diff --git a/venv/lib/python3.10/site-packages/PIL/AvifImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/AvifImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..366e0c864bf6ece4c401fe827430e07fd7fc4a09 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/AvifImagePlugin.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import os +from io import BytesIO +from typing import IO + +from . import ExifTags, Image, ImageFile + +try: + from . import _avif + + SUPPORTED = True +except ImportError: + SUPPORTED = False + +# Decoder options as module globals, until there is a way to pass parameters +# to Image.open (see https://github.com/python-pillow/Pillow/issues/569) +DECODE_CODEC_CHOICE = "auto" +DEFAULT_MAX_THREADS = 0 + + +def get_codec_version(codec_name: str) -> str | None: + versions = _avif.codec_versions() + for version in versions.split(", "): + if version.split(" [")[0] == codec_name: + return version.split(":")[-1].split(" ")[0] + return None + + +def _accept(prefix: bytes) -> bool | str: + if prefix[4:8] != b"ftyp": + return False + major_brand = prefix[8:12] + if major_brand in ( + # coding brands + b"avif", + b"avis", + # We accept files with AVIF container brands; we can't yet know if + # the ftyp box has the correct compatible brands, but if it doesn't + # then the plugin will raise a SyntaxError which Pillow will catch + # before moving on to the next plugin that accepts the file. + # + # Also, because this file might not actually be an AVIF file, we + # don't raise an error if AVIF support isn't properly compiled. + b"mif1", + b"msf1", + ): + if not SUPPORTED: + return ( + "image file could not be identified because AVIF support not installed" + ) + return True + return False + + +def _get_default_max_threads() -> int: + if DEFAULT_MAX_THREADS: + return DEFAULT_MAX_THREADS + if hasattr(os, "sched_getaffinity"): + return len(os.sched_getaffinity(0)) + else: + return os.cpu_count() or 1 + + +class AvifImageFile(ImageFile.ImageFile): + format = "AVIF" + format_description = "AVIF image" + __frame = -1 + + def _open(self) -> None: + if not SUPPORTED: + msg = "image file could not be opened because AVIF support not installed" + raise SyntaxError(msg) + + if DECODE_CODEC_CHOICE != "auto" and not _avif.decoder_codec_available( + DECODE_CODEC_CHOICE + ): + msg = "Invalid opening codec" + raise ValueError(msg) + self._decoder = _avif.AvifDecoder( + self.fp.read(), + DECODE_CODEC_CHOICE, + _get_default_max_threads(), + ) + + # Get info from decoder + self._size, self.n_frames, self._mode, icc, exif, exif_orientation, xmp = ( + self._decoder.get_info() + ) + self.is_animated = self.n_frames > 1 + + if icc: + self.info["icc_profile"] = icc + if xmp: + self.info["xmp"] = xmp + + if exif_orientation != 1 or exif: + exif_data = Image.Exif() + if exif: + exif_data.load(exif) + original_orientation = exif_data.get(ExifTags.Base.Orientation, 1) + else: + original_orientation = 1 + if exif_orientation != original_orientation: + exif_data[ExifTags.Base.Orientation] = exif_orientation + exif = exif_data.tobytes() + if exif: + self.info["exif"] = exif + self.seek(0) + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + + # Set tile + self.__frame = frame + self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)] + + def load(self) -> Image.core.PixelAccess | None: + if self.tile: + # We need to load the image data for this frame + data, timescale, pts_in_timescales, duration_in_timescales = ( + self._decoder.get_frame(self.__frame) + ) + self.info["timestamp"] = round(1000 * (pts_in_timescales / timescale)) + self.info["duration"] = round(1000 * (duration_in_timescales / timescale)) + + if self.fp and self._exclusive_fp: + self.fp.close() + self.fp = BytesIO(data) + + return super().load() + + def load_seek(self, pos: int) -> None: + pass + + def tell(self) -> int: + return self.__frame + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + _save(im, fp, filename, save_all=True) + + +def _save( + im: Image.Image, fp: IO[bytes], filename: str | bytes, save_all: bool = False +) -> None: + info = im.encoderinfo.copy() + if save_all: + append_images = list(info.get("append_images", [])) + else: + append_images = [] + + total = 0 + for ims in [im] + append_images: + total += getattr(ims, "n_frames", 1) + + quality = info.get("quality", 75) + if not isinstance(quality, int) or quality < 0 or quality > 100: + msg = "Invalid quality setting" + raise ValueError(msg) + + duration = info.get("duration", 0) + subsampling = info.get("subsampling", "4:2:0") + speed = info.get("speed", 6) + max_threads = info.get("max_threads", _get_default_max_threads()) + codec = info.get("codec", "auto") + if codec != "auto" and not _avif.encoder_codec_available(codec): + msg = "Invalid saving codec" + raise ValueError(msg) + range_ = info.get("range", "full") + tile_rows_log2 = info.get("tile_rows", 0) + tile_cols_log2 = info.get("tile_cols", 0) + alpha_premultiplied = bool(info.get("alpha_premultiplied", False)) + autotiling = bool(info.get("autotiling", tile_rows_log2 == tile_cols_log2 == 0)) + + icc_profile = info.get("icc_profile", im.info.get("icc_profile")) + exif_orientation = 1 + if exif := info.get("exif"): + if isinstance(exif, Image.Exif): + exif_data = exif + else: + exif_data = Image.Exif() + exif_data.load(exif) + if ExifTags.Base.Orientation in exif_data: + exif_orientation = exif_data.pop(ExifTags.Base.Orientation) + exif = exif_data.tobytes() if exif_data else b"" + elif isinstance(exif, Image.Exif): + exif = exif_data.tobytes() + + xmp = info.get("xmp") + + if isinstance(xmp, str): + xmp = xmp.encode("utf-8") + + advanced = info.get("advanced") + if advanced is not None: + if isinstance(advanced, dict): + advanced = advanced.items() + try: + advanced = tuple(advanced) + except TypeError: + invalid = True + else: + invalid = any(not isinstance(v, tuple) or len(v) != 2 for v in advanced) + if invalid: + msg = ( + "advanced codec options must be a dict of key-value string " + "pairs or a series of key-value two-tuples" + ) + raise ValueError(msg) + + # Setup the AVIF encoder + enc = _avif.AvifEncoder( + im.size, + subsampling, + quality, + speed, + max_threads, + codec, + range_, + tile_rows_log2, + tile_cols_log2, + alpha_premultiplied, + autotiling, + icc_profile or b"", + exif or b"", + exif_orientation, + xmp or b"", + advanced, + ) + + # Add each frame + frame_idx = 0 + frame_duration = 0 + cur_idx = im.tell() + is_single_frame = total == 1 + try: + for ims in [im] + append_images: + # Get number of frames in this image + nfr = getattr(ims, "n_frames", 1) + + for idx in range(nfr): + ims.seek(idx) + + # Make sure image mode is supported + frame = ims + rawmode = ims.mode + if ims.mode not in {"RGB", "RGBA"}: + rawmode = "RGBA" if ims.has_transparency_data else "RGB" + frame = ims.convert(rawmode) + + # Update frame duration + if isinstance(duration, (list, tuple)): + frame_duration = duration[frame_idx] + else: + frame_duration = duration + + # Append the frame to the animation encoder + enc.add( + frame.tobytes("raw", rawmode), + frame_duration, + frame.size, + rawmode, + is_single_frame, + ) + + # Update frame index + frame_idx += 1 + + if not save_all: + break + + finally: + im.seek(cur_idx) + + # Get the final output from the encoder + data = enc.finish() + if data is None: + msg = "cannot write file as AVIF (encoder returned None)" + raise OSError(msg) + + fp.write(data) + + +Image.register_open(AvifImageFile.format, AvifImageFile, _accept) +if SUPPORTED: + Image.register_save(AvifImageFile.format, _save) + Image.register_save_all(AvifImageFile.format, _save_all) + Image.register_extensions(AvifImageFile.format, [".avif", ".avifs"]) + Image.register_mime(AvifImageFile.format, "image/avif") diff --git a/venv/lib/python3.10/site-packages/PIL/BdfFontFile.py b/venv/lib/python3.10/site-packages/PIL/BdfFontFile.py new file mode 100644 index 0000000000000000000000000000000000000000..f175e2f4f80b1b232d79f15a6db0667296917c97 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/BdfFontFile.py @@ -0,0 +1,122 @@ +# +# The Python Imaging Library +# $Id$ +# +# bitmap distribution font (bdf) file parser +# +# history: +# 1996-05-16 fl created (as bdf2pil) +# 1997-08-25 fl converted to FontFile driver +# 2001-05-25 fl removed bogus __init__ call +# 2002-11-20 fl robustification (from Kevin Cazabon, Dmitry Vasiliev) +# 2003-04-22 fl more robustification (from Graham Dumpleton) +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1997-2003 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + +""" +Parse X Bitmap Distribution Format (BDF) +""" +from __future__ import annotations + +from typing import BinaryIO + +from . import FontFile, Image + + +def bdf_char( + f: BinaryIO, +) -> ( + tuple[ + str, + int, + tuple[tuple[int, int], tuple[int, int, int, int], tuple[int, int, int, int]], + Image.Image, + ] + | None +): + # skip to STARTCHAR + while True: + s = f.readline() + if not s: + return None + if s.startswith(b"STARTCHAR"): + break + id = s[9:].strip().decode("ascii") + + # load symbol properties + props = {} + while True: + s = f.readline() + if not s or s.startswith(b"BITMAP"): + break + i = s.find(b" ") + props[s[:i].decode("ascii")] = s[i + 1 : -1].decode("ascii") + + # load bitmap + bitmap = bytearray() + while True: + s = f.readline() + if not s or s.startswith(b"ENDCHAR"): + break + bitmap += s[:-1] + + # The word BBX + # followed by the width in x (BBw), height in y (BBh), + # and x and y displacement (BBxoff0, BByoff0) + # of the lower left corner from the origin of the character. + width, height, x_disp, y_disp = (int(p) for p in props["BBX"].split()) + + # The word DWIDTH + # followed by the width in x and y of the character in device pixels. + dwx, dwy = (int(p) for p in props["DWIDTH"].split()) + + bbox = ( + (dwx, dwy), + (x_disp, -y_disp - height, width + x_disp, -y_disp), + (0, 0, width, height), + ) + + try: + im = Image.frombytes("1", (width, height), bitmap, "hex", "1") + except ValueError: + # deal with zero-width characters + im = Image.new("1", (width, height)) + + return id, int(props["ENCODING"]), bbox, im + + +class BdfFontFile(FontFile.FontFile): + """Font file plugin for the X11 BDF format.""" + + def __init__(self, fp: BinaryIO) -> None: + super().__init__() + + s = fp.readline() + if not s.startswith(b"STARTFONT 2.1"): + msg = "not a valid BDF file" + raise SyntaxError(msg) + + props = {} + comments = [] + + while True: + s = fp.readline() + if not s or s.startswith(b"ENDPROPERTIES"): + break + i = s.find(b" ") + props[s[:i].decode("ascii")] = s[i + 1 : -1].decode("ascii") + if s[:i] in [b"COMMENT", b"COPYRIGHT"]: + if s.find(b"LogicalFontDescription") < 0: + comments.append(s[i + 1 : -1].decode("ascii")) + + while True: + c = bdf_char(fp) + if not c: + break + id, ch, (xy, dst, src), im = c + if 0 <= ch < len(self.glyph): + self.glyph[ch] = xy, dst, src, im diff --git a/venv/lib/python3.10/site-packages/PIL/BlpImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/BlpImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..f7be7746d84bbc076e0a41124a903a8b2b05ae61 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/BlpImagePlugin.py @@ -0,0 +1,497 @@ +""" +Blizzard Mipmap Format (.blp) +Jerome Leclanche + +The contents of this file are hereby released in the public domain (CC0) +Full text of the CC0 license: + https://creativecommons.org/publicdomain/zero/1.0/ + +BLP1 files, used mostly in Warcraft III, are not fully supported. +All types of BLP2 files used in World of Warcraft are supported. + +The BLP file structure consists of a header, up to 16 mipmaps of the +texture + +Texture sizes must be powers of two, though the two dimensions do +not have to be equal; 512x256 is valid, but 512x200 is not. +The first mipmap (mipmap #0) is the full size image; each subsequent +mipmap halves both dimensions. The final mipmap should be 1x1. + +BLP files come in many different flavours: +* JPEG-compressed (type == 0) - only supported for BLP1. +* RAW images (type == 1, encoding == 1). Each mipmap is stored as an + array of 8-bit values, one per pixel, left to right, top to bottom. + Each value is an index to the palette. +* DXT-compressed (type == 1, encoding == 2): +- DXT1 compression is used if alpha_encoding == 0. + - An additional alpha bit is used if alpha_depth == 1. + - DXT3 compression is used if alpha_encoding == 1. + - DXT5 compression is used if alpha_encoding == 7. +""" + +from __future__ import annotations + +import abc +import os +import struct +from enum import IntEnum +from io import BytesIO +from typing import IO + +from . import Image, ImageFile + + +class Format(IntEnum): + JPEG = 0 + + +class Encoding(IntEnum): + UNCOMPRESSED = 1 + DXT = 2 + UNCOMPRESSED_RAW_BGRA = 3 + + +class AlphaEncoding(IntEnum): + DXT1 = 0 + DXT3 = 1 + DXT5 = 7 + + +def unpack_565(i: int) -> tuple[int, int, int]: + return ((i >> 11) & 0x1F) << 3, ((i >> 5) & 0x3F) << 2, (i & 0x1F) << 3 + + +def decode_dxt1( + data: bytes, alpha: bool = False +) -> tuple[bytearray, bytearray, bytearray, bytearray]: + """ + input: one "row" of data (i.e. will produce 4*width pixels) + """ + + blocks = len(data) // 8 # number of blocks in row + ret = (bytearray(), bytearray(), bytearray(), bytearray()) + + for block_index in range(blocks): + # Decode next 8-byte block. + idx = block_index * 8 + color0, color1, bits = struct.unpack_from("> 2 + + a = 0xFF + if control == 0: + r, g, b = r0, g0, b0 + elif control == 1: + r, g, b = r1, g1, b1 + elif control == 2: + if color0 > color1: + r = (2 * r0 + r1) // 3 + g = (2 * g0 + g1) // 3 + b = (2 * b0 + b1) // 3 + else: + r = (r0 + r1) // 2 + g = (g0 + g1) // 2 + b = (b0 + b1) // 2 + elif control == 3: + if color0 > color1: + r = (2 * r1 + r0) // 3 + g = (2 * g1 + g0) // 3 + b = (2 * b1 + b0) // 3 + else: + r, g, b, a = 0, 0, 0, 0 + + if alpha: + ret[j].extend([r, g, b, a]) + else: + ret[j].extend([r, g, b]) + + return ret + + +def decode_dxt3(data: bytes) -> tuple[bytearray, bytearray, bytearray, bytearray]: + """ + input: one "row" of data (i.e. will produce 4*width pixels) + """ + + blocks = len(data) // 16 # number of blocks in row + ret = (bytearray(), bytearray(), bytearray(), bytearray()) + + for block_index in range(blocks): + idx = block_index * 16 + block = data[idx : idx + 16] + # Decode next 16-byte block. + bits = struct.unpack_from("<8B", block) + color0, color1 = struct.unpack_from(">= 4 + else: + high = True + a &= 0xF + a *= 17 # We get a value between 0 and 15 + + color_code = (code >> 2 * (4 * j + i)) & 0x03 + + if color_code == 0: + r, g, b = r0, g0, b0 + elif color_code == 1: + r, g, b = r1, g1, b1 + elif color_code == 2: + r = (2 * r0 + r1) // 3 + g = (2 * g0 + g1) // 3 + b = (2 * b0 + b1) // 3 + elif color_code == 3: + r = (2 * r1 + r0) // 3 + g = (2 * g1 + g0) // 3 + b = (2 * b1 + b0) // 3 + + ret[j].extend([r, g, b, a]) + + return ret + + +def decode_dxt5(data: bytes) -> tuple[bytearray, bytearray, bytearray, bytearray]: + """ + input: one "row" of data (i.e. will produce 4 * width pixels) + """ + + blocks = len(data) // 16 # number of blocks in row + ret = (bytearray(), bytearray(), bytearray(), bytearray()) + + for block_index in range(blocks): + idx = block_index * 16 + block = data[idx : idx + 16] + # Decode next 16-byte block. + a0, a1 = struct.unpack_from("> alphacode_index) & 0x07 + elif alphacode_index == 15: + alphacode = (alphacode2 >> 15) | ((alphacode1 << 1) & 0x06) + else: # alphacode_index >= 18 and alphacode_index <= 45 + alphacode = (alphacode1 >> (alphacode_index - 16)) & 0x07 + + if alphacode == 0: + a = a0 + elif alphacode == 1: + a = a1 + elif a0 > a1: + a = ((8 - alphacode) * a0 + (alphacode - 1) * a1) // 7 + elif alphacode == 6: + a = 0 + elif alphacode == 7: + a = 255 + else: + a = ((6 - alphacode) * a0 + (alphacode - 1) * a1) // 5 + + color_code = (code >> 2 * (4 * j + i)) & 0x03 + + if color_code == 0: + r, g, b = r0, g0, b0 + elif color_code == 1: + r, g, b = r1, g1, b1 + elif color_code == 2: + r = (2 * r0 + r1) // 3 + g = (2 * g0 + g1) // 3 + b = (2 * b0 + b1) // 3 + elif color_code == 3: + r = (2 * r1 + r0) // 3 + g = (2 * g1 + g0) // 3 + b = (2 * b1 + b0) // 3 + + ret[j].extend([r, g, b, a]) + + return ret + + +class BLPFormatError(NotImplementedError): + pass + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith((b"BLP1", b"BLP2")) + + +class BlpImageFile(ImageFile.ImageFile): + """ + Blizzard Mipmap Format + """ + + format = "BLP" + format_description = "Blizzard Mipmap Format" + + def _open(self) -> None: + self.magic = self.fp.read(4) + if not _accept(self.magic): + msg = f"Bad BLP magic {repr(self.magic)}" + raise BLPFormatError(msg) + + compression = struct.unpack(" tuple[int, int]: + try: + self._read_header() + self._load() + except struct.error as e: + msg = "Truncated BLP file" + raise OSError(msg) from e + return -1, 0 + + @abc.abstractmethod + def _load(self) -> None: + pass + + def _read_header(self) -> None: + self._offsets = struct.unpack("<16I", self._safe_read(16 * 4)) + self._lengths = struct.unpack("<16I", self._safe_read(16 * 4)) + + def _safe_read(self, length: int) -> bytes: + assert self.fd is not None + return ImageFile._safe_read(self.fd, length) + + def _read_palette(self) -> list[tuple[int, int, int, int]]: + ret = [] + for i in range(256): + try: + b, g, r, a = struct.unpack("<4B", self._safe_read(4)) + except struct.error: + break + ret.append((b, g, r, a)) + return ret + + def _read_bgra( + self, palette: list[tuple[int, int, int, int]], alpha: bool + ) -> bytearray: + data = bytearray() + _data = BytesIO(self._safe_read(self._lengths[0])) + while True: + try: + (offset,) = struct.unpack(" None: + self._compression, self._encoding, alpha = self.args + + if self._compression == Format.JPEG: + self._decode_jpeg_stream() + + elif self._compression == 1: + if self._encoding in (4, 5): + palette = self._read_palette() + data = self._read_bgra(palette, alpha) + self.set_as_raw(data) + else: + msg = f"Unsupported BLP encoding {repr(self._encoding)}" + raise BLPFormatError(msg) + else: + msg = f"Unsupported BLP compression {repr(self._encoding)}" + raise BLPFormatError(msg) + + def _decode_jpeg_stream(self) -> None: + from .JpegImagePlugin import JpegImageFile + + (jpeg_header_size,) = struct.unpack(" None: + self._compression, self._encoding, alpha, self._alpha_encoding = self.args + + palette = self._read_palette() + + assert self.fd is not None + self.fd.seek(self._offsets[0]) + + if self._compression == 1: + # Uncompressed or DirectX compression + + if self._encoding == Encoding.UNCOMPRESSED: + data = self._read_bgra(palette, alpha) + + elif self._encoding == Encoding.DXT: + data = bytearray() + if self._alpha_encoding == AlphaEncoding.DXT1: + linesize = (self.state.xsize + 3) // 4 * 8 + for yb in range((self.state.ysize + 3) // 4): + for d in decode_dxt1(self._safe_read(linesize), alpha): + data += d + + elif self._alpha_encoding == AlphaEncoding.DXT3: + linesize = (self.state.xsize + 3) // 4 * 16 + for yb in range((self.state.ysize + 3) // 4): + for d in decode_dxt3(self._safe_read(linesize)): + data += d + + elif self._alpha_encoding == AlphaEncoding.DXT5: + linesize = (self.state.xsize + 3) // 4 * 16 + for yb in range((self.state.ysize + 3) // 4): + for d in decode_dxt5(self._safe_read(linesize)): + data += d + else: + msg = f"Unsupported alpha encoding {repr(self._alpha_encoding)}" + raise BLPFormatError(msg) + else: + msg = f"Unknown BLP encoding {repr(self._encoding)}" + raise BLPFormatError(msg) + + else: + msg = f"Unknown BLP compression {repr(self._compression)}" + raise BLPFormatError(msg) + + self.set_as_raw(data) + + +class BLPEncoder(ImageFile.PyEncoder): + _pushes_fd = True + + def _write_palette(self) -> bytes: + data = b"" + assert self.im is not None + palette = self.im.getpalette("RGBA", "RGBA") + for i in range(len(palette) // 4): + r, g, b, a = palette[i * 4 : (i + 1) * 4] + data += struct.pack("<4B", b, g, r, a) + while len(data) < 256 * 4: + data += b"\x00" * 4 + return data + + def encode(self, bufsize: int) -> tuple[int, int, bytes]: + palette_data = self._write_palette() + + offset = 20 + 16 * 4 * 2 + len(palette_data) + data = struct.pack("<16I", offset, *((0,) * 15)) + + assert self.im is not None + w, h = self.im.size + data += struct.pack("<16I", w * h, *((0,) * 15)) + + data += palette_data + + for y in range(h): + for x in range(w): + data += struct.pack(" None: + if im.mode != "P": + msg = "Unsupported BLP image mode" + raise ValueError(msg) + + magic = b"BLP1" if im.encoderinfo.get("blp_version") == "BLP1" else b"BLP2" + fp.write(magic) + + assert im.palette is not None + fp.write(struct.pack(" mode, rawmode + 1: ("P", "P;1"), + 4: ("P", "P;4"), + 8: ("P", "P"), + 16: ("RGB", "BGR;15"), + 24: ("RGB", "BGR"), + 32: ("RGB", "BGRX"), +} + +USE_RAW_ALPHA = False + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"BM") + + +def _dib_accept(prefix: bytes) -> bool: + return i32(prefix) in [12, 40, 52, 56, 64, 108, 124] + + +# ============================================================================= +# Image plugin for the Windows BMP format. +# ============================================================================= +class BmpImageFile(ImageFile.ImageFile): + """Image plugin for the Windows Bitmap format (BMP)""" + + # ------------------------------------------------------------- Description + format_description = "Windows Bitmap" + format = "BMP" + + # -------------------------------------------------- BMP Compression values + COMPRESSIONS = {"RAW": 0, "RLE8": 1, "RLE4": 2, "BITFIELDS": 3, "JPEG": 4, "PNG": 5} + for k, v in COMPRESSIONS.items(): + vars()[k] = v + + def _bitmap(self, header: int = 0, offset: int = 0) -> None: + """Read relevant info about the BMP""" + read, seek = self.fp.read, self.fp.seek + if header: + seek(header) + # read bmp header size @offset 14 (this is part of the header size) + file_info: dict[str, bool | int | tuple[int, ...]] = { + "header_size": i32(read(4)), + "direction": -1, + } + + # -------------------- If requested, read header at a specific position + # read the rest of the bmp header, without its size + assert isinstance(file_info["header_size"], int) + header_data = ImageFile._safe_read(self.fp, file_info["header_size"] - 4) + + # ------------------------------- Windows Bitmap v2, IBM OS/2 Bitmap v1 + # ----- This format has different offsets because of width/height types + # 12: BITMAPCOREHEADER/OS21XBITMAPHEADER + if file_info["header_size"] == 12: + file_info["width"] = i16(header_data, 0) + file_info["height"] = i16(header_data, 2) + file_info["planes"] = i16(header_data, 4) + file_info["bits"] = i16(header_data, 6) + file_info["compression"] = self.COMPRESSIONS["RAW"] + file_info["palette_padding"] = 3 + + # --------------------------------------------- Windows Bitmap v3 to v5 + # 40: BITMAPINFOHEADER + # 52: BITMAPV2HEADER + # 56: BITMAPV3HEADER + # 64: BITMAPCOREHEADER2/OS22XBITMAPHEADER + # 108: BITMAPV4HEADER + # 124: BITMAPV5HEADER + elif file_info["header_size"] in (40, 52, 56, 64, 108, 124): + file_info["y_flip"] = header_data[7] == 0xFF + file_info["direction"] = 1 if file_info["y_flip"] else -1 + file_info["width"] = i32(header_data, 0) + file_info["height"] = ( + i32(header_data, 4) + if not file_info["y_flip"] + else 2**32 - i32(header_data, 4) + ) + file_info["planes"] = i16(header_data, 8) + file_info["bits"] = i16(header_data, 10) + file_info["compression"] = i32(header_data, 12) + # byte size of pixel data + file_info["data_size"] = i32(header_data, 16) + file_info["pixels_per_meter"] = ( + i32(header_data, 20), + i32(header_data, 24), + ) + file_info["colors"] = i32(header_data, 28) + file_info["palette_padding"] = 4 + assert isinstance(file_info["pixels_per_meter"], tuple) + self.info["dpi"] = tuple(x / 39.3701 for x in file_info["pixels_per_meter"]) + if file_info["compression"] == self.COMPRESSIONS["BITFIELDS"]: + masks = ["r_mask", "g_mask", "b_mask"] + if len(header_data) >= 48: + if len(header_data) >= 52: + masks.append("a_mask") + else: + file_info["a_mask"] = 0x0 + for idx, mask in enumerate(masks): + file_info[mask] = i32(header_data, 36 + idx * 4) + else: + # 40 byte headers only have the three components in the + # bitfields masks, ref: + # https://msdn.microsoft.com/en-us/library/windows/desktop/dd183376(v=vs.85).aspx + # See also + # https://github.com/python-pillow/Pillow/issues/1293 + # There is a 4th component in the RGBQuad, in the alpha + # location, but it is listed as a reserved component, + # and it is not generally an alpha channel + file_info["a_mask"] = 0x0 + for mask in masks: + file_info[mask] = i32(read(4)) + assert isinstance(file_info["r_mask"], int) + assert isinstance(file_info["g_mask"], int) + assert isinstance(file_info["b_mask"], int) + assert isinstance(file_info["a_mask"], int) + file_info["rgb_mask"] = ( + file_info["r_mask"], + file_info["g_mask"], + file_info["b_mask"], + ) + file_info["rgba_mask"] = ( + file_info["r_mask"], + file_info["g_mask"], + file_info["b_mask"], + file_info["a_mask"], + ) + else: + msg = f"Unsupported BMP header type ({file_info['header_size']})" + raise OSError(msg) + + # ------------------ Special case : header is reported 40, which + # ---------------------- is shorter than real size for bpp >= 16 + assert isinstance(file_info["width"], int) + assert isinstance(file_info["height"], int) + self._size = file_info["width"], file_info["height"] + + # ------- If color count was not found in the header, compute from bits + assert isinstance(file_info["bits"], int) + file_info["colors"] = ( + file_info["colors"] + if file_info.get("colors", 0) + else (1 << file_info["bits"]) + ) + assert isinstance(file_info["colors"], int) + if offset == 14 + file_info["header_size"] and file_info["bits"] <= 8: + offset += 4 * file_info["colors"] + + # ---------------------- Check bit depth for unusual unsupported values + self._mode, raw_mode = BIT2MODE.get(file_info["bits"], ("", "")) + if not self.mode: + msg = f"Unsupported BMP pixel depth ({file_info['bits']})" + raise OSError(msg) + + # ---------------- Process BMP with Bitfields compression (not palette) + decoder_name = "raw" + if file_info["compression"] == self.COMPRESSIONS["BITFIELDS"]: + SUPPORTED: dict[int, list[tuple[int, ...]]] = { + 32: [ + (0xFF0000, 0xFF00, 0xFF, 0x0), + (0xFF000000, 0xFF0000, 0xFF00, 0x0), + (0xFF000000, 0xFF00, 0xFF, 0x0), + (0xFF000000, 0xFF0000, 0xFF00, 0xFF), + (0xFF, 0xFF00, 0xFF0000, 0xFF000000), + (0xFF0000, 0xFF00, 0xFF, 0xFF000000), + (0xFF000000, 0xFF00, 0xFF, 0xFF0000), + (0x0, 0x0, 0x0, 0x0), + ], + 24: [(0xFF0000, 0xFF00, 0xFF)], + 16: [(0xF800, 0x7E0, 0x1F), (0x7C00, 0x3E0, 0x1F)], + } + MASK_MODES = { + (32, (0xFF0000, 0xFF00, 0xFF, 0x0)): "BGRX", + (32, (0xFF000000, 0xFF0000, 0xFF00, 0x0)): "XBGR", + (32, (0xFF000000, 0xFF00, 0xFF, 0x0)): "BGXR", + (32, (0xFF000000, 0xFF0000, 0xFF00, 0xFF)): "ABGR", + (32, (0xFF, 0xFF00, 0xFF0000, 0xFF000000)): "RGBA", + (32, (0xFF0000, 0xFF00, 0xFF, 0xFF000000)): "BGRA", + (32, (0xFF000000, 0xFF00, 0xFF, 0xFF0000)): "BGAR", + (32, (0x0, 0x0, 0x0, 0x0)): "BGRA", + (24, (0xFF0000, 0xFF00, 0xFF)): "BGR", + (16, (0xF800, 0x7E0, 0x1F)): "BGR;16", + (16, (0x7C00, 0x3E0, 0x1F)): "BGR;15", + } + if file_info["bits"] in SUPPORTED: + if ( + file_info["bits"] == 32 + and file_info["rgba_mask"] in SUPPORTED[file_info["bits"]] + ): + assert isinstance(file_info["rgba_mask"], tuple) + raw_mode = MASK_MODES[(file_info["bits"], file_info["rgba_mask"])] + self._mode = "RGBA" if "A" in raw_mode else self.mode + elif ( + file_info["bits"] in (24, 16) + and file_info["rgb_mask"] in SUPPORTED[file_info["bits"]] + ): + assert isinstance(file_info["rgb_mask"], tuple) + raw_mode = MASK_MODES[(file_info["bits"], file_info["rgb_mask"])] + else: + msg = "Unsupported BMP bitfields layout" + raise OSError(msg) + else: + msg = "Unsupported BMP bitfields layout" + raise OSError(msg) + elif file_info["compression"] == self.COMPRESSIONS["RAW"]: + if file_info["bits"] == 32 and ( + header == 22 or USE_RAW_ALPHA # 32-bit .cur offset + ): + raw_mode, self._mode = "BGRA", "RGBA" + elif file_info["compression"] in ( + self.COMPRESSIONS["RLE8"], + self.COMPRESSIONS["RLE4"], + ): + decoder_name = "bmp_rle" + else: + msg = f"Unsupported BMP compression ({file_info['compression']})" + raise OSError(msg) + + # --------------- Once the header is processed, process the palette/LUT + if self.mode == "P": # Paletted for 1, 4 and 8 bit images + # ---------------------------------------------------- 1-bit images + if not (0 < file_info["colors"] <= 65536): + msg = f"Unsupported BMP Palette size ({file_info['colors']})" + raise OSError(msg) + else: + assert isinstance(file_info["palette_padding"], int) + padding = file_info["palette_padding"] + palette = read(padding * file_info["colors"]) + grayscale = True + indices = ( + (0, 255) + if file_info["colors"] == 2 + else list(range(file_info["colors"])) + ) + + # ----------------- Check if grayscale and ignore palette if so + for ind, val in enumerate(indices): + rgb = palette[ind * padding : ind * padding + 3] + if rgb != o8(val) * 3: + grayscale = False + + # ------- If all colors are gray, white or black, ditch palette + if grayscale: + self._mode = "1" if file_info["colors"] == 2 else "L" + raw_mode = self.mode + else: + self._mode = "P" + self.palette = ImagePalette.raw( + "BGRX" if padding == 4 else "BGR", palette + ) + + # ---------------------------- Finally set the tile data for the plugin + self.info["compression"] = file_info["compression"] + args: list[Any] = [raw_mode] + if decoder_name == "bmp_rle": + args.append(file_info["compression"] == self.COMPRESSIONS["RLE4"]) + else: + assert isinstance(file_info["width"], int) + args.append(((file_info["width"] * file_info["bits"] + 31) >> 3) & (~3)) + args.append(file_info["direction"]) + self.tile = [ + ImageFile._Tile( + decoder_name, + (0, 0, file_info["width"], file_info["height"]), + offset or self.fp.tell(), + tuple(args), + ) + ] + + def _open(self) -> None: + """Open file, check magic number and read header""" + # read 14 bytes: magic number, filesize, reserved, header final offset + head_data = self.fp.read(14) + # choke if the file does not have the required magic bytes + if not _accept(head_data): + msg = "Not a BMP file" + raise SyntaxError(msg) + # read the start position of the BMP image data (u32) + offset = i32(head_data, 10) + # load bitmap information (offset=raster info) + self._bitmap(offset=offset) + + +class BmpRleDecoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + rle4 = self.args[1] + data = bytearray() + x = 0 + dest_length = self.state.xsize * self.state.ysize + while len(data) < dest_length: + pixels = self.fd.read(1) + byte = self.fd.read(1) + if not pixels or not byte: + break + num_pixels = pixels[0] + if num_pixels: + # encoded mode + if x + num_pixels > self.state.xsize: + # Too much data for row + num_pixels = max(0, self.state.xsize - x) + if rle4: + first_pixel = o8(byte[0] >> 4) + second_pixel = o8(byte[0] & 0x0F) + for index in range(num_pixels): + if index % 2 == 0: + data += first_pixel + else: + data += second_pixel + else: + data += byte * num_pixels + x += num_pixels + else: + if byte[0] == 0: + # end of line + while len(data) % self.state.xsize != 0: + data += b"\x00" + x = 0 + elif byte[0] == 1: + # end of bitmap + break + elif byte[0] == 2: + # delta + bytes_read = self.fd.read(2) + if len(bytes_read) < 2: + break + right, up = self.fd.read(2) + data += b"\x00" * (right + up * self.state.xsize) + x = len(data) % self.state.xsize + else: + # absolute mode + if rle4: + # 2 pixels per byte + byte_count = byte[0] // 2 + bytes_read = self.fd.read(byte_count) + for byte_read in bytes_read: + data += o8(byte_read >> 4) + data += o8(byte_read & 0x0F) + else: + byte_count = byte[0] + bytes_read = self.fd.read(byte_count) + data += bytes_read + if len(bytes_read) < byte_count: + break + x += byte[0] + + # align to 16-bit word boundary + if self.fd.tell() % 2 != 0: + self.fd.seek(1, os.SEEK_CUR) + rawmode = "L" if self.mode == "L" else "P" + self.set_as_raw(bytes(data), rawmode, (0, self.args[-1])) + return -1, 0 + + +# ============================================================================= +# Image plugin for the DIB format (BMP alias) +# ============================================================================= +class DibImageFile(BmpImageFile): + format = "DIB" + format_description = "Windows Bitmap" + + def _open(self) -> None: + self._bitmap() + + +# +# -------------------------------------------------------------------- +# Write BMP file + + +SAVE = { + "1": ("1", 1, 2), + "L": ("L", 8, 256), + "P": ("P", 8, 256), + "RGB": ("BGR", 24, 0), + "RGBA": ("BGRA", 32, 0), +} + + +def _dib_save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + _save(im, fp, filename, False) + + +def _save( + im: Image.Image, fp: IO[bytes], filename: str | bytes, bitmap_header: bool = True +) -> None: + try: + rawmode, bits, colors = SAVE[im.mode] + except KeyError as e: + msg = f"cannot write mode {im.mode} as BMP" + raise OSError(msg) from e + + info = im.encoderinfo + + dpi = info.get("dpi", (96, 96)) + + # 1 meter == 39.3701 inches + ppm = tuple(int(x * 39.3701 + 0.5) for x in dpi) + + stride = ((im.size[0] * bits + 7) // 8 + 3) & (~3) + header = 40 # or 64 for OS/2 version 2 + image = stride * im.size[1] + + if im.mode == "1": + palette = b"".join(o8(i) * 3 + b"\x00" for i in (0, 255)) + elif im.mode == "L": + palette = b"".join(o8(i) * 3 + b"\x00" for i in range(256)) + elif im.mode == "P": + palette = im.im.getpalette("RGB", "BGRX") + colors = len(palette) // 4 + else: + palette = None + + # bitmap header + if bitmap_header: + offset = 14 + header + colors * 4 + file_size = offset + image + if file_size > 2**32 - 1: + msg = "File size is too large for the BMP format" + raise ValueError(msg) + fp.write( + b"BM" # file type (magic) + + o32(file_size) # file size + + o32(0) # reserved + + o32(offset) # image data offset + ) + + # bitmap info header + fp.write( + o32(header) # info header size + + o32(im.size[0]) # width + + o32(im.size[1]) # height + + o16(1) # planes + + o16(bits) # depth + + o32(0) # compression (0=uncompressed) + + o32(image) # size of bitmap + + o32(ppm[0]) # resolution + + o32(ppm[1]) # resolution + + o32(colors) # colors used + + o32(colors) # colors important + ) + + fp.write(b"\0" * (header - 40)) # padding (for OS/2 format) + + if palette: + fp.write(palette) + + ImageFile._save( + im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 0, (rawmode, stride, -1))] + ) + + +# +# -------------------------------------------------------------------- +# Registry + + +Image.register_open(BmpImageFile.format, BmpImageFile, _accept) +Image.register_save(BmpImageFile.format, _save) + +Image.register_extension(BmpImageFile.format, ".bmp") + +Image.register_mime(BmpImageFile.format, "image/bmp") + +Image.register_decoder("bmp_rle", BmpRleDecoder) + +Image.register_open(DibImageFile.format, DibImageFile, _dib_accept) +Image.register_save(DibImageFile.format, _dib_save) + +Image.register_extension(DibImageFile.format, ".dib") + +Image.register_mime(DibImageFile.format, "image/bmp") diff --git a/venv/lib/python3.10/site-packages/PIL/BufrStubImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/BufrStubImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..8c5da14f5f6769127e0ed23d36fa8f3dd0b086c2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/BufrStubImagePlugin.py @@ -0,0 +1,75 @@ +# +# The Python Imaging Library +# $Id$ +# +# BUFR stub adapter +# +# Copyright (c) 1996-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +from typing import IO + +from . import Image, ImageFile + +_handler = None + + +def register_handler(handler: ImageFile.StubHandler | None) -> None: + """ + Install application-specific BUFR image handler. + + :param handler: Handler object. + """ + global _handler + _handler = handler + + +# -------------------------------------------------------------------- +# Image adapter + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith((b"BUFR", b"ZCZC")) + + +class BufrStubImageFile(ImageFile.StubImageFile): + format = "BUFR" + format_description = "BUFR" + + def _open(self) -> None: + if not _accept(self.fp.read(4)): + msg = "Not a BUFR file" + raise SyntaxError(msg) + + self.fp.seek(-4, os.SEEK_CUR) + + # make something up + self._mode = "F" + self._size = 1, 1 + + loader = self._load() + if loader: + loader.open(self) + + def _load(self) -> ImageFile.StubHandler | None: + return _handler + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if _handler is None or not hasattr(_handler, "save"): + msg = "BUFR save handler not installed" + raise OSError(msg) + _handler.save(im, fp, filename) + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(BufrStubImageFile.format, BufrStubImageFile, _accept) +Image.register_save(BufrStubImageFile.format, _save) + +Image.register_extension(BufrStubImageFile.format, ".bufr") diff --git a/venv/lib/python3.10/site-packages/PIL/ContainerIO.py b/venv/lib/python3.10/site-packages/PIL/ContainerIO.py new file mode 100644 index 0000000000000000000000000000000000000000..ec9e66c714fbbfec8c597f6127e5d932b0da521f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ContainerIO.py @@ -0,0 +1,173 @@ +# +# The Python Imaging Library. +# $Id$ +# +# a class to read from a container file +# +# History: +# 1995-06-18 fl Created +# 1995-09-07 fl Added readline(), readlines() +# +# Copyright (c) 1997-2001 by Secret Labs AB +# Copyright (c) 1995 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +from collections.abc import Iterable +from typing import IO, AnyStr, NoReturn + + +class ContainerIO(IO[AnyStr]): + """ + A file object that provides read access to a part of an existing + file (for example a TAR file). + """ + + def __init__(self, file: IO[AnyStr], offset: int, length: int) -> None: + """ + Create file object. + + :param file: Existing file. + :param offset: Start of region, in bytes. + :param length: Size of region, in bytes. + """ + self.fh: IO[AnyStr] = file + self.pos = 0 + self.offset = offset + self.length = length + self.fh.seek(offset) + + ## + # Always false. + + def isatty(self) -> bool: + return False + + def seekable(self) -> bool: + return True + + def seek(self, offset: int, mode: int = io.SEEK_SET) -> int: + """ + Move file pointer. + + :param offset: Offset in bytes. + :param mode: Starting position. Use 0 for beginning of region, 1 + for current offset, and 2 for end of region. You cannot move + the pointer outside the defined region. + :returns: Offset from start of region, in bytes. + """ + if mode == 1: + self.pos = self.pos + offset + elif mode == 2: + self.pos = self.length + offset + else: + self.pos = offset + # clamp + self.pos = max(0, min(self.pos, self.length)) + self.fh.seek(self.offset + self.pos) + return self.pos + + def tell(self) -> int: + """ + Get current file pointer. + + :returns: Offset from start of region, in bytes. + """ + return self.pos + + def readable(self) -> bool: + return True + + def read(self, n: int = -1) -> AnyStr: + """ + Read data. + + :param n: Number of bytes to read. If omitted, zero or negative, + read until end of region. + :returns: An 8-bit string. + """ + if n > 0: + n = min(n, self.length - self.pos) + else: + n = self.length - self.pos + if n <= 0: # EOF + return b"" if "b" in self.fh.mode else "" # type: ignore[return-value] + self.pos = self.pos + n + return self.fh.read(n) + + def readline(self, n: int = -1) -> AnyStr: + """ + Read a line of text. + + :param n: Number of bytes to read. If omitted, zero or negative, + read until end of line. + :returns: An 8-bit string. + """ + s: AnyStr = b"" if "b" in self.fh.mode else "" # type: ignore[assignment] + newline_character = b"\n" if "b" in self.fh.mode else "\n" + while True: + c = self.read(1) + if not c: + break + s = s + c + if c == newline_character or len(s) == n: + break + return s + + def readlines(self, n: int | None = -1) -> list[AnyStr]: + """ + Read multiple lines of text. + + :param n: Number of lines to read. If omitted, zero, negative or None, + read until end of region. + :returns: A list of 8-bit strings. + """ + lines = [] + while True: + s = self.readline() + if not s: + break + lines.append(s) + if len(lines) == n: + break + return lines + + def writable(self) -> bool: + return False + + def write(self, b: AnyStr) -> NoReturn: + raise NotImplementedError() + + def writelines(self, lines: Iterable[AnyStr]) -> NoReturn: + raise NotImplementedError() + + def truncate(self, size: int | None = None) -> int: + raise NotImplementedError() + + def __enter__(self) -> ContainerIO[AnyStr]: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def __iter__(self) -> ContainerIO[AnyStr]: + return self + + def __next__(self) -> AnyStr: + line = self.readline() + if not line: + msg = "end of region" + raise StopIteration(msg) + return line + + def fileno(self) -> int: + return self.fh.fileno() + + def flush(self) -> None: + self.fh.flush() + + def close(self) -> None: + self.fh.close() diff --git a/venv/lib/python3.10/site-packages/PIL/CurImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/CurImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..9c188e084463dd0607934196ff9ad0dd1e51b210 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/CurImagePlugin.py @@ -0,0 +1,75 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Windows Cursor support for PIL +# +# notes: +# uses BmpImagePlugin.py to read the bitmap data. +# +# history: +# 96-05-27 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import BmpImagePlugin, Image +from ._binary import i16le as i16 +from ._binary import i32le as i32 + +# +# -------------------------------------------------------------------- + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"\0\0\2\0") + + +## +# Image plugin for Windows Cursor files. + + +class CurImageFile(BmpImagePlugin.BmpImageFile): + format = "CUR" + format_description = "Windows Cursor" + + def _open(self) -> None: + assert self.fp is not None + offset = self.fp.tell() + + # check magic + s = self.fp.read(6) + if not _accept(s): + msg = "not a CUR file" + raise SyntaxError(msg) + + # pick the largest cursor in the file + m = b"" + for i in range(i16(s, 4)): + s = self.fp.read(16) + if not m: + m = s + elif s[0] > m[0] and s[1] > m[1]: + m = s + if not m: + msg = "No cursors were found" + raise TypeError(msg) + + # load as bitmap + self._bitmap(i32(m, 12) + offset) + + # patch up the bitmap height + self._size = self.size[0], self.size[1] // 2 + self.tile = [self.tile[0]._replace(extents=(0, 0) + self.size)] + + +# +# -------------------------------------------------------------------- + +Image.register_open(CurImageFile.format, CurImageFile, _accept) + +Image.register_extension(CurImageFile.format, ".cur") diff --git a/venv/lib/python3.10/site-packages/PIL/DcxImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/DcxImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..aea661b9cb6eef184696e377678ee69f66c5f772 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/DcxImagePlugin.py @@ -0,0 +1,83 @@ +# +# The Python Imaging Library. +# $Id$ +# +# DCX file handling +# +# DCX is a container file format defined by Intel, commonly used +# for fax applications. Each DCX file consists of a directory +# (a list of file offsets) followed by a set of (usually 1-bit) +# PCX files. +# +# History: +# 1995-09-09 fl Created +# 1996-03-20 fl Properly derived from PcxImageFile. +# 1998-07-15 fl Renamed offset attribute to avoid name clash +# 2002-07-30 fl Fixed file handling +# +# Copyright (c) 1997-98 by Secret Labs AB. +# Copyright (c) 1995-96 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image +from ._binary import i32le as i32 +from ._util import DeferredError +from .PcxImagePlugin import PcxImageFile + +MAGIC = 0x3ADE68B1 # QUIZ: what's this value, then? + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 4 and i32(prefix) == MAGIC + + +## +# Image plugin for the Intel DCX format. + + +class DcxImageFile(PcxImageFile): + format = "DCX" + format_description = "Intel DCX" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + # Header + s = self.fp.read(4) + if not _accept(s): + msg = "not a DCX file" + raise SyntaxError(msg) + + # Component directory + self._offset = [] + for i in range(1024): + offset = i32(self.fp.read(4)) + if not offset: + break + self._offset.append(offset) + + self._fp = self.fp + self.frame = -1 + self.n_frames = len(self._offset) + self.is_animated = self.n_frames > 1 + self.seek(0) + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self.frame = frame + self.fp = self._fp + self.fp.seek(self._offset[frame]) + PcxImageFile._open(self) + + def tell(self) -> int: + return self.frame + + +Image.register_open(DcxImageFile.format, DcxImageFile, _accept) + +Image.register_extension(DcxImageFile.format, ".dcx") diff --git a/venv/lib/python3.10/site-packages/PIL/DdsImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/DdsImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..f9ade18f9a1edf431524dd86a238f6b0445e6ab0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/DdsImagePlugin.py @@ -0,0 +1,624 @@ +""" +A Pillow plugin for .dds files (S3TC-compressed aka DXTC) +Jerome Leclanche + +Documentation: +https://web.archive.org/web/20170802060935/http://oss.sgi.com/projects/ogl-sample/registry/EXT/texture_compression_s3tc.txt + +The contents of this file are hereby released in the public domain (CC0) +Full text of the CC0 license: +https://creativecommons.org/publicdomain/zero/1.0/ +""" + +from __future__ import annotations + +import io +import struct +import sys +from enum import IntEnum, IntFlag +from typing import IO + +from . import Image, ImageFile, ImagePalette +from ._binary import i32le as i32 +from ._binary import o8 +from ._binary import o32le as o32 + +# Magic ("DDS ") +DDS_MAGIC = 0x20534444 + + +# DDS flags +class DDSD(IntFlag): + CAPS = 0x1 + HEIGHT = 0x2 + WIDTH = 0x4 + PITCH = 0x8 + PIXELFORMAT = 0x1000 + MIPMAPCOUNT = 0x20000 + LINEARSIZE = 0x80000 + DEPTH = 0x800000 + + +# DDS caps +class DDSCAPS(IntFlag): + COMPLEX = 0x8 + TEXTURE = 0x1000 + MIPMAP = 0x400000 + + +class DDSCAPS2(IntFlag): + CUBEMAP = 0x200 + CUBEMAP_POSITIVEX = 0x400 + CUBEMAP_NEGATIVEX = 0x800 + CUBEMAP_POSITIVEY = 0x1000 + CUBEMAP_NEGATIVEY = 0x2000 + CUBEMAP_POSITIVEZ = 0x4000 + CUBEMAP_NEGATIVEZ = 0x8000 + VOLUME = 0x200000 + + +# Pixel Format +class DDPF(IntFlag): + ALPHAPIXELS = 0x1 + ALPHA = 0x2 + FOURCC = 0x4 + PALETTEINDEXED8 = 0x20 + RGB = 0x40 + LUMINANCE = 0x20000 + + +# dxgiformat.h +class DXGI_FORMAT(IntEnum): + UNKNOWN = 0 + R32G32B32A32_TYPELESS = 1 + R32G32B32A32_FLOAT = 2 + R32G32B32A32_UINT = 3 + R32G32B32A32_SINT = 4 + R32G32B32_TYPELESS = 5 + R32G32B32_FLOAT = 6 + R32G32B32_UINT = 7 + R32G32B32_SINT = 8 + R16G16B16A16_TYPELESS = 9 + R16G16B16A16_FLOAT = 10 + R16G16B16A16_UNORM = 11 + R16G16B16A16_UINT = 12 + R16G16B16A16_SNORM = 13 + R16G16B16A16_SINT = 14 + R32G32_TYPELESS = 15 + R32G32_FLOAT = 16 + R32G32_UINT = 17 + R32G32_SINT = 18 + R32G8X24_TYPELESS = 19 + D32_FLOAT_S8X24_UINT = 20 + R32_FLOAT_X8X24_TYPELESS = 21 + X32_TYPELESS_G8X24_UINT = 22 + R10G10B10A2_TYPELESS = 23 + R10G10B10A2_UNORM = 24 + R10G10B10A2_UINT = 25 + R11G11B10_FLOAT = 26 + R8G8B8A8_TYPELESS = 27 + R8G8B8A8_UNORM = 28 + R8G8B8A8_UNORM_SRGB = 29 + R8G8B8A8_UINT = 30 + R8G8B8A8_SNORM = 31 + R8G8B8A8_SINT = 32 + R16G16_TYPELESS = 33 + R16G16_FLOAT = 34 + R16G16_UNORM = 35 + R16G16_UINT = 36 + R16G16_SNORM = 37 + R16G16_SINT = 38 + R32_TYPELESS = 39 + D32_FLOAT = 40 + R32_FLOAT = 41 + R32_UINT = 42 + R32_SINT = 43 + R24G8_TYPELESS = 44 + D24_UNORM_S8_UINT = 45 + R24_UNORM_X8_TYPELESS = 46 + X24_TYPELESS_G8_UINT = 47 + R8G8_TYPELESS = 48 + R8G8_UNORM = 49 + R8G8_UINT = 50 + R8G8_SNORM = 51 + R8G8_SINT = 52 + R16_TYPELESS = 53 + R16_FLOAT = 54 + D16_UNORM = 55 + R16_UNORM = 56 + R16_UINT = 57 + R16_SNORM = 58 + R16_SINT = 59 + R8_TYPELESS = 60 + R8_UNORM = 61 + R8_UINT = 62 + R8_SNORM = 63 + R8_SINT = 64 + A8_UNORM = 65 + R1_UNORM = 66 + R9G9B9E5_SHAREDEXP = 67 + R8G8_B8G8_UNORM = 68 + G8R8_G8B8_UNORM = 69 + BC1_TYPELESS = 70 + BC1_UNORM = 71 + BC1_UNORM_SRGB = 72 + BC2_TYPELESS = 73 + BC2_UNORM = 74 + BC2_UNORM_SRGB = 75 + BC3_TYPELESS = 76 + BC3_UNORM = 77 + BC3_UNORM_SRGB = 78 + BC4_TYPELESS = 79 + BC4_UNORM = 80 + BC4_SNORM = 81 + BC5_TYPELESS = 82 + BC5_UNORM = 83 + BC5_SNORM = 84 + B5G6R5_UNORM = 85 + B5G5R5A1_UNORM = 86 + B8G8R8A8_UNORM = 87 + B8G8R8X8_UNORM = 88 + R10G10B10_XR_BIAS_A2_UNORM = 89 + B8G8R8A8_TYPELESS = 90 + B8G8R8A8_UNORM_SRGB = 91 + B8G8R8X8_TYPELESS = 92 + B8G8R8X8_UNORM_SRGB = 93 + BC6H_TYPELESS = 94 + BC6H_UF16 = 95 + BC6H_SF16 = 96 + BC7_TYPELESS = 97 + BC7_UNORM = 98 + BC7_UNORM_SRGB = 99 + AYUV = 100 + Y410 = 101 + Y416 = 102 + NV12 = 103 + P010 = 104 + P016 = 105 + OPAQUE_420 = 106 + YUY2 = 107 + Y210 = 108 + Y216 = 109 + NV11 = 110 + AI44 = 111 + IA44 = 112 + P8 = 113 + A8P8 = 114 + B4G4R4A4_UNORM = 115 + P208 = 130 + V208 = 131 + V408 = 132 + SAMPLER_FEEDBACK_MIN_MIP_OPAQUE = 189 + SAMPLER_FEEDBACK_MIP_REGION_USED_OPAQUE = 190 + + +class D3DFMT(IntEnum): + UNKNOWN = 0 + R8G8B8 = 20 + A8R8G8B8 = 21 + X8R8G8B8 = 22 + R5G6B5 = 23 + X1R5G5B5 = 24 + A1R5G5B5 = 25 + A4R4G4B4 = 26 + R3G3B2 = 27 + A8 = 28 + A8R3G3B2 = 29 + X4R4G4B4 = 30 + A2B10G10R10 = 31 + A8B8G8R8 = 32 + X8B8G8R8 = 33 + G16R16 = 34 + A2R10G10B10 = 35 + A16B16G16R16 = 36 + A8P8 = 40 + P8 = 41 + L8 = 50 + A8L8 = 51 + A4L4 = 52 + V8U8 = 60 + L6V5U5 = 61 + X8L8V8U8 = 62 + Q8W8V8U8 = 63 + V16U16 = 64 + A2W10V10U10 = 67 + D16_LOCKABLE = 70 + D32 = 71 + D15S1 = 73 + D24S8 = 75 + D24X8 = 77 + D24X4S4 = 79 + D16 = 80 + D32F_LOCKABLE = 82 + D24FS8 = 83 + D32_LOCKABLE = 84 + S8_LOCKABLE = 85 + L16 = 81 + VERTEXDATA = 100 + INDEX16 = 101 + INDEX32 = 102 + Q16W16V16U16 = 110 + R16F = 111 + G16R16F = 112 + A16B16G16R16F = 113 + R32F = 114 + G32R32F = 115 + A32B32G32R32F = 116 + CxV8U8 = 117 + A1 = 118 + A2B10G10R10_XR_BIAS = 119 + BINARYBUFFER = 199 + + UYVY = i32(b"UYVY") + R8G8_B8G8 = i32(b"RGBG") + YUY2 = i32(b"YUY2") + G8R8_G8B8 = i32(b"GRGB") + DXT1 = i32(b"DXT1") + DXT2 = i32(b"DXT2") + DXT3 = i32(b"DXT3") + DXT4 = i32(b"DXT4") + DXT5 = i32(b"DXT5") + DX10 = i32(b"DX10") + BC4S = i32(b"BC4S") + BC4U = i32(b"BC4U") + BC5S = i32(b"BC5S") + BC5U = i32(b"BC5U") + ATI1 = i32(b"ATI1") + ATI2 = i32(b"ATI2") + MULTI2_ARGB8 = i32(b"MET1") + + +# Backward compatibility layer +module = sys.modules[__name__] +for item in DDSD: + assert item.name is not None + setattr(module, f"DDSD_{item.name}", item.value) +for item1 in DDSCAPS: + assert item1.name is not None + setattr(module, f"DDSCAPS_{item1.name}", item1.value) +for item2 in DDSCAPS2: + assert item2.name is not None + setattr(module, f"DDSCAPS2_{item2.name}", item2.value) +for item3 in DDPF: + assert item3.name is not None + setattr(module, f"DDPF_{item3.name}", item3.value) + +DDS_FOURCC = DDPF.FOURCC +DDS_RGB = DDPF.RGB +DDS_RGBA = DDPF.RGB | DDPF.ALPHAPIXELS +DDS_LUMINANCE = DDPF.LUMINANCE +DDS_LUMINANCEA = DDPF.LUMINANCE | DDPF.ALPHAPIXELS +DDS_ALPHA = DDPF.ALPHA +DDS_PAL8 = DDPF.PALETTEINDEXED8 + +DDS_HEADER_FLAGS_TEXTURE = DDSD.CAPS | DDSD.HEIGHT | DDSD.WIDTH | DDSD.PIXELFORMAT +DDS_HEADER_FLAGS_MIPMAP = DDSD.MIPMAPCOUNT +DDS_HEADER_FLAGS_VOLUME = DDSD.DEPTH +DDS_HEADER_FLAGS_PITCH = DDSD.PITCH +DDS_HEADER_FLAGS_LINEARSIZE = DDSD.LINEARSIZE + +DDS_HEIGHT = DDSD.HEIGHT +DDS_WIDTH = DDSD.WIDTH + +DDS_SURFACE_FLAGS_TEXTURE = DDSCAPS.TEXTURE +DDS_SURFACE_FLAGS_MIPMAP = DDSCAPS.COMPLEX | DDSCAPS.MIPMAP +DDS_SURFACE_FLAGS_CUBEMAP = DDSCAPS.COMPLEX + +DDS_CUBEMAP_POSITIVEX = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_POSITIVEX +DDS_CUBEMAP_NEGATIVEX = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_NEGATIVEX +DDS_CUBEMAP_POSITIVEY = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_POSITIVEY +DDS_CUBEMAP_NEGATIVEY = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_NEGATIVEY +DDS_CUBEMAP_POSITIVEZ = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_POSITIVEZ +DDS_CUBEMAP_NEGATIVEZ = DDSCAPS2.CUBEMAP | DDSCAPS2.CUBEMAP_NEGATIVEZ + +DXT1_FOURCC = D3DFMT.DXT1 +DXT3_FOURCC = D3DFMT.DXT3 +DXT5_FOURCC = D3DFMT.DXT5 + +DXGI_FORMAT_R8G8B8A8_TYPELESS = DXGI_FORMAT.R8G8B8A8_TYPELESS +DXGI_FORMAT_R8G8B8A8_UNORM = DXGI_FORMAT.R8G8B8A8_UNORM +DXGI_FORMAT_R8G8B8A8_UNORM_SRGB = DXGI_FORMAT.R8G8B8A8_UNORM_SRGB +DXGI_FORMAT_BC5_TYPELESS = DXGI_FORMAT.BC5_TYPELESS +DXGI_FORMAT_BC5_UNORM = DXGI_FORMAT.BC5_UNORM +DXGI_FORMAT_BC5_SNORM = DXGI_FORMAT.BC5_SNORM +DXGI_FORMAT_BC6H_UF16 = DXGI_FORMAT.BC6H_UF16 +DXGI_FORMAT_BC6H_SF16 = DXGI_FORMAT.BC6H_SF16 +DXGI_FORMAT_BC7_TYPELESS = DXGI_FORMAT.BC7_TYPELESS +DXGI_FORMAT_BC7_UNORM = DXGI_FORMAT.BC7_UNORM +DXGI_FORMAT_BC7_UNORM_SRGB = DXGI_FORMAT.BC7_UNORM_SRGB + + +class DdsImageFile(ImageFile.ImageFile): + format = "DDS" + format_description = "DirectDraw Surface" + + def _open(self) -> None: + if not _accept(self.fp.read(4)): + msg = "not a DDS file" + raise SyntaxError(msg) + (header_size,) = struct.unpack(" None: + pass + + +class DdsRgbDecoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + bitcount, masks = self.args + + # Some masks will be padded with zeros, e.g. R 0b11 G 0b1100 + # Calculate how many zeros each mask is padded with + mask_offsets = [] + # And the maximum value of each channel without the padding + mask_totals = [] + for mask in masks: + offset = 0 + if mask != 0: + while mask >> (offset + 1) << (offset + 1) == mask: + offset += 1 + mask_offsets.append(offset) + mask_totals.append(mask >> offset) + + data = bytearray() + bytecount = bitcount // 8 + dest_length = self.state.xsize * self.state.ysize * len(masks) + while len(data) < dest_length: + value = int.from_bytes(self.fd.read(bytecount), "little") + for i, mask in enumerate(masks): + masked_value = value & mask + # Remove the zero padding, and scale it to 8 bits + data += o8( + int(((masked_value >> mask_offsets[i]) / mask_totals[i]) * 255) + ) + self.set_as_raw(data) + return -1, 0 + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode not in ("RGB", "RGBA", "L", "LA"): + msg = f"cannot write mode {im.mode} as DDS" + raise OSError(msg) + + flags = DDSD.CAPS | DDSD.HEIGHT | DDSD.WIDTH | DDSD.PIXELFORMAT + bitcount = len(im.getbands()) * 8 + pixel_format = im.encoderinfo.get("pixel_format") + args: tuple[int] | str + if pixel_format: + codec_name = "bcn" + flags |= DDSD.LINEARSIZE + pitch = (im.width + 3) * 4 + rgba_mask = [0, 0, 0, 0] + pixel_flags = DDPF.FOURCC + if pixel_format == "DXT1": + fourcc = D3DFMT.DXT1 + args = (1,) + elif pixel_format == "DXT3": + fourcc = D3DFMT.DXT3 + args = (2,) + elif pixel_format == "DXT5": + fourcc = D3DFMT.DXT5 + args = (3,) + else: + fourcc = D3DFMT.DX10 + if pixel_format == "BC2": + args = (2,) + dxgi_format = DXGI_FORMAT.BC2_TYPELESS + elif pixel_format == "BC3": + args = (3,) + dxgi_format = DXGI_FORMAT.BC3_TYPELESS + elif pixel_format == "BC5": + args = (5,) + dxgi_format = DXGI_FORMAT.BC5_TYPELESS + if im.mode != "RGB": + msg = "only RGB mode can be written as BC5" + raise OSError(msg) + else: + msg = f"cannot write pixel format {pixel_format}" + raise OSError(msg) + else: + codec_name = "raw" + flags |= DDSD.PITCH + pitch = (im.width * bitcount + 7) // 8 + + alpha = im.mode[-1] == "A" + if im.mode[0] == "L": + pixel_flags = DDPF.LUMINANCE + args = im.mode + if alpha: + rgba_mask = [0x000000FF, 0x000000FF, 0x000000FF] + else: + rgba_mask = [0xFF000000, 0xFF000000, 0xFF000000] + else: + pixel_flags = DDPF.RGB + args = im.mode[::-1] + rgba_mask = [0x00FF0000, 0x0000FF00, 0x000000FF] + + if alpha: + r, g, b, a = im.split() + im = Image.merge("RGBA", (a, r, g, b)) + if alpha: + pixel_flags |= DDPF.ALPHAPIXELS + rgba_mask.append(0xFF000000 if alpha else 0) + + fourcc = D3DFMT.UNKNOWN + fp.write( + o32(DDS_MAGIC) + + struct.pack( + "<7I", + 124, # header size + flags, # flags + im.height, + im.width, + pitch, + 0, # depth + 0, # mipmaps + ) + + struct.pack("11I", *((0,) * 11)) # reserved + # pfsize, pfflags, fourcc, bitcount + + struct.pack("<4I", 32, pixel_flags, fourcc, bitcount) + + struct.pack("<4I", *rgba_mask) # dwRGBABitMask + + struct.pack("<5I", DDSCAPS.TEXTURE, 0, 0, 0, 0) + ) + if fourcc == D3DFMT.DX10: + fp.write( + # dxgi_format, 2D resource, misc, array size, straight alpha + struct.pack("<5I", dxgi_format, 3, 0, 0, 1) + ) + ImageFile._save(im, fp, [ImageFile._Tile(codec_name, (0, 0) + im.size, 0, args)]) + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"DDS ") + + +Image.register_open(DdsImageFile.format, DdsImageFile, _accept) +Image.register_decoder("dds_rgb", DdsRgbDecoder) +Image.register_save(DdsImageFile.format, _save) +Image.register_extension(DdsImageFile.format, ".dds") diff --git a/venv/lib/python3.10/site-packages/PIL/EpsImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/EpsImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..69f3062b4d4c18887970ab1980bbc9880b1d6621 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/EpsImagePlugin.py @@ -0,0 +1,479 @@ +# +# The Python Imaging Library. +# $Id$ +# +# EPS file handling +# +# History: +# 1995-09-01 fl Created (0.1) +# 1996-05-18 fl Don't choke on "atend" fields, Ghostscript interface (0.2) +# 1996-08-22 fl Don't choke on floating point BoundingBox values +# 1996-08-23 fl Handle files from Macintosh (0.3) +# 2001-02-17 fl Use 're' instead of 'regex' (Python 2.1) (0.4) +# 2003-09-07 fl Check gs.close status (from Federico Di Gregorio) (0.5) +# 2014-05-07 e Handling of EPS with binary preview and fixed resolution +# resizing +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1995-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import os +import re +import subprocess +import sys +import tempfile +from typing import IO + +from . import Image, ImageFile +from ._binary import i32le as i32 + +# -------------------------------------------------------------------- + + +split = re.compile(r"^%%([^:]*):[ \t]*(.*)[ \t]*$") +field = re.compile(r"^%[%!\w]([^:]*)[ \t]*$") + +gs_binary: str | bool | None = None +gs_windows_binary = None + + +def has_ghostscript() -> bool: + global gs_binary, gs_windows_binary + if gs_binary is None: + if sys.platform.startswith("win"): + if gs_windows_binary is None: + import shutil + + for binary in ("gswin32c", "gswin64c", "gs"): + if shutil.which(binary) is not None: + gs_windows_binary = binary + break + else: + gs_windows_binary = False + gs_binary = gs_windows_binary + else: + try: + subprocess.check_call(["gs", "--version"], stdout=subprocess.DEVNULL) + gs_binary = "gs" + except OSError: + gs_binary = False + return gs_binary is not False + + +def Ghostscript( + tile: list[ImageFile._Tile], + size: tuple[int, int], + fp: IO[bytes], + scale: int = 1, + transparency: bool = False, +) -> Image.core.ImagingCore: + """Render an image using Ghostscript""" + global gs_binary + if not has_ghostscript(): + msg = "Unable to locate Ghostscript on paths" + raise OSError(msg) + assert isinstance(gs_binary, str) + + # Unpack decoder tile + args = tile[0].args + assert isinstance(args, tuple) + length, bbox = args + + # Hack to support hi-res rendering + scale = int(scale) or 1 + width = size[0] * scale + height = size[1] * scale + # resolution is dependent on bbox and size + res_x = 72.0 * width / (bbox[2] - bbox[0]) + res_y = 72.0 * height / (bbox[3] - bbox[1]) + + out_fd, outfile = tempfile.mkstemp() + os.close(out_fd) + + infile_temp = None + if hasattr(fp, "name") and os.path.exists(fp.name): + infile = fp.name + else: + in_fd, infile_temp = tempfile.mkstemp() + os.close(in_fd) + infile = infile_temp + + # Ignore length and offset! + # Ghostscript can read it + # Copy whole file to read in Ghostscript + with open(infile_temp, "wb") as f: + # fetch length of fp + fp.seek(0, io.SEEK_END) + fsize = fp.tell() + # ensure start position + # go back + fp.seek(0) + lengthfile = fsize + while lengthfile > 0: + s = fp.read(min(lengthfile, 100 * 1024)) + if not s: + break + lengthfile -= len(s) + f.write(s) + + if transparency: + # "RGBA" + device = "pngalpha" + else: + # "pnmraw" automatically chooses between + # PBM ("1"), PGM ("L"), and PPM ("RGB"). + device = "pnmraw" + + # Build Ghostscript command + command = [ + gs_binary, + "-q", # quiet mode + f"-g{width:d}x{height:d}", # set output geometry (pixels) + f"-r{res_x:f}x{res_y:f}", # set input DPI (dots per inch) + "-dBATCH", # exit after processing + "-dNOPAUSE", # don't pause between pages + "-dSAFER", # safe mode + f"-sDEVICE={device}", + f"-sOutputFile={outfile}", # output file + # adjust for image origin + "-c", + f"{-bbox[0]} {-bbox[1]} translate", + "-f", + infile, # input file + # showpage (see https://bugs.ghostscript.com/show_bug.cgi?id=698272) + "-c", + "showpage", + ] + + # push data through Ghostscript + try: + startupinfo = None + if sys.platform.startswith("win"): + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + subprocess.check_call(command, startupinfo=startupinfo) + with Image.open(outfile) as out_im: + out_im.load() + return out_im.im.copy() + finally: + try: + os.unlink(outfile) + if infile_temp: + os.unlink(infile_temp) + except OSError: + pass + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"%!PS") or ( + len(prefix) >= 4 and i32(prefix) == 0xC6D3D0C5 + ) + + +## +# Image plugin for Encapsulated PostScript. This plugin supports only +# a few variants of this format. + + +class EpsImageFile(ImageFile.ImageFile): + """EPS File Parser for the Python Imaging Library""" + + format = "EPS" + format_description = "Encapsulated Postscript" + + mode_map = {1: "L", 2: "LAB", 3: "RGB", 4: "CMYK"} + + def _open(self) -> None: + (length, offset) = self._find_offset(self.fp) + + # go to offset - start of "%!PS" + self.fp.seek(offset) + + self._mode = "RGB" + + # When reading header comments, the first comment is used. + # When reading trailer comments, the last comment is used. + bounding_box: list[int] | None = None + imagedata_size: tuple[int, int] | None = None + + byte_arr = bytearray(255) + bytes_mv = memoryview(byte_arr) + bytes_read = 0 + reading_header_comments = True + reading_trailer_comments = False + trailer_reached = False + + def check_required_header_comments() -> None: + """ + The EPS specification requires that some headers exist. + This should be checked when the header comments formally end, + when image data starts, or when the file ends, whichever comes first. + """ + if "PS-Adobe" not in self.info: + msg = 'EPS header missing "%!PS-Adobe" comment' + raise SyntaxError(msg) + if "BoundingBox" not in self.info: + msg = 'EPS header missing "%%BoundingBox" comment' + raise SyntaxError(msg) + + def read_comment(s: str) -> bool: + nonlocal bounding_box, reading_trailer_comments + try: + m = split.match(s) + except re.error as e: + msg = "not an EPS file" + raise SyntaxError(msg) from e + + if not m: + return False + + k, v = m.group(1, 2) + self.info[k] = v + if k == "BoundingBox": + if v == "(atend)": + reading_trailer_comments = True + elif not bounding_box or (trailer_reached and reading_trailer_comments): + try: + # Note: The DSC spec says that BoundingBox + # fields should be integers, but some drivers + # put floating point values there anyway. + bounding_box = [int(float(i)) for i in v.split()] + except Exception: + pass + return True + + while True: + byte = self.fp.read(1) + if byte == b"": + # if we didn't read a byte we must be at the end of the file + if bytes_read == 0: + if reading_header_comments: + check_required_header_comments() + break + elif byte in b"\r\n": + # if we read a line ending character, ignore it and parse what + # we have already read. if we haven't read any other characters, + # continue reading + if bytes_read == 0: + continue + else: + # ASCII/hexadecimal lines in an EPS file must not exceed + # 255 characters, not including line ending characters + if bytes_read >= 255: + # only enforce this for lines starting with a "%", + # otherwise assume it's binary data + if byte_arr[0] == ord("%"): + msg = "not an EPS file" + raise SyntaxError(msg) + else: + if reading_header_comments: + check_required_header_comments() + reading_header_comments = False + # reset bytes_read so we can keep reading + # data until the end of the line + bytes_read = 0 + byte_arr[bytes_read] = byte[0] + bytes_read += 1 + continue + + if reading_header_comments: + # Load EPS header + + # if this line doesn't start with a "%", + # or does start with "%%EndComments", + # then we've reached the end of the header/comments + if byte_arr[0] != ord("%") or bytes_mv[:13] == b"%%EndComments": + check_required_header_comments() + reading_header_comments = False + continue + + s = str(bytes_mv[:bytes_read], "latin-1") + if not read_comment(s): + m = field.match(s) + if m: + k = m.group(1) + if k.startswith("PS-Adobe"): + self.info["PS-Adobe"] = k[9:] + else: + self.info[k] = "" + elif s[0] == "%": + # handle non-DSC PostScript comments that some + # tools mistakenly put in the Comments section + pass + else: + msg = "bad EPS header" + raise OSError(msg) + elif bytes_mv[:11] == b"%ImageData:": + # Check for an "ImageData" descriptor + # https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#50577413_pgfId-1035096 + + # If we've already read an "ImageData" descriptor, + # don't read another one. + if imagedata_size: + bytes_read = 0 + continue + + # Values: + # columns + # rows + # bit depth (1 or 8) + # mode (1: L, 2: LAB, 3: RGB, 4: CMYK) + # number of padding channels + # block size (number of bytes per row per channel) + # binary/ascii (1: binary, 2: ascii) + # data start identifier (the image data follows after a single line + # consisting only of this quoted value) + image_data_values = byte_arr[11:bytes_read].split(None, 7) + columns, rows, bit_depth, mode_id = ( + int(value) for value in image_data_values[:4] + ) + + if bit_depth == 1: + self._mode = "1" + elif bit_depth == 8: + try: + self._mode = self.mode_map[mode_id] + except ValueError: + break + else: + break + + # Parse the columns and rows after checking the bit depth and mode + # in case the bit depth and/or mode are invalid. + imagedata_size = columns, rows + elif bytes_mv[:5] == b"%%EOF": + break + elif trailer_reached and reading_trailer_comments: + # Load EPS trailer + s = str(bytes_mv[:bytes_read], "latin-1") + read_comment(s) + elif bytes_mv[:9] == b"%%Trailer": + trailer_reached = True + elif bytes_mv[:14] == b"%%BeginBinary:": + bytecount = int(byte_arr[14:bytes_read]) + self.fp.seek(bytecount, os.SEEK_CUR) + bytes_read = 0 + + # A "BoundingBox" is always required, + # even if an "ImageData" descriptor size exists. + if not bounding_box: + msg = "cannot determine EPS bounding box" + raise OSError(msg) + + # An "ImageData" size takes precedence over the "BoundingBox". + self._size = imagedata_size or ( + bounding_box[2] - bounding_box[0], + bounding_box[3] - bounding_box[1], + ) + + self.tile = [ + ImageFile._Tile("eps", (0, 0) + self.size, offset, (length, bounding_box)) + ] + + def _find_offset(self, fp: IO[bytes]) -> tuple[int, int]: + s = fp.read(4) + + if s == b"%!PS": + # for HEAD without binary preview + fp.seek(0, io.SEEK_END) + length = fp.tell() + offset = 0 + elif i32(s) == 0xC6D3D0C5: + # FIX for: Some EPS file not handled correctly / issue #302 + # EPS can contain binary data + # or start directly with latin coding + # more info see: + # https://web.archive.org/web/20160528181353/http://partners.adobe.com/public/developer/en/ps/5002.EPSF_Spec.pdf + s = fp.read(8) + offset = i32(s) + length = i32(s, 4) + else: + msg = "not an EPS file" + raise SyntaxError(msg) + + return length, offset + + def load( + self, scale: int = 1, transparency: bool = False + ) -> Image.core.PixelAccess | None: + # Load EPS via Ghostscript + if self.tile: + self.im = Ghostscript(self.tile, self.size, self.fp, scale, transparency) + self._mode = self.im.mode + self._size = self.im.size + self.tile = [] + return Image.Image.load(self) + + def load_seek(self, pos: int) -> None: + # we can't incrementally load, so force ImageFile.parser to + # use our custom load method by defining this method. + pass + + +# -------------------------------------------------------------------- + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes, eps: int = 1) -> None: + """EPS Writer for the Python Imaging Library.""" + + # make sure image data is available + im.load() + + # determine PostScript image mode + if im.mode == "L": + operator = (8, 1, b"image") + elif im.mode == "RGB": + operator = (8, 3, b"false 3 colorimage") + elif im.mode == "CMYK": + operator = (8, 4, b"false 4 colorimage") + else: + msg = "image mode is not supported" + raise ValueError(msg) + + if eps: + # write EPS header + fp.write(b"%!PS-Adobe-3.0 EPSF-3.0\n") + fp.write(b"%%Creator: PIL 0.1 EpsEncode\n") + # fp.write("%%CreationDate: %s"...) + fp.write(b"%%%%BoundingBox: 0 0 %d %d\n" % im.size) + fp.write(b"%%Pages: 1\n") + fp.write(b"%%EndComments\n") + fp.write(b"%%Page: 1 1\n") + fp.write(b"%%ImageData: %d %d " % im.size) + fp.write(b'%d %d 0 1 1 "%s"\n' % operator) + + # image header + fp.write(b"gsave\n") + fp.write(b"10 dict begin\n") + fp.write(b"/buf %d string def\n" % (im.size[0] * operator[1])) + fp.write(b"%d %d scale\n" % im.size) + fp.write(b"%d %d 8\n" % im.size) # <= bits + fp.write(b"[%d 0 0 -%d 0 %d]\n" % (im.size[0], im.size[1], im.size[1])) + fp.write(b"{ currentfile buf readhexstring pop } bind\n") + fp.write(operator[2] + b"\n") + if hasattr(fp, "flush"): + fp.flush() + + ImageFile._save(im, fp, [ImageFile._Tile("eps", (0, 0) + im.size)]) + + fp.write(b"\n%%%%EndBinary\n") + fp.write(b"grestore end\n") + if hasattr(fp, "flush"): + fp.flush() + + +# -------------------------------------------------------------------- + + +Image.register_open(EpsImageFile.format, EpsImageFile, _accept) + +Image.register_save(EpsImageFile.format, _save) + +Image.register_extensions(EpsImageFile.format, [".ps", ".eps"]) + +Image.register_mime(EpsImageFile.format, "application/postscript") diff --git a/venv/lib/python3.10/site-packages/PIL/ExifTags.py b/venv/lib/python3.10/site-packages/PIL/ExifTags.py new file mode 100644 index 0000000000000000000000000000000000000000..2280d5ce84b9badabe16cfb0db37f739d50d51c6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ExifTags.py @@ -0,0 +1,382 @@ +# +# The Python Imaging Library. +# $Id$ +# +# EXIF tags +# +# Copyright (c) 2003 by Secret Labs AB +# +# See the README file for information on usage and redistribution. +# + +""" +This module provides constants and clear-text names for various +well-known EXIF tags. +""" +from __future__ import annotations + +from enum import IntEnum + + +class Base(IntEnum): + # possibly incomplete + InteropIndex = 0x0001 + ProcessingSoftware = 0x000B + NewSubfileType = 0x00FE + SubfileType = 0x00FF + ImageWidth = 0x0100 + ImageLength = 0x0101 + BitsPerSample = 0x0102 + Compression = 0x0103 + PhotometricInterpretation = 0x0106 + Thresholding = 0x0107 + CellWidth = 0x0108 + CellLength = 0x0109 + FillOrder = 0x010A + DocumentName = 0x010D + ImageDescription = 0x010E + Make = 0x010F + Model = 0x0110 + StripOffsets = 0x0111 + Orientation = 0x0112 + SamplesPerPixel = 0x0115 + RowsPerStrip = 0x0116 + StripByteCounts = 0x0117 + MinSampleValue = 0x0118 + MaxSampleValue = 0x0119 + XResolution = 0x011A + YResolution = 0x011B + PlanarConfiguration = 0x011C + PageName = 0x011D + FreeOffsets = 0x0120 + FreeByteCounts = 0x0121 + GrayResponseUnit = 0x0122 + GrayResponseCurve = 0x0123 + T4Options = 0x0124 + T6Options = 0x0125 + ResolutionUnit = 0x0128 + PageNumber = 0x0129 + TransferFunction = 0x012D + Software = 0x0131 + DateTime = 0x0132 + Artist = 0x013B + HostComputer = 0x013C + Predictor = 0x013D + WhitePoint = 0x013E + PrimaryChromaticities = 0x013F + ColorMap = 0x0140 + HalftoneHints = 0x0141 + TileWidth = 0x0142 + TileLength = 0x0143 + TileOffsets = 0x0144 + TileByteCounts = 0x0145 + SubIFDs = 0x014A + InkSet = 0x014C + InkNames = 0x014D + NumberOfInks = 0x014E + DotRange = 0x0150 + TargetPrinter = 0x0151 + ExtraSamples = 0x0152 + SampleFormat = 0x0153 + SMinSampleValue = 0x0154 + SMaxSampleValue = 0x0155 + TransferRange = 0x0156 + ClipPath = 0x0157 + XClipPathUnits = 0x0158 + YClipPathUnits = 0x0159 + Indexed = 0x015A + JPEGTables = 0x015B + OPIProxy = 0x015F + JPEGProc = 0x0200 + JpegIFOffset = 0x0201 + JpegIFByteCount = 0x0202 + JpegRestartInterval = 0x0203 + JpegLosslessPredictors = 0x0205 + JpegPointTransforms = 0x0206 + JpegQTables = 0x0207 + JpegDCTables = 0x0208 + JpegACTables = 0x0209 + YCbCrCoefficients = 0x0211 + YCbCrSubSampling = 0x0212 + YCbCrPositioning = 0x0213 + ReferenceBlackWhite = 0x0214 + XMLPacket = 0x02BC + RelatedImageFileFormat = 0x1000 + RelatedImageWidth = 0x1001 + RelatedImageLength = 0x1002 + Rating = 0x4746 + RatingPercent = 0x4749 + ImageID = 0x800D + CFARepeatPatternDim = 0x828D + BatteryLevel = 0x828F + Copyright = 0x8298 + ExposureTime = 0x829A + FNumber = 0x829D + IPTCNAA = 0x83BB + ImageResources = 0x8649 + ExifOffset = 0x8769 + InterColorProfile = 0x8773 + ExposureProgram = 0x8822 + SpectralSensitivity = 0x8824 + GPSInfo = 0x8825 + ISOSpeedRatings = 0x8827 + OECF = 0x8828 + Interlace = 0x8829 + TimeZoneOffset = 0x882A + SelfTimerMode = 0x882B + SensitivityType = 0x8830 + StandardOutputSensitivity = 0x8831 + RecommendedExposureIndex = 0x8832 + ISOSpeed = 0x8833 + ISOSpeedLatitudeyyy = 0x8834 + ISOSpeedLatitudezzz = 0x8835 + ExifVersion = 0x9000 + DateTimeOriginal = 0x9003 + DateTimeDigitized = 0x9004 + OffsetTime = 0x9010 + OffsetTimeOriginal = 0x9011 + OffsetTimeDigitized = 0x9012 + ComponentsConfiguration = 0x9101 + CompressedBitsPerPixel = 0x9102 + ShutterSpeedValue = 0x9201 + ApertureValue = 0x9202 + BrightnessValue = 0x9203 + ExposureBiasValue = 0x9204 + MaxApertureValue = 0x9205 + SubjectDistance = 0x9206 + MeteringMode = 0x9207 + LightSource = 0x9208 + Flash = 0x9209 + FocalLength = 0x920A + Noise = 0x920D + ImageNumber = 0x9211 + SecurityClassification = 0x9212 + ImageHistory = 0x9213 + TIFFEPStandardID = 0x9216 + MakerNote = 0x927C + UserComment = 0x9286 + SubsecTime = 0x9290 + SubsecTimeOriginal = 0x9291 + SubsecTimeDigitized = 0x9292 + AmbientTemperature = 0x9400 + Humidity = 0x9401 + Pressure = 0x9402 + WaterDepth = 0x9403 + Acceleration = 0x9404 + CameraElevationAngle = 0x9405 + XPTitle = 0x9C9B + XPComment = 0x9C9C + XPAuthor = 0x9C9D + XPKeywords = 0x9C9E + XPSubject = 0x9C9F + FlashPixVersion = 0xA000 + ColorSpace = 0xA001 + ExifImageWidth = 0xA002 + ExifImageHeight = 0xA003 + RelatedSoundFile = 0xA004 + ExifInteroperabilityOffset = 0xA005 + FlashEnergy = 0xA20B + SpatialFrequencyResponse = 0xA20C + FocalPlaneXResolution = 0xA20E + FocalPlaneYResolution = 0xA20F + FocalPlaneResolutionUnit = 0xA210 + SubjectLocation = 0xA214 + ExposureIndex = 0xA215 + SensingMethod = 0xA217 + FileSource = 0xA300 + SceneType = 0xA301 + CFAPattern = 0xA302 + CustomRendered = 0xA401 + ExposureMode = 0xA402 + WhiteBalance = 0xA403 + DigitalZoomRatio = 0xA404 + FocalLengthIn35mmFilm = 0xA405 + SceneCaptureType = 0xA406 + GainControl = 0xA407 + Contrast = 0xA408 + Saturation = 0xA409 + Sharpness = 0xA40A + DeviceSettingDescription = 0xA40B + SubjectDistanceRange = 0xA40C + ImageUniqueID = 0xA420 + CameraOwnerName = 0xA430 + BodySerialNumber = 0xA431 + LensSpecification = 0xA432 + LensMake = 0xA433 + LensModel = 0xA434 + LensSerialNumber = 0xA435 + CompositeImage = 0xA460 + CompositeImageCount = 0xA461 + CompositeImageExposureTimes = 0xA462 + Gamma = 0xA500 + PrintImageMatching = 0xC4A5 + DNGVersion = 0xC612 + DNGBackwardVersion = 0xC613 + UniqueCameraModel = 0xC614 + LocalizedCameraModel = 0xC615 + CFAPlaneColor = 0xC616 + CFALayout = 0xC617 + LinearizationTable = 0xC618 + BlackLevelRepeatDim = 0xC619 + BlackLevel = 0xC61A + BlackLevelDeltaH = 0xC61B + BlackLevelDeltaV = 0xC61C + WhiteLevel = 0xC61D + DefaultScale = 0xC61E + DefaultCropOrigin = 0xC61F + DefaultCropSize = 0xC620 + ColorMatrix1 = 0xC621 + ColorMatrix2 = 0xC622 + CameraCalibration1 = 0xC623 + CameraCalibration2 = 0xC624 + ReductionMatrix1 = 0xC625 + ReductionMatrix2 = 0xC626 + AnalogBalance = 0xC627 + AsShotNeutral = 0xC628 + AsShotWhiteXY = 0xC629 + BaselineExposure = 0xC62A + BaselineNoise = 0xC62B + BaselineSharpness = 0xC62C + BayerGreenSplit = 0xC62D + LinearResponseLimit = 0xC62E + CameraSerialNumber = 0xC62F + LensInfo = 0xC630 + ChromaBlurRadius = 0xC631 + AntiAliasStrength = 0xC632 + ShadowScale = 0xC633 + DNGPrivateData = 0xC634 + MakerNoteSafety = 0xC635 + CalibrationIlluminant1 = 0xC65A + CalibrationIlluminant2 = 0xC65B + BestQualityScale = 0xC65C + RawDataUniqueID = 0xC65D + OriginalRawFileName = 0xC68B + OriginalRawFileData = 0xC68C + ActiveArea = 0xC68D + MaskedAreas = 0xC68E + AsShotICCProfile = 0xC68F + AsShotPreProfileMatrix = 0xC690 + CurrentICCProfile = 0xC691 + CurrentPreProfileMatrix = 0xC692 + ColorimetricReference = 0xC6BF + CameraCalibrationSignature = 0xC6F3 + ProfileCalibrationSignature = 0xC6F4 + AsShotProfileName = 0xC6F6 + NoiseReductionApplied = 0xC6F7 + ProfileName = 0xC6F8 + ProfileHueSatMapDims = 0xC6F9 + ProfileHueSatMapData1 = 0xC6FA + ProfileHueSatMapData2 = 0xC6FB + ProfileToneCurve = 0xC6FC + ProfileEmbedPolicy = 0xC6FD + ProfileCopyright = 0xC6FE + ForwardMatrix1 = 0xC714 + ForwardMatrix2 = 0xC715 + PreviewApplicationName = 0xC716 + PreviewApplicationVersion = 0xC717 + PreviewSettingsName = 0xC718 + PreviewSettingsDigest = 0xC719 + PreviewColorSpace = 0xC71A + PreviewDateTime = 0xC71B + RawImageDigest = 0xC71C + OriginalRawFileDigest = 0xC71D + SubTileBlockSize = 0xC71E + RowInterleaveFactor = 0xC71F + ProfileLookTableDims = 0xC725 + ProfileLookTableData = 0xC726 + OpcodeList1 = 0xC740 + OpcodeList2 = 0xC741 + OpcodeList3 = 0xC74E + NoiseProfile = 0xC761 + + +"""Maps EXIF tags to tag names.""" +TAGS = { + **{i.value: i.name for i in Base}, + 0x920C: "SpatialFrequencyResponse", + 0x9214: "SubjectLocation", + 0x9215: "ExposureIndex", + 0x828E: "CFAPattern", + 0x920B: "FlashEnergy", + 0x9216: "TIFF/EPStandardID", +} + + +class GPS(IntEnum): + GPSVersionID = 0x00 + GPSLatitudeRef = 0x01 + GPSLatitude = 0x02 + GPSLongitudeRef = 0x03 + GPSLongitude = 0x04 + GPSAltitudeRef = 0x05 + GPSAltitude = 0x06 + GPSTimeStamp = 0x07 + GPSSatellites = 0x08 + GPSStatus = 0x09 + GPSMeasureMode = 0x0A + GPSDOP = 0x0B + GPSSpeedRef = 0x0C + GPSSpeed = 0x0D + GPSTrackRef = 0x0E + GPSTrack = 0x0F + GPSImgDirectionRef = 0x10 + GPSImgDirection = 0x11 + GPSMapDatum = 0x12 + GPSDestLatitudeRef = 0x13 + GPSDestLatitude = 0x14 + GPSDestLongitudeRef = 0x15 + GPSDestLongitude = 0x16 + GPSDestBearingRef = 0x17 + GPSDestBearing = 0x18 + GPSDestDistanceRef = 0x19 + GPSDestDistance = 0x1A + GPSProcessingMethod = 0x1B + GPSAreaInformation = 0x1C + GPSDateStamp = 0x1D + GPSDifferential = 0x1E + GPSHPositioningError = 0x1F + + +"""Maps EXIF GPS tags to tag names.""" +GPSTAGS = {i.value: i.name for i in GPS} + + +class Interop(IntEnum): + InteropIndex = 0x0001 + InteropVersion = 0x0002 + RelatedImageFileFormat = 0x1000 + RelatedImageWidth = 0x1001 + RelatedImageHeight = 0x1002 + + +class IFD(IntEnum): + Exif = 0x8769 + GPSInfo = 0x8825 + MakerNote = 0x927C + Makernote = 0x927C # Deprecated + Interop = 0xA005 + IFD1 = -1 + + +class LightSource(IntEnum): + Unknown = 0x00 + Daylight = 0x01 + Fluorescent = 0x02 + Tungsten = 0x03 + Flash = 0x04 + Fine = 0x09 + Cloudy = 0x0A + Shade = 0x0B + DaylightFluorescent = 0x0C + DayWhiteFluorescent = 0x0D + CoolWhiteFluorescent = 0x0E + WhiteFluorescent = 0x0F + StandardLightA = 0x11 + StandardLightB = 0x12 + StandardLightC = 0x13 + D55 = 0x14 + D65 = 0x15 + D75 = 0x16 + D50 = 0x17 + ISO = 0x18 + Other = 0xFF diff --git a/venv/lib/python3.10/site-packages/PIL/FitsImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/FitsImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..a3fdc0efeec6f7ec195112ded41d8ff1e248a6a0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/FitsImagePlugin.py @@ -0,0 +1,152 @@ +# +# The Python Imaging Library +# $Id$ +# +# FITS file handling +# +# Copyright (c) 1998-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import gzip +import math + +from . import Image, ImageFile + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"SIMPLE") + + +class FitsImageFile(ImageFile.ImageFile): + format = "FITS" + format_description = "FITS" + + def _open(self) -> None: + assert self.fp is not None + + headers: dict[bytes, bytes] = {} + header_in_progress = False + decoder_name = "" + while True: + header = self.fp.read(80) + if not header: + msg = "Truncated FITS file" + raise OSError(msg) + keyword = header[:8].strip() + if keyword in (b"SIMPLE", b"XTENSION"): + header_in_progress = True + elif headers and not header_in_progress: + # This is now a data unit + break + elif keyword == b"END": + # Seek to the end of the header unit + self.fp.seek(math.ceil(self.fp.tell() / 2880) * 2880) + if not decoder_name: + decoder_name, offset, args = self._parse_headers(headers) + + header_in_progress = False + continue + + if decoder_name: + # Keep going to read past the headers + continue + + value = header[8:].split(b"/")[0].strip() + if value.startswith(b"="): + value = value[1:].strip() + if not headers and (not _accept(keyword) or value != b"T"): + msg = "Not a FITS file" + raise SyntaxError(msg) + headers[keyword] = value + + if not decoder_name: + msg = "No image data" + raise ValueError(msg) + + offset += self.fp.tell() - 80 + self.tile = [ImageFile._Tile(decoder_name, (0, 0) + self.size, offset, args)] + + def _get_size( + self, headers: dict[bytes, bytes], prefix: bytes + ) -> tuple[int, int] | None: + naxis = int(headers[prefix + b"NAXIS"]) + if naxis == 0: + return None + + if naxis == 1: + return 1, int(headers[prefix + b"NAXIS1"]) + else: + return int(headers[prefix + b"NAXIS1"]), int(headers[prefix + b"NAXIS2"]) + + def _parse_headers( + self, headers: dict[bytes, bytes] + ) -> tuple[str, int, tuple[str | int, ...]]: + prefix = b"" + decoder_name = "raw" + offset = 0 + if ( + headers.get(b"XTENSION") == b"'BINTABLE'" + and headers.get(b"ZIMAGE") == b"T" + and headers[b"ZCMPTYPE"] == b"'GZIP_1 '" + ): + no_prefix_size = self._get_size(headers, prefix) or (0, 0) + number_of_bits = int(headers[b"BITPIX"]) + offset = no_prefix_size[0] * no_prefix_size[1] * (number_of_bits // 8) + + prefix = b"Z" + decoder_name = "fits_gzip" + + size = self._get_size(headers, prefix) + if not size: + return "", 0, () + + self._size = size + + number_of_bits = int(headers[prefix + b"BITPIX"]) + if number_of_bits == 8: + self._mode = "L" + elif number_of_bits == 16: + self._mode = "I;16" + elif number_of_bits == 32: + self._mode = "I" + elif number_of_bits in (-32, -64): + self._mode = "F" + + args: tuple[str | int, ...] + if decoder_name == "raw": + args = (self.mode, 0, -1) + else: + args = (number_of_bits,) + return decoder_name, offset, args + + +class FitsGzipDecoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + value = gzip.decompress(self.fd.read()) + + rows = [] + offset = 0 + number_of_bits = min(self.args[0] // 8, 4) + for y in range(self.state.ysize): + row = bytearray() + for x in range(self.state.xsize): + row += value[offset + (4 - number_of_bits) : offset + 4] + offset += 4 + rows.append(row) + self.set_as_raw(bytes([pixel for row in rows[::-1] for pixel in row])) + return -1, 0 + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(FitsImageFile.format, FitsImageFile, _accept) +Image.register_decoder("fits_gzip", FitsGzipDecoder) + +Image.register_extensions(FitsImageFile.format, [".fit", ".fits"]) diff --git a/venv/lib/python3.10/site-packages/PIL/FliImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/FliImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..da1e8e95cf3263bb2e22ba45d6d2934fc4f2b54f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/FliImagePlugin.py @@ -0,0 +1,184 @@ +# +# The Python Imaging Library. +# $Id$ +# +# FLI/FLC file handling. +# +# History: +# 95-09-01 fl Created +# 97-01-03 fl Fixed parser, setup decoder tile +# 98-07-15 fl Renamed offset attribute to avoid name clash +# +# Copyright (c) Secret Labs AB 1997-98. +# Copyright (c) Fredrik Lundh 1995-97. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os + +from . import Image, ImageFile, ImagePalette +from ._binary import i16le as i16 +from ._binary import i32le as i32 +from ._binary import o8 +from ._util import DeferredError + +# +# decoder + + +def _accept(prefix: bytes) -> bool: + return ( + len(prefix) >= 16 + and i16(prefix, 4) in [0xAF11, 0xAF12] + and i16(prefix, 14) in [0, 3] # flags + ) + + +## +# Image plugin for the FLI/FLC animation format. Use the seek +# method to load individual frames. + + +class FliImageFile(ImageFile.ImageFile): + format = "FLI" + format_description = "Autodesk FLI/FLC Animation" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + # HEAD + assert self.fp is not None + s = self.fp.read(128) + if not ( + _accept(s) + and s[20:22] == b"\x00" * 2 + and s[42:80] == b"\x00" * 38 + and s[88:] == b"\x00" * 40 + ): + msg = "not an FLI/FLC file" + raise SyntaxError(msg) + + # frames + self.n_frames = i16(s, 6) + self.is_animated = self.n_frames > 1 + + # image characteristics + self._mode = "P" + self._size = i16(s, 8), i16(s, 10) + + # animation speed + duration = i32(s, 16) + magic = i16(s, 4) + if magic == 0xAF11: + duration = (duration * 1000) // 70 + self.info["duration"] = duration + + # look for palette + palette = [(a, a, a) for a in range(256)] + + s = self.fp.read(16) + + self.__offset = 128 + + if i16(s, 4) == 0xF100: + # prefix chunk; ignore it + self.fp.seek(self.__offset + i32(s)) + s = self.fp.read(16) + + if i16(s, 4) == 0xF1FA: + # look for palette chunk + number_of_subchunks = i16(s, 6) + chunk_size: int | None = None + for _ in range(number_of_subchunks): + if chunk_size is not None: + self.fp.seek(chunk_size - 6, os.SEEK_CUR) + s = self.fp.read(6) + chunk_type = i16(s, 4) + if chunk_type in (4, 11): + self._palette(palette, 2 if chunk_type == 11 else 0) + break + chunk_size = i32(s) + if not chunk_size: + break + + self.palette = ImagePalette.raw( + "RGB", b"".join(o8(r) + o8(g) + o8(b) for (r, g, b) in palette) + ) + + # set things up to decode first frame + self.__frame = -1 + self._fp = self.fp + self.__rewind = self.fp.tell() + self.seek(0) + + def _palette(self, palette: list[tuple[int, int, int]], shift: int) -> None: + # load palette + + i = 0 + assert self.fp is not None + for e in range(i16(self.fp.read(2))): + s = self.fp.read(2) + i = i + s[0] + n = s[1] + if n == 0: + n = 256 + s = self.fp.read(n * 3) + for n in range(0, len(s), 3): + r = s[n] << shift + g = s[n + 1] << shift + b = s[n + 2] << shift + palette[i] = (r, g, b) + i += 1 + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if frame < self.__frame: + self._seek(0) + + for f in range(self.__frame + 1, frame + 1): + self._seek(f) + + def _seek(self, frame: int) -> None: + if isinstance(self._fp, DeferredError): + raise self._fp.ex + if frame == 0: + self.__frame = -1 + self._fp.seek(self.__rewind) + self.__offset = 128 + else: + # ensure that the previous frame was loaded + self.load() + + if frame != self.__frame + 1: + msg = f"cannot seek to frame {frame}" + raise ValueError(msg) + self.__frame = frame + + # move to next frame + self.fp = self._fp + self.fp.seek(self.__offset) + + s = self.fp.read(4) + if not s: + msg = "missing frame size" + raise EOFError(msg) + + framesize = i32(s) + + self.decodermaxblock = framesize + self.tile = [ImageFile._Tile("fli", (0, 0) + self.size, self.__offset)] + + self.__offset += framesize + + def tell(self) -> int: + return self.__frame + + +# +# registry + +Image.register_open(FliImageFile.format, FliImageFile, _accept) + +Image.register_extensions(FliImageFile.format, [".fli", ".flc"]) diff --git a/venv/lib/python3.10/site-packages/PIL/FontFile.py b/venv/lib/python3.10/site-packages/PIL/FontFile.py new file mode 100644 index 0000000000000000000000000000000000000000..1e0c1c166b5932a7621e510eba047586465e03d8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/FontFile.py @@ -0,0 +1,134 @@ +# +# The Python Imaging Library +# $Id$ +# +# base class for raster font file parsers +# +# history: +# 1997-06-05 fl created +# 1997-08-19 fl restrict image width +# +# Copyright (c) 1997-1998 by Secret Labs AB +# Copyright (c) 1997-1998 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +from typing import BinaryIO + +from . import Image, _binary + +WIDTH = 800 + + +def puti16( + fp: BinaryIO, values: tuple[int, int, int, int, int, int, int, int, int, int] +) -> None: + """Write network order (big-endian) 16-bit sequence""" + for v in values: + if v < 0: + v += 65536 + fp.write(_binary.o16be(v)) + + +class FontFile: + """Base class for raster font file handlers.""" + + bitmap: Image.Image | None = None + + def __init__(self) -> None: + self.info: dict[bytes, bytes | int] = {} + self.glyph: list[ + tuple[ + tuple[int, int], + tuple[int, int, int, int], + tuple[int, int, int, int], + Image.Image, + ] + | None + ] = [None] * 256 + + def __getitem__(self, ix: int) -> ( + tuple[ + tuple[int, int], + tuple[int, int, int, int], + tuple[int, int, int, int], + Image.Image, + ] + | None + ): + return self.glyph[ix] + + def compile(self) -> None: + """Create metrics and bitmap""" + + if self.bitmap: + return + + # create bitmap large enough to hold all data + h = w = maxwidth = 0 + lines = 1 + for glyph in self.glyph: + if glyph: + d, dst, src, im = glyph + h = max(h, src[3] - src[1]) + w = w + (src[2] - src[0]) + if w > WIDTH: + lines += 1 + w = src[2] - src[0] + maxwidth = max(maxwidth, w) + + xsize = maxwidth + ysize = lines * h + + if xsize == 0 and ysize == 0: + return + + self.ysize = h + + # paste glyphs into bitmap + self.bitmap = Image.new("1", (xsize, ysize)) + self.metrics: list[ + tuple[tuple[int, int], tuple[int, int, int, int], tuple[int, int, int, int]] + | None + ] = [None] * 256 + x = y = 0 + for i in range(256): + glyph = self[i] + if glyph: + d, dst, src, im = glyph + xx = src[2] - src[0] + x0, y0 = x, y + x = x + xx + if x > WIDTH: + x, y = 0, y + h + x0, y0 = x, y + x = xx + s = src[0] + x0, src[1] + y0, src[2] + x0, src[3] + y0 + self.bitmap.paste(im.crop(src), s) + self.metrics[i] = d, dst, s + + def save(self, filename: str) -> None: + """Save font""" + + self.compile() + + # font data + if not self.bitmap: + msg = "No bitmap created" + raise ValueError(msg) + self.bitmap.save(os.path.splitext(filename)[0] + ".pbm", "PNG") + + # font metrics + with open(os.path.splitext(filename)[0] + ".pil", "wb") as fp: + fp.write(b"PILfont\n") + fp.write(f";;;;;;{self.ysize};\n".encode("ascii")) # HACK!!! + fp.write(b"DATA\n") + for id in range(256): + m = self.metrics[id] + if not m: + puti16(fp, (0,) * 10) + else: + puti16(fp, m[0] + m[1] + m[2]) diff --git a/venv/lib/python3.10/site-packages/PIL/FpxImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/FpxImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..fd992cd9e20eb7cf0c6de347ac0a76f928ffc238 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/FpxImagePlugin.py @@ -0,0 +1,257 @@ +# +# THIS IS WORK IN PROGRESS +# +# The Python Imaging Library. +# $Id$ +# +# FlashPix support for PIL +# +# History: +# 97-01-25 fl Created (reads uncompressed RGB images only) +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import olefile + +from . import Image, ImageFile +from ._binary import i32le as i32 + +# we map from colour field tuples to (mode, rawmode) descriptors +MODES = { + # opacity + (0x00007FFE,): ("A", "L"), + # monochrome + (0x00010000,): ("L", "L"), + (0x00018000, 0x00017FFE): ("RGBA", "LA"), + # photo YCC + (0x00020000, 0x00020001, 0x00020002): ("RGB", "YCC;P"), + (0x00028000, 0x00028001, 0x00028002, 0x00027FFE): ("RGBA", "YCCA;P"), + # standard RGB (NIFRGB) + (0x00030000, 0x00030001, 0x00030002): ("RGB", "RGB"), + (0x00038000, 0x00038001, 0x00038002, 0x00037FFE): ("RGBA", "RGBA"), +} + + +# +# -------------------------------------------------------------------- + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(olefile.MAGIC) + + +## +# Image plugin for the FlashPix images. + + +class FpxImageFile(ImageFile.ImageFile): + format = "FPX" + format_description = "FlashPix" + + def _open(self) -> None: + # + # read the OLE directory and see if this is a likely + # to be a FlashPix file + + try: + self.ole = olefile.OleFileIO(self.fp) + except OSError as e: + msg = "not an FPX file; invalid OLE file" + raise SyntaxError(msg) from e + + root = self.ole.root + if not root or root.clsid != "56616700-C154-11CE-8553-00AA00A1F95B": + msg = "not an FPX file; bad root CLSID" + raise SyntaxError(msg) + + self._open_index(1) + + def _open_index(self, index: int = 1) -> None: + # + # get the Image Contents Property Set + + prop = self.ole.getproperties( + [f"Data Object Store {index:06d}", "\005Image Contents"] + ) + + # size (highest resolution) + + assert isinstance(prop[0x1000002], int) + assert isinstance(prop[0x1000003], int) + self._size = prop[0x1000002], prop[0x1000003] + + size = max(self.size) + i = 1 + while size > 64: + size = size // 2 + i += 1 + self.maxid = i - 1 + + # mode. instead of using a single field for this, flashpix + # requires you to specify the mode for each channel in each + # resolution subimage, and leaves it to the decoder to make + # sure that they all match. for now, we'll cheat and assume + # that this is always the case. + + id = self.maxid << 16 + + s = prop[0x2000002 | id] + + if not isinstance(s, bytes) or (bands := i32(s, 4)) > 4: + msg = "Invalid number of bands" + raise OSError(msg) + + # note: for now, we ignore the "uncalibrated" flag + colors = tuple(i32(s, 8 + i * 4) & 0x7FFFFFFF for i in range(bands)) + + self._mode, self.rawmode = MODES[colors] + + # load JPEG tables, if any + self.jpeg = {} + for i in range(256): + id = 0x3000001 | (i << 16) + if id in prop: + self.jpeg[i] = prop[id] + + self._open_subimage(1, self.maxid) + + def _open_subimage(self, index: int = 1, subimage: int = 0) -> None: + # + # setup tile descriptors for a given subimage + + stream = [ + f"Data Object Store {index:06d}", + f"Resolution {subimage:04d}", + "Subimage 0000 Header", + ] + + fp = self.ole.openstream(stream) + + # skip prefix + fp.read(28) + + # header stream + s = fp.read(36) + + size = i32(s, 4), i32(s, 8) + # tilecount = i32(s, 12) + tilesize = i32(s, 16), i32(s, 20) + # channels = i32(s, 24) + offset = i32(s, 28) + length = i32(s, 32) + + if size != self.size: + msg = "subimage mismatch" + raise OSError(msg) + + # get tile descriptors + fp.seek(28 + offset) + s = fp.read(i32(s, 12) * length) + + x = y = 0 + xsize, ysize = size + xtile, ytile = tilesize + self.tile = [] + + for i in range(0, len(s), length): + x1 = min(xsize, x + xtile) + y1 = min(ysize, y + ytile) + + compression = i32(s, i + 8) + + if compression == 0: + self.tile.append( + ImageFile._Tile( + "raw", + (x, y, x1, y1), + i32(s, i) + 28, + self.rawmode, + ) + ) + + elif compression == 1: + # FIXME: the fill decoder is not implemented + self.tile.append( + ImageFile._Tile( + "fill", + (x, y, x1, y1), + i32(s, i) + 28, + (self.rawmode, s[12:16]), + ) + ) + + elif compression == 2: + internal_color_conversion = s[14] + jpeg_tables = s[15] + rawmode = self.rawmode + + if internal_color_conversion: + # The image is stored as usual (usually YCbCr). + if rawmode == "RGBA": + # For "RGBA", data is stored as YCbCrA based on + # negative RGB. The following trick works around + # this problem : + jpegmode, rawmode = "YCbCrK", "CMYK" + else: + jpegmode = None # let the decoder decide + + else: + # The image is stored as defined by rawmode + jpegmode = rawmode + + self.tile.append( + ImageFile._Tile( + "jpeg", + (x, y, x1, y1), + i32(s, i) + 28, + (rawmode, jpegmode), + ) + ) + + # FIXME: jpeg tables are tile dependent; the prefix + # data must be placed in the tile descriptor itself! + + if jpeg_tables: + self.tile_prefix = self.jpeg[jpeg_tables] + + else: + msg = "unknown/invalid compression" + raise OSError(msg) + + x = x + xtile + if x >= xsize: + x, y = 0, y + ytile + if y >= ysize: + break # isn't really required + + self.stream = stream + self._fp = self.fp + self.fp = None + + def load(self) -> Image.core.PixelAccess | None: + if not self.fp: + self.fp = self.ole.openstream(self.stream[:2] + ["Subimage 0000 Data"]) + + return ImageFile.ImageFile.load(self) + + def close(self) -> None: + self.ole.close() + super().close() + + def __exit__(self, *args: object) -> None: + self.ole.close() + super().__exit__() + + +# +# -------------------------------------------------------------------- + + +Image.register_open(FpxImageFile.format, FpxImageFile, _accept) + +Image.register_extension(FpxImageFile.format, ".fpx") diff --git a/venv/lib/python3.10/site-packages/PIL/FtexImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/FtexImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..d60e75bb60bdb5113c7cb3c48840918207ced694 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/FtexImagePlugin.py @@ -0,0 +1,114 @@ +""" +A Pillow loader for .ftc and .ftu files (FTEX) +Jerome Leclanche + +The contents of this file are hereby released in the public domain (CC0) +Full text of the CC0 license: + https://creativecommons.org/publicdomain/zero/1.0/ + +Independence War 2: Edge Of Chaos - Texture File Format - 16 October 2001 + +The textures used for 3D objects in Independence War 2: Edge Of Chaos are in a +packed custom format called FTEX. This file format uses file extensions FTC +and FTU. +* FTC files are compressed textures (using standard texture compression). +* FTU files are not compressed. +Texture File Format +The FTC and FTU texture files both use the same format. This +has the following structure: +{header} +{format_directory} +{data} +Where: +{header} = { + u32:magic, + u32:version, + u32:width, + u32:height, + u32:mipmap_count, + u32:format_count +} + +* The "magic" number is "FTEX". +* "width" and "height" are the dimensions of the texture. +* "mipmap_count" is the number of mipmaps in the texture. +* "format_count" is the number of texture formats (different versions of the +same texture) in this file. + +{format_directory} = format_count * { u32:format, u32:where } + +The format value is 0 for DXT1 compressed textures and 1 for 24-bit RGB +uncompressed textures. +The texture data for a format starts at the position "where" in the file. + +Each set of texture data in the file has the following structure: +{data} = format_count * { u32:mipmap_size, mipmap_size * { u8 } } +* "mipmap_size" is the number of bytes in that mip level. For compressed +textures this is the size of the texture data compressed with DXT1. For 24 bit +uncompressed textures, this is 3 * width * height. Following this are the image +bytes for that mipmap level. + +Note: All data is stored in little-Endian (Intel) byte order. +""" + +from __future__ import annotations + +import struct +from enum import IntEnum +from io import BytesIO + +from . import Image, ImageFile + +MAGIC = b"FTEX" + + +class Format(IntEnum): + DXT1 = 0 + UNCOMPRESSED = 1 + + +class FtexImageFile(ImageFile.ImageFile): + format = "FTEX" + format_description = "Texture File Format (IW2:EOC)" + + def _open(self) -> None: + if not _accept(self.fp.read(4)): + msg = "not an FTEX file" + raise SyntaxError(msg) + struct.unpack(" None: + pass + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(MAGIC) + + +Image.register_open(FtexImageFile.format, FtexImageFile, _accept) +Image.register_extensions(FtexImageFile.format, [".ftc", ".ftu"]) diff --git a/venv/lib/python3.10/site-packages/PIL/GbrImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/GbrImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..d69295363f372240593d1bc44a7a11e9b076718d --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GbrImagePlugin.py @@ -0,0 +1,101 @@ +# +# The Python Imaging Library +# +# load a GIMP brush file +# +# History: +# 96-03-14 fl Created +# 16-01-08 es Version 2 +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# Copyright (c) Eric Soroos 2016. +# +# See the README file for information on usage and redistribution. +# +# +# See https://github.com/GNOME/gimp/blob/mainline/devel-docs/gbr.txt for +# format documentation. +# +# This code Interprets version 1 and 2 .gbr files. +# Version 1 files are obsolete, and should not be used for new +# brushes. +# Version 2 files are saved by GIMP v2.8 (at least) +# Version 3 files have a format specifier of 18 for 16bit floats in +# the color depth field. This is currently unsupported by Pillow. +from __future__ import annotations + +from . import Image, ImageFile +from ._binary import i32be as i32 + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 8 and i32(prefix, 0) >= 20 and i32(prefix, 4) in (1, 2) + + +## +# Image plugin for the GIMP brush format. + + +class GbrImageFile(ImageFile.ImageFile): + format = "GBR" + format_description = "GIMP brush file" + + def _open(self) -> None: + header_size = i32(self.fp.read(4)) + if header_size < 20: + msg = "not a GIMP brush" + raise SyntaxError(msg) + version = i32(self.fp.read(4)) + if version not in (1, 2): + msg = f"Unsupported GIMP brush version: {version}" + raise SyntaxError(msg) + + width = i32(self.fp.read(4)) + height = i32(self.fp.read(4)) + color_depth = i32(self.fp.read(4)) + if width == 0 or height == 0: + msg = "not a GIMP brush" + raise SyntaxError(msg) + if color_depth not in (1, 4): + msg = f"Unsupported GIMP brush color depth: {color_depth}" + raise SyntaxError(msg) + + if version == 1: + comment_length = header_size - 20 + else: + comment_length = header_size - 28 + magic_number = self.fp.read(4) + if magic_number != b"GIMP": + msg = "not a GIMP brush, bad magic number" + raise SyntaxError(msg) + self.info["spacing"] = i32(self.fp.read(4)) + + self.info["comment"] = self.fp.read(comment_length)[:-1] + + if color_depth == 1: + self._mode = "L" + else: + self._mode = "RGBA" + + self._size = width, height + + # Image might not be small + Image._decompression_bomb_check(self.size) + + # Data is an uncompressed block of w * h * bytes/pixel + self._data_size = width * height * color_depth + + def load(self) -> Image.core.PixelAccess | None: + if self._im is None: + self.im = Image.core.new(self.mode, self.size) + self.frombytes(self.fp.read(self._data_size)) + return Image.Image.load(self) + + +# +# registry + + +Image.register_open(GbrImageFile.format, GbrImageFile, _accept) +Image.register_extension(GbrImageFile.format, ".gbr") diff --git a/venv/lib/python3.10/site-packages/PIL/GdImageFile.py b/venv/lib/python3.10/site-packages/PIL/GdImageFile.py new file mode 100644 index 0000000000000000000000000000000000000000..891225ce2fd034a11963bb64212cfa7311190441 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GdImageFile.py @@ -0,0 +1,102 @@ +# +# The Python Imaging Library. +# $Id$ +# +# GD file handling +# +# History: +# 1996-04-12 fl Created +# +# Copyright (c) 1997 by Secret Labs AB. +# Copyright (c) 1996 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + + +""" +.. note:: + This format cannot be automatically recognized, so the + class is not registered for use with :py:func:`PIL.Image.open()`. To open a + gd file, use the :py:func:`PIL.GdImageFile.open()` function instead. + +.. warning:: + THE GD FORMAT IS NOT DESIGNED FOR DATA INTERCHANGE. This + implementation is provided for convenience and demonstrational + purposes only. +""" +from __future__ import annotations + +from typing import IO + +from . import ImageFile, ImagePalette, UnidentifiedImageError +from ._binary import i16be as i16 +from ._binary import i32be as i32 +from ._typing import StrOrBytesPath + + +class GdImageFile(ImageFile.ImageFile): + """ + Image plugin for the GD uncompressed format. Note that this format + is not supported by the standard :py:func:`PIL.Image.open()` function. To use + this plugin, you have to import the :py:mod:`PIL.GdImageFile` module and + use the :py:func:`PIL.GdImageFile.open()` function. + """ + + format = "GD" + format_description = "GD uncompressed images" + + def _open(self) -> None: + # Header + assert self.fp is not None + + s = self.fp.read(1037) + + if i16(s) not in [65534, 65535]: + msg = "Not a valid GD 2.x .gd file" + raise SyntaxError(msg) + + self._mode = "P" + self._size = i16(s, 2), i16(s, 4) + + true_color = s[6] + true_color_offset = 2 if true_color else 0 + + # transparency index + tindex = i32(s, 7 + true_color_offset) + if tindex < 256: + self.info["transparency"] = tindex + + self.palette = ImagePalette.raw( + "RGBX", s[7 + true_color_offset + 6 : 7 + true_color_offset + 6 + 256 * 4] + ) + + self.tile = [ + ImageFile._Tile( + "raw", + (0, 0) + self.size, + 7 + true_color_offset + 6 + 256 * 4, + "L", + ) + ] + + +def open(fp: StrOrBytesPath | IO[bytes], mode: str = "r") -> GdImageFile: + """ + Load texture from a GD image file. + + :param fp: GD file name, or an opened file handle. + :param mode: Optional mode. In this version, if the mode argument + is given, it must be "r". + :returns: An image instance. + :raises OSError: If the image could not be read. + """ + if mode != "r": + msg = "bad mode" + raise ValueError(msg) + + try: + return GdImageFile(fp) + except SyntaxError as e: + msg = "cannot identify this image file" + raise UnidentifiedImageError(msg) from e diff --git a/venv/lib/python3.10/site-packages/PIL/GifImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/GifImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..58c460ef3db7d2d4e130095d4f64a2586a9dadbb --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GifImagePlugin.py @@ -0,0 +1,1215 @@ +# +# The Python Imaging Library. +# $Id$ +# +# GIF file handling +# +# History: +# 1995-09-01 fl Created +# 1996-12-14 fl Added interlace support +# 1996-12-30 fl Added animation support +# 1997-01-05 fl Added write support, fixed local colour map bug +# 1997-02-23 fl Make sure to load raster data in getdata() +# 1997-07-05 fl Support external decoder (0.4) +# 1998-07-09 fl Handle all modes when saving (0.5) +# 1998-07-15 fl Renamed offset attribute to avoid name clash +# 2001-04-16 fl Added rewind support (seek to frame 0) (0.6) +# 2001-04-17 fl Added palette optimization (0.7) +# 2002-06-06 fl Added transparency support for save (0.8) +# 2004-02-24 fl Disable interlacing for small images +# +# Copyright (c) 1997-2004 by Secret Labs AB +# Copyright (c) 1995-2004 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import itertools +import math +import os +import subprocess +from enum import IntEnum +from functools import cached_property +from typing import Any, NamedTuple, cast + +from . import ( + Image, + ImageChops, + ImageFile, + ImageMath, + ImageOps, + ImagePalette, + ImageSequence, +) +from ._binary import i16le as i16 +from ._binary import o8 +from ._binary import o16le as o16 +from ._util import DeferredError + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import IO, Literal + + from . import _imaging + from ._typing import Buffer + + +class LoadingStrategy(IntEnum): + """.. versionadded:: 9.1.0""" + + RGB_AFTER_FIRST = 0 + RGB_AFTER_DIFFERENT_PALETTE_ONLY = 1 + RGB_ALWAYS = 2 + + +#: .. versionadded:: 9.1.0 +LOADING_STRATEGY = LoadingStrategy.RGB_AFTER_FIRST + +# -------------------------------------------------------------------- +# Identify/read GIF files + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith((b"GIF87a", b"GIF89a")) + + +## +# Image plugin for GIF images. This plugin supports both GIF87 and +# GIF89 images. + + +class GifImageFile(ImageFile.ImageFile): + format = "GIF" + format_description = "Compuserve GIF" + _close_exclusive_fp_after_loading = False + + global_palette = None + + def data(self) -> bytes | None: + s = self.fp.read(1) + if s and s[0]: + return self.fp.read(s[0]) + return None + + def _is_palette_needed(self, p: bytes) -> bool: + for i in range(0, len(p), 3): + if not (i // 3 == p[i] == p[i + 1] == p[i + 2]): + return True + return False + + def _open(self) -> None: + # Screen + s = self.fp.read(13) + if not _accept(s): + msg = "not a GIF file" + raise SyntaxError(msg) + + self.info["version"] = s[:6] + self._size = i16(s, 6), i16(s, 8) + flags = s[10] + bits = (flags & 7) + 1 + + if flags & 128: + # get global palette + self.info["background"] = s[11] + # check if palette contains colour indices + p = self.fp.read(3 << bits) + if self._is_palette_needed(p): + p = ImagePalette.raw("RGB", p) + self.global_palette = self.palette = p + + self._fp = self.fp # FIXME: hack + self.__rewind = self.fp.tell() + self._n_frames: int | None = None + self._seek(0) # get ready to read first frame + + @property + def n_frames(self) -> int: + if self._n_frames is None: + current = self.tell() + try: + while True: + self._seek(self.tell() + 1, False) + except EOFError: + self._n_frames = self.tell() + 1 + self.seek(current) + return self._n_frames + + @cached_property + def is_animated(self) -> bool: + if self._n_frames is not None: + return self._n_frames != 1 + + current = self.tell() + if current: + return True + + try: + self._seek(1, False) + is_animated = True + except EOFError: + is_animated = False + + self.seek(current) + return is_animated + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if frame < self.__frame: + self._im = None + self._seek(0) + + last_frame = self.__frame + for f in range(self.__frame + 1, frame + 1): + try: + self._seek(f) + except EOFError as e: + self.seek(last_frame) + msg = "no more images in GIF file" + raise EOFError(msg) from e + + def _seek(self, frame: int, update_image: bool = True) -> None: + if isinstance(self._fp, DeferredError): + raise self._fp.ex + if frame == 0: + # rewind + self.__offset = 0 + self.dispose: _imaging.ImagingCore | None = None + self.__frame = -1 + self._fp.seek(self.__rewind) + self.disposal_method = 0 + if "comment" in self.info: + del self.info["comment"] + else: + # ensure that the previous frame was loaded + if self.tile and update_image: + self.load() + + if frame != self.__frame + 1: + msg = f"cannot seek to frame {frame}" + raise ValueError(msg) + + self.fp = self._fp + if self.__offset: + # backup to last frame + self.fp.seek(self.__offset) + while self.data(): + pass + self.__offset = 0 + + s = self.fp.read(1) + if not s or s == b";": + msg = "no more images in GIF file" + raise EOFError(msg) + + palette: ImagePalette.ImagePalette | Literal[False] | None = None + + info: dict[str, Any] = {} + frame_transparency = None + interlace = None + frame_dispose_extent = None + while True: + if not s: + s = self.fp.read(1) + if not s or s == b";": + break + + elif s == b"!": + # + # extensions + # + s = self.fp.read(1) + block = self.data() + if s[0] == 249 and block is not None: + # + # graphic control extension + # + flags = block[0] + if flags & 1: + frame_transparency = block[3] + info["duration"] = i16(block, 1) * 10 + + # disposal method - find the value of bits 4 - 6 + dispose_bits = 0b00011100 & flags + dispose_bits = dispose_bits >> 2 + if dispose_bits: + # only set the dispose if it is not + # unspecified. I'm not sure if this is + # correct, but it seems to prevent the last + # frame from looking odd for some animations + self.disposal_method = dispose_bits + elif s[0] == 254: + # + # comment extension + # + comment = b"" + + # Read this comment block + while block: + comment += block + block = self.data() + + if "comment" in info: + # If multiple comment blocks in frame, separate with \n + info["comment"] += b"\n" + comment + else: + info["comment"] = comment + s = None + continue + elif s[0] == 255 and frame == 0 and block is not None: + # + # application extension + # + info["extension"] = block, self.fp.tell() + if block.startswith(b"NETSCAPE2.0"): + block = self.data() + if block and len(block) >= 3 and block[0] == 1: + self.info["loop"] = i16(block, 1) + while self.data(): + pass + + elif s == b",": + # + # local image + # + s = self.fp.read(9) + + # extent + x0, y0 = i16(s, 0), i16(s, 2) + x1, y1 = x0 + i16(s, 4), y0 + i16(s, 6) + if (x1 > self.size[0] or y1 > self.size[1]) and update_image: + self._size = max(x1, self.size[0]), max(y1, self.size[1]) + Image._decompression_bomb_check(self._size) + frame_dispose_extent = x0, y0, x1, y1 + flags = s[8] + + interlace = (flags & 64) != 0 + + if flags & 128: + bits = (flags & 7) + 1 + p = self.fp.read(3 << bits) + if self._is_palette_needed(p): + palette = ImagePalette.raw("RGB", p) + else: + palette = False + + # image data + bits = self.fp.read(1)[0] + self.__offset = self.fp.tell() + break + s = None + + if interlace is None: + msg = "image not found in GIF frame" + raise EOFError(msg) + + self.__frame = frame + if not update_image: + return + + self.tile = [] + + if self.dispose: + self.im.paste(self.dispose, self.dispose_extent) + + self._frame_palette = palette if palette is not None else self.global_palette + self._frame_transparency = frame_transparency + if frame == 0: + if self._frame_palette: + if LOADING_STRATEGY == LoadingStrategy.RGB_ALWAYS: + self._mode = "RGBA" if frame_transparency is not None else "RGB" + else: + self._mode = "P" + else: + self._mode = "L" + + if palette: + self.palette = palette + elif self.global_palette: + from copy import copy + + self.palette = copy(self.global_palette) + else: + self.palette = None + else: + if self.mode == "P": + if ( + LOADING_STRATEGY != LoadingStrategy.RGB_AFTER_DIFFERENT_PALETTE_ONLY + or palette + ): + if "transparency" in self.info: + self.im.putpalettealpha(self.info["transparency"], 0) + self.im = self.im.convert("RGBA", Image.Dither.FLOYDSTEINBERG) + self._mode = "RGBA" + del self.info["transparency"] + else: + self._mode = "RGB" + self.im = self.im.convert("RGB", Image.Dither.FLOYDSTEINBERG) + + def _rgb(color: int) -> tuple[int, int, int]: + if self._frame_palette: + if color * 3 + 3 > len(self._frame_palette.palette): + color = 0 + return cast( + tuple[int, int, int], + tuple(self._frame_palette.palette[color * 3 : color * 3 + 3]), + ) + else: + return (color, color, color) + + self.dispose = None + self.dispose_extent: tuple[int, int, int, int] | None = frame_dispose_extent + if self.dispose_extent and self.disposal_method >= 2: + try: + if self.disposal_method == 2: + # replace with background colour + + # only dispose the extent in this frame + x0, y0, x1, y1 = self.dispose_extent + dispose_size = (x1 - x0, y1 - y0) + + Image._decompression_bomb_check(dispose_size) + + # by convention, attempt to use transparency first + dispose_mode = "P" + color = self.info.get("transparency", frame_transparency) + if color is not None: + if self.mode in ("RGB", "RGBA"): + dispose_mode = "RGBA" + color = _rgb(color) + (0,) + else: + color = self.info.get("background", 0) + if self.mode in ("RGB", "RGBA"): + dispose_mode = "RGB" + color = _rgb(color) + self.dispose = Image.core.fill(dispose_mode, dispose_size, color) + else: + # replace with previous contents + if self._im is not None: + # only dispose the extent in this frame + self.dispose = self._crop(self.im, self.dispose_extent) + elif frame_transparency is not None: + x0, y0, x1, y1 = self.dispose_extent + dispose_size = (x1 - x0, y1 - y0) + + Image._decompression_bomb_check(dispose_size) + dispose_mode = "P" + color = frame_transparency + if self.mode in ("RGB", "RGBA"): + dispose_mode = "RGBA" + color = _rgb(frame_transparency) + (0,) + self.dispose = Image.core.fill( + dispose_mode, dispose_size, color + ) + except AttributeError: + pass + + if interlace is not None: + transparency = -1 + if frame_transparency is not None: + if frame == 0: + if LOADING_STRATEGY != LoadingStrategy.RGB_ALWAYS: + self.info["transparency"] = frame_transparency + elif self.mode not in ("RGB", "RGBA"): + transparency = frame_transparency + self.tile = [ + ImageFile._Tile( + "gif", + (x0, y0, x1, y1), + self.__offset, + (bits, interlace, transparency), + ) + ] + + if info.get("comment"): + self.info["comment"] = info["comment"] + for k in ["duration", "extension"]: + if k in info: + self.info[k] = info[k] + elif k in self.info: + del self.info[k] + + def load_prepare(self) -> None: + temp_mode = "P" if self._frame_palette else "L" + self._prev_im = None + if self.__frame == 0: + if self._frame_transparency is not None: + self.im = Image.core.fill( + temp_mode, self.size, self._frame_transparency + ) + elif self.mode in ("RGB", "RGBA"): + self._prev_im = self.im + if self._frame_palette: + self.im = Image.core.fill("P", self.size, self._frame_transparency or 0) + self.im.putpalette("RGB", *self._frame_palette.getdata()) + else: + self._im = None + if not self._prev_im and self._im is not None and self.size != self.im.size: + expanded_im = Image.core.fill(self.im.mode, self.size) + if self._frame_palette: + expanded_im.putpalette("RGB", *self._frame_palette.getdata()) + expanded_im.paste(self.im, (0, 0) + self.im.size) + + self.im = expanded_im + self._mode = temp_mode + self._frame_palette = None + + super().load_prepare() + + def load_end(self) -> None: + if self.__frame == 0: + if self.mode == "P" and LOADING_STRATEGY == LoadingStrategy.RGB_ALWAYS: + if self._frame_transparency is not None: + self.im.putpalettealpha(self._frame_transparency, 0) + self._mode = "RGBA" + else: + self._mode = "RGB" + self.im = self.im.convert(self.mode, Image.Dither.FLOYDSTEINBERG) + return + if not self._prev_im: + return + if self.size != self._prev_im.size: + if self._frame_transparency is not None: + expanded_im = Image.core.fill("RGBA", self.size) + else: + expanded_im = Image.core.fill("P", self.size) + expanded_im.putpalette("RGB", "RGB", self.im.getpalette()) + expanded_im = expanded_im.convert("RGB") + expanded_im.paste(self._prev_im, (0, 0) + self._prev_im.size) + + self._prev_im = expanded_im + assert self._prev_im is not None + if self._frame_transparency is not None: + if self.mode == "L": + frame_im = self.im.convert_transparent("LA", self._frame_transparency) + else: + self.im.putpalettealpha(self._frame_transparency, 0) + frame_im = self.im.convert("RGBA") + else: + frame_im = self.im.convert("RGB") + + assert self.dispose_extent is not None + frame_im = self._crop(frame_im, self.dispose_extent) + + self.im = self._prev_im + self._mode = self.im.mode + if frame_im.mode in ("LA", "RGBA"): + self.im.paste(frame_im, self.dispose_extent, frame_im) + else: + self.im.paste(frame_im, self.dispose_extent) + + def tell(self) -> int: + return self.__frame + + +# -------------------------------------------------------------------- +# Write GIF files + + +RAWMODE = {"1": "L", "L": "L", "P": "P"} + + +def _normalize_mode(im: Image.Image) -> Image.Image: + """ + Takes an image (or frame), returns an image in a mode that is appropriate + for saving in a Gif. + + It may return the original image, or it may return an image converted to + palette or 'L' mode. + + :param im: Image object + :returns: Image object + """ + if im.mode in RAWMODE: + im.load() + return im + if Image.getmodebase(im.mode) == "RGB": + im = im.convert("P", palette=Image.Palette.ADAPTIVE) + assert im.palette is not None + if im.palette.mode == "RGBA": + for rgba in im.palette.colors: + if rgba[3] == 0: + im.info["transparency"] = im.palette.colors[rgba] + break + return im + return im.convert("L") + + +_Palette = bytes | bytearray | list[int] | ImagePalette.ImagePalette + + +def _normalize_palette( + im: Image.Image, palette: _Palette | None, info: dict[str, Any] +) -> Image.Image: + """ + Normalizes the palette for image. + - Sets the palette to the incoming palette, if provided. + - Ensures that there's a palette for L mode images + - Optimizes the palette if necessary/desired. + + :param im: Image object + :param palette: bytes object containing the source palette, or .... + :param info: encoderinfo + :returns: Image object + """ + source_palette = None + if palette: + # a bytes palette + if isinstance(palette, (bytes, bytearray, list)): + source_palette = bytearray(palette[:768]) + if isinstance(palette, ImagePalette.ImagePalette): + source_palette = bytearray(palette.palette) + + if im.mode == "P": + if not source_palette: + im_palette = im.getpalette(None) + assert im_palette is not None + source_palette = bytearray(im_palette) + else: # L-mode + if not source_palette: + source_palette = bytearray(i // 3 for i in range(768)) + im.palette = ImagePalette.ImagePalette("RGB", palette=source_palette) + assert source_palette is not None + + if palette: + used_palette_colors: list[int | None] = [] + assert im.palette is not None + for i in range(0, len(source_palette), 3): + source_color = tuple(source_palette[i : i + 3]) + index = im.palette.colors.get(source_color) + if index in used_palette_colors: + index = None + used_palette_colors.append(index) + for i, index in enumerate(used_palette_colors): + if index is None: + for j in range(len(used_palette_colors)): + if j not in used_palette_colors: + used_palette_colors[i] = j + break + dest_map: list[int] = [] + for index in used_palette_colors: + assert index is not None + dest_map.append(index) + im = im.remap_palette(dest_map) + else: + optimized_palette_colors = _get_optimize(im, info) + if optimized_palette_colors is not None: + im = im.remap_palette(optimized_palette_colors, source_palette) + if "transparency" in info: + try: + info["transparency"] = optimized_palette_colors.index( + info["transparency"] + ) + except ValueError: + del info["transparency"] + return im + + assert im.palette is not None + im.palette.palette = source_palette + return im + + +def _write_single_frame( + im: Image.Image, + fp: IO[bytes], + palette: _Palette | None, +) -> None: + im_out = _normalize_mode(im) + for k, v in im_out.info.items(): + if isinstance(k, str): + im.encoderinfo.setdefault(k, v) + im_out = _normalize_palette(im_out, palette, im.encoderinfo) + + for s in _get_global_header(im_out, im.encoderinfo): + fp.write(s) + + # local image header + flags = 0 + if get_interlace(im): + flags = flags | 64 + _write_local_header(fp, im, (0, 0), flags) + + im_out.encoderconfig = (8, get_interlace(im)) + ImageFile._save( + im_out, fp, [ImageFile._Tile("gif", (0, 0) + im.size, 0, RAWMODE[im_out.mode])] + ) + + fp.write(b"\0") # end of image data + + +def _getbbox( + base_im: Image.Image, im_frame: Image.Image +) -> tuple[Image.Image, tuple[int, int, int, int] | None]: + palette_bytes = [ + bytes(im.palette.palette) if im.palette else b"" for im in (base_im, im_frame) + ] + if palette_bytes[0] != palette_bytes[1]: + im_frame = im_frame.convert("RGBA") + base_im = base_im.convert("RGBA") + delta = ImageChops.subtract_modulo(im_frame, base_im) + return delta, delta.getbbox(alpha_only=False) + + +class _Frame(NamedTuple): + im: Image.Image + bbox: tuple[int, int, int, int] | None + encoderinfo: dict[str, Any] + + +def _write_multiple_frames( + im: Image.Image, fp: IO[bytes], palette: _Palette | None +) -> bool: + duration = im.encoderinfo.get("duration") + disposal = im.encoderinfo.get("disposal", im.info.get("disposal")) + + im_frames: list[_Frame] = [] + previous_im: Image.Image | None = None + frame_count = 0 + background_im = None + for imSequence in itertools.chain([im], im.encoderinfo.get("append_images", [])): + for im_frame in ImageSequence.Iterator(imSequence): + # a copy is required here since seek can still mutate the image + im_frame = _normalize_mode(im_frame.copy()) + if frame_count == 0: + for k, v in im_frame.info.items(): + if k == "transparency": + continue + if isinstance(k, str): + im.encoderinfo.setdefault(k, v) + + encoderinfo = im.encoderinfo.copy() + if "transparency" in im_frame.info: + encoderinfo.setdefault("transparency", im_frame.info["transparency"]) + im_frame = _normalize_palette(im_frame, palette, encoderinfo) + if isinstance(duration, (list, tuple)): + encoderinfo["duration"] = duration[frame_count] + elif duration is None and "duration" in im_frame.info: + encoderinfo["duration"] = im_frame.info["duration"] + if isinstance(disposal, (list, tuple)): + encoderinfo["disposal"] = disposal[frame_count] + frame_count += 1 + + diff_frame = None + if im_frames and previous_im: + # delta frame + delta, bbox = _getbbox(previous_im, im_frame) + if not bbox: + # This frame is identical to the previous frame + if encoderinfo.get("duration"): + im_frames[-1].encoderinfo["duration"] += encoderinfo["duration"] + continue + if im_frames[-1].encoderinfo.get("disposal") == 2: + # To appear correctly in viewers using a convention, + # only consider transparency, and not background color + color = im.encoderinfo.get( + "transparency", im.info.get("transparency") + ) + if color is not None: + if background_im is None: + background = _get_background(im_frame, color) + background_im = Image.new("P", im_frame.size, background) + first_palette = im_frames[0].im.palette + assert first_palette is not None + background_im.putpalette(first_palette, first_palette.mode) + bbox = _getbbox(background_im, im_frame)[1] + else: + bbox = (0, 0) + im_frame.size + elif encoderinfo.get("optimize") and im_frame.mode != "1": + if "transparency" not in encoderinfo: + assert im_frame.palette is not None + try: + encoderinfo["transparency"] = ( + im_frame.palette._new_color_index(im_frame) + ) + except ValueError: + pass + if "transparency" in encoderinfo: + # When the delta is zero, fill the image with transparency + diff_frame = im_frame.copy() + fill = Image.new("P", delta.size, encoderinfo["transparency"]) + if delta.mode == "RGBA": + r, g, b, a = delta.split() + mask = ImageMath.lambda_eval( + lambda args: args["convert"]( + args["max"]( + args["max"]( + args["max"](args["r"], args["g"]), args["b"] + ), + args["a"], + ) + * 255, + "1", + ), + r=r, + g=g, + b=b, + a=a, + ) + else: + if delta.mode == "P": + # Convert to L without considering palette + delta_l = Image.new("L", delta.size) + delta_l.putdata(delta.getdata()) + delta = delta_l + mask = ImageMath.lambda_eval( + lambda args: args["convert"](args["im"] * 255, "1"), + im=delta, + ) + diff_frame.paste(fill, mask=ImageOps.invert(mask)) + else: + bbox = None + previous_im = im_frame + im_frames.append(_Frame(diff_frame or im_frame, bbox, encoderinfo)) + + if len(im_frames) == 1: + if "duration" in im.encoderinfo: + # Since multiple frames will not be written, use the combined duration + im.encoderinfo["duration"] = im_frames[0].encoderinfo["duration"] + return False + + for frame_data in im_frames: + im_frame = frame_data.im + if not frame_data.bbox: + # global header + for s in _get_global_header(im_frame, frame_data.encoderinfo): + fp.write(s) + offset = (0, 0) + else: + # compress difference + if not palette: + frame_data.encoderinfo["include_color_table"] = True + + if frame_data.bbox != (0, 0) + im_frame.size: + im_frame = im_frame.crop(frame_data.bbox) + offset = frame_data.bbox[:2] + _write_frame_data(fp, im_frame, offset, frame_data.encoderinfo) + return True + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + _save(im, fp, filename, save_all=True) + + +def _save( + im: Image.Image, fp: IO[bytes], filename: str | bytes, save_all: bool = False +) -> None: + # header + if "palette" in im.encoderinfo or "palette" in im.info: + palette = im.encoderinfo.get("palette", im.info.get("palette")) + else: + palette = None + im.encoderinfo.setdefault("optimize", True) + + if not save_all or not _write_multiple_frames(im, fp, palette): + _write_single_frame(im, fp, palette) + + fp.write(b";") # end of file + + if hasattr(fp, "flush"): + fp.flush() + + +def get_interlace(im: Image.Image) -> int: + interlace = im.encoderinfo.get("interlace", 1) + + # workaround for @PIL153 + if min(im.size) < 16: + interlace = 0 + + return interlace + + +def _write_local_header( + fp: IO[bytes], im: Image.Image, offset: tuple[int, int], flags: int +) -> None: + try: + transparency = im.encoderinfo["transparency"] + except KeyError: + transparency = None + + if "duration" in im.encoderinfo: + duration = int(im.encoderinfo["duration"] / 10) + else: + duration = 0 + + disposal = int(im.encoderinfo.get("disposal", 0)) + + if transparency is not None or duration != 0 or disposal: + packed_flag = 1 if transparency is not None else 0 + packed_flag |= disposal << 2 + + fp.write( + b"!" + + o8(249) # extension intro + + o8(4) # length + + o8(packed_flag) # packed fields + + o16(duration) # duration + + o8(transparency or 0) # transparency index + + o8(0) + ) + + include_color_table = im.encoderinfo.get("include_color_table") + if include_color_table: + palette_bytes = _get_palette_bytes(im) + color_table_size = _get_color_table_size(palette_bytes) + if color_table_size: + flags = flags | 128 # local color table flag + flags = flags | color_table_size + + fp.write( + b"," + + o16(offset[0]) # offset + + o16(offset[1]) + + o16(im.size[0]) # size + + o16(im.size[1]) + + o8(flags) # flags + ) + if include_color_table and color_table_size: + fp.write(_get_header_palette(palette_bytes)) + fp.write(o8(8)) # bits + + +def _save_netpbm(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + # Unused by default. + # To use, uncomment the register_save call at the end of the file. + # + # If you need real GIF compression and/or RGB quantization, you + # can use the external NETPBM/PBMPLUS utilities. See comments + # below for information on how to enable this. + tempfile = im._dump() + + try: + with open(filename, "wb") as f: + if im.mode != "RGB": + subprocess.check_call( + ["ppmtogif", tempfile], stdout=f, stderr=subprocess.DEVNULL + ) + else: + # Pipe ppmquant output into ppmtogif + # "ppmquant 256 %s | ppmtogif > %s" % (tempfile, filename) + quant_cmd = ["ppmquant", "256", tempfile] + togif_cmd = ["ppmtogif"] + quant_proc = subprocess.Popen( + quant_cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL + ) + togif_proc = subprocess.Popen( + togif_cmd, + stdin=quant_proc.stdout, + stdout=f, + stderr=subprocess.DEVNULL, + ) + + # Allow ppmquant to receive SIGPIPE if ppmtogif exits + assert quant_proc.stdout is not None + quant_proc.stdout.close() + + retcode = quant_proc.wait() + if retcode: + raise subprocess.CalledProcessError(retcode, quant_cmd) + + retcode = togif_proc.wait() + if retcode: + raise subprocess.CalledProcessError(retcode, togif_cmd) + finally: + try: + os.unlink(tempfile) + except OSError: + pass + + +# Force optimization so that we can test performance against +# cases where it took lots of memory and time previously. +_FORCE_OPTIMIZE = False + + +def _get_optimize(im: Image.Image, info: dict[str, Any]) -> list[int] | None: + """ + Palette optimization is a potentially expensive operation. + + This function determines if the palette should be optimized using + some heuristics, then returns the list of palette entries in use. + + :param im: Image object + :param info: encoderinfo + :returns: list of indexes of palette entries in use, or None + """ + if im.mode in ("P", "L") and info and info.get("optimize"): + # Potentially expensive operation. + + # The palette saves 3 bytes per color not used, but palette + # lengths are restricted to 3*(2**N) bytes. Max saving would + # be 768 -> 6 bytes if we went all the way down to 2 colors. + # * If we're over 128 colors, we can't save any space. + # * If there aren't any holes, it's not worth collapsing. + # * If we have a 'large' image, the palette is in the noise. + + # create the new palette if not every color is used + optimise = _FORCE_OPTIMIZE or im.mode == "L" + if optimise or im.width * im.height < 512 * 512: + # check which colors are used + used_palette_colors = [] + for i, count in enumerate(im.histogram()): + if count: + used_palette_colors.append(i) + + if optimise or max(used_palette_colors) >= len(used_palette_colors): + return used_palette_colors + + assert im.palette is not None + num_palette_colors = len(im.palette.palette) // Image.getmodebands( + im.palette.mode + ) + current_palette_size = 1 << (num_palette_colors - 1).bit_length() + if ( + # check that the palette would become smaller when saved + len(used_palette_colors) <= current_palette_size // 2 + # check that the palette is not already the smallest possible size + and current_palette_size > 2 + ): + return used_palette_colors + return None + + +def _get_color_table_size(palette_bytes: bytes) -> int: + # calculate the palette size for the header + if not palette_bytes: + return 0 + elif len(palette_bytes) < 9: + return 1 + else: + return math.ceil(math.log(len(palette_bytes) // 3, 2)) - 1 + + +def _get_header_palette(palette_bytes: bytes) -> bytes: + """ + Returns the palette, null padded to the next power of 2 (*3) bytes + suitable for direct inclusion in the GIF header + + :param palette_bytes: Unpadded palette bytes, in RGBRGB form + :returns: Null padded palette + """ + color_table_size = _get_color_table_size(palette_bytes) + + # add the missing amount of bytes + # the palette has to be 2< 0: + palette_bytes += o8(0) * 3 * actual_target_size_diff + return palette_bytes + + +def _get_palette_bytes(im: Image.Image) -> bytes: + """ + Gets the palette for inclusion in the gif header + + :param im: Image object + :returns: Bytes, len<=768 suitable for inclusion in gif header + """ + if not im.palette: + return b"" + + palette = bytes(im.palette.palette) + if im.palette.mode == "RGBA": + palette = b"".join(palette[i * 4 : i * 4 + 3] for i in range(len(palette) // 3)) + return palette + + +def _get_background( + im: Image.Image, + info_background: int | tuple[int, int, int] | tuple[int, int, int, int] | None, +) -> int: + background = 0 + if info_background: + if isinstance(info_background, tuple): + # WebPImagePlugin stores an RGBA value in info["background"] + # So it must be converted to the same format as GifImagePlugin's + # info["background"] - a global color table index + assert im.palette is not None + try: + background = im.palette.getcolor(info_background, im) + except ValueError as e: + if str(e) not in ( + # If all 256 colors are in use, + # then there is no need for the background color + "cannot allocate more than 256 colors", + # Ignore non-opaque WebP background + "cannot add non-opaque RGBA color to RGB palette", + ): + raise + else: + background = info_background + return background + + +def _get_global_header(im: Image.Image, info: dict[str, Any]) -> list[bytes]: + """Return a list of strings representing a GIF header""" + + # Header Block + # https://www.matthewflickinger.com/lab/whatsinagif/bits_and_bytes.asp + + version = b"87a" + if im.info.get("version") == b"89a" or ( + info + and ( + "transparency" in info + or info.get("loop") is not None + or info.get("duration") + or info.get("comment") + ) + ): + version = b"89a" + + background = _get_background(im, info.get("background")) + + palette_bytes = _get_palette_bytes(im) + color_table_size = _get_color_table_size(palette_bytes) + + header = [ + b"GIF" # signature + + version # version + + o16(im.size[0]) # canvas width + + o16(im.size[1]), # canvas height + # Logical Screen Descriptor + # size of global color table + global color table flag + o8(color_table_size + 128), # packed fields + # background + reserved/aspect + o8(background) + o8(0), + # Global Color Table + _get_header_palette(palette_bytes), + ] + if info.get("loop") is not None: + header.append( + b"!" + + o8(255) # extension intro + + o8(11) + + b"NETSCAPE2.0" + + o8(3) + + o8(1) + + o16(info["loop"]) # number of loops + + o8(0) + ) + if info.get("comment"): + comment_block = b"!" + o8(254) # extension intro + + comment = info["comment"] + if isinstance(comment, str): + comment = comment.encode() + for i in range(0, len(comment), 255): + subblock = comment[i : i + 255] + comment_block += o8(len(subblock)) + subblock + + comment_block += o8(0) + header.append(comment_block) + return header + + +def _write_frame_data( + fp: IO[bytes], + im_frame: Image.Image, + offset: tuple[int, int], + params: dict[str, Any], +) -> None: + try: + im_frame.encoderinfo = params + + # local image header + _write_local_header(fp, im_frame, offset, 0) + + ImageFile._save( + im_frame, + fp, + [ImageFile._Tile("gif", (0, 0) + im_frame.size, 0, RAWMODE[im_frame.mode])], + ) + + fp.write(b"\0") # end of image data + finally: + del im_frame.encoderinfo + + +# -------------------------------------------------------------------- +# Legacy GIF utilities + + +def getheader( + im: Image.Image, palette: _Palette | None = None, info: dict[str, Any] | None = None +) -> tuple[list[bytes], list[int] | None]: + """ + Legacy Method to get Gif data from image. + + Warning:: May modify image data. + + :param im: Image object + :param palette: bytes object containing the source palette, or .... + :param info: encoderinfo + :returns: tuple of(list of header items, optimized palette) + + """ + if info is None: + info = {} + + used_palette_colors = _get_optimize(im, info) + + if "background" not in info and "background" in im.info: + info["background"] = im.info["background"] + + im_mod = _normalize_palette(im, palette, info) + im.palette = im_mod.palette + im.im = im_mod.im + header = _get_global_header(im, info) + + return header, used_palette_colors + + +def getdata( + im: Image.Image, offset: tuple[int, int] = (0, 0), **params: Any +) -> list[bytes]: + """ + Legacy Method + + Return a list of strings representing this image. + The first string is a local image header, the rest contains + encoded image data. + + To specify duration, add the time in milliseconds, + e.g. ``getdata(im_frame, duration=1000)`` + + :param im: Image object + :param offset: Tuple of (x, y) pixels. Defaults to (0, 0) + :param \\**params: e.g. duration or other encoder info parameters + :returns: List of bytes containing GIF encoded frame data + + """ + from io import BytesIO + + class Collector(BytesIO): + data = [] + + def write(self, data: Buffer) -> int: + self.data.append(data) + return len(data) + + im.load() # make sure raster data is available + + fp = Collector() + + _write_frame_data(fp, im, offset, params) + + return fp.data + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(GifImageFile.format, GifImageFile, _accept) +Image.register_save(GifImageFile.format, _save) +Image.register_save_all(GifImageFile.format, _save_all) +Image.register_extension(GifImageFile.format, ".gif") +Image.register_mime(GifImageFile.format, "image/gif") + +# +# Uncomment the following line if you wish to use NETPBM/PBMPLUS +# instead of the built-in "uncompressed" GIF encoder + +# Image.register_save(GifImageFile.format, _save_netpbm) diff --git a/venv/lib/python3.10/site-packages/PIL/GimpGradientFile.py b/venv/lib/python3.10/site-packages/PIL/GimpGradientFile.py new file mode 100644 index 0000000000000000000000000000000000000000..5f2691882c46130fc2f83c45f01db34e6ce1efe6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GimpGradientFile.py @@ -0,0 +1,153 @@ +# +# Python Imaging Library +# $Id$ +# +# stuff to read (and render) GIMP gradient files +# +# History: +# 97-08-23 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# + +""" +Stuff to translate curve segments to palette values (derived from +the corresponding code in GIMP, written by Federico Mena Quintero. +See the GIMP distribution for more information.) +""" +from __future__ import annotations + +from math import log, pi, sin, sqrt + +from ._binary import o8 + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from typing import IO + +EPSILON = 1e-10 +"""""" # Enable auto-doc for data member + + +def linear(middle: float, pos: float) -> float: + if pos <= middle: + if middle < EPSILON: + return 0.0 + else: + return 0.5 * pos / middle + else: + pos = pos - middle + middle = 1.0 - middle + if middle < EPSILON: + return 1.0 + else: + return 0.5 + 0.5 * pos / middle + + +def curved(middle: float, pos: float) -> float: + return pos ** (log(0.5) / log(max(middle, EPSILON))) + + +def sine(middle: float, pos: float) -> float: + return (sin((-pi / 2.0) + pi * linear(middle, pos)) + 1.0) / 2.0 + + +def sphere_increasing(middle: float, pos: float) -> float: + return sqrt(1.0 - (linear(middle, pos) - 1.0) ** 2) + + +def sphere_decreasing(middle: float, pos: float) -> float: + return 1.0 - sqrt(1.0 - linear(middle, pos) ** 2) + + +SEGMENTS = [linear, curved, sine, sphere_increasing, sphere_decreasing] +"""""" # Enable auto-doc for data member + + +class GradientFile: + gradient: ( + list[ + tuple[ + float, + float, + float, + list[float], + list[float], + Callable[[float, float], float], + ] + ] + | None + ) = None + + def getpalette(self, entries: int = 256) -> tuple[bytes, str]: + assert self.gradient is not None + palette = [] + + ix = 0 + x0, x1, xm, rgb0, rgb1, segment = self.gradient[ix] + + for i in range(entries): + x = i / (entries - 1) + + while x1 < x: + ix += 1 + x0, x1, xm, rgb0, rgb1, segment = self.gradient[ix] + + w = x1 - x0 + + if w < EPSILON: + scale = segment(0.5, 0.5) + else: + scale = segment((xm - x0) / w, (x - x0) / w) + + # expand to RGBA + r = o8(int(255 * ((rgb1[0] - rgb0[0]) * scale + rgb0[0]) + 0.5)) + g = o8(int(255 * ((rgb1[1] - rgb0[1]) * scale + rgb0[1]) + 0.5)) + b = o8(int(255 * ((rgb1[2] - rgb0[2]) * scale + rgb0[2]) + 0.5)) + a = o8(int(255 * ((rgb1[3] - rgb0[3]) * scale + rgb0[3]) + 0.5)) + + # add to palette + palette.append(r + g + b + a) + + return b"".join(palette), "RGBA" + + +class GimpGradientFile(GradientFile): + """File handler for GIMP's gradient format.""" + + def __init__(self, fp: IO[bytes]) -> None: + if not fp.readline().startswith(b"GIMP Gradient"): + msg = "not a GIMP gradient file" + raise SyntaxError(msg) + + line = fp.readline() + + # GIMP 1.2 gradient files don't contain a name, but GIMP 1.3 files do + if line.startswith(b"Name: "): + line = fp.readline().strip() + + count = int(line) + + self.gradient = [] + + for i in range(count): + s = fp.readline().split() + w = [float(x) for x in s[:11]] + + x0, x1 = w[0], w[2] + xm = w[1] + rgb0 = w[3:7] + rgb1 = w[7:11] + + segment = SEGMENTS[int(s[11])] + cspace = int(s[12]) + + if cspace != 0: + msg = "cannot handle HSV colour space" + raise OSError(msg) + + self.gradient.append((x0, x1, xm, rgb0, rgb1, segment)) diff --git a/venv/lib/python3.10/site-packages/PIL/GimpPaletteFile.py b/venv/lib/python3.10/site-packages/PIL/GimpPaletteFile.py new file mode 100644 index 0000000000000000000000000000000000000000..016257d3dd29e83ed8d68f89363965dc39b93811 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GimpPaletteFile.py @@ -0,0 +1,75 @@ +# +# Python Imaging Library +# $Id$ +# +# stuff to read GIMP palette files +# +# History: +# 1997-08-23 fl Created +# 2004-09-07 fl Support GIMP 2.0 palette files. +# +# Copyright (c) Secret Labs AB 1997-2004. All rights reserved. +# Copyright (c) Fredrik Lundh 1997-2004. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import re +from io import BytesIO + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import IO + + +class GimpPaletteFile: + """File handler for GIMP's palette format.""" + + rawmode = "RGB" + + def _read(self, fp: IO[bytes], limit: bool = True) -> None: + if not fp.readline().startswith(b"GIMP Palette"): + msg = "not a GIMP palette file" + raise SyntaxError(msg) + + palette: list[int] = [] + i = 0 + while True: + if limit and i == 256 + 3: + break + + i += 1 + s = fp.readline() + if not s: + break + + # skip fields and comment lines + if re.match(rb"\w+:|#", s): + continue + if limit and len(s) > 100: + msg = "bad palette file" + raise SyntaxError(msg) + + v = s.split(maxsplit=3) + if len(v) < 3: + msg = "bad palette entry" + raise ValueError(msg) + + palette += (int(v[i]) for i in range(3)) + if limit and len(palette) == 768: + break + + self.palette = bytes(palette) + + def __init__(self, fp: IO[bytes]) -> None: + self._read(fp) + + @classmethod + def frombytes(cls, data: bytes) -> GimpPaletteFile: + self = cls.__new__(cls) + self._read(BytesIO(data), False) + return self + + def getpalette(self) -> tuple[bytes, str]: + return self.palette, self.rawmode diff --git a/venv/lib/python3.10/site-packages/PIL/GribStubImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/GribStubImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..dfa798893cbd6c0adc2c787ea210a8a0d6433f8f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/GribStubImagePlugin.py @@ -0,0 +1,75 @@ +# +# The Python Imaging Library +# $Id$ +# +# GRIB stub adapter +# +# Copyright (c) 1996-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +from typing import IO + +from . import Image, ImageFile + +_handler = None + + +def register_handler(handler: ImageFile.StubHandler | None) -> None: + """ + Install application-specific GRIB image handler. + + :param handler: Handler object. + """ + global _handler + _handler = handler + + +# -------------------------------------------------------------------- +# Image adapter + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 8 and prefix.startswith(b"GRIB") and prefix[7] == 1 + + +class GribStubImageFile(ImageFile.StubImageFile): + format = "GRIB" + format_description = "GRIB" + + def _open(self) -> None: + if not _accept(self.fp.read(8)): + msg = "Not a GRIB file" + raise SyntaxError(msg) + + self.fp.seek(-8, os.SEEK_CUR) + + # make something up + self._mode = "F" + self._size = 1, 1 + + loader = self._load() + if loader: + loader.open(self) + + def _load(self) -> ImageFile.StubHandler | None: + return _handler + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if _handler is None or not hasattr(_handler, "save"): + msg = "GRIB save handler not installed" + raise OSError(msg) + _handler.save(im, fp, filename) + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(GribStubImageFile.format, GribStubImageFile, _accept) +Image.register_save(GribStubImageFile.format, _save) + +Image.register_extension(GribStubImageFile.format, ".grib") diff --git a/venv/lib/python3.10/site-packages/PIL/Hdf5StubImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/Hdf5StubImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..76e640f15abfe60a56a571380133a0463c104035 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/Hdf5StubImagePlugin.py @@ -0,0 +1,75 @@ +# +# The Python Imaging Library +# $Id$ +# +# HDF5 stub adapter +# +# Copyright (c) 2000-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +from typing import IO + +from . import Image, ImageFile + +_handler = None + + +def register_handler(handler: ImageFile.StubHandler | None) -> None: + """ + Install application-specific HDF5 image handler. + + :param handler: Handler object. + """ + global _handler + _handler = handler + + +# -------------------------------------------------------------------- +# Image adapter + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"\x89HDF\r\n\x1a\n") + + +class HDF5StubImageFile(ImageFile.StubImageFile): + format = "HDF5" + format_description = "HDF5" + + def _open(self) -> None: + if not _accept(self.fp.read(8)): + msg = "Not an HDF file" + raise SyntaxError(msg) + + self.fp.seek(-8, os.SEEK_CUR) + + # make something up + self._mode = "F" + self._size = 1, 1 + + loader = self._load() + if loader: + loader.open(self) + + def _load(self) -> ImageFile.StubHandler | None: + return _handler + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if _handler is None or not hasattr(_handler, "save"): + msg = "HDF5 save handler not installed" + raise OSError(msg) + _handler.save(im, fp, filename) + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(HDF5StubImageFile.format, HDF5StubImageFile, _accept) +Image.register_save(HDF5StubImageFile.format, _save) + +Image.register_extensions(HDF5StubImageFile.format, [".h5", ".hdf"]) diff --git a/venv/lib/python3.10/site-packages/PIL/IcnsImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/IcnsImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..197ea7a2bb27a8b98d3331550b68a9d68d3e23a5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/IcnsImagePlugin.py @@ -0,0 +1,401 @@ +# +# The Python Imaging Library. +# $Id$ +# +# macOS icns file decoder, based on icns.py by Bob Ippolito. +# +# history: +# 2004-10-09 fl Turned into a PIL plugin; removed 2.3 dependencies. +# 2020-04-04 Allow saving on all operating systems. +# +# Copyright (c) 2004 by Bob Ippolito. +# Copyright (c) 2004 by Secret Labs. +# Copyright (c) 2004 by Fredrik Lundh. +# Copyright (c) 2014 by Alastair Houghton. +# Copyright (c) 2020 by Pan Jing. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import os +import struct +import sys +from typing import IO + +from . import Image, ImageFile, PngImagePlugin, features + +enable_jpeg2k = features.check_codec("jpg_2000") +if enable_jpeg2k: + from . import Jpeg2KImagePlugin + +MAGIC = b"icns" +HEADERSIZE = 8 + + +def nextheader(fobj: IO[bytes]) -> tuple[bytes, int]: + return struct.unpack(">4sI", fobj.read(HEADERSIZE)) + + +def read_32t( + fobj: IO[bytes], start_length: tuple[int, int], size: tuple[int, int, int] +) -> dict[str, Image.Image]: + # The 128x128 icon seems to have an extra header for some reason. + (start, length) = start_length + fobj.seek(start) + sig = fobj.read(4) + if sig != b"\x00\x00\x00\x00": + msg = "Unknown signature, expecting 0x00000000" + raise SyntaxError(msg) + return read_32(fobj, (start + 4, length - 4), size) + + +def read_32( + fobj: IO[bytes], start_length: tuple[int, int], size: tuple[int, int, int] +) -> dict[str, Image.Image]: + """ + Read a 32bit RGB icon resource. Seems to be either uncompressed or + an RLE packbits-like scheme. + """ + (start, length) = start_length + fobj.seek(start) + pixel_size = (size[0] * size[2], size[1] * size[2]) + sizesq = pixel_size[0] * pixel_size[1] + if length == sizesq * 3: + # uncompressed ("RGBRGBGB") + indata = fobj.read(length) + im = Image.frombuffer("RGB", pixel_size, indata, "raw", "RGB", 0, 1) + else: + # decode image + im = Image.new("RGB", pixel_size, None) + for band_ix in range(3): + data = [] + bytesleft = sizesq + while bytesleft > 0: + byte = fobj.read(1) + if not byte: + break + byte_int = byte[0] + if byte_int & 0x80: + blocksize = byte_int - 125 + byte = fobj.read(1) + for i in range(blocksize): + data.append(byte) + else: + blocksize = byte_int + 1 + data.append(fobj.read(blocksize)) + bytesleft -= blocksize + if bytesleft <= 0: + break + if bytesleft != 0: + msg = f"Error reading channel [{repr(bytesleft)} left]" + raise SyntaxError(msg) + band = Image.frombuffer("L", pixel_size, b"".join(data), "raw", "L", 0, 1) + im.im.putband(band.im, band_ix) + return {"RGB": im} + + +def read_mk( + fobj: IO[bytes], start_length: tuple[int, int], size: tuple[int, int, int] +) -> dict[str, Image.Image]: + # Alpha masks seem to be uncompressed + start = start_length[0] + fobj.seek(start) + pixel_size = (size[0] * size[2], size[1] * size[2]) + sizesq = pixel_size[0] * pixel_size[1] + band = Image.frombuffer("L", pixel_size, fobj.read(sizesq), "raw", "L", 0, 1) + return {"A": band} + + +def read_png_or_jpeg2000( + fobj: IO[bytes], start_length: tuple[int, int], size: tuple[int, int, int] +) -> dict[str, Image.Image]: + (start, length) = start_length + fobj.seek(start) + sig = fobj.read(12) + + im: Image.Image + if sig.startswith(b"\x89PNG\x0d\x0a\x1a\x0a"): + fobj.seek(start) + im = PngImagePlugin.PngImageFile(fobj) + Image._decompression_bomb_check(im.size) + return {"RGBA": im} + elif ( + sig.startswith((b"\xff\x4f\xff\x51", b"\x0d\x0a\x87\x0a")) + or sig == b"\x00\x00\x00\x0cjP \x0d\x0a\x87\x0a" + ): + if not enable_jpeg2k: + msg = ( + "Unsupported icon subimage format (rebuild PIL " + "with JPEG 2000 support to fix this)" + ) + raise ValueError(msg) + # j2k, jpc or j2c + fobj.seek(start) + jp2kstream = fobj.read(length) + f = io.BytesIO(jp2kstream) + im = Jpeg2KImagePlugin.Jpeg2KImageFile(f) + Image._decompression_bomb_check(im.size) + if im.mode != "RGBA": + im = im.convert("RGBA") + return {"RGBA": im} + else: + msg = "Unsupported icon subimage format" + raise ValueError(msg) + + +class IcnsFile: + SIZES = { + (512, 512, 2): [(b"ic10", read_png_or_jpeg2000)], + (512, 512, 1): [(b"ic09", read_png_or_jpeg2000)], + (256, 256, 2): [(b"ic14", read_png_or_jpeg2000)], + (256, 256, 1): [(b"ic08", read_png_or_jpeg2000)], + (128, 128, 2): [(b"ic13", read_png_or_jpeg2000)], + (128, 128, 1): [ + (b"ic07", read_png_or_jpeg2000), + (b"it32", read_32t), + (b"t8mk", read_mk), + ], + (64, 64, 1): [(b"icp6", read_png_or_jpeg2000)], + (32, 32, 2): [(b"ic12", read_png_or_jpeg2000)], + (48, 48, 1): [(b"ih32", read_32), (b"h8mk", read_mk)], + (32, 32, 1): [ + (b"icp5", read_png_or_jpeg2000), + (b"il32", read_32), + (b"l8mk", read_mk), + ], + (16, 16, 2): [(b"ic11", read_png_or_jpeg2000)], + (16, 16, 1): [ + (b"icp4", read_png_or_jpeg2000), + (b"is32", read_32), + (b"s8mk", read_mk), + ], + } + + def __init__(self, fobj: IO[bytes]) -> None: + """ + fobj is a file-like object as an icns resource + """ + # signature : (start, length) + self.dct = {} + self.fobj = fobj + sig, filesize = nextheader(fobj) + if not _accept(sig): + msg = "not an icns file" + raise SyntaxError(msg) + i = HEADERSIZE + while i < filesize: + sig, blocksize = nextheader(fobj) + if blocksize <= 0: + msg = "invalid block header" + raise SyntaxError(msg) + i += HEADERSIZE + blocksize -= HEADERSIZE + self.dct[sig] = (i, blocksize) + fobj.seek(blocksize, io.SEEK_CUR) + i += blocksize + + def itersizes(self) -> list[tuple[int, int, int]]: + sizes = [] + for size, fmts in self.SIZES.items(): + for fmt, reader in fmts: + if fmt in self.dct: + sizes.append(size) + break + return sizes + + def bestsize(self) -> tuple[int, int, int]: + sizes = self.itersizes() + if not sizes: + msg = "No 32bit icon resources found" + raise SyntaxError(msg) + return max(sizes) + + def dataforsize(self, size: tuple[int, int, int]) -> dict[str, Image.Image]: + """ + Get an icon resource as {channel: array}. Note that + the arrays are bottom-up like windows bitmaps and will likely + need to be flipped or transposed in some way. + """ + dct = {} + for code, reader in self.SIZES[size]: + desc = self.dct.get(code) + if desc is not None: + dct.update(reader(self.fobj, desc, size)) + return dct + + def getimage( + self, size: tuple[int, int] | tuple[int, int, int] | None = None + ) -> Image.Image: + if size is None: + size = self.bestsize() + elif len(size) == 2: + size = (size[0], size[1], 1) + channels = self.dataforsize(size) + + im = channels.get("RGBA") + if im: + return im + + im = channels["RGB"].copy() + try: + im.putalpha(channels["A"]) + except KeyError: + pass + return im + + +## +# Image plugin for Mac OS icons. + + +class IcnsImageFile(ImageFile.ImageFile): + """ + PIL image support for Mac OS .icns files. + Chooses the best resolution, but will possibly load + a different size image if you mutate the size attribute + before calling 'load'. + + The info dictionary has a key 'sizes' that is a list + of sizes that the icns file has. + """ + + format = "ICNS" + format_description = "Mac OS icns resource" + + def _open(self) -> None: + self.icns = IcnsFile(self.fp) + self._mode = "RGBA" + self.info["sizes"] = self.icns.itersizes() + self.best_size = self.icns.bestsize() + self.size = ( + self.best_size[0] * self.best_size[2], + self.best_size[1] * self.best_size[2], + ) + + @property + def size(self) -> tuple[int, int]: + return self._size + + @size.setter + def size(self, value: tuple[int, int]) -> None: + # Check that a matching size exists, + # or that there is a scale that would create a size that matches + for size in self.info["sizes"]: + simple_size = size[0] * size[2], size[1] * size[2] + scale = simple_size[0] // value[0] + if simple_size[1] / value[1] == scale: + self._size = value + return + msg = "This is not one of the allowed sizes of this image" + raise ValueError(msg) + + def load(self, scale: int | None = None) -> Image.core.PixelAccess | None: + if scale is not None: + width, height = self.size[:2] + self.size = width * scale, height * scale + self.best_size = width, height, scale + + px = Image.Image.load(self) + if self._im is not None and self.im.size == self.size: + # Already loaded + return px + self.load_prepare() + # This is likely NOT the best way to do it, but whatever. + im = self.icns.getimage(self.best_size) + + # If this is a PNG or JPEG 2000, it won't be loaded yet + px = im.load() + + self.im = im.im + self._mode = im.mode + self.size = im.size + + return px + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + """ + Saves the image as a series of PNG files, + that are then combined into a .icns file. + """ + if hasattr(fp, "flush"): + fp.flush() + + sizes = { + b"ic07": 128, + b"ic08": 256, + b"ic09": 512, + b"ic10": 1024, + b"ic11": 32, + b"ic12": 64, + b"ic13": 256, + b"ic14": 512, + } + provided_images = {im.width: im for im in im.encoderinfo.get("append_images", [])} + size_streams = {} + for size in set(sizes.values()): + image = ( + provided_images[size] + if size in provided_images + else im.resize((size, size)) + ) + + temp = io.BytesIO() + image.save(temp, "png") + size_streams[size] = temp.getvalue() + + entries = [] + for type, size in sizes.items(): + stream = size_streams[size] + entries.append((type, HEADERSIZE + len(stream), stream)) + + # Header + fp.write(MAGIC) + file_length = HEADERSIZE # Header + file_length += HEADERSIZE + 8 * len(entries) # TOC + file_length += sum(entry[1] for entry in entries) + fp.write(struct.pack(">i", file_length)) + + # TOC + fp.write(b"TOC ") + fp.write(struct.pack(">i", HEADERSIZE + len(entries) * HEADERSIZE)) + for entry in entries: + fp.write(entry[0]) + fp.write(struct.pack(">i", entry[1])) + + # Data + for entry in entries: + fp.write(entry[0]) + fp.write(struct.pack(">i", entry[1])) + fp.write(entry[2]) + + if hasattr(fp, "flush"): + fp.flush() + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(MAGIC) + + +Image.register_open(IcnsImageFile.format, IcnsImageFile, _accept) +Image.register_extension(IcnsImageFile.format, ".icns") + +Image.register_save(IcnsImageFile.format, _save) +Image.register_mime(IcnsImageFile.format, "image/icns") + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Syntax: python3 IcnsImagePlugin.py [file]") + sys.exit() + + with open(sys.argv[1], "rb") as fp: + imf = IcnsImageFile(fp) + for size in imf.info["sizes"]: + width, height, scale = imf.size = size + imf.save(f"out-{width}-{height}-{scale}.png") + with Image.open(sys.argv[1]) as im: + im.save("out.png") + if sys.platform == "windows": + os.startfile("out.png") diff --git a/venv/lib/python3.10/site-packages/PIL/IcoImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/IcoImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..bd35ac890e6cf824e9c890404416d871e5b94f7c --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/IcoImagePlugin.py @@ -0,0 +1,381 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Windows Icon support for PIL +# +# History: +# 96-05-27 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# + +# This plugin is a refactored version of Win32IconImagePlugin by Bryan Davis +# . +# https://code.google.com/archive/p/casadebender/wikis/Win32IconImagePlugin.wiki +# +# Icon format references: +# * https://en.wikipedia.org/wiki/ICO_(file_format) +# * https://msdn.microsoft.com/en-us/library/ms997538.aspx +from __future__ import annotations + +import warnings +from io import BytesIO +from math import ceil, log +from typing import IO, NamedTuple + +from . import BmpImagePlugin, Image, ImageFile, PngImagePlugin +from ._binary import i16le as i16 +from ._binary import i32le as i32 +from ._binary import o8 +from ._binary import o16le as o16 +from ._binary import o32le as o32 + +# +# -------------------------------------------------------------------- + +_MAGIC = b"\0\0\1\0" + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + fp.write(_MAGIC) # (2+2) + bmp = im.encoderinfo.get("bitmap_format") == "bmp" + sizes = im.encoderinfo.get( + "sizes", + [(16, 16), (24, 24), (32, 32), (48, 48), (64, 64), (128, 128), (256, 256)], + ) + frames = [] + provided_ims = [im] + im.encoderinfo.get("append_images", []) + width, height = im.size + for size in sorted(set(sizes)): + if size[0] > width or size[1] > height or size[0] > 256 or size[1] > 256: + continue + + for provided_im in provided_ims: + if provided_im.size != size: + continue + frames.append(provided_im) + if bmp: + bits = BmpImagePlugin.SAVE[provided_im.mode][1] + bits_used = [bits] + for other_im in provided_ims: + if other_im.size != size: + continue + bits = BmpImagePlugin.SAVE[other_im.mode][1] + if bits not in bits_used: + # Another image has been supplied for this size + # with a different bit depth + frames.append(other_im) + bits_used.append(bits) + break + else: + # TODO: invent a more convenient method for proportional scalings + frame = provided_im.copy() + frame.thumbnail(size, Image.Resampling.LANCZOS, reducing_gap=None) + frames.append(frame) + fp.write(o16(len(frames))) # idCount(2) + offset = fp.tell() + len(frames) * 16 + for frame in frames: + width, height = frame.size + # 0 means 256 + fp.write(o8(width if width < 256 else 0)) # bWidth(1) + fp.write(o8(height if height < 256 else 0)) # bHeight(1) + + bits, colors = BmpImagePlugin.SAVE[frame.mode][1:] if bmp else (32, 0) + fp.write(o8(colors)) # bColorCount(1) + fp.write(b"\0") # bReserved(1) + fp.write(b"\0\0") # wPlanes(2) + fp.write(o16(bits)) # wBitCount(2) + + image_io = BytesIO() + if bmp: + frame.save(image_io, "dib") + + if bits != 32: + and_mask = Image.new("1", size) + ImageFile._save( + and_mask, + image_io, + [ImageFile._Tile("raw", (0, 0) + size, 0, ("1", 0, -1))], + ) + else: + frame.save(image_io, "png") + image_io.seek(0) + image_bytes = image_io.read() + if bmp: + image_bytes = image_bytes[:8] + o32(height * 2) + image_bytes[12:] + bytes_len = len(image_bytes) + fp.write(o32(bytes_len)) # dwBytesInRes(4) + fp.write(o32(offset)) # dwImageOffset(4) + current = fp.tell() + fp.seek(offset) + fp.write(image_bytes) + offset = offset + bytes_len + fp.seek(current) + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(_MAGIC) + + +class IconHeader(NamedTuple): + width: int + height: int + nb_color: int + reserved: int + planes: int + bpp: int + size: int + offset: int + dim: tuple[int, int] + square: int + color_depth: int + + +class IcoFile: + def __init__(self, buf: IO[bytes]) -> None: + """ + Parse image from file-like object containing ico file data + """ + + # check magic + s = buf.read(6) + if not _accept(s): + msg = "not an ICO file" + raise SyntaxError(msg) + + self.buf = buf + self.entry = [] + + # Number of items in file + self.nb_items = i16(s, 4) + + # Get headers for each item + for i in range(self.nb_items): + s = buf.read(16) + + # See Wikipedia + width = s[0] or 256 + height = s[1] or 256 + + # No. of colors in image (0 if >=8bpp) + nb_color = s[2] + bpp = i16(s, 6) + icon_header = IconHeader( + width=width, + height=height, + nb_color=nb_color, + reserved=s[3], + planes=i16(s, 4), + bpp=i16(s, 6), + size=i32(s, 8), + offset=i32(s, 12), + dim=(width, height), + square=width * height, + # See Wikipedia notes about color depth. + # We need this just to differ images with equal sizes + color_depth=bpp or (nb_color != 0 and ceil(log(nb_color, 2))) or 256, + ) + + self.entry.append(icon_header) + + self.entry = sorted(self.entry, key=lambda x: x.color_depth) + # ICO images are usually squares + self.entry = sorted(self.entry, key=lambda x: x.square, reverse=True) + + def sizes(self) -> set[tuple[int, int]]: + """ + Get a set of all available icon sizes and color depths. + """ + return {(h.width, h.height) for h in self.entry} + + def getentryindex(self, size: tuple[int, int], bpp: int | bool = False) -> int: + for i, h in enumerate(self.entry): + if size == h.dim and (bpp is False or bpp == h.color_depth): + return i + return 0 + + def getimage(self, size: tuple[int, int], bpp: int | bool = False) -> Image.Image: + """ + Get an image from the icon + """ + return self.frame(self.getentryindex(size, bpp)) + + def frame(self, idx: int) -> Image.Image: + """ + Get an image from frame idx + """ + + header = self.entry[idx] + + self.buf.seek(header.offset) + data = self.buf.read(8) + self.buf.seek(header.offset) + + im: Image.Image + if data[:8] == PngImagePlugin._MAGIC: + # png frame + im = PngImagePlugin.PngImageFile(self.buf) + Image._decompression_bomb_check(im.size) + else: + # XOR + AND mask bmp frame + im = BmpImagePlugin.DibImageFile(self.buf) + Image._decompression_bomb_check(im.size) + + # change tile dimension to only encompass XOR image + im._size = (im.size[0], int(im.size[1] / 2)) + d, e, o, a = im.tile[0] + im.tile[0] = ImageFile._Tile(d, (0, 0) + im.size, o, a) + + # figure out where AND mask image starts + if header.bpp == 32: + # 32-bit color depth icon image allows semitransparent areas + # PIL's DIB format ignores transparency bits, recover them. + # The DIB is packed in BGRX byte order where X is the alpha + # channel. + + # Back up to start of bmp data + self.buf.seek(o) + # extract every 4th byte (eg. 3,7,11,15,...) + alpha_bytes = self.buf.read(im.size[0] * im.size[1] * 4)[3::4] + + # convert to an 8bpp grayscale image + try: + mask = Image.frombuffer( + "L", # 8bpp + im.size, # (w, h) + alpha_bytes, # source chars + "raw", # raw decoder + ("L", 0, -1), # 8bpp inverted, unpadded, reversed + ) + except ValueError: + if ImageFile.LOAD_TRUNCATED_IMAGES: + mask = None + else: + raise + else: + # get AND image from end of bitmap + w = im.size[0] + if (w % 32) > 0: + # bitmap row data is aligned to word boundaries + w += 32 - (im.size[0] % 32) + + # the total mask data is + # padded row size * height / bits per char + + total_bytes = int((w * im.size[1]) / 8) + and_mask_offset = header.offset + header.size - total_bytes + + self.buf.seek(and_mask_offset) + mask_data = self.buf.read(total_bytes) + + # convert raw data to image + try: + mask = Image.frombuffer( + "1", # 1 bpp + im.size, # (w, h) + mask_data, # source chars + "raw", # raw decoder + ("1;I", int(w / 8), -1), # 1bpp inverted, padded, reversed + ) + except ValueError: + if ImageFile.LOAD_TRUNCATED_IMAGES: + mask = None + else: + raise + + # now we have two images, im is XOR image and mask is AND image + + # apply mask image as alpha channel + if mask: + im = im.convert("RGBA") + im.putalpha(mask) + + return im + + +## +# Image plugin for Windows Icon files. + + +class IcoImageFile(ImageFile.ImageFile): + """ + PIL read-only image support for Microsoft Windows .ico files. + + By default the largest resolution image in the file will be loaded. This + can be changed by altering the 'size' attribute before calling 'load'. + + The info dictionary has a key 'sizes' that is a list of the sizes available + in the icon file. + + Handles classic, XP and Vista icon formats. + + When saving, PNG compression is used. Support for this was only added in + Windows Vista. If you are unable to view the icon in Windows, convert the + image to "RGBA" mode before saving. + + This plugin is a refactored version of Win32IconImagePlugin by Bryan Davis + . + https://code.google.com/archive/p/casadebender/wikis/Win32IconImagePlugin.wiki + """ + + format = "ICO" + format_description = "Windows Icon" + + def _open(self) -> None: + self.ico = IcoFile(self.fp) + self.info["sizes"] = self.ico.sizes() + self.size = self.ico.entry[0].dim + self.load() + + @property + def size(self) -> tuple[int, int]: + return self._size + + @size.setter + def size(self, value: tuple[int, int]) -> None: + if value not in self.info["sizes"]: + msg = "This is not one of the allowed sizes of this image" + raise ValueError(msg) + self._size = value + + def load(self) -> Image.core.PixelAccess | None: + if self._im is not None and self.im.size == self.size: + # Already loaded + return Image.Image.load(self) + im = self.ico.getimage(self.size) + # if tile is PNG, it won't really be loaded yet + im.load() + self.im = im.im + self._mode = im.mode + if im.palette: + self.palette = im.palette + if im.size != self.size: + warnings.warn("Image was not the expected size") + + index = self.ico.getentryindex(self.size) + sizes = list(self.info["sizes"]) + sizes[index] = im.size + self.info["sizes"] = set(sizes) + + self.size = im.size + return Image.Image.load(self) + + def load_seek(self, pos: int) -> None: + # Flag the ImageFile.Parser so that it + # just does all the decode at the end. + pass + + +# +# -------------------------------------------------------------------- + + +Image.register_open(IcoImageFile.format, IcoImageFile, _accept) +Image.register_save(IcoImageFile.format, _save) +Image.register_extension(IcoImageFile.format, ".ico") + +Image.register_mime(IcoImageFile.format, "image/x-icon") diff --git a/venv/lib/python3.10/site-packages/PIL/ImImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/ImImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..71b9996780ce8dfc420670b5732216f934a1f677 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImImagePlugin.py @@ -0,0 +1,389 @@ +# +# The Python Imaging Library. +# $Id$ +# +# IFUNC IM file handling for PIL +# +# history: +# 1995-09-01 fl Created. +# 1997-01-03 fl Save palette images +# 1997-01-08 fl Added sequence support +# 1997-01-23 fl Added P and RGB save support +# 1997-05-31 fl Read floating point images +# 1997-06-22 fl Save floating point images +# 1997-08-27 fl Read and save 1-bit images +# 1998-06-25 fl Added support for RGB+LUT images +# 1998-07-02 fl Added support for YCC images +# 1998-07-15 fl Renamed offset attribute to avoid name clash +# 1998-12-29 fl Added I;16 support +# 2001-02-17 fl Use 're' instead of 'regex' (Python 2.1) (0.7) +# 2003-09-26 fl Added LA/PA support +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1995-2001 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +import re +from typing import IO, Any + +from . import Image, ImageFile, ImagePalette +from ._util import DeferredError + +# -------------------------------------------------------------------- +# Standard tags + +COMMENT = "Comment" +DATE = "Date" +EQUIPMENT = "Digitalization equipment" +FRAMES = "File size (no of images)" +LUT = "Lut" +NAME = "Name" +SCALE = "Scale (x,y)" +SIZE = "Image size (x*y)" +MODE = "Image type" + +TAGS = { + COMMENT: 0, + DATE: 0, + EQUIPMENT: 0, + FRAMES: 0, + LUT: 0, + NAME: 0, + SCALE: 0, + SIZE: 0, + MODE: 0, +} + +OPEN = { + # ifunc93/p3cfunc formats + "0 1 image": ("1", "1"), + "L 1 image": ("1", "1"), + "Greyscale image": ("L", "L"), + "Grayscale image": ("L", "L"), + "RGB image": ("RGB", "RGB;L"), + "RLB image": ("RGB", "RLB"), + "RYB image": ("RGB", "RLB"), + "B1 image": ("1", "1"), + "B2 image": ("P", "P;2"), + "B4 image": ("P", "P;4"), + "X 24 image": ("RGB", "RGB"), + "L 32 S image": ("I", "I;32"), + "L 32 F image": ("F", "F;32"), + # old p3cfunc formats + "RGB3 image": ("RGB", "RGB;T"), + "RYB3 image": ("RGB", "RYB;T"), + # extensions + "LA image": ("LA", "LA;L"), + "PA image": ("LA", "PA;L"), + "RGBA image": ("RGBA", "RGBA;L"), + "RGBX image": ("RGB", "RGBX;L"), + "CMYK image": ("CMYK", "CMYK;L"), + "YCC image": ("YCbCr", "YCbCr;L"), +} + +# ifunc95 extensions +for i in ["8", "8S", "16", "16S", "32", "32F"]: + OPEN[f"L {i} image"] = ("F", f"F;{i}") + OPEN[f"L*{i} image"] = ("F", f"F;{i}") +for i in ["16", "16L", "16B"]: + OPEN[f"L {i} image"] = (f"I;{i}", f"I;{i}") + OPEN[f"L*{i} image"] = (f"I;{i}", f"I;{i}") +for i in ["32S"]: + OPEN[f"L {i} image"] = ("I", f"I;{i}") + OPEN[f"L*{i} image"] = ("I", f"I;{i}") +for j in range(2, 33): + OPEN[f"L*{j} image"] = ("F", f"F;{j}") + + +# -------------------------------------------------------------------- +# Read IM directory + +split = re.compile(rb"^([A-Za-z][^:]*):[ \t]*(.*)[ \t]*$") + + +def number(s: Any) -> float: + try: + return int(s) + except ValueError: + return float(s) + + +## +# Image plugin for the IFUNC IM file format. + + +class ImImageFile(ImageFile.ImageFile): + format = "IM" + format_description = "IFUNC Image Memory" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + # Quick rejection: if there's not an LF among the first + # 100 bytes, this is (probably) not a text header. + + if b"\n" not in self.fp.read(100): + msg = "not an IM file" + raise SyntaxError(msg) + self.fp.seek(0) + + n = 0 + + # Default values + self.info[MODE] = "L" + self.info[SIZE] = (512, 512) + self.info[FRAMES] = 1 + + self.rawmode = "L" + + while True: + s = self.fp.read(1) + + # Some versions of IFUNC uses \n\r instead of \r\n... + if s == b"\r": + continue + + if not s or s == b"\0" or s == b"\x1a": + break + + # FIXME: this may read whole file if not a text file + s = s + self.fp.readline() + + if len(s) > 100: + msg = "not an IM file" + raise SyntaxError(msg) + + if s.endswith(b"\r\n"): + s = s[:-2] + elif s.endswith(b"\n"): + s = s[:-1] + + try: + m = split.match(s) + except re.error as e: + msg = "not an IM file" + raise SyntaxError(msg) from e + + if m: + k, v = m.group(1, 2) + + # Don't know if this is the correct encoding, + # but a decent guess (I guess) + k = k.decode("latin-1", "replace") + v = v.decode("latin-1", "replace") + + # Convert value as appropriate + if k in [FRAMES, SCALE, SIZE]: + v = v.replace("*", ",") + v = tuple(map(number, v.split(","))) + if len(v) == 1: + v = v[0] + elif k == MODE and v in OPEN: + v, self.rawmode = OPEN[v] + + # Add to dictionary. Note that COMMENT tags are + # combined into a list of strings. + if k == COMMENT: + if k in self.info: + self.info[k].append(v) + else: + self.info[k] = [v] + else: + self.info[k] = v + + if k in TAGS: + n += 1 + + else: + msg = f"Syntax error in IM header: {s.decode('ascii', 'replace')}" + raise SyntaxError(msg) + + if not n: + msg = "Not an IM file" + raise SyntaxError(msg) + + # Basic attributes + self._size = self.info[SIZE] + self._mode = self.info[MODE] + + # Skip forward to start of image data + while s and not s.startswith(b"\x1a"): + s = self.fp.read(1) + if not s: + msg = "File truncated" + raise SyntaxError(msg) + + if LUT in self.info: + # convert lookup table to palette or lut attribute + palette = self.fp.read(768) + greyscale = 1 # greyscale palette + linear = 1 # linear greyscale palette + for i in range(256): + if palette[i] == palette[i + 256] == palette[i + 512]: + if palette[i] != i: + linear = 0 + else: + greyscale = 0 + if self.mode in ["L", "LA", "P", "PA"]: + if greyscale: + if not linear: + self.lut = list(palette[:256]) + else: + if self.mode in ["L", "P"]: + self._mode = self.rawmode = "P" + elif self.mode in ["LA", "PA"]: + self._mode = "PA" + self.rawmode = "PA;L" + self.palette = ImagePalette.raw("RGB;L", palette) + elif self.mode == "RGB": + if not greyscale or not linear: + self.lut = list(palette) + + self.frame = 0 + + self.__offset = offs = self.fp.tell() + + self._fp = self.fp # FIXME: hack + + if self.rawmode.startswith("F;"): + # ifunc95 formats + try: + # use bit decoder (if necessary) + bits = int(self.rawmode[2:]) + if bits not in [8, 16, 32]: + self.tile = [ + ImageFile._Tile( + "bit", (0, 0) + self.size, offs, (bits, 8, 3, 0, -1) + ) + ] + return + except ValueError: + pass + + if self.rawmode in ["RGB;T", "RYB;T"]: + # Old LabEye/3PC files. Would be very surprised if anyone + # ever stumbled upon such a file ;-) + size = self.size[0] * self.size[1] + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, offs, ("G", 0, -1)), + ImageFile._Tile("raw", (0, 0) + self.size, offs + size, ("R", 0, -1)), + ImageFile._Tile( + "raw", (0, 0) + self.size, offs + 2 * size, ("B", 0, -1) + ), + ] + else: + # LabEye/IFUNC files + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, offs, (self.rawmode, 0, -1)) + ] + + @property + def n_frames(self) -> int: + return self.info[FRAMES] + + @property + def is_animated(self) -> bool: + return self.info[FRAMES] > 1 + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if isinstance(self._fp, DeferredError): + raise self._fp.ex + + self.frame = frame + + if self.mode == "1": + bits = 1 + else: + bits = 8 * len(self.mode) + + size = ((self.size[0] * bits + 7) // 8) * self.size[1] + offs = self.__offset + frame * size + + self.fp = self._fp + + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, offs, (self.rawmode, 0, -1)) + ] + + def tell(self) -> int: + return self.frame + + +# +# -------------------------------------------------------------------- +# Save IM files + + +SAVE = { + # mode: (im type, raw mode) + "1": ("0 1", "1"), + "L": ("Greyscale", "L"), + "LA": ("LA", "LA;L"), + "P": ("Greyscale", "P"), + "PA": ("LA", "PA;L"), + "I": ("L 32S", "I;32S"), + "I;16": ("L 16", "I;16"), + "I;16L": ("L 16L", "I;16L"), + "I;16B": ("L 16B", "I;16B"), + "F": ("L 32F", "F;32F"), + "RGB": ("RGB", "RGB;L"), + "RGBA": ("RGBA", "RGBA;L"), + "RGBX": ("RGBX", "RGBX;L"), + "CMYK": ("CMYK", "CMYK;L"), + "YCbCr": ("YCC", "YCbCr;L"), +} + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + try: + image_type, rawmode = SAVE[im.mode] + except KeyError as e: + msg = f"Cannot save {im.mode} images as IM" + raise ValueError(msg) from e + + frames = im.encoderinfo.get("frames", 1) + + fp.write(f"Image type: {image_type} image\r\n".encode("ascii")) + if filename: + # Each line must be 100 characters or less, + # or: SyntaxError("not an IM file") + # 8 characters are used for "Name: " and "\r\n" + # Keep just the filename, ditch the potentially overlong path + if isinstance(filename, bytes): + filename = filename.decode("ascii") + name, ext = os.path.splitext(os.path.basename(filename)) + name = "".join([name[: 92 - len(ext)], ext]) + + fp.write(f"Name: {name}\r\n".encode("ascii")) + fp.write(f"Image size (x*y): {im.size[0]}*{im.size[1]}\r\n".encode("ascii")) + fp.write(f"File size (no of images): {frames}\r\n".encode("ascii")) + if im.mode in ["P", "PA"]: + fp.write(b"Lut: 1\r\n") + fp.write(b"\000" * (511 - fp.tell()) + b"\032") + if im.mode in ["P", "PA"]: + im_palette = im.im.getpalette("RGB", "RGB;L") + colors = len(im_palette) // 3 + palette = b"" + for i in range(3): + palette += im_palette[colors * i : colors * (i + 1)] + palette += b"\x00" * (256 - colors) + fp.write(palette) # 768 bytes + ImageFile._save( + im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 0, (rawmode, 0, -1))] + ) + + +# +# -------------------------------------------------------------------- +# Registry + + +Image.register_open(ImImageFile.format, ImImageFile) +Image.register_save(ImImageFile.format, _save) + +Image.register_extension(ImImageFile.format, ".im") diff --git a/venv/lib/python3.10/site-packages/PIL/Image.py b/venv/lib/python3.10/site-packages/PIL/Image.py new file mode 100644 index 0000000000000000000000000000000000000000..9d50812eb3e0f9c1dbf6541cf6b8ffcf8797dcdd --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/Image.py @@ -0,0 +1,4227 @@ +# +# The Python Imaging Library. +# $Id$ +# +# the Image class wrapper +# +# partial release history: +# 1995-09-09 fl Created +# 1996-03-11 fl PIL release 0.0 (proof of concept) +# 1996-04-30 fl PIL release 0.1b1 +# 1999-07-28 fl PIL release 1.0 final +# 2000-06-07 fl PIL release 1.1 +# 2000-10-20 fl PIL release 1.1.1 +# 2001-05-07 fl PIL release 1.1.2 +# 2002-03-15 fl PIL release 1.1.3 +# 2003-05-10 fl PIL release 1.1.4 +# 2005-03-28 fl PIL release 1.1.5 +# 2006-12-02 fl PIL release 1.1.6 +# 2009-11-15 fl PIL release 1.1.7 +# +# Copyright (c) 1997-2009 by Secret Labs AB. All rights reserved. +# Copyright (c) 1995-2009 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + +from __future__ import annotations + +import abc +import atexit +import builtins +import io +import logging +import math +import os +import re +import struct +import sys +import tempfile +import warnings +from collections.abc import MutableMapping +from enum import IntEnum +from typing import IO, Protocol, cast + +# VERSION was removed in Pillow 6.0.0. +# PILLOW_VERSION was removed in Pillow 9.0.0. +# Use __version__ instead. +from . import ( + ExifTags, + ImageMode, + TiffTags, + UnidentifiedImageError, + __version__, + _plugins, +) +from ._binary import i32le, o32be, o32le +from ._deprecate import deprecate +from ._util import DeferredError, is_path + +ElementTree: ModuleType | None +try: + from defusedxml import ElementTree +except ImportError: + ElementTree = None + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + from types import ModuleType + from typing import Any, Literal + +logger = logging.getLogger(__name__) + + +class DecompressionBombWarning(RuntimeWarning): + pass + + +class DecompressionBombError(Exception): + pass + + +WARN_POSSIBLE_FORMATS: bool = False + +# Limit to around a quarter gigabyte for a 24-bit (3 bpp) image +MAX_IMAGE_PIXELS: int | None = int(1024 * 1024 * 1024 // 4 // 3) + + +try: + # If the _imaging C module is not present, Pillow will not load. + # Note that other modules should not refer to _imaging directly; + # import Image and use the Image.core variable instead. + # Also note that Image.core is not a publicly documented interface, + # and should be considered private and subject to change. + from . import _imaging as core + + if __version__ != getattr(core, "PILLOW_VERSION", None): + msg = ( + "The _imaging extension was built for another version of Pillow or PIL:\n" + f"Core version: {getattr(core, 'PILLOW_VERSION', None)}\n" + f"Pillow version: {__version__}" + ) + raise ImportError(msg) + +except ImportError as v: + # Explanations for ways that we know we might have an import error + if str(v).startswith("Module use of python"): + # The _imaging C module is present, but not compiled for + # the right version (windows only). Print a warning, if + # possible. + warnings.warn( + "The _imaging extension was built for another version of Python.", + RuntimeWarning, + ) + elif str(v).startswith("The _imaging extension"): + warnings.warn(str(v), RuntimeWarning) + # Fail here anyway. Don't let people run with a mostly broken Pillow. + # see docs/porting.rst + raise + + +# +# Constants + + +# transpose +class Transpose(IntEnum): + FLIP_LEFT_RIGHT = 0 + FLIP_TOP_BOTTOM = 1 + ROTATE_90 = 2 + ROTATE_180 = 3 + ROTATE_270 = 4 + TRANSPOSE = 5 + TRANSVERSE = 6 + + +# transforms (also defined in Imaging.h) +class Transform(IntEnum): + AFFINE = 0 + EXTENT = 1 + PERSPECTIVE = 2 + QUAD = 3 + MESH = 4 + + +# resampling filters (also defined in Imaging.h) +class Resampling(IntEnum): + NEAREST = 0 + BOX = 4 + BILINEAR = 2 + HAMMING = 5 + BICUBIC = 3 + LANCZOS = 1 + + +_filters_support = { + Resampling.BOX: 0.5, + Resampling.BILINEAR: 1.0, + Resampling.HAMMING: 1.0, + Resampling.BICUBIC: 2.0, + Resampling.LANCZOS: 3.0, +} + + +# dithers +class Dither(IntEnum): + NONE = 0 + ORDERED = 1 # Not yet implemented + RASTERIZE = 2 # Not yet implemented + FLOYDSTEINBERG = 3 # default + + +# palettes/quantizers +class Palette(IntEnum): + WEB = 0 + ADAPTIVE = 1 + + +class Quantize(IntEnum): + MEDIANCUT = 0 + MAXCOVERAGE = 1 + FASTOCTREE = 2 + LIBIMAGEQUANT = 3 + + +module = sys.modules[__name__] +for enum in (Transpose, Transform, Resampling, Dither, Palette, Quantize): + for item in enum: + setattr(module, item.name, item.value) + + +if hasattr(core, "DEFAULT_STRATEGY"): + DEFAULT_STRATEGY = core.DEFAULT_STRATEGY + FILTERED = core.FILTERED + HUFFMAN_ONLY = core.HUFFMAN_ONLY + RLE = core.RLE + FIXED = core.FIXED + + +# -------------------------------------------------------------------- +# Registries + +TYPE_CHECKING = False +if TYPE_CHECKING: + import mmap + from xml.etree.ElementTree import Element + + from IPython.lib.pretty import PrettyPrinter + + from . import ImageFile, ImageFilter, ImagePalette, ImageQt, TiffImagePlugin + from ._typing import CapsuleType, NumpyArray, StrOrBytesPath +ID: list[str] = [] +OPEN: dict[ + str, + tuple[ + Callable[[IO[bytes], str | bytes], ImageFile.ImageFile], + Callable[[bytes], bool | str] | None, + ], +] = {} +MIME: dict[str, str] = {} +SAVE: dict[str, Callable[[Image, IO[bytes], str | bytes], None]] = {} +SAVE_ALL: dict[str, Callable[[Image, IO[bytes], str | bytes], None]] = {} +EXTENSION: dict[str, str] = {} +DECODERS: dict[str, type[ImageFile.PyDecoder]] = {} +ENCODERS: dict[str, type[ImageFile.PyEncoder]] = {} + +# -------------------------------------------------------------------- +# Modes + +_ENDIAN = "<" if sys.byteorder == "little" else ">" + + +def _conv_type_shape(im: Image) -> tuple[tuple[int, ...], str]: + m = ImageMode.getmode(im.mode) + shape: tuple[int, ...] = (im.height, im.width) + extra = len(m.bands) + if extra != 1: + shape += (extra,) + return shape, m.typestr + + +MODES = [ + "1", + "CMYK", + "F", + "HSV", + "I", + "I;16", + "I;16B", + "I;16L", + "I;16N", + "L", + "LA", + "La", + "LAB", + "P", + "PA", + "RGB", + "RGBA", + "RGBa", + "RGBX", + "YCbCr", +] + +# raw modes that may be memory mapped. NOTE: if you change this, you +# may have to modify the stride calculation in map.c too! +_MAPMODES = ("L", "P", "RGBX", "RGBA", "CMYK", "I;16", "I;16L", "I;16B") + + +def getmodebase(mode: str) -> str: + """ + Gets the "base" mode for given mode. This function returns "L" for + images that contain grayscale data, and "RGB" for images that + contain color data. + + :param mode: Input mode. + :returns: "L" or "RGB". + :exception KeyError: If the input mode was not a standard mode. + """ + return ImageMode.getmode(mode).basemode + + +def getmodetype(mode: str) -> str: + """ + Gets the storage type mode. Given a mode, this function returns a + single-layer mode suitable for storing individual bands. + + :param mode: Input mode. + :returns: "L", "I", or "F". + :exception KeyError: If the input mode was not a standard mode. + """ + return ImageMode.getmode(mode).basetype + + +def getmodebandnames(mode: str) -> tuple[str, ...]: + """ + Gets a list of individual band names. Given a mode, this function returns + a tuple containing the names of individual bands (use + :py:method:`~PIL.Image.getmodetype` to get the mode used to store each + individual band. + + :param mode: Input mode. + :returns: A tuple containing band names. The length of the tuple + gives the number of bands in an image of the given mode. + :exception KeyError: If the input mode was not a standard mode. + """ + return ImageMode.getmode(mode).bands + + +def getmodebands(mode: str) -> int: + """ + Gets the number of individual bands for this mode. + + :param mode: Input mode. + :returns: The number of bands in this mode. + :exception KeyError: If the input mode was not a standard mode. + """ + return len(ImageMode.getmode(mode).bands) + + +# -------------------------------------------------------------------- +# Helpers + +_initialized = 0 + + +def preinit() -> None: + """ + Explicitly loads BMP, GIF, JPEG, PPM and PPM file format drivers. + + It is called when opening or saving images. + """ + + global _initialized + if _initialized >= 1: + return + + try: + from . import BmpImagePlugin + + assert BmpImagePlugin + except ImportError: + pass + try: + from . import GifImagePlugin + + assert GifImagePlugin + except ImportError: + pass + try: + from . import JpegImagePlugin + + assert JpegImagePlugin + except ImportError: + pass + try: + from . import PpmImagePlugin + + assert PpmImagePlugin + except ImportError: + pass + try: + from . import PngImagePlugin + + assert PngImagePlugin + except ImportError: + pass + + _initialized = 1 + + +def init() -> bool: + """ + Explicitly initializes the Python Imaging Library. This function + loads all available file format drivers. + + It is called when opening or saving images if :py:meth:`~preinit()` is + insufficient, and by :py:meth:`~PIL.features.pilinfo`. + """ + + global _initialized + if _initialized >= 2: + return False + + parent_name = __name__.rpartition(".")[0] + for plugin in _plugins: + try: + logger.debug("Importing %s", plugin) + __import__(f"{parent_name}.{plugin}", globals(), locals(), []) + except ImportError as e: + logger.debug("Image: failed to import %s: %s", plugin, e) + + if OPEN or SAVE: + _initialized = 2 + return True + return False + + +# -------------------------------------------------------------------- +# Codec factories (used by tobytes/frombytes and ImageFile.load) + + +def _getdecoder( + mode: str, decoder_name: str, args: Any, extra: tuple[Any, ...] = () +) -> core.ImagingDecoder | ImageFile.PyDecoder: + # tweak arguments + if args is None: + args = () + elif not isinstance(args, tuple): + args = (args,) + + try: + decoder = DECODERS[decoder_name] + except KeyError: + pass + else: + return decoder(mode, *args + extra) + + try: + # get decoder + decoder = getattr(core, f"{decoder_name}_decoder") + except AttributeError as e: + msg = f"decoder {decoder_name} not available" + raise OSError(msg) from e + return decoder(mode, *args + extra) + + +def _getencoder( + mode: str, encoder_name: str, args: Any, extra: tuple[Any, ...] = () +) -> core.ImagingEncoder | ImageFile.PyEncoder: + # tweak arguments + if args is None: + args = () + elif not isinstance(args, tuple): + args = (args,) + + try: + encoder = ENCODERS[encoder_name] + except KeyError: + pass + else: + return encoder(mode, *args + extra) + + try: + # get encoder + encoder = getattr(core, f"{encoder_name}_encoder") + except AttributeError as e: + msg = f"encoder {encoder_name} not available" + raise OSError(msg) from e + return encoder(mode, *args + extra) + + +# -------------------------------------------------------------------- +# Simple expression analyzer + + +class ImagePointTransform: + """ + Used with :py:meth:`~PIL.Image.Image.point` for single band images with more than + 8 bits, this represents an affine transformation, where the value is multiplied by + ``scale`` and ``offset`` is added. + """ + + def __init__(self, scale: float, offset: float) -> None: + self.scale = scale + self.offset = offset + + def __neg__(self) -> ImagePointTransform: + return ImagePointTransform(-self.scale, -self.offset) + + def __add__(self, other: ImagePointTransform | float) -> ImagePointTransform: + if isinstance(other, ImagePointTransform): + return ImagePointTransform( + self.scale + other.scale, self.offset + other.offset + ) + return ImagePointTransform(self.scale, self.offset + other) + + __radd__ = __add__ + + def __sub__(self, other: ImagePointTransform | float) -> ImagePointTransform: + return self + -other + + def __rsub__(self, other: ImagePointTransform | float) -> ImagePointTransform: + return other + -self + + def __mul__(self, other: ImagePointTransform | float) -> ImagePointTransform: + if isinstance(other, ImagePointTransform): + return NotImplemented + return ImagePointTransform(self.scale * other, self.offset * other) + + __rmul__ = __mul__ + + def __truediv__(self, other: ImagePointTransform | float) -> ImagePointTransform: + if isinstance(other, ImagePointTransform): + return NotImplemented + return ImagePointTransform(self.scale / other, self.offset / other) + + +def _getscaleoffset( + expr: Callable[[ImagePointTransform], ImagePointTransform | float], +) -> tuple[float, float]: + a = expr(ImagePointTransform(1, 0)) + return (a.scale, a.offset) if isinstance(a, ImagePointTransform) else (0, a) + + +# -------------------------------------------------------------------- +# Implementation wrapper + + +class SupportsGetData(Protocol): + def getdata( + self, + ) -> tuple[Transform, Sequence[int]]: ... + + +class Image: + """ + This class represents an image object. To create + :py:class:`~PIL.Image.Image` objects, use the appropriate factory + functions. There's hardly ever any reason to call the Image constructor + directly. + + * :py:func:`~PIL.Image.open` + * :py:func:`~PIL.Image.new` + * :py:func:`~PIL.Image.frombytes` + """ + + format: str | None = None + format_description: str | None = None + _close_exclusive_fp_after_loading = True + + def __init__(self) -> None: + # FIXME: take "new" parameters / other image? + self._im: core.ImagingCore | DeferredError | None = None + self._mode = "" + self._size = (0, 0) + self.palette: ImagePalette.ImagePalette | None = None + self.info: dict[str | tuple[int, int], Any] = {} + self.readonly = 0 + self._exif: Exif | None = None + + @property + def im(self) -> core.ImagingCore: + if isinstance(self._im, DeferredError): + raise self._im.ex + assert self._im is not None + return self._im + + @im.setter + def im(self, im: core.ImagingCore) -> None: + self._im = im + + @property + def width(self) -> int: + return self.size[0] + + @property + def height(self) -> int: + return self.size[1] + + @property + def size(self) -> tuple[int, int]: + return self._size + + @property + def mode(self) -> str: + return self._mode + + @property + def readonly(self) -> int: + return (self._im and self._im.readonly) or self._readonly + + @readonly.setter + def readonly(self, readonly: int) -> None: + self._readonly = readonly + + def _new(self, im: core.ImagingCore) -> Image: + new = Image() + new.im = im + new._mode = im.mode + new._size = im.size + if im.mode in ("P", "PA"): + if self.palette: + new.palette = self.palette.copy() + else: + from . import ImagePalette + + new.palette = ImagePalette.ImagePalette() + new.info = self.info.copy() + return new + + # Context manager support + def __enter__(self): + return self + + def __exit__(self, *args): + from . import ImageFile + + if isinstance(self, ImageFile.ImageFile): + if getattr(self, "_exclusive_fp", False): + self._close_fp() + self.fp = None + + def close(self) -> None: + """ + This operation will destroy the image core and release its memory. + The image data will be unusable afterward. + + This function is required to close images that have multiple frames or + have not had their file read and closed by the + :py:meth:`~PIL.Image.Image.load` method. See :ref:`file-handling` for + more information. + """ + if getattr(self, "map", None): + if sys.platform == "win32" and hasattr(sys, "pypy_version_info"): + self.map.close() + self.map: mmap.mmap | None = None + + # Instead of simply setting to None, we're setting up a + # deferred error that will better explain that the core image + # object is gone. + self._im = DeferredError(ValueError("Operation on closed image")) + + def _copy(self) -> None: + self.load() + self.im = self.im.copy() + self.readonly = 0 + + def _ensure_mutable(self) -> None: + if self.readonly: + self._copy() + else: + self.load() + + def _dump( + self, file: str | None = None, format: str | None = None, **options: Any + ) -> str: + suffix = "" + if format: + suffix = f".{format}" + + if not file: + f, filename = tempfile.mkstemp(suffix) + os.close(f) + else: + filename = file + if not filename.endswith(suffix): + filename = filename + suffix + + self.load() + + if not format or format == "PPM": + self.im.save_ppm(filename) + else: + self.save(filename, format, **options) + + return filename + + def __eq__(self, other: object) -> bool: + if self.__class__ is not other.__class__: + return False + assert isinstance(other, Image) + return ( + self.mode == other.mode + and self.size == other.size + and self.info == other.info + and self.getpalette() == other.getpalette() + and self.tobytes() == other.tobytes() + ) + + def __repr__(self) -> str: + return ( + f"<{self.__class__.__module__}.{self.__class__.__name__} " + f"image mode={self.mode} size={self.size[0]}x{self.size[1]} " + f"at 0x{id(self):X}>" + ) + + def _repr_pretty_(self, p: PrettyPrinter, cycle: bool) -> None: + """IPython plain text display support""" + + # Same as __repr__ but without unpredictable id(self), + # to keep Jupyter notebook `text/plain` output stable. + p.text( + f"<{self.__class__.__module__}.{self.__class__.__name__} " + f"image mode={self.mode} size={self.size[0]}x{self.size[1]}>" + ) + + def _repr_image(self, image_format: str, **kwargs: Any) -> bytes | None: + """Helper function for iPython display hook. + + :param image_format: Image format. + :returns: image as bytes, saved into the given format. + """ + b = io.BytesIO() + try: + self.save(b, image_format, **kwargs) + except Exception: + return None + return b.getvalue() + + def _repr_png_(self) -> bytes | None: + """iPython display hook support for PNG format. + + :returns: PNG version of the image as bytes + """ + return self._repr_image("PNG", compress_level=1) + + def _repr_jpeg_(self) -> bytes | None: + """iPython display hook support for JPEG format. + + :returns: JPEG version of the image as bytes + """ + return self._repr_image("JPEG") + + @property + def __array_interface__(self) -> dict[str, str | bytes | int | tuple[int, ...]]: + # numpy array interface support + new: dict[str, str | bytes | int | tuple[int, ...]] = {"version": 3} + if self.mode == "1": + # Binary images need to be extended from bits to bytes + # See: https://github.com/python-pillow/Pillow/issues/350 + new["data"] = self.tobytes("raw", "L") + else: + new["data"] = self.tobytes() + new["shape"], new["typestr"] = _conv_type_shape(self) + return new + + def __arrow_c_schema__(self) -> object: + self.load() + return self.im.__arrow_c_schema__() + + def __arrow_c_array__( + self, requested_schema: object | None = None + ) -> tuple[object, object]: + self.load() + return (self.im.__arrow_c_schema__(), self.im.__arrow_c_array__()) + + def __getstate__(self) -> list[Any]: + im_data = self.tobytes() # load image first + return [self.info, self.mode, self.size, self.getpalette(), im_data] + + def __setstate__(self, state: list[Any]) -> None: + Image.__init__(self) + info, mode, size, palette, data = state[:5] + self.info = info + self._mode = mode + self._size = size + self.im = core.new(mode, size) + if mode in ("L", "LA", "P", "PA") and palette: + self.putpalette(palette) + self.frombytes(data) + + def tobytes(self, encoder_name: str = "raw", *args: Any) -> bytes: + """ + Return image as a bytes object. + + .. warning:: + + This method returns raw image data derived from Pillow's internal + storage. For compressed image data (e.g. PNG, JPEG) use + :meth:`~.save`, with a BytesIO parameter for in-memory data. + + :param encoder_name: What encoder to use. + + The default is to use the standard "raw" encoder. + To see how this packs pixel data into the returned + bytes, see :file:`libImaging/Pack.c`. + + A list of C encoders can be seen under codecs + section of the function array in + :file:`_imaging.c`. Python encoders are registered + within the relevant plugins. + :param args: Extra arguments to the encoder. + :returns: A :py:class:`bytes` object. + """ + + encoder_args: Any = args + if len(encoder_args) == 1 and isinstance(encoder_args[0], tuple): + # may pass tuple instead of argument list + encoder_args = encoder_args[0] + + if encoder_name == "raw" and encoder_args == (): + encoder_args = self.mode + + self.load() + + if self.width == 0 or self.height == 0: + return b"" + + # unpack data + e = _getencoder(self.mode, encoder_name, encoder_args) + e.setimage(self.im) + + from . import ImageFile + + bufsize = max(ImageFile.MAXBLOCK, self.size[0] * 4) # see RawEncode.c + + output = [] + while True: + bytes_consumed, errcode, data = e.encode(bufsize) + output.append(data) + if errcode: + break + if errcode < 0: + msg = f"encoder error {errcode} in tobytes" + raise RuntimeError(msg) + + return b"".join(output) + + def tobitmap(self, name: str = "image") -> bytes: + """ + Returns the image converted to an X11 bitmap. + + .. note:: This method only works for mode "1" images. + + :param name: The name prefix to use for the bitmap variables. + :returns: A string containing an X11 bitmap. + :raises ValueError: If the mode is not "1" + """ + + self.load() + if self.mode != "1": + msg = "not a bitmap" + raise ValueError(msg) + data = self.tobytes("xbm") + return b"".join( + [ + f"#define {name}_width {self.size[0]}\n".encode("ascii"), + f"#define {name}_height {self.size[1]}\n".encode("ascii"), + f"static char {name}_bits[] = {{\n".encode("ascii"), + data, + b"};", + ] + ) + + def frombytes( + self, + data: bytes | bytearray | SupportsArrayInterface, + decoder_name: str = "raw", + *args: Any, + ) -> None: + """ + Loads this image with pixel data from a bytes object. + + This method is similar to the :py:func:`~PIL.Image.frombytes` function, + but loads data into this image instead of creating a new image object. + """ + + if self.width == 0 or self.height == 0: + return + + decoder_args: Any = args + if len(decoder_args) == 1 and isinstance(decoder_args[0], tuple): + # may pass tuple instead of argument list + decoder_args = decoder_args[0] + + # default format + if decoder_name == "raw" and decoder_args == (): + decoder_args = self.mode + + # unpack data + d = _getdecoder(self.mode, decoder_name, decoder_args) + d.setimage(self.im) + s = d.decode(data) + + if s[0] >= 0: + msg = "not enough image data" + raise ValueError(msg) + if s[1] != 0: + msg = "cannot decode image data" + raise ValueError(msg) + + def load(self) -> core.PixelAccess | None: + """ + Allocates storage for the image and loads the pixel data. In + normal cases, you don't need to call this method, since the + Image class automatically loads an opened image when it is + accessed for the first time. + + If the file associated with the image was opened by Pillow, then this + method will close it. The exception to this is if the image has + multiple frames, in which case the file will be left open for seek + operations. See :ref:`file-handling` for more information. + + :returns: An image access object. + :rtype: :py:class:`.PixelAccess` + """ + if self._im is not None and self.palette and self.palette.dirty: + # realize palette + mode, arr = self.palette.getdata() + self.im.putpalette(self.palette.mode, mode, arr) + self.palette.dirty = 0 + self.palette.rawmode = None + if "transparency" in self.info and mode in ("LA", "PA"): + if isinstance(self.info["transparency"], int): + self.im.putpalettealpha(self.info["transparency"], 0) + else: + self.im.putpalettealphas(self.info["transparency"]) + self.palette.mode = "RGBA" + else: + self.palette.palette = self.im.getpalette( + self.palette.mode, self.palette.mode + ) + + if self._im is not None: + return self.im.pixel_access(self.readonly) + return None + + def verify(self) -> None: + """ + Verifies the contents of a file. For data read from a file, this + method attempts to determine if the file is broken, without + actually decoding the image data. If this method finds any + problems, it raises suitable exceptions. If you need to load + the image after using this method, you must reopen the image + file. + """ + pass + + def convert( + self, + mode: str | None = None, + matrix: tuple[float, ...] | None = None, + dither: Dither | None = None, + palette: Palette = Palette.WEB, + colors: int = 256, + ) -> Image: + """ + Returns a converted copy of this image. For the "P" mode, this + method translates pixels through the palette. If mode is + omitted, a mode is chosen so that all information in the image + and the palette can be represented without a palette. + + This supports all possible conversions between "L", "RGB" and "CMYK". The + ``matrix`` argument only supports "L" and "RGB". + + When translating a color image to grayscale (mode "L"), + the library uses the ITU-R 601-2 luma transform:: + + L = R * 299/1000 + G * 587/1000 + B * 114/1000 + + The default method of converting a grayscale ("L") or "RGB" + image into a bilevel (mode "1") image uses Floyd-Steinberg + dither to approximate the original image luminosity levels. If + dither is ``None``, all values larger than 127 are set to 255 (white), + all other values to 0 (black). To use other thresholds, use the + :py:meth:`~PIL.Image.Image.point` method. + + When converting from "RGBA" to "P" without a ``matrix`` argument, + this passes the operation to :py:meth:`~PIL.Image.Image.quantize`, + and ``dither`` and ``palette`` are ignored. + + When converting from "PA", if an "RGBA" palette is present, the alpha + channel from the image will be used instead of the values from the palette. + + :param mode: The requested mode. See: :ref:`concept-modes`. + :param matrix: An optional conversion matrix. If given, this + should be 4- or 12-tuple containing floating point values. + :param dither: Dithering method, used when converting from + mode "RGB" to "P" or from "RGB" or "L" to "1". + Available methods are :data:`Dither.NONE` or :data:`Dither.FLOYDSTEINBERG` + (default). Note that this is not used when ``matrix`` is supplied. + :param palette: Palette to use when converting from mode "RGB" + to "P". Available palettes are :data:`Palette.WEB` or + :data:`Palette.ADAPTIVE`. + :param colors: Number of colors to use for the :data:`Palette.ADAPTIVE` + palette. Defaults to 256. + :rtype: :py:class:`~PIL.Image.Image` + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + self.load() + + has_transparency = "transparency" in self.info + if not mode and self.mode == "P": + # determine default mode + if self.palette: + mode = self.palette.mode + else: + mode = "RGB" + if mode == "RGB" and has_transparency: + mode = "RGBA" + if not mode or (mode == self.mode and not matrix): + return self.copy() + + if matrix: + # matrix conversion + if mode not in ("L", "RGB"): + msg = "illegal conversion" + raise ValueError(msg) + im = self.im.convert_matrix(mode, matrix) + new_im = self._new(im) + if has_transparency and self.im.bands == 3: + transparency = new_im.info["transparency"] + + def convert_transparency( + m: tuple[float, ...], v: tuple[int, int, int] + ) -> int: + value = m[0] * v[0] + m[1] * v[1] + m[2] * v[2] + m[3] * 0.5 + return max(0, min(255, int(value))) + + if mode == "L": + transparency = convert_transparency(matrix, transparency) + elif len(mode) == 3: + transparency = tuple( + convert_transparency(matrix[i * 4 : i * 4 + 4], transparency) + for i in range(len(transparency)) + ) + new_im.info["transparency"] = transparency + return new_im + + if self.mode == "RGBA": + if mode == "P": + return self.quantize(colors) + elif mode == "PA": + r, g, b, a = self.split() + rgb = merge("RGB", (r, g, b)) + p = rgb.quantize(colors) + return merge("PA", (p, a)) + + trns = None + delete_trns = False + # transparency handling + if has_transparency: + if (self.mode in ("1", "L", "I", "I;16") and mode in ("LA", "RGBA")) or ( + self.mode == "RGB" and mode in ("La", "LA", "RGBa", "RGBA") + ): + # Use transparent conversion to promote from transparent + # color to an alpha channel. + new_im = self._new( + self.im.convert_transparent(mode, self.info["transparency"]) + ) + del new_im.info["transparency"] + return new_im + elif self.mode in ("L", "RGB", "P") and mode in ("L", "RGB", "P"): + t = self.info["transparency"] + if isinstance(t, bytes): + # Dragons. This can't be represented by a single color + warnings.warn( + "Palette images with Transparency expressed in bytes should be " + "converted to RGBA images" + ) + delete_trns = True + else: + # get the new transparency color. + # use existing conversions + trns_im = new(self.mode, (1, 1)) + if self.mode == "P": + assert self.palette is not None + trns_im.putpalette(self.palette, self.palette.mode) + if isinstance(t, tuple): + err = "Couldn't allocate a palette color for transparency" + assert trns_im.palette is not None + try: + t = trns_im.palette.getcolor(t, self) + except ValueError as e: + if str(e) == "cannot allocate more than 256 colors": + # If all 256 colors are in use, + # then there is no need for transparency + t = None + else: + raise ValueError(err) from e + if t is None: + trns = None + else: + trns_im.putpixel((0, 0), t) + + if mode in ("L", "RGB"): + trns_im = trns_im.convert(mode) + else: + # can't just retrieve the palette number, got to do it + # after quantization. + trns_im = trns_im.convert("RGB") + trns = trns_im.getpixel((0, 0)) + + elif self.mode == "P" and mode in ("LA", "PA", "RGBA"): + t = self.info["transparency"] + delete_trns = True + + if isinstance(t, bytes): + self.im.putpalettealphas(t) + elif isinstance(t, int): + self.im.putpalettealpha(t, 0) + else: + msg = "Transparency for P mode should be bytes or int" + raise ValueError(msg) + + if mode == "P" and palette == Palette.ADAPTIVE: + im = self.im.quantize(colors) + new_im = self._new(im) + from . import ImagePalette + + new_im.palette = ImagePalette.ImagePalette( + "RGB", new_im.im.getpalette("RGB") + ) + if delete_trns: + # This could possibly happen if we requantize to fewer colors. + # The transparency would be totally off in that case. + del new_im.info["transparency"] + if trns is not None: + try: + new_im.info["transparency"] = new_im.palette.getcolor( + cast(tuple[int, ...], trns), # trns was converted to RGB + new_im, + ) + except Exception: + # if we can't make a transparent color, don't leave the old + # transparency hanging around to mess us up. + del new_im.info["transparency"] + warnings.warn("Couldn't allocate palette entry for transparency") + return new_im + + if "LAB" in (self.mode, mode): + im = self + if mode == "LAB": + if im.mode not in ("RGB", "RGBA", "RGBX"): + im = im.convert("RGBA") + other_mode = im.mode + else: + other_mode = mode + if other_mode in ("RGB", "RGBA", "RGBX"): + from . import ImageCms + + srgb = ImageCms.createProfile("sRGB") + lab = ImageCms.createProfile("LAB") + profiles = [lab, srgb] if im.mode == "LAB" else [srgb, lab] + transform = ImageCms.buildTransform( + profiles[0], profiles[1], im.mode, mode + ) + return transform.apply(im) + + # colorspace conversion + if dither is None: + dither = Dither.FLOYDSTEINBERG + + try: + im = self.im.convert(mode, dither) + except ValueError: + try: + # normalize source image and try again + modebase = getmodebase(self.mode) + if modebase == self.mode: + raise + im = self.im.convert(modebase) + im = im.convert(mode, dither) + except KeyError as e: + msg = "illegal conversion" + raise ValueError(msg) from e + + new_im = self._new(im) + if mode in ("P", "PA") and palette != Palette.ADAPTIVE: + from . import ImagePalette + + new_im.palette = ImagePalette.ImagePalette("RGB", im.getpalette("RGB")) + if delete_trns: + # crash fail if we leave a bytes transparency in an rgb/l mode. + del new_im.info["transparency"] + if trns is not None: + if new_im.mode == "P" and new_im.palette: + try: + new_im.info["transparency"] = new_im.palette.getcolor( + cast(tuple[int, ...], trns), new_im # trns was converted to RGB + ) + except ValueError as e: + del new_im.info["transparency"] + if str(e) != "cannot allocate more than 256 colors": + # If all 256 colors are in use, + # then there is no need for transparency + warnings.warn( + "Couldn't allocate palette entry for transparency" + ) + else: + new_im.info["transparency"] = trns + return new_im + + def quantize( + self, + colors: int = 256, + method: int | None = None, + kmeans: int = 0, + palette: Image | None = None, + dither: Dither = Dither.FLOYDSTEINBERG, + ) -> Image: + """ + Convert the image to 'P' mode with the specified number + of colors. + + :param colors: The desired number of colors, <= 256 + :param method: :data:`Quantize.MEDIANCUT` (median cut), + :data:`Quantize.MAXCOVERAGE` (maximum coverage), + :data:`Quantize.FASTOCTREE` (fast octree), + :data:`Quantize.LIBIMAGEQUANT` (libimagequant; check support + using :py:func:`PIL.features.check_feature` with + ``feature="libimagequant"``). + + By default, :data:`Quantize.MEDIANCUT` will be used. + + The exception to this is RGBA images. :data:`Quantize.MEDIANCUT` + and :data:`Quantize.MAXCOVERAGE` do not support RGBA images, so + :data:`Quantize.FASTOCTREE` is used by default instead. + :param kmeans: Integer greater than or equal to zero. + :param palette: Quantize to the palette of given + :py:class:`PIL.Image.Image`. + :param dither: Dithering method, used when converting from + mode "RGB" to "P" or from "RGB" or "L" to "1". + Available methods are :data:`Dither.NONE` or :data:`Dither.FLOYDSTEINBERG` + (default). + :returns: A new image + """ + + self.load() + + if method is None: + # defaults: + method = Quantize.MEDIANCUT + if self.mode == "RGBA": + method = Quantize.FASTOCTREE + + if self.mode == "RGBA" and method not in ( + Quantize.FASTOCTREE, + Quantize.LIBIMAGEQUANT, + ): + # Caller specified an invalid mode. + msg = ( + "Fast Octree (method == 2) and libimagequant (method == 3) " + "are the only valid methods for quantizing RGBA images" + ) + raise ValueError(msg) + + if palette: + # use palette from reference image + palette.load() + if palette.mode != "P": + msg = "bad mode for palette image" + raise ValueError(msg) + if self.mode not in {"RGB", "L"}: + msg = "only RGB or L mode images can be quantized to a palette" + raise ValueError(msg) + im = self.im.convert("P", dither, palette.im) + new_im = self._new(im) + assert palette.palette is not None + new_im.palette = palette.palette.copy() + return new_im + + if kmeans < 0: + msg = "kmeans must not be negative" + raise ValueError(msg) + + im = self._new(self.im.quantize(colors, method, kmeans)) + + from . import ImagePalette + + mode = im.im.getpalettemode() + palette_data = im.im.getpalette(mode, mode)[: colors * len(mode)] + im.palette = ImagePalette.ImagePalette(mode, palette_data) + + return im + + def copy(self) -> Image: + """ + Copies this image. Use this method if you wish to paste things + into an image, but still retain the original. + + :rtype: :py:class:`~PIL.Image.Image` + :returns: An :py:class:`~PIL.Image.Image` object. + """ + self.load() + return self._new(self.im.copy()) + + __copy__ = copy + + def crop(self, box: tuple[float, float, float, float] | None = None) -> Image: + """ + Returns a rectangular region from this image. The box is a + 4-tuple defining the left, upper, right, and lower pixel + coordinate. See :ref:`coordinate-system`. + + Note: Prior to Pillow 3.4.0, this was a lazy operation. + + :param box: The crop rectangle, as a (left, upper, right, lower)-tuple. + :rtype: :py:class:`~PIL.Image.Image` + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + if box is None: + return self.copy() + + if box[2] < box[0]: + msg = "Coordinate 'right' is less than 'left'" + raise ValueError(msg) + elif box[3] < box[1]: + msg = "Coordinate 'lower' is less than 'upper'" + raise ValueError(msg) + + self.load() + return self._new(self._crop(self.im, box)) + + def _crop( + self, im: core.ImagingCore, box: tuple[float, float, float, float] + ) -> core.ImagingCore: + """ + Returns a rectangular region from the core image object im. + + This is equivalent to calling im.crop((x0, y0, x1, y1)), but + includes additional sanity checks. + + :param im: a core image object + :param box: The crop rectangle, as a (left, upper, right, lower)-tuple. + :returns: A core image object. + """ + + x0, y0, x1, y1 = map(int, map(round, box)) + + absolute_values = (abs(x1 - x0), abs(y1 - y0)) + + _decompression_bomb_check(absolute_values) + + return im.crop((x0, y0, x1, y1)) + + def draft( + self, mode: str | None, size: tuple[int, int] | None + ) -> tuple[str, tuple[int, int, float, float]] | None: + """ + Configures the image file loader so it returns a version of the + image that as closely as possible matches the given mode and + size. For example, you can use this method to convert a color + JPEG to grayscale while loading it. + + If any changes are made, returns a tuple with the chosen ``mode`` and + ``box`` with coordinates of the original image within the altered one. + + Note that this method modifies the :py:class:`~PIL.Image.Image` object + in place. If the image has already been loaded, this method has no + effect. + + Note: This method is not implemented for most images. It is + currently implemented only for JPEG and MPO images. + + :param mode: The requested mode. + :param size: The requested size in pixels, as a 2-tuple: + (width, height). + """ + pass + + def filter(self, filter: ImageFilter.Filter | type[ImageFilter.Filter]) -> Image: + """ + Filters this image using the given filter. For a list of + available filters, see the :py:mod:`~PIL.ImageFilter` module. + + :param filter: Filter kernel. + :returns: An :py:class:`~PIL.Image.Image` object.""" + + from . import ImageFilter + + self.load() + + if callable(filter): + filter = filter() + if not hasattr(filter, "filter"): + msg = "filter argument should be ImageFilter.Filter instance or class" + raise TypeError(msg) + + multiband = isinstance(filter, ImageFilter.MultibandFilter) + if self.im.bands == 1 or multiband: + return self._new(filter.filter(self.im)) + + ims = [ + self._new(filter.filter(self.im.getband(c))) for c in range(self.im.bands) + ] + return merge(self.mode, ims) + + def getbands(self) -> tuple[str, ...]: + """ + Returns a tuple containing the name of each band in this image. + For example, ``getbands`` on an RGB image returns ("R", "G", "B"). + + :returns: A tuple containing band names. + :rtype: tuple + """ + return ImageMode.getmode(self.mode).bands + + def getbbox(self, *, alpha_only: bool = True) -> tuple[int, int, int, int] | None: + """ + Calculates the bounding box of the non-zero regions in the + image. + + :param alpha_only: Optional flag, defaulting to ``True``. + If ``True`` and the image has an alpha channel, trim transparent pixels. + Otherwise, trim pixels when all channels are zero. + Keyword-only argument. + :returns: The bounding box is returned as a 4-tuple defining the + left, upper, right, and lower pixel coordinate. See + :ref:`coordinate-system`. If the image is completely empty, this + method returns None. + + """ + + self.load() + return self.im.getbbox(alpha_only) + + def getcolors( + self, maxcolors: int = 256 + ) -> list[tuple[int, tuple[int, ...]]] | list[tuple[int, float]] | None: + """ + Returns a list of colors used in this image. + + The colors will be in the image's mode. For example, an RGB image will + return a tuple of (red, green, blue) color values, and a P image will + return the index of the color in the palette. + + :param maxcolors: Maximum number of colors. If this number is + exceeded, this method returns None. The default limit is + 256 colors. + :returns: An unsorted list of (count, pixel) values. + """ + + self.load() + if self.mode in ("1", "L", "P"): + h = self.im.histogram() + out: list[tuple[int, float]] = [(h[i], i) for i in range(256) if h[i]] + if len(out) > maxcolors: + return None + return out + return self.im.getcolors(maxcolors) + + def getdata(self, band: int | None = None) -> core.ImagingCore: + """ + Returns the contents of this image as a sequence object + containing pixel values. The sequence object is flattened, so + that values for line one follow directly after the values of + line zero, and so on. + + Note that the sequence object returned by this method is an + internal PIL data type, which only supports certain sequence + operations. To convert it to an ordinary sequence (e.g. for + printing), use ``list(im.getdata())``. + + :param band: What band to return. The default is to return + all bands. To return a single band, pass in the index + value (e.g. 0 to get the "R" band from an "RGB" image). + :returns: A sequence-like object. + """ + + self.load() + if band is not None: + return self.im.getband(band) + return self.im # could be abused + + def getextrema(self) -> tuple[float, float] | tuple[tuple[int, int], ...]: + """ + Gets the minimum and maximum pixel values for each band in + the image. + + :returns: For a single-band image, a 2-tuple containing the + minimum and maximum pixel value. For a multi-band image, + a tuple containing one 2-tuple for each band. + """ + + self.load() + if self.im.bands > 1: + return tuple(self.im.getband(i).getextrema() for i in range(self.im.bands)) + return self.im.getextrema() + + def getxmp(self) -> dict[str, Any]: + """ + Returns a dictionary containing the XMP tags. + Requires defusedxml to be installed. + + :returns: XMP tags in a dictionary. + """ + + def get_name(tag: str) -> str: + return re.sub("^{[^}]+}", "", tag) + + def get_value(element: Element) -> str | dict[str, Any] | None: + value: dict[str, Any] = {get_name(k): v for k, v in element.attrib.items()} + children = list(element) + if children: + for child in children: + name = get_name(child.tag) + child_value = get_value(child) + if name in value: + if not isinstance(value[name], list): + value[name] = [value[name]] + value[name].append(child_value) + else: + value[name] = child_value + elif value: + if element.text: + value["text"] = element.text + else: + return element.text + return value + + if ElementTree is None: + warnings.warn("XMP data cannot be read without defusedxml dependency") + return {} + if "xmp" not in self.info: + return {} + root = ElementTree.fromstring(self.info["xmp"].rstrip(b"\x00 ")) + return {get_name(root.tag): get_value(root)} + + def getexif(self) -> Exif: + """ + Gets EXIF data from the image. + + :returns: an :py:class:`~PIL.Image.Exif` object. + """ + if self._exif is None: + self._exif = Exif() + elif self._exif._loaded: + return self._exif + self._exif._loaded = True + + exif_info = self.info.get("exif") + if exif_info is None: + if "Raw profile type exif" in self.info: + exif_info = bytes.fromhex( + "".join(self.info["Raw profile type exif"].split("\n")[3:]) + ) + elif hasattr(self, "tag_v2"): + self._exif.bigtiff = self.tag_v2._bigtiff + self._exif.endian = self.tag_v2._endian + self._exif.load_from_fp(self.fp, self.tag_v2._offset) + if exif_info is not None: + self._exif.load(exif_info) + + # XMP tags + if ExifTags.Base.Orientation not in self._exif: + xmp_tags = self.info.get("XML:com.adobe.xmp") + pattern: str | bytes = r'tiff:Orientation(="|>)([0-9])' + if not xmp_tags and (xmp_tags := self.info.get("xmp")): + pattern = rb'tiff:Orientation(="|>)([0-9])' + if xmp_tags: + match = re.search(pattern, xmp_tags) + if match: + self._exif[ExifTags.Base.Orientation] = int(match[2]) + + return self._exif + + def _reload_exif(self) -> None: + if self._exif is None or not self._exif._loaded: + return + self._exif._loaded = False + self.getexif() + + def get_child_images(self) -> list[ImageFile.ImageFile]: + from . import ImageFile + + deprecate("Image.Image.get_child_images", 13) + return ImageFile.ImageFile.get_child_images(self) # type: ignore[arg-type] + + def getim(self) -> CapsuleType: + """ + Returns a capsule that points to the internal image memory. + + :returns: A capsule object. + """ + + self.load() + return self.im.ptr + + def getpalette(self, rawmode: str | None = "RGB") -> list[int] | None: + """ + Returns the image palette as a list. + + :param rawmode: The mode in which to return the palette. ``None`` will + return the palette in its current mode. + + .. versionadded:: 9.1.0 + + :returns: A list of color values [r, g, b, ...], or None if the + image has no palette. + """ + + self.load() + try: + mode = self.im.getpalettemode() + except ValueError: + return None # no palette + if rawmode is None: + rawmode = mode + return list(self.im.getpalette(mode, rawmode)) + + @property + def has_transparency_data(self) -> bool: + """ + Determine if an image has transparency data, whether in the form of an + alpha channel, a palette with an alpha channel, or a "transparency" key + in the info dictionary. + + Note the image might still appear solid, if all of the values shown + within are opaque. + + :returns: A boolean. + """ + if ( + self.mode in ("LA", "La", "PA", "RGBA", "RGBa") + or "transparency" in self.info + ): + return True + if self.mode == "P": + assert self.palette is not None + return self.palette.mode.endswith("A") + return False + + def apply_transparency(self) -> None: + """ + If a P mode image has a "transparency" key in the info dictionary, + remove the key and instead apply the transparency to the palette. + Otherwise, the image is unchanged. + """ + if self.mode != "P" or "transparency" not in self.info: + return + + from . import ImagePalette + + palette = self.getpalette("RGBA") + assert palette is not None + transparency = self.info["transparency"] + if isinstance(transparency, bytes): + for i, alpha in enumerate(transparency): + palette[i * 4 + 3] = alpha + else: + palette[transparency * 4 + 3] = 0 + self.palette = ImagePalette.ImagePalette("RGBA", bytes(palette)) + self.palette.dirty = 1 + + del self.info["transparency"] + + def getpixel( + self, xy: tuple[int, int] | list[int] + ) -> float | tuple[int, ...] | None: + """ + Returns the pixel value at a given position. + + :param xy: The coordinate, given as (x, y). See + :ref:`coordinate-system`. + :returns: The pixel value. If the image is a multi-layer image, + this method returns a tuple. + """ + + self.load() + return self.im.getpixel(tuple(xy)) + + def getprojection(self) -> tuple[list[int], list[int]]: + """ + Get projection to x and y axes + + :returns: Two sequences, indicating where there are non-zero + pixels along the X-axis and the Y-axis, respectively. + """ + + self.load() + x, y = self.im.getprojection() + return list(x), list(y) + + def histogram( + self, mask: Image | None = None, extrema: tuple[float, float] | None = None + ) -> list[int]: + """ + Returns a histogram for the image. The histogram is returned as a + list of pixel counts, one for each pixel value in the source + image. Counts are grouped into 256 bins for each band, even if + the image has more than 8 bits per band. If the image has more + than one band, the histograms for all bands are concatenated (for + example, the histogram for an "RGB" image contains 768 values). + + A bilevel image (mode "1") is treated as a grayscale ("L") image + by this method. + + If a mask is provided, the method returns a histogram for those + parts of the image where the mask image is non-zero. The mask + image must have the same size as the image, and be either a + bi-level image (mode "1") or a grayscale image ("L"). + + :param mask: An optional mask. + :param extrema: An optional tuple of manually-specified extrema. + :returns: A list containing pixel counts. + """ + self.load() + if mask: + mask.load() + return self.im.histogram((0, 0), mask.im) + if self.mode in ("I", "F"): + return self.im.histogram( + extrema if extrema is not None else self.getextrema() + ) + return self.im.histogram() + + def entropy( + self, mask: Image | None = None, extrema: tuple[float, float] | None = None + ) -> float: + """ + Calculates and returns the entropy for the image. + + A bilevel image (mode "1") is treated as a grayscale ("L") + image by this method. + + If a mask is provided, the method employs the histogram for + those parts of the image where the mask image is non-zero. + The mask image must have the same size as the image, and be + either a bi-level image (mode "1") or a grayscale image ("L"). + + :param mask: An optional mask. + :param extrema: An optional tuple of manually-specified extrema. + :returns: A float value representing the image entropy + """ + self.load() + if mask: + mask.load() + return self.im.entropy((0, 0), mask.im) + if self.mode in ("I", "F"): + return self.im.entropy( + extrema if extrema is not None else self.getextrema() + ) + return self.im.entropy() + + def paste( + self, + im: Image | str | float | tuple[float, ...], + box: Image | tuple[int, int, int, int] | tuple[int, int] | None = None, + mask: Image | None = None, + ) -> None: + """ + Pastes another image into this image. The box argument is either + a 2-tuple giving the upper left corner, a 4-tuple defining the + left, upper, right, and lower pixel coordinate, or None (same as + (0, 0)). See :ref:`coordinate-system`. If a 4-tuple is given, the size + of the pasted image must match the size of the region. + + If the modes don't match, the pasted image is converted to the mode of + this image (see the :py:meth:`~PIL.Image.Image.convert` method for + details). + + Instead of an image, the source can be a integer or tuple + containing pixel values. The method then fills the region + with the given color. When creating RGB images, you can + also use color strings as supported by the ImageColor module. See + :ref:`colors` for more information. + + If a mask is given, this method updates only the regions + indicated by the mask. You can use either "1", "L", "LA", "RGBA" + or "RGBa" images (if present, the alpha band is used as mask). + Where the mask is 255, the given image is copied as is. Where + the mask is 0, the current value is preserved. Intermediate + values will mix the two images together, including their alpha + channels if they have them. + + See :py:meth:`~PIL.Image.Image.alpha_composite` if you want to + combine images with respect to their alpha channels. + + :param im: Source image or pixel value (integer, float or tuple). + :param box: An optional 4-tuple giving the region to paste into. + If a 2-tuple is used instead, it's treated as the upper left + corner. If omitted or None, the source is pasted into the + upper left corner. + + If an image is given as the second argument and there is no + third, the box defaults to (0, 0), and the second argument + is interpreted as a mask image. + :param mask: An optional mask image. + """ + + if isinstance(box, Image): + if mask is not None: + msg = "If using second argument as mask, third argument must be None" + raise ValueError(msg) + # abbreviated paste(im, mask) syntax + mask = box + box = None + + if box is None: + box = (0, 0) + + if len(box) == 2: + # upper left corner given; get size from image or mask + if isinstance(im, Image): + size = im.size + elif isinstance(mask, Image): + size = mask.size + else: + # FIXME: use self.size here? + msg = "cannot determine region size; use 4-item box" + raise ValueError(msg) + box += (box[0] + size[0], box[1] + size[1]) + + source: core.ImagingCore | str | float | tuple[float, ...] + if isinstance(im, str): + from . import ImageColor + + source = ImageColor.getcolor(im, self.mode) + elif isinstance(im, Image): + im.load() + if self.mode != im.mode: + if self.mode != "RGB" or im.mode not in ("LA", "RGBA", "RGBa"): + # should use an adapter for this! + im = im.convert(self.mode) + source = im.im + else: + source = im + + self._ensure_mutable() + + if mask: + mask.load() + self.im.paste(source, box, mask.im) + else: + self.im.paste(source, box) + + def alpha_composite( + self, im: Image, dest: Sequence[int] = (0, 0), source: Sequence[int] = (0, 0) + ) -> None: + """'In-place' analog of Image.alpha_composite. Composites an image + onto this image. + + :param im: image to composite over this one + :param dest: Optional 2 tuple (left, top) specifying the upper + left corner in this (destination) image. + :param source: Optional 2 (left, top) tuple for the upper left + corner in the overlay source image, or 4 tuple (left, top, right, + bottom) for the bounds of the source rectangle + + Performance Note: Not currently implemented in-place in the core layer. + """ + + if not isinstance(source, (list, tuple)): + msg = "Source must be a list or tuple" + raise ValueError(msg) + if not isinstance(dest, (list, tuple)): + msg = "Destination must be a list or tuple" + raise ValueError(msg) + + if len(source) == 4: + overlay_crop_box = tuple(source) + elif len(source) == 2: + overlay_crop_box = tuple(source) + im.size + else: + msg = "Source must be a sequence of length 2 or 4" + raise ValueError(msg) + + if not len(dest) == 2: + msg = "Destination must be a sequence of length 2" + raise ValueError(msg) + if min(source) < 0: + msg = "Source must be non-negative" + raise ValueError(msg) + + # over image, crop if it's not the whole image. + if overlay_crop_box == (0, 0) + im.size: + overlay = im + else: + overlay = im.crop(overlay_crop_box) + + # target for the paste + box = tuple(dest) + (dest[0] + overlay.width, dest[1] + overlay.height) + + # destination image. don't copy if we're using the whole image. + if box == (0, 0) + self.size: + background = self + else: + background = self.crop(box) + + result = alpha_composite(background, overlay) + self.paste(result, box) + + def point( + self, + lut: ( + Sequence[float] + | NumpyArray + | Callable[[int], float] + | Callable[[ImagePointTransform], ImagePointTransform | float] + | ImagePointHandler + ), + mode: str | None = None, + ) -> Image: + """ + Maps this image through a lookup table or function. + + :param lut: A lookup table, containing 256 (or 65536 if + self.mode=="I" and mode == "L") values per band in the + image. A function can be used instead, it should take a + single argument. The function is called once for each + possible pixel value, and the resulting table is applied to + all bands of the image. + + It may also be an :py:class:`~PIL.Image.ImagePointHandler` + object:: + + class Example(Image.ImagePointHandler): + def point(self, im: Image) -> Image: + # Return result + :param mode: Output mode (default is same as input). This can only be used if + the source image has mode "L" or "P", and the output has mode "1" or the + source image mode is "I" and the output mode is "L". + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + self.load() + + if isinstance(lut, ImagePointHandler): + return lut.point(self) + + if callable(lut): + # if it isn't a list, it should be a function + if self.mode in ("I", "I;16", "F"): + # check if the function can be used with point_transform + # UNDONE wiredfool -- I think this prevents us from ever doing + # a gamma function point transform on > 8bit images. + scale, offset = _getscaleoffset(lut) # type: ignore[arg-type] + return self._new(self.im.point_transform(scale, offset)) + # for other modes, convert the function to a table + flatLut = [lut(i) for i in range(256)] * self.im.bands # type: ignore[arg-type] + else: + flatLut = lut + + if self.mode == "F": + # FIXME: _imaging returns a confusing error message for this case + msg = "point operation not supported for this mode" + raise ValueError(msg) + + if mode != "F": + flatLut = [round(i) for i in flatLut] + return self._new(self.im.point(flatLut, mode)) + + def putalpha(self, alpha: Image | int) -> None: + """ + Adds or replaces the alpha layer in this image. If the image + does not have an alpha layer, it's converted to "LA" or "RGBA". + The new layer must be either "L" or "1". + + :param alpha: The new alpha layer. This can either be an "L" or "1" + image having the same size as this image, or an integer. + """ + + self._ensure_mutable() + + if self.mode not in ("LA", "PA", "RGBA"): + # attempt to promote self to a matching alpha mode + try: + mode = getmodebase(self.mode) + "A" + try: + self.im.setmode(mode) + except (AttributeError, ValueError) as e: + # do things the hard way + im = self.im.convert(mode) + if im.mode not in ("LA", "PA", "RGBA"): + msg = "alpha channel could not be added" + raise ValueError(msg) from e # sanity check + self.im = im + self._mode = self.im.mode + except KeyError as e: + msg = "illegal image mode" + raise ValueError(msg) from e + + if self.mode in ("LA", "PA"): + band = 1 + else: + band = 3 + + if isinstance(alpha, Image): + # alpha layer + if alpha.mode not in ("1", "L"): + msg = "illegal image mode" + raise ValueError(msg) + alpha.load() + if alpha.mode == "1": + alpha = alpha.convert("L") + else: + # constant alpha + try: + self.im.fillband(band, alpha) + except (AttributeError, ValueError): + # do things the hard way + alpha = new("L", self.size, alpha) + else: + return + + self.im.putband(alpha.im, band) + + def putdata( + self, + data: Sequence[float] | Sequence[Sequence[int]] | core.ImagingCore | NumpyArray, + scale: float = 1.0, + offset: float = 0.0, + ) -> None: + """ + Copies pixel data from a flattened sequence object into the image. The + values should start at the upper left corner (0, 0), continue to the + end of the line, followed directly by the first value of the second + line, and so on. Data will be read until either the image or the + sequence ends. The scale and offset values are used to adjust the + sequence values: **pixel = value*scale + offset**. + + :param data: A flattened sequence object. See :ref:`colors` for more + information about values. + :param scale: An optional scale value. The default is 1.0. + :param offset: An optional offset value. The default is 0.0. + """ + + self._ensure_mutable() + + self.im.putdata(data, scale, offset) + + def putpalette( + self, + data: ImagePalette.ImagePalette | bytes | Sequence[int], + rawmode: str = "RGB", + ) -> None: + """ + Attaches a palette to this image. The image must be a "P", "PA", "L" + or "LA" image. + + The palette sequence must contain at most 256 colors, made up of one + integer value for each channel in the raw mode. + For example, if the raw mode is "RGB", then it can contain at most 768 + values, made up of red, green and blue values for the corresponding pixel + index in the 256 colors. + If the raw mode is "RGBA", then it can contain at most 1024 values, + containing red, green, blue and alpha values. + + Alternatively, an 8-bit string may be used instead of an integer sequence. + + :param data: A palette sequence (either a list or a string). + :param rawmode: The raw mode of the palette. Either "RGB", "RGBA", or a mode + that can be transformed to "RGB" or "RGBA" (e.g. "R", "BGR;15", "RGBA;L"). + """ + from . import ImagePalette + + if self.mode not in ("L", "LA", "P", "PA"): + msg = "illegal image mode" + raise ValueError(msg) + if isinstance(data, ImagePalette.ImagePalette): + if data.rawmode is not None: + palette = ImagePalette.raw(data.rawmode, data.palette) + else: + palette = ImagePalette.ImagePalette(palette=data.palette) + palette.dirty = 1 + else: + if not isinstance(data, bytes): + data = bytes(data) + palette = ImagePalette.raw(rawmode, data) + self._mode = "PA" if "A" in self.mode else "P" + self.palette = palette + self.palette.mode = "RGBA" if "A" in rawmode else "RGB" + self.load() # install new palette + + def putpixel( + self, xy: tuple[int, int], value: float | tuple[int, ...] | list[int] + ) -> None: + """ + Modifies the pixel at the given position. The color is given as + a single numerical value for single-band images, and a tuple for + multi-band images. In addition to this, RGB and RGBA tuples are + accepted for P and PA images. See :ref:`colors` for more information. + + Note that this method is relatively slow. For more extensive changes, + use :py:meth:`~PIL.Image.Image.paste` or the :py:mod:`~PIL.ImageDraw` + module instead. + + See: + + * :py:meth:`~PIL.Image.Image.paste` + * :py:meth:`~PIL.Image.Image.putdata` + * :py:mod:`~PIL.ImageDraw` + + :param xy: The pixel coordinate, given as (x, y). See + :ref:`coordinate-system`. + :param value: The pixel value. + """ + + self._ensure_mutable() + + if ( + self.mode in ("P", "PA") + and isinstance(value, (list, tuple)) + and len(value) in [3, 4] + ): + # RGB or RGBA value for a P or PA image + if self.mode == "PA": + alpha = value[3] if len(value) == 4 else 255 + value = value[:3] + assert self.palette is not None + palette_index = self.palette.getcolor(tuple(value), self) + value = (palette_index, alpha) if self.mode == "PA" else palette_index + return self.im.putpixel(xy, value) + + def remap_palette( + self, dest_map: list[int], source_palette: bytes | bytearray | None = None + ) -> Image: + """ + Rewrites the image to reorder the palette. + + :param dest_map: A list of indexes into the original palette. + e.g. ``[1,0]`` would swap a two item palette, and ``list(range(256))`` + is the identity transform. + :param source_palette: Bytes or None. + :returns: An :py:class:`~PIL.Image.Image` object. + + """ + from . import ImagePalette + + if self.mode not in ("L", "P"): + msg = "illegal image mode" + raise ValueError(msg) + + bands = 3 + palette_mode = "RGB" + if source_palette is None: + if self.mode == "P": + self.load() + palette_mode = self.im.getpalettemode() + if palette_mode == "RGBA": + bands = 4 + source_palette = self.im.getpalette(palette_mode, palette_mode) + else: # L-mode + source_palette = bytearray(i // 3 for i in range(768)) + elif len(source_palette) > 768: + bands = 4 + palette_mode = "RGBA" + + palette_bytes = b"" + new_positions = [0] * 256 + + # pick only the used colors from the palette + for i, oldPosition in enumerate(dest_map): + palette_bytes += source_palette[ + oldPosition * bands : oldPosition * bands + bands + ] + new_positions[oldPosition] = i + + # replace the palette color id of all pixel with the new id + + # Palette images are [0..255], mapped through a 1 or 3 + # byte/color map. We need to remap the whole image + # from palette 1 to palette 2. New_positions is + # an array of indexes into palette 1. Palette 2 is + # palette 1 with any holes removed. + + # We're going to leverage the convert mechanism to use the + # C code to remap the image from palette 1 to palette 2, + # by forcing the source image into 'L' mode and adding a + # mapping 'L' mode palette, then converting back to 'L' + # sans palette thus converting the image bytes, then + # assigning the optimized RGB palette. + + # perf reference, 9500x4000 gif, w/~135 colors + # 14 sec prepatch, 1 sec postpatch with optimization forced. + + mapping_palette = bytearray(new_positions) + + m_im = self.copy() + m_im._mode = "P" + + m_im.palette = ImagePalette.ImagePalette( + palette_mode, palette=mapping_palette * bands + ) + # possibly set palette dirty, then + # m_im.putpalette(mapping_palette, 'L') # converts to 'P' + # or just force it. + # UNDONE -- this is part of the general issue with palettes + m_im.im.putpalette(palette_mode, palette_mode + ";L", m_im.palette.tobytes()) + + m_im = m_im.convert("L") + + m_im.putpalette(palette_bytes, palette_mode) + m_im.palette = ImagePalette.ImagePalette(palette_mode, palette=palette_bytes) + + if "transparency" in self.info: + try: + m_im.info["transparency"] = dest_map.index(self.info["transparency"]) + except ValueError: + if "transparency" in m_im.info: + del m_im.info["transparency"] + + return m_im + + def _get_safe_box( + self, + size: tuple[int, int], + resample: Resampling, + box: tuple[float, float, float, float], + ) -> tuple[int, int, int, int]: + """Expands the box so it includes adjacent pixels + that may be used by resampling with the given resampling filter. + """ + filter_support = _filters_support[resample] - 0.5 + scale_x = (box[2] - box[0]) / size[0] + scale_y = (box[3] - box[1]) / size[1] + support_x = filter_support * scale_x + support_y = filter_support * scale_y + + return ( + max(0, int(box[0] - support_x)), + max(0, int(box[1] - support_y)), + min(self.size[0], math.ceil(box[2] + support_x)), + min(self.size[1], math.ceil(box[3] + support_y)), + ) + + def resize( + self, + size: tuple[int, int] | list[int] | NumpyArray, + resample: int | None = None, + box: tuple[float, float, float, float] | None = None, + reducing_gap: float | None = None, + ) -> Image: + """ + Returns a resized copy of this image. + + :param size: The requested size in pixels, as a tuple or array: + (width, height). + :param resample: An optional resampling filter. This can be + one of :py:data:`Resampling.NEAREST`, :py:data:`Resampling.BOX`, + :py:data:`Resampling.BILINEAR`, :py:data:`Resampling.HAMMING`, + :py:data:`Resampling.BICUBIC` or :py:data:`Resampling.LANCZOS`. + If the image has mode "1" or "P", it is always set to + :py:data:`Resampling.NEAREST`. Otherwise, the default filter is + :py:data:`Resampling.BICUBIC`. See: :ref:`concept-filters`. + :param box: An optional 4-tuple of floats providing + the source image region to be scaled. + The values must be within (0, 0, width, height) rectangle. + If omitted or None, the entire source is used. + :param reducing_gap: Apply optimization by resizing the image + in two steps. First, reducing the image by integer times + using :py:meth:`~PIL.Image.Image.reduce`. + Second, resizing using regular resampling. The last step + changes size no less than by ``reducing_gap`` times. + ``reducing_gap`` may be None (no first step is performed) + or should be greater than 1.0. The bigger ``reducing_gap``, + the closer the result to the fair resampling. + The smaller ``reducing_gap``, the faster resizing. + With ``reducing_gap`` greater or equal to 3.0, the result is + indistinguishable from fair resampling in most cases. + The default value is None (no optimization). + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + if resample is None: + resample = Resampling.BICUBIC + elif resample not in ( + Resampling.NEAREST, + Resampling.BILINEAR, + Resampling.BICUBIC, + Resampling.LANCZOS, + Resampling.BOX, + Resampling.HAMMING, + ): + msg = f"Unknown resampling filter ({resample})." + + filters = [ + f"{filter[1]} ({filter[0]})" + for filter in ( + (Resampling.NEAREST, "Image.Resampling.NEAREST"), + (Resampling.LANCZOS, "Image.Resampling.LANCZOS"), + (Resampling.BILINEAR, "Image.Resampling.BILINEAR"), + (Resampling.BICUBIC, "Image.Resampling.BICUBIC"), + (Resampling.BOX, "Image.Resampling.BOX"), + (Resampling.HAMMING, "Image.Resampling.HAMMING"), + ) + ] + msg += f" Use {', '.join(filters[:-1])} or {filters[-1]}" + raise ValueError(msg) + + if reducing_gap is not None and reducing_gap < 1.0: + msg = "reducing_gap must be 1.0 or greater" + raise ValueError(msg) + + if box is None: + box = (0, 0) + self.size + + size = tuple(size) + if self.size == size and box == (0, 0) + self.size: + return self.copy() + + if self.mode in ("1", "P"): + resample = Resampling.NEAREST + + if self.mode in ["LA", "RGBA"] and resample != Resampling.NEAREST: + im = self.convert({"LA": "La", "RGBA": "RGBa"}[self.mode]) + im = im.resize(size, resample, box) + return im.convert(self.mode) + + self.load() + + if reducing_gap is not None and resample != Resampling.NEAREST: + factor_x = int((box[2] - box[0]) / size[0] / reducing_gap) or 1 + factor_y = int((box[3] - box[1]) / size[1] / reducing_gap) or 1 + if factor_x > 1 or factor_y > 1: + reduce_box = self._get_safe_box(size, cast(Resampling, resample), box) + factor = (factor_x, factor_y) + self = ( + self.reduce(factor, box=reduce_box) + if callable(self.reduce) + else Image.reduce(self, factor, box=reduce_box) + ) + box = ( + (box[0] - reduce_box[0]) / factor_x, + (box[1] - reduce_box[1]) / factor_y, + (box[2] - reduce_box[0]) / factor_x, + (box[3] - reduce_box[1]) / factor_y, + ) + + return self._new(self.im.resize(size, resample, box)) + + def reduce( + self, + factor: int | tuple[int, int], + box: tuple[int, int, int, int] | None = None, + ) -> Image: + """ + Returns a copy of the image reduced ``factor`` times. + If the size of the image is not dividable by ``factor``, + the resulting size will be rounded up. + + :param factor: A greater than 0 integer or tuple of two integers + for width and height separately. + :param box: An optional 4-tuple of ints providing + the source image region to be reduced. + The values must be within ``(0, 0, width, height)`` rectangle. + If omitted or ``None``, the entire source is used. + """ + if not isinstance(factor, (list, tuple)): + factor = (factor, factor) + + if box is None: + box = (0, 0) + self.size + + if factor == (1, 1) and box == (0, 0) + self.size: + return self.copy() + + if self.mode in ["LA", "RGBA"]: + im = self.convert({"LA": "La", "RGBA": "RGBa"}[self.mode]) + im = im.reduce(factor, box) + return im.convert(self.mode) + + self.load() + + return self._new(self.im.reduce(factor, box)) + + def rotate( + self, + angle: float, + resample: Resampling = Resampling.NEAREST, + expand: int | bool = False, + center: tuple[float, float] | None = None, + translate: tuple[int, int] | None = None, + fillcolor: float | tuple[float, ...] | str | None = None, + ) -> Image: + """ + Returns a rotated copy of this image. This method returns a + copy of this image, rotated the given number of degrees counter + clockwise around its centre. + + :param angle: In degrees counter clockwise. + :param resample: An optional resampling filter. This can be + one of :py:data:`Resampling.NEAREST` (use nearest neighbour), + :py:data:`Resampling.BILINEAR` (linear interpolation in a 2x2 + environment), or :py:data:`Resampling.BICUBIC` (cubic spline + interpolation in a 4x4 environment). If omitted, or if the image has + mode "1" or "P", it is set to :py:data:`Resampling.NEAREST`. + See :ref:`concept-filters`. + :param expand: Optional expansion flag. If true, expands the output + image to make it large enough to hold the entire rotated image. + If false or omitted, make the output image the same size as the + input image. Note that the expand flag assumes rotation around + the center and no translation. + :param center: Optional center of rotation (a 2-tuple). Origin is + the upper left corner. Default is the center of the image. + :param translate: An optional post-rotate translation (a 2-tuple). + :param fillcolor: An optional color for area outside the rotated image. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + angle = angle % 360.0 + + # Fast paths regardless of filter, as long as we're not + # translating or changing the center. + if not (center or translate): + if angle == 0: + return self.copy() + if angle == 180: + return self.transpose(Transpose.ROTATE_180) + if angle in (90, 270) and (expand or self.width == self.height): + return self.transpose( + Transpose.ROTATE_90 if angle == 90 else Transpose.ROTATE_270 + ) + + # Calculate the affine matrix. Note that this is the reverse + # transformation (from destination image to source) because we + # want to interpolate the (discrete) destination pixel from + # the local area around the (floating) source pixel. + + # The matrix we actually want (note that it operates from the right): + # (1, 0, tx) (1, 0, cx) ( cos a, sin a, 0) (1, 0, -cx) + # (0, 1, ty) * (0, 1, cy) * (-sin a, cos a, 0) * (0, 1, -cy) + # (0, 0, 1) (0, 0, 1) ( 0, 0, 1) (0, 0, 1) + + # The reverse matrix is thus: + # (1, 0, cx) ( cos -a, sin -a, 0) (1, 0, -cx) (1, 0, -tx) + # (0, 1, cy) * (-sin -a, cos -a, 0) * (0, 1, -cy) * (0, 1, -ty) + # (0, 0, 1) ( 0, 0, 1) (0, 0, 1) (0, 0, 1) + + # In any case, the final translation may be updated at the end to + # compensate for the expand flag. + + w, h = self.size + + if translate is None: + post_trans = (0, 0) + else: + post_trans = translate + if center is None: + center = (w / 2, h / 2) + + angle = -math.radians(angle) + matrix = [ + round(math.cos(angle), 15), + round(math.sin(angle), 15), + 0.0, + round(-math.sin(angle), 15), + round(math.cos(angle), 15), + 0.0, + ] + + def transform(x: float, y: float, matrix: list[float]) -> tuple[float, float]: + (a, b, c, d, e, f) = matrix + return a * x + b * y + c, d * x + e * y + f + + matrix[2], matrix[5] = transform( + -center[0] - post_trans[0], -center[1] - post_trans[1], matrix + ) + matrix[2] += center[0] + matrix[5] += center[1] + + if expand: + # calculate output size + xx = [] + yy = [] + for x, y in ((0, 0), (w, 0), (w, h), (0, h)): + transformed_x, transformed_y = transform(x, y, matrix) + xx.append(transformed_x) + yy.append(transformed_y) + nw = math.ceil(max(xx)) - math.floor(min(xx)) + nh = math.ceil(max(yy)) - math.floor(min(yy)) + + # We multiply a translation matrix from the right. Because of its + # special form, this is the same as taking the image of the + # translation vector as new translation vector. + matrix[2], matrix[5] = transform(-(nw - w) / 2.0, -(nh - h) / 2.0, matrix) + w, h = nw, nh + + return self.transform( + (w, h), Transform.AFFINE, matrix, resample, fillcolor=fillcolor + ) + + def save( + self, fp: StrOrBytesPath | IO[bytes], format: str | None = None, **params: Any + ) -> None: + """ + Saves this image under the given filename. If no format is + specified, the format to use is determined from the filename + extension, if possible. + + Keyword options can be used to provide additional instructions + to the writer. If a writer doesn't recognise an option, it is + silently ignored. The available options are described in the + :doc:`image format documentation + <../handbook/image-file-formats>` for each writer. + + You can use a file object instead of a filename. In this case, + you must always specify the format. The file object must + implement the ``seek``, ``tell``, and ``write`` + methods, and be opened in binary mode. + + :param fp: A filename (string), os.PathLike object or file object. + :param format: Optional format override. If omitted, the + format to use is determined from the filename extension. + If a file object was used instead of a filename, this + parameter should always be used. + :param params: Extra parameters to the image writer. These can also be + set on the image itself through ``encoderinfo``. This is useful when + saving multiple images:: + + # Saving XMP data to a single image + from PIL import Image + red = Image.new("RGB", (1, 1), "#f00") + red.save("out.mpo", xmp=b"test") + + # Saving XMP data to the second frame of an image + from PIL import Image + black = Image.new("RGB", (1, 1)) + red = Image.new("RGB", (1, 1), "#f00") + red.encoderinfo = {"xmp": b"test"} + black.save("out.mpo", save_all=True, append_images=[red]) + :returns: None + :exception ValueError: If the output format could not be determined + from the file name. Use the format option to solve this. + :exception OSError: If the file could not be written. The file + may have been created, and may contain partial data. + """ + + filename: str | bytes = "" + open_fp = False + if is_path(fp): + filename = os.fspath(fp) + open_fp = True + elif fp == sys.stdout: + try: + fp = sys.stdout.buffer + except AttributeError: + pass + if not filename and hasattr(fp, "name") and is_path(fp.name): + # only set the name for metadata purposes + filename = os.fspath(fp.name) + + preinit() + + filename_ext = os.path.splitext(filename)[1].lower() + ext = filename_ext.decode() if isinstance(filename_ext, bytes) else filename_ext + + if not format: + if ext not in EXTENSION: + init() + try: + format = EXTENSION[ext] + except KeyError as e: + msg = f"unknown file extension: {ext}" + raise ValueError(msg) from e + + from . import ImageFile + + # may mutate self! + if isinstance(self, ImageFile.ImageFile) and os.path.abspath( + filename + ) == os.path.abspath(self.filename): + self._ensure_mutable() + else: + self.load() + + save_all = params.pop("save_all", None) + self._default_encoderinfo = params + encoderinfo = getattr(self, "encoderinfo", {}) + self._attach_default_encoderinfo(self) + self.encoderconfig: tuple[Any, ...] = () + + if format.upper() not in SAVE: + init() + if save_all or ( + save_all is None + and params.get("append_images") + and format.upper() in SAVE_ALL + ): + save_handler = SAVE_ALL[format.upper()] + else: + save_handler = SAVE[format.upper()] + + created = False + if open_fp: + created = not os.path.exists(filename) + if params.get("append", False): + # Open also for reading ("+"), because TIFF save_all + # writer needs to go back and edit the written data. + fp = builtins.open(filename, "r+b") + else: + fp = builtins.open(filename, "w+b") + else: + fp = cast(IO[bytes], fp) + + try: + save_handler(self, fp, filename) + except Exception: + if open_fp: + fp.close() + if created: + try: + os.remove(filename) + except PermissionError: + pass + raise + finally: + self.encoderinfo = encoderinfo + if open_fp: + fp.close() + + def _attach_default_encoderinfo(self, im: Image) -> dict[str, Any]: + encoderinfo = getattr(self, "encoderinfo", {}) + self.encoderinfo = {**im._default_encoderinfo, **encoderinfo} + return encoderinfo + + def seek(self, frame: int) -> None: + """ + Seeks to the given frame in this sequence file. If you seek + beyond the end of the sequence, the method raises an + ``EOFError`` exception. When a sequence file is opened, the + library automatically seeks to frame 0. + + See :py:meth:`~PIL.Image.Image.tell`. + + If defined, :attr:`~PIL.Image.Image.n_frames` refers to the + number of available frames. + + :param frame: Frame number, starting at 0. + :exception EOFError: If the call attempts to seek beyond the end + of the sequence. + """ + + # overridden by file handlers + if frame != 0: + msg = "no more images in file" + raise EOFError(msg) + + def show(self, title: str | None = None) -> None: + """ + Displays this image. This method is mainly intended for debugging purposes. + + This method calls :py:func:`PIL.ImageShow.show` internally. You can use + :py:func:`PIL.ImageShow.register` to override its default behaviour. + + The image is first saved to a temporary file. By default, it will be in + PNG format. + + On Unix, the image is then opened using the **xdg-open**, **display**, + **gm**, **eog** or **xv** utility, depending on which one can be found. + + On macOS, the image is opened with the native Preview application. + + On Windows, the image is opened with the standard PNG display utility. + + :param title: Optional title to use for the image window, where possible. + """ + + from . import ImageShow + + ImageShow.show(self, title) + + def split(self) -> tuple[Image, ...]: + """ + Split this image into individual bands. This method returns a + tuple of individual image bands from an image. For example, + splitting an "RGB" image creates three new images each + containing a copy of one of the original bands (red, green, + blue). + + If you need only one band, :py:meth:`~PIL.Image.Image.getchannel` + method can be more convenient and faster. + + :returns: A tuple containing bands. + """ + + self.load() + if self.im.bands == 1: + return (self.copy(),) + return tuple(map(self._new, self.im.split())) + + def getchannel(self, channel: int | str) -> Image: + """ + Returns an image containing a single channel of the source image. + + :param channel: What channel to return. Could be index + (0 for "R" channel of "RGB") or channel name + ("A" for alpha channel of "RGBA"). + :returns: An image in "L" mode. + + .. versionadded:: 4.3.0 + """ + self.load() + + if isinstance(channel, str): + try: + channel = self.getbands().index(channel) + except ValueError as e: + msg = f'The image has no channel "{channel}"' + raise ValueError(msg) from e + + return self._new(self.im.getband(channel)) + + def tell(self) -> int: + """ + Returns the current frame number. See :py:meth:`~PIL.Image.Image.seek`. + + If defined, :attr:`~PIL.Image.Image.n_frames` refers to the + number of available frames. + + :returns: Frame number, starting with 0. + """ + return 0 + + def thumbnail( + self, + size: tuple[float, float], + resample: Resampling = Resampling.BICUBIC, + reducing_gap: float | None = 2.0, + ) -> None: + """ + Make this image into a thumbnail. This method modifies the + image to contain a thumbnail version of itself, no larger than + the given size. This method calculates an appropriate thumbnail + size to preserve the aspect of the image, calls the + :py:meth:`~PIL.Image.Image.draft` method to configure the file reader + (where applicable), and finally resizes the image. + + Note that this function modifies the :py:class:`~PIL.Image.Image` + object in place. If you need to use the full resolution image as well, + apply this method to a :py:meth:`~PIL.Image.Image.copy` of the original + image. + + :param size: The requested size in pixels, as a 2-tuple: + (width, height). + :param resample: Optional resampling filter. This can be one + of :py:data:`Resampling.NEAREST`, :py:data:`Resampling.BOX`, + :py:data:`Resampling.BILINEAR`, :py:data:`Resampling.HAMMING`, + :py:data:`Resampling.BICUBIC` or :py:data:`Resampling.LANCZOS`. + If omitted, it defaults to :py:data:`Resampling.BICUBIC`. + (was :py:data:`Resampling.NEAREST` prior to version 2.5.0). + See: :ref:`concept-filters`. + :param reducing_gap: Apply optimization by resizing the image + in two steps. First, reducing the image by integer times + using :py:meth:`~PIL.Image.Image.reduce` or + :py:meth:`~PIL.Image.Image.draft` for JPEG images. + Second, resizing using regular resampling. The last step + changes size no less than by ``reducing_gap`` times. + ``reducing_gap`` may be None (no first step is performed) + or should be greater than 1.0. The bigger ``reducing_gap``, + the closer the result to the fair resampling. + The smaller ``reducing_gap``, the faster resizing. + With ``reducing_gap`` greater or equal to 3.0, the result is + indistinguishable from fair resampling in most cases. + The default value is 2.0 (very close to fair resampling + while still being faster in many cases). + :returns: None + """ + + provided_size = tuple(map(math.floor, size)) + + def preserve_aspect_ratio() -> tuple[int, int] | None: + def round_aspect(number: float, key: Callable[[int], float]) -> int: + return max(min(math.floor(number), math.ceil(number), key=key), 1) + + x, y = provided_size + if x >= self.width and y >= self.height: + return None + + aspect = self.width / self.height + if x / y >= aspect: + x = round_aspect(y * aspect, key=lambda n: abs(aspect - n / y)) + else: + y = round_aspect( + x / aspect, key=lambda n: 0 if n == 0 else abs(aspect - x / n) + ) + return x, y + + preserved_size = preserve_aspect_ratio() + if preserved_size is None: + return + final_size = preserved_size + + box = None + if reducing_gap is not None: + res = self.draft( + None, (int(size[0] * reducing_gap), int(size[1] * reducing_gap)) + ) + if res is not None: + box = res[1] + + if self.size != final_size: + im = self.resize(final_size, resample, box=box, reducing_gap=reducing_gap) + + self.im = im.im + self._size = final_size + self._mode = self.im.mode + + self.readonly = 0 + + # FIXME: the different transform methods need further explanation + # instead of bloating the method docs, add a separate chapter. + def transform( + self, + size: tuple[int, int], + method: Transform | ImageTransformHandler | SupportsGetData, + data: Sequence[Any] | None = None, + resample: int = Resampling.NEAREST, + fill: int = 1, + fillcolor: float | tuple[float, ...] | str | None = None, + ) -> Image: + """ + Transforms this image. This method creates a new image with the + given size, and the same mode as the original, and copies data + to the new image using the given transform. + + :param size: The output size in pixels, as a 2-tuple: + (width, height). + :param method: The transformation method. This is one of + :py:data:`Transform.EXTENT` (cut out a rectangular subregion), + :py:data:`Transform.AFFINE` (affine transform), + :py:data:`Transform.PERSPECTIVE` (perspective transform), + :py:data:`Transform.QUAD` (map a quadrilateral to a rectangle), or + :py:data:`Transform.MESH` (map a number of source quadrilaterals + in one operation). + + It may also be an :py:class:`~PIL.Image.ImageTransformHandler` + object:: + + class Example(Image.ImageTransformHandler): + def transform(self, size, data, resample, fill=1): + # Return result + + Implementations of :py:class:`~PIL.Image.ImageTransformHandler` + for some of the :py:class:`Transform` methods are provided + in :py:mod:`~PIL.ImageTransform`. + + It may also be an object with a ``method.getdata`` method + that returns a tuple supplying new ``method`` and ``data`` values:: + + class Example: + def getdata(self): + method = Image.Transform.EXTENT + data = (0, 0, 100, 100) + return method, data + :param data: Extra data to the transformation method. + :param resample: Optional resampling filter. It can be one of + :py:data:`Resampling.NEAREST` (use nearest neighbour), + :py:data:`Resampling.BILINEAR` (linear interpolation in a 2x2 + environment), or :py:data:`Resampling.BICUBIC` (cubic spline + interpolation in a 4x4 environment). If omitted, or if the image + has mode "1" or "P", it is set to :py:data:`Resampling.NEAREST`. + See: :ref:`concept-filters`. + :param fill: If ``method`` is an + :py:class:`~PIL.Image.ImageTransformHandler` object, this is one of + the arguments passed to it. Otherwise, it is unused. + :param fillcolor: Optional fill color for the area outside the + transform in the output image. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + if self.mode in ("LA", "RGBA") and resample != Resampling.NEAREST: + return ( + self.convert({"LA": "La", "RGBA": "RGBa"}[self.mode]) + .transform(size, method, data, resample, fill, fillcolor) + .convert(self.mode) + ) + + if isinstance(method, ImageTransformHandler): + return method.transform(size, self, resample=resample, fill=fill) + + if hasattr(method, "getdata"): + # compatibility w. old-style transform objects + method, data = method.getdata() + + if data is None: + msg = "missing method data" + raise ValueError(msg) + + im = new(self.mode, size, fillcolor) + if self.mode == "P" and self.palette: + im.palette = self.palette.copy() + im.info = self.info.copy() + if method == Transform.MESH: + # list of quads + for box, quad in data: + im.__transformer( + box, self, Transform.QUAD, quad, resample, fillcolor is None + ) + else: + im.__transformer( + (0, 0) + size, self, method, data, resample, fillcolor is None + ) + + return im + + def __transformer( + self, + box: tuple[int, int, int, int], + image: Image, + method: Transform, + data: Sequence[float], + resample: int = Resampling.NEAREST, + fill: bool = True, + ) -> None: + w = box[2] - box[0] + h = box[3] - box[1] + + if method == Transform.AFFINE: + data = data[:6] + + elif method == Transform.EXTENT: + # convert extent to an affine transform + x0, y0, x1, y1 = data + xs = (x1 - x0) / w + ys = (y1 - y0) / h + method = Transform.AFFINE + data = (xs, 0, x0, 0, ys, y0) + + elif method == Transform.PERSPECTIVE: + data = data[:8] + + elif method == Transform.QUAD: + # quadrilateral warp. data specifies the four corners + # given as NW, SW, SE, and NE. + nw = data[:2] + sw = data[2:4] + se = data[4:6] + ne = data[6:8] + x0, y0 = nw + As = 1.0 / w + At = 1.0 / h + data = ( + x0, + (ne[0] - x0) * As, + (sw[0] - x0) * At, + (se[0] - sw[0] - ne[0] + x0) * As * At, + y0, + (ne[1] - y0) * As, + (sw[1] - y0) * At, + (se[1] - sw[1] - ne[1] + y0) * As * At, + ) + + else: + msg = "unknown transformation method" + raise ValueError(msg) + + if resample not in ( + Resampling.NEAREST, + Resampling.BILINEAR, + Resampling.BICUBIC, + ): + if resample in (Resampling.BOX, Resampling.HAMMING, Resampling.LANCZOS): + unusable: dict[int, str] = { + Resampling.BOX: "Image.Resampling.BOX", + Resampling.HAMMING: "Image.Resampling.HAMMING", + Resampling.LANCZOS: "Image.Resampling.LANCZOS", + } + msg = unusable[resample] + f" ({resample}) cannot be used." + else: + msg = f"Unknown resampling filter ({resample})." + + filters = [ + f"{filter[1]} ({filter[0]})" + for filter in ( + (Resampling.NEAREST, "Image.Resampling.NEAREST"), + (Resampling.BILINEAR, "Image.Resampling.BILINEAR"), + (Resampling.BICUBIC, "Image.Resampling.BICUBIC"), + ) + ] + msg += f" Use {', '.join(filters[:-1])} or {filters[-1]}" + raise ValueError(msg) + + image.load() + + self.load() + + if image.mode in ("1", "P"): + resample = Resampling.NEAREST + + self.im.transform(box, image.im, method, data, resample, fill) + + def transpose(self, method: Transpose) -> Image: + """ + Transpose image (flip or rotate in 90 degree steps) + + :param method: One of :py:data:`Transpose.FLIP_LEFT_RIGHT`, + :py:data:`Transpose.FLIP_TOP_BOTTOM`, :py:data:`Transpose.ROTATE_90`, + :py:data:`Transpose.ROTATE_180`, :py:data:`Transpose.ROTATE_270`, + :py:data:`Transpose.TRANSPOSE` or :py:data:`Transpose.TRANSVERSE`. + :returns: Returns a flipped or rotated copy of this image. + """ + + self.load() + return self._new(self.im.transpose(method)) + + def effect_spread(self, distance: int) -> Image: + """ + Randomly spread pixels in an image. + + :param distance: Distance to spread pixels. + """ + self.load() + return self._new(self.im.effect_spread(distance)) + + def toqimage(self) -> ImageQt.ImageQt: + """Returns a QImage copy of this image""" + from . import ImageQt + + if not ImageQt.qt_is_installed: + msg = "Qt bindings are not installed" + raise ImportError(msg) + return ImageQt.toqimage(self) + + def toqpixmap(self) -> ImageQt.QPixmap: + """Returns a QPixmap copy of this image""" + from . import ImageQt + + if not ImageQt.qt_is_installed: + msg = "Qt bindings are not installed" + raise ImportError(msg) + return ImageQt.toqpixmap(self) + + +# -------------------------------------------------------------------- +# Abstract handlers. + + +class ImagePointHandler(abc.ABC): + """ + Used as a mixin by point transforms + (for use with :py:meth:`~PIL.Image.Image.point`) + """ + + @abc.abstractmethod + def point(self, im: Image) -> Image: + pass + + +class ImageTransformHandler(abc.ABC): + """ + Used as a mixin by geometry transforms + (for use with :py:meth:`~PIL.Image.Image.transform`) + """ + + @abc.abstractmethod + def transform( + self, + size: tuple[int, int], + image: Image, + **options: Any, + ) -> Image: + pass + + +# -------------------------------------------------------------------- +# Factories + + +def _check_size(size: Any) -> None: + """ + Common check to enforce type and sanity check on size tuples + + :param size: Should be a 2 tuple of (width, height) + :returns: None, or raises a ValueError + """ + + if not isinstance(size, (list, tuple)): + msg = "Size must be a list or tuple" + raise ValueError(msg) + if len(size) != 2: + msg = "Size must be a sequence of length 2" + raise ValueError(msg) + if size[0] < 0 or size[1] < 0: + msg = "Width and height must be >= 0" + raise ValueError(msg) + + +def new( + mode: str, + size: tuple[int, int] | list[int], + color: float | tuple[float, ...] | str | None = 0, +) -> Image: + """ + Creates a new image with the given mode and size. + + :param mode: The mode to use for the new image. See: + :ref:`concept-modes`. + :param size: A 2-tuple, containing (width, height) in pixels. + :param color: What color to use for the image. Default is black. If given, + this should be a single integer or floating point value for single-band + modes, and a tuple for multi-band modes (one value per band). When + creating RGB or HSV images, you can also use color strings as supported + by the ImageColor module. See :ref:`colors` for more information. If the + color is None, the image is not initialised. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + _check_size(size) + + if color is None: + # don't initialize + return Image()._new(core.new(mode, size)) + + if isinstance(color, str): + # css3-style specifier + + from . import ImageColor + + color = ImageColor.getcolor(color, mode) + + im = Image() + if ( + mode == "P" + and isinstance(color, (list, tuple)) + and all(isinstance(i, int) for i in color) + ): + color_ints: tuple[int, ...] = cast(tuple[int, ...], tuple(color)) + if len(color_ints) == 3 or len(color_ints) == 4: + # RGB or RGBA value for a P image + from . import ImagePalette + + im.palette = ImagePalette.ImagePalette() + color = im.palette.getcolor(color_ints) + return im._new(core.fill(mode, size, color)) + + +def frombytes( + mode: str, + size: tuple[int, int], + data: bytes | bytearray | SupportsArrayInterface, + decoder_name: str = "raw", + *args: Any, +) -> Image: + """ + Creates a copy of an image memory from pixel data in a buffer. + + In its simplest form, this function takes three arguments + (mode, size, and unpacked pixel data). + + You can also use any pixel decoder supported by PIL. For more + information on available decoders, see the section + :ref:`Writing Your Own File Codec `. + + Note that this function decodes pixel data only, not entire images. + If you have an entire image in a string, wrap it in a + :py:class:`~io.BytesIO` object, and use :py:func:`~PIL.Image.open` to load + it. + + :param mode: The image mode. See: :ref:`concept-modes`. + :param size: The image size. + :param data: A byte buffer containing raw data for the given mode. + :param decoder_name: What decoder to use. + :param args: Additional parameters for the given decoder. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + _check_size(size) + + im = new(mode, size) + if im.width != 0 and im.height != 0: + decoder_args: Any = args + if len(decoder_args) == 1 and isinstance(decoder_args[0], tuple): + # may pass tuple instead of argument list + decoder_args = decoder_args[0] + + if decoder_name == "raw" and decoder_args == (): + decoder_args = mode + + im.frombytes(data, decoder_name, decoder_args) + return im + + +def frombuffer( + mode: str, + size: tuple[int, int], + data: bytes | SupportsArrayInterface, + decoder_name: str = "raw", + *args: Any, +) -> Image: + """ + Creates an image memory referencing pixel data in a byte buffer. + + This function is similar to :py:func:`~PIL.Image.frombytes`, but uses data + in the byte buffer, where possible. This means that changes to the + original buffer object are reflected in this image). Not all modes can + share memory; supported modes include "L", "RGBX", "RGBA", and "CMYK". + + Note that this function decodes pixel data only, not entire images. + If you have an entire image file in a string, wrap it in a + :py:class:`~io.BytesIO` object, and use :py:func:`~PIL.Image.open` to load it. + + The default parameters used for the "raw" decoder differs from that used for + :py:func:`~PIL.Image.frombytes`. This is a bug, and will probably be fixed in a + future release. The current release issues a warning if you do this; to disable + the warning, you should provide the full set of parameters. See below for details. + + :param mode: The image mode. See: :ref:`concept-modes`. + :param size: The image size. + :param data: A bytes or other buffer object containing raw + data for the given mode. + :param decoder_name: What decoder to use. + :param args: Additional parameters for the given decoder. For the + default encoder ("raw"), it's recommended that you provide the + full set of parameters:: + + frombuffer(mode, size, data, "raw", mode, 0, 1) + + :returns: An :py:class:`~PIL.Image.Image` object. + + .. versionadded:: 1.1.4 + """ + + _check_size(size) + + # may pass tuple instead of argument list + if len(args) == 1 and isinstance(args[0], tuple): + args = args[0] + + if decoder_name == "raw": + if args == (): + args = mode, 0, 1 + if args[0] in _MAPMODES: + im = new(mode, (0, 0)) + im = im._new(core.map_buffer(data, size, decoder_name, 0, args)) + if mode == "P": + from . import ImagePalette + + im.palette = ImagePalette.ImagePalette("RGB", im.im.getpalette("RGB")) + im.readonly = 1 + return im + + return frombytes(mode, size, data, decoder_name, args) + + +class SupportsArrayInterface(Protocol): + """ + An object that has an ``__array_interface__`` dictionary. + """ + + @property + def __array_interface__(self) -> dict[str, Any]: + raise NotImplementedError() + + +class SupportsArrowArrayInterface(Protocol): + """ + An object that has an ``__arrow_c_array__`` method corresponding to the arrow c + data interface. + """ + + def __arrow_c_array__( + self, requested_schema: "PyCapsule" = None # type: ignore[name-defined] # noqa: F821, UP037 + ) -> tuple["PyCapsule", "PyCapsule"]: # type: ignore[name-defined] # noqa: F821, UP037 + raise NotImplementedError() + + +def fromarray(obj: SupportsArrayInterface, mode: str | None = None) -> Image: + """ + Creates an image memory from an object exporting the array interface + (using the buffer protocol):: + + from PIL import Image + import numpy as np + a = np.zeros((5, 5)) + im = Image.fromarray(a) + + If ``obj`` is not contiguous, then the ``tobytes`` method is called + and :py:func:`~PIL.Image.frombuffer` is used. + + In the case of NumPy, be aware that Pillow modes do not always correspond + to NumPy dtypes. Pillow modes only offer 1-bit pixels, 8-bit pixels, + 32-bit signed integer pixels, and 32-bit floating point pixels. + + Pillow images can also be converted to arrays:: + + from PIL import Image + import numpy as np + im = Image.open("hopper.jpg") + a = np.asarray(im) + + When converting Pillow images to arrays however, only pixel values are + transferred. This means that P and PA mode images will lose their palette. + + :param obj: Object with array interface + :param mode: Optional mode to use when reading ``obj``. Since pixel values do not + contain information about palettes or color spaces, this can be used to place + grayscale L mode data within a P mode image, or read RGB data as YCbCr for + example. + + See: :ref:`concept-modes` for general information about modes. + :returns: An image object. + + .. versionadded:: 1.1.6 + """ + arr = obj.__array_interface__ + shape = arr["shape"] + ndim = len(shape) + strides = arr.get("strides", None) + try: + typekey = (1, 1) + shape[2:], arr["typestr"] + except KeyError as e: + if mode is not None: + typekey = None + color_modes: list[str] = [] + else: + msg = "Cannot handle this data type" + raise TypeError(msg) from e + if typekey is not None: + try: + typemode, rawmode, color_modes = _fromarray_typemap[typekey] + except KeyError as e: + typekey_shape, typestr = typekey + msg = f"Cannot handle this data type: {typekey_shape}, {typestr}" + raise TypeError(msg) from e + if mode is not None: + if mode != typemode and mode not in color_modes: + deprecate("'mode' parameter for changing data types", 13) + rawmode = mode + else: + mode = typemode + if mode in ["1", "L", "I", "P", "F"]: + ndmax = 2 + elif mode == "RGB": + ndmax = 3 + else: + ndmax = 4 + if ndim > ndmax: + msg = f"Too many dimensions: {ndim} > {ndmax}." + raise ValueError(msg) + + size = 1 if ndim == 1 else shape[1], shape[0] + if strides is not None: + if hasattr(obj, "tobytes"): + obj = obj.tobytes() + elif hasattr(obj, "tostring"): + obj = obj.tostring() + else: + msg = "'strides' requires either tobytes() or tostring()" + raise ValueError(msg) + + return frombuffer(mode, size, obj, "raw", rawmode, 0, 1) + + +def fromarrow( + obj: SupportsArrowArrayInterface, mode: str, size: tuple[int, int] +) -> Image: + """Creates an image with zero-copy shared memory from an object exporting + the arrow_c_array interface protocol:: + + from PIL import Image + import pyarrow as pa + arr = pa.array([0]*(5*5*4), type=pa.uint8()) + im = Image.fromarrow(arr, 'RGBA', (5, 5)) + + If the data representation of the ``obj`` is not compatible with + Pillow internal storage, a ValueError is raised. + + Pillow images can also be converted to Arrow objects:: + + from PIL import Image + import pyarrow as pa + im = Image.open('hopper.jpg') + arr = pa.array(im) + + As with array support, when converting Pillow images to arrays, + only pixel values are transferred. This means that P and PA mode + images will lose their palette. + + :param obj: Object with an arrow_c_array interface + :param mode: Image mode. + :param size: Image size. This must match the storage of the arrow object. + :returns: An Image object + + Note that according to the Arrow spec, both the producer and the + consumer should consider the exported array to be immutable, as + unsynchronized updates will potentially cause inconsistent data. + + See: :ref:`arrow-support` for more detailed information + + .. versionadded:: 11.2.1 + + """ + if not hasattr(obj, "__arrow_c_array__"): + msg = "arrow_c_array interface not found" + raise ValueError(msg) + + (schema_capsule, array_capsule) = obj.__arrow_c_array__() + _im = core.new_arrow(mode, size, schema_capsule, array_capsule) + if _im: + return Image()._new(_im) + + msg = "new_arrow returned None without an exception" + raise ValueError(msg) + + +def fromqimage(im: ImageQt.QImage) -> ImageFile.ImageFile: + """Creates an image instance from a QImage image""" + from . import ImageQt + + if not ImageQt.qt_is_installed: + msg = "Qt bindings are not installed" + raise ImportError(msg) + return ImageQt.fromqimage(im) + + +def fromqpixmap(im: ImageQt.QPixmap) -> ImageFile.ImageFile: + """Creates an image instance from a QPixmap image""" + from . import ImageQt + + if not ImageQt.qt_is_installed: + msg = "Qt bindings are not installed" + raise ImportError(msg) + return ImageQt.fromqpixmap(im) + + +_fromarray_typemap = { + # (shape, typestr) => mode, rawmode, color modes + # first two members of shape are set to one + ((1, 1), "|b1"): ("1", "1;8", []), + ((1, 1), "|u1"): ("L", "L", ["P"]), + ((1, 1), "|i1"): ("I", "I;8", []), + ((1, 1), "u2"): ("I", "I;16B", []), + ((1, 1), "i2"): ("I", "I;16BS", []), + ((1, 1), "u4"): ("I", "I;32B", []), + ((1, 1), "i4"): ("I", "I;32BS", []), + ((1, 1), "f4"): ("F", "F;32BF", []), + ((1, 1), "f8"): ("F", "F;64BF", []), + ((1, 1, 2), "|u1"): ("LA", "LA", ["La", "PA"]), + ((1, 1, 3), "|u1"): ("RGB", "RGB", ["YCbCr", "LAB", "HSV"]), + ((1, 1, 4), "|u1"): ("RGBA", "RGBA", ["RGBa", "RGBX", "CMYK"]), + # shortcuts: + ((1, 1), f"{_ENDIAN}i4"): ("I", "I", []), + ((1, 1), f"{_ENDIAN}f4"): ("F", "F", []), +} + + +def _decompression_bomb_check(size: tuple[int, int]) -> None: + if MAX_IMAGE_PIXELS is None: + return + + pixels = max(1, size[0]) * max(1, size[1]) + + if pixels > 2 * MAX_IMAGE_PIXELS: + msg = ( + f"Image size ({pixels} pixels) exceeds limit of {2 * MAX_IMAGE_PIXELS} " + "pixels, could be decompression bomb DOS attack." + ) + raise DecompressionBombError(msg) + + if pixels > MAX_IMAGE_PIXELS: + warnings.warn( + f"Image size ({pixels} pixels) exceeds limit of {MAX_IMAGE_PIXELS} pixels, " + "could be decompression bomb DOS attack.", + DecompressionBombWarning, + ) + + +def open( + fp: StrOrBytesPath | IO[bytes], + mode: Literal["r"] = "r", + formats: list[str] | tuple[str, ...] | None = None, +) -> ImageFile.ImageFile: + """ + Opens and identifies the given image file. + + This is a lazy operation; this function identifies the file, but + the file remains open and the actual image data is not read from + the file until you try to process the data (or call the + :py:meth:`~PIL.Image.Image.load` method). See + :py:func:`~PIL.Image.new`. See :ref:`file-handling`. + + :param fp: A filename (string), os.PathLike object or a file object. + The file object must implement ``file.read``, + ``file.seek``, and ``file.tell`` methods, + and be opened in binary mode. The file object will also seek to zero + before reading. + :param mode: The mode. If given, this argument must be "r". + :param formats: A list or tuple of formats to attempt to load the file in. + This can be used to restrict the set of formats checked. + Pass ``None`` to try all supported formats. You can print the set of + available formats by running ``python3 -m PIL`` or using + the :py:func:`PIL.features.pilinfo` function. + :returns: An :py:class:`~PIL.Image.Image` object. + :exception FileNotFoundError: If the file cannot be found. + :exception PIL.UnidentifiedImageError: If the image cannot be opened and + identified. + :exception ValueError: If the ``mode`` is not "r", or if a ``StringIO`` + instance is used for ``fp``. + :exception TypeError: If ``formats`` is not ``None``, a list or a tuple. + """ + + if mode != "r": + msg = f"bad mode {repr(mode)}" # type: ignore[unreachable] + raise ValueError(msg) + elif isinstance(fp, io.StringIO): + msg = ( # type: ignore[unreachable] + "StringIO cannot be used to open an image. " + "Binary data must be used instead." + ) + raise ValueError(msg) + + if formats is None: + formats = ID + elif not isinstance(formats, (list, tuple)): + msg = "formats must be a list or tuple" # type: ignore[unreachable] + raise TypeError(msg) + + exclusive_fp = False + filename: str | bytes = "" + if is_path(fp): + filename = os.fspath(fp) + fp = builtins.open(filename, "rb") + exclusive_fp = True + else: + fp = cast(IO[bytes], fp) + + try: + fp.seek(0) + except (AttributeError, io.UnsupportedOperation): + fp = io.BytesIO(fp.read()) + exclusive_fp = True + + prefix = fp.read(16) + + preinit() + + warning_messages: list[str] = [] + + def _open_core( + fp: IO[bytes], + filename: str | bytes, + prefix: bytes, + formats: list[str] | tuple[str, ...], + ) -> ImageFile.ImageFile | None: + for i in formats: + i = i.upper() + if i not in OPEN: + init() + try: + factory, accept = OPEN[i] + result = not accept or accept(prefix) + if isinstance(result, str): + warning_messages.append(result) + elif result: + fp.seek(0) + im = factory(fp, filename) + _decompression_bomb_check(im.size) + return im + except (SyntaxError, IndexError, TypeError, struct.error) as e: + if WARN_POSSIBLE_FORMATS: + warning_messages.append(i + " opening failed. " + str(e)) + except BaseException: + if exclusive_fp: + fp.close() + raise + return None + + im = _open_core(fp, filename, prefix, formats) + + if im is None and formats is ID: + checked_formats = ID.copy() + if init(): + im = _open_core( + fp, + filename, + prefix, + tuple(format for format in formats if format not in checked_formats), + ) + + if im: + im._exclusive_fp = exclusive_fp + return im + + if exclusive_fp: + fp.close() + for message in warning_messages: + warnings.warn(message) + msg = "cannot identify image file %r" % (filename if filename else fp) + raise UnidentifiedImageError(msg) + + +# +# Image processing. + + +def alpha_composite(im1: Image, im2: Image) -> Image: + """ + Alpha composite im2 over im1. + + :param im1: The first image. Must have mode RGBA or LA. + :param im2: The second image. Must have the same mode and size as the first image. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + im1.load() + im2.load() + return im1._new(core.alpha_composite(im1.im, im2.im)) + + +def blend(im1: Image, im2: Image, alpha: float) -> Image: + """ + Creates a new image by interpolating between two input images, using + a constant alpha:: + + out = image1 * (1.0 - alpha) + image2 * alpha + + :param im1: The first image. + :param im2: The second image. Must have the same mode and size as + the first image. + :param alpha: The interpolation alpha factor. If alpha is 0.0, a + copy of the first image is returned. If alpha is 1.0, a copy of + the second image is returned. There are no restrictions on the + alpha value. If necessary, the result is clipped to fit into + the allowed output range. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + im1.load() + im2.load() + return im1._new(core.blend(im1.im, im2.im, alpha)) + + +def composite(image1: Image, image2: Image, mask: Image) -> Image: + """ + Create composite image by blending images using a transparency mask. + + :param image1: The first image. + :param image2: The second image. Must have the same mode and + size as the first image. + :param mask: A mask image. This image can have mode + "1", "L", or "RGBA", and must have the same size as the + other two images. + """ + + image = image2.copy() + image.paste(image1, None, mask) + return image + + +def eval(image: Image, *args: Callable[[int], float]) -> Image: + """ + Applies the function (which should take one argument) to each pixel + in the given image. If the image has more than one band, the same + function is applied to each band. Note that the function is + evaluated once for each possible pixel value, so you cannot use + random components or other generators. + + :param image: The input image. + :param function: A function object, taking one integer argument. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + return image.point(args[0]) + + +def merge(mode: str, bands: Sequence[Image]) -> Image: + """ + Merge a set of single band images into a new multiband image. + + :param mode: The mode to use for the output image. See: + :ref:`concept-modes`. + :param bands: A sequence containing one single-band image for + each band in the output image. All bands must have the + same size. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + + if getmodebands(mode) != len(bands) or "*" in mode: + msg = "wrong number of bands" + raise ValueError(msg) + for band in bands[1:]: + if band.mode != getmodetype(mode): + msg = "mode mismatch" + raise ValueError(msg) + if band.size != bands[0].size: + msg = "size mismatch" + raise ValueError(msg) + for band in bands: + band.load() + return bands[0]._new(core.merge(mode, *[b.im for b in bands])) + + +# -------------------------------------------------------------------- +# Plugin registry + + +def register_open( + id: str, + factory: ( + Callable[[IO[bytes], str | bytes], ImageFile.ImageFile] + | type[ImageFile.ImageFile] + ), + accept: Callable[[bytes], bool | str] | None = None, +) -> None: + """ + Register an image file plugin. This function should not be used + in application code. + + :param id: An image format identifier. + :param factory: An image file factory method. + :param accept: An optional function that can be used to quickly + reject images having another format. + """ + id = id.upper() + if id not in ID: + ID.append(id) + OPEN[id] = factory, accept + + +def register_mime(id: str, mimetype: str) -> None: + """ + Registers an image MIME type by populating ``Image.MIME``. This function + should not be used in application code. + + ``Image.MIME`` provides a mapping from image format identifiers to mime + formats, but :py:meth:`~PIL.ImageFile.ImageFile.get_format_mimetype` can + provide a different result for specific images. + + :param id: An image format identifier. + :param mimetype: The image MIME type for this format. + """ + MIME[id.upper()] = mimetype + + +def register_save( + id: str, driver: Callable[[Image, IO[bytes], str | bytes], None] +) -> None: + """ + Registers an image save function. This function should not be + used in application code. + + :param id: An image format identifier. + :param driver: A function to save images in this format. + """ + SAVE[id.upper()] = driver + + +def register_save_all( + id: str, driver: Callable[[Image, IO[bytes], str | bytes], None] +) -> None: + """ + Registers an image function to save all the frames + of a multiframe format. This function should not be + used in application code. + + :param id: An image format identifier. + :param driver: A function to save images in this format. + """ + SAVE_ALL[id.upper()] = driver + + +def register_extension(id: str, extension: str) -> None: + """ + Registers an image extension. This function should not be + used in application code. + + :param id: An image format identifier. + :param extension: An extension used for this format. + """ + EXTENSION[extension.lower()] = id.upper() + + +def register_extensions(id: str, extensions: list[str]) -> None: + """ + Registers image extensions. This function should not be + used in application code. + + :param id: An image format identifier. + :param extensions: A list of extensions used for this format. + """ + for extension in extensions: + register_extension(id, extension) + + +def registered_extensions() -> dict[str, str]: + """ + Returns a dictionary containing all file extensions belonging + to registered plugins + """ + init() + return EXTENSION + + +def register_decoder(name: str, decoder: type[ImageFile.PyDecoder]) -> None: + """ + Registers an image decoder. This function should not be + used in application code. + + :param name: The name of the decoder + :param decoder: An ImageFile.PyDecoder object + + .. versionadded:: 4.1.0 + """ + DECODERS[name] = decoder + + +def register_encoder(name: str, encoder: type[ImageFile.PyEncoder]) -> None: + """ + Registers an image encoder. This function should not be + used in application code. + + :param name: The name of the encoder + :param encoder: An ImageFile.PyEncoder object + + .. versionadded:: 4.1.0 + """ + ENCODERS[name] = encoder + + +# -------------------------------------------------------------------- +# Simple display support. + + +def _show(image: Image, **options: Any) -> None: + from . import ImageShow + + deprecate("Image._show", 13, "ImageShow.show") + ImageShow.show(image, **options) + + +# -------------------------------------------------------------------- +# Effects + + +def effect_mandelbrot( + size: tuple[int, int], extent: tuple[float, float, float, float], quality: int +) -> Image: + """ + Generate a Mandelbrot set covering the given extent. + + :param size: The requested size in pixels, as a 2-tuple: + (width, height). + :param extent: The extent to cover, as a 4-tuple: + (x0, y0, x1, y1). + :param quality: Quality. + """ + return Image()._new(core.effect_mandelbrot(size, extent, quality)) + + +def effect_noise(size: tuple[int, int], sigma: float) -> Image: + """ + Generate Gaussian noise centered around 128. + + :param size: The requested size in pixels, as a 2-tuple: + (width, height). + :param sigma: Standard deviation of noise. + """ + return Image()._new(core.effect_noise(size, sigma)) + + +def linear_gradient(mode: str) -> Image: + """ + Generate 256x256 linear gradient from black to white, top to bottom. + + :param mode: Input mode. + """ + return Image()._new(core.linear_gradient(mode)) + + +def radial_gradient(mode: str) -> Image: + """ + Generate 256x256 radial gradient from black to white, centre to edge. + + :param mode: Input mode. + """ + return Image()._new(core.radial_gradient(mode)) + + +# -------------------------------------------------------------------- +# Resources + + +def _apply_env_variables(env: dict[str, str] | None = None) -> None: + env_dict = env if env is not None else os.environ + + for var_name, setter in [ + ("PILLOW_ALIGNMENT", core.set_alignment), + ("PILLOW_BLOCK_SIZE", core.set_block_size), + ("PILLOW_BLOCKS_MAX", core.set_blocks_max), + ]: + if var_name not in env_dict: + continue + + var = env_dict[var_name].lower() + + units = 1 + for postfix, mul in [("k", 1024), ("m", 1024 * 1024)]: + if var.endswith(postfix): + units = mul + var = var[: -len(postfix)] + + try: + var_int = int(var) * units + except ValueError: + warnings.warn(f"{var_name} is not int") + continue + + try: + setter(var_int) + except ValueError as e: + warnings.warn(f"{var_name}: {e}") + + +_apply_env_variables() +atexit.register(core.clear_cache) + + +if TYPE_CHECKING: + _ExifBase = MutableMapping[int, Any] +else: + _ExifBase = MutableMapping + + +class Exif(_ExifBase): + """ + This class provides read and write access to EXIF image data:: + + from PIL import Image + im = Image.open("exif.png") + exif = im.getexif() # Returns an instance of this class + + Information can be read and written, iterated over or deleted:: + + print(exif[274]) # 1 + exif[274] = 2 + for k, v in exif.items(): + print("Tag", k, "Value", v) # Tag 274 Value 2 + del exif[274] + + To access information beyond IFD0, :py:meth:`~PIL.Image.Exif.get_ifd` + returns a dictionary:: + + from PIL import ExifTags + im = Image.open("exif_gps.jpg") + exif = im.getexif() + gps_ifd = exif.get_ifd(ExifTags.IFD.GPSInfo) + print(gps_ifd) + + Other IFDs include ``ExifTags.IFD.Exif``, ``ExifTags.IFD.MakerNote``, + ``ExifTags.IFD.Interop`` and ``ExifTags.IFD.IFD1``. + + :py:mod:`~PIL.ExifTags` also has enum classes to provide names for data:: + + print(exif[ExifTags.Base.Software]) # PIL + print(gps_ifd[ExifTags.GPS.GPSDateStamp]) # 1999:99:99 99:99:99 + """ + + endian: str | None = None + bigtiff = False + _loaded = False + + def __init__(self) -> None: + self._data: dict[int, Any] = {} + self._hidden_data: dict[int, Any] = {} + self._ifds: dict[int, dict[int, Any]] = {} + self._info: TiffImagePlugin.ImageFileDirectory_v2 | None = None + self._loaded_exif: bytes | None = None + + def _fixup(self, value: Any) -> Any: + try: + if len(value) == 1 and isinstance(value, tuple): + return value[0] + except Exception: + pass + return value + + def _fixup_dict(self, src_dict: dict[int, Any]) -> dict[int, Any]: + # Helper function + # returns a dict with any single item tuples/lists as individual values + return {k: self._fixup(v) for k, v in src_dict.items()} + + def _get_ifd_dict( + self, offset: int, group: int | None = None + ) -> dict[int, Any] | None: + try: + # an offset pointer to the location of the nested embedded IFD. + # It should be a long, but may be corrupted. + self.fp.seek(offset) + except (KeyError, TypeError): + return None + else: + from . import TiffImagePlugin + + info = TiffImagePlugin.ImageFileDirectory_v2(self.head, group=group) + info.load(self.fp) + return self._fixup_dict(dict(info)) + + def _get_head(self) -> bytes: + version = b"\x2b" if self.bigtiff else b"\x2a" + if self.endian == "<": + head = b"II" + version + b"\x00" + o32le(8) + else: + head = b"MM\x00" + version + o32be(8) + if self.bigtiff: + head += o32le(8) if self.endian == "<" else o32be(8) + head += b"\x00\x00\x00\x00" + return head + + def load(self, data: bytes) -> None: + # Extract EXIF information. This is highly experimental, + # and is likely to be replaced with something better in a future + # version. + + # The EXIF record consists of a TIFF file embedded in a JPEG + # application marker (!). + if data == self._loaded_exif: + return + self._loaded_exif = data + self._data.clear() + self._hidden_data.clear() + self._ifds.clear() + while data and data.startswith(b"Exif\x00\x00"): + data = data[6:] + if not data: + self._info = None + return + + self.fp: IO[bytes] = io.BytesIO(data) + self.head = self.fp.read(8) + # process dictionary + from . import TiffImagePlugin + + self._info = TiffImagePlugin.ImageFileDirectory_v2(self.head) + self.endian = self._info._endian + self.fp.seek(self._info.next) + self._info.load(self.fp) + + def load_from_fp(self, fp: IO[bytes], offset: int | None = None) -> None: + self._loaded_exif = None + self._data.clear() + self._hidden_data.clear() + self._ifds.clear() + + # process dictionary + from . import TiffImagePlugin + + self.fp = fp + if offset is not None: + self.head = self._get_head() + else: + self.head = self.fp.read(8) + self._info = TiffImagePlugin.ImageFileDirectory_v2(self.head) + if self.endian is None: + self.endian = self._info._endian + if offset is None: + offset = self._info.next + self.fp.tell() + self.fp.seek(offset) + self._info.load(self.fp) + + def _get_merged_dict(self) -> dict[int, Any]: + merged_dict = dict(self) + + # get EXIF extension + if ExifTags.IFD.Exif in self: + ifd = self._get_ifd_dict(self[ExifTags.IFD.Exif], ExifTags.IFD.Exif) + if ifd: + merged_dict.update(ifd) + + # GPS + if ExifTags.IFD.GPSInfo in self: + merged_dict[ExifTags.IFD.GPSInfo] = self._get_ifd_dict( + self[ExifTags.IFD.GPSInfo], ExifTags.IFD.GPSInfo + ) + + return merged_dict + + def tobytes(self, offset: int = 8) -> bytes: + from . import TiffImagePlugin + + head = self._get_head() + ifd = TiffImagePlugin.ImageFileDirectory_v2(ifh=head) + for tag, ifd_dict in self._ifds.items(): + if tag not in self: + ifd[tag] = ifd_dict + for tag, value in self.items(): + if tag in [ + ExifTags.IFD.Exif, + ExifTags.IFD.GPSInfo, + ] and not isinstance(value, dict): + value = self.get_ifd(tag) + if ( + tag == ExifTags.IFD.Exif + and ExifTags.IFD.Interop in value + and not isinstance(value[ExifTags.IFD.Interop], dict) + ): + value = value.copy() + value[ExifTags.IFD.Interop] = self.get_ifd(ExifTags.IFD.Interop) + ifd[tag] = value + return b"Exif\x00\x00" + head + ifd.tobytes(offset) + + def get_ifd(self, tag: int) -> dict[int, Any]: + if tag not in self._ifds: + if tag == ExifTags.IFD.IFD1: + if self._info is not None and self._info.next != 0: + ifd = self._get_ifd_dict(self._info.next) + if ifd is not None: + self._ifds[tag] = ifd + elif tag in [ExifTags.IFD.Exif, ExifTags.IFD.GPSInfo]: + offset = self._hidden_data.get(tag, self.get(tag)) + if offset is not None: + ifd = self._get_ifd_dict(offset, tag) + if ifd is not None: + self._ifds[tag] = ifd + elif tag in [ExifTags.IFD.Interop, ExifTags.IFD.MakerNote]: + if ExifTags.IFD.Exif not in self._ifds: + self.get_ifd(ExifTags.IFD.Exif) + tag_data = self._ifds[ExifTags.IFD.Exif][tag] + if tag == ExifTags.IFD.MakerNote: + from .TiffImagePlugin import ImageFileDirectory_v2 + + if tag_data.startswith(b"FUJIFILM"): + ifd_offset = i32le(tag_data, 8) + ifd_data = tag_data[ifd_offset:] + + makernote = {} + for i in range(struct.unpack(" 4: + (offset,) = struct.unpack("H", tag_data[:2])[0]): + ifd_tag, typ, count, data = struct.unpack( + ">HHL4s", tag_data[i * 12 + 2 : (i + 1) * 12 + 2] + ) + if ifd_tag == 0x1101: + # CameraInfo + (offset,) = struct.unpack(">L", data) + self.fp.seek(offset) + + camerainfo: dict[str, int | bytes] = { + "ModelID": self.fp.read(4) + } + + self.fp.read(4) + # Seconds since 2000 + camerainfo["TimeStamp"] = i32le(self.fp.read(12)) + + self.fp.read(4) + camerainfo["InternalSerialNumber"] = self.fp.read(4) + + self.fp.read(12) + parallax = self.fp.read(4) + handler = ImageFileDirectory_v2._load_dispatch[ + TiffTags.FLOAT + ][1] + camerainfo["Parallax"] = handler( + ImageFileDirectory_v2(), parallax, False + )[0] + + self.fp.read(4) + camerainfo["Category"] = self.fp.read(2) + + makernote = {0x1101: camerainfo} + self._ifds[tag] = makernote + else: + # Interop + ifd = self._get_ifd_dict(tag_data, tag) + if ifd is not None: + self._ifds[tag] = ifd + ifd = self._ifds.setdefault(tag, {}) + if tag == ExifTags.IFD.Exif and self._hidden_data: + ifd = { + k: v + for (k, v) in ifd.items() + if k not in (ExifTags.IFD.Interop, ExifTags.IFD.MakerNote) + } + return ifd + + def hide_offsets(self) -> None: + for tag in (ExifTags.IFD.Exif, ExifTags.IFD.GPSInfo): + if tag in self: + self._hidden_data[tag] = self[tag] + del self[tag] + + def __str__(self) -> str: + if self._info is not None: + # Load all keys into self._data + for tag in self._info: + self[tag] + + return str(self._data) + + def __len__(self) -> int: + keys = set(self._data) + if self._info is not None: + keys.update(self._info) + return len(keys) + + def __getitem__(self, tag: int) -> Any: + if self._info is not None and tag not in self._data and tag in self._info: + self._data[tag] = self._fixup(self._info[tag]) + del self._info[tag] + return self._data[tag] + + def __contains__(self, tag: object) -> bool: + return tag in self._data or (self._info is not None and tag in self._info) + + def __setitem__(self, tag: int, value: Any) -> None: + if self._info is not None and tag in self._info: + del self._info[tag] + self._data[tag] = value + + def __delitem__(self, tag: int) -> None: + if self._info is not None and tag in self._info: + del self._info[tag] + else: + del self._data[tag] + if tag in self._ifds: + del self._ifds[tag] + + def __iter__(self) -> Iterator[int]: + keys = set(self._data) + if self._info is not None: + keys.update(self._info) + return iter(keys) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageChops.py b/venv/lib/python3.10/site-packages/PIL/ImageChops.py new file mode 100644 index 0000000000000000000000000000000000000000..29a5c995fd802c9be16784f80707cfecb88b2002 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageChops.py @@ -0,0 +1,311 @@ +# +# The Python Imaging Library. +# $Id$ +# +# standard channel operations +# +# History: +# 1996-03-24 fl Created +# 1996-08-13 fl Added logical operations (for "1" images) +# 2000-10-12 fl Added offset method (from Image.py) +# +# Copyright (c) 1997-2000 by Secret Labs AB +# Copyright (c) 1996-2000 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# + +from __future__ import annotations + +from . import Image + + +def constant(image: Image.Image, value: int) -> Image.Image: + """Fill a channel with a given gray level. + + :rtype: :py:class:`~PIL.Image.Image` + """ + + return Image.new("L", image.size, value) + + +def duplicate(image: Image.Image) -> Image.Image: + """Copy a channel. Alias for :py:meth:`PIL.Image.Image.copy`. + + :rtype: :py:class:`~PIL.Image.Image` + """ + + return image.copy() + + +def invert(image: Image.Image) -> Image.Image: + """ + Invert an image (channel). :: + + out = MAX - image + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image.load() + return image._new(image.im.chop_invert()) + + +def lighter(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Compares the two images, pixel by pixel, and returns a new image containing + the lighter values. :: + + out = max(image1, image2) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_lighter(image2.im)) + + +def darker(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Compares the two images, pixel by pixel, and returns a new image containing + the darker values. :: + + out = min(image1, image2) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_darker(image2.im)) + + +def difference(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Returns the absolute value of the pixel-by-pixel difference between the two + images. :: + + out = abs(image1 - image2) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_difference(image2.im)) + + +def multiply(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Superimposes two images on top of each other. + + If you multiply an image with a solid black image, the result is black. If + you multiply with a solid white image, the image is unaffected. :: + + out = image1 * image2 / MAX + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_multiply(image2.im)) + + +def screen(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Superimposes two inverted images on top of each other. :: + + out = MAX - ((MAX - image1) * (MAX - image2) / MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_screen(image2.im)) + + +def soft_light(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Superimposes two images on top of each other using the Soft Light algorithm + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_soft_light(image2.im)) + + +def hard_light(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Superimposes two images on top of each other using the Hard Light algorithm + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_hard_light(image2.im)) + + +def overlay(image1: Image.Image, image2: Image.Image) -> Image.Image: + """ + Superimposes two images on top of each other using the Overlay algorithm + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_overlay(image2.im)) + + +def add( + image1: Image.Image, image2: Image.Image, scale: float = 1.0, offset: float = 0 +) -> Image.Image: + """ + Adds two images, dividing the result by scale and adding the + offset. If omitted, scale defaults to 1.0, and offset to 0.0. :: + + out = ((image1 + image2) / scale + offset) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_add(image2.im, scale, offset)) + + +def subtract( + image1: Image.Image, image2: Image.Image, scale: float = 1.0, offset: float = 0 +) -> Image.Image: + """ + Subtracts two images, dividing the result by scale and adding the offset. + If omitted, scale defaults to 1.0, and offset to 0.0. :: + + out = ((image1 - image2) / scale + offset) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_subtract(image2.im, scale, offset)) + + +def add_modulo(image1: Image.Image, image2: Image.Image) -> Image.Image: + """Add two images, without clipping the result. :: + + out = ((image1 + image2) % MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_add_modulo(image2.im)) + + +def subtract_modulo(image1: Image.Image, image2: Image.Image) -> Image.Image: + """Subtract two images, without clipping the result. :: + + out = ((image1 - image2) % MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_subtract_modulo(image2.im)) + + +def logical_and(image1: Image.Image, image2: Image.Image) -> Image.Image: + """Logical AND between two images. + + Both of the images must have mode "1". If you would like to perform a + logical AND on an image with a mode other than "1", try + :py:meth:`~PIL.ImageChops.multiply` instead, using a black-and-white mask + as the second image. :: + + out = ((image1 and image2) % MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_and(image2.im)) + + +def logical_or(image1: Image.Image, image2: Image.Image) -> Image.Image: + """Logical OR between two images. + + Both of the images must have mode "1". :: + + out = ((image1 or image2) % MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_or(image2.im)) + + +def logical_xor(image1: Image.Image, image2: Image.Image) -> Image.Image: + """Logical XOR between two images. + + Both of the images must have mode "1". :: + + out = ((bool(image1) != bool(image2)) % MAX) + + :rtype: :py:class:`~PIL.Image.Image` + """ + + image1.load() + image2.load() + return image1._new(image1.im.chop_xor(image2.im)) + + +def blend(image1: Image.Image, image2: Image.Image, alpha: float) -> Image.Image: + """Blend images using constant transparency weight. Alias for + :py:func:`PIL.Image.blend`. + + :rtype: :py:class:`~PIL.Image.Image` + """ + + return Image.blend(image1, image2, alpha) + + +def composite( + image1: Image.Image, image2: Image.Image, mask: Image.Image +) -> Image.Image: + """Create composite using transparency mask. Alias for + :py:func:`PIL.Image.composite`. + + :rtype: :py:class:`~PIL.Image.Image` + """ + + return Image.composite(image1, image2, mask) + + +def offset(image: Image.Image, xoffset: int, yoffset: int | None = None) -> Image.Image: + """Returns a copy of the image where data has been offset by the given + distances. Data wraps around the edges. If ``yoffset`` is omitted, it + is assumed to be equal to ``xoffset``. + + :param image: Input image. + :param xoffset: The horizontal distance. + :param yoffset: The vertical distance. If omitted, both + distances are set to the same value. + :rtype: :py:class:`~PIL.Image.Image` + """ + + if yoffset is None: + yoffset = xoffset + image.load() + return image._new(image.im.offset(xoffset, yoffset)) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageCms.py b/venv/lib/python3.10/site-packages/PIL/ImageCms.py new file mode 100644 index 0000000000000000000000000000000000000000..513e28acf33867bd4f6d046ac45ce3f71d35a433 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageCms.py @@ -0,0 +1,1076 @@ +# The Python Imaging Library. +# $Id$ + +# Optional color management support, based on Kevin Cazabon's PyCMS +# library. + +# Originally released under LGPL. Graciously donated to PIL in +# March 2009, for distribution under the standard PIL license + +# History: + +# 2009-03-08 fl Added to PIL. + +# Copyright (C) 2002-2003 Kevin Cazabon +# Copyright (c) 2009 by Fredrik Lundh +# Copyright (c) 2013 by Eric Soroos + +# See the README file for information on usage and redistribution. See +# below for the original description. +from __future__ import annotations + +import operator +import sys +from enum import IntEnum, IntFlag +from functools import reduce +from typing import Any, Literal, SupportsFloat, SupportsInt, Union + +from . import Image +from ._deprecate import deprecate +from ._typing import SupportsRead + +try: + from . import _imagingcms as core + + _CmsProfileCompatible = Union[ + str, SupportsRead[bytes], core.CmsProfile, "ImageCmsProfile" + ] +except ImportError as ex: + # Allow error import for doc purposes, but error out when accessing + # anything in core. + from ._util import DeferredError + + core = DeferredError.new(ex) + +_DESCRIPTION = """ +pyCMS + + a Python / PIL interface to the littleCMS ICC Color Management System + Copyright (C) 2002-2003 Kevin Cazabon + kevin@cazabon.com + https://www.cazabon.com + + pyCMS home page: https://www.cazabon.com/pyCMS + littleCMS home page: https://www.littlecms.com + (littleCMS is Copyright (C) 1998-2001 Marti Maria) + + Originally released under LGPL. Graciously donated to PIL in + March 2009, for distribution under the standard PIL license + + The pyCMS.py module provides a "clean" interface between Python/PIL and + pyCMSdll, taking care of some of the more complex handling of the direct + pyCMSdll functions, as well as error-checking and making sure that all + relevant data is kept together. + + While it is possible to call pyCMSdll functions directly, it's not highly + recommended. + + Version History: + + 1.0.0 pil Oct 2013 Port to LCMS 2. + + 0.1.0 pil mod March 10, 2009 + + Renamed display profile to proof profile. The proof + profile is the profile of the device that is being + simulated, not the profile of the device which is + actually used to display/print the final simulation + (that'd be the output profile) - also see LCMSAPI.txt + input colorspace -> using 'renderingIntent' -> proof + colorspace -> using 'proofRenderingIntent' -> output + colorspace + + Added LCMS FLAGS support. + Added FLAGS["SOFTPROOFING"] as default flag for + buildProofTransform (otherwise the proof profile/intent + would be ignored). + + 0.1.0 pil March 2009 - added to PIL, as PIL.ImageCms + + 0.0.2 alpha Jan 6, 2002 + + Added try/except statements around type() checks of + potential CObjects... Python won't let you use type() + on them, and raises a TypeError (stupid, if you ask + me!) + + Added buildProofTransformFromOpenProfiles() function. + Additional fixes in DLL, see DLL code for details. + + 0.0.1 alpha first public release, Dec. 26, 2002 + + Known to-do list with current version (of Python interface, not pyCMSdll): + + none + +""" + +_VERSION = "1.0.0 pil" + + +# --------------------------------------------------------------------. + + +# +# intent/direction values + + +class Intent(IntEnum): + PERCEPTUAL = 0 + RELATIVE_COLORIMETRIC = 1 + SATURATION = 2 + ABSOLUTE_COLORIMETRIC = 3 + + +class Direction(IntEnum): + INPUT = 0 + OUTPUT = 1 + PROOF = 2 + + +# +# flags + + +class Flags(IntFlag): + """Flags and documentation are taken from ``lcms2.h``.""" + + NONE = 0 + NOCACHE = 0x0040 + """Inhibit 1-pixel cache""" + NOOPTIMIZE = 0x0100 + """Inhibit optimizations""" + NULLTRANSFORM = 0x0200 + """Don't transform anyway""" + GAMUTCHECK = 0x1000 + """Out of Gamut alarm""" + SOFTPROOFING = 0x4000 + """Do softproofing""" + BLACKPOINTCOMPENSATION = 0x2000 + NOWHITEONWHITEFIXUP = 0x0004 + """Don't fix scum dot""" + HIGHRESPRECALC = 0x0400 + """Use more memory to give better accuracy""" + LOWRESPRECALC = 0x0800 + """Use less memory to minimize resources""" + # this should be 8BITS_DEVICELINK, but that is not a valid name in Python: + USE_8BITS_DEVICELINK = 0x0008 + """Create 8 bits devicelinks""" + GUESSDEVICECLASS = 0x0020 + """Guess device class (for ``transform2devicelink``)""" + KEEP_SEQUENCE = 0x0080 + """Keep profile sequence for devicelink creation""" + FORCE_CLUT = 0x0002 + """Force CLUT optimization""" + CLUT_POST_LINEARIZATION = 0x0001 + """create postlinearization tables if possible""" + CLUT_PRE_LINEARIZATION = 0x0010 + """create prelinearization tables if possible""" + NONEGATIVES = 0x8000 + """Prevent negative numbers in floating point transforms""" + COPY_ALPHA = 0x04000000 + """Alpha channels are copied on ``cmsDoTransform()``""" + NODEFAULTRESOURCEDEF = 0x01000000 + + _GRIDPOINTS_1 = 1 << 16 + _GRIDPOINTS_2 = 2 << 16 + _GRIDPOINTS_4 = 4 << 16 + _GRIDPOINTS_8 = 8 << 16 + _GRIDPOINTS_16 = 16 << 16 + _GRIDPOINTS_32 = 32 << 16 + _GRIDPOINTS_64 = 64 << 16 + _GRIDPOINTS_128 = 128 << 16 + + @staticmethod + def GRIDPOINTS(n: int) -> Flags: + """ + Fine-tune control over number of gridpoints + + :param n: :py:class:`int` in range ``0 <= n <= 255`` + """ + return Flags.NONE | ((n & 0xFF) << 16) + + +_MAX_FLAG = reduce(operator.or_, Flags) + + +_FLAGS = { + "MATRIXINPUT": 1, + "MATRIXOUTPUT": 2, + "MATRIXONLY": (1 | 2), + "NOWHITEONWHITEFIXUP": 4, # Don't hot fix scum dot + # Don't create prelinearization tables on precalculated transforms + # (internal use): + "NOPRELINEARIZATION": 16, + "GUESSDEVICECLASS": 32, # Guess device class (for transform2devicelink) + "NOTCACHE": 64, # Inhibit 1-pixel cache + "NOTPRECALC": 256, + "NULLTRANSFORM": 512, # Don't transform anyway + "HIGHRESPRECALC": 1024, # Use more memory to give better accuracy + "LOWRESPRECALC": 2048, # Use less memory to minimize resources + "WHITEBLACKCOMPENSATION": 8192, + "BLACKPOINTCOMPENSATION": 8192, + "GAMUTCHECK": 4096, # Out of Gamut alarm + "SOFTPROOFING": 16384, # Do softproofing + "PRESERVEBLACK": 32768, # Black preservation + "NODEFAULTRESOURCEDEF": 16777216, # CRD special + "GRIDPOINTS": lambda n: (n & 0xFF) << 16, # Gridpoints +} + + +# --------------------------------------------------------------------. +# Experimental PIL-level API +# --------------------------------------------------------------------. + +## +# Profile. + + +class ImageCmsProfile: + def __init__(self, profile: str | SupportsRead[bytes] | core.CmsProfile) -> None: + """ + :param profile: Either a string representing a filename, + a file like object containing a profile or a + low-level profile object + + """ + self.filename: str | None = None + + if isinstance(profile, str): + if sys.platform == "win32": + profile_bytes_path = profile.encode() + try: + profile_bytes_path.decode("ascii") + except UnicodeDecodeError: + with open(profile, "rb") as f: + self.profile = core.profile_frombytes(f.read()) + return + self.filename = profile + self.profile = core.profile_open(profile) + elif hasattr(profile, "read"): + self.profile = core.profile_frombytes(profile.read()) + elif isinstance(profile, core.CmsProfile): + self.profile = profile + else: + msg = "Invalid type for Profile" # type: ignore[unreachable] + raise TypeError(msg) + + def __getattr__(self, name: str) -> Any: + if name in ("product_name", "product_info"): + deprecate(f"ImageCms.ImageCmsProfile.{name}", 13) + return None + msg = f"'{self.__class__.__name__}' object has no attribute '{name}'" + raise AttributeError(msg) + + def tobytes(self) -> bytes: + """ + Returns the profile in a format suitable for embedding in + saved images. + + :returns: a bytes object containing the ICC profile. + """ + + return core.profile_tobytes(self.profile) + + +class ImageCmsTransform(Image.ImagePointHandler): + """ + Transform. This can be used with the procedural API, or with the standard + :py:func:`~PIL.Image.Image.point` method. + + Will return the output profile in the ``output.info['icc_profile']``. + """ + + def __init__( + self, + input: ImageCmsProfile, + output: ImageCmsProfile, + input_mode: str, + output_mode: str, + intent: Intent = Intent.PERCEPTUAL, + proof: ImageCmsProfile | None = None, + proof_intent: Intent = Intent.ABSOLUTE_COLORIMETRIC, + flags: Flags = Flags.NONE, + ): + if proof is None: + self.transform = core.buildTransform( + input.profile, output.profile, input_mode, output_mode, intent, flags + ) + else: + self.transform = core.buildProofTransform( + input.profile, + output.profile, + proof.profile, + input_mode, + output_mode, + intent, + proof_intent, + flags, + ) + # Note: inputMode and outputMode are for pyCMS compatibility only + self.input_mode = self.inputMode = input_mode + self.output_mode = self.outputMode = output_mode + + self.output_profile = output + + def point(self, im: Image.Image) -> Image.Image: + return self.apply(im) + + def apply(self, im: Image.Image, imOut: Image.Image | None = None) -> Image.Image: + if imOut is None: + imOut = Image.new(self.output_mode, im.size, None) + self.transform.apply(im.getim(), imOut.getim()) + imOut.info["icc_profile"] = self.output_profile.tobytes() + return imOut + + def apply_in_place(self, im: Image.Image) -> Image.Image: + if im.mode != self.output_mode: + msg = "mode mismatch" + raise ValueError(msg) # wrong output mode + self.transform.apply(im.getim(), im.getim()) + im.info["icc_profile"] = self.output_profile.tobytes() + return im + + +def get_display_profile(handle: SupportsInt | None = None) -> ImageCmsProfile | None: + """ + (experimental) Fetches the profile for the current display device. + + :returns: ``None`` if the profile is not known. + """ + + if sys.platform != "win32": + return None + + from . import ImageWin # type: ignore[unused-ignore, unreachable] + + if isinstance(handle, ImageWin.HDC): + profile = core.get_display_profile_win32(int(handle), 1) + else: + profile = core.get_display_profile_win32(int(handle or 0)) + if profile is None: + return None + return ImageCmsProfile(profile) + + +# --------------------------------------------------------------------. +# pyCMS compatible layer +# --------------------------------------------------------------------. + + +class PyCMSError(Exception): + """(pyCMS) Exception class. + This is used for all errors in the pyCMS API.""" + + pass + + +def profileToProfile( + im: Image.Image, + inputProfile: _CmsProfileCompatible, + outputProfile: _CmsProfileCompatible, + renderingIntent: Intent = Intent.PERCEPTUAL, + outputMode: str | None = None, + inPlace: bool = False, + flags: Flags = Flags.NONE, +) -> Image.Image | None: + """ + (pyCMS) Applies an ICC transformation to a given image, mapping from + ``inputProfile`` to ``outputProfile``. + + If the input or output profiles specified are not valid filenames, a + :exc:`PyCMSError` will be raised. If ``inPlace`` is ``True`` and + ``outputMode != im.mode``, a :exc:`PyCMSError` will be raised. + If an error occurs during application of the profiles, + a :exc:`PyCMSError` will be raised. + If ``outputMode`` is not a mode supported by the ``outputProfile`` (or by pyCMS), + a :exc:`PyCMSError` will be raised. + + This function applies an ICC transformation to im from ``inputProfile``'s + color space to ``outputProfile``'s color space using the specified rendering + intent to decide how to handle out-of-gamut colors. + + ``outputMode`` can be used to specify that a color mode conversion is to + be done using these profiles, but the specified profiles must be able + to handle that mode. I.e., if converting im from RGB to CMYK using + profiles, the input profile must handle RGB data, and the output + profile must handle CMYK data. + + :param im: An open :py:class:`~PIL.Image.Image` object (i.e. Image.new(...) + or Image.open(...), etc.) + :param inputProfile: String, as a valid filename path to the ICC input + profile you wish to use for this image, or a profile object + :param outputProfile: String, as a valid filename path to the ICC output + profile you wish to use for this image, or a profile object + :param renderingIntent: Integer (0-3) specifying the rendering intent you + wish to use for the transform + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :param outputMode: A valid PIL mode for the output image (i.e. "RGB", + "CMYK", etc.). Note: if rendering the image "inPlace", outputMode + MUST be the same mode as the input, or omitted completely. If + omitted, the outputMode will be the same as the mode of the input + image (im.mode) + :param inPlace: Boolean. If ``True``, the original image is modified in-place, + and ``None`` is returned. If ``False`` (default), a new + :py:class:`~PIL.Image.Image` object is returned with the transform applied. + :param flags: Integer (0-...) specifying additional flags + :returns: Either None or a new :py:class:`~PIL.Image.Image` object, depending on + the value of ``inPlace`` + :exception PyCMSError: + """ + + if outputMode is None: + outputMode = im.mode + + if not isinstance(renderingIntent, int) or not (0 <= renderingIntent <= 3): + msg = "renderingIntent must be an integer between 0 and 3" + raise PyCMSError(msg) + + if not isinstance(flags, int) or not (0 <= flags <= _MAX_FLAG): + msg = f"flags must be an integer between 0 and {_MAX_FLAG}" + raise PyCMSError(msg) + + try: + if not isinstance(inputProfile, ImageCmsProfile): + inputProfile = ImageCmsProfile(inputProfile) + if not isinstance(outputProfile, ImageCmsProfile): + outputProfile = ImageCmsProfile(outputProfile) + transform = ImageCmsTransform( + inputProfile, + outputProfile, + im.mode, + outputMode, + renderingIntent, + flags=flags, + ) + if inPlace: + transform.apply_in_place(im) + imOut = None + else: + imOut = transform.apply(im) + except (OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + return imOut + + +def getOpenProfile( + profileFilename: str | SupportsRead[bytes] | core.CmsProfile, +) -> ImageCmsProfile: + """ + (pyCMS) Opens an ICC profile file. + + The PyCMSProfile object can be passed back into pyCMS for use in creating + transforms and such (as in ImageCms.buildTransformFromOpenProfiles()). + + If ``profileFilename`` is not a valid filename for an ICC profile, + a :exc:`PyCMSError` will be raised. + + :param profileFilename: String, as a valid filename path to the ICC profile + you wish to open, or a file-like object. + :returns: A CmsProfile class object. + :exception PyCMSError: + """ + + try: + return ImageCmsProfile(profileFilename) + except (OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def buildTransform( + inputProfile: _CmsProfileCompatible, + outputProfile: _CmsProfileCompatible, + inMode: str, + outMode: str, + renderingIntent: Intent = Intent.PERCEPTUAL, + flags: Flags = Flags.NONE, +) -> ImageCmsTransform: + """ + (pyCMS) Builds an ICC transform mapping from the ``inputProfile`` to the + ``outputProfile``. Use applyTransform to apply the transform to a given + image. + + If the input or output profiles specified are not valid filenames, a + :exc:`PyCMSError` will be raised. If an error occurs during creation + of the transform, a :exc:`PyCMSError` will be raised. + + If ``inMode`` or ``outMode`` are not a mode supported by the ``outputProfile`` + (or by pyCMS), a :exc:`PyCMSError` will be raised. + + This function builds and returns an ICC transform from the ``inputProfile`` + to the ``outputProfile`` using the ``renderingIntent`` to determine what to do + with out-of-gamut colors. It will ONLY work for converting images that + are in ``inMode`` to images that are in ``outMode`` color format (PIL mode, + i.e. "RGB", "RGBA", "CMYK", etc.). + + Building the transform is a fair part of the overhead in + ImageCms.profileToProfile(), so if you're planning on converting multiple + images using the same input/output settings, this can save you time. + Once you have a transform object, it can be used with + ImageCms.applyProfile() to convert images without the need to re-compute + the lookup table for the transform. + + The reason pyCMS returns a class object rather than a handle directly + to the transform is that it needs to keep track of the PIL input/output + modes that the transform is meant for. These attributes are stored in + the ``inMode`` and ``outMode`` attributes of the object (which can be + manually overridden if you really want to, but I don't know of any + time that would be of use, or would even work). + + :param inputProfile: String, as a valid filename path to the ICC input + profile you wish to use for this transform, or a profile object + :param outputProfile: String, as a valid filename path to the ICC output + profile you wish to use for this transform, or a profile object + :param inMode: String, as a valid PIL mode that the appropriate profile + also supports (i.e. "RGB", "RGBA", "CMYK", etc.) + :param outMode: String, as a valid PIL mode that the appropriate profile + also supports (i.e. "RGB", "RGBA", "CMYK", etc.) + :param renderingIntent: Integer (0-3) specifying the rendering intent you + wish to use for the transform + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :param flags: Integer (0-...) specifying additional flags + :returns: A CmsTransform class object. + :exception PyCMSError: + """ + + if not isinstance(renderingIntent, int) or not (0 <= renderingIntent <= 3): + msg = "renderingIntent must be an integer between 0 and 3" + raise PyCMSError(msg) + + if not isinstance(flags, int) or not (0 <= flags <= _MAX_FLAG): + msg = f"flags must be an integer between 0 and {_MAX_FLAG}" + raise PyCMSError(msg) + + try: + if not isinstance(inputProfile, ImageCmsProfile): + inputProfile = ImageCmsProfile(inputProfile) + if not isinstance(outputProfile, ImageCmsProfile): + outputProfile = ImageCmsProfile(outputProfile) + return ImageCmsTransform( + inputProfile, outputProfile, inMode, outMode, renderingIntent, flags=flags + ) + except (OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def buildProofTransform( + inputProfile: _CmsProfileCompatible, + outputProfile: _CmsProfileCompatible, + proofProfile: _CmsProfileCompatible, + inMode: str, + outMode: str, + renderingIntent: Intent = Intent.PERCEPTUAL, + proofRenderingIntent: Intent = Intent.ABSOLUTE_COLORIMETRIC, + flags: Flags = Flags.SOFTPROOFING, +) -> ImageCmsTransform: + """ + (pyCMS) Builds an ICC transform mapping from the ``inputProfile`` to the + ``outputProfile``, but tries to simulate the result that would be + obtained on the ``proofProfile`` device. + + If the input, output, or proof profiles specified are not valid + filenames, a :exc:`PyCMSError` will be raised. + + If an error occurs during creation of the transform, + a :exc:`PyCMSError` will be raised. + + If ``inMode`` or ``outMode`` are not a mode supported by the ``outputProfile`` + (or by pyCMS), a :exc:`PyCMSError` will be raised. + + This function builds and returns an ICC transform from the ``inputProfile`` + to the ``outputProfile``, but tries to simulate the result that would be + obtained on the ``proofProfile`` device using ``renderingIntent`` and + ``proofRenderingIntent`` to determine what to do with out-of-gamut + colors. This is known as "soft-proofing". It will ONLY work for + converting images that are in ``inMode`` to images that are in outMode + color format (PIL mode, i.e. "RGB", "RGBA", "CMYK", etc.). + + Usage of the resulting transform object is exactly the same as with + ImageCms.buildTransform(). + + Proof profiling is generally used when using an output device to get a + good idea of what the final printed/displayed image would look like on + the ``proofProfile`` device when it's quicker and easier to use the + output device for judging color. Generally, this means that the + output device is a monitor, or a dye-sub printer (etc.), and the simulated + device is something more expensive, complicated, or time consuming + (making it difficult to make a real print for color judgement purposes). + + Soft-proofing basically functions by adjusting the colors on the + output device to match the colors of the device being simulated. However, + when the simulated device has a much wider gamut than the output + device, you may obtain marginal results. + + :param inputProfile: String, as a valid filename path to the ICC input + profile you wish to use for this transform, or a profile object + :param outputProfile: String, as a valid filename path to the ICC output + (monitor, usually) profile you wish to use for this transform, or a + profile object + :param proofProfile: String, as a valid filename path to the ICC proof + profile you wish to use for this transform, or a profile object + :param inMode: String, as a valid PIL mode that the appropriate profile + also supports (i.e. "RGB", "RGBA", "CMYK", etc.) + :param outMode: String, as a valid PIL mode that the appropriate profile + also supports (i.e. "RGB", "RGBA", "CMYK", etc.) + :param renderingIntent: Integer (0-3) specifying the rendering intent you + wish to use for the input->proof (simulated) transform + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :param proofRenderingIntent: Integer (0-3) specifying the rendering intent + you wish to use for proof->output transform + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :param flags: Integer (0-...) specifying additional flags + :returns: A CmsTransform class object. + :exception PyCMSError: + """ + + if not isinstance(renderingIntent, int) or not (0 <= renderingIntent <= 3): + msg = "renderingIntent must be an integer between 0 and 3" + raise PyCMSError(msg) + + if not isinstance(flags, int) or not (0 <= flags <= _MAX_FLAG): + msg = f"flags must be an integer between 0 and {_MAX_FLAG}" + raise PyCMSError(msg) + + try: + if not isinstance(inputProfile, ImageCmsProfile): + inputProfile = ImageCmsProfile(inputProfile) + if not isinstance(outputProfile, ImageCmsProfile): + outputProfile = ImageCmsProfile(outputProfile) + if not isinstance(proofProfile, ImageCmsProfile): + proofProfile = ImageCmsProfile(proofProfile) + return ImageCmsTransform( + inputProfile, + outputProfile, + inMode, + outMode, + renderingIntent, + proofProfile, + proofRenderingIntent, + flags, + ) + except (OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +buildTransformFromOpenProfiles = buildTransform +buildProofTransformFromOpenProfiles = buildProofTransform + + +def applyTransform( + im: Image.Image, transform: ImageCmsTransform, inPlace: bool = False +) -> Image.Image | None: + """ + (pyCMS) Applies a transform to a given image. + + If ``im.mode != transform.input_mode``, a :exc:`PyCMSError` is raised. + + If ``inPlace`` is ``True`` and ``transform.input_mode != transform.output_mode``, a + :exc:`PyCMSError` is raised. + + If ``im.mode``, ``transform.input_mode`` or ``transform.output_mode`` is not + supported by pyCMSdll or the profiles you used for the transform, a + :exc:`PyCMSError` is raised. + + If an error occurs while the transform is being applied, + a :exc:`PyCMSError` is raised. + + This function applies a pre-calculated transform (from + ImageCms.buildTransform() or ImageCms.buildTransformFromOpenProfiles()) + to an image. The transform can be used for multiple images, saving + considerable calculation time if doing the same conversion multiple times. + + If you want to modify im in-place instead of receiving a new image as + the return value, set ``inPlace`` to ``True``. This can only be done if + ``transform.input_mode`` and ``transform.output_mode`` are the same, because we + can't change the mode in-place (the buffer sizes for some modes are + different). The default behavior is to return a new :py:class:`~PIL.Image.Image` + object of the same dimensions in mode ``transform.output_mode``. + + :param im: An :py:class:`~PIL.Image.Image` object, and ``im.mode`` must be the same + as the ``input_mode`` supported by the transform. + :param transform: A valid CmsTransform class object + :param inPlace: Bool. If ``True``, ``im`` is modified in place and ``None`` is + returned, if ``False``, a new :py:class:`~PIL.Image.Image` object with the + transform applied is returned (and ``im`` is not changed). The default is + ``False``. + :returns: Either ``None``, or a new :py:class:`~PIL.Image.Image` object, + depending on the value of ``inPlace``. The profile will be returned in + the image's ``info['icc_profile']``. + :exception PyCMSError: + """ + + try: + if inPlace: + transform.apply_in_place(im) + imOut = None + else: + imOut = transform.apply(im) + except (TypeError, ValueError) as v: + raise PyCMSError(v) from v + + return imOut + + +def createProfile( + colorSpace: Literal["LAB", "XYZ", "sRGB"], colorTemp: SupportsFloat = 0 +) -> core.CmsProfile: + """ + (pyCMS) Creates a profile. + + If colorSpace not in ``["LAB", "XYZ", "sRGB"]``, + a :exc:`PyCMSError` is raised. + + If using LAB and ``colorTemp`` is not a positive integer, + a :exc:`PyCMSError` is raised. + + If an error occurs while creating the profile, + a :exc:`PyCMSError` is raised. + + Use this function to create common profiles on-the-fly instead of + having to supply a profile on disk and knowing the path to it. It + returns a normal CmsProfile object that can be passed to + ImageCms.buildTransformFromOpenProfiles() to create a transform to apply + to images. + + :param colorSpace: String, the color space of the profile you wish to + create. + Currently only "LAB", "XYZ", and "sRGB" are supported. + :param colorTemp: Positive number for the white point for the profile, in + degrees Kelvin (i.e. 5000, 6500, 9600, etc.). The default is for D50 + illuminant if omitted (5000k). colorTemp is ONLY applied to LAB + profiles, and is ignored for XYZ and sRGB. + :returns: A CmsProfile class object + :exception PyCMSError: + """ + + if colorSpace not in ["LAB", "XYZ", "sRGB"]: + msg = ( + f"Color space not supported for on-the-fly profile creation ({colorSpace})" + ) + raise PyCMSError(msg) + + if colorSpace == "LAB": + try: + colorTemp = float(colorTemp) + except (TypeError, ValueError) as e: + msg = f'Color temperature must be numeric, "{colorTemp}" not valid' + raise PyCMSError(msg) from e + + try: + return core.createProfile(colorSpace, colorTemp) + except (TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileName(profile: _CmsProfileCompatible) -> str: + """ + + (pyCMS) Gets the internal product name for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, + a :exc:`PyCMSError` is raised If an error occurs while trying + to obtain the name tag, a :exc:`PyCMSError` is raised. + + Use this function to obtain the INTERNAL name of the profile (stored + in an ICC tag in the profile itself), usually the one used when the + profile was originally created. Sometimes this tag also contains + additional information supplied by the creator. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal name of the profile as stored + in an ICC tag. + :exception PyCMSError: + """ + + try: + # add an extra newline to preserve pyCMS compatibility + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + # do it in python, not c. + # // name was "%s - %s" (model, manufacturer) || Description , + # // but if the Model and Manufacturer were the same or the model + # // was long, Just the model, in 1.x + model = profile.profile.model + manufacturer = profile.profile.manufacturer + + if not (model or manufacturer): + return (profile.profile.profile_description or "") + "\n" + if not manufacturer or (model and len(model) > 30): + return f"{model}\n" + return f"{model} - {manufacturer}\n" + + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileInfo(profile: _CmsProfileCompatible) -> str: + """ + (pyCMS) Gets the internal product information for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, + a :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the info tag, + a :exc:`PyCMSError` is raised. + + Use this function to obtain the information stored in the profile's + info tag. This often contains details about the profile, and how it + was created, as supplied by the creator. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal profile information stored in + an ICC tag. + :exception PyCMSError: + """ + + try: + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + # add an extra newline to preserve pyCMS compatibility + # Python, not C. the white point bits weren't working well, + # so skipping. + # info was description \r\n\r\n copyright \r\n\r\n K007 tag \r\n\r\n whitepoint + description = profile.profile.profile_description + cpright = profile.profile.copyright + elements = [element for element in (description, cpright) if element] + return "\r\n\r\n".join(elements) + "\r\n\r\n" + + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileCopyright(profile: _CmsProfileCompatible) -> str: + """ + (pyCMS) Gets the copyright for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, a + :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the copyright tag, + a :exc:`PyCMSError` is raised. + + Use this function to obtain the information stored in the profile's + copyright tag. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal profile information stored in + an ICC tag. + :exception PyCMSError: + """ + try: + # add an extra newline to preserve pyCMS compatibility + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + return (profile.profile.copyright or "") + "\n" + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileManufacturer(profile: _CmsProfileCompatible) -> str: + """ + (pyCMS) Gets the manufacturer for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, a + :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the manufacturer tag, a + :exc:`PyCMSError` is raised. + + Use this function to obtain the information stored in the profile's + manufacturer tag. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal profile information stored in + an ICC tag. + :exception PyCMSError: + """ + try: + # add an extra newline to preserve pyCMS compatibility + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + return (profile.profile.manufacturer or "") + "\n" + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileModel(profile: _CmsProfileCompatible) -> str: + """ + (pyCMS) Gets the model for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, a + :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the model tag, + a :exc:`PyCMSError` is raised. + + Use this function to obtain the information stored in the profile's + model tag. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal profile information stored in + an ICC tag. + :exception PyCMSError: + """ + + try: + # add an extra newline to preserve pyCMS compatibility + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + return (profile.profile.model or "") + "\n" + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getProfileDescription(profile: _CmsProfileCompatible) -> str: + """ + (pyCMS) Gets the description for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, a + :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the description tag, + a :exc:`PyCMSError` is raised. + + Use this function to obtain the information stored in the profile's + description tag. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: A string containing the internal profile information stored in an + ICC tag. + :exception PyCMSError: + """ + + try: + # add an extra newline to preserve pyCMS compatibility + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + return (profile.profile.profile_description or "") + "\n" + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def getDefaultIntent(profile: _CmsProfileCompatible) -> int: + """ + (pyCMS) Gets the default intent name for the given profile. + + If ``profile`` isn't a valid CmsProfile object or filename to a profile, a + :exc:`PyCMSError` is raised. + + If an error occurs while trying to obtain the default intent, a + :exc:`PyCMSError` is raised. + + Use this function to determine the default (and usually best optimized) + rendering intent for this profile. Most profiles support multiple + rendering intents, but are intended mostly for one type of conversion. + If you wish to use a different intent than returned, use + ImageCms.isIntentSupported() to verify it will work first. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :returns: Integer 0-3 specifying the default rendering intent for this + profile. + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :exception PyCMSError: + """ + + try: + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + return profile.profile.rendering_intent + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v + + +def isIntentSupported( + profile: _CmsProfileCompatible, intent: Intent, direction: Direction +) -> Literal[-1, 1]: + """ + (pyCMS) Checks if a given intent is supported. + + Use this function to verify that you can use your desired + ``intent`` with ``profile``, and that ``profile`` can be used for the + input/output/proof profile as you desire. + + Some profiles are created specifically for one "direction", can cannot + be used for others. Some profiles can only be used for certain + rendering intents, so it's best to either verify this before trying + to create a transform with them (using this function), or catch the + potential :exc:`PyCMSError` that will occur if they don't + support the modes you select. + + :param profile: EITHER a valid CmsProfile object, OR a string of the + filename of an ICC profile. + :param intent: Integer (0-3) specifying the rendering intent you wish to + use with this profile + + ImageCms.Intent.PERCEPTUAL = 0 (DEFAULT) + ImageCms.Intent.RELATIVE_COLORIMETRIC = 1 + ImageCms.Intent.SATURATION = 2 + ImageCms.Intent.ABSOLUTE_COLORIMETRIC = 3 + + see the pyCMS documentation for details on rendering intents and what + they do. + :param direction: Integer specifying if the profile is to be used for + input, output, or proof + + INPUT = 0 (or use ImageCms.Direction.INPUT) + OUTPUT = 1 (or use ImageCms.Direction.OUTPUT) + PROOF = 2 (or use ImageCms.Direction.PROOF) + + :returns: 1 if the intent/direction are supported, -1 if they are not. + :exception PyCMSError: + """ + + try: + if not isinstance(profile, ImageCmsProfile): + profile = ImageCmsProfile(profile) + # FIXME: I get different results for the same data w. different + # compilers. Bug in LittleCMS or in the binding? + if profile.profile.is_intent_supported(intent, direction): + return 1 + else: + return -1 + except (AttributeError, OSError, TypeError, ValueError) as v: + raise PyCMSError(v) from v diff --git a/venv/lib/python3.10/site-packages/PIL/ImageColor.py b/venv/lib/python3.10/site-packages/PIL/ImageColor.py new file mode 100644 index 0000000000000000000000000000000000000000..9a15a8eb7597998f1bc9a01e8eae3588c087838b --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageColor.py @@ -0,0 +1,320 @@ +# +# The Python Imaging Library +# $Id$ +# +# map CSS3-style colour description strings to RGB +# +# History: +# 2002-10-24 fl Added support for CSS-style color strings +# 2002-12-15 fl Added RGBA support +# 2004-03-27 fl Fixed remaining int() problems for Python 1.5.2 +# 2004-07-19 fl Fixed gray/grey spelling issues +# 2009-03-05 fl Fixed rounding error in grayscale calculation +# +# Copyright (c) 2002-2004 by Secret Labs AB +# Copyright (c) 2002-2004 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import re +from functools import lru_cache + +from . import Image + + +@lru_cache +def getrgb(color: str) -> tuple[int, int, int] | tuple[int, int, int, int]: + """ + Convert a color string to an RGB or RGBA tuple. If the string cannot be + parsed, this function raises a :py:exc:`ValueError` exception. + + .. versionadded:: 1.1.4 + + :param color: A color string + :return: ``(red, green, blue[, alpha])`` + """ + if len(color) > 100: + msg = "color specifier is too long" + raise ValueError(msg) + color = color.lower() + + rgb = colormap.get(color, None) + if rgb: + if isinstance(rgb, tuple): + return rgb + rgb_tuple = getrgb(rgb) + assert len(rgb_tuple) == 3 + colormap[color] = rgb_tuple + return rgb_tuple + + # check for known string formats + if re.match("#[a-f0-9]{3}$", color): + return int(color[1] * 2, 16), int(color[2] * 2, 16), int(color[3] * 2, 16) + + if re.match("#[a-f0-9]{4}$", color): + return ( + int(color[1] * 2, 16), + int(color[2] * 2, 16), + int(color[3] * 2, 16), + int(color[4] * 2, 16), + ) + + if re.match("#[a-f0-9]{6}$", color): + return int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16) + + if re.match("#[a-f0-9]{8}$", color): + return ( + int(color[1:3], 16), + int(color[3:5], 16), + int(color[5:7], 16), + int(color[7:9], 16), + ) + + m = re.match(r"rgb\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)$", color) + if m: + return int(m.group(1)), int(m.group(2)), int(m.group(3)) + + m = re.match(r"rgb\(\s*(\d+)%\s*,\s*(\d+)%\s*,\s*(\d+)%\s*\)$", color) + if m: + return ( + int((int(m.group(1)) * 255) / 100.0 + 0.5), + int((int(m.group(2)) * 255) / 100.0 + 0.5), + int((int(m.group(3)) * 255) / 100.0 + 0.5), + ) + + m = re.match( + r"hsl\(\s*(\d+\.?\d*)\s*,\s*(\d+\.?\d*)%\s*,\s*(\d+\.?\d*)%\s*\)$", color + ) + if m: + from colorsys import hls_to_rgb + + rgb_floats = hls_to_rgb( + float(m.group(1)) / 360.0, + float(m.group(3)) / 100.0, + float(m.group(2)) / 100.0, + ) + return ( + int(rgb_floats[0] * 255 + 0.5), + int(rgb_floats[1] * 255 + 0.5), + int(rgb_floats[2] * 255 + 0.5), + ) + + m = re.match( + r"hs[bv]\(\s*(\d+\.?\d*)\s*,\s*(\d+\.?\d*)%\s*,\s*(\d+\.?\d*)%\s*\)$", color + ) + if m: + from colorsys import hsv_to_rgb + + rgb_floats = hsv_to_rgb( + float(m.group(1)) / 360.0, + float(m.group(2)) / 100.0, + float(m.group(3)) / 100.0, + ) + return ( + int(rgb_floats[0] * 255 + 0.5), + int(rgb_floats[1] * 255 + 0.5), + int(rgb_floats[2] * 255 + 0.5), + ) + + m = re.match(r"rgba\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)$", color) + if m: + return int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4)) + msg = f"unknown color specifier: {repr(color)}" + raise ValueError(msg) + + +@lru_cache +def getcolor(color: str, mode: str) -> int | tuple[int, ...]: + """ + Same as :py:func:`~PIL.ImageColor.getrgb` for most modes. However, if + ``mode`` is HSV, converts the RGB value to a HSV value, or if ``mode`` is + not color or a palette image, converts the RGB value to a grayscale value. + If the string cannot be parsed, this function raises a :py:exc:`ValueError` + exception. + + .. versionadded:: 1.1.4 + + :param color: A color string + :param mode: Convert result to this mode + :return: ``graylevel, (graylevel, alpha) or (red, green, blue[, alpha])`` + """ + # same as getrgb, but converts the result to the given mode + rgb, alpha = getrgb(color), 255 + if len(rgb) == 4: + alpha = rgb[3] + rgb = rgb[:3] + + if mode == "HSV": + from colorsys import rgb_to_hsv + + r, g, b = rgb + h, s, v = rgb_to_hsv(r / 255, g / 255, b / 255) + return int(h * 255), int(s * 255), int(v * 255) + elif Image.getmodebase(mode) == "L": + r, g, b = rgb + # ITU-R Recommendation 601-2 for nonlinear RGB + # scaled to 24 bits to match the convert's implementation. + graylevel = (r * 19595 + g * 38470 + b * 7471 + 0x8000) >> 16 + if mode[-1] == "A": + return graylevel, alpha + return graylevel + elif mode[-1] == "A": + return rgb + (alpha,) + return rgb + + +colormap: dict[str, str | tuple[int, int, int]] = { + # X11 colour table from https://drafts.csswg.org/css-color-4/, with + # gray/grey spelling issues fixed. This is a superset of HTML 4.0 + # colour names used in CSS 1. + "aliceblue": "#f0f8ff", + "antiquewhite": "#faebd7", + "aqua": "#00ffff", + "aquamarine": "#7fffd4", + "azure": "#f0ffff", + "beige": "#f5f5dc", + "bisque": "#ffe4c4", + "black": "#000000", + "blanchedalmond": "#ffebcd", + "blue": "#0000ff", + "blueviolet": "#8a2be2", + "brown": "#a52a2a", + "burlywood": "#deb887", + "cadetblue": "#5f9ea0", + "chartreuse": "#7fff00", + "chocolate": "#d2691e", + "coral": "#ff7f50", + "cornflowerblue": "#6495ed", + "cornsilk": "#fff8dc", + "crimson": "#dc143c", + "cyan": "#00ffff", + "darkblue": "#00008b", + "darkcyan": "#008b8b", + "darkgoldenrod": "#b8860b", + "darkgray": "#a9a9a9", + "darkgrey": "#a9a9a9", + "darkgreen": "#006400", + "darkkhaki": "#bdb76b", + "darkmagenta": "#8b008b", + "darkolivegreen": "#556b2f", + "darkorange": "#ff8c00", + "darkorchid": "#9932cc", + "darkred": "#8b0000", + "darksalmon": "#e9967a", + "darkseagreen": "#8fbc8f", + "darkslateblue": "#483d8b", + "darkslategray": "#2f4f4f", + "darkslategrey": "#2f4f4f", + "darkturquoise": "#00ced1", + "darkviolet": "#9400d3", + "deeppink": "#ff1493", + "deepskyblue": "#00bfff", + "dimgray": "#696969", + "dimgrey": "#696969", + "dodgerblue": "#1e90ff", + "firebrick": "#b22222", + "floralwhite": "#fffaf0", + "forestgreen": "#228b22", + "fuchsia": "#ff00ff", + "gainsboro": "#dcdcdc", + "ghostwhite": "#f8f8ff", + "gold": "#ffd700", + "goldenrod": "#daa520", + "gray": "#808080", + "grey": "#808080", + "green": "#008000", + "greenyellow": "#adff2f", + "honeydew": "#f0fff0", + "hotpink": "#ff69b4", + "indianred": "#cd5c5c", + "indigo": "#4b0082", + "ivory": "#fffff0", + "khaki": "#f0e68c", + "lavender": "#e6e6fa", + "lavenderblush": "#fff0f5", + "lawngreen": "#7cfc00", + "lemonchiffon": "#fffacd", + "lightblue": "#add8e6", + "lightcoral": "#f08080", + "lightcyan": "#e0ffff", + "lightgoldenrodyellow": "#fafad2", + "lightgreen": "#90ee90", + "lightgray": "#d3d3d3", + "lightgrey": "#d3d3d3", + "lightpink": "#ffb6c1", + "lightsalmon": "#ffa07a", + "lightseagreen": "#20b2aa", + "lightskyblue": "#87cefa", + "lightslategray": "#778899", + "lightslategrey": "#778899", + "lightsteelblue": "#b0c4de", + "lightyellow": "#ffffe0", + "lime": "#00ff00", + "limegreen": "#32cd32", + "linen": "#faf0e6", + "magenta": "#ff00ff", + "maroon": "#800000", + "mediumaquamarine": "#66cdaa", + "mediumblue": "#0000cd", + "mediumorchid": "#ba55d3", + "mediumpurple": "#9370db", + "mediumseagreen": "#3cb371", + "mediumslateblue": "#7b68ee", + "mediumspringgreen": "#00fa9a", + "mediumturquoise": "#48d1cc", + "mediumvioletred": "#c71585", + "midnightblue": "#191970", + "mintcream": "#f5fffa", + "mistyrose": "#ffe4e1", + "moccasin": "#ffe4b5", + "navajowhite": "#ffdead", + "navy": "#000080", + "oldlace": "#fdf5e6", + "olive": "#808000", + "olivedrab": "#6b8e23", + "orange": "#ffa500", + "orangered": "#ff4500", + "orchid": "#da70d6", + "palegoldenrod": "#eee8aa", + "palegreen": "#98fb98", + "paleturquoise": "#afeeee", + "palevioletred": "#db7093", + "papayawhip": "#ffefd5", + "peachpuff": "#ffdab9", + "peru": "#cd853f", + "pink": "#ffc0cb", + "plum": "#dda0dd", + "powderblue": "#b0e0e6", + "purple": "#800080", + "rebeccapurple": "#663399", + "red": "#ff0000", + "rosybrown": "#bc8f8f", + "royalblue": "#4169e1", + "saddlebrown": "#8b4513", + "salmon": "#fa8072", + "sandybrown": "#f4a460", + "seagreen": "#2e8b57", + "seashell": "#fff5ee", + "sienna": "#a0522d", + "silver": "#c0c0c0", + "skyblue": "#87ceeb", + "slateblue": "#6a5acd", + "slategray": "#708090", + "slategrey": "#708090", + "snow": "#fffafa", + "springgreen": "#00ff7f", + "steelblue": "#4682b4", + "tan": "#d2b48c", + "teal": "#008080", + "thistle": "#d8bfd8", + "tomato": "#ff6347", + "turquoise": "#40e0d0", + "violet": "#ee82ee", + "wheat": "#f5deb3", + "white": "#ffffff", + "whitesmoke": "#f5f5f5", + "yellow": "#ffff00", + "yellowgreen": "#9acd32", +} diff --git a/venv/lib/python3.10/site-packages/PIL/ImageDraw.py b/venv/lib/python3.10/site-packages/PIL/ImageDraw.py new file mode 100644 index 0000000000000000000000000000000000000000..8bcf2d8ee06ff87e966cf7bab9e1f448848cedc8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageDraw.py @@ -0,0 +1,1036 @@ +# +# The Python Imaging Library +# $Id$ +# +# drawing interface operations +# +# History: +# 1996-04-13 fl Created (experimental) +# 1996-08-07 fl Filled polygons, ellipses. +# 1996-08-13 fl Added text support +# 1998-06-28 fl Handle I and F images +# 1998-12-29 fl Added arc; use arc primitive to draw ellipses +# 1999-01-10 fl Added shape stuff (experimental) +# 1999-02-06 fl Added bitmap support +# 1999-02-11 fl Changed all primitives to take options +# 1999-02-20 fl Fixed backwards compatibility +# 2000-10-12 fl Copy on write, when necessary +# 2001-02-18 fl Use default ink for bitmap/text also in fill mode +# 2002-10-24 fl Added support for CSS-style color strings +# 2002-12-10 fl Added experimental support for RGBA-on-RGB drawing +# 2002-12-11 fl Refactored low-level drawing API (work in progress) +# 2004-08-26 fl Made Draw() a factory function, added getdraw() support +# 2004-09-04 fl Added width support to line primitive +# 2004-09-10 fl Added font mode handling +# 2006-06-19 fl Added font bearing support (getmask2) +# +# Copyright (c) 1997-2006 by Secret Labs AB +# Copyright (c) 1996-2006 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import math +import struct +from collections.abc import Sequence +from typing import cast + +from . import Image, ImageColor, ImageText + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from types import ModuleType + from typing import Any, AnyStr + + from . import ImageDraw2, ImageFont + from ._typing import Coords, _Ink + +# experimental access to the outline API +Outline: Callable[[], Image.core._Outline] = Image.core.outline + +""" +A simple 2D drawing interface for PIL images. +

+Application code should use the Draw factory, instead of +directly. +""" + + +class ImageDraw: + font: ( + ImageFont.ImageFont | ImageFont.FreeTypeFont | ImageFont.TransposedFont | None + ) = None + + def __init__(self, im: Image.Image, mode: str | None = None) -> None: + """ + Create a drawing instance. + + :param im: The image to draw in. + :param mode: Optional mode to use for color values. For RGB + images, this argument can be RGB or RGBA (to blend the + drawing into the image). For all other modes, this argument + must be the same as the image mode. If omitted, the mode + defaults to the mode of the image. + """ + im._ensure_mutable() + blend = 0 + if mode is None: + mode = im.mode + if mode != im.mode: + if mode == "RGBA" and im.mode == "RGB": + blend = 1 + else: + msg = "mode mismatch" + raise ValueError(msg) + if mode == "P": + self.palette = im.palette + else: + self.palette = None + self._image = im + self.im = im.im + self.draw = Image.core.draw(self.im, blend) + self.mode = mode + if mode in ("I", "F"): + self.ink = self.draw.draw_ink(1) + else: + self.ink = self.draw.draw_ink(-1) + if mode in ("1", "P", "I", "F"): + # FIXME: fix Fill2 to properly support matte for I+F images + self.fontmode = "1" + else: + self.fontmode = "L" # aliasing is okay for other modes + self.fill = False + + def getfont( + self, + ) -> ImageFont.ImageFont | ImageFont.FreeTypeFont | ImageFont.TransposedFont: + """ + Get the current default font. + + To set the default font for this ImageDraw instance:: + + from PIL import ImageDraw, ImageFont + draw.font = ImageFont.truetype("Tests/fonts/FreeMono.ttf") + + To set the default font for all future ImageDraw instances:: + + from PIL import ImageDraw, ImageFont + ImageDraw.ImageDraw.font = ImageFont.truetype("Tests/fonts/FreeMono.ttf") + + If the current default font is ``None``, + it is initialized with ``ImageFont.load_default()``. + + :returns: An image font.""" + if not self.font: + # FIXME: should add a font repository + from . import ImageFont + + self.font = ImageFont.load_default() + return self.font + + def _getfont( + self, font_size: float | None + ) -> ImageFont.ImageFont | ImageFont.FreeTypeFont | ImageFont.TransposedFont: + if font_size is not None: + from . import ImageFont + + return ImageFont.load_default(font_size) + else: + return self.getfont() + + def _getink( + self, ink: _Ink | None, fill: _Ink | None = None + ) -> tuple[int | None, int | None]: + result_ink = None + result_fill = None + if ink is None and fill is None: + if self.fill: + result_fill = self.ink + else: + result_ink = self.ink + else: + if ink is not None: + if isinstance(ink, str): + ink = ImageColor.getcolor(ink, self.mode) + if self.palette and isinstance(ink, tuple): + ink = self.palette.getcolor(ink, self._image) + result_ink = self.draw.draw_ink(ink) + if fill is not None: + if isinstance(fill, str): + fill = ImageColor.getcolor(fill, self.mode) + if self.palette and isinstance(fill, tuple): + fill = self.palette.getcolor(fill, self._image) + result_fill = self.draw.draw_ink(fill) + return result_ink, result_fill + + def arc( + self, + xy: Coords, + start: float, + end: float, + fill: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw an arc.""" + ink, fill = self._getink(fill) + if ink is not None: + self.draw.draw_arc(xy, start, end, ink, width) + + def bitmap( + self, xy: Sequence[int], bitmap: Image.Image, fill: _Ink | None = None + ) -> None: + """Draw a bitmap.""" + bitmap.load() + ink, fill = self._getink(fill) + if ink is None: + ink = fill + if ink is not None: + self.draw.draw_bitmap(xy, bitmap.im, ink) + + def chord( + self, + xy: Coords, + start: float, + end: float, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a chord.""" + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_chord(xy, start, end, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + self.draw.draw_chord(xy, start, end, ink, 0, width) + + def ellipse( + self, + xy: Coords, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw an ellipse.""" + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_ellipse(xy, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + self.draw.draw_ellipse(xy, ink, 0, width) + + def circle( + self, + xy: Sequence[float], + radius: float, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a circle given center coordinates and a radius.""" + ellipse_xy = (xy[0] - radius, xy[1] - radius, xy[0] + radius, xy[1] + radius) + self.ellipse(ellipse_xy, fill, outline, width) + + def line( + self, + xy: Coords, + fill: _Ink | None = None, + width: int = 0, + joint: str | None = None, + ) -> None: + """Draw a line, or a connected sequence of line segments.""" + ink = self._getink(fill)[0] + if ink is not None: + self.draw.draw_lines(xy, ink, width) + if joint == "curve" and width > 4: + points: Sequence[Sequence[float]] + if isinstance(xy[0], (list, tuple)): + points = cast(Sequence[Sequence[float]], xy) + else: + points = [ + cast(Sequence[float], tuple(xy[i : i + 2])) + for i in range(0, len(xy), 2) + ] + for i in range(1, len(points) - 1): + point = points[i] + angles = [ + math.degrees(math.atan2(end[0] - start[0], start[1] - end[1])) + % 360 + for start, end in ( + (points[i - 1], point), + (point, points[i + 1]), + ) + ] + if angles[0] == angles[1]: + # This is a straight line, so no joint is required + continue + + def coord_at_angle( + coord: Sequence[float], angle: float + ) -> tuple[float, ...]: + x, y = coord + angle -= 90 + distance = width / 2 - 1 + return tuple( + p + (math.floor(p_d) if p_d > 0 else math.ceil(p_d)) + for p, p_d in ( + (x, distance * math.cos(math.radians(angle))), + (y, distance * math.sin(math.radians(angle))), + ) + ) + + flipped = ( + angles[1] > angles[0] and angles[1] - 180 > angles[0] + ) or (angles[1] < angles[0] and angles[1] + 180 > angles[0]) + coords = [ + (point[0] - width / 2 + 1, point[1] - width / 2 + 1), + (point[0] + width / 2 - 1, point[1] + width / 2 - 1), + ] + if flipped: + start, end = (angles[1] + 90, angles[0] + 90) + else: + start, end = (angles[0] - 90, angles[1] - 90) + self.pieslice(coords, start - 90, end - 90, fill) + + if width > 8: + # Cover potential gaps between the line and the joint + if flipped: + gap_coords = [ + coord_at_angle(point, angles[0] + 90), + point, + coord_at_angle(point, angles[1] + 90), + ] + else: + gap_coords = [ + coord_at_angle(point, angles[0] - 90), + point, + coord_at_angle(point, angles[1] - 90), + ] + self.line(gap_coords, fill, width=3) + + def shape( + self, + shape: Image.core._Outline, + fill: _Ink | None = None, + outline: _Ink | None = None, + ) -> None: + """(Experimental) Draw a shape.""" + shape.close() + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_outline(shape, fill_ink, 1) + if ink is not None and ink != fill_ink: + self.draw.draw_outline(shape, ink, 0) + + def pieslice( + self, + xy: Coords, + start: float, + end: float, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a pieslice.""" + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_pieslice(xy, start, end, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + self.draw.draw_pieslice(xy, start, end, ink, 0, width) + + def point(self, xy: Coords, fill: _Ink | None = None) -> None: + """Draw one or more individual pixels.""" + ink, fill = self._getink(fill) + if ink is not None: + self.draw.draw_points(xy, ink) + + def polygon( + self, + xy: Coords, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a polygon.""" + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_polygon(xy, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + if width == 1: + self.draw.draw_polygon(xy, ink, 0, width) + elif self.im is not None: + # To avoid expanding the polygon outwards, + # use the fill as a mask + mask = Image.new("1", self.im.size) + mask_ink = self._getink(1)[0] + draw = Draw(mask) + draw.draw.draw_polygon(xy, mask_ink, 1) + + self.draw.draw_polygon(xy, ink, 0, width * 2 - 1, mask.im) + + def regular_polygon( + self, + bounding_circle: Sequence[Sequence[float] | float], + n_sides: int, + rotation: float = 0, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a regular polygon.""" + xy = _compute_regular_polygon_vertices(bounding_circle, n_sides, rotation) + self.polygon(xy, fill, outline, width) + + def rectangle( + self, + xy: Coords, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + ) -> None: + """Draw a rectangle.""" + ink, fill_ink = self._getink(outline, fill) + if fill_ink is not None: + self.draw.draw_rectangle(xy, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + self.draw.draw_rectangle(xy, ink, 0, width) + + def rounded_rectangle( + self, + xy: Coords, + radius: float = 0, + fill: _Ink | None = None, + outline: _Ink | None = None, + width: int = 1, + *, + corners: tuple[bool, bool, bool, bool] | None = None, + ) -> None: + """Draw a rounded rectangle.""" + if isinstance(xy[0], (list, tuple)): + (x0, y0), (x1, y1) = cast(Sequence[Sequence[float]], xy) + else: + x0, y0, x1, y1 = cast(Sequence[float], xy) + if x1 < x0: + msg = "x1 must be greater than or equal to x0" + raise ValueError(msg) + if y1 < y0: + msg = "y1 must be greater than or equal to y0" + raise ValueError(msg) + if corners is None: + corners = (True, True, True, True) + + d = radius * 2 + + x0 = round(x0) + y0 = round(y0) + x1 = round(x1) + y1 = round(y1) + full_x, full_y = False, False + if all(corners): + full_x = d >= x1 - x0 - 1 + if full_x: + # The two left and two right corners are joined + d = x1 - x0 + full_y = d >= y1 - y0 - 1 + if full_y: + # The two top and two bottom corners are joined + d = y1 - y0 + if full_x and full_y: + # If all corners are joined, that is a circle + return self.ellipse(xy, fill, outline, width) + + if d == 0 or not any(corners): + # If the corners have no curve, + # or there are no corners, + # that is a rectangle + return self.rectangle(xy, fill, outline, width) + + r = int(d // 2) + ink, fill_ink = self._getink(outline, fill) + + def draw_corners(pieslice: bool) -> None: + parts: tuple[tuple[tuple[float, float, float, float], int, int], ...] + if full_x: + # Draw top and bottom halves + parts = ( + ((x0, y0, x0 + d, y0 + d), 180, 360), + ((x0, y1 - d, x0 + d, y1), 0, 180), + ) + elif full_y: + # Draw left and right halves + parts = ( + ((x0, y0, x0 + d, y0 + d), 90, 270), + ((x1 - d, y0, x1, y0 + d), 270, 90), + ) + else: + # Draw four separate corners + parts = tuple( + part + for i, part in enumerate( + ( + ((x0, y0, x0 + d, y0 + d), 180, 270), + ((x1 - d, y0, x1, y0 + d), 270, 360), + ((x1 - d, y1 - d, x1, y1), 0, 90), + ((x0, y1 - d, x0 + d, y1), 90, 180), + ) + ) + if corners[i] + ) + for part in parts: + if pieslice: + self.draw.draw_pieslice(*(part + (fill_ink, 1))) + else: + self.draw.draw_arc(*(part + (ink, width))) + + if fill_ink is not None: + draw_corners(True) + + if full_x: + self.draw.draw_rectangle((x0, y0 + r + 1, x1, y1 - r - 1), fill_ink, 1) + elif x1 - r - 1 > x0 + r + 1: + self.draw.draw_rectangle((x0 + r + 1, y0, x1 - r - 1, y1), fill_ink, 1) + if not full_x and not full_y: + left = [x0, y0, x0 + r, y1] + if corners[0]: + left[1] += r + 1 + if corners[3]: + left[3] -= r + 1 + self.draw.draw_rectangle(left, fill_ink, 1) + + right = [x1 - r, y0, x1, y1] + if corners[1]: + right[1] += r + 1 + if corners[2]: + right[3] -= r + 1 + self.draw.draw_rectangle(right, fill_ink, 1) + if ink is not None and ink != fill_ink and width != 0: + draw_corners(False) + + if not full_x: + top = [x0, y0, x1, y0 + width - 1] + if corners[0]: + top[0] += r + 1 + if corners[1]: + top[2] -= r + 1 + self.draw.draw_rectangle(top, ink, 1) + + bottom = [x0, y1 - width + 1, x1, y1] + if corners[3]: + bottom[0] += r + 1 + if corners[2]: + bottom[2] -= r + 1 + self.draw.draw_rectangle(bottom, ink, 1) + if not full_y: + left = [x0, y0, x0 + width - 1, y1] + if corners[0]: + left[1] += r + 1 + if corners[3]: + left[3] -= r + 1 + self.draw.draw_rectangle(left, ink, 1) + + right = [x1 - width + 1, y0, x1, y1] + if corners[1]: + right[1] += r + 1 + if corners[2]: + right[3] -= r + 1 + self.draw.draw_rectangle(right, ink, 1) + + def text( + self, + xy: tuple[float, float], + text: AnyStr | ImageText.Text, + fill: _Ink | None = None, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + anchor: str | None = None, + spacing: float = 4, + align: str = "left", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + stroke_fill: _Ink | None = None, + embedded_color: bool = False, + *args: Any, + **kwargs: Any, + ) -> None: + """Draw text.""" + if isinstance(text, ImageText.Text): + image_text = text + else: + if font is None: + font = self._getfont(kwargs.get("font_size")) + image_text = ImageText.Text( + text, font, self.mode, spacing, direction, features, language + ) + if embedded_color: + image_text.embed_color() + if stroke_width: + image_text.stroke(stroke_width, stroke_fill) + + def getink(fill: _Ink | None) -> int: + ink, fill_ink = self._getink(fill) + if ink is None: + assert fill_ink is not None + return fill_ink + return ink + + ink = getink(fill) + if ink is None: + return + + stroke_ink = None + if image_text.stroke_width: + stroke_ink = ( + getink(image_text.stroke_fill) + if image_text.stroke_fill is not None + else ink + ) + + for xy, anchor, line in image_text._split(xy, anchor, align): + + def draw_text(ink: int, stroke_width: float = 0) -> None: + mode = self.fontmode + if stroke_width == 0 and embedded_color: + mode = "RGBA" + coord = [] + for i in range(2): + coord.append(int(xy[i])) + start = (math.modf(xy[0])[0], math.modf(xy[1])[0]) + try: + mask, offset = image_text.font.getmask2( # type: ignore[union-attr,misc] + line, + mode, + direction=direction, + features=features, + language=language, + stroke_width=stroke_width, + stroke_filled=True, + anchor=anchor, + ink=ink, + start=start, + *args, + **kwargs, + ) + coord = [coord[0] + offset[0], coord[1] + offset[1]] + except AttributeError: + try: + mask = image_text.font.getmask( # type: ignore[misc] + line, + mode, + direction, + features, + language, + stroke_width, + anchor, + ink, + start=start, + *args, + **kwargs, + ) + except TypeError: + mask = image_text.font.getmask(line) + if mode == "RGBA": + # image_text.font.getmask2(mode="RGBA") + # returns color in RGB bands and mask in A + # extract mask and set text alpha + color, mask = mask, mask.getband(3) + ink_alpha = struct.pack("i", ink)[3] + color.fillband(3, ink_alpha) + x, y = coord + if self.im is not None: + self.im.paste( + color, (x, y, x + mask.size[0], y + mask.size[1]), mask + ) + else: + self.draw.draw_bitmap(coord, mask, ink) + + if stroke_ink is not None: + # Draw stroked text + draw_text(stroke_ink, image_text.stroke_width) + + # Draw normal text + if ink != stroke_ink: + draw_text(ink) + else: + # Only draw normal text + draw_text(ink) + + def multiline_text( + self, + xy: tuple[float, float], + text: AnyStr, + fill: _Ink | None = None, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + anchor: str | None = None, + spacing: float = 4, + align: str = "left", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + stroke_fill: _Ink | None = None, + embedded_color: bool = False, + *, + font_size: float | None = None, + ) -> None: + return self.text( + xy, + text, + fill, + font, + anchor, + spacing, + align, + direction, + features, + language, + stroke_width, + stroke_fill, + embedded_color, + font_size=font_size, + ) + + def textlength( + self, + text: AnyStr, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + embedded_color: bool = False, + *, + font_size: float | None = None, + ) -> float: + """Get the length of a given string, in pixels with 1/64 precision.""" + if font is None: + font = self._getfont(font_size) + image_text = ImageText.Text( + text, + font, + self.mode, + direction=direction, + features=features, + language=language, + ) + if embedded_color: + image_text.embed_color() + return image_text.get_length() + + def textbbox( + self, + xy: tuple[float, float], + text: AnyStr, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + anchor: str | None = None, + spacing: float = 4, + align: str = "left", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + embedded_color: bool = False, + *, + font_size: float | None = None, + ) -> tuple[float, float, float, float]: + """Get the bounding box of a given string, in pixels.""" + if font is None: + font = self._getfont(font_size) + image_text = ImageText.Text( + text, font, self.mode, spacing, direction, features, language + ) + if embedded_color: + image_text.embed_color() + if stroke_width: + image_text.stroke(stroke_width) + return image_text.get_bbox(xy, anchor, align) + + def multiline_textbbox( + self, + xy: tuple[float, float], + text: AnyStr, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + anchor: str | None = None, + spacing: float = 4, + align: str = "left", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + embedded_color: bool = False, + *, + font_size: float | None = None, + ) -> tuple[float, float, float, float]: + return self.textbbox( + xy, + text, + font, + anchor, + spacing, + align, + direction, + features, + language, + stroke_width, + embedded_color, + font_size=font_size, + ) + + +def Draw(im: Image.Image, mode: str | None = None) -> ImageDraw: + """ + A simple 2D drawing interface for PIL images. + + :param im: The image to draw in. + :param mode: Optional mode to use for color values. For RGB + images, this argument can be RGB or RGBA (to blend the + drawing into the image). For all other modes, this argument + must be the same as the image mode. If omitted, the mode + defaults to the mode of the image. + """ + try: + return getattr(im, "getdraw")(mode) + except AttributeError: + return ImageDraw(im, mode) + + +def getdraw(im: Image.Image | None = None) -> tuple[ImageDraw2.Draw | None, ModuleType]: + """ + :param im: The image to draw in. + :returns: A (drawing context, drawing resource factory) tuple. + """ + from . import ImageDraw2 + + draw = ImageDraw2.Draw(im) if im is not None else None + return draw, ImageDraw2 + + +def floodfill( + image: Image.Image, + xy: tuple[int, int], + value: float | tuple[int, ...], + border: float | tuple[int, ...] | None = None, + thresh: float = 0, +) -> None: + """ + .. warning:: This method is experimental. + + Fills a bounded region with a given color. + + :param image: Target image. + :param xy: Seed position (a 2-item coordinate tuple). See + :ref:`coordinate-system`. + :param value: Fill color. + :param border: Optional border value. If given, the region consists of + pixels with a color different from the border color. If not given, + the region consists of pixels having the same color as the seed + pixel. + :param thresh: Optional threshold value which specifies a maximum + tolerable difference of a pixel value from the 'background' in + order for it to be replaced. Useful for filling regions of + non-homogeneous, but similar, colors. + """ + # based on an implementation by Eric S. Raymond + # amended by yo1995 @20180806 + pixel = image.load() + assert pixel is not None + x, y = xy + try: + background = pixel[x, y] + if _color_diff(value, background) <= thresh: + return # seed point already has fill color + pixel[x, y] = value + except (ValueError, IndexError): + return # seed point outside image + edge = {(x, y)} + # use a set to keep record of current and previous edge pixels + # to reduce memory consumption + full_edge = set() + while edge: + new_edge = set() + for x, y in edge: # 4 adjacent method + for s, t in ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)): + # If already processed, or if a coordinate is negative, skip + if (s, t) in full_edge or s < 0 or t < 0: + continue + try: + p = pixel[s, t] + except (ValueError, IndexError): + pass + else: + full_edge.add((s, t)) + if border is None: + fill = _color_diff(p, background) <= thresh + else: + fill = p not in (value, border) + if fill: + pixel[s, t] = value + new_edge.add((s, t)) + full_edge = edge # discard pixels processed + edge = new_edge + + +def _compute_regular_polygon_vertices( + bounding_circle: Sequence[Sequence[float] | float], n_sides: int, rotation: float +) -> list[tuple[float, float]]: + """ + Generate a list of vertices for a 2D regular polygon. + + :param bounding_circle: The bounding circle is a sequence defined + by a point and radius. The polygon is inscribed in this circle. + (e.g. ``bounding_circle=(x, y, r)`` or ``((x, y), r)``) + :param n_sides: Number of sides + (e.g. ``n_sides=3`` for a triangle, ``6`` for a hexagon) + :param rotation: Apply an arbitrary rotation to the polygon + (e.g. ``rotation=90``, applies a 90 degree rotation) + :return: List of regular polygon vertices + (e.g. ``[(25, 50), (50, 50), (50, 25), (25, 25)]``) + + How are the vertices computed? + 1. Compute the following variables + - theta: Angle between the apothem & the nearest polygon vertex + - side_length: Length of each polygon edge + - centroid: Center of bounding circle (1st, 2nd elements of bounding_circle) + - polygon_radius: Polygon radius (last element of bounding_circle) + - angles: Location of each polygon vertex in polar grid + (e.g. A square with 0 degree rotation => [225.0, 315.0, 45.0, 135.0]) + + 2. For each angle in angles, get the polygon vertex at that angle + The vertex is computed using the equation below. + X= xcos(φ) + ysin(φ) + Y= −xsin(φ) + ycos(φ) + + Note: + φ = angle in degrees + x = 0 + y = polygon_radius + + The formula above assumes rotation around the origin. + In our case, we are rotating around the centroid. + To account for this, we use the formula below + X = xcos(φ) + ysin(φ) + centroid_x + Y = −xsin(φ) + ycos(φ) + centroid_y + """ + # 1. Error Handling + # 1.1 Check `n_sides` has an appropriate value + if not isinstance(n_sides, int): + msg = "n_sides should be an int" # type: ignore[unreachable] + raise TypeError(msg) + if n_sides < 3: + msg = "n_sides should be an int > 2" + raise ValueError(msg) + + # 1.2 Check `bounding_circle` has an appropriate value + if not isinstance(bounding_circle, (list, tuple)): + msg = "bounding_circle should be a sequence" + raise TypeError(msg) + + if len(bounding_circle) == 3: + if not all(isinstance(i, (int, float)) for i in bounding_circle): + msg = "bounding_circle should only contain numeric data" + raise ValueError(msg) + + *centroid, polygon_radius = cast(list[float], list(bounding_circle)) + elif len(bounding_circle) == 2 and isinstance(bounding_circle[0], (list, tuple)): + if not all( + isinstance(i, (int, float)) for i in bounding_circle[0] + ) or not isinstance(bounding_circle[1], (int, float)): + msg = "bounding_circle should only contain numeric data" + raise ValueError(msg) + + if len(bounding_circle[0]) != 2: + msg = "bounding_circle centre should contain 2D coordinates (e.g. (x, y))" + raise ValueError(msg) + + centroid = cast(list[float], list(bounding_circle[0])) + polygon_radius = cast(float, bounding_circle[1]) + else: + msg = ( + "bounding_circle should contain 2D coordinates " + "and a radius (e.g. (x, y, r) or ((x, y), r) )" + ) + raise ValueError(msg) + + if polygon_radius <= 0: + msg = "bounding_circle radius should be > 0" + raise ValueError(msg) + + # 1.3 Check `rotation` has an appropriate value + if not isinstance(rotation, (int, float)): + msg = "rotation should be an int or float" # type: ignore[unreachable] + raise ValueError(msg) + + # 2. Define Helper Functions + def _apply_rotation(point: list[float], degrees: float) -> tuple[float, float]: + return ( + round( + point[0] * math.cos(math.radians(360 - degrees)) + - point[1] * math.sin(math.radians(360 - degrees)) + + centroid[0], + 2, + ), + round( + point[1] * math.cos(math.radians(360 - degrees)) + + point[0] * math.sin(math.radians(360 - degrees)) + + centroid[1], + 2, + ), + ) + + def _compute_polygon_vertex(angle: float) -> tuple[float, float]: + start_point = [polygon_radius, 0] + return _apply_rotation(start_point, angle) + + def _get_angles(n_sides: int, rotation: float) -> list[float]: + angles = [] + degrees = 360 / n_sides + # Start with the bottom left polygon vertex + current_angle = (270 - 0.5 * degrees) + rotation + for _ in range(n_sides): + angles.append(current_angle) + current_angle += degrees + if current_angle > 360: + current_angle -= 360 + return angles + + # 3. Variable Declarations + angles = _get_angles(n_sides, rotation) + + # 4. Compute Vertices + return [_compute_polygon_vertex(angle) for angle in angles] + + +def _color_diff( + color1: float | tuple[int, ...], color2: float | tuple[int, ...] +) -> float: + """ + Uses 1-norm distance to calculate difference between two values. + """ + first = color1 if isinstance(color1, tuple) else (color1,) + second = color2 if isinstance(color2, tuple) else (color2,) + + return sum(abs(first[i] - second[i]) for i in range(len(second))) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageDraw2.py b/venv/lib/python3.10/site-packages/PIL/ImageDraw2.py new file mode 100644 index 0000000000000000000000000000000000000000..3d68658ed5b79a36597e4953b888c41aa82fc7da --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageDraw2.py @@ -0,0 +1,243 @@ +# +# The Python Imaging Library +# $Id$ +# +# WCK-style drawing interface operations +# +# History: +# 2003-12-07 fl created +# 2005-05-15 fl updated; added to PIL as ImageDraw2 +# 2005-05-15 fl added text support +# 2005-05-20 fl added arc/chord/pieslice support +# +# Copyright (c) 2003-2005 by Secret Labs AB +# Copyright (c) 2003-2005 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# + + +""" +(Experimental) WCK-style drawing interface operations + +.. seealso:: :py:mod:`PIL.ImageDraw` +""" +from __future__ import annotations + +from typing import Any, AnyStr, BinaryIO + +from . import Image, ImageColor, ImageDraw, ImageFont, ImagePath +from ._typing import Coords, StrOrBytesPath + + +class Pen: + """Stores an outline color and width.""" + + def __init__(self, color: str, width: int = 1, opacity: int = 255) -> None: + self.color = ImageColor.getrgb(color) + self.width = width + + +class Brush: + """Stores a fill color""" + + def __init__(self, color: str, opacity: int = 255) -> None: + self.color = ImageColor.getrgb(color) + + +class Font: + """Stores a TrueType font and color""" + + def __init__( + self, color: str, file: StrOrBytesPath | BinaryIO, size: float = 12 + ) -> None: + # FIXME: add support for bitmap fonts + self.color = ImageColor.getrgb(color) + self.font = ImageFont.truetype(file, size) + + +class Draw: + """ + (Experimental) WCK-style drawing interface + """ + + def __init__( + self, + image: Image.Image | str, + size: tuple[int, int] | list[int] | None = None, + color: float | tuple[float, ...] | str | None = None, + ) -> None: + if isinstance(image, str): + if size is None: + msg = "If image argument is mode string, size must be a list or tuple" + raise ValueError(msg) + image = Image.new(image, size, color) + self.draw = ImageDraw.Draw(image) + self.image = image + self.transform: tuple[float, float, float, float, float, float] | None = None + + def flush(self) -> Image.Image: + return self.image + + def render( + self, + op: str, + xy: Coords, + pen: Pen | Brush | None, + brush: Brush | Pen | None = None, + **kwargs: Any, + ) -> None: + # handle color arguments + outline = fill = None + width = 1 + if isinstance(pen, Pen): + outline = pen.color + width = pen.width + elif isinstance(brush, Pen): + outline = brush.color + width = brush.width + if isinstance(brush, Brush): + fill = brush.color + elif isinstance(pen, Brush): + fill = pen.color + # handle transformation + if self.transform: + path = ImagePath.Path(xy) + path.transform(self.transform) + xy = path + # render the item + if op in ("arc", "line"): + kwargs.setdefault("fill", outline) + else: + kwargs.setdefault("fill", fill) + kwargs.setdefault("outline", outline) + if op == "line": + kwargs.setdefault("width", width) + getattr(self.draw, op)(xy, **kwargs) + + def settransform(self, offset: tuple[float, float]) -> None: + """Sets a transformation offset.""" + (xoffset, yoffset) = offset + self.transform = (1, 0, xoffset, 0, 1, yoffset) + + def arc( + self, + xy: Coords, + pen: Pen | Brush | None, + start: float, + end: float, + *options: Any, + ) -> None: + """ + Draws an arc (a portion of a circle outline) between the start and end + angles, inside the given bounding box. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.arc` + """ + self.render("arc", xy, pen, *options, start=start, end=end) + + def chord( + self, + xy: Coords, + pen: Pen | Brush | None, + start: float, + end: float, + *options: Any, + ) -> None: + """ + Same as :py:meth:`~PIL.ImageDraw2.Draw.arc`, but connects the end points + with a straight line. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.chord` + """ + self.render("chord", xy, pen, *options, start=start, end=end) + + def ellipse(self, xy: Coords, pen: Pen | Brush | None, *options: Any) -> None: + """ + Draws an ellipse inside the given bounding box. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.ellipse` + """ + self.render("ellipse", xy, pen, *options) + + def line(self, xy: Coords, pen: Pen | Brush | None, *options: Any) -> None: + """ + Draws a line between the coordinates in the ``xy`` list. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.line` + """ + self.render("line", xy, pen, *options) + + def pieslice( + self, + xy: Coords, + pen: Pen | Brush | None, + start: float, + end: float, + *options: Any, + ) -> None: + """ + Same as arc, but also draws straight lines between the end points and the + center of the bounding box. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.pieslice` + """ + self.render("pieslice", xy, pen, *options, start=start, end=end) + + def polygon(self, xy: Coords, pen: Pen | Brush | None, *options: Any) -> None: + """ + Draws a polygon. + + The polygon outline consists of straight lines between the given + coordinates, plus a straight line between the last and the first + coordinate. + + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.polygon` + """ + self.render("polygon", xy, pen, *options) + + def rectangle(self, xy: Coords, pen: Pen | Brush | None, *options: Any) -> None: + """ + Draws a rectangle. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.rectangle` + """ + self.render("rectangle", xy, pen, *options) + + def text(self, xy: tuple[float, float], text: AnyStr, font: Font) -> None: + """ + Draws the string at the given position. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.text` + """ + if self.transform: + path = ImagePath.Path(xy) + path.transform(self.transform) + xy = path + self.draw.text(xy, text, font=font.font, fill=font.color) + + def textbbox( + self, xy: tuple[float, float], text: AnyStr, font: Font + ) -> tuple[float, float, float, float]: + """ + Returns bounding box (in pixels) of given text. + + :return: ``(left, top, right, bottom)`` bounding box + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.textbbox` + """ + if self.transform: + path = ImagePath.Path(xy) + path.transform(self.transform) + xy = path + return self.draw.textbbox(xy, text, font=font.font) + + def textlength(self, text: AnyStr, font: Font) -> float: + """ + Returns length (in pixels) of given text. + This is the amount by which following text should be offset. + + .. seealso:: :py:meth:`PIL.ImageDraw.ImageDraw.textlength` + """ + return self.draw.textlength(text, font=font.font) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageEnhance.py b/venv/lib/python3.10/site-packages/PIL/ImageEnhance.py new file mode 100644 index 0000000000000000000000000000000000000000..0e7e6dd8ae631ad3577bda1d3e823bd2a3227536 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageEnhance.py @@ -0,0 +1,113 @@ +# +# The Python Imaging Library. +# $Id$ +# +# image enhancement classes +# +# For a background, see "Image Processing By Interpolation and +# Extrapolation", Paul Haeberli and Douglas Voorhies. Available +# at http://www.graficaobscura.com/interp/index.html +# +# History: +# 1996-03-23 fl Created +# 2009-06-16 fl Fixed mean calculation +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image, ImageFilter, ImageStat + + +class _Enhance: + image: Image.Image + degenerate: Image.Image + + def enhance(self, factor: float) -> Image.Image: + """ + Returns an enhanced image. + + :param factor: A floating point value controlling the enhancement. + Factor 1.0 always returns a copy of the original image, + lower factors mean less color (brightness, contrast, + etc), and higher values more. There are no restrictions + on this value. + :rtype: :py:class:`~PIL.Image.Image` + """ + return Image.blend(self.degenerate, self.image, factor) + + +class Color(_Enhance): + """Adjust image color balance. + + This class can be used to adjust the colour balance of an image, in + a manner similar to the controls on a colour TV set. An enhancement + factor of 0.0 gives a black and white image. A factor of 1.0 gives + the original image. + """ + + def __init__(self, image: Image.Image) -> None: + self.image = image + self.intermediate_mode = "L" + if "A" in image.getbands(): + self.intermediate_mode = "LA" + + if self.intermediate_mode != image.mode: + image = image.convert(self.intermediate_mode).convert(image.mode) + self.degenerate = image + + +class Contrast(_Enhance): + """Adjust image contrast. + + This class can be used to control the contrast of an image, similar + to the contrast control on a TV set. An enhancement factor of 0.0 + gives a solid gray image. A factor of 1.0 gives the original image. + """ + + def __init__(self, image: Image.Image) -> None: + self.image = image + if image.mode != "L": + image = image.convert("L") + mean = int(ImageStat.Stat(image).mean[0] + 0.5) + self.degenerate = Image.new("L", image.size, mean) + if self.degenerate.mode != self.image.mode: + self.degenerate = self.degenerate.convert(self.image.mode) + + if "A" in self.image.getbands(): + self.degenerate.putalpha(self.image.getchannel("A")) + + +class Brightness(_Enhance): + """Adjust image brightness. + + This class can be used to control the brightness of an image. An + enhancement factor of 0.0 gives a black image. A factor of 1.0 gives the + original image. + """ + + def __init__(self, image: Image.Image) -> None: + self.image = image + self.degenerate = Image.new(image.mode, image.size, 0) + + if "A" in image.getbands(): + self.degenerate.putalpha(image.getchannel("A")) + + +class Sharpness(_Enhance): + """Adjust image sharpness. + + This class can be used to adjust the sharpness of an image. An + enhancement factor of 0.0 gives a blurred image, a factor of 1.0 gives the + original image, and a factor of 2.0 gives a sharpened image. + """ + + def __init__(self, image: Image.Image) -> None: + self.image = image + self.degenerate = image.filter(ImageFilter.SMOOTH) + + if "A" in image.getbands(): + self.degenerate.putalpha(image.getchannel("A")) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageFile.py b/venv/lib/python3.10/site-packages/PIL/ImageFile.py new file mode 100644 index 0000000000000000000000000000000000000000..a1d98bd5103c181c682e4fafe8cb2c3093f21586 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageFile.py @@ -0,0 +1,926 @@ +# +# The Python Imaging Library. +# $Id$ +# +# base class for image file handlers +# +# history: +# 1995-09-09 fl Created +# 1996-03-11 fl Fixed load mechanism. +# 1996-04-15 fl Added pcx/xbm decoders. +# 1996-04-30 fl Added encoders. +# 1996-12-14 fl Added load helpers +# 1997-01-11 fl Use encode_to_file where possible +# 1997-08-27 fl Flush output in _save +# 1998-03-05 fl Use memory mapping for some modes +# 1999-02-04 fl Use memory mapping also for "I;16" and "I;16B" +# 1999-05-31 fl Added image parser +# 2000-10-12 fl Set readonly flag on memory-mapped images +# 2002-03-20 fl Use better messages for common decoder errors +# 2003-04-21 fl Fall back on mmap/map_buffer if map is not available +# 2003-10-30 fl Added StubImageFile class +# 2004-02-25 fl Made incremental parser more robust +# +# Copyright (c) 1997-2004 by Secret Labs AB +# Copyright (c) 1995-2004 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import abc +import io +import itertools +import logging +import os +import struct +from typing import IO, Any, NamedTuple, cast + +from . import ExifTags, Image +from ._util import DeferredError, is_path + +TYPE_CHECKING = False +if TYPE_CHECKING: + from ._typing import StrOrBytesPath + +logger = logging.getLogger(__name__) + +MAXBLOCK = 65536 +""" +By default, Pillow processes image data in blocks. This helps to prevent excessive use +of resources. Codecs may disable this behaviour with ``_pulls_fd`` or ``_pushes_fd``. + +When reading an image, this is the number of bytes to read at once. + +When writing an image, this is the number of bytes to write at once. +If the image width times 4 is greater, then that will be used instead. +Plugins may also set a greater number. + +User code may set this to another number. +""" + +SAFEBLOCK = 1024 * 1024 + +LOAD_TRUNCATED_IMAGES = False +"""Whether or not to load truncated image files. User code may change this.""" + +ERRORS = { + -1: "image buffer overrun error", + -2: "decoding error", + -3: "unknown error", + -8: "bad configuration", + -9: "out of memory error", +} +""" +Dict of known error codes returned from :meth:`.PyDecoder.decode`, +:meth:`.PyEncoder.encode` :meth:`.PyEncoder.encode_to_pyfd` and +:meth:`.PyEncoder.encode_to_file`. +""" + + +# +# -------------------------------------------------------------------- +# Helpers + + +def _get_oserror(error: int, *, encoder: bool) -> OSError: + try: + msg = Image.core.getcodecstatus(error) + except AttributeError: + msg = ERRORS.get(error) + if not msg: + msg = f"{'encoder' if encoder else 'decoder'} error {error}" + msg += f" when {'writing' if encoder else 'reading'} image file" + return OSError(msg) + + +def _tilesort(t: _Tile) -> int: + # sort on offset + return t[2] + + +class _Tile(NamedTuple): + codec_name: str + extents: tuple[int, int, int, int] | None + offset: int = 0 + args: tuple[Any, ...] | str | None = None + + +# +# -------------------------------------------------------------------- +# ImageFile base class + + +class ImageFile(Image.Image): + """Base class for image file format handlers.""" + + def __init__( + self, fp: StrOrBytesPath | IO[bytes], filename: str | bytes | None = None + ) -> None: + super().__init__() + + self._min_frame = 0 + + self.custom_mimetype: str | None = None + + self.tile: list[_Tile] = [] + """ A list of tile descriptors """ + + self.readonly = 1 # until we know better + + self.decoderconfig: tuple[Any, ...] = () + self.decodermaxblock = MAXBLOCK + + if is_path(fp): + # filename + self.fp = open(fp, "rb") + self.filename = os.fspath(fp) + self._exclusive_fp = True + else: + # stream + self.fp = cast(IO[bytes], fp) + self.filename = filename if filename is not None else "" + # can be overridden + self._exclusive_fp = False + + try: + try: + self._open() + except ( + IndexError, # end of data + TypeError, # end of data (ord) + KeyError, # unsupported mode + EOFError, # got header but not the first frame + struct.error, + ) as v: + raise SyntaxError(v) from v + + if not self.mode or self.size[0] <= 0 or self.size[1] <= 0: + msg = "not identified by this driver" + raise SyntaxError(msg) + except BaseException: + # close the file only if we have opened it this constructor + if self._exclusive_fp: + self.fp.close() + raise + + def _open(self) -> None: + pass + + def _close_fp(self): + if getattr(self, "_fp", False) and not isinstance(self._fp, DeferredError): + if self._fp != self.fp: + self._fp.close() + self._fp = DeferredError(ValueError("Operation on closed image")) + if self.fp: + self.fp.close() + + def close(self) -> None: + """ + Closes the file pointer, if possible. + + This operation will destroy the image core and release its memory. + The image data will be unusable afterward. + + This function is required to close images that have multiple frames or + have not had their file read and closed by the + :py:meth:`~PIL.Image.Image.load` method. See :ref:`file-handling` for + more information. + """ + try: + self._close_fp() + self.fp = None + except Exception as msg: + logger.debug("Error closing: %s", msg) + + super().close() + + def get_child_images(self) -> list[ImageFile]: + child_images = [] + exif = self.getexif() + ifds = [] + if ExifTags.Base.SubIFDs in exif: + subifd_offsets = exif[ExifTags.Base.SubIFDs] + if subifd_offsets: + if not isinstance(subifd_offsets, tuple): + subifd_offsets = (subifd_offsets,) + for subifd_offset in subifd_offsets: + ifds.append((exif._get_ifd_dict(subifd_offset), subifd_offset)) + ifd1 = exif.get_ifd(ExifTags.IFD.IFD1) + if ifd1 and ifd1.get(ExifTags.Base.JpegIFOffset): + assert exif._info is not None + ifds.append((ifd1, exif._info.next)) + + offset = None + for ifd, ifd_offset in ifds: + assert self.fp is not None + current_offset = self.fp.tell() + if offset is None: + offset = current_offset + + fp = self.fp + if ifd is not None: + thumbnail_offset = ifd.get(ExifTags.Base.JpegIFOffset) + if thumbnail_offset is not None: + thumbnail_offset += getattr(self, "_exif_offset", 0) + self.fp.seek(thumbnail_offset) + + length = ifd.get(ExifTags.Base.JpegIFByteCount) + assert isinstance(length, int) + data = self.fp.read(length) + fp = io.BytesIO(data) + + with Image.open(fp) as im: + from . import TiffImagePlugin + + if thumbnail_offset is None and isinstance( + im, TiffImagePlugin.TiffImageFile + ): + im._frame_pos = [ifd_offset] + im._seek(0) + im.load() + child_images.append(im) + + if offset is not None: + assert self.fp is not None + self.fp.seek(offset) + return child_images + + def get_format_mimetype(self) -> str | None: + if self.custom_mimetype: + return self.custom_mimetype + if self.format is not None: + return Image.MIME.get(self.format.upper()) + return None + + def __getstate__(self) -> list[Any]: + return super().__getstate__() + [self.filename] + + def __setstate__(self, state: list[Any]) -> None: + self.tile = [] + if len(state) > 5: + self.filename = state[5] + super().__setstate__(state) + + def verify(self) -> None: + """Check file integrity""" + + # raise exception if something's wrong. must be called + # directly after open, and closes file when finished. + if self._exclusive_fp: + self.fp.close() + self.fp = None + + def load(self) -> Image.core.PixelAccess | None: + """Load image data based on tile list""" + + if not self.tile and self._im is None: + msg = "cannot load this image" + raise OSError(msg) + + pixel = Image.Image.load(self) + if not self.tile: + return pixel + + self.map: mmap.mmap | None = None + use_mmap = self.filename and len(self.tile) == 1 + + readonly = 0 + + # look for read/seek overrides + if hasattr(self, "load_read"): + read = self.load_read + # don't use mmap if there are custom read/seek functions + use_mmap = False + else: + read = self.fp.read + + if hasattr(self, "load_seek"): + seek = self.load_seek + use_mmap = False + else: + seek = self.fp.seek + + if use_mmap: + # try memory mapping + decoder_name, extents, offset, args = self.tile[0] + if isinstance(args, str): + args = (args, 0, 1) + if ( + decoder_name == "raw" + and isinstance(args, tuple) + and len(args) >= 3 + and args[0] == self.mode + and args[0] in Image._MAPMODES + ): + if offset < 0: + msg = "Tile offset cannot be negative" + raise ValueError(msg) + try: + # use mmap, if possible + import mmap + + with open(self.filename) as fp: + self.map = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ) + if offset + self.size[1] * args[1] > self.map.size(): + msg = "buffer is not large enough" + raise OSError(msg) + self.im = Image.core.map_buffer( + self.map, self.size, decoder_name, offset, args + ) + readonly = 1 + # After trashing self.im, + # we might need to reload the palette data. + if self.palette: + self.palette.dirty = 1 + except (AttributeError, OSError, ImportError): + self.map = None + + self.load_prepare() + err_code = -3 # initialize to unknown error + if not self.map: + # sort tiles in file order + self.tile.sort(key=_tilesort) + + # FIXME: This is a hack to handle TIFF's JpegTables tag. + prefix = getattr(self, "tile_prefix", b"") + + # Remove consecutive duplicates that only differ by their offset + self.tile = [ + list(tiles)[-1] + for _, tiles in itertools.groupby( + self.tile, lambda tile: (tile[0], tile[1], tile[3]) + ) + ] + for i, (decoder_name, extents, offset, args) in enumerate(self.tile): + seek(offset) + decoder = Image._getdecoder( + self.mode, decoder_name, args, self.decoderconfig + ) + try: + decoder.setimage(self.im, extents) + if decoder.pulls_fd: + decoder.setfd(self.fp) + err_code = decoder.decode(b"")[1] + else: + b = prefix + while True: + read_bytes = self.decodermaxblock + if i + 1 < len(self.tile): + next_offset = self.tile[i + 1].offset + if next_offset > offset: + read_bytes = next_offset - offset + try: + s = read(read_bytes) + except (IndexError, struct.error) as e: + # truncated png/gif + if LOAD_TRUNCATED_IMAGES: + break + else: + msg = "image file is truncated" + raise OSError(msg) from e + + if not s: # truncated jpeg + if LOAD_TRUNCATED_IMAGES: + break + else: + msg = ( + "image file is truncated " + f"({len(b)} bytes not processed)" + ) + raise OSError(msg) + + b = b + s + n, err_code = decoder.decode(b) + if n < 0: + break + b = b[n:] + finally: + # Need to cleanup here to prevent leaks + decoder.cleanup() + + self.tile = [] + self.readonly = readonly + + self.load_end() + + if self._exclusive_fp and self._close_exclusive_fp_after_loading: + self.fp.close() + self.fp = None + + if not self.map and not LOAD_TRUNCATED_IMAGES and err_code < 0: + # still raised if decoder fails to return anything + raise _get_oserror(err_code, encoder=False) + + return Image.Image.load(self) + + def load_prepare(self) -> None: + # create image memory if necessary + if self._im is None: + self.im = Image.core.new(self.mode, self.size) + # create palette (optional) + if self.mode == "P": + Image.Image.load(self) + + def load_end(self) -> None: + # may be overridden + pass + + # may be defined for contained formats + # def load_seek(self, pos: int) -> None: + # pass + + # may be defined for blocked formats (e.g. PNG) + # def load_read(self, read_bytes: int) -> bytes: + # pass + + def _seek_check(self, frame: int) -> bool: + if ( + frame < self._min_frame + # Only check upper limit on frames if additional seek operations + # are not required to do so + or ( + not (hasattr(self, "_n_frames") and self._n_frames is None) + and frame >= getattr(self, "n_frames") + self._min_frame + ) + ): + msg = "attempt to seek outside sequence" + raise EOFError(msg) + + return self.tell() != frame + + +class StubHandler(abc.ABC): + def open(self, im: StubImageFile) -> None: + pass + + @abc.abstractmethod + def load(self, im: StubImageFile) -> Image.Image: + pass + + +class StubImageFile(ImageFile, metaclass=abc.ABCMeta): + """ + Base class for stub image loaders. + + A stub loader is an image loader that can identify files of a + certain format, but relies on external code to load the file. + """ + + @abc.abstractmethod + def _open(self) -> None: + pass + + def load(self) -> Image.core.PixelAccess | None: + loader = self._load() + if loader is None: + msg = f"cannot find loader for this {self.format} file" + raise OSError(msg) + image = loader.load(self) + assert image is not None + # become the other object (!) + self.__class__ = image.__class__ # type: ignore[assignment] + self.__dict__ = image.__dict__ + return image.load() + + @abc.abstractmethod + def _load(self) -> StubHandler | None: + """(Hook) Find actual image loader.""" + pass + + +class Parser: + """ + Incremental image parser. This class implements the standard + feed/close consumer interface. + """ + + incremental = None + image: Image.Image | None = None + data: bytes | None = None + decoder: Image.core.ImagingDecoder | PyDecoder | None = None + offset = 0 + finished = 0 + + def reset(self) -> None: + """ + (Consumer) Reset the parser. Note that you can only call this + method immediately after you've created a parser; parser + instances cannot be reused. + """ + assert self.data is None, "cannot reuse parsers" + + def feed(self, data: bytes) -> None: + """ + (Consumer) Feed data to the parser. + + :param data: A string buffer. + :exception OSError: If the parser failed to parse the image file. + """ + # collect data + + if self.finished: + return + + if self.data is None: + self.data = data + else: + self.data = self.data + data + + # parse what we have + if self.decoder: + if self.offset > 0: + # skip header + skip = min(len(self.data), self.offset) + self.data = self.data[skip:] + self.offset = self.offset - skip + if self.offset > 0 or not self.data: + return + + n, e = self.decoder.decode(self.data) + + if n < 0: + # end of stream + self.data = None + self.finished = 1 + if e < 0: + # decoding error + self.image = None + raise _get_oserror(e, encoder=False) + else: + # end of image + return + self.data = self.data[n:] + + elif self.image: + # if we end up here with no decoder, this file cannot + # be incrementally parsed. wait until we've gotten all + # available data + pass + + else: + # attempt to open this file + try: + with io.BytesIO(self.data) as fp: + im = Image.open(fp) + except OSError: + pass # not enough data + else: + flag = hasattr(im, "load_seek") or hasattr(im, "load_read") + if flag or len(im.tile) != 1: + # custom load code, or multiple tiles + self.decode = None + else: + # initialize decoder + im.load_prepare() + d, e, o, a = im.tile[0] + im.tile = [] + self.decoder = Image._getdecoder(im.mode, d, a, im.decoderconfig) + self.decoder.setimage(im.im, e) + + # calculate decoder offset + self.offset = o + if self.offset <= len(self.data): + self.data = self.data[self.offset :] + self.offset = 0 + + self.image = im + + def __enter__(self) -> Parser: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def close(self) -> Image.Image: + """ + (Consumer) Close the stream. + + :returns: An image object. + :exception OSError: If the parser failed to parse the image file either + because it cannot be identified or cannot be + decoded. + """ + # finish decoding + if self.decoder: + # get rid of what's left in the buffers + self.feed(b"") + self.data = self.decoder = None + if not self.finished: + msg = "image was incomplete" + raise OSError(msg) + if not self.image: + msg = "cannot parse this image" + raise OSError(msg) + if self.data: + # incremental parsing not possible; reopen the file + # not that we have all data + with io.BytesIO(self.data) as fp: + try: + self.image = Image.open(fp) + finally: + self.image.load() + return self.image + + +# -------------------------------------------------------------------- + + +def _save(im: Image.Image, fp: IO[bytes], tile: list[_Tile], bufsize: int = 0) -> None: + """Helper to save image based on tile list + + :param im: Image object. + :param fp: File object. + :param tile: Tile list. + :param bufsize: Optional buffer size + """ + + im.load() + if not hasattr(im, "encoderconfig"): + im.encoderconfig = () + tile.sort(key=_tilesort) + # FIXME: make MAXBLOCK a configuration parameter + # It would be great if we could have the encoder specify what it needs + # But, it would need at least the image size in most cases. RawEncode is + # a tricky case. + bufsize = max(MAXBLOCK, bufsize, im.size[0] * 4) # see RawEncode.c + try: + fh = fp.fileno() + fp.flush() + _encode_tile(im, fp, tile, bufsize, fh) + except (AttributeError, io.UnsupportedOperation) as exc: + _encode_tile(im, fp, tile, bufsize, None, exc) + if hasattr(fp, "flush"): + fp.flush() + + +def _encode_tile( + im: Image.Image, + fp: IO[bytes], + tile: list[_Tile], + bufsize: int, + fh: int | None, + exc: BaseException | None = None, +) -> None: + for encoder_name, extents, offset, args in tile: + if offset > 0: + fp.seek(offset) + encoder = Image._getencoder(im.mode, encoder_name, args, im.encoderconfig) + try: + encoder.setimage(im.im, extents) + if encoder.pushes_fd: + encoder.setfd(fp) + errcode = encoder.encode_to_pyfd()[1] + else: + if exc: + # compress to Python file-compatible object + while True: + errcode, data = encoder.encode(bufsize)[1:] + fp.write(data) + if errcode: + break + else: + # slight speedup: compress to real file object + assert fh is not None + errcode = encoder.encode_to_file(fh, bufsize) + if errcode < 0: + raise _get_oserror(errcode, encoder=True) from exc + finally: + encoder.cleanup() + + +def _safe_read(fp: IO[bytes], size: int) -> bytes: + """ + Reads large blocks in a safe way. Unlike fp.read(n), this function + doesn't trust the user. If the requested size is larger than + SAFEBLOCK, the file is read block by block. + + :param fp: File handle. Must implement a read method. + :param size: Number of bytes to read. + :returns: A string containing size bytes of data. + + Raises an OSError if the file is truncated and the read cannot be completed + + """ + if size <= 0: + return b"" + if size <= SAFEBLOCK: + data = fp.read(size) + if len(data) < size: + msg = "Truncated File Read" + raise OSError(msg) + return data + blocks: list[bytes] = [] + remaining_size = size + while remaining_size > 0: + block = fp.read(min(remaining_size, SAFEBLOCK)) + if not block: + break + blocks.append(block) + remaining_size -= len(block) + if sum(len(block) for block in blocks) < size: + msg = "Truncated File Read" + raise OSError(msg) + return b"".join(blocks) + + +class PyCodecState: + def __init__(self) -> None: + self.xsize = 0 + self.ysize = 0 + self.xoff = 0 + self.yoff = 0 + + def extents(self) -> tuple[int, int, int, int]: + return self.xoff, self.yoff, self.xoff + self.xsize, self.yoff + self.ysize + + +class PyCodec: + fd: IO[bytes] | None + + def __init__(self, mode: str, *args: Any) -> None: + self.im: Image.core.ImagingCore | None = None + self.state = PyCodecState() + self.fd = None + self.mode = mode + self.init(args) + + def init(self, args: tuple[Any, ...]) -> None: + """ + Override to perform codec specific initialization + + :param args: Tuple of arg items from the tile entry + :returns: None + """ + self.args = args + + def cleanup(self) -> None: + """ + Override to perform codec specific cleanup + + :returns: None + """ + pass + + def setfd(self, fd: IO[bytes]) -> None: + """ + Called from ImageFile to set the Python file-like object + + :param fd: A Python file-like object + :returns: None + """ + self.fd = fd + + def setimage( + self, + im: Image.core.ImagingCore, + extents: tuple[int, int, int, int] | None = None, + ) -> None: + """ + Called from ImageFile to set the core output image for the codec + + :param im: A core image object + :param extents: a 4 tuple of (x0, y0, x1, y1) defining the rectangle + for this tile + :returns: None + """ + + # following c code + self.im = im + + if extents: + (x0, y0, x1, y1) = extents + else: + (x0, y0, x1, y1) = (0, 0, 0, 0) + + if x0 == 0 and x1 == 0: + self.state.xsize, self.state.ysize = self.im.size + else: + self.state.xoff = x0 + self.state.yoff = y0 + self.state.xsize = x1 - x0 + self.state.ysize = y1 - y0 + + if self.state.xsize <= 0 or self.state.ysize <= 0: + msg = "Size cannot be negative" + raise ValueError(msg) + + if ( + self.state.xsize + self.state.xoff > self.im.size[0] + or self.state.ysize + self.state.yoff > self.im.size[1] + ): + msg = "Tile cannot extend outside image" + raise ValueError(msg) + + +class PyDecoder(PyCodec): + """ + Python implementation of a format decoder. Override this class and + add the decoding logic in the :meth:`decode` method. + + See :ref:`Writing Your Own File Codec in Python` + """ + + _pulls_fd = False + + @property + def pulls_fd(self) -> bool: + return self._pulls_fd + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + """ + Override to perform the decoding process. + + :param buffer: A bytes object with the data to be decoded. + :returns: A tuple of ``(bytes consumed, errcode)``. + If finished with decoding return -1 for the bytes consumed. + Err codes are from :data:`.ImageFile.ERRORS`. + """ + msg = "unavailable in base decoder" + raise NotImplementedError(msg) + + def set_as_raw( + self, data: bytes, rawmode: str | None = None, extra: tuple[Any, ...] = () + ) -> None: + """ + Convenience method to set the internal image from a stream of raw data + + :param data: Bytes to be set + :param rawmode: The rawmode to be used for the decoder. + If not specified, it will default to the mode of the image + :param extra: Extra arguments for the decoder. + :returns: None + """ + + if not rawmode: + rawmode = self.mode + d = Image._getdecoder(self.mode, "raw", rawmode, extra) + assert self.im is not None + d.setimage(self.im, self.state.extents()) + s = d.decode(data) + + if s[0] >= 0: + msg = "not enough image data" + raise ValueError(msg) + if s[1] != 0: + msg = "cannot decode image data" + raise ValueError(msg) + + +class PyEncoder(PyCodec): + """ + Python implementation of a format encoder. Override this class and + add the decoding logic in the :meth:`encode` method. + + See :ref:`Writing Your Own File Codec in Python` + """ + + _pushes_fd = False + + @property + def pushes_fd(self) -> bool: + return self._pushes_fd + + def encode(self, bufsize: int) -> tuple[int, int, bytes]: + """ + Override to perform the encoding process. + + :param bufsize: Buffer size. + :returns: A tuple of ``(bytes encoded, errcode, bytes)``. + If finished with encoding return 1 for the error code. + Err codes are from :data:`.ImageFile.ERRORS`. + """ + msg = "unavailable in base encoder" + raise NotImplementedError(msg) + + def encode_to_pyfd(self) -> tuple[int, int]: + """ + If ``pushes_fd`` is ``True``, then this method will be used, + and ``encode()`` will only be called once. + + :returns: A tuple of ``(bytes consumed, errcode)``. + Err codes are from :data:`.ImageFile.ERRORS`. + """ + if not self.pushes_fd: + return 0, -8 # bad configuration + bytes_consumed, errcode, data = self.encode(0) + if data: + assert self.fd is not None + self.fd.write(data) + return bytes_consumed, errcode + + def encode_to_file(self, fh: int, bufsize: int) -> int: + """ + :param fh: File handle. + :param bufsize: Buffer size. + + :returns: If finished successfully, return 0. + Otherwise, return an error code. Err codes are from + :data:`.ImageFile.ERRORS`. + """ + errcode = 0 + while errcode == 0: + status, errcode, buf = self.encode(bufsize) + if status > 0: + os.write(fh, buf[status:]) + return errcode diff --git a/venv/lib/python3.10/site-packages/PIL/ImageFilter.py b/venv/lib/python3.10/site-packages/PIL/ImageFilter.py new file mode 100644 index 0000000000000000000000000000000000000000..9326eeeda9de4d479ae4c8c056d19caec12a576f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageFilter.py @@ -0,0 +1,607 @@ +# +# The Python Imaging Library. +# $Id$ +# +# standard filters +# +# History: +# 1995-11-27 fl Created +# 2002-06-08 fl Added rank and mode filters +# 2003-09-15 fl Fixed rank calculation in rank filter; added expand call +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1995-2002 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import abc +import functools +from collections.abc import Sequence +from typing import cast + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from types import ModuleType + from typing import Any + + from . import _imaging + from ._typing import NumpyArray + + +class Filter(abc.ABC): + @abc.abstractmethod + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + pass + + +class MultibandFilter(Filter): + pass + + +class BuiltinFilter(MultibandFilter): + filterargs: tuple[Any, ...] + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + if image.mode == "P": + msg = "cannot filter palette images" + raise ValueError(msg) + return image.filter(*self.filterargs) + + +class Kernel(BuiltinFilter): + """ + Create a convolution kernel. This only supports 3x3 and 5x5 integer and floating + point kernels. + + Kernels can only be applied to "L" and "RGB" images. + + :param size: Kernel size, given as (width, height). This must be (3,3) or (5,5). + :param kernel: A sequence containing kernel weights. The kernel will be flipped + vertically before being applied to the image. + :param scale: Scale factor. If given, the result for each pixel is divided by this + value. The default is the sum of the kernel weights. + :param offset: Offset. If given, this value is added to the result, after it has + been divided by the scale factor. + """ + + name = "Kernel" + + def __init__( + self, + size: tuple[int, int], + kernel: Sequence[float], + scale: float | None = None, + offset: float = 0, + ) -> None: + if scale is None: + # default scale is sum of kernel + scale = functools.reduce(lambda a, b: a + b, kernel) + if size[0] * size[1] != len(kernel): + msg = "not enough coefficients in kernel" + raise ValueError(msg) + self.filterargs = size, scale, offset, kernel + + +class RankFilter(Filter): + """ + Create a rank filter. The rank filter sorts all pixels in + a window of the given size, and returns the ``rank``'th value. + + :param size: The kernel size, in pixels. + :param rank: What pixel value to pick. Use 0 for a min filter, + ``size * size / 2`` for a median filter, ``size * size - 1`` + for a max filter, etc. + """ + + name = "Rank" + + def __init__(self, size: int, rank: int) -> None: + self.size = size + self.rank = rank + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + if image.mode == "P": + msg = "cannot filter palette images" + raise ValueError(msg) + image = image.expand(self.size // 2, self.size // 2) + return image.rankfilter(self.size, self.rank) + + +class MedianFilter(RankFilter): + """ + Create a median filter. Picks the median pixel value in a window with the + given size. + + :param size: The kernel size, in pixels. + """ + + name = "Median" + + def __init__(self, size: int = 3) -> None: + self.size = size + self.rank = size * size // 2 + + +class MinFilter(RankFilter): + """ + Create a min filter. Picks the lowest pixel value in a window with the + given size. + + :param size: The kernel size, in pixels. + """ + + name = "Min" + + def __init__(self, size: int = 3) -> None: + self.size = size + self.rank = 0 + + +class MaxFilter(RankFilter): + """ + Create a max filter. Picks the largest pixel value in a window with the + given size. + + :param size: The kernel size, in pixels. + """ + + name = "Max" + + def __init__(self, size: int = 3) -> None: + self.size = size + self.rank = size * size - 1 + + +class ModeFilter(Filter): + """ + Create a mode filter. Picks the most frequent pixel value in a box with the + given size. Pixel values that occur only once or twice are ignored; if no + pixel value occurs more than twice, the original pixel value is preserved. + + :param size: The kernel size, in pixels. + """ + + name = "Mode" + + def __init__(self, size: int = 3) -> None: + self.size = size + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + return image.modefilter(self.size) + + +class GaussianBlur(MultibandFilter): + """Blurs the image with a sequence of extended box filters, which + approximates a Gaussian kernel. For details on accuracy see + + + :param radius: Standard deviation of the Gaussian kernel. Either a sequence of two + numbers for x and y, or a single number for both. + """ + + name = "GaussianBlur" + + def __init__(self, radius: float | Sequence[float] = 2) -> None: + self.radius = radius + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + xy = self.radius + if isinstance(xy, (int, float)): + xy = (xy, xy) + if xy == (0, 0): + return image.copy() + return image.gaussian_blur(xy) + + +class BoxBlur(MultibandFilter): + """Blurs the image by setting each pixel to the average value of the pixels + in a square box extending radius pixels in each direction. + Supports float radius of arbitrary size. Uses an optimized implementation + which runs in linear time relative to the size of the image + for any radius value. + + :param radius: Size of the box in a direction. Either a sequence of two numbers for + x and y, or a single number for both. + + Radius 0 does not blur, returns an identical image. + Radius 1 takes 1 pixel in each direction, i.e. 9 pixels in total. + """ + + name = "BoxBlur" + + def __init__(self, radius: float | Sequence[float]) -> None: + xy = radius if isinstance(radius, (tuple, list)) else (radius, radius) + if xy[0] < 0 or xy[1] < 0: + msg = "radius must be >= 0" + raise ValueError(msg) + self.radius = radius + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + xy = self.radius + if isinstance(xy, (int, float)): + xy = (xy, xy) + if xy == (0, 0): + return image.copy() + return image.box_blur(xy) + + +class UnsharpMask(MultibandFilter): + """Unsharp mask filter. + + See Wikipedia's entry on `digital unsharp masking`_ for an explanation of + the parameters. + + :param radius: Blur Radius + :param percent: Unsharp strength, in percent + :param threshold: Threshold controls the minimum brightness change that + will be sharpened + + .. _digital unsharp masking: https://en.wikipedia.org/wiki/Unsharp_masking#Digital_unsharp_masking + + """ + + name = "UnsharpMask" + + def __init__( + self, radius: float = 2, percent: int = 150, threshold: int = 3 + ) -> None: + self.radius = radius + self.percent = percent + self.threshold = threshold + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + return image.unsharp_mask(self.radius, self.percent, self.threshold) + + +class BLUR(BuiltinFilter): + name = "Blur" + # fmt: off + filterargs = (5, 5), 16, 0, ( + 1, 1, 1, 1, 1, + 1, 0, 0, 0, 1, + 1, 0, 0, 0, 1, + 1, 0, 0, 0, 1, + 1, 1, 1, 1, 1, + ) + # fmt: on + + +class CONTOUR(BuiltinFilter): + name = "Contour" + # fmt: off + filterargs = (3, 3), 1, 255, ( + -1, -1, -1, + -1, 8, -1, + -1, -1, -1, + ) + # fmt: on + + +class DETAIL(BuiltinFilter): + name = "Detail" + # fmt: off + filterargs = (3, 3), 6, 0, ( + 0, -1, 0, + -1, 10, -1, + 0, -1, 0, + ) + # fmt: on + + +class EDGE_ENHANCE(BuiltinFilter): + name = "Edge-enhance" + # fmt: off + filterargs = (3, 3), 2, 0, ( + -1, -1, -1, + -1, 10, -1, + -1, -1, -1, + ) + # fmt: on + + +class EDGE_ENHANCE_MORE(BuiltinFilter): + name = "Edge-enhance More" + # fmt: off + filterargs = (3, 3), 1, 0, ( + -1, -1, -1, + -1, 9, -1, + -1, -1, -1, + ) + # fmt: on + + +class EMBOSS(BuiltinFilter): + name = "Emboss" + # fmt: off + filterargs = (3, 3), 1, 128, ( + -1, 0, 0, + 0, 1, 0, + 0, 0, 0, + ) + # fmt: on + + +class FIND_EDGES(BuiltinFilter): + name = "Find Edges" + # fmt: off + filterargs = (3, 3), 1, 0, ( + -1, -1, -1, + -1, 8, -1, + -1, -1, -1, + ) + # fmt: on + + +class SHARPEN(BuiltinFilter): + name = "Sharpen" + # fmt: off + filterargs = (3, 3), 16, 0, ( + -2, -2, -2, + -2, 32, -2, + -2, -2, -2, + ) + # fmt: on + + +class SMOOTH(BuiltinFilter): + name = "Smooth" + # fmt: off + filterargs = (3, 3), 13, 0, ( + 1, 1, 1, + 1, 5, 1, + 1, 1, 1, + ) + # fmt: on + + +class SMOOTH_MORE(BuiltinFilter): + name = "Smooth More" + # fmt: off + filterargs = (5, 5), 100, 0, ( + 1, 1, 1, 1, 1, + 1, 5, 5, 5, 1, + 1, 5, 44, 5, 1, + 1, 5, 5, 5, 1, + 1, 1, 1, 1, 1, + ) + # fmt: on + + +class Color3DLUT(MultibandFilter): + """Three-dimensional color lookup table. + + Transforms 3-channel pixels using the values of the channels as coordinates + in the 3D lookup table and interpolating the nearest elements. + + This method allows you to apply almost any color transformation + in constant time by using pre-calculated decimated tables. + + .. versionadded:: 5.2.0 + + :param size: Size of the table. One int or tuple of (int, int, int). + Minimal size in any dimension is 2, maximum is 65. + :param table: Flat lookup table. A list of ``channels * size**3`` + float elements or a list of ``size**3`` channels-sized + tuples with floats. Channels are changed first, + then first dimension, then second, then third. + Value 0.0 corresponds lowest value of output, 1.0 highest. + :param channels: Number of channels in the table. Could be 3 or 4. + Default is 3. + :param target_mode: A mode for the result image. Should have not less + than ``channels`` channels. Default is ``None``, + which means that mode wouldn't be changed. + """ + + name = "Color 3D LUT" + + def __init__( + self, + size: int | tuple[int, int, int], + table: Sequence[float] | Sequence[Sequence[int]] | NumpyArray, + channels: int = 3, + target_mode: str | None = None, + **kwargs: bool, + ) -> None: + if channels not in (3, 4): + msg = "Only 3 or 4 output channels are supported" + raise ValueError(msg) + self.size = size = self._check_size(size) + self.channels = channels + self.mode = target_mode + + # Hidden flag `_copy_table=False` could be used to avoid extra copying + # of the table if the table is specially made for the constructor. + copy_table = kwargs.get("_copy_table", True) + items = size[0] * size[1] * size[2] + wrong_size = False + + numpy: ModuleType | None = None + if hasattr(table, "shape"): + try: + import numpy + except ImportError: + pass + + if numpy and isinstance(table, numpy.ndarray): + numpy_table: NumpyArray = table + if copy_table: + numpy_table = numpy_table.copy() + + if numpy_table.shape in [ + (items * channels,), + (items, channels), + (size[2], size[1], size[0], channels), + ]: + table = numpy_table.reshape(items * channels) + else: + wrong_size = True + + else: + if copy_table: + table = list(table) + + # Convert to a flat list + if table and isinstance(table[0], (list, tuple)): + raw_table = cast(Sequence[Sequence[int]], table) + flat_table: list[int] = [] + for pixel in raw_table: + if len(pixel) != channels: + msg = ( + "The elements of the table should " + f"have a length of {channels}." + ) + raise ValueError(msg) + flat_table.extend(pixel) + table = flat_table + + if wrong_size or len(table) != items * channels: + msg = ( + "The table should have either channels * size**3 float items " + "or size**3 items of channels-sized tuples with floats. " + f"Table should be: {channels}x{size[0]}x{size[1]}x{size[2]}. " + f"Actual length: {len(table)}" + ) + raise ValueError(msg) + self.table = table + + @staticmethod + def _check_size(size: Any) -> tuple[int, int, int]: + try: + _, _, _ = size + except ValueError as e: + msg = "Size should be either an integer or a tuple of three integers." + raise ValueError(msg) from e + except TypeError: + size = (size, size, size) + size = tuple(int(x) for x in size) + for size_1d in size: + if not 2 <= size_1d <= 65: + msg = "Size should be in [2, 65] range." + raise ValueError(msg) + return size + + @classmethod + def generate( + cls, + size: int | tuple[int, int, int], + callback: Callable[[float, float, float], tuple[float, ...]], + channels: int = 3, + target_mode: str | None = None, + ) -> Color3DLUT: + """Generates new LUT using provided callback. + + :param size: Size of the table. Passed to the constructor. + :param callback: Function with three parameters which correspond + three color channels. Will be called ``size**3`` + times with values from 0.0 to 1.0 and should return + a tuple with ``channels`` elements. + :param channels: The number of channels which should return callback. + :param target_mode: Passed to the constructor of the resulting + lookup table. + """ + size_1d, size_2d, size_3d = cls._check_size(size) + if channels not in (3, 4): + msg = "Only 3 or 4 output channels are supported" + raise ValueError(msg) + + table: list[float] = [0] * (size_1d * size_2d * size_3d * channels) + idx_out = 0 + for b in range(size_3d): + for g in range(size_2d): + for r in range(size_1d): + table[idx_out : idx_out + channels] = callback( + r / (size_1d - 1), g / (size_2d - 1), b / (size_3d - 1) + ) + idx_out += channels + + return cls( + (size_1d, size_2d, size_3d), + table, + channels=channels, + target_mode=target_mode, + _copy_table=False, + ) + + def transform( + self, + callback: Callable[..., tuple[float, ...]], + with_normals: bool = False, + channels: int | None = None, + target_mode: str | None = None, + ) -> Color3DLUT: + """Transforms the table values using provided callback and returns + a new LUT with altered values. + + :param callback: A function which takes old lookup table values + and returns a new set of values. The number + of arguments which function should take is + ``self.channels`` or ``3 + self.channels`` + if ``with_normals`` flag is set. + Should return a tuple of ``self.channels`` or + ``channels`` elements if it is set. + :param with_normals: If true, ``callback`` will be called with + coordinates in the color cube as the first + three arguments. Otherwise, ``callback`` + will be called only with actual color values. + :param channels: The number of channels in the resulting lookup table. + :param target_mode: Passed to the constructor of the resulting + lookup table. + """ + if channels not in (None, 3, 4): + msg = "Only 3 or 4 output channels are supported" + raise ValueError(msg) + ch_in = self.channels + ch_out = channels or ch_in + size_1d, size_2d, size_3d = self.size + + table: list[float] = [0] * (size_1d * size_2d * size_3d * ch_out) + idx_in = 0 + idx_out = 0 + for b in range(size_3d): + for g in range(size_2d): + for r in range(size_1d): + values = self.table[idx_in : idx_in + ch_in] + if with_normals: + values = callback( + r / (size_1d - 1), + g / (size_2d - 1), + b / (size_3d - 1), + *values, + ) + else: + values = callback(*values) + table[idx_out : idx_out + ch_out] = values + idx_in += ch_in + idx_out += ch_out + + return type(self)( + self.size, + table, + channels=ch_out, + target_mode=target_mode or self.mode, + _copy_table=False, + ) + + def __repr__(self) -> str: + r = [ + f"{self.__class__.__name__} from {self.table.__class__.__name__}", + "size={:d}x{:d}x{:d}".format(*self.size), + f"channels={self.channels:d}", + ] + if self.mode: + r.append(f"target_mode={self.mode}") + return "<{}>".format(" ".join(r)) + + def filter(self, image: _imaging.ImagingCore) -> _imaging.ImagingCore: + from . import Image + + return image.color_lut_3d( + self.mode or image.mode, + Image.Resampling.BILINEAR, + self.channels, + self.size, + self.table, + ) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageFont.py b/venv/lib/python3.10/site-packages/PIL/ImageFont.py new file mode 100644 index 0000000000000000000000000000000000000000..92eb763a51e15102926f892ca85f74ebc6a4d7f4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageFont.py @@ -0,0 +1,1312 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PIL raster font management +# +# History: +# 1996-08-07 fl created (experimental) +# 1997-08-25 fl minor adjustments to handle fonts from pilfont 0.3 +# 1999-02-06 fl rewrote most font management stuff in C +# 1999-03-17 fl take pth files into account in load_path (from Richard Jones) +# 2001-02-17 fl added freetype support +# 2001-05-09 fl added TransposedFont wrapper class +# 2002-03-04 fl make sure we have a "L" or "1" font +# 2002-12-04 fl skip non-directory entries in the system path +# 2003-04-29 fl add embedded default font +# 2003-09-27 fl added support for truetype charmap encodings +# +# Todo: +# Adapt to PILFONT2 format (16-bit fonts, compressed, single file) +# +# Copyright (c) 1997-2003 by Secret Labs AB +# Copyright (c) 1996-2003 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# + +from __future__ import annotations + +import base64 +import os +import sys +import warnings +from enum import IntEnum +from io import BytesIO +from types import ModuleType +from typing import IO, Any, BinaryIO, TypedDict, cast + +from . import Image +from ._typing import StrOrBytesPath +from ._util import DeferredError, is_path + +TYPE_CHECKING = False +if TYPE_CHECKING: + from . import ImageFile + from ._imaging import ImagingFont + from ._imagingft import Font + + +class Axis(TypedDict): + minimum: int | None + default: int | None + maximum: int | None + name: bytes | None + + +class Layout(IntEnum): + BASIC = 0 + RAQM = 1 + + +MAX_STRING_LENGTH = 1_000_000 + + +core: ModuleType | DeferredError +try: + from . import _imagingft as core +except ImportError as ex: + core = DeferredError.new(ex) + + +def _string_length_check(text: str | bytes | bytearray) -> None: + if MAX_STRING_LENGTH is not None and len(text) > MAX_STRING_LENGTH: + msg = "too many characters in string" + raise ValueError(msg) + + +# FIXME: add support for pilfont2 format (see FontFile.py) + +# -------------------------------------------------------------------- +# Font metrics format: +# "PILfont" LF +# fontdescriptor LF +# (optional) key=value... LF +# "DATA" LF +# binary data: 256*10*2 bytes (dx, dy, dstbox, srcbox) +# +# To place a character, cut out srcbox and paste at dstbox, +# relative to the character position. Then move the character +# position according to dx, dy. +# -------------------------------------------------------------------- + + +class ImageFont: + """PIL font wrapper""" + + font: ImagingFont + + def _load_pilfont(self, filename: str) -> None: + with open(filename, "rb") as fp: + image: ImageFile.ImageFile | None = None + root = os.path.splitext(filename)[0] + + for ext in (".png", ".gif", ".pbm"): + if image: + image.close() + try: + fullname = root + ext + image = Image.open(fullname) + except Exception: + pass + else: + if image and image.mode in ("1", "L"): + break + else: + if image: + image.close() + + msg = f"cannot find glyph data file {root}.{{gif|pbm|png}}" + raise OSError(msg) + + self.file = fullname + + self._load_pilfont_data(fp, image) + image.close() + + def _load_pilfont_data(self, file: IO[bytes], image: Image.Image) -> None: + # check image + if image.mode not in ("1", "L"): + msg = "invalid font image mode" + raise TypeError(msg) + + # read PILfont header + if file.read(8) != b"PILfont\n": + msg = "Not a PILfont file" + raise SyntaxError(msg) + file.readline() + self.info = [] # FIXME: should be a dictionary + while True: + s = file.readline() + if not s or s == b"DATA\n": + break + self.info.append(s) + + # read PILfont metrics + data = file.read(256 * 20) + + image.load() + + self.font = Image.core.font(image.im, data) + + def getmask( + self, text: str | bytes, mode: str = "", *args: Any, **kwargs: Any + ) -> Image.core.ImagingCore: + """ + Create a bitmap for the text. + + If the font uses antialiasing, the bitmap should have mode ``L`` and use a + maximum value of 255. Otherwise, it should have mode ``1``. + + :param text: Text to render. + :param mode: Used by some graphics drivers to indicate what mode the + driver prefers; if empty, the renderer may return either + mode. Note that the mode is always a string, to simplify + C-level implementations. + + .. versionadded:: 1.1.5 + + :return: An internal PIL storage memory instance as defined by the + :py:mod:`PIL.Image.core` interface module. + """ + _string_length_check(text) + Image._decompression_bomb_check(self.font.getsize(text)) + return self.font.getmask(text, mode) + + def getbbox( + self, text: str | bytes | bytearray, *args: Any, **kwargs: Any + ) -> tuple[int, int, int, int]: + """ + Returns bounding box (in pixels) of given text. + + .. versionadded:: 9.2.0 + + :param text: Text to render. + + :return: ``(left, top, right, bottom)`` bounding box + """ + _string_length_check(text) + width, height = self.font.getsize(text) + return 0, 0, width, height + + def getlength( + self, text: str | bytes | bytearray, *args: Any, **kwargs: Any + ) -> int: + """ + Returns length (in pixels) of given text. + This is the amount by which following text should be offset. + + .. versionadded:: 9.2.0 + """ + _string_length_check(text) + width, height = self.font.getsize(text) + return width + + +## +# Wrapper for FreeType fonts. Application code should use the +# truetype factory function to create font objects. + + +class FreeTypeFont: + """FreeType font wrapper (requires _imagingft service)""" + + font: Font + font_bytes: bytes + + def __init__( + self, + font: StrOrBytesPath | BinaryIO, + size: float = 10, + index: int = 0, + encoding: str = "", + layout_engine: Layout | None = None, + ) -> None: + # FIXME: use service provider instead + + if isinstance(core, DeferredError): + raise core.ex + + if size <= 0: + msg = f"font size must be greater than 0, not {size}" + raise ValueError(msg) + + self.path = font + self.size = size + self.index = index + self.encoding = encoding + + if layout_engine not in (Layout.BASIC, Layout.RAQM): + layout_engine = Layout.BASIC + if core.HAVE_RAQM: + layout_engine = Layout.RAQM + elif layout_engine == Layout.RAQM and not core.HAVE_RAQM: + warnings.warn( + "Raqm layout was requested, but Raqm is not available. " + "Falling back to basic layout." + ) + layout_engine = Layout.BASIC + + self.layout_engine = layout_engine + + def load_from_bytes(f: IO[bytes]) -> None: + self.font_bytes = f.read() + self.font = core.getfont( + "", size, index, encoding, self.font_bytes, layout_engine + ) + + if is_path(font): + font = os.fspath(font) + if sys.platform == "win32": + font_bytes_path = font if isinstance(font, bytes) else font.encode() + try: + font_bytes_path.decode("ascii") + except UnicodeDecodeError: + # FreeType cannot load fonts with non-ASCII characters on Windows + # So load it into memory first + with open(font, "rb") as f: + load_from_bytes(f) + return + self.font = core.getfont( + font, size, index, encoding, layout_engine=layout_engine + ) + else: + load_from_bytes(cast(IO[bytes], font)) + + def __getstate__(self) -> list[Any]: + return [self.path, self.size, self.index, self.encoding, self.layout_engine] + + def __setstate__(self, state: list[Any]) -> None: + path, size, index, encoding, layout_engine = state + FreeTypeFont.__init__(self, path, size, index, encoding, layout_engine) + + def getname(self) -> tuple[str | None, str | None]: + """ + :return: A tuple of the font family (e.g. Helvetica) and the font style + (e.g. Bold) + """ + return self.font.family, self.font.style + + def getmetrics(self) -> tuple[int, int]: + """ + :return: A tuple of the font ascent (the distance from the baseline to + the highest outline point) and descent (the distance from the + baseline to the lowest outline point, a negative value) + """ + return self.font.ascent, self.font.descent + + def getlength( + self, + text: str | bytes, + mode: str = "", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + ) -> float: + """ + Returns length (in pixels with 1/64 precision) of given text when rendered + in font with provided direction, features, and language. + + This is the amount by which following text should be offset. + Text bounding box may extend past the length in some fonts, + e.g. when using italics or accents. + + The result is returned as a float; it is a whole number if using basic layout. + + Note that the sum of two lengths may not equal the length of a concatenated + string due to kerning. If you need to adjust for kerning, include the following + character and subtract its length. + + For example, instead of :: + + hello = font.getlength("Hello") + world = font.getlength("World") + hello_world = hello + world # not adjusted for kerning + assert hello_world == font.getlength("HelloWorld") # may fail + + use :: + + hello = font.getlength("HelloW") - font.getlength("W") # adjusted for kerning + world = font.getlength("World") + hello_world = hello + world # adjusted for kerning + assert hello_world == font.getlength("HelloWorld") # True + + or disable kerning with (requires libraqm) :: + + hello = draw.textlength("Hello", font, features=["-kern"]) + world = draw.textlength("World", font, features=["-kern"]) + hello_world = hello + world # kerning is disabled, no need to adjust + assert hello_world == draw.textlength("HelloWorld", font, features=["-kern"]) + + .. versionadded:: 8.0.0 + + :param text: Text to measure. + :param mode: Used by some graphics drivers to indicate what mode the + driver prefers; if empty, the renderer may return either + mode. Note that the mode is always a string, to simplify + C-level implementations. + + :param direction: Direction of the text. It can be 'rtl' (right to + left), 'ltr' (left to right) or 'ttb' (top to bottom). + Requires libraqm. + + :param features: A list of OpenType font features to be used during text + layout. This is usually used to turn on optional + font features that are not enabled by default, + for example 'dlig' or 'ss01', but can be also + used to turn off default font features for + example '-liga' to disable ligatures or '-kern' + to disable kerning. To get all supported + features, see + https://learn.microsoft.com/en-us/typography/opentype/spec/featurelist + Requires libraqm. + + :param language: Language of the text. Different languages may use + different glyph shapes or ligatures. This parameter tells + the font which language the text is in, and to apply the + correct substitutions as appropriate, if available. + It should be a `BCP 47 language code + `_ + Requires libraqm. + + :return: Either width for horizontal text, or height for vertical text. + """ + _string_length_check(text) + return self.font.getlength(text, mode, direction, features, language) / 64 + + def getbbox( + self, + text: str | bytes, + mode: str = "", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + anchor: str | None = None, + ) -> tuple[float, float, float, float]: + """ + Returns bounding box (in pixels) of given text relative to given anchor + when rendered in font with provided direction, features, and language. + + Use :py:meth:`getlength()` to get the offset of following text with + 1/64 pixel precision. The bounding box includes extra margins for + some fonts, e.g. italics or accents. + + .. versionadded:: 8.0.0 + + :param text: Text to render. + :param mode: Used by some graphics drivers to indicate what mode the + driver prefers; if empty, the renderer may return either + mode. Note that the mode is always a string, to simplify + C-level implementations. + + :param direction: Direction of the text. It can be 'rtl' (right to + left), 'ltr' (left to right) or 'ttb' (top to bottom). + Requires libraqm. + + :param features: A list of OpenType font features to be used during text + layout. This is usually used to turn on optional + font features that are not enabled by default, + for example 'dlig' or 'ss01', but can be also + used to turn off default font features for + example '-liga' to disable ligatures or '-kern' + to disable kerning. To get all supported + features, see + https://learn.microsoft.com/en-us/typography/opentype/spec/featurelist + Requires libraqm. + + :param language: Language of the text. Different languages may use + different glyph shapes or ligatures. This parameter tells + the font which language the text is in, and to apply the + correct substitutions as appropriate, if available. + It should be a `BCP 47 language code + `_ + Requires libraqm. + + :param stroke_width: The width of the text stroke. + + :param anchor: The text anchor alignment. Determines the relative location of + the anchor to the text. The default alignment is top left, + specifically ``la`` for horizontal text and ``lt`` for + vertical text. See :ref:`text-anchors` for details. + + :return: ``(left, top, right, bottom)`` bounding box + """ + _string_length_check(text) + size, offset = self.font.getsize( + text, mode, direction, features, language, anchor + ) + left, top = offset[0] - stroke_width, offset[1] - stroke_width + width, height = size[0] + 2 * stroke_width, size[1] + 2 * stroke_width + return left, top, left + width, top + height + + def getmask( + self, + text: str | bytes, + mode: str = "", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + anchor: str | None = None, + ink: int = 0, + start: tuple[float, float] | None = None, + ) -> Image.core.ImagingCore: + """ + Create a bitmap for the text. + + If the font uses antialiasing, the bitmap should have mode ``L`` and use a + maximum value of 255. If the font has embedded color data, the bitmap + should have mode ``RGBA``. Otherwise, it should have mode ``1``. + + :param text: Text to render. + :param mode: Used by some graphics drivers to indicate what mode the + driver prefers; if empty, the renderer may return either + mode. Note that the mode is always a string, to simplify + C-level implementations. + + .. versionadded:: 1.1.5 + + :param direction: Direction of the text. It can be 'rtl' (right to + left), 'ltr' (left to right) or 'ttb' (top to bottom). + Requires libraqm. + + .. versionadded:: 4.2.0 + + :param features: A list of OpenType font features to be used during text + layout. This is usually used to turn on optional + font features that are not enabled by default, + for example 'dlig' or 'ss01', but can be also + used to turn off default font features for + example '-liga' to disable ligatures or '-kern' + to disable kerning. To get all supported + features, see + https://learn.microsoft.com/en-us/typography/opentype/spec/featurelist + Requires libraqm. + + .. versionadded:: 4.2.0 + + :param language: Language of the text. Different languages may use + different glyph shapes or ligatures. This parameter tells + the font which language the text is in, and to apply the + correct substitutions as appropriate, if available. + It should be a `BCP 47 language code + `_ + Requires libraqm. + + .. versionadded:: 6.0.0 + + :param stroke_width: The width of the text stroke. + + .. versionadded:: 6.2.0 + + :param anchor: The text anchor alignment. Determines the relative location of + the anchor to the text. The default alignment is top left, + specifically ``la`` for horizontal text and ``lt`` for + vertical text. See :ref:`text-anchors` for details. + + .. versionadded:: 8.0.0 + + :param ink: Foreground ink for rendering in RGBA mode. + + .. versionadded:: 8.0.0 + + :param start: Tuple of horizontal and vertical offset, as text may render + differently when starting at fractional coordinates. + + .. versionadded:: 9.4.0 + + :return: An internal PIL storage memory instance as defined by the + :py:mod:`PIL.Image.core` interface module. + """ + return self.getmask2( + text, + mode, + direction=direction, + features=features, + language=language, + stroke_width=stroke_width, + anchor=anchor, + ink=ink, + start=start, + )[0] + + def getmask2( + self, + text: str | bytes, + mode: str = "", + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + stroke_width: float = 0, + anchor: str | None = None, + ink: int = 0, + start: tuple[float, float] | None = None, + *args: Any, + **kwargs: Any, + ) -> tuple[Image.core.ImagingCore, tuple[int, int]]: + """ + Create a bitmap for the text. + + If the font uses antialiasing, the bitmap should have mode ``L`` and use a + maximum value of 255. If the font has embedded color data, the bitmap + should have mode ``RGBA``. Otherwise, it should have mode ``1``. + + :param text: Text to render. + :param mode: Used by some graphics drivers to indicate what mode the + driver prefers; if empty, the renderer may return either + mode. Note that the mode is always a string, to simplify + C-level implementations. + + .. versionadded:: 1.1.5 + + :param direction: Direction of the text. It can be 'rtl' (right to + left), 'ltr' (left to right) or 'ttb' (top to bottom). + Requires libraqm. + + .. versionadded:: 4.2.0 + + :param features: A list of OpenType font features to be used during text + layout. This is usually used to turn on optional + font features that are not enabled by default, + for example 'dlig' or 'ss01', but can be also + used to turn off default font features for + example '-liga' to disable ligatures or '-kern' + to disable kerning. To get all supported + features, see + https://learn.microsoft.com/en-us/typography/opentype/spec/featurelist + Requires libraqm. + + .. versionadded:: 4.2.0 + + :param language: Language of the text. Different languages may use + different glyph shapes or ligatures. This parameter tells + the font which language the text is in, and to apply the + correct substitutions as appropriate, if available. + It should be a `BCP 47 language code + `_ + Requires libraqm. + + .. versionadded:: 6.0.0 + + :param stroke_width: The width of the text stroke. + + .. versionadded:: 6.2.0 + + :param anchor: The text anchor alignment. Determines the relative location of + the anchor to the text. The default alignment is top left, + specifically ``la`` for horizontal text and ``lt`` for + vertical text. See :ref:`text-anchors` for details. + + .. versionadded:: 8.0.0 + + :param ink: Foreground ink for rendering in RGBA mode. + + .. versionadded:: 8.0.0 + + :param start: Tuple of horizontal and vertical offset, as text may render + differently when starting at fractional coordinates. + + .. versionadded:: 9.4.0 + + :return: A tuple of an internal PIL storage memory instance as defined by the + :py:mod:`PIL.Image.core` interface module, and the text offset, the + gap between the starting coordinate and the first marking + """ + _string_length_check(text) + if start is None: + start = (0, 0) + + def fill(width: int, height: int) -> Image.core.ImagingCore: + size = (width, height) + Image._decompression_bomb_check(size) + return Image.core.fill("RGBA" if mode == "RGBA" else "L", size) + + return self.font.render( + text, + fill, + mode, + direction, + features, + language, + stroke_width, + kwargs.get("stroke_filled", False), + anchor, + ink, + start, + ) + + def font_variant( + self, + font: StrOrBytesPath | BinaryIO | None = None, + size: float | None = None, + index: int | None = None, + encoding: str | None = None, + layout_engine: Layout | None = None, + ) -> FreeTypeFont: + """ + Create a copy of this FreeTypeFont object, + using any specified arguments to override the settings. + + Parameters are identical to the parameters used to initialize this + object. + + :return: A FreeTypeFont object. + """ + if font is None: + try: + font = BytesIO(self.font_bytes) + except AttributeError: + font = self.path + return FreeTypeFont( + font=font, + size=self.size if size is None else size, + index=self.index if index is None else index, + encoding=self.encoding if encoding is None else encoding, + layout_engine=layout_engine or self.layout_engine, + ) + + def get_variation_names(self) -> list[bytes]: + """ + :returns: A list of the named styles in a variation font. + :exception OSError: If the font is not a variation font. + """ + names = self.font.getvarnames() + return [name.replace(b"\x00", b"") for name in names] + + def set_variation_by_name(self, name: str | bytes) -> None: + """ + :param name: The name of the style. + :exception OSError: If the font is not a variation font. + """ + names = self.get_variation_names() + if not isinstance(name, bytes): + name = name.encode() + index = names.index(name) + 1 + + if index == getattr(self, "_last_variation_index", None): + # When the same name is set twice in a row, + # there is an 'unknown freetype error' + # https://savannah.nongnu.org/bugs/?56186 + return + self._last_variation_index = index + + self.font.setvarname(index) + + def get_variation_axes(self) -> list[Axis]: + """ + :returns: A list of the axes in a variation font. + :exception OSError: If the font is not a variation font. + """ + axes = self.font.getvaraxes() + for axis in axes: + if axis["name"]: + axis["name"] = axis["name"].replace(b"\x00", b"") + return axes + + def set_variation_by_axes(self, axes: list[float]) -> None: + """ + :param axes: A list of values for each axis. + :exception OSError: If the font is not a variation font. + """ + self.font.setvaraxes(axes) + + +class TransposedFont: + """Wrapper for writing rotated or mirrored text""" + + def __init__( + self, font: ImageFont | FreeTypeFont, orientation: Image.Transpose | None = None + ): + """ + Wrapper that creates a transposed font from any existing font + object. + + :param font: A font object. + :param orientation: An optional orientation. If given, this should + be one of Image.Transpose.FLIP_LEFT_RIGHT, Image.Transpose.FLIP_TOP_BOTTOM, + Image.Transpose.ROTATE_90, Image.Transpose.ROTATE_180, or + Image.Transpose.ROTATE_270. + """ + self.font = font + self.orientation = orientation # any 'transpose' argument, or None + + def getmask( + self, text: str | bytes, mode: str = "", *args: Any, **kwargs: Any + ) -> Image.core.ImagingCore: + im = self.font.getmask(text, mode, *args, **kwargs) + if self.orientation is not None: + return im.transpose(self.orientation) + return im + + def getbbox( + self, text: str | bytes, *args: Any, **kwargs: Any + ) -> tuple[int, int, float, float]: + # TransposedFont doesn't support getmask2, move top-left point to (0, 0) + # this has no effect on ImageFont and simulates anchor="lt" for FreeTypeFont + left, top, right, bottom = self.font.getbbox(text, *args, **kwargs) + width = right - left + height = bottom - top + if self.orientation in (Image.Transpose.ROTATE_90, Image.Transpose.ROTATE_270): + return 0, 0, height, width + return 0, 0, width, height + + def getlength(self, text: str | bytes, *args: Any, **kwargs: Any) -> float: + if self.orientation in (Image.Transpose.ROTATE_90, Image.Transpose.ROTATE_270): + msg = "text length is undefined for text rotated by 90 or 270 degrees" + raise ValueError(msg) + return self.font.getlength(text, *args, **kwargs) + + +def load(filename: str) -> ImageFont: + """ + Load a font file. This function loads a font object from the given + bitmap font file, and returns the corresponding font object. For loading TrueType + or OpenType fonts instead, see :py:func:`~PIL.ImageFont.truetype`. + + :param filename: Name of font file. + :return: A font object. + :exception OSError: If the file could not be read. + """ + f = ImageFont() + f._load_pilfont(filename) + return f + + +def truetype( + font: StrOrBytesPath | BinaryIO, + size: float = 10, + index: int = 0, + encoding: str = "", + layout_engine: Layout | None = None, +) -> FreeTypeFont: + """ + Load a TrueType or OpenType font from a file or file-like object, + and create a font object. This function loads a font object from the given + file or file-like object, and creates a font object for a font of the given + size. For loading bitmap fonts instead, see :py:func:`~PIL.ImageFont.load` + and :py:func:`~PIL.ImageFont.load_path`. + + Pillow uses FreeType to open font files. On Windows, be aware that FreeType + will keep the file open as long as the FreeTypeFont object exists. Windows + limits the number of files that can be open in C at once to 512, so if many + fonts are opened simultaneously and that limit is approached, an + ``OSError`` may be thrown, reporting that FreeType "cannot open resource". + A workaround would be to copy the file(s) into memory, and open that instead. + + This function requires the _imagingft service. + + :param font: A filename or file-like object containing a TrueType font. + If the file is not found in this filename, the loader may also + search in other directories, such as: + + * The :file:`fonts/` directory on Windows, + * :file:`/Library/Fonts/`, :file:`/System/Library/Fonts/` + and :file:`~/Library/Fonts/` on macOS. + * :file:`~/.local/share/fonts`, :file:`/usr/local/share/fonts`, + and :file:`/usr/share/fonts` on Linux; or those specified by + the ``XDG_DATA_HOME`` and ``XDG_DATA_DIRS`` environment variables + for user-installed and system-wide fonts, respectively. + + :param size: The requested size, in pixels. + :param index: Which font face to load (default is first available face). + :param encoding: Which font encoding to use (default is Unicode). Possible + encodings include (see the FreeType documentation for more + information): + + * "unic" (Unicode) + * "symb" (Microsoft Symbol) + * "ADOB" (Adobe Standard) + * "ADBE" (Adobe Expert) + * "ADBC" (Adobe Custom) + * "armn" (Apple Roman) + * "sjis" (Shift JIS) + * "gb " (PRC) + * "big5" + * "wans" (Extended Wansung) + * "joha" (Johab) + * "lat1" (Latin-1) + + This specifies the character set to use. It does not alter the + encoding of any text provided in subsequent operations. + :param layout_engine: Which layout engine to use, if available: + :attr:`.ImageFont.Layout.BASIC` or :attr:`.ImageFont.Layout.RAQM`. + If it is available, Raqm layout will be used by default. + Otherwise, basic layout will be used. + + Raqm layout is recommended for all non-English text. If Raqm layout + is not required, basic layout will have better performance. + + You can check support for Raqm layout using + :py:func:`PIL.features.check_feature` with ``feature="raqm"``. + + .. versionadded:: 4.2.0 + :return: A font object. + :exception OSError: If the file could not be read. + :exception ValueError: If the font size is not greater than zero. + """ + + def freetype(font: StrOrBytesPath | BinaryIO) -> FreeTypeFont: + return FreeTypeFont(font, size, index, encoding, layout_engine) + + try: + return freetype(font) + except OSError: + if not is_path(font): + raise + ttf_filename = os.path.basename(font) + + dirs = [] + if sys.platform == "win32": + # check the windows font repository + # NOTE: must use uppercase WINDIR, to work around bugs in + # 1.5.2's os.environ.get() + windir = os.environ.get("WINDIR") + if windir: + dirs.append(os.path.join(windir, "fonts")) + elif sys.platform in ("linux", "linux2"): + data_home = os.environ.get("XDG_DATA_HOME") + if not data_home: + # The freedesktop spec defines the following default directory for + # when XDG_DATA_HOME is unset or empty. This user-level directory + # takes precedence over system-level directories. + data_home = os.path.expanduser("~/.local/share") + xdg_dirs = [data_home] + + data_dirs = os.environ.get("XDG_DATA_DIRS") + if not data_dirs: + # Similarly, defaults are defined for the system-level directories + data_dirs = "/usr/local/share:/usr/share" + xdg_dirs += data_dirs.split(":") + + dirs += [os.path.join(xdg_dir, "fonts") for xdg_dir in xdg_dirs] + elif sys.platform == "darwin": + dirs += [ + "/Library/Fonts", + "/System/Library/Fonts", + os.path.expanduser("~/Library/Fonts"), + ] + + ext = os.path.splitext(ttf_filename)[1] + first_font_with_a_different_extension = None + for directory in dirs: + for walkroot, walkdir, walkfilenames in os.walk(directory): + for walkfilename in walkfilenames: + if ext and walkfilename == ttf_filename: + return freetype(os.path.join(walkroot, walkfilename)) + elif not ext and os.path.splitext(walkfilename)[0] == ttf_filename: + fontpath = os.path.join(walkroot, walkfilename) + if os.path.splitext(fontpath)[1] == ".ttf": + return freetype(fontpath) + if not ext and first_font_with_a_different_extension is None: + first_font_with_a_different_extension = fontpath + if first_font_with_a_different_extension: + return freetype(first_font_with_a_different_extension) + raise + + +def load_path(filename: str | bytes) -> ImageFont: + """ + Load font file. Same as :py:func:`~PIL.ImageFont.load`, but searches for a + bitmap font along the Python path. + + :param filename: Name of font file. + :return: A font object. + :exception OSError: If the file could not be read. + """ + if not isinstance(filename, str): + filename = filename.decode("utf-8") + for directory in sys.path: + try: + return load(os.path.join(directory, filename)) + except OSError: + pass + msg = f'cannot find font file "{filename}" in sys.path' + if os.path.exists(filename): + msg += f', did you mean ImageFont.load("{filename}") instead?' + + raise OSError(msg) + + +def load_default_imagefont() -> ImageFont: + f = ImageFont() + f._load_pilfont_data( + # courB08 + BytesIO( + base64.b64decode( + b""" +UElMZm9udAo7Ozs7OzsxMDsKREFUQQoAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAYAAAAA//8AAQAAAAAAAAABAAEA +BgAAAAH/+gADAAAAAQAAAAMABgAGAAAAAf/6AAT//QADAAAABgADAAYAAAAA//kABQABAAYAAAAL +AAgABgAAAAD/+AAFAAEACwAAABAACQAGAAAAAP/5AAUAAAAQAAAAFQAHAAYAAP////oABQAAABUA +AAAbAAYABgAAAAH/+QAE//wAGwAAAB4AAwAGAAAAAf/5AAQAAQAeAAAAIQAIAAYAAAAB//kABAAB +ACEAAAAkAAgABgAAAAD/+QAE//0AJAAAACgABAAGAAAAAP/6AAX//wAoAAAALQAFAAYAAAAB//8A +BAACAC0AAAAwAAMABgAAAAD//AAF//0AMAAAADUAAQAGAAAAAf//AAMAAAA1AAAANwABAAYAAAAB +//kABQABADcAAAA7AAgABgAAAAD/+QAFAAAAOwAAAEAABwAGAAAAAP/5AAYAAABAAAAARgAHAAYA +AAAA//kABQAAAEYAAABLAAcABgAAAAD/+QAFAAAASwAAAFAABwAGAAAAAP/5AAYAAABQAAAAVgAH +AAYAAAAA//kABQAAAFYAAABbAAcABgAAAAD/+QAFAAAAWwAAAGAABwAGAAAAAP/5AAUAAABgAAAA +ZQAHAAYAAAAA//kABQAAAGUAAABqAAcABgAAAAD/+QAFAAAAagAAAG8ABwAGAAAAAf/8AAMAAABv +AAAAcQAEAAYAAAAA//wAAwACAHEAAAB0AAYABgAAAAD/+gAE//8AdAAAAHgABQAGAAAAAP/7AAT/ +/gB4AAAAfAADAAYAAAAB//oABf//AHwAAACAAAUABgAAAAD/+gAFAAAAgAAAAIUABgAGAAAAAP/5 +AAYAAQCFAAAAiwAIAAYAAP////oABgAAAIsAAACSAAYABgAA////+gAFAAAAkgAAAJgABgAGAAAA +AP/6AAUAAACYAAAAnQAGAAYAAP////oABQAAAJ0AAACjAAYABgAA////+gAFAAAAowAAAKkABgAG +AAD////6AAUAAACpAAAArwAGAAYAAAAA//oABQAAAK8AAAC0AAYABgAA////+gAGAAAAtAAAALsA +BgAGAAAAAP/6AAQAAAC7AAAAvwAGAAYAAP////oABQAAAL8AAADFAAYABgAA////+gAGAAAAxQAA +AMwABgAGAAD////6AAUAAADMAAAA0gAGAAYAAP////oABQAAANIAAADYAAYABgAA////+gAGAAAA +2AAAAN8ABgAGAAAAAP/6AAUAAADfAAAA5AAGAAYAAP////oABQAAAOQAAADqAAYABgAAAAD/+gAF +AAEA6gAAAO8ABwAGAAD////6AAYAAADvAAAA9gAGAAYAAAAA//oABQAAAPYAAAD7AAYABgAA//// ++gAFAAAA+wAAAQEABgAGAAD////6AAYAAAEBAAABCAAGAAYAAP////oABgAAAQgAAAEPAAYABgAA +////+gAGAAABDwAAARYABgAGAAAAAP/6AAYAAAEWAAABHAAGAAYAAP////oABgAAARwAAAEjAAYA +BgAAAAD/+gAFAAABIwAAASgABgAGAAAAAf/5AAQAAQEoAAABKwAIAAYAAAAA//kABAABASsAAAEv +AAgABgAAAAH/+QAEAAEBLwAAATIACAAGAAAAAP/5AAX//AEyAAABNwADAAYAAAAAAAEABgACATcA +AAE9AAEABgAAAAH/+QAE//wBPQAAAUAAAwAGAAAAAP/7AAYAAAFAAAABRgAFAAYAAP////kABQAA +AUYAAAFMAAcABgAAAAD/+wAFAAABTAAAAVEABQAGAAAAAP/5AAYAAAFRAAABVwAHAAYAAAAA//sA +BQAAAVcAAAFcAAUABgAAAAD/+QAFAAABXAAAAWEABwAGAAAAAP/7AAYAAgFhAAABZwAHAAYAAP// +//kABQAAAWcAAAFtAAcABgAAAAD/+QAGAAABbQAAAXMABwAGAAAAAP/5AAQAAgFzAAABdwAJAAYA +AP////kABgAAAXcAAAF+AAcABgAAAAD/+QAGAAABfgAAAYQABwAGAAD////7AAUAAAGEAAABigAF +AAYAAP////sABQAAAYoAAAGQAAUABgAAAAD/+wAFAAABkAAAAZUABQAGAAD////7AAUAAgGVAAAB +mwAHAAYAAAAA//sABgACAZsAAAGhAAcABgAAAAD/+wAGAAABoQAAAacABQAGAAAAAP/7AAYAAAGn +AAABrQAFAAYAAAAA//kABgAAAa0AAAGzAAcABgAA////+wAGAAABswAAAboABQAGAAD////7AAUA +AAG6AAABwAAFAAYAAP////sABgAAAcAAAAHHAAUABgAAAAD/+wAGAAABxwAAAc0ABQAGAAD////7 +AAYAAgHNAAAB1AAHAAYAAAAA//sABQAAAdQAAAHZAAUABgAAAAH/+QAFAAEB2QAAAd0ACAAGAAAA +Av/6AAMAAQHdAAAB3gAHAAYAAAAA//kABAABAd4AAAHiAAgABgAAAAD/+wAF//0B4gAAAecAAgAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAYAAAAB +//sAAwACAecAAAHpAAcABgAAAAD/+QAFAAEB6QAAAe4ACAAGAAAAAP/5AAYAAAHuAAAB9AAHAAYA +AAAA//oABf//AfQAAAH5AAUABgAAAAD/+QAGAAAB+QAAAf8ABwAGAAAAAv/5AAMAAgH/AAACAAAJ +AAYAAAAA//kABQABAgAAAAIFAAgABgAAAAH/+gAE//sCBQAAAggAAQAGAAAAAP/5AAYAAAIIAAAC +DgAHAAYAAAAB//kABf/+Ag4AAAISAAUABgAA////+wAGAAACEgAAAhkABQAGAAAAAP/7AAX//gIZ +AAACHgADAAYAAAAA//wABf/9Ah4AAAIjAAEABgAAAAD/+QAHAAACIwAAAioABwAGAAAAAP/6AAT/ ++wIqAAACLgABAAYAAAAA//kABP/8Ai4AAAIyAAMABgAAAAD/+gAFAAACMgAAAjcABgAGAAAAAf/5 +AAT//QI3AAACOgAEAAYAAAAB//kABP/9AjoAAAI9AAQABgAAAAL/+QAE//sCPQAAAj8AAgAGAAD/ +///7AAYAAgI/AAACRgAHAAYAAAAA//kABgABAkYAAAJMAAgABgAAAAH//AAD//0CTAAAAk4AAQAG +AAAAAf//AAQAAgJOAAACUQADAAYAAAAB//kABP/9AlEAAAJUAAQABgAAAAH/+QAF//4CVAAAAlgA +BQAGAAD////7AAYAAAJYAAACXwAFAAYAAP////kABgAAAl8AAAJmAAcABgAA////+QAGAAACZgAA +Am0ABwAGAAD////5AAYAAAJtAAACdAAHAAYAAAAA//sABQACAnQAAAJ5AAcABgAA////9wAGAAAC +eQAAAoAACQAGAAD////3AAYAAAKAAAAChwAJAAYAAP////cABgAAAocAAAKOAAkABgAA////9wAG +AAACjgAAApUACQAGAAD////4AAYAAAKVAAACnAAIAAYAAP////cABgAAApwAAAKjAAkABgAA//// ++gAGAAACowAAAqoABgAGAAAAAP/6AAUAAgKqAAACrwAIAAYAAP////cABQAAAq8AAAK1AAkABgAA +////9wAFAAACtQAAArsACQAGAAD////3AAUAAAK7AAACwQAJAAYAAP////gABQAAAsEAAALHAAgA +BgAAAAD/9wAEAAACxwAAAssACQAGAAAAAP/3AAQAAALLAAACzwAJAAYAAAAA//cABAAAAs8AAALT +AAkABgAAAAD/+AAEAAAC0wAAAtcACAAGAAD////6AAUAAALXAAAC3QAGAAYAAP////cABgAAAt0A +AALkAAkABgAAAAD/9wAFAAAC5AAAAukACQAGAAAAAP/3AAUAAALpAAAC7gAJAAYAAAAA//cABQAA +Au4AAALzAAkABgAAAAD/9wAFAAAC8wAAAvgACQAGAAAAAP/4AAUAAAL4AAAC/QAIAAYAAAAA//oA +Bf//Av0AAAMCAAUABgAA////+gAGAAADAgAAAwkABgAGAAD////3AAYAAAMJAAADEAAJAAYAAP// +//cABgAAAxAAAAMXAAkABgAA////9wAGAAADFwAAAx4ACQAGAAD////4AAYAAAAAAAoABwASAAYA +AP////cABgAAAAcACgAOABMABgAA////+gAFAAAADgAKABQAEAAGAAD////6AAYAAAAUAAoAGwAQ +AAYAAAAA//gABgAAABsACgAhABIABgAAAAD/+AAGAAAAIQAKACcAEgAGAAAAAP/4AAYAAAAnAAoA +LQASAAYAAAAA//gABgAAAC0ACgAzABIABgAAAAD/+QAGAAAAMwAKADkAEQAGAAAAAP/3AAYAAAA5 +AAoAPwATAAYAAP////sABQAAAD8ACgBFAA8ABgAAAAD/+wAFAAIARQAKAEoAEQAGAAAAAP/4AAUA +AABKAAoATwASAAYAAAAA//gABQAAAE8ACgBUABIABgAAAAD/+AAFAAAAVAAKAFkAEgAGAAAAAP/5 +AAUAAABZAAoAXgARAAYAAAAA//gABgAAAF4ACgBkABIABgAAAAD/+AAGAAAAZAAKAGoAEgAGAAAA +AP/4AAYAAABqAAoAcAASAAYAAAAA//kABgAAAHAACgB2ABEABgAAAAD/+AAFAAAAdgAKAHsAEgAG +AAD////4AAYAAAB7AAoAggASAAYAAAAA//gABQAAAIIACgCHABIABgAAAAD/+AAFAAAAhwAKAIwA +EgAGAAAAAP/4AAUAAACMAAoAkQASAAYAAAAA//gABQAAAJEACgCWABIABgAAAAD/+QAFAAAAlgAK +AJsAEQAGAAAAAP/6AAX//wCbAAoAoAAPAAYAAAAA//oABQABAKAACgClABEABgAA////+AAGAAAA +pQAKAKwAEgAGAAD////4AAYAAACsAAoAswASAAYAAP////gABgAAALMACgC6ABIABgAA////+QAG +AAAAugAKAMEAEQAGAAD////4AAYAAgDBAAoAyAAUAAYAAP////kABQACAMgACgDOABMABgAA//// ++QAGAAIAzgAKANUAEw== +""" + ) + ), + Image.open( + BytesIO( + base64.b64decode( + b""" +iVBORw0KGgoAAAANSUhEUgAAAx4AAAAUAQAAAAArMtZoAAAEwElEQVR4nABlAJr/AHVE4czCI/4u +Mc4b7vuds/xzjz5/3/7u/n9vMe7vnfH/9++vPn/xyf5zhxzjt8GHw8+2d83u8x27199/nxuQ6Od9 +M43/5z2I+9n9ZtmDBwMQECDRQw/eQIQohJXxpBCNVE6QCCAAAAD//wBlAJr/AgALyj1t/wINwq0g +LeNZUworuN1cjTPIzrTX6ofHWeo3v336qPzfEwRmBnHTtf95/fglZK5N0PDgfRTslpGBvz7LFc4F +IUXBWQGjQ5MGCx34EDFPwXiY4YbYxavpnhHFrk14CDAAAAD//wBlAJr/AgKqRooH2gAgPeggvUAA +Bu2WfgPoAwzRAABAAAAAAACQgLz/3Uv4Gv+gX7BJgDeeGP6AAAD1NMDzKHD7ANWr3loYbxsAD791 +NAADfcoIDyP44K/jv4Y63/Z+t98Ovt+ub4T48LAAAAD//wBlAJr/AuplMlADJAAAAGuAphWpqhMx +in0A/fRvAYBABPgBwBUgABBQ/sYAyv9g0bCHgOLoGAAAAAAAREAAwI7nr0ArYpow7aX8//9LaP/9 +SjdavWA8ePHeBIKB//81/83ndznOaXx379wAAAD//wBlAJr/AqDxW+D3AABAAbUh/QMnbQag/gAY +AYDAAACgtgD/gOqAAAB5IA/8AAAk+n9w0AAA8AAAmFRJuPo27ciC0cD5oeW4E7KA/wD3ECMAn2tt +y8PgwH8AfAxFzC0JzeAMtratAsC/ffwAAAD//wBlAJr/BGKAyCAA4AAAAvgeYTAwHd1kmQF5chkG +ABoMIHcL5xVpTfQbUqzlAAAErwAQBgAAEOClA5D9il08AEh/tUzdCBsXkbgACED+woQg8Si9VeqY +lODCn7lmF6NhnAEYgAAA/NMIAAAAAAD//2JgjLZgVGBg5Pv/Tvpc8hwGBjYGJADjHDrAwPzAjv/H +/Wf3PzCwtzcwHmBgYGcwbZz8wHaCAQMDOwMDQ8MCBgYOC3W7mp+f0w+wHOYxO3OG+e376hsMZjk3 +AAAAAP//YmCMY2A4wMAIN5e5gQETPD6AZisDAwMDgzSDAAPjByiHcQMDAwMDg1nOze1lByRu5/47 +c4859311AYNZzg0AAAAA//9iYGDBYihOIIMuwIjGL39/fwffA8b//xv/P2BPtzzHwCBjUQAAAAD/ +/yLFBrIBAAAA//9i1HhcwdhizX7u8NZNzyLbvT97bfrMf/QHI8evOwcSqGUJAAAA//9iYBB81iSw +pEE170Qrg5MIYydHqwdDQRMrAwcVrQAAAAD//2J4x7j9AAMDn8Q/BgYLBoaiAwwMjPdvMDBYM1Tv +oJodAAAAAP//Yqo/83+dxePWlxl3npsel9lvLfPcqlE9725C+acfVLMEAAAA//9i+s9gwCoaaGMR +evta/58PTEWzr21hufPjA8N+qlnBwAAAAAD//2JiWLci5v1+HmFXDqcnULE/MxgYGBj+f6CaJQAA +AAD//2Ji2FrkY3iYpYC5qDeGgeEMAwPDvwQBBoYvcTwOVLMEAAAA//9isDBgkP///0EOg9z35v// +Gc/eeW7BwPj5+QGZhANUswMAAAD//2JgqGBgYGBgqEMXlvhMPUsAAAAA//8iYDd1AAAAAP//AwDR +w7IkEbzhVQAAAABJRU5ErkJggg== +""" + ) + ) + ), + ) + return f + + +def load_default(size: float | None = None) -> FreeTypeFont | ImageFont: + """If FreeType support is available, load a version of Aileron Regular, + https://dotcolon.net/fonts/aileron, with a more limited character set. + + Otherwise, load a "better than nothing" font. + + .. versionadded:: 1.1.4 + + :param size: The font size of Aileron Regular. + + .. versionadded:: 10.1.0 + + :return: A font object. + """ + if isinstance(core, ModuleType) or size is not None: + return truetype( + BytesIO( + base64.b64decode( + b""" +AAEAAAAPAIAAAwBwRkZUTYwDlUAAADFoAAAAHEdERUYAqADnAAAo8AAAACRHUE9ThhmITwAAKfgAA +AduR1NVQnHxefoAACkUAAAA4k9TLzJovoHLAAABeAAAAGBjbWFw5lFQMQAAA6gAAAGqZ2FzcP//AA +MAACjoAAAACGdseWYmRXoPAAAGQAAAHfhoZWFkE18ayQAAAPwAAAA2aGhlYQboArEAAAE0AAAAJGh +tdHjjERZ8AAAB2AAAAdBsb2NhuOexrgAABVQAAADqbWF4cAC7AEYAAAFYAAAAIG5hbWUr+h5lAAAk +OAAAA6Jwb3N0D3oPTQAAJ9wAAAEKAAEAAAABGhxJDqIhXw889QALA+gAAAAA0Bqf2QAAAADhCh2h/ +2r/LgOxAyAAAAAIAAIAAAAAAAAAAQAAA8r/GgAAA7j/av9qA7EAAQAAAAAAAAAAAAAAAAAAAHQAAQ +AAAHQAQwAFAAAAAAACAAAAAQABAAAAQAAAAAAAAAADAfoBkAAFAAgCigJYAAAASwKKAlgAAAFeADI +BPgAAAAAFAAAAAAAAAAAAAAcAAAAAAAAAAAAAAABVS1dOAEAAIPsCAwL/GgDIA8oA5iAAAJMAAAAA +AhICsgAAACAAAwH0AAAAAAAAAU0AAADYAAAA8gA5AVMAVgJEAEYCRAA1AuQAKQKOAEAAsAArATsAZ +AE7AB4CMABVAkQAUADc/+EBEgAgANwAJQEv//sCRAApAkQAggJEADwCRAAtAkQAIQJEADkCRAArAk +QAMgJEACwCRAAxANwAJQDc/+ECRABnAkQAUAJEAEQB8wAjA1QANgJ/AB0CcwBkArsALwLFAGQCSwB +kAjcAZALGAC8C2gBkAQgAZAIgADcCYQBkAj8AZANiAGQCzgBkAuEALwJWAGQC3QAvAmsAZAJJADQC +ZAAiAqoAXgJuACADuAAaAnEAGQJFABMCTwAuATMAYgEv//sBJwAiAkQAUAH0ADIBLAApAhMAJAJjA +EoCEQAeAmcAHgIlAB4BIgAVAmcAHgJRAEoA7gA+AOn/8wIKAEoA9wBGA1cASgJRAEoCSgAeAmMASg +JnAB4BSgBKAcsAGAE5ABQCUABCAgIAAQMRAAEB4v/6AgEAAQHOABQBLwBAAPoAYAEvACECRABNA0Y +AJAItAHgBKgAcAkQAUAEsAHQAygAgAi0AOQD3ADYA9wAWAaEANgGhABYCbAAlAYMAeAGDADkA6/9q +AhsAFAIKABUB/QAVAAAAAwAAAAMAAAAcAAEAAAAAAKQAAwABAAAAHAAEAIgAAAAeABAAAwAOAH4Aq +QCrALEAtAC3ALsgGSAdICYgOiBEISL7Av//AAAAIACpAKsAsAC0ALcAuyAYIBwgJiA5IEQhIvsB// +//4/+5/7j/tP+y/7D/reBR4E/gR+A14CzfTwVxAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAEGAAABAAAAAAAAAAECAAAAAgAAAAAAAAAAAAAAAAAAAAEAAAMEBQYHCAkKCwwNDg8QERIT +FBUWFxgZGhscHR4fICEiIyQlJicoKSorLC0uLzAxMjM0NTY3ODk6Ozw9Pj9AQUJDREVGR0hJSktMT +U5PUFFSU1RVVldYWVpbXF1eX2BhAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAGQAAA +AAAAAAYnFmAAAAAABlAAAAAAAAAAAAAAAAAAAAAAAAAAAAY2htAAAAAAAAAABrbGlqAAAAAHAAbm9 +ycwBnAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAmACYAJgAmAD4AUgCCAMoBCgFO +AVwBcgGIAaYBvAHKAdYB6AH2AgwCIAJKAogCpgLWAw4DIgNkA5wDugPUA+gD/AQQBEYEogS8BPoFJ +gVSBWoFgAWwBcoF1gX6BhQGJAZMBmgGiga0BuIHGgdUB2YHkAeiB8AH3AfyCAoIHAgqCDoITghcCG +oIogjSCPoJKglYCXwJwgnqCgIKKApACl4Klgq8CtwLDAs8C1YLjAuyC9oL7gwMDCYMSAxgDKAMrAz +qDQoNTA1mDYQNoA2uDcAN2g3oDfYODA4iDkoOXA5sDnoOnA7EDvwAAAAFAAAAAAH0ArwAAwAGAAkA +DAAPAAAxESERAxMhExcRASELARETAfT6qv6syKr+jgFUqsiqArz9RAGLAP/+1P8B/v3VAP8BLP4CA +P8AAgA5//IAuQKyAAMACwAANyMDMwIyFhQGIiY0oE4MZk84JCQ4JLQB/v3AJDgkJDgAAgBWAeUBPA +LfAAMABwAAEyMnMxcjJzOmRgpagkYKWgHl+vr6AAAAAAIARgAAAf4CsgAbAB8AAAEHMxUjByM3Iwc +jNyM1MzcjNTM3MwczNzMHMxUrAQczAZgdZXEvOi9bLzovWmYdZXEvOi9bLzovWp9bHlsBn4w429vb +2ziMONvb29s4jAAAAAMANf+mAg4DDAAfACYALAAAJRQGBxUjNS4BJzMeARcRLgE0Njc1MxUeARcjJ +icVHgEBFBYXNQ4BExU+ATU0Ag5xWDpgcgRcBz41Xl9oVTpVYwpcC1ttXP6cLTQuM5szOrVRZwlOTQ +ZqVzZECAEAGlukZAlOTQdrUG8O7iNlAQgxNhDlCDj+8/YGOjReAAAAAAUAKf/yArsCvAAHAAsAFQA +dACcAABIyFhQGIiY0EyMBMwQiBhUUFjI2NTQSMhYUBiImNDYiBhUUFjI2NTR5iFBQiFCVVwHAV/5c +OiMjOiPmiFBQiFCxOiMjOiMCvFaSVlaS/ZoCsjIzMC80NC8w/uNWklZWkhozMC80NC8wAAAAAgBA/ +/ICbgLAACIALgAAARUjEQYjIiY1NDY3LgE1NDYzMhcVJiMiBhUUFhcWOwE1MxUFFBYzMjc1IyIHDg +ECbmBcYYOOVkg7R4hsQjY4Q0RNRD4SLDxW/pJUXzksPCkUUk0BgUb+zBVUZ0BkDw5RO1huCkULQzp +COAMBcHDHRz0J/AIHRQAAAAEAKwHlAIUC3wADAAATIycze0YKWgHl+gAAAAABAGT/sAEXAwwACQAA +EzMGEBcjLgE0Nt06dXU6OUBAAwzG/jDGVePs4wAAAAEAHv+wANEDDAAJAAATMx4BFAYHIzYQHjo5Q +EA5OnUDDFXj7ONVxgHQAAAAAQBVAFIB2wHbAA4AAAE3FwcXBycHJzcnNxcnMwEtmxOfcTJjYzJxnx +ObCj4BKD07KYolmZkliik7PbMAAQBQAFUB9AIlAAsAAAEjFSM1IzUzNTMVMwH0tTq1tTq1AR/Kyjj +OzgAAAAAB/+H/iACMAGQABAAANwcjNzOMWlFOXVrS3AAAAQAgAP8A8gE3AAMAABMjNTPy0tIA/zgA +AQAl//IApQByAAcAADYyFhQGIiY0STgkJDgkciQ4JCQ4AAAAAf/7/+IBNALQAAMAABcjEzM5Pvs+H +gLuAAAAAAIAKf/yAhsCwAADAAcAABIgECA2IBAgKQHy/g5gATL+zgLA/TJEAkYAAAAAAQCCAAABlg +KyAAgAAAERIxEHNTc2MwGWVr6SIygCsv1OAldxW1sWAAEAPAAAAg4CwAAZAAA3IRUhNRM+ATU0JiM +iDwEjNz4BMzIWFRQGB7kBUv4x+kI2QTt+EAFWAQp8aGVtSl5GRjEA/0RVLzlLmAoKa3FsUkNxXQAA +AAEALf/yAhYCwAAqAAABHgEVFAYjIi8BMxceATMyNjU0KwE1MzI2NTQmIyIGDwEjNz4BMzIWFRQGA +YxBSZJo2RUBVgEHV0JBUaQREUBUQzc5TQcBVgEKfGhfcEMBbxJbQl1x0AoKRkZHPn9GSD80QUVCCg +pfbGBPOlgAAAACACEAAAIkArIACgAPAAAlIxUjNSE1ATMRMyMRBg8BAiRXVv6qAVZWV60dHLCurq4 +rAdn+QgFLMibzAAABADn/8gIZArIAHQAAATIWFRQGIyIvATMXFjMyNjU0JiMiByMTIRUhBzc2ATNv +d5Fl1RQBVgIad0VSTkVhL1IwAYj+vh8rMAHHgGdtgcUKCoFXTU5bYgGRRvAuHQAAAAACACv/8gITA +sAAFwAjAAABMhYVFAYjIhE0NjMyFh8BIycmIyIDNzYTMjY1NCYjIgYVFBYBLmp7imr0l3RZdAgBXA +IYZ5wKJzU6QVNJSz5SUAHSgWltiQFGxcNlVQoKdv7sPiz+ZF1LTmJbU0lhAAAAAQAyAAACGgKyAAY +AAAEVASMBITUCGv6oXAFL/oECsij9dgJsRgAAAAMALP/xAhgCwAAWACAALAAAAR4BFRQGIyImNTQ2 +Ny4BNTQ2MhYVFAYmIgYVFBYyNjU0AzI2NTQmIyIGFRQWAZQ5S5BmbIpPOjA7ecp5P2F8Q0J8RIVJS +0pLTEtOAW0TXTxpZ2ZqPF0SE1A3VWVlVTdQ/UU0N0RENzT9/ko+Ok1NOj1LAAIAMf/yAhkCwAAXAC +MAAAEyERQGIyImLwEzFxYzMhMHBiMiJjU0NhMyNjU0JiMiBhUUFgEl9Jd0WXQIAVwCGGecCic1SWp +7imo+UlBAQVNJAsD+usXDZVUKCnYBFD4sgWltif5kW1NJYV1LTmIAAAACACX/8gClAiAABwAPAAAS +MhYUBiImNBIyFhQGIiY0STgkJDgkJDgkJDgkAiAkOCQkOP52JDgkJDgAAAAC/+H/iAClAiAABwAMA +AASMhYUBiImNBMHIzczSTgkJDgkaFpSTl4CICQ4JCQ4/mba5gAAAQBnAB4B+AH0AAYAAAENARUlNS +UB+P6qAVb+bwGRAbCmpkbJRMkAAAIAUAC7AfQBuwADAAcAAAEhNSERITUhAfT+XAGk/lwBpAGDOP8 +AOAABAEQAHgHVAfQABgAAARUFNS0BNQHV/m8BVv6qAStEyUSmpkYAAAAAAgAj//IB1ALAABgAIAAA +ATIWFRQHDgEHIz4BNz4BNTQmIyIGByM+ARIyFhQGIiY0AQRibmktIAJWBSEqNig+NTlHBFoDezQ4J +CQ4JALAZ1BjaS03JS1DMD5LLDQ/SUVgcv2yJDgkJDgAAAAAAgA2/5gDFgKYADYAQgAAAQMGFRQzMj +Y1NCYjIg4CFRQWMzI2NxcGIyImNTQ+AjMyFhUUBiMiJwcGIyImNTQ2MzIfATcHNzYmIyIGFRQzMjY +Cej8EJjJJlnBAfGQ+oHtAhjUYg5OPx0h2k06Os3xRWQsVLjY5VHtdPBwJETcJDyUoOkZEJz8B0f74 +EQ8kZl6EkTFZjVOLlyknMVm1pmCiaTq4lX6CSCknTVRmmR8wPdYnQzxuSWVGAAIAHQAAAncCsgAHA +AoAACUjByMTMxMjATMDAcj+UVz4dO5d/sjPZPT0ArL9TgE6ATQAAAADAGQAAAJMArIAEAAbACcAAA +EeARUUBgcGKwERMzIXFhUUJRUzMjc2NTQnJiMTPgE1NCcmKwEVMzIBvkdHZkwiNt7LOSGq/oeFHBt +hahIlSTM+cB8Yj5UWAW8QT0VYYgwFArIEF5Fv1eMED2NfDAL93AU+N24PBP0AAAAAAQAv//ICjwLA +ABsAAAEyFh8BIycmIyIGFRQWMzI/ATMHDgEjIiY1NDYBdX+PCwFWAiKiaHx5ZaIiAlYBCpWBk6a0A +sCAagoKpqN/gaOmCgplhcicn8sAAAIAZAAAAp8CsgAMABkAAAEeARUUBgcGKwERMzITPgE1NCYnJi +sBETMyAY59lJp8IzXN0jUVWmdjWRs5d3I4Aq4QqJWUug8EArL9mQ+PeHGHDgX92gAAAAABAGQAAAI +vArIACwAAJRUhESEVIRUhFSEVAi/+NQHB/pUBTf6zRkYCskbwRvAAAAABAGQAAAIlArIACQAAExUh +FSERIxEhFboBQ/69VgHBAmzwRv7KArJGAAAAAAEAL//yAo8CwAAfAAABMxEjNQcGIyImNTQ2MzIWH +wEjJyYjIgYVFBYzMjY1IwGP90wfPnWTprSSf48LAVYCIqJofHllVG+hAU3+s3hARsicn8uAagoKpq +N/gaN1XAAAAAEAZAAAAowCsgALAAABESMRIREjETMRIRECjFb+hFZWAXwCsv1OAS7+0gKy/sQBPAA +AAAABAGQAAAC6ArIAAwAAMyMRM7pWVgKyAAABADf/8gHoArIAEwAAAREUBw4BIyImLwEzFxYzMjc2 +NREB6AIFcGpgbQIBVgIHfXQKAQKy/lYxIltob2EpKYyEFD0BpwAAAAABAGQAAAJ0ArIACwAACQEjA +wcVIxEzEQEzATsBJ3ntQlZWAVVlAWH+nwEnR+ACsv6RAW8AAQBkAAACLwKyAAUAACUVIREzEQIv/j +VWRkYCsv2UAAABAGQAAAMUArIAFAAAAREjETQ3BgcDIwMmJxYVESMRMxsBAxRWAiMxemx8NxsCVo7 +MywKy/U4BY7ZLco7+nAFmoFxLtP6dArL9lwJpAAAAAAEAZAAAAoACsgANAAAhIwEWFREjETMBJjUR +MwKAhP67A1aEAUUDVAJeeov+pwKy/aJ5jAFZAAAAAgAv//ICuwLAAAkAEwAAEiAWFRQGICY1NBIyN +jU0JiIGFRTbATSsrP7MrNrYenrYegLAxaKhxsahov47nIeIm5uIhwACAGQAAAJHArIADgAYAAABHg +EVFAYHBisBESMRMzITNjQnJisBETMyAZRUX2VOHzuAVtY7GlxcGDWIiDUCrgtnVlVpCgT+5gKy/rU +V1BUF/vgAAAACAC//zAK9AsAAEgAcAAAlFhcHJiMiBwYjIiY1NDYgFhUUJRQWMjY1NCYiBgI9PUMx +UDcfKh8omqysATSs/dR62Hp62HpICTg7NgkHxqGixcWitbWHnJyHiJubAAIAZAAAAlgCsgAXACMAA +CUWFyMmJyYnJisBESMRMzIXHgEVFAYHFiUzMjc+ATU0JyYrAQIqDCJfGQwNWhAhglbiOx9QXEY1Tv +6bhDATMj1lGSyMtYgtOXR0BwH+1wKyBApbU0BSESRAAgVAOGoQBAABADT/8gIoAsAAJQAAATIWFyM +uASMiBhUUFhceARUUBiMiJiczHgEzMjY1NCYnLgE1NDYBOmd2ClwGS0E6SUNRdW+HZnKKC1wPWkQ9 +Uk1cZGuEAsBwXUJHNjQ3OhIbZVZZbm5kREo+NT5DFRdYUFdrAAAAAAEAIgAAAmQCsgAHAAABIxEjE +SM1IQJk9lb2AkICbP2UAmxGAAEAXv/yAmQCsgAXAAABERQHDgEiJicmNREzERQXHgEyNjc2NRECZA +IIgfCBCAJWAgZYmlgGAgKy/k0qFFxzc1wUKgGz/lUrEkRQUEQSKwGrAAAAAAEAIAAAAnoCsgAGAAA +hIwMzGwEzAYJ07l3N1FwCsv2PAnEAAAEAGgAAA7ECsgAMAAABAyMLASMDMxsBMxsBA7HAcZyicrZi +kaB0nJkCsv1OAlP9rQKy/ZsCW/2kAmYAAAEAGQAAAm8CsgALAAAhCwEjEwMzGwEzAxMCCsrEY/bkY +re+Y/D6AST+3AFcAVb+5gEa/q3+oQAAAQATAAACUQKyAAgAAAERIxEDMxsBMwFdVvRjwLphARD+8A +EQAaL+sQFPAAABAC4AAAI5ArIACQAAJRUhNQEhNSEVAQI5/fUBof57Aen+YUZGQgIqRkX92QAAAAA +BAGL/sAEFAwwABwAAARUjETMVIxEBBWlpowMMOP0UOANcAAAB//v/4gE0AtAAAwAABSMDMwE0Pvs+ +HgLuAAAAAQAi/7AAxQMMAAcAABcjNTMRIzUzxaNpaaNQOALsOAABAFAA1wH0AmgABgAAJQsBIxMzE +wGwjY1GsESw1wFZ/qcBkf5vAAAAAQAy/6oBwv/iAAMAAAUhNSEBwv5wAZBWOAAAAAEAKQJEALYCsg +ADAAATIycztjhVUAJEbgAAAAACACT/8gHQAiAAHQAlAAAhJwcGIyImNTQ2OwE1NCcmIyIHIz4BMzI +XFh0BFBcnMjY9ASYVFAF6CR0wVUtgkJoiAgdgaQlaBm1Zrg4DCuQ9R+5MOSFQR1tbDiwUUXBUXowf +J8c9SjRORzYSgVwAAAAAAgBK//ICRQLfABEAHgAAATIWFRQGIyImLwEVIxEzETc2EzI2NTQmIyIGH +QEUFgFUcYCVbiNJEyNWVigySElcU01JXmECIJd4i5QTEDRJAt/+3jkq/hRuZV55ZWsdX14AAQAe// +IB9wIgABgAAAEyFhcjJiMiBhUUFjMyNjczDgEjIiY1NDYBF152DFocbEJXU0A1Rw1aE3pbaoKQAiB +oWH5qZm1tPDlaXYuLgZcAAAACAB7/8gIZAt8AEQAeAAABESM1BwYjIiY1NDYzMhYfAREDMjY9ATQm +IyIGFRQWAhlWKDJacYCVbiNJEyOnSV5hQUlcUwLf/SFVOSqXeIuUExA0ARb9VWVrHV9ebmVeeQACA +B7/8gH9AiAAFQAbAAABFAchHgEzMjY3Mw4BIyImNTQ2MzIWJyIGByEmAf0C/oAGUkA1SwlaD4FXbI +WObmt45UBVBwEqDQEYFhNjWD84W16Oh3+akU9aU60AAAEAFQAAARoC8gAWAAATBh0BMxUjESMRIzU +zNTQ3PgEzMhcVJqcDbW1WOTkDB0k8Hx5oAngVITRC/jQBzEIsJRs5PwVHEwAAAAIAHv8uAhkCIAAi +AC8AAAERFAcOASMiLwEzFx4BMzI2NzY9AQcGIyImNTQ2MzIWHwE1AzI2PQE0JiMiBhUUFgIZAQSEd +NwRAVcBBU5DTlUDASgyWnGAlW4jSRMjp0leYUFJXFMCEv5wSh1zeq8KCTI8VU0ZIQk5Kpd4i5QTED +RJ/iJlax1fXm5lXnkAAQBKAAACCgLkABcAAAEWFREjETQnLgEHDgEdASMRMxE3NjMyFgIIAlYCBDs +6RVRWViE5UVViAYUbQP7WASQxGzI7AQJyf+kC5P7TPSxUAAACAD4AAACsAsAABwALAAASMhYUBiIm +NBMjETNeLiAgLiBiVlYCwCAuICAu/WACEgAC//P/LgCnAsAABwAVAAASMhYUBiImNBcRFAcGIyInN +RY3NjURWS4gIC4gYgMLcRwNSgYCAsAgLiAgLo79wCUbZAJGBzMOHgJEAAAAAQBKAAACCALfAAsAAC +EnBxUjETMREzMHEwGTwTJWVvdu9/rgN6kC3/4oAQv6/ugAAQBG//wA3gLfAA8AABMRFBceATcVBiM +iJicmNRGcAQIcIxkkKi4CAQLf/bkhERoSBD4EJC8SNAJKAAAAAQBKAAADEAIgACQAAAEWFREjETQn +JiMiFREjETQnJiMiFREjETMVNzYzMhYXNzYzMhYDCwVWBAxedFYEDF50VlYiJko7ThAvJkpEVAGfI +jn+vAEcQyRZ1v76ARxDJFnW/voCEk08HzYtRB9HAAAAAAEASgAAAgoCIAAWAAABFhURIxE0JyYjIg +YdASMRMxU3NjMyFgIIAlYCCXBEVVZWITlRVWIBhRtA/tYBJDEbbHR/6QISWz0sVAAAAAACAB7/8gI +sAiAABwARAAASIBYUBiAmNBIyNjU0JiIGFRSlAQCHh/8Ah7ieWlqeWgIgn/Cfn/D+s3ZfYHV1YF8A +AgBK/zwCRQIgABEAHgAAATIWFRQGIyImLwERIxEzFTc2EzI2NTQmIyIGHQEUFgFUcYCVbiNJEyNWV +igySElcU01JXmECIJd4i5QTEDT+8wLWVTkq/hRuZV55ZWsdX14AAgAe/zwCGQIgABEAHgAAAREjEQ +cGIyImNTQ2MzIWHwE1AzI2PQE0JiMiBhUUFgIZVigyWnGAlW4jSRMjp0leYUFJXFMCEv0qARk5Kpd +4i5QTEDRJ/iJlax1fXm5lXnkAAQBKAAABPgIeAA0AAAEyFxUmBhURIxEzFTc2ARoWDkdXVlYwIwIe +B0EFVlf+0gISU0cYAAEAGP/yAa0CIAAjAAATMhYXIyYjIgYVFBYXHgEVFAYjIiYnMxYzMjY1NCYnL +gE1NDbkV2MJWhNdKy04PF1XbVhWbgxaE2ktOjlEUllkAiBaS2MrJCUoEBlPQkhOVFZoKCUmLhIWSE +BIUwAAAAEAFP/4ARQCiQAXAAATERQXHgE3FQYjIiYnJjURIzUzNTMVMxWxAQMmMx8qMjMEAUdHVmM +BzP7PGw4mFgY/BSwxDjQBNUJ7e0IAAAABAEL/8gICAhIAFwAAAREjNQcGIyImJyY1ETMRFBceATMy +Nj0BAgJWITlRT2EKBVYEBkA1RFECEv3uWj4qTToiOQE+/tIlJC43c4DpAAAAAAEAAQAAAfwCEgAGA +AABAyMDMxsBAfzJaclfop8CEv3uAhL+LQHTAAABAAEAAAMLAhIADAAAAQMjCwEjAzMbATMbAQMLqW +Z2dmapY3t0a3Z7AhL97gG+/kICEv5AAcD+QwG9AAAB//oAAAHWAhIACwAAARMjJwcjEwMzFzczARq +8ZIuKY763ZoWFYwEO/vLV1QEMAQbNzQAAAQAB/y4B+wISABEAAAEDDgEjIic1FjMyNj8BAzMbAQH7 +2iFZQB8NDRIpNhQH02GenQIS/cFVUAJGASozEwIt/i4B0gABABQAAAGxAg4ACQAAJRUhNQEhNSEVA +QGx/mMBNP7iAYL+zkREQgGIREX+ewAAAAABAED/sAEOAwwALAAAASMiBhUUFxYVFAYHHgEVFAcGFR +QWOwEVIyImNTQ3NjU0JzU2NTQnJjU0NjsBAQ4MKiMLDS4pKS4NCyMqDAtERAwLUlILDERECwLUGBk +WTlsgKzUFBTcrIFtOFhkYOC87GFVMIkUIOAhFIkxVGDsvAAAAAAEAYP84AJoDIAADAAAXIxEzmjo6 +yAPoAAEAIf+wAO8DDAAsAAATFQYVFBcWFRQGKwE1MzI2NTQnJjU0NjcuATU0NzY1NCYrATUzMhYVF +AcGFRTvUgsMREQLDCojCw0uKSkuDQsjKgwLREQMCwF6OAhFIkxVGDsvOBgZFk5bICs1BQU3KyBbTh +YZGDgvOxhVTCJFAAABAE0A3wH2AWQAEwAAATMUIyImJyYjIhUjNDMyFhcWMzIBvjhuGywtQR0xOG4 +bLC1BHTEBZIURGCNMhREYIwAAAwAk/94DIgLoAAcAEQApAAAAIBYQBiAmECQgBhUUFiA2NTQlMhYX +IyYjIgYUFjMyNjczDgEjIiY1NDYBAQFE3d3+vN0CB/7wubkBELn+xVBnD1wSWDo+QTcqOQZcEmZWX +HN2Aujg/rbg4AFKpr+Mjb6+jYxbWEldV5ZZNShLVn5na34AAgB4AFIB9AGeAAUACwAAAQcXIyc3Mw +cXIyc3AUqJiUmJifOJiUmJiQGepqampqampqYAAAIAHAHSAQ4CwAAHAA8AABIyFhQGIiY0NiIGFBY +yNjRgakREakSTNCEhNCECwEJqQkJqCiM4IyM4AAAAAAIAUAAAAfQCCwALAA8AAAEzFSMVIzUjNTM1 +MxMhNSEBP7W1OrW1OrX+XAGkAVs4tLQ4sP31OAAAAQB0AkQBAQKyAAMAABMjNzOsOD1QAkRuAAAAA +AEAIADsAKoBdgAHAAASMhYUBiImNEg6KCg6KAF2KDooKDoAAAIAOQBSAbUBngAFAAsAACUHIzcnMw +UHIzcnMwELiUmJiUkBM4lJiYlJ+KampqampqYAAAABADYB5QDhAt8ABAAAEzczByM2Xk1OXQHv8Po +AAQAWAeUAwQLfAAQAABMHIzczwV5NTl0C1fD6AAIANgHlAYsC3wAEAAkAABM3MwcjPwEzByM2Xk1O +XapeTU5dAe/w+grw+gAAAgAWAeUBawLfAAQACQAAEwcjNzMXByM3M8FeTU5dql5NTl0C1fD6CvD6A +AADACX/8gI1AHIABwAPABcAADYyFhQGIiY0NjIWFAYiJjQ2MhYUBiImNEk4JCQ4JOw4JCQ4JOw4JC +Q4JHIkOCQkOCQkOCQkOCQkOCQkOAAAAAEAeABSAUoBngAFAAABBxcjJzcBSomJSYmJAZ6mpqamAAA +AAAEAOQBSAQsBngAFAAAlByM3JzMBC4lJiYlJ+KampgAAAf9qAAABgQKyAAMAACsBATM/VwHAVwKy +AAAAAAIAFAHIAdwClAAHABQAABMVIxUjNSM1BRUjNwcjJxcjNTMXN9pKMkoByDICKzQqATJLKysCl +CmjoykBy46KiY3Lm5sAAQAVAAABvALyABgAAAERIxEjESMRIzUzNTQ3NjMyFxUmBgcGHQEBvFbCVj +k5AxHHHx5iVgcDAg798gHM/jQBzEIOJRuWBUcIJDAVIRYAAAABABX//AHkAvIAJQAAJR4BNxUGIyI +mJyY1ESYjIgcGHQEzFSMRIxEjNTM1NDc2MzIXERQBowIcIxkkKi4CAR4nXgwDbW1WLy8DEbNdOmYa +EQQ/BCQvEjQCFQZWFSEWQv40AcxCDiUblhP9uSEAAAAAAAAWAQ4AAQAAAAAAAAATACgAAQAAAAAAA +QAHAEwAAQAAAAAAAgAHAGQAAQAAAAAAAwAaAKIAAQAAAAAABAAHAM0AAQAAAAAABQA8AU8AAQAAAA +AABgAPAawAAQAAAAAACAALAdQAAQAAAAAACQALAfgAAQAAAAAACwAXAjQAAQAAAAAADAAXAnwAAwA +BBAkAAAAmAAAAAwABBAkAAQAOADwAAwABBAkAAgAOAFQAAwABBAkAAwA0AGwAAwABBAkABAAOAL0A +AwABBAkABQB4ANUAAwABBAkABgAeAYwAAwABBAkACAAWAbwAAwABBAkACQAWAeAAAwABBAkACwAuA +gQAAwABBAkADAAuAkwATgBvACAAUgBpAGcAaAB0AHMAIABSAGUAcwBlAHIAdgBlAGQALgAATm8gUm +lnaHRzIFJlc2VydmVkLgAAQQBpAGwAZQByAG8AbgAAQWlsZXJvbgAAUgBlAGcAdQBsAGEAcgAAUmV +ndWxhcgAAMQAuADEAMAAyADsAVQBLAFcATgA7AEEAaQBsAGUAcgBvAG4ALQBSAGUAZwB1AGwAYQBy +AAAxLjEwMjtVS1dOO0FpbGVyb24tUmVndWxhcgAAQQBpAGwAZQByAG8AbgAAQWlsZXJvbgAAVgBlA +HIAcwBpAG8AbgAgADEALgAxADAAMgA7AFAAUwAgADAAMAAxAC4AMQAwADIAOwBoAG8AdABjAG8Abg +B2ACAAMQAuADAALgA3ADAAOwBtAGEAawBlAG8AdABmAC4AbABpAGIAMgAuADUALgA1ADgAMwAyADk +AAFZlcnNpb24gMS4xMDI7UFMgMDAxLjEwMjtob3Rjb252IDEuMC43MDttYWtlb3RmLmxpYjIuNS41 +ODMyOQAAQQBpAGwAZQByAG8AbgAtAFIAZQBnAHUAbABhAHIAAEFpbGVyb24tUmVndWxhcgAAUwBvA +HIAYQAgAFMAYQBnAGEAbgBvAABTb3JhIFNhZ2FubwAAUwBvAHIAYQAgAFMAYQBnAGEAbgBvAABTb3 +JhIFNhZ2FubwAAaAB0AHQAcAA6AC8ALwB3AHcAdwAuAGQAbwB0AGMAbwBsAG8AbgAuAG4AZQB0AAB +odHRwOi8vd3d3LmRvdGNvbG9uLm5ldAAAaAB0AHQAcAA6AC8ALwB3AHcAdwAuAGQAbwB0AGMAbwBs +AG8AbgAuAG4AZQB0AABodHRwOi8vd3d3LmRvdGNvbG9uLm5ldAAAAAACAAAAAAAA/4MAMgAAAAAAA +AAAAAAAAAAAAAAAAAAAAHQAAAABAAIAAwAEAAUABgAHAAgACQAKAAsADAANAA4ADwAQABEAEgATAB +QAFQAWABcAGAAZABoAGwAcAB0AHgAfACAAIQAiACMAJAAlACYAJwAoACkAKgArACwALQAuAC8AMAA +xADIAMwA0ADUANgA3ADgAOQA6ADsAPAA9AD4APwBAAEEAQgBDAEQARQBGAEcASABJAEoASwBMAE0A +TgBPAFAAUQBSAFMAVABVAFYAVwBYAFkAWgBbAFwAXQBeAF8AYABhAIsAqQCDAJMAjQDDAKoAtgC3A +LQAtQCrAL4AvwC8AIwAwADBAAAAAAAB//8AAgABAAAADAAAABwAAAACAAIAAwBxAAEAcgBzAAIABA +AAAAIAAAABAAAACgBMAGYAAkRGTFQADmxhdG4AGgAEAAAAAP//AAEAAAAWAANDQVQgAB5NT0wgABZ +ST00gABYAAP//AAEAAAAA//8AAgAAAAEAAmxpZ2EADmxvY2wAFAAAAAEAAQAAAAEAAAACAAYAEAAG +AAAAAgASADQABAAAAAEATAADAAAAAgAQABYAAQAcAAAAAQABAE8AAQABAGcAAQABAE8AAwAAAAIAE +AAWAAEAHAAAAAEAAQAvAAEAAQBnAAEAAQAvAAEAGgABAAgAAgAGAAwAcwACAE8AcgACAEwAAQABAE +kAAAABAAAACgBGAGAAAkRGTFQADmxhdG4AHAAEAAAAAP//AAIAAAABABYAA0NBVCAAFk1PTCAAFlJ +PTSAAFgAA//8AAgAAAAEAAmNwc3AADmtlcm4AFAAAAAEAAAAAAAEAAQACAAYADgABAAAAAQASAAIA +AAACAB4ANgABAAoABQAFAAoAAgABACQAPQAAAAEAEgAEAAAAAQAMAAEAOP/nAAEAAQAkAAIGigAEA +AAFJAXKABoAGQAA//gAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAD/sv+4/+z/7v/MAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAD/xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/9T/6AAAAAD/8QAA +ABD/vQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/7gAAAAAAAAAAAAAAAAAA//MAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABIAAAAAAAAAAP/5AAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP/gAAD/4AAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA//L/9AAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAA/+gAAAAAAAkAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP/zAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP/mAAAAAAAAAAAAAAAAAAD +/4gAA//AAAAAA//YAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/+AAAAAAAAP/OAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/zv/qAAAAAP/0AAAACAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP/ZAAD/egAA/1kAAAAA/5D/rgAAAAAAAAAAAA +AAAAAAAAAAAAAAAAD/9AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAD/8AAA/7b/8P+wAAD/8P/E/98AAAAA/8P/+P/0//oAAAAAAAAAAAAA//gA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/+AAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/w//C/9MAAP/SAAD/9wAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAD/yAAA/+kAAAAA//QAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/9wAAAAD//QAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAP/2AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAP/cAAAAAAAAAAAAAAAA/7YAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAP/8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/6AAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAkAFAAEAAAAAQACwAAABcA +BgAAAAAAAAAIAA4AAAAAAAsAEgAAAAAAAAATABkAAwANAAAAAQAJAAAAAAAAAAAAAAAAAAAAGAAAA +AAABwAAAAAAAAAAAAAAFQAFAAAAAAAYABgAAAAUAAAACgAAAAwAAgAPABEAFgAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFAAEAEQBdAAYAAAAAAAAAAAAAAAAAAAAAAAA +AAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAcAAAAAAAAABwAAAAAACAAAAAAAAAAAAAcAAAAHAAAAEwAJ +ABUADgAPAAAACwAQAAAAAAAAAAAAAAAAAAUAGAACAAIAAgAAAAIAGAAXAAAAGAAAABYAFgACABYAA +gAWAAAAEQADAAoAFAAMAA0ABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAASAAAAEgAGAAEAHgAkAC +YAJwApACoALQAuAC8AMgAzADcAOAA5ADoAPAA9AEUASABOAE8AUgBTAFUAVwBZAFoAWwBcAF0AcwA +AAAAAAQAAAADa3tfFAAAAANAan9kAAAAA4QodoQ== +""" + ) + ), + 10 if size is None else size, + layout_engine=Layout.BASIC, + ) + return load_default_imagefont() diff --git a/venv/lib/python3.10/site-packages/PIL/ImageGrab.py b/venv/lib/python3.10/site-packages/PIL/ImageGrab.py new file mode 100644 index 0000000000000000000000000000000000000000..1eb4507344cd412bed5c113d4c52e39f26cecab2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageGrab.py @@ -0,0 +1,196 @@ +# +# The Python Imaging Library +# $Id$ +# +# screen grabber +# +# History: +# 2001-04-26 fl created +# 2001-09-17 fl use builtin driver, if present +# 2002-11-19 fl added grabclipboard support +# +# Copyright (c) 2001-2002 by Secret Labs AB +# Copyright (c) 2001-2002 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import os +import shutil +import subprocess +import sys +import tempfile + +from . import Image + +TYPE_CHECKING = False +if TYPE_CHECKING: + from . import ImageWin + + +def grab( + bbox: tuple[int, int, int, int] | None = None, + include_layered_windows: bool = False, + all_screens: bool = False, + xdisplay: str | None = None, + window: int | ImageWin.HWND | None = None, +) -> Image.Image: + im: Image.Image + if xdisplay is None: + if sys.platform == "darwin": + fh, filepath = tempfile.mkstemp(".png") + os.close(fh) + args = ["screencapture"] + if bbox: + left, top, right, bottom = bbox + args += ["-R", f"{left},{top},{right-left},{bottom-top}"] + subprocess.call(args + ["-x", filepath]) + im = Image.open(filepath) + im.load() + os.unlink(filepath) + if bbox: + im_resized = im.resize((right - left, bottom - top)) + im.close() + return im_resized + return im + elif sys.platform == "win32": + if window is not None: + all_screens = -1 + offset, size, data = Image.core.grabscreen_win32( + include_layered_windows, + all_screens, + int(window) if window is not None else 0, + ) + im = Image.frombytes( + "RGB", + size, + data, + # RGB, 32-bit line padding, origin lower left corner + "raw", + "BGR", + (size[0] * 3 + 3) & -4, + -1, + ) + if bbox: + x0, y0 = offset + left, top, right, bottom = bbox + im = im.crop((left - x0, top - y0, right - x0, bottom - y0)) + return im + # Cast to Optional[str] needed for Windows and macOS. + display_name: str | None = xdisplay + try: + if not Image.core.HAVE_XCB: + msg = "Pillow was built without XCB support" + raise OSError(msg) + size, data = Image.core.grabscreen_x11(display_name) + except OSError: + if display_name is None and sys.platform not in ("darwin", "win32"): + if shutil.which("gnome-screenshot"): + args = ["gnome-screenshot", "-f"] + elif shutil.which("grim"): + args = ["grim"] + elif shutil.which("spectacle"): + args = ["spectacle", "-n", "-b", "-f", "-o"] + else: + raise + fh, filepath = tempfile.mkstemp(".png") + os.close(fh) + subprocess.call(args + [filepath]) + im = Image.open(filepath) + im.load() + os.unlink(filepath) + if bbox: + im_cropped = im.crop(bbox) + im.close() + return im_cropped + return im + else: + raise + else: + im = Image.frombytes("RGB", size, data, "raw", "BGRX", size[0] * 4, 1) + if bbox: + im = im.crop(bbox) + return im + + +def grabclipboard() -> Image.Image | list[str] | None: + if sys.platform == "darwin": + p = subprocess.run( + ["osascript", "-e", "get the clipboard as «class PNGf»"], + capture_output=True, + ) + if p.returncode != 0: + return None + + import binascii + + data = io.BytesIO(binascii.unhexlify(p.stdout[11:-3])) + return Image.open(data) + elif sys.platform == "win32": + fmt, data = Image.core.grabclipboard_win32() + if fmt == "file": # CF_HDROP + import struct + + o = struct.unpack_from("I", data)[0] + if data[16] == 0: + files = data[o:].decode("mbcs").split("\0") + else: + files = data[o:].decode("utf-16le").split("\0") + return files[: files.index("")] + if isinstance(data, bytes): + data = io.BytesIO(data) + if fmt == "png": + from . import PngImagePlugin + + return PngImagePlugin.PngImageFile(data) + elif fmt == "DIB": + from . import BmpImagePlugin + + return BmpImagePlugin.DibImageFile(data) + return None + else: + if os.getenv("WAYLAND_DISPLAY"): + session_type = "wayland" + elif os.getenv("DISPLAY"): + session_type = "x11" + else: # Session type check failed + session_type = None + + if shutil.which("wl-paste") and session_type in ("wayland", None): + args = ["wl-paste", "-t", "image"] + elif shutil.which("xclip") and session_type in ("x11", None): + args = ["xclip", "-selection", "clipboard", "-t", "image/png", "-o"] + else: + msg = "wl-paste or xclip is required for ImageGrab.grabclipboard() on Linux" + raise NotImplementedError(msg) + + p = subprocess.run(args, capture_output=True) + if p.returncode != 0: + err = p.stderr + for silent_error in [ + # wl-paste, when the clipboard is empty + b"Nothing is copied", + # Ubuntu/Debian wl-paste, when the clipboard is empty + b"No selection", + # Ubuntu/Debian wl-paste, when an image isn't available + b"No suitable type of content copied", + # wl-paste or Ubuntu/Debian xclip, when an image isn't available + b" not available", + # xclip, when an image isn't available + b"cannot convert ", + # xclip, when the clipboard isn't initialized + b"xclip: Error: There is no owner for the ", + ]: + if silent_error in err: + return None + msg = f"{args[0]} error" + if err: + msg += f": {err.strip().decode()}" + raise ChildProcessError(msg) + + data = io.BytesIO(p.stdout) + im = Image.open(data) + im.load() + return im diff --git a/venv/lib/python3.10/site-packages/PIL/ImageMath.py b/venv/lib/python3.10/site-packages/PIL/ImageMath.py new file mode 100644 index 0000000000000000000000000000000000000000..dfdc50c0552cc0a3d325aa6b5479acd90893c2de --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageMath.py @@ -0,0 +1,314 @@ +# +# The Python Imaging Library +# $Id$ +# +# a simple math add-on for the Python Imaging Library +# +# History: +# 1999-02-15 fl Original PIL Plus release +# 2005-05-05 fl Simplified and cleaned up for PIL 1.1.6 +# 2005-09-12 fl Fixed int() and float() for Python 2.4.1 +# +# Copyright (c) 1999-2005 by Secret Labs AB +# Copyright (c) 2005 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import builtins + +from . import Image, _imagingmath + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from types import CodeType + from typing import Any + + +class _Operand: + """Wraps an image operand, providing standard operators""" + + def __init__(self, im: Image.Image): + self.im = im + + def __fixup(self, im1: _Operand | float) -> Image.Image: + # convert image to suitable mode + if isinstance(im1, _Operand): + # argument was an image. + if im1.im.mode in ("1", "L"): + return im1.im.convert("I") + elif im1.im.mode in ("I", "F"): + return im1.im + else: + msg = f"unsupported mode: {im1.im.mode}" + raise ValueError(msg) + else: + # argument was a constant + if isinstance(im1, (int, float)) and self.im.mode in ("1", "L", "I"): + return Image.new("I", self.im.size, im1) + else: + return Image.new("F", self.im.size, im1) + + def apply( + self, + op: str, + im1: _Operand | float, + im2: _Operand | float | None = None, + mode: str | None = None, + ) -> _Operand: + im_1 = self.__fixup(im1) + if im2 is None: + # unary operation + out = Image.new(mode or im_1.mode, im_1.size, None) + try: + op = getattr(_imagingmath, f"{op}_{im_1.mode}") + except AttributeError as e: + msg = f"bad operand type for '{op}'" + raise TypeError(msg) from e + _imagingmath.unop(op, out.getim(), im_1.getim()) + else: + # binary operation + im_2 = self.__fixup(im2) + if im_1.mode != im_2.mode: + # convert both arguments to floating point + if im_1.mode != "F": + im_1 = im_1.convert("F") + if im_2.mode != "F": + im_2 = im_2.convert("F") + if im_1.size != im_2.size: + # crop both arguments to a common size + size = ( + min(im_1.size[0], im_2.size[0]), + min(im_1.size[1], im_2.size[1]), + ) + if im_1.size != size: + im_1 = im_1.crop((0, 0) + size) + if im_2.size != size: + im_2 = im_2.crop((0, 0) + size) + out = Image.new(mode or im_1.mode, im_1.size, None) + try: + op = getattr(_imagingmath, f"{op}_{im_1.mode}") + except AttributeError as e: + msg = f"bad operand type for '{op}'" + raise TypeError(msg) from e + _imagingmath.binop(op, out.getim(), im_1.getim(), im_2.getim()) + return _Operand(out) + + # unary operators + def __bool__(self) -> bool: + # an image is "true" if it contains at least one non-zero pixel + return self.im.getbbox() is not None + + def __abs__(self) -> _Operand: + return self.apply("abs", self) + + def __pos__(self) -> _Operand: + return self + + def __neg__(self) -> _Operand: + return self.apply("neg", self) + + # binary operators + def __add__(self, other: _Operand | float) -> _Operand: + return self.apply("add", self, other) + + def __radd__(self, other: _Operand | float) -> _Operand: + return self.apply("add", other, self) + + def __sub__(self, other: _Operand | float) -> _Operand: + return self.apply("sub", self, other) + + def __rsub__(self, other: _Operand | float) -> _Operand: + return self.apply("sub", other, self) + + def __mul__(self, other: _Operand | float) -> _Operand: + return self.apply("mul", self, other) + + def __rmul__(self, other: _Operand | float) -> _Operand: + return self.apply("mul", other, self) + + def __truediv__(self, other: _Operand | float) -> _Operand: + return self.apply("div", self, other) + + def __rtruediv__(self, other: _Operand | float) -> _Operand: + return self.apply("div", other, self) + + def __mod__(self, other: _Operand | float) -> _Operand: + return self.apply("mod", self, other) + + def __rmod__(self, other: _Operand | float) -> _Operand: + return self.apply("mod", other, self) + + def __pow__(self, other: _Operand | float) -> _Operand: + return self.apply("pow", self, other) + + def __rpow__(self, other: _Operand | float) -> _Operand: + return self.apply("pow", other, self) + + # bitwise + def __invert__(self) -> _Operand: + return self.apply("invert", self) + + def __and__(self, other: _Operand | float) -> _Operand: + return self.apply("and", self, other) + + def __rand__(self, other: _Operand | float) -> _Operand: + return self.apply("and", other, self) + + def __or__(self, other: _Operand | float) -> _Operand: + return self.apply("or", self, other) + + def __ror__(self, other: _Operand | float) -> _Operand: + return self.apply("or", other, self) + + def __xor__(self, other: _Operand | float) -> _Operand: + return self.apply("xor", self, other) + + def __rxor__(self, other: _Operand | float) -> _Operand: + return self.apply("xor", other, self) + + def __lshift__(self, other: _Operand | float) -> _Operand: + return self.apply("lshift", self, other) + + def __rshift__(self, other: _Operand | float) -> _Operand: + return self.apply("rshift", self, other) + + # logical + def __eq__(self, other: _Operand | float) -> _Operand: # type: ignore[override] + return self.apply("eq", self, other) + + def __ne__(self, other: _Operand | float) -> _Operand: # type: ignore[override] + return self.apply("ne", self, other) + + def __lt__(self, other: _Operand | float) -> _Operand: + return self.apply("lt", self, other) + + def __le__(self, other: _Operand | float) -> _Operand: + return self.apply("le", self, other) + + def __gt__(self, other: _Operand | float) -> _Operand: + return self.apply("gt", self, other) + + def __ge__(self, other: _Operand | float) -> _Operand: + return self.apply("ge", self, other) + + +# conversions +def imagemath_int(self: _Operand) -> _Operand: + return _Operand(self.im.convert("I")) + + +def imagemath_float(self: _Operand) -> _Operand: + return _Operand(self.im.convert("F")) + + +# logical +def imagemath_equal(self: _Operand, other: _Operand | float | None) -> _Operand: + return self.apply("eq", self, other, mode="I") + + +def imagemath_notequal(self: _Operand, other: _Operand | float | None) -> _Operand: + return self.apply("ne", self, other, mode="I") + + +def imagemath_min(self: _Operand, other: _Operand | float | None) -> _Operand: + return self.apply("min", self, other) + + +def imagemath_max(self: _Operand, other: _Operand | float | None) -> _Operand: + return self.apply("max", self, other) + + +def imagemath_convert(self: _Operand, mode: str) -> _Operand: + return _Operand(self.im.convert(mode)) + + +ops = { + "int": imagemath_int, + "float": imagemath_float, + "equal": imagemath_equal, + "notequal": imagemath_notequal, + "min": imagemath_min, + "max": imagemath_max, + "convert": imagemath_convert, +} + + +def lambda_eval(expression: Callable[[dict[str, Any]], Any], **kw: Any) -> Any: + """ + Returns the result of an image function. + + :py:mod:`~PIL.ImageMath` only supports single-layer images. To process multi-band + images, use the :py:meth:`~PIL.Image.Image.split` method or + :py:func:`~PIL.Image.merge` function. + + :param expression: A function that receives a dictionary. + :param **kw: Values to add to the function's dictionary. + :return: The expression result. This is usually an image object, but can + also be an integer, a floating point value, or a pixel tuple, + depending on the expression. + """ + + args: dict[str, Any] = ops.copy() + args.update(kw) + for k, v in args.items(): + if isinstance(v, Image.Image): + args[k] = _Operand(v) + + out = expression(args) + try: + return out.im + except AttributeError: + return out + + +def unsafe_eval(expression: str, **kw: Any) -> Any: + """ + Evaluates an image expression. This uses Python's ``eval()`` function to process + the expression string, and carries the security risks of doing so. It is not + recommended to process expressions without considering this. + :py:meth:`~lambda_eval` is a more secure alternative. + + :py:mod:`~PIL.ImageMath` only supports single-layer images. To process multi-band + images, use the :py:meth:`~PIL.Image.Image.split` method or + :py:func:`~PIL.Image.merge` function. + + :param expression: A string containing a Python-style expression. + :param **kw: Values to add to the evaluation context. + :return: The evaluated expression. This is usually an image object, but can + also be an integer, a floating point value, or a pixel tuple, + depending on the expression. + """ + + # build execution namespace + args: dict[str, Any] = ops.copy() + for k in kw: + if "__" in k or hasattr(builtins, k): + msg = f"'{k}' not allowed" + raise ValueError(msg) + + args.update(kw) + for k, v in args.items(): + if isinstance(v, Image.Image): + args[k] = _Operand(v) + + compiled_code = compile(expression, "", "eval") + + def scan(code: CodeType) -> None: + for const in code.co_consts: + if type(const) is type(compiled_code): + scan(const) + + for name in code.co_names: + if name not in args and name != "abs": + msg = f"'{name}' not allowed" + raise ValueError(msg) + + scan(compiled_code) + out = builtins.eval(expression, {"__builtins": {"abs": abs}}, args) + try: + return out.im + except AttributeError: + return out diff --git a/venv/lib/python3.10/site-packages/PIL/ImageMode.py b/venv/lib/python3.10/site-packages/PIL/ImageMode.py new file mode 100644 index 0000000000000000000000000000000000000000..b7c6c863659b93416cac4acd8d8d31587566a9e6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageMode.py @@ -0,0 +1,85 @@ +# +# The Python Imaging Library. +# $Id$ +# +# standard mode descriptors +# +# History: +# 2006-03-20 fl Added +# +# Copyright (c) 2006 by Secret Labs AB. +# Copyright (c) 2006 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import sys +from functools import lru_cache +from typing import NamedTuple + + +class ModeDescriptor(NamedTuple): + """Wrapper for mode strings.""" + + mode: str + bands: tuple[str, ...] + basemode: str + basetype: str + typestr: str + + def __str__(self) -> str: + return self.mode + + +@lru_cache +def getmode(mode: str) -> ModeDescriptor: + """Gets a mode descriptor for the given mode.""" + endian = "<" if sys.byteorder == "little" else ">" + + modes = { + # core modes + # Bits need to be extended to bytes + "1": ("L", "L", ("1",), "|b1"), + "L": ("L", "L", ("L",), "|u1"), + "I": ("L", "I", ("I",), f"{endian}i4"), + "F": ("L", "F", ("F",), f"{endian}f4"), + "P": ("P", "L", ("P",), "|u1"), + "RGB": ("RGB", "L", ("R", "G", "B"), "|u1"), + "RGBX": ("RGB", "L", ("R", "G", "B", "X"), "|u1"), + "RGBA": ("RGB", "L", ("R", "G", "B", "A"), "|u1"), + "CMYK": ("RGB", "L", ("C", "M", "Y", "K"), "|u1"), + "YCbCr": ("RGB", "L", ("Y", "Cb", "Cr"), "|u1"), + # UNDONE - unsigned |u1i1i1 + "LAB": ("RGB", "L", ("L", "A", "B"), "|u1"), + "HSV": ("RGB", "L", ("H", "S", "V"), "|u1"), + # extra experimental modes + "RGBa": ("RGB", "L", ("R", "G", "B", "a"), "|u1"), + "LA": ("L", "L", ("L", "A"), "|u1"), + "La": ("L", "L", ("L", "a"), "|u1"), + "PA": ("RGB", "L", ("P", "A"), "|u1"), + } + if mode in modes: + base_mode, base_type, bands, type_str = modes[mode] + return ModeDescriptor(mode, bands, base_mode, base_type, type_str) + + mapping_modes = { + # I;16 == I;16L, and I;32 == I;32L + "I;16": "u2", + "I;16BS": ">i2", + "I;16N": f"{endian}u2", + "I;16NS": f"{endian}i2", + "I;32": "u4", + "I;32L": "i4", + "I;32LS": " +from __future__ import annotations + +import re + +from . import Image, _imagingmorph + +LUT_SIZE = 1 << 9 + +# fmt: off +ROTATION_MATRIX = [ + 6, 3, 0, + 7, 4, 1, + 8, 5, 2, +] +MIRROR_MATRIX = [ + 2, 1, 0, + 5, 4, 3, + 8, 7, 6, +] +# fmt: on + + +class LutBuilder: + """A class for building a MorphLut from a descriptive language + + The input patterns is a list of a strings sequences like these:: + + 4:(... + .1. + 111)->1 + + (whitespaces including linebreaks are ignored). The option 4 + describes a series of symmetry operations (in this case a + 4-rotation), the pattern is described by: + + - . or X - Ignore + - 1 - Pixel is on + - 0 - Pixel is off + + The result of the operation is described after "->" string. + + The default is to return the current pixel value, which is + returned if no other match is found. + + Operations: + + - 4 - 4 way rotation + - N - Negate + - 1 - Dummy op for no other operation (an op must always be given) + - M - Mirroring + + Example:: + + lb = LutBuilder(patterns = ["4:(... .1. 111)->1"]) + lut = lb.build_lut() + + """ + + def __init__( + self, patterns: list[str] | None = None, op_name: str | None = None + ) -> None: + if patterns is not None: + self.patterns = patterns + else: + self.patterns = [] + self.lut: bytearray | None = None + if op_name is not None: + known_patterns = { + "corner": ["1:(... ... ...)->0", "4:(00. 01. ...)->1"], + "dilation4": ["4:(... .0. .1.)->1"], + "dilation8": ["4:(... .0. .1.)->1", "4:(... .0. ..1)->1"], + "erosion4": ["4:(... .1. .0.)->0"], + "erosion8": ["4:(... .1. .0.)->0", "4:(... .1. ..0)->0"], + "edge": [ + "1:(... ... ...)->0", + "4:(.0. .1. ...)->1", + "4:(01. .1. ...)->1", + ], + } + if op_name not in known_patterns: + msg = f"Unknown pattern {op_name}!" + raise Exception(msg) + + self.patterns = known_patterns[op_name] + + def add_patterns(self, patterns: list[str]) -> None: + self.patterns += patterns + + def build_default_lut(self) -> None: + symbols = [0, 1] + m = 1 << 4 # pos of current pixel + self.lut = bytearray(symbols[(i & m) > 0] for i in range(LUT_SIZE)) + + def get_lut(self) -> bytearray | None: + return self.lut + + def _string_permute(self, pattern: str, permutation: list[int]) -> str: + """string_permute takes a pattern and a permutation and returns the + string permuted according to the permutation list. + """ + assert len(permutation) == 9 + return "".join(pattern[p] for p in permutation) + + def _pattern_permute( + self, basic_pattern: str, options: str, basic_result: int + ) -> list[tuple[str, int]]: + """pattern_permute takes a basic pattern and its result and clones + the pattern according to the modifications described in the $options + parameter. It returns a list of all cloned patterns.""" + patterns = [(basic_pattern, basic_result)] + + # rotations + if "4" in options: + res = patterns[-1][1] + for i in range(4): + patterns.append( + (self._string_permute(patterns[-1][0], ROTATION_MATRIX), res) + ) + # mirror + if "M" in options: + n = len(patterns) + for pattern, res in patterns[:n]: + patterns.append((self._string_permute(pattern, MIRROR_MATRIX), res)) + + # negate + if "N" in options: + n = len(patterns) + for pattern, res in patterns[:n]: + # Swap 0 and 1 + pattern = pattern.replace("0", "Z").replace("1", "0").replace("Z", "1") + res = 1 - int(res) + patterns.append((pattern, res)) + + return patterns + + def build_lut(self) -> bytearray: + """Compile all patterns into a morphology lut. + + TBD :Build based on (file) morphlut:modify_lut + """ + self.build_default_lut() + assert self.lut is not None + patterns = [] + + # Parse and create symmetries of the patterns strings + for p in self.patterns: + m = re.search(r"(\w):?\s*\((.+?)\)\s*->\s*(\d)", p.replace("\n", "")) + if not m: + msg = 'Syntax error in pattern "' + p + '"' + raise Exception(msg) + options = m.group(1) + pattern = m.group(2) + result = int(m.group(3)) + + # Get rid of spaces + pattern = pattern.replace(" ", "").replace("\n", "") + + patterns += self._pattern_permute(pattern, options, result) + + # compile the patterns into regular expressions for speed + compiled_patterns = [] + for pattern in patterns: + p = pattern[0].replace(".", "X").replace("X", "[01]") + compiled_patterns.append((re.compile(p), pattern[1])) + + # Step through table and find patterns that match. + # Note that all the patterns are searched. The last one + # caught overrides + for i in range(LUT_SIZE): + # Build the bit pattern + bitpattern = bin(i)[2:] + bitpattern = ("0" * (9 - len(bitpattern)) + bitpattern)[::-1] + + for pattern, r in compiled_patterns: + if pattern.match(bitpattern): + self.lut[i] = [0, 1][r] + + return self.lut + + +class MorphOp: + """A class for binary morphological operators""" + + def __init__( + self, + lut: bytearray | None = None, + op_name: str | None = None, + patterns: list[str] | None = None, + ) -> None: + """Create a binary morphological operator""" + self.lut = lut + if op_name is not None: + self.lut = LutBuilder(op_name=op_name).build_lut() + elif patterns is not None: + self.lut = LutBuilder(patterns=patterns).build_lut() + + def apply(self, image: Image.Image) -> tuple[int, Image.Image]: + """Run a single morphological operation on an image + + Returns a tuple of the number of changed pixels and the + morphed image""" + if self.lut is None: + msg = "No operator loaded" + raise Exception(msg) + + if image.mode != "L": + msg = "Image mode must be L" + raise ValueError(msg) + outimage = Image.new(image.mode, image.size, None) + count = _imagingmorph.apply(bytes(self.lut), image.getim(), outimage.getim()) + return count, outimage + + def match(self, image: Image.Image) -> list[tuple[int, int]]: + """Get a list of coordinates matching the morphological operation on + an image. + + Returns a list of tuples of (x,y) coordinates + of all matching pixels. See :ref:`coordinate-system`.""" + if self.lut is None: + msg = "No operator loaded" + raise Exception(msg) + + if image.mode != "L": + msg = "Image mode must be L" + raise ValueError(msg) + return _imagingmorph.match(bytes(self.lut), image.getim()) + + def get_on_pixels(self, image: Image.Image) -> list[tuple[int, int]]: + """Get a list of all turned on pixels in a binary image + + Returns a list of tuples of (x,y) coordinates + of all matching pixels. See :ref:`coordinate-system`.""" + + if image.mode != "L": + msg = "Image mode must be L" + raise ValueError(msg) + return _imagingmorph.get_on_pixels(image.getim()) + + def load_lut(self, filename: str) -> None: + """Load an operator from an mrl file""" + with open(filename, "rb") as f: + self.lut = bytearray(f.read()) + + if len(self.lut) != LUT_SIZE: + self.lut = None + msg = "Wrong size operator file!" + raise Exception(msg) + + def save_lut(self, filename: str) -> None: + """Save an operator to an mrl file""" + if self.lut is None: + msg = "No operator loaded" + raise Exception(msg) + with open(filename, "wb") as f: + f.write(self.lut) + + def set_lut(self, lut: bytearray | None) -> None: + """Set the lut from an external source""" + self.lut = lut diff --git a/venv/lib/python3.10/site-packages/PIL/ImageOps.py b/venv/lib/python3.10/site-packages/PIL/ImageOps.py new file mode 100644 index 0000000000000000000000000000000000000000..42b10bd7bc8d40bc6aa5ea77e1097461d5196b96 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageOps.py @@ -0,0 +1,746 @@ +# +# The Python Imaging Library. +# $Id$ +# +# standard image operations +# +# History: +# 2001-10-20 fl Created +# 2001-10-23 fl Added autocontrast operator +# 2001-12-18 fl Added Kevin's fit operator +# 2004-03-14 fl Fixed potential division by zero in equalize +# 2005-05-05 fl Fixed equalize for low number of values +# +# Copyright (c) 2001-2004 by Secret Labs AB +# Copyright (c) 2001-2004 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import functools +import operator +import re +from collections.abc import Sequence +from typing import Literal, Protocol, cast, overload + +from . import ExifTags, Image, ImagePalette + +# +# helpers + + +def _border(border: int | tuple[int, ...]) -> tuple[int, int, int, int]: + if isinstance(border, tuple): + if len(border) == 2: + left, top = right, bottom = border + elif len(border) == 4: + left, top, right, bottom = border + else: + left = top = right = bottom = border + return left, top, right, bottom + + +def _color(color: str | int | tuple[int, ...], mode: str) -> int | tuple[int, ...]: + if isinstance(color, str): + from . import ImageColor + + color = ImageColor.getcolor(color, mode) + return color + + +def _lut(image: Image.Image, lut: list[int]) -> Image.Image: + if image.mode == "P": + # FIXME: apply to lookup table, not image data + msg = "mode P support coming soon" + raise NotImplementedError(msg) + elif image.mode in ("L", "RGB"): + if image.mode == "RGB" and len(lut) == 256: + lut = lut + lut + lut + return image.point(lut) + else: + msg = f"not supported for mode {image.mode}" + raise OSError(msg) + + +# +# actions + + +def autocontrast( + image: Image.Image, + cutoff: float | tuple[float, float] = 0, + ignore: int | Sequence[int] | None = None, + mask: Image.Image | None = None, + preserve_tone: bool = False, +) -> Image.Image: + """ + Maximize (normalize) image contrast. This function calculates a + histogram of the input image (or mask region), removes ``cutoff`` percent of the + lightest and darkest pixels from the histogram, and remaps the image + so that the darkest pixel becomes black (0), and the lightest + becomes white (255). + + :param image: The image to process. + :param cutoff: The percent to cut off from the histogram on the low and + high ends. Either a tuple of (low, high), or a single + number for both. + :param ignore: The background pixel value (use None for no background). + :param mask: Histogram used in contrast operation is computed using pixels + within the mask. If no mask is given the entire image is used + for histogram computation. + :param preserve_tone: Preserve image tone in Photoshop-like style autocontrast. + + .. versionadded:: 8.2.0 + + :return: An image. + """ + if preserve_tone: + histogram = image.convert("L").histogram(mask) + else: + histogram = image.histogram(mask) + + lut = [] + for layer in range(0, len(histogram), 256): + h = histogram[layer : layer + 256] + if ignore is not None: + # get rid of outliers + if isinstance(ignore, int): + h[ignore] = 0 + else: + for ix in ignore: + h[ix] = 0 + if cutoff: + # cut off pixels from both ends of the histogram + if not isinstance(cutoff, tuple): + cutoff = (cutoff, cutoff) + # get number of pixels + n = 0 + for ix in range(256): + n = n + h[ix] + # remove cutoff% pixels from the low end + cut = int(n * cutoff[0] // 100) + for lo in range(256): + if cut > h[lo]: + cut = cut - h[lo] + h[lo] = 0 + else: + h[lo] -= cut + cut = 0 + if cut <= 0: + break + # remove cutoff% samples from the high end + cut = int(n * cutoff[1] // 100) + for hi in range(255, -1, -1): + if cut > h[hi]: + cut = cut - h[hi] + h[hi] = 0 + else: + h[hi] -= cut + cut = 0 + if cut <= 0: + break + # find lowest/highest samples after preprocessing + for lo in range(256): + if h[lo]: + break + for hi in range(255, -1, -1): + if h[hi]: + break + if hi <= lo: + # don't bother + lut.extend(list(range(256))) + else: + scale = 255.0 / (hi - lo) + offset = -lo * scale + for ix in range(256): + ix = int(ix * scale + offset) + if ix < 0: + ix = 0 + elif ix > 255: + ix = 255 + lut.append(ix) + return _lut(image, lut) + + +def colorize( + image: Image.Image, + black: str | tuple[int, ...], + white: str | tuple[int, ...], + mid: str | int | tuple[int, ...] | None = None, + blackpoint: int = 0, + whitepoint: int = 255, + midpoint: int = 127, +) -> Image.Image: + """ + Colorize grayscale image. + This function calculates a color wedge which maps all black pixels in + the source image to the first color and all white pixels to the + second color. If ``mid`` is specified, it uses three-color mapping. + The ``black`` and ``white`` arguments should be RGB tuples or color names; + optionally you can use three-color mapping by also specifying ``mid``. + Mapping positions for any of the colors can be specified + (e.g. ``blackpoint``), where these parameters are the integer + value corresponding to where the corresponding color should be mapped. + These parameters must have logical order, such that + ``blackpoint <= midpoint <= whitepoint`` (if ``mid`` is specified). + + :param image: The image to colorize. + :param black: The color to use for black input pixels. + :param white: The color to use for white input pixels. + :param mid: The color to use for midtone input pixels. + :param blackpoint: an int value [0, 255] for the black mapping. + :param whitepoint: an int value [0, 255] for the white mapping. + :param midpoint: an int value [0, 255] for the midtone mapping. + :return: An image. + """ + + # Initial asserts + assert image.mode == "L" + if mid is None: + assert 0 <= blackpoint <= whitepoint <= 255 + else: + assert 0 <= blackpoint <= midpoint <= whitepoint <= 255 + + # Define colors from arguments + rgb_black = cast(Sequence[int], _color(black, "RGB")) + rgb_white = cast(Sequence[int], _color(white, "RGB")) + rgb_mid = cast(Sequence[int], _color(mid, "RGB")) if mid is not None else None + + # Empty lists for the mapping + red = [] + green = [] + blue = [] + + # Create the low-end values + for i in range(blackpoint): + red.append(rgb_black[0]) + green.append(rgb_black[1]) + blue.append(rgb_black[2]) + + # Create the mapping (2-color) + if rgb_mid is None: + range_map = range(whitepoint - blackpoint) + + for i in range_map: + red.append( + rgb_black[0] + i * (rgb_white[0] - rgb_black[0]) // len(range_map) + ) + green.append( + rgb_black[1] + i * (rgb_white[1] - rgb_black[1]) // len(range_map) + ) + blue.append( + rgb_black[2] + i * (rgb_white[2] - rgb_black[2]) // len(range_map) + ) + + # Create the mapping (3-color) + else: + range_map1 = range(midpoint - blackpoint) + range_map2 = range(whitepoint - midpoint) + + for i in range_map1: + red.append( + rgb_black[0] + i * (rgb_mid[0] - rgb_black[0]) // len(range_map1) + ) + green.append( + rgb_black[1] + i * (rgb_mid[1] - rgb_black[1]) // len(range_map1) + ) + blue.append( + rgb_black[2] + i * (rgb_mid[2] - rgb_black[2]) // len(range_map1) + ) + for i in range_map2: + red.append(rgb_mid[0] + i * (rgb_white[0] - rgb_mid[0]) // len(range_map2)) + green.append( + rgb_mid[1] + i * (rgb_white[1] - rgb_mid[1]) // len(range_map2) + ) + blue.append(rgb_mid[2] + i * (rgb_white[2] - rgb_mid[2]) // len(range_map2)) + + # Create the high-end values + for i in range(256 - whitepoint): + red.append(rgb_white[0]) + green.append(rgb_white[1]) + blue.append(rgb_white[2]) + + # Return converted image + image = image.convert("RGB") + return _lut(image, red + green + blue) + + +def contain( + image: Image.Image, size: tuple[int, int], method: int = Image.Resampling.BICUBIC +) -> Image.Image: + """ + Returns a resized version of the image, set to the maximum width and height + within the requested size, while maintaining the original aspect ratio. + + :param image: The image to resize. + :param size: The requested output size in pixels, given as a + (width, height) tuple. + :param method: Resampling method to use. Default is + :py:attr:`~PIL.Image.Resampling.BICUBIC`. + See :ref:`concept-filters`. + :return: An image. + """ + + im_ratio = image.width / image.height + dest_ratio = size[0] / size[1] + + if im_ratio != dest_ratio: + if im_ratio > dest_ratio: + new_height = round(image.height / image.width * size[0]) + if new_height != size[1]: + size = (size[0], new_height) + else: + new_width = round(image.width / image.height * size[1]) + if new_width != size[0]: + size = (new_width, size[1]) + return image.resize(size, resample=method) + + +def cover( + image: Image.Image, size: tuple[int, int], method: int = Image.Resampling.BICUBIC +) -> Image.Image: + """ + Returns a resized version of the image, so that the requested size is + covered, while maintaining the original aspect ratio. + + :param image: The image to resize. + :param size: The requested output size in pixels, given as a + (width, height) tuple. + :param method: Resampling method to use. Default is + :py:attr:`~PIL.Image.Resampling.BICUBIC`. + See :ref:`concept-filters`. + :return: An image. + """ + + im_ratio = image.width / image.height + dest_ratio = size[0] / size[1] + + if im_ratio != dest_ratio: + if im_ratio < dest_ratio: + new_height = round(image.height / image.width * size[0]) + if new_height != size[1]: + size = (size[0], new_height) + else: + new_width = round(image.width / image.height * size[1]) + if new_width != size[0]: + size = (new_width, size[1]) + return image.resize(size, resample=method) + + +def pad( + image: Image.Image, + size: tuple[int, int], + method: int = Image.Resampling.BICUBIC, + color: str | int | tuple[int, ...] | None = None, + centering: tuple[float, float] = (0.5, 0.5), +) -> Image.Image: + """ + Returns a resized and padded version of the image, expanded to fill the + requested aspect ratio and size. + + :param image: The image to resize and crop. + :param size: The requested output size in pixels, given as a + (width, height) tuple. + :param method: Resampling method to use. Default is + :py:attr:`~PIL.Image.Resampling.BICUBIC`. + See :ref:`concept-filters`. + :param color: The background color of the padded image. + :param centering: Control the position of the original image within the + padded version. + + (0.5, 0.5) will keep the image centered + (0, 0) will keep the image aligned to the top left + (1, 1) will keep the image aligned to the bottom + right + :return: An image. + """ + + resized = contain(image, size, method) + if resized.size == size: + out = resized + else: + out = Image.new(image.mode, size, color) + if resized.palette: + palette = resized.getpalette() + if palette is not None: + out.putpalette(palette) + if resized.width != size[0]: + x = round((size[0] - resized.width) * max(0, min(centering[0], 1))) + out.paste(resized, (x, 0)) + else: + y = round((size[1] - resized.height) * max(0, min(centering[1], 1))) + out.paste(resized, (0, y)) + return out + + +def crop(image: Image.Image, border: int = 0) -> Image.Image: + """ + Remove border from image. The same amount of pixels are removed + from all four sides. This function works on all image modes. + + .. seealso:: :py:meth:`~PIL.Image.Image.crop` + + :param image: The image to crop. + :param border: The number of pixels to remove. + :return: An image. + """ + left, top, right, bottom = _border(border) + return image.crop((left, top, image.size[0] - right, image.size[1] - bottom)) + + +def scale( + image: Image.Image, factor: float, resample: int = Image.Resampling.BICUBIC +) -> Image.Image: + """ + Returns a rescaled image by a specific factor given in parameter. + A factor greater than 1 expands the image, between 0 and 1 contracts the + image. + + :param image: The image to rescale. + :param factor: The expansion factor, as a float. + :param resample: Resampling method to use. Default is + :py:attr:`~PIL.Image.Resampling.BICUBIC`. + See :ref:`concept-filters`. + :returns: An :py:class:`~PIL.Image.Image` object. + """ + if factor == 1: + return image.copy() + elif factor <= 0: + msg = "the factor must be greater than 0" + raise ValueError(msg) + else: + size = (round(factor * image.width), round(factor * image.height)) + return image.resize(size, resample) + + +class SupportsGetMesh(Protocol): + """ + An object that supports the ``getmesh`` method, taking an image as an + argument, and returning a list of tuples. Each tuple contains two tuples, + the source box as a tuple of 4 integers, and a tuple of 8 integers for the + final quadrilateral, in order of top left, bottom left, bottom right, top + right. + """ + + def getmesh( + self, image: Image.Image + ) -> list[ + tuple[tuple[int, int, int, int], tuple[int, int, int, int, int, int, int, int]] + ]: ... + + +def deform( + image: Image.Image, + deformer: SupportsGetMesh, + resample: int = Image.Resampling.BILINEAR, +) -> Image.Image: + """ + Deform the image. + + :param image: The image to deform. + :param deformer: A deformer object. Any object that implements a + ``getmesh`` method can be used. + :param resample: An optional resampling filter. Same values possible as + in the PIL.Image.transform function. + :return: An image. + """ + return image.transform( + image.size, Image.Transform.MESH, deformer.getmesh(image), resample + ) + + +def equalize(image: Image.Image, mask: Image.Image | None = None) -> Image.Image: + """ + Equalize the image histogram. This function applies a non-linear + mapping to the input image, in order to create a uniform + distribution of grayscale values in the output image. + + :param image: The image to equalize. + :param mask: An optional mask. If given, only the pixels selected by + the mask are included in the analysis. + :return: An image. + """ + if image.mode == "P": + image = image.convert("RGB") + h = image.histogram(mask) + lut = [] + for b in range(0, len(h), 256): + histo = [_f for _f in h[b : b + 256] if _f] + if len(histo) <= 1: + lut.extend(list(range(256))) + else: + step = (functools.reduce(operator.add, histo) - histo[-1]) // 255 + if not step: + lut.extend(list(range(256))) + else: + n = step // 2 + for i in range(256): + lut.append(n // step) + n = n + h[i + b] + return _lut(image, lut) + + +def expand( + image: Image.Image, + border: int | tuple[int, ...] = 0, + fill: str | int | tuple[int, ...] = 0, +) -> Image.Image: + """ + Add border to the image + + :param image: The image to expand. + :param border: Border width, in pixels. + :param fill: Pixel fill value (a color value). Default is 0 (black). + :return: An image. + """ + left, top, right, bottom = _border(border) + width = left + image.size[0] + right + height = top + image.size[1] + bottom + color = _color(fill, image.mode) + if image.palette: + mode = image.palette.mode + palette = ImagePalette.ImagePalette(mode, image.getpalette(mode)) + if isinstance(color, tuple) and (len(color) == 3 or len(color) == 4): + color = palette.getcolor(color) + else: + palette = None + out = Image.new(image.mode, (width, height), color) + if palette: + out.putpalette(palette.palette, mode) + out.paste(image, (left, top)) + return out + + +def fit( + image: Image.Image, + size: tuple[int, int], + method: int = Image.Resampling.BICUBIC, + bleed: float = 0.0, + centering: tuple[float, float] = (0.5, 0.5), +) -> Image.Image: + """ + Returns a resized and cropped version of the image, cropped to the + requested aspect ratio and size. + + This function was contributed by Kevin Cazabon. + + :param image: The image to resize and crop. + :param size: The requested output size in pixels, given as a + (width, height) tuple. + :param method: Resampling method to use. Default is + :py:attr:`~PIL.Image.Resampling.BICUBIC`. + See :ref:`concept-filters`. + :param bleed: Remove a border around the outside of the image from all + four edges. The value is a decimal percentage (use 0.01 for + one percent). The default value is 0 (no border). + Cannot be greater than or equal to 0.5. + :param centering: Control the cropping position. Use (0.5, 0.5) for + center cropping (e.g. if cropping the width, take 50% off + of the left side, and therefore 50% off the right side). + (0.0, 0.0) will crop from the top left corner (i.e. if + cropping the width, take all of the crop off of the right + side, and if cropping the height, take all of it off the + bottom). (1.0, 0.0) will crop from the bottom left + corner, etc. (i.e. if cropping the width, take all of the + crop off the left side, and if cropping the height take + none from the top, and therefore all off the bottom). + :return: An image. + """ + + # by Kevin Cazabon, Feb 17/2000 + # kevin@cazabon.com + # https://www.cazabon.com + + centering_x, centering_y = centering + + if not 0.0 <= centering_x <= 1.0: + centering_x = 0.5 + if not 0.0 <= centering_y <= 1.0: + centering_y = 0.5 + + if not 0.0 <= bleed < 0.5: + bleed = 0.0 + + # calculate the area to use for resizing and cropping, subtracting + # the 'bleed' around the edges + + # number of pixels to trim off on Top and Bottom, Left and Right + bleed_pixels = (bleed * image.size[0], bleed * image.size[1]) + + live_size = ( + image.size[0] - bleed_pixels[0] * 2, + image.size[1] - bleed_pixels[1] * 2, + ) + + # calculate the aspect ratio of the live_size + live_size_ratio = live_size[0] / live_size[1] + + # calculate the aspect ratio of the output image + output_ratio = size[0] / size[1] + + # figure out if the sides or top/bottom will be cropped off + if live_size_ratio == output_ratio: + # live_size is already the needed ratio + crop_width = live_size[0] + crop_height = live_size[1] + elif live_size_ratio >= output_ratio: + # live_size is wider than what's needed, crop the sides + crop_width = output_ratio * live_size[1] + crop_height = live_size[1] + else: + # live_size is taller than what's needed, crop the top and bottom + crop_width = live_size[0] + crop_height = live_size[0] / output_ratio + + # make the crop + crop_left = bleed_pixels[0] + (live_size[0] - crop_width) * centering_x + crop_top = bleed_pixels[1] + (live_size[1] - crop_height) * centering_y + + crop = (crop_left, crop_top, crop_left + crop_width, crop_top + crop_height) + + # resize the image and return it + return image.resize(size, method, box=crop) + + +def flip(image: Image.Image) -> Image.Image: + """ + Flip the image vertically (top to bottom). + + :param image: The image to flip. + :return: An image. + """ + return image.transpose(Image.Transpose.FLIP_TOP_BOTTOM) + + +def grayscale(image: Image.Image) -> Image.Image: + """ + Convert the image to grayscale. + + :param image: The image to convert. + :return: An image. + """ + return image.convert("L") + + +def invert(image: Image.Image) -> Image.Image: + """ + Invert (negate) the image. + + :param image: The image to invert. + :return: An image. + """ + lut = list(range(255, -1, -1)) + return image.point(lut) if image.mode == "1" else _lut(image, lut) + + +def mirror(image: Image.Image) -> Image.Image: + """ + Flip image horizontally (left to right). + + :param image: The image to mirror. + :return: An image. + """ + return image.transpose(Image.Transpose.FLIP_LEFT_RIGHT) + + +def posterize(image: Image.Image, bits: int) -> Image.Image: + """ + Reduce the number of bits for each color channel. + + :param image: The image to posterize. + :param bits: The number of bits to keep for each channel (1-8). + :return: An image. + """ + mask = ~(2 ** (8 - bits) - 1) + lut = [i & mask for i in range(256)] + return _lut(image, lut) + + +def solarize(image: Image.Image, threshold: int = 128) -> Image.Image: + """ + Invert all pixel values above a threshold. + + :param image: The image to solarize. + :param threshold: All pixels above this grayscale level are inverted. + :return: An image. + """ + lut = [] + for i in range(256): + if i < threshold: + lut.append(i) + else: + lut.append(255 - i) + return _lut(image, lut) + + +@overload +def exif_transpose(image: Image.Image, *, in_place: Literal[True]) -> None: ... + + +@overload +def exif_transpose( + image: Image.Image, *, in_place: Literal[False] = False +) -> Image.Image: ... + + +def exif_transpose(image: Image.Image, *, in_place: bool = False) -> Image.Image | None: + """ + If an image has an EXIF Orientation tag, other than 1, transpose the image + accordingly, and remove the orientation data. + + :param image: The image to transpose. + :param in_place: Boolean. Keyword-only argument. + If ``True``, the original image is modified in-place, and ``None`` is returned. + If ``False`` (default), a new :py:class:`~PIL.Image.Image` object is returned + with the transposition applied. If there is no transposition, a copy of the + image will be returned. + """ + image.load() + image_exif = image.getexif() + orientation = image_exif.get(ExifTags.Base.Orientation, 1) + method = { + 2: Image.Transpose.FLIP_LEFT_RIGHT, + 3: Image.Transpose.ROTATE_180, + 4: Image.Transpose.FLIP_TOP_BOTTOM, + 5: Image.Transpose.TRANSPOSE, + 6: Image.Transpose.ROTATE_270, + 7: Image.Transpose.TRANSVERSE, + 8: Image.Transpose.ROTATE_90, + }.get(orientation) + if method is not None: + if in_place: + image.im = image.im.transpose(method) + image._size = image.im.size + else: + transposed_image = image.transpose(method) + exif_image = image if in_place else transposed_image + + exif = exif_image.getexif() + if ExifTags.Base.Orientation in exif: + del exif[ExifTags.Base.Orientation] + if "exif" in exif_image.info: + exif_image.info["exif"] = exif.tobytes() + elif "Raw profile type exif" in exif_image.info: + exif_image.info["Raw profile type exif"] = exif.tobytes().hex() + for key in ("XML:com.adobe.xmp", "xmp"): + if key in exif_image.info: + for pattern in ( + r'tiff:Orientation="([0-9])"', + r"([0-9])", + ): + value = exif_image.info[key] + if isinstance(value, str): + value = re.sub(pattern, "", value) + elif isinstance(value, tuple): + value = tuple( + re.sub(pattern.encode(), b"", v) for v in value + ) + else: + value = re.sub(pattern.encode(), b"", value) + exif_image.info[key] = value + if not in_place: + return transposed_image + elif not in_place: + return image.copy() + return None diff --git a/venv/lib/python3.10/site-packages/PIL/ImagePalette.py b/venv/lib/python3.10/site-packages/PIL/ImagePalette.py new file mode 100644 index 0000000000000000000000000000000000000000..103697117b92a3dca9794fbea5b9c92306b9b198 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImagePalette.py @@ -0,0 +1,286 @@ +# +# The Python Imaging Library. +# $Id$ +# +# image palette object +# +# History: +# 1996-03-11 fl Rewritten. +# 1997-01-03 fl Up and running. +# 1997-08-23 fl Added load hack +# 2001-04-16 fl Fixed randint shadow bug in random() +# +# Copyright (c) 1997-2001 by Secret Labs AB +# Copyright (c) 1996-1997 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import array +from collections.abc import Sequence +from typing import IO + +from . import GimpGradientFile, GimpPaletteFile, ImageColor, PaletteFile + +TYPE_CHECKING = False +if TYPE_CHECKING: + from . import Image + + +class ImagePalette: + """ + Color palette for palette mapped images + + :param mode: The mode to use for the palette. See: + :ref:`concept-modes`. Defaults to "RGB" + :param palette: An optional palette. If given, it must be a bytearray, + an array or a list of ints between 0-255. The list must consist of + all channels for one color followed by the next color (e.g. RGBRGBRGB). + Defaults to an empty palette. + """ + + def __init__( + self, + mode: str = "RGB", + palette: Sequence[int] | bytes | bytearray | None = None, + ) -> None: + self.mode = mode + self.rawmode: str | None = None # if set, palette contains raw data + self.palette = palette or bytearray() + self.dirty: int | None = None + + @property + def palette(self) -> Sequence[int] | bytes | bytearray: + return self._palette + + @palette.setter + def palette(self, palette: Sequence[int] | bytes | bytearray) -> None: + self._colors: dict[tuple[int, ...], int] | None = None + self._palette = palette + + @property + def colors(self) -> dict[tuple[int, ...], int]: + if self._colors is None: + mode_len = len(self.mode) + self._colors = {} + for i in range(0, len(self.palette), mode_len): + color = tuple(self.palette[i : i + mode_len]) + if color in self._colors: + continue + self._colors[color] = i // mode_len + return self._colors + + @colors.setter + def colors(self, colors: dict[tuple[int, ...], int]) -> None: + self._colors = colors + + def copy(self) -> ImagePalette: + new = ImagePalette() + + new.mode = self.mode + new.rawmode = self.rawmode + if self.palette is not None: + new.palette = self.palette[:] + new.dirty = self.dirty + + return new + + def getdata(self) -> tuple[str, Sequence[int] | bytes | bytearray]: + """ + Get palette contents in format suitable for the low-level + ``im.putpalette`` primitive. + + .. warning:: This method is experimental. + """ + if self.rawmode: + return self.rawmode, self.palette + return self.mode, self.tobytes() + + def tobytes(self) -> bytes: + """Convert palette to bytes. + + .. warning:: This method is experimental. + """ + if self.rawmode: + msg = "palette contains raw palette data" + raise ValueError(msg) + if isinstance(self.palette, bytes): + return self.palette + arr = array.array("B", self.palette) + return arr.tobytes() + + # Declare tostring as an alias for tobytes + tostring = tobytes + + def _new_color_index( + self, image: Image.Image | None = None, e: Exception | None = None + ) -> int: + if not isinstance(self.palette, bytearray): + self._palette = bytearray(self.palette) + index = len(self.palette) // 3 + special_colors: tuple[int | tuple[int, ...] | None, ...] = () + if image: + special_colors = ( + image.info.get("background"), + image.info.get("transparency"), + ) + while index in special_colors: + index += 1 + if index >= 256: + if image: + # Search for an unused index + for i, count in reversed(list(enumerate(image.histogram()))): + if count == 0 and i not in special_colors: + index = i + break + if index >= 256: + msg = "cannot allocate more than 256 colors" + raise ValueError(msg) from e + return index + + def getcolor( + self, + color: tuple[int, ...], + image: Image.Image | None = None, + ) -> int: + """Given an rgb tuple, allocate palette entry. + + .. warning:: This method is experimental. + """ + if self.rawmode: + msg = "palette contains raw palette data" + raise ValueError(msg) + if isinstance(color, tuple): + if self.mode == "RGB": + if len(color) == 4: + if color[3] != 255: + msg = "cannot add non-opaque RGBA color to RGB palette" + raise ValueError(msg) + color = color[:3] + elif self.mode == "RGBA": + if len(color) == 3: + color += (255,) + try: + return self.colors[color] + except KeyError as e: + # allocate new color slot + index = self._new_color_index(image, e) + assert isinstance(self._palette, bytearray) + self.colors[color] = index + if index * 3 < len(self.palette): + self._palette = ( + self._palette[: index * 3] + + bytes(color) + + self._palette[index * 3 + 3 :] + ) + else: + self._palette += bytes(color) + self.dirty = 1 + return index + else: + msg = f"unknown color specifier: {repr(color)}" # type: ignore[unreachable] + raise ValueError(msg) + + def save(self, fp: str | IO[str]) -> None: + """Save palette to text file. + + .. warning:: This method is experimental. + """ + if self.rawmode: + msg = "palette contains raw palette data" + raise ValueError(msg) + if isinstance(fp, str): + fp = open(fp, "w") + fp.write("# Palette\n") + fp.write(f"# Mode: {self.mode}\n") + for i in range(256): + fp.write(f"{i}") + for j in range(i * len(self.mode), (i + 1) * len(self.mode)): + try: + fp.write(f" {self.palette[j]}") + except IndexError: + fp.write(" 0") + fp.write("\n") + fp.close() + + +# -------------------------------------------------------------------- +# Internal + + +def raw(rawmode: str, data: Sequence[int] | bytes | bytearray) -> ImagePalette: + palette = ImagePalette() + palette.rawmode = rawmode + palette.palette = data + palette.dirty = 1 + return palette + + +# -------------------------------------------------------------------- +# Factories + + +def make_linear_lut(black: int, white: float) -> list[int]: + if black == 0: + return [int(white * i // 255) for i in range(256)] + + msg = "unavailable when black is non-zero" + raise NotImplementedError(msg) # FIXME + + +def make_gamma_lut(exp: float) -> list[int]: + return [int(((i / 255.0) ** exp) * 255.0 + 0.5) for i in range(256)] + + +def negative(mode: str = "RGB") -> ImagePalette: + palette = list(range(256 * len(mode))) + palette.reverse() + return ImagePalette(mode, [i // len(mode) for i in palette]) + + +def random(mode: str = "RGB") -> ImagePalette: + from random import randint + + palette = [randint(0, 255) for _ in range(256 * len(mode))] + return ImagePalette(mode, palette) + + +def sepia(white: str = "#fff0c0") -> ImagePalette: + bands = [make_linear_lut(0, band) for band in ImageColor.getrgb(white)] + return ImagePalette("RGB", [bands[i % 3][i // 3] for i in range(256 * 3)]) + + +def wedge(mode: str = "RGB") -> ImagePalette: + palette = list(range(256 * len(mode))) + return ImagePalette(mode, [i // len(mode) for i in palette]) + + +def load(filename: str) -> tuple[bytes, str]: + # FIXME: supports GIMP gradients only + + with open(filename, "rb") as fp: + paletteHandlers: list[ + type[ + GimpPaletteFile.GimpPaletteFile + | GimpGradientFile.GimpGradientFile + | PaletteFile.PaletteFile + ] + ] = [ + GimpPaletteFile.GimpPaletteFile, + GimpGradientFile.GimpGradientFile, + PaletteFile.PaletteFile, + ] + for paletteHandler in paletteHandlers: + try: + fp.seek(0) + lut = paletteHandler(fp).getpalette() + if lut: + break + except (SyntaxError, ValueError): + pass + else: + msg = "cannot load palette" + raise OSError(msg) + + return lut # data, rawmode diff --git a/venv/lib/python3.10/site-packages/PIL/ImagePath.py b/venv/lib/python3.10/site-packages/PIL/ImagePath.py new file mode 100644 index 0000000000000000000000000000000000000000..77e8a609a552ae7d8c6b87e78a36ecbfc1cdce89 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImagePath.py @@ -0,0 +1,20 @@ +# +# The Python Imaging Library +# $Id$ +# +# path interface +# +# History: +# 1996-11-04 fl Created +# 2002-04-14 fl Added documentation stub class +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image + +Path = Image.core.path diff --git a/venv/lib/python3.10/site-packages/PIL/ImageQt.py b/venv/lib/python3.10/site-packages/PIL/ImageQt.py new file mode 100644 index 0000000000000000000000000000000000000000..af4d0742d6bc5f434cf47ae7a629bf08e23fd3a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageQt.py @@ -0,0 +1,219 @@ +# +# The Python Imaging Library. +# $Id$ +# +# a simple Qt image interface. +# +# history: +# 2006-06-03 fl: created +# 2006-06-04 fl: inherit from QImage instead of wrapping it +# 2006-06-05 fl: removed toimage helper; move string support to ImageQt +# 2013-11-13 fl: add support for Qt5 (aurelien.ballier@cyclonit.com) +# +# Copyright (c) 2006 by Secret Labs AB +# Copyright (c) 2006 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import sys +from io import BytesIO + +from . import Image +from ._util import is_path + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from typing import Any + + from . import ImageFile + + QBuffer: type + +qt_version: str | None +qt_versions = [ + ["6", "PyQt6"], + ["side6", "PySide6"], +] + +# If a version has already been imported, attempt it first +qt_versions.sort(key=lambda version: version[1] in sys.modules, reverse=True) +for version, qt_module in qt_versions: + try: + qRgba: Callable[[int, int, int, int], int] + if qt_module == "PyQt6": + from PyQt6.QtCore import QBuffer, QByteArray, QIODevice + from PyQt6.QtGui import QImage, QPixmap, qRgba + elif qt_module == "PySide6": + from PySide6.QtCore import ( # type: ignore[assignment] + QBuffer, + QByteArray, + QIODevice, + ) + from PySide6.QtGui import QImage, QPixmap, qRgba # type: ignore[assignment] + except (ImportError, RuntimeError): + continue + qt_is_installed = True + qt_version = version + break +else: + qt_is_installed = False + qt_version = None + + +def rgb(r: int, g: int, b: int, a: int = 255) -> int: + """(Internal) Turns an RGB color into a Qt compatible color integer.""" + # use qRgb to pack the colors, and then turn the resulting long + # into a negative integer with the same bitpattern. + return qRgba(r, g, b, a) & 0xFFFFFFFF + + +def fromqimage(im: QImage | QPixmap) -> ImageFile.ImageFile: + """ + :param im: QImage or PIL ImageQt object + """ + buffer = QBuffer() + qt_openmode: object + if qt_version == "6": + try: + qt_openmode = getattr(QIODevice, "OpenModeFlag") + except AttributeError: + qt_openmode = getattr(QIODevice, "OpenMode") + else: + qt_openmode = QIODevice + buffer.open(getattr(qt_openmode, "ReadWrite")) + # preserve alpha channel with png + # otherwise ppm is more friendly with Image.open + if im.hasAlphaChannel(): + im.save(buffer, "png") + else: + im.save(buffer, "ppm") + + b = BytesIO() + b.write(buffer.data()) + buffer.close() + b.seek(0) + + return Image.open(b) + + +def fromqpixmap(im: QPixmap) -> ImageFile.ImageFile: + return fromqimage(im) + + +def align8to32(bytes: bytes, width: int, mode: str) -> bytes: + """ + converts each scanline of data from 8 bit to 32 bit aligned + """ + + bits_per_pixel = {"1": 1, "L": 8, "P": 8, "I;16": 16}[mode] + + # calculate bytes per line and the extra padding if needed + bits_per_line = bits_per_pixel * width + full_bytes_per_line, remaining_bits_per_line = divmod(bits_per_line, 8) + bytes_per_line = full_bytes_per_line + (1 if remaining_bits_per_line else 0) + + extra_padding = -bytes_per_line % 4 + + # already 32 bit aligned by luck + if not extra_padding: + return bytes + + new_data = [ + bytes[i * bytes_per_line : (i + 1) * bytes_per_line] + b"\x00" * extra_padding + for i in range(len(bytes) // bytes_per_line) + ] + + return b"".join(new_data) + + +def _toqclass_helper(im: Image.Image | str | QByteArray) -> dict[str, Any]: + data = None + colortable = None + exclusive_fp = False + + # handle filename, if given instead of image name + if hasattr(im, "toUtf8"): + # FIXME - is this really the best way to do this? + im = str(im.toUtf8(), "utf-8") + if is_path(im): + im = Image.open(im) + exclusive_fp = True + assert isinstance(im, Image.Image) + + qt_format = getattr(QImage, "Format") if qt_version == "6" else QImage + if im.mode == "1": + format = getattr(qt_format, "Format_Mono") + elif im.mode == "L": + format = getattr(qt_format, "Format_Indexed8") + colortable = [rgb(i, i, i) for i in range(256)] + elif im.mode == "P": + format = getattr(qt_format, "Format_Indexed8") + palette = im.getpalette() + assert palette is not None + colortable = [rgb(*palette[i : i + 3]) for i in range(0, len(palette), 3)] + elif im.mode == "RGB": + # Populate the 4th channel with 255 + im = im.convert("RGBA") + + data = im.tobytes("raw", "BGRA") + format = getattr(qt_format, "Format_RGB32") + elif im.mode == "RGBA": + data = im.tobytes("raw", "BGRA") + format = getattr(qt_format, "Format_ARGB32") + elif im.mode == "I;16": + im = im.point(lambda i: i * 256) + + format = getattr(qt_format, "Format_Grayscale16") + else: + if exclusive_fp: + im.close() + msg = f"unsupported image mode {repr(im.mode)}" + raise ValueError(msg) + + size = im.size + __data = data or align8to32(im.tobytes(), size[0], im.mode) + if exclusive_fp: + im.close() + return {"data": __data, "size": size, "format": format, "colortable": colortable} + + +if qt_is_installed: + + class ImageQt(QImage): + def __init__(self, im: Image.Image | str | QByteArray) -> None: + """ + An PIL image wrapper for Qt. This is a subclass of PyQt's QImage + class. + + :param im: A PIL Image object, or a file name (given either as + Python string or a PyQt string object). + """ + im_data = _toqclass_helper(im) + # must keep a reference, or Qt will crash! + # All QImage constructors that take data operate on an existing + # buffer, so this buffer has to hang on for the life of the image. + # Fixes https://github.com/python-pillow/Pillow/issues/1370 + self.__data = im_data["data"] + super().__init__( + self.__data, + im_data["size"][0], + im_data["size"][1], + im_data["format"], + ) + if im_data["colortable"]: + self.setColorTable(im_data["colortable"]) + + +def toqimage(im: Image.Image | str | QByteArray) -> ImageQt: + return ImageQt(im) + + +def toqpixmap(im: Image.Image | str | QByteArray) -> QPixmap: + qimage = toqimage(im) + pixmap = getattr(QPixmap, "fromImage")(qimage) + if qt_version == "6": + pixmap.detach() + return pixmap diff --git a/venv/lib/python3.10/site-packages/PIL/ImageSequence.py b/venv/lib/python3.10/site-packages/PIL/ImageSequence.py new file mode 100644 index 0000000000000000000000000000000000000000..361be48971e0446ef39e2428b613f75951e04e87 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageSequence.py @@ -0,0 +1,88 @@ +# +# The Python Imaging Library. +# $Id$ +# +# sequence support classes +# +# history: +# 1997-02-20 fl Created +# +# Copyright (c) 1997 by Secret Labs AB. +# Copyright (c) 1997 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + +## +from __future__ import annotations + +from . import Image + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + + +class Iterator: + """ + This class implements an iterator object that can be used to loop + over an image sequence. + + You can use the ``[]`` operator to access elements by index. This operator + will raise an :py:exc:`IndexError` if you try to access a nonexistent + frame. + + :param im: An image object. + """ + + def __init__(self, im: Image.Image) -> None: + if not hasattr(im, "seek"): + msg = "im must have seek method" + raise AttributeError(msg) + self.im = im + self.position = getattr(self.im, "_min_frame", 0) + + def __getitem__(self, ix: int) -> Image.Image: + try: + self.im.seek(ix) + return self.im + except EOFError as e: + msg = "end of sequence" + raise IndexError(msg) from e + + def __iter__(self) -> Iterator: + return self + + def __next__(self) -> Image.Image: + try: + self.im.seek(self.position) + self.position += 1 + return self.im + except EOFError as e: + msg = "end of sequence" + raise StopIteration(msg) from e + + +def all_frames( + im: Image.Image | list[Image.Image], + func: Callable[[Image.Image], Image.Image] | None = None, +) -> list[Image.Image]: + """ + Applies a given function to all frames in an image or a list of images. + The frames are returned as a list of separate images. + + :param im: An image, or a list of images. + :param func: The function to apply to all of the image frames. + :returns: A list of images. + """ + if not isinstance(im, list): + im = [im] + + ims = [] + for imSequence in im: + current = imSequence.tell() + + ims += [im_frame.copy() for im_frame in Iterator(imSequence)] + + imSequence.seek(current) + return [func(im) for im in ims] if func else ims diff --git a/venv/lib/python3.10/site-packages/PIL/ImageShow.py b/venv/lib/python3.10/site-packages/PIL/ImageShow.py new file mode 100644 index 0000000000000000000000000000000000000000..7705608e3eccd5e82cfca87daa1264df2c81dacd --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageShow.py @@ -0,0 +1,362 @@ +# +# The Python Imaging Library. +# $Id$ +# +# im.show() drivers +# +# History: +# 2008-04-06 fl Created +# +# Copyright (c) Secret Labs AB 2008. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import abc +import os +import shutil +import subprocess +import sys +from shlex import quote +from typing import Any + +from . import Image + +_viewers = [] + + +def register(viewer: type[Viewer] | Viewer, order: int = 1) -> None: + """ + The :py:func:`register` function is used to register additional viewers:: + + from PIL import ImageShow + ImageShow.register(MyViewer()) # MyViewer will be used as a last resort + ImageShow.register(MySecondViewer(), 0) # MySecondViewer will be prioritised + ImageShow.register(ImageShow.XVViewer(), 0) # XVViewer will be prioritised + + :param viewer: The viewer to be registered. + :param order: + Zero or a negative integer to prepend this viewer to the list, + a positive integer to append it. + """ + if isinstance(viewer, type) and issubclass(viewer, Viewer): + viewer = viewer() + if order > 0: + _viewers.append(viewer) + else: + _viewers.insert(0, viewer) + + +def show(image: Image.Image, title: str | None = None, **options: Any) -> bool: + r""" + Display a given image. + + :param image: An image object. + :param title: Optional title. Not all viewers can display the title. + :param \**options: Additional viewer options. + :returns: ``True`` if a suitable viewer was found, ``False`` otherwise. + """ + for viewer in _viewers: + if viewer.show(image, title=title, **options): + return True + return False + + +class Viewer: + """Base class for viewers.""" + + # main api + + def show(self, image: Image.Image, **options: Any) -> int: + """ + The main function for displaying an image. + Converts the given image to the target format and displays it. + """ + + if not ( + image.mode in ("1", "RGBA") + or (self.format == "PNG" and image.mode in ("I;16", "LA")) + ): + base = Image.getmodebase(image.mode) + if image.mode != base: + image = image.convert(base) + + return self.show_image(image, **options) + + # hook methods + + format: str | None = None + """The format to convert the image into.""" + options: dict[str, Any] = {} + """Additional options used to convert the image.""" + + def get_format(self, image: Image.Image) -> str | None: + """Return format name, or ``None`` to save as PGM/PPM.""" + return self.format + + def get_command(self, file: str, **options: Any) -> str: + """ + Returns the command used to display the file. + Not implemented in the base class. + """ + msg = "unavailable in base viewer" + raise NotImplementedError(msg) + + def save_image(self, image: Image.Image) -> str: + """Save to temporary file and return filename.""" + return image._dump(format=self.get_format(image), **self.options) + + def show_image(self, image: Image.Image, **options: Any) -> int: + """Display the given image.""" + return self.show_file(self.save_image(image), **options) + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + os.system(self.get_command(path, **options)) # nosec + return 1 + + +# -------------------------------------------------------------------- + + +class WindowsViewer(Viewer): + """The default viewer on Windows is the default system application for PNG files.""" + + format = "PNG" + options = {"compress_level": 1, "save_all": True} + + def get_command(self, file: str, **options: Any) -> str: + return ( + f'start "Pillow" /WAIT "{file}" ' + "&& ping -n 4 127.0.0.1 >NUL " + f'&& del /f "{file}"' + ) + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + subprocess.Popen( + self.get_command(path, **options), + shell=True, + creationflags=getattr(subprocess, "CREATE_NO_WINDOW"), + ) # nosec + return 1 + + +if sys.platform == "win32": + register(WindowsViewer) + + +class MacViewer(Viewer): + """The default viewer on macOS using ``Preview.app``.""" + + format = "PNG" + options = {"compress_level": 1, "save_all": True} + + def get_command(self, file: str, **options: Any) -> str: + # on darwin open returns immediately resulting in the temp + # file removal while app is opening + command = "open -a Preview.app" + command = f"({command} {quote(file)}; sleep 20; rm -f {quote(file)})&" + return command + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + subprocess.call(["open", "-a", "Preview.app", path]) + + pyinstaller = getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS") + executable = (not pyinstaller and sys.executable) or shutil.which("python3") + if executable: + subprocess.Popen( + [ + executable, + "-c", + "import os, sys, time; time.sleep(20); os.remove(sys.argv[1])", + path, + ] + ) + return 1 + + +if sys.platform == "darwin": + register(MacViewer) + + +class UnixViewer(abc.ABC, Viewer): + format = "PNG" + options = {"compress_level": 1, "save_all": True} + + @abc.abstractmethod + def get_command_ex(self, file: str, **options: Any) -> tuple[str, str]: + pass + + def get_command(self, file: str, **options: Any) -> str: + command = self.get_command_ex(file, **options)[0] + return f"{command} {quote(file)}" + + +class XDGViewer(UnixViewer): + """ + The freedesktop.org ``xdg-open`` command. + """ + + def get_command_ex(self, file: str, **options: Any) -> tuple[str, str]: + command = executable = "xdg-open" + return command, executable + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + subprocess.Popen(["xdg-open", path]) + return 1 + + +class DisplayViewer(UnixViewer): + """ + The ImageMagick ``display`` command. + This viewer supports the ``title`` parameter. + """ + + def get_command_ex( + self, file: str, title: str | None = None, **options: Any + ) -> tuple[str, str]: + command = executable = "display" + if title: + command += f" -title {quote(title)}" + return command, executable + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + args = ["display"] + title = options.get("title") + if title: + args += ["-title", title] + args.append(path) + + subprocess.Popen(args) + return 1 + + +class GmDisplayViewer(UnixViewer): + """The GraphicsMagick ``gm display`` command.""" + + def get_command_ex(self, file: str, **options: Any) -> tuple[str, str]: + executable = "gm" + command = "gm display" + return command, executable + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + subprocess.Popen(["gm", "display", path]) + return 1 + + +class EogViewer(UnixViewer): + """The GNOME Image Viewer ``eog`` command.""" + + def get_command_ex(self, file: str, **options: Any) -> tuple[str, str]: + executable = "eog" + command = "eog -n" + return command, executable + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + subprocess.Popen(["eog", "-n", path]) + return 1 + + +class XVViewer(UnixViewer): + """ + The X Viewer ``xv`` command. + This viewer supports the ``title`` parameter. + """ + + def get_command_ex( + self, file: str, title: str | None = None, **options: Any + ) -> tuple[str, str]: + # note: xv is pretty outdated. most modern systems have + # imagemagick's display command instead. + command = executable = "xv" + if title: + command += f" -name {quote(title)}" + return command, executable + + def show_file(self, path: str, **options: Any) -> int: + """ + Display given file. + """ + if not os.path.exists(path): + raise FileNotFoundError + args = ["xv"] + title = options.get("title") + if title: + args += ["-name", title] + args.append(path) + + subprocess.Popen(args) + return 1 + + +if sys.platform not in ("win32", "darwin"): # unixoids + if shutil.which("xdg-open"): + register(XDGViewer) + if shutil.which("display"): + register(DisplayViewer) + if shutil.which("gm"): + register(GmDisplayViewer) + if shutil.which("eog"): + register(EogViewer) + if shutil.which("xv"): + register(XVViewer) + + +class IPythonViewer(Viewer): + """The viewer for IPython frontends.""" + + def show_image(self, image: Image.Image, **options: Any) -> int: + ipython_display(image) + return 1 + + +try: + from IPython.display import display as ipython_display +except ImportError: + pass +else: + register(IPythonViewer) + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Syntax: python3 ImageShow.py imagefile [title]") + sys.exit() + + with Image.open(sys.argv[1]) as im: + print(show(im, *sys.argv[2:])) diff --git a/venv/lib/python3.10/site-packages/PIL/ImageStat.py b/venv/lib/python3.10/site-packages/PIL/ImageStat.py new file mode 100644 index 0000000000000000000000000000000000000000..3a1044ba449408cd038e81bb4f6cc45a2149e3ce --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageStat.py @@ -0,0 +1,167 @@ +# +# The Python Imaging Library. +# $Id$ +# +# global image statistics +# +# History: +# 1996-04-05 fl Created +# 1997-05-21 fl Added mask; added rms, var, stddev attributes +# 1997-08-05 fl Added median +# 1998-07-05 hk Fixed integer overflow error +# +# Notes: +# This class shows how to implement delayed evaluation of attributes. +# To get a certain value, simply access the corresponding attribute. +# The __getattr__ dispatcher takes care of the rest. +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996-97. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import math +from functools import cached_property + +from . import Image + + +class Stat: + def __init__( + self, image_or_list: Image.Image | list[int], mask: Image.Image | None = None + ) -> None: + """ + Calculate statistics for the given image. If a mask is included, + only the regions covered by that mask are included in the + statistics. You can also pass in a previously calculated histogram. + + :param image: A PIL image, or a precalculated histogram. + + .. note:: + + For a PIL image, calculations rely on the + :py:meth:`~PIL.Image.Image.histogram` method. The pixel counts are + grouped into 256 bins, even if the image has more than 8 bits per + channel. So ``I`` and ``F`` mode images have a maximum ``mean``, + ``median`` and ``rms`` of 255, and cannot have an ``extrema`` maximum + of more than 255. + + :param mask: An optional mask. + """ + if isinstance(image_or_list, Image.Image): + self.h = image_or_list.histogram(mask) + elif isinstance(image_or_list, list): + self.h = image_or_list + else: + msg = "first argument must be image or list" # type: ignore[unreachable] + raise TypeError(msg) + self.bands = list(range(len(self.h) // 256)) + + @cached_property + def extrema(self) -> list[tuple[int, int]]: + """ + Min/max values for each band in the image. + + .. note:: + This relies on the :py:meth:`~PIL.Image.Image.histogram` method, and + simply returns the low and high bins used. This is correct for + images with 8 bits per channel, but fails for other modes such as + ``I`` or ``F``. Instead, use :py:meth:`~PIL.Image.Image.getextrema` to + return per-band extrema for the image. This is more correct and + efficient because, for non-8-bit modes, the histogram method uses + :py:meth:`~PIL.Image.Image.getextrema` to determine the bins used. + """ + + def minmax(histogram: list[int]) -> tuple[int, int]: + res_min, res_max = 255, 0 + for i in range(256): + if histogram[i]: + res_min = i + break + for i in range(255, -1, -1): + if histogram[i]: + res_max = i + break + return res_min, res_max + + return [minmax(self.h[i:]) for i in range(0, len(self.h), 256)] + + @cached_property + def count(self) -> list[int]: + """Total number of pixels for each band in the image.""" + return [sum(self.h[i : i + 256]) for i in range(0, len(self.h), 256)] + + @cached_property + def sum(self) -> list[float]: + """Sum of all pixels for each band in the image.""" + + v = [] + for i in range(0, len(self.h), 256): + layer_sum = 0.0 + for j in range(256): + layer_sum += j * self.h[i + j] + v.append(layer_sum) + return v + + @cached_property + def sum2(self) -> list[float]: + """Squared sum of all pixels for each band in the image.""" + + v = [] + for i in range(0, len(self.h), 256): + sum2 = 0.0 + for j in range(256): + sum2 += (j**2) * float(self.h[i + j]) + v.append(sum2) + return v + + @cached_property + def mean(self) -> list[float]: + """Average (arithmetic mean) pixel level for each band in the image.""" + return [self.sum[i] / self.count[i] if self.count[i] else 0 for i in self.bands] + + @cached_property + def median(self) -> list[int]: + """Median pixel level for each band in the image.""" + + v = [] + for i in self.bands: + s = 0 + half = self.count[i] // 2 + b = i * 256 + for j in range(256): + s = s + self.h[b + j] + if s > half: + break + v.append(j) + return v + + @cached_property + def rms(self) -> list[float]: + """RMS (root-mean-square) for each band in the image.""" + return [ + math.sqrt(self.sum2[i] / self.count[i]) if self.count[i] else 0 + for i in self.bands + ] + + @cached_property + def var(self) -> list[float]: + """Variance for each band in the image.""" + return [ + ( + (self.sum2[i] - (self.sum[i] ** 2.0) / self.count[i]) / self.count[i] + if self.count[i] + else 0 + ) + for i in self.bands + ] + + @cached_property + def stddev(self) -> list[float]: + """Standard deviation for each band in the image.""" + return [math.sqrt(self.var[i]) for i in self.bands] + + +Global = Stat # compatibility diff --git a/venv/lib/python3.10/site-packages/PIL/ImageText.py b/venv/lib/python3.10/site-packages/PIL/ImageText.py new file mode 100644 index 0000000000000000000000000000000000000000..c74570e693989b739c85b98044debff2c2c8ddc7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageText.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +from . import ImageFont +from ._typing import _Ink + + +class Text: + def __init__( + self, + text: str | bytes, + font: ( + ImageFont.ImageFont + | ImageFont.FreeTypeFont + | ImageFont.TransposedFont + | None + ) = None, + mode: str = "RGB", + spacing: float = 4, + direction: str | None = None, + features: list[str] | None = None, + language: str | None = None, + ) -> None: + """ + :param text: String to be drawn. + :param font: Either an :py:class:`~PIL.ImageFont.ImageFont` instance, + :py:class:`~PIL.ImageFont.FreeTypeFont` instance, + :py:class:`~PIL.ImageFont.TransposedFont` instance or ``None``. If + ``None``, the default font from :py:meth:`.ImageFont.load_default` + will be used. + :param mode: The image mode this will be used with. + :param spacing: The number of pixels between lines. + :param direction: Direction of the text. It can be ``"rtl"`` (right to left), + ``"ltr"`` (left to right) or ``"ttb"`` (top to bottom). + Requires libraqm. + :param features: A list of OpenType font features to be used during text + layout. This is usually used to turn on optional font features + that are not enabled by default, for example ``"dlig"`` or + ``"ss01"``, but can be also used to turn off default font + features, for example ``"-liga"`` to disable ligatures or + ``"-kern"`` to disable kerning. To get all supported + features, see `OpenType docs`_. + Requires libraqm. + :param language: Language of the text. Different languages may use + different glyph shapes or ligatures. This parameter tells + the font which language the text is in, and to apply the + correct substitutions as appropriate, if available. + It should be a `BCP 47 language code`_. + Requires libraqm. + """ + self.text = text + self.font = font or ImageFont.load_default() + + self.mode = mode + self.spacing = spacing + self.direction = direction + self.features = features + self.language = language + + self.embedded_color = False + + self.stroke_width: float = 0 + self.stroke_fill: _Ink | None = None + + def embed_color(self) -> None: + """ + Use embedded color glyphs (COLR, CBDT, SBIX). + """ + if self.mode not in ("RGB", "RGBA"): + msg = "Embedded color supported only in RGB and RGBA modes" + raise ValueError(msg) + self.embedded_color = True + + def stroke(self, width: float = 0, fill: _Ink | None = None) -> None: + """ + :param width: The width of the text stroke. + :param fill: Color to use for the text stroke when drawing. If not given, will + default to the ``fill`` parameter from + :py:meth:`.ImageDraw.ImageDraw.text`. + """ + self.stroke_width = width + self.stroke_fill = fill + + def _get_fontmode(self) -> str: + if self.mode in ("1", "P", "I", "F"): + return "1" + elif self.embedded_color: + return "RGBA" + else: + return "L" + + def get_length(self): + """ + Returns length (in pixels with 1/64 precision) of text. + + This is the amount by which following text should be offset. + Text bounding box may extend past the length in some fonts, + e.g. when using italics or accents. + + The result is returned as a float; it is a whole number if using basic layout. + + Note that the sum of two lengths may not equal the length of a concatenated + string due to kerning. If you need to adjust for kerning, include the following + character and subtract its length. + + For example, instead of:: + + hello = ImageText.Text("Hello", font).get_length() + world = ImageText.Text("World", font).get_length() + helloworld = ImageText.Text("HelloWorld", font).get_length() + assert hello + world == helloworld + + use:: + + hello = ( + ImageText.Text("HelloW", font).get_length() - + ImageText.Text("W", font).get_length() + ) # adjusted for kerning + world = ImageText.Text("World", font).get_length() + helloworld = ImageText.Text("HelloWorld", font).get_length() + assert hello + world == helloworld + + or disable kerning with (requires libraqm):: + + hello = ImageText.Text("Hello", font, features=["-kern"]).get_length() + world = ImageText.Text("World", font, features=["-kern"]).get_length() + helloworld = ImageText.Text( + "HelloWorld", font, features=["-kern"] + ).get_length() + assert hello + world == helloworld + + :return: Either width for horizontal text, or height for vertical text. + """ + split_character = "\n" if isinstance(self.text, str) else b"\n" + if split_character in self.text: + msg = "can't measure length of multiline text" + raise ValueError(msg) + return self.font.getlength( + self.text, + self._get_fontmode(), + self.direction, + self.features, + self.language, + ) + + def _split( + self, xy: tuple[float, float], anchor: str | None, align: str + ) -> list[tuple[tuple[float, float], str, str | bytes]]: + if anchor is None: + anchor = "lt" if self.direction == "ttb" else "la" + elif len(anchor) != 2: + msg = "anchor must be a 2 character string" + raise ValueError(msg) + + lines = ( + self.text.split("\n") + if isinstance(self.text, str) + else self.text.split(b"\n") + ) + if len(lines) == 1: + return [(xy, anchor, self.text)] + + if anchor[1] in "tb" and self.direction != "ttb": + msg = "anchor not supported for multiline text" + raise ValueError(msg) + + fontmode = self._get_fontmode() + line_spacing = ( + self.font.getbbox( + "A", + fontmode, + None, + self.features, + self.language, + self.stroke_width, + )[3] + + self.stroke_width + + self.spacing + ) + + top = xy[1] + parts = [] + if self.direction == "ttb": + left = xy[0] + for line in lines: + parts.append(((left, top), anchor, line)) + left += line_spacing + else: + widths = [] + max_width: float = 0 + for line in lines: + line_width = self.font.getlength( + line, fontmode, self.direction, self.features, self.language + ) + widths.append(line_width) + max_width = max(max_width, line_width) + + if anchor[1] == "m": + top -= (len(lines) - 1) * line_spacing / 2.0 + elif anchor[1] == "d": + top -= (len(lines) - 1) * line_spacing + + idx = -1 + for line in lines: + left = xy[0] + idx += 1 + width_difference = max_width - widths[idx] + + # align by align parameter + if align in ("left", "justify"): + pass + elif align == "center": + left += width_difference / 2.0 + elif align == "right": + left += width_difference + else: + msg = 'align must be "left", "center", "right" or "justify"' + raise ValueError(msg) + + if ( + align == "justify" + and width_difference != 0 + and idx != len(lines) - 1 + ): + words = ( + line.split(" ") if isinstance(line, str) else line.split(b" ") + ) + if len(words) > 1: + # align left by anchor + if anchor[0] == "m": + left -= max_width / 2.0 + elif anchor[0] == "r": + left -= max_width + + word_widths = [ + self.font.getlength( + word, + fontmode, + self.direction, + self.features, + self.language, + ) + for word in words + ] + word_anchor = "l" + anchor[1] + width_difference = max_width - sum(word_widths) + i = 0 + for word in words: + parts.append(((left, top), word_anchor, word)) + left += word_widths[i] + width_difference / (len(words) - 1) + i += 1 + top += line_spacing + continue + + # align left by anchor + if anchor[0] == "m": + left -= width_difference / 2.0 + elif anchor[0] == "r": + left -= width_difference + parts.append(((left, top), anchor, line)) + top += line_spacing + + return parts + + def get_bbox( + self, + xy: tuple[float, float] = (0, 0), + anchor: str | None = None, + align: str = "left", + ) -> tuple[float, float, float, float]: + """ + Returns bounding box (in pixels) of text. + + Use :py:meth:`get_length` to get the offset of following text with 1/64 pixel + precision. The bounding box includes extra margins for some fonts, e.g. italics + or accents. + + :param xy: The anchor coordinates of the text. + :param anchor: The text anchor alignment. Determines the relative location of + the anchor to the text. The default alignment is top left, + specifically ``la`` for horizontal text and ``lt`` for + vertical text. See :ref:`text-anchors` for details. + :param align: For multiline text, ``"left"``, ``"center"``, ``"right"`` or + ``"justify"`` determines the relative alignment of lines. Use the + ``anchor`` parameter to specify the alignment to ``xy``. + + :return: ``(left, top, right, bottom)`` bounding box + """ + bbox: tuple[float, float, float, float] | None = None + fontmode = self._get_fontmode() + for xy, anchor, line in self._split(xy, anchor, align): + bbox_line = self.font.getbbox( + line, + fontmode, + self.direction, + self.features, + self.language, + self.stroke_width, + anchor, + ) + bbox_line = ( + bbox_line[0] + xy[0], + bbox_line[1] + xy[1], + bbox_line[2] + xy[0], + bbox_line[3] + xy[1], + ) + if bbox is None: + bbox = bbox_line + else: + bbox = ( + min(bbox[0], bbox_line[0]), + min(bbox[1], bbox_line[1]), + max(bbox[2], bbox_line[2]), + max(bbox[3], bbox_line[3]), + ) + + if bbox is None: + return xy[0], xy[1], xy[0], xy[1] + return bbox diff --git a/venv/lib/python3.10/site-packages/PIL/ImageTk.py b/venv/lib/python3.10/site-packages/PIL/ImageTk.py new file mode 100644 index 0000000000000000000000000000000000000000..3a4cb81e9ef5ef4abe617d4a364074c2203571ad --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageTk.py @@ -0,0 +1,266 @@ +# +# The Python Imaging Library. +# $Id$ +# +# a Tk display interface +# +# History: +# 96-04-08 fl Created +# 96-09-06 fl Added getimage method +# 96-11-01 fl Rewritten, removed image attribute and crop method +# 97-05-09 fl Use PyImagingPaste method instead of image type +# 97-05-12 fl Minor tweaks to match the IFUNC95 interface +# 97-05-17 fl Support the "pilbitmap" booster patch +# 97-06-05 fl Added file= and data= argument to image constructors +# 98-03-09 fl Added width and height methods to Image classes +# 98-07-02 fl Use default mode for "P" images without palette attribute +# 98-07-02 fl Explicitly destroy Tkinter image objects +# 99-07-24 fl Support multiple Tk interpreters (from Greg Couch) +# 99-07-26 fl Automatically hook into Tkinter (if possible) +# 99-08-15 fl Hook uses _imagingtk instead of _imaging +# +# Copyright (c) 1997-1999 by Secret Labs AB +# Copyright (c) 1996-1997 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import tkinter +from io import BytesIO +from typing import Any + +from . import Image, ImageFile + +TYPE_CHECKING = False +if TYPE_CHECKING: + from ._typing import CapsuleType + +# -------------------------------------------------------------------- +# Check for Tkinter interface hooks + + +def _get_image_from_kw(kw: dict[str, Any]) -> ImageFile.ImageFile | None: + source = None + if "file" in kw: + source = kw.pop("file") + elif "data" in kw: + source = BytesIO(kw.pop("data")) + if not source: + return None + return Image.open(source) + + +def _pyimagingtkcall( + command: str, photo: PhotoImage | tkinter.PhotoImage, ptr: CapsuleType +) -> None: + tk = photo.tk + try: + tk.call(command, photo, repr(ptr)) + except tkinter.TclError: + # activate Tkinter hook + # may raise an error if it cannot attach to Tkinter + from . import _imagingtk + + _imagingtk.tkinit(tk.interpaddr()) + tk.call(command, photo, repr(ptr)) + + +# -------------------------------------------------------------------- +# PhotoImage + + +class PhotoImage: + """ + A Tkinter-compatible photo image. This can be used + everywhere Tkinter expects an image object. If the image is an RGBA + image, pixels having alpha 0 are treated as transparent. + + The constructor takes either a PIL image, or a mode and a size. + Alternatively, you can use the ``file`` or ``data`` options to initialize + the photo image object. + + :param image: Either a PIL image, or a mode string. If a mode string is + used, a size must also be given. + :param size: If the first argument is a mode string, this defines the size + of the image. + :keyword file: A filename to load the image from (using + ``Image.open(file)``). + :keyword data: An 8-bit string containing image data (as loaded from an + image file). + """ + + def __init__( + self, + image: Image.Image | str | None = None, + size: tuple[int, int] | None = None, + **kw: Any, + ) -> None: + # Tk compatibility: file or data + if image is None: + image = _get_image_from_kw(kw) + + if image is None: + msg = "Image is required" + raise ValueError(msg) + elif isinstance(image, str): + mode = image + image = None + + if size is None: + msg = "If first argument is mode, size is required" + raise ValueError(msg) + else: + # got an image instead of a mode + mode = image.mode + if mode == "P": + # palette mapped data + image.apply_transparency() + image.load() + mode = image.palette.mode if image.palette else "RGB" + size = image.size + kw["width"], kw["height"] = size + + if mode not in ["1", "L", "RGB", "RGBA"]: + mode = Image.getmodebase(mode) + + self.__mode = mode + self.__size = size + self.__photo = tkinter.PhotoImage(**kw) + self.tk = self.__photo.tk + if image: + self.paste(image) + + def __del__(self) -> None: + try: + name = self.__photo.name + except AttributeError: + return + self.__photo.name = None + try: + self.__photo.tk.call("image", "delete", name) + except Exception: + pass # ignore internal errors + + def __str__(self) -> str: + """ + Get the Tkinter photo image identifier. This method is automatically + called by Tkinter whenever a PhotoImage object is passed to a Tkinter + method. + + :return: A Tkinter photo image identifier (a string). + """ + return str(self.__photo) + + def width(self) -> int: + """ + Get the width of the image. + + :return: The width, in pixels. + """ + return self.__size[0] + + def height(self) -> int: + """ + Get the height of the image. + + :return: The height, in pixels. + """ + return self.__size[1] + + def paste(self, im: Image.Image) -> None: + """ + Paste a PIL image into the photo image. Note that this can + be very slow if the photo image is displayed. + + :param im: A PIL image. The size must match the target region. If the + mode does not match, the image is converted to the mode of + the bitmap image. + """ + # convert to blittable + ptr = im.getim() + image = im.im + if not image.isblock() or im.mode != self.__mode: + block = Image.core.new_block(self.__mode, im.size) + image.convert2(block, image) # convert directly between buffers + ptr = block.ptr + + _pyimagingtkcall("PyImagingPhoto", self.__photo, ptr) + + +# -------------------------------------------------------------------- +# BitmapImage + + +class BitmapImage: + """ + A Tkinter-compatible bitmap image. This can be used everywhere Tkinter + expects an image object. + + The given image must have mode "1". Pixels having value 0 are treated as + transparent. Options, if any, are passed on to Tkinter. The most commonly + used option is ``foreground``, which is used to specify the color for the + non-transparent parts. See the Tkinter documentation for information on + how to specify colours. + + :param image: A PIL image. + """ + + def __init__(self, image: Image.Image | None = None, **kw: Any) -> None: + # Tk compatibility: file or data + if image is None: + image = _get_image_from_kw(kw) + + if image is None: + msg = "Image is required" + raise ValueError(msg) + self.__mode = image.mode + self.__size = image.size + + self.__photo = tkinter.BitmapImage(data=image.tobitmap(), **kw) + + def __del__(self) -> None: + try: + name = self.__photo.name + except AttributeError: + return + self.__photo.name = None + try: + self.__photo.tk.call("image", "delete", name) + except Exception: + pass # ignore internal errors + + def width(self) -> int: + """ + Get the width of the image. + + :return: The width, in pixels. + """ + return self.__size[0] + + def height(self) -> int: + """ + Get the height of the image. + + :return: The height, in pixels. + """ + return self.__size[1] + + def __str__(self) -> str: + """ + Get the Tkinter bitmap image identifier. This method is automatically + called by Tkinter whenever a BitmapImage object is passed to a Tkinter + method. + + :return: A Tkinter bitmap image identifier (a string). + """ + return str(self.__photo) + + +def getimage(photo: PhotoImage) -> Image.Image: + """Copies the contents of a PhotoImage to a PIL image memory.""" + im = Image.new("RGBA", (photo.width(), photo.height())) + + _pyimagingtkcall("PyImagingPhotoGet", photo, im.getim()) + + return im diff --git a/venv/lib/python3.10/site-packages/PIL/ImageTransform.py b/venv/lib/python3.10/site-packages/PIL/ImageTransform.py new file mode 100644 index 0000000000000000000000000000000000000000..fb144ff38a1ee7ff77cc01f3b941756a60b2b4cd --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageTransform.py @@ -0,0 +1,136 @@ +# +# The Python Imaging Library. +# $Id$ +# +# transform wrappers +# +# History: +# 2002-04-08 fl Created +# +# Copyright (c) 2002 by Secret Labs AB +# Copyright (c) 2002 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from . import Image + + +class Transform(Image.ImageTransformHandler): + """Base class for other transforms defined in :py:mod:`~PIL.ImageTransform`.""" + + method: Image.Transform + + def __init__(self, data: Sequence[Any]) -> None: + self.data = data + + def getdata(self) -> tuple[Image.Transform, Sequence[int]]: + return self.method, self.data + + def transform( + self, + size: tuple[int, int], + image: Image.Image, + **options: Any, + ) -> Image.Image: + """Perform the transform. Called from :py:meth:`.Image.transform`.""" + # can be overridden + method, data = self.getdata() + return image.transform(size, method, data, **options) + + +class AffineTransform(Transform): + """ + Define an affine image transform. + + This function takes a 6-tuple (a, b, c, d, e, f) which contain the first + two rows from the inverse of an affine transform matrix. For each pixel + (x, y) in the output image, the new value is taken from a position (a x + + b y + c, d x + e y + f) in the input image, rounded to nearest pixel. + + This function can be used to scale, translate, rotate, and shear the + original image. + + See :py:meth:`.Image.transform` + + :param matrix: A 6-tuple (a, b, c, d, e, f) containing the first two rows + from the inverse of an affine transform matrix. + """ + + method = Image.Transform.AFFINE + + +class PerspectiveTransform(Transform): + """ + Define a perspective image transform. + + This function takes an 8-tuple (a, b, c, d, e, f, g, h). For each pixel + (x, y) in the output image, the new value is taken from a position + ((a x + b y + c) / (g x + h y + 1), (d x + e y + f) / (g x + h y + 1)) in + the input image, rounded to nearest pixel. + + This function can be used to scale, translate, rotate, and shear the + original image. + + See :py:meth:`.Image.transform` + + :param matrix: An 8-tuple (a, b, c, d, e, f, g, h). + """ + + method = Image.Transform.PERSPECTIVE + + +class ExtentTransform(Transform): + """ + Define a transform to extract a subregion from an image. + + Maps a rectangle (defined by two corners) from the image to a rectangle of + the given size. The resulting image will contain data sampled from between + the corners, such that (x0, y0) in the input image will end up at (0,0) in + the output image, and (x1, y1) at size. + + This method can be used to crop, stretch, shrink, or mirror an arbitrary + rectangle in the current image. It is slightly slower than crop, but about + as fast as a corresponding resize operation. + + See :py:meth:`.Image.transform` + + :param bbox: A 4-tuple (x0, y0, x1, y1) which specifies two points in the + input image's coordinate system. See :ref:`coordinate-system`. + """ + + method = Image.Transform.EXTENT + + +class QuadTransform(Transform): + """ + Define a quad image transform. + + Maps a quadrilateral (a region defined by four corners) from the image to a + rectangle of the given size. + + See :py:meth:`.Image.transform` + + :param xy: An 8-tuple (x0, y0, x1, y1, x2, y2, x3, y3) which contain the + upper left, lower left, lower right, and upper right corner of the + source quadrilateral. + """ + + method = Image.Transform.QUAD + + +class MeshTransform(Transform): + """ + Define a mesh image transform. A mesh transform consists of one or more + individual quad transforms. + + See :py:meth:`.Image.transform` + + :param data: A list of (bbox, quad) tuples. + """ + + method = Image.Transform.MESH diff --git a/venv/lib/python3.10/site-packages/PIL/ImageWin.py b/venv/lib/python3.10/site-packages/PIL/ImageWin.py new file mode 100644 index 0000000000000000000000000000000000000000..98c28f29f1dbbb069b68dc9359051b6629148f0d --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImageWin.py @@ -0,0 +1,247 @@ +# +# The Python Imaging Library. +# $Id$ +# +# a Windows DIB display interface +# +# History: +# 1996-05-20 fl Created +# 1996-09-20 fl Fixed subregion exposure +# 1997-09-21 fl Added draw primitive (for tzPrint) +# 2003-05-21 fl Added experimental Window/ImageWindow classes +# 2003-09-05 fl Added fromstring/tostring methods +# +# Copyright (c) Secret Labs AB 1997-2003. +# Copyright (c) Fredrik Lundh 1996-2003. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image + + +class HDC: + """ + Wraps an HDC integer. The resulting object can be passed to the + :py:meth:`~PIL.ImageWin.Dib.draw` and :py:meth:`~PIL.ImageWin.Dib.expose` + methods. + """ + + def __init__(self, dc: int) -> None: + self.dc = dc + + def __int__(self) -> int: + return self.dc + + +class HWND: + """ + Wraps an HWND integer. The resulting object can be passed to the + :py:meth:`~PIL.ImageWin.Dib.draw` and :py:meth:`~PIL.ImageWin.Dib.expose` + methods, instead of a DC. + """ + + def __init__(self, wnd: int) -> None: + self.wnd = wnd + + def __int__(self) -> int: + return self.wnd + + +class Dib: + """ + A Windows bitmap with the given mode and size. The mode can be one of "1", + "L", "P", or "RGB". + + If the display requires a palette, this constructor creates a suitable + palette and associates it with the image. For an "L" image, 128 graylevels + are allocated. For an "RGB" image, a 6x6x6 colour cube is used, together + with 20 graylevels. + + To make sure that palettes work properly under Windows, you must call the + ``palette`` method upon certain events from Windows. + + :param image: Either a PIL image, or a mode string. If a mode string is + used, a size must also be given. The mode can be one of "1", + "L", "P", or "RGB". + :param size: If the first argument is a mode string, this + defines the size of the image. + """ + + def __init__( + self, image: Image.Image | str, size: tuple[int, int] | None = None + ) -> None: + if isinstance(image, str): + mode = image + image = "" + if size is None: + msg = "If first argument is mode, size is required" + raise ValueError(msg) + else: + mode = image.mode + size = image.size + if mode not in ["1", "L", "P", "RGB"]: + mode = Image.getmodebase(mode) + self.image = Image.core.display(mode, size) + self.mode = mode + self.size = size + if image: + assert not isinstance(image, str) + self.paste(image) + + def expose(self, handle: int | HDC | HWND) -> None: + """ + Copy the bitmap contents to a device context. + + :param handle: Device context (HDC), cast to a Python integer, or an + HDC or HWND instance. In PythonWin, you can use + ``CDC.GetHandleAttrib()`` to get a suitable handle. + """ + handle_int = int(handle) + if isinstance(handle, HWND): + dc = self.image.getdc(handle_int) + try: + self.image.expose(dc) + finally: + self.image.releasedc(handle_int, dc) + else: + self.image.expose(handle_int) + + def draw( + self, + handle: int | HDC | HWND, + dst: tuple[int, int, int, int], + src: tuple[int, int, int, int] | None = None, + ) -> None: + """ + Same as expose, but allows you to specify where to draw the image, and + what part of it to draw. + + The destination and source areas are given as 4-tuple rectangles. If + the source is omitted, the entire image is copied. If the source and + the destination have different sizes, the image is resized as + necessary. + """ + if src is None: + src = (0, 0) + self.size + handle_int = int(handle) + if isinstance(handle, HWND): + dc = self.image.getdc(handle_int) + try: + self.image.draw(dc, dst, src) + finally: + self.image.releasedc(handle_int, dc) + else: + self.image.draw(handle_int, dst, src) + + def query_palette(self, handle: int | HDC | HWND) -> int: + """ + Installs the palette associated with the image in the given device + context. + + This method should be called upon **QUERYNEWPALETTE** and + **PALETTECHANGED** events from Windows. If this method returns a + non-zero value, one or more display palette entries were changed, and + the image should be redrawn. + + :param handle: Device context (HDC), cast to a Python integer, or an + HDC or HWND instance. + :return: The number of entries that were changed (if one or more entries, + this indicates that the image should be redrawn). + """ + handle_int = int(handle) + if isinstance(handle, HWND): + handle = self.image.getdc(handle_int) + try: + result = self.image.query_palette(handle) + finally: + self.image.releasedc(handle, handle) + else: + result = self.image.query_palette(handle_int) + return result + + def paste( + self, im: Image.Image, box: tuple[int, int, int, int] | None = None + ) -> None: + """ + Paste a PIL image into the bitmap image. + + :param im: A PIL image. The size must match the target region. + If the mode does not match, the image is converted to the + mode of the bitmap image. + :param box: A 4-tuple defining the left, upper, right, and + lower pixel coordinate. See :ref:`coordinate-system`. If + None is given instead of a tuple, all of the image is + assumed. + """ + im.load() + if self.mode != im.mode: + im = im.convert(self.mode) + if box: + self.image.paste(im.im, box) + else: + self.image.paste(im.im) + + def frombytes(self, buffer: bytes) -> None: + """ + Load display memory contents from byte data. + + :param buffer: A buffer containing display data (usually + data returned from :py:func:`~PIL.ImageWin.Dib.tobytes`) + """ + self.image.frombytes(buffer) + + def tobytes(self) -> bytes: + """ + Copy display memory contents to bytes object. + + :return: A bytes object containing display data. + """ + return self.image.tobytes() + + +class Window: + """Create a Window with the given title size.""" + + def __init__( + self, title: str = "PIL", width: int | None = None, height: int | None = None + ) -> None: + self.hwnd = Image.core.createwindow( + title, self.__dispatcher, width or 0, height or 0 + ) + + def __dispatcher(self, action: str, *args: int) -> None: + getattr(self, f"ui_handle_{action}")(*args) + + def ui_handle_clear(self, dc: int, x0: int, y0: int, x1: int, y1: int) -> None: + pass + + def ui_handle_damage(self, x0: int, y0: int, x1: int, y1: int) -> None: + pass + + def ui_handle_destroy(self) -> None: + pass + + def ui_handle_repair(self, dc: int, x0: int, y0: int, x1: int, y1: int) -> None: + pass + + def ui_handle_resize(self, width: int, height: int) -> None: + pass + + def mainloop(self) -> None: + Image.core.eventloop() + + +class ImageWindow(Window): + """Create an image window which displays the given image.""" + + def __init__(self, image: Image.Image | Dib, title: str = "PIL") -> None: + if not isinstance(image, Dib): + image = Dib(image) + self.image = image + width, height = image.size + super().__init__(title, width=width, height=height) + + def ui_handle_repair(self, dc: int, x0: int, y0: int, x1: int, y1: int) -> None: + self.image.draw(dc, (x0, y0, x1, y1)) diff --git a/venv/lib/python3.10/site-packages/PIL/ImtImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/ImtImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..c4eccee3423dc6c273bdc1ea88eda5ef4e17cf7d --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/ImtImagePlugin.py @@ -0,0 +1,103 @@ +# +# The Python Imaging Library. +# $Id$ +# +# IM Tools support for PIL +# +# history: +# 1996-05-27 fl Created (read 8-bit images only) +# 2001-02-17 fl Use 're' instead of 'regex' (Python 2.1) (0.2) +# +# Copyright (c) Secret Labs AB 1997-2001. +# Copyright (c) Fredrik Lundh 1996-2001. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import re + +from . import Image, ImageFile + +# +# -------------------------------------------------------------------- + +field = re.compile(rb"([a-z]*) ([^ \r\n]*)") + + +## +# Image plugin for IM Tools images. + + +class ImtImageFile(ImageFile.ImageFile): + format = "IMT" + format_description = "IM Tools" + + def _open(self) -> None: + # Quick rejection: if there's not a LF among the first + # 100 bytes, this is (probably) not a text header. + + assert self.fp is not None + + buffer = self.fp.read(100) + if b"\n" not in buffer: + msg = "not an IM file" + raise SyntaxError(msg) + + xsize = ysize = 0 + + while True: + if buffer: + s = buffer[:1] + buffer = buffer[1:] + else: + s = self.fp.read(1) + if not s: + break + + if s == b"\x0c": + # image data begins + self.tile = [ + ImageFile._Tile( + "raw", + (0, 0) + self.size, + self.fp.tell() - len(buffer), + self.mode, + ) + ] + + break + + else: + # read key/value pair + if b"\n" not in buffer: + buffer += self.fp.read(100) + lines = buffer.split(b"\n") + s += lines.pop(0) + buffer = b"\n".join(lines) + if len(s) == 1 or len(s) > 100: + break + if s[0] == ord(b"*"): + continue # comment + + m = field.match(s) + if not m: + break + k, v = m.group(1, 2) + if k == b"width": + xsize = int(v) + self._size = xsize, ysize + elif k == b"height": + ysize = int(v) + self._size = xsize, ysize + elif k == b"pixel" and v == b"n8": + self._mode = "L" + + +# +# -------------------------------------------------------------------- + +Image.register_open(ImtImageFile.format, ImtImageFile) + +# +# no extension registered (".im" is simply too common) diff --git a/venv/lib/python3.10/site-packages/PIL/IptcImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/IptcImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..c28f4dcc797796e19d164be473b679e848b3e790 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/IptcImagePlugin.py @@ -0,0 +1,229 @@ +# +# The Python Imaging Library. +# $Id$ +# +# IPTC/NAA file handling +# +# history: +# 1995-10-01 fl Created +# 1998-03-09 fl Cleaned up and added to PIL +# 2002-06-18 fl Added getiptcinfo helper +# +# Copyright (c) Secret Labs AB 1997-2002. +# Copyright (c) Fredrik Lundh 1995. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from io import BytesIO +from typing import cast + +from . import Image, ImageFile +from ._binary import i16be as i16 +from ._binary import i32be as i32 + +COMPRESSION = {1: "raw", 5: "jpeg"} + + +# +# Helpers + + +def _i(c: bytes) -> int: + return i32((b"\0\0\0\0" + c)[-4:]) + + +## +# Image plugin for IPTC/NAA datastreams. To read IPTC/NAA fields +# from TIFF and JPEG files, use the getiptcinfo function. + + +class IptcImageFile(ImageFile.ImageFile): + format = "IPTC" + format_description = "IPTC/NAA" + + def getint(self, key: tuple[int, int]) -> int: + return _i(self.info[key]) + + def field(self) -> tuple[tuple[int, int] | None, int]: + # + # get a IPTC field header + s = self.fp.read(5) + if not s.strip(b"\x00"): + return None, 0 + + tag = s[1], s[2] + + # syntax + if s[0] != 0x1C or tag[0] not in [1, 2, 3, 4, 5, 6, 7, 8, 9, 240]: + msg = "invalid IPTC/NAA file" + raise SyntaxError(msg) + + # field size + size = s[3] + if size > 132: + msg = "illegal field length in IPTC/NAA file" + raise OSError(msg) + elif size == 128: + size = 0 + elif size > 128: + size = _i(self.fp.read(size - 128)) + else: + size = i16(s, 3) + + return tag, size + + def _open(self) -> None: + # load descriptive fields + while True: + offset = self.fp.tell() + tag, size = self.field() + if not tag or tag == (8, 10): + break + if size: + tagdata = self.fp.read(size) + else: + tagdata = None + if tag in self.info: + if isinstance(self.info[tag], list): + self.info[tag].append(tagdata) + else: + self.info[tag] = [self.info[tag], tagdata] + else: + self.info[tag] = tagdata + + # mode + layers = self.info[(3, 60)][0] + component = self.info[(3, 60)][1] + if layers == 1 and not component: + self._mode = "L" + band = None + else: + if layers == 3 and component: + self._mode = "RGB" + elif layers == 4 and component: + self._mode = "CMYK" + if (3, 65) in self.info: + band = self.info[(3, 65)][0] - 1 + else: + band = 0 + + # size + self._size = self.getint((3, 20)), self.getint((3, 30)) + + # compression + try: + compression = COMPRESSION[self.getint((3, 120))] + except KeyError as e: + msg = "Unknown IPTC image compression" + raise OSError(msg) from e + + # tile + if tag == (8, 10): + self.tile = [ + ImageFile._Tile("iptc", (0, 0) + self.size, offset, (compression, band)) + ] + + def load(self) -> Image.core.PixelAccess | None: + if self.tile: + args = self.tile[0].args + assert isinstance(args, tuple) + compression, band = args + + self.fp.seek(self.tile[0].offset) + + # Copy image data to temporary file + o = BytesIO() + if compression == "raw": + # To simplify access to the extracted file, + # prepend a PPM header + o.write(b"P5\n%d %d\n255\n" % self.size) + while True: + type, size = self.field() + if type != (8, 10): + break + while size > 0: + s = self.fp.read(min(size, 8192)) + if not s: + break + o.write(s) + size -= len(s) + + with Image.open(o) as _im: + if band is not None: + bands = [Image.new("L", _im.size)] * Image.getmodebands(self.mode) + bands[band] = _im + _im = Image.merge(self.mode, bands) + else: + _im.load() + self.im = _im.im + self.tile = [] + return ImageFile.ImageFile.load(self) + + +Image.register_open(IptcImageFile.format, IptcImageFile) + +Image.register_extension(IptcImageFile.format, ".iim") + + +def getiptcinfo( + im: ImageFile.ImageFile, +) -> dict[tuple[int, int], bytes | list[bytes]] | None: + """ + Get IPTC information from TIFF, JPEG, or IPTC file. + + :param im: An image containing IPTC data. + :returns: A dictionary containing IPTC information, or None if + no IPTC information block was found. + """ + from . import JpegImagePlugin, TiffImagePlugin + + data = None + + info: dict[tuple[int, int], bytes | list[bytes]] = {} + if isinstance(im, IptcImageFile): + # return info dictionary right away + for k, v in im.info.items(): + if isinstance(k, tuple): + info[k] = v + return info + + elif isinstance(im, JpegImagePlugin.JpegImageFile): + # extract the IPTC/NAA resource + photoshop = im.info.get("photoshop") + if photoshop: + data = photoshop.get(0x0404) + + elif isinstance(im, TiffImagePlugin.TiffImageFile): + # get raw data from the IPTC/NAA tag (PhotoShop tags the data + # as 4-byte integers, so we cannot use the get method...) + try: + data = im.tag_v2._tagdata[TiffImagePlugin.IPTC_NAA_CHUNK] + except KeyError: + pass + + if data is None: + return None # no properties + + # create an IptcImagePlugin object without initializing it + class FakeImage: + pass + + fake_im = FakeImage() + fake_im.__class__ = IptcImageFile # type: ignore[assignment] + iptc_im = cast(IptcImageFile, fake_im) + + # parse the IPTC information chunk + iptc_im.info = {} + iptc_im.fp = BytesIO(data) + + try: + iptc_im._open() + except (IndexError, KeyError): + pass # expected failure + + for k, v in iptc_im.info.items(): + if isinstance(k, tuple): + info[k] = v + return info diff --git a/venv/lib/python3.10/site-packages/PIL/Jpeg2KImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/Jpeg2KImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..4c85dd4e2818ee7df837f3f4834eb1fe54b4ddf8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/Jpeg2KImagePlugin.py @@ -0,0 +1,446 @@ +# +# The Python Imaging Library +# $Id$ +# +# JPEG2000 file handling +# +# History: +# 2014-03-12 ajh Created +# 2021-06-30 rogermb Extract dpi information from the 'resc' header box +# +# Copyright (c) 2014 Coriolis Systems Limited +# Copyright (c) 2014 Alastair Houghton +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import os +import struct +from typing import cast + +from . import Image, ImageFile, ImagePalette, _binary + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from typing import IO + + +class BoxReader: + """ + A small helper class to read fields stored in JPEG2000 header boxes + and to easily step into and read sub-boxes. + """ + + def __init__(self, fp: IO[bytes], length: int = -1) -> None: + self.fp = fp + self.has_length = length >= 0 + self.length = length + self.remaining_in_box = -1 + + def _can_read(self, num_bytes: int) -> bool: + if self.has_length and self.fp.tell() + num_bytes > self.length: + # Outside box: ensure we don't read past the known file length + return False + if self.remaining_in_box >= 0: + # Inside box contents: ensure read does not go past box boundaries + return num_bytes <= self.remaining_in_box + else: + return True # No length known, just read + + def _read_bytes(self, num_bytes: int) -> bytes: + if not self._can_read(num_bytes): + msg = "Not enough data in header" + raise SyntaxError(msg) + + data = self.fp.read(num_bytes) + if len(data) < num_bytes: + msg = f"Expected to read {num_bytes} bytes but only got {len(data)}." + raise OSError(msg) + + if self.remaining_in_box > 0: + self.remaining_in_box -= num_bytes + return data + + def read_fields(self, field_format: str) -> tuple[int | bytes, ...]: + size = struct.calcsize(field_format) + data = self._read_bytes(size) + return struct.unpack(field_format, data) + + def read_boxes(self) -> BoxReader: + size = self.remaining_in_box + data = self._read_bytes(size) + return BoxReader(io.BytesIO(data), size) + + def has_next_box(self) -> bool: + if self.has_length: + return self.fp.tell() + self.remaining_in_box < self.length + else: + return True + + def next_box_type(self) -> bytes: + # Skip the rest of the box if it has not been read + if self.remaining_in_box > 0: + self.fp.seek(self.remaining_in_box, os.SEEK_CUR) + self.remaining_in_box = -1 + + # Read the length and type of the next box + lbox, tbox = cast(tuple[int, bytes], self.read_fields(">I4s")) + if lbox == 1: + lbox = cast(int, self.read_fields(">Q")[0]) + hlen = 16 + else: + hlen = 8 + + if lbox < hlen or not self._can_read(lbox - hlen): + msg = "Invalid header length" + raise SyntaxError(msg) + + self.remaining_in_box = lbox - hlen + return tbox + + +def _parse_codestream(fp: IO[bytes]) -> tuple[tuple[int, int], str]: + """Parse the JPEG 2000 codestream to extract the size and component + count from the SIZ marker segment, returning a PIL (size, mode) tuple.""" + + hdr = fp.read(2) + lsiz = _binary.i16be(hdr) + siz = hdr + fp.read(lsiz - 2) + lsiz, rsiz, xsiz, ysiz, xosiz, yosiz, _, _, _, _, csiz = struct.unpack_from( + ">HHIIIIIIIIH", siz + ) + + size = (xsiz - xosiz, ysiz - yosiz) + if csiz == 1: + ssiz = struct.unpack_from(">B", siz, 38) + if (ssiz[0] & 0x7F) + 1 > 8: + mode = "I;16" + else: + mode = "L" + elif csiz == 2: + mode = "LA" + elif csiz == 3: + mode = "RGB" + elif csiz == 4: + mode = "RGBA" + else: + msg = "unable to determine J2K image mode" + raise SyntaxError(msg) + + return size, mode + + +def _res_to_dpi(num: int, denom: int, exp: int) -> float | None: + """Convert JPEG2000's (numerator, denominator, exponent-base-10) resolution, + calculated as (num / denom) * 10^exp and stored in dots per meter, + to floating-point dots per inch.""" + if denom == 0: + return None + return (254 * num * (10**exp)) / (10000 * denom) + + +def _parse_jp2_header( + fp: IO[bytes], +) -> tuple[ + tuple[int, int], + str, + str | None, + tuple[float, float] | None, + ImagePalette.ImagePalette | None, +]: + """Parse the JP2 header box to extract size, component count, + color space information, and optionally DPI information, + returning a (size, mode, mimetype, dpi) tuple.""" + + # Find the JP2 header box + reader = BoxReader(fp) + header = None + mimetype = None + while reader.has_next_box(): + tbox = reader.next_box_type() + + if tbox == b"jp2h": + header = reader.read_boxes() + break + elif tbox == b"ftyp": + if reader.read_fields(">4s")[0] == b"jpx ": + mimetype = "image/jpx" + assert header is not None + + size = None + mode = None + bpc = None + nc = None + dpi = None # 2-tuple of DPI info, or None + palette = None + + while header.has_next_box(): + tbox = header.next_box_type() + + if tbox == b"ihdr": + height, width, nc, bpc = header.read_fields(">IIHB") + assert isinstance(height, int) + assert isinstance(width, int) + assert isinstance(bpc, int) + size = (width, height) + if nc == 1 and (bpc & 0x7F) > 8: + mode = "I;16" + elif nc == 1: + mode = "L" + elif nc == 2: + mode = "LA" + elif nc == 3: + mode = "RGB" + elif nc == 4: + mode = "RGBA" + elif tbox == b"colr" and nc == 4: + meth, _, _, enumcs = header.read_fields(">BBBI") + if meth == 1 and enumcs == 12: + mode = "CMYK" + elif tbox == b"pclr" and mode in ("L", "LA"): + ne, npc = header.read_fields(">HB") + assert isinstance(ne, int) + assert isinstance(npc, int) + max_bitdepth = 0 + for bitdepth in header.read_fields(">" + ("B" * npc)): + assert isinstance(bitdepth, int) + if bitdepth > max_bitdepth: + max_bitdepth = bitdepth + if max_bitdepth <= 8: + palette = ImagePalette.ImagePalette("RGBA" if npc == 4 else "RGB") + for i in range(ne): + color: list[int] = [] + for value in header.read_fields(">" + ("B" * npc)): + assert isinstance(value, int) + color.append(value) + palette.getcolor(tuple(color)) + mode = "P" if mode == "L" else "PA" + elif tbox == b"res ": + res = header.read_boxes() + while res.has_next_box(): + tres = res.next_box_type() + if tres == b"resc": + vrcn, vrcd, hrcn, hrcd, vrce, hrce = res.read_fields(">HHHHBB") + assert isinstance(vrcn, int) + assert isinstance(vrcd, int) + assert isinstance(hrcn, int) + assert isinstance(hrcd, int) + assert isinstance(vrce, int) + assert isinstance(hrce, int) + hres = _res_to_dpi(hrcn, hrcd, hrce) + vres = _res_to_dpi(vrcn, vrcd, vrce) + if hres is not None and vres is not None: + dpi = (hres, vres) + break + + if size is None or mode is None: + msg = "Malformed JP2 header" + raise SyntaxError(msg) + + return size, mode, mimetype, dpi, palette + + +## +# Image plugin for JPEG2000 images. + + +class Jpeg2KImageFile(ImageFile.ImageFile): + format = "JPEG2000" + format_description = "JPEG 2000 (ISO 15444)" + + def _open(self) -> None: + sig = self.fp.read(4) + if sig == b"\xff\x4f\xff\x51": + self.codec = "j2k" + self._size, self._mode = _parse_codestream(self.fp) + self._parse_comment() + else: + sig = sig + self.fp.read(8) + + if sig == b"\x00\x00\x00\x0cjP \x0d\x0a\x87\x0a": + self.codec = "jp2" + header = _parse_jp2_header(self.fp) + self._size, self._mode, self.custom_mimetype, dpi, self.palette = header + if dpi is not None: + self.info["dpi"] = dpi + if self.fp.read(12).endswith(b"jp2c\xff\x4f\xff\x51"): + hdr = self.fp.read(2) + length = _binary.i16be(hdr) + self.fp.seek(length - 2, os.SEEK_CUR) + self._parse_comment() + else: + msg = "not a JPEG 2000 file" + raise SyntaxError(msg) + + self._reduce = 0 + self.layers = 0 + + fd = -1 + length = -1 + + try: + fd = self.fp.fileno() + length = os.fstat(fd).st_size + except Exception: + fd = -1 + try: + pos = self.fp.tell() + self.fp.seek(0, io.SEEK_END) + length = self.fp.tell() + self.fp.seek(pos) + except Exception: + length = -1 + + self.tile = [ + ImageFile._Tile( + "jpeg2k", + (0, 0) + self.size, + 0, + (self.codec, self._reduce, self.layers, fd, length), + ) + ] + + def _parse_comment(self) -> None: + while True: + marker = self.fp.read(2) + if not marker: + break + typ = marker[1] + if typ in (0x90, 0xD9): + # Start of tile or end of codestream + break + hdr = self.fp.read(2) + length = _binary.i16be(hdr) + if typ == 0x64: + # Comment + self.info["comment"] = self.fp.read(length - 2)[2:] + break + else: + self.fp.seek(length - 2, os.SEEK_CUR) + + @property # type: ignore[override] + def reduce( + self, + ) -> ( + Callable[[int | tuple[int, int], tuple[int, int, int, int] | None], Image.Image] + | int + ): + # https://github.com/python-pillow/Pillow/issues/4343 found that the + # new Image 'reduce' method was shadowed by this plugin's 'reduce' + # property. This attempts to allow for both scenarios + return self._reduce or super().reduce + + @reduce.setter + def reduce(self, value: int) -> None: + self._reduce = value + + def load(self) -> Image.core.PixelAccess | None: + if self.tile and self._reduce: + power = 1 << self._reduce + adjust = power >> 1 + self._size = ( + int((self.size[0] + adjust) / power), + int((self.size[1] + adjust) / power), + ) + + # Update the reduce and layers settings + t = self.tile[0] + assert isinstance(t[3], tuple) + t3 = (t[3][0], self._reduce, self.layers, t[3][3], t[3][4]) + self.tile = [ImageFile._Tile(t[0], (0, 0) + self.size, t[2], t3)] + + return ImageFile.ImageFile.load(self) + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith( + (b"\xff\x4f\xff\x51", b"\x00\x00\x00\x0cjP \x0d\x0a\x87\x0a") + ) + + +# ------------------------------------------------------------ +# Save support + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + # Get the keyword arguments + info = im.encoderinfo + + if isinstance(filename, str): + filename = filename.encode() + if filename.endswith(b".j2k") or info.get("no_jp2", False): + kind = "j2k" + else: + kind = "jp2" + + offset = info.get("offset", None) + tile_offset = info.get("tile_offset", None) + tile_size = info.get("tile_size", None) + quality_mode = info.get("quality_mode", "rates") + quality_layers = info.get("quality_layers", None) + if quality_layers is not None and not ( + isinstance(quality_layers, (list, tuple)) + and all( + isinstance(quality_layer, (int, float)) for quality_layer in quality_layers + ) + ): + msg = "quality_layers must be a sequence of numbers" + raise ValueError(msg) + + num_resolutions = info.get("num_resolutions", 0) + cblk_size = info.get("codeblock_size", None) + precinct_size = info.get("precinct_size", None) + irreversible = info.get("irreversible", False) + progression = info.get("progression", "LRCP") + cinema_mode = info.get("cinema_mode", "no") + mct = info.get("mct", 0) + signed = info.get("signed", False) + comment = info.get("comment") + if isinstance(comment, str): + comment = comment.encode() + plt = info.get("plt", False) + + fd = -1 + if hasattr(fp, "fileno"): + try: + fd = fp.fileno() + except Exception: + fd = -1 + + im.encoderconfig = ( + offset, + tile_offset, + tile_size, + quality_mode, + quality_layers, + num_resolutions, + cblk_size, + precinct_size, + irreversible, + progression, + cinema_mode, + mct, + signed, + fd, + comment, + plt, + ) + + ImageFile._save(im, fp, [ImageFile._Tile("jpeg2k", (0, 0) + im.size, 0, kind)]) + + +# ------------------------------------------------------------ +# Registry stuff + + +Image.register_open(Jpeg2KImageFile.format, Jpeg2KImageFile, _accept) +Image.register_save(Jpeg2KImageFile.format, _save) + +Image.register_extensions( + Jpeg2KImageFile.format, [".jp2", ".j2k", ".jpc", ".jpf", ".jpx", ".j2c"] +) + +Image.register_mime(Jpeg2KImageFile.format, "image/jp2") diff --git a/venv/lib/python3.10/site-packages/PIL/JpegImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/JpegImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..755ca648e5567675d35fef291d5ede4e8b89ba7a --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/JpegImagePlugin.py @@ -0,0 +1,888 @@ +# +# The Python Imaging Library. +# $Id$ +# +# JPEG (JFIF) file handling +# +# See "Digital Compression and Coding of Continuous-Tone Still Images, +# Part 1, Requirements and Guidelines" (CCITT T.81 / ISO 10918-1) +# +# History: +# 1995-09-09 fl Created +# 1995-09-13 fl Added full parser +# 1996-03-25 fl Added hack to use the IJG command line utilities +# 1996-05-05 fl Workaround Photoshop 2.5 CMYK polarity bug +# 1996-05-28 fl Added draft support, JFIF version (0.1) +# 1996-12-30 fl Added encoder options, added progression property (0.2) +# 1997-08-27 fl Save mode 1 images as BW (0.3) +# 1998-07-12 fl Added YCbCr to draft and save methods (0.4) +# 1998-10-19 fl Don't hang on files using 16-bit DQT's (0.4.1) +# 2001-04-16 fl Extract DPI settings from JFIF files (0.4.2) +# 2002-07-01 fl Skip pad bytes before markers; identify Exif files (0.4.3) +# 2003-04-25 fl Added experimental EXIF decoder (0.5) +# 2003-06-06 fl Added experimental EXIF GPSinfo decoder +# 2003-09-13 fl Extract COM markers +# 2009-09-06 fl Added icc_profile support (from Florian Hoech) +# 2009-03-06 fl Changed CMYK handling; always use Adobe polarity (0.6) +# 2009-03-08 fl Added subsampling support (from Justin Huff). +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1995-1996 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import array +import io +import math +import os +import struct +import subprocess +import sys +import tempfile +import warnings + +from . import Image, ImageFile +from ._binary import i16be as i16 +from ._binary import i32be as i32 +from ._binary import o8 +from ._binary import o16be as o16 +from .JpegPresets import presets + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import IO, Any + + from .MpoImagePlugin import MpoImageFile + +# +# Parser + + +def Skip(self: JpegImageFile, marker: int) -> None: + n = i16(self.fp.read(2)) - 2 + ImageFile._safe_read(self.fp, n) + + +def APP(self: JpegImageFile, marker: int) -> None: + # + # Application marker. Store these in the APP dictionary. + # Also look for well-known application markers. + + n = i16(self.fp.read(2)) - 2 + s = ImageFile._safe_read(self.fp, n) + + app = f"APP{marker & 15}" + + self.app[app] = s # compatibility + self.applist.append((app, s)) + + if marker == 0xFFE0 and s.startswith(b"JFIF"): + # extract JFIF information + self.info["jfif"] = version = i16(s, 5) # version + self.info["jfif_version"] = divmod(version, 256) + # extract JFIF properties + try: + jfif_unit = s[7] + jfif_density = i16(s, 8), i16(s, 10) + except Exception: + pass + else: + if jfif_unit == 1: + self.info["dpi"] = jfif_density + elif jfif_unit == 2: # cm + # 1 dpcm = 2.54 dpi + self.info["dpi"] = tuple(d * 2.54 for d in jfif_density) + self.info["jfif_unit"] = jfif_unit + self.info["jfif_density"] = jfif_density + elif marker == 0xFFE1 and s.startswith(b"Exif\0\0"): + # extract EXIF information + if "exif" in self.info: + self.info["exif"] += s[6:] + else: + self.info["exif"] = s + self._exif_offset = self.fp.tell() - n + 6 + elif marker == 0xFFE1 and s.startswith(b"http://ns.adobe.com/xap/1.0/\x00"): + self.info["xmp"] = s.split(b"\x00", 1)[1] + elif marker == 0xFFE2 and s.startswith(b"FPXR\0"): + # extract FlashPix information (incomplete) + self.info["flashpix"] = s # FIXME: value will change + elif marker == 0xFFE2 and s.startswith(b"ICC_PROFILE\0"): + # Since an ICC profile can be larger than the maximum size of + # a JPEG marker (64K), we need provisions to split it into + # multiple markers. The format defined by the ICC specifies + # one or more APP2 markers containing the following data: + # Identifying string ASCII "ICC_PROFILE\0" (12 bytes) + # Marker sequence number 1, 2, etc (1 byte) + # Number of markers Total of APP2's used (1 byte) + # Profile data (remainder of APP2 data) + # Decoders should use the marker sequence numbers to + # reassemble the profile, rather than assuming that the APP2 + # markers appear in the correct sequence. + self.icclist.append(s) + elif marker == 0xFFED and s.startswith(b"Photoshop 3.0\x00"): + # parse the image resource block + offset = 14 + photoshop = self.info.setdefault("photoshop", {}) + while s[offset : offset + 4] == b"8BIM": + try: + offset += 4 + # resource code + code = i16(s, offset) + offset += 2 + # resource name (usually empty) + name_len = s[offset] + # name = s[offset+1:offset+1+name_len] + offset += 1 + name_len + offset += offset & 1 # align + # resource data block + size = i32(s, offset) + offset += 4 + data = s[offset : offset + size] + if code == 0x03ED: # ResolutionInfo + photoshop[code] = { + "XResolution": i32(data, 0) / 65536, + "DisplayedUnitsX": i16(data, 4), + "YResolution": i32(data, 8) / 65536, + "DisplayedUnitsY": i16(data, 12), + } + else: + photoshop[code] = data + offset += size + offset += offset & 1 # align + except struct.error: + break # insufficient data + + elif marker == 0xFFEE and s.startswith(b"Adobe"): + self.info["adobe"] = i16(s, 5) + # extract Adobe custom properties + try: + adobe_transform = s[11] + except IndexError: + pass + else: + self.info["adobe_transform"] = adobe_transform + elif marker == 0xFFE2 and s.startswith(b"MPF\0"): + # extract MPO information + self.info["mp"] = s[4:] + # offset is current location minus buffer size + # plus constant header size + self.info["mpoffset"] = self.fp.tell() - n + 4 + + +def COM(self: JpegImageFile, marker: int) -> None: + # + # Comment marker. Store these in the APP dictionary. + n = i16(self.fp.read(2)) - 2 + s = ImageFile._safe_read(self.fp, n) + + self.info["comment"] = s + self.app["COM"] = s # compatibility + self.applist.append(("COM", s)) + + +def SOF(self: JpegImageFile, marker: int) -> None: + # + # Start of frame marker. Defines the size and mode of the + # image. JPEG is colour blind, so we use some simple + # heuristics to map the number of layers to an appropriate + # mode. Note that this could be made a bit brighter, by + # looking for JFIF and Adobe APP markers. + + n = i16(self.fp.read(2)) - 2 + s = ImageFile._safe_read(self.fp, n) + self._size = i16(s, 3), i16(s, 1) + if self._im is not None and self.size != self.im.size: + self._im = None + + self.bits = s[0] + if self.bits != 8: + msg = f"cannot handle {self.bits}-bit layers" + raise SyntaxError(msg) + + self.layers = s[5] + if self.layers == 1: + self._mode = "L" + elif self.layers == 3: + self._mode = "RGB" + elif self.layers == 4: + self._mode = "CMYK" + else: + msg = f"cannot handle {self.layers}-layer images" + raise SyntaxError(msg) + + if marker in [0xFFC2, 0xFFC6, 0xFFCA, 0xFFCE]: + self.info["progressive"] = self.info["progression"] = 1 + + if self.icclist: + # fixup icc profile + self.icclist.sort() # sort by sequence number + if self.icclist[0][13] == len(self.icclist): + profile = [p[14:] for p in self.icclist] + icc_profile = b"".join(profile) + else: + icc_profile = None # wrong number of fragments + self.info["icc_profile"] = icc_profile + self.icclist = [] + + for i in range(6, len(s), 3): + t = s[i : i + 3] + # 4-tuples: id, vsamp, hsamp, qtable + self.layer.append((t[0], t[1] // 16, t[1] & 15, t[2])) + + +def DQT(self: JpegImageFile, marker: int) -> None: + # + # Define quantization table. Note that there might be more + # than one table in each marker. + + # FIXME: The quantization tables can be used to estimate the + # compression quality. + + n = i16(self.fp.read(2)) - 2 + s = ImageFile._safe_read(self.fp, n) + while len(s): + v = s[0] + precision = 1 if (v // 16 == 0) else 2 # in bytes + qt_length = 1 + precision * 64 + if len(s) < qt_length: + msg = "bad quantization table marker" + raise SyntaxError(msg) + data = array.array("B" if precision == 1 else "H", s[1:qt_length]) + if sys.byteorder == "little" and precision > 1: + data.byteswap() # the values are always big-endian + self.quantization[v & 15] = [data[i] for i in zigzag_index] + s = s[qt_length:] + + +# +# JPEG marker table + +MARKER = { + 0xFFC0: ("SOF0", "Baseline DCT", SOF), + 0xFFC1: ("SOF1", "Extended Sequential DCT", SOF), + 0xFFC2: ("SOF2", "Progressive DCT", SOF), + 0xFFC3: ("SOF3", "Spatial lossless", SOF), + 0xFFC4: ("DHT", "Define Huffman table", Skip), + 0xFFC5: ("SOF5", "Differential sequential DCT", SOF), + 0xFFC6: ("SOF6", "Differential progressive DCT", SOF), + 0xFFC7: ("SOF7", "Differential spatial", SOF), + 0xFFC8: ("JPG", "Extension", None), + 0xFFC9: ("SOF9", "Extended sequential DCT (AC)", SOF), + 0xFFCA: ("SOF10", "Progressive DCT (AC)", SOF), + 0xFFCB: ("SOF11", "Spatial lossless DCT (AC)", SOF), + 0xFFCC: ("DAC", "Define arithmetic coding conditioning", Skip), + 0xFFCD: ("SOF13", "Differential sequential DCT (AC)", SOF), + 0xFFCE: ("SOF14", "Differential progressive DCT (AC)", SOF), + 0xFFCF: ("SOF15", "Differential spatial (AC)", SOF), + 0xFFD0: ("RST0", "Restart 0", None), + 0xFFD1: ("RST1", "Restart 1", None), + 0xFFD2: ("RST2", "Restart 2", None), + 0xFFD3: ("RST3", "Restart 3", None), + 0xFFD4: ("RST4", "Restart 4", None), + 0xFFD5: ("RST5", "Restart 5", None), + 0xFFD6: ("RST6", "Restart 6", None), + 0xFFD7: ("RST7", "Restart 7", None), + 0xFFD8: ("SOI", "Start of image", None), + 0xFFD9: ("EOI", "End of image", None), + 0xFFDA: ("SOS", "Start of scan", Skip), + 0xFFDB: ("DQT", "Define quantization table", DQT), + 0xFFDC: ("DNL", "Define number of lines", Skip), + 0xFFDD: ("DRI", "Define restart interval", Skip), + 0xFFDE: ("DHP", "Define hierarchical progression", SOF), + 0xFFDF: ("EXP", "Expand reference component", Skip), + 0xFFE0: ("APP0", "Application segment 0", APP), + 0xFFE1: ("APP1", "Application segment 1", APP), + 0xFFE2: ("APP2", "Application segment 2", APP), + 0xFFE3: ("APP3", "Application segment 3", APP), + 0xFFE4: ("APP4", "Application segment 4", APP), + 0xFFE5: ("APP5", "Application segment 5", APP), + 0xFFE6: ("APP6", "Application segment 6", APP), + 0xFFE7: ("APP7", "Application segment 7", APP), + 0xFFE8: ("APP8", "Application segment 8", APP), + 0xFFE9: ("APP9", "Application segment 9", APP), + 0xFFEA: ("APP10", "Application segment 10", APP), + 0xFFEB: ("APP11", "Application segment 11", APP), + 0xFFEC: ("APP12", "Application segment 12", APP), + 0xFFED: ("APP13", "Application segment 13", APP), + 0xFFEE: ("APP14", "Application segment 14", APP), + 0xFFEF: ("APP15", "Application segment 15", APP), + 0xFFF0: ("JPG0", "Extension 0", None), + 0xFFF1: ("JPG1", "Extension 1", None), + 0xFFF2: ("JPG2", "Extension 2", None), + 0xFFF3: ("JPG3", "Extension 3", None), + 0xFFF4: ("JPG4", "Extension 4", None), + 0xFFF5: ("JPG5", "Extension 5", None), + 0xFFF6: ("JPG6", "Extension 6", None), + 0xFFF7: ("JPG7", "Extension 7", None), + 0xFFF8: ("JPG8", "Extension 8", None), + 0xFFF9: ("JPG9", "Extension 9", None), + 0xFFFA: ("JPG10", "Extension 10", None), + 0xFFFB: ("JPG11", "Extension 11", None), + 0xFFFC: ("JPG12", "Extension 12", None), + 0xFFFD: ("JPG13", "Extension 13", None), + 0xFFFE: ("COM", "Comment", COM), +} + + +def _accept(prefix: bytes) -> bool: + # Magic number was taken from https://en.wikipedia.org/wiki/JPEG + return prefix.startswith(b"\xff\xd8\xff") + + +## +# Image plugin for JPEG and JFIF images. + + +class JpegImageFile(ImageFile.ImageFile): + format = "JPEG" + format_description = "JPEG (ISO 10918)" + + def _open(self) -> None: + s = self.fp.read(3) + + if not _accept(s): + msg = "not a JPEG file" + raise SyntaxError(msg) + s = b"\xff" + + # Create attributes + self.bits = self.layers = 0 + self._exif_offset = 0 + + # JPEG specifics (internal) + self.layer: list[tuple[int, int, int, int]] = [] + self._huffman_dc: dict[Any, Any] = {} + self._huffman_ac: dict[Any, Any] = {} + self.quantization: dict[int, list[int]] = {} + self.app: dict[str, bytes] = {} # compatibility + self.applist: list[tuple[str, bytes]] = [] + self.icclist: list[bytes] = [] + + while True: + i = s[0] + if i == 0xFF: + s = s + self.fp.read(1) + i = i16(s) + else: + # Skip non-0xFF junk + s = self.fp.read(1) + continue + + if i in MARKER: + name, description, handler = MARKER[i] + if handler is not None: + handler(self, i) + if i == 0xFFDA: # start of scan + rawmode = self.mode + if self.mode == "CMYK": + rawmode = "CMYK;I" # assume adobe conventions + self.tile = [ + ImageFile._Tile("jpeg", (0, 0) + self.size, 0, (rawmode, "")) + ] + # self.__offset = self.fp.tell() + break + s = self.fp.read(1) + elif i in {0, 0xFFFF}: + # padded marker or junk; move on + s = b"\xff" + elif i == 0xFF00: # Skip extraneous data (escaped 0xFF) + s = self.fp.read(1) + else: + msg = "no marker found" + raise SyntaxError(msg) + + self._read_dpi_from_exif() + + def __getstate__(self) -> list[Any]: + return super().__getstate__() + [self.layers, self.layer] + + def __setstate__(self, state: list[Any]) -> None: + self.layers, self.layer = state[6:] + super().__setstate__(state) + + def load_read(self, read_bytes: int) -> bytes: + """ + internal: read more image data + For premature EOF and LOAD_TRUNCATED_IMAGES adds EOI marker + so libjpeg can finish decoding + """ + s = self.fp.read(read_bytes) + + if not s and ImageFile.LOAD_TRUNCATED_IMAGES and not hasattr(self, "_ended"): + # Premature EOF. + # Pretend file is finished adding EOI marker + self._ended = True + return b"\xff\xd9" + + return s + + def draft( + self, mode: str | None, size: tuple[int, int] | None + ) -> tuple[str, tuple[int, int, float, float]] | None: + if len(self.tile) != 1: + return None + + # Protect from second call + if self.decoderconfig: + return None + + d, e, o, a = self.tile[0] + scale = 1 + original_size = self.size + + assert isinstance(a, tuple) + if a[0] == "RGB" and mode in ["L", "YCbCr"]: + self._mode = mode + a = mode, "" + + if size: + scale = min(self.size[0] // size[0], self.size[1] // size[1]) + for s in [8, 4, 2, 1]: + if scale >= s: + break + assert e is not None + e = ( + e[0], + e[1], + (e[2] - e[0] + s - 1) // s + e[0], + (e[3] - e[1] + s - 1) // s + e[1], + ) + self._size = ((self.size[0] + s - 1) // s, (self.size[1] + s - 1) // s) + scale = s + + self.tile = [ImageFile._Tile(d, e, o, a)] + self.decoderconfig = (scale, 0) + + box = (0, 0, original_size[0] / scale, original_size[1] / scale) + return self.mode, box + + def load_djpeg(self) -> None: + # ALTERNATIVE: handle JPEGs via the IJG command line utilities + + f, path = tempfile.mkstemp() + os.close(f) + if os.path.exists(self.filename): + subprocess.check_call(["djpeg", "-outfile", path, self.filename]) + else: + try: + os.unlink(path) + except OSError: + pass + + msg = "Invalid Filename" + raise ValueError(msg) + + try: + with Image.open(path) as _im: + _im.load() + self.im = _im.im + finally: + try: + os.unlink(path) + except OSError: + pass + + self._mode = self.im.mode + self._size = self.im.size + + self.tile = [] + + def _getexif(self) -> dict[int, Any] | None: + return _getexif(self) + + def _read_dpi_from_exif(self) -> None: + # If DPI isn't in JPEG header, fetch from EXIF + if "dpi" in self.info or "exif" not in self.info: + return + try: + exif = self.getexif() + resolution_unit = exif[0x0128] + x_resolution = exif[0x011A] + try: + dpi = float(x_resolution[0]) / x_resolution[1] + except TypeError: + dpi = x_resolution + if math.isnan(dpi): + msg = "DPI is not a number" + raise ValueError(msg) + if resolution_unit == 3: # cm + # 1 dpcm = 2.54 dpi + dpi *= 2.54 + self.info["dpi"] = dpi, dpi + except ( + struct.error, # truncated EXIF + KeyError, # dpi not included + SyntaxError, # invalid/unreadable EXIF + TypeError, # dpi is an invalid float + ValueError, # dpi is an invalid float + ZeroDivisionError, # invalid dpi rational value + ): + self.info["dpi"] = 72, 72 + + def _getmp(self) -> dict[int, Any] | None: + return _getmp(self) + + +def _getexif(self: JpegImageFile) -> dict[int, Any] | None: + if "exif" not in self.info: + return None + return self.getexif()._get_merged_dict() + + +def _getmp(self: JpegImageFile) -> dict[int, Any] | None: + # Extract MP information. This method was inspired by the "highly + # experimental" _getexif version that's been in use for years now, + # itself based on the ImageFileDirectory class in the TIFF plugin. + + # The MP record essentially consists of a TIFF file embedded in a JPEG + # application marker. + try: + data = self.info["mp"] + except KeyError: + return None + file_contents = io.BytesIO(data) + head = file_contents.read(8) + endianness = ">" if head.startswith(b"\x4d\x4d\x00\x2a") else "<" + # process dictionary + from . import TiffImagePlugin + + try: + info = TiffImagePlugin.ImageFileDirectory_v2(head) + file_contents.seek(info.next) + info.load(file_contents) + mp = dict(info) + except Exception as e: + msg = "malformed MP Index (unreadable directory)" + raise SyntaxError(msg) from e + # it's an error not to have a number of images + try: + quant = mp[0xB001] + except KeyError as e: + msg = "malformed MP Index (no number of images)" + raise SyntaxError(msg) from e + # get MP entries + mpentries = [] + try: + rawmpentries = mp[0xB002] + for entrynum in range(quant): + unpackedentry = struct.unpack_from( + f"{endianness}LLLHH", rawmpentries, entrynum * 16 + ) + labels = ("Attribute", "Size", "DataOffset", "EntryNo1", "EntryNo2") + mpentry = dict(zip(labels, unpackedentry)) + mpentryattr = { + "DependentParentImageFlag": bool(mpentry["Attribute"] & (1 << 31)), + "DependentChildImageFlag": bool(mpentry["Attribute"] & (1 << 30)), + "RepresentativeImageFlag": bool(mpentry["Attribute"] & (1 << 29)), + "Reserved": (mpentry["Attribute"] & (3 << 27)) >> 27, + "ImageDataFormat": (mpentry["Attribute"] & (7 << 24)) >> 24, + "MPType": mpentry["Attribute"] & 0x00FFFFFF, + } + if mpentryattr["ImageDataFormat"] == 0: + mpentryattr["ImageDataFormat"] = "JPEG" + else: + msg = "unsupported picture format in MPO" + raise SyntaxError(msg) + mptypemap = { + 0x000000: "Undefined", + 0x010001: "Large Thumbnail (VGA Equivalent)", + 0x010002: "Large Thumbnail (Full HD Equivalent)", + 0x020001: "Multi-Frame Image (Panorama)", + 0x020002: "Multi-Frame Image: (Disparity)", + 0x020003: "Multi-Frame Image: (Multi-Angle)", + 0x030000: "Baseline MP Primary Image", + } + mpentryattr["MPType"] = mptypemap.get(mpentryattr["MPType"], "Unknown") + mpentry["Attribute"] = mpentryattr + mpentries.append(mpentry) + mp[0xB002] = mpentries + except KeyError as e: + msg = "malformed MP Index (bad MP Entry)" + raise SyntaxError(msg) from e + # Next we should try and parse the individual image unique ID list; + # we don't because I've never seen this actually used in a real MPO + # file and so can't test it. + return mp + + +# -------------------------------------------------------------------- +# stuff to save JPEG files + +RAWMODE = { + "1": "L", + "L": "L", + "RGB": "RGB", + "RGBX": "RGB", + "CMYK": "CMYK;I", # assume adobe conventions + "YCbCr": "YCbCr", +} + +# fmt: off +zigzag_index = ( + 0, 1, 5, 6, 14, 15, 27, 28, + 2, 4, 7, 13, 16, 26, 29, 42, + 3, 8, 12, 17, 25, 30, 41, 43, + 9, 11, 18, 24, 31, 40, 44, 53, + 10, 19, 23, 32, 39, 45, 52, 54, + 20, 22, 33, 38, 46, 51, 55, 60, + 21, 34, 37, 47, 50, 56, 59, 61, + 35, 36, 48, 49, 57, 58, 62, 63, +) + +samplings = { + (1, 1, 1, 1, 1, 1): 0, + (2, 1, 1, 1, 1, 1): 1, + (2, 2, 1, 1, 1, 1): 2, +} +# fmt: on + + +def get_sampling(im: Image.Image) -> int: + # There's no subsampling when images have only 1 layer + # (grayscale images) or when they are CMYK (4 layers), + # so set subsampling to the default value. + # + # NOTE: currently Pillow can't encode JPEG to YCCK format. + # If YCCK support is added in the future, subsampling code will have + # to be updated (here and in JpegEncode.c) to deal with 4 layers. + if not isinstance(im, JpegImageFile) or im.layers in (1, 4): + return -1 + sampling = im.layer[0][1:3] + im.layer[1][1:3] + im.layer[2][1:3] + return samplings.get(sampling, -1) + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.width == 0 or im.height == 0: + msg = "cannot write empty image as JPEG" + raise ValueError(msg) + + try: + rawmode = RAWMODE[im.mode] + except KeyError as e: + msg = f"cannot write mode {im.mode} as JPEG" + raise OSError(msg) from e + + info = im.encoderinfo + + dpi = [round(x) for x in info.get("dpi", (0, 0))] + + quality = info.get("quality", -1) + subsampling = info.get("subsampling", -1) + qtables = info.get("qtables") + + if quality == "keep": + quality = -1 + subsampling = "keep" + qtables = "keep" + elif quality in presets: + preset = presets[quality] + quality = -1 + subsampling = preset.get("subsampling", -1) + qtables = preset.get("quantization") + elif not isinstance(quality, int): + msg = "Invalid quality setting" + raise ValueError(msg) + else: + if subsampling in presets: + subsampling = presets[subsampling].get("subsampling", -1) + if isinstance(qtables, str) and qtables in presets: + qtables = presets[qtables].get("quantization") + + if subsampling == "4:4:4": + subsampling = 0 + elif subsampling == "4:2:2": + subsampling = 1 + elif subsampling == "4:2:0": + subsampling = 2 + elif subsampling == "4:1:1": + # For compatibility. Before Pillow 4.3, 4:1:1 actually meant 4:2:0. + # Set 4:2:0 if someone is still using that value. + subsampling = 2 + elif subsampling == "keep": + if im.format != "JPEG": + msg = "Cannot use 'keep' when original image is not a JPEG" + raise ValueError(msg) + subsampling = get_sampling(im) + + def validate_qtables( + qtables: ( + str | tuple[list[int], ...] | list[list[int]] | dict[int, list[int]] | None + ), + ) -> list[list[int]] | None: + if qtables is None: + return qtables + if isinstance(qtables, str): + try: + lines = [ + int(num) + for line in qtables.splitlines() + for num in line.split("#", 1)[0].split() + ] + except ValueError as e: + msg = "Invalid quantization table" + raise ValueError(msg) from e + else: + qtables = [lines[s : s + 64] for s in range(0, len(lines), 64)] + if isinstance(qtables, (tuple, list, dict)): + if isinstance(qtables, dict): + qtables = [ + qtables[key] for key in range(len(qtables)) if key in qtables + ] + elif isinstance(qtables, tuple): + qtables = list(qtables) + if not (0 < len(qtables) < 5): + msg = "None or too many quantization tables" + raise ValueError(msg) + for idx, table in enumerate(qtables): + try: + if len(table) != 64: + msg = "Invalid quantization table" + raise TypeError(msg) + table_array = array.array("H", table) + except TypeError as e: + msg = "Invalid quantization table" + raise ValueError(msg) from e + else: + qtables[idx] = list(table_array) + return qtables + + if qtables == "keep": + if im.format != "JPEG": + msg = "Cannot use 'keep' when original image is not a JPEG" + raise ValueError(msg) + qtables = getattr(im, "quantization", None) + qtables = validate_qtables(qtables) + + extra = info.get("extra", b"") + + MAX_BYTES_IN_MARKER = 65533 + if xmp := info.get("xmp"): + overhead_len = 29 # b"http://ns.adobe.com/xap/1.0/\x00" + max_data_bytes_in_marker = MAX_BYTES_IN_MARKER - overhead_len + if len(xmp) > max_data_bytes_in_marker: + msg = "XMP data is too long" + raise ValueError(msg) + size = o16(2 + overhead_len + len(xmp)) + extra += b"\xff\xe1" + size + b"http://ns.adobe.com/xap/1.0/\x00" + xmp + + if icc_profile := info.get("icc_profile"): + overhead_len = 14 # b"ICC_PROFILE\0" + o8(i) + o8(len(markers)) + max_data_bytes_in_marker = MAX_BYTES_IN_MARKER - overhead_len + markers = [] + while icc_profile: + markers.append(icc_profile[:max_data_bytes_in_marker]) + icc_profile = icc_profile[max_data_bytes_in_marker:] + i = 1 + for marker in markers: + size = o16(2 + overhead_len + len(marker)) + extra += ( + b"\xff\xe2" + + size + + b"ICC_PROFILE\0" + + o8(i) + + o8(len(markers)) + + marker + ) + i += 1 + + comment = info.get("comment", im.info.get("comment")) + + # "progressive" is the official name, but older documentation + # says "progression" + # FIXME: issue a warning if the wrong form is used (post-1.1.7) + progressive = info.get("progressive", False) or info.get("progression", False) + + optimize = info.get("optimize", False) + + exif = info.get("exif", b"") + if isinstance(exif, Image.Exif): + exif = exif.tobytes() + if len(exif) > MAX_BYTES_IN_MARKER: + msg = "EXIF data is too long" + raise ValueError(msg) + + # get keyword arguments + im.encoderconfig = ( + quality, + progressive, + info.get("smooth", 0), + optimize, + info.get("keep_rgb", False), + info.get("streamtype", 0), + dpi, + subsampling, + info.get("restart_marker_blocks", 0), + info.get("restart_marker_rows", 0), + qtables, + comment, + extra, + exif, + ) + + # if we optimize, libjpeg needs a buffer big enough to hold the whole image + # in a shot. Guessing on the size, at im.size bytes. (raw pixel size is + # channels*size, this is a value that's been used in a django patch. + # https://github.com/matthewwithanm/django-imagekit/issues/50 + if optimize or progressive: + # CMYK can be bigger + if im.mode == "CMYK": + bufsize = 4 * im.size[0] * im.size[1] + # keep sets quality to -1, but the actual value may be high. + elif quality >= 95 or quality == -1: + bufsize = 2 * im.size[0] * im.size[1] + else: + bufsize = im.size[0] * im.size[1] + if exif: + bufsize += len(exif) + 5 + if extra: + bufsize += len(extra) + 1 + else: + # The EXIF info needs to be written as one block, + APP1, + one spare byte. + # Ensure that our buffer is big enough. Same with the icc_profile block. + bufsize = max(len(exif) + 5, len(extra) + 1) + + ImageFile._save( + im, fp, [ImageFile._Tile("jpeg", (0, 0) + im.size, 0, rawmode)], bufsize + ) + + +## +# Factory for making JPEG and MPO instances +def jpeg_factory( + fp: IO[bytes], filename: str | bytes | None = None +) -> JpegImageFile | MpoImageFile: + im = JpegImageFile(fp, filename) + try: + mpheader = im._getmp() + if mpheader is not None and mpheader[45057] > 1: + for segment, content in im.applist: + if segment == "APP1" and b' hdrgm:Version="' in content: + # Ultra HDR images are not yet supported + return im + # It's actually an MPO + from .MpoImagePlugin import MpoImageFile + + # Don't reload everything, just convert it. + im = MpoImageFile.adopt(im, mpheader) + except (TypeError, IndexError): + # It is really a JPEG + pass + except SyntaxError: + warnings.warn( + "Image appears to be a malformed MPO file, it will be " + "interpreted as a base JPEG file" + ) + return im + + +# --------------------------------------------------------------------- +# Registry stuff + +Image.register_open(JpegImageFile.format, jpeg_factory, _accept) +Image.register_save(JpegImageFile.format, _save) + +Image.register_extensions(JpegImageFile.format, [".jfif", ".jpe", ".jpg", ".jpeg"]) + +Image.register_mime(JpegImageFile.format, "image/jpeg") diff --git a/venv/lib/python3.10/site-packages/PIL/JpegPresets.py b/venv/lib/python3.10/site-packages/PIL/JpegPresets.py new file mode 100644 index 0000000000000000000000000000000000000000..d0e64a35ee1b6fe3ac6da792682a3129253993bb --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/JpegPresets.py @@ -0,0 +1,242 @@ +""" +JPEG quality settings equivalent to the Photoshop settings. +Can be used when saving JPEG files. + +The following presets are available by default: +``web_low``, ``web_medium``, ``web_high``, ``web_very_high``, ``web_maximum``, +``low``, ``medium``, ``high``, ``maximum``. +More presets can be added to the :py:data:`presets` dict if needed. + +To apply the preset, specify:: + + quality="preset_name" + +To apply only the quantization table:: + + qtables="preset_name" + +To apply only the subsampling setting:: + + subsampling="preset_name" + +Example:: + + im.save("image_name.jpg", quality="web_high") + +Subsampling +----------- + +Subsampling is the practice of encoding images by implementing less resolution +for chroma information than for luma information. +(ref.: https://en.wikipedia.org/wiki/Chroma_subsampling) + +Possible subsampling values are 0, 1 and 2 that correspond to 4:4:4, 4:2:2 and +4:2:0. + +You can get the subsampling of a JPEG with the +:func:`.JpegImagePlugin.get_sampling` function. + +In JPEG compressed data a JPEG marker is used instead of an EXIF tag. +(ref.: https://exiv2.org/tags.html) + + +Quantization tables +------------------- + +They are values use by the DCT (Discrete cosine transform) to remove +*unnecessary* information from the image (the lossy part of the compression). +(ref.: https://en.wikipedia.org/wiki/Quantization_matrix#Quantization_matrices, +https://en.wikipedia.org/wiki/JPEG#Quantization) + +You can get the quantization tables of a JPEG with:: + + im.quantization + +This will return a dict with a number of lists. You can pass this dict +directly as the qtables argument when saving a JPEG. + +The quantization table format in presets is a list with sublists. These formats +are interchangeable. + +Libjpeg ref.: +https://web.archive.org/web/20120328125543/http://www.jpegcameras.com/libjpeg/libjpeg-3.html + +""" + +from __future__ import annotations + +# fmt: off +presets = { + 'web_low': {'subsampling': 2, # "4:2:0" + 'quantization': [ + [20, 16, 25, 39, 50, 46, 62, 68, + 16, 18, 23, 38, 38, 53, 65, 68, + 25, 23, 31, 38, 53, 65, 68, 68, + 39, 38, 38, 53, 65, 68, 68, 68, + 50, 38, 53, 65, 68, 68, 68, 68, + 46, 53, 65, 68, 68, 68, 68, 68, + 62, 65, 68, 68, 68, 68, 68, 68, + 68, 68, 68, 68, 68, 68, 68, 68], + [21, 25, 32, 38, 54, 68, 68, 68, + 25, 28, 24, 38, 54, 68, 68, 68, + 32, 24, 32, 43, 66, 68, 68, 68, + 38, 38, 43, 53, 68, 68, 68, 68, + 54, 54, 66, 68, 68, 68, 68, 68, + 68, 68, 68, 68, 68, 68, 68, 68, + 68, 68, 68, 68, 68, 68, 68, 68, + 68, 68, 68, 68, 68, 68, 68, 68] + ]}, + 'web_medium': {'subsampling': 2, # "4:2:0" + 'quantization': [ + [16, 11, 11, 16, 23, 27, 31, 30, + 11, 12, 12, 15, 20, 23, 23, 30, + 11, 12, 13, 16, 23, 26, 35, 47, + 16, 15, 16, 23, 26, 37, 47, 64, + 23, 20, 23, 26, 39, 51, 64, 64, + 27, 23, 26, 37, 51, 64, 64, 64, + 31, 23, 35, 47, 64, 64, 64, 64, + 30, 30, 47, 64, 64, 64, 64, 64], + [17, 15, 17, 21, 20, 26, 38, 48, + 15, 19, 18, 17, 20, 26, 35, 43, + 17, 18, 20, 22, 26, 30, 46, 53, + 21, 17, 22, 28, 30, 39, 53, 64, + 20, 20, 26, 30, 39, 48, 64, 64, + 26, 26, 30, 39, 48, 63, 64, 64, + 38, 35, 46, 53, 64, 64, 64, 64, + 48, 43, 53, 64, 64, 64, 64, 64] + ]}, + 'web_high': {'subsampling': 0, # "4:4:4" + 'quantization': [ + [6, 4, 4, 6, 9, 11, 12, 16, + 4, 5, 5, 6, 8, 10, 12, 12, + 4, 5, 5, 6, 10, 12, 14, 19, + 6, 6, 6, 11, 12, 15, 19, 28, + 9, 8, 10, 12, 16, 20, 27, 31, + 11, 10, 12, 15, 20, 27, 31, 31, + 12, 12, 14, 19, 27, 31, 31, 31, + 16, 12, 19, 28, 31, 31, 31, 31], + [7, 7, 13, 24, 26, 31, 31, 31, + 7, 12, 16, 21, 31, 31, 31, 31, + 13, 16, 17, 31, 31, 31, 31, 31, + 24, 21, 31, 31, 31, 31, 31, 31, + 26, 31, 31, 31, 31, 31, 31, 31, + 31, 31, 31, 31, 31, 31, 31, 31, + 31, 31, 31, 31, 31, 31, 31, 31, + 31, 31, 31, 31, 31, 31, 31, 31] + ]}, + 'web_very_high': {'subsampling': 0, # "4:4:4" + 'quantization': [ + [2, 2, 2, 2, 3, 4, 5, 6, + 2, 2, 2, 2, 3, 4, 5, 6, + 2, 2, 2, 2, 4, 5, 7, 9, + 2, 2, 2, 4, 5, 7, 9, 12, + 3, 3, 4, 5, 8, 10, 12, 12, + 4, 4, 5, 7, 10, 12, 12, 12, + 5, 5, 7, 9, 12, 12, 12, 12, + 6, 6, 9, 12, 12, 12, 12, 12], + [3, 3, 5, 9, 13, 15, 15, 15, + 3, 4, 6, 11, 14, 12, 12, 12, + 5, 6, 9, 14, 12, 12, 12, 12, + 9, 11, 14, 12, 12, 12, 12, 12, + 13, 14, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12] + ]}, + 'web_maximum': {'subsampling': 0, # "4:4:4" + 'quantization': [ + [1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 2, + 1, 1, 1, 1, 1, 1, 2, 2, + 1, 1, 1, 1, 1, 2, 2, 3, + 1, 1, 1, 1, 2, 2, 3, 3, + 1, 1, 1, 2, 2, 3, 3, 3, + 1, 1, 2, 2, 3, 3, 3, 3], + [1, 1, 1, 2, 2, 3, 3, 3, + 1, 1, 1, 2, 3, 3, 3, 3, + 1, 1, 1, 3, 3, 3, 3, 3, + 2, 2, 3, 3, 3, 3, 3, 3, + 2, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3] + ]}, + 'low': {'subsampling': 2, # "4:2:0" + 'quantization': [ + [18, 14, 14, 21, 30, 35, 34, 17, + 14, 16, 16, 19, 26, 23, 12, 12, + 14, 16, 17, 21, 23, 12, 12, 12, + 21, 19, 21, 23, 12, 12, 12, 12, + 30, 26, 23, 12, 12, 12, 12, 12, + 35, 23, 12, 12, 12, 12, 12, 12, + 34, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12], + [20, 19, 22, 27, 20, 20, 17, 17, + 19, 25, 23, 14, 14, 12, 12, 12, + 22, 23, 14, 14, 12, 12, 12, 12, + 27, 14, 14, 12, 12, 12, 12, 12, + 20, 14, 12, 12, 12, 12, 12, 12, + 20, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12] + ]}, + 'medium': {'subsampling': 2, # "4:2:0" + 'quantization': [ + [12, 8, 8, 12, 17, 21, 24, 17, + 8, 9, 9, 11, 15, 19, 12, 12, + 8, 9, 10, 12, 19, 12, 12, 12, + 12, 11, 12, 21, 12, 12, 12, 12, + 17, 15, 19, 12, 12, 12, 12, 12, + 21, 19, 12, 12, 12, 12, 12, 12, + 24, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12], + [13, 11, 13, 16, 20, 20, 17, 17, + 11, 14, 14, 14, 14, 12, 12, 12, + 13, 14, 14, 14, 12, 12, 12, 12, + 16, 14, 14, 12, 12, 12, 12, 12, + 20, 14, 12, 12, 12, 12, 12, 12, + 20, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12] + ]}, + 'high': {'subsampling': 0, # "4:4:4" + 'quantization': [ + [6, 4, 4, 6, 9, 11, 12, 16, + 4, 5, 5, 6, 8, 10, 12, 12, + 4, 5, 5, 6, 10, 12, 12, 12, + 6, 6, 6, 11, 12, 12, 12, 12, + 9, 8, 10, 12, 12, 12, 12, 12, + 11, 10, 12, 12, 12, 12, 12, 12, + 12, 12, 12, 12, 12, 12, 12, 12, + 16, 12, 12, 12, 12, 12, 12, 12], + [7, 7, 13, 24, 20, 20, 17, 17, + 7, 12, 16, 14, 14, 12, 12, 12, + 13, 16, 14, 14, 12, 12, 12, 12, + 24, 14, 14, 12, 12, 12, 12, 12, + 20, 14, 12, 12, 12, 12, 12, 12, + 20, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12, + 17, 12, 12, 12, 12, 12, 12, 12] + ]}, + 'maximum': {'subsampling': 0, # "4:4:4" + 'quantization': [ + [2, 2, 2, 2, 3, 4, 5, 6, + 2, 2, 2, 2, 3, 4, 5, 6, + 2, 2, 2, 2, 4, 5, 7, 9, + 2, 2, 2, 4, 5, 7, 9, 12, + 3, 3, 4, 5, 8, 10, 12, 12, + 4, 4, 5, 7, 10, 12, 12, 12, + 5, 5, 7, 9, 12, 12, 12, 12, + 6, 6, 9, 12, 12, 12, 12, 12], + [3, 3, 5, 9, 13, 15, 15, 15, + 3, 4, 6, 10, 14, 12, 12, 12, + 5, 6, 9, 14, 12, 12, 12, 12, + 9, 10, 14, 12, 12, 12, 12, 12, + 13, 14, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12, + 15, 12, 12, 12, 12, 12, 12, 12] + ]}, +} +# fmt: on diff --git a/venv/lib/python3.10/site-packages/PIL/McIdasImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/McIdasImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..9a47933b69cbdc628faafb67b2fca8de703abfc1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/McIdasImagePlugin.py @@ -0,0 +1,78 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Basic McIdas support for PIL +# +# History: +# 1997-05-05 fl Created (8-bit images only) +# 2009-03-08 fl Added 16/32-bit support. +# +# Thanks to Richard Jones and Craig Swank for specs and samples. +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import struct + +from . import Image, ImageFile + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"\x00\x00\x00\x00\x00\x00\x00\x04") + + +## +# Image plugin for McIdas area images. + + +class McIdasImageFile(ImageFile.ImageFile): + format = "MCIDAS" + format_description = "McIdas area file" + + def _open(self) -> None: + # parse area file directory + assert self.fp is not None + + s = self.fp.read(256) + if not _accept(s) or len(s) != 256: + msg = "not an McIdas area file" + raise SyntaxError(msg) + + self.area_descriptor_raw = s + self.area_descriptor = w = [0, *struct.unpack("!64i", s)] + + # get mode + if w[11] == 1: + mode = rawmode = "L" + elif w[11] == 2: + mode = rawmode = "I;16B" + elif w[11] == 4: + # FIXME: add memory map support + mode = "I" + rawmode = "I;32B" + else: + msg = "unsupported McIdas format" + raise SyntaxError(msg) + + self._mode = mode + self._size = w[10], w[9] + + offset = w[34] + w[15] + stride = w[15] + w[10] * w[11] * w[14] + + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, offset, (rawmode, stride, 1)) + ] + + +# -------------------------------------------------------------------- +# registry + +Image.register_open(McIdasImageFile.format, McIdasImageFile, _accept) + +# no default extension diff --git a/venv/lib/python3.10/site-packages/PIL/MicImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/MicImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..9ce38c427b6c19be9e0c5092181a54b936a7a2f3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/MicImagePlugin.py @@ -0,0 +1,102 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Microsoft Image Composer support for PIL +# +# Notes: +# uses TiffImagePlugin.py to read the actual image streams +# +# History: +# 97-01-20 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import olefile + +from . import Image, TiffImagePlugin + +# +# -------------------------------------------------------------------- + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(olefile.MAGIC) + + +## +# Image plugin for Microsoft's Image Composer file format. + + +class MicImageFile(TiffImagePlugin.TiffImageFile): + format = "MIC" + format_description = "Microsoft Image Composer" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + # read the OLE directory and see if this is a likely + # to be a Microsoft Image Composer file + + try: + self.ole = olefile.OleFileIO(self.fp) + except OSError as e: + msg = "not an MIC file; invalid OLE file" + raise SyntaxError(msg) from e + + # find ACI subfiles with Image members (maybe not the + # best way to identify MIC files, but what the... ;-) + + self.images = [ + path + for path in self.ole.listdir() + if path[1:] and path[0].endswith(".ACI") and path[1] == "Image" + ] + + # if we didn't find any images, this is probably not + # an MIC file. + if not self.images: + msg = "not an MIC file; no image entries" + raise SyntaxError(msg) + + self.frame = -1 + self._n_frames = len(self.images) + self.is_animated = self._n_frames > 1 + + self.__fp = self.fp + self.seek(0) + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + filename = self.images[frame] + self.fp = self.ole.openstream(filename) + + TiffImagePlugin.TiffImageFile._open(self) + + self.frame = frame + + def tell(self) -> int: + return self.frame + + def close(self) -> None: + self.__fp.close() + self.ole.close() + super().close() + + def __exit__(self, *args: object) -> None: + self.__fp.close() + self.ole.close() + super().__exit__() + + +# +# -------------------------------------------------------------------- + +Image.register_open(MicImageFile.format, MicImageFile, _accept) + +Image.register_extension(MicImageFile.format, ".mic") diff --git a/venv/lib/python3.10/site-packages/PIL/MpegImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/MpegImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..47ebe9d62c4edd3b5e97f760ff7e9b0417e5b5ab --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/MpegImagePlugin.py @@ -0,0 +1,84 @@ +# +# The Python Imaging Library. +# $Id$ +# +# MPEG file handling +# +# History: +# 95-09-09 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1995. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image, ImageFile +from ._binary import i8 +from ._typing import SupportsRead + +# +# Bitstream parser + + +class BitStream: + def __init__(self, fp: SupportsRead[bytes]) -> None: + self.fp = fp + self.bits = 0 + self.bitbuffer = 0 + + def next(self) -> int: + return i8(self.fp.read(1)) + + def peek(self, bits: int) -> int: + while self.bits < bits: + self.bitbuffer = (self.bitbuffer << 8) + self.next() + self.bits += 8 + return self.bitbuffer >> (self.bits - bits) & (1 << bits) - 1 + + def skip(self, bits: int) -> None: + while self.bits < bits: + self.bitbuffer = (self.bitbuffer << 8) + i8(self.fp.read(1)) + self.bits += 8 + self.bits = self.bits - bits + + def read(self, bits: int) -> int: + v = self.peek(bits) + self.bits = self.bits - bits + return v + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"\x00\x00\x01\xb3") + + +## +# Image plugin for MPEG streams. This plugin can identify a stream, +# but it cannot read it. + + +class MpegImageFile(ImageFile.ImageFile): + format = "MPEG" + format_description = "MPEG" + + def _open(self) -> None: + assert self.fp is not None + + s = BitStream(self.fp) + if s.read(32) != 0x1B3: + msg = "not an MPEG file" + raise SyntaxError(msg) + + self._mode = "RGB" + self._size = s.read(12), s.read(12) + + +# -------------------------------------------------------------------- +# Registry stuff + +Image.register_open(MpegImageFile.format, MpegImageFile, _accept) + +Image.register_extensions(MpegImageFile.format, [".mpg", ".mpeg"]) + +Image.register_mime(MpegImageFile.format, "video/mpeg") diff --git a/venv/lib/python3.10/site-packages/PIL/MpoImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/MpoImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..b1ae07873ac215b7abeeed9fe32d0f17db45d124 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/MpoImagePlugin.py @@ -0,0 +1,202 @@ +# +# The Python Imaging Library. +# $Id$ +# +# MPO file handling +# +# See "Multi-Picture Format" (CIPA DC-007-Translation 2009, Standard of the +# Camera & Imaging Products Association) +# +# The multi-picture object combines multiple JPEG images (with a modified EXIF +# data format) into a single file. While it can theoretically be used much like +# a GIF animation, it is commonly used to represent 3D photographs and is (as +# of this writing) the most commonly used format by 3D cameras. +# +# History: +# 2014-03-13 Feneric Created +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +import struct +from typing import IO, Any, cast + +from . import ( + Image, + ImageFile, + ImageSequence, + JpegImagePlugin, + TiffImagePlugin, +) +from ._binary import o32le +from ._util import DeferredError + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + JpegImagePlugin._save(im, fp, filename) + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + append_images = im.encoderinfo.get("append_images", []) + if not append_images and not getattr(im, "is_animated", False): + _save(im, fp, filename) + return + + mpf_offset = 28 + offsets: list[int] = [] + im_sequences = [im, *append_images] + total = sum(getattr(seq, "n_frames", 1) for seq in im_sequences) + for im_sequence in im_sequences: + for im_frame in ImageSequence.Iterator(im_sequence): + if not offsets: + # APP2 marker + ifd_length = 66 + 16 * total + im_frame.encoderinfo["extra"] = ( + b"\xff\xe2" + + struct.pack(">H", 6 + ifd_length) + + b"MPF\0" + + b" " * ifd_length + ) + exif = im_frame.encoderinfo.get("exif") + if isinstance(exif, Image.Exif): + exif = exif.tobytes() + im_frame.encoderinfo["exif"] = exif + if exif: + mpf_offset += 4 + len(exif) + + JpegImagePlugin._save(im_frame, fp, filename) + offsets.append(fp.tell()) + else: + encoderinfo = im_frame._attach_default_encoderinfo(im) + im_frame.save(fp, "JPEG") + im_frame.encoderinfo = encoderinfo + offsets.append(fp.tell() - offsets[-1]) + + ifd = TiffImagePlugin.ImageFileDirectory_v2() + ifd[0xB000] = b"0100" + ifd[0xB001] = len(offsets) + + mpentries = b"" + data_offset = 0 + for i, size in enumerate(offsets): + if i == 0: + mptype = 0x030000 # Baseline MP Primary Image + else: + mptype = 0x000000 # Undefined + mpentries += struct.pack(" None: + self.fp.seek(0) # prep the fp in order to pass the JPEG test + JpegImagePlugin.JpegImageFile._open(self) + self._after_jpeg_open() + + def _after_jpeg_open(self, mpheader: dict[int, Any] | None = None) -> None: + self.mpinfo = mpheader if mpheader is not None else self._getmp() + if self.mpinfo is None: + msg = "Image appears to be a malformed MPO file" + raise ValueError(msg) + self.n_frames = self.mpinfo[0xB001] + self.__mpoffsets = [ + mpent["DataOffset"] + self.info["mpoffset"] for mpent in self.mpinfo[0xB002] + ] + self.__mpoffsets[0] = 0 + # Note that the following assertion will only be invalid if something + # gets broken within JpegImagePlugin. + assert self.n_frames == len(self.__mpoffsets) + del self.info["mpoffset"] # no longer needed + self.is_animated = self.n_frames > 1 + self._fp = self.fp # FIXME: hack + self._fp.seek(self.__mpoffsets[0]) # get ready to read first frame + self.__frame = 0 + self.offset = 0 + # for now we can only handle reading and individual frame extraction + self.readonly = 1 + + def load_seek(self, pos: int) -> None: + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self._fp.seek(pos) + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self.fp = self._fp + self.offset = self.__mpoffsets[frame] + + original_exif = self.info.get("exif") + if "exif" in self.info: + del self.info["exif"] + + self.fp.seek(self.offset + 2) # skip SOI marker + if not self.fp.read(2): + msg = "No data found for frame" + raise ValueError(msg) + self.fp.seek(self.offset) + JpegImagePlugin.JpegImageFile._open(self) + if self.info.get("exif") != original_exif: + self._reload_exif() + + self.tile = [ + ImageFile._Tile("jpeg", (0, 0) + self.size, self.offset, self.tile[0][-1]) + ] + self.__frame = frame + + def tell(self) -> int: + return self.__frame + + @staticmethod + def adopt( + jpeg_instance: JpegImagePlugin.JpegImageFile, + mpheader: dict[int, Any] | None = None, + ) -> MpoImageFile: + """ + Transform the instance of JpegImageFile into + an instance of MpoImageFile. + After the call, the JpegImageFile is extended + to be an MpoImageFile. + + This is essentially useful when opening a JPEG + file that reveals itself as an MPO, to avoid + double call to _open. + """ + jpeg_instance.__class__ = MpoImageFile + mpo_instance = cast(MpoImageFile, jpeg_instance) + mpo_instance._after_jpeg_open(mpheader) + return mpo_instance + + +# --------------------------------------------------------------------- +# Registry stuff + +# Note that since MPO shares a factory with JPEG, we do not need to do a +# separate registration for it here. +# Image.register_open(MpoImageFile.format, +# JpegImagePlugin.jpeg_factory, _accept) +Image.register_save(MpoImageFile.format, _save) +Image.register_save_all(MpoImageFile.format, _save_all) + +Image.register_extension(MpoImageFile.format, ".mpo") + +Image.register_mime(MpoImageFile.format, "image/mpo") diff --git a/venv/lib/python3.10/site-packages/PIL/MspImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/MspImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..277087a8677708a3a5fe21a3f6d2c3b27f880d03 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/MspImagePlugin.py @@ -0,0 +1,200 @@ +# +# The Python Imaging Library. +# +# MSP file handling +# +# This is the format used by the Paint program in Windows 1 and 2. +# +# History: +# 95-09-05 fl Created +# 97-01-03 fl Read/write MSP images +# 17-02-21 es Fixed RLE interpretation +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1995-97. +# Copyright (c) Eric Soroos 2017. +# +# See the README file for information on usage and redistribution. +# +# More info on this format: https://archive.org/details/gg243631 +# Page 313: +# Figure 205. Windows Paint Version 1: "DanM" Format +# Figure 206. Windows Paint Version 2: "LinS" Format. Used in Windows V2.03 +# +# See also: https://www.fileformat.info/format/mspaint/egff.htm +from __future__ import annotations + +import io +import struct +from typing import IO + +from . import Image, ImageFile +from ._binary import i16le as i16 +from ._binary import o16le as o16 + +# +# read MSP files + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith((b"DanM", b"LinS")) + + +## +# Image plugin for Windows MSP images. This plugin supports both +# uncompressed (Windows 1.0). + + +class MspImageFile(ImageFile.ImageFile): + format = "MSP" + format_description = "Windows Paint" + + def _open(self) -> None: + # Header + assert self.fp is not None + + s = self.fp.read(32) + if not _accept(s): + msg = "not an MSP file" + raise SyntaxError(msg) + + # Header checksum + checksum = 0 + for i in range(0, 32, 2): + checksum = checksum ^ i16(s, i) + if checksum != 0: + msg = "bad MSP checksum" + raise SyntaxError(msg) + + self._mode = "1" + self._size = i16(s, 4), i16(s, 6) + + if s.startswith(b"DanM"): + self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 32, "1")] + else: + self.tile = [ImageFile._Tile("MSP", (0, 0) + self.size, 32)] + + +class MspDecoder(ImageFile.PyDecoder): + # The algo for the MSP decoder is from + # https://www.fileformat.info/format/mspaint/egff.htm + # cc-by-attribution -- That page references is taken from the + # Encyclopedia of Graphics File Formats and is licensed by + # O'Reilly under the Creative Common/Attribution license + # + # For RLE encoded files, the 32byte header is followed by a scan + # line map, encoded as one 16bit word of encoded byte length per + # line. + # + # NOTE: the encoded length of the line can be 0. This was not + # handled in the previous version of this encoder, and there's no + # mention of how to handle it in the documentation. From the few + # examples I've seen, I've assumed that it is a fill of the + # background color, in this case, white. + # + # + # Pseudocode of the decoder: + # Read a BYTE value as the RunType + # If the RunType value is zero + # Read next byte as the RunCount + # Read the next byte as the RunValue + # Write the RunValue byte RunCount times + # If the RunType value is non-zero + # Use this value as the RunCount + # Read and write the next RunCount bytes literally + # + # e.g.: + # 0x00 03 ff 05 00 01 02 03 04 + # would yield the bytes: + # 0xff ff ff 00 01 02 03 04 + # + # which are then interpreted as a bit packed mode '1' image + + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + + img = io.BytesIO() + blank_line = bytearray((0xFF,) * ((self.state.xsize + 7) // 8)) + try: + self.fd.seek(32) + rowmap = struct.unpack_from( + f"<{self.state.ysize}H", self.fd.read(self.state.ysize * 2) + ) + except struct.error as e: + msg = "Truncated MSP file in row map" + raise OSError(msg) from e + + for x, rowlen in enumerate(rowmap): + try: + if rowlen == 0: + img.write(blank_line) + continue + row = self.fd.read(rowlen) + if len(row) != rowlen: + msg = f"Truncated MSP file, expected {rowlen} bytes on row {x}" + raise OSError(msg) + idx = 0 + while idx < rowlen: + runtype = row[idx] + idx += 1 + if runtype == 0: + (runcount, runval) = struct.unpack_from("Bc", row, idx) + img.write(runval * runcount) + idx += 2 + else: + runcount = runtype + img.write(row[idx : idx + runcount]) + idx += runcount + + except struct.error as e: + msg = f"Corrupted MSP file in row {x}" + raise OSError(msg) from e + + self.set_as_raw(img.getvalue(), "1") + + return -1, 0 + + +Image.register_decoder("MSP", MspDecoder) + + +# +# write MSP files (uncompressed only) + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode != "1": + msg = f"cannot write mode {im.mode} as MSP" + raise OSError(msg) + + # create MSP header + header = [0] * 16 + + header[0], header[1] = i16(b"Da"), i16(b"nM") # version 1 + header[2], header[3] = im.size + header[4], header[5] = 1, 1 + header[6], header[7] = 1, 1 + header[8], header[9] = im.size + + checksum = 0 + for h in header: + checksum = checksum ^ h + header[12] = checksum # FIXME: is this the right field? + + # header + for h in header: + fp.write(o16(h)) + + # image body + ImageFile._save(im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 32, "1")]) + + +# +# registry + +Image.register_open(MspImageFile.format, MspImageFile, _accept) +Image.register_save(MspImageFile.format, _save) + +Image.register_extension(MspImageFile.format, ".msp") diff --git a/venv/lib/python3.10/site-packages/PIL/PSDraw.py b/venv/lib/python3.10/site-packages/PIL/PSDraw.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd4c5c94cfa7ec46332f4da78f3e402fd5b311b --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PSDraw.py @@ -0,0 +1,237 @@ +# +# The Python Imaging Library +# $Id$ +# +# Simple PostScript graphics interface +# +# History: +# 1996-04-20 fl Created +# 1999-01-10 fl Added gsave/grestore to image method +# 2005-05-04 fl Fixed floating point issue in image (from Eric Etheridge) +# +# Copyright (c) 1997-2005 by Secret Labs AB. All rights reserved. +# Copyright (c) 1996 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import sys +from typing import IO + +from . import EpsImagePlugin + +TYPE_CHECKING = False + + +## +# Simple PostScript graphics interface. + + +class PSDraw: + """ + Sets up printing to the given file. If ``fp`` is omitted, + ``sys.stdout.buffer`` is assumed. + """ + + def __init__(self, fp: IO[bytes] | None = None) -> None: + if not fp: + fp = sys.stdout.buffer + self.fp = fp + + def begin_document(self, id: str | None = None) -> None: + """Set up printing of a document. (Write PostScript DSC header.)""" + # FIXME: incomplete + self.fp.write( + b"%!PS-Adobe-3.0\n" + b"save\n" + b"/showpage { } def\n" + b"%%EndComments\n" + b"%%BeginDocument\n" + ) + # self.fp.write(ERROR_PS) # debugging! + self.fp.write(EDROFF_PS) + self.fp.write(VDI_PS) + self.fp.write(b"%%EndProlog\n") + self.isofont: dict[bytes, int] = {} + + def end_document(self) -> None: + """Ends printing. (Write PostScript DSC footer.)""" + self.fp.write(b"%%EndDocument\nrestore showpage\n%%End\n") + if hasattr(self.fp, "flush"): + self.fp.flush() + + def setfont(self, font: str, size: int) -> None: + """ + Selects which font to use. + + :param font: A PostScript font name + :param size: Size in points. + """ + font_bytes = bytes(font, "UTF-8") + if font_bytes not in self.isofont: + # reencode font + self.fp.write( + b"/PSDraw-%s ISOLatin1Encoding /%s E\n" % (font_bytes, font_bytes) + ) + self.isofont[font_bytes] = 1 + # rough + self.fp.write(b"/F0 %d /PSDraw-%s F\n" % (size, font_bytes)) + + def line(self, xy0: tuple[int, int], xy1: tuple[int, int]) -> None: + """ + Draws a line between the two points. Coordinates are given in + PostScript point coordinates (72 points per inch, (0, 0) is the lower + left corner of the page). + """ + self.fp.write(b"%d %d %d %d Vl\n" % (*xy0, *xy1)) + + def rectangle(self, box: tuple[int, int, int, int]) -> None: + """ + Draws a rectangle. + + :param box: A tuple of four integers, specifying left, bottom, width and + height. + """ + self.fp.write(b"%d %d M 0 %d %d Vr\n" % box) + + def text(self, xy: tuple[int, int], text: str) -> None: + """ + Draws text at the given position. You must use + :py:meth:`~PIL.PSDraw.PSDraw.setfont` before calling this method. + """ + text_bytes = bytes(text, "UTF-8") + text_bytes = b"\\(".join(text_bytes.split(b"(")) + text_bytes = b"\\)".join(text_bytes.split(b")")) + self.fp.write(b"%d %d M (%s) S\n" % (xy + (text_bytes,))) + + if TYPE_CHECKING: + from . import Image + + def image( + self, box: tuple[int, int, int, int], im: Image.Image, dpi: int | None = None + ) -> None: + """Draw a PIL image, centered in the given box.""" + # default resolution depends on mode + if not dpi: + if im.mode == "1": + dpi = 200 # fax + else: + dpi = 100 # grayscale + # image size (on paper) + x = im.size[0] * 72 / dpi + y = im.size[1] * 72 / dpi + # max allowed size + xmax = float(box[2] - box[0]) + ymax = float(box[3] - box[1]) + if x > xmax: + y = y * xmax / x + x = xmax + if y > ymax: + x = x * ymax / y + y = ymax + dx = (xmax - x) / 2 + box[0] + dy = (ymax - y) / 2 + box[1] + self.fp.write(b"gsave\n%f %f translate\n" % (dx, dy)) + if (x, y) != im.size: + # EpsImagePlugin._save prints the image at (0,0,xsize,ysize) + sx = x / im.size[0] + sy = y / im.size[1] + self.fp.write(b"%f %f scale\n" % (sx, sy)) + EpsImagePlugin._save(im, self.fp, "", 0) + self.fp.write(b"\ngrestore\n") + + +# -------------------------------------------------------------------- +# PostScript driver + +# +# EDROFF.PS -- PostScript driver for Edroff 2 +# +# History: +# 94-01-25 fl: created (edroff 2.04) +# +# Copyright (c) Fredrik Lundh 1994. +# + + +EDROFF_PS = b"""\ +/S { show } bind def +/P { moveto show } bind def +/M { moveto } bind def +/X { 0 rmoveto } bind def +/Y { 0 exch rmoveto } bind def +/E { findfont + dup maxlength dict begin + { + 1 index /FID ne { def } { pop pop } ifelse + } forall + /Encoding exch def + dup /FontName exch def + currentdict end definefont pop +} bind def +/F { findfont exch scalefont dup setfont + [ exch /setfont cvx ] cvx bind def +} bind def +""" + +# +# VDI.PS -- PostScript driver for VDI meta commands +# +# History: +# 94-01-25 fl: created (edroff 2.04) +# +# Copyright (c) Fredrik Lundh 1994. +# + +VDI_PS = b"""\ +/Vm { moveto } bind def +/Va { newpath arcn stroke } bind def +/Vl { moveto lineto stroke } bind def +/Vc { newpath 0 360 arc closepath } bind def +/Vr { exch dup 0 rlineto + exch dup 0 exch rlineto + exch neg 0 rlineto + 0 exch neg rlineto + setgray fill } bind def +/Tm matrix def +/Ve { Tm currentmatrix pop + translate scale newpath 0 0 .5 0 360 arc closepath + Tm setmatrix +} bind def +/Vf { currentgray exch setgray fill setgray } bind def +""" + +# +# ERROR.PS -- Error handler +# +# History: +# 89-11-21 fl: created (pslist 1.10) +# + +ERROR_PS = b"""\ +/landscape false def +/errorBUF 200 string def +/errorNL { currentpoint 10 sub exch pop 72 exch moveto } def +errordict begin /handleerror { + initmatrix /Courier findfont 10 scalefont setfont + newpath 72 720 moveto $error begin /newerror false def + (PostScript Error) show errorNL errorNL + (Error: ) show + /errorname load errorBUF cvs show errorNL errorNL + (Command: ) show + /command load dup type /stringtype ne { errorBUF cvs } if show + errorNL errorNL + (VMstatus: ) show + vmstatus errorBUF cvs show ( bytes available, ) show + errorBUF cvs show ( bytes used at level ) show + errorBUF cvs show errorNL errorNL + (Operand stargck: ) show errorNL /ostargck load { + dup type /stringtype ne { errorBUF cvs } if 72 0 rmoveto show errorNL + } forall errorNL + (Execution stargck: ) show errorNL /estargck load { + dup type /stringtype ne { errorBUF cvs } if 72 0 rmoveto show errorNL + } forall + end showpage +} def end +""" diff --git a/venv/lib/python3.10/site-packages/PIL/PaletteFile.py b/venv/lib/python3.10/site-packages/PIL/PaletteFile.py new file mode 100644 index 0000000000000000000000000000000000000000..2a26e5d4e223ba0bc80ad1bfb37b4c3927e222ac --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PaletteFile.py @@ -0,0 +1,54 @@ +# +# Python Imaging Library +# $Id$ +# +# stuff to read simple, teragon-style palette files +# +# History: +# 97-08-23 fl Created +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from typing import IO + +from ._binary import o8 + + +class PaletteFile: + """File handler for Teragon-style palette files.""" + + rawmode = "RGB" + + def __init__(self, fp: IO[bytes]) -> None: + palette = [o8(i) * 3 for i in range(256)] + + while True: + s = fp.readline() + + if not s: + break + if s.startswith(b"#"): + continue + if len(s) > 100: + msg = "bad palette file" + raise SyntaxError(msg) + + v = [int(x) for x in s.split()] + try: + [i, r, g, b] = v + except ValueError: + [i, r] = v + g = b = r + + if 0 <= i <= 255: + palette[i] = o8(r) + o8(g) + o8(b) + + self.palette = b"".join(palette) + + def getpalette(self) -> tuple[bytes, str]: + return self.palette, self.rawmode diff --git a/venv/lib/python3.10/site-packages/PIL/PalmImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PalmImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..15f71290816c5fa6a5178842260a1520eb0b372f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PalmImagePlugin.py @@ -0,0 +1,217 @@ +# +# The Python Imaging Library. +# $Id$ +# + +## +# Image plugin for Palm pixmap images (output only). +## +from __future__ import annotations + +from typing import IO + +from . import Image, ImageFile +from ._binary import o8 +from ._binary import o16be as o16b + +# fmt: off +_Palm8BitColormapValues = ( + (255, 255, 255), (255, 204, 255), (255, 153, 255), (255, 102, 255), + (255, 51, 255), (255, 0, 255), (255, 255, 204), (255, 204, 204), + (255, 153, 204), (255, 102, 204), (255, 51, 204), (255, 0, 204), + (255, 255, 153), (255, 204, 153), (255, 153, 153), (255, 102, 153), + (255, 51, 153), (255, 0, 153), (204, 255, 255), (204, 204, 255), + (204, 153, 255), (204, 102, 255), (204, 51, 255), (204, 0, 255), + (204, 255, 204), (204, 204, 204), (204, 153, 204), (204, 102, 204), + (204, 51, 204), (204, 0, 204), (204, 255, 153), (204, 204, 153), + (204, 153, 153), (204, 102, 153), (204, 51, 153), (204, 0, 153), + (153, 255, 255), (153, 204, 255), (153, 153, 255), (153, 102, 255), + (153, 51, 255), (153, 0, 255), (153, 255, 204), (153, 204, 204), + (153, 153, 204), (153, 102, 204), (153, 51, 204), (153, 0, 204), + (153, 255, 153), (153, 204, 153), (153, 153, 153), (153, 102, 153), + (153, 51, 153), (153, 0, 153), (102, 255, 255), (102, 204, 255), + (102, 153, 255), (102, 102, 255), (102, 51, 255), (102, 0, 255), + (102, 255, 204), (102, 204, 204), (102, 153, 204), (102, 102, 204), + (102, 51, 204), (102, 0, 204), (102, 255, 153), (102, 204, 153), + (102, 153, 153), (102, 102, 153), (102, 51, 153), (102, 0, 153), + (51, 255, 255), (51, 204, 255), (51, 153, 255), (51, 102, 255), + (51, 51, 255), (51, 0, 255), (51, 255, 204), (51, 204, 204), + (51, 153, 204), (51, 102, 204), (51, 51, 204), (51, 0, 204), + (51, 255, 153), (51, 204, 153), (51, 153, 153), (51, 102, 153), + (51, 51, 153), (51, 0, 153), (0, 255, 255), (0, 204, 255), + (0, 153, 255), (0, 102, 255), (0, 51, 255), (0, 0, 255), + (0, 255, 204), (0, 204, 204), (0, 153, 204), (0, 102, 204), + (0, 51, 204), (0, 0, 204), (0, 255, 153), (0, 204, 153), + (0, 153, 153), (0, 102, 153), (0, 51, 153), (0, 0, 153), + (255, 255, 102), (255, 204, 102), (255, 153, 102), (255, 102, 102), + (255, 51, 102), (255, 0, 102), (255, 255, 51), (255, 204, 51), + (255, 153, 51), (255, 102, 51), (255, 51, 51), (255, 0, 51), + (255, 255, 0), (255, 204, 0), (255, 153, 0), (255, 102, 0), + (255, 51, 0), (255, 0, 0), (204, 255, 102), (204, 204, 102), + (204, 153, 102), (204, 102, 102), (204, 51, 102), (204, 0, 102), + (204, 255, 51), (204, 204, 51), (204, 153, 51), (204, 102, 51), + (204, 51, 51), (204, 0, 51), (204, 255, 0), (204, 204, 0), + (204, 153, 0), (204, 102, 0), (204, 51, 0), (204, 0, 0), + (153, 255, 102), (153, 204, 102), (153, 153, 102), (153, 102, 102), + (153, 51, 102), (153, 0, 102), (153, 255, 51), (153, 204, 51), + (153, 153, 51), (153, 102, 51), (153, 51, 51), (153, 0, 51), + (153, 255, 0), (153, 204, 0), (153, 153, 0), (153, 102, 0), + (153, 51, 0), (153, 0, 0), (102, 255, 102), (102, 204, 102), + (102, 153, 102), (102, 102, 102), (102, 51, 102), (102, 0, 102), + (102, 255, 51), (102, 204, 51), (102, 153, 51), (102, 102, 51), + (102, 51, 51), (102, 0, 51), (102, 255, 0), (102, 204, 0), + (102, 153, 0), (102, 102, 0), (102, 51, 0), (102, 0, 0), + (51, 255, 102), (51, 204, 102), (51, 153, 102), (51, 102, 102), + (51, 51, 102), (51, 0, 102), (51, 255, 51), (51, 204, 51), + (51, 153, 51), (51, 102, 51), (51, 51, 51), (51, 0, 51), + (51, 255, 0), (51, 204, 0), (51, 153, 0), (51, 102, 0), + (51, 51, 0), (51, 0, 0), (0, 255, 102), (0, 204, 102), + (0, 153, 102), (0, 102, 102), (0, 51, 102), (0, 0, 102), + (0, 255, 51), (0, 204, 51), (0, 153, 51), (0, 102, 51), + (0, 51, 51), (0, 0, 51), (0, 255, 0), (0, 204, 0), + (0, 153, 0), (0, 102, 0), (0, 51, 0), (17, 17, 17), + (34, 34, 34), (68, 68, 68), (85, 85, 85), (119, 119, 119), + (136, 136, 136), (170, 170, 170), (187, 187, 187), (221, 221, 221), + (238, 238, 238), (192, 192, 192), (128, 0, 0), (128, 0, 128), + (0, 128, 0), (0, 128, 128), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), + (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)) +# fmt: on + + +# so build a prototype image to be used for palette resampling +def build_prototype_image() -> Image.Image: + image = Image.new("L", (1, len(_Palm8BitColormapValues))) + image.putdata(list(range(len(_Palm8BitColormapValues)))) + palettedata: tuple[int, ...] = () + for colormapValue in _Palm8BitColormapValues: + palettedata += colormapValue + palettedata += (0, 0, 0) * (256 - len(_Palm8BitColormapValues)) + image.putpalette(palettedata) + return image + + +Palm8BitColormapImage = build_prototype_image() + +# OK, we now have in Palm8BitColormapImage, +# a "P"-mode image with the right palette +# +# -------------------------------------------------------------------- + +_FLAGS = {"custom-colormap": 0x4000, "is-compressed": 0x8000, "has-transparent": 0x2000} + +_COMPRESSION_TYPES = {"none": 0xFF, "rle": 0x01, "scanline": 0x00} + + +# +# -------------------------------------------------------------------- + +## +# (Internal) Image save plugin for the Palm format. + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode == "P": + rawmode = "P" + bpp = 8 + version = 1 + + elif im.mode == "L": + if im.encoderinfo.get("bpp") in (1, 2, 4): + # this is 8-bit grayscale, so we shift it to get the high-order bits, + # and invert it because + # Palm does grayscale from white (0) to black (1) + bpp = im.encoderinfo["bpp"] + maxval = (1 << bpp) - 1 + shift = 8 - bpp + im = im.point(lambda x: maxval - (x >> shift)) + elif im.info.get("bpp") in (1, 2, 4): + # here we assume that even though the inherent mode is 8-bit grayscale, + # only the lower bpp bits are significant. + # We invert them to match the Palm. + bpp = im.info["bpp"] + maxval = (1 << bpp) - 1 + im = im.point(lambda x: maxval - (x & maxval)) + else: + msg = f"cannot write mode {im.mode} as Palm" + raise OSError(msg) + + # we ignore the palette here + im._mode = "P" + rawmode = f"P;{bpp}" + version = 1 + + elif im.mode == "1": + # monochrome -- write it inverted, as is the Palm standard + rawmode = "1;I" + bpp = 1 + version = 0 + + else: + msg = f"cannot write mode {im.mode} as Palm" + raise OSError(msg) + + # + # make sure image data is available + im.load() + + # write header + + cols = im.size[0] + rows = im.size[1] + + rowbytes = int((cols + (16 // bpp - 1)) / (16 // bpp)) * 2 + transparent_index = 0 + compression_type = _COMPRESSION_TYPES["none"] + + flags = 0 + if im.mode == "P": + flags |= _FLAGS["custom-colormap"] + colormap = im.im.getpalette() + colors = len(colormap) // 3 + colormapsize = 4 * colors + 2 + else: + colormapsize = 0 + + if "offset" in im.info: + offset = (rowbytes * rows + 16 + 3 + colormapsize) // 4 + else: + offset = 0 + + fp.write(o16b(cols) + o16b(rows) + o16b(rowbytes) + o16b(flags)) + fp.write(o8(bpp)) + fp.write(o8(version)) + fp.write(o16b(offset)) + fp.write(o8(transparent_index)) + fp.write(o8(compression_type)) + fp.write(o16b(0)) # reserved by Palm + + # now write colormap if necessary + + if colormapsize: + fp.write(o16b(colors)) + for i in range(colors): + fp.write(o8(i)) + fp.write(colormap[3 * i : 3 * i + 3]) + + # now convert data to raw form + ImageFile._save( + im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 0, (rawmode, rowbytes, 1))] + ) + + if hasattr(fp, "flush"): + fp.flush() + + +# +# -------------------------------------------------------------------- + +Image.register_save("Palm", _save) + +Image.register_extension("Palm", ".palm") + +Image.register_mime("Palm", "image/palm") diff --git a/venv/lib/python3.10/site-packages/PIL/PcdImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PcdImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..296f3775b0d064627888df8958786731643989cd --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PcdImagePlugin.py @@ -0,0 +1,68 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PCD file handling +# +# History: +# 96-05-10 fl Created +# 96-05-27 fl Added draft mode (128x192, 256x384) +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image, ImageFile + +## +# Image plugin for PhotoCD images. This plugin only reads the 768x512 +# image from the file; higher resolutions are encoded in a proprietary +# encoding. + + +class PcdImageFile(ImageFile.ImageFile): + format = "PCD" + format_description = "Kodak PhotoCD" + + def _open(self) -> None: + # rough + assert self.fp is not None + + self.fp.seek(2048) + s = self.fp.read(1539) + + if not s.startswith(b"PCD_"): + msg = "not a PCD file" + raise SyntaxError(msg) + + orientation = s[1538] & 3 + self.tile_post_rotate = None + if orientation == 1: + self.tile_post_rotate = 90 + elif orientation == 3: + self.tile_post_rotate = 270 + + self._mode = "RGB" + self._size = (512, 768) if orientation in (1, 3) else (768, 512) + self.tile = [ImageFile._Tile("pcd", (0, 0, 768, 512), 96 * 2048)] + + def load_prepare(self) -> None: + if self._im is None and self.tile_post_rotate: + self.im = Image.core.new(self.mode, (768, 512)) + ImageFile.ImageFile.load_prepare(self) + + def load_end(self) -> None: + if self.tile_post_rotate: + # Handle rotated PCDs + self.im = self.rotate(self.tile_post_rotate, expand=True).im + + +# +# registry + +Image.register_open(PcdImageFile.format, PcdImageFile) + +Image.register_extension(PcdImageFile.format, ".pcd") diff --git a/venv/lib/python3.10/site-packages/PIL/PcfFontFile.py b/venv/lib/python3.10/site-packages/PIL/PcfFontFile.py new file mode 100644 index 0000000000000000000000000000000000000000..a00e9b9198430c28b2aeb49df57ab17f78369995 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PcfFontFile.py @@ -0,0 +1,258 @@ +# +# THIS IS WORK IN PROGRESS +# +# The Python Imaging Library +# $Id$ +# +# portable compiled font file parser +# +# history: +# 1997-08-19 fl created +# 2003-09-13 fl fixed loading of unicode fonts +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1997-2003 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io + +from . import FontFile, Image +from ._binary import i8 +from ._binary import i16be as b16 +from ._binary import i16le as l16 +from ._binary import i32be as b32 +from ._binary import i32le as l32 + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from typing import BinaryIO + +# -------------------------------------------------------------------- +# declarations + +PCF_MAGIC = 0x70636601 # "\x01fcp" + +PCF_PROPERTIES = 1 << 0 +PCF_ACCELERATORS = 1 << 1 +PCF_METRICS = 1 << 2 +PCF_BITMAPS = 1 << 3 +PCF_INK_METRICS = 1 << 4 +PCF_BDF_ENCODINGS = 1 << 5 +PCF_SWIDTHS = 1 << 6 +PCF_GLYPH_NAMES = 1 << 7 +PCF_BDF_ACCELERATORS = 1 << 8 + +BYTES_PER_ROW: list[Callable[[int], int]] = [ + lambda bits: ((bits + 7) >> 3), + lambda bits: ((bits + 15) >> 3) & ~1, + lambda bits: ((bits + 31) >> 3) & ~3, + lambda bits: ((bits + 63) >> 3) & ~7, +] + + +def sz(s: bytes, o: int) -> bytes: + return s[o : s.index(b"\0", o)] + + +class PcfFontFile(FontFile.FontFile): + """Font file plugin for the X11 PCF format.""" + + name = "name" + + def __init__(self, fp: BinaryIO, charset_encoding: str = "iso8859-1"): + self.charset_encoding = charset_encoding + + magic = l32(fp.read(4)) + if magic != PCF_MAGIC: + msg = "not a PCF file" + raise SyntaxError(msg) + + super().__init__() + + count = l32(fp.read(4)) + self.toc = {} + for i in range(count): + type = l32(fp.read(4)) + self.toc[type] = l32(fp.read(4)), l32(fp.read(4)), l32(fp.read(4)) + + self.fp = fp + + self.info = self._load_properties() + + metrics = self._load_metrics() + bitmaps = self._load_bitmaps(metrics) + encoding = self._load_encoding() + + # + # create glyph structure + + for ch, ix in enumerate(encoding): + if ix is not None: + ( + xsize, + ysize, + left, + right, + width, + ascent, + descent, + attributes, + ) = metrics[ix] + self.glyph[ch] = ( + (width, 0), + (left, descent - ysize, xsize + left, descent), + (0, 0, xsize, ysize), + bitmaps[ix], + ) + + def _getformat( + self, tag: int + ) -> tuple[BinaryIO, int, Callable[[bytes], int], Callable[[bytes], int]]: + format, size, offset = self.toc[tag] + + fp = self.fp + fp.seek(offset) + + format = l32(fp.read(4)) + + if format & 4: + i16, i32 = b16, b32 + else: + i16, i32 = l16, l32 + + return fp, format, i16, i32 + + def _load_properties(self) -> dict[bytes, bytes | int]: + # + # font properties + + properties = {} + + fp, format, i16, i32 = self._getformat(PCF_PROPERTIES) + + nprops = i32(fp.read(4)) + + # read property description + p = [(i32(fp.read(4)), i8(fp.read(1)), i32(fp.read(4))) for _ in range(nprops)] + + if nprops & 3: + fp.seek(4 - (nprops & 3), io.SEEK_CUR) # pad + + data = fp.read(i32(fp.read(4))) + + for k, s, v in p: + property_value: bytes | int = sz(data, v) if s else v + properties[sz(data, k)] = property_value + + return properties + + def _load_metrics(self) -> list[tuple[int, int, int, int, int, int, int, int]]: + # + # font metrics + + metrics: list[tuple[int, int, int, int, int, int, int, int]] = [] + + fp, format, i16, i32 = self._getformat(PCF_METRICS) + + append = metrics.append + + if (format & 0xFF00) == 0x100: + # "compressed" metrics + for i in range(i16(fp.read(2))): + left = i8(fp.read(1)) - 128 + right = i8(fp.read(1)) - 128 + width = i8(fp.read(1)) - 128 + ascent = i8(fp.read(1)) - 128 + descent = i8(fp.read(1)) - 128 + xsize = right - left + ysize = ascent + descent + append((xsize, ysize, left, right, width, ascent, descent, 0)) + + else: + # "jumbo" metrics + for i in range(i32(fp.read(4))): + left = i16(fp.read(2)) + right = i16(fp.read(2)) + width = i16(fp.read(2)) + ascent = i16(fp.read(2)) + descent = i16(fp.read(2)) + attributes = i16(fp.read(2)) + xsize = right - left + ysize = ascent + descent + append((xsize, ysize, left, right, width, ascent, descent, attributes)) + + return metrics + + def _load_bitmaps( + self, metrics: list[tuple[int, int, int, int, int, int, int, int]] + ) -> list[Image.Image]: + # + # bitmap data + + fp, format, i16, i32 = self._getformat(PCF_BITMAPS) + + nbitmaps = i32(fp.read(4)) + + if nbitmaps != len(metrics): + msg = "Wrong number of bitmaps" + raise OSError(msg) + + offsets = [i32(fp.read(4)) for _ in range(nbitmaps)] + + bitmap_sizes = [i32(fp.read(4)) for _ in range(4)] + + # byteorder = format & 4 # non-zero => MSB + bitorder = format & 8 # non-zero => MSB + padindex = format & 3 + + bitmapsize = bitmap_sizes[padindex] + offsets.append(bitmapsize) + + data = fp.read(bitmapsize) + + pad = BYTES_PER_ROW[padindex] + mode = "1;R" + if bitorder: + mode = "1" + + bitmaps = [] + for i in range(nbitmaps): + xsize, ysize = metrics[i][:2] + b, e = offsets[i : i + 2] + bitmaps.append( + Image.frombytes("1", (xsize, ysize), data[b:e], "raw", mode, pad(xsize)) + ) + + return bitmaps + + def _load_encoding(self) -> list[int | None]: + fp, format, i16, i32 = self._getformat(PCF_BDF_ENCODINGS) + + first_col, last_col = i16(fp.read(2)), i16(fp.read(2)) + first_row, last_row = i16(fp.read(2)), i16(fp.read(2)) + + i16(fp.read(2)) # default + + nencoding = (last_col - first_col + 1) * (last_row - first_row + 1) + + # map character code to bitmap index + encoding: list[int | None] = [None] * min(256, nencoding) + + encoding_offsets = [i16(fp.read(2)) for _ in range(nencoding)] + + for i in range(first_col, len(encoding)): + try: + encoding_offset = encoding_offsets[ + ord(bytearray([i]).decode(self.charset_encoding)) + ] + if encoding_offset != 0xFFFF: + encoding[i] = encoding_offset + except UnicodeDecodeError: + # character is not supported in selected encoding + pass + + return encoding diff --git a/venv/lib/python3.10/site-packages/PIL/PcxImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PcxImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..6b16d5385379ea3549f04459b862e60abe85bc93 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PcxImagePlugin.py @@ -0,0 +1,228 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PCX file handling +# +# This format was originally used by ZSoft's popular PaintBrush +# program for the IBM PC. It is also supported by many MS-DOS and +# Windows applications, including the Windows PaintBrush program in +# Windows 3. +# +# history: +# 1995-09-01 fl Created +# 1996-05-20 fl Fixed RGB support +# 1997-01-03 fl Fixed 2-bit and 4-bit support +# 1999-02-03 fl Fixed 8-bit support (broken in 1.0b1) +# 1999-02-07 fl Added write support +# 2002-06-09 fl Made 2-bit and 4-bit support a bit more robust +# 2002-07-30 fl Seek from to current position, not beginning of file +# 2003-06-03 fl Extract DPI settings (info["dpi"]) +# +# Copyright (c) 1997-2003 by Secret Labs AB. +# Copyright (c) 1995-2003 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import logging +from typing import IO + +from . import Image, ImageFile, ImagePalette +from ._binary import i16le as i16 +from ._binary import o8 +from ._binary import o16le as o16 + +logger = logging.getLogger(__name__) + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 2 and prefix[0] == 10 and prefix[1] in [0, 2, 3, 5] + + +## +# Image plugin for Paintbrush images. + + +class PcxImageFile(ImageFile.ImageFile): + format = "PCX" + format_description = "Paintbrush" + + def _open(self) -> None: + # header + assert self.fp is not None + + s = self.fp.read(68) + if not _accept(s): + msg = "not a PCX file" + raise SyntaxError(msg) + + # image + bbox = i16(s, 4), i16(s, 6), i16(s, 8) + 1, i16(s, 10) + 1 + if bbox[2] <= bbox[0] or bbox[3] <= bbox[1]: + msg = "bad PCX image size" + raise SyntaxError(msg) + logger.debug("BBox: %s %s %s %s", *bbox) + + offset = self.fp.tell() + 60 + + # format + version = s[1] + bits = s[3] + planes = s[65] + provided_stride = i16(s, 66) + logger.debug( + "PCX version %s, bits %s, planes %s, stride %s", + version, + bits, + planes, + provided_stride, + ) + + self.info["dpi"] = i16(s, 12), i16(s, 14) + + if bits == 1 and planes == 1: + mode = rawmode = "1" + + elif bits == 1 and planes in (2, 4): + mode = "P" + rawmode = f"P;{planes}L" + self.palette = ImagePalette.raw("RGB", s[16:64]) + + elif version == 5 and bits == 8 and planes == 1: + mode = rawmode = "L" + # FIXME: hey, this doesn't work with the incremental loader !!! + self.fp.seek(-769, io.SEEK_END) + s = self.fp.read(769) + if len(s) == 769 and s[0] == 12: + # check if the palette is linear grayscale + for i in range(256): + if s[i * 3 + 1 : i * 3 + 4] != o8(i) * 3: + mode = rawmode = "P" + break + if mode == "P": + self.palette = ImagePalette.raw("RGB", s[1:]) + + elif version == 5 and bits == 8 and planes == 3: + mode = "RGB" + rawmode = "RGB;L" + + else: + msg = "unknown PCX mode" + raise OSError(msg) + + self._mode = mode + self._size = bbox[2] - bbox[0], bbox[3] - bbox[1] + + # Don't trust the passed in stride. + # Calculate the approximate position for ourselves. + # CVE-2020-35653 + stride = (self._size[0] * bits + 7) // 8 + + # While the specification states that this must be even, + # not all images follow this + if provided_stride != stride: + stride += stride % 2 + + bbox = (0, 0) + self.size + logger.debug("size: %sx%s", *self.size) + + self.tile = [ImageFile._Tile("pcx", bbox, offset, (rawmode, planes * stride))] + + +# -------------------------------------------------------------------- +# save PCX files + + +SAVE = { + # mode: (version, bits, planes, raw mode) + "1": (2, 1, 1, "1"), + "L": (5, 8, 1, "L"), + "P": (5, 8, 1, "P"), + "RGB": (5, 8, 3, "RGB;L"), +} + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + try: + version, bits, planes, rawmode = SAVE[im.mode] + except KeyError as e: + msg = f"Cannot save {im.mode} images as PCX" + raise ValueError(msg) from e + + # bytes per plane + stride = (im.size[0] * bits + 7) // 8 + # stride should be even + stride += stride % 2 + # Stride needs to be kept in sync with the PcxEncode.c version. + # Ideally it should be passed in in the state, but the bytes value + # gets overwritten. + + logger.debug( + "PcxImagePlugin._save: xwidth: %d, bits: %d, stride: %d", + im.size[0], + bits, + stride, + ) + + # under windows, we could determine the current screen size with + # "Image.core.display_mode()[1]", but I think that's overkill... + + screen = im.size + + dpi = 100, 100 + + # PCX header + fp.write( + o8(10) + + o8(version) + + o8(1) + + o8(bits) + + o16(0) + + o16(0) + + o16(im.size[0] - 1) + + o16(im.size[1] - 1) + + o16(dpi[0]) + + o16(dpi[1]) + + b"\0" * 24 + + b"\xff" * 24 + + b"\0" + + o8(planes) + + o16(stride) + + o16(1) + + o16(screen[0]) + + o16(screen[1]) + + b"\0" * 54 + ) + + assert fp.tell() == 128 + + ImageFile._save( + im, fp, [ImageFile._Tile("pcx", (0, 0) + im.size, 0, (rawmode, bits * planes))] + ) + + if im.mode == "P": + # colour palette + fp.write(o8(12)) + palette = im.im.getpalette("RGB", "RGB") + palette += b"\x00" * (768 - len(palette)) + fp.write(palette) # 768 bytes + elif im.mode == "L": + # grayscale palette + fp.write(o8(12)) + for i in range(256): + fp.write(o8(i) * 3) + + +# -------------------------------------------------------------------- +# registry + + +Image.register_open(PcxImageFile.format, PcxImageFile, _accept) +Image.register_save(PcxImageFile.format, _save) + +Image.register_extension(PcxImageFile.format, ".pcx") + +Image.register_mime(PcxImageFile.format, "image/x-pcx") diff --git a/venv/lib/python3.10/site-packages/PIL/PdfImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PdfImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..5594c7e0f2b17b431c27e51ccf2bd66cf45c48b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PdfImagePlugin.py @@ -0,0 +1,311 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PDF (Acrobat) file handling +# +# History: +# 1996-07-16 fl Created +# 1997-01-18 fl Fixed header +# 2004-02-21 fl Fixes for 1/L/CMYK images, etc. +# 2004-02-24 fl Fixes for 1 and P images. +# +# Copyright (c) 1997-2004 by Secret Labs AB. All rights reserved. +# Copyright (c) 1996-1997 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + +## +# Image plugin for PDF images (output only). +## +from __future__ import annotations + +import io +import math +import os +import time +from typing import IO, Any + +from . import Image, ImageFile, ImageSequence, PdfParser, features + +# +# -------------------------------------------------------------------- + +# object ids: +# 1. catalogue +# 2. pages +# 3. image +# 4. page +# 5. page contents + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + _save(im, fp, filename, save_all=True) + + +## +# (Internal) Image save plugin for the PDF format. + + +def _write_image( + im: Image.Image, + filename: str | bytes, + existing_pdf: PdfParser.PdfParser, + image_refs: list[PdfParser.IndirectReference], +) -> tuple[PdfParser.IndirectReference, str]: + # FIXME: Should replace ASCIIHexDecode with RunLengthDecode + # (packbits) or LZWDecode (tiff/lzw compression). Note that + # PDF 1.2 also supports Flatedecode (zip compression). + + params = None + decode = None + + # + # Get image characteristics + + width, height = im.size + + dict_obj: dict[str, Any] = {"BitsPerComponent": 8} + if im.mode == "1": + if features.check("libtiff"): + decode_filter = "CCITTFaxDecode" + dict_obj["BitsPerComponent"] = 1 + params = PdfParser.PdfArray( + [ + PdfParser.PdfDict( + { + "K": -1, + "BlackIs1": True, + "Columns": width, + "Rows": height, + } + ) + ] + ) + else: + decode_filter = "DCTDecode" + dict_obj["ColorSpace"] = PdfParser.PdfName("DeviceGray") + procset = "ImageB" # grayscale + elif im.mode == "L": + decode_filter = "DCTDecode" + # params = f"<< /Predictor 15 /Columns {width-2} >>" + dict_obj["ColorSpace"] = PdfParser.PdfName("DeviceGray") + procset = "ImageB" # grayscale + elif im.mode == "LA": + decode_filter = "JPXDecode" + # params = f"<< /Predictor 15 /Columns {width-2} >>" + procset = "ImageB" # grayscale + dict_obj["SMaskInData"] = 1 + elif im.mode == "P": + decode_filter = "ASCIIHexDecode" + palette = im.getpalette() + assert palette is not None + dict_obj["ColorSpace"] = [ + PdfParser.PdfName("Indexed"), + PdfParser.PdfName("DeviceRGB"), + len(palette) // 3 - 1, + PdfParser.PdfBinary(palette), + ] + procset = "ImageI" # indexed color + + if "transparency" in im.info: + smask = im.convert("LA").getchannel("A") + smask.encoderinfo = {} + + image_ref = _write_image(smask, filename, existing_pdf, image_refs)[0] + dict_obj["SMask"] = image_ref + elif im.mode == "RGB": + decode_filter = "DCTDecode" + dict_obj["ColorSpace"] = PdfParser.PdfName("DeviceRGB") + procset = "ImageC" # color images + elif im.mode == "RGBA": + decode_filter = "JPXDecode" + procset = "ImageC" # color images + dict_obj["SMaskInData"] = 1 + elif im.mode == "CMYK": + decode_filter = "DCTDecode" + dict_obj["ColorSpace"] = PdfParser.PdfName("DeviceCMYK") + procset = "ImageC" # color images + decode = [1, 0, 1, 0, 1, 0, 1, 0] + else: + msg = f"cannot save mode {im.mode}" + raise ValueError(msg) + + # + # image + + op = io.BytesIO() + + if decode_filter == "ASCIIHexDecode": + ImageFile._save(im, op, [ImageFile._Tile("hex", (0, 0) + im.size, 0, im.mode)]) + elif decode_filter == "CCITTFaxDecode": + im.save( + op, + "TIFF", + compression="group4", + # use a single strip + strip_size=math.ceil(width / 8) * height, + ) + elif decode_filter == "DCTDecode": + Image.SAVE["JPEG"](im, op, filename) + elif decode_filter == "JPXDecode": + del dict_obj["BitsPerComponent"] + Image.SAVE["JPEG2000"](im, op, filename) + else: + msg = f"unsupported PDF filter ({decode_filter})" + raise ValueError(msg) + + stream = op.getvalue() + filter: PdfParser.PdfArray | PdfParser.PdfName + if decode_filter == "CCITTFaxDecode": + stream = stream[8:] + filter = PdfParser.PdfArray([PdfParser.PdfName(decode_filter)]) + else: + filter = PdfParser.PdfName(decode_filter) + + image_ref = image_refs.pop(0) + existing_pdf.write_obj( + image_ref, + stream=stream, + Type=PdfParser.PdfName("XObject"), + Subtype=PdfParser.PdfName("Image"), + Width=width, # * 72.0 / x_resolution, + Height=height, # * 72.0 / y_resolution, + Filter=filter, + Decode=decode, + DecodeParms=params, + **dict_obj, + ) + + return image_ref, procset + + +def _save( + im: Image.Image, fp: IO[bytes], filename: str | bytes, save_all: bool = False +) -> None: + is_appending = im.encoderinfo.get("append", False) + filename_str = filename.decode() if isinstance(filename, bytes) else filename + if is_appending: + existing_pdf = PdfParser.PdfParser(f=fp, filename=filename_str, mode="r+b") + else: + existing_pdf = PdfParser.PdfParser(f=fp, filename=filename_str, mode="w+b") + + dpi = im.encoderinfo.get("dpi") + if dpi: + x_resolution = dpi[0] + y_resolution = dpi[1] + else: + x_resolution = y_resolution = im.encoderinfo.get("resolution", 72.0) + + info = { + "title": ( + None if is_appending else os.path.splitext(os.path.basename(filename))[0] + ), + "author": None, + "subject": None, + "keywords": None, + "creator": None, + "producer": None, + "creationDate": None if is_appending else time.gmtime(), + "modDate": None if is_appending else time.gmtime(), + } + for k, default in info.items(): + v = im.encoderinfo.get(k) if k in im.encoderinfo else default + if v: + existing_pdf.info[k[0].upper() + k[1:]] = v + + # + # make sure image data is available + im.load() + + existing_pdf.start_writing() + existing_pdf.write_header() + existing_pdf.write_comment("created by Pillow PDF driver") + + # + # pages + ims = [im] + if save_all: + append_images = im.encoderinfo.get("append_images", []) + for append_im in append_images: + append_im.encoderinfo = im.encoderinfo.copy() + ims.append(append_im) + number_of_pages = 0 + image_refs = [] + page_refs = [] + contents_refs = [] + for im in ims: + im_number_of_pages = 1 + if save_all: + im_number_of_pages = getattr(im, "n_frames", 1) + number_of_pages += im_number_of_pages + for i in range(im_number_of_pages): + image_refs.append(existing_pdf.next_object_id(0)) + if im.mode == "P" and "transparency" in im.info: + image_refs.append(existing_pdf.next_object_id(0)) + + page_refs.append(existing_pdf.next_object_id(0)) + contents_refs.append(existing_pdf.next_object_id(0)) + existing_pdf.pages.append(page_refs[-1]) + + # + # catalog and list of pages + existing_pdf.write_catalog() + + page_number = 0 + for im_sequence in ims: + im_pages: ImageSequence.Iterator | list[Image.Image] = ( + ImageSequence.Iterator(im_sequence) if save_all else [im_sequence] + ) + for im in im_pages: + image_ref, procset = _write_image(im, filename, existing_pdf, image_refs) + + # + # page + + existing_pdf.write_page( + page_refs[page_number], + Resources=PdfParser.PdfDict( + ProcSet=[PdfParser.PdfName("PDF"), PdfParser.PdfName(procset)], + XObject=PdfParser.PdfDict(image=image_ref), + ), + MediaBox=[ + 0, + 0, + im.width * 72.0 / x_resolution, + im.height * 72.0 / y_resolution, + ], + Contents=contents_refs[page_number], + ) + + # + # page contents + + page_contents = b"q %f 0 0 %f 0 0 cm /image Do Q\n" % ( + im.width * 72.0 / x_resolution, + im.height * 72.0 / y_resolution, + ) + + existing_pdf.write_obj(contents_refs[page_number], stream=page_contents) + + page_number += 1 + + # + # trailer + existing_pdf.write_xref_and_trailer() + if hasattr(fp, "flush"): + fp.flush() + existing_pdf.close() + + +# +# -------------------------------------------------------------------- + + +Image.register_save("PDF", _save) +Image.register_save_all("PDF", _save_all) + +Image.register_extension("PDF", ".pdf") + +Image.register_mime("PDF", "application/pdf") diff --git a/venv/lib/python3.10/site-packages/PIL/PdfParser.py b/venv/lib/python3.10/site-packages/PIL/PdfParser.py new file mode 100644 index 0000000000000000000000000000000000000000..2c9031469ad5a4374e859e389df18c3a3248a2b4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PdfParser.py @@ -0,0 +1,1075 @@ +from __future__ import annotations + +import calendar +import codecs +import collections +import mmap +import os +import re +import time +import zlib +from typing import Any, NamedTuple + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import IO + + _DictBase = collections.UserDict[str | bytes, Any] +else: + _DictBase = collections.UserDict + + +# see 7.9.2.2 Text String Type on page 86 and D.3 PDFDocEncoding Character Set +# on page 656 +def encode_text(s: str) -> bytes: + return codecs.BOM_UTF16_BE + s.encode("utf_16_be") + + +PDFDocEncoding = { + 0x16: "\u0017", + 0x18: "\u02d8", + 0x19: "\u02c7", + 0x1A: "\u02c6", + 0x1B: "\u02d9", + 0x1C: "\u02dd", + 0x1D: "\u02db", + 0x1E: "\u02da", + 0x1F: "\u02dc", + 0x80: "\u2022", + 0x81: "\u2020", + 0x82: "\u2021", + 0x83: "\u2026", + 0x84: "\u2014", + 0x85: "\u2013", + 0x86: "\u0192", + 0x87: "\u2044", + 0x88: "\u2039", + 0x89: "\u203a", + 0x8A: "\u2212", + 0x8B: "\u2030", + 0x8C: "\u201e", + 0x8D: "\u201c", + 0x8E: "\u201d", + 0x8F: "\u2018", + 0x90: "\u2019", + 0x91: "\u201a", + 0x92: "\u2122", + 0x93: "\ufb01", + 0x94: "\ufb02", + 0x95: "\u0141", + 0x96: "\u0152", + 0x97: "\u0160", + 0x98: "\u0178", + 0x99: "\u017d", + 0x9A: "\u0131", + 0x9B: "\u0142", + 0x9C: "\u0153", + 0x9D: "\u0161", + 0x9E: "\u017e", + 0xA0: "\u20ac", +} + + +def decode_text(b: bytes) -> str: + if b[: len(codecs.BOM_UTF16_BE)] == codecs.BOM_UTF16_BE: + return b[len(codecs.BOM_UTF16_BE) :].decode("utf_16_be") + else: + return "".join(PDFDocEncoding.get(byte, chr(byte)) for byte in b) + + +class PdfFormatError(RuntimeError): + """An error that probably indicates a syntactic or semantic error in the + PDF file structure""" + + pass + + +def check_format_condition(condition: bool, error_message: str) -> None: + if not condition: + raise PdfFormatError(error_message) + + +class IndirectReferenceTuple(NamedTuple): + object_id: int + generation: int + + +class IndirectReference(IndirectReferenceTuple): + def __str__(self) -> str: + return f"{self.object_id} {self.generation} R" + + def __bytes__(self) -> bytes: + return self.__str__().encode("us-ascii") + + def __eq__(self, other: object) -> bool: + if self.__class__ is not other.__class__: + return False + assert isinstance(other, IndirectReference) + return other.object_id == self.object_id and other.generation == self.generation + + def __ne__(self, other: object) -> bool: + return not (self == other) + + def __hash__(self) -> int: + return hash((self.object_id, self.generation)) + + +class IndirectObjectDef(IndirectReference): + def __str__(self) -> str: + return f"{self.object_id} {self.generation} obj" + + +class XrefTable: + def __init__(self) -> None: + self.existing_entries: dict[int, tuple[int, int]] = ( + {} + ) # object ID => (offset, generation) + self.new_entries: dict[int, tuple[int, int]] = ( + {} + ) # object ID => (offset, generation) + self.deleted_entries = {0: 65536} # object ID => generation + self.reading_finished = False + + def __setitem__(self, key: int, value: tuple[int, int]) -> None: + if self.reading_finished: + self.new_entries[key] = value + else: + self.existing_entries[key] = value + if key in self.deleted_entries: + del self.deleted_entries[key] + + def __getitem__(self, key: int) -> tuple[int, int]: + try: + return self.new_entries[key] + except KeyError: + return self.existing_entries[key] + + def __delitem__(self, key: int) -> None: + if key in self.new_entries: + generation = self.new_entries[key][1] + 1 + del self.new_entries[key] + self.deleted_entries[key] = generation + elif key in self.existing_entries: + generation = self.existing_entries[key][1] + 1 + self.deleted_entries[key] = generation + elif key in self.deleted_entries: + generation = self.deleted_entries[key] + else: + msg = f"object ID {key} cannot be deleted because it doesn't exist" + raise IndexError(msg) + + def __contains__(self, key: int) -> bool: + return key in self.existing_entries or key in self.new_entries + + def __len__(self) -> int: + return len( + set(self.existing_entries.keys()) + | set(self.new_entries.keys()) + | set(self.deleted_entries.keys()) + ) + + def keys(self) -> set[int]: + return ( + set(self.existing_entries.keys()) - set(self.deleted_entries.keys()) + ) | set(self.new_entries.keys()) + + def write(self, f: IO[bytes]) -> int: + keys = sorted(set(self.new_entries.keys()) | set(self.deleted_entries.keys())) + deleted_keys = sorted(set(self.deleted_entries.keys())) + startxref = f.tell() + f.write(b"xref\n") + while keys: + # find a contiguous sequence of object IDs + prev: int | None = None + for index, key in enumerate(keys): + if prev is None or prev + 1 == key: + prev = key + else: + contiguous_keys = keys[:index] + keys = keys[index:] + break + else: + contiguous_keys = keys + keys = [] + f.write(b"%d %d\n" % (contiguous_keys[0], len(contiguous_keys))) + for object_id in contiguous_keys: + if object_id in self.new_entries: + f.write(b"%010d %05d n \n" % self.new_entries[object_id]) + else: + this_deleted_object_id = deleted_keys.pop(0) + check_format_condition( + object_id == this_deleted_object_id, + f"expected the next deleted object ID to be {object_id}, " + f"instead found {this_deleted_object_id}", + ) + try: + next_in_linked_list = deleted_keys[0] + except IndexError: + next_in_linked_list = 0 + f.write( + b"%010d %05d f \n" + % (next_in_linked_list, self.deleted_entries[object_id]) + ) + return startxref + + +class PdfName: + name: bytes + + def __init__(self, name: PdfName | bytes | str) -> None: + if isinstance(name, PdfName): + self.name = name.name + elif isinstance(name, bytes): + self.name = name + else: + self.name = name.encode("us-ascii") + + def name_as_str(self) -> str: + return self.name.decode("us-ascii") + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, PdfName) and other.name == self.name + ) or other == self.name + + def __hash__(self) -> int: + return hash(self.name) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({repr(self.name)})" + + @classmethod + def from_pdf_stream(cls, data: bytes) -> PdfName: + return cls(PdfParser.interpret_name(data)) + + allowed_chars = set(range(33, 127)) - {ord(c) for c in "#%/()<>[]{}"} + + def __bytes__(self) -> bytes: + result = bytearray(b"/") + for b in self.name: + if b in self.allowed_chars: + result.append(b) + else: + result.extend(b"#%02X" % b) + return bytes(result) + + +class PdfArray(list[Any]): + def __bytes__(self) -> bytes: + return b"[ " + b" ".join(pdf_repr(x) for x in self) + b" ]" + + +class PdfDict(_DictBase): + def __setattr__(self, key: str, value: Any) -> None: + if key == "data": + collections.UserDict.__setattr__(self, key, value) + else: + self[key.encode("us-ascii")] = value + + def __getattr__(self, key: str) -> str | time.struct_time: + try: + value = self[key.encode("us-ascii")] + except KeyError as e: + raise AttributeError(key) from e + if isinstance(value, bytes): + value = decode_text(value) + if key.endswith("Date"): + if value.startswith("D:"): + value = value[2:] + + relationship = "Z" + if len(value) > 17: + relationship = value[14] + offset = int(value[15:17]) * 60 + if len(value) > 20: + offset += int(value[18:20]) + + format = "%Y%m%d%H%M%S"[: len(value) - 2] + value = time.strptime(value[: len(format) + 2], format) + if relationship in ["+", "-"]: + offset *= 60 + if relationship == "+": + offset *= -1 + value = time.gmtime(calendar.timegm(value) + offset) + return value + + def __bytes__(self) -> bytes: + out = bytearray(b"<<") + for key, value in self.items(): + if value is None: + continue + value = pdf_repr(value) + out.extend(b"\n") + out.extend(bytes(PdfName(key))) + out.extend(b" ") + out.extend(value) + out.extend(b"\n>>") + return bytes(out) + + +class PdfBinary: + def __init__(self, data: list[int] | bytes) -> None: + self.data = data + + def __bytes__(self) -> bytes: + return b"<%s>" % b"".join(b"%02X" % b for b in self.data) + + +class PdfStream: + def __init__(self, dictionary: PdfDict, buf: bytes) -> None: + self.dictionary = dictionary + self.buf = buf + + def decode(self) -> bytes: + try: + filter = self.dictionary[b"Filter"] + except KeyError: + return self.buf + if filter == b"FlateDecode": + try: + expected_length = self.dictionary[b"DL"] + except KeyError: + expected_length = self.dictionary[b"Length"] + return zlib.decompress(self.buf, bufsize=int(expected_length)) + else: + msg = f"stream filter {repr(filter)} unknown/unsupported" + raise NotImplementedError(msg) + + +def pdf_repr(x: Any) -> bytes: + if x is True: + return b"true" + elif x is False: + return b"false" + elif x is None: + return b"null" + elif isinstance(x, (PdfName, PdfDict, PdfArray, PdfBinary)): + return bytes(x) + elif isinstance(x, (int, float)): + return str(x).encode("us-ascii") + elif isinstance(x, time.struct_time): + return b"(D:" + time.strftime("%Y%m%d%H%M%SZ", x).encode("us-ascii") + b")" + elif isinstance(x, dict): + return bytes(PdfDict(x)) + elif isinstance(x, list): + return bytes(PdfArray(x)) + elif isinstance(x, str): + return pdf_repr(encode_text(x)) + elif isinstance(x, bytes): + # XXX escape more chars? handle binary garbage + x = x.replace(b"\\", b"\\\\") + x = x.replace(b"(", b"\\(") + x = x.replace(b")", b"\\)") + return b"(" + x + b")" + else: + return bytes(x) + + +class PdfParser: + """Based on + https://www.adobe.com/content/dam/acom/en/devnet/acrobat/pdfs/PDF32000_2008.pdf + Supports PDF up to 1.4 + """ + + def __init__( + self, + filename: str | None = None, + f: IO[bytes] | None = None, + buf: bytes | bytearray | None = None, + start_offset: int = 0, + mode: str = "rb", + ) -> None: + if buf and f: + msg = "specify buf or f or filename, but not both buf and f" + raise RuntimeError(msg) + self.filename = filename + self.buf: bytes | bytearray | mmap.mmap | None = buf + self.f = f + self.start_offset = start_offset + self.should_close_buf = False + self.should_close_file = False + if filename is not None and f is None: + self.f = f = open(filename, mode) + self.should_close_file = True + if f is not None: + self.buf = self.get_buf_from_file(f) + self.should_close_buf = True + if not filename and hasattr(f, "name"): + self.filename = f.name + self.cached_objects: dict[IndirectReference, Any] = {} + self.root_ref: IndirectReference | None + self.info_ref: IndirectReference | None + self.pages_ref: IndirectReference | None + self.last_xref_section_offset: int | None + if self.buf: + self.read_pdf_info() + else: + self.file_size_total = self.file_size_this = 0 + self.root = PdfDict() + self.root_ref = None + self.info = PdfDict() + self.info_ref = None + self.page_tree_root = PdfDict() + self.pages: list[IndirectReference] = [] + self.orig_pages: list[IndirectReference] = [] + self.pages_ref = None + self.last_xref_section_offset = None + self.trailer_dict: dict[bytes, Any] = {} + self.xref_table = XrefTable() + self.xref_table.reading_finished = True + if f: + self.seek_end() + + def __enter__(self) -> PdfParser: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def start_writing(self) -> None: + self.close_buf() + self.seek_end() + + def close_buf(self) -> None: + if isinstance(self.buf, mmap.mmap): + self.buf.close() + self.buf = None + + def close(self) -> None: + if self.should_close_buf: + self.close_buf() + if self.f is not None and self.should_close_file: + self.f.close() + self.f = None + + def seek_end(self) -> None: + assert self.f is not None + self.f.seek(0, os.SEEK_END) + + def write_header(self) -> None: + assert self.f is not None + self.f.write(b"%PDF-1.4\n") + + def write_comment(self, s: str) -> None: + assert self.f is not None + self.f.write(f"% {s}\n".encode()) + + def write_catalog(self) -> IndirectReference: + assert self.f is not None + self.del_root() + self.root_ref = self.next_object_id(self.f.tell()) + self.pages_ref = self.next_object_id(0) + self.rewrite_pages() + self.write_obj(self.root_ref, Type=PdfName(b"Catalog"), Pages=self.pages_ref) + self.write_obj( + self.pages_ref, + Type=PdfName(b"Pages"), + Count=len(self.pages), + Kids=self.pages, + ) + return self.root_ref + + def rewrite_pages(self) -> None: + pages_tree_nodes_to_delete = [] + for i, page_ref in enumerate(self.orig_pages): + page_info = self.cached_objects[page_ref] + del self.xref_table[page_ref.object_id] + pages_tree_nodes_to_delete.append(page_info[PdfName(b"Parent")]) + if page_ref not in self.pages: + # the page has been deleted + continue + # make dict keys into strings for passing to write_page + stringified_page_info = {} + for key, value in page_info.items(): + # key should be a PdfName + stringified_page_info[key.name_as_str()] = value + stringified_page_info["Parent"] = self.pages_ref + new_page_ref = self.write_page(None, **stringified_page_info) + for j, cur_page_ref in enumerate(self.pages): + if cur_page_ref == page_ref: + # replace the page reference with the new one + self.pages[j] = new_page_ref + # delete redundant Pages tree nodes from xref table + for pages_tree_node_ref in pages_tree_nodes_to_delete: + while pages_tree_node_ref: + pages_tree_node = self.cached_objects[pages_tree_node_ref] + if pages_tree_node_ref.object_id in self.xref_table: + del self.xref_table[pages_tree_node_ref.object_id] + pages_tree_node_ref = pages_tree_node.get(b"Parent", None) + self.orig_pages = [] + + def write_xref_and_trailer( + self, new_root_ref: IndirectReference | None = None + ) -> None: + assert self.f is not None + if new_root_ref: + self.del_root() + self.root_ref = new_root_ref + if self.info: + self.info_ref = self.write_obj(None, self.info) + start_xref = self.xref_table.write(self.f) + num_entries = len(self.xref_table) + trailer_dict: dict[str | bytes, Any] = { + b"Root": self.root_ref, + b"Size": num_entries, + } + if self.last_xref_section_offset is not None: + trailer_dict[b"Prev"] = self.last_xref_section_offset + if self.info: + trailer_dict[b"Info"] = self.info_ref + self.last_xref_section_offset = start_xref + self.f.write( + b"trailer\n" + + bytes(PdfDict(trailer_dict)) + + b"\nstartxref\n%d\n%%%%EOF" % start_xref + ) + + def write_page( + self, ref: int | IndirectReference | None, *objs: Any, **dict_obj: Any + ) -> IndirectReference: + obj_ref = self.pages[ref] if isinstance(ref, int) else ref + if "Type" not in dict_obj: + dict_obj["Type"] = PdfName(b"Page") + if "Parent" not in dict_obj: + dict_obj["Parent"] = self.pages_ref + return self.write_obj(obj_ref, *objs, **dict_obj) + + def write_obj( + self, ref: IndirectReference | None, *objs: Any, **dict_obj: Any + ) -> IndirectReference: + assert self.f is not None + f = self.f + if ref is None: + ref = self.next_object_id(f.tell()) + else: + self.xref_table[ref.object_id] = (f.tell(), ref.generation) + f.write(bytes(IndirectObjectDef(*ref))) + stream = dict_obj.pop("stream", None) + if stream is not None: + dict_obj["Length"] = len(stream) + if dict_obj: + f.write(pdf_repr(dict_obj)) + for obj in objs: + f.write(pdf_repr(obj)) + if stream is not None: + f.write(b"stream\n") + f.write(stream) + f.write(b"\nendstream\n") + f.write(b"endobj\n") + return ref + + def del_root(self) -> None: + if self.root_ref is None: + return + del self.xref_table[self.root_ref.object_id] + del self.xref_table[self.root[b"Pages"].object_id] + + @staticmethod + def get_buf_from_file(f: IO[bytes]) -> bytes | mmap.mmap: + if hasattr(f, "getbuffer"): + return f.getbuffer() + elif hasattr(f, "getvalue"): + return f.getvalue() + else: + try: + return mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) + except ValueError: # cannot mmap an empty file + return b"" + + def read_pdf_info(self) -> None: + assert self.buf is not None + self.file_size_total = len(self.buf) + self.file_size_this = self.file_size_total - self.start_offset + self.read_trailer() + check_format_condition( + self.trailer_dict.get(b"Root") is not None, "Root is missing" + ) + self.root_ref = self.trailer_dict[b"Root"] + assert self.root_ref is not None + self.info_ref = self.trailer_dict.get(b"Info", None) + self.root = PdfDict(self.read_indirect(self.root_ref)) + if self.info_ref is None: + self.info = PdfDict() + else: + self.info = PdfDict(self.read_indirect(self.info_ref)) + check_format_condition(b"Type" in self.root, "/Type missing in Root") + check_format_condition( + self.root[b"Type"] == b"Catalog", "/Type in Root is not /Catalog" + ) + check_format_condition( + self.root.get(b"Pages") is not None, "/Pages missing in Root" + ) + check_format_condition( + isinstance(self.root[b"Pages"], IndirectReference), + "/Pages in Root is not an indirect reference", + ) + self.pages_ref = self.root[b"Pages"] + assert self.pages_ref is not None + self.page_tree_root = self.read_indirect(self.pages_ref) + self.pages = self.linearize_page_tree(self.page_tree_root) + # save the original list of page references + # in case the user modifies, adds or deletes some pages + # and we need to rewrite the pages and their list + self.orig_pages = self.pages[:] + + def next_object_id(self, offset: int | None = None) -> IndirectReference: + try: + # TODO: support reuse of deleted objects + reference = IndirectReference(max(self.xref_table.keys()) + 1, 0) + except ValueError: + reference = IndirectReference(1, 0) + if offset is not None: + self.xref_table[reference.object_id] = (offset, 0) + return reference + + delimiter = rb"[][()<>{}/%]" + delimiter_or_ws = rb"[][()<>{}/%\000\011\012\014\015\040]" + whitespace = rb"[\000\011\012\014\015\040]" + whitespace_or_hex = rb"[\000\011\012\014\015\0400-9a-fA-F]" + whitespace_optional = whitespace + b"*" + whitespace_mandatory = whitespace + b"+" + # No "\012" aka "\n" or "\015" aka "\r": + whitespace_optional_no_nl = rb"[\000\011\014\040]*" + newline_only = rb"[\r\n]+" + newline = whitespace_optional_no_nl + newline_only + whitespace_optional_no_nl + re_trailer_end = re.compile( + whitespace_mandatory + + rb"trailer" + + whitespace_optional + + rb"<<(.*>>)" + + newline + + rb"startxref" + + newline + + rb"([0-9]+)" + + newline + + rb"%%EOF" + + whitespace_optional + + rb"$", + re.DOTALL, + ) + re_trailer_prev = re.compile( + whitespace_optional + + rb"trailer" + + whitespace_optional + + rb"<<(.*?>>)" + + newline + + rb"startxref" + + newline + + rb"([0-9]+)" + + newline + + rb"%%EOF" + + whitespace_optional, + re.DOTALL, + ) + + def read_trailer(self) -> None: + assert self.buf is not None + search_start_offset = len(self.buf) - 16384 + if search_start_offset < self.start_offset: + search_start_offset = self.start_offset + m = self.re_trailer_end.search(self.buf, search_start_offset) + check_format_condition(m is not None, "trailer end not found") + # make sure we found the LAST trailer + last_match = m + while m: + last_match = m + m = self.re_trailer_end.search(self.buf, m.start() + 16) + if not m: + m = last_match + assert m is not None + trailer_data = m.group(1) + self.last_xref_section_offset = int(m.group(2)) + self.trailer_dict = self.interpret_trailer(trailer_data) + self.xref_table = XrefTable() + self.read_xref_table(xref_section_offset=self.last_xref_section_offset) + if b"Prev" in self.trailer_dict: + self.read_prev_trailer(self.trailer_dict[b"Prev"]) + + def read_prev_trailer(self, xref_section_offset: int) -> None: + assert self.buf is not None + trailer_offset = self.read_xref_table(xref_section_offset=xref_section_offset) + m = self.re_trailer_prev.search( + self.buf[trailer_offset : trailer_offset + 16384] + ) + check_format_condition(m is not None, "previous trailer not found") + assert m is not None + trailer_data = m.group(1) + check_format_condition( + int(m.group(2)) == xref_section_offset, + "xref section offset in previous trailer doesn't match what was expected", + ) + trailer_dict = self.interpret_trailer(trailer_data) + if b"Prev" in trailer_dict: + self.read_prev_trailer(trailer_dict[b"Prev"]) + + re_whitespace_optional = re.compile(whitespace_optional) + re_name = re.compile( + whitespace_optional + + rb"/([!-$&'*-.0-;=?-Z\\^-z|~]+)(?=" + + delimiter_or_ws + + rb")" + ) + re_dict_start = re.compile(whitespace_optional + rb"<<") + re_dict_end = re.compile(whitespace_optional + rb">>" + whitespace_optional) + + @classmethod + def interpret_trailer(cls, trailer_data: bytes) -> dict[bytes, Any]: + trailer = {} + offset = 0 + while True: + m = cls.re_name.match(trailer_data, offset) + if not m: + m = cls.re_dict_end.match(trailer_data, offset) + check_format_condition( + m is not None and m.end() == len(trailer_data), + "name not found in trailer, remaining data: " + + repr(trailer_data[offset:]), + ) + break + key = cls.interpret_name(m.group(1)) + assert isinstance(key, bytes) + value, value_offset = cls.get_value(trailer_data, m.end()) + trailer[key] = value + if value_offset is None: + break + offset = value_offset + check_format_condition( + b"Size" in trailer and isinstance(trailer[b"Size"], int), + "/Size not in trailer or not an integer", + ) + check_format_condition( + b"Root" in trailer and isinstance(trailer[b"Root"], IndirectReference), + "/Root not in trailer or not an indirect reference", + ) + return trailer + + re_hashes_in_name = re.compile(rb"([^#]*)(#([0-9a-fA-F]{2}))?") + + @classmethod + def interpret_name(cls, raw: bytes, as_text: bool = False) -> str | bytes: + name = b"" + for m in cls.re_hashes_in_name.finditer(raw): + if m.group(3): + name += m.group(1) + bytearray.fromhex(m.group(3).decode("us-ascii")) + else: + name += m.group(1) + if as_text: + return name.decode("utf-8") + else: + return bytes(name) + + re_null = re.compile(whitespace_optional + rb"null(?=" + delimiter_or_ws + rb")") + re_true = re.compile(whitespace_optional + rb"true(?=" + delimiter_or_ws + rb")") + re_false = re.compile(whitespace_optional + rb"false(?=" + delimiter_or_ws + rb")") + re_int = re.compile( + whitespace_optional + rb"([-+]?[0-9]+)(?=" + delimiter_or_ws + rb")" + ) + re_real = re.compile( + whitespace_optional + + rb"([-+]?([0-9]+\.[0-9]*|[0-9]*\.[0-9]+))(?=" + + delimiter_or_ws + + rb")" + ) + re_array_start = re.compile(whitespace_optional + rb"\[") + re_array_end = re.compile(whitespace_optional + rb"]") + re_string_hex = re.compile( + whitespace_optional + rb"<(" + whitespace_or_hex + rb"*)>" + ) + re_string_lit = re.compile(whitespace_optional + rb"\(") + re_indirect_reference = re.compile( + whitespace_optional + + rb"([-+]?[0-9]+)" + + whitespace_mandatory + + rb"([-+]?[0-9]+)" + + whitespace_mandatory + + rb"R(?=" + + delimiter_or_ws + + rb")" + ) + re_indirect_def_start = re.compile( + whitespace_optional + + rb"([-+]?[0-9]+)" + + whitespace_mandatory + + rb"([-+]?[0-9]+)" + + whitespace_mandatory + + rb"obj(?=" + + delimiter_or_ws + + rb")" + ) + re_indirect_def_end = re.compile( + whitespace_optional + rb"endobj(?=" + delimiter_or_ws + rb")" + ) + re_comment = re.compile( + rb"(" + whitespace_optional + rb"%[^\r\n]*" + newline + rb")*" + ) + re_stream_start = re.compile(whitespace_optional + rb"stream\r?\n") + re_stream_end = re.compile( + whitespace_optional + rb"endstream(?=" + delimiter_or_ws + rb")" + ) + + @classmethod + def get_value( + cls, + data: bytes | bytearray | mmap.mmap, + offset: int, + expect_indirect: IndirectReference | None = None, + max_nesting: int = -1, + ) -> tuple[Any, int | None]: + if max_nesting == 0: + return None, None + m = cls.re_comment.match(data, offset) + if m: + offset = m.end() + m = cls.re_indirect_def_start.match(data, offset) + if m: + check_format_condition( + int(m.group(1)) > 0, + "indirect object definition: object ID must be greater than 0", + ) + check_format_condition( + int(m.group(2)) >= 0, + "indirect object definition: generation must be non-negative", + ) + check_format_condition( + expect_indirect is None + or expect_indirect + == IndirectReference(int(m.group(1)), int(m.group(2))), + "indirect object definition different than expected", + ) + object, object_offset = cls.get_value( + data, m.end(), max_nesting=max_nesting - 1 + ) + if object_offset is None: + return object, None + m = cls.re_indirect_def_end.match(data, object_offset) + check_format_condition( + m is not None, "indirect object definition end not found" + ) + assert m is not None + return object, m.end() + check_format_condition( + not expect_indirect, "indirect object definition not found" + ) + m = cls.re_indirect_reference.match(data, offset) + if m: + check_format_condition( + int(m.group(1)) > 0, + "indirect object reference: object ID must be greater than 0", + ) + check_format_condition( + int(m.group(2)) >= 0, + "indirect object reference: generation must be non-negative", + ) + return IndirectReference(int(m.group(1)), int(m.group(2))), m.end() + m = cls.re_dict_start.match(data, offset) + if m: + offset = m.end() + result: dict[Any, Any] = {} + m = cls.re_dict_end.match(data, offset) + current_offset: int | None = offset + while not m: + assert current_offset is not None + key, current_offset = cls.get_value( + data, current_offset, max_nesting=max_nesting - 1 + ) + if current_offset is None: + return result, None + value, current_offset = cls.get_value( + data, current_offset, max_nesting=max_nesting - 1 + ) + result[key] = value + if current_offset is None: + return result, None + m = cls.re_dict_end.match(data, current_offset) + current_offset = m.end() + m = cls.re_stream_start.match(data, current_offset) + if m: + stream_len = result.get(b"Length") + if stream_len is None or not isinstance(stream_len, int): + msg = f"bad or missing Length in stream dict ({stream_len})" + raise PdfFormatError(msg) + stream_data = data[m.end() : m.end() + stream_len] + m = cls.re_stream_end.match(data, m.end() + stream_len) + check_format_condition(m is not None, "stream end not found") + assert m is not None + current_offset = m.end() + return PdfStream(PdfDict(result), stream_data), current_offset + return PdfDict(result), current_offset + m = cls.re_array_start.match(data, offset) + if m: + offset = m.end() + results = [] + m = cls.re_array_end.match(data, offset) + current_offset = offset + while not m: + assert current_offset is not None + value, current_offset = cls.get_value( + data, current_offset, max_nesting=max_nesting - 1 + ) + results.append(value) + if current_offset is None: + return results, None + m = cls.re_array_end.match(data, current_offset) + return results, m.end() + m = cls.re_null.match(data, offset) + if m: + return None, m.end() + m = cls.re_true.match(data, offset) + if m: + return True, m.end() + m = cls.re_false.match(data, offset) + if m: + return False, m.end() + m = cls.re_name.match(data, offset) + if m: + return PdfName(cls.interpret_name(m.group(1))), m.end() + m = cls.re_int.match(data, offset) + if m: + return int(m.group(1)), m.end() + m = cls.re_real.match(data, offset) + if m: + # XXX Decimal instead of float??? + return float(m.group(1)), m.end() + m = cls.re_string_hex.match(data, offset) + if m: + # filter out whitespace + hex_string = bytearray( + b for b in m.group(1) if b in b"0123456789abcdefABCDEF" + ) + if len(hex_string) % 2 == 1: + # append a 0 if the length is not even - yes, at the end + hex_string.append(ord(b"0")) + return bytearray.fromhex(hex_string.decode("us-ascii")), m.end() + m = cls.re_string_lit.match(data, offset) + if m: + return cls.get_literal_string(data, m.end()) + # return None, offset # fallback (only for debugging) + msg = f"unrecognized object: {repr(data[offset : offset + 32])}" + raise PdfFormatError(msg) + + re_lit_str_token = re.compile( + rb"(\\[nrtbf()\\])|(\\[0-9]{1,3})|(\\(\r\n|\r|\n))|(\r\n|\r|\n)|(\()|(\))" + ) + escaped_chars = { + b"n": b"\n", + b"r": b"\r", + b"t": b"\t", + b"b": b"\b", + b"f": b"\f", + b"(": b"(", + b")": b")", + b"\\": b"\\", + ord(b"n"): b"\n", + ord(b"r"): b"\r", + ord(b"t"): b"\t", + ord(b"b"): b"\b", + ord(b"f"): b"\f", + ord(b"("): b"(", + ord(b")"): b")", + ord(b"\\"): b"\\", + } + + @classmethod + def get_literal_string( + cls, data: bytes | bytearray | mmap.mmap, offset: int + ) -> tuple[bytes, int]: + nesting_depth = 0 + result = bytearray() + for m in cls.re_lit_str_token.finditer(data, offset): + result.extend(data[offset : m.start()]) + if m.group(1): + result.extend(cls.escaped_chars[m.group(1)[1]]) + elif m.group(2): + result.append(int(m.group(2)[1:], 8)) + elif m.group(3): + pass + elif m.group(5): + result.extend(b"\n") + elif m.group(6): + result.extend(b"(") + nesting_depth += 1 + elif m.group(7): + if nesting_depth == 0: + return bytes(result), m.end() + result.extend(b")") + nesting_depth -= 1 + offset = m.end() + msg = "unfinished literal string" + raise PdfFormatError(msg) + + re_xref_section_start = re.compile(whitespace_optional + rb"xref" + newline) + re_xref_subsection_start = re.compile( + whitespace_optional + + rb"([0-9]+)" + + whitespace_mandatory + + rb"([0-9]+)" + + whitespace_optional + + newline_only + ) + re_xref_entry = re.compile(rb"([0-9]{10}) ([0-9]{5}) ([fn])( \r| \n|\r\n)") + + def read_xref_table(self, xref_section_offset: int) -> int: + assert self.buf is not None + subsection_found = False + m = self.re_xref_section_start.match( + self.buf, xref_section_offset + self.start_offset + ) + check_format_condition(m is not None, "xref section start not found") + assert m is not None + offset = m.end() + while True: + m = self.re_xref_subsection_start.match(self.buf, offset) + if not m: + check_format_condition( + subsection_found, "xref subsection start not found" + ) + break + subsection_found = True + offset = m.end() + first_object = int(m.group(1)) + num_objects = int(m.group(2)) + for i in range(first_object, first_object + num_objects): + m = self.re_xref_entry.match(self.buf, offset) + check_format_condition(m is not None, "xref entry not found") + assert m is not None + offset = m.end() + is_free = m.group(3) == b"f" + if not is_free: + generation = int(m.group(2)) + new_entry = (int(m.group(1)), generation) + if i not in self.xref_table: + self.xref_table[i] = new_entry + return offset + + def read_indirect(self, ref: IndirectReference, max_nesting: int = -1) -> Any: + offset, generation = self.xref_table[ref[0]] + check_format_condition( + generation == ref[1], + f"expected to find generation {ref[1]} for object ID {ref[0]} in xref " + f"table, instead found generation {generation} at offset {offset}", + ) + assert self.buf is not None + value = self.get_value( + self.buf, + offset + self.start_offset, + expect_indirect=IndirectReference(*ref), + max_nesting=max_nesting, + )[0] + self.cached_objects[ref] = value + return value + + def linearize_page_tree( + self, node: PdfDict | None = None + ) -> list[IndirectReference]: + page_node = node if node is not None else self.page_tree_root + check_format_condition( + page_node[b"Type"] == b"Pages", "/Type of page tree node is not /Pages" + ) + pages = [] + for kid in page_node[b"Kids"]: + kid_object = self.read_indirect(kid) + if kid_object[b"Type"] == b"Page": + pages.append(kid) + else: + pages.extend(self.linearize_page_tree(node=kid_object)) + return pages diff --git a/venv/lib/python3.10/site-packages/PIL/PixarImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PixarImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..d2b6d0a97e4bd230134d4741fc997baca5b4507f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PixarImagePlugin.py @@ -0,0 +1,72 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PIXAR raster support for PIL +# +# history: +# 97-01-29 fl Created +# +# notes: +# This is incomplete; it is based on a few samples created with +# Photoshop 2.5 and 3.0, and a summary description provided by +# Greg Coats . Hopefully, "L" and +# "RGBA" support will be added in future versions. +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1997. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image, ImageFile +from ._binary import i16le as i16 + +# +# helpers + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"\200\350\000\000") + + +## +# Image plugin for PIXAR raster images. + + +class PixarImageFile(ImageFile.ImageFile): + format = "PIXAR" + format_description = "PIXAR raster image" + + def _open(self) -> None: + # assuming a 4-byte magic label + assert self.fp is not None + + s = self.fp.read(4) + if not _accept(s): + msg = "not a PIXAR file" + raise SyntaxError(msg) + + # read rest of header + s = s + self.fp.read(508) + + self._size = i16(s, 418), i16(s, 416) + + # get channel/depth descriptions + mode = i16(s, 424), i16(s, 426) + + if mode == (14, 2): + self._mode = "RGB" + # FIXME: to be continued... + + # create tile descriptor (assuming "dumped") + self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 1024, self.mode)] + + +# +# -------------------------------------------------------------------- + +Image.register_open(PixarImageFile.format, PixarImageFile, _accept) + +Image.register_extension(PixarImageFile.format, ".pxr") diff --git a/venv/lib/python3.10/site-packages/PIL/PngImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PngImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..d0f22f812e50ab5a5d5722444bbeabd08eb5e28f --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PngImagePlugin.py @@ -0,0 +1,1553 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PNG support code +# +# See "PNG (Portable Network Graphics) Specification, version 1.0; +# W3C Recommendation", 1996-10-01, Thomas Boutell (ed.). +# +# history: +# 1996-05-06 fl Created (couldn't resist it) +# 1996-12-14 fl Upgraded, added read and verify support (0.2) +# 1996-12-15 fl Separate PNG stream parser +# 1996-12-29 fl Added write support, added getchunks +# 1996-12-30 fl Eliminated circular references in decoder (0.3) +# 1998-07-12 fl Read/write 16-bit images as mode I (0.4) +# 2001-02-08 fl Added transparency support (from Zircon) (0.5) +# 2001-04-16 fl Don't close data source in "open" method (0.6) +# 2004-02-24 fl Don't even pretend to support interlaced files (0.7) +# 2004-08-31 fl Do basic sanity check on chunk identifiers (0.8) +# 2004-09-20 fl Added PngInfo chunk container +# 2004-12-18 fl Added DPI read support (based on code by Niki Spahiev) +# 2008-08-13 fl Added tRNS support for RGB images +# 2009-03-06 fl Support for preserving ICC profiles (by Florian Hoech) +# 2009-03-08 fl Added zTXT support (from Lowell Alleman) +# 2009-03-29 fl Read interlaced PNG files (from Conrado Porto Lopes Gouvua) +# +# Copyright (c) 1997-2009 by Secret Labs AB +# Copyright (c) 1996 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import itertools +import logging +import re +import struct +import warnings +import zlib +from enum import IntEnum +from typing import IO, NamedTuple, cast + +from . import Image, ImageChops, ImageFile, ImagePalette, ImageSequence +from ._binary import i16be as i16 +from ._binary import i32be as i32 +from ._binary import o8 +from ._binary import o16be as o16 +from ._binary import o32be as o32 +from ._deprecate import deprecate +from ._util import DeferredError + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Callable + from typing import Any, NoReturn + + from . import _imaging + +logger = logging.getLogger(__name__) + +is_cid = re.compile(rb"\w\w\w\w").match + + +_MAGIC = b"\211PNG\r\n\032\n" + + +_MODES = { + # supported bits/color combinations, and corresponding modes/rawmodes + # Grayscale + (1, 0): ("1", "1"), + (2, 0): ("L", "L;2"), + (4, 0): ("L", "L;4"), + (8, 0): ("L", "L"), + (16, 0): ("I;16", "I;16B"), + # Truecolour + (8, 2): ("RGB", "RGB"), + (16, 2): ("RGB", "RGB;16B"), + # Indexed-colour + (1, 3): ("P", "P;1"), + (2, 3): ("P", "P;2"), + (4, 3): ("P", "P;4"), + (8, 3): ("P", "P"), + # Grayscale with alpha + (8, 4): ("LA", "LA"), + (16, 4): ("RGBA", "LA;16B"), # LA;16B->LA not yet available + # Truecolour with alpha + (8, 6): ("RGBA", "RGBA"), + (16, 6): ("RGBA", "RGBA;16B"), +} + + +_simple_palette = re.compile(b"^\xff*\x00\xff*$") + +MAX_TEXT_CHUNK = ImageFile.SAFEBLOCK +""" +Maximum decompressed size for a iTXt or zTXt chunk. +Eliminates decompression bombs where compressed chunks can expand 1000x. +See :ref:`Text in PNG File Format`. +""" +MAX_TEXT_MEMORY = 64 * MAX_TEXT_CHUNK +""" +Set the maximum total text chunk size. +See :ref:`Text in PNG File Format`. +""" + + +# APNG frame disposal modes +class Disposal(IntEnum): + OP_NONE = 0 + """ + No disposal is done on this frame before rendering the next frame. + See :ref:`Saving APNG sequences`. + """ + OP_BACKGROUND = 1 + """ + This frame’s modified region is cleared to fully transparent black before rendering + the next frame. + See :ref:`Saving APNG sequences`. + """ + OP_PREVIOUS = 2 + """ + This frame’s modified region is reverted to the previous frame’s contents before + rendering the next frame. + See :ref:`Saving APNG sequences`. + """ + + +# APNG frame blend modes +class Blend(IntEnum): + OP_SOURCE = 0 + """ + All color components of this frame, including alpha, overwrite the previous output + image contents. + See :ref:`Saving APNG sequences`. + """ + OP_OVER = 1 + """ + This frame should be alpha composited with the previous output image contents. + See :ref:`Saving APNG sequences`. + """ + + +def _safe_zlib_decompress(s: bytes) -> bytes: + dobj = zlib.decompressobj() + plaintext = dobj.decompress(s, MAX_TEXT_CHUNK) + if dobj.unconsumed_tail: + msg = "Decompressed data too large for PngImagePlugin.MAX_TEXT_CHUNK" + raise ValueError(msg) + return plaintext + + +def _crc32(data: bytes, seed: int = 0) -> int: + return zlib.crc32(data, seed) & 0xFFFFFFFF + + +# -------------------------------------------------------------------- +# Support classes. Suitable for PNG and related formats like MNG etc. + + +class ChunkStream: + def __init__(self, fp: IO[bytes]) -> None: + self.fp: IO[bytes] | None = fp + self.queue: list[tuple[bytes, int, int]] | None = [] + + def read(self) -> tuple[bytes, int, int]: + """Fetch a new chunk. Returns header information.""" + cid = None + + assert self.fp is not None + if self.queue: + cid, pos, length = self.queue.pop() + self.fp.seek(pos) + else: + s = self.fp.read(8) + cid = s[4:] + pos = self.fp.tell() + length = i32(s) + + if not is_cid(cid): + if not ImageFile.LOAD_TRUNCATED_IMAGES: + msg = f"broken PNG file (chunk {repr(cid)})" + raise SyntaxError(msg) + + return cid, pos, length + + def __enter__(self) -> ChunkStream: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def close(self) -> None: + self.queue = self.fp = None + + def push(self, cid: bytes, pos: int, length: int) -> None: + assert self.queue is not None + self.queue.append((cid, pos, length)) + + def call(self, cid: bytes, pos: int, length: int) -> bytes: + """Call the appropriate chunk handler""" + + logger.debug("STREAM %r %s %s", cid, pos, length) + return getattr(self, f"chunk_{cid.decode('ascii')}")(pos, length) + + def crc(self, cid: bytes, data: bytes) -> None: + """Read and verify checksum""" + + # Skip CRC checks for ancillary chunks if allowed to load truncated + # images + # 5th byte of first char is 1 [specs, section 5.4] + if ImageFile.LOAD_TRUNCATED_IMAGES and (cid[0] >> 5 & 1): + self.crc_skip(cid, data) + return + + assert self.fp is not None + try: + crc1 = _crc32(data, _crc32(cid)) + crc2 = i32(self.fp.read(4)) + if crc1 != crc2: + msg = f"broken PNG file (bad header checksum in {repr(cid)})" + raise SyntaxError(msg) + except struct.error as e: + msg = f"broken PNG file (incomplete checksum in {repr(cid)})" + raise SyntaxError(msg) from e + + def crc_skip(self, cid: bytes, data: bytes) -> None: + """Read checksum""" + + assert self.fp is not None + self.fp.read(4) + + def verify(self, endchunk: bytes = b"IEND") -> list[bytes]: + # Simple approach; just calculate checksum for all remaining + # blocks. Must be called directly after open. + + cids = [] + + assert self.fp is not None + while True: + try: + cid, pos, length = self.read() + except struct.error as e: + msg = "truncated PNG file" + raise OSError(msg) from e + + if cid == endchunk: + break + self.crc(cid, ImageFile._safe_read(self.fp, length)) + cids.append(cid) + + return cids + + +class iTXt(str): + """ + Subclass of string to allow iTXt chunks to look like strings while + keeping their extra information + + """ + + lang: str | bytes | None + tkey: str | bytes | None + + @staticmethod + def __new__( + cls, text: str, lang: str | None = None, tkey: str | None = None + ) -> iTXt: + """ + :param cls: the class to use when creating the instance + :param text: value for this key + :param lang: language code + :param tkey: UTF-8 version of the key name + """ + + self = str.__new__(cls, text) + self.lang = lang + self.tkey = tkey + return self + + +class PngInfo: + """ + PNG chunk container (for use with save(pnginfo=)) + + """ + + def __init__(self) -> None: + self.chunks: list[tuple[bytes, bytes, bool]] = [] + + def add(self, cid: bytes, data: bytes, after_idat: bool = False) -> None: + """Appends an arbitrary chunk. Use with caution. + + :param cid: a byte string, 4 bytes long. + :param data: a byte string of the encoded data + :param after_idat: for use with private chunks. Whether the chunk + should be written after IDAT + + """ + + self.chunks.append((cid, data, after_idat)) + + def add_itxt( + self, + key: str | bytes, + value: str | bytes, + lang: str | bytes = "", + tkey: str | bytes = "", + zip: bool = False, + ) -> None: + """Appends an iTXt chunk. + + :param key: latin-1 encodable text key name + :param value: value for this key + :param lang: language code + :param tkey: UTF-8 version of the key name + :param zip: compression flag + + """ + + if not isinstance(key, bytes): + key = key.encode("latin-1", "strict") + if not isinstance(value, bytes): + value = value.encode("utf-8", "strict") + if not isinstance(lang, bytes): + lang = lang.encode("utf-8", "strict") + if not isinstance(tkey, bytes): + tkey = tkey.encode("utf-8", "strict") + + if zip: + self.add( + b"iTXt", + key + b"\0\x01\0" + lang + b"\0" + tkey + b"\0" + zlib.compress(value), + ) + else: + self.add(b"iTXt", key + b"\0\0\0" + lang + b"\0" + tkey + b"\0" + value) + + def add_text( + self, key: str | bytes, value: str | bytes | iTXt, zip: bool = False + ) -> None: + """Appends a text chunk. + + :param key: latin-1 encodable text key name + :param value: value for this key, text or an + :py:class:`PIL.PngImagePlugin.iTXt` instance + :param zip: compression flag + + """ + if isinstance(value, iTXt): + return self.add_itxt( + key, + value, + value.lang if value.lang is not None else b"", + value.tkey if value.tkey is not None else b"", + zip=zip, + ) + + # The tEXt chunk stores latin-1 text + if not isinstance(value, bytes): + try: + value = value.encode("latin-1", "strict") + except UnicodeError: + return self.add_itxt(key, value, zip=zip) + + if not isinstance(key, bytes): + key = key.encode("latin-1", "strict") + + if zip: + self.add(b"zTXt", key + b"\0\0" + zlib.compress(value)) + else: + self.add(b"tEXt", key + b"\0" + value) + + +# -------------------------------------------------------------------- +# PNG image stream (IHDR/IEND) + + +class _RewindState(NamedTuple): + info: dict[str | tuple[int, int], Any] + tile: list[ImageFile._Tile] + seq_num: int | None + + +class PngStream(ChunkStream): + def __init__(self, fp: IO[bytes]) -> None: + super().__init__(fp) + + # local copies of Image attributes + self.im_info: dict[str | tuple[int, int], Any] = {} + self.im_text: dict[str, str | iTXt] = {} + self.im_size = (0, 0) + self.im_mode = "" + self.im_tile: list[ImageFile._Tile] = [] + self.im_palette: tuple[str, bytes] | None = None + self.im_custom_mimetype: str | None = None + self.im_n_frames: int | None = None + self._seq_num: int | None = None + self.rewind_state = _RewindState({}, [], None) + + self.text_memory = 0 + + def check_text_memory(self, chunklen: int) -> None: + self.text_memory += chunklen + if self.text_memory > MAX_TEXT_MEMORY: + msg = ( + "Too much memory used in text chunks: " + f"{self.text_memory}>MAX_TEXT_MEMORY" + ) + raise ValueError(msg) + + def save_rewind(self) -> None: + self.rewind_state = _RewindState( + self.im_info.copy(), + self.im_tile, + self._seq_num, + ) + + def rewind(self) -> None: + self.im_info = self.rewind_state.info.copy() + self.im_tile = self.rewind_state.tile + self._seq_num = self.rewind_state.seq_num + + def chunk_iCCP(self, pos: int, length: int) -> bytes: + # ICC profile + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + # according to PNG spec, the iCCP chunk contains: + # Profile name 1-79 bytes (character string) + # Null separator 1 byte (null character) + # Compression method 1 byte (0) + # Compressed profile n bytes (zlib with deflate compression) + i = s.find(b"\0") + logger.debug("iCCP profile name %r", s[:i]) + comp_method = s[i + 1] + logger.debug("Compression method %s", comp_method) + if comp_method != 0: + msg = f"Unknown compression method {comp_method} in iCCP chunk" + raise SyntaxError(msg) + try: + icc_profile = _safe_zlib_decompress(s[i + 2 :]) + except ValueError: + if ImageFile.LOAD_TRUNCATED_IMAGES: + icc_profile = None + else: + raise + except zlib.error: + icc_profile = None # FIXME + self.im_info["icc_profile"] = icc_profile + return s + + def chunk_IHDR(self, pos: int, length: int) -> bytes: + # image header + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if length < 13: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + msg = "Truncated IHDR chunk" + raise ValueError(msg) + self.im_size = i32(s, 0), i32(s, 4) + try: + self.im_mode, self.im_rawmode = _MODES[(s[8], s[9])] + except Exception: + pass + if s[12]: + self.im_info["interlace"] = 1 + if s[11]: + msg = "unknown filter category" + raise SyntaxError(msg) + return s + + def chunk_IDAT(self, pos: int, length: int) -> NoReturn: + # image data + if "bbox" in self.im_info: + tile = [ImageFile._Tile("zip", self.im_info["bbox"], pos, self.im_rawmode)] + else: + if self.im_n_frames is not None: + self.im_info["default_image"] = True + tile = [ImageFile._Tile("zip", (0, 0) + self.im_size, pos, self.im_rawmode)] + self.im_tile = tile + self.im_idat = length + msg = "image data found" + raise EOFError(msg) + + def chunk_IEND(self, pos: int, length: int) -> NoReturn: + msg = "end of PNG image" + raise EOFError(msg) + + def chunk_PLTE(self, pos: int, length: int) -> bytes: + # palette + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if self.im_mode == "P": + self.im_palette = "RGB", s + return s + + def chunk_tRNS(self, pos: int, length: int) -> bytes: + # transparency + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if self.im_mode == "P": + if _simple_palette.match(s): + # tRNS contains only one full-transparent entry, + # other entries are full opaque + i = s.find(b"\0") + if i >= 0: + self.im_info["transparency"] = i + else: + # otherwise, we have a byte string with one alpha value + # for each palette entry + self.im_info["transparency"] = s + elif self.im_mode in ("1", "L", "I;16"): + self.im_info["transparency"] = i16(s) + elif self.im_mode == "RGB": + self.im_info["transparency"] = i16(s), i16(s, 2), i16(s, 4) + return s + + def chunk_gAMA(self, pos: int, length: int) -> bytes: + # gamma setting + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + self.im_info["gamma"] = i32(s) / 100000.0 + return s + + def chunk_cHRM(self, pos: int, length: int) -> bytes: + # chromaticity, 8 unsigned ints, actual value is scaled by 100,000 + # WP x,y, Red x,y, Green x,y Blue x,y + + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + raw_vals = struct.unpack(f">{len(s) // 4}I", s) + self.im_info["chromaticity"] = tuple(elt / 100000.0 for elt in raw_vals) + return s + + def chunk_sRGB(self, pos: int, length: int) -> bytes: + # srgb rendering intent, 1 byte + # 0 perceptual + # 1 relative colorimetric + # 2 saturation + # 3 absolute colorimetric + + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if length < 1: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + msg = "Truncated sRGB chunk" + raise ValueError(msg) + self.im_info["srgb"] = s[0] + return s + + def chunk_pHYs(self, pos: int, length: int) -> bytes: + # pixels per unit + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if length < 9: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + msg = "Truncated pHYs chunk" + raise ValueError(msg) + px, py = i32(s, 0), i32(s, 4) + unit = s[8] + if unit == 1: # meter + dpi = px * 0.0254, py * 0.0254 + self.im_info["dpi"] = dpi + elif unit == 0: + self.im_info["aspect"] = px, py + return s + + def chunk_tEXt(self, pos: int, length: int) -> bytes: + # text + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + try: + k, v = s.split(b"\0", 1) + except ValueError: + # fallback for broken tEXt tags + k = s + v = b"" + if k: + k_str = k.decode("latin-1", "strict") + v_str = v.decode("latin-1", "replace") + + self.im_info[k_str] = v if k == b"exif" else v_str + self.im_text[k_str] = v_str + self.check_text_memory(len(v_str)) + + return s + + def chunk_zTXt(self, pos: int, length: int) -> bytes: + # compressed text + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + try: + k, v = s.split(b"\0", 1) + except ValueError: + k = s + v = b"" + if v: + comp_method = v[0] + else: + comp_method = 0 + if comp_method != 0: + msg = f"Unknown compression method {comp_method} in zTXt chunk" + raise SyntaxError(msg) + try: + v = _safe_zlib_decompress(v[1:]) + except ValueError: + if ImageFile.LOAD_TRUNCATED_IMAGES: + v = b"" + else: + raise + except zlib.error: + v = b"" + + if k: + k_str = k.decode("latin-1", "strict") + v_str = v.decode("latin-1", "replace") + + self.im_info[k_str] = self.im_text[k_str] = v_str + self.check_text_memory(len(v_str)) + + return s + + def chunk_iTXt(self, pos: int, length: int) -> bytes: + # international text + assert self.fp is not None + r = s = ImageFile._safe_read(self.fp, length) + try: + k, r = r.split(b"\0", 1) + except ValueError: + return s + if len(r) < 2: + return s + cf, cm, r = r[0], r[1], r[2:] + try: + lang, tk, v = r.split(b"\0", 2) + except ValueError: + return s + if cf != 0: + if cm == 0: + try: + v = _safe_zlib_decompress(v) + except ValueError: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + else: + raise + except zlib.error: + return s + else: + return s + if k == b"XML:com.adobe.xmp": + self.im_info["xmp"] = v + try: + k_str = k.decode("latin-1", "strict") + lang_str = lang.decode("utf-8", "strict") + tk_str = tk.decode("utf-8", "strict") + v_str = v.decode("utf-8", "strict") + except UnicodeError: + return s + + self.im_info[k_str] = self.im_text[k_str] = iTXt(v_str, lang_str, tk_str) + self.check_text_memory(len(v_str)) + + return s + + def chunk_eXIf(self, pos: int, length: int) -> bytes: + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + self.im_info["exif"] = b"Exif\x00\x00" + s + return s + + # APNG chunks + def chunk_acTL(self, pos: int, length: int) -> bytes: + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if length < 8: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + msg = "APNG contains truncated acTL chunk" + raise ValueError(msg) + if self.im_n_frames is not None: + self.im_n_frames = None + warnings.warn("Invalid APNG, will use default PNG image if possible") + return s + n_frames = i32(s) + if n_frames == 0 or n_frames > 0x80000000: + warnings.warn("Invalid APNG, will use default PNG image if possible") + return s + self.im_n_frames = n_frames + self.im_info["loop"] = i32(s, 4) + self.im_custom_mimetype = "image/apng" + return s + + def chunk_fcTL(self, pos: int, length: int) -> bytes: + assert self.fp is not None + s = ImageFile._safe_read(self.fp, length) + if length < 26: + if ImageFile.LOAD_TRUNCATED_IMAGES: + return s + msg = "APNG contains truncated fcTL chunk" + raise ValueError(msg) + seq = i32(s) + if (self._seq_num is None and seq != 0) or ( + self._seq_num is not None and self._seq_num != seq - 1 + ): + msg = "APNG contains frame sequence errors" + raise SyntaxError(msg) + self._seq_num = seq + width, height = i32(s, 4), i32(s, 8) + px, py = i32(s, 12), i32(s, 16) + im_w, im_h = self.im_size + if px + width > im_w or py + height > im_h: + msg = "APNG contains invalid frames" + raise SyntaxError(msg) + self.im_info["bbox"] = (px, py, px + width, py + height) + delay_num, delay_den = i16(s, 20), i16(s, 22) + if delay_den == 0: + delay_den = 100 + self.im_info["duration"] = float(delay_num) / float(delay_den) * 1000 + self.im_info["disposal"] = s[24] + self.im_info["blend"] = s[25] + return s + + def chunk_fdAT(self, pos: int, length: int) -> bytes: + assert self.fp is not None + if length < 4: + if ImageFile.LOAD_TRUNCATED_IMAGES: + s = ImageFile._safe_read(self.fp, length) + return s + msg = "APNG contains truncated fDAT chunk" + raise ValueError(msg) + s = ImageFile._safe_read(self.fp, 4) + seq = i32(s) + if self._seq_num != seq - 1: + msg = "APNG contains frame sequence errors" + raise SyntaxError(msg) + self._seq_num = seq + return self.chunk_IDAT(pos + 4, length - 4) + + +# -------------------------------------------------------------------- +# PNG reader + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(_MAGIC) + + +## +# Image plugin for PNG images. + + +class PngImageFile(ImageFile.ImageFile): + format = "PNG" + format_description = "Portable network graphics" + + def _open(self) -> None: + if not _accept(self.fp.read(8)): + msg = "not a PNG file" + raise SyntaxError(msg) + self._fp = self.fp + self.__frame = 0 + + # + # Parse headers up to the first IDAT or fDAT chunk + + self.private_chunks: list[tuple[bytes, bytes] | tuple[bytes, bytes, bool]] = [] + self.png: PngStream | None = PngStream(self.fp) + + while True: + # + # get next chunk + + cid, pos, length = self.png.read() + + try: + s = self.png.call(cid, pos, length) + except EOFError: + break + except AttributeError: + logger.debug("%r %s %s (unknown)", cid, pos, length) + s = ImageFile._safe_read(self.fp, length) + if cid[1:2].islower(): + self.private_chunks.append((cid, s)) + + self.png.crc(cid, s) + + # + # Copy relevant attributes from the PngStream. An alternative + # would be to let the PngStream class modify these attributes + # directly, but that introduces circular references which are + # difficult to break if things go wrong in the decoder... + # (believe me, I've tried ;-) + + self._mode = self.png.im_mode + self._size = self.png.im_size + self.info = self.png.im_info + self._text: dict[str, str | iTXt] | None = None + self.tile = self.png.im_tile + self.custom_mimetype = self.png.im_custom_mimetype + self.n_frames = self.png.im_n_frames or 1 + self.default_image = self.info.get("default_image", False) + + if self.png.im_palette: + rawmode, data = self.png.im_palette + self.palette = ImagePalette.raw(rawmode, data) + + if cid == b"fdAT": + self.__prepare_idat = length - 4 + else: + self.__prepare_idat = length # used by load_prepare() + + if self.png.im_n_frames is not None: + self._close_exclusive_fp_after_loading = False + self.png.save_rewind() + self.__rewind_idat = self.__prepare_idat + self.__rewind = self._fp.tell() + if self.default_image: + # IDAT chunk contains default image and not first animation frame + self.n_frames += 1 + self._seek(0) + self.is_animated = self.n_frames > 1 + + @property + def text(self) -> dict[str, str | iTXt]: + # experimental + if self._text is None: + # iTxt, tEXt and zTXt chunks may appear at the end of the file + # So load the file to ensure that they are read + if self.is_animated: + frame = self.__frame + # for APNG, seek to the final frame before loading + self.seek(self.n_frames - 1) + self.load() + if self.is_animated: + self.seek(frame) + assert self._text is not None + return self._text + + def verify(self) -> None: + """Verify PNG file""" + + if self.fp is None: + msg = "verify must be called directly after open" + raise RuntimeError(msg) + + # back up to beginning of IDAT block + self.fp.seek(self.tile[0][2] - 8) + + assert self.png is not None + self.png.verify() + self.png.close() + + if self._exclusive_fp: + self.fp.close() + self.fp = None + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + if frame < self.__frame: + self._seek(0, True) + + last_frame = self.__frame + for f in range(self.__frame + 1, frame + 1): + try: + self._seek(f) + except EOFError as e: + self.seek(last_frame) + msg = "no more images in APNG file" + raise EOFError(msg) from e + + def _seek(self, frame: int, rewind: bool = False) -> None: + assert self.png is not None + if isinstance(self._fp, DeferredError): + raise self._fp.ex + + self.dispose: _imaging.ImagingCore | None + dispose_extent = None + if frame == 0: + if rewind: + self._fp.seek(self.__rewind) + self.png.rewind() + self.__prepare_idat = self.__rewind_idat + self._im = None + self.info = self.png.im_info + self.tile = self.png.im_tile + self.fp = self._fp + self._prev_im = None + self.dispose = None + self.default_image = self.info.get("default_image", False) + self.dispose_op = self.info.get("disposal") + self.blend_op = self.info.get("blend") + dispose_extent = self.info.get("bbox") + self.__frame = 0 + else: + if frame != self.__frame + 1: + msg = f"cannot seek to frame {frame}" + raise ValueError(msg) + + # ensure previous frame was loaded + self.load() + + if self.dispose: + self.im.paste(self.dispose, self.dispose_extent) + self._prev_im = self.im.copy() + + self.fp = self._fp + + # advance to the next frame + if self.__prepare_idat: + ImageFile._safe_read(self.fp, self.__prepare_idat) + self.__prepare_idat = 0 + frame_start = False + while True: + self.fp.read(4) # CRC + + try: + cid, pos, length = self.png.read() + except (struct.error, SyntaxError): + break + + if cid == b"IEND": + msg = "No more images in APNG file" + raise EOFError(msg) + if cid == b"fcTL": + if frame_start: + # there must be at least one fdAT chunk between fcTL chunks + msg = "APNG missing frame data" + raise SyntaxError(msg) + frame_start = True + + try: + self.png.call(cid, pos, length) + except UnicodeDecodeError: + break + except EOFError: + if cid == b"fdAT": + length -= 4 + if frame_start: + self.__prepare_idat = length + break + ImageFile._safe_read(self.fp, length) + except AttributeError: + logger.debug("%r %s %s (unknown)", cid, pos, length) + ImageFile._safe_read(self.fp, length) + + self.__frame = frame + self.tile = self.png.im_tile + self.dispose_op = self.info.get("disposal") + self.blend_op = self.info.get("blend") + dispose_extent = self.info.get("bbox") + + if not self.tile: + msg = "image not found in APNG frame" + raise EOFError(msg) + if dispose_extent: + self.dispose_extent: tuple[float, float, float, float] = dispose_extent + + # setup frame disposal (actual disposal done when needed in the next _seek()) + if self._prev_im is None and self.dispose_op == Disposal.OP_PREVIOUS: + self.dispose_op = Disposal.OP_BACKGROUND + + self.dispose = None + if self.dispose_op == Disposal.OP_PREVIOUS: + if self._prev_im: + self.dispose = self._prev_im.copy() + self.dispose = self._crop(self.dispose, self.dispose_extent) + elif self.dispose_op == Disposal.OP_BACKGROUND: + self.dispose = Image.core.fill(self.mode, self.size) + self.dispose = self._crop(self.dispose, self.dispose_extent) + + def tell(self) -> int: + return self.__frame + + def load_prepare(self) -> None: + """internal: prepare to read PNG file""" + + if self.info.get("interlace"): + self.decoderconfig = self.decoderconfig + (1,) + + self.__idat = self.__prepare_idat # used by load_read() + ImageFile.ImageFile.load_prepare(self) + + def load_read(self, read_bytes: int) -> bytes: + """internal: read more image data""" + + assert self.png is not None + while self.__idat == 0: + # end of chunk, skip forward to next one + + self.fp.read(4) # CRC + + cid, pos, length = self.png.read() + + if cid not in [b"IDAT", b"DDAT", b"fdAT"]: + self.png.push(cid, pos, length) + return b"" + + if cid == b"fdAT": + try: + self.png.call(cid, pos, length) + except EOFError: + pass + self.__idat = length - 4 # sequence_num has already been read + else: + self.__idat = length # empty chunks are allowed + + # read more data from this chunk + if read_bytes <= 0: + read_bytes = self.__idat + else: + read_bytes = min(read_bytes, self.__idat) + + self.__idat = self.__idat - read_bytes + + return self.fp.read(read_bytes) + + def load_end(self) -> None: + """internal: finished reading image data""" + assert self.png is not None + if self.__idat != 0: + self.fp.read(self.__idat) + while True: + self.fp.read(4) # CRC + + try: + cid, pos, length = self.png.read() + except (struct.error, SyntaxError): + break + + if cid == b"IEND": + break + elif cid == b"fcTL" and self.is_animated: + # start of the next frame, stop reading + self.__prepare_idat = 0 + self.png.push(cid, pos, length) + break + + try: + self.png.call(cid, pos, length) + except UnicodeDecodeError: + break + except EOFError: + if cid == b"fdAT": + length -= 4 + try: + ImageFile._safe_read(self.fp, length) + except OSError as e: + if ImageFile.LOAD_TRUNCATED_IMAGES: + break + else: + raise e + except AttributeError: + logger.debug("%r %s %s (unknown)", cid, pos, length) + s = ImageFile._safe_read(self.fp, length) + if cid[1:2].islower(): + self.private_chunks.append((cid, s, True)) + self._text = self.png.im_text + if not self.is_animated: + self.png.close() + self.png = None + else: + if self._prev_im and self.blend_op == Blend.OP_OVER: + updated = self._crop(self.im, self.dispose_extent) + if self.im.mode == "RGB" and "transparency" in self.info: + mask = updated.convert_transparent( + "RGBA", self.info["transparency"] + ) + else: + if self.im.mode == "P" and "transparency" in self.info: + t = self.info["transparency"] + if isinstance(t, bytes): + updated.putpalettealphas(t) + elif isinstance(t, int): + updated.putpalettealpha(t) + mask = updated.convert("RGBA") + self._prev_im.paste(updated, self.dispose_extent, mask) + self.im = self._prev_im + + def _getexif(self) -> dict[int, Any] | None: + if "exif" not in self.info: + self.load() + if "exif" not in self.info and "Raw profile type exif" not in self.info: + return None + return self.getexif()._get_merged_dict() + + def getexif(self) -> Image.Exif: + if "exif" not in self.info: + self.load() + + return super().getexif() + + +# -------------------------------------------------------------------- +# PNG writer + +_OUTMODES = { + # supported PIL modes, and corresponding rawmode, bit depth and color type + "1": ("1", b"\x01", b"\x00"), + "L;1": ("L;1", b"\x01", b"\x00"), + "L;2": ("L;2", b"\x02", b"\x00"), + "L;4": ("L;4", b"\x04", b"\x00"), + "L": ("L", b"\x08", b"\x00"), + "LA": ("LA", b"\x08", b"\x04"), + "I": ("I;16B", b"\x10", b"\x00"), + "I;16": ("I;16B", b"\x10", b"\x00"), + "I;16B": ("I;16B", b"\x10", b"\x00"), + "P;1": ("P;1", b"\x01", b"\x03"), + "P;2": ("P;2", b"\x02", b"\x03"), + "P;4": ("P;4", b"\x04", b"\x03"), + "P": ("P", b"\x08", b"\x03"), + "RGB": ("RGB", b"\x08", b"\x02"), + "RGBA": ("RGBA", b"\x08", b"\x06"), +} + + +def putchunk(fp: IO[bytes], cid: bytes, *data: bytes) -> None: + """Write a PNG chunk (including CRC field)""" + + byte_data = b"".join(data) + + fp.write(o32(len(byte_data)) + cid) + fp.write(byte_data) + crc = _crc32(byte_data, _crc32(cid)) + fp.write(o32(crc)) + + +class _idat: + # wrap output from the encoder in IDAT chunks + + def __init__(self, fp: IO[bytes], chunk: Callable[..., None]) -> None: + self.fp = fp + self.chunk = chunk + + def write(self, data: bytes) -> None: + self.chunk(self.fp, b"IDAT", data) + + +class _fdat: + # wrap encoder output in fdAT chunks + + def __init__(self, fp: IO[bytes], chunk: Callable[..., None], seq_num: int) -> None: + self.fp = fp + self.chunk = chunk + self.seq_num = seq_num + + def write(self, data: bytes) -> None: + self.chunk(self.fp, b"fdAT", o32(self.seq_num), data) + self.seq_num += 1 + + +class _Frame(NamedTuple): + im: Image.Image + bbox: tuple[int, int, int, int] | None + encoderinfo: dict[str, Any] + + +def _write_multiple_frames( + im: Image.Image, + fp: IO[bytes], + chunk: Callable[..., None], + mode: str, + rawmode: str, + default_image: Image.Image | None, + append_images: list[Image.Image], +) -> Image.Image | None: + duration = im.encoderinfo.get("duration") + loop = im.encoderinfo.get("loop", im.info.get("loop", 0)) + disposal = im.encoderinfo.get("disposal", im.info.get("disposal", Disposal.OP_NONE)) + blend = im.encoderinfo.get("blend", im.info.get("blend", Blend.OP_SOURCE)) + + if default_image: + chain = itertools.chain(append_images) + else: + chain = itertools.chain([im], append_images) + + im_frames: list[_Frame] = [] + frame_count = 0 + for im_seq in chain: + for im_frame in ImageSequence.Iterator(im_seq): + if im_frame.mode == mode: + im_frame = im_frame.copy() + else: + im_frame = im_frame.convert(mode) + encoderinfo = im.encoderinfo.copy() + if isinstance(duration, (list, tuple)): + encoderinfo["duration"] = duration[frame_count] + elif duration is None and "duration" in im_frame.info: + encoderinfo["duration"] = im_frame.info["duration"] + if isinstance(disposal, (list, tuple)): + encoderinfo["disposal"] = disposal[frame_count] + if isinstance(blend, (list, tuple)): + encoderinfo["blend"] = blend[frame_count] + frame_count += 1 + + if im_frames: + previous = im_frames[-1] + prev_disposal = previous.encoderinfo.get("disposal") + prev_blend = previous.encoderinfo.get("blend") + if prev_disposal == Disposal.OP_PREVIOUS and len(im_frames) < 2: + prev_disposal = Disposal.OP_BACKGROUND + + if prev_disposal == Disposal.OP_BACKGROUND: + base_im = previous.im.copy() + dispose = Image.core.fill("RGBA", im.size, (0, 0, 0, 0)) + bbox = previous.bbox + if bbox: + dispose = dispose.crop(bbox) + else: + bbox = (0, 0) + im.size + base_im.paste(dispose, bbox) + elif prev_disposal == Disposal.OP_PREVIOUS: + base_im = im_frames[-2].im + else: + base_im = previous.im + delta = ImageChops.subtract_modulo( + im_frame.convert("RGBA"), base_im.convert("RGBA") + ) + bbox = delta.getbbox(alpha_only=False) + if ( + not bbox + and prev_disposal == encoderinfo.get("disposal") + and prev_blend == encoderinfo.get("blend") + and "duration" in encoderinfo + ): + previous.encoderinfo["duration"] += encoderinfo["duration"] + continue + else: + bbox = None + im_frames.append(_Frame(im_frame, bbox, encoderinfo)) + + if len(im_frames) == 1 and not default_image: + return im_frames[0].im + + # animation control + chunk( + fp, + b"acTL", + o32(len(im_frames)), # 0: num_frames + o32(loop), # 4: num_plays + ) + + # default image IDAT (if it exists) + if default_image: + if im.mode != mode: + im = im.convert(mode) + ImageFile._save( + im, + cast(IO[bytes], _idat(fp, chunk)), + [ImageFile._Tile("zip", (0, 0) + im.size, 0, rawmode)], + ) + + seq_num = 0 + for frame, frame_data in enumerate(im_frames): + im_frame = frame_data.im + if not frame_data.bbox: + bbox = (0, 0) + im_frame.size + else: + bbox = frame_data.bbox + im_frame = im_frame.crop(bbox) + size = im_frame.size + encoderinfo = frame_data.encoderinfo + frame_duration = int(round(encoderinfo.get("duration", 0))) + frame_disposal = encoderinfo.get("disposal", disposal) + frame_blend = encoderinfo.get("blend", blend) + # frame control + chunk( + fp, + b"fcTL", + o32(seq_num), # sequence_number + o32(size[0]), # width + o32(size[1]), # height + o32(bbox[0]), # x_offset + o32(bbox[1]), # y_offset + o16(frame_duration), # delay_numerator + o16(1000), # delay_denominator + o8(frame_disposal), # dispose_op + o8(frame_blend), # blend_op + ) + seq_num += 1 + # frame data + if frame == 0 and not default_image: + # first frame must be in IDAT chunks for backwards compatibility + ImageFile._save( + im_frame, + cast(IO[bytes], _idat(fp, chunk)), + [ImageFile._Tile("zip", (0, 0) + im_frame.size, 0, rawmode)], + ) + else: + fdat_chunks = _fdat(fp, chunk, seq_num) + ImageFile._save( + im_frame, + cast(IO[bytes], fdat_chunks), + [ImageFile._Tile("zip", (0, 0) + im_frame.size, 0, rawmode)], + ) + seq_num = fdat_chunks.seq_num + return None + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + _save(im, fp, filename, save_all=True) + + +def _save( + im: Image.Image, + fp: IO[bytes], + filename: str | bytes, + chunk: Callable[..., None] = putchunk, + save_all: bool = False, +) -> None: + # save an image to disk (called by the save method) + + if save_all: + default_image = im.encoderinfo.get( + "default_image", im.info.get("default_image") + ) + modes = set() + sizes = set() + append_images = im.encoderinfo.get("append_images", []) + for im_seq in itertools.chain([im], append_images): + for im_frame in ImageSequence.Iterator(im_seq): + modes.add(im_frame.mode) + sizes.add(im_frame.size) + for mode in ("RGBA", "RGB", "P"): + if mode in modes: + break + else: + mode = modes.pop() + size = tuple(max(frame_size[i] for frame_size in sizes) for i in range(2)) + else: + size = im.size + mode = im.mode + + outmode = mode + if mode == "P": + # + # attempt to minimize storage requirements for palette images + if "bits" in im.encoderinfo: + # number of bits specified by user + colors = min(1 << im.encoderinfo["bits"], 256) + else: + # check palette contents + if im.palette: + colors = max(min(len(im.palette.getdata()[1]) // 3, 256), 1) + else: + colors = 256 + + if colors <= 16: + if colors <= 2: + bits = 1 + elif colors <= 4: + bits = 2 + else: + bits = 4 + outmode += f";{bits}" + + # encoder options + im.encoderconfig = ( + im.encoderinfo.get("optimize", False), + im.encoderinfo.get("compress_level", -1), + im.encoderinfo.get("compress_type", -1), + im.encoderinfo.get("dictionary", b""), + ) + + # get the corresponding PNG mode + try: + rawmode, bit_depth, color_type = _OUTMODES[outmode] + except KeyError as e: + msg = f"cannot write mode {mode} as PNG" + raise OSError(msg) from e + if outmode == "I": + deprecate("Saving I mode images as PNG", 13, stacklevel=4) + + # + # write minimal PNG file + + fp.write(_MAGIC) + + chunk( + fp, + b"IHDR", + o32(size[0]), # 0: size + o32(size[1]), + bit_depth, + color_type, + b"\0", # 10: compression + b"\0", # 11: filter category + b"\0", # 12: interlace flag + ) + + chunks = [b"cHRM", b"cICP", b"gAMA", b"sBIT", b"sRGB", b"tIME"] + + icc = im.encoderinfo.get("icc_profile", im.info.get("icc_profile")) + if icc: + # ICC profile + # according to PNG spec, the iCCP chunk contains: + # Profile name 1-79 bytes (character string) + # Null separator 1 byte (null character) + # Compression method 1 byte (0) + # Compressed profile n bytes (zlib with deflate compression) + name = b"ICC Profile" + data = name + b"\0\0" + zlib.compress(icc) + chunk(fp, b"iCCP", data) + + # You must either have sRGB or iCCP. + # Disallow sRGB chunks when an iCCP-chunk has been emitted. + chunks.remove(b"sRGB") + + info = im.encoderinfo.get("pnginfo") + if info: + chunks_multiple_allowed = [b"sPLT", b"iTXt", b"tEXt", b"zTXt"] + for info_chunk in info.chunks: + cid, data = info_chunk[:2] + if cid in chunks: + chunks.remove(cid) + chunk(fp, cid, data) + elif cid in chunks_multiple_allowed: + chunk(fp, cid, data) + elif cid[1:2].islower(): + # Private chunk + after_idat = len(info_chunk) == 3 and info_chunk[2] + if not after_idat: + chunk(fp, cid, data) + + if im.mode == "P": + palette_byte_number = colors * 3 + palette_bytes = im.im.getpalette("RGB")[:palette_byte_number] + while len(palette_bytes) < palette_byte_number: + palette_bytes += b"\0" + chunk(fp, b"PLTE", palette_bytes) + + transparency = im.encoderinfo.get("transparency", im.info.get("transparency", None)) + + if transparency or transparency == 0: + if im.mode == "P": + # limit to actual palette size + alpha_bytes = colors + if isinstance(transparency, bytes): + chunk(fp, b"tRNS", transparency[:alpha_bytes]) + else: + transparency = max(0, min(255, transparency)) + alpha = b"\xff" * transparency + b"\0" + chunk(fp, b"tRNS", alpha[:alpha_bytes]) + elif im.mode in ("1", "L", "I", "I;16"): + transparency = max(0, min(65535, transparency)) + chunk(fp, b"tRNS", o16(transparency)) + elif im.mode == "RGB": + red, green, blue = transparency + chunk(fp, b"tRNS", o16(red) + o16(green) + o16(blue)) + else: + if "transparency" in im.encoderinfo: + # don't bother with transparency if it's an RGBA + # and it's in the info dict. It's probably just stale. + msg = "cannot use transparency for this mode" + raise OSError(msg) + else: + if im.mode == "P" and im.im.getpalettemode() == "RGBA": + alpha = im.im.getpalette("RGBA", "A") + alpha_bytes = colors + chunk(fp, b"tRNS", alpha[:alpha_bytes]) + + dpi = im.encoderinfo.get("dpi") + if dpi: + chunk( + fp, + b"pHYs", + o32(int(dpi[0] / 0.0254 + 0.5)), + o32(int(dpi[1] / 0.0254 + 0.5)), + b"\x01", + ) + + if info: + chunks = [b"bKGD", b"hIST"] + for info_chunk in info.chunks: + cid, data = info_chunk[:2] + if cid in chunks: + chunks.remove(cid) + chunk(fp, cid, data) + + exif = im.encoderinfo.get("exif") + if exif: + if isinstance(exif, Image.Exif): + exif = exif.tobytes(8) + if exif.startswith(b"Exif\x00\x00"): + exif = exif[6:] + chunk(fp, b"eXIf", exif) + + single_im: Image.Image | None = im + if save_all: + single_im = _write_multiple_frames( + im, fp, chunk, mode, rawmode, default_image, append_images + ) + if single_im: + ImageFile._save( + single_im, + cast(IO[bytes], _idat(fp, chunk)), + [ImageFile._Tile("zip", (0, 0) + single_im.size, 0, rawmode)], + ) + + if info: + for info_chunk in info.chunks: + cid, data = info_chunk[:2] + if cid[1:2].islower(): + # Private chunk + after_idat = len(info_chunk) == 3 and info_chunk[2] + if after_idat: + chunk(fp, cid, data) + + chunk(fp, b"IEND", b"") + + if hasattr(fp, "flush"): + fp.flush() + + +# -------------------------------------------------------------------- +# PNG chunk converter + + +def getchunks(im: Image.Image, **params: Any) -> list[tuple[bytes, bytes, bytes]]: + """Return a list of PNG chunks representing this image.""" + from io import BytesIO + + chunks = [] + + def append(fp: IO[bytes], cid: bytes, *data: bytes) -> None: + byte_data = b"".join(data) + crc = o32(_crc32(byte_data, _crc32(cid))) + chunks.append((cid, byte_data, crc)) + + fp = BytesIO() + + try: + im.encoderinfo = params + _save(im, fp, "", append) + finally: + del im.encoderinfo + + return chunks + + +# -------------------------------------------------------------------- +# Registry + +Image.register_open(PngImageFile.format, PngImageFile, _accept) +Image.register_save(PngImageFile.format, _save) +Image.register_save_all(PngImageFile.format, _save_all) + +Image.register_extensions(PngImageFile.format, [".png", ".apng"]) + +Image.register_mime(PngImageFile.format, "image/png") diff --git a/venv/lib/python3.10/site-packages/PIL/PpmImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PpmImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..307bc97ff65160fabf533d12e96ef4382ed5f1e7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PpmImagePlugin.py @@ -0,0 +1,375 @@ +# +# The Python Imaging Library. +# $Id$ +# +# PPM support for PIL +# +# History: +# 96-03-24 fl Created +# 98-03-06 fl Write RGBA images (as RGB, that is) +# +# Copyright (c) Secret Labs AB 1997-98. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import math +from typing import IO + +from . import Image, ImageFile +from ._binary import i16be as i16 +from ._binary import o8 +from ._binary import o32le as o32 + +# +# -------------------------------------------------------------------- + +b_whitespace = b"\x20\x09\x0a\x0b\x0c\x0d" + +MODES = { + # standard + b"P1": "1", + b"P2": "L", + b"P3": "RGB", + b"P4": "1", + b"P5": "L", + b"P6": "RGB", + # extensions + b"P0CMYK": "CMYK", + b"Pf": "F", + # PIL extensions (for test purposes only) + b"PyP": "P", + b"PyRGBA": "RGBA", + b"PyCMYK": "CMYK", +} + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 2 and prefix.startswith(b"P") and prefix[1] in b"0123456fy" + + +## +# Image plugin for PBM, PGM, and PPM images. + + +class PpmImageFile(ImageFile.ImageFile): + format = "PPM" + format_description = "Pbmplus image" + + def _read_magic(self) -> bytes: + assert self.fp is not None + + magic = b"" + # read until whitespace or longest available magic number + for _ in range(6): + c = self.fp.read(1) + if not c or c in b_whitespace: + break + magic += c + return magic + + def _read_token(self) -> bytes: + assert self.fp is not None + + token = b"" + while len(token) <= 10: # read until next whitespace or limit of 10 characters + c = self.fp.read(1) + if not c: + break + elif c in b_whitespace: # token ended + if not token: + # skip whitespace at start + continue + break + elif c == b"#": + # ignores rest of the line; stops at CR, LF or EOF + while self.fp.read(1) not in b"\r\n": + pass + continue + token += c + if not token: + # Token was not even 1 byte + msg = "Reached EOF while reading header" + raise ValueError(msg) + elif len(token) > 10: + msg_too_long = b"Token too long in file header: %s" % token + raise ValueError(msg_too_long) + return token + + def _open(self) -> None: + assert self.fp is not None + + magic_number = self._read_magic() + try: + mode = MODES[magic_number] + except KeyError: + msg = "not a PPM file" + raise SyntaxError(msg) + self._mode = mode + + if magic_number in (b"P1", b"P4"): + self.custom_mimetype = "image/x-portable-bitmap" + elif magic_number in (b"P2", b"P5"): + self.custom_mimetype = "image/x-portable-graymap" + elif magic_number in (b"P3", b"P6"): + self.custom_mimetype = "image/x-portable-pixmap" + + self._size = int(self._read_token()), int(self._read_token()) + + decoder_name = "raw" + if magic_number in (b"P1", b"P2", b"P3"): + decoder_name = "ppm_plain" + + args: str | tuple[str | int, ...] + if mode == "1": + args = "1;I" + elif mode == "F": + scale = float(self._read_token()) + if scale == 0.0 or not math.isfinite(scale): + msg = "scale must be finite and non-zero" + raise ValueError(msg) + self.info["scale"] = abs(scale) + + rawmode = "F;32F" if scale < 0 else "F;32BF" + args = (rawmode, 0, -1) + else: + maxval = int(self._read_token()) + if not 0 < maxval < 65536: + msg = "maxval must be greater than 0 and less than 65536" + raise ValueError(msg) + if maxval > 255 and mode == "L": + self._mode = "I" + + rawmode = mode + if decoder_name != "ppm_plain": + # If maxval matches a bit depth, use the raw decoder directly + if maxval == 65535 and mode == "L": + rawmode = "I;16B" + elif maxval != 255: + decoder_name = "ppm" + + args = rawmode if decoder_name == "raw" else (rawmode, maxval) + self.tile = [ + ImageFile._Tile(decoder_name, (0, 0) + self.size, self.fp.tell(), args) + ] + + +# +# -------------------------------------------------------------------- + + +class PpmPlainDecoder(ImageFile.PyDecoder): + _pulls_fd = True + _comment_spans: bool + + def _read_block(self) -> bytes: + assert self.fd is not None + + return self.fd.read(ImageFile.SAFEBLOCK) + + def _find_comment_end(self, block: bytes, start: int = 0) -> int: + a = block.find(b"\n", start) + b = block.find(b"\r", start) + return min(a, b) if a * b > 0 else max(a, b) # lowest nonnegative index (or -1) + + def _ignore_comments(self, block: bytes) -> bytes: + if self._comment_spans: + # Finish current comment + while block: + comment_end = self._find_comment_end(block) + if comment_end != -1: + # Comment ends in this block + # Delete tail of comment + block = block[comment_end + 1 :] + break + else: + # Comment spans whole block + # So read the next block, looking for the end + block = self._read_block() + + # Search for any further comments + self._comment_spans = False + while True: + comment_start = block.find(b"#") + if comment_start == -1: + # No comment found + break + comment_end = self._find_comment_end(block, comment_start) + if comment_end != -1: + # Comment ends in this block + # Delete comment + block = block[:comment_start] + block[comment_end + 1 :] + else: + # Comment continues to next block(s) + block = block[:comment_start] + self._comment_spans = True + break + return block + + def _decode_bitonal(self) -> bytearray: + """ + This is a separate method because in the plain PBM format, all data tokens are + exactly one byte, so the inter-token whitespace is optional. + """ + data = bytearray() + total_bytes = self.state.xsize * self.state.ysize + + while len(data) != total_bytes: + block = self._read_block() # read next block + if not block: + # eof + break + + block = self._ignore_comments(block) + + tokens = b"".join(block.split()) + for token in tokens: + if token not in (48, 49): + msg = b"Invalid token for this mode: %s" % bytes([token]) + raise ValueError(msg) + data = (data + tokens)[:total_bytes] + invert = bytes.maketrans(b"01", b"\xff\x00") + return data.translate(invert) + + def _decode_blocks(self, maxval: int) -> bytearray: + data = bytearray() + max_len = 10 + out_byte_count = 4 if self.mode == "I" else 1 + out_max = 65535 if self.mode == "I" else 255 + bands = Image.getmodebands(self.mode) + total_bytes = self.state.xsize * self.state.ysize * bands * out_byte_count + + half_token = b"" + while len(data) != total_bytes: + block = self._read_block() # read next block + if not block: + if half_token: + block = bytearray(b" ") # flush half_token + else: + # eof + break + + block = self._ignore_comments(block) + + if half_token: + block = half_token + block # stitch half_token to new block + half_token = b"" + + tokens = block.split() + + if block and not block[-1:].isspace(): # block might split token + half_token = tokens.pop() # save half token for later + if len(half_token) > max_len: # prevent buildup of half_token + msg = ( + b"Token too long found in data: %s" % half_token[: max_len + 1] + ) + raise ValueError(msg) + + for token in tokens: + if len(token) > max_len: + msg = b"Token too long found in data: %s" % token[: max_len + 1] + raise ValueError(msg) + value = int(token) + if value < 0: + msg_str = f"Channel value is negative: {value}" + raise ValueError(msg_str) + if value > maxval: + msg_str = f"Channel value too large for this mode: {value}" + raise ValueError(msg_str) + value = round(value / maxval * out_max) + data += o32(value) if self.mode == "I" else o8(value) + if len(data) == total_bytes: # finished! + break + return data + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + self._comment_spans = False + if self.mode == "1": + data = self._decode_bitonal() + rawmode = "1;8" + else: + maxval = self.args[-1] + data = self._decode_blocks(maxval) + rawmode = "I;32" if self.mode == "I" else self.mode + self.set_as_raw(bytes(data), rawmode) + return -1, 0 + + +class PpmDecoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + + data = bytearray() + maxval = self.args[-1] + in_byte_count = 1 if maxval < 256 else 2 + out_byte_count = 4 if self.mode == "I" else 1 + out_max = 65535 if self.mode == "I" else 255 + bands = Image.getmodebands(self.mode) + dest_length = self.state.xsize * self.state.ysize * bands * out_byte_count + while len(data) < dest_length: + pixels = self.fd.read(in_byte_count * bands) + if len(pixels) < in_byte_count * bands: + # eof + break + for b in range(bands): + value = ( + pixels[b] if in_byte_count == 1 else i16(pixels, b * in_byte_count) + ) + value = min(out_max, round(value / maxval * out_max)) + data += o32(value) if self.mode == "I" else o8(value) + rawmode = "I;32" if self.mode == "I" else self.mode + self.set_as_raw(bytes(data), rawmode) + return -1, 0 + + +# +# -------------------------------------------------------------------- + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode == "1": + rawmode, head = "1;I", b"P4" + elif im.mode == "L": + rawmode, head = "L", b"P5" + elif im.mode in ("I", "I;16"): + rawmode, head = "I;16B", b"P5" + elif im.mode in ("RGB", "RGBA"): + rawmode, head = "RGB", b"P6" + elif im.mode == "F": + rawmode, head = "F;32F", b"Pf" + else: + msg = f"cannot write mode {im.mode} as PPM" + raise OSError(msg) + fp.write(head + b"\n%d %d\n" % im.size) + if head == b"P6": + fp.write(b"255\n") + elif head == b"P5": + if rawmode == "L": + fp.write(b"255\n") + else: + fp.write(b"65535\n") + elif head == b"Pf": + fp.write(b"-1.0\n") + row_order = -1 if im.mode == "F" else 1 + ImageFile._save( + im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 0, (rawmode, 0, row_order))] + ) + + +# +# -------------------------------------------------------------------- + + +Image.register_open(PpmImageFile.format, PpmImageFile, _accept) +Image.register_save(PpmImageFile.format, _save) + +Image.register_decoder("ppm", PpmDecoder) +Image.register_decoder("ppm_plain", PpmPlainDecoder) + +Image.register_extensions(PpmImageFile.format, [".pbm", ".pgm", ".ppm", ".pnm", ".pfm"]) + +Image.register_mime(PpmImageFile.format, "image/x-portable-anymap") diff --git a/venv/lib/python3.10/site-packages/PIL/PsdImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/PsdImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..f49aaeeb1f55cd2b6b17d9aa10f68876384fd410 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/PsdImagePlugin.py @@ -0,0 +1,333 @@ +# +# The Python Imaging Library +# $Id$ +# +# Adobe PSD 2.5/3.0 file handling +# +# History: +# 1995-09-01 fl Created +# 1997-01-03 fl Read most PSD images +# 1997-01-18 fl Fixed P and CMYK support +# 2001-10-21 fl Added seek/tell support (for layers) +# +# Copyright (c) 1997-2001 by Secret Labs AB. +# Copyright (c) 1995-2001 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +from functools import cached_property +from typing import IO + +from . import Image, ImageFile, ImagePalette +from ._binary import i8 +from ._binary import i16be as i16 +from ._binary import i32be as i32 +from ._binary import si16be as si16 +from ._binary import si32be as si32 +from ._util import DeferredError + +MODES = { + # (photoshop mode, bits) -> (pil mode, required channels) + (0, 1): ("1", 1), + (0, 8): ("L", 1), + (1, 8): ("L", 1), + (2, 8): ("P", 1), + (3, 8): ("RGB", 3), + (4, 8): ("CMYK", 4), + (7, 8): ("L", 1), # FIXME: multilayer + (8, 8): ("L", 1), # duotone + (9, 8): ("LAB", 3), +} + + +# --------------------------------------------------------------------. +# read PSD images + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"8BPS") + + +## +# Image plugin for Photoshop images. + + +class PsdImageFile(ImageFile.ImageFile): + format = "PSD" + format_description = "Adobe Photoshop" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + read = self.fp.read + + # + # header + + s = read(26) + if not _accept(s) or i16(s, 4) != 1: + msg = "not a PSD file" + raise SyntaxError(msg) + + psd_bits = i16(s, 22) + psd_channels = i16(s, 12) + psd_mode = i16(s, 24) + + mode, channels = MODES[(psd_mode, psd_bits)] + + if channels > psd_channels: + msg = "not enough channels" + raise OSError(msg) + if mode == "RGB" and psd_channels == 4: + mode = "RGBA" + channels = 4 + + self._mode = mode + self._size = i32(s, 18), i32(s, 14) + + # + # color mode data + + size = i32(read(4)) + if size: + data = read(size) + if mode == "P" and size == 768: + self.palette = ImagePalette.raw("RGB;L", data) + + # + # image resources + + self.resources = [] + + size = i32(read(4)) + if size: + # load resources + end = self.fp.tell() + size + while self.fp.tell() < end: + read(4) # signature + id = i16(read(2)) + name = read(i8(read(1))) + if not (len(name) & 1): + read(1) # padding + data = read(i32(read(4))) + if len(data) & 1: + read(1) # padding + self.resources.append((id, name, data)) + if id == 1039: # ICC profile + self.info["icc_profile"] = data + + # + # layer and mask information + + self._layers_position = None + + size = i32(read(4)) + if size: + end = self.fp.tell() + size + size = i32(read(4)) + if size: + self._layers_position = self.fp.tell() + self._layers_size = size + self.fp.seek(end) + self._n_frames: int | None = None + + # + # image descriptor + + self.tile = _maketile(self.fp, mode, (0, 0) + self.size, channels) + + # keep the file open + self._fp = self.fp + self.frame = 1 + self._min_frame = 1 + + @cached_property + def layers( + self, + ) -> list[tuple[str, str, tuple[int, int, int, int], list[ImageFile._Tile]]]: + layers = [] + if self._layers_position is not None: + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self._fp.seek(self._layers_position) + _layer_data = io.BytesIO(ImageFile._safe_read(self._fp, self._layers_size)) + layers = _layerinfo(_layer_data, self._layers_size) + self._n_frames = len(layers) + return layers + + @property + def n_frames(self) -> int: + if self._n_frames is None: + self._n_frames = len(self.layers) + return self._n_frames + + @property + def is_animated(self) -> bool: + return len(self.layers) > 1 + + def seek(self, layer: int) -> None: + if not self._seek_check(layer): + return + if isinstance(self._fp, DeferredError): + raise self._fp.ex + + # seek to given layer (1..max) + _, mode, _, tile = self.layers[layer - 1] + self._mode = mode + self.tile = tile + self.frame = layer + self.fp = self._fp + + def tell(self) -> int: + # return layer number (0=image, 1..max=layers) + return self.frame + + +def _layerinfo( + fp: IO[bytes], ct_bytes: int +) -> list[tuple[str, str, tuple[int, int, int, int], list[ImageFile._Tile]]]: + # read layerinfo block + layers = [] + + def read(size: int) -> bytes: + return ImageFile._safe_read(fp, size) + + ct = si16(read(2)) + + # sanity check + if ct_bytes < (abs(ct) * 20): + msg = "Layer block too short for number of layers requested" + raise SyntaxError(msg) + + for _ in range(abs(ct)): + # bounding box + y0 = si32(read(4)) + x0 = si32(read(4)) + y1 = si32(read(4)) + x1 = si32(read(4)) + + # image info + bands = [] + ct_types = i16(read(2)) + if ct_types > 4: + fp.seek(ct_types * 6 + 12, io.SEEK_CUR) + size = i32(read(4)) + fp.seek(size, io.SEEK_CUR) + continue + + for _ in range(ct_types): + type = i16(read(2)) + + if type == 65535: + b = "A" + else: + b = "RGBA"[type] + + bands.append(b) + read(4) # size + + # figure out the image mode + bands.sort() + if bands == ["R"]: + mode = "L" + elif bands == ["B", "G", "R"]: + mode = "RGB" + elif bands == ["A", "B", "G", "R"]: + mode = "RGBA" + else: + mode = "" # unknown + + # skip over blend flags and extra information + read(12) # filler + name = "" + size = i32(read(4)) # length of the extra data field + if size: + data_end = fp.tell() + size + + length = i32(read(4)) + if length: + fp.seek(length - 16, io.SEEK_CUR) + + length = i32(read(4)) + if length: + fp.seek(length, io.SEEK_CUR) + + length = i8(read(1)) + if length: + # Don't know the proper encoding, + # Latin-1 should be a good guess + name = read(length).decode("latin-1", "replace") + + fp.seek(data_end) + layers.append((name, mode, (x0, y0, x1, y1))) + + # get tiles + layerinfo = [] + for i, (name, mode, bbox) in enumerate(layers): + tile = [] + for m in mode: + t = _maketile(fp, m, bbox, 1) + if t: + tile.extend(t) + layerinfo.append((name, mode, bbox, tile)) + + return layerinfo + + +def _maketile( + file: IO[bytes], mode: str, bbox: tuple[int, int, int, int], channels: int +) -> list[ImageFile._Tile]: + tiles = [] + read = file.read + + compression = i16(read(2)) + + xsize = bbox[2] - bbox[0] + ysize = bbox[3] - bbox[1] + + offset = file.tell() + + if compression == 0: + # + # raw compression + for channel in range(channels): + layer = mode[channel] + if mode == "CMYK": + layer += ";I" + tiles.append(ImageFile._Tile("raw", bbox, offset, layer)) + offset = offset + xsize * ysize + + elif compression == 1: + # + # packbits compression + i = 0 + bytecount = read(channels * ysize * 2) + offset = file.tell() + for channel in range(channels): + layer = mode[channel] + if mode == "CMYK": + layer += ";I" + tiles.append(ImageFile._Tile("packbits", bbox, offset, layer)) + for y in range(ysize): + offset = offset + i16(bytecount, i) + i += 2 + + file.seek(offset) + + if offset & 1: + read(1) # padding + + return tiles + + +# -------------------------------------------------------------------- +# registry + + +Image.register_open(PsdImageFile.format, PsdImageFile, _accept) + +Image.register_extension(PsdImageFile.format, ".psd") + +Image.register_mime(PsdImageFile.format, "image/vnd.adobe.photoshop") diff --git a/venv/lib/python3.10/site-packages/PIL/QoiImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/QoiImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..dba5d809fef75e281ac10f92f1868c58b1b4508c --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/QoiImagePlugin.py @@ -0,0 +1,234 @@ +# +# The Python Imaging Library. +# +# QOI support for PIL +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +from typing import IO + +from . import Image, ImageFile +from ._binary import i32be as i32 +from ._binary import o8 +from ._binary import o32be as o32 + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"qoif") + + +class QoiImageFile(ImageFile.ImageFile): + format = "QOI" + format_description = "Quite OK Image" + + def _open(self) -> None: + if not _accept(self.fp.read(4)): + msg = "not a QOI file" + raise SyntaxError(msg) + + self._size = i32(self.fp.read(4)), i32(self.fp.read(4)) + + channels = self.fp.read(1)[0] + self._mode = "RGB" if channels == 3 else "RGBA" + + self.fp.seek(1, os.SEEK_CUR) # colorspace + self.tile = [ImageFile._Tile("qoi", (0, 0) + self._size, self.fp.tell())] + + +class QoiDecoder(ImageFile.PyDecoder): + _pulls_fd = True + _previous_pixel: bytes | bytearray | None = None + _previously_seen_pixels: dict[int, bytes | bytearray] = {} + + def _add_to_previous_pixels(self, value: bytes | bytearray) -> None: + self._previous_pixel = value + + r, g, b, a = value + hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64 + self._previously_seen_pixels[hash_value] = value + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + + self._previously_seen_pixels = {} + self._previous_pixel = bytearray((0, 0, 0, 255)) + + data = bytearray() + bands = Image.getmodebands(self.mode) + dest_length = self.state.xsize * self.state.ysize * bands + while len(data) < dest_length: + byte = self.fd.read(1)[0] + value: bytes | bytearray + if byte == 0b11111110 and self._previous_pixel: # QOI_OP_RGB + value = bytearray(self.fd.read(3)) + self._previous_pixel[3:] + elif byte == 0b11111111: # QOI_OP_RGBA + value = self.fd.read(4) + else: + op = byte >> 6 + if op == 0: # QOI_OP_INDEX + op_index = byte & 0b00111111 + value = self._previously_seen_pixels.get( + op_index, bytearray((0, 0, 0, 0)) + ) + elif op == 1 and self._previous_pixel: # QOI_OP_DIFF + value = bytearray( + ( + (self._previous_pixel[0] + ((byte & 0b00110000) >> 4) - 2) + % 256, + (self._previous_pixel[1] + ((byte & 0b00001100) >> 2) - 2) + % 256, + (self._previous_pixel[2] + (byte & 0b00000011) - 2) % 256, + self._previous_pixel[3], + ) + ) + elif op == 2 and self._previous_pixel: # QOI_OP_LUMA + second_byte = self.fd.read(1)[0] + diff_green = (byte & 0b00111111) - 32 + diff_red = ((second_byte & 0b11110000) >> 4) - 8 + diff_blue = (second_byte & 0b00001111) - 8 + + value = bytearray( + tuple( + (self._previous_pixel[i] + diff_green + diff) % 256 + for i, diff in enumerate((diff_red, 0, diff_blue)) + ) + ) + value += self._previous_pixel[3:] + elif op == 3 and self._previous_pixel: # QOI_OP_RUN + run_length = (byte & 0b00111111) + 1 + value = self._previous_pixel + if bands == 3: + value = value[:3] + data += value * run_length + continue + self._add_to_previous_pixels(value) + + if bands == 3: + value = value[:3] + data += value + self.set_as_raw(data) + return -1, 0 + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode == "RGB": + channels = 3 + elif im.mode == "RGBA": + channels = 4 + else: + msg = "Unsupported QOI image mode" + raise ValueError(msg) + + colorspace = 0 if im.encoderinfo.get("colorspace") == "sRGB" else 1 + + fp.write(b"qoif") + fp.write(o32(im.size[0])) + fp.write(o32(im.size[1])) + fp.write(o8(channels)) + fp.write(o8(colorspace)) + + ImageFile._save(im, fp, [ImageFile._Tile("qoi", (0, 0) + im.size)]) + + +class QoiEncoder(ImageFile.PyEncoder): + _pushes_fd = True + _previous_pixel: tuple[int, int, int, int] | None = None + _previously_seen_pixels: dict[int, tuple[int, int, int, int]] = {} + _run = 0 + + def _write_run(self) -> bytes: + data = o8(0b11000000 | (self._run - 1)) # QOI_OP_RUN + self._run = 0 + return data + + def _delta(self, left: int, right: int) -> int: + result = (left - right) & 255 + if result >= 128: + result -= 256 + return result + + def encode(self, bufsize: int) -> tuple[int, int, bytes]: + assert self.im is not None + + self._previously_seen_pixels = {0: (0, 0, 0, 0)} + self._previous_pixel = (0, 0, 0, 255) + + data = bytearray() + w, h = self.im.size + bands = Image.getmodebands(self.mode) + + for y in range(h): + for x in range(w): + pixel = self.im.getpixel((x, y)) + if bands == 3: + pixel = (*pixel, 255) + + if pixel == self._previous_pixel: + self._run += 1 + if self._run == 62: + data += self._write_run() + else: + if self._run: + data += self._write_run() + + r, g, b, a = pixel + hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64 + if self._previously_seen_pixels.get(hash_value) == pixel: + data += o8(hash_value) # QOI_OP_INDEX + elif self._previous_pixel: + self._previously_seen_pixels[hash_value] = pixel + + prev_r, prev_g, prev_b, prev_a = self._previous_pixel + if prev_a == a: + delta_r = self._delta(r, prev_r) + delta_g = self._delta(g, prev_g) + delta_b = self._delta(b, prev_b) + + if ( + -2 <= delta_r < 2 + and -2 <= delta_g < 2 + and -2 <= delta_b < 2 + ): + data += o8( + 0b01000000 + | (delta_r + 2) << 4 + | (delta_g + 2) << 2 + | (delta_b + 2) + ) # QOI_OP_DIFF + else: + delta_gr = self._delta(delta_r, delta_g) + delta_gb = self._delta(delta_b, delta_g) + if ( + -8 <= delta_gr < 8 + and -32 <= delta_g < 32 + and -8 <= delta_gb < 8 + ): + data += o8( + 0b10000000 | (delta_g + 32) + ) # QOI_OP_LUMA + data += o8((delta_gr + 8) << 4 | (delta_gb + 8)) + else: + data += o8(0b11111110) # QOI_OP_RGB + data += bytes(pixel[:3]) + else: + data += o8(0b11111111) # QOI_OP_RGBA + data += bytes(pixel) + + self._previous_pixel = pixel + + if self._run: + data += self._write_run() + data += bytes((0, 0, 0, 0, 0, 0, 0, 1)) # padding + + return len(data), 0, data + + +Image.register_open(QoiImageFile.format, QoiImageFile, _accept) +Image.register_decoder("qoi", QoiDecoder) +Image.register_extension(QoiImageFile.format, ".qoi") + +Image.register_save(QoiImageFile.format, _save) +Image.register_encoder("qoi", QoiEncoder) diff --git a/venv/lib/python3.10/site-packages/PIL/SgiImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/SgiImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..853022150ae849e490378e41e831897050c207a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/SgiImagePlugin.py @@ -0,0 +1,231 @@ +# +# The Python Imaging Library. +# $Id$ +# +# SGI image file handling +# +# See "The SGI Image File Format (Draft version 0.97)", Paul Haeberli. +# +# +# +# History: +# 2017-22-07 mb Add RLE decompression +# 2016-16-10 mb Add save method without compression +# 1995-09-10 fl Created +# +# Copyright (c) 2016 by Mickael Bonfill. +# Copyright (c) 2008 by Karsten Hiddemann. +# Copyright (c) 1997 by Secret Labs AB. +# Copyright (c) 1995 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import os +import struct +from typing import IO + +from . import Image, ImageFile +from ._binary import i16be as i16 +from ._binary import o8 + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 2 and i16(prefix) == 474 + + +MODES = { + (1, 1, 1): "L", + (1, 2, 1): "L", + (2, 1, 1): "L;16B", + (2, 2, 1): "L;16B", + (1, 3, 3): "RGB", + (2, 3, 3): "RGB;16B", + (1, 3, 4): "RGBA", + (2, 3, 4): "RGBA;16B", +} + + +## +# Image plugin for SGI images. +class SgiImageFile(ImageFile.ImageFile): + format = "SGI" + format_description = "SGI Image File Format" + + def _open(self) -> None: + # HEAD + assert self.fp is not None + + headlen = 512 + s = self.fp.read(headlen) + + if not _accept(s): + msg = "Not an SGI image file" + raise ValueError(msg) + + # compression : verbatim or RLE + compression = s[2] + + # bpc : 1 or 2 bytes (8bits or 16bits) + bpc = s[3] + + # dimension : 1, 2 or 3 (depending on xsize, ysize and zsize) + dimension = i16(s, 4) + + # xsize : width + xsize = i16(s, 6) + + # ysize : height + ysize = i16(s, 8) + + # zsize : channels count + zsize = i16(s, 10) + + # determine mode from bits/zsize + try: + rawmode = MODES[(bpc, dimension, zsize)] + except KeyError: + msg = "Unsupported SGI image mode" + raise ValueError(msg) + + self._size = xsize, ysize + self._mode = rawmode.split(";")[0] + if self.mode == "RGB": + self.custom_mimetype = "image/rgb" + + # orientation -1 : scanlines begins at the bottom-left corner + orientation = -1 + + # decoder info + if compression == 0: + pagesize = xsize * ysize * bpc + if bpc == 2: + self.tile = [ + ImageFile._Tile( + "SGI16", + (0, 0) + self.size, + headlen, + (self.mode, 0, orientation), + ) + ] + else: + self.tile = [] + offset = headlen + for layer in self.mode: + self.tile.append( + ImageFile._Tile( + "raw", (0, 0) + self.size, offset, (layer, 0, orientation) + ) + ) + offset += pagesize + elif compression == 1: + self.tile = [ + ImageFile._Tile( + "sgi_rle", (0, 0) + self.size, headlen, (rawmode, orientation, bpc) + ) + ] + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode not in {"RGB", "RGBA", "L"}: + msg = "Unsupported SGI image mode" + raise ValueError(msg) + + # Get the keyword arguments + info = im.encoderinfo + + # Byte-per-pixel precision, 1 = 8bits per pixel + bpc = info.get("bpc", 1) + + if bpc not in (1, 2): + msg = "Unsupported number of bytes per pixel" + raise ValueError(msg) + + # Flip the image, since the origin of SGI file is the bottom-left corner + orientation = -1 + # Define the file as SGI File Format + magic_number = 474 + # Run-Length Encoding Compression - Unsupported at this time + rle = 0 + + # X Dimension = width / Y Dimension = height + x, y = im.size + # Z Dimension: Number of channels + z = len(im.mode) + # Number of dimensions (x,y,z) + if im.mode == "L": + dimension = 1 if y == 1 else 2 + else: + dimension = 3 + + # Minimum Byte value + pinmin = 0 + # Maximum Byte value (255 = 8bits per pixel) + pinmax = 255 + # Image name (79 characters max, truncated below in write) + img_name = os.path.splitext(os.path.basename(filename))[0] + if isinstance(img_name, str): + img_name = img_name.encode("ascii", "ignore") + # Standard representation of pixel in the file + colormap = 0 + fp.write(struct.pack(">h", magic_number)) + fp.write(o8(rle)) + fp.write(o8(bpc)) + fp.write(struct.pack(">H", dimension)) + fp.write(struct.pack(">H", x)) + fp.write(struct.pack(">H", y)) + fp.write(struct.pack(">H", z)) + fp.write(struct.pack(">l", pinmin)) + fp.write(struct.pack(">l", pinmax)) + fp.write(struct.pack("4s", b"")) # dummy + fp.write(struct.pack("79s", img_name)) # truncates to 79 chars + fp.write(struct.pack("s", b"")) # force null byte after img_name + fp.write(struct.pack(">l", colormap)) + fp.write(struct.pack("404s", b"")) # dummy + + rawmode = "L" + if bpc == 2: + rawmode = "L;16B" + + for channel in im.split(): + fp.write(channel.tobytes("raw", rawmode, 0, orientation)) + + if hasattr(fp, "flush"): + fp.flush() + + +class SGI16Decoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + assert self.im is not None + + rawmode, stride, orientation = self.args + pagesize = self.state.xsize * self.state.ysize + zsize = len(self.mode) + self.fd.seek(512) + + for band in range(zsize): + channel = Image.new("L", (self.state.xsize, self.state.ysize)) + channel.frombytes( + self.fd.read(2 * pagesize), "raw", "L;16B", stride, orientation + ) + self.im.putband(channel.im, band) + + return -1, 0 + + +# +# registry + + +Image.register_decoder("SGI16", SGI16Decoder) +Image.register_open(SgiImageFile.format, SgiImageFile, _accept) +Image.register_save(SgiImageFile.format, _save) +Image.register_mime(SgiImageFile.format, "image/sgi") + +Image.register_extensions(SgiImageFile.format, [".bw", ".rgb", ".rgba", ".sgi"]) + +# End of file diff --git a/venv/lib/python3.10/site-packages/PIL/SpiderImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/SpiderImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..868019e80a80cffc5e9f193ddbf96a0ba64ad9ea --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/SpiderImagePlugin.py @@ -0,0 +1,331 @@ +# +# The Python Imaging Library. +# +# SPIDER image file handling +# +# History: +# 2004-08-02 Created BB +# 2006-03-02 added save method +# 2006-03-13 added support for stack images +# +# Copyright (c) 2004 by Health Research Inc. (HRI) RENSSELAER, NY 12144. +# Copyright (c) 2004 by William Baxter. +# Copyright (c) 2004 by Secret Labs AB. +# Copyright (c) 2004 by Fredrik Lundh. +# + +## +# Image plugin for the Spider image format. This format is used +# by the SPIDER software, in processing image data from electron +# microscopy and tomography. +## + +# +# SpiderImagePlugin.py +# +# The Spider image format is used by SPIDER software, in processing +# image data from electron microscopy and tomography. +# +# Spider home page: +# https://spider.wadsworth.org/spider_doc/spider/docs/spider.html +# +# Details about the Spider image format: +# https://spider.wadsworth.org/spider_doc/spider/docs/image_doc.html +# +from __future__ import annotations + +import os +import struct +import sys +from typing import IO, Any, cast + +from . import Image, ImageFile +from ._util import DeferredError + +TYPE_CHECKING = False + + +def isInt(f: Any) -> int: + try: + i = int(f) + if f - i == 0: + return 1 + else: + return 0 + except (ValueError, OverflowError): + return 0 + + +iforms = [1, 3, -11, -12, -21, -22] + + +# There is no magic number to identify Spider files, so just check a +# series of header locations to see if they have reasonable values. +# Returns no. of bytes in the header, if it is a valid Spider header, +# otherwise returns 0 + + +def isSpiderHeader(t: tuple[float, ...]) -> int: + h = (99,) + t # add 1 value so can use spider header index start=1 + # header values 1,2,5,12,13,22,23 should be integers + for i in [1, 2, 5, 12, 13, 22, 23]: + if not isInt(h[i]): + return 0 + # check iform + iform = int(h[5]) + if iform not in iforms: + return 0 + # check other header values + labrec = int(h[13]) # no. records in file header + labbyt = int(h[22]) # total no. of bytes in header + lenbyt = int(h[23]) # record length in bytes + if labbyt != (labrec * lenbyt): + return 0 + # looks like a valid header + return labbyt + + +def isSpiderImage(filename: str) -> int: + with open(filename, "rb") as fp: + f = fp.read(92) # read 23 * 4 bytes + t = struct.unpack(">23f", f) # try big-endian first + hdrlen = isSpiderHeader(t) + if hdrlen == 0: + t = struct.unpack("<23f", f) # little-endian + hdrlen = isSpiderHeader(t) + return hdrlen + + +class SpiderImageFile(ImageFile.ImageFile): + format = "SPIDER" + format_description = "Spider 2D image" + _close_exclusive_fp_after_loading = False + + def _open(self) -> None: + # check header + n = 27 * 4 # read 27 float values + f = self.fp.read(n) + + try: + self.bigendian = 1 + t = struct.unpack(">27f", f) # try big-endian first + hdrlen = isSpiderHeader(t) + if hdrlen == 0: + self.bigendian = 0 + t = struct.unpack("<27f", f) # little-endian + hdrlen = isSpiderHeader(t) + if hdrlen == 0: + msg = "not a valid Spider file" + raise SyntaxError(msg) + except struct.error as e: + msg = "not a valid Spider file" + raise SyntaxError(msg) from e + + h = (99,) + t # add 1 value : spider header index starts at 1 + iform = int(h[5]) + if iform != 1: + msg = "not a Spider 2D image" + raise SyntaxError(msg) + + self._size = int(h[12]), int(h[2]) # size in pixels (width, height) + self.istack = int(h[24]) + self.imgnumber = int(h[27]) + + if self.istack == 0 and self.imgnumber == 0: + # stk=0, img=0: a regular 2D image + offset = hdrlen + self._nimages = 1 + elif self.istack > 0 and self.imgnumber == 0: + # stk>0, img=0: Opening the stack for the first time + self.imgbytes = int(h[12]) * int(h[2]) * 4 + self.hdrlen = hdrlen + self._nimages = int(h[26]) + # Point to the first image in the stack + offset = hdrlen * 2 + self.imgnumber = 1 + elif self.istack == 0 and self.imgnumber > 0: + # stk=0, img>0: an image within the stack + offset = hdrlen + self.stkoffset + self.istack = 2 # So Image knows it's still a stack + else: + msg = "inconsistent stack header values" + raise SyntaxError(msg) + + if self.bigendian: + self.rawmode = "F;32BF" + else: + self.rawmode = "F;32F" + self._mode = "F" + + self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, offset, self.rawmode)] + self._fp = self.fp # FIXME: hack + + @property + def n_frames(self) -> int: + return self._nimages + + @property + def is_animated(self) -> bool: + return self._nimages > 1 + + # 1st image index is zero (although SPIDER imgnumber starts at 1) + def tell(self) -> int: + if self.imgnumber < 1: + return 0 + else: + return self.imgnumber - 1 + + def seek(self, frame: int) -> None: + if self.istack == 0: + msg = "attempt to seek in a non-stack file" + raise EOFError(msg) + if not self._seek_check(frame): + return + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self.stkoffset = self.hdrlen + frame * (self.hdrlen + self.imgbytes) + self.fp = self._fp + self.fp.seek(self.stkoffset) + self._open() + + # returns a byte image after rescaling to 0..255 + def convert2byte(self, depth: int = 255) -> Image.Image: + extrema = self.getextrema() + assert isinstance(extrema[0], float) + minimum, maximum = cast(tuple[float, float], extrema) + m: float = 1 + if maximum != minimum: + m = depth / (maximum - minimum) + b = -m * minimum + return self.point(lambda i: i * m + b).convert("L") + + if TYPE_CHECKING: + from . import ImageTk + + # returns a ImageTk.PhotoImage object, after rescaling to 0..255 + def tkPhotoImage(self) -> ImageTk.PhotoImage: + from . import ImageTk + + return ImageTk.PhotoImage(self.convert2byte(), palette=256) + + +# -------------------------------------------------------------------- +# Image series + + +# given a list of filenames, return a list of images +def loadImageSeries(filelist: list[str] | None = None) -> list[Image.Image] | None: + """create a list of :py:class:`~PIL.Image.Image` objects for use in a montage""" + if filelist is None or len(filelist) < 1: + return None + + byte_imgs = [] + for img in filelist: + if not os.path.exists(img): + print(f"unable to find {img}") + continue + try: + with Image.open(img) as im: + assert isinstance(im, SpiderImageFile) + byte_im = im.convert2byte() + except Exception: + if not isSpiderImage(img): + print(f"{img} is not a Spider image file") + continue + byte_im.info["filename"] = img + byte_imgs.append(byte_im) + return byte_imgs + + +# -------------------------------------------------------------------- +# For saving images in Spider format + + +def makeSpiderHeader(im: Image.Image) -> list[bytes]: + nsam, nrow = im.size + lenbyt = nsam * 4 # There are labrec records in the header + labrec = int(1024 / lenbyt) + if 1024 % lenbyt != 0: + labrec += 1 + labbyt = labrec * lenbyt + nvalues = int(labbyt / 4) + if nvalues < 23: + return [] + + hdr = [0.0] * nvalues + + # NB these are Fortran indices + hdr[1] = 1.0 # nslice (=1 for an image) + hdr[2] = float(nrow) # number of rows per slice + hdr[3] = float(nrow) # number of records in the image + hdr[5] = 1.0 # iform for 2D image + hdr[12] = float(nsam) # number of pixels per line + hdr[13] = float(labrec) # number of records in file header + hdr[22] = float(labbyt) # total number of bytes in header + hdr[23] = float(lenbyt) # record length in bytes + + # adjust for Fortran indexing + hdr = hdr[1:] + hdr.append(0.0) + # pack binary data into a string + return [struct.pack("f", v) for v in hdr] + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode != "F": + im = im.convert("F") + + hdr = makeSpiderHeader(im) + if len(hdr) < 256: + msg = "Error creating Spider header" + raise OSError(msg) + + # write the SPIDER header + fp.writelines(hdr) + + rawmode = "F;32NF" # 32-bit native floating point + ImageFile._save(im, fp, [ImageFile._Tile("raw", (0, 0) + im.size, 0, rawmode)]) + + +def _save_spider(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + # get the filename extension and register it with Image + filename_ext = os.path.splitext(filename)[1] + ext = filename_ext.decode() if isinstance(filename_ext, bytes) else filename_ext + Image.register_extension(SpiderImageFile.format, ext) + _save(im, fp, filename) + + +# -------------------------------------------------------------------- + + +Image.register_open(SpiderImageFile.format, SpiderImageFile) +Image.register_save(SpiderImageFile.format, _save_spider) + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Syntax: python3 SpiderImagePlugin.py [infile] [outfile]") + sys.exit() + + filename = sys.argv[1] + if not isSpiderImage(filename): + print("input image must be in Spider format") + sys.exit() + + with Image.open(filename) as im: + print(f"image: {im}") + print(f"format: {im.format}") + print(f"size: {im.size}") + print(f"mode: {im.mode}") + print("max, min: ", end=" ") + print(im.getextrema()) + + if len(sys.argv) > 2: + outfile = sys.argv[2] + + # perform some image operation + im = im.transpose(Image.Transpose.FLIP_LEFT_RIGHT) + print( + f"saving a flipped version of {os.path.basename(filename)} " + f"as {outfile} " + ) + im.save(outfile, SpiderImageFile.format) diff --git a/venv/lib/python3.10/site-packages/PIL/SunImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/SunImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..8912379ea3e7801cdac9a557d2bc0c557bce8991 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/SunImagePlugin.py @@ -0,0 +1,145 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Sun image file handling +# +# History: +# 1995-09-10 fl Created +# 1996-05-28 fl Fixed 32-bit alignment +# 1998-12-29 fl Import ImagePalette module +# 2001-12-18 fl Fixed palette loading (from Jean-Claude Rimbault) +# +# Copyright (c) 1997-2001 by Secret Labs AB +# Copyright (c) 1995-1996 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +from . import Image, ImageFile, ImagePalette +from ._binary import i32be as i32 + + +def _accept(prefix: bytes) -> bool: + return len(prefix) >= 4 and i32(prefix) == 0x59A66A95 + + +## +# Image plugin for Sun raster files. + + +class SunImageFile(ImageFile.ImageFile): + format = "SUN" + format_description = "Sun Raster File" + + def _open(self) -> None: + # The Sun Raster file header is 32 bytes in length + # and has the following format: + + # typedef struct _SunRaster + # { + # DWORD MagicNumber; /* Magic (identification) number */ + # DWORD Width; /* Width of image in pixels */ + # DWORD Height; /* Height of image in pixels */ + # DWORD Depth; /* Number of bits per pixel */ + # DWORD Length; /* Size of image data in bytes */ + # DWORD Type; /* Type of raster file */ + # DWORD ColorMapType; /* Type of color map */ + # DWORD ColorMapLength; /* Size of the color map in bytes */ + # } SUNRASTER; + + assert self.fp is not None + + # HEAD + s = self.fp.read(32) + if not _accept(s): + msg = "not an SUN raster file" + raise SyntaxError(msg) + + offset = 32 + + self._size = i32(s, 4), i32(s, 8) + + depth = i32(s, 12) + # data_length = i32(s, 16) # unreliable, ignore. + file_type = i32(s, 20) + palette_type = i32(s, 24) # 0: None, 1: RGB, 2: Raw/arbitrary + palette_length = i32(s, 28) + + if depth == 1: + self._mode, rawmode = "1", "1;I" + elif depth == 4: + self._mode, rawmode = "L", "L;4" + elif depth == 8: + self._mode = rawmode = "L" + elif depth == 24: + if file_type == 3: + self._mode, rawmode = "RGB", "RGB" + else: + self._mode, rawmode = "RGB", "BGR" + elif depth == 32: + if file_type == 3: + self._mode, rawmode = "RGB", "RGBX" + else: + self._mode, rawmode = "RGB", "BGRX" + else: + msg = "Unsupported Mode/Bit Depth" + raise SyntaxError(msg) + + if palette_length: + if palette_length > 1024: + msg = "Unsupported Color Palette Length" + raise SyntaxError(msg) + + if palette_type != 1: + msg = "Unsupported Palette Type" + raise SyntaxError(msg) + + offset = offset + palette_length + self.palette = ImagePalette.raw("RGB;L", self.fp.read(palette_length)) + if self.mode == "L": + self._mode = "P" + rawmode = rawmode.replace("L", "P") + + # 16 bit boundaries on stride + stride = ((self.size[0] * depth + 15) // 16) * 2 + + # file type: Type is the version (or flavor) of the bitmap + # file. The following values are typically found in the Type + # field: + # 0000h Old + # 0001h Standard + # 0002h Byte-encoded + # 0003h RGB format + # 0004h TIFF format + # 0005h IFF format + # FFFFh Experimental + + # Old and standard are the same, except for the length tag. + # byte-encoded is run-length-encoded + # RGB looks similar to standard, but RGB byte order + # TIFF and IFF mean that they were converted from T/IFF + # Experimental means that it's something else. + # (https://www.fileformat.info/format/sunraster/egff.htm) + + if file_type in (0, 1, 3, 4, 5): + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, offset, (rawmode, stride)) + ] + elif file_type == 2: + self.tile = [ + ImageFile._Tile("sun_rle", (0, 0) + self.size, offset, rawmode) + ] + else: + msg = "Unsupported Sun Raster file type" + raise SyntaxError(msg) + + +# +# registry + + +Image.register_open(SunImageFile.format, SunImageFile, _accept) + +Image.register_extension(SunImageFile.format, ".ras") diff --git a/venv/lib/python3.10/site-packages/PIL/TarIO.py b/venv/lib/python3.10/site-packages/PIL/TarIO.py new file mode 100644 index 0000000000000000000000000000000000000000..86490a496f3f106fcc042c03fb235ed5fb41f3a7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/TarIO.py @@ -0,0 +1,61 @@ +# +# The Python Imaging Library. +# $Id$ +# +# read files from within a tar file +# +# History: +# 95-06-18 fl Created +# 96-05-28 fl Open files in binary mode +# +# Copyright (c) Secret Labs AB 1997. +# Copyright (c) Fredrik Lundh 1995-96. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io + +from . import ContainerIO + + +class TarIO(ContainerIO.ContainerIO[bytes]): + """A file object that provides read access to a given member of a TAR file.""" + + def __init__(self, tarfile: str, file: str) -> None: + """ + Create file object. + + :param tarfile: Name of TAR file. + :param file: Name of member file. + """ + self.fh = open(tarfile, "rb") + + while True: + s = self.fh.read(512) + if len(s) != 512: + self.fh.close() + + msg = "unexpected end of tar file" + raise OSError(msg) + + name = s[:100].decode("utf-8") + i = name.find("\0") + if i == 0: + self.fh.close() + + msg = "cannot find subfile" + raise OSError(msg) + if i > 0: + name = name[:i] + + size = int(s[124:135], 8) + + if file == name: + break + + self.fh.seek((size + 511) & (~511), io.SEEK_CUR) + + # Open region + super().__init__(self.fh, self.fh.tell(), size) diff --git a/venv/lib/python3.10/site-packages/PIL/TgaImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/TgaImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..90d5b5cf4ee17fc050784bf591adae3247407b56 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/TgaImagePlugin.py @@ -0,0 +1,264 @@ +# +# The Python Imaging Library. +# $Id$ +# +# TGA file handling +# +# History: +# 95-09-01 fl created (reads 24-bit files only) +# 97-01-04 fl support more TGA versions, including compressed images +# 98-07-04 fl fixed orientation and alpha layer bugs +# 98-09-11 fl fixed orientation for runlength decoder +# +# Copyright (c) Secret Labs AB 1997-98. +# Copyright (c) Fredrik Lundh 1995-97. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import warnings +from typing import IO + +from . import Image, ImageFile, ImagePalette +from ._binary import i16le as i16 +from ._binary import o8 +from ._binary import o16le as o16 + +# +# -------------------------------------------------------------------- +# Read RGA file + + +MODES = { + # map imagetype/depth to rawmode + (1, 8): "P", + (3, 1): "1", + (3, 8): "L", + (3, 16): "LA", + (2, 16): "BGRA;15Z", + (2, 24): "BGR", + (2, 32): "BGRA", +} + + +## +# Image plugin for Targa files. + + +class TgaImageFile(ImageFile.ImageFile): + format = "TGA" + format_description = "Targa" + + def _open(self) -> None: + # process header + assert self.fp is not None + + s = self.fp.read(18) + + id_len = s[0] + + colormaptype = s[1] + imagetype = s[2] + + depth = s[16] + + flags = s[17] + + self._size = i16(s, 12), i16(s, 14) + + # validate header fields + if ( + colormaptype not in (0, 1) + or self.size[0] <= 0 + or self.size[1] <= 0 + or depth not in (1, 8, 16, 24, 32) + ): + msg = "not a TGA file" + raise SyntaxError(msg) + + # image mode + if imagetype in (3, 11): + self._mode = "L" + if depth == 1: + self._mode = "1" # ??? + elif depth == 16: + self._mode = "LA" + elif imagetype in (1, 9): + self._mode = "P" if colormaptype else "L" + elif imagetype in (2, 10): + self._mode = "RGB" if depth == 24 else "RGBA" + else: + msg = "unknown TGA mode" + raise SyntaxError(msg) + + # orientation + orientation = flags & 0x30 + self._flip_horizontally = orientation in [0x10, 0x30] + if orientation in [0x20, 0x30]: + orientation = 1 + elif orientation in [0, 0x10]: + orientation = -1 + else: + msg = "unknown TGA orientation" + raise SyntaxError(msg) + + self.info["orientation"] = orientation + + if imagetype & 8: + self.info["compression"] = "tga_rle" + + if id_len: + self.info["id_section"] = self.fp.read(id_len) + + if colormaptype: + # read palette + start, size, mapdepth = i16(s, 3), i16(s, 5), s[7] + if mapdepth == 16: + self.palette = ImagePalette.raw( + "BGRA;15Z", bytes(2 * start) + self.fp.read(2 * size) + ) + self.palette.mode = "RGBA" + elif mapdepth == 24: + self.palette = ImagePalette.raw( + "BGR", bytes(3 * start) + self.fp.read(3 * size) + ) + elif mapdepth == 32: + self.palette = ImagePalette.raw( + "BGRA", bytes(4 * start) + self.fp.read(4 * size) + ) + else: + msg = "unknown TGA map depth" + raise SyntaxError(msg) + + # setup tile descriptor + try: + rawmode = MODES[(imagetype & 7, depth)] + if imagetype & 8: + # compressed + self.tile = [ + ImageFile._Tile( + "tga_rle", + (0, 0) + self.size, + self.fp.tell(), + (rawmode, orientation, depth), + ) + ] + else: + self.tile = [ + ImageFile._Tile( + "raw", + (0, 0) + self.size, + self.fp.tell(), + (rawmode, 0, orientation), + ) + ] + except KeyError: + pass # cannot decode + + def load_end(self) -> None: + if self._flip_horizontally: + self.im = self.im.transpose(Image.Transpose.FLIP_LEFT_RIGHT) + + +# +# -------------------------------------------------------------------- +# Write TGA file + + +SAVE = { + "1": ("1", 1, 0, 3), + "L": ("L", 8, 0, 3), + "LA": ("LA", 16, 0, 3), + "P": ("P", 8, 1, 1), + "RGB": ("BGR", 24, 0, 2), + "RGBA": ("BGRA", 32, 0, 2), +} + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + try: + rawmode, bits, colormaptype, imagetype = SAVE[im.mode] + except KeyError as e: + msg = f"cannot write mode {im.mode} as TGA" + raise OSError(msg) from e + + if "rle" in im.encoderinfo: + rle = im.encoderinfo["rle"] + else: + compression = im.encoderinfo.get("compression", im.info.get("compression")) + rle = compression == "tga_rle" + if rle: + imagetype += 8 + + id_section = im.encoderinfo.get("id_section", im.info.get("id_section", "")) + id_len = len(id_section) + if id_len > 255: + id_len = 255 + id_section = id_section[:255] + warnings.warn("id_section has been trimmed to 255 characters") + + if colormaptype: + palette = im.im.getpalette("RGB", "BGR") + colormaplength, colormapentry = len(palette) // 3, 24 + else: + colormaplength, colormapentry = 0, 0 + + if im.mode in ("LA", "RGBA"): + flags = 8 + else: + flags = 0 + + orientation = im.encoderinfo.get("orientation", im.info.get("orientation", -1)) + if orientation > 0: + flags = flags | 0x20 + + fp.write( + o8(id_len) + + o8(colormaptype) + + o8(imagetype) + + o16(0) # colormapfirst + + o16(colormaplength) + + o8(colormapentry) + + o16(0) + + o16(0) + + o16(im.size[0]) + + o16(im.size[1]) + + o8(bits) + + o8(flags) + ) + + if id_section: + fp.write(id_section) + + if colormaptype: + fp.write(palette) + + if rle: + ImageFile._save( + im, + fp, + [ImageFile._Tile("tga_rle", (0, 0) + im.size, 0, (rawmode, orientation))], + ) + else: + ImageFile._save( + im, + fp, + [ImageFile._Tile("raw", (0, 0) + im.size, 0, (rawmode, 0, orientation))], + ) + + # write targa version 2 footer + fp.write(b"\000" * 8 + b"TRUEVISION-XFILE." + b"\000") + + +# +# -------------------------------------------------------------------- +# Registry + + +Image.register_open(TgaImageFile.format, TgaImageFile) +Image.register_save(TgaImageFile.format, _save) + +Image.register_extensions(TgaImageFile.format, [".tga", ".icb", ".vda", ".vst"]) + +Image.register_mime(TgaImageFile.format, "image/x-tga") diff --git a/venv/lib/python3.10/site-packages/PIL/TiffImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/TiffImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..de2ce066ebf551335c572fdde8b9176585eb0f30 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/TiffImagePlugin.py @@ -0,0 +1,2338 @@ +# +# The Python Imaging Library. +# $Id$ +# +# TIFF file handling +# +# TIFF is a flexible, if somewhat aged, image file format originally +# defined by Aldus. Although TIFF supports a wide variety of pixel +# layouts and compression methods, the name doesn't really stand for +# "thousands of incompatible file formats," it just feels that way. +# +# To read TIFF data from a stream, the stream must be seekable. For +# progressive decoding, make sure to use TIFF files where the tag +# directory is placed first in the file. +# +# History: +# 1995-09-01 fl Created +# 1996-05-04 fl Handle JPEGTABLES tag +# 1996-05-18 fl Fixed COLORMAP support +# 1997-01-05 fl Fixed PREDICTOR support +# 1997-08-27 fl Added support for rational tags (from Perry Stoll) +# 1998-01-10 fl Fixed seek/tell (from Jan Blom) +# 1998-07-15 fl Use private names for internal variables +# 1999-06-13 fl Rewritten for PIL 1.0 (1.0) +# 2000-10-11 fl Additional fixes for Python 2.0 (1.1) +# 2001-04-17 fl Fixed rewind support (seek to frame 0) (1.2) +# 2001-05-12 fl Added write support for more tags (from Greg Couch) (1.3) +# 2001-12-18 fl Added workaround for broken Matrox library +# 2002-01-18 fl Don't mess up if photometric tag is missing (D. Alan Stewart) +# 2003-05-19 fl Check FILLORDER tag +# 2003-09-26 fl Added RGBa support +# 2004-02-24 fl Added DPI support; fixed rational write support +# 2005-02-07 fl Added workaround for broken Corel Draw 10 files +# 2006-01-09 fl Added support for float/double tags (from Russell Nelson) +# +# Copyright (c) 1997-2006 by Secret Labs AB. All rights reserved. +# Copyright (c) 1995-1997 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import io +import itertools +import logging +import math +import os +import struct +import warnings +from collections.abc import Callable, MutableMapping +from fractions import Fraction +from numbers import Number, Rational +from typing import IO, Any, cast + +from . import ExifTags, Image, ImageFile, ImageOps, ImagePalette, TiffTags +from ._binary import i16be as i16 +from ._binary import i32be as i32 +from ._binary import o8 +from ._util import DeferredError, is_path +from .TiffTags import TYPES + +TYPE_CHECKING = False +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import NoReturn + + from ._typing import Buffer, IntegralLike, StrOrBytesPath + +logger = logging.getLogger(__name__) + +# Set these to true to force use of libtiff for reading or writing. +READ_LIBTIFF = False +WRITE_LIBTIFF = False +STRIP_SIZE = 65536 + +II = b"II" # little-endian (Intel style) +MM = b"MM" # big-endian (Motorola style) + +# +# -------------------------------------------------------------------- +# Read TIFF files + +# a few tag names, just to make the code below a bit more readable +OSUBFILETYPE = 255 +IMAGEWIDTH = 256 +IMAGELENGTH = 257 +BITSPERSAMPLE = 258 +COMPRESSION = 259 +PHOTOMETRIC_INTERPRETATION = 262 +FILLORDER = 266 +IMAGEDESCRIPTION = 270 +STRIPOFFSETS = 273 +SAMPLESPERPIXEL = 277 +ROWSPERSTRIP = 278 +STRIPBYTECOUNTS = 279 +X_RESOLUTION = 282 +Y_RESOLUTION = 283 +PLANAR_CONFIGURATION = 284 +RESOLUTION_UNIT = 296 +TRANSFERFUNCTION = 301 +SOFTWARE = 305 +DATE_TIME = 306 +ARTIST = 315 +PREDICTOR = 317 +COLORMAP = 320 +TILEWIDTH = 322 +TILELENGTH = 323 +TILEOFFSETS = 324 +TILEBYTECOUNTS = 325 +SUBIFD = 330 +EXTRASAMPLES = 338 +SAMPLEFORMAT = 339 +JPEGTABLES = 347 +YCBCRSUBSAMPLING = 530 +REFERENCEBLACKWHITE = 532 +COPYRIGHT = 33432 +IPTC_NAA_CHUNK = 33723 # newsphoto properties +PHOTOSHOP_CHUNK = 34377 # photoshop properties +ICCPROFILE = 34675 +EXIFIFD = 34665 +XMP = 700 +JPEGQUALITY = 65537 # pseudo-tag by libtiff + +# https://github.com/imagej/ImageJA/blob/master/src/main/java/ij/io/TiffDecoder.java +IMAGEJ_META_DATA_BYTE_COUNTS = 50838 +IMAGEJ_META_DATA = 50839 + +COMPRESSION_INFO = { + # Compression => pil compression name + 1: "raw", + 2: "tiff_ccitt", + 3: "group3", + 4: "group4", + 5: "tiff_lzw", + 6: "tiff_jpeg", # obsolete + 7: "jpeg", + 8: "tiff_adobe_deflate", + 32771: "tiff_raw_16", # 16-bit padding + 32773: "packbits", + 32809: "tiff_thunderscan", + 32946: "tiff_deflate", + 34676: "tiff_sgilog", + 34677: "tiff_sgilog24", + 34925: "lzma", + 50000: "zstd", + 50001: "webp", +} + +COMPRESSION_INFO_REV = {v: k for k, v in COMPRESSION_INFO.items()} + +OPEN_INFO = { + # (ByteOrder, PhotoInterpretation, SampleFormat, FillOrder, BitsPerSample, + # ExtraSamples) => mode, rawmode + (II, 0, (1,), 1, (1,), ()): ("1", "1;I"), + (MM, 0, (1,), 1, (1,), ()): ("1", "1;I"), + (II, 0, (1,), 2, (1,), ()): ("1", "1;IR"), + (MM, 0, (1,), 2, (1,), ()): ("1", "1;IR"), + (II, 1, (1,), 1, (1,), ()): ("1", "1"), + (MM, 1, (1,), 1, (1,), ()): ("1", "1"), + (II, 1, (1,), 2, (1,), ()): ("1", "1;R"), + (MM, 1, (1,), 2, (1,), ()): ("1", "1;R"), + (II, 0, (1,), 1, (2,), ()): ("L", "L;2I"), + (MM, 0, (1,), 1, (2,), ()): ("L", "L;2I"), + (II, 0, (1,), 2, (2,), ()): ("L", "L;2IR"), + (MM, 0, (1,), 2, (2,), ()): ("L", "L;2IR"), + (II, 1, (1,), 1, (2,), ()): ("L", "L;2"), + (MM, 1, (1,), 1, (2,), ()): ("L", "L;2"), + (II, 1, (1,), 2, (2,), ()): ("L", "L;2R"), + (MM, 1, (1,), 2, (2,), ()): ("L", "L;2R"), + (II, 0, (1,), 1, (4,), ()): ("L", "L;4I"), + (MM, 0, (1,), 1, (4,), ()): ("L", "L;4I"), + (II, 0, (1,), 2, (4,), ()): ("L", "L;4IR"), + (MM, 0, (1,), 2, (4,), ()): ("L", "L;4IR"), + (II, 1, (1,), 1, (4,), ()): ("L", "L;4"), + (MM, 1, (1,), 1, (4,), ()): ("L", "L;4"), + (II, 1, (1,), 2, (4,), ()): ("L", "L;4R"), + (MM, 1, (1,), 2, (4,), ()): ("L", "L;4R"), + (II, 0, (1,), 1, (8,), ()): ("L", "L;I"), + (MM, 0, (1,), 1, (8,), ()): ("L", "L;I"), + (II, 0, (1,), 2, (8,), ()): ("L", "L;IR"), + (MM, 0, (1,), 2, (8,), ()): ("L", "L;IR"), + (II, 1, (1,), 1, (8,), ()): ("L", "L"), + (MM, 1, (1,), 1, (8,), ()): ("L", "L"), + (II, 1, (2,), 1, (8,), ()): ("L", "L"), + (MM, 1, (2,), 1, (8,), ()): ("L", "L"), + (II, 1, (1,), 2, (8,), ()): ("L", "L;R"), + (MM, 1, (1,), 2, (8,), ()): ("L", "L;R"), + (II, 1, (1,), 1, (12,), ()): ("I;16", "I;12"), + (II, 0, (1,), 1, (16,), ()): ("I;16", "I;16"), + (II, 1, (1,), 1, (16,), ()): ("I;16", "I;16"), + (MM, 1, (1,), 1, (16,), ()): ("I;16B", "I;16B"), + (II, 1, (1,), 2, (16,), ()): ("I;16", "I;16R"), + (II, 1, (2,), 1, (16,), ()): ("I", "I;16S"), + (MM, 1, (2,), 1, (16,), ()): ("I", "I;16BS"), + (II, 0, (3,), 1, (32,), ()): ("F", "F;32F"), + (MM, 0, (3,), 1, (32,), ()): ("F", "F;32BF"), + (II, 1, (1,), 1, (32,), ()): ("I", "I;32N"), + (II, 1, (2,), 1, (32,), ()): ("I", "I;32S"), + (MM, 1, (2,), 1, (32,), ()): ("I", "I;32BS"), + (II, 1, (3,), 1, (32,), ()): ("F", "F;32F"), + (MM, 1, (3,), 1, (32,), ()): ("F", "F;32BF"), + (II, 1, (1,), 1, (8, 8), (2,)): ("LA", "LA"), + (MM, 1, (1,), 1, (8, 8), (2,)): ("LA", "LA"), + (II, 2, (1,), 1, (8, 8, 8), ()): ("RGB", "RGB"), + (MM, 2, (1,), 1, (8, 8, 8), ()): ("RGB", "RGB"), + (II, 2, (1,), 2, (8, 8, 8), ()): ("RGB", "RGB;R"), + (MM, 2, (1,), 2, (8, 8, 8), ()): ("RGB", "RGB;R"), + (II, 2, (1,), 1, (8, 8, 8, 8), ()): ("RGBA", "RGBA"), # missing ExtraSamples + (MM, 2, (1,), 1, (8, 8, 8, 8), ()): ("RGBA", "RGBA"), # missing ExtraSamples + (II, 2, (1,), 1, (8, 8, 8, 8), (0,)): ("RGB", "RGBX"), + (MM, 2, (1,), 1, (8, 8, 8, 8), (0,)): ("RGB", "RGBX"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8), (0, 0)): ("RGB", "RGBXX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8), (0, 0)): ("RGB", "RGBXX"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (0, 0, 0)): ("RGB", "RGBXXX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (0, 0, 0)): ("RGB", "RGBXXX"), + (II, 2, (1,), 1, (8, 8, 8, 8), (1,)): ("RGBA", "RGBa"), + (MM, 2, (1,), 1, (8, 8, 8, 8), (1,)): ("RGBA", "RGBa"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8), (1, 0)): ("RGBA", "RGBaX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8), (1, 0)): ("RGBA", "RGBaX"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (1, 0, 0)): ("RGBA", "RGBaXX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (1, 0, 0)): ("RGBA", "RGBaXX"), + (II, 2, (1,), 1, (8, 8, 8, 8), (2,)): ("RGBA", "RGBA"), + (MM, 2, (1,), 1, (8, 8, 8, 8), (2,)): ("RGBA", "RGBA"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8), (2, 0)): ("RGBA", "RGBAX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8), (2, 0)): ("RGBA", "RGBAX"), + (II, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (2, 0, 0)): ("RGBA", "RGBAXX"), + (MM, 2, (1,), 1, (8, 8, 8, 8, 8, 8), (2, 0, 0)): ("RGBA", "RGBAXX"), + (II, 2, (1,), 1, (8, 8, 8, 8), (999,)): ("RGBA", "RGBA"), # Corel Draw 10 + (MM, 2, (1,), 1, (8, 8, 8, 8), (999,)): ("RGBA", "RGBA"), # Corel Draw 10 + (II, 2, (1,), 1, (16, 16, 16), ()): ("RGB", "RGB;16L"), + (MM, 2, (1,), 1, (16, 16, 16), ()): ("RGB", "RGB;16B"), + (II, 2, (1,), 1, (16, 16, 16, 16), ()): ("RGBA", "RGBA;16L"), + (MM, 2, (1,), 1, (16, 16, 16, 16), ()): ("RGBA", "RGBA;16B"), + (II, 2, (1,), 1, (16, 16, 16, 16), (0,)): ("RGB", "RGBX;16L"), + (MM, 2, (1,), 1, (16, 16, 16, 16), (0,)): ("RGB", "RGBX;16B"), + (II, 2, (1,), 1, (16, 16, 16, 16), (1,)): ("RGBA", "RGBa;16L"), + (MM, 2, (1,), 1, (16, 16, 16, 16), (1,)): ("RGBA", "RGBa;16B"), + (II, 2, (1,), 1, (16, 16, 16, 16), (2,)): ("RGBA", "RGBA;16L"), + (MM, 2, (1,), 1, (16, 16, 16, 16), (2,)): ("RGBA", "RGBA;16B"), + (II, 3, (1,), 1, (1,), ()): ("P", "P;1"), + (MM, 3, (1,), 1, (1,), ()): ("P", "P;1"), + (II, 3, (1,), 2, (1,), ()): ("P", "P;1R"), + (MM, 3, (1,), 2, (1,), ()): ("P", "P;1R"), + (II, 3, (1,), 1, (2,), ()): ("P", "P;2"), + (MM, 3, (1,), 1, (2,), ()): ("P", "P;2"), + (II, 3, (1,), 2, (2,), ()): ("P", "P;2R"), + (MM, 3, (1,), 2, (2,), ()): ("P", "P;2R"), + (II, 3, (1,), 1, (4,), ()): ("P", "P;4"), + (MM, 3, (1,), 1, (4,), ()): ("P", "P;4"), + (II, 3, (1,), 2, (4,), ()): ("P", "P;4R"), + (MM, 3, (1,), 2, (4,), ()): ("P", "P;4R"), + (II, 3, (1,), 1, (8,), ()): ("P", "P"), + (MM, 3, (1,), 1, (8,), ()): ("P", "P"), + (II, 3, (1,), 1, (8, 8), (0,)): ("P", "PX"), + (MM, 3, (1,), 1, (8, 8), (0,)): ("P", "PX"), + (II, 3, (1,), 1, (8, 8), (2,)): ("PA", "PA"), + (MM, 3, (1,), 1, (8, 8), (2,)): ("PA", "PA"), + (II, 3, (1,), 2, (8,), ()): ("P", "P;R"), + (MM, 3, (1,), 2, (8,), ()): ("P", "P;R"), + (II, 5, (1,), 1, (8, 8, 8, 8), ()): ("CMYK", "CMYK"), + (MM, 5, (1,), 1, (8, 8, 8, 8), ()): ("CMYK", "CMYK"), + (II, 5, (1,), 1, (8, 8, 8, 8, 8), (0,)): ("CMYK", "CMYKX"), + (MM, 5, (1,), 1, (8, 8, 8, 8, 8), (0,)): ("CMYK", "CMYKX"), + (II, 5, (1,), 1, (8, 8, 8, 8, 8, 8), (0, 0)): ("CMYK", "CMYKXX"), + (MM, 5, (1,), 1, (8, 8, 8, 8, 8, 8), (0, 0)): ("CMYK", "CMYKXX"), + (II, 5, (1,), 1, (16, 16, 16, 16), ()): ("CMYK", "CMYK;16L"), + (MM, 5, (1,), 1, (16, 16, 16, 16), ()): ("CMYK", "CMYK;16B"), + (II, 6, (1,), 1, (8,), ()): ("L", "L"), + (MM, 6, (1,), 1, (8,), ()): ("L", "L"), + # JPEG compressed images handled by LibTiff and auto-converted to RGBX + # Minimal Baseline TIFF requires YCbCr images to have 3 SamplesPerPixel + (II, 6, (1,), 1, (8, 8, 8), ()): ("RGB", "RGBX"), + (MM, 6, (1,), 1, (8, 8, 8), ()): ("RGB", "RGBX"), + (II, 8, (1,), 1, (8, 8, 8), ()): ("LAB", "LAB"), + (MM, 8, (1,), 1, (8, 8, 8), ()): ("LAB", "LAB"), +} + +MAX_SAMPLESPERPIXEL = max(len(key_tp[4]) for key_tp in OPEN_INFO) + +PREFIXES = [ + b"MM\x00\x2a", # Valid TIFF header with big-endian byte order + b"II\x2a\x00", # Valid TIFF header with little-endian byte order + b"MM\x2a\x00", # Invalid TIFF header, assume big-endian + b"II\x00\x2a", # Invalid TIFF header, assume little-endian + b"MM\x00\x2b", # BigTIFF with big-endian byte order + b"II\x2b\x00", # BigTIFF with little-endian byte order +] + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(tuple(PREFIXES)) + + +def _limit_rational( + val: float | Fraction | IFDRational, max_val: int +) -> tuple[IntegralLike, IntegralLike]: + inv = abs(val) > 1 + n_d = IFDRational(1 / val if inv else val).limit_rational(max_val) + return n_d[::-1] if inv else n_d + + +def _limit_signed_rational( + val: IFDRational, max_val: int, min_val: int +) -> tuple[IntegralLike, IntegralLike]: + frac = Fraction(val) + n_d: tuple[IntegralLike, IntegralLike] = frac.numerator, frac.denominator + + if min(float(i) for i in n_d) < min_val: + n_d = _limit_rational(val, abs(min_val)) + + n_d_float = tuple(float(i) for i in n_d) + if max(n_d_float) > max_val: + n_d = _limit_rational(n_d_float[0] / n_d_float[1], max_val) + + return n_d + + +## +# Wrapper for TIFF IFDs. + +_load_dispatch = {} +_write_dispatch = {} + + +def _delegate(op: str) -> Any: + def delegate( + self: IFDRational, *args: tuple[float, ...] + ) -> bool | float | Fraction: + return getattr(self._val, op)(*args) + + return delegate + + +class IFDRational(Rational): + """Implements a rational class where 0/0 is a legal value to match + the in the wild use of exif rationals. + + e.g., DigitalZoomRatio - 0.00/0.00 indicates that no digital zoom was used + """ + + """ If the denominator is 0, store this as a float('nan'), otherwise store + as a fractions.Fraction(). Delegate as appropriate + + """ + + __slots__ = ("_numerator", "_denominator", "_val") + + def __init__( + self, value: float | Fraction | IFDRational, denominator: int = 1 + ) -> None: + """ + :param value: either an integer numerator, a + float/rational/other number, or an IFDRational + :param denominator: Optional integer denominator + """ + self._val: Fraction | float + if isinstance(value, IFDRational): + self._numerator = value.numerator + self._denominator = value.denominator + self._val = value._val + return + + if isinstance(value, Fraction): + self._numerator = value.numerator + self._denominator = value.denominator + else: + if TYPE_CHECKING: + self._numerator = cast(IntegralLike, value) + else: + self._numerator = value + self._denominator = denominator + + if denominator == 0: + self._val = float("nan") + elif denominator == 1: + self._val = Fraction(value) + elif int(value) == value: + self._val = Fraction(int(value), denominator) + else: + self._val = Fraction(value / denominator) + + @property + def numerator(self) -> IntegralLike: + return self._numerator + + @property + def denominator(self) -> int: + return self._denominator + + def limit_rational(self, max_denominator: int) -> tuple[IntegralLike, int]: + """ + + :param max_denominator: Integer, the maximum denominator value + :returns: Tuple of (numerator, denominator) + """ + + if self.denominator == 0: + return self.numerator, self.denominator + + assert isinstance(self._val, Fraction) + f = self._val.limit_denominator(max_denominator) + return f.numerator, f.denominator + + def __repr__(self) -> str: + return str(float(self._val)) + + def __hash__(self) -> int: # type: ignore[override] + return self._val.__hash__() + + def __eq__(self, other: object) -> bool: + val = self._val + if isinstance(other, IFDRational): + other = other._val + if isinstance(other, float): + val = float(val) + return val == other + + def __getstate__(self) -> list[float | Fraction | IntegralLike]: + return [self._val, self._numerator, self._denominator] + + def __setstate__(self, state: list[float | Fraction | IntegralLike]) -> None: + IFDRational.__init__(self, 0) + _val, _numerator, _denominator = state + assert isinstance(_val, (float, Fraction)) + self._val = _val + if TYPE_CHECKING: + self._numerator = cast(IntegralLike, _numerator) + else: + self._numerator = _numerator + assert isinstance(_denominator, int) + self._denominator = _denominator + + """ a = ['add','radd', 'sub', 'rsub', 'mul', 'rmul', + 'truediv', 'rtruediv', 'floordiv', 'rfloordiv', + 'mod','rmod', 'pow','rpow', 'pos', 'neg', + 'abs', 'trunc', 'lt', 'gt', 'le', 'ge', 'bool', + 'ceil', 'floor', 'round'] + print("\n".join("__%s__ = _delegate('__%s__')" % (s,s) for s in a)) + """ + + __add__ = _delegate("__add__") + __radd__ = _delegate("__radd__") + __sub__ = _delegate("__sub__") + __rsub__ = _delegate("__rsub__") + __mul__ = _delegate("__mul__") + __rmul__ = _delegate("__rmul__") + __truediv__ = _delegate("__truediv__") + __rtruediv__ = _delegate("__rtruediv__") + __floordiv__ = _delegate("__floordiv__") + __rfloordiv__ = _delegate("__rfloordiv__") + __mod__ = _delegate("__mod__") + __rmod__ = _delegate("__rmod__") + __pow__ = _delegate("__pow__") + __rpow__ = _delegate("__rpow__") + __pos__ = _delegate("__pos__") + __neg__ = _delegate("__neg__") + __abs__ = _delegate("__abs__") + __trunc__ = _delegate("__trunc__") + __lt__ = _delegate("__lt__") + __gt__ = _delegate("__gt__") + __le__ = _delegate("__le__") + __ge__ = _delegate("__ge__") + __bool__ = _delegate("__bool__") + __ceil__ = _delegate("__ceil__") + __floor__ = _delegate("__floor__") + __round__ = _delegate("__round__") + # Python >= 3.11 + if hasattr(Fraction, "__int__"): + __int__ = _delegate("__int__") + + +_LoaderFunc = Callable[["ImageFileDirectory_v2", bytes, bool], Any] + + +def _register_loader(idx: int, size: int) -> Callable[[_LoaderFunc], _LoaderFunc]: + def decorator(func: _LoaderFunc) -> _LoaderFunc: + from .TiffTags import TYPES + + if func.__name__.startswith("load_"): + TYPES[idx] = func.__name__[5:].replace("_", " ") + _load_dispatch[idx] = size, func # noqa: F821 + return func + + return decorator + + +def _register_writer(idx: int) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + _write_dispatch[idx] = func # noqa: F821 + return func + + return decorator + + +def _register_basic(idx_fmt_name: tuple[int, str, str]) -> None: + from .TiffTags import TYPES + + idx, fmt, name = idx_fmt_name + TYPES[idx] = name + size = struct.calcsize(f"={fmt}") + + def basic_handler( + self: ImageFileDirectory_v2, data: bytes, legacy_api: bool = True + ) -> tuple[Any, ...]: + return self._unpack(f"{len(data) // size}{fmt}", data) + + _load_dispatch[idx] = size, basic_handler # noqa: F821 + _write_dispatch[idx] = lambda self, *values: ( # noqa: F821 + b"".join(self._pack(fmt, value) for value in values) + ) + + +if TYPE_CHECKING: + _IFDv2Base = MutableMapping[int, Any] +else: + _IFDv2Base = MutableMapping + + +class ImageFileDirectory_v2(_IFDv2Base): + """This class represents a TIFF tag directory. To speed things up, we + don't decode tags unless they're asked for. + + Exposes a dictionary interface of the tags in the directory:: + + ifd = ImageFileDirectory_v2() + ifd[key] = 'Some Data' + ifd.tagtype[key] = TiffTags.ASCII + print(ifd[key]) + 'Some Data' + + Individual values are returned as the strings or numbers, sequences are + returned as tuples of the values. + + The tiff metadata type of each item is stored in a dictionary of + tag types in + :attr:`~PIL.TiffImagePlugin.ImageFileDirectory_v2.tagtype`. The types + are read from a tiff file, guessed from the type added, or added + manually. + + Data Structures: + + * ``self.tagtype = {}`` + + * Key: numerical TIFF tag number + * Value: integer corresponding to the data type from + :py:data:`.TiffTags.TYPES` + + .. versionadded:: 3.0.0 + + 'Internal' data structures: + + * ``self._tags_v2 = {}`` + + * Key: numerical TIFF tag number + * Value: decoded data, as tuple for multiple values + + * ``self._tagdata = {}`` + + * Key: numerical TIFF tag number + * Value: undecoded byte string from file + + * ``self._tags_v1 = {}`` + + * Key: numerical TIFF tag number + * Value: decoded data in the v1 format + + Tags will be found in the private attributes ``self._tagdata``, and in + ``self._tags_v2`` once decoded. + + ``self.legacy_api`` is a value for internal use, and shouldn't be changed + from outside code. In cooperation with + :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v1`, if ``legacy_api`` + is true, then decoded tags will be populated into both ``_tags_v1`` and + ``_tags_v2``. ``_tags_v2`` will be used if this IFD is used in the TIFF + save routine. Tags should be read from ``_tags_v1`` if + ``legacy_api == true``. + + """ + + _load_dispatch: dict[int, tuple[int, _LoaderFunc]] = {} + _write_dispatch: dict[int, Callable[..., Any]] = {} + + def __init__( + self, + ifh: bytes = b"II\x2a\x00\x00\x00\x00\x00", + prefix: bytes | None = None, + group: int | None = None, + ) -> None: + """Initialize an ImageFileDirectory. + + To construct an ImageFileDirectory from a real file, pass the 8-byte + magic header to the constructor. To only set the endianness, pass it + as the 'prefix' keyword argument. + + :param ifh: One of the accepted magic headers (cf. PREFIXES); also sets + endianness. + :param prefix: Override the endianness of the file. + """ + if not _accept(ifh): + msg = f"not a TIFF file (header {repr(ifh)} not valid)" + raise SyntaxError(msg) + self._prefix = prefix if prefix is not None else ifh[:2] + if self._prefix == MM: + self._endian = ">" + elif self._prefix == II: + self._endian = "<" + else: + msg = "not a TIFF IFD" + raise SyntaxError(msg) + self._bigtiff = ifh[2] == 43 + self.group = group + self.tagtype: dict[int, int] = {} + """ Dictionary of tag types """ + self.reset() + self.next = ( + self._unpack("Q", ifh[8:])[0] + if self._bigtiff + else self._unpack("L", ifh[4:])[0] + ) + self._legacy_api = False + + prefix = property(lambda self: self._prefix) + offset = property(lambda self: self._offset) + + @property + def legacy_api(self) -> bool: + return self._legacy_api + + @legacy_api.setter + def legacy_api(self, value: bool) -> NoReturn: + msg = "Not allowing setting of legacy api" + raise Exception(msg) + + def reset(self) -> None: + self._tags_v1: dict[int, Any] = {} # will remain empty if legacy_api is false + self._tags_v2: dict[int, Any] = {} # main tag storage + self._tagdata: dict[int, bytes] = {} + self.tagtype = {} # added 2008-06-05 by Florian Hoech + self._next = None + self._offset: int | None = None + + def __str__(self) -> str: + return str(dict(self)) + + def named(self) -> dict[str, Any]: + """ + :returns: dict of name|key: value + + Returns the complete tag dictionary, with named tags where possible. + """ + return { + TiffTags.lookup(code, self.group).name: value + for code, value in self.items() + } + + def __len__(self) -> int: + return len(set(self._tagdata) | set(self._tags_v2)) + + def __getitem__(self, tag: int) -> Any: + if tag not in self._tags_v2: # unpack on the fly + data = self._tagdata[tag] + typ = self.tagtype[tag] + size, handler = self._load_dispatch[typ] + self[tag] = handler(self, data, self.legacy_api) # check type + val = self._tags_v2[tag] + if self.legacy_api and not isinstance(val, (tuple, bytes)): + val = (val,) + return val + + def __contains__(self, tag: object) -> bool: + return tag in self._tags_v2 or tag in self._tagdata + + def __setitem__(self, tag: int, value: Any) -> None: + self._setitem(tag, value, self.legacy_api) + + def _setitem(self, tag: int, value: Any, legacy_api: bool) -> None: + basetypes = (Number, bytes, str) + + info = TiffTags.lookup(tag, self.group) + values = [value] if isinstance(value, basetypes) else value + + if tag not in self.tagtype: + if info.type: + self.tagtype[tag] = info.type + else: + self.tagtype[tag] = TiffTags.UNDEFINED + if all(isinstance(v, IFDRational) for v in values): + for v in values: + assert isinstance(v, IFDRational) + if v < 0: + self.tagtype[tag] = TiffTags.SIGNED_RATIONAL + break + else: + self.tagtype[tag] = TiffTags.RATIONAL + elif all(isinstance(v, int) for v in values): + short = True + signed_short = True + long = True + for v in values: + assert isinstance(v, int) + if short and not (0 <= v < 2**16): + short = False + if signed_short and not (-(2**15) < v < 2**15): + signed_short = False + if long and v < 0: + long = False + if short: + self.tagtype[tag] = TiffTags.SHORT + elif signed_short: + self.tagtype[tag] = TiffTags.SIGNED_SHORT + elif long: + self.tagtype[tag] = TiffTags.LONG + else: + self.tagtype[tag] = TiffTags.SIGNED_LONG + elif all(isinstance(v, float) for v in values): + self.tagtype[tag] = TiffTags.DOUBLE + elif all(isinstance(v, str) for v in values): + self.tagtype[tag] = TiffTags.ASCII + elif all(isinstance(v, bytes) for v in values): + self.tagtype[tag] = TiffTags.BYTE + + if self.tagtype[tag] == TiffTags.UNDEFINED: + values = [ + v.encode("ascii", "replace") if isinstance(v, str) else v + for v in values + ] + elif self.tagtype[tag] == TiffTags.RATIONAL: + values = [float(v) if isinstance(v, int) else v for v in values] + + is_ifd = self.tagtype[tag] == TiffTags.LONG and isinstance(values, dict) + if not is_ifd: + values = tuple( + info.cvt_enum(value) if isinstance(value, str) else value + for value in values + ) + + dest = self._tags_v1 if legacy_api else self._tags_v2 + + # Three branches: + # Spec'd length == 1, Actual length 1, store as element + # Spec'd length == 1, Actual > 1, Warn and truncate. Formerly barfed. + # No Spec, Actual length 1, Formerly (<4.2) returned a 1 element tuple. + # Don't mess with the legacy api, since it's frozen. + if not is_ifd and ( + (info.length == 1) + or self.tagtype[tag] == TiffTags.BYTE + or (info.length is None and len(values) == 1 and not legacy_api) + ): + # Don't mess with the legacy api, since it's frozen. + if legacy_api and self.tagtype[tag] in [ + TiffTags.RATIONAL, + TiffTags.SIGNED_RATIONAL, + ]: # rationals + values = (values,) + try: + (dest[tag],) = values + except ValueError: + # We've got a builtin tag with 1 expected entry + warnings.warn( + f"Metadata Warning, tag {tag} had too many entries: " + f"{len(values)}, expected 1" + ) + dest[tag] = values[0] + + else: + # Spec'd length > 1 or undefined + # Unspec'd, and length > 1 + dest[tag] = values + + def __delitem__(self, tag: int) -> None: + self._tags_v2.pop(tag, None) + self._tags_v1.pop(tag, None) + self._tagdata.pop(tag, None) + + def __iter__(self) -> Iterator[int]: + return iter(set(self._tagdata) | set(self._tags_v2)) + + def _unpack(self, fmt: str, data: bytes) -> tuple[Any, ...]: + return struct.unpack(self._endian + fmt, data) + + def _pack(self, fmt: str, *values: Any) -> bytes: + return struct.pack(self._endian + fmt, *values) + + list( + map( + _register_basic, + [ + (TiffTags.SHORT, "H", "short"), + (TiffTags.LONG, "L", "long"), + (TiffTags.SIGNED_BYTE, "b", "signed byte"), + (TiffTags.SIGNED_SHORT, "h", "signed short"), + (TiffTags.SIGNED_LONG, "l", "signed long"), + (TiffTags.FLOAT, "f", "float"), + (TiffTags.DOUBLE, "d", "double"), + (TiffTags.IFD, "L", "long"), + (TiffTags.LONG8, "Q", "long8"), + ], + ) + ) + + @_register_loader(1, 1) # Basic type, except for the legacy API. + def load_byte(self, data: bytes, legacy_api: bool = True) -> bytes: + return data + + @_register_writer(1) # Basic type, except for the legacy API. + def write_byte(self, data: bytes | int | IFDRational) -> bytes: + if isinstance(data, IFDRational): + data = int(data) + if isinstance(data, int): + data = bytes((data,)) + return data + + @_register_loader(2, 1) + def load_string(self, data: bytes, legacy_api: bool = True) -> str: + if data.endswith(b"\0"): + data = data[:-1] + return data.decode("latin-1", "replace") + + @_register_writer(2) + def write_string(self, value: str | bytes | int) -> bytes: + # remerge of https://github.com/python-pillow/Pillow/pull/1416 + if isinstance(value, int): + value = str(value) + if not isinstance(value, bytes): + value = value.encode("ascii", "replace") + return value + b"\0" + + @_register_loader(5, 8) + def load_rational( + self, data: bytes, legacy_api: bool = True + ) -> tuple[tuple[int, int] | IFDRational, ...]: + vals = self._unpack(f"{len(data) // 4}L", data) + + def combine(a: int, b: int) -> tuple[int, int] | IFDRational: + return (a, b) if legacy_api else IFDRational(a, b) + + return tuple(combine(num, denom) for num, denom in zip(vals[::2], vals[1::2])) + + @_register_writer(5) + def write_rational(self, *values: IFDRational) -> bytes: + return b"".join( + self._pack("2L", *_limit_rational(frac, 2**32 - 1)) for frac in values + ) + + @_register_loader(7, 1) + def load_undefined(self, data: bytes, legacy_api: bool = True) -> bytes: + return data + + @_register_writer(7) + def write_undefined(self, value: bytes | int | IFDRational) -> bytes: + if isinstance(value, IFDRational): + value = int(value) + if isinstance(value, int): + value = str(value).encode("ascii", "replace") + return value + + @_register_loader(10, 8) + def load_signed_rational( + self, data: bytes, legacy_api: bool = True + ) -> tuple[tuple[int, int] | IFDRational, ...]: + vals = self._unpack(f"{len(data) // 4}l", data) + + def combine(a: int, b: int) -> tuple[int, int] | IFDRational: + return (a, b) if legacy_api else IFDRational(a, b) + + return tuple(combine(num, denom) for num, denom in zip(vals[::2], vals[1::2])) + + @_register_writer(10) + def write_signed_rational(self, *values: IFDRational) -> bytes: + return b"".join( + self._pack("2l", *_limit_signed_rational(frac, 2**31 - 1, -(2**31))) + for frac in values + ) + + def _ensure_read(self, fp: IO[bytes], size: int) -> bytes: + ret = fp.read(size) + if len(ret) != size: + msg = ( + "Corrupt EXIF data. " + f"Expecting to read {size} bytes but only got {len(ret)}. " + ) + raise OSError(msg) + return ret + + def load(self, fp: IO[bytes]) -> None: + self.reset() + self._offset = fp.tell() + + try: + tag_count = ( + self._unpack("Q", self._ensure_read(fp, 8)) + if self._bigtiff + else self._unpack("H", self._ensure_read(fp, 2)) + )[0] + for i in range(tag_count): + tag, typ, count, data = ( + self._unpack("HHQ8s", self._ensure_read(fp, 20)) + if self._bigtiff + else self._unpack("HHL4s", self._ensure_read(fp, 12)) + ) + + tagname = TiffTags.lookup(tag, self.group).name + typname = TYPES.get(typ, "unknown") + msg = f"tag: {tagname} ({tag}) - type: {typname} ({typ})" + + try: + unit_size, handler = self._load_dispatch[typ] + except KeyError: + logger.debug("%s - unsupported type %s", msg, typ) + continue # ignore unsupported type + size = count * unit_size + if size > (8 if self._bigtiff else 4): + here = fp.tell() + (offset,) = self._unpack("Q" if self._bigtiff else "L", data) + msg += f" Tag Location: {here} - Data Location: {offset}" + fp.seek(offset) + data = ImageFile._safe_read(fp, size) + fp.seek(here) + else: + data = data[:size] + + if len(data) != size: + warnings.warn( + "Possibly corrupt EXIF data. " + f"Expecting to read {size} bytes but only got {len(data)}." + f" Skipping tag {tag}" + ) + logger.debug(msg) + continue + + if not data: + logger.debug(msg) + continue + + self._tagdata[tag] = data + self.tagtype[tag] = typ + + msg += " - value: " + msg += f"" if size > 32 else repr(data) + + logger.debug(msg) + + (self.next,) = ( + self._unpack("Q", self._ensure_read(fp, 8)) + if self._bigtiff + else self._unpack("L", self._ensure_read(fp, 4)) + ) + except OSError as msg: + warnings.warn(str(msg)) + return + + def _get_ifh(self) -> bytes: + ifh = self._prefix + self._pack("H", 43 if self._bigtiff else 42) + if self._bigtiff: + ifh += self._pack("HH", 8, 0) + ifh += self._pack("Q", 16) if self._bigtiff else self._pack("L", 8) + + return ifh + + def tobytes(self, offset: int = 0) -> bytes: + # FIXME What about tagdata? + result = self._pack("Q" if self._bigtiff else "H", len(self._tags_v2)) + + entries: list[tuple[int, int, int, bytes, bytes]] = [] + + fmt = "Q" if self._bigtiff else "L" + fmt_size = 8 if self._bigtiff else 4 + offset += ( + len(result) + len(self._tags_v2) * (20 if self._bigtiff else 12) + fmt_size + ) + stripoffsets = None + + # pass 1: convert tags to binary format + # always write tags in ascending order + for tag, value in sorted(self._tags_v2.items()): + if tag == STRIPOFFSETS: + stripoffsets = len(entries) + typ = self.tagtype[tag] + logger.debug("Tag %s, Type: %s, Value: %s", tag, typ, repr(value)) + is_ifd = typ == TiffTags.LONG and isinstance(value, dict) + if is_ifd: + ifd = ImageFileDirectory_v2(self._get_ifh(), group=tag) + values = self._tags_v2[tag] + for ifd_tag, ifd_value in values.items(): + ifd[ifd_tag] = ifd_value + data = ifd.tobytes(offset) + else: + values = value if isinstance(value, tuple) else (value,) + data = self._write_dispatch[typ](self, *values) + + tagname = TiffTags.lookup(tag, self.group).name + typname = "ifd" if is_ifd else TYPES.get(typ, "unknown") + msg = f"save: {tagname} ({tag}) - type: {typname} ({typ}) - value: " + msg += f"" if len(data) >= 16 else str(values) + logger.debug(msg) + + # count is sum of lengths for string and arbitrary data + if is_ifd: + count = 1 + elif typ in [TiffTags.BYTE, TiffTags.ASCII, TiffTags.UNDEFINED]: + count = len(data) + else: + count = len(values) + # figure out if data fits into the entry + if len(data) <= fmt_size: + entries.append((tag, typ, count, data.ljust(fmt_size, b"\0"), b"")) + else: + entries.append((tag, typ, count, self._pack(fmt, offset), data)) + offset += (len(data) + 1) // 2 * 2 # pad to word + + # update strip offset data to point beyond auxiliary data + if stripoffsets is not None: + tag, typ, count, value, data = entries[stripoffsets] + if data: + size, handler = self._load_dispatch[typ] + values = [val + offset for val in handler(self, data, self.legacy_api)] + data = self._write_dispatch[typ](self, *values) + else: + value = self._pack(fmt, self._unpack(fmt, value)[0] + offset) + entries[stripoffsets] = tag, typ, count, value, data + + # pass 2: write entries to file + for tag, typ, count, value, data in entries: + logger.debug("%s %s %s %s %s", tag, typ, count, repr(value), repr(data)) + result += self._pack( + "HHQ8s" if self._bigtiff else "HHL4s", tag, typ, count, value + ) + + # -- overwrite here for multi-page -- + result += self._pack(fmt, 0) # end of entries + + # pass 3: write auxiliary data to file + for tag, typ, count, value, data in entries: + result += data + if len(data) & 1: + result += b"\0" + + return result + + def save(self, fp: IO[bytes]) -> int: + if fp.tell() == 0: # skip TIFF header on subsequent pages + fp.write(self._get_ifh()) + + offset = fp.tell() + result = self.tobytes(offset) + fp.write(result) + return offset + len(result) + + +ImageFileDirectory_v2._load_dispatch = _load_dispatch +ImageFileDirectory_v2._write_dispatch = _write_dispatch +for idx, name in TYPES.items(): + name = name.replace(" ", "_") + setattr(ImageFileDirectory_v2, f"load_{name}", _load_dispatch[idx][1]) + setattr(ImageFileDirectory_v2, f"write_{name}", _write_dispatch[idx]) +del _load_dispatch, _write_dispatch, idx, name + + +# Legacy ImageFileDirectory support. +class ImageFileDirectory_v1(ImageFileDirectory_v2): + """This class represents the **legacy** interface to a TIFF tag directory. + + Exposes a dictionary interface of the tags in the directory:: + + ifd = ImageFileDirectory_v1() + ifd[key] = 'Some Data' + ifd.tagtype[key] = TiffTags.ASCII + print(ifd[key]) + ('Some Data',) + + Also contains a dictionary of tag types as read from the tiff image file, + :attr:`~PIL.TiffImagePlugin.ImageFileDirectory_v1.tagtype`. + + Values are returned as a tuple. + + .. deprecated:: 3.0.0 + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._legacy_api = True + + tags = property(lambda self: self._tags_v1) + tagdata = property(lambda self: self._tagdata) + + # defined in ImageFileDirectory_v2 + tagtype: dict[int, int] + """Dictionary of tag types""" + + @classmethod + def from_v2(cls, original: ImageFileDirectory_v2) -> ImageFileDirectory_v1: + """Returns an + :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v1` + instance with the same data as is contained in the original + :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v2` + instance. + + :returns: :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v1` + + """ + + ifd = cls(prefix=original.prefix) + ifd._tagdata = original._tagdata + ifd.tagtype = original.tagtype + ifd.next = original.next # an indicator for multipage tiffs + return ifd + + def to_v2(self) -> ImageFileDirectory_v2: + """Returns an + :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v2` + instance with the same data as is contained in the original + :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v1` + instance. + + :returns: :py:class:`~PIL.TiffImagePlugin.ImageFileDirectory_v2` + + """ + + ifd = ImageFileDirectory_v2(prefix=self.prefix) + ifd._tagdata = dict(self._tagdata) + ifd.tagtype = dict(self.tagtype) + ifd._tags_v2 = dict(self._tags_v2) + return ifd + + def __contains__(self, tag: object) -> bool: + return tag in self._tags_v1 or tag in self._tagdata + + def __len__(self) -> int: + return len(set(self._tagdata) | set(self._tags_v1)) + + def __iter__(self) -> Iterator[int]: + return iter(set(self._tagdata) | set(self._tags_v1)) + + def __setitem__(self, tag: int, value: Any) -> None: + for legacy_api in (False, True): + self._setitem(tag, value, legacy_api) + + def __getitem__(self, tag: int) -> Any: + if tag not in self._tags_v1: # unpack on the fly + data = self._tagdata[tag] + typ = self.tagtype[tag] + size, handler = self._load_dispatch[typ] + for legacy in (False, True): + self._setitem(tag, handler(self, data, legacy), legacy) + val = self._tags_v1[tag] + if not isinstance(val, (tuple, bytes)): + val = (val,) + return val + + +# undone -- switch this pointer +ImageFileDirectory = ImageFileDirectory_v1 + + +## +# Image plugin for TIFF files. + + +class TiffImageFile(ImageFile.ImageFile): + format = "TIFF" + format_description = "Adobe TIFF" + _close_exclusive_fp_after_loading = False + + def __init__( + self, + fp: StrOrBytesPath | IO[bytes], + filename: str | bytes | None = None, + ) -> None: + self.tag_v2: ImageFileDirectory_v2 + """ Image file directory (tag dictionary) """ + + self.tag: ImageFileDirectory_v1 + """ Legacy tag entries """ + + super().__init__(fp, filename) + + def _open(self) -> None: + """Open the first image in a TIFF file""" + + # Header + assert self.fp is not None + ifh = self.fp.read(8) + if ifh[2] == 43: + ifh += self.fp.read(8) + + self.tag_v2 = ImageFileDirectory_v2(ifh) + + # setup frame pointers + self.__first = self.__next = self.tag_v2.next + self.__frame = -1 + self._fp = self.fp + self._frame_pos: list[int] = [] + self._n_frames: int | None = None + + logger.debug("*** TiffImageFile._open ***") + logger.debug("- __first: %s", self.__first) + logger.debug("- ifh: %s", repr(ifh)) # Use repr to avoid str(bytes) + + # and load the first frame + self._seek(0) + + @property + def n_frames(self) -> int: + current_n_frames = self._n_frames + if current_n_frames is None: + current = self.tell() + self._seek(len(self._frame_pos)) + while self._n_frames is None: + self._seek(self.tell() + 1) + self.seek(current) + assert self._n_frames is not None + return self._n_frames + + def seek(self, frame: int) -> None: + """Select a given frame as current image""" + if not self._seek_check(frame): + return + self._seek(frame) + if self._im is not None and ( + self.im.size != self._tile_size + or self.im.mode != self.mode + or self.readonly + ): + self._im = None + + def _seek(self, frame: int) -> None: + if isinstance(self._fp, DeferredError): + raise self._fp.ex + self.fp = self._fp + + while len(self._frame_pos) <= frame: + if not self.__next: + msg = "no more images in TIFF file" + raise EOFError(msg) + logger.debug( + "Seeking to frame %s, on frame %s, __next %s, location: %s", + frame, + self.__frame, + self.__next, + self.fp.tell(), + ) + if self.__next >= 2**63: + msg = "Unable to seek to frame" + raise ValueError(msg) + self.fp.seek(self.__next) + self._frame_pos.append(self.__next) + logger.debug("Loading tags, location: %s", self.fp.tell()) + self.tag_v2.load(self.fp) + if self.tag_v2.next in self._frame_pos: + # This IFD has already been processed + # Declare this to be the end of the image + self.__next = 0 + else: + self.__next = self.tag_v2.next + if self.__next == 0: + self._n_frames = frame + 1 + if len(self._frame_pos) == 1: + self.is_animated = self.__next != 0 + self.__frame += 1 + self.fp.seek(self._frame_pos[frame]) + self.tag_v2.load(self.fp) + if XMP in self.tag_v2: + xmp = self.tag_v2[XMP] + if isinstance(xmp, tuple) and len(xmp) == 1: + xmp = xmp[0] + self.info["xmp"] = xmp + elif "xmp" in self.info: + del self.info["xmp"] + self._reload_exif() + # fill the legacy tag/ifd entries + self.tag = self.ifd = ImageFileDirectory_v1.from_v2(self.tag_v2) + self.__frame = frame + self._setup() + + def tell(self) -> int: + """Return the current frame number""" + return self.__frame + + def get_photoshop_blocks(self) -> dict[int, dict[str, bytes]]: + """ + Returns a dictionary of Photoshop "Image Resource Blocks". + The keys are the image resource ID. For more information, see + https://www.adobe.com/devnet-apps/photoshop/fileformatashtml/#50577409_pgfId-1037727 + + :returns: Photoshop "Image Resource Blocks" in a dictionary. + """ + blocks = {} + val = self.tag_v2.get(ExifTags.Base.ImageResources) + if val: + while val.startswith(b"8BIM"): + id = i16(val[4:6]) + n = math.ceil((val[6] + 1) / 2) * 2 + size = i32(val[6 + n : 10 + n]) + data = val[10 + n : 10 + n + size] + blocks[id] = {"data": data} + + val = val[math.ceil((10 + n + size) / 2) * 2 :] + return blocks + + def load(self) -> Image.core.PixelAccess | None: + if self.tile and self.use_load_libtiff: + return self._load_libtiff() + return super().load() + + def load_prepare(self) -> None: + if self._im is None: + Image._decompression_bomb_check(self._tile_size) + self.im = Image.core.new(self.mode, self._tile_size) + ImageFile.ImageFile.load_prepare(self) + + def load_end(self) -> None: + # allow closing if we're on the first frame, there's no next + # This is the ImageFile.load path only, libtiff specific below. + if not self.is_animated: + self._close_exclusive_fp_after_loading = True + + # load IFD data from fp before it is closed + exif = self.getexif() + for key in TiffTags.TAGS_V2_GROUPS: + if key not in exif: + continue + exif.get_ifd(key) + + ImageOps.exif_transpose(self, in_place=True) + if ExifTags.Base.Orientation in self.tag_v2: + del self.tag_v2[ExifTags.Base.Orientation] + + def _load_libtiff(self) -> Image.core.PixelAccess | None: + """Overload method triggered when we detect a compressed tiff + Calls out to libtiff""" + + Image.Image.load(self) + + self.load_prepare() + + if not len(self.tile) == 1: + msg = "Not exactly one tile" + raise OSError(msg) + + # (self._compression, (extents tuple), + # 0, (rawmode, self._compression, fp)) + extents = self.tile[0][1] + args = self.tile[0][3] + + # To be nice on memory footprint, if there's a + # file descriptor, use that instead of reading + # into a string in python. + assert self.fp is not None + try: + fp = hasattr(self.fp, "fileno") and self.fp.fileno() + # flush the file descriptor, prevents error on pypy 2.4+ + # should also eliminate the need for fp.tell + # in _seek + if hasattr(self.fp, "flush"): + self.fp.flush() + except OSError: + # io.BytesIO have a fileno, but returns an OSError if + # it doesn't use a file descriptor. + fp = False + + if fp: + assert isinstance(args, tuple) + args_list = list(args) + args_list[2] = fp + args = tuple(args_list) + + decoder = Image._getdecoder(self.mode, "libtiff", args, self.decoderconfig) + try: + decoder.setimage(self.im, extents) + except ValueError as e: + msg = "Couldn't set the image" + raise OSError(msg) from e + + close_self_fp = self._exclusive_fp and not self.is_animated + if hasattr(self.fp, "getvalue"): + # We've got a stringio like thing passed in. Yay for all in memory. + # The decoder needs the entire file in one shot, so there's not + # a lot we can do here other than give it the entire file. + # unless we could do something like get the address of the + # underlying string for stringio. + # + # Rearranging for supporting byteio items, since they have a fileno + # that returns an OSError if there's no underlying fp. Easier to + # deal with here by reordering. + logger.debug("have getvalue. just sending in a string from getvalue") + n, err = decoder.decode(self.fp.getvalue()) + elif fp: + # we've got a actual file on disk, pass in the fp. + logger.debug("have fileno, calling fileno version of the decoder.") + if not close_self_fp: + self.fp.seek(0) + # Save and restore the file position, because libtiff will move it + # outside of the Python runtime, and that will confuse + # io.BufferedReader and possible others. + # NOTE: This must use os.lseek(), and not fp.tell()/fp.seek(), + # because the buffer read head already may not equal the actual + # file position, and fp.seek() may just adjust it's internal + # pointer and not actually seek the OS file handle. + pos = os.lseek(fp, 0, os.SEEK_CUR) + # 4 bytes, otherwise the trace might error out + n, err = decoder.decode(b"fpfp") + os.lseek(fp, pos, os.SEEK_SET) + else: + # we have something else. + logger.debug("don't have fileno or getvalue. just reading") + self.fp.seek(0) + # UNDONE -- so much for that buffer size thing. + n, err = decoder.decode(self.fp.read()) + + self.tile = [] + self.readonly = 0 + + self.load_end() + + if close_self_fp: + self.fp.close() + self.fp = None # might be shared + + if err < 0: + msg = f"decoder error {err}" + raise OSError(msg) + + return Image.Image.load(self) + + def _setup(self) -> None: + """Setup this image object based on current tags""" + + if 0xBC01 in self.tag_v2: + msg = "Windows Media Photo files not yet supported" + raise OSError(msg) + + # extract relevant tags + self._compression = COMPRESSION_INFO[self.tag_v2.get(COMPRESSION, 1)] + self._planar_configuration = self.tag_v2.get(PLANAR_CONFIGURATION, 1) + + # photometric is a required tag, but not everyone is reading + # the specification + photo = self.tag_v2.get(PHOTOMETRIC_INTERPRETATION, 0) + + # old style jpeg compression images most certainly are YCbCr + if self._compression == "tiff_jpeg": + photo = 6 + + fillorder = self.tag_v2.get(FILLORDER, 1) + + logger.debug("*** Summary ***") + logger.debug("- compression: %s", self._compression) + logger.debug("- photometric_interpretation: %s", photo) + logger.debug("- planar_configuration: %s", self._planar_configuration) + logger.debug("- fill_order: %s", fillorder) + logger.debug("- YCbCr subsampling: %s", self.tag_v2.get(YCBCRSUBSAMPLING)) + + # size + try: + xsize = self.tag_v2[IMAGEWIDTH] + ysize = self.tag_v2[IMAGELENGTH] + except KeyError as e: + msg = "Missing dimensions" + raise TypeError(msg) from e + if not isinstance(xsize, int) or not isinstance(ysize, int): + msg = "Invalid dimensions" + raise ValueError(msg) + self._tile_size = xsize, ysize + orientation = self.tag_v2.get(ExifTags.Base.Orientation) + if orientation in (5, 6, 7, 8): + self._size = ysize, xsize + else: + self._size = xsize, ysize + + logger.debug("- size: %s", self.size) + + sample_format = self.tag_v2.get(SAMPLEFORMAT, (1,)) + if len(sample_format) > 1 and max(sample_format) == min(sample_format) == 1: + # SAMPLEFORMAT is properly per band, so an RGB image will + # be (1,1,1). But, we don't support per band pixel types, + # and anything more than one band is a uint8. So, just + # take the first element. Revisit this if adding support + # for more exotic images. + sample_format = (1,) + + bps_tuple = self.tag_v2.get(BITSPERSAMPLE, (1,)) + extra_tuple = self.tag_v2.get(EXTRASAMPLES, ()) + if photo in (2, 6, 8): # RGB, YCbCr, LAB + bps_count = 3 + elif photo == 5: # CMYK + bps_count = 4 + else: + bps_count = 1 + bps_count += len(extra_tuple) + bps_actual_count = len(bps_tuple) + samples_per_pixel = self.tag_v2.get( + SAMPLESPERPIXEL, + 3 if self._compression == "tiff_jpeg" and photo in (2, 6) else 1, + ) + + if samples_per_pixel > MAX_SAMPLESPERPIXEL: + # DOS check, samples_per_pixel can be a Long, and we extend the tuple below + logger.error( + "More samples per pixel than can be decoded: %s", samples_per_pixel + ) + msg = "Invalid value for samples per pixel" + raise SyntaxError(msg) + + if samples_per_pixel < bps_actual_count: + # If a file has more values in bps_tuple than expected, + # remove the excess. + bps_tuple = bps_tuple[:samples_per_pixel] + elif samples_per_pixel > bps_actual_count and bps_actual_count == 1: + # If a file has only one value in bps_tuple, when it should have more, + # presume it is the same number of bits for all of the samples. + bps_tuple = bps_tuple * samples_per_pixel + + if len(bps_tuple) != samples_per_pixel: + msg = "unknown data organization" + raise SyntaxError(msg) + + # mode: check photometric interpretation and bits per pixel + key = ( + self.tag_v2.prefix, + photo, + sample_format, + fillorder, + bps_tuple, + extra_tuple, + ) + logger.debug("format key: %s", key) + try: + self._mode, rawmode = OPEN_INFO[key] + except KeyError as e: + logger.debug("- unsupported format") + msg = "unknown pixel mode" + raise SyntaxError(msg) from e + + logger.debug("- raw mode: %s", rawmode) + logger.debug("- pil mode: %s", self.mode) + + self.info["compression"] = self._compression + + xres = self.tag_v2.get(X_RESOLUTION, 1) + yres = self.tag_v2.get(Y_RESOLUTION, 1) + + if xres and yres: + resunit = self.tag_v2.get(RESOLUTION_UNIT) + if resunit == 2: # dots per inch + self.info["dpi"] = (xres, yres) + elif resunit == 3: # dots per centimeter. convert to dpi + self.info["dpi"] = (xres * 2.54, yres * 2.54) + elif resunit is None: # used to default to 1, but now 2) + self.info["dpi"] = (xres, yres) + # For backward compatibility, + # we also preserve the old behavior + self.info["resolution"] = xres, yres + else: # No absolute unit of measurement + self.info["resolution"] = xres, yres + + # build tile descriptors + x = y = layer = 0 + self.tile = [] + self.use_load_libtiff = READ_LIBTIFF or self._compression != "raw" + if self.use_load_libtiff: + # Decoder expects entire file as one tile. + # There's a buffer size limit in load (64k) + # so large g4 images will fail if we use that + # function. + # + # Setup the one tile for the whole image, then + # use the _load_libtiff function. + + # libtiff handles the fillmode for us, so 1;IR should + # actually be 1;I. Including the R double reverses the + # bits, so stripes of the image are reversed. See + # https://github.com/python-pillow/Pillow/issues/279 + if fillorder == 2: + # Replace fillorder with fillorder=1 + key = key[:3] + (1,) + key[4:] + logger.debug("format key: %s", key) + # this should always work, since all the + # fillorder==2 modes have a corresponding + # fillorder=1 mode + self._mode, rawmode = OPEN_INFO[key] + # YCbCr images with new jpeg compression with pixels in one plane + # unpacked straight into RGB values + if ( + photo == 6 + and self._compression == "jpeg" + and self._planar_configuration == 1 + ): + rawmode = "RGB" + # libtiff always returns the bytes in native order. + # we're expecting image byte order. So, if the rawmode + # contains I;16, we need to convert from native to image + # byte order. + elif rawmode == "I;16": + rawmode = "I;16N" + elif rawmode.endswith((";16B", ";16L")): + rawmode = rawmode[:-1] + "N" + + # Offset in the tile tuple is 0, we go from 0,0 to + # w,h, and we only do this once -- eds + a = (rawmode, self._compression, False, self.tag_v2.offset) + self.tile.append(ImageFile._Tile("libtiff", (0, 0, xsize, ysize), 0, a)) + + elif STRIPOFFSETS in self.tag_v2 or TILEOFFSETS in self.tag_v2: + # striped image + if STRIPOFFSETS in self.tag_v2: + offsets = self.tag_v2[STRIPOFFSETS] + h = self.tag_v2.get(ROWSPERSTRIP, ysize) + w = xsize + else: + # tiled image + offsets = self.tag_v2[TILEOFFSETS] + tilewidth = self.tag_v2.get(TILEWIDTH) + h = self.tag_v2.get(TILELENGTH) + if not isinstance(tilewidth, int) or not isinstance(h, int): + msg = "Invalid tile dimensions" + raise ValueError(msg) + w = tilewidth + + if w == xsize and h == ysize and self._planar_configuration != 2: + # Every tile covers the image. Only use the last offset + offsets = offsets[-1:] + + for offset in offsets: + if x + w > xsize: + stride = w * sum(bps_tuple) / 8 # bytes per line + else: + stride = 0 + + tile_rawmode = rawmode + if self._planar_configuration == 2: + # each band on it's own layer + tile_rawmode = rawmode[layer] + # adjust stride width accordingly + stride /= bps_count + + args = (tile_rawmode, int(stride), 1) + self.tile.append( + ImageFile._Tile( + self._compression, + (x, y, min(x + w, xsize), min(y + h, ysize)), + offset, + args, + ) + ) + x += w + if x >= xsize: + x, y = 0, y + h + if y >= ysize: + y = 0 + layer += 1 + else: + logger.debug("- unsupported data organization") + msg = "unknown data organization" + raise SyntaxError(msg) + + # Fix up info. + if ICCPROFILE in self.tag_v2: + self.info["icc_profile"] = self.tag_v2[ICCPROFILE] + + # fixup palette descriptor + + if self.mode in ["P", "PA"]: + palette = [o8(b // 256) for b in self.tag_v2[COLORMAP]] + self.palette = ImagePalette.raw("RGB;L", b"".join(palette)) + + +# +# -------------------------------------------------------------------- +# Write TIFF files + +# little endian is default except for image modes with +# explicit big endian byte-order + +SAVE_INFO = { + # mode => rawmode, byteorder, photometrics, + # sampleformat, bitspersample, extra + "1": ("1", II, 1, 1, (1,), None), + "L": ("L", II, 1, 1, (8,), None), + "LA": ("LA", II, 1, 1, (8, 8), 2), + "P": ("P", II, 3, 1, (8,), None), + "PA": ("PA", II, 3, 1, (8, 8), 2), + "I": ("I;32S", II, 1, 2, (32,), None), + "I;16": ("I;16", II, 1, 1, (16,), None), + "I;16L": ("I;16L", II, 1, 1, (16,), None), + "F": ("F;32F", II, 1, 3, (32,), None), + "RGB": ("RGB", II, 2, 1, (8, 8, 8), None), + "RGBX": ("RGBX", II, 2, 1, (8, 8, 8, 8), 0), + "RGBA": ("RGBA", II, 2, 1, (8, 8, 8, 8), 2), + "CMYK": ("CMYK", II, 5, 1, (8, 8, 8, 8), None), + "YCbCr": ("YCbCr", II, 6, 1, (8, 8, 8), None), + "LAB": ("LAB", II, 8, 1, (8, 8, 8), None), + "I;16B": ("I;16B", MM, 1, 1, (16,), None), +} + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + try: + rawmode, prefix, photo, format, bits, extra = SAVE_INFO[im.mode] + except KeyError as e: + msg = f"cannot write mode {im.mode} as TIFF" + raise OSError(msg) from e + + encoderinfo = im.encoderinfo + encoderconfig = im.encoderconfig + + ifd = ImageFileDirectory_v2(prefix=prefix) + if encoderinfo.get("big_tiff"): + ifd._bigtiff = True + + try: + compression = encoderinfo["compression"] + except KeyError: + compression = im.info.get("compression") + if isinstance(compression, int): + # compression value may be from BMP. Ignore it + compression = None + if compression is None: + compression = "raw" + elif compression == "tiff_jpeg": + # OJPEG is obsolete, so use new-style JPEG compression instead + compression = "jpeg" + elif compression == "tiff_deflate": + compression = "tiff_adobe_deflate" + + libtiff = WRITE_LIBTIFF or compression != "raw" + + # required for color libtiff images + ifd[PLANAR_CONFIGURATION] = 1 + + ifd[IMAGEWIDTH] = im.size[0] + ifd[IMAGELENGTH] = im.size[1] + + # write any arbitrary tags passed in as an ImageFileDirectory + if "tiffinfo" in encoderinfo: + info = encoderinfo["tiffinfo"] + elif "exif" in encoderinfo: + info = encoderinfo["exif"] + if isinstance(info, bytes): + exif = Image.Exif() + exif.load(info) + info = exif + else: + info = {} + logger.debug("Tiffinfo Keys: %s", list(info)) + if isinstance(info, ImageFileDirectory_v1): + info = info.to_v2() + for key in info: + if isinstance(info, Image.Exif) and key in TiffTags.TAGS_V2_GROUPS: + ifd[key] = info.get_ifd(key) + else: + ifd[key] = info.get(key) + try: + ifd.tagtype[key] = info.tagtype[key] + except Exception: + pass # might not be an IFD. Might not have populated type + + legacy_ifd = {} + if hasattr(im, "tag"): + legacy_ifd = im.tag.to_v2() + + supplied_tags = {**legacy_ifd, **getattr(im, "tag_v2", {})} + for tag in ( + # IFD offset that may not be correct in the saved image + EXIFIFD, + # Determined by the image format and should not be copied from legacy_ifd. + SAMPLEFORMAT, + ): + if tag in supplied_tags: + del supplied_tags[tag] + + # additions written by Greg Couch, gregc@cgl.ucsf.edu + # inspired by image-sig posting from Kevin Cazabon, kcazabon@home.com + if hasattr(im, "tag_v2"): + # preserve tags from original TIFF image file + for key in ( + RESOLUTION_UNIT, + X_RESOLUTION, + Y_RESOLUTION, + IPTC_NAA_CHUNK, + PHOTOSHOP_CHUNK, + XMP, + ): + if key in im.tag_v2: + if key == IPTC_NAA_CHUNK and im.tag_v2.tagtype[key] not in ( + TiffTags.BYTE, + TiffTags.UNDEFINED, + ): + del supplied_tags[key] + else: + ifd[key] = im.tag_v2[key] + ifd.tagtype[key] = im.tag_v2.tagtype[key] + + # preserve ICC profile (should also work when saving other formats + # which support profiles as TIFF) -- 2008-06-06 Florian Hoech + icc = encoderinfo.get("icc_profile", im.info.get("icc_profile")) + if icc: + ifd[ICCPROFILE] = icc + + for key, name in [ + (IMAGEDESCRIPTION, "description"), + (X_RESOLUTION, "resolution"), + (Y_RESOLUTION, "resolution"), + (X_RESOLUTION, "x_resolution"), + (Y_RESOLUTION, "y_resolution"), + (RESOLUTION_UNIT, "resolution_unit"), + (SOFTWARE, "software"), + (DATE_TIME, "date_time"), + (ARTIST, "artist"), + (COPYRIGHT, "copyright"), + ]: + if name in encoderinfo: + ifd[key] = encoderinfo[name] + + dpi = encoderinfo.get("dpi") + if dpi: + ifd[RESOLUTION_UNIT] = 2 + ifd[X_RESOLUTION] = dpi[0] + ifd[Y_RESOLUTION] = dpi[1] + + if bits != (1,): + ifd[BITSPERSAMPLE] = bits + if len(bits) != 1: + ifd[SAMPLESPERPIXEL] = len(bits) + if extra is not None: + ifd[EXTRASAMPLES] = extra + if format != 1: + ifd[SAMPLEFORMAT] = format + + if PHOTOMETRIC_INTERPRETATION not in ifd: + ifd[PHOTOMETRIC_INTERPRETATION] = photo + elif im.mode in ("1", "L") and ifd[PHOTOMETRIC_INTERPRETATION] == 0: + if im.mode == "1": + inverted_im = im.copy() + px = inverted_im.load() + if px is not None: + for y in range(inverted_im.height): + for x in range(inverted_im.width): + px[x, y] = 0 if px[x, y] == 255 else 255 + im = inverted_im + else: + im = ImageOps.invert(im) + + if im.mode in ["P", "PA"]: + lut = im.im.getpalette("RGB", "RGB;L") + colormap = [] + colors = len(lut) // 3 + for i in range(3): + colormap += [v * 256 for v in lut[colors * i : colors * (i + 1)]] + colormap += [0] * (256 - colors) + ifd[COLORMAP] = colormap + # data orientation + w, h = ifd[IMAGEWIDTH], ifd[IMAGELENGTH] + stride = len(bits) * ((w * bits[0] + 7) // 8) + if ROWSPERSTRIP not in ifd: + # aim for given strip size (64 KB by default) when using libtiff writer + if libtiff: + im_strip_size = encoderinfo.get("strip_size", STRIP_SIZE) + rows_per_strip = 1 if stride == 0 else min(im_strip_size // stride, h) + # JPEG encoder expects multiple of 8 rows + if compression == "jpeg": + rows_per_strip = min(((rows_per_strip + 7) // 8) * 8, h) + else: + rows_per_strip = h + if rows_per_strip == 0: + rows_per_strip = 1 + ifd[ROWSPERSTRIP] = rows_per_strip + strip_byte_counts = 1 if stride == 0 else stride * ifd[ROWSPERSTRIP] + strips_per_image = (h + ifd[ROWSPERSTRIP] - 1) // ifd[ROWSPERSTRIP] + if strip_byte_counts >= 2**16: + ifd.tagtype[STRIPBYTECOUNTS] = TiffTags.LONG + ifd[STRIPBYTECOUNTS] = (strip_byte_counts,) * (strips_per_image - 1) + ( + stride * h - strip_byte_counts * (strips_per_image - 1), + ) + ifd[STRIPOFFSETS] = tuple( + range(0, strip_byte_counts * strips_per_image, strip_byte_counts) + ) # this is adjusted by IFD writer + # no compression by default: + ifd[COMPRESSION] = COMPRESSION_INFO_REV.get(compression, 1) + + if im.mode == "YCbCr": + for tag, default_value in { + YCBCRSUBSAMPLING: (1, 1), + REFERENCEBLACKWHITE: (0, 255, 128, 255, 128, 255), + }.items(): + ifd.setdefault(tag, default_value) + + blocklist = [TILEWIDTH, TILELENGTH, TILEOFFSETS, TILEBYTECOUNTS] + if libtiff: + if "quality" in encoderinfo: + quality = encoderinfo["quality"] + if not isinstance(quality, int) or quality < 0 or quality > 100: + msg = "Invalid quality setting" + raise ValueError(msg) + if compression != "jpeg": + msg = "quality setting only supported for 'jpeg' compression" + raise ValueError(msg) + ifd[JPEGQUALITY] = quality + + logger.debug("Saving using libtiff encoder") + logger.debug("Items: %s", sorted(ifd.items())) + _fp = 0 + if hasattr(fp, "fileno"): + try: + fp.seek(0) + _fp = fp.fileno() + except io.UnsupportedOperation: + pass + + # optional types for non core tags + types = {} + # STRIPOFFSETS and STRIPBYTECOUNTS are added by the library + # based on the data in the strip. + # OSUBFILETYPE is deprecated. + # The other tags expect arrays with a certain length (fixed or depending on + # BITSPERSAMPLE, etc), passing arrays with a different length will result in + # segfaults. Block these tags until we add extra validation. + # SUBIFD may also cause a segfault. + blocklist += [ + OSUBFILETYPE, + REFERENCEBLACKWHITE, + STRIPBYTECOUNTS, + STRIPOFFSETS, + TRANSFERFUNCTION, + SUBIFD, + ] + + # bits per sample is a single short in the tiff directory, not a list. + atts: dict[int, Any] = {BITSPERSAMPLE: bits[0]} + # Merge the ones that we have with (optional) more bits from + # the original file, e.g x,y resolution so that we can + # save(load('')) == original file. + for tag, value in itertools.chain(ifd.items(), supplied_tags.items()): + # Libtiff can only process certain core items without adding + # them to the custom dictionary. + # Custom items are supported for int, float, unicode, string and byte + # values. Other types and tuples require a tagtype. + if tag not in TiffTags.LIBTIFF_CORE: + if tag in TiffTags.TAGS_V2_GROUPS: + types[tag] = TiffTags.LONG8 + elif tag in ifd.tagtype: + types[tag] = ifd.tagtype[tag] + elif isinstance(value, (int, float, str, bytes)) or ( + isinstance(value, tuple) + and all(isinstance(v, (int, float, IFDRational)) for v in value) + ): + type = TiffTags.lookup(tag).type + if type: + types[tag] = type + if tag not in atts and tag not in blocklist: + if isinstance(value, str): + atts[tag] = value.encode("ascii", "replace") + b"\0" + elif isinstance(value, IFDRational): + atts[tag] = float(value) + else: + atts[tag] = value + + if SAMPLEFORMAT in atts and len(atts[SAMPLEFORMAT]) == 1: + atts[SAMPLEFORMAT] = atts[SAMPLEFORMAT][0] + + logger.debug("Converted items: %s", sorted(atts.items())) + + # libtiff always expects the bytes in native order. + # we're storing image byte order. So, if the rawmode + # contains I;16, we need to convert from native to image + # byte order. + if im.mode in ("I;16", "I;16B", "I;16L"): + rawmode = "I;16N" + + # Pass tags as sorted list so that the tags are set in a fixed order. + # This is required by libtiff for some tags. For example, the JPEGQUALITY + # pseudo tag requires that the COMPRESS tag was already set. + tags = list(atts.items()) + tags.sort() + a = (rawmode, compression, _fp, filename, tags, types) + encoder = Image._getencoder(im.mode, "libtiff", a, encoderconfig) + encoder.setimage(im.im, (0, 0) + im.size) + while True: + errcode, data = encoder.encode(ImageFile.MAXBLOCK)[1:] + if not _fp: + fp.write(data) + if errcode: + break + if errcode < 0: + msg = f"encoder error {errcode} when writing image file" + raise OSError(msg) + + else: + for tag in blocklist: + del ifd[tag] + offset = ifd.save(fp) + + ImageFile._save( + im, + fp, + [ImageFile._Tile("raw", (0, 0) + im.size, offset, (rawmode, stride, 1))], + ) + + # -- helper for multi-page save -- + if "_debug_multipage" in encoderinfo: + # just to access o32 and o16 (using correct byte order) + setattr(im, "_debug_multipage", ifd) + + +class AppendingTiffWriter(io.BytesIO): + fieldSizes = [ + 0, # None + 1, # byte + 1, # ascii + 2, # short + 4, # long + 8, # rational + 1, # sbyte + 1, # undefined + 2, # sshort + 4, # slong + 8, # srational + 4, # float + 8, # double + 4, # ifd + 2, # unicode + 4, # complex + 8, # long8 + ] + + Tags = { + 273, # StripOffsets + 288, # FreeOffsets + 324, # TileOffsets + 519, # JPEGQTables + 520, # JPEGDCTables + 521, # JPEGACTables + } + + def __init__(self, fn: StrOrBytesPath | IO[bytes], new: bool = False) -> None: + self.f: IO[bytes] + if is_path(fn): + self.name = fn + self.close_fp = True + try: + self.f = open(fn, "w+b" if new else "r+b") + except OSError: + self.f = open(fn, "w+b") + else: + self.f = cast(IO[bytes], fn) + self.close_fp = False + self.beginning = self.f.tell() + self.setup() + + def setup(self) -> None: + # Reset everything. + self.f.seek(self.beginning, os.SEEK_SET) + + self.whereToWriteNewIFDOffset: int | None = None + self.offsetOfNewPage = 0 + + self.IIMM = iimm = self.f.read(4) + self._bigtiff = b"\x2b" in iimm + if not iimm: + # empty file - first page + self.isFirst = True + return + + self.isFirst = False + if iimm not in PREFIXES: + msg = "Invalid TIFF file header" + raise RuntimeError(msg) + + self.setEndian("<" if iimm.startswith(II) else ">") + + if self._bigtiff: + self.f.seek(4, os.SEEK_CUR) + self.skipIFDs() + self.goToEnd() + + def finalize(self) -> None: + if self.isFirst: + return + + # fix offsets + self.f.seek(self.offsetOfNewPage) + + iimm = self.f.read(4) + if not iimm: + # Make it easy to finish a frame without committing to a new one. + return + + if iimm != self.IIMM: + msg = "IIMM of new page doesn't match IIMM of first page" + raise RuntimeError(msg) + + if self._bigtiff: + self.f.seek(4, os.SEEK_CUR) + ifd_offset = self._read(8 if self._bigtiff else 4) + ifd_offset += self.offsetOfNewPage + assert self.whereToWriteNewIFDOffset is not None + self.f.seek(self.whereToWriteNewIFDOffset) + self._write(ifd_offset, 8 if self._bigtiff else 4) + self.f.seek(ifd_offset) + self.fixIFD() + + def newFrame(self) -> None: + # Call this to finish a frame. + self.finalize() + self.setup() + + def __enter__(self) -> AppendingTiffWriter: + return self + + def __exit__(self, *args: object) -> None: + if self.close_fp: + self.close() + + def tell(self) -> int: + return self.f.tell() - self.offsetOfNewPage + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + """ + :param offset: Distance to seek. + :param whence: Whether the distance is relative to the start, + end or current position. + :returns: The resulting position, relative to the start. + """ + if whence == os.SEEK_SET: + offset += self.offsetOfNewPage + + self.f.seek(offset, whence) + return self.tell() + + def goToEnd(self) -> None: + self.f.seek(0, os.SEEK_END) + pos = self.f.tell() + + # pad to 16 byte boundary + pad_bytes = 16 - pos % 16 + if 0 < pad_bytes < 16: + self.f.write(bytes(pad_bytes)) + self.offsetOfNewPage = self.f.tell() + + def setEndian(self, endian: str) -> None: + self.endian = endian + self.longFmt = f"{self.endian}L" + self.shortFmt = f"{self.endian}H" + self.tagFormat = f"{self.endian}HH" + ("Q" if self._bigtiff else "L") + + def skipIFDs(self) -> None: + while True: + ifd_offset = self._read(8 if self._bigtiff else 4) + if ifd_offset == 0: + self.whereToWriteNewIFDOffset = self.f.tell() - ( + 8 if self._bigtiff else 4 + ) + break + + self.f.seek(ifd_offset) + num_tags = self._read(8 if self._bigtiff else 2) + self.f.seek(num_tags * (20 if self._bigtiff else 12), os.SEEK_CUR) + + def write(self, data: Buffer, /) -> int: + return self.f.write(data) + + def _fmt(self, field_size: int) -> str: + try: + return {2: "H", 4: "L", 8: "Q"}[field_size] + except KeyError: + msg = "offset is not supported" + raise RuntimeError(msg) + + def _read(self, field_size: int) -> int: + (value,) = struct.unpack( + self.endian + self._fmt(field_size), self.f.read(field_size) + ) + return value + + def readShort(self) -> int: + return self._read(2) + + def readLong(self) -> int: + return self._read(4) + + @staticmethod + def _verify_bytes_written(bytes_written: int | None, expected: int) -> None: + if bytes_written is not None and bytes_written != expected: + msg = f"wrote only {bytes_written} bytes but wanted {expected}" + raise RuntimeError(msg) + + def _rewriteLast( + self, value: int, field_size: int, new_field_size: int = 0 + ) -> None: + self.f.seek(-field_size, os.SEEK_CUR) + if not new_field_size: + new_field_size = field_size + bytes_written = self.f.write( + struct.pack(self.endian + self._fmt(new_field_size), value) + ) + self._verify_bytes_written(bytes_written, new_field_size) + + def rewriteLastShortToLong(self, value: int) -> None: + self._rewriteLast(value, 2, 4) + + def rewriteLastShort(self, value: int) -> None: + return self._rewriteLast(value, 2) + + def rewriteLastLong(self, value: int) -> None: + return self._rewriteLast(value, 4) + + def _write(self, value: int, field_size: int) -> None: + bytes_written = self.f.write( + struct.pack(self.endian + self._fmt(field_size), value) + ) + self._verify_bytes_written(bytes_written, field_size) + + def writeShort(self, value: int) -> None: + self._write(value, 2) + + def writeLong(self, value: int) -> None: + self._write(value, 4) + + def close(self) -> None: + self.finalize() + if self.close_fp: + self.f.close() + + def fixIFD(self) -> None: + num_tags = self._read(8 if self._bigtiff else 2) + + for i in range(num_tags): + tag, field_type, count = struct.unpack( + self.tagFormat, self.f.read(12 if self._bigtiff else 8) + ) + + field_size = self.fieldSizes[field_type] + total_size = field_size * count + fmt_size = 8 if self._bigtiff else 4 + is_local = total_size <= fmt_size + if not is_local: + offset = self._read(fmt_size) + self.offsetOfNewPage + self._rewriteLast(offset, fmt_size) + + if tag in self.Tags: + cur_pos = self.f.tell() + + logger.debug( + "fixIFD: %s (%d) - type: %s (%d) - type size: %d - count: %d", + TiffTags.lookup(tag).name, + tag, + TYPES.get(field_type, "unknown"), + field_type, + field_size, + count, + ) + + if is_local: + self._fixOffsets(count, field_size) + self.f.seek(cur_pos + fmt_size) + else: + self.f.seek(offset) + self._fixOffsets(count, field_size) + self.f.seek(cur_pos) + + elif is_local: + # skip the locally stored value that is not an offset + self.f.seek(fmt_size, os.SEEK_CUR) + + def _fixOffsets(self, count: int, field_size: int) -> None: + for i in range(count): + offset = self._read(field_size) + offset += self.offsetOfNewPage + + new_field_size = 0 + if self._bigtiff and field_size in (2, 4) and offset >= 2**32: + # offset is now too large - we must convert long to long8 + new_field_size = 8 + elif field_size == 2 and offset >= 2**16: + # offset is now too large - we must convert short to long + new_field_size = 4 + if new_field_size: + if count != 1: + msg = "not implemented" + raise RuntimeError(msg) # XXX TODO + + # simple case - the offset is just one and therefore it is + # local (not referenced with another offset) + self._rewriteLast(offset, field_size, new_field_size) + # Move back past the new offset, past 'count', and before 'field_type' + rewind = -new_field_size - 4 - 2 + self.f.seek(rewind, os.SEEK_CUR) + self.writeShort(new_field_size) # rewrite the type + self.f.seek(2 - rewind, os.SEEK_CUR) + else: + self._rewriteLast(offset, field_size) + + def fixOffsets( + self, count: int, isShort: bool = False, isLong: bool = False + ) -> None: + if isShort: + field_size = 2 + elif isLong: + field_size = 4 + else: + field_size = 0 + return self._fixOffsets(count, field_size) + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + append_images = list(im.encoderinfo.get("append_images", [])) + if not hasattr(im, "n_frames") and not append_images: + return _save(im, fp, filename) + + cur_idx = im.tell() + try: + with AppendingTiffWriter(fp) as tf: + for ims in [im] + append_images: + encoderinfo = ims._attach_default_encoderinfo(im) + if not hasattr(ims, "encoderconfig"): + ims.encoderconfig = () + nfr = getattr(ims, "n_frames", 1) + + for idx in range(nfr): + ims.seek(idx) + ims.load() + _save(ims, tf, filename) + tf.newFrame() + ims.encoderinfo = encoderinfo + finally: + im.seek(cur_idx) + + +# +# -------------------------------------------------------------------- +# Register + +Image.register_open(TiffImageFile.format, TiffImageFile, _accept) +Image.register_save(TiffImageFile.format, _save) +Image.register_save_all(TiffImageFile.format, _save_all) + +Image.register_extensions(TiffImageFile.format, [".tif", ".tiff"]) + +Image.register_mime(TiffImageFile.format, "image/tiff") diff --git a/venv/lib/python3.10/site-packages/PIL/TiffTags.py b/venv/lib/python3.10/site-packages/PIL/TiffTags.py new file mode 100644 index 0000000000000000000000000000000000000000..761aa3f6b0869cd1dc36f0a632d5e59ecbd931c4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/TiffTags.py @@ -0,0 +1,567 @@ +# +# The Python Imaging Library. +# $Id$ +# +# TIFF tags +# +# This module provides clear-text names for various well-known +# TIFF tags. the TIFF codec works just fine without it. +# +# Copyright (c) Secret Labs AB 1999. +# +# See the README file for information on usage and redistribution. +# + +## +# This module provides constants and clear-text names for various +# well-known TIFF tags. +## +from __future__ import annotations + +from typing import NamedTuple + + +class _TagInfo(NamedTuple): + value: int | None + name: str + type: int | None + length: int | None + enum: dict[str, int] + + +class TagInfo(_TagInfo): + __slots__: list[str] = [] + + def __new__( + cls, + value: int | None = None, + name: str = "unknown", + type: int | None = None, + length: int | None = None, + enum: dict[str, int] | None = None, + ) -> TagInfo: + return super().__new__(cls, value, name, type, length, enum or {}) + + def cvt_enum(self, value: str) -> int | str: + # Using get will call hash(value), which can be expensive + # for some types (e.g. Fraction). Since self.enum is rarely + # used, it's usually better to test it first. + return self.enum.get(value, value) if self.enum else value + + +def lookup(tag: int, group: int | None = None) -> TagInfo: + """ + :param tag: Integer tag number + :param group: Which :py:data:`~PIL.TiffTags.TAGS_V2_GROUPS` to look in + + .. versionadded:: 8.3.0 + + :returns: Taginfo namedtuple, From the ``TAGS_V2`` info if possible, + otherwise just populating the value and name from ``TAGS``. + If the tag is not recognized, "unknown" is returned for the name + + """ + + if group is not None: + info = TAGS_V2_GROUPS[group].get(tag) if group in TAGS_V2_GROUPS else None + else: + info = TAGS_V2.get(tag) + return info or TagInfo(tag, TAGS.get(tag, "unknown")) + + +## +# Map tag numbers to tag info. +# +# id: (Name, Type, Length[, enum_values]) +# +# The length here differs from the length in the tiff spec. For +# numbers, the tiff spec is for the number of fields returned. We +# agree here. For string-like types, the tiff spec uses the length of +# field in bytes. In Pillow, we are using the number of expected +# fields, in general 1 for string-like types. + + +BYTE = 1 +ASCII = 2 +SHORT = 3 +LONG = 4 +RATIONAL = 5 +SIGNED_BYTE = 6 +UNDEFINED = 7 +SIGNED_SHORT = 8 +SIGNED_LONG = 9 +SIGNED_RATIONAL = 10 +FLOAT = 11 +DOUBLE = 12 +IFD = 13 +LONG8 = 16 + +_tags_v2: dict[int, tuple[str, int, int] | tuple[str, int, int, dict[str, int]]] = { + 254: ("NewSubfileType", LONG, 1), + 255: ("SubfileType", SHORT, 1), + 256: ("ImageWidth", LONG, 1), + 257: ("ImageLength", LONG, 1), + 258: ("BitsPerSample", SHORT, 0), + 259: ( + "Compression", + SHORT, + 1, + { + "Uncompressed": 1, + "CCITT 1d": 2, + "Group 3 Fax": 3, + "Group 4 Fax": 4, + "LZW": 5, + "JPEG": 6, + "PackBits": 32773, + }, + ), + 262: ( + "PhotometricInterpretation", + SHORT, + 1, + { + "WhiteIsZero": 0, + "BlackIsZero": 1, + "RGB": 2, + "RGB Palette": 3, + "Transparency Mask": 4, + "CMYK": 5, + "YCbCr": 6, + "CieLAB": 8, + "CFA": 32803, # TIFF/EP, Adobe DNG + "LinearRaw": 32892, # Adobe DNG + }, + ), + 263: ("Threshholding", SHORT, 1), + 264: ("CellWidth", SHORT, 1), + 265: ("CellLength", SHORT, 1), + 266: ("FillOrder", SHORT, 1), + 269: ("DocumentName", ASCII, 1), + 270: ("ImageDescription", ASCII, 1), + 271: ("Make", ASCII, 1), + 272: ("Model", ASCII, 1), + 273: ("StripOffsets", LONG, 0), + 274: ("Orientation", SHORT, 1), + 277: ("SamplesPerPixel", SHORT, 1), + 278: ("RowsPerStrip", LONG, 1), + 279: ("StripByteCounts", LONG, 0), + 280: ("MinSampleValue", SHORT, 0), + 281: ("MaxSampleValue", SHORT, 0), + 282: ("XResolution", RATIONAL, 1), + 283: ("YResolution", RATIONAL, 1), + 284: ("PlanarConfiguration", SHORT, 1, {"Contiguous": 1, "Separate": 2}), + 285: ("PageName", ASCII, 1), + 286: ("XPosition", RATIONAL, 1), + 287: ("YPosition", RATIONAL, 1), + 288: ("FreeOffsets", LONG, 1), + 289: ("FreeByteCounts", LONG, 1), + 290: ("GrayResponseUnit", SHORT, 1), + 291: ("GrayResponseCurve", SHORT, 0), + 292: ("T4Options", LONG, 1), + 293: ("T6Options", LONG, 1), + 296: ("ResolutionUnit", SHORT, 1, {"none": 1, "inch": 2, "cm": 3}), + 297: ("PageNumber", SHORT, 2), + 301: ("TransferFunction", SHORT, 0), + 305: ("Software", ASCII, 1), + 306: ("DateTime", ASCII, 1), + 315: ("Artist", ASCII, 1), + 316: ("HostComputer", ASCII, 1), + 317: ("Predictor", SHORT, 1, {"none": 1, "Horizontal Differencing": 2}), + 318: ("WhitePoint", RATIONAL, 2), + 319: ("PrimaryChromaticities", RATIONAL, 6), + 320: ("ColorMap", SHORT, 0), + 321: ("HalftoneHints", SHORT, 2), + 322: ("TileWidth", LONG, 1), + 323: ("TileLength", LONG, 1), + 324: ("TileOffsets", LONG, 0), + 325: ("TileByteCounts", LONG, 0), + 330: ("SubIFDs", LONG, 0), + 332: ("InkSet", SHORT, 1), + 333: ("InkNames", ASCII, 1), + 334: ("NumberOfInks", SHORT, 1), + 336: ("DotRange", SHORT, 0), + 337: ("TargetPrinter", ASCII, 1), + 338: ("ExtraSamples", SHORT, 0), + 339: ("SampleFormat", SHORT, 0), + 340: ("SMinSampleValue", DOUBLE, 0), + 341: ("SMaxSampleValue", DOUBLE, 0), + 342: ("TransferRange", SHORT, 6), + 347: ("JPEGTables", UNDEFINED, 1), + # obsolete JPEG tags + 512: ("JPEGProc", SHORT, 1), + 513: ("JPEGInterchangeFormat", LONG, 1), + 514: ("JPEGInterchangeFormatLength", LONG, 1), + 515: ("JPEGRestartInterval", SHORT, 1), + 517: ("JPEGLosslessPredictors", SHORT, 0), + 518: ("JPEGPointTransforms", SHORT, 0), + 519: ("JPEGQTables", LONG, 0), + 520: ("JPEGDCTables", LONG, 0), + 521: ("JPEGACTables", LONG, 0), + 529: ("YCbCrCoefficients", RATIONAL, 3), + 530: ("YCbCrSubSampling", SHORT, 2), + 531: ("YCbCrPositioning", SHORT, 1), + 532: ("ReferenceBlackWhite", RATIONAL, 6), + 700: ("XMP", BYTE, 0), + # Four private SGI tags + 32995: ("Matteing", SHORT, 1), + 32996: ("DataType", SHORT, 0), + 32997: ("ImageDepth", LONG, 1), + 32998: ("TileDepth", LONG, 1), + 33432: ("Copyright", ASCII, 1), + 33723: ("IptcNaaInfo", UNDEFINED, 1), + 34377: ("PhotoshopInfo", BYTE, 0), + # FIXME add more tags here + 34665: ("ExifIFD", LONG, 1), + 34675: ("ICCProfile", UNDEFINED, 1), + 34853: ("GPSInfoIFD", LONG, 1), + 36864: ("ExifVersion", UNDEFINED, 1), + 37724: ("ImageSourceData", UNDEFINED, 1), + 40965: ("InteroperabilityIFD", LONG, 1), + 41730: ("CFAPattern", UNDEFINED, 1), + # MPInfo + 45056: ("MPFVersion", UNDEFINED, 1), + 45057: ("NumberOfImages", LONG, 1), + 45058: ("MPEntry", UNDEFINED, 1), + 45059: ("ImageUIDList", UNDEFINED, 0), # UNDONE, check + 45060: ("TotalFrames", LONG, 1), + 45313: ("MPIndividualNum", LONG, 1), + 45569: ("PanOrientation", LONG, 1), + 45570: ("PanOverlap_H", RATIONAL, 1), + 45571: ("PanOverlap_V", RATIONAL, 1), + 45572: ("BaseViewpointNum", LONG, 1), + 45573: ("ConvergenceAngle", SIGNED_RATIONAL, 1), + 45574: ("BaselineLength", RATIONAL, 1), + 45575: ("VerticalDivergence", SIGNED_RATIONAL, 1), + 45576: ("AxisDistance_X", SIGNED_RATIONAL, 1), + 45577: ("AxisDistance_Y", SIGNED_RATIONAL, 1), + 45578: ("AxisDistance_Z", SIGNED_RATIONAL, 1), + 45579: ("YawAngle", SIGNED_RATIONAL, 1), + 45580: ("PitchAngle", SIGNED_RATIONAL, 1), + 45581: ("RollAngle", SIGNED_RATIONAL, 1), + 40960: ("FlashPixVersion", UNDEFINED, 1), + 50741: ("MakerNoteSafety", SHORT, 1, {"Unsafe": 0, "Safe": 1}), + 50780: ("BestQualityScale", RATIONAL, 1), + 50838: ("ImageJMetaDataByteCounts", LONG, 0), # Can be more than one + 50839: ("ImageJMetaData", UNDEFINED, 1), # see Issue #2006 +} +_tags_v2_groups = { + # ExifIFD + 34665: { + 36864: ("ExifVersion", UNDEFINED, 1), + 40960: ("FlashPixVersion", UNDEFINED, 1), + 40965: ("InteroperabilityIFD", LONG, 1), + 41730: ("CFAPattern", UNDEFINED, 1), + }, + # GPSInfoIFD + 34853: { + 0: ("GPSVersionID", BYTE, 4), + 1: ("GPSLatitudeRef", ASCII, 2), + 2: ("GPSLatitude", RATIONAL, 3), + 3: ("GPSLongitudeRef", ASCII, 2), + 4: ("GPSLongitude", RATIONAL, 3), + 5: ("GPSAltitudeRef", BYTE, 1), + 6: ("GPSAltitude", RATIONAL, 1), + 7: ("GPSTimeStamp", RATIONAL, 3), + 8: ("GPSSatellites", ASCII, 0), + 9: ("GPSStatus", ASCII, 2), + 10: ("GPSMeasureMode", ASCII, 2), + 11: ("GPSDOP", RATIONAL, 1), + 12: ("GPSSpeedRef", ASCII, 2), + 13: ("GPSSpeed", RATIONAL, 1), + 14: ("GPSTrackRef", ASCII, 2), + 15: ("GPSTrack", RATIONAL, 1), + 16: ("GPSImgDirectionRef", ASCII, 2), + 17: ("GPSImgDirection", RATIONAL, 1), + 18: ("GPSMapDatum", ASCII, 0), + 19: ("GPSDestLatitudeRef", ASCII, 2), + 20: ("GPSDestLatitude", RATIONAL, 3), + 21: ("GPSDestLongitudeRef", ASCII, 2), + 22: ("GPSDestLongitude", RATIONAL, 3), + 23: ("GPSDestBearingRef", ASCII, 2), + 24: ("GPSDestBearing", RATIONAL, 1), + 25: ("GPSDestDistanceRef", ASCII, 2), + 26: ("GPSDestDistance", RATIONAL, 1), + 27: ("GPSProcessingMethod", UNDEFINED, 0), + 28: ("GPSAreaInformation", UNDEFINED, 0), + 29: ("GPSDateStamp", ASCII, 11), + 30: ("GPSDifferential", SHORT, 1), + }, + # InteroperabilityIFD + 40965: {1: ("InteropIndex", ASCII, 1), 2: ("InteropVersion", UNDEFINED, 1)}, +} + +# Legacy Tags structure +# these tags aren't included above, but were in the previous versions +TAGS: dict[int | tuple[int, int], str] = { + 347: "JPEGTables", + 700: "XMP", + # Additional Exif Info + 32932: "Wang Annotation", + 33434: "ExposureTime", + 33437: "FNumber", + 33445: "MD FileTag", + 33446: "MD ScalePixel", + 33447: "MD ColorTable", + 33448: "MD LabName", + 33449: "MD SampleInfo", + 33450: "MD PrepDate", + 33451: "MD PrepTime", + 33452: "MD FileUnits", + 33550: "ModelPixelScaleTag", + 33723: "IptcNaaInfo", + 33918: "INGR Packet Data Tag", + 33919: "INGR Flag Registers", + 33920: "IrasB Transformation Matrix", + 33922: "ModelTiepointTag", + 34264: "ModelTransformationTag", + 34377: "PhotoshopInfo", + 34735: "GeoKeyDirectoryTag", + 34736: "GeoDoubleParamsTag", + 34737: "GeoAsciiParamsTag", + 34850: "ExposureProgram", + 34852: "SpectralSensitivity", + 34855: "ISOSpeedRatings", + 34856: "OECF", + 34864: "SensitivityType", + 34865: "StandardOutputSensitivity", + 34866: "RecommendedExposureIndex", + 34867: "ISOSpeed", + 34868: "ISOSpeedLatitudeyyy", + 34869: "ISOSpeedLatitudezzz", + 34908: "HylaFAX FaxRecvParams", + 34909: "HylaFAX FaxSubAddress", + 34910: "HylaFAX FaxRecvTime", + 36864: "ExifVersion", + 36867: "DateTimeOriginal", + 36868: "DateTimeDigitized", + 37121: "ComponentsConfiguration", + 37122: "CompressedBitsPerPixel", + 37724: "ImageSourceData", + 37377: "ShutterSpeedValue", + 37378: "ApertureValue", + 37379: "BrightnessValue", + 37380: "ExposureBiasValue", + 37381: "MaxApertureValue", + 37382: "SubjectDistance", + 37383: "MeteringMode", + 37384: "LightSource", + 37385: "Flash", + 37386: "FocalLength", + 37396: "SubjectArea", + 37500: "MakerNote", + 37510: "UserComment", + 37520: "SubSec", + 37521: "SubSecTimeOriginal", + 37522: "SubsecTimeDigitized", + 40960: "FlashPixVersion", + 40961: "ColorSpace", + 40962: "PixelXDimension", + 40963: "PixelYDimension", + 40964: "RelatedSoundFile", + 40965: "InteroperabilityIFD", + 41483: "FlashEnergy", + 41484: "SpatialFrequencyResponse", + 41486: "FocalPlaneXResolution", + 41487: "FocalPlaneYResolution", + 41488: "FocalPlaneResolutionUnit", + 41492: "SubjectLocation", + 41493: "ExposureIndex", + 41495: "SensingMethod", + 41728: "FileSource", + 41729: "SceneType", + 41730: "CFAPattern", + 41985: "CustomRendered", + 41986: "ExposureMode", + 41987: "WhiteBalance", + 41988: "DigitalZoomRatio", + 41989: "FocalLengthIn35mmFilm", + 41990: "SceneCaptureType", + 41991: "GainControl", + 41992: "Contrast", + 41993: "Saturation", + 41994: "Sharpness", + 41995: "DeviceSettingDescription", + 41996: "SubjectDistanceRange", + 42016: "ImageUniqueID", + 42032: "CameraOwnerName", + 42033: "BodySerialNumber", + 42034: "LensSpecification", + 42035: "LensMake", + 42036: "LensModel", + 42037: "LensSerialNumber", + 42112: "GDAL_METADATA", + 42113: "GDAL_NODATA", + 42240: "Gamma", + 50215: "Oce Scanjob Description", + 50216: "Oce Application Selector", + 50217: "Oce Identification Number", + 50218: "Oce ImageLogic Characteristics", + # Adobe DNG + 50706: "DNGVersion", + 50707: "DNGBackwardVersion", + 50708: "UniqueCameraModel", + 50709: "LocalizedCameraModel", + 50710: "CFAPlaneColor", + 50711: "CFALayout", + 50712: "LinearizationTable", + 50713: "BlackLevelRepeatDim", + 50714: "BlackLevel", + 50715: "BlackLevelDeltaH", + 50716: "BlackLevelDeltaV", + 50717: "WhiteLevel", + 50718: "DefaultScale", + 50719: "DefaultCropOrigin", + 50720: "DefaultCropSize", + 50721: "ColorMatrix1", + 50722: "ColorMatrix2", + 50723: "CameraCalibration1", + 50724: "CameraCalibration2", + 50725: "ReductionMatrix1", + 50726: "ReductionMatrix2", + 50727: "AnalogBalance", + 50728: "AsShotNeutral", + 50729: "AsShotWhiteXY", + 50730: "BaselineExposure", + 50731: "BaselineNoise", + 50732: "BaselineSharpness", + 50733: "BayerGreenSplit", + 50734: "LinearResponseLimit", + 50735: "CameraSerialNumber", + 50736: "LensInfo", + 50737: "ChromaBlurRadius", + 50738: "AntiAliasStrength", + 50740: "DNGPrivateData", + 50778: "CalibrationIlluminant1", + 50779: "CalibrationIlluminant2", + 50784: "Alias Layer Metadata", +} + +TAGS_V2: dict[int, TagInfo] = {} +TAGS_V2_GROUPS: dict[int, dict[int, TagInfo]] = {} + + +def _populate() -> None: + for k, v in _tags_v2.items(): + # Populate legacy structure. + TAGS[k] = v[0] + if len(v) == 4: + for sk, sv in v[3].items(): + TAGS[(k, sv)] = sk + + TAGS_V2[k] = TagInfo(k, *v) + + for group, tags in _tags_v2_groups.items(): + TAGS_V2_GROUPS[group] = {k: TagInfo(k, *v) for k, v in tags.items()} + + +_populate() +## +# Map type numbers to type names -- defined in ImageFileDirectory. + +TYPES: dict[int, str] = {} + +# +# These tags are handled by default in libtiff, without +# adding to the custom dictionary. From tif_dir.c, searching for +# case TIFFTAG in the _TIFFVSetField function: +# Line: item. +# 148: case TIFFTAG_SUBFILETYPE: +# 151: case TIFFTAG_IMAGEWIDTH: +# 154: case TIFFTAG_IMAGELENGTH: +# 157: case TIFFTAG_BITSPERSAMPLE: +# 181: case TIFFTAG_COMPRESSION: +# 202: case TIFFTAG_PHOTOMETRIC: +# 205: case TIFFTAG_THRESHHOLDING: +# 208: case TIFFTAG_FILLORDER: +# 214: case TIFFTAG_ORIENTATION: +# 221: case TIFFTAG_SAMPLESPERPIXEL: +# 228: case TIFFTAG_ROWSPERSTRIP: +# 238: case TIFFTAG_MINSAMPLEVALUE: +# 241: case TIFFTAG_MAXSAMPLEVALUE: +# 244: case TIFFTAG_SMINSAMPLEVALUE: +# 247: case TIFFTAG_SMAXSAMPLEVALUE: +# 250: case TIFFTAG_XRESOLUTION: +# 256: case TIFFTAG_YRESOLUTION: +# 262: case TIFFTAG_PLANARCONFIG: +# 268: case TIFFTAG_XPOSITION: +# 271: case TIFFTAG_YPOSITION: +# 274: case TIFFTAG_RESOLUTIONUNIT: +# 280: case TIFFTAG_PAGENUMBER: +# 284: case TIFFTAG_HALFTONEHINTS: +# 288: case TIFFTAG_COLORMAP: +# 294: case TIFFTAG_EXTRASAMPLES: +# 298: case TIFFTAG_MATTEING: +# 305: case TIFFTAG_TILEWIDTH: +# 316: case TIFFTAG_TILELENGTH: +# 327: case TIFFTAG_TILEDEPTH: +# 333: case TIFFTAG_DATATYPE: +# 344: case TIFFTAG_SAMPLEFORMAT: +# 361: case TIFFTAG_IMAGEDEPTH: +# 364: case TIFFTAG_SUBIFD: +# 376: case TIFFTAG_YCBCRPOSITIONING: +# 379: case TIFFTAG_YCBCRSUBSAMPLING: +# 383: case TIFFTAG_TRANSFERFUNCTION: +# 389: case TIFFTAG_REFERENCEBLACKWHITE: +# 393: case TIFFTAG_INKNAMES: + +# Following pseudo-tags are also handled by default in libtiff: +# TIFFTAG_JPEGQUALITY 65537 + +# some of these are not in our TAGS_V2 dict and were included from tiff.h + +# This list also exists in encode.c +LIBTIFF_CORE = { + 255, + 256, + 257, + 258, + 259, + 262, + 263, + 266, + 274, + 277, + 278, + 280, + 281, + 340, + 341, + 282, + 283, + 284, + 286, + 287, + 296, + 297, + 321, + 320, + 338, + 32995, + 322, + 323, + 32998, + 32996, + 339, + 32997, + 330, + 531, + 530, + 301, + 532, + 333, + # as above + 269, # this has been in our tests forever, and works + 65537, +} + +LIBTIFF_CORE.remove(255) # We don't have support for subfiletypes +LIBTIFF_CORE.remove(322) # We don't have support for writing tiled images with libtiff +LIBTIFF_CORE.remove(323) # Tiled images +LIBTIFF_CORE.remove(333) # Ink Names either + +# Note to advanced users: There may be combinations of these +# parameters and values that when added properly, will work and +# produce valid tiff images that may work in your application. +# It is safe to add and remove tags from this set from Pillow's point +# of view so long as you test against libtiff. diff --git a/venv/lib/python3.10/site-packages/PIL/WalImageFile.py b/venv/lib/python3.10/site-packages/PIL/WalImageFile.py new file mode 100644 index 0000000000000000000000000000000000000000..5494f62e89294882ff2e650cde23be5e70104680 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/WalImageFile.py @@ -0,0 +1,126 @@ +# +# The Python Imaging Library. +# $Id$ +# +# WAL file handling +# +# History: +# 2003-04-23 fl created +# +# Copyright (c) 2003 by Fredrik Lundh. +# +# See the README file for information on usage and redistribution. +# + +""" +This reader is based on the specification available from: +https://www.flipcode.com/archives/Quake_2_BSP_File_Format.shtml +and has been tested with a few sample files found using google. + +.. note:: + This format cannot be automatically recognized, so the reader + is not registered for use with :py:func:`PIL.Image.open()`. + To open a WAL file, use the :py:func:`PIL.WalImageFile.open()` function instead. +""" +from __future__ import annotations + +from typing import IO + +from . import Image, ImageFile +from ._binary import i32le as i32 +from ._typing import StrOrBytesPath + + +class WalImageFile(ImageFile.ImageFile): + format = "WAL" + format_description = "Quake2 Texture" + + def _open(self) -> None: + self._mode = "P" + + # read header fields + header = self.fp.read(32 + 24 + 32 + 12) + self._size = i32(header, 32), i32(header, 36) + Image._decompression_bomb_check(self.size) + + # load pixel data + offset = i32(header, 40) + self.fp.seek(offset) + + # strings are null-terminated + self.info["name"] = header[:32].split(b"\0", 1)[0] + if next_name := header[56 : 56 + 32].split(b"\0", 1)[0]: + self.info["next_name"] = next_name + + def load(self) -> Image.core.PixelAccess | None: + if self._im is None: + self.im = Image.core.new(self.mode, self.size) + self.frombytes(self.fp.read(self.size[0] * self.size[1])) + self.putpalette(quake2palette) + return Image.Image.load(self) + + +def open(filename: StrOrBytesPath | IO[bytes]) -> WalImageFile: + """ + Load texture from a Quake2 WAL texture file. + + By default, a Quake2 standard palette is attached to the texture. + To override the palette, use the :py:func:`PIL.Image.Image.putpalette()` method. + + :param filename: WAL file name, or an opened file handle. + :returns: An image instance. + """ + return WalImageFile(filename) + + +quake2palette = ( + # default palette taken from piffo 0.93 by Hans Häggström + b"\x01\x01\x01\x0b\x0b\x0b\x12\x12\x12\x17\x17\x17\x1b\x1b\x1b\x1e" + b"\x1e\x1e\x22\x22\x22\x26\x26\x26\x29\x29\x29\x2c\x2c\x2c\x2f\x2f" + b"\x2f\x32\x32\x32\x35\x35\x35\x37\x37\x37\x3a\x3a\x3a\x3c\x3c\x3c" + b"\x24\x1e\x13\x22\x1c\x12\x20\x1b\x12\x1f\x1a\x10\x1d\x19\x10\x1b" + b"\x17\x0f\x1a\x16\x0f\x18\x14\x0d\x17\x13\x0d\x16\x12\x0d\x14\x10" + b"\x0b\x13\x0f\x0b\x10\x0d\x0a\x0f\x0b\x0a\x0d\x0b\x07\x0b\x0a\x07" + b"\x23\x23\x26\x22\x22\x25\x22\x20\x23\x21\x1f\x22\x20\x1e\x20\x1f" + b"\x1d\x1e\x1d\x1b\x1c\x1b\x1a\x1a\x1a\x19\x19\x18\x17\x17\x17\x16" + b"\x16\x14\x14\x14\x13\x13\x13\x10\x10\x10\x0f\x0f\x0f\x0d\x0d\x0d" + b"\x2d\x28\x20\x29\x24\x1c\x27\x22\x1a\x25\x1f\x17\x38\x2e\x1e\x31" + b"\x29\x1a\x2c\x25\x17\x26\x20\x14\x3c\x30\x14\x37\x2c\x13\x33\x28" + b"\x12\x2d\x24\x10\x28\x1f\x0f\x22\x1a\x0b\x1b\x14\x0a\x13\x0f\x07" + b"\x31\x1a\x16\x30\x17\x13\x2e\x16\x10\x2c\x14\x0d\x2a\x12\x0b\x27" + b"\x0f\x0a\x25\x0f\x07\x21\x0d\x01\x1e\x0b\x01\x1c\x0b\x01\x1a\x0b" + b"\x01\x18\x0a\x01\x16\x0a\x01\x13\x0a\x01\x10\x07\x01\x0d\x07\x01" + b"\x29\x23\x1e\x27\x21\x1c\x26\x20\x1b\x25\x1f\x1a\x23\x1d\x19\x21" + b"\x1c\x18\x20\x1b\x17\x1e\x19\x16\x1c\x18\x14\x1b\x17\x13\x19\x14" + b"\x10\x17\x13\x0f\x14\x10\x0d\x12\x0f\x0b\x0f\x0b\x0a\x0b\x0a\x07" + b"\x26\x1a\x0f\x23\x19\x0f\x20\x17\x0f\x1c\x16\x0f\x19\x13\x0d\x14" + b"\x10\x0b\x10\x0d\x0a\x0b\x0a\x07\x33\x22\x1f\x35\x29\x26\x37\x2f" + b"\x2d\x39\x35\x34\x37\x39\x3a\x33\x37\x39\x30\x34\x36\x2b\x31\x34" + b"\x27\x2e\x31\x22\x2b\x2f\x1d\x28\x2c\x17\x25\x2a\x0f\x20\x26\x0d" + b"\x1e\x25\x0b\x1c\x22\x0a\x1b\x20\x07\x19\x1e\x07\x17\x1b\x07\x14" + b"\x18\x01\x12\x16\x01\x0f\x12\x01\x0b\x0d\x01\x07\x0a\x01\x01\x01" + b"\x2c\x21\x21\x2a\x1f\x1f\x29\x1d\x1d\x27\x1c\x1c\x26\x1a\x1a\x24" + b"\x18\x18\x22\x17\x17\x21\x16\x16\x1e\x13\x13\x1b\x12\x12\x18\x10" + b"\x10\x16\x0d\x0d\x12\x0b\x0b\x0d\x0a\x0a\x0a\x07\x07\x01\x01\x01" + b"\x2e\x30\x29\x2d\x2e\x27\x2b\x2c\x26\x2a\x2a\x24\x28\x29\x23\x27" + b"\x27\x21\x26\x26\x1f\x24\x24\x1d\x22\x22\x1c\x1f\x1f\x1a\x1c\x1c" + b"\x18\x19\x19\x16\x17\x17\x13\x13\x13\x10\x0f\x0f\x0d\x0b\x0b\x0a" + b"\x30\x1e\x1b\x2d\x1c\x19\x2c\x1a\x17\x2a\x19\x14\x28\x17\x13\x26" + b"\x16\x10\x24\x13\x0f\x21\x12\x0d\x1f\x10\x0b\x1c\x0f\x0a\x19\x0d" + b"\x0a\x16\x0b\x07\x12\x0a\x07\x0f\x07\x01\x0a\x01\x01\x01\x01\x01" + b"\x28\x29\x38\x26\x27\x36\x25\x26\x34\x24\x24\x31\x22\x22\x2f\x20" + b"\x21\x2d\x1e\x1f\x2a\x1d\x1d\x27\x1b\x1b\x25\x19\x19\x21\x17\x17" + b"\x1e\x14\x14\x1b\x13\x12\x17\x10\x0f\x13\x0d\x0b\x0f\x0a\x07\x07" + b"\x2f\x32\x29\x2d\x30\x26\x2b\x2e\x24\x29\x2c\x21\x27\x2a\x1e\x25" + b"\x28\x1c\x23\x26\x1a\x21\x25\x18\x1e\x22\x14\x1b\x1f\x10\x19\x1c" + b"\x0d\x17\x1a\x0a\x13\x17\x07\x10\x13\x01\x0d\x0f\x01\x0a\x0b\x01" + b"\x01\x3f\x01\x13\x3c\x0b\x1b\x39\x10\x20\x35\x14\x23\x31\x17\x23" + b"\x2d\x18\x23\x29\x18\x3f\x3f\x3f\x3f\x3f\x39\x3f\x3f\x31\x3f\x3f" + b"\x2a\x3f\x3f\x20\x3f\x3f\x14\x3f\x3c\x12\x3f\x39\x0f\x3f\x35\x0b" + b"\x3f\x32\x07\x3f\x2d\x01\x3d\x2a\x01\x3b\x26\x01\x39\x21\x01\x37" + b"\x1d\x01\x34\x1a\x01\x32\x16\x01\x2f\x12\x01\x2d\x0f\x01\x2a\x0b" + b"\x01\x27\x07\x01\x23\x01\x01\x1d\x01\x01\x17\x01\x01\x10\x01\x01" + b"\x3d\x01\x01\x19\x19\x3f\x3f\x01\x01\x01\x01\x3f\x16\x16\x13\x10" + b"\x10\x0f\x0d\x0d\x0b\x3c\x2e\x2a\x36\x27\x20\x30\x21\x18\x29\x1b" + b"\x10\x3c\x39\x37\x37\x32\x2f\x31\x2c\x28\x2b\x26\x21\x30\x22\x20" +) diff --git a/venv/lib/python3.10/site-packages/PIL/WebPImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/WebPImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..2847fed20a0f08299185a71b8058a9cbc58ef9c5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/WebPImagePlugin.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +from io import BytesIO + +from . import Image, ImageFile + +try: + from . import _webp + + SUPPORTED = True +except ImportError: + SUPPORTED = False + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import IO, Any + +_VP8_MODES_BY_IDENTIFIER = { + b"VP8 ": "RGB", + b"VP8X": "RGBA", + b"VP8L": "RGBA", # lossless +} + + +def _accept(prefix: bytes) -> bool | str: + is_riff_file_format = prefix.startswith(b"RIFF") + is_webp_file = prefix[8:12] == b"WEBP" + is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER + + if is_riff_file_format and is_webp_file and is_valid_vp8_mode: + if not SUPPORTED: + return ( + "image file could not be identified because WEBP support not installed" + ) + return True + return False + + +class WebPImageFile(ImageFile.ImageFile): + format = "WEBP" + format_description = "WebP image" + __loaded = 0 + __logical_frame = 0 + + def _open(self) -> None: + # Use the newer AnimDecoder API to parse the (possibly) animated file, + # and access muxed chunks like ICC/EXIF/XMP. + self._decoder = _webp.WebPAnimDecoder(self.fp.read()) + + # Get info from decoder + self._size, loop_count, bgcolor, frame_count, mode = self._decoder.get_info() + self.info["loop"] = loop_count + bg_a, bg_r, bg_g, bg_b = ( + (bgcolor >> 24) & 0xFF, + (bgcolor >> 16) & 0xFF, + (bgcolor >> 8) & 0xFF, + bgcolor & 0xFF, + ) + self.info["background"] = (bg_r, bg_g, bg_b, bg_a) + self.n_frames = frame_count + self.is_animated = self.n_frames > 1 + self._mode = "RGB" if mode == "RGBX" else mode + self.rawmode = mode + + # Attempt to read ICC / EXIF / XMP chunks from file + icc_profile = self._decoder.get_chunk("ICCP") + exif = self._decoder.get_chunk("EXIF") + xmp = self._decoder.get_chunk("XMP ") + if icc_profile: + self.info["icc_profile"] = icc_profile + if exif: + self.info["exif"] = exif + if xmp: + self.info["xmp"] = xmp + + # Initialize seek state + self._reset(reset=False) + + def _getexif(self) -> dict[int, Any] | None: + if "exif" not in self.info: + return None + return self.getexif()._get_merged_dict() + + def seek(self, frame: int) -> None: + if not self._seek_check(frame): + return + + # Set logical frame to requested position + self.__logical_frame = frame + + def _reset(self, reset: bool = True) -> None: + if reset: + self._decoder.reset() + self.__physical_frame = 0 + self.__loaded = -1 + self.__timestamp = 0 + + def _get_next(self) -> tuple[bytes, int, int]: + # Get next frame + ret = self._decoder.get_next() + self.__physical_frame += 1 + + # Check if an error occurred + if ret is None: + self._reset() # Reset just to be safe + self.seek(0) + msg = "failed to decode next frame in WebP file" + raise EOFError(msg) + + # Compute duration + data, timestamp = ret + duration = timestamp - self.__timestamp + self.__timestamp = timestamp + + # libwebp gives frame end, adjust to start of frame + timestamp -= duration + return data, timestamp, duration + + def _seek(self, frame: int) -> None: + if self.__physical_frame == frame: + return # Nothing to do + if frame < self.__physical_frame: + self._reset() # Rewind to beginning + while self.__physical_frame < frame: + self._get_next() # Advance to the requested frame + + def load(self) -> Image.core.PixelAccess | None: + if self.__loaded != self.__logical_frame: + self._seek(self.__logical_frame) + + # We need to load the image data for this frame + data, timestamp, duration = self._get_next() + self.info["timestamp"] = timestamp + self.info["duration"] = duration + self.__loaded = self.__logical_frame + + # Set tile + if self.fp and self._exclusive_fp: + self.fp.close() + self.fp = BytesIO(data) + self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.rawmode)] + + return super().load() + + def load_seek(self, pos: int) -> None: + pass + + def tell(self) -> int: + return self.__logical_frame + + +def _convert_frame(im: Image.Image) -> Image.Image: + # Make sure image mode is supported + if im.mode not in ("RGBX", "RGBA", "RGB"): + im = im.convert("RGBA" if im.has_transparency_data else "RGB") + return im + + +def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + encoderinfo = im.encoderinfo.copy() + append_images = list(encoderinfo.get("append_images", [])) + + # If total frame count is 1, then save using the legacy API, which + # will preserve non-alpha modes + total = 0 + for ims in [im] + append_images: + total += getattr(ims, "n_frames", 1) + if total == 1: + _save(im, fp, filename) + return + + background: int | tuple[int, ...] = (0, 0, 0, 0) + if "background" in encoderinfo: + background = encoderinfo["background"] + elif "background" in im.info: + background = im.info["background"] + if isinstance(background, int): + # GifImagePlugin stores a global color table index in + # info["background"]. So it must be converted to an RGBA value + palette = im.getpalette() + if palette: + r, g, b = palette[background * 3 : (background + 1) * 3] + background = (r, g, b, 255) + else: + background = (background, background, background, 255) + + duration = im.encoderinfo.get("duration", im.info.get("duration", 0)) + loop = im.encoderinfo.get("loop", 0) + minimize_size = im.encoderinfo.get("minimize_size", False) + kmin = im.encoderinfo.get("kmin", None) + kmax = im.encoderinfo.get("kmax", None) + allow_mixed = im.encoderinfo.get("allow_mixed", False) + verbose = False + lossless = im.encoderinfo.get("lossless", False) + quality = im.encoderinfo.get("quality", 80) + alpha_quality = im.encoderinfo.get("alpha_quality", 100) + method = im.encoderinfo.get("method", 0) + icc_profile = im.encoderinfo.get("icc_profile") or "" + exif = im.encoderinfo.get("exif", "") + if isinstance(exif, Image.Exif): + exif = exif.tobytes() + xmp = im.encoderinfo.get("xmp", "") + if allow_mixed: + lossless = False + + # Sensible keyframe defaults are from gif2webp.c script + if kmin is None: + kmin = 9 if lossless else 3 + if kmax is None: + kmax = 17 if lossless else 5 + + # Validate background color + if ( + not isinstance(background, (list, tuple)) + or len(background) != 4 + or not all(0 <= v < 256 for v in background) + ): + msg = f"Background color is not an RGBA tuple clamped to (0-255): {background}" + raise OSError(msg) + + # Convert to packed uint + bg_r, bg_g, bg_b, bg_a = background + background = (bg_a << 24) | (bg_r << 16) | (bg_g << 8) | (bg_b << 0) + + # Setup the WebP animation encoder + enc = _webp.WebPAnimEncoder( + im.size, + background, + loop, + minimize_size, + kmin, + kmax, + allow_mixed, + verbose, + ) + + # Add each frame + frame_idx = 0 + timestamp = 0 + cur_idx = im.tell() + try: + for ims in [im] + append_images: + # Get number of frames in this image + nfr = getattr(ims, "n_frames", 1) + + for idx in range(nfr): + ims.seek(idx) + + frame = _convert_frame(ims) + + # Append the frame to the animation encoder + enc.add( + frame.getim(), + round(timestamp), + lossless, + quality, + alpha_quality, + method, + ) + + # Update timestamp and frame index + if isinstance(duration, (list, tuple)): + timestamp += duration[frame_idx] + else: + timestamp += duration + frame_idx += 1 + + finally: + im.seek(cur_idx) + + # Force encoder to flush frames + enc.add(None, round(timestamp), lossless, quality, alpha_quality, 0) + + # Get the final output from the encoder + data = enc.assemble(icc_profile, exif, xmp) + if data is None: + msg = "cannot write file as WebP (encoder returned None)" + raise OSError(msg) + + fp.write(data) + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + lossless = im.encoderinfo.get("lossless", False) + quality = im.encoderinfo.get("quality", 80) + alpha_quality = im.encoderinfo.get("alpha_quality", 100) + icc_profile = im.encoderinfo.get("icc_profile") or "" + exif = im.encoderinfo.get("exif", b"") + if isinstance(exif, Image.Exif): + exif = exif.tobytes() + if exif.startswith(b"Exif\x00\x00"): + exif = exif[6:] + xmp = im.encoderinfo.get("xmp", "") + method = im.encoderinfo.get("method", 4) + exact = 1 if im.encoderinfo.get("exact") else 0 + + im = _convert_frame(im) + + data = _webp.WebPEncode( + im.getim(), + lossless, + float(quality), + float(alpha_quality), + icc_profile, + method, + exact, + exif, + xmp, + ) + if data is None: + msg = "cannot write file as WebP (encoder returned None)" + raise OSError(msg) + + fp.write(data) + + +Image.register_open(WebPImageFile.format, WebPImageFile, _accept) +if SUPPORTED: + Image.register_save(WebPImageFile.format, _save) + Image.register_save_all(WebPImageFile.format, _save_all) + Image.register_extension(WebPImageFile.format, ".webp") + Image.register_mime(WebPImageFile.format, "image/webp") diff --git a/venv/lib/python3.10/site-packages/PIL/WmfImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/WmfImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..de714d337945f2d41046f9962f0fe034effab45e --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/WmfImagePlugin.py @@ -0,0 +1,186 @@ +# +# The Python Imaging Library +# $Id$ +# +# WMF stub codec +# +# history: +# 1996-12-14 fl Created +# 2004-02-22 fl Turned into a stub driver +# 2004-02-23 fl Added EMF support +# +# Copyright (c) Secret Labs AB 1997-2004. All rights reserved. +# Copyright (c) Fredrik Lundh 1996. +# +# See the README file for information on usage and redistribution. +# +# WMF/EMF reference documentation: +# https://winprotocoldoc.blob.core.windows.net/productionwindowsarchives/MS-WMF/[MS-WMF].pdf +# http://wvware.sourceforge.net/caolan/index.html +# http://wvware.sourceforge.net/caolan/ora-wmf.html +from __future__ import annotations + +from typing import IO + +from . import Image, ImageFile +from ._binary import i16le as word +from ._binary import si16le as short +from ._binary import si32le as _long + +_handler = None + + +def register_handler(handler: ImageFile.StubHandler | None) -> None: + """ + Install application-specific WMF image handler. + + :param handler: Handler object. + """ + global _handler + _handler = handler + + +if hasattr(Image.core, "drawwmf"): + # install default handler (windows only) + + class WmfHandler(ImageFile.StubHandler): + def open(self, im: ImageFile.StubImageFile) -> None: + im._mode = "RGB" + self.bbox = im.info["wmf_bbox"] + + def load(self, im: ImageFile.StubImageFile) -> Image.Image: + im.fp.seek(0) # rewind + return Image.frombytes( + "RGB", + im.size, + Image.core.drawwmf(im.fp.read(), im.size, self.bbox), + "raw", + "BGR", + (im.size[0] * 3 + 3) & -4, + -1, + ) + + register_handler(WmfHandler()) + +# +# -------------------------------------------------------------------- +# Read WMF file + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith((b"\xd7\xcd\xc6\x9a\x00\x00", b"\x01\x00\x00\x00")) + + +## +# Image plugin for Windows metafiles. + + +class WmfStubImageFile(ImageFile.StubImageFile): + format = "WMF" + format_description = "Windows Metafile" + + def _open(self) -> None: + # check placeable header + s = self.fp.read(44) + + if s.startswith(b"\xd7\xcd\xc6\x9a\x00\x00"): + # placeable windows metafile + + # get units per inch + inch = word(s, 14) + if inch == 0: + msg = "Invalid inch" + raise ValueError(msg) + self._inch: tuple[float, float] = inch, inch + + # get bounding box + x0 = short(s, 6) + y0 = short(s, 8) + x1 = short(s, 10) + y1 = short(s, 12) + + # normalize size to 72 dots per inch + self.info["dpi"] = 72 + size = ( + (x1 - x0) * self.info["dpi"] // inch, + (y1 - y0) * self.info["dpi"] // inch, + ) + + self.info["wmf_bbox"] = x0, y0, x1, y1 + + # sanity check (standard metafile header) + if s[22:26] != b"\x01\x00\t\x00": + msg = "Unsupported WMF file format" + raise SyntaxError(msg) + + elif s.startswith(b"\x01\x00\x00\x00") and s[40:44] == b" EMF": + # enhanced metafile + + # get bounding box + x0 = _long(s, 8) + y0 = _long(s, 12) + x1 = _long(s, 16) + y1 = _long(s, 20) + + # get frame (in 0.01 millimeter units) + frame = _long(s, 24), _long(s, 28), _long(s, 32), _long(s, 36) + + size = x1 - x0, y1 - y0 + + # calculate dots per inch from bbox and frame + xdpi = 2540.0 * (x1 - x0) / (frame[2] - frame[0]) + ydpi = 2540.0 * (y1 - y0) / (frame[3] - frame[1]) + + self.info["wmf_bbox"] = x0, y0, x1, y1 + + if xdpi == ydpi: + self.info["dpi"] = xdpi + else: + self.info["dpi"] = xdpi, ydpi + self._inch = xdpi, ydpi + + else: + msg = "Unsupported file format" + raise SyntaxError(msg) + + self._mode = "RGB" + self._size = size + + loader = self._load() + if loader: + loader.open(self) + + def _load(self) -> ImageFile.StubHandler | None: + return _handler + + def load( + self, dpi: float | tuple[float, float] | None = None + ) -> Image.core.PixelAccess | None: + if dpi is not None: + self.info["dpi"] = dpi + x0, y0, x1, y1 = self.info["wmf_bbox"] + if not isinstance(dpi, tuple): + dpi = dpi, dpi + self._size = ( + int((x1 - x0) * dpi[0] / self._inch[0]), + int((y1 - y0) * dpi[1] / self._inch[1]), + ) + return super().load() + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if _handler is None or not hasattr(_handler, "save"): + msg = "WMF save handler not installed" + raise OSError(msg) + _handler.save(im, fp, filename) + + +# +# -------------------------------------------------------------------- +# Registry stuff + + +Image.register_open(WmfStubImageFile.format, WmfStubImageFile, _accept) +Image.register_save(WmfStubImageFile.format, _save) + +Image.register_extensions(WmfStubImageFile.format, [".wmf", ".emf"]) diff --git a/venv/lib/python3.10/site-packages/PIL/XVThumbImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/XVThumbImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..cde28388ff0535262770dd0336ee2b48db2178ba --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/XVThumbImagePlugin.py @@ -0,0 +1,83 @@ +# +# The Python Imaging Library. +# $Id$ +# +# XV Thumbnail file handler by Charles E. "Gene" Cash +# (gcash@magicnet.net) +# +# see xvcolor.c and xvbrowse.c in the sources to John Bradley's XV, +# available from ftp://ftp.cis.upenn.edu/pub/xv/ +# +# history: +# 98-08-15 cec created (b/w only) +# 98-12-09 cec added color palette +# 98-12-28 fl added to PIL (with only a few very minor modifications) +# +# To do: +# FIXME: make save work (this requires quantization support) +# +from __future__ import annotations + +from . import Image, ImageFile, ImagePalette +from ._binary import o8 + +_MAGIC = b"P7 332" + +# standard color palette for thumbnails (RGB332) +PALETTE = b"" +for r in range(8): + for g in range(8): + for b in range(4): + PALETTE = PALETTE + ( + o8((r * 255) // 7) + o8((g * 255) // 7) + o8((b * 255) // 3) + ) + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(_MAGIC) + + +## +# Image plugin for XV thumbnail images. + + +class XVThumbImageFile(ImageFile.ImageFile): + format = "XVThumb" + format_description = "XV thumbnail image" + + def _open(self) -> None: + # check magic + assert self.fp is not None + + if not _accept(self.fp.read(6)): + msg = "not an XV thumbnail file" + raise SyntaxError(msg) + + # Skip to beginning of next line + self.fp.readline() + + # skip info comments + while True: + s = self.fp.readline() + if not s: + msg = "Unexpected EOF reading XV thumbnail file" + raise SyntaxError(msg) + if s[0] != 35: # ie. when not a comment: '#' + break + + # parse header line (already read) + s = s.strip().split() + + self._mode = "P" + self._size = int(s[0]), int(s[1]) + + self.palette = ImagePalette.raw("RGB", PALETTE) + + self.tile = [ + ImageFile._Tile("raw", (0, 0) + self.size, self.fp.tell(), self.mode) + ] + + +# -------------------------------------------------------------------- + +Image.register_open(XVThumbImageFile.format, XVThumbImageFile, _accept) diff --git a/venv/lib/python3.10/site-packages/PIL/XbmImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/XbmImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..1e57aa162ea4f8618dac66cf042352f73d2199c8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/XbmImagePlugin.py @@ -0,0 +1,98 @@ +# +# The Python Imaging Library. +# $Id$ +# +# XBM File handling +# +# History: +# 1995-09-08 fl Created +# 1996-11-01 fl Added save support +# 1997-07-07 fl Made header parser more tolerant +# 1997-07-22 fl Fixed yet another parser bug +# 2001-02-17 fl Use 're' instead of 'regex' (Python 2.1) (0.4) +# 2001-05-13 fl Added hotspot handling (based on code from Bernhard Herzog) +# 2004-02-24 fl Allow some whitespace before first #define +# +# Copyright (c) 1997-2004 by Secret Labs AB +# Copyright (c) 1996-1997 by Fredrik Lundh +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import re +from typing import IO + +from . import Image, ImageFile + +# XBM header +xbm_head = re.compile( + rb"\s*#define[ \t]+.*_width[ \t]+(?P[0-9]+)[\r\n]+" + b"#define[ \t]+.*_height[ \t]+(?P[0-9]+)[\r\n]+" + b"(?P" + b"#define[ \t]+[^_]*_x_hot[ \t]+(?P[0-9]+)[\r\n]+" + b"#define[ \t]+[^_]*_y_hot[ \t]+(?P[0-9]+)[\r\n]+" + b")?" + rb"[\000-\377]*_bits\[]" +) + + +def _accept(prefix: bytes) -> bool: + return prefix.lstrip().startswith(b"#define") + + +## +# Image plugin for X11 bitmaps. + + +class XbmImageFile(ImageFile.ImageFile): + format = "XBM" + format_description = "X11 Bitmap" + + def _open(self) -> None: + assert self.fp is not None + + m = xbm_head.match(self.fp.read(512)) + + if not m: + msg = "not a XBM file" + raise SyntaxError(msg) + + xsize = int(m.group("width")) + ysize = int(m.group("height")) + + if m.group("hotspot"): + self.info["hotspot"] = (int(m.group("xhot")), int(m.group("yhot"))) + + self._mode = "1" + self._size = xsize, ysize + + self.tile = [ImageFile._Tile("xbm", (0, 0) + self.size, m.end())] + + +def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: + if im.mode != "1": + msg = f"cannot write mode {im.mode} as XBM" + raise OSError(msg) + + fp.write(f"#define im_width {im.size[0]}\n".encode("ascii")) + fp.write(f"#define im_height {im.size[1]}\n".encode("ascii")) + + hotspot = im.encoderinfo.get("hotspot") + if hotspot: + fp.write(f"#define im_x_hot {hotspot[0]}\n".encode("ascii")) + fp.write(f"#define im_y_hot {hotspot[1]}\n".encode("ascii")) + + fp.write(b"static char im_bits[] = {\n") + + ImageFile._save(im, fp, [ImageFile._Tile("xbm", (0, 0) + im.size)]) + + fp.write(b"};\n") + + +Image.register_open(XbmImageFile.format, XbmImageFile, _accept) +Image.register_save(XbmImageFile.format, _save) + +Image.register_extension(XbmImageFile.format, ".xbm") + +Image.register_mime(XbmImageFile.format, "image/xbm") diff --git a/venv/lib/python3.10/site-packages/PIL/XpmImagePlugin.py b/venv/lib/python3.10/site-packages/PIL/XpmImagePlugin.py new file mode 100644 index 0000000000000000000000000000000000000000..3be240fbc1aeb7660de46fbd4f99f309ce9915dd --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/XpmImagePlugin.py @@ -0,0 +1,157 @@ +# +# The Python Imaging Library. +# $Id$ +# +# XPM File handling +# +# History: +# 1996-12-29 fl Created +# 2001-02-17 fl Use 're' instead of 'regex' (Python 2.1) (0.7) +# +# Copyright (c) Secret Labs AB 1997-2001. +# Copyright (c) Fredrik Lundh 1996-2001. +# +# See the README file for information on usage and redistribution. +# +from __future__ import annotations + +import re + +from . import Image, ImageFile, ImagePalette +from ._binary import o8 + +# XPM header +xpm_head = re.compile(b'"([0-9]*) ([0-9]*) ([0-9]*) ([0-9]*)') + + +def _accept(prefix: bytes) -> bool: + return prefix.startswith(b"/* XPM */") + + +## +# Image plugin for X11 pixel maps. + + +class XpmImageFile(ImageFile.ImageFile): + format = "XPM" + format_description = "X11 Pixel Map" + + def _open(self) -> None: + assert self.fp is not None + if not _accept(self.fp.read(9)): + msg = "not an XPM file" + raise SyntaxError(msg) + + # skip forward to next string + while True: + line = self.fp.readline() + if not line: + msg = "broken XPM file" + raise SyntaxError(msg) + m = xpm_head.match(line) + if m: + break + + self._size = int(m.group(1)), int(m.group(2)) + + palette_length = int(m.group(3)) + bpp = int(m.group(4)) + + # + # load palette description + + palette = {} + + for _ in range(palette_length): + line = self.fp.readline().rstrip() + + c = line[1 : bpp + 1] + s = line[bpp + 1 : -2].split() + + for i in range(0, len(s), 2): + if s[i] == b"c": + # process colour key + rgb = s[i + 1] + if rgb == b"None": + self.info["transparency"] = c + elif rgb.startswith(b"#"): + rgb_int = int(rgb[1:], 16) + palette[c] = ( + o8((rgb_int >> 16) & 255) + + o8((rgb_int >> 8) & 255) + + o8(rgb_int & 255) + ) + else: + # unknown colour + msg = "cannot read this XPM file" + raise ValueError(msg) + break + + else: + # missing colour key + msg = "cannot read this XPM file" + raise ValueError(msg) + + args: tuple[int, dict[bytes, bytes] | tuple[bytes, ...]] + if palette_length > 256: + self._mode = "RGB" + args = (bpp, palette) + else: + self._mode = "P" + self.palette = ImagePalette.raw("RGB", b"".join(palette.values())) + args = (bpp, tuple(palette.keys())) + + self.tile = [ImageFile._Tile("xpm", (0, 0) + self.size, self.fp.tell(), args)] + + def load_read(self, read_bytes: int) -> bytes: + # + # load all image data in one chunk + + xsize, ysize = self.size + + assert self.fp is not None + s = [self.fp.readline()[1 : xsize + 1].ljust(xsize) for i in range(ysize)] + + return b"".join(s) + + +class XpmDecoder(ImageFile.PyDecoder): + _pulls_fd = True + + def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]: + assert self.fd is not None + + data = bytearray() + bpp, palette = self.args + dest_length = self.state.xsize * self.state.ysize + if self.mode == "RGB": + dest_length *= 3 + pixel_header = False + while len(data) < dest_length: + line = self.fd.readline() + if not line: + break + if line.rstrip() == b"/* pixels */" and not pixel_header: + pixel_header = True + continue + line = b'"'.join(line.split(b'"')[1:-1]) + for i in range(0, len(line), bpp): + key = line[i : i + bpp] + if self.mode == "RGB": + data += palette[key] + else: + data += o8(palette.index(key)) + self.set_as_raw(bytes(data)) + return -1, 0 + + +# +# Registry + + +Image.register_open(XpmImageFile.format, XpmImageFile, _accept) +Image.register_decoder("xpm", XpmDecoder) + +Image.register_extension(XpmImageFile.format, ".xpm") + +Image.register_mime(XpmImageFile.format, "image/xpm") diff --git a/venv/lib/python3.10/site-packages/PIL/__init__.py b/venv/lib/python3.10/site-packages/PIL/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6e4c23f897f83ef72fc10070bd22e9dc70614cf9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/__init__.py @@ -0,0 +1,87 @@ +"""Pillow (Fork of the Python Imaging Library) + +Pillow is the friendly PIL fork by Jeffrey A. Clark and contributors. + https://github.com/python-pillow/Pillow/ + +Pillow is forked from PIL 1.1.7. + +PIL is the Python Imaging Library by Fredrik Lundh and contributors. +Copyright (c) 1999 by Secret Labs AB. + +Use PIL.__version__ for this Pillow version. + +;-) +""" + +from __future__ import annotations + +from . import _version + +# VERSION was removed in Pillow 6.0.0. +# PILLOW_VERSION was removed in Pillow 9.0.0. +# Use __version__ instead. +__version__ = _version.__version__ +del _version + + +_plugins = [ + "AvifImagePlugin", + "BlpImagePlugin", + "BmpImagePlugin", + "BufrStubImagePlugin", + "CurImagePlugin", + "DcxImagePlugin", + "DdsImagePlugin", + "EpsImagePlugin", + "FitsImagePlugin", + "FliImagePlugin", + "FpxImagePlugin", + "FtexImagePlugin", + "GbrImagePlugin", + "GifImagePlugin", + "GribStubImagePlugin", + "Hdf5StubImagePlugin", + "IcnsImagePlugin", + "IcoImagePlugin", + "ImImagePlugin", + "ImtImagePlugin", + "IptcImagePlugin", + "JpegImagePlugin", + "Jpeg2KImagePlugin", + "McIdasImagePlugin", + "MicImagePlugin", + "MpegImagePlugin", + "MpoImagePlugin", + "MspImagePlugin", + "PalmImagePlugin", + "PcdImagePlugin", + "PcxImagePlugin", + "PdfImagePlugin", + "PixarImagePlugin", + "PngImagePlugin", + "PpmImagePlugin", + "PsdImagePlugin", + "QoiImagePlugin", + "SgiImagePlugin", + "SpiderImagePlugin", + "SunImagePlugin", + "TgaImagePlugin", + "TiffImagePlugin", + "WebPImagePlugin", + "WmfImagePlugin", + "XbmImagePlugin", + "XpmImagePlugin", + "XVThumbImagePlugin", +] + + +class UnidentifiedImageError(OSError): + """ + Raised in :py:meth:`PIL.Image.open` if an image cannot be opened and identified. + + If a PNG image raises this error, setting :data:`.ImageFile.LOAD_TRUNCATED_IMAGES` + to true may allow the image to be opened after all. The setting will ignore missing + data and checksum failures. + """ + + pass diff --git a/venv/lib/python3.10/site-packages/PIL/__main__.py b/venv/lib/python3.10/site-packages/PIL/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..043156e892dadc4fb1222b33f5eda33251cd15aa --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/__main__.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +import sys + +from .features import pilinfo + +pilinfo(supported_formats="--report" not in sys.argv) diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/AvifImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/AvifImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f91682fe9fe7b63123c04f0e815fc802ef633ac Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/AvifImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/BdfFontFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/BdfFontFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a0a0909606fb0f71ff563de6ede1d2ab82c0b4a Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/BdfFontFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/BlpImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/BlpImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e50167387f9fd535a08c8324e6443d6bbec9e4c Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/BlpImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/BmpImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/BmpImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e989e5d597a27507a342761e891c6a5aa9ae536 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/BmpImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/BufrStubImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/BufrStubImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be725a47b48febb63b7ade8755ab86da115b7593 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/BufrStubImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ContainerIO.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ContainerIO.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da12a6cf069ba957a5e29dfe60c1007481156a97 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ContainerIO.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/CurImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/CurImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0bf53467fb8439b0a5e4860d9c860235a63457e2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/CurImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/DcxImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/DcxImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93b86ac286531252fda41007c95bb8c221f5e065 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/DcxImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/DdsImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/DdsImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..093b182763fdf54988f7de6ceca9c21343039735 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/DdsImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/EpsImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/EpsImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36ffaea05086c912bcbb357bf113e808d2846a8c Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/EpsImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ExifTags.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ExifTags.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..436cbb96485fb17051cd50dddbe23db06a559030 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ExifTags.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/FitsImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/FitsImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8b38894ff14f66715e0e7b6ae36fb7bf66e8cee Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/FitsImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/FliImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/FliImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0f491d74721b0745d51c84f2220e23df5164ef8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/FliImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/FontFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/FontFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e214099bba9fae82485fe7d0f5eb011bd681a94 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/FontFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/FpxImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/FpxImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b95a4d5dd6f6e26ad335f8076e26a601e4c9b6d9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/FpxImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/FtexImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/FtexImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f3474cabb123a9eb1fa818a777b84fd194bfe02 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/FtexImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GbrImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GbrImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7677df524fc6f849f8b161f0975fb9712ccf129e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GbrImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GdImageFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GdImageFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b79229cd0c238f33136a85643667a6f394df773 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GdImageFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GifImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GifImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d67ea442e3a2376ee19a1e1f25b34d1e494517c Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GifImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpGradientFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpGradientFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f2f44de6522d8b6d23979e2bc585b3ed494cc74 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpGradientFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpPaletteFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpPaletteFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6f92ba6c98d887e8c762f41496e08e35807da1c Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GimpPaletteFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/GribStubImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/GribStubImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39bbe71680e0738ce0c6142fd1805e4595a318e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/GribStubImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/Hdf5StubImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/Hdf5StubImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4973638335ed6da91279b9024daf6a99f85dbcdd Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/Hdf5StubImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/IcnsImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/IcnsImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50b3a6600275ecc7c64850c1cd3323b12d41f9c6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/IcnsImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/IcoImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/IcoImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95e72e1fcfa88d07f6e4469d66aec43249d75f7b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/IcoImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ccdb7ac19e02752ad7e5f10e6308b6910836c42 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageChops.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageChops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94de868be0168df27cae6eb47b2194b7082a1035 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageChops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageCms.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageCms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56b2786af878434c144a81b25d614aab836a9eef Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageCms.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageColor.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageColor.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a5897513d33d196fc5e60fea158de77c79ef137 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageColor.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d68909a1b960f1bf09abd02715e6434eb720af1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw2.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw2.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4bfc7cb899a686ad8f8a3e2a73ea37a334590f1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageDraw2.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageEnhance.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageEnhance.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71f9030a60a3b5db98b7c10f47ef779ed725f7cb Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageEnhance.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6155fc3657a653d3741b9b55e12300378eac1d9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFilter.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFilter.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7034b203a599c67e4ba09649c9cd7a1242dae5e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFilter.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFont.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFont.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f90900739cb793bead082fb134e7ed152721cab8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageFont.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageGrab.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageGrab.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34490a10561876cde144d4cb192a81227dfa7247 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageGrab.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMath.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMath.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f9fafd72479f30fdc3c41188a9f94109bb80eac Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMath.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMode.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMode.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b49096c126e49daa4f2c1bbd1874f95ff4862ff Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMode.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMorph.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMorph.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38e102c2f4bb575008e9160c8a72cdd589f5c204 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageMorph.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageOps.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageOps.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e4d02e102f20efd18c54fe294ad5b2ba0612016 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageOps.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePalette.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePalette.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2a8036f9e716fd138e16568950fa1130a743d2a Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePalette.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePath.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePath.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0dbee43d696e8de03233c5a6be0ed912dcc983b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImagePath.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageQt.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageQt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2760bbc538dc336a36a5436f27e86fd76039ef6d Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageQt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageSequence.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageSequence.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a3236a253d40900ed1b42edcb663761084e6988 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageSequence.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageShow.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageShow.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..504af2dbceff8f7ed3fdcb21abff06d4e3f949c0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageShow.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageStat.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageStat.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b93e0f1d152d3bb6dc2733a8ff14dd4858751c9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageStat.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageText.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageText.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b64d2a34e463830943406e607247dd1af93cfcc4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageText.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTk.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTk.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e95693471d6feba480732d4f7190189dc28c76d Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTk.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTransform.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTransform.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..312a78b20e6b43342fb1c2d5851048de90f465cc Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageTransform.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageWin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageWin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0df710c45a65b6487979b790495953bec8da89e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImageWin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/ImtImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImtImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ab3000d2d79e564d353dd9960a5c0ca0d9e37ca Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/ImtImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/IptcImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/IptcImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..efe4805580f9088387c5504b0c02749d450aa1f1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/IptcImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/Jpeg2KImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/Jpeg2KImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..041611588913f8feff010da0dc53fd1409b09301 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/Jpeg2KImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3474bfce2e0b58f830eab99fa8047d79c10fffd3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegPresets.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegPresets.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96e0763599b140558a081595e41cfee0a34d25c3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/JpegPresets.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/McIdasImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/McIdasImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9cfcca4847a2deb697652c120f204fbbc6126c7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/McIdasImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/MicImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/MicImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34acffc956bfb4616084ecf18b9e1bc424a94d55 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/MicImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/MpegImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/MpegImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4dd2e20ece29b068799df78f96257b867ff248a5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/MpegImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/MpoImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/MpoImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e88272773dae7ced24d7b75100b5f061cab2304c Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/MpoImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/MspImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/MspImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1129cfcd72e260babdfebb5369bba54d202af74 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/MspImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PSDraw.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PSDraw.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c923ae4a74cafdd77bb64436b2dbb4edc60f99db Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PSDraw.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PaletteFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PaletteFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb5c5aa5f52a288806a3b01d1510f4e8dd391bc4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PaletteFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PalmImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PalmImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f95f63fdc769d20fc778f970f75f88fc6ac560d3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PalmImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PcdImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcdImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64641027bc86907567155e725e751e48efa6d7ca Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcdImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PcfFontFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcfFontFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa8f6140701cc024efa99b701651a9bc0c4fd9b2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcfFontFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PcxImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcxImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..faafc2b54a891df376bfca1f01c10c7a71aae16b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PcxImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a30c87f7b7dcb2ab226e67c79d90bcc6b1f9c6e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfParser.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfParser.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e048bc65f0976f01a961b13426ef7ec73c5a0ed8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PdfParser.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PixarImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PixarImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4a7e21a046fe203d0ddbe1c4a521216d3b98ae1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PixarImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PngImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PngImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..213dc144538612accc5621104b2e8e1611b9ae53 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PngImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PpmImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PpmImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a06d79ef1e264498f24a083c6968aacd28224ac Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PpmImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/PsdImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/PsdImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ea275be1767f6c5eb09a568f477df63eac1bace Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/PsdImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/QoiImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/QoiImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..339deee0b2f970f5c05755c9938300d163adaa8e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/QoiImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/SgiImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/SgiImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf2fc6f0980cead6dbf137f5679aafa312271950 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/SgiImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/SpiderImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/SpiderImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9aa0c834059d7973bf712b5bbf74414020f60dd Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/SpiderImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/SunImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/SunImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ef91433c522caef971efe9923e3b7459f8042a4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/SunImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/TarIO.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/TarIO.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..153d9fb84b2a5b1ea9d3838eb0062c4d74c7891d Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/TarIO.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/TgaImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/TgaImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6982cd75f7dd2c3caec1ee3187a0c0d462243171 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/TgaImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0fbf3cbbb2879eb6b046dd879323da2edf031cec Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffTags.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffTags.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9f568d3040892b072f794da482a59dd6341327b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/TiffTags.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/WalImageFile.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/WalImageFile.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43dd375a5c6e9fa9152fe542d295413ec7b0a069 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/WalImageFile.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/WebPImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/WebPImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..105ce5b070fed70a4738dcf38e2f49b722aa8246 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/WebPImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/WmfImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/WmfImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..669896b9e1dd6ec033a2a916af223eaf62c54f9b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/WmfImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/XVThumbImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/XVThumbImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..276ac607e454ea89097c36c94a73ca487d0c2b8e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/XVThumbImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/XbmImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/XbmImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05f78ab65fa231265bd9dad98bb542efc7b9140a Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/XbmImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/XpmImagePlugin.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/XpmImagePlugin.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17c59f55387d7b9516f5092a888b03f5f156a7d8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/XpmImagePlugin.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..158a0e8eddafacc44293412c5bb31f2930a839b5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/__main__.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/__main__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f074f987f9b9e3687144d3474b09577a262fd9ae Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/__main__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_binary.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_binary.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce9baa50eade3648d30b22b59634a9ae61370d3f Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_binary.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_deprecate.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_deprecate.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25be2d214f63fb09832f1f2d9215625a30c0ad0d Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_deprecate.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_tkinter_finder.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_tkinter_finder.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5bca340d989d67c6887e6a05b3a796083e6a9a7e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_tkinter_finder.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_typing.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_typing.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad0639e920aab7d0190944edbbea74dc8dad0d6e Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_typing.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_util.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cc9a4fef82d4b12b2457cbcbcc2d98a94711621 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_util.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/_version.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/_version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ae1711c5d14eaf93b6e1476287b1a990d214fdf Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/_version.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/features.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/features.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65bae7c344aac86e19d97ee4c2839b58bcdfb165 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/features.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/__pycache__/report.cpython-310.pyc b/venv/lib/python3.10/site-packages/PIL/__pycache__/report.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec643c8e43d79ae942128783819e179c50a62c6b Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/__pycache__/report.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/PIL/_avif.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/PIL/_avif.cpython-310-x86_64-linux-gnu.so new file mode 100644 index 0000000000000000000000000000000000000000..fa636a7129fdfb4a491e5ab881d4469203f204e9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/_avif.cpython-310-x86_64-linux-gnu.so differ diff --git a/venv/lib/python3.10/site-packages/PIL/_avif.pyi b/venv/lib/python3.10/site-packages/PIL/_avif.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e27843e5338213713e26973127c738c14313ff98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_avif.pyi @@ -0,0 +1,3 @@ +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_binary.py b/venv/lib/python3.10/site-packages/PIL/_binary.py new file mode 100644 index 0000000000000000000000000000000000000000..4594ccce361168cf77e630cb88ffb09bb4362831 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_binary.py @@ -0,0 +1,112 @@ +# +# The Python Imaging Library. +# $Id$ +# +# Binary input/output support routines. +# +# Copyright (c) 1997-2003 by Secret Labs AB +# Copyright (c) 1995-2003 by Fredrik Lundh +# Copyright (c) 2012 by Brian Crowell +# +# See the README file for information on usage and redistribution. +# + + +"""Binary input/output support routines.""" +from __future__ import annotations + +from struct import pack, unpack_from + + +def i8(c: bytes) -> int: + return c[0] + + +def o8(i: int) -> bytes: + return bytes((i & 255,)) + + +# Input, le = little endian, be = big endian +def i16le(c: bytes, o: int = 0) -> int: + """ + Converts a 2-bytes (16 bits) string to an unsigned integer. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(" int: + """ + Converts a 2-bytes (16 bits) string to a signed integer. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(" int: + """ + Converts a 2-bytes (16 bits) string to a signed integer, big endian. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(">h", c, o)[0] + + +def i32le(c: bytes, o: int = 0) -> int: + """ + Converts a 4-bytes (32 bits) string to an unsigned integer. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(" int: + """ + Converts a 4-bytes (32 bits) string to a signed integer. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(" int: + """ + Converts a 4-bytes (32 bits) string to a signed integer, big endian. + + :param c: string containing bytes to convert + :param o: offset of bytes to convert in string + """ + return unpack_from(">i", c, o)[0] + + +def i16be(c: bytes, o: int = 0) -> int: + return unpack_from(">H", c, o)[0] + + +def i32be(c: bytes, o: int = 0) -> int: + return unpack_from(">I", c, o)[0] + + +# Output, le = little endian, be = big endian +def o16le(i: int) -> bytes: + return pack(" bytes: + return pack(" bytes: + return pack(">H", i) + + +def o32be(i: int) -> bytes: + return pack(">I", i) diff --git a/venv/lib/python3.10/site-packages/PIL/_deprecate.py b/venv/lib/python3.10/site-packages/PIL/_deprecate.py new file mode 100644 index 0000000000000000000000000000000000000000..616a9aace9f4d05e32093c832cc437cb2fa67962 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_deprecate.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import warnings + +from . import __version__ + + +def deprecate( + deprecated: str, + when: int | None, + replacement: str | None = None, + *, + action: str | None = None, + plural: bool = False, + stacklevel: int = 3, +) -> None: + """ + Deprecations helper. + + :param deprecated: Name of thing to be deprecated. + :param when: Pillow major version to be removed in. + :param replacement: Name of replacement. + :param action: Instead of "replacement", give a custom call to action + e.g. "Upgrade to new thing". + :param plural: if the deprecated thing is plural, needing "are" instead of "is". + + Usually of the form: + + "[deprecated] is deprecated and will be removed in Pillow [when] (yyyy-mm-dd). + Use [replacement] instead." + + You can leave out the replacement sentence: + + "[deprecated] is deprecated and will be removed in Pillow [when] (yyyy-mm-dd)" + + Or with another call to action: + + "[deprecated] is deprecated and will be removed in Pillow [when] (yyyy-mm-dd). + [action]." + """ + + is_ = "are" if plural else "is" + + if when is None: + removed = "a future version" + elif when <= int(__version__.split(".")[0]): + msg = f"{deprecated} {is_} deprecated and should be removed." + raise RuntimeError(msg) + elif when == 13: + removed = "Pillow 13 (2026-10-15)" + else: + msg = f"Unknown removal version: {when}. Update {__name__}?" + raise ValueError(msg) + + if replacement and action: + msg = "Use only one of 'replacement' and 'action'" + raise ValueError(msg) + + if replacement: + action = f". Use {replacement} instead." + elif action: + action = f". {action.rstrip('.')}." + else: + action = "" + + warnings.warn( + f"{deprecated} {is_} deprecated and will be removed in {removed}{action}", + DeprecationWarning, + stacklevel=stacklevel, + ) diff --git a/venv/lib/python3.10/site-packages/PIL/_imaging.pyi b/venv/lib/python3.10/site-packages/PIL/_imaging.pyi new file mode 100644 index 0000000000000000000000000000000000000000..998bc52eb8a73b5ee5868cd2c8e5c87c4e6d3037 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imaging.pyi @@ -0,0 +1,31 @@ +from typing import Any + +class ImagingCore: + def __getitem__(self, index: int) -> float: ... + def __getattr__(self, name: str) -> Any: ... + +class ImagingFont: + def __getattr__(self, name: str) -> Any: ... + +class ImagingDraw: + def __getattr__(self, name: str) -> Any: ... + +class PixelAccess: + def __getitem__(self, xy: tuple[int, int]) -> float | tuple[int, ...]: ... + def __setitem__( + self, xy: tuple[int, int], color: float | tuple[int, ...] + ) -> None: ... + +class ImagingDecoder: + def __getattr__(self, name: str) -> Any: ... + +class ImagingEncoder: + def __getattr__(self, name: str) -> Any: ... + +class _Outline: + def close(self) -> None: ... + def __getattr__(self, name: str) -> Any: ... + +def font(image: ImagingCore, glyphdata: bytes) -> ImagingFont: ... +def outline() -> _Outline: ... +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingcms.pyi b/venv/lib/python3.10/site-packages/PIL/_imagingcms.pyi new file mode 100644 index 0000000000000000000000000000000000000000..4fc0d60ab79375281abb6d53d9d707cf18f09815 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imagingcms.pyi @@ -0,0 +1,143 @@ +import datetime +import sys +from typing import Literal, SupportsFloat, TypeAlias, TypedDict + +from ._typing import CapsuleType + +littlecms_version: str | None + +_Tuple3f: TypeAlias = tuple[float, float, float] +_Tuple2x3f: TypeAlias = tuple[_Tuple3f, _Tuple3f] +_Tuple3x3f: TypeAlias = tuple[_Tuple3f, _Tuple3f, _Tuple3f] + +class _IccMeasurementCondition(TypedDict): + observer: int + backing: _Tuple3f + geo: str + flare: float + illuminant_type: str + +class _IccViewingCondition(TypedDict): + illuminant: _Tuple3f + surround: _Tuple3f + illuminant_type: str + +class CmsProfile: + @property + def rendering_intent(self) -> int: ... + @property + def creation_date(self) -> datetime.datetime | None: ... + @property + def copyright(self) -> str | None: ... + @property + def target(self) -> str | None: ... + @property + def manufacturer(self) -> str | None: ... + @property + def model(self) -> str | None: ... + @property + def profile_description(self) -> str | None: ... + @property + def screening_description(self) -> str | None: ... + @property + def viewing_condition(self) -> str | None: ... + @property + def version(self) -> float: ... + @property + def icc_version(self) -> int: ... + @property + def attributes(self) -> int: ... + @property + def header_flags(self) -> int: ... + @property + def header_manufacturer(self) -> str: ... + @property + def header_model(self) -> str: ... + @property + def device_class(self) -> str: ... + @property + def connection_space(self) -> str: ... + @property + def xcolor_space(self) -> str: ... + @property + def profile_id(self) -> bytes: ... + @property + def is_matrix_shaper(self) -> bool: ... + @property + def technology(self) -> str | None: ... + @property + def colorimetric_intent(self) -> str | None: ... + @property + def perceptual_rendering_intent_gamut(self) -> str | None: ... + @property + def saturation_rendering_intent_gamut(self) -> str | None: ... + @property + def red_colorant(self) -> _Tuple2x3f | None: ... + @property + def green_colorant(self) -> _Tuple2x3f | None: ... + @property + def blue_colorant(self) -> _Tuple2x3f | None: ... + @property + def red_primary(self) -> _Tuple2x3f | None: ... + @property + def green_primary(self) -> _Tuple2x3f | None: ... + @property + def blue_primary(self) -> _Tuple2x3f | None: ... + @property + def media_white_point_temperature(self) -> float | None: ... + @property + def media_white_point(self) -> _Tuple2x3f | None: ... + @property + def media_black_point(self) -> _Tuple2x3f | None: ... + @property + def luminance(self) -> _Tuple2x3f | None: ... + @property + def chromatic_adaptation(self) -> tuple[_Tuple3x3f, _Tuple3x3f] | None: ... + @property + def chromaticity(self) -> _Tuple3x3f | None: ... + @property + def colorant_table(self) -> list[str] | None: ... + @property + def colorant_table_out(self) -> list[str] | None: ... + @property + def intent_supported(self) -> dict[int, tuple[bool, bool, bool]] | None: ... + @property + def clut(self) -> dict[int, tuple[bool, bool, bool]] | None: ... + @property + def icc_measurement_condition(self) -> _IccMeasurementCondition | None: ... + @property + def icc_viewing_condition(self) -> _IccViewingCondition | None: ... + def is_intent_supported(self, intent: int, direction: int, /) -> int: ... + +class CmsTransform: + def apply(self, id_in: CapsuleType, id_out: CapsuleType) -> int: ... + +def profile_open(profile: str, /) -> CmsProfile: ... +def profile_frombytes(profile: bytes, /) -> CmsProfile: ... +def profile_tobytes(profile: CmsProfile, /) -> bytes: ... +def buildTransform( + input_profile: CmsProfile, + output_profile: CmsProfile, + in_mode: str, + out_mode: str, + rendering_intent: int = 0, + cms_flags: int = 0, + /, +) -> CmsTransform: ... +def buildProofTransform( + input_profile: CmsProfile, + output_profile: CmsProfile, + proof_profile: CmsProfile, + in_mode: str, + out_mode: str, + rendering_intent: int = 0, + proof_intent: int = 0, + cms_flags: int = 0, + /, +) -> CmsTransform: ... +def createProfile( + color_space: Literal["LAB", "XYZ", "sRGB"], color_temp: SupportsFloat = 0.0, / +) -> CmsProfile: ... + +if sys.platform == "win32": + def get_display_profile_win32(handle: int = 0, is_dc: int = 0, /) -> str | None: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingft.pyi b/venv/lib/python3.10/site-packages/PIL/_imagingft.pyi new file mode 100644 index 0000000000000000000000000000000000000000..2136810ba6a95e0c025a71cf7fbbc5b60490fed7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imagingft.pyi @@ -0,0 +1,70 @@ +from collections.abc import Callable +from typing import Any + +from . import ImageFont, _imaging + +class Font: + @property + def family(self) -> str | None: ... + @property + def style(self) -> str | None: ... + @property + def ascent(self) -> int: ... + @property + def descent(self) -> int: ... + @property + def height(self) -> int: ... + @property + def x_ppem(self) -> int: ... + @property + def y_ppem(self) -> int: ... + @property + def glyphs(self) -> int: ... + def render( + self, + string: str | bytes, + fill: Callable[[int, int], _imaging.ImagingCore], + mode: str, + dir: str | None, + features: list[str] | None, + lang: str | None, + stroke_width: float, + stroke_filled: bool, + anchor: str | None, + foreground_ink_long: int, + start: tuple[float, float], + /, + ) -> tuple[_imaging.ImagingCore, tuple[int, int]]: ... + def getsize( + self, + string: str | bytes | bytearray, + mode: str, + dir: str | None, + features: list[str] | None, + lang: str | None, + anchor: str | None, + /, + ) -> tuple[tuple[int, int], tuple[int, int]]: ... + def getlength( + self, + string: str | bytes, + mode: str, + dir: str | None, + features: list[str] | None, + lang: str | None, + /, + ) -> float: ... + def getvarnames(self) -> list[bytes]: ... + def getvaraxes(self) -> list[ImageFont.Axis]: ... + def setvarname(self, instance_index: int, /) -> None: ... + def setvaraxes(self, axes: list[float], /) -> None: ... + +def getfont( + filename: str | bytes, + size: float, + index: int, + encoding: str, + font_bytes: bytes, + layout_engine: int, +) -> Font: ... +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingmath.pyi b/venv/lib/python3.10/site-packages/PIL/_imagingmath.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e27843e5338213713e26973127c738c14313ff98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imagingmath.pyi @@ -0,0 +1,3 @@ +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingmorph.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/PIL/_imagingmorph.cpython-310-x86_64-linux-gnu.so new file mode 100644 index 0000000000000000000000000000000000000000..cde33213c524eabb80d9f620033c519b21c99cc9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/_imagingmorph.cpython-310-x86_64-linux-gnu.so differ diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingmorph.pyi b/venv/lib/python3.10/site-packages/PIL/_imagingmorph.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e27843e5338213713e26973127c738c14313ff98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imagingmorph.pyi @@ -0,0 +1,3 @@ +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingtk.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/PIL/_imagingtk.cpython-310-x86_64-linux-gnu.so new file mode 100644 index 0000000000000000000000000000000000000000..1bb0a2850302be3549e226199e1c0368fd501646 Binary files /dev/null and b/venv/lib/python3.10/site-packages/PIL/_imagingtk.cpython-310-x86_64-linux-gnu.so differ diff --git a/venv/lib/python3.10/site-packages/PIL/_imagingtk.pyi b/venv/lib/python3.10/site-packages/PIL/_imagingtk.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e27843e5338213713e26973127c738c14313ff98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_imagingtk.pyi @@ -0,0 +1,3 @@ +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/_tkinter_finder.py b/venv/lib/python3.10/site-packages/PIL/_tkinter_finder.py new file mode 100644 index 0000000000000000000000000000000000000000..9c0143003a7320dd475cfcd168168b82e4f64964 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_tkinter_finder.py @@ -0,0 +1,20 @@ +"""Find compiled module linking to Tcl / Tk libraries""" + +from __future__ import annotations + +import sys +import tkinter + +tk = getattr(tkinter, "_tkinter") + +try: + if hasattr(sys, "pypy_find_executable"): + TKINTER_LIB = tk.tklib_cffi.__file__ + else: + TKINTER_LIB = tk.__file__ +except AttributeError: + # _tkinter may be compiled directly into Python, in which case __file__ is + # not available. load_tkinter_funcs will check the binary first in any case. + TKINTER_LIB = None + +tk_version = str(tkinter.TkVersion) diff --git a/venv/lib/python3.10/site-packages/PIL/_typing.py b/venv/lib/python3.10/site-packages/PIL/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..a941f89806faabe61486fc164d8bd3f60fb48626 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_typing.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import os +import sys +from collections.abc import Sequence +from typing import Any, Protocol, TypeVar + +TYPE_CHECKING = False +if TYPE_CHECKING: + from numbers import _IntegralLike as IntegralLike + + try: + import numpy.typing as npt + + NumpyArray = npt.NDArray[Any] + except ImportError: + pass + +if sys.version_info >= (3, 13): + from types import CapsuleType +else: + CapsuleType = object + +if sys.version_info >= (3, 12): + from collections.abc import Buffer +else: + Buffer = Any + + +_Ink = float | tuple[int, ...] | str + +Coords = Sequence[float] | Sequence[Sequence[float]] + + +_T_co = TypeVar("_T_co", covariant=True) + + +class SupportsRead(Protocol[_T_co]): + def read(self, length: int = ..., /) -> _T_co: ... + + +StrOrBytesPath = str | bytes | os.PathLike[str] | os.PathLike[bytes] + + +__all__ = ["Buffer", "IntegralLike", "StrOrBytesPath", "SupportsRead"] diff --git a/venv/lib/python3.10/site-packages/PIL/_util.py b/venv/lib/python3.10/site-packages/PIL/_util.py new file mode 100644 index 0000000000000000000000000000000000000000..b1fa6a0f39ed2847723f20fc81462adc9245d970 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_util.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import os + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Any, NoReturn, TypeGuard + + from ._typing import StrOrBytesPath + + +def is_path(f: Any) -> TypeGuard[StrOrBytesPath]: + return isinstance(f, (bytes, str, os.PathLike)) + + +class DeferredError: + def __init__(self, ex: BaseException): + self.ex = ex + + def __getattr__(self, elt: str) -> NoReturn: + raise self.ex + + @staticmethod + def new(ex: BaseException) -> Any: + """ + Creates an object that raises the wrapped exception ``ex`` when used, + and casts it to :py:obj:`~typing.Any` type. + """ + return DeferredError(ex) diff --git a/venv/lib/python3.10/site-packages/PIL/_version.py b/venv/lib/python3.10/site-packages/PIL/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..79ce194c33476aa657567e25d65d1179557373a7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_version.py @@ -0,0 +1,4 @@ +# Master version for Pillow +from __future__ import annotations + +__version__ = "12.0.0" diff --git a/venv/lib/python3.10/site-packages/PIL/_webp.pyi b/venv/lib/python3.10/site-packages/PIL/_webp.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e27843e5338213713e26973127c738c14313ff98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/_webp.pyi @@ -0,0 +1,3 @@ +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/PIL/features.py b/venv/lib/python3.10/site-packages/PIL/features.py new file mode 100644 index 0000000000000000000000000000000000000000..ff32c2510453aebb44a0b60d52db29d4718acee6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/features.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import collections +import os +import sys +import warnings +from typing import IO + +import PIL + +from . import Image + +modules = { + "pil": ("PIL._imaging", "PILLOW_VERSION"), + "tkinter": ("PIL._tkinter_finder", "tk_version"), + "freetype2": ("PIL._imagingft", "freetype2_version"), + "littlecms2": ("PIL._imagingcms", "littlecms_version"), + "webp": ("PIL._webp", "webpdecoder_version"), + "avif": ("PIL._avif", "libavif_version"), +} + + +def check_module(feature: str) -> bool: + """ + Checks if a module is available. + + :param feature: The module to check for. + :returns: ``True`` if available, ``False`` otherwise. + :raises ValueError: If the module is not defined in this version of Pillow. + """ + if feature not in modules: + msg = f"Unknown module {feature}" + raise ValueError(msg) + + module, ver = modules[feature] + + try: + __import__(module) + return True + except ModuleNotFoundError: + return False + except ImportError as ex: + warnings.warn(str(ex)) + return False + + +def version_module(feature: str) -> str | None: + """ + :param feature: The module to check for. + :returns: + The loaded version number as a string, or ``None`` if unknown or not available. + :raises ValueError: If the module is not defined in this version of Pillow. + """ + if not check_module(feature): + return None + + module, ver = modules[feature] + + return getattr(__import__(module, fromlist=[ver]), ver) + + +def get_supported_modules() -> list[str]: + """ + :returns: A list of all supported modules. + """ + return [f for f in modules if check_module(f)] + + +codecs = { + "jpg": ("jpeg", "jpeglib"), + "jpg_2000": ("jpeg2k", "jp2klib"), + "zlib": ("zip", "zlib"), + "libtiff": ("libtiff", "libtiff"), +} + + +def check_codec(feature: str) -> bool: + """ + Checks if a codec is available. + + :param feature: The codec to check for. + :returns: ``True`` if available, ``False`` otherwise. + :raises ValueError: If the codec is not defined in this version of Pillow. + """ + if feature not in codecs: + msg = f"Unknown codec {feature}" + raise ValueError(msg) + + codec, lib = codecs[feature] + + return f"{codec}_encoder" in dir(Image.core) + + +def version_codec(feature: str) -> str | None: + """ + :param feature: The codec to check for. + :returns: + The version number as a string, or ``None`` if not available. + Checked at compile time for ``jpg``, run-time otherwise. + :raises ValueError: If the codec is not defined in this version of Pillow. + """ + if not check_codec(feature): + return None + + codec, lib = codecs[feature] + + version = getattr(Image.core, f"{lib}_version") + + if feature == "libtiff": + return version.split("\n")[0].split("Version ")[1] + + return version + + +def get_supported_codecs() -> list[str]: + """ + :returns: A list of all supported codecs. + """ + return [f for f in codecs if check_codec(f)] + + +features: dict[str, tuple[str, str, str | None]] = { + "raqm": ("PIL._imagingft", "HAVE_RAQM", "raqm_version"), + "fribidi": ("PIL._imagingft", "HAVE_FRIBIDI", "fribidi_version"), + "harfbuzz": ("PIL._imagingft", "HAVE_HARFBUZZ", "harfbuzz_version"), + "libjpeg_turbo": ("PIL._imaging", "HAVE_LIBJPEGTURBO", "libjpeg_turbo_version"), + "mozjpeg": ("PIL._imaging", "HAVE_MOZJPEG", "libjpeg_turbo_version"), + "zlib_ng": ("PIL._imaging", "HAVE_ZLIBNG", "zlib_ng_version"), + "libimagequant": ("PIL._imaging", "HAVE_LIBIMAGEQUANT", "imagequant_version"), + "xcb": ("PIL._imaging", "HAVE_XCB", None), +} + + +def check_feature(feature: str) -> bool | None: + """ + Checks if a feature is available. + + :param feature: The feature to check for. + :returns: ``True`` if available, ``False`` if unavailable, ``None`` if unknown. + :raises ValueError: If the feature is not defined in this version of Pillow. + """ + if feature not in features: + msg = f"Unknown feature {feature}" + raise ValueError(msg) + + module, flag, ver = features[feature] + + try: + imported_module = __import__(module, fromlist=["PIL"]) + return getattr(imported_module, flag) + except ModuleNotFoundError: + return None + except ImportError as ex: + warnings.warn(str(ex)) + return None + + +def version_feature(feature: str) -> str | None: + """ + :param feature: The feature to check for. + :returns: The version number as a string, or ``None`` if not available. + :raises ValueError: If the feature is not defined in this version of Pillow. + """ + if not check_feature(feature): + return None + + module, flag, ver = features[feature] + + if ver is None: + return None + + return getattr(__import__(module, fromlist=[ver]), ver) + + +def get_supported_features() -> list[str]: + """ + :returns: A list of all supported features. + """ + return [f for f in features if check_feature(f)] + + +def check(feature: str) -> bool | None: + """ + :param feature: A module, codec, or feature name. + :returns: + ``True`` if the module, codec, or feature is available, + ``False`` or ``None`` otherwise. + """ + + if feature in modules: + return check_module(feature) + if feature in codecs: + return check_codec(feature) + if feature in features: + return check_feature(feature) + warnings.warn(f"Unknown feature '{feature}'.", stacklevel=2) + return False + + +def version(feature: str) -> str | None: + """ + :param feature: + The module, codec, or feature to check for. + :returns: + The version number as a string, or ``None`` if unknown or not available. + """ + if feature in modules: + return version_module(feature) + if feature in codecs: + return version_codec(feature) + if feature in features: + return version_feature(feature) + return None + + +def get_supported() -> list[str]: + """ + :returns: A list of all supported modules, features, and codecs. + """ + + ret = get_supported_modules() + ret.extend(get_supported_features()) + ret.extend(get_supported_codecs()) + return ret + + +def pilinfo(out: IO[str] | None = None, supported_formats: bool = True) -> None: + """ + Prints information about this installation of Pillow. + This function can be called with ``python3 -m PIL``. + It can also be called with ``python3 -m PIL.report`` or ``python3 -m PIL --report`` + to have "supported_formats" set to ``False``, omitting the list of all supported + image file formats. + + :param out: + The output stream to print to. Defaults to ``sys.stdout`` if ``None``. + :param supported_formats: + If ``True``, a list of all supported image file formats will be printed. + """ + + if out is None: + out = sys.stdout + + Image.init() + + print("-" * 68, file=out) + print(f"Pillow {PIL.__version__}", file=out) + py_version_lines = sys.version.splitlines() + print(f"Python {py_version_lines[0].strip()}", file=out) + for py_version in py_version_lines[1:]: + print(f" {py_version.strip()}", file=out) + print("-" * 68, file=out) + print(f"Python executable is {sys.executable or 'unknown'}", file=out) + if sys.prefix != sys.base_prefix: + print(f"Environment Python files loaded from {sys.prefix}", file=out) + print(f"System Python files loaded from {sys.base_prefix}", file=out) + print("-" * 68, file=out) + print( + f"Python Pillow modules loaded from {os.path.dirname(Image.__file__)}", + file=out, + ) + print( + f"Binary Pillow modules loaded from {os.path.dirname(Image.core.__file__)}", + file=out, + ) + print("-" * 68, file=out) + + for name, feature in [ + ("pil", "PIL CORE"), + ("tkinter", "TKINTER"), + ("freetype2", "FREETYPE2"), + ("littlecms2", "LITTLECMS2"), + ("webp", "WEBP"), + ("avif", "AVIF"), + ("jpg", "JPEG"), + ("jpg_2000", "OPENJPEG (JPEG2000)"), + ("zlib", "ZLIB (PNG/ZIP)"), + ("libtiff", "LIBTIFF"), + ("raqm", "RAQM (Bidirectional Text)"), + ("libimagequant", "LIBIMAGEQUANT (Quantization method)"), + ("xcb", "XCB (X protocol)"), + ]: + if check(name): + v: str | None = None + if name == "jpg": + libjpeg_turbo_version = version_feature("libjpeg_turbo") + if libjpeg_turbo_version is not None: + v = "mozjpeg" if check_feature("mozjpeg") else "libjpeg-turbo" + v += " " + libjpeg_turbo_version + if v is None: + v = version(name) + if v is not None: + version_static = name in ("pil", "jpg") + if name == "littlecms2": + # this check is also in src/_imagingcms.c:setup_module() + version_static = tuple(int(x) for x in v.split(".")) < (2, 7) + t = "compiled for" if version_static else "loaded" + if name == "zlib": + zlib_ng_version = version_feature("zlib_ng") + if zlib_ng_version is not None: + v += ", compiled for zlib-ng " + zlib_ng_version + elif name == "raqm": + for f in ("fribidi", "harfbuzz"): + v2 = version_feature(f) + if v2 is not None: + v += f", {f} {v2}" + print("---", feature, "support ok,", t, v, file=out) + else: + print("---", feature, "support ok", file=out) + else: + print("***", feature, "support not installed", file=out) + print("-" * 68, file=out) + + if supported_formats: + extensions = collections.defaultdict(list) + for ext, i in Image.EXTENSION.items(): + extensions[i].append(ext) + + for i in sorted(Image.ID): + line = f"{i}" + if i in Image.MIME: + line = f"{line} {Image.MIME[i]}" + print(line, file=out) + + if i in extensions: + print( + "Extensions: {}".format(", ".join(sorted(extensions[i]))), file=out + ) + + features = [] + if i in Image.OPEN: + features.append("open") + if i in Image.SAVE: + features.append("save") + if i in Image.SAVE_ALL: + features.append("save_all") + if i in Image.DECODERS: + features.append("decode") + if i in Image.ENCODERS: + features.append("encode") + + print("Features: {}".format(", ".join(features)), file=out) + print("-" * 68, file=out) diff --git a/venv/lib/python3.10/site-packages/PIL/py.typed b/venv/lib/python3.10/site-packages/PIL/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/PIL/report.py b/venv/lib/python3.10/site-packages/PIL/report.py new file mode 100644 index 0000000000000000000000000000000000000000..d2815e8455e2ead803de4417314987ce7e9b7598 --- /dev/null +++ b/venv/lib/python3.10/site-packages/PIL/report.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .features import pilinfo + +pilinfo(supported_formats=False) diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/LICENSE b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..033c86b7a40a331f281bd406e991ef1db597c208 --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/LICENSE @@ -0,0 +1,13 @@ +Copyright 2016-2020 aio-libs collaboration. + +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. diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..1eecd7d586cde8ccf4b4a460422d2566b665937b --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/METADATA @@ -0,0 +1,163 @@ +Metadata-Version: 2.1 +Name: async-timeout +Version: 5.0.1 +Summary: Timeout context manager for asyncio programs +Home-page: https://github.com/aio-libs/async-timeout +Author: Andrew Svetlov +Author-email: andrew.svetlov@gmail.com +License: Apache 2 +Project-URL: Chat: Gitter, https://gitter.im/aio-libs/Lobby +Project-URL: CI: GitHub Actions, https://github.com/aio-libs/async-timeout/actions +Project-URL: Coverage: codecov, https://codecov.io/github/aio-libs/async-timeout +Project-URL: GitHub: issues, https://github.com/aio-libs/async-timeout/issues +Project-URL: GitHub: repo, https://github.com/aio-libs/async-timeout +Classifier: Development Status :: 5 - Production/Stable +Classifier: Topic :: Software Development :: Libraries +Classifier: Framework :: AsyncIO +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Requires-Python: >=3.8 +Description-Content-Type: text/x-rst +License-File: LICENSE + +async-timeout +============= +.. image:: https://travis-ci.com/aio-libs/async-timeout.svg?branch=master + :target: https://travis-ci.com/aio-libs/async-timeout +.. image:: https://codecov.io/gh/aio-libs/async-timeout/branch/master/graph/badge.svg + :target: https://codecov.io/gh/aio-libs/async-timeout +.. image:: https://img.shields.io/pypi/v/async-timeout.svg + :target: https://pypi.python.org/pypi/async-timeout +.. image:: https://badges.gitter.im/Join%20Chat.svg + :target: https://gitter.im/aio-libs/Lobby + :alt: Chat on Gitter + +asyncio-compatible timeout context manager. + + + +DEPRECATED +---------- + +This library has effectively been upstreamed into Python 3.11+. + +Therefore this library is considered deprecated and no longer actively supported. + +Version 5.0+ provides dual-mode when executed on Python 3.11+: +``asyncio_timeout.Timeout`` is fully compatible with ``asyncio.Timeout`` *and* old +versions of the library. + +Anyway, using upstream is highly recommended. ``asyncio_timeout`` exists only for the +sake of backward compatibility, easy supporting both old and new Python by the same +code, and easy misgration. + +If rescheduling API is not important and only ``async with timeout(...): ...`` functionality is required, +a user could apply conditional import:: + + if sys.version_info >= (3, 11): + from asyncio import timeout, timeout_at + else: + from async_timeout import timeout, timeout_at + + +Usage example +------------- + + +The context manager is useful in cases when you want to apply timeout +logic around block of code or in cases when ``asyncio.wait_for()`` is +not suitable. Also it's much faster than ``asyncio.wait_for()`` +because ``timeout`` doesn't create a new task. + +The ``timeout(delay, *, loop=None)`` call returns a context manager +that cancels a block on *timeout* expiring:: + + from async_timeout import timeout + async with timeout(1.5): + await inner() + +1. If ``inner()`` is executed faster than in ``1.5`` seconds nothing + happens. +2. Otherwise ``inner()`` is cancelled internally by sending + ``asyncio.CancelledError`` into but ``asyncio.TimeoutError`` is + raised outside of context manager scope. + +*timeout* parameter could be ``None`` for skipping timeout functionality. + + +Alternatively, ``timeout_at(when)`` can be used for scheduling +at the absolute time:: + + loop = asyncio.get_event_loop() + now = loop.time() + + async with timeout_at(now + 1.5): + await inner() + + +Please note: it is not POSIX time but a time with +undefined starting base, e.g. the time of the system power on. + + +Context manager has ``.expired()`` / ``.expired`` for check if timeout happens +exactly in context manager:: + + async with timeout(1.5) as cm: + await inner() + print(cm.expired()) # recommended api + print(cm.expired) # compatible api + +The property is ``True`` if ``inner()`` execution is cancelled by +timeout context manager. + +If ``inner()`` call explicitly raises ``TimeoutError`` ``cm.expired`` +is ``False``. + +The scheduled deadline time is available as ``.when()`` / ``.deadline``:: + + async with timeout(1.5) as cm: + cm.when() # recommended api + cm.deadline # compatible api + +Not finished yet timeout can be rescheduled by ``shift()`` +or ``update()`` methods:: + + async with timeout(1.5) as cm: + # recommended api + cm.reschedule(cm.when() + 1) # add another second on waiting + # compatible api + cm.shift(1) # add another second on waiting + cm.update(loop.time() + 5) # reschedule to now+5 seconds + +Rescheduling is forbidden if the timeout is expired or after exit from ``async with`` +code block. + + +Disable scheduled timeout:: + + async with timeout(1.5) as cm: + cm.reschedule(None) # recommended api + cm.reject() # compatible api + + + +Installation +------------ + +:: + + $ pip install async-timeout + +The library is Python 3 only! + + + +Authors and License +------------------- + +The module is written by Andrew Svetlov. + +It's *Apache 2* licensed and freely available. diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..cd953bfdd7e53a5721fe35923786008b9412901c --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/RECORD @@ -0,0 +1,11 @@ +async_timeout-5.0.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +async_timeout-5.0.1.dist-info/LICENSE,sha256=4Y17uPUT4sRrtYXJS1hb0wcg3TzLId2weG9y0WZY-Sw,568 +async_timeout-5.0.1.dist-info/METADATA,sha256=RVDNEIPYIBJKPsjThJDaKRX1h79-4QYQNuBLSXPItU8,5131 +async_timeout-5.0.1.dist-info/RECORD,, +async_timeout-5.0.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +async_timeout-5.0.1.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91 +async_timeout-5.0.1.dist-info/top_level.txt,sha256=9oM4e7Twq8iD_7_Q3Mz0E6GPIB6vJvRFo-UBwUQtBDU,14 +async_timeout-5.0.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 +async_timeout/__init__.py,sha256=QF0zpfX1vGmxib7kAqNPm9YehPV0oBVozxJ--Mxq9dI,9186 +async_timeout/__pycache__/__init__.cpython-310.pyc,, +async_timeout/py.typed,sha256=tyozzRT1fziXETDxokmuyt6jhOmtjUbnVNJdZcG7ik0,12 diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..9b78c44519d450e4036a1ea44361a2ac7b1ee300 --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (75.3.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..ad29955ef909f5f38e96b6d8a6c9ba54d9bccd53 --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/top_level.txt @@ -0,0 +1 @@ +async_timeout diff --git a/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/zip-safe b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/venv/lib/python3.10/site-packages/async_timeout-5.0.1.dist-info/zip-safe @@ -0,0 +1 @@ + diff --git a/venv/lib/python3.10/site-packages/filelock/__init__.py b/venv/lib/python3.10/site-packages/filelock/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9d8c5b8ebe565a652b3671b3dfa066f7346af45 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/__init__.py @@ -0,0 +1,70 @@ +""" +A platform independent file lock that supports the with-statement. + +.. autodata:: filelock.__version__ + :no-value: + +""" + +from __future__ import annotations + +import sys +import warnings +from typing import TYPE_CHECKING + +from ._api import AcquireReturnProxy, BaseFileLock +from ._error import Timeout +from ._soft import SoftFileLock +from ._unix import UnixFileLock, has_fcntl +from ._windows import WindowsFileLock +from .asyncio import ( + AsyncAcquireReturnProxy, + AsyncSoftFileLock, + AsyncUnixFileLock, + AsyncWindowsFileLock, + BaseAsyncFileLock, +) +from .version import version + +#: version of the project as a string +__version__: str = version + + +if sys.platform == "win32": # pragma: win32 cover + _FileLock: type[BaseFileLock] = WindowsFileLock + _AsyncFileLock: type[BaseAsyncFileLock] = AsyncWindowsFileLock +else: # pragma: win32 no cover # noqa: PLR5501 + if has_fcntl: + _FileLock: type[BaseFileLock] = UnixFileLock + _AsyncFileLock: type[BaseAsyncFileLock] = AsyncUnixFileLock + else: + _FileLock = SoftFileLock + _AsyncFileLock = AsyncSoftFileLock + if warnings is not None: + warnings.warn("only soft file lock is available", stacklevel=2) + +if TYPE_CHECKING: + FileLock = SoftFileLock + AsyncFileLock = AsyncSoftFileLock +else: + #: Alias for the lock, which should be used for the current platform. + FileLock = _FileLock + AsyncFileLock = _AsyncFileLock + + +__all__ = [ + "AcquireReturnProxy", + "AsyncAcquireReturnProxy", + "AsyncFileLock", + "AsyncSoftFileLock", + "AsyncUnixFileLock", + "AsyncWindowsFileLock", + "BaseAsyncFileLock", + "BaseFileLock", + "FileLock", + "SoftFileLock", + "Timeout", + "UnixFileLock", + "WindowsFileLock", + "__version__", +] diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67155036f35fccd72b5cba2b0c4f7bc4fc200144 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_api.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_api.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aefc5c2408966f6819054e7dd892f747721fa507 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_api.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_error.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_error.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45070effb74ae582fbf2a044ee72802c1665963a Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_error.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_soft.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_soft.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..021bc3267ca3ebd3786d8e3520f5a55272f9c0a4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_soft.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_unix.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_unix.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a982985071a0fdfb491d5cfbc02382c7dd3448b5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_unix.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_util.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8229097685c8b017aa4bd329214afe9f35d3e062 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_util.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/_windows.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/_windows.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b5632fb842fb548439897882f6a981c2599366d Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/_windows.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/asyncio.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/asyncio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2c2d8fa720c2dfaae8cad8b3a054846c4e3e453 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/asyncio.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/__pycache__/version.cpython-310.pyc b/venv/lib/python3.10/site-packages/filelock/__pycache__/version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14e646ba6d635395fe84b2a1137d9647ec4cc208 Binary files /dev/null and b/venv/lib/python3.10/site-packages/filelock/__pycache__/version.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/filelock/_api.py b/venv/lib/python3.10/site-packages/filelock/_api.py new file mode 100644 index 0000000000000000000000000000000000000000..8fde69a0fef7badcc123d17735cd784a99baed52 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_api.py @@ -0,0 +1,403 @@ +from __future__ import annotations + +import contextlib +import inspect +import logging +import os +import time +import warnings +from abc import ABCMeta, abstractmethod +from dataclasses import dataclass +from threading import local +from typing import TYPE_CHECKING, Any, cast +from weakref import WeakValueDictionary + +from ._error import Timeout + +if TYPE_CHECKING: + import sys + from types import TracebackType + + if sys.version_info >= (3, 11): # pragma: no cover (py311+) + from typing import Self + else: # pragma: no cover ( None: + self.lock = lock + + def __enter__(self) -> BaseFileLock: + return self.lock + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.lock.release() + + +@dataclass +class FileLockContext: + """A dataclass which holds the context for a ``BaseFileLock`` object.""" + + # The context is held in a separate class to allow optional use of thread local storage via the + # ThreadLocalFileContext class. + + #: The path to the lock file. + lock_file: str + + #: The default timeout value. + timeout: float + + #: The mode for the lock files + mode: int + + #: Whether the lock should be blocking or not + blocking: bool + + #: The file descriptor for the *_lock_file* as it is returned by the os.open() function, not None when lock held + lock_file_fd: int | None = None + + #: The lock counter is used for implementing the nested locking mechanism. + lock_counter: int = 0 # When the lock is acquired is increased and the lock is only released, when this value is 0 + + +class ThreadLocalFileContext(FileLockContext, local): + """A thread local version of the ``FileLockContext`` class.""" + + +class FileLockMeta(ABCMeta): + def __call__( # noqa: PLR0913 + cls, + lock_file: str | os.PathLike[str], + timeout: float = -1, + mode: int = 0o644, + thread_local: bool = True, # noqa: FBT001, FBT002 + *, + blocking: bool = True, + is_singleton: bool = False, + **kwargs: Any, # capture remaining kwargs for subclasses # noqa: ANN401 + ) -> BaseFileLock: + if is_singleton: + instance = cls._instances.get(str(lock_file)) # type: ignore[attr-defined] + if instance: + params_to_check = { + "thread_local": (thread_local, instance.is_thread_local()), + "timeout": (timeout, instance.timeout), + "mode": (mode, instance.mode), + "blocking": (blocking, instance.blocking), + } + + non_matching_params = { + name: (passed_param, set_param) + for name, (passed_param, set_param) in params_to_check.items() + if passed_param != set_param + } + if not non_matching_params: + return cast("BaseFileLock", instance) + + # parameters do not match; raise error + msg = "Singleton lock instances cannot be initialized with differing arguments" + msg += "\nNon-matching arguments: " + for param_name, (passed_param, set_param) in non_matching_params.items(): + msg += f"\n\t{param_name} (existing lock has {set_param} but {passed_param} was passed)" + raise ValueError(msg) + + # Workaround to make `__init__`'s params optional in subclasses + # E.g. virtualenv changes the signature of the `__init__` method in the `BaseFileLock` class descendant + # (https://github.com/tox-dev/filelock/pull/340) + + all_params = { + "timeout": timeout, + "mode": mode, + "thread_local": thread_local, + "blocking": blocking, + "is_singleton": is_singleton, + **kwargs, + } + + present_params = inspect.signature(cls.__init__).parameters # type: ignore[misc] + init_params = {key: value for key, value in all_params.items() if key in present_params} + + instance = super().__call__(lock_file, **init_params) + + if is_singleton: + cls._instances[str(lock_file)] = instance # type: ignore[attr-defined] + + return cast("BaseFileLock", instance) + + +class BaseFileLock(contextlib.ContextDecorator, metaclass=FileLockMeta): + """Abstract base class for a file lock object.""" + + _instances: WeakValueDictionary[str, BaseFileLock] + + def __init_subclass__(cls, **kwargs: dict[str, Any]) -> None: + """Setup unique state for lock subclasses.""" + super().__init_subclass__(**kwargs) + cls._instances = WeakValueDictionary() + + def __init__( # noqa: PLR0913 + self, + lock_file: str | os.PathLike[str], + timeout: float = -1, + mode: int = 0o644, + thread_local: bool = True, # noqa: FBT001, FBT002 + *, + blocking: bool = True, + is_singleton: bool = False, + ) -> None: + """ + Create a new lock object. + + :param lock_file: path to the file + :param timeout: default timeout when acquiring the lock, in seconds. It will be used as fallback value in \ + the acquire method, if no timeout value (``None``) is given. If you want to disable the timeout, set it \ + to a negative value. A timeout of 0 means that there is exactly one attempt to acquire the file lock. + :param mode: file permissions for the lockfile + :param thread_local: Whether this object's internal context should be thread local or not. If this is set to \ + ``False`` then the lock will be reentrant across threads. + :param blocking: whether the lock should be blocking or not + :param is_singleton: If this is set to ``True`` then only one instance of this class will be created \ + per lock file. This is useful if you want to use the lock object for reentrant locking without needing \ + to pass the same object around. + + """ + self._is_thread_local = thread_local + self._is_singleton = is_singleton + + # Create the context. Note that external code should not work with the context directly and should instead use + # properties of this class. + kwargs: dict[str, Any] = { + "lock_file": os.fspath(lock_file), + "timeout": timeout, + "mode": mode, + "blocking": blocking, + } + self._context: FileLockContext = (ThreadLocalFileContext if thread_local else FileLockContext)(**kwargs) + + def is_thread_local(self) -> bool: + """:return: a flag indicating if this lock is thread local or not""" + return self._is_thread_local + + @property + def is_singleton(self) -> bool: + """:return: a flag indicating if this lock is singleton or not""" + return self._is_singleton + + @property + def lock_file(self) -> str: + """:return: path to the lock file""" + return self._context.lock_file + + @property + def timeout(self) -> float: + """ + :return: the default timeout value, in seconds + + .. versionadded:: 2.0.0 + """ + return self._context.timeout + + @timeout.setter + def timeout(self, value: float | str) -> None: + """ + Change the default timeout value. + + :param value: the new value, in seconds + + """ + self._context.timeout = float(value) + + @property + def blocking(self) -> bool: + """:return: whether the locking is blocking or not""" + return self._context.blocking + + @blocking.setter + def blocking(self, value: bool) -> None: + """ + Change the default blocking value. + + :param value: the new value as bool + + """ + self._context.blocking = value + + @property + def mode(self) -> int: + """:return: the file permissions for the lockfile""" + return self._context.mode + + @abstractmethod + def _acquire(self) -> None: + """If the file lock could be acquired, self._context.lock_file_fd holds the file descriptor of the lock file.""" + raise NotImplementedError + + @abstractmethod + def _release(self) -> None: + """Releases the lock and sets self._context.lock_file_fd to None.""" + raise NotImplementedError + + @property + def is_locked(self) -> bool: + """ + + :return: A boolean indicating if the lock file is holding the lock currently. + + .. versionchanged:: 2.0.0 + + This was previously a method and is now a property. + """ + return self._context.lock_file_fd is not None + + @property + def lock_counter(self) -> int: + """:return: The number of times this lock has been acquired (but not yet released).""" + return self._context.lock_counter + + def acquire( + self, + timeout: float | None = None, + poll_interval: float = 0.05, + *, + poll_intervall: float | None = None, + blocking: bool | None = None, + ) -> AcquireReturnProxy: + """ + Try to acquire the file lock. + + :param timeout: maximum wait time for acquiring the lock, ``None`` means use the default :attr:`~timeout` is and + if ``timeout < 0``, there is no timeout and this method will block until the lock could be acquired + :param poll_interval: interval of trying to acquire the lock file + :param poll_intervall: deprecated, kept for backwards compatibility, use ``poll_interval`` instead + :param blocking: defaults to True. If False, function will return immediately if it cannot obtain a lock on the + first attempt. Otherwise, this method will block until the timeout expires or the lock is acquired. + :raises Timeout: if fails to acquire lock within the timeout period + :return: a context object that will unlock the file when the context is exited + + .. code-block:: python + + # You can use this method in the context manager (recommended) + with lock.acquire(): + pass + + # Or use an equivalent try-finally construct: + lock.acquire() + try: + pass + finally: + lock.release() + + .. versionchanged:: 2.0.0 + + This method returns now a *proxy* object instead of *self*, + so that it can be used in a with statement without side effects. + + """ + # Use the default timeout, if no timeout is provided. + if timeout is None: + timeout = self._context.timeout + + if blocking is None: + blocking = self._context.blocking + + if poll_intervall is not None: + msg = "use poll_interval instead of poll_intervall" + warnings.warn(msg, DeprecationWarning, stacklevel=2) + poll_interval = poll_intervall + + # Increment the number right at the beginning. We can still undo it, if something fails. + self._context.lock_counter += 1 + + lock_id = id(self) + lock_filename = self.lock_file + start_time = time.perf_counter() + try: + while True: + if not self.is_locked: + _LOGGER.debug("Attempting to acquire lock %s on %s", lock_id, lock_filename) + self._acquire() + if self.is_locked: + _LOGGER.debug("Lock %s acquired on %s", lock_id, lock_filename) + break + if blocking is False: + _LOGGER.debug("Failed to immediately acquire lock %s on %s", lock_id, lock_filename) + raise Timeout(lock_filename) # noqa: TRY301 + if 0 <= timeout < time.perf_counter() - start_time: + _LOGGER.debug("Timeout on acquiring lock %s on %s", lock_id, lock_filename) + raise Timeout(lock_filename) # noqa: TRY301 + msg = "Lock %s not acquired on %s, waiting %s seconds ..." + _LOGGER.debug(msg, lock_id, lock_filename, poll_interval) + time.sleep(poll_interval) + except BaseException: # Something did go wrong, so decrement the counter. + self._context.lock_counter = max(0, self._context.lock_counter - 1) + raise + return AcquireReturnProxy(lock=self) + + def release(self, force: bool = False) -> None: # noqa: FBT001, FBT002 + """ + Releases the file lock. Please note, that the lock is only completely released, if the lock counter is 0. + Also note, that the lock file itself is not automatically deleted. + + :param force: If true, the lock counter is ignored and the lock is released in every case/ + + """ + if self.is_locked: + self._context.lock_counter -= 1 + + if self._context.lock_counter == 0 or force: + lock_id, lock_filename = id(self), self.lock_file + + _LOGGER.debug("Attempting to release lock %s on %s", lock_id, lock_filename) + self._release() + self._context.lock_counter = 0 + _LOGGER.debug("Lock %s released on %s", lock_id, lock_filename) + + def __enter__(self) -> Self: + """ + Acquire the lock. + + :return: the lock object + + """ + self.acquire() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + """ + Release the lock. + + :param exc_type: the exception type if raised + :param exc_value: the exception value if raised + :param traceback: the exception traceback if raised + + """ + self.release() + + def __del__(self) -> None: + """Called when the lock object is deleted.""" + self.release(force=True) + + +__all__ = [ + "AcquireReturnProxy", + "BaseFileLock", +] diff --git a/venv/lib/python3.10/site-packages/filelock/_error.py b/venv/lib/python3.10/site-packages/filelock/_error.py new file mode 100644 index 0000000000000000000000000000000000000000..f7ff08c0f508ad7077eb6ed1990898840c952b3a --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_error.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import Any + + +class Timeout(TimeoutError): # noqa: N818 + """Raised when the lock could not be acquired in *timeout* seconds.""" + + def __init__(self, lock_file: str) -> None: + super().__init__() + self._lock_file = lock_file + + def __reduce__(self) -> str | tuple[Any, ...]: + return self.__class__, (self._lock_file,) # Properly pickle the exception + + def __str__(self) -> str: + return f"The file lock '{self._lock_file}' could not be acquired." + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.lock_file!r})" + + @property + def lock_file(self) -> str: + """:return: The path of the file lock.""" + return self._lock_file + + +__all__ = [ + "Timeout", +] diff --git a/venv/lib/python3.10/site-packages/filelock/_soft.py b/venv/lib/python3.10/site-packages/filelock/_soft.py new file mode 100644 index 0000000000000000000000000000000000000000..28c67f74cc82b8f55e47afd6a71972cc1fb95eb6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_soft.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import os +import sys +from contextlib import suppress +from errno import EACCES, EEXIST +from pathlib import Path + +from ._api import BaseFileLock +from ._util import ensure_directory_exists, raise_on_not_writable_file + + +class SoftFileLock(BaseFileLock): + """Simply watches the existence of the lock file.""" + + def _acquire(self) -> None: + raise_on_not_writable_file(self.lock_file) + ensure_directory_exists(self.lock_file) + # first check for exists and read-only mode as the open will mask this case as EEXIST + flags = ( + os.O_WRONLY # open for writing only + | os.O_CREAT + | os.O_EXCL # together with above raise EEXIST if the file specified by filename exists + | os.O_TRUNC # truncate the file to zero byte + ) + try: + file_handler = os.open(self.lock_file, flags, self._context.mode) + except OSError as exception: # re-raise unless expected exception + if not ( + exception.errno == EEXIST # lock already exist + or (exception.errno == EACCES and sys.platform == "win32") # has no access to this lock + ): # pragma: win32 no cover + raise + else: + self._context.lock_file_fd = file_handler + + def _release(self) -> None: + assert self._context.lock_file_fd is not None # noqa: S101 + os.close(self._context.lock_file_fd) # the lock file is definitely not None + self._context.lock_file_fd = None + with suppress(OSError): # the file is already deleted and that's what we want + Path(self.lock_file).unlink() + + +__all__ = [ + "SoftFileLock", +] diff --git a/venv/lib/python3.10/site-packages/filelock/_unix.py b/venv/lib/python3.10/site-packages/filelock/_unix.py new file mode 100644 index 0000000000000000000000000000000000000000..b2fd0f33d25d2bdf4a2a883380154771b4a25f9b --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_unix.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import os +import sys +from contextlib import suppress +from errno import ENOSYS +from pathlib import Path +from typing import cast + +from ._api import BaseFileLock +from ._util import ensure_directory_exists + +#: a flag to indicate if the fcntl API is available +has_fcntl = False +if sys.platform == "win32": # pragma: win32 cover + + class UnixFileLock(BaseFileLock): + """Uses the :func:`fcntl.flock` to hard lock the lock file on unix systems.""" + + def _acquire(self) -> None: + raise NotImplementedError + + def _release(self) -> None: + raise NotImplementedError + +else: # pragma: win32 no cover + try: + import fcntl + + _ = (fcntl.flock, fcntl.LOCK_EX, fcntl.LOCK_NB, fcntl.LOCK_UN) + except (ImportError, AttributeError): + pass + else: + has_fcntl = True + + class UnixFileLock(BaseFileLock): + """Uses the :func:`fcntl.flock` to hard lock the lock file on unix systems.""" + + def _acquire(self) -> None: + ensure_directory_exists(self.lock_file) + open_flags = os.O_RDWR | os.O_TRUNC + if not Path(self.lock_file).exists(): + open_flags |= os.O_CREAT + fd = os.open(self.lock_file, open_flags, self._context.mode) + with suppress(PermissionError): # This locked is not owned by this UID + os.fchmod(fd, self._context.mode) + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exception: + os.close(fd) + if exception.errno == ENOSYS: # NotImplemented error + msg = "FileSystem does not appear to support flock; use SoftFileLock instead" + raise NotImplementedError(msg) from exception + else: + self._context.lock_file_fd = fd + + def _release(self) -> None: + # Do not remove the lockfile: + # https://github.com/tox-dev/py-filelock/issues/31 + # https://stackoverflow.com/questions/17708885/flock-removing-locked-file-without-race-condition + fd = cast("int", self._context.lock_file_fd) + self._context.lock_file_fd = None + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +__all__ = [ + "UnixFileLock", + "has_fcntl", +] diff --git a/venv/lib/python3.10/site-packages/filelock/_util.py b/venv/lib/python3.10/site-packages/filelock/_util.py new file mode 100644 index 0000000000000000000000000000000000000000..c671e8533873948f0e1b5575ff952c722019f067 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_util.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import os +import stat +import sys +from errno import EACCES, EISDIR +from pathlib import Path + + +def raise_on_not_writable_file(filename: str) -> None: + """ + Raise an exception if attempting to open the file for writing would fail. + + This is done so files that will never be writable can be separated from files that are writable but currently + locked. + + :param filename: file to check + :raises OSError: as if the file was opened for writing. + + """ + try: # use stat to do exists + can write to check without race condition + file_stat = os.stat(filename) # noqa: PTH116 + except OSError: + return # swallow does not exist or other errors + + if file_stat.st_mtime != 0: # if os.stat returns but modification is zero that's an invalid os.stat - ignore it + if not (file_stat.st_mode & stat.S_IWUSR): + raise PermissionError(EACCES, "Permission denied", filename) + + if stat.S_ISDIR(file_stat.st_mode): + if sys.platform == "win32": # pragma: win32 cover + # On Windows, this is PermissionError + raise PermissionError(EACCES, "Permission denied", filename) + else: # pragma: win32 no cover # noqa: RET506 + # On linux / macOS, this is IsADirectoryError + raise IsADirectoryError(EISDIR, "Is a directory", filename) + + +def ensure_directory_exists(filename: Path | str) -> None: + """ + Ensure the directory containing the file exists (create it if necessary). + + :param filename: file. + + """ + Path(filename).parent.mkdir(parents=True, exist_ok=True) + + +__all__ = [ + "ensure_directory_exists", + "raise_on_not_writable_file", +] diff --git a/venv/lib/python3.10/site-packages/filelock/_windows.py b/venv/lib/python3.10/site-packages/filelock/_windows.py new file mode 100644 index 0000000000000000000000000000000000000000..348251d1067c28c55a6a267f8d11337abfae837f --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/_windows.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import os +import sys +from contextlib import suppress +from errno import EACCES +from pathlib import Path +from typing import cast + +from ._api import BaseFileLock +from ._util import ensure_directory_exists, raise_on_not_writable_file + +if sys.platform == "win32": # pragma: win32 cover + import msvcrt + + class WindowsFileLock(BaseFileLock): + """Uses the :func:`msvcrt.locking` function to hard lock the lock file on Windows systems.""" + + def _acquire(self) -> None: + raise_on_not_writable_file(self.lock_file) + ensure_directory_exists(self.lock_file) + flags = ( + os.O_RDWR # open for read and write + | os.O_CREAT # create file if not exists + | os.O_TRUNC # truncate file if not empty + ) + try: + fd = os.open(self.lock_file, flags, self._context.mode) + except OSError as exception: + if exception.errno != EACCES: # has no access to this lock + raise + else: + try: + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + except OSError as exception: + os.close(fd) # close file first + if exception.errno != EACCES: # file is already locked + raise + else: + self._context.lock_file_fd = fd + + def _release(self) -> None: + fd = cast("int", self._context.lock_file_fd) + self._context.lock_file_fd = None + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + os.close(fd) + + with suppress(OSError): # Probably another instance of the application hat acquired the file lock. + Path(self.lock_file).unlink() + +else: # pragma: win32 no cover + + class WindowsFileLock(BaseFileLock): + """Uses the :func:`msvcrt.locking` function to hard lock the lock file on Windows systems.""" + + def _acquire(self) -> None: + raise NotImplementedError + + def _release(self) -> None: + raise NotImplementedError + + +__all__ = [ + "WindowsFileLock", +] diff --git a/venv/lib/python3.10/site-packages/filelock/asyncio.py b/venv/lib/python3.10/site-packages/filelock/asyncio.py new file mode 100644 index 0000000000000000000000000000000000000000..022d0ef6995dbcf34764366313f4f3f1e9ce8ed6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/asyncio.py @@ -0,0 +1,344 @@ +"""An asyncio-based implementation of the file lock.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import time +from dataclasses import dataclass +from inspect import iscoroutinefunction +from threading import local +from typing import TYPE_CHECKING, Any, NoReturn, cast + +from ._api import BaseFileLock, FileLockContext, FileLockMeta +from ._error import Timeout +from ._soft import SoftFileLock +from ._unix import UnixFileLock +from ._windows import WindowsFileLock + +if TYPE_CHECKING: + import sys + from collections.abc import Callable + from concurrent import futures + from types import TracebackType + + if sys.version_info >= (3, 11): # pragma: no cover (py311+) + from typing import Self + else: # pragma: no cover ( None: # noqa: D107 + self.lock = lock + + async def __aenter__(self) -> BaseAsyncFileLock: # noqa: D105 + return self.lock + + async def __aexit__( # noqa: D105 + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.lock.release() + + +class AsyncFileLockMeta(FileLockMeta): + def __call__( # type: ignore[override] # noqa: PLR0913 + cls, # noqa: N805 + lock_file: str | os.PathLike[str], + timeout: float = -1, + mode: int = 0o644, + thread_local: bool = False, # noqa: FBT001, FBT002 + *, + blocking: bool = True, + is_singleton: bool = False, + loop: asyncio.AbstractEventLoop | None = None, + run_in_executor: bool = True, + executor: futures.Executor | None = None, + ) -> BaseAsyncFileLock: + if thread_local and run_in_executor: + msg = "run_in_executor is not supported when thread_local is True" + raise ValueError(msg) + instance = super().__call__( + lock_file=lock_file, + timeout=timeout, + mode=mode, + thread_local=thread_local, + blocking=blocking, + is_singleton=is_singleton, + loop=loop, + run_in_executor=run_in_executor, + executor=executor, + ) + return cast("BaseAsyncFileLock", instance) + + +class BaseAsyncFileLock(BaseFileLock, metaclass=AsyncFileLockMeta): + """Base class for asynchronous file locks.""" + + def __init__( # noqa: PLR0913 + self, + lock_file: str | os.PathLike[str], + timeout: float = -1, + mode: int = 0o644, + thread_local: bool = False, # noqa: FBT001, FBT002 + *, + blocking: bool = True, + is_singleton: bool = False, + loop: asyncio.AbstractEventLoop | None = None, + run_in_executor: bool = True, + executor: futures.Executor | None = None, + ) -> None: + """ + Create a new lock object. + + :param lock_file: path to the file + :param timeout: default timeout when acquiring the lock, in seconds. It will be used as fallback value in \ + the acquire method, if no timeout value (``None``) is given. If you want to disable the timeout, set it \ + to a negative value. A timeout of 0 means that there is exactly one attempt to acquire the file lock. + :param mode: file permissions for the lockfile + :param thread_local: Whether this object's internal context should be thread local or not. If this is set to \ + ``False`` then the lock will be reentrant across threads. + :param blocking: whether the lock should be blocking or not + :param is_singleton: If this is set to ``True`` then only one instance of this class will be created \ + per lock file. This is useful if you want to use the lock object for reentrant locking without needing \ + to pass the same object around. + :param loop: The event loop to use. If not specified, the running event loop will be used. + :param run_in_executor: If this is set to ``True`` then the lock will be acquired in an executor. + :param executor: The executor to use. If not specified, the default executor will be used. + + """ + self._is_thread_local = thread_local + self._is_singleton = is_singleton + + # Create the context. Note that external code should not work with the context directly and should instead use + # properties of this class. + kwargs: dict[str, Any] = { + "lock_file": os.fspath(lock_file), + "timeout": timeout, + "mode": mode, + "blocking": blocking, + "loop": loop, + "run_in_executor": run_in_executor, + "executor": executor, + } + self._context: AsyncFileLockContext = (AsyncThreadLocalFileContext if thread_local else AsyncFileLockContext)( + **kwargs + ) + + @property + def run_in_executor(self) -> bool: + """::return: whether run in executor.""" + return self._context.run_in_executor + + @property + def executor(self) -> futures.Executor | None: + """::return: the executor.""" + return self._context.executor + + @executor.setter + def executor(self, value: futures.Executor | None) -> None: # pragma: no cover + """ + Change the executor. + + :param value: the new executor or ``None`` + :type value: futures.Executor | None + + """ + self._context.executor = value + + @property + def loop(self) -> asyncio.AbstractEventLoop | None: + """::return: the event loop.""" + return self._context.loop + + async def acquire( # type: ignore[override] + self, + timeout: float | None = None, + poll_interval: float = 0.05, + *, + blocking: bool | None = None, + ) -> AsyncAcquireReturnProxy: + """ + Try to acquire the file lock. + + :param timeout: maximum wait time for acquiring the lock, ``None`` means use the default + :attr:`~BaseFileLock.timeout` is and if ``timeout < 0``, there is no timeout and + this method will block until the lock could be acquired + :param poll_interval: interval of trying to acquire the lock file + :param blocking: defaults to True. If False, function will return immediately if it cannot obtain a lock on the + first attempt. Otherwise, this method will block until the timeout expires or the lock is acquired. + :raises Timeout: if fails to acquire lock within the timeout period + :return: a context object that will unlock the file when the context is exited + + .. code-block:: python + + # You can use this method in the context manager (recommended) + with lock.acquire(): + pass + + # Or use an equivalent try-finally construct: + lock.acquire() + try: + pass + finally: + lock.release() + + """ + # Use the default timeout, if no timeout is provided. + if timeout is None: + timeout = self._context.timeout + + if blocking is None: + blocking = self._context.blocking + + # Increment the number right at the beginning. We can still undo it, if something fails. + self._context.lock_counter += 1 + + lock_id = id(self) + lock_filename = self.lock_file + start_time = time.perf_counter() + try: + while True: + if not self.is_locked: + _LOGGER.debug("Attempting to acquire lock %s on %s", lock_id, lock_filename) + await self._run_internal_method(self._acquire) + if self.is_locked: + _LOGGER.debug("Lock %s acquired on %s", lock_id, lock_filename) + break + if blocking is False: + _LOGGER.debug("Failed to immediately acquire lock %s on %s", lock_id, lock_filename) + raise Timeout(lock_filename) # noqa: TRY301 + if 0 <= timeout < time.perf_counter() - start_time: + _LOGGER.debug("Timeout on acquiring lock %s on %s", lock_id, lock_filename) + raise Timeout(lock_filename) # noqa: TRY301 + msg = "Lock %s not acquired on %s, waiting %s seconds ..." + _LOGGER.debug(msg, lock_id, lock_filename, poll_interval) + await asyncio.sleep(poll_interval) + except BaseException: # Something did go wrong, so decrement the counter. + self._context.lock_counter = max(0, self._context.lock_counter - 1) + raise + return AsyncAcquireReturnProxy(lock=self) + + async def release(self, force: bool = False) -> None: # type: ignore[override] # noqa: FBT001, FBT002 + """ + Releases the file lock. Please note, that the lock is only completely released, if the lock counter is 0. + Also note, that the lock file itself is not automatically deleted. + + :param force: If true, the lock counter is ignored and the lock is released in every case/ + + """ + if self.is_locked: + self._context.lock_counter -= 1 + + if self._context.lock_counter == 0 or force: + lock_id, lock_filename = id(self), self.lock_file + + _LOGGER.debug("Attempting to release lock %s on %s", lock_id, lock_filename) + await self._run_internal_method(self._release) + self._context.lock_counter = 0 + _LOGGER.debug("Lock %s released on %s", lock_id, lock_filename) + + async def _run_internal_method(self, method: Callable[[], Any]) -> None: + if iscoroutinefunction(method): + await method() + elif self.run_in_executor: + loop = self.loop or asyncio.get_running_loop() + await loop.run_in_executor(self.executor, method) + else: + method() + + def __enter__(self) -> NoReturn: + """ + Replace old __enter__ method to avoid using it. + + NOTE: DO NOT USE `with` FOR ASYNCIO LOCKS, USE `async with` INSTEAD. + + :return: none + :rtype: NoReturn + """ + msg = "Do not use `with` for asyncio locks, use `async with` instead." + raise NotImplementedError(msg) + + async def __aenter__(self) -> Self: + """ + Acquire the lock. + + :return: the lock object + + """ + await self.acquire() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + """ + Release the lock. + + :param exc_type: the exception type if raised + :param exc_value: the exception value if raised + :param traceback: the exception traceback if raised + + """ + await self.release() + + def __del__(self) -> None: + """Called when the lock object is deleted.""" + with contextlib.suppress(RuntimeError): + loop = self.loop or asyncio.get_running_loop() + if not loop.is_running(): # pragma: no cover + loop.run_until_complete(self.release(force=True)) + else: + loop.create_task(self.release(force=True)) + + +class AsyncSoftFileLock(SoftFileLock, BaseAsyncFileLock): + """Simply watches the existence of the lock file.""" + + +class AsyncUnixFileLock(UnixFileLock, BaseAsyncFileLock): + """Uses the :func:`fcntl.flock` to hard lock the lock file on unix systems.""" + + +class AsyncWindowsFileLock(WindowsFileLock, BaseAsyncFileLock): + """Uses the :func:`msvcrt.locking` to hard lock the lock file on windows systems.""" + + +__all__ = [ + "AsyncAcquireReturnProxy", + "AsyncSoftFileLock", + "AsyncUnixFileLock", + "AsyncWindowsFileLock", + "BaseAsyncFileLock", +] diff --git a/venv/lib/python3.10/site-packages/filelock/py.typed b/venv/lib/python3.10/site-packages/filelock/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/filelock/version.py b/venv/lib/python3.10/site-packages/filelock/version.py new file mode 100644 index 0000000000000000000000000000000000000000..093125cd6ae5a94fa087734b389ccbb3082495a1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/filelock/version.py @@ -0,0 +1,34 @@ +# file generated by setuptools-scm +# don't change, don't track in version control + +__all__ = [ + "__version__", + "__version_tuple__", + "version", + "version_tuple", + "__commit_id__", + "commit_id", +] + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Tuple + from typing import Union + + VERSION_TUPLE = Tuple[Union[int, str], ...] + COMMIT_ID = Union[str, None] +else: + VERSION_TUPLE = object + COMMIT_ID = object + +version: str +__version__: str +__version_tuple__: VERSION_TUPLE +version_tuple: VERSION_TUPLE +commit_id: COMMIT_ID +__commit_id__: COMMIT_ID + +__version__ = version = '3.20.0' +__version_tuple__ = version_tuple = (3, 20, 0) + +__commit_id__ = commit_id = None diff --git a/venv/lib/python3.10/site-packages/hf_xet/__init__.py b/venv/lib/python3.10/site-packages/hf_xet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..96ed54a8a066d681e4973e5841a0f5577b619698 --- /dev/null +++ b/venv/lib/python3.10/site-packages/hf_xet/__init__.py @@ -0,0 +1,5 @@ +from .hf_xet import * + +__doc__ = hf_xet.__doc__ +if hasattr(hf_xet, "__all__"): + __all__ = hf_xet.__all__ \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/hf_xet/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/hf_xet/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00ebd06a1e257ad00a97eeec9021845a369f31fb Binary files /dev/null and b/venv/lib/python3.10/site-packages/hf_xet/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/kiwisolver/__init__.py b/venv/lib/python3.10/site-packages/kiwisolver/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa3bcfe3be9877b5e17b072cf20d00085e6c8efb --- /dev/null +++ b/venv/lib/python3.10/site-packages/kiwisolver/__init__.py @@ -0,0 +1,42 @@ +# -------------------------------------------------------------------------------------- +# Copyright (c) 2013-2025, Nucleic Development Team. +# +# Distributed under the terms of the Modified BSD License. +# +# The full license is in the file LICENSE, distributed with this software. +# -------------------------------------------------------------------------------------- +from ._cext import ( + Constraint, + Expression, + Solver, + Term, + Variable, + __kiwi_version__, + __version__, + strength, +) +from .exceptions import ( + BadRequiredStrength, + DuplicateConstraint, + DuplicateEditVariable, + UnknownConstraint, + UnknownEditVariable, + UnsatisfiableConstraint, +) + +__all__ = [ + "BadRequiredStrength", + "Constraint", + "DuplicateConstraint", + "DuplicateEditVariable", + "Expression", + "Solver", + "Term", + "UnknownConstraint", + "UnknownEditVariable", + "UnsatisfiableConstraint", + "Variable", + "__kiwi_version__", + "__version__", + "strength", +] diff --git a/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c75b34c92ba38c400cc6b7429a4c04fe7a1b4934 Binary files /dev/null and b/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/exceptions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54a4ad75869e331385ab313b06535c64170780fe Binary files /dev/null and b/venv/lib/python3.10/site-packages/kiwisolver/__pycache__/exceptions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/kiwisolver/_cext.pyi b/venv/lib/python3.10/site-packages/kiwisolver/_cext.pyi new file mode 100644 index 0000000000000000000000000000000000000000..33fd65693d42ebd04078263394f1b43802ae1b46 --- /dev/null +++ b/venv/lib/python3.10/site-packages/kiwisolver/_cext.pyi @@ -0,0 +1,228 @@ +# -------------------------------------------------------------------------------------- +# Copyright (c) 2021-2024, Nucleic Development Team. +# +# Distributed under the terms of the Modified BSD License. +# +# The full license is in the file LICENSE, distributed with this software. +# -------------------------------------------------------------------------------------- + +from typing import Any, Iterable, NoReturn, Tuple, type_check_only + +try: + from typing import Literal +except ImportError: + from typing_extensions import Literal # type: ignore + +__version__: str +__kiwi_version__: str + +# Types +@type_check_only +class Strength: + @property + def weak(self) -> float: ... + @property + def medium(self) -> float: ... + @property + def strong(self) -> float: ... + @property + def required(self) -> float: ... + def create( + self, + a: int | float, + b: int | float, + c: int | float, + weight: int | float = 1.0, + /, + ) -> float: ... + +# This is meant as a singleton and users should not access the Strength type. +strength: Strength + +class Variable: + """Variable to express a constraint in a solver.""" + + __hash__: None # type: ignore + def __init__(self, name: str = "", context: Any = None, /) -> None: ... + def name(self) -> str: + """Get the name of the variable.""" + ... + def setName(self, name: str, /) -> Any: + """Set the name of the variable.""" + ... + def value(self) -> float: + """Get the current value of the variable.""" + ... + def context(self) -> Any: + """Get the context object associated with the variable.""" + ... + def setContext(self, context: Any, /) -> Any: + """Set the context object associated with the variable.""" + ... + def __neg__(self) -> Term: ... + def __add__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __radd__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __sub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __rsub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __mul__(self, other: float) -> Term: ... + def __rmul__(self, other: float) -> Term: ... + def __truediv__(self, other: float) -> Term: ... + def __rtruediv__(self, other: float) -> Term: ... + def __eq__(self, other: float | Variable | Term | Expression) -> Constraint: ... # type: ignore + def __ge__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __le__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __ne__(self, other: Any) -> NoReturn: ... + def __gt__(self, other: Any) -> NoReturn: ... + def __lt__(self, other: Any) -> NoReturn: ... + +class Term: + """Product of a variable by a constant pre-factor.""" + + __hash__: None # type: ignore + def __init__( + self, variable: Variable, coefficient: int | float = 1.0, / + ) -> None: ... + def coefficient(self) -> float: + """Get the coefficient for the term.""" + ... + def variable(self) -> Variable: + """Get the variable for the term.""" + ... + def value(self) -> float: + """Get the value for the term.""" + ... + def __neg__(self) -> Term: ... + def __add__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __radd__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __sub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __rsub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __mul__(self, other: float) -> Term: ... + def __rmul__(self, other: float) -> Term: ... + def __truediv__(self, other: float) -> Term: ... + def __rtruediv__(self, other: float) -> Term: ... + def __eq__(self, other: float | Variable | Term | Expression) -> Constraint: ... # type: ignore + def __ge__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __le__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __ne__(self, other: Any) -> NoReturn: ... + def __gt__(self, other: Any) -> NoReturn: ... + def __lt__(self, other: Any) -> NoReturn: ... + +class Expression: + """Sum of terms and an additional constant.""" + + __hash__: None # type: ignore + def __init__( + self, terms: Iterable[Term], constant: int | float = 0.0, / + ) -> None: ... + def constant(self) -> float: + """Get the constant for the expression.""" + ... + def terms(self) -> Tuple[Term, ...]: + """Get the tuple of terms for the expression.""" + ... + def value(self) -> float: + """Get the value for the expression.""" + ... + def __neg__(self) -> Expression: ... + def __add__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __radd__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __sub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __rsub__(self, other: float | Variable | Term | Expression) -> Expression: ... + def __mul__(self, other: float) -> Expression: ... + def __rmul__(self, other: float) -> Expression: ... + def __truediv__(self, other: float) -> Expression: ... + def __rtruediv__(self, other: float) -> Expression: ... + def __eq__(self, other: float | Variable | Term | Expression) -> Constraint: ... # type: ignore + def __ge__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __le__(self, other: float | Variable | Term | Expression) -> Constraint: ... + def __ne__(self, other: Any) -> NoReturn: ... + def __gt__(self, other: Any) -> NoReturn: ... + def __lt__(self, other: Any) -> NoReturn: ... + +class Constraint: + def __init__( + self, + expression: Expression, + op: Literal["=="] | Literal["<="] | Literal[">="], + strength: float + | Literal["weak"] + | Literal["medium"] + | Literal["strong"] + | Literal["required"] = "required", + /, + ) -> None: ... + def expression(self) -> Expression: + """Get the expression object for the constraint.""" + ... + def op(self) -> Literal["=="] | Literal["<="] | Literal[">="]: + """Get the relational operator for the constraint.""" + ... + def strength(self) -> float: + """Get the strength for the constraint.""" + ... + def violated(self) -> bool: + """Indicate if the constraint is violated in teh current state of the solver.""" + ... + def __or__( + self, + other: float + | Literal["weak"] + | Literal["medium"] + | Literal["strong"] + | Literal["required"], + ) -> Constraint: ... + def __ror__( + self, + other: float + | Literal["weak"] + | Literal["medium"] + | Literal["strong"] + | Literal["required"], + ) -> Constraint: ... + +class Solver: + """Kiwi solver class.""" + + def __init__(self) -> None: ... + def addConstraint(self, constraint: Constraint, /) -> None: + """Add a constraint to the solver.""" + ... + def removeConstraint(self, constraint: Constraint, /) -> None: + """Remove a constraint from the solver.""" + ... + def hasConstraint(self, constraint: Constraint, /) -> bool: + """Check whether the solver contains a constraint.""" + ... + def addEditVariable( + self, + variable: Variable, + strength: float + | Literal["weak"] + | Literal["medium"] + | Literal["strong"] + | Literal["required"], + /, + ) -> None: + """Add an edit variable to the solver.""" + ... + def removeEditVariable(self, variable: Variable, /) -> None: + """Remove an edit variable from the solver.""" + ... + def hasEditVariable(self, variable: Variable, /) -> bool: + """Check whether the solver contains an edit variable.""" + ... + def suggestValue(self, variable: Variable, value: int | float, /) -> None: + """Suggest a desired value for an edit variable.""" + ... + def updateVariables(self) -> None: + """Update the values of the solver variables.""" + ... + def reset(self) -> None: + """Reset the solver to the initial empty starting condition.""" + ... + def dump(self) -> None: + """Dump a representation of the solver internals to stdout.""" + ... + def dumps(self) -> str: + """Dump a representation of the solver internals to a string.""" + ... diff --git a/venv/lib/python3.10/site-packages/kiwisolver/exceptions.py b/venv/lib/python3.10/site-packages/kiwisolver/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..fd06f23f3a76dae34f03708a831b7f34ff175343 --- /dev/null +++ b/venv/lib/python3.10/site-packages/kiwisolver/exceptions.py @@ -0,0 +1,51 @@ +# -------------------------------------------------------------------------------------- +# Copyright (c) 2023-2024, Nucleic Development Team. +# +# Distributed under the terms of the Modified BSD License. +# +# The full license is in the file LICENSE, distributed with this software. +# -------------------------------------------------------------------------------------- +"""Kiwi exceptions. + +Imported by the kiwisolver C extension. + +""" + + +class BadRequiredStrength(Exception): + pass + + +class DuplicateConstraint(Exception): + __slots__ = ("constraint",) + + def __init__(self, constraint): + self.constraint = constraint + + +class DuplicateEditVariable(Exception): + __slots__ = ("edit_variable",) + + def __init__(self, edit_variable): + self.edit_variable = edit_variable + + +class UnknownConstraint(Exception): + __slots__ = ("constraint",) + + def __init__(self, constraint): + self.constraint = constraint + + +class UnknownEditVariable(Exception): + __slots__ = ("edit_variable",) + + def __init__(self, edit_variable): + self.edit_variable = edit_variable + + +class UnsatisfiableConstraint(Exception): + __slots__ = ("constraint",) + + def __init__(self, constraint): + self.constraint = constraint diff --git a/venv/lib/python3.10/site-packages/kiwisolver/py.typed b/venv/lib/python3.10/site-packages/kiwisolver/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/LICENSE b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..ec51537db27dd4d9c9ed3cd39fd96485f3cfddea --- /dev/null +++ b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/LICENSE @@ -0,0 +1,99 @@ +License agreement for matplotlib versions 1.3.0 and later +========================================================= + +1. This LICENSE AGREEMENT is between the Matplotlib Development Team +("MDT"), and the Individual or Organization ("Licensee") accessing and +otherwise using matplotlib software in source or binary form and its +associated documentation. + +2. Subject to the terms and conditions of this License Agreement, MDT +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that MDT's +License Agreement and MDT's notice of copyright, i.e., "Copyright (c) +2012- Matplotlib Development Team; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib . + +4. MDT is making matplotlib available to Licensee on an "AS +IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between MDT and +Licensee. This License Agreement does not grant permission to use MDT +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib , +Licensee agrees to be bound by the terms and conditions of this License +Agreement. + +License agreement for matplotlib versions prior to 1.3.0 +======================================================== + +1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the +Individual or Organization ("Licensee") accessing and otherwise using +matplotlib software in source or binary form and its associated +documentation. + +2. Subject to the terms and conditions of this License Agreement, JDH +hereby grants Licensee a nonexclusive, royalty-free, world-wide license +to reproduce, analyze, test, perform and/or display publicly, prepare +derivative works, distribute, and otherwise use matplotlib +alone or in any derivative version, provided, however, that JDH's +License Agreement and JDH's notice of copyright, i.e., "Copyright (c) +2002-2011 John D. Hunter; All Rights Reserved" are retained in +matplotlib alone or in any derivative version prepared by +Licensee. + +3. In the event Licensee prepares a derivative work that is based on or +incorporates matplotlib or any part thereof, and wants to +make the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to matplotlib. + +4. JDH is making matplotlib available to Licensee on an "AS +IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB +WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + +5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR +LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING +MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF +THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between JDH and +Licensee. This License Agreement does not grant permission to use JDH +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using matplotlib, +Licensee agrees to be bound by the terms and conditions of this License +Agreement. \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/METADATA b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c8608ba08a5fcd1fa649eee9765dc174e619f01f --- /dev/null +++ b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/METADATA @@ -0,0 +1,215 @@ +Metadata-Version: 2.1 +Name: matplotlib +Version: 3.10.7 +Summary: Python plotting package +Author: John D. Hunter, Michael Droettboom +Author-Email: Unknown +License: License agreement for matplotlib versions 1.3.0 and later + ========================================================= + + 1. This LICENSE AGREEMENT is between the Matplotlib Development Team + ("MDT"), and the Individual or Organization ("Licensee") accessing and + otherwise using matplotlib software in source or binary form and its + associated documentation. + + 2. Subject to the terms and conditions of this License Agreement, MDT + hereby grants Licensee a nonexclusive, royalty-free, world-wide license + to reproduce, analyze, test, perform and/or display publicly, prepare + derivative works, distribute, and otherwise use matplotlib + alone or in any derivative version, provided, however, that MDT's + License Agreement and MDT's notice of copyright, i.e., "Copyright (c) + 2012- Matplotlib Development Team; All Rights Reserved" are retained in + matplotlib alone or in any derivative version prepared by + Licensee. + + 3. In the event Licensee prepares a derivative work that is based on or + incorporates matplotlib or any part thereof, and wants to + make the derivative work available to others as provided herein, then + Licensee hereby agrees to include in any such work a brief summary of + the changes made to matplotlib . + + 4. MDT is making matplotlib available to Licensee on an "AS + IS" basis. MDT MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR + IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, MDT MAKES NO AND + DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS + FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB + WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + + 5. MDT SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR + LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING + MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF + THE POSSIBILITY THEREOF. + + 6. This License Agreement will automatically terminate upon a material + breach of its terms and conditions. + + 7. Nothing in this License Agreement shall be deemed to create any + relationship of agency, partnership, or joint venture between MDT and + Licensee. This License Agreement does not grant permission to use MDT + trademarks or trade name in a trademark sense to endorse or promote + products or services of Licensee, or any third party. + + 8. By copying, installing or otherwise using matplotlib , + Licensee agrees to be bound by the terms and conditions of this License + Agreement. + + License agreement for matplotlib versions prior to 1.3.0 + ======================================================== + + 1. This LICENSE AGREEMENT is between John D. Hunter ("JDH"), and the + Individual or Organization ("Licensee") accessing and otherwise using + matplotlib software in source or binary form and its associated + documentation. + + 2. Subject to the terms and conditions of this License Agreement, JDH + hereby grants Licensee a nonexclusive, royalty-free, world-wide license + to reproduce, analyze, test, perform and/or display publicly, prepare + derivative works, distribute, and otherwise use matplotlib + alone or in any derivative version, provided, however, that JDH's + License Agreement and JDH's notice of copyright, i.e., "Copyright (c) + 2002-2011 John D. Hunter; All Rights Reserved" are retained in + matplotlib alone or in any derivative version prepared by + Licensee. + + 3. In the event Licensee prepares a derivative work that is based on or + incorporates matplotlib or any part thereof, and wants to + make the derivative work available to others as provided herein, then + Licensee hereby agrees to include in any such work a brief summary of + the changes made to matplotlib. + + 4. JDH is making matplotlib available to Licensee on an "AS + IS" basis. JDH MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR + IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, JDH MAKES NO AND + DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS + FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF MATPLOTLIB + WILL NOT INFRINGE ANY THIRD PARTY RIGHTS. + + 5. JDH SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF MATPLOTLIB + FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR + LOSS AS A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING + MATPLOTLIB , OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF + THE POSSIBILITY THEREOF. + + 6. This License Agreement will automatically terminate upon a material + breach of its terms and conditions. + + 7. Nothing in this License Agreement shall be deemed to create any + relationship of agency, partnership, or joint venture between JDH and + Licensee. This License Agreement does not grant permission to use JDH + trademarks or trade name in a trademark sense to endorse or promote + products or services of Licensee, or any third party. + + 8. By copying, installing or otherwise using matplotlib, + Licensee agrees to be bound by the terms and conditions of this License + Agreement. +Classifier: Development Status :: 5 - Production/Stable +Classifier: Framework :: Matplotlib +Classifier: Intended Audience :: Science/Research +Classifier: Intended Audience :: Education +Classifier: License :: OSI Approved :: Python Software Foundation License +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Topic :: Scientific/Engineering :: Visualization +Project-URL: Homepage, https://matplotlib.org +Project-URL: Download, https://matplotlib.org/stable/install/index.html +Project-URL: Documentation, https://matplotlib.org +Project-URL: Source Code, https://github.com/matplotlib/matplotlib +Project-URL: Bug Tracker, https://github.com/matplotlib/matplotlib/issues +Project-URL: Forum, https://discourse.matplotlib.org/ +Project-URL: Donate, https://numfocus.org/donate-to-matplotlib +Requires-Python: >=3.10 +Requires-Dist: contourpy>=1.0.1 +Requires-Dist: cycler>=0.10 +Requires-Dist: fonttools>=4.22.0 +Requires-Dist: kiwisolver>=1.3.1 +Requires-Dist: numpy>=1.23 +Requires-Dist: packaging>=20.0 +Requires-Dist: pillow>=8 +Requires-Dist: pyparsing>=3 +Requires-Dist: python-dateutil>=2.7 +Provides-Extra: dev +Requires-Dist: meson-python<0.17.0,>=0.13.1; extra == "dev" +Requires-Dist: pybind11!=2.13.3,>=2.13.2; extra == "dev" +Requires-Dist: setuptools_scm>=7; extra == "dev" +Requires-Dist: setuptools>=64; extra == "dev" +Description-Content-Type: text/markdown + +[![PyPi](https://img.shields.io/pypi/v/matplotlib)](https://pypi.org/project/matplotlib/) +[![Conda](https://img.shields.io/conda/vn/conda-forge/matplotlib)](https://anaconda.org/conda-forge/matplotlib) +[![Downloads](https://img.shields.io/pypi/dm/matplotlib)](https://pypi.org/project/matplotlib) +[![NUMFocus](https://img.shields.io/badge/powered%20by-NumFOCUS-orange.svg?style=flat&colorA=E1523D&colorB=007D8A)](https://numfocus.org) + +[![Discourse help forum](https://img.shields.io/badge/help_forum-discourse-blue.svg)](https://discourse.matplotlib.org) +[![Gitter](https://badges.gitter.im/matplotlib/matplotlib.svg)](https://gitter.im/matplotlib/matplotlib) +[![GitHub issues](https://img.shields.io/badge/issue_tracking-github-blue.svg)](https://github.com/matplotlib/matplotlib/issues) +[![Contributing](https://img.shields.io/badge/PR-Welcome-%23FF8300.svg?)](https://matplotlib.org/stable/devel/index.html) + +[![GitHub actions status](https://github.com/matplotlib/matplotlib/workflows/Tests/badge.svg)](https://github.com/matplotlib/matplotlib/actions?query=workflow%3ATests) +[![Azure pipelines status](https://dev.azure.com/matplotlib/matplotlib/_apis/build/status/matplotlib.matplotlib?branchName=main)](https://dev.azure.com/matplotlib/matplotlib/_build/latest?definitionId=1&branchName=main) +[![AppVeyor status](https://ci.appveyor.com/api/projects/status/github/matplotlib/matplotlib?branch=main&svg=true)](https://ci.appveyor.com/project/matplotlib/matplotlib) +[![Codecov status](https://codecov.io/github/matplotlib/matplotlib/badge.svg?branch=main&service=github)](https://app.codecov.io/gh/matplotlib/matplotlib) +[![EffVer Versioning](https://img.shields.io/badge/version_scheme-EffVer-0097a7)](https://jacobtomlinson.dev/effver) + +![Matplotlib logotype](https://matplotlib.org/_static/logo2.svg) + +Matplotlib is a comprehensive library for creating static, animated, and +interactive visualizations in Python. + +Check out our [home page](https://matplotlib.org/) for more information. + +![image](https://matplotlib.org/_static/readme_preview.png) + +Matplotlib produces publication-quality figures in a variety of hardcopy +formats and interactive environments across platforms. Matplotlib can be +used in Python scripts, Python/IPython shells, web application servers, +and various graphical user interface toolkits. + +## Install + +See the [install +documentation](https://matplotlib.org/stable/users/installing/index.html), +which is generated from `/doc/install/index.rst` + +## Contribute + +You've discovered a bug or something else you want to change — excellent! + +You've worked out a way to fix it — even better! + +You want to tell us about it — best of all! + +Start at the [contributing +guide](https://matplotlib.org/devdocs/devel/contribute.html)! + +## Contact + +[Discourse](https://discourse.matplotlib.org/) is the discussion forum +for general questions and discussions and our recommended starting +point. + +Our active mailing lists (which are mirrored on Discourse) are: + +- [Users](https://mail.python.org/mailman/listinfo/matplotlib-users) + mailing list: +- [Announcement](https://mail.python.org/mailman/listinfo/matplotlib-announce) + mailing list: +- [Development](https://mail.python.org/mailman/listinfo/matplotlib-devel) + mailing list: + +[Gitter](https://gitter.im/matplotlib/matplotlib) is for coordinating +development and asking questions directly related to contributing to +matplotlib. + +## Citing Matplotlib + +If Matplotlib contributes to a project that leads to publication, please +acknowledge this by citing Matplotlib. + +[A ready-made citation +entry](https://matplotlib.org/stable/users/project/citing.html) is +available. diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/RECORD b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..e2088cafc6fa52582d9153aff5c8355a87cafe11 --- /dev/null +++ b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/RECORD @@ -0,0 +1,884 @@ +__pycache__/pylab.cpython-310.pyc,, +matplotlib-3.10.7.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +matplotlib-3.10.7.dist-info/LICENSE,sha256=WhqB6jAXKMi7opM9qDLAzWIina8giToCSrPVMkRGjbw,4830 +matplotlib-3.10.7.dist-info/METADATA,sha256=_Yvvshgf49EfZHfDWdj0UrT7yQ9paaaJPgkEbZNeFnc,11264 +matplotlib-3.10.7.dist-info/RECORD,, +matplotlib-3.10.7.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib-3.10.7.dist-info/WHEEL,sha256=sZM_NeUMz2G4fDenMf11eikcCxcLaQWiYRmjwQBavQs,137 +matplotlib/__init__.py,sha256=xsR9Hs8l15NPOMilXDuOEW1BqgjY7PqC9q6K9Ykhajo,55823 +matplotlib/__init__.pyi,sha256=s0YJ8fkGVioZL9bcBtyYgCq_MbY8LRrQUtmaD3XaIxQ,3439 +matplotlib/__pycache__/__init__.cpython-310.pyc,, +matplotlib/__pycache__/_afm.cpython-310.pyc,, +matplotlib/__pycache__/_animation_data.cpython-310.pyc,, +matplotlib/__pycache__/_blocking_input.cpython-310.pyc,, +matplotlib/__pycache__/_cm.cpython-310.pyc,, +matplotlib/__pycache__/_cm_bivar.cpython-310.pyc,, +matplotlib/__pycache__/_cm_listed.cpython-310.pyc,, +matplotlib/__pycache__/_cm_multivar.cpython-310.pyc,, +matplotlib/__pycache__/_color_data.cpython-310.pyc,, +matplotlib/__pycache__/_constrained_layout.cpython-310.pyc,, +matplotlib/__pycache__/_docstring.cpython-310.pyc,, +matplotlib/__pycache__/_enums.cpython-310.pyc,, +matplotlib/__pycache__/_fontconfig_pattern.cpython-310.pyc,, +matplotlib/__pycache__/_internal_utils.cpython-310.pyc,, +matplotlib/__pycache__/_layoutgrid.cpython-310.pyc,, +matplotlib/__pycache__/_mathtext.cpython-310.pyc,, +matplotlib/__pycache__/_mathtext_data.cpython-310.pyc,, +matplotlib/__pycache__/_pylab_helpers.cpython-310.pyc,, +matplotlib/__pycache__/_text_helpers.cpython-310.pyc,, +matplotlib/__pycache__/_tight_bbox.cpython-310.pyc,, +matplotlib/__pycache__/_tight_layout.cpython-310.pyc,, +matplotlib/__pycache__/_type1font.cpython-310.pyc,, +matplotlib/__pycache__/_version.cpython-310.pyc,, +matplotlib/__pycache__/animation.cpython-310.pyc,, +matplotlib/__pycache__/artist.cpython-310.pyc,, +matplotlib/__pycache__/axis.cpython-310.pyc,, +matplotlib/__pycache__/backend_bases.cpython-310.pyc,, +matplotlib/__pycache__/backend_managers.cpython-310.pyc,, +matplotlib/__pycache__/backend_tools.cpython-310.pyc,, +matplotlib/__pycache__/bezier.cpython-310.pyc,, +matplotlib/__pycache__/category.cpython-310.pyc,, +matplotlib/__pycache__/cbook.cpython-310.pyc,, +matplotlib/__pycache__/cm.cpython-310.pyc,, +matplotlib/__pycache__/collections.cpython-310.pyc,, +matplotlib/__pycache__/colorbar.cpython-310.pyc,, +matplotlib/__pycache__/colorizer.cpython-310.pyc,, +matplotlib/__pycache__/colors.cpython-310.pyc,, +matplotlib/__pycache__/container.cpython-310.pyc,, +matplotlib/__pycache__/contour.cpython-310.pyc,, +matplotlib/__pycache__/dates.cpython-310.pyc,, +matplotlib/__pycache__/dviread.cpython-310.pyc,, +matplotlib/__pycache__/figure.cpython-310.pyc,, +matplotlib/__pycache__/font_manager.cpython-310.pyc,, +matplotlib/__pycache__/gridspec.cpython-310.pyc,, +matplotlib/__pycache__/hatch.cpython-310.pyc,, +matplotlib/__pycache__/image.cpython-310.pyc,, +matplotlib/__pycache__/inset.cpython-310.pyc,, +matplotlib/__pycache__/layout_engine.cpython-310.pyc,, +matplotlib/__pycache__/legend.cpython-310.pyc,, +matplotlib/__pycache__/legend_handler.cpython-310.pyc,, +matplotlib/__pycache__/lines.cpython-310.pyc,, +matplotlib/__pycache__/markers.cpython-310.pyc,, +matplotlib/__pycache__/mathtext.cpython-310.pyc,, +matplotlib/__pycache__/mlab.cpython-310.pyc,, +matplotlib/__pycache__/offsetbox.cpython-310.pyc,, +matplotlib/__pycache__/patches.cpython-310.pyc,, +matplotlib/__pycache__/path.cpython-310.pyc,, +matplotlib/__pycache__/patheffects.cpython-310.pyc,, +matplotlib/__pycache__/pylab.cpython-310.pyc,, +matplotlib/__pycache__/pyplot.cpython-310.pyc,, +matplotlib/__pycache__/quiver.cpython-310.pyc,, +matplotlib/__pycache__/rcsetup.cpython-310.pyc,, +matplotlib/__pycache__/sankey.cpython-310.pyc,, +matplotlib/__pycache__/scale.cpython-310.pyc,, +matplotlib/__pycache__/spines.cpython-310.pyc,, +matplotlib/__pycache__/stackplot.cpython-310.pyc,, +matplotlib/__pycache__/streamplot.cpython-310.pyc,, +matplotlib/__pycache__/table.cpython-310.pyc,, +matplotlib/__pycache__/texmanager.cpython-310.pyc,, +matplotlib/__pycache__/text.cpython-310.pyc,, +matplotlib/__pycache__/textpath.cpython-310.pyc,, +matplotlib/__pycache__/ticker.cpython-310.pyc,, +matplotlib/__pycache__/transforms.cpython-310.pyc,, +matplotlib/__pycache__/typing.cpython-310.pyc,, +matplotlib/__pycache__/units.cpython-310.pyc,, +matplotlib/__pycache__/widgets.cpython-310.pyc,, +matplotlib/_afm.py,sha256=cWe1Ib37T6ZyHbR6_hPuzAjotMmi32y-kDB-i28iyqE,16692 +matplotlib/_animation_data.py,sha256=JJJbbc-fMdPjkbQ7ng9BHL5i91VTDHQVTtEdWOvWBAI,7986 +matplotlib/_api/__init__.py,sha256=7Fs57FtnCmbyr_NwuGWe0EhfXOYrBw22iKMU9onGc_k,13799 +matplotlib/_api/__init__.pyi,sha256=XNL-oGkk1MZPtSXk3rHlC3jLWZxsElmpNpOongKq8qA,2246 +matplotlib/_api/__pycache__/__init__.cpython-310.pyc,, +matplotlib/_api/__pycache__/deprecation.cpython-310.pyc,, +matplotlib/_api/deprecation.py,sha256=XJxksSV8ukS8kSVmsGPS0zgv1ffFa8WNaRLAk31dmOM,20091 +matplotlib/_api/deprecation.pyi,sha256=A8De57amX2GlZSrDYwVYCMxHFb8AsYGrH0k_bCxuJus,2217 +matplotlib/_blocking_input.py,sha256=VHNsxvX2mTx_xBknd30MSicVlRXS4dCDe9hDctbV5rk,1224 +matplotlib/_c_internal_utils.cpython-310-x86_64-linux-gnu.so,sha256=qK4dPJ3AidxVzcIq03fohY1N8lcXOWGWMdPdWVgpX68,267216 +matplotlib/_c_internal_utils.pyi,sha256=Z3bLs9pMGXrmZjt-4_A-x4321bLP-B54xDbr4PIgUfc,377 +matplotlib/_cm.py,sha256=PuYIAkUpz4u4aiUjvdV5njIfG0J_MQY9pM2Yz3j3KXs,68014 +matplotlib/_cm_bivar.py,sha256=gpmKiSxsWoVGWIifseh0goND7Y7zTWJKsznr_QkEtDg,97461 +matplotlib/_cm_listed.py,sha256=3a02mPUSnOsXkqRFNxKGwIvJCAKY5lezBUrqGcllnvk,135004 +matplotlib/_cm_multivar.py,sha256=0UjNFW7Sytj1t31QD99ebfr-al7NmCpWhKUOyq9VzK8,6630 +matplotlib/_color_data.py,sha256=k-wdTi6ArJxksqBfMT-7Uy2qWz8XX4Th5gsjf32CwmM,34780 +matplotlib/_color_data.pyi,sha256=RdBRk01yuf3jYVlCwG351tIBCxehizkZMnKs9c8gnOw,170 +matplotlib/_constrained_layout.py,sha256=XX_2elqHukF3toeEDupwQPlN_8UEv3-yGKlR0A92n74,31485 +matplotlib/_docstring.py,sha256=u9yJorJidI8k1W8S01SIMqNmh5VOMVc-AqyES66EEdk,4435 +matplotlib/_docstring.pyi,sha256=6ze5DoqZaFy6BQ5Z8vtDiJwSOXhBcpZJlkM5enI9cuE,800 +matplotlib/_enums.py,sha256=euD2sj2FIbQMzIPA4rCrII_y8RVzaMyEBKsSdBQaUb4,6175 +matplotlib/_enums.pyi,sha256=K7j_kDwGOnx37CYnXhwfJ9NJBGeet45KvAj0om05RUs,326 +matplotlib/_fontconfig_pattern.py,sha256=hCpupToheqJJpISVEN3FNv6WZZPVaoKOuqGMmD3Wxh4,4365 +matplotlib/_image.cpython-310-x86_64-linux-gnu.so,sha256=0rOaO_cd-qmhzmPGUUAsVQ5jYsdTuEas-14Gzz-OCtc,586232 +matplotlib/_image.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/_internal_utils.py,sha256=nhK6LLWYW93fBcsFiO09JmqFj2rgHEsGYFOeaC7HRKw,2140 +matplotlib/_layoutgrid.py,sha256=mFB9asZVol2aV9hLBpdntuG5x1isKUrajF-3TUhCgso,21676 +matplotlib/_mathtext.py,sha256=KnNxS6r_YRzato7XwQ17SMCmVxHPI_91hSAubJJRKow,107636 +matplotlib/_mathtext_data.py,sha256=9y__7jf3bzgOmD2lEYW3QpgsT-z9QDuiwRLVH7Wq8Pw,65067 +matplotlib/_path.cpython-310-x86_64-linux-gnu.so,sha256=YrBD-cBXUrieGFfFpJ9VvVF16XPxOoznxsd6nSfKic0,506808 +matplotlib/_path.pyi,sha256=yznyfzoUogH9vvi0vK68ga4Shlbrn5UBhAnLX8Ght1o,325 +matplotlib/_pylab_helpers.py,sha256=pJERytHDmXo2VP3sH9Qw6NwJP0hefNPQMUXCVntnoHI,4307 +matplotlib/_pylab_helpers.pyi,sha256=7OZKr-OL3ipVt1EDZ6e-tRwHASz-ijYfcIdlPczXhvQ,1012 +matplotlib/_qhull.cpython-310-x86_64-linux-gnu.so,sha256=Ezc0F3Flzd7JzWsmYJxK7HsLHam6HAkQxQRqdz_ER6I,764784 +matplotlib/_qhull.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/_text_helpers.py,sha256=nz1pvEp8756PCrn6KIY2M6vVA-beYUNODvvf4NnGF0g,2538 +matplotlib/_tight_bbox.py,sha256=ddJ5ViXulbPLocHr-RkWK27WJSuV2WXUx74jZyL0NOg,2787 +matplotlib/_tight_layout.py,sha256=A3vZKBmxci7w35OO9lxgfUqrRKgTr-N_dUza0M05WxE,12675 +matplotlib/_tri.cpython-310-x86_64-linux-gnu.so,sha256=-ZecvM0Oogw83NeR1B6Ud6XwExtSGHxs0JVXp_OV95M,441352 +matplotlib/_tri.pyi,sha256=bFnu97pIlSboJT3ijQBw4YlZdlR1UzXYtI3DkEPVzeM,1429 +matplotlib/_type1font.py,sha256=c_PXvyOFQ4YCRQJ23snWNkdEhuZ4k0JbLe55zaeoXfQ,28409 +matplotlib/_version.py,sha256=qbaaDQaXfJSjlv6Dopir4jYnWHPR6yfkLYEwTRXzgQs,19 +matplotlib/animation.py,sha256=58vUSy_0M2Ve2GO5JcXsCHxCwr7TqaaMed-KNxN6Nno,73114 +matplotlib/animation.pyi,sha256=C8Fgmswh4um-qYgqPZNBEpXVaefRXi6e026pYjonLHg,6566 +matplotlib/artist.py,sha256=hfCGETRkHLfP-i4PmmNsc7syXyFTZi4HYvzyE8q9ak0,63293 +matplotlib/artist.pyi,sha256=9sR0w8Kvi1GaFV6IbUzOw9LvQkPEB8FP_uRvUcI6_hg,7336 +matplotlib/axes/__init__.py,sha256=aHE_zIjphIJkW4_1fyoFuEWTb7gCRewK5jNZWJHEdgM,351 +matplotlib/axes/__init__.pyi,sha256=HP1z2v-PboHQS4dQjvJ7XjUjX-zw6taZRTTB9oVKwYE,303 +matplotlib/axes/__pycache__/__init__.cpython-310.pyc,, +matplotlib/axes/__pycache__/_axes.cpython-310.pyc,, +matplotlib/axes/__pycache__/_base.cpython-310.pyc,, +matplotlib/axes/__pycache__/_secondary_axes.cpython-310.pyc,, +matplotlib/axes/_axes.py,sha256=CwBUFh4fIlJF4Km80apFeH9sPVBi-3OPTOnxqKveZ5Q,353841 +matplotlib/axes/_axes.pyi,sha256=fvlVj8AnKt52EOCgxJ5N_6_lb9n0gmXlOP7WGf2puNI,26116 +matplotlib/axes/_base.py,sha256=eXarXIkk1SiQRPa7Rdu_R6gDgHdm2l60Hh2RPtPTf7E,186445 +matplotlib/axes/_base.pyi,sha256=EX6xYlHGEMRWXoE85w-UfCkN7m990Gg5VFZJ44D0XmE,17058 +matplotlib/axes/_secondary_axes.py,sha256=Sew-LK57mDndUXb-9_VHzs5dJLQ6fh4nstZtfpZjui0,11887 +matplotlib/axes/_secondary_axes.pyi,sha256=GtU55YzLNN7XHMRQcAmcAxcdcN-FzHw1RAJpOiTcZFk,1414 +matplotlib/axis.py,sha256=bsPEyn4eMAnkTCbW4rxSm8ysPNYH2p_ur7RN_ZfZRps,104714 +matplotlib/axis.pyi,sha256=UO1-oCeHyB7Mrxxnuo0Qys2u_k8e6qTmz4B6-rHrWdw,10181 +matplotlib/backend_bases.py,sha256=5EaIYhLC7LdKrk0SYSNBiYW3G16KbbexndFhoOX7_jY,131953 +matplotlib/backend_bases.pyi,sha256=G6MqvmWf-y2SwhuhNF2qN8hmovxAkV1TJMbg4LzU9cM,16270 +matplotlib/backend_managers.py,sha256=RQheCO_cQBlaWsYMbAmswu0UPKU7bmLTI5LEFgotklA,11795 +matplotlib/backend_managers.pyi,sha256=agnuM0wiZRqSgqti2AgbKJijRLvEPNLOrSY8PEwLjFE,2253 +matplotlib/backend_tools.py,sha256=eCx84l8JBYayju5Ry1LUFyKLIrdgxmGUTJw72VNP1yo,33186 +matplotlib/backend_tools.pyi,sha256=lc3W6FcY0XhTb8Z8XAqQ6MBk_SV7WH41t7zjwAlQfmE,4122 +matplotlib/backends/__init__.py,sha256=JowJe-tDrUBMNJTiJATgiEuACpgdxsKnRCYa-nC255A,206 +matplotlib/backends/__pycache__/__init__.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_gtk.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_pdf_ps.cpython-310.pyc,, +matplotlib/backends/__pycache__/_backend_tk.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk3cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_gtk4cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_macosx.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_mixed.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_nbagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_pdf.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_pgf.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_ps.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5agg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qt5cairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qtagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_qtcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_svg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_template.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_tkagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_tkcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_webagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_webagg_core.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wx.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wxagg.cpython-310.pyc,, +matplotlib/backends/__pycache__/backend_wxcairo.cpython-310.pyc,, +matplotlib/backends/__pycache__/qt_compat.cpython-310.pyc,, +matplotlib/backends/__pycache__/registry.cpython-310.pyc,, +matplotlib/backends/_backend_agg.cpython-310-x86_64-linux-gnu.so,sha256=jael5ah2SIAY6xwIDADqGx6c-vCdLNf_hJ8u8WIaFNM,771424 +matplotlib/backends/_backend_agg.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/_backend_gtk.py,sha256=zNCmMOgjgZGfAIlxDU4ph02XMSSn0wIUoPQ4wsaeeAg,11274 +matplotlib/backends/_backend_pdf_ps.py,sha256=M0kjQpMlBoAim6Pdw8u3FqUIoXz4UwRv4mSzXPJPI44,5968 +matplotlib/backends/_backend_tk.py,sha256=zx9GZYgrcmTwDgUT3rPvU7YgW79sUbJIWQ-SJtEsH5s,44581 +matplotlib/backends/_macosx.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/_tkagg.cpython-310-x86_64-linux-gnu.so,sha256=9Eksl57ceMRw2gdtCzhSyyhwl95NstKEuwwCabeSQI8,298400 +matplotlib/backends/_tkagg.pyi,sha256=1yKnYKSgNFesrb0ZI5RQCyTO4zONgHW5MrJ9vq4HzyI,379 +matplotlib/backends/backend_agg.py,sha256=Py_rXRbDTP0ZSrNB2O5Tmh5V8G5V9ViXH_YV5PTX9xw,19888 +matplotlib/backends/backend_cairo.py,sha256=Ah5M85Ppa8zYGOtLOizW0gzltjs3b4nD56qFeQR795U,18618 +matplotlib/backends/backend_gtk3.py,sha256=ZHf4CKhDGJvl8J3p6Sx_ZJBTgVetZBbaKWfnXKvNe0c,22277 +matplotlib/backends/backend_gtk3agg.py,sha256=aFJw05L5V-FyFSj_20l_FuC5piPf8fhKzft_m1NMhYc,2463 +matplotlib/backends/backend_gtk3cairo.py,sha256=zPsJzVm750if2LiQ9ybzsbX0rBhkP05XVvR9Lhz65so,1392 +matplotlib/backends/backend_gtk4.py,sha256=lkcW3cEKdASQ4vD6VHBUXvZQZEGvj5uenCVg4fHof_g,23541 +matplotlib/backends/backend_gtk4agg.py,sha256=00i3qpIt9Tcf_S74GOWbeckiPlfVJoQ2pBbhXDMthF0,1262 +matplotlib/backends/backend_gtk4cairo.py,sha256=sqWm3WgfNO8EsBjqzsD4U4cAc0f4q5SufYx7ZJacDf0,1125 +matplotlib/backends/backend_macosx.py,sha256=Gjbu1MKXIZpCM_luMG6U7qDgwotcGSegCvYw-NNWxeQ,7397 +matplotlib/backends/backend_mixed.py,sha256=Gf_2BDjy94FZRXALm0X3xRpK_1l8bZGSYCqnrSjrjME,4696 +matplotlib/backends/backend_nbagg.py,sha256=Au9RHfRufpI0ngT4R0K0CUVtAMFi9Bg-YhDunlj_Lko,8000 +matplotlib/backends/backend_pdf.py,sha256=DfFKdD4-wc7PPpfjXh8MAaaDppXUvLwicaMhXKgRhGw,105320 +matplotlib/backends/backend_pgf.py,sha256=GT7UHkAtHtnZcyV0GxHtUL9TTre9dji7Q6qexyi4AUU,39567 +matplotlib/backends/backend_ps.py,sha256=JRZZ6hsPcCfZmb-HczqyBulgow15yvoaHQe3Zcl1Zd4,51991 +matplotlib/backends/backend_qt.py,sha256=FjwUd9E2Fu2gCUhAVcJjXC--sAJsAoLUEHpMpa-IueE,42312 +matplotlib/backends/backend_qt5.py,sha256=kzfoo2ksEGsiWAa2LGtZYzKvfzqJJWyGOohohcRAu1g,787 +matplotlib/backends/backend_qt5agg.py,sha256=Vh7H8kqWH4X8a3VX2XZ2Vze9srwJavkNHAZxdJUz_bk,352 +matplotlib/backends/backend_qt5cairo.py,sha256=Go2Y0GVkXh1xh6x4F255_e5Xbwwws-OiD1Fc0805E78,292 +matplotlib/backends/backend_qtagg.py,sha256=ZjPtp5wR6tZGjbngPXRdVXYRhiPPrc5C0q2DmtdRkpY,3413 +matplotlib/backends/backend_qtcairo.py,sha256=e3SUG50VGqo68eS_8ebTCVQPa4AaxLxuo1JiWX4TIWg,1770 +matplotlib/backends/backend_svg.py,sha256=RQvPE9pR0xSu55NIMajX1WBqmmXbUWWwdWr2qMG-63k,51000 +matplotlib/backends/backend_template.py,sha256=qhWvXGiPeO2jvH61L3NCRYLReo--d4WJrQYanIaEmiE,8012 +matplotlib/backends/backend_tkagg.py,sha256=z9gB16fq2d-DUNpbeSDDLWaYmc0Jz3cDqNlBKhnQg0c,592 +matplotlib/backends/backend_tkcairo.py,sha256=JaGGXh8Y5FwVZtgryIucN941Olf_Pn6f4Re7Vuxl1-c,845 +matplotlib/backends/backend_webagg.py,sha256=jQhB9JA5tXcRwm6_7qlBg2-Q8_y-5hnv2IUNBPpCSbA,11014 +matplotlib/backends/backend_webagg_core.py,sha256=RZisxLS_pocuBRr5_lqBnF91flQA5upI8PiP8-Zitl0,18748 +matplotlib/backends/backend_wx.py,sha256=A4ztnMlVo-p1z77uZvJ-sJiyE21_QrujaTxo5P1glOY,51310 +matplotlib/backends/backend_wxagg.py,sha256=tzcwYyW34j4LPfHm9uhuHwepwZIcspi3y8oPC8FJkdk,1468 +matplotlib/backends/backend_wxcairo.py,sha256=TK-m3S0c1WipfKE2IpIPNeE4hoXPjfMvnWAzHpCXpFs,848 +matplotlib/backends/qt_compat.py,sha256=t_4aD4rg5maItXABLtVG6JIfuQwZnHwhMcjtz6AUwQs,5346 +matplotlib/backends/qt_editor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/backends/qt_editor/__pycache__/__init__.cpython-310.pyc,, +matplotlib/backends/qt_editor/__pycache__/_formlayout.cpython-310.pyc,, +matplotlib/backends/qt_editor/__pycache__/figureoptions.cpython-310.pyc,, +matplotlib/backends/qt_editor/_formlayout.py,sha256=QmqqqLO6waqeSGOKjDNUwjvon53Z7yqil5AfqfftDWY,20953 +matplotlib/backends/qt_editor/figureoptions.py,sha256=ljd7fCBiE0kk2ATNTaB5Dg8U6Gn7Lz4Vhi3XX-VbXeE,9851 +matplotlib/backends/registry.py,sha256=M2BHhBKbUJ8E2B2QExfVf8NbAK4ZFhJtmKRelBZZDh4,15480 +matplotlib/backends/web_backend/all_figures.html,sha256=44Y-GvIJbNlqQaKSW3kwVKpTxBSG1WsdYz3ZmYHlUsA,1753 +matplotlib/backends/web_backend/css/boilerplate.css,sha256=qui16QXRnQFNJDbcMasfH6KtN9hLjv8883U9cJmsVCE,2310 +matplotlib/backends/web_backend/css/fbm.css,sha256=wa4vNkNv7fQ_TufjJjecFZEzPMR6W8x6uXJga_wQILw,1456 +matplotlib/backends/web_backend/css/mpl.css,sha256=ruca_aA5kNnP-MZmLkriu8teVP1nIgwcFEpoB16j8Z4,1611 +matplotlib/backends/web_backend/css/page.css,sha256=ca3nO3TaPw7865PN5SlGJBTc2H3rBXQCMaFPywX29y4,1623 +matplotlib/backends/web_backend/ipython_inline_figure.html,sha256=wgSxUh3xpPAxOnZgSMnrhDM5hYncOfWRGgaCUezvedY,1311 +matplotlib/backends/web_backend/js/mpl.js,sha256=FAFf8huEmvymJbTHFUQ5pbQt48NtXmN_DNACP2dCDMw,24432 +matplotlib/backends/web_backend/js/mpl_tornado.js,sha256=Zs2Uzs7YUilG765nYvanCo-IK8HkHDtIum1KAq6bQ_w,302 +matplotlib/backends/web_backend/js/nbagg_mpl.js,sha256=F-By4ZjOSmwpNAkxUUxUk35qCjGlf0B28Z0aOyfpxDM,9514 +matplotlib/backends/web_backend/single_figure.html,sha256=wEBwF602JLHErBUEXiS6jXqmxYAIzHpa3MMFrnev6rs,1357 +matplotlib/bezier.py,sha256=ahJ1BRJoTbcxJoNmF1mCYc96mWs5jSM4DOK58jYaXkU,19049 +matplotlib/bezier.pyi,sha256=itubQt_U1IKnQ508bhy6J8LSpaspvZIeOmJqg7JvOmU,2586 +matplotlib/category.py,sha256=UsZ8rbADAH96YfpjD-N1DLTcS_cQwlszIIQ2ZRaiWm4,7377 +matplotlib/cbook.py,sha256=nF2hNC_ryiJcec58Cnp_0mAGSD0LfsnjbzUuCH4p-Os,80020 +matplotlib/cbook.pyi,sha256=YkPZz_bsTYQJQ_wKx-qbsJYLhy2t2nCyvLNjJyiRtpw,6037 +matplotlib/cm.py,sha256=5xav-BvJ4pWSVdvGql4oYc4EpA-R6OuUP5xjrZCMKvg,10350 +matplotlib/cm.pyi,sha256=Ot5j2zDHYTra4kKuTwxYp7X5TsNwg_cJuYa7BxmQz-A,939 +matplotlib/collections.py,sha256=5AcX3lTsKD1ObSOkLr-irzYPfdJxsNrgw5dAD886vCE,96311 +matplotlib/collections.pyi,sha256=uZ4fPdq34oghdEtS1UkY7P0oOE2oPUzvS0Cr09yktaQ,10775 +matplotlib/colorbar.py,sha256=-PwV3xkczcyPq9Ki_zfpQdP_1EqOFBzoeUSsMlUF2cE,60687 +matplotlib/colorbar.pyi,sha256=C1NzjMWFKm01rPUUr15Yjz1WB0pGakJQeatJO7G0mNs,4966 +matplotlib/colorizer.py,sha256=RzE0zjqB6elpketKidy6n-RRDbYF98speEIOIMfIFb4,25180 +matplotlib/colorizer.pyi,sha256=B_5-rEvNo4DblSifLodf3OVVEc3bTGsIYqqaQdU_4Tw,3308 +matplotlib/colors.py,sha256=h1YV59y4WvstOPC0x8eZzwHYItal7VjjVcyVS3niaSM,137303 +matplotlib/colors.pyi,sha256=XhCi3MtgSozOY5FATtPUHyKO4xMrvi7TvbCAl714Uxc,14908 +matplotlib/container.py,sha256=Y6v4j79gMk8QDYfrdOqbbJHH0BoOIO9enz4dtlaBSJU,4565 +matplotlib/container.pyi,sha256=DdthHVj1-bQQV3VcpMD1uIPLXVUKyKhWMXy8yCniK1I,1805 +matplotlib/contour.py,sha256=wpMtrMZEiDYPRuFbZ592HvQqn4NOBySfM17detKoenY,68391 +matplotlib/contour.pyi,sha256=X4DSH0u1YFUvKxuoAv8YbmVOJTJ-wbHig4lznzhsI0U,5300 +matplotlib/dates.py,sha256=xfcFY3g3AsW9isnEIKFGj4gXq0Xeoxox_3fxNCDQ01c,66306 +matplotlib/dviread.py,sha256=BqQg0fqXD_TmABJPSWDC-H2azt7B53NfOE1SC6mbXiI,42590 +matplotlib/dviread.pyi,sha256=5fps3GiPf1ibKY1TH6__kNEiJ9xhpVgzgait6jo09Hc,2139 +matplotlib/figure.py,sha256=CNPMl0QCxhr3h3WGcd271pE4Xtv4rapRWuUSFSNJSuc,141862 +matplotlib/figure.pyi,sha256=rROP8wvPmeSwTKoJwFaWsCHFpZjfgDk_7HE0wPfZQRg,14914 +matplotlib/font_manager.py,sha256=tSpOlZqiNUdIn3gQCFN_-nOLZ08WXBqj_unP6o1JyIw,57651 +matplotlib/font_manager.pyi,sha256=Jse1IshOdjhXHdjcvTxtos3rgsuECbRJzc264A-ZGus,5052 +matplotlib/ft2font.cpython-310-x86_64-linux-gnu.so,sha256=6HlilDHYRzRplaxd8etug5dy8XEuh44OWfLENJzeNRE,1438856 +matplotlib/ft2font.pyi,sha256=yuDEfdw-7VyFuAHts-sMRm4VKV7IxYresBxEQgNOaPM,9253 +matplotlib/gridspec.py,sha256=s6RkXEk4ORWsMHjXkK98ZvosZ5D-qgenVNk0hIvMwOw,29786 +matplotlib/gridspec.pyi,sha256=3IQi_Q5KdwyPT_J_FWcngJ7iUOOCH4PSm42USm16Ofg,5099 +matplotlib/hatch.py,sha256=tfEM0DQxyLkUE87lHduMNlE55GKINtAZo4Y_YsAmSJo,7453 +matplotlib/hatch.pyi,sha256=OEFkNe8TgrBli_MDMW99Tvhr44S4ySPaYQomjVXpUUA,2098 +matplotlib/image.py,sha256=MLFg06J2xllm1tmgDnlLYw8gTy7__NGnkf4SZP-3mSs,69765 +matplotlib/image.pyi,sha256=R88aGGu6wP0lFKRudaM7iNxGgq4kh9uX9QRaZSorTLQ,7066 +matplotlib/inset.py,sha256=fJV8VMMY07_97iJR3_RO892WAKE7myfV_Gqg-qK99vg,10154 +matplotlib/inset.pyi,sha256=c-IVW14bpqb3qu7bz0EabvNXJnZwJAXEXoQoIcrL8ns,968 +matplotlib/layout_engine.py,sha256=mwRNh13mIQ-wpjQVgBaB_LWKN1AiQTWldg-X9nLmVok,11433 +matplotlib/layout_engine.pyi,sha256=9cKFTFPCvwmKIygTUvrvOq1iWnFYme9Shniiv2QbC74,1788 +matplotlib/legend.py,sha256=5NDPbWqkiIkP1FJ3sLNqmew2jkNjawQf2G_BqZNG_u4,55343 +matplotlib/legend.pyi,sha256=hn-MNF3SPHtSUqIhwVXTebU_Nzk_wIh5iKgf7AEAZRg,5364 +matplotlib/legend_handler.py,sha256=ekhZTT1X9yQd7TNIRdu_vui5FR2CTmI0NWTFJTqF1N4,29872 +matplotlib/legend_handler.pyi,sha256=3VEfeioGIAxhd3mhg4PXETZjCKf4OlXL0jz1MAFGtos,7655 +matplotlib/lines.py,sha256=BG3xhTuzYDFjzf-k5CcQc_V7W4HlQl4CL7RCQQz0EC8,57920 +matplotlib/lines.pyi,sha256=3tG7tD8GZ8YPphw4nmkgS9SfdLV8icW36bv61jb9RlU,6081 +matplotlib/markers.py,sha256=g6ukerZ5n_sD_Enk0eJA4kkTvOOp8AfSGBq60PHot-E,33708 +matplotlib/markers.pyi,sha256=FFFBsvilnbd8-5L04U70kfVoyBwc18w_fZ6DysTj9p0,1678 +matplotlib/mathtext.py,sha256=yTlpUfnKxzrFl7XNjyuC-xMBBnxUufU6GVwQyhQAI2o,5104 +matplotlib/mathtext.pyi,sha256=RCVxYGQ_CJ6wC7v_HkqoouU2PhtcvlJ1ffIyxzAC-so,1045 +matplotlib/mlab.py,sha256=T71p4x0KJNVGrwUuoDdgcjWQucIQTuVrocOC98Zdwsc,30210 +matplotlib/mlab.pyi,sha256=mkR7wbJS9eCQfCFsUWoXnqrAy0kcE3cVqpFZCeOGC_Q,3583 +matplotlib/mpl-data/fonts/afm/cmex10.afm,sha256=blR3ERmrVBV5XKkAnDCj4NMeYVgzH7cXtJ3u59u9GuE,12070 +matplotlib/mpl-data/fonts/afm/cmmi10.afm,sha256=5qwEOpedEo76bDUahyuuF1q0cD84tRrX-VQ4p3MlfBo,10416 +matplotlib/mpl-data/fonts/afm/cmr10.afm,sha256=WDvgC_D3UkGJg9u-J0U6RaT02lF4oz3lQxHtg1r3lYw,10101 +matplotlib/mpl-data/fonts/afm/cmsy10.afm,sha256=AbmzvCVWBceHRfmRfeJ9E6xzOQTFLk0U1zDfpf3_MaM,8295 +matplotlib/mpl-data/fonts/afm/cmtt10.afm,sha256=4ji7_mTpeWMa93o_UHBWPKCnqsBfhJJNllat1lJArP4,6501 +matplotlib/mpl-data/fonts/afm/pagd8a.afm,sha256=jjFrigwkTpYLqa26cpzZvKQNBo-PuF4bmDVqaM4pMWw,17183 +matplotlib/mpl-data/fonts/afm/pagdo8a.afm,sha256=sgNQdeYyx8J-itGw9h31y95aMBiTCRvmNSPTXwwS7xg,17255 +matplotlib/mpl-data/fonts/afm/pagk8a.afm,sha256=ZUtfHPloNqcvGMHMxaKDSlshhOcjwheUx143RwpGdIU,17241 +matplotlib/mpl-data/fonts/afm/pagko8a.afm,sha256=Yj1wBg6Jsqqz1KBfhRoJ3ACR-CMQol8Fj_ZM5NZ1gDk,17346 +matplotlib/mpl-data/fonts/afm/pbkd8a.afm,sha256=Zl5o6J_di9Y5j2EpHtjew-_sfg7-WoeVmO9PzOYSTUc,15157 +matplotlib/mpl-data/fonts/afm/pbkdi8a.afm,sha256=JAOno930iTyfZILMf11vWtiaTgrJcPpP6FRTRhEMMD4,15278 +matplotlib/mpl-data/fonts/afm/pbkl8a.afm,sha256=UJqJjOJ6xQDgDBLX157mKpohIJFVmHM-N6x2-DiGv14,15000 +matplotlib/mpl-data/fonts/afm/pbkli8a.afm,sha256=AWislZ2hDbs0ox_qOWREugsbS8_8lpL48LPMR40qpi0,15181 +matplotlib/mpl-data/fonts/afm/pcrb8a.afm,sha256=6j1TS2Uc7DWSc-8l42TGDc1u0Fg8JspeWfxFayjUwi8,15352 +matplotlib/mpl-data/fonts/afm/pcrbo8a.afm,sha256=smg3mjl9QaBDtQIt06ko5GvaxLsO9QtTvYANuE5hfG0,15422 +matplotlib/mpl-data/fonts/afm/pcrr8a.afm,sha256=7nxFr0Ehz4E5KG_zSE5SZOhxRH8MyfnCbw-7x5wu7tw,15339 +matplotlib/mpl-data/fonts/afm/pcrro8a.afm,sha256=NKEz7XtdFkh9cA8MvY-S3UOZlV2Y_J3tMEWFFxj7QSg,15443 +matplotlib/mpl-data/fonts/afm/phvb8a.afm,sha256=NAx4M4HjL7vANCJbc-tk04Vkol-T0oaXeQ3T2h-XUvM,17155 +matplotlib/mpl-data/fonts/afm/phvb8an.afm,sha256=8e_myD-AQkNF7q9XNLb2m76_lX2TUr3a5wog_LIE1sk,17086 +matplotlib/mpl-data/fonts/afm/phvbo8a.afm,sha256=8fkBRmJ-SWY2YrBg8fFyjJyrJp8daQ6JPO6LvhM8xPI,17230 +matplotlib/mpl-data/fonts/afm/phvbo8an.afm,sha256=aeVRvV4r15BBvxuRJ0MG8ZHuH2HViuIiCYkvuapmkmM,17195 +matplotlib/mpl-data/fonts/afm/phvl8a.afm,sha256=IyMYM-bgl-gI6rG0EuZZ2OLzlxJfGeSh8xqsh0t-eJQ,15627 +matplotlib/mpl-data/fonts/afm/phvlo8a.afm,sha256=s12C-eNnIDHJ_UVbuiprjxBjCiHIbS3Y8ORTC-qTpuI,15729 +matplotlib/mpl-data/fonts/afm/phvr8a.afm,sha256=Kt8KaRidts89EBIK29X2JomDUEDxvroeaJz_RNTi6r4,17839 +matplotlib/mpl-data/fonts/afm/phvr8an.afm,sha256=lL5fAHTRwODl-sB5mH7IfsD1tnnea4yRUK-_Ca2bQHM,17781 +matplotlib/mpl-data/fonts/afm/phvro8a.afm,sha256=3KqK3eejiR4hIFBUynuSX_4lMdE2V2T58xOF8lX-fwc,17919 +matplotlib/mpl-data/fonts/afm/phvro8an.afm,sha256=Vx9rRf3YfasMY7tz-njSxz67xHKk-fNkN7yBi0X2IP0,17877 +matplotlib/mpl-data/fonts/afm/pncb8a.afm,sha256=aoXepTcDQtQa_mspflMJkEFKefzXHoyjz6ioJVI0YNc,16028 +matplotlib/mpl-data/fonts/afm/pncbi8a.afm,sha256=pCWW1MYgy0EmvwaYsaYJaAI_LfrsKmDANHu7Pk0RaiU,17496 +matplotlib/mpl-data/fonts/afm/pncr8a.afm,sha256=0CIB2BLe9r-6_Wl5ObRTTf98UOrezmGQ8ZOuBX5kLks,16665 +matplotlib/mpl-data/fonts/afm/pncri8a.afm,sha256=5R-pLZOnaHNG8pjV6MP3Ai-d2OTQYR_cYCb5zQhzfSU,16920 +matplotlib/mpl-data/fonts/afm/pplb8a.afm,sha256=3EzUbNnXr5Ft5eFLY00W9oWu59rHORgDXUuJaOoKN58,15662 +matplotlib/mpl-data/fonts/afm/pplbi8a.afm,sha256=X_9tVspvrcMer3OS8qvdwjFFqpAXYZneyCL2NHA902g,15810 +matplotlib/mpl-data/fonts/afm/pplr8a.afm,sha256=ijMb497FDJ9nVdVMb21F7W3-cu9sb_9nF0oriFpSn8k,15752 +matplotlib/mpl-data/fonts/afm/pplri8a.afm,sha256=8KITbarcUUMi_hdoRLLmNHtlqs0TtOSKqtPFft7X5nY,15733 +matplotlib/mpl-data/fonts/afm/psyr.afm,sha256=Iyt8ajE4B2Tm34oBj2pKtctIf9kPfq05suQefq8p3Ro,9644 +matplotlib/mpl-data/fonts/afm/ptmb8a.afm,sha256=bL1fA1NC4_nW14Zrnxz4nHlXJb4dzELJPvodqKnYeMg,17983 +matplotlib/mpl-data/fonts/afm/ptmbi8a.afm,sha256=-_Ui6XlKaFTHEnkoS_-1GtIr5VtGa3gFQ2ezLOYHs08,18070 +matplotlib/mpl-data/fonts/afm/ptmr8a.afm,sha256=IEcsWcmzJyjCwkgsw4o6hIMmzlyXUglJat9s1PZNnEU,17942 +matplotlib/mpl-data/fonts/afm/ptmri8a.afm,sha256=49fQMg5fIGguZ7rgc_2styMK55Pv5bPTs7wCzqpcGpk,18068 +matplotlib/mpl-data/fonts/afm/putb8a.afm,sha256=qMaHTdpkrNL-m4DWhjpxJCSmgYkCv1qIzLlFfM0rl40,21532 +matplotlib/mpl-data/fonts/afm/putbi8a.afm,sha256=g7AVJyiTxeMpNk_1cSfmYgM09uNUfPlZyWGv3D1vcAk,21931 +matplotlib/mpl-data/fonts/afm/putr8a.afm,sha256=XYmNC5GQgSVAZKTIYdYeNksE6znNm9GF_0SmQlriqx0,22148 +matplotlib/mpl-data/fonts/afm/putri8a.afm,sha256=i7fVe-iLyLtQxCfAa4IxdxH-ufcHmMk7hbCGG5TxAY4,21891 +matplotlib/mpl-data/fonts/afm/pzcmi8a.afm,sha256=wyuoIWEZOcoXrSl1tPzLkEahik7kGi91JJj-tkFRG4A,16250 +matplotlib/mpl-data/fonts/afm/pzdr.afm,sha256=MyjLAnzKYRdQBfof1W3k_hf30MvqOkqL__G22mQ5xww,9467 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-Bold.afm,sha256=sIDDI-B82VZ3C0mI_mHFITCZ7PVn37AIYMv1CrHX4sE,15333 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-BoldOblique.afm,sha256=zg61QobD3YU9UBfCXmvmhBNaFKno-xj8sY0b2RpgfLw,15399 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier-Oblique.afm,sha256=vRQm5j1sTUN4hicT1PcVZ9P9DTTUHhEzfPXqUUzVZhE,15441 +matplotlib/mpl-data/fonts/pdfcorefonts/Courier.afm,sha256=Mdcq2teZEBJrIqVXnsnhee7oZnTs6-P8_292kWGTrw4,15335 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-Bold.afm,sha256=i2l4gcjuYXoXf28uK7yIVwuf0rnw6J7PwPVQeHj5iPw,69269 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-BoldOblique.afm,sha256=Um5O6qK11DXLt8uj_0IoWkc84TKqHK3bObSKUswQqvY,69365 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica-Oblique.afm,sha256=hVYDg2b52kqtbVeCzmiv25bW1yYdpkZS-LXlGREN2Rs,74392 +matplotlib/mpl-data/fonts/pdfcorefonts/Helvetica.afm,sha256=23cvKDD7bQAJB3kdjSahJSTZaUOppznlIO6FXGslyW8,74292 +matplotlib/mpl-data/fonts/pdfcorefonts/Symbol.afm,sha256=P5UaoXr4y0qh4SiMa5uqijDT6ZDr2-jPmj1ayry593E,9740 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Bold.afm,sha256=cQTmr2LFPwKQE_sGQageMcmFicjye16mKJslsJLHQyE,64251 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-BoldItalic.afm,sha256=pzWOdycm6RqocBWgAVY5Jq0z3Fp7LuqWgLNMx4q6OFw,59642 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Italic.afm,sha256=bK5puSMpGT_YUILwyJrXoxjfj7XJOdfv5TQ_iKsJRzw,66328 +matplotlib/mpl-data/fonts/pdfcorefonts/Times-Roman.afm,sha256=hhNrUnpazuDDKD1WpraPxqPWCYLrO7D7bMVOg-zI13o,60460 +matplotlib/mpl-data/fonts/pdfcorefonts/ZapfDingbats.afm,sha256=ZuOmt9GcKofjdOq8kqhPhtAIhOwkL2rTJTmZxAjFakA,9527 +matplotlib/mpl-data/fonts/pdfcorefonts/readme.txt,sha256=MRv8ppSITYYAb7lt5EOw9DWWNZIblfxsFhu5TQE7cpI,828 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-Bold.ttf,sha256=sYS4njwQdfIva3FXW2_CDUlys8_TsjMiym_Vltyu8Wc,704128 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-BoldOblique.ttf,sha256=bt8CgxYBhq9FHL7nHnuEXy5Mq_Jku5ks5mjIPCVGXm8,641720 +matplotlib/mpl-data/fonts/ttf/DejaVuSans-Oblique.ttf,sha256=zN90s1DxH9PdV3TeUOXmNGoaXaH1t9X7g1kGZel6UhM,633840 +matplotlib/mpl-data/fonts/ttf/DejaVuSans.ttf,sha256=P99pyr8GBJ6nCgC1kZNA4s4ebQKwzDxLRPtoAb0eDSI,756072 +matplotlib/mpl-data/fonts/ttf/DejaVuSansDisplay.ttf,sha256=ggmdz7paqGjN_CdFGYlSX-MpL3N_s8ngMozpzvWWUvY,25712 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Bold.ttf,sha256=uq2ppRcv4giGJRr_BDP8OEYZEtXa8HKH577lZiCo2pY,331536 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-BoldOblique.ttf,sha256=ppCBwVx2yCfgonpaf1x0thNchDSZlVSV_6jCDTqYKIs,253116 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono-Oblique.ttf,sha256=KAUoE_enCfyJ9S0ZLcmV708P3Fw9e3OknWhJsZFtDNA,251472 +matplotlib/mpl-data/fonts/ttf/DejaVuSansMono.ttf,sha256=YC7Ia4lIz82VZIL-ZPlMNshndwFJ7y95HUYT9EO87LM,340240 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-Bold.ttf,sha256=w3U_Lta8Zz8VhG3EWt2-s7nIcvMvsY_VOiHxvvHtdnY,355692 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-BoldItalic.ttf,sha256=2T7-x6nS6CZ2jRou6VuVhw4V4pWZqE80hK8d4c7C4YE,347064 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif-Italic.ttf,sha256=PnmU-8VPoQzjNSpC1Uj63X2crbacsRCbydlg9trFfwQ,345612 +matplotlib/mpl-data/fonts/ttf/DejaVuSerif.ttf,sha256=EHJElW6ZYrnpb6zNxVGCXgrgiYrhNzcTPhuSGi_TX_o,379740 +matplotlib/mpl-data/fonts/ttf/DejaVuSerifDisplay.ttf,sha256=KRTzLkfHd8J75Wd6-ufbTeefnkXeb8kJfZlJwjwU99U,14300 +matplotlib/mpl-data/fonts/ttf/LICENSE_DEJAVU,sha256=11k43sCY8G8Kw8AIUwZdlPAgvhw8Yu8dwpdboVtNmw4,4816 +matplotlib/mpl-data/fonts/ttf/LICENSE_STIX,sha256=urPTHf7wf0g2JPL2XycR52BluOcnMnixwHHt4QQcmVk,5476 +matplotlib/mpl-data/fonts/ttf/STIXGeneral.ttf,sha256=FnN4Ax4t3cYhbWeBnJJg6aBv_ExHjk4jy5im_USxg8I,448228 +matplotlib/mpl-data/fonts/ttf/STIXGeneralBol.ttf,sha256=6FM9xwg_o0a9oZM9YOpKg7Z9CUW86vGzVB-CtKDixqA,237360 +matplotlib/mpl-data/fonts/ttf/STIXGeneralBolIta.ttf,sha256=mHiP1LpI37sr0CbA4gokeosGxzcoeWKLemuw1bsJc2w,181152 +matplotlib/mpl-data/fonts/ttf/STIXGeneralItalic.ttf,sha256=bPyzM9IrfDxiO9_UAXTxTIXD1nMcphZsHtyAFA6uhSc,175040 +matplotlib/mpl-data/fonts/ttf/STIXNonUni.ttf,sha256=Ulb34CEzWsSFTRgPDovxmJZOwvyCAXYnbhaqvGU3u1c,59108 +matplotlib/mpl-data/fonts/ttf/STIXNonUniBol.ttf,sha256=XRBqW3jR_8MBdFU0ObhiV7-kXwiBIMs7QVClHcT5tgs,30512 +matplotlib/mpl-data/fonts/ttf/STIXNonUniBolIta.ttf,sha256=pb22DnbDf2yQqizotc3wBDqFGC_g27YcCGJivH9-Le8,41272 +matplotlib/mpl-data/fonts/ttf/STIXNonUniIta.ttf,sha256=BMr9pWiBv2YIZdq04X4c3CgL6NPLUPrl64aV1N4w9Ug,46752 +matplotlib/mpl-data/fonts/ttf/STIXSizFiveSymReg.ttf,sha256=wYuH1gYUpCuusqItRH5kf9p_s6mUD-9X3L5RvRtKSxs,13656 +matplotlib/mpl-data/fonts/ttf/STIXSizFourSymBol.ttf,sha256=yNdvjUoSmsZCULmD7SVq9HabndG9P4dPhboL1JpAf0s,12228 +matplotlib/mpl-data/fonts/ttf/STIXSizFourSymReg.ttf,sha256=-9xVMYL4_1rcO8FiCKrCfR4PaSmKtA42ddLGqwtei1w,15972 +matplotlib/mpl-data/fonts/ttf/STIXSizOneSymBol.ttf,sha256=cYexyo8rZcdqMlpa9fNF5a2IoXLUTZuIvh0JD1Qp0i4,12556 +matplotlib/mpl-data/fonts/ttf/STIXSizOneSymReg.ttf,sha256=0lbHzpndzJmO8S42mlkhsz5NbvJLQCaH5Mcc7QZRDzc,19760 +matplotlib/mpl-data/fonts/ttf/STIXSizThreeSymBol.ttf,sha256=3eBc-VtYbhQU3BnxiypfO6eAzEu8BdDvtIJSFbkS2oY,12192 +matplotlib/mpl-data/fonts/ttf/STIXSizThreeSymReg.ttf,sha256=XFSKCptbESM8uxHtUFSAV2cybwxhSjd8dWVByq6f3w0,15836 +matplotlib/mpl-data/fonts/ttf/STIXSizTwoSymBol.ttf,sha256=MUCYHrA0ZqFiSE_PjIGlJZgMuv79aUgQqE7Dtu3kuo0,12116 +matplotlib/mpl-data/fonts/ttf/STIXSizTwoSymReg.ttf,sha256=_sdxDuEwBDtADpu9CyIXQxV7sIqA2TZVBCUiUjq5UCk,15704 +matplotlib/mpl-data/fonts/ttf/cmb10.ttf,sha256=B0SXtQxD6ldZcYFZH5iT04_BKofpUQT1ZX_CSB9hojo,25680 +matplotlib/mpl-data/fonts/ttf/cmex10.ttf,sha256=ryjwwXByOsd2pxv6WVrKCemNFa5cPVTOGa_VYZyWqQU,21092 +matplotlib/mpl-data/fonts/ttf/cmmi10.ttf,sha256=MJKWW4gR_WpnZXmWZIRRgfwd0TMLk3-RWAjEhdMWI00,32560 +matplotlib/mpl-data/fonts/ttf/cmr10.ttf,sha256=Tdl2GwWMAJ25shRfVe5mF9CTwnPdPWxbPkP_YRD6m_Y,26348 +matplotlib/mpl-data/fonts/ttf/cmss10.ttf,sha256=ffkag9BbLkcexjjLC0NaNgo8eSsJ_EKn2mfpHy55EVo,20376 +matplotlib/mpl-data/fonts/ttf/cmsy10.ttf,sha256=uyJu2TLz8QDNDlL15JEu5VO0G2nnv9uNOFTbDrZgUjI,29396 +matplotlib/mpl-data/fonts/ttf/cmtt10.ttf,sha256=YhHwmuk1mZka_alwwkZp2tGnfiU9kVYk-_IS9wLwcdc,28136 +matplotlib/mpl-data/images/back-symbolic.svg,sha256=Okj_ressZkfe6Ewv_o7GF5toc5qWCeFkQ2cHQ25BdVE,1532 +matplotlib/mpl-data/images/back.pdf,sha256=ZR7CJo_dAeCM-KlaGvskgtHQyRtrPIolc8REOmcoqJk,1623 +matplotlib/mpl-data/images/back.png,sha256=E4dGf4Gnz1xJ1v2tMygHV0YNQgShreDeVApaMb-74mU,380 +matplotlib/mpl-data/images/back.svg,sha256=Okj_ressZkfe6Ewv_o7GF5toc5qWCeFkQ2cHQ25BdVE,1532 +matplotlib/mpl-data/images/back_large.png,sha256=9A6hUSQeszhYONE4ZuH3kvOItM0JfDVu6tkfromCbsQ,620 +matplotlib/mpl-data/images/filesave-symbolic.svg,sha256=dMGXvLSOHPu44kiWgZx-B_My_tLWaP6J6GgxJfL4FW0,2049 +matplotlib/mpl-data/images/filesave.pdf,sha256=P1EPPV2g50WTt8UaX-6kFoTZM1xVqo6S2H6FJ6Zd1ec,1734 +matplotlib/mpl-data/images/filesave.png,sha256=b7ctucrM_F2mG-DycTedG_a_y4pHkx3F-zM7l18GLhk,458 +matplotlib/mpl-data/images/filesave.svg,sha256=dMGXvLSOHPu44kiWgZx-B_My_tLWaP6J6GgxJfL4FW0,2049 +matplotlib/mpl-data/images/filesave_large.png,sha256=LNbRD5KZ3Kf7nbp-stx_a1_6XfGBSWUfDdpgmnzoRvk,720 +matplotlib/mpl-data/images/forward-symbolic.svg,sha256=kOiKq3a4mieMRLVCwQBdOMTRrWG2NOX_5-rbAFHpdmQ,1551 +matplotlib/mpl-data/images/forward.pdf,sha256=KIqIL4YId43LkcOxV_TT5uvz1SP8k5iUNUeJmAElMV8,1630 +matplotlib/mpl-data/images/forward.png,sha256=pKbLepgGiGeyY2TCBl8svjvm7Z4CS3iysFxcq4GR-wk,357 +matplotlib/mpl-data/images/forward.svg,sha256=kOiKq3a4mieMRLVCwQBdOMTRrWG2NOX_5-rbAFHpdmQ,1551 +matplotlib/mpl-data/images/forward_large.png,sha256=36h7m7DZDHql6kkdpNPckyi2LKCe_xhhyavWARz_2kQ,593 +matplotlib/mpl-data/images/hand.pdf,sha256=hspwkNY915KPD7AMWnVQs7LFPOtlcj0VUiLu76dMabQ,4172 +matplotlib/mpl-data/images/hand.png,sha256=2cchRETGKa0hYNKUxnJABwkyYXEBPqJy_VqSPlT0W2Q,979 +matplotlib/mpl-data/images/hand.svg,sha256=hxxBtakaVFA7mpZOGakvo0QUcb2x06rojeS5gnVmyuc,4906 +matplotlib/mpl-data/images/help-symbolic.svg,sha256=XVcFcuzcL3SQ3LjfSbtdLYDjoB5YUkj2jk2Gk8vaZF8,1890 +matplotlib/mpl-data/images/help.pdf,sha256=CeE978IMi0YWznWKjIT1R8IrP4KhZ0S7usPUvreSgcA,1813 +matplotlib/mpl-data/images/help.png,sha256=s4pQrqaQ0py8I7vc9hv3BI3DO_tky-7YBMpaHuBDCBY,472 +matplotlib/mpl-data/images/help.svg,sha256=XVcFcuzcL3SQ3LjfSbtdLYDjoB5YUkj2jk2Gk8vaZF8,1890 +matplotlib/mpl-data/images/help_large.png,sha256=1IwEyWfGRgnoCWM-r9CJHEogTJVD5n1c8LXTK4AJ4RE,747 +matplotlib/mpl-data/images/home-symbolic.svg,sha256=ptrus8h5PZTi9ahYfnaz-uZ8MAHCr72aPeMW48TBR9Q,1911 +matplotlib/mpl-data/images/home.pdf,sha256=e0e0pI-XRtPmvUCW2VTKL1DeYu1pvPmUUeRSgEbWmik,1737 +matplotlib/mpl-data/images/home.png,sha256=IcFdAAUa6_A0qt8IO3I8p4rpXpQgAlJ8ndBECCh7C1w,468 +matplotlib/mpl-data/images/home.svg,sha256=ptrus8h5PZTi9ahYfnaz-uZ8MAHCr72aPeMW48TBR9Q,1911 +matplotlib/mpl-data/images/home_large.png,sha256=uxS2O3tWOHh1iau7CaVV4ermIJaZ007ibm5Z3i8kXYg,790 +matplotlib/mpl-data/images/matplotlib.pdf,sha256=BkSUf-2xoij-eXfpV2t7y1JFKG1zD1gtV6aAg3Xi_wE,22852 +matplotlib/mpl-data/images/matplotlib.png,sha256=w8KLRYVa-voUZXa41hgJauQuoois23f3NFfdc72pUYY,1283 +matplotlib/mpl-data/images/matplotlib.svg,sha256=QiTIcqlQwGaVPtHsEk-vtmJk1wxwZSvijhqBe_b9VCI,62087 +matplotlib/mpl-data/images/matplotlib_large.png,sha256=ElRoue9grUqkZXJngk-nvh4GKfpvJ4gE69WryjCbX5U,3088 +matplotlib/mpl-data/images/move-symbolic.svg,sha256=_uamLnjQ20iwSuKbd8JvTXUFaRq4206MrpFWvtErr8I,2529 +matplotlib/mpl-data/images/move.pdf,sha256=CXk3PGK9WL5t-5J-G2X5Tl-nb6lcErTBS5oUj2St6aU,1867 +matplotlib/mpl-data/images/move.png,sha256=TmjR41IzSzxGbhiUcV64X0zx2BjrxbWH3cSKvnG2vzc,481 +matplotlib/mpl-data/images/move.svg,sha256=_uamLnjQ20iwSuKbd8JvTXUFaRq4206MrpFWvtErr8I,2529 +matplotlib/mpl-data/images/move_large.png,sha256=Skjz2nW_RTA5s_0g88gdq2hrVbm6DOcfYW4Fu42Fn9U,767 +matplotlib/mpl-data/images/qt4_editor_options.pdf,sha256=2qu6GVyBrJvVHxychQoJUiXPYxBylbH2j90QnytXs_w,1568 +matplotlib/mpl-data/images/qt4_editor_options.png,sha256=EryQjQ5hh2dwmIxtzCFiMN1U6Tnd11p1CDfgH5ZHjNM,380 +matplotlib/mpl-data/images/qt4_editor_options.svg,sha256=sdrNIxYT-BLvJ30ASnaRQ5PxF3SB41-pgdaIJT0KqBg,1264 +matplotlib/mpl-data/images/qt4_editor_options_large.png,sha256=-Pd-9Vh5aIr3PZa8O6Ge_BLo41kiEnpmkdDj8a11JkY,619 +matplotlib/mpl-data/images/subplots-symbolic.svg,sha256=Gq4fDSS99Rv5rbR8_nenV6jcY5VsKPARWeH-BZBk9CU,2150 +matplotlib/mpl-data/images/subplots.pdf,sha256=Q0syPMI5EvtgM-CE-YXKOkL9eFUAZnj_X2Ihoj6R4p4,1714 +matplotlib/mpl-data/images/subplots.png,sha256=MUfCItq3_yzb9yRieGOglpn0Y74h8IA7m5i70B63iRc,445 +matplotlib/mpl-data/images/subplots.svg,sha256=Gq4fDSS99Rv5rbR8_nenV6jcY5VsKPARWeH-BZBk9CU,2150 +matplotlib/mpl-data/images/subplots_large.png,sha256=Edu9SwVMQEXJZ5ogU5cyW7VLcwXJdhdf-EtxxmxdkIs,662 +matplotlib/mpl-data/images/zoom_to_rect-symbolic.svg,sha256=uMmdGkO43ZHlezkpieR3_MiqlEc5vROffRDOhY4sxm4,1499 +matplotlib/mpl-data/images/zoom_to_rect.pdf,sha256=SEvPc24gfZRpl-dHv7nx8KkxPyU66Kq4zgQTvGFm9KA,1609 +matplotlib/mpl-data/images/zoom_to_rect.png,sha256=aNz3QZBrIgxu9E-fFfaQweCVNitGuDUFoC27e5NU2L4,530 +matplotlib/mpl-data/images/zoom_to_rect.svg,sha256=uMmdGkO43ZHlezkpieR3_MiqlEc5vROffRDOhY4sxm4,1499 +matplotlib/mpl-data/images/zoom_to_rect_large.png,sha256=V6pkxmm6VwFExdg_PEJWdK37HB7k3cE_corLa7RbUMk,1016 +matplotlib/mpl-data/kpsewhich.lua,sha256=RdyYaBnBLy3NsB5c2R5FGrKu-V-WBcZim24NWilsTfw,139 +matplotlib/mpl-data/matplotlibrc,sha256=CXIWz2WeGv3kyoQsMBDOO0m3yF9M_pK1EihXl01w0bw,43413 +matplotlib/mpl-data/plot_directive/plot_directive.css,sha256=utSJ1oETz0UG6AC9hU134J_JY78ENijqMZXN0JMBUfk,318 +matplotlib/mpl-data/sample_data/Minduka_Present_Blue_Pack.png,sha256=XnKGiCanpDKalQ5anvo5NZSAeDP7fyflzQAaivuc0IE,13634 +matplotlib/mpl-data/sample_data/README.txt,sha256=ABz19VBKfGewdY39QInG9Qccgn1MTYV3bT5Ph7TCy2Y,128 +matplotlib/mpl-data/sample_data/Stocks.csv,sha256=72878aZNXGxd5wLvFUw_rnj-nfg4gqtrucZji-w830c,67924 +matplotlib/mpl-data/sample_data/axes_grid/bivariate_normal.npy,sha256=DpWZ9udAh6ospYqneEa27D6EkRgORFwHosacZXVu98U,1880 +matplotlib/mpl-data/sample_data/data_x_x2_x3.csv,sha256=A0SU3buOUGhT-NI_6LQ6p70fFSIU3iLFdgzvzrKR6SE,132 +matplotlib/mpl-data/sample_data/eeg.dat,sha256=KGVjFt8ABKz7p6XZirNfcxSTOpGGNuyA8JYErRKLRBc,25600 +matplotlib/mpl-data/sample_data/embedding_in_wx3.xrc,sha256=cUqVw5vDHNSZoaO4J0ebZUf5SrJP36775abs7R9Bclg,2186 +matplotlib/mpl-data/sample_data/goog.npz,sha256=QAkXzzDmtmT3sNqT18dFhg06qQCNqLfxYNLdEuajGLE,22845 +matplotlib/mpl-data/sample_data/grace_hopper.jpg,sha256=qMptc0dlcDsJcoq0f-WfRz2Trjln_CTHwCiMPHrbcTA,61306 +matplotlib/mpl-data/sample_data/jacksboro_fault_dem.npz,sha256=1JP1CjPoKkQgSUxU0fyhU50Xe9wnqxkLxf5ukvYvtjc,174061 +matplotlib/mpl-data/sample_data/logo2.png,sha256=DXNx4FXeyqxHy26AmvNELpwezQLxweLQY9HP7ktKIdc,22279 +matplotlib/mpl-data/sample_data/membrane.dat,sha256=q3lbQpIBpbtXXGNw1eFwkN_PwxdDGqk4L46IE2b0M1c,48000 +matplotlib/mpl-data/sample_data/msft.csv,sha256=GArKb0O3DgKZRsKdJf6lX3rMSf-PCekIiBoLNdgF7Mk,3211 +matplotlib/mpl-data/sample_data/s1045.ima.gz,sha256=MrQk1k9it-ccsk0p_VOTitVmTWCAVaZ6srKvQ2n4uJ4,33229 +matplotlib/mpl-data/sample_data/topobathy.npz,sha256=AkTgMpFwLfRQJNy1ysvE89TLMNct-n_TccSsYcQrT78,45224 +matplotlib/mpl-data/stylelib/Solarize_Light2.mplstyle,sha256=aytOm4eT_SPvs7HC28ZY4GukeN44q-SE0JEMCR8kVOk,1257 +matplotlib/mpl-data/stylelib/_classic_test_patch.mplstyle,sha256=iopHpMaM3im_AK2aiHGuM2DKM5i9Kc84v6NQEoSb10Q,167 +matplotlib/mpl-data/stylelib/_mpl-gallery-nogrid.mplstyle,sha256=1VOL3USqD6iuGQaSynNg1QhyUwvKLnkLyUKdbBMnnqg,489 +matplotlib/mpl-data/stylelib/_mpl-gallery.mplstyle,sha256=MN-q59CiDqHXB8xFKXxzCbJJbJmNDhBe9lDJJAoMTPA,504 +matplotlib/mpl-data/stylelib/bmh.mplstyle,sha256=-KbhaI859BITHIoyUZIfpQDjfckgLAlDAS_ydKsm6mc,712 +matplotlib/mpl-data/stylelib/classic.mplstyle,sha256=1o5b47VD_RIZv3unnG9Gm2tbprTvOeNGXM8hJCmGuYI,24670 +matplotlib/mpl-data/stylelib/dark_background.mplstyle,sha256=GzSBD06jvlRYOqu7D5Z5a5x25l9JnTgtObn7S4D9zug,607 +matplotlib/mpl-data/stylelib/fast.mplstyle,sha256=yTa2YEIIP9xi5V_G0p2vSlxghuhNwjRi9gPECMxyRiM,288 +matplotlib/mpl-data/stylelib/fivethirtyeight.mplstyle,sha256=IcBt2DSDz9CtuyCis9JZtT3Nqoh6LVO-dB66AxObUcA,780 +matplotlib/mpl-data/stylelib/ggplot.mplstyle,sha256=u2oPHMLWFtZcpIjHk2swi2Nrt4NgnEtof5lxcwM0RD0,956 +matplotlib/mpl-data/stylelib/grayscale.mplstyle,sha256=KCLg-pXpns9cnKDXKN2WH6mV41OH-6cbT-5zKQotSdw,526 +matplotlib/mpl-data/stylelib/petroff10.mplstyle,sha256=Qj7pPbHh3L25kdsPt1ypwvKR3dPzuOrFzzQTP3Mfilg,298 +matplotlib/mpl-data/stylelib/seaborn-v0_8-bright.mplstyle,sha256=pDqn3-NUyVLvlfkYs8n8HzNZvmslVMChkeH-HtZuJIc,144 +matplotlib/mpl-data/stylelib/seaborn-v0_8-colorblind.mplstyle,sha256=eCSzFj5_2vR6n5qu1rHE46wvSVGZcdVqz85ov40ZsH8,148 +matplotlib/mpl-data/stylelib/seaborn-v0_8-dark-palette.mplstyle,sha256=p5ABKNQHRG7bk4HXqMQrRBjDlxGAo3RCXHdQmP7g-Ng,142 +matplotlib/mpl-data/stylelib/seaborn-v0_8-dark.mplstyle,sha256=I4xQ75vE5_9X4k0cNDiqhhnF3OcrZ2xlPX8Ll7OCkoE,667 +matplotlib/mpl-data/stylelib/seaborn-v0_8-darkgrid.mplstyle,sha256=2bXOSzS5gmPzRBrRmzVWyhg_7ZaBRQ6t_-O-cRuyZoA,670 +matplotlib/mpl-data/stylelib/seaborn-v0_8-deep.mplstyle,sha256=44dLcXjjRgR-6yaopgGRInaVgz3jk8VJVQTbBIcxRB0,142 +matplotlib/mpl-data/stylelib/seaborn-v0_8-muted.mplstyle,sha256=T4o3jvqKD_ImXDkp66XFOV_xrBVFUolJU34JDFk1Xkk,143 +matplotlib/mpl-data/stylelib/seaborn-v0_8-notebook.mplstyle,sha256=PcvZQbYrDdducrNlavBPmQ1g2minio_9GkUUFRdgtoM,382 +matplotlib/mpl-data/stylelib/seaborn-v0_8-paper.mplstyle,sha256=n0mboUp2C4Usq2j6tNWcu4TZ_YT4-kKgrYO0t-rz1yw,393 +matplotlib/mpl-data/stylelib/seaborn-v0_8-pastel.mplstyle,sha256=8nV8qRpbUrnFZeyE6VcQ1oRuZPLil2W74M2U37DNMOE,144 +matplotlib/mpl-data/stylelib/seaborn-v0_8-poster.mplstyle,sha256=dUaKqTE4MRfUq2rWVXbbou7kzD7Z9PE9Ko8aXLza8JA,403 +matplotlib/mpl-data/stylelib/seaborn-v0_8-talk.mplstyle,sha256=7FnBaBEdWBbncTm6_ER-EQVa_bZgU7dncgez-ez8R74,403 +matplotlib/mpl-data/stylelib/seaborn-v0_8-ticks.mplstyle,sha256=CITZmZFUFp40MK2Oz8tI8a7WRoCizQU9Z4J172YWfWw,665 +matplotlib/mpl-data/stylelib/seaborn-v0_8-white.mplstyle,sha256=WjJ6LEU6rlCwUugToawciAbKP9oERFHr9rfFlUrdTx0,665 +matplotlib/mpl-data/stylelib/seaborn-v0_8-whitegrid.mplstyle,sha256=ec4BjsNzmOvHptcJ3mdPxULF3S1_U1EUocuqfIpw-Nk,664 +matplotlib/mpl-data/stylelib/seaborn-v0_8.mplstyle,sha256=_Xu6qXKzi4b3GymCOB1b1-ykKTQ8xhDliZ8ezHGTiAs,1130 +matplotlib/mpl-data/stylelib/tableau-colorblind10.mplstyle,sha256=BsirZVd1LmPWT4tBIz6loZPjZcInoQrIGfC7rvzqmJw,190 +matplotlib/offsetbox.py,sha256=hcBN7g2UBFLu_WhBg_bDlnHcIdjc6Pj-kHicfhOpTpI,54490 +matplotlib/offsetbox.pyi,sha256=ISf7gRrOhvefNKfutbHIPh6QYS0kAgtA97bYa6N-82o,9909 +matplotlib/patches.py,sha256=uXew19bSegJ1vxVIhZMhmf8f4XXyB7wKIwgAnFPfwnk,165064 +matplotlib/patches.pyi,sha256=d-sCqVXw_ANsvoUpliH7dllMbTy7TcyRm2IV97Ll0nM,22649 +matplotlib/path.py,sha256=QiQRSMjHbKCBpjWyJ3KPIJgK1Eq07qXUYjULaDHgt5g,43269 +matplotlib/path.pyi,sha256=VfIuc2KpkcB28uzRjLvkC3Hq5HwH4Zp0Uyj9H9bwcWU,4808 +matplotlib/patheffects.py,sha256=4paQgPGKc_xdBz_BZ_MaVhKJ2FxZDGeh5zHwzn3evDY,18387 +matplotlib/patheffects.pyi,sha256=7-FhuxGGrer94GtJ1sZ0YxOmK6Nv4oixTmsKb6-ijOg,3664 +matplotlib/projections/__init__.py,sha256=gICKgNfjJ-tioWEFa75BKKJZG6JJjx46_RpEqt5fk94,4438 +matplotlib/projections/__init__.pyi,sha256=D28dSYmwZcSBFBtNDer-QqE_lqXAhHKersuAlvi89jE,673 +matplotlib/projections/__pycache__/__init__.cpython-310.pyc,, +matplotlib/projections/__pycache__/geo.cpython-310.pyc,, +matplotlib/projections/__pycache__/polar.cpython-310.pyc,, +matplotlib/projections/geo.py,sha256=feR2dybCylhk6xe-MzxBr_e8MdbpVc8tBtCQQJ0V2co,17605 +matplotlib/projections/geo.pyi,sha256=vPfhvj7_e0ZnKjyfDUNC89RGCktycJBPnn5D8w0A7N8,3775 +matplotlib/projections/polar.py,sha256=vEy7FB4FnEu7COWyjZT-zYDt4tbTwkW8BvZCr7kqYiQ,57241 +matplotlib/projections/polar.pyi,sha256=IWNSLQIo5cXFBJur5VFVDQBwAtq5n1XFZMnrj85mCBg,6636 +matplotlib/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/pylab.py,sha256=VqUqd2J2-dKtltZtsYP8ePKoX1jsNDpTyqlcmpfC-lo,2369 +matplotlib/pyplot.py,sha256=0-anFye98wVBms7QS-Bd7adPaWVNJWEVRqzKKnor3Uw,149717 +matplotlib/quiver.py,sha256=gx-mX4BE2DV1EjRQMTR_3jqZM8NDh1npUOizdACy4iQ,48676 +matplotlib/quiver.pyi,sha256=Kq1FvQP-DRX1N9Xlp6XF6OJfKkJ7tc4dYq4CEndxEzg,5640 +matplotlib/rcsetup.py,sha256=87KrQeDFg4sRu4obxF4mc0eGSCuClGyuzXsxoZ_kme4,51606 +matplotlib/rcsetup.pyi,sha256=0VTIhzfKgBcKxOeOdoKq1oa7ZubalOj9Ks0efE0-cms,4337 +matplotlib/sankey.py,sha256=eKA9DrX96-bpoWRdeNnvC73aJfPtCnB78WmMPCXQ9rc,36151 +matplotlib/sankey.pyi,sha256=P2AGkhdg3ZLqXaQdf8yodblrBpdgEQZBLKPL1mbxmFc,1451 +matplotlib/scale.py,sha256=n3JqLU_d3dAKrP1uO5lTipVEaxgiLhDR2sSk8TAXE_M,26924 +matplotlib/scale.pyi,sha256=-ptRptcqiAuzfKwrjgSWWOxFmjRUTOKGwIoWtuBXKgY,5057 +matplotlib/sphinxext/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +matplotlib/sphinxext/__pycache__/__init__.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/figmpl_directive.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/mathmpl.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/plot_directive.cpython-310.pyc,, +matplotlib/sphinxext/__pycache__/roles.cpython-310.pyc,, +matplotlib/sphinxext/figmpl_directive.py,sha256=rgoJcGRfLdp4jiOlyBIwG83wcyfzbFVfrlrSCmDHL9g,9308 +matplotlib/sphinxext/mathmpl.py,sha256=CKIocBSV4prRMPR50NL6WxBTULBa1zReoqo7ZtbvCsc,7871 +matplotlib/sphinxext/plot_directive.py,sha256=6GEOX2y0K1Izc98Aojt8pCXS21Ucru6E-T_33lRU92s,32542 +matplotlib/sphinxext/roles.py,sha256=U7P4t06foLF2PsZJrTN4mfaspKGRoYwFZ3ijIhIUVyo,4882 +matplotlib/spines.py,sha256=_9Dc4oIAgDs3MA6rZCwr1scPGCSzbAnmYp3PQzpKFeU,21923 +matplotlib/spines.pyi,sha256=eceCS47bOawgkWdLUMIPt6XX_H-EtWS2t-euz5Bgkbg,2951 +matplotlib/stackplot.py,sha256=a6IUAujDjGpaFqiTzGr85eDD8l6baWLPrM5YNvb6g9I,4997 +matplotlib/stackplot.pyi,sha256=BX2g3-6dJ4RhjjIguo-gwWhvKZNYqFCQYVw5SKvpM4c,561 +matplotlib/streamplot.py,sha256=rXVKYEi4MfSg3HnAaWb8w4qFp95mRf6dosMVqtXZEpE,24011 +matplotlib/streamplot.pyi,sha256=fkdkew8u807qd7WXFd_ZjF7FUbx8okOI8A7u4x_ktvg,2690 +matplotlib/style/__init__.py,sha256=fraQtyBC3TY2ZTsLdxL7zNs9tJYZtje7tbiqEf3M56M,140 +matplotlib/style/__pycache__/__init__.cpython-310.pyc,, +matplotlib/style/__pycache__/core.cpython-310.pyc,, +matplotlib/style/core.py,sha256=RSNH2NO0xK_Y37dbpMWyaWMMjvcMa_NHaKXDa5I4_tA,8362 +matplotlib/style/core.pyi,sha256=mIlyz6eChdMjZ_A7S05iJgcWwrWE2NXg2rhm1C1eYlQ,521 +matplotlib/table.py,sha256=MkF9s8u2X-9_Cei9sJixljb2_I_B_4UOqVx0iReKIU0,27744 +matplotlib/table.pyi,sha256=tcR40hoCWCRLlL-8PtFMa7P5sq8qajDSMFxd25Q90z4,3098 +matplotlib/testing/__init__.py,sha256=pTaqH1qgGNBeWCmvK_vjqFJXJ9dvZXCaKQ-8jHSpaEc,6942 +matplotlib/testing/__init__.pyi,sha256=ffqfetWzyCVrSx7BlnoxCmbgIaZg2x57nDrq9eucRk0,1752 +matplotlib/testing/__pycache__/__init__.cpython-310.pyc,, +matplotlib/testing/__pycache__/_markers.cpython-310.pyc,, +matplotlib/testing/__pycache__/compare.cpython-310.pyc,, +matplotlib/testing/__pycache__/conftest.cpython-310.pyc,, +matplotlib/testing/__pycache__/decorators.cpython-310.pyc,, +matplotlib/testing/__pycache__/exceptions.cpython-310.pyc,, +matplotlib/testing/__pycache__/widgets.cpython-310.pyc,, +matplotlib/testing/_markers.py,sha256=0iNyOi25XLv_gTfSUqiRizdSqJzozePPBMRo72H2Je4,1419 +matplotlib/testing/compare.py,sha256=8dVcuP7FauHY1JZQXBD7fNBfQKb4IoXM6pVeq6zJgBs,19780 +matplotlib/testing/compare.pyi,sha256=xlJ4chgXKe567NUavlu-dalyPr4wAQUbd2Fz6aK_JII,1192 +matplotlib/testing/conftest.py,sha256=CcwfUtIkUOgGBlXXNplQE6dg2U_UBKdH6AajROPJ5uw,6683 +matplotlib/testing/conftest.pyi,sha256=zOF_MM2GjCkNi9nZ8vldV9VI41bD67nVhrIgLVuMNsk,416 +matplotlib/testing/decorators.py,sha256=1Xiyb2exeN0nDEO4he8MW_VKwFE34cB7EtRl6bcGRzg,18022 +matplotlib/testing/decorators.pyi,sha256=0fSpdLBtEH7ZP_trVJ7RPxNtOX9sJ_z-MkNsbUxF8nM,872 +matplotlib/testing/exceptions.py,sha256=72QmjiHG7DwxSvlJf8mei-hRit5AH3NKh0-osBo4YbY,138 +matplotlib/testing/jpl_units/Duration.py,sha256=9FMBu9uj6orCWtf23cf6_9HCFUC50xAHrCzaxATwQfM,3966 +matplotlib/testing/jpl_units/Epoch.py,sha256=-FGxeq-VvCS9GVPwOEE5ind_G4Tl9ztD-gYcW9CWzjo,6100 +matplotlib/testing/jpl_units/EpochConverter.py,sha256=fhWjyP567bzcTU_oNuJJpucoolqS88Nt-yEFg1-3yEk,2944 +matplotlib/testing/jpl_units/StrConverter.py,sha256=codGw9b_Zc-MG_YK4CiyMrnMR8ahR9hw836O2SsV8QI,2865 +matplotlib/testing/jpl_units/UnitDbl.py,sha256=EABjyEK4MVouyvlwi_9KdYDg-qbYY3aLHoUjRw37Fb0,5882 +matplotlib/testing/jpl_units/UnitDblConverter.py,sha256=B8DssrQVyC4mwvSFP78cGL0vCnZgVhDaAbZE-jsXLUg,2828 +matplotlib/testing/jpl_units/UnitDblFormatter.py,sha256=246hgA4_pCfJm-P94hEsxqnTS9t0XlvLC8p1v_bw2pU,657 +matplotlib/testing/jpl_units/__init__.py,sha256=p__9RUwrt2LJ2eoT2JPM-42XLxSJrfA4az3rN5uP6d4,2684 +matplotlib/testing/jpl_units/__pycache__/Duration.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/Epoch.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/EpochConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/StrConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDbl.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDblConverter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/UnitDblFormatter.cpython-310.pyc,, +matplotlib/testing/jpl_units/__pycache__/__init__.cpython-310.pyc,, +matplotlib/testing/widgets.py,sha256=IaoPHgDguJSHd7t97kSVWUZvqTQ-Tsr9tU9kfnmJrAM,3480 +matplotlib/testing/widgets.pyi,sha256=Ioau7Q2aPRDZLx8hze2DOe3E1vn7QPxePC74WMR7tFc,831 +matplotlib/tests/__init__.py,sha256=XyXveEAxafB87gnbx0jkC0MggzKO8FvORq_6RtJRwo4,366 +matplotlib/tests/__pycache__/__init__.cpython-310.pyc,, +matplotlib/tests/__pycache__/conftest.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_afm.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_agg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_agg_filter.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_animation.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_api.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_arrow_patches.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_artist.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_axes.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_axis.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_bases.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_cairo.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_gtk3.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_inline.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_macosx.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_nbagg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_pdf.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_pgf.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_ps.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_qt.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_registry.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_svg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_template.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_tk.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_tools.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backend_webagg.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_backends_interactive.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_basic.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_bbox_tight.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_bezier.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_category.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_cbook.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_collections.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_colorbar.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_colors.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_compare_images.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_constrainedlayout.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_container.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_contour.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_cycles.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_dates.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_datetime.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_determinism.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_doc.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_dviread.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_figure.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_font_manager.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_fontconfig_pattern.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_ft2font.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_getattr.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_gridspec.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_image.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_legend.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_lines.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_marker.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_mathtext.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_matplotlib.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_mlab.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_multivariate_colormaps.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_offsetbox.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_patches.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_path.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_patheffects.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_pickle.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_png.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_polar.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_preprocess_data.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_pyplot.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_quiver.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_rcparams.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_sankey.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_scale.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_simplification.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_skew.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_sphinxext.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_spines.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_streamplot.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_style.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_subplots.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_table.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_testing.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_texmanager.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_text.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_textpath.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_ticker.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_tightlayout.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_transforms.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_triangulation.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_type1font.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_units.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_usetex.cpython-310.pyc,, +matplotlib/tests/__pycache__/test_widgets.cpython-310.pyc,, +matplotlib/tests/conftest.py,sha256=0cKxWVUC25UfQJQcAjFJa73MQE9eSYqKSXGDx7kAq48,138 +matplotlib/tests/test_afm.py,sha256=A7jm2o-QaQH9SSpiFxGtZkbVU0LJZE69jfPv7RczOD4,3701 +matplotlib/tests/test_agg.py,sha256=xAKrTzvT1dNM_49efiWisnGNrJtINOBenYSOtaRXvIg,10884 +matplotlib/tests/test_agg_filter.py,sha256=3c_Smtb4OHEOfdMFCOb2qKhzMXSbNoaUtsJ0pW47Q44,1067 +matplotlib/tests/test_animation.py,sha256=z0FJ0gx15rJB_khEPSg1E-vLGQyTj85Ki-uTjFsMEPE,18357 +matplotlib/tests/test_api.py,sha256=UJU_qoi7Y8Rbz1S89NYKvOkNwv_RBVPNFajrcXixNds,5736 +matplotlib/tests/test_arrow_patches.py,sha256=qgbS9OrlHa4TJihAd2ekW4AupCVgKDd2B5YyF704Pak,6577 +matplotlib/tests/test_artist.py,sha256=oblCTkCPZzufghCp19AKwIJftV69tZwEOjeRLn8soAQ,18599 +matplotlib/tests/test_axes.py,sha256=PbsmnNA9f3GRtSjCEf4HUvZ5onrBEUIASd1bcalhil8,329934 +matplotlib/tests/test_axis.py,sha256=Ljc_Yc7dAFkTuVW0PD0i3vEnjkGOuAehkpcNpWg9A7Y,1446 +matplotlib/tests/test_backend_bases.py,sha256=FwUgearYkC7UgqpRpaKk51o5F62ipDL4kUvFROwGVLo,22775 +matplotlib/tests/test_backend_cairo.py,sha256=O2LTYjsfPn__bKtTz4MGGBodpSshoPkzu0INsc18xmI,1821 +matplotlib/tests/test_backend_gtk3.py,sha256=Zb29s85nbRzwLVmChCA5_LphAeC_yKyP45kCcHxfgPE,2849 +matplotlib/tests/test_backend_inline.py,sha256=b7GrmMCrJfI5kXGJKbxdMM9kWBd-JKUdCkRq-vk8JXo,1608 +matplotlib/tests/test_backend_macosx.py,sha256=s4Kd1fnwYCuamGRLtkOAYtzer7jVOjLaA2LH52hhPqM,2233 +matplotlib/tests/test_backend_nbagg.py,sha256=o6ON7tt_LOW7wzg4StIMYuTgMFSPNWt1FOv2FFMIefk,1459 +matplotlib/tests/test_backend_pdf.py,sha256=CxVYTlz8_AQ7o90yh5cavpSnNBLozIL8UhscRYShUVM,14602 +matplotlib/tests/test_backend_pgf.py,sha256=WlGdZThz5uQ-x8-HqpMcB5IWmYijflNCC-rVNgvz0yw,13028 +matplotlib/tests/test_backend_ps.py,sha256=veHbm4_dpBDcczgomEjQxdz-Sv8wHij3LVSKWST9Bd0,12600 +matplotlib/tests/test_backend_qt.py,sha256=JfHFhQ1WRw6nMFCRBOPBzAvdvHGMXfoxjiwaUm4GS6Q,12928 +matplotlib/tests/test_backend_registry.py,sha256=SiyyRRMT1tIXKxwn7IXq0SaIrYt3DEvBBe8TvNxa644,5950 +matplotlib/tests/test_backend_svg.py,sha256=L73hpBa1P-y978Y3Gpd_eSJAFJMvUvrOqiYaJ-CkQUw,22930 +matplotlib/tests/test_backend_template.py,sha256=uuE7oZ9pSBVDWrPL05B0WgCFsgv6HlXyetuPTfJn6a8,2184 +matplotlib/tests/test_backend_tk.py,sha256=A1h3AfPhEMbG5PLdDr1gbnwt_DlaQOBgFux5QsIpCpM,8850 +matplotlib/tests/test_backend_tools.py,sha256=C-B7NCkyWsQ5KzQEnI5Be16DsAHHZJU9P5v9--wsF-o,501 +matplotlib/tests/test_backend_webagg.py,sha256=4Oc-z7w-mTGETB1x0FQ_gZP9qHfyWh5lwWc9qPkkisc,938 +matplotlib/tests/test_backends_interactive.py,sha256=oo-vRlDTrLQZOm53FfWECF_JSKDDBJ-BKigekZIO8XE,28880 +matplotlib/tests/test_basic.py,sha256=ubAnlE-lFQzMhoBlWYZJKCAsox9Y3jBXS_IMn29Zi84,1141 +matplotlib/tests/test_bbox_tight.py,sha256=5yo6yy3HTUIf2-oGZllieaLynmOupzWcuPpeGmgS5BQ,6314 +matplotlib/tests/test_bezier.py,sha256=IrrDWAV5O6ELMdziaaOeBI7LfyEaTBVsT4B1T9KMzD8,692 +matplotlib/tests/test_category.py,sha256=jWXOAGOR_gpvrjarvJB6Tyot0yOIWaO6d1F9mrNcVPo,12043 +matplotlib/tests/test_cbook.py,sha256=xESpD9VE9dXInGp_QyzRCeGvmsgY2j9IyNMSUBae-WA,33656 +matplotlib/tests/test_collections.py,sha256=C4jbvAq9t7wcsXY-0SH666aP2t8w4dyO8nRnu7TRdYc,48896 +matplotlib/tests/test_colorbar.py,sha256=tDdmWz_LxCbD80ksM6iqrHYqpombyL8QtqpKQSKOS14,46711 +matplotlib/tests/test_colors.py,sha256=-Q5VZfQWI-NSMOa4tusxDgWe705a7uymMHLI8h10C8A,61215 +matplotlib/tests/test_compare_images.py,sha256=NcBoT8HAxtzoR1zZBu0V3MFpGWtdFaDblQ8o3OTO6zM,3260 +matplotlib/tests/test_constrainedlayout.py,sha256=Mx86XCAfw0XEG67dU73GjZZZi38EZMJErwvgrm7-ThU,24270 +matplotlib/tests/test_container.py,sha256=FTewckOd3dJqLOzEUa29Itjqusk7Mx7MK5xPVMhzMmc,694 +matplotlib/tests/test_contour.py,sha256=O_VIhjDgsb0A-brYR-gJ-WUcqzYzH0XLQD4aI_5SocM,30289 +matplotlib/tests/test_cycles.py,sha256=3KyRmWH29WUgvIXUT06tKVDNCfDWqxuxlLueIp-FIl0,5996 +matplotlib/tests/test_dates.py,sha256=jOK-KIAMBRCcnlgpzknpLJb3JvO3xFa5Qe3O5mVTeGM,56533 +matplotlib/tests/test_datetime.py,sha256=n2TYVItNQcUFBLUivuTXBnc0fElIeSmy9hu4ZNMVurc,32635 +matplotlib/tests/test_determinism.py,sha256=harG-hnOpD5HZFKPtHebT0snHu8V6C6BScVFml1Mqyg,7958 +matplotlib/tests/test_doc.py,sha256=K6HhdRcHRUNNC0iIhxBLCq5tOqxqHBSvanHxXg5w5iI,1015 +matplotlib/tests/test_dviread.py,sha256=JeTuA2FMUj1FddxDVBXUtnvZYTgztE-CyRXL_mI20P0,2764 +matplotlib/tests/test_figure.py,sha256=bGyAW9fwJc3gVVul_Vd5f8TEmbrEoDAMWV4Rm5-QXQQ,60289 +matplotlib/tests/test_font_manager.py,sha256=Qkewu4tMylgVciamVwWL8TFoDon8YyzN-9po6usN6Mk,14136 +matplotlib/tests/test_fontconfig_pattern.py,sha256=LSR6lWF_0cVOshkGiflYaTCLcISRIJM6mjXm5QtDjX4,2168 +matplotlib/tests/test_ft2font.py,sha256=lz13lRGmxdQ9eRjQlnigdfQWBwaST1A3_AgEAR5dJA0,40838 +matplotlib/tests/test_getattr.py,sha256=Tl_H1zpwLdSIVutc4vi-QwDCeWPzBGpN31N9ItzTkeQ,1090 +matplotlib/tests/test_gridspec.py,sha256=SYJuJucA_PyQ2Rsov1RaNhadOEWGDcMbQcVFrWIoy3I,1580 +matplotlib/tests/test_image.py,sha256=O9pUxNrGyewPievlC-WDhJ3hkCTsiAq4id6E6hJ22Is,59562 +matplotlib/tests/test_legend.py,sha256=QeMjTc20_IZeVShlxEZQQsv2suFAPqpvD86TBHESyks,55122 +matplotlib/tests/test_lines.py,sha256=KElNi2GlJCtGLhNpf8rl8bkDhMcQ9ittyRdsX4rdSxo,15035 +matplotlib/tests/test_marker.py,sha256=w0WVHoaD-6iybjUSENoVFFdUOOR13943JcE4sgz3qhI,11410 +matplotlib/tests/test_mathtext.py,sha256=2S1i3BLNffEqz1t-mehhWBROBcsdBZsdgrticxW6Xrw,24549 +matplotlib/tests/test_matplotlib.py,sha256=ZQgD9x7UVWFnVDOm9K5iKfobz6BtMtFkgyVRQ5QezYY,3368 +matplotlib/tests/test_mlab.py,sha256=d4qMyogTFMrvlRZEpDs7SjhSmmCnBUMNgSX2bJU6eDk,42269 +matplotlib/tests/test_multivariate_colormaps.py,sha256=wGc08cdUTBDJvhUiwVqTrBS7o19Ei-wQRnWlkny7h5M,20785 +matplotlib/tests/test_offsetbox.py,sha256=VXFfgputVnLodHhK5ofEWqge42trRSvIm2areAbv1ZQ,16656 +matplotlib/tests/test_patches.py,sha256=FeQGNDS0G8DJNWCAQVit2vN4JPN0ARNougGOTCNa3Fk,33592 +matplotlib/tests/test_path.py,sha256=R3oTrr6kS1Nu_q8zeK1uxrrbnQI8NMb5B7yTk8R1eRs,23254 +matplotlib/tests/test_patheffects.py,sha256=q7JA5VbZU5Zsqvc2eGM7YhB9TClkNp3akQCZVsHrNQU,8109 +matplotlib/tests/test_pickle.py,sha256=2VcsJ6QwN9vP4cvkuUCZD7uwKyd1rOCDIaYPEBr3u6E,10014 +matplotlib/tests/test_png.py,sha256=d6u6UkU71T6ULxDVSdMkIT3CUAL46dEuDuNrnjZe5BE,1407 +matplotlib/tests/test_polar.py,sha256=pPXBMmeUIUZhykYyZa5mrrjFL9GJE3vsaTX94cRjstU,18306 +matplotlib/tests/test_preprocess_data.py,sha256=cIVICUi1iahMQS30sqI5IlT2RYJRH2gJ0z60FyuYUFk,11363 +matplotlib/tests/test_pyplot.py,sha256=3sdLSjSnZtXrag2BFH5ah-AjYcrFzzQwgxnLZ-jx7ak,13929 +matplotlib/tests/test_quiver.py,sha256=5zSF_0DpoDXPBvGWD6iu1VfIUna15npLd2UyCG6_Ksg,11953 +matplotlib/tests/test_rcparams.py,sha256=PLos0asOBQMn0hjrVP_itywcbZjT8og_TaB3tqSFHac,26506 +matplotlib/tests/test_sankey.py,sha256=yg4i-qHT5MljFzFGViOnHbMkumM8bhPwRgoQ5M6CUEs,3900 +matplotlib/tests/test_scale.py,sha256=fqLu88lVy7-JZFNslpgiid11_TIivSkM9DxVnHwZJYo,8429 +matplotlib/tests/test_simplification.py,sha256=IJ1yEz2Y58QtKRRHz41NaXVS4kkRipaE_Un0HgM4K2g,21570 +matplotlib/tests/test_skew.py,sha256=_Mjwfgce6WxLH3dafI7irv2Cw0ANifytCtQXdyOStCk,6349 +matplotlib/tests/test_sphinxext.py,sha256=9hzVcUgMhtfo_6A_c-SvCYWDQY45Oq-lD6Os4PnRg3s,9937 +matplotlib/tests/test_spines.py,sha256=g2CeF0S73XoGA-T04NvYFalIyZvcq7CpbTvN-ikVFn0,5274 +matplotlib/tests/test_streamplot.py,sha256=mSLuZ7E2eEtcI5nKHBZZ_30OYte5awqsXWFWROJYtCQ,5731 +matplotlib/tests/test_style.py,sha256=sd6rMLpPBS1hinR2svI1rXZmnXUnMQk1szxSNBUXtzU,6509 +matplotlib/tests/test_subplots.py,sha256=-YVCEob48FyLUmdInE8vQ2E2b1txTbqChFb1tDUu16s,11082 +matplotlib/tests/test_table.py,sha256=QsrDRpe864LP3zkQBtTciKub5raYz3i3Vrt_Ve3kdzQ,8659 +matplotlib/tests/test_testing.py,sha256=eh-1r4PIXcM7hfSKCNTU899QxRYWhZBg6W7x0TDVUxo,1057 +matplotlib/tests/test_texmanager.py,sha256=j_5ZG02ztIO-m0oXzd31a13YFw_G7JR7Ldt3fwf-phI,2647 +matplotlib/tests/test_text.py,sha256=3RxKrfJJBWKFGSrEvxeCG5rxcvlkIivir8wCj1ls5-M,38473 +matplotlib/tests/test_textpath.py,sha256=WLXvT5OzE6Tew4pr87LH-cgioCzM7srgMNRemiMEC5o,271 +matplotlib/tests/test_ticker.py,sha256=Bj3hiEgXMBXzDi0Lmg8k-Qg4ChPZrQVqGU0Qfh7Sv84,74759 +matplotlib/tests/test_tightlayout.py,sha256=KU0m9BnpbKV-jBdaI4EbDLE5Cf_BLtF0Y9hcFfQ4zlQ,13978 +matplotlib/tests/test_transforms.py,sha256=NfWkPWOKAmY1sq_R1T85ct4CYdFmYUjmjtg2uwOYIIY,48671 +matplotlib/tests/test_triangulation.py,sha256=hstVYhH09BoVl8L0n6vHXBabniZLUu2KYl6wqZn2Xac,55289 +matplotlib/tests/test_type1font.py,sha256=gZZDFi0buFus_aAgOrmtwKqbC36-88UkvM8Afbcngs4,6369 +matplotlib/tests/test_units.py,sha256=0v7JNqjkA8MVnvlOeFoYqQ3sW1KjRY_ijp8BGM8ZW_U,11675 +matplotlib/tests/test_usetex.py,sha256=a-Y6NuyROPHDGP2ELsOZNSVwBUoGtAv1xQZfisl9lSE,6405 +matplotlib/tests/test_widgets.py,sha256=U2FMCCJyKD-pcLbzGzEBAL91-3Fx80wsYdinF7eCOrg,66435 +matplotlib/texmanager.py,sha256=k3Ud7vHbtY0CjXmzHP-ZnEAxg6bHVxo2ALve7_e1rKg,15319 +matplotlib/texmanager.pyi,sha256=di3gbC9muXKTo5VCrW5ye-e19A1mrOl8e8lvI3b904A,1116 +matplotlib/text.py,sha256=xYSkLZiOpew4G1jkLA5pUjtF7qUhOs3yp8IEFMqE6Gg,70856 +matplotlib/text.pyi,sha256=EL_-kiqz88kuPs0hss6IucRrOCsLJ4zaKOYRFV5Wbk8,7019 +matplotlib/textpath.py,sha256=QmrZQ2Mtfyao6qWg-wmEGAbTwkuJ54AnAzFf9NzZHg4,13254 +matplotlib/textpath.pyi,sha256=rqOeTAeQYgm2b2NpetrEX0gMF8PzzW43xS5mNfUA98M,2529 +matplotlib/ticker.py,sha256=YoIJyWVlYk3DIZbBl-w6fzkuVES3j1dN4nyf-SHKVDg,107381 +matplotlib/ticker.pyi,sha256=YvnQDHxJT9H-iomDZpvBCoYkfRRYSot2MR0H0pjZM78,10604 +matplotlib/transforms.py,sha256=ctnOTqgPP4dgHCmSkUayQWL9KkjzWSzQEZALcc9oEzA,99763 +matplotlib/transforms.pyi,sha256=0I70FR4ZWmRoo6w5KjbXRrH0dVDmQyPFbGt_sNTYq10,12102 +matplotlib/tri/__init__.py,sha256=asnfefKRpJv7sGbfddCMybnJInVDPwgph7g0mpoh2u4,820 +matplotlib/tri/__pycache__/__init__.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triangulation.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tricontour.cpython-310.pyc,, +matplotlib/tri/__pycache__/_trifinder.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triinterpolate.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tripcolor.cpython-310.pyc,, +matplotlib/tri/__pycache__/_triplot.cpython-310.pyc,, +matplotlib/tri/__pycache__/_trirefine.cpython-310.pyc,, +matplotlib/tri/__pycache__/_tritools.cpython-310.pyc,, +matplotlib/tri/_triangulation.py,sha256=Ur2lKMOx4NrZxwyi0hBeBnVzicuKaCke0NkrZneSklM,9784 +matplotlib/tri/_triangulation.pyi,sha256=pVw1rvpIcl00p7V7E9GcvJSqQWyoxlZXX_p0_VSxTiY,1017 +matplotlib/tri/_tricontour.py,sha256=yxeH8QgBub1nTq9qbO5Iiiz81bcLXDeMcJHUS2Ahmns,10220 +matplotlib/tri/_tricontour.pyi,sha256=jnsAmVRX0-FOUw9ptUgci9J4T4JQRloKeH8fh8aAi-o,1155 +matplotlib/tri/_trifinder.py,sha256=3gUzJZDIwfdsSJUE8hIKso9e1-UGvynUN9HxaqC1EEc,3522 +matplotlib/tri/_trifinder.pyi,sha256=dXcZucacAS3Ch6nrDBPh2e3LYZLfZ7VwqpBUBb-vMPo,405 +matplotlib/tri/_triinterpolate.py,sha256=4FtyJSoJpHcFxkSkZHZ1aNengVNWqVKF4l78PgCH8O0,62445 +matplotlib/tri/_triinterpolate.pyi,sha256=nR0o0Jm0uPH-6l0ft1WRHqYlq-6o84X6M0aX1UJ9IvE,1044 +matplotlib/tri/_tripcolor.py,sha256=Z6urFv-naDM3_rhv_pkgIZI53h88aOo9qWc3MH7GmBA,6705 +matplotlib/tri/_tripcolor.pyi,sha256=QsA-A2ohj3r_tAElt2-9pzi47JiU01tNlRPDIptqnh4,1781 +matplotlib/tri/_triplot.py,sha256=jlHSz36Z5S18zBKc639PlSqdhfl7jHol8ExlddJuDI4,3102 +matplotlib/tri/_triplot.pyi,sha256=9USU-BfitrcdQE8yWOUlBX59QBNoHCWivDon9JbDQ0k,446 +matplotlib/tri/_trirefine.py,sha256=NG8bsDhZ5EOxMT-MsEWzJm11ZR3_8CAYHlG53IGi0ps,13178 +matplotlib/tri/_trirefine.pyi,sha256=J_PmjbeI6UbLaeecgj1OCvGe_sr9UUsNK9NGBSlQ320,1056 +matplotlib/tri/_tritools.py,sha256=wC9KVE6UqkWVHpyW9FU4hQdqRVRVmJlhaBF1EXsaD8U,10575 +matplotlib/tri/_tritools.pyi,sha256=XWwwvH2nIAmH8k59aRjnLBVQbTwKvd_FzdsRNASCJMw,402 +matplotlib/typing.py,sha256=ZxpQ5jAqiprq3UkqWParWJITgIuOi4ZJ6WEM_q3QOis,2439 +matplotlib/units.py,sha256=7O-llc8k3GpdotUs2tWcEGgoUHHX-Y7o0R7f-1Jve3k,6429 +matplotlib/widgets.py,sha256=zUdNQ_unqHLh76ucCxLKFRUGDiFB1Yi8MmEod5Im5H8,152754 +matplotlib/widgets.pyi,sha256=BpbjCZec7n-uOj_cImu-cB1imVC4oeL8rJUwjxbmN4c,15383 +mpl_toolkits/axes_grid1/__init__.py,sha256=wiuUCQo1g20SW5T3mFOmI9dGCJY6aDmglpQw5DfszEU,371 +mpl_toolkits/axes_grid1/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/anchored_artists.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_divider.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_grid.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_rgb.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/axes_size.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/inset_locator.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/mpl_axes.cpython-310.pyc,, +mpl_toolkits/axes_grid1/__pycache__/parasite_axes.cpython-310.pyc,, +mpl_toolkits/axes_grid1/anchored_artists.py,sha256=ZO5bIF_29sWrk-lWocKJ_45OQdluf_gu7nLrCtfKtdY,17161 +mpl_toolkits/axes_grid1/axes_divider.py,sha256=JD3jGBAqRHRYb1jablrpq_RFQvYtYfwtIIPJgQpbKPc,21892 +mpl_toolkits/axes_grid1/axes_grid.py,sha256=sjEbq6hu7YFimoOYVIcaccEGXsSljNPaTLfJjFHWVtM,22344 +mpl_toolkits/axes_grid1/axes_rgb.py,sha256=pabgaWJuLTCPw2FlT6Zfy5d0_95CEvaLeosWRTElR98,5227 +mpl_toolkits/axes_grid1/axes_size.py,sha256=UAEQ-t0qeJlQcwbbrY4tK2NuSWheFyUXIkk4BUCauHc,7713 +mpl_toolkits/axes_grid1/inset_locator.py,sha256=_2U8ZAj_x7gjKhPk40VSbzE7L6wO2VY-OoOMfd6xOJs,19649 +mpl_toolkits/axes_grid1/mpl_axes.py,sha256=vFCttnj9JIgY3Mt2eOi-O_FVvdZ6SW_sBtIBFib6bz4,4251 +mpl_toolkits/axes_grid1/parasite_axes.py,sha256=809Uy3bLgXIDIGztefcvS9EoZcu3KfMV1QwLzX8qVT4,9404 +mpl_toolkits/axes_grid1/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/axes_grid1/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/__pycache__/test_axes_grid1.cpython-310.pyc,, +mpl_toolkits/axes_grid1/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/axes_grid1/tests/test_axes_grid1.py,sha256=x-W845Cvl2O-eGGhpxKDayFg7-QQXNA94nODLCFPpJc,29093 +mpl_toolkits/axisartist/__init__.py,sha256=RPaNDl22FbmDP7ZRsku1yCqpoNqcclCk0a3rXj3G7fE,631 +mpl_toolkits/axisartist/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/angle_helper.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axes_divider.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axis_artist.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axisline_style.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/axislines.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/floating_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/grid_finder.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/grid_helper_curvelinear.cpython-310.pyc,, +mpl_toolkits/axisartist/__pycache__/parasite_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/angle_helper.py,sha256=-mjKpaR1pLMJuoc0sx0_V3bv0iRPMrpS7r_WI0UYrCc,12952 +mpl_toolkits/axisartist/axes_divider.py,sha256=65xSCQ9cHSC3KE7J7HxS4bfsDTAbPmwyz1jJ43BqnBs,122 +mpl_toolkits/axisartist/axis_artist.py,sha256=9FY9yXl8eF5QuBvF8-Vipv_rmdpcOlA3Q48a5v7FoeA,38328 +mpl_toolkits/axisartist/axisline_style.py,sha256=9jbDkXEzMQiDHR-lDYKZEvTADtJwt2qlN1cErVUUdx0,6723 +mpl_toolkits/axisartist/axislines.py,sha256=QxKvChTaRPj0ovvxdrfr3pOzExEIf5svDwxRE7enEXg,16556 +mpl_toolkits/axisartist/floating_axes.py,sha256=kfhWKkmiy8tkXoVHvooTidXIlicU861VncElsUcBhLo,10337 +mpl_toolkits/axisartist/grid_finder.py,sha256=Hi2zwnQilavgrqWKScuSRba_WBPqqrbmErGU0XyAsOo,12265 +mpl_toolkits/axisartist/grid_helper_curvelinear.py,sha256=ofN7pPqEMh3r6bSf9eOHkAZEOrmhiNZV-Yig3jozkC4,12349 +mpl_toolkits/axisartist/parasite_axes.py,sha256=Ydi4-0Lbczr6K7Sz1-fRwK4Tm8KlHrOIumx67Xbo_9c,244 +mpl_toolkits/axisartist/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/axisartist/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_angle_helper.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_axis_artist.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_axislines.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_floating_axes.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_grid_finder.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/__pycache__/test_grid_helper_curvelinear.cpython-310.pyc,, +mpl_toolkits/axisartist/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/axisartist/tests/test_angle_helper.py,sha256=PwhJwBm2kk4uMyhdO5arQs8IlqSX2vN0hvUzI7YHqrw,5670 +mpl_toolkits/axisartist/tests/test_axis_artist.py,sha256=wt3bicVgUPnBX48-dH0Z6hboHgutIgwVpaGkcUZDeVU,2980 +mpl_toolkits/axisartist/tests/test_axislines.py,sha256=NXegrvEzVWovshta-qjbUKA2tpQcAbjYbfwkf6tKT6Y,4353 +mpl_toolkits/axisartist/tests/test_floating_axes.py,sha256=l24VB1SLrsJZOMMH2jmBny9ETha4AqAM5KokdGOa5Wk,4083 +mpl_toolkits/axisartist/tests/test_grid_finder.py,sha256=cwQLDOdcJbAY2E7dr8595yzuNh1_Yh80r_O8WGT2hMY,1156 +mpl_toolkits/axisartist/tests/test_grid_helper_curvelinear.py,sha256=OhHej0vCfCjJJknT7yIt4OxZd6OMJCXnFoT3pzqUtTo,7216 +mpl_toolkits/mplot3d/__init__.py,sha256=fH9HdMfFMvjbIWqy2gjQnm2m3ae1CvLiuH6LwKHo0kI,49 +mpl_toolkits/mplot3d/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/art3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/axes3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/axis3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/__pycache__/proj3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/art3d.py,sha256=ovygkCV4P7dQ9xIicczbJyyAbaEy2C1HYnzjuV3COus,50548 +mpl_toolkits/mplot3d/axes3d.py,sha256=hCBda2QB6_OYyKEJOG0paPEgDPmLIcrEMkZuKqOUFho,157921 +mpl_toolkits/mplot3d/axis3d.py,sha256=eQPWo2TKbRsPY8JUt5drjBbyRdVKBgeT2EF4HGg51oU,29327 +mpl_toolkits/mplot3d/proj3d.py,sha256=6Hm6WPzeu_wjfeR8afrQ1nCfjS0p3wjvoSIxJWVlS0s,6349 +mpl_toolkits/mplot3d/tests/__init__.py,sha256=sKLxL9jEJBX7eh5OumtXSOnTriPrJUkujTHFtnJVFrM,365 +mpl_toolkits/mplot3d/tests/__pycache__/__init__.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/conftest.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_art3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_axes3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/__pycache__/test_legend3d.cpython-310.pyc,, +mpl_toolkits/mplot3d/tests/conftest.py,sha256=zB61sy90X97YJ16mIGiuaEAaBIjBEzRAK_qfSCichQM,147 +mpl_toolkits/mplot3d/tests/test_art3d.py,sha256=AdFcJf0vz1_b-PIn2cVonPHF1vZw2EOHvmjdgBRa5Yo,3317 +mpl_toolkits/mplot3d/tests/test_axes3d.py,sha256=Xf5M1DipyVfsBBF8dA3AXg5U0Hts194EfAh57VSupE4,92348 +mpl_toolkits/mplot3d/tests/test_legend3d.py,sha256=QGPaoaucJP9KIC68g8zmk4divB_w5PtQc4DMIHMpcA8,4343 +pylab.py,sha256=zUXU0l7e7C5jmDSJbM0GLQxBun3xzuXNf1tuoZYA6Xk,110 diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/WHEEL b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..4e4c38ae320920b8f083b87f408214cdecd350d2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/matplotlib-3.10.7.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: meson +Root-Is-Purelib: false +Tag: cp310-cp310-manylinux_2_17_x86_64 +Tag: cp310-cp310-manylinux2014_x86_64 + diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52e150f1da2497f40c8fed79f235effdf1720aa6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/anchored_artists.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/anchored_artists.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c47cb12fe7b8eba6ae4a7f777c7109a65de17ab Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/anchored_artists.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_divider.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_divider.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27fa4ce9f5fc7594442b7845badb88f6c73d6779 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_divider.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_grid.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_grid.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f80dacce61d1d5a4ac9bc01309c35851d719673b Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_grid.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_rgb.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_rgb.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae921781fdc48ff62eab5c2b63540e3cae94c317 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_rgb.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_size.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_size.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79805e15b248a602544ddd73040734631cac4adb Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/axes_size.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/inset_locator.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/inset_locator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ac8e3f4fb2e4d68b7d7b7931f2e4de1ed0ac650 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/inset_locator.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/mpl_axes.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/mpl_axes.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21bcfab0660f56a678ed2bf9eef99e3e31b7a0fb Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/mpl_axes.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/parasite_axes.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/parasite_axes.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90e4d30b42749f356f9ab037046564e3e73eb1c9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/__pycache__/parasite_axes.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/anchored_artists.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/anchored_artists.py new file mode 100644 index 0000000000000000000000000000000000000000..214b15843ebfaf2e2ee26e100a2194dfe092382d --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/anchored_artists.py @@ -0,0 +1,414 @@ +from matplotlib import transforms +from matplotlib.offsetbox import (AnchoredOffsetbox, AuxTransformBox, + DrawingArea, TextArea, VPacker) +from matplotlib.patches import (Rectangle, ArrowStyle, + FancyArrowPatch, PathPatch) +from matplotlib.text import TextPath + +__all__ = ['AnchoredDrawingArea', 'AnchoredAuxTransformBox', + 'AnchoredSizeBar', 'AnchoredDirectionArrows'] + + +class AnchoredDrawingArea(AnchoredOffsetbox): + def __init__(self, width, height, xdescent, ydescent, + loc, pad=0.4, borderpad=0.5, prop=None, frameon=True, + **kwargs): + """ + An anchored container with a fixed size and fillable `.DrawingArea`. + + Artists added to the *drawing_area* will have their coordinates + interpreted as pixels. Any transformations set on the artists will be + overridden. + + Parameters + ---------- + width, height : float + Width and height of the container, in pixels. + xdescent, ydescent : float + Descent of the container in the x- and y- direction, in pixels. + loc : str + Location of this artist. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + pad : float, default: 0.4 + Padding around the child objects, in fraction of the font size. + borderpad : float, default: 0.5 + Border padding, in fraction of the font size. + prop : `~matplotlib.font_manager.FontProperties`, optional + Font property used as a reference for paddings. + frameon : bool, default: True + If True, draw a box around this artist. + **kwargs + Keyword arguments forwarded to `.AnchoredOffsetbox`. + + Attributes + ---------- + drawing_area : `~matplotlib.offsetbox.DrawingArea` + A container for artists to display. + + Examples + -------- + To display blue and red circles of different sizes in the upper right + of an Axes *ax*: + + >>> ada = AnchoredDrawingArea(20, 20, 0, 0, + ... loc='upper right', frameon=False) + >>> ada.drawing_area.add_artist(Circle((10, 10), 10, fc="b")) + >>> ada.drawing_area.add_artist(Circle((30, 10), 5, fc="r")) + >>> ax.add_artist(ada) + """ + self.da = DrawingArea(width, height, xdescent, ydescent) + self.drawing_area = self.da + + super().__init__( + loc, pad=pad, borderpad=borderpad, child=self.da, prop=None, + frameon=frameon, **kwargs + ) + + +class AnchoredAuxTransformBox(AnchoredOffsetbox): + def __init__(self, transform, loc, + pad=0.4, borderpad=0.5, prop=None, frameon=True, **kwargs): + """ + An anchored container with transformed coordinates. + + Artists added to the *drawing_area* are scaled according to the + coordinates of the transformation used. The dimensions of this artist + will scale to contain the artists added. + + Parameters + ---------- + transform : `~matplotlib.transforms.Transform` + The transformation object for the coordinate system in use, i.e., + :attr:`matplotlib.axes.Axes.transData`. + loc : str + Location of this artist. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + pad : float, default: 0.4 + Padding around the child objects, in fraction of the font size. + borderpad : float, default: 0.5 + Border padding, in fraction of the font size. + prop : `~matplotlib.font_manager.FontProperties`, optional + Font property used as a reference for paddings. + frameon : bool, default: True + If True, draw a box around this artist. + **kwargs + Keyword arguments forwarded to `.AnchoredOffsetbox`. + + Attributes + ---------- + drawing_area : `~matplotlib.offsetbox.AuxTransformBox` + A container for artists to display. + + Examples + -------- + To display an ellipse in the upper left, with a width of 0.1 and + height of 0.4 in data coordinates: + + >>> box = AnchoredAuxTransformBox(ax.transData, loc='upper left') + >>> el = Ellipse((0, 0), width=0.1, height=0.4, angle=30) + >>> box.drawing_area.add_artist(el) + >>> ax.add_artist(box) + """ + self.drawing_area = AuxTransformBox(transform) + + super().__init__(loc, pad=pad, borderpad=borderpad, + child=self.drawing_area, prop=prop, frameon=frameon, + **kwargs) + + +class AnchoredSizeBar(AnchoredOffsetbox): + def __init__(self, transform, size, label, loc, + pad=0.1, borderpad=0.1, sep=2, + frameon=True, size_vertical=0, color='black', + label_top=False, fontproperties=None, fill_bar=None, + **kwargs): + """ + Draw a horizontal scale bar with a center-aligned label underneath. + + Parameters + ---------- + transform : `~matplotlib.transforms.Transform` + The transformation object for the coordinate system in use, i.e., + :attr:`matplotlib.axes.Axes.transData`. + size : float + Horizontal length of the size bar, given in coordinates of + *transform*. + label : str + Label to display. + loc : str + Location of the size bar. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + pad : float, default: 0.1 + Padding around the label and size bar, in fraction of the font + size. + borderpad : float, default: 0.1 + Border padding, in fraction of the font size. + sep : float, default: 2 + Separation between the label and the size bar, in points. + frameon : bool, default: True + If True, draw a box around the horizontal bar and label. + size_vertical : float, default: 0 + Vertical length of the size bar, given in coordinates of + *transform*. + color : str, default: 'black' + Color for the size bar and label. + label_top : bool, default: False + If True, the label will be over the size bar. + fontproperties : `~matplotlib.font_manager.FontProperties`, optional + Font properties for the label text. + fill_bar : bool, optional + If True and if *size_vertical* is nonzero, the size bar will + be filled in with the color specified by the size bar. + Defaults to True if *size_vertical* is greater than + zero and False otherwise. + **kwargs + Keyword arguments forwarded to `.AnchoredOffsetbox`. + + Attributes + ---------- + size_bar : `~matplotlib.offsetbox.AuxTransformBox` + Container for the size bar. + txt_label : `~matplotlib.offsetbox.TextArea` + Container for the label of the size bar. + + Notes + ----- + If *prop* is passed as a keyword argument, but *fontproperties* is + not, then *prop* is assumed to be the intended *fontproperties*. + Using both *prop* and *fontproperties* is not supported. + + Examples + -------- + >>> import matplotlib.pyplot as plt + >>> import numpy as np + >>> from mpl_toolkits.axes_grid1.anchored_artists import ( + ... AnchoredSizeBar) + >>> fig, ax = plt.subplots() + >>> ax.imshow(np.random.random((10, 10))) + >>> bar = AnchoredSizeBar(ax.transData, 3, '3 data units', 4) + >>> ax.add_artist(bar) + >>> fig.show() + + Using all the optional parameters + + >>> import matplotlib.font_manager as fm + >>> fontprops = fm.FontProperties(size=14, family='monospace') + >>> bar = AnchoredSizeBar(ax.transData, 3, '3 units', 4, pad=0.5, + ... sep=5, borderpad=0.5, frameon=False, + ... size_vertical=0.5, color='white', + ... fontproperties=fontprops) + """ + if fill_bar is None: + fill_bar = size_vertical > 0 + + self.size_bar = AuxTransformBox(transform) + self.size_bar.add_artist(Rectangle((0, 0), size, size_vertical, + fill=fill_bar, facecolor=color, + edgecolor=color)) + + if fontproperties is None and 'prop' in kwargs: + fontproperties = kwargs.pop('prop') + + if fontproperties is None: + textprops = {'color': color} + else: + textprops = {'color': color, 'fontproperties': fontproperties} + + self.txt_label = TextArea(label, textprops=textprops) + + if label_top: + _box_children = [self.txt_label, self.size_bar] + else: + _box_children = [self.size_bar, self.txt_label] + + self._box = VPacker(children=_box_children, + align="center", + pad=0, sep=sep) + + super().__init__(loc, pad=pad, borderpad=borderpad, child=self._box, + prop=fontproperties, frameon=frameon, **kwargs) + + +class AnchoredDirectionArrows(AnchoredOffsetbox): + def __init__(self, transform, label_x, label_y, length=0.15, + fontsize=0.08, loc='upper left', angle=0, aspect_ratio=1, + pad=0.4, borderpad=0.4, frameon=False, color='w', alpha=1, + sep_x=0.01, sep_y=0, fontproperties=None, back_length=0.15, + head_width=10, head_length=15, tail_width=2, + text_props=None, arrow_props=None, + **kwargs): + """ + Draw two perpendicular arrows to indicate directions. + + Parameters + ---------- + transform : `~matplotlib.transforms.Transform` + The transformation object for the coordinate system in use, i.e., + :attr:`matplotlib.axes.Axes.transAxes`. + label_x, label_y : str + Label text for the x and y arrows + length : float, default: 0.15 + Length of the arrow, given in coordinates of *transform*. + fontsize : float, default: 0.08 + Size of label strings, given in coordinates of *transform*. + loc : str, default: 'upper left' + Location of the arrow. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + angle : float, default: 0 + The angle of the arrows in degrees. + aspect_ratio : float, default: 1 + The ratio of the length of arrow_x and arrow_y. + Negative numbers can be used to change the direction. + pad : float, default: 0.4 + Padding around the labels and arrows, in fraction of the font size. + borderpad : float, default: 0.4 + Border padding, in fraction of the font size. + frameon : bool, default: False + If True, draw a box around the arrows and labels. + color : str, default: 'white' + Color for the arrows and labels. + alpha : float, default: 1 + Alpha values of the arrows and labels + sep_x, sep_y : float, default: 0.01 and 0 respectively + Separation between the arrows and labels in coordinates of + *transform*. + fontproperties : `~matplotlib.font_manager.FontProperties`, optional + Font properties for the label text. + back_length : float, default: 0.15 + Fraction of the arrow behind the arrow crossing. + head_width : float, default: 10 + Width of arrow head, sent to `.ArrowStyle`. + head_length : float, default: 15 + Length of arrow head, sent to `.ArrowStyle`. + tail_width : float, default: 2 + Width of arrow tail, sent to `.ArrowStyle`. + text_props, arrow_props : dict + Properties of the text and arrows, passed to `.TextPath` and + `.FancyArrowPatch`. + **kwargs + Keyword arguments forwarded to `.AnchoredOffsetbox`. + + Attributes + ---------- + arrow_x, arrow_y : `~matplotlib.patches.FancyArrowPatch` + Arrow x and y + text_path_x, text_path_y : `~matplotlib.text.TextPath` + Path for arrow labels + p_x, p_y : `~matplotlib.patches.PathPatch` + Patch for arrow labels + box : `~matplotlib.offsetbox.AuxTransformBox` + Container for the arrows and labels. + + Notes + ----- + If *prop* is passed as a keyword argument, but *fontproperties* is + not, then *prop* is assumed to be the intended *fontproperties*. + Using both *prop* and *fontproperties* is not supported. + + Examples + -------- + >>> import matplotlib.pyplot as plt + >>> import numpy as np + >>> from mpl_toolkits.axes_grid1.anchored_artists import ( + ... AnchoredDirectionArrows) + >>> fig, ax = plt.subplots() + >>> ax.imshow(np.random.random((10, 10))) + >>> arrows = AnchoredDirectionArrows(ax.transAxes, '111', '110') + >>> ax.add_artist(arrows) + >>> fig.show() + + Using several of the optional parameters, creating downward pointing + arrow and high contrast text labels. + + >>> import matplotlib.font_manager as fm + >>> fontprops = fm.FontProperties(family='monospace') + >>> arrows = AnchoredDirectionArrows(ax.transAxes, 'East', 'South', + ... loc='lower left', color='k', + ... aspect_ratio=-1, sep_x=0.02, + ... sep_y=-0.01, + ... text_props={'ec':'w', 'fc':'k'}, + ... fontproperties=fontprops) + """ + if arrow_props is None: + arrow_props = {} + + if text_props is None: + text_props = {} + + arrowstyle = ArrowStyle("Simple", + head_width=head_width, + head_length=head_length, + tail_width=tail_width) + + if fontproperties is None and 'prop' in kwargs: + fontproperties = kwargs.pop('prop') + + if 'color' not in arrow_props: + arrow_props['color'] = color + + if 'alpha' not in arrow_props: + arrow_props['alpha'] = alpha + + if 'color' not in text_props: + text_props['color'] = color + + if 'alpha' not in text_props: + text_props['alpha'] = alpha + + t_start = transform + t_end = t_start + transforms.Affine2D().rotate_deg(angle) + + self.box = AuxTransformBox(t_end) + + length_x = length + length_y = length*aspect_ratio + + self.arrow_x = FancyArrowPatch( + (0, back_length*length_y), + (length_x, back_length*length_y), + arrowstyle=arrowstyle, + shrinkA=0.0, + shrinkB=0.0, + **arrow_props) + + self.arrow_y = FancyArrowPatch( + (back_length*length_x, 0), + (back_length*length_x, length_y), + arrowstyle=arrowstyle, + shrinkA=0.0, + shrinkB=0.0, + **arrow_props) + + self.box.add_artist(self.arrow_x) + self.box.add_artist(self.arrow_y) + + text_path_x = TextPath(( + length_x+sep_x, back_length*length_y+sep_y), label_x, + size=fontsize, prop=fontproperties) + self.p_x = PathPatch(text_path_x, transform=t_start, **text_props) + self.box.add_artist(self.p_x) + + text_path_y = TextPath(( + length_x*back_length+sep_x, length_y*(1-back_length)+sep_y), + label_y, size=fontsize, prop=fontproperties) + self.p_y = PathPatch(text_path_y, **text_props) + self.box.add_artist(self.p_y) + + super().__init__(loc, pad=pad, borderpad=borderpad, child=self.box, + frameon=frameon, **kwargs) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_grid.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..20abf18ea79c32ccdfe678d81ad975e83fe4bc0b --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_grid.py @@ -0,0 +1,563 @@ +from numbers import Number +import functools +from types import MethodType + +import numpy as np + +from matplotlib import _api, cbook +from matplotlib.gridspec import SubplotSpec + +from .axes_divider import Size, SubplotDivider, Divider +from .mpl_axes import Axes, SimpleAxisArtist + + +class CbarAxesBase: + def __init__(self, *args, orientation, **kwargs): + self.orientation = orientation + super().__init__(*args, **kwargs) + + def colorbar(self, mappable, **kwargs): + return self.get_figure(root=False).colorbar( + mappable, cax=self, location=self.orientation, **kwargs) + + +_cbaraxes_class_factory = cbook._make_class_factory(CbarAxesBase, "Cbar{}") + + +class Grid: + """ + A grid of Axes. + + In Matplotlib, the Axes location (and size) is specified in normalized + figure coordinates. This may not be ideal for images that needs to be + displayed with a given aspect ratio; for example, it is difficult to + display multiple images of a same size with some fixed padding between + them. AxesGrid can be used in such case. + + Attributes + ---------- + axes_all : list of Axes + A flat list of Axes. Note that you can also access this directly + from the grid. The following is equivalent :: + + grid[i] == grid.axes_all[i] + len(grid) == len(grid.axes_all) + + axes_column : list of list of Axes + A 2D list of Axes where the first index is the column. This results + in the usage pattern ``grid.axes_column[col][row]``. + axes_row : list of list of Axes + A 2D list of Axes where the first index is the row. This results + in the usage pattern ``grid.axes_row[row][col]``. + axes_llc : Axes + The Axes in the lower left corner. + ngrids : int + Number of Axes in the grid. + """ + + _defaultAxesClass = Axes + + def __init__(self, fig, + rect, + nrows_ncols, + ngrids=None, + direction="row", + axes_pad=0.02, + *, + share_all=False, + share_x=True, + share_y=True, + label_mode="L", + axes_class=None, + aspect=False, + ): + """ + Parameters + ---------- + fig : `.Figure` + The parent figure. + rect : (float, float, float, float), (int, int, int), int, or \ + `~.SubplotSpec` + The axes position, as a ``(left, bottom, width, height)`` tuple, + as a three-digit subplot position code (e.g., ``(1, 2, 1)`` or + ``121``), or as a `~.SubplotSpec`. + nrows_ncols : (int, int) + Number of rows and columns in the grid. + ngrids : int or None, default: None + If not None, only the first *ngrids* axes in the grid are created. + direction : {"row", "column"}, default: "row" + Whether axes are created in row-major ("row by row") or + column-major order ("column by column"). This also affects the + order in which axes are accessed using indexing (``grid[index]``). + axes_pad : float or (float, float), default: 0.02 + Padding or (horizontal padding, vertical padding) between axes, in + inches. + share_all : bool, default: False + Whether all axes share their x- and y-axis. Overrides *share_x* + and *share_y*. + share_x : bool, default: True + Whether all axes of a column share their x-axis. + share_y : bool, default: True + Whether all axes of a row share their y-axis. + label_mode : {"L", "1", "all", "keep"}, default: "L" + Determines which axes will get tick labels: + + - "L": All axes on the left column get vertical tick labels; + all axes on the bottom row get horizontal tick labels. + - "1": Only the bottom left axes is labelled. + - "all": All axes are labelled. + - "keep": Do not do anything. + + axes_class : subclass of `matplotlib.axes.Axes`, default: `.mpl_axes.Axes` + The type of Axes to create. + aspect : bool, default: False + Whether the axes aspect ratio follows the aspect ratio of the data + limits. + """ + self._nrows, self._ncols = nrows_ncols + + if ngrids is None: + ngrids = self._nrows * self._ncols + else: + if not 0 < ngrids <= self._nrows * self._ncols: + raise ValueError( + "ngrids must be positive and not larger than nrows*ncols") + + self.ngrids = ngrids + + self._horiz_pad_size, self._vert_pad_size = map( + Size.Fixed, np.broadcast_to(axes_pad, 2)) + + _api.check_in_list(["column", "row"], direction=direction) + self._direction = direction + + if axes_class is None: + axes_class = self._defaultAxesClass + elif isinstance(axes_class, (list, tuple)): + cls, kwargs = axes_class + axes_class = functools.partial(cls, **kwargs) + + kw = dict(horizontal=[], vertical=[], aspect=aspect) + if isinstance(rect, (Number, SubplotSpec)): + self._divider = SubplotDivider(fig, rect, **kw) + elif len(rect) == 3: + self._divider = SubplotDivider(fig, *rect, **kw) + elif len(rect) == 4: + self._divider = Divider(fig, rect, **kw) + else: + raise TypeError("Incorrect rect format") + + rect = self._divider.get_position() + + axes_array = np.full((self._nrows, self._ncols), None, dtype=object) + for i in range(self.ngrids): + col, row = self._get_col_row(i) + if share_all: + sharex = sharey = axes_array[0, 0] + else: + sharex = axes_array[0, col] if share_x else None + sharey = axes_array[row, 0] if share_y else None + axes_array[row, col] = axes_class( + fig, rect, sharex=sharex, sharey=sharey) + self.axes_all = axes_array.ravel( + order="C" if self._direction == "row" else "F").tolist() + self.axes_column = axes_array.T.tolist() + self.axes_row = axes_array.tolist() + self.axes_llc = self.axes_column[0][-1] + + self._init_locators() + + for ax in self.axes_all: + fig.add_axes(ax) + + self.set_label_mode(label_mode) + + def _init_locators(self): + self._divider.set_horizontal( + [Size.Scaled(1), self._horiz_pad_size] * (self._ncols-1) + [Size.Scaled(1)]) + self._divider.set_vertical( + [Size.Scaled(1), self._vert_pad_size] * (self._nrows-1) + [Size.Scaled(1)]) + for i in range(self.ngrids): + col, row = self._get_col_row(i) + self.axes_all[i].set_axes_locator( + self._divider.new_locator(nx=2 * col, ny=2 * (self._nrows - 1 - row))) + + def _get_col_row(self, n): + if self._direction == "column": + col, row = divmod(n, self._nrows) + else: + row, col = divmod(n, self._ncols) + + return col, row + + # Good to propagate __len__ if we have __getitem__ + def __len__(self): + return len(self.axes_all) + + def __getitem__(self, i): + return self.axes_all[i] + + def get_geometry(self): + """ + Return the number of rows and columns of the grid as (nrows, ncols). + """ + return self._nrows, self._ncols + + def set_axes_pad(self, axes_pad): + """ + Set the padding between the axes. + + Parameters + ---------- + axes_pad : (float, float) + The padding (horizontal pad, vertical pad) in inches. + """ + self._horiz_pad_size.fixed_size = axes_pad[0] + self._vert_pad_size.fixed_size = axes_pad[1] + + def get_axes_pad(self): + """ + Return the axes padding. + + Returns + ------- + hpad, vpad + Padding (horizontal pad, vertical pad) in inches. + """ + return (self._horiz_pad_size.fixed_size, + self._vert_pad_size.fixed_size) + + def set_aspect(self, aspect): + """Set the aspect of the SubplotDivider.""" + self._divider.set_aspect(aspect) + + def get_aspect(self): + """Return the aspect of the SubplotDivider.""" + return self._divider.get_aspect() + + def set_label_mode(self, mode): + """ + Define which axes have tick labels. + + Parameters + ---------- + mode : {"L", "1", "all", "keep"} + The label mode: + + - "L": All axes on the left column get vertical tick labels; + all axes on the bottom row get horizontal tick labels. + - "1": Only the bottom left axes is labelled. + - "all": All axes are labelled. + - "keep": Do not do anything. + """ + _api.check_in_list(["all", "L", "1", "keep"], mode=mode) + is_last_row, is_first_col = ( + np.mgrid[:self._nrows, :self._ncols] == [[[self._nrows - 1]], [[0]]]) + if mode == "all": + bottom = left = np.full((self._nrows, self._ncols), True) + elif mode == "L": + bottom = is_last_row + left = is_first_col + elif mode == "1": + bottom = left = is_last_row & is_first_col + else: + return + for i in range(self._nrows): + for j in range(self._ncols): + ax = self.axes_row[i][j] + if isinstance(ax.axis, MethodType): + bottom_axis = SimpleAxisArtist(ax.xaxis, 1, ax.spines["bottom"]) + left_axis = SimpleAxisArtist(ax.yaxis, 1, ax.spines["left"]) + else: + bottom_axis = ax.axis["bottom"] + left_axis = ax.axis["left"] + bottom_axis.toggle(ticklabels=bottom[i, j], label=bottom[i, j]) + left_axis.toggle(ticklabels=left[i, j], label=left[i, j]) + + def get_divider(self): + return self._divider + + def set_axes_locator(self, locator): + self._divider.set_locator(locator) + + def get_axes_locator(self): + return self._divider.get_locator() + + +class ImageGrid(Grid): + """ + A grid of Axes for Image display. + + This class is a specialization of `~.axes_grid1.axes_grid.Grid` for displaying a + grid of images. In particular, it forces all axes in a column to share their x-axis + and all axes in a row to share their y-axis. It further provides helpers to add + colorbars to some or all axes. + """ + + def __init__(self, fig, + rect, + nrows_ncols, + ngrids=None, + direction="row", + axes_pad=0.02, + *, + share_all=False, + aspect=True, + label_mode="L", + cbar_mode=None, + cbar_location="right", + cbar_pad=None, + cbar_size="5%", + cbar_set_cax=True, + axes_class=None, + ): + """ + Parameters + ---------- + fig : `.Figure` + The parent figure. + rect : (float, float, float, float) or int + The axes position, as a ``(left, bottom, width, height)`` tuple or + as a three-digit subplot position code (e.g., "121"). + nrows_ncols : (int, int) + Number of rows and columns in the grid. + ngrids : int or None, default: None + If not None, only the first *ngrids* axes in the grid are created. + direction : {"row", "column"}, default: "row" + Whether axes are created in row-major ("row by row") or + column-major order ("column by column"). This also affects the + order in which axes are accessed using indexing (``grid[index]``). + axes_pad : float or (float, float), default: 0.02in + Padding or (horizontal padding, vertical padding) between axes, in + inches. + share_all : bool, default: False + Whether all axes share their x- and y-axis. Note that in any case, + all axes in a column share their x-axis and all axes in a row share + their y-axis. + aspect : bool, default: True + Whether the axes aspect ratio follows the aspect ratio of the data + limits. + label_mode : {"L", "1", "all"}, default: "L" + Determines which axes will get tick labels: + + - "L": All axes on the left column get vertical tick labels; + all axes on the bottom row get horizontal tick labels. + - "1": Only the bottom left axes is labelled. + - "all": all axes are labelled. + + cbar_mode : {"each", "single", "edge", None}, default: None + Whether to create a colorbar for "each" axes, a "single" colorbar + for the entire grid, colorbars only for axes on the "edge" + determined by *cbar_location*, or no colorbars. The colorbars are + stored in the :attr:`cbar_axes` attribute. + cbar_location : {"left", "right", "bottom", "top"}, default: "right" + cbar_pad : float, default: None + Padding between the image axes and the colorbar axes. + + .. versionchanged:: 3.10 + ``cbar_mode="single"`` no longer adds *axes_pad* between the axes + and the colorbar if the *cbar_location* is "left" or "bottom". + + cbar_size : size specification (see `.Size.from_any`), default: "5%" + Colorbar size. + cbar_set_cax : bool, default: True + If True, each axes in the grid has a *cax* attribute that is bound + to associated *cbar_axes*. + axes_class : subclass of `matplotlib.axes.Axes`, default: None + """ + _api.check_in_list(["each", "single", "edge", None], + cbar_mode=cbar_mode) + _api.check_in_list(["left", "right", "bottom", "top"], + cbar_location=cbar_location) + self._colorbar_mode = cbar_mode + self._colorbar_location = cbar_location + self._colorbar_pad = cbar_pad + self._colorbar_size = cbar_size + # The colorbar axes are created in _init_locators(). + + super().__init__( + fig, rect, nrows_ncols, ngrids, + direction=direction, axes_pad=axes_pad, + share_all=share_all, share_x=True, share_y=True, aspect=aspect, + label_mode=label_mode, axes_class=axes_class) + + for ax in self.cbar_axes: + fig.add_axes(ax) + + if cbar_set_cax: + if self._colorbar_mode == "single": + for ax in self.axes_all: + ax.cax = self.cbar_axes[0] + elif self._colorbar_mode == "edge": + for index, ax in enumerate(self.axes_all): + col, row = self._get_col_row(index) + if self._colorbar_location in ("left", "right"): + ax.cax = self.cbar_axes[row] + else: + ax.cax = self.cbar_axes[col] + else: + for ax, cax in zip(self.axes_all, self.cbar_axes): + ax.cax = cax + + def _init_locators(self): + # Slightly abusing this method to inject colorbar creation into init. + + if self._colorbar_pad is None: + # horizontal or vertical arrangement? + if self._colorbar_location in ("left", "right"): + self._colorbar_pad = self._horiz_pad_size.fixed_size + else: + self._colorbar_pad = self._vert_pad_size.fixed_size + self.cbar_axes = [ + _cbaraxes_class_factory(self._defaultAxesClass)( + self.axes_all[0].get_figure(root=False), self._divider.get_position(), + orientation=self._colorbar_location) + for _ in range(self.ngrids)] + + cb_mode = self._colorbar_mode + cb_location = self._colorbar_location + + h = [] + v = [] + + h_ax_pos = [] + h_cb_pos = [] + if cb_mode == "single" and cb_location in ("left", "bottom"): + if cb_location == "left": + sz = self._nrows * Size.AxesX(self.axes_llc) + h.append(Size.from_any(self._colorbar_size, sz)) + h.append(Size.from_any(self._colorbar_pad, sz)) + locator = self._divider.new_locator(nx=0, ny=0, ny1=-1) + elif cb_location == "bottom": + sz = self._ncols * Size.AxesY(self.axes_llc) + v.append(Size.from_any(self._colorbar_size, sz)) + v.append(Size.from_any(self._colorbar_pad, sz)) + locator = self._divider.new_locator(nx=0, nx1=-1, ny=0) + for i in range(self.ngrids): + self.cbar_axes[i].set_visible(False) + self.cbar_axes[0].set_axes_locator(locator) + self.cbar_axes[0].set_visible(True) + + for col, ax in enumerate(self.axes_row[0]): + if col != 0: + h.append(self._horiz_pad_size) + + if ax: + sz = Size.AxesX(ax, aspect="axes", ref_ax=self.axes_all[0]) + else: + sz = Size.AxesX(self.axes_all[0], + aspect="axes", ref_ax=self.axes_all[0]) + + if (cb_location == "left" + and (cb_mode == "each" + or (cb_mode == "edge" and col == 0))): + h_cb_pos.append(len(h)) + h.append(Size.from_any(self._colorbar_size, sz)) + h.append(Size.from_any(self._colorbar_pad, sz)) + + h_ax_pos.append(len(h)) + h.append(sz) + + if (cb_location == "right" + and (cb_mode == "each" + or (cb_mode == "edge" and col == self._ncols - 1))): + h.append(Size.from_any(self._colorbar_pad, sz)) + h_cb_pos.append(len(h)) + h.append(Size.from_any(self._colorbar_size, sz)) + + v_ax_pos = [] + v_cb_pos = [] + for row, ax in enumerate(self.axes_column[0][::-1]): + if row != 0: + v.append(self._vert_pad_size) + + if ax: + sz = Size.AxesY(ax, aspect="axes", ref_ax=self.axes_all[0]) + else: + sz = Size.AxesY(self.axes_all[0], + aspect="axes", ref_ax=self.axes_all[0]) + + if (cb_location == "bottom" + and (cb_mode == "each" + or (cb_mode == "edge" and row == 0))): + v_cb_pos.append(len(v)) + v.append(Size.from_any(self._colorbar_size, sz)) + v.append(Size.from_any(self._colorbar_pad, sz)) + + v_ax_pos.append(len(v)) + v.append(sz) + + if (cb_location == "top" + and (cb_mode == "each" + or (cb_mode == "edge" and row == self._nrows - 1))): + v.append(Size.from_any(self._colorbar_pad, sz)) + v_cb_pos.append(len(v)) + v.append(Size.from_any(self._colorbar_size, sz)) + + for i in range(self.ngrids): + col, row = self._get_col_row(i) + locator = self._divider.new_locator(nx=h_ax_pos[col], + ny=v_ax_pos[self._nrows-1-row]) + self.axes_all[i].set_axes_locator(locator) + + if cb_mode == "each": + if cb_location in ("right", "left"): + locator = self._divider.new_locator( + nx=h_cb_pos[col], ny=v_ax_pos[self._nrows - 1 - row]) + + elif cb_location in ("top", "bottom"): + locator = self._divider.new_locator( + nx=h_ax_pos[col], ny=v_cb_pos[self._nrows - 1 - row]) + + self.cbar_axes[i].set_axes_locator(locator) + elif cb_mode == "edge": + if (cb_location == "left" and col == 0 + or cb_location == "right" and col == self._ncols - 1): + locator = self._divider.new_locator( + nx=h_cb_pos[0], ny=v_ax_pos[self._nrows - 1 - row]) + self.cbar_axes[row].set_axes_locator(locator) + elif (cb_location == "bottom" and row == self._nrows - 1 + or cb_location == "top" and row == 0): + locator = self._divider.new_locator(nx=h_ax_pos[col], + ny=v_cb_pos[0]) + self.cbar_axes[col].set_axes_locator(locator) + + if cb_mode == "single": + if cb_location == "right": + sz = self._nrows * Size.AxesX(self.axes_llc) + h.append(Size.from_any(self._colorbar_pad, sz)) + h.append(Size.from_any(self._colorbar_size, sz)) + locator = self._divider.new_locator(nx=-2, ny=0, ny1=-1) + elif cb_location == "top": + sz = self._ncols * Size.AxesY(self.axes_llc) + v.append(Size.from_any(self._colorbar_pad, sz)) + v.append(Size.from_any(self._colorbar_size, sz)) + locator = self._divider.new_locator(nx=0, nx1=-1, ny=-2) + if cb_location in ("right", "top"): + for i in range(self.ngrids): + self.cbar_axes[i].set_visible(False) + self.cbar_axes[0].set_axes_locator(locator) + self.cbar_axes[0].set_visible(True) + elif cb_mode == "each": + for i in range(self.ngrids): + self.cbar_axes[i].set_visible(True) + elif cb_mode == "edge": + if cb_location in ("right", "left"): + count = self._nrows + else: + count = self._ncols + for i in range(count): + self.cbar_axes[i].set_visible(True) + for j in range(i + 1, self.ngrids): + self.cbar_axes[j].set_visible(False) + else: + for i in range(self.ngrids): + self.cbar_axes[i].set_visible(False) + self.cbar_axes[i].set_position([1., 1., 0.001, 0.001], + which="active") + + self._divider.set_horizontal(h) + self._divider.set_vertical(v) + + +AxesGrid = ImageGrid diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_size.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_size.py new file mode 100644 index 0000000000000000000000000000000000000000..86e5f70d9824392a8c49bce9a18aaf45cf6f0232 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/axes_size.py @@ -0,0 +1,271 @@ +""" +Provides classes of simple units that will be used with `.AxesDivider` +class (or others) to determine the size of each Axes. The unit +classes define `get_size` method that returns a tuple of two floats, +meaning relative and absolute sizes, respectively. + +Note that this class is nothing more than a simple tuple of two +floats. Take a look at the Divider class to see how these two +values are used. + +Once created, the unit classes can be modified by simple arithmetic +operations: addition /subtraction with another unit type or a real number and scaling +(multiplication or division) by a real number. +""" + +from numbers import Real + +from matplotlib import _api +from matplotlib.axes import Axes + + +class _Base: + def __rmul__(self, other): + return self * other + + def __mul__(self, other): + if not isinstance(other, Real): + return NotImplemented + return Fraction(other, self) + + def __div__(self, other): + return (1 / other) * self + + def __add__(self, other): + if isinstance(other, _Base): + return Add(self, other) + else: + return Add(self, Fixed(other)) + + def __neg__(self): + return -1 * self + + def __radd__(self, other): + # other cannot be a _Base instance, because A + B would trigger + # A.__add__(B) first. + return Add(self, Fixed(other)) + + def __sub__(self, other): + return self + (-other) + + def get_size(self, renderer): + """ + Return two-float tuple with relative and absolute sizes. + """ + raise NotImplementedError("Subclasses must implement") + + +class Add(_Base): + """ + Sum of two sizes. + """ + + def __init__(self, a, b): + self._a = a + self._b = b + + def get_size(self, renderer): + a_rel_size, a_abs_size = self._a.get_size(renderer) + b_rel_size, b_abs_size = self._b.get_size(renderer) + return a_rel_size + b_rel_size, a_abs_size + b_abs_size + + +class Fixed(_Base): + """ + Simple fixed size with absolute part = *fixed_size* and relative part = 0. + """ + + def __init__(self, fixed_size): + _api.check_isinstance(Real, fixed_size=fixed_size) + self.fixed_size = fixed_size + + def get_size(self, renderer): + rel_size = 0. + abs_size = self.fixed_size + return rel_size, abs_size + + +class Scaled(_Base): + """ + Simple scaled(?) size with absolute part = 0 and + relative part = *scalable_size*. + """ + + def __init__(self, scalable_size): + self._scalable_size = scalable_size + + def get_size(self, renderer): + rel_size = self._scalable_size + abs_size = 0. + return rel_size, abs_size + +Scalable = Scaled + + +def _get_axes_aspect(ax): + aspect = ax.get_aspect() + if aspect == "auto": + aspect = 1. + return aspect + + +class AxesX(_Base): + """ + Scaled size whose relative part corresponds to the data width + of the *axes* multiplied by the *aspect*. + """ + + def __init__(self, axes, aspect=1., ref_ax=None): + self._axes = axes + self._aspect = aspect + if aspect == "axes" and ref_ax is None: + raise ValueError("ref_ax must be set when aspect='axes'") + self._ref_ax = ref_ax + + def get_size(self, renderer): + l1, l2 = self._axes.get_xlim() + if self._aspect == "axes": + ref_aspect = _get_axes_aspect(self._ref_ax) + aspect = ref_aspect / _get_axes_aspect(self._axes) + else: + aspect = self._aspect + + rel_size = abs(l2-l1)*aspect + abs_size = 0. + return rel_size, abs_size + + +class AxesY(_Base): + """ + Scaled size whose relative part corresponds to the data height + of the *axes* multiplied by the *aspect*. + """ + + def __init__(self, axes, aspect=1., ref_ax=None): + self._axes = axes + self._aspect = aspect + if aspect == "axes" and ref_ax is None: + raise ValueError("ref_ax must be set when aspect='axes'") + self._ref_ax = ref_ax + + def get_size(self, renderer): + l1, l2 = self._axes.get_ylim() + + if self._aspect == "axes": + ref_aspect = _get_axes_aspect(self._ref_ax) + aspect = _get_axes_aspect(self._axes) + else: + aspect = self._aspect + + rel_size = abs(l2-l1)*aspect + abs_size = 0. + return rel_size, abs_size + + +class MaxExtent(_Base): + """ + Size whose absolute part is either the largest width or the largest height + of the given *artist_list*. + """ + + def __init__(self, artist_list, w_or_h): + self._artist_list = artist_list + _api.check_in_list(["width", "height"], w_or_h=w_or_h) + self._w_or_h = w_or_h + + def add_artist(self, a): + self._artist_list.append(a) + + def get_size(self, renderer): + rel_size = 0. + extent_list = [ + getattr(a.get_window_extent(renderer), self._w_or_h) / a.figure.dpi + for a in self._artist_list] + abs_size = max(extent_list, default=0) + return rel_size, abs_size + + +class MaxWidth(MaxExtent): + """ + Size whose absolute part is the largest width of the given *artist_list*. + """ + + def __init__(self, artist_list): + super().__init__(artist_list, "width") + + +class MaxHeight(MaxExtent): + """ + Size whose absolute part is the largest height of the given *artist_list*. + """ + + def __init__(self, artist_list): + super().__init__(artist_list, "height") + + +class Fraction(_Base): + """ + An instance whose size is a *fraction* of the *ref_size*. + + >>> s = Fraction(0.3, AxesX(ax)) + """ + + def __init__(self, fraction, ref_size): + _api.check_isinstance(Real, fraction=fraction) + self._fraction_ref = ref_size + self._fraction = fraction + + def get_size(self, renderer): + if self._fraction_ref is None: + return self._fraction, 0. + else: + r, a = self._fraction_ref.get_size(renderer) + rel_size = r*self._fraction + abs_size = a*self._fraction + return rel_size, abs_size + + +def from_any(size, fraction_ref=None): + """ + Create a Fixed unit when the first argument is a float, or a + Fraction unit if that is a string that ends with %. The second + argument is only meaningful when Fraction unit is created. + + >>> from mpl_toolkits.axes_grid1.axes_size import from_any + >>> a = from_any(1.2) # => Fixed(1.2) + >>> from_any("50%", a) # => Fraction(0.5, a) + """ + if isinstance(size, Real): + return Fixed(size) + elif isinstance(size, str): + if size[-1] == "%": + return Fraction(float(size[:-1]) / 100, fraction_ref) + raise ValueError("Unknown format") + + +class _AxesDecorationsSize(_Base): + """ + Fixed size, corresponding to the size of decorations on a given Axes side. + """ + + _get_size_map = { + "left": lambda tight_bb, axes_bb: axes_bb.xmin - tight_bb.xmin, + "right": lambda tight_bb, axes_bb: tight_bb.xmax - axes_bb.xmax, + "bottom": lambda tight_bb, axes_bb: axes_bb.ymin - tight_bb.ymin, + "top": lambda tight_bb, axes_bb: tight_bb.ymax - axes_bb.ymax, + } + + def __init__(self, ax, direction): + _api.check_in_list(self._get_size_map, direction=direction) + self._direction = direction + self._ax_list = [ax] if isinstance(ax, Axes) else ax + + def get_size(self, renderer): + sz = max([ + self._get_size_map[self._direction]( + ax.get_tightbbox(renderer, call_axes_locator=False), ax.bbox) + for ax in self._ax_list]) + dpi = renderer.points_to_pixels(72) + abs_size = sz / dpi + rel_size = 0 + return rel_size, abs_size diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/inset_locator.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/inset_locator.py new file mode 100644 index 0000000000000000000000000000000000000000..52fe6efc061856bd33939c31efe404f619faf033 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/inset_locator.py @@ -0,0 +1,519 @@ +""" +A collection of functions and objects for creating or placing inset axes. +""" + +from matplotlib import _api, _docstring +from matplotlib.offsetbox import AnchoredOffsetbox +from matplotlib.patches import Patch, Rectangle +from matplotlib.path import Path +from matplotlib.transforms import Bbox +from matplotlib.transforms import IdentityTransform, TransformedBbox + +from . import axes_size as Size +from .parasite_axes import HostAxes + + +class AnchoredLocatorBase(AnchoredOffsetbox): + def __init__(self, bbox_to_anchor, offsetbox, loc, + borderpad=0.5, bbox_transform=None): + super().__init__( + loc, pad=0., child=None, borderpad=borderpad, + bbox_to_anchor=bbox_to_anchor, bbox_transform=bbox_transform + ) + + def draw(self, renderer): + raise RuntimeError("No draw method should be called") + + def __call__(self, ax, renderer): + fig = ax.get_figure(root=False) + if renderer is None: + renderer = fig._get_renderer() + self.axes = ax + bbox = self.get_window_extent(renderer) + px, py = self.get_offset(bbox.width, bbox.height, 0, 0, renderer) + bbox_canvas = Bbox.from_bounds(px, py, bbox.width, bbox.height) + tr = fig.transSubfigure.inverted() + return TransformedBbox(bbox_canvas, tr) + + +class AnchoredSizeLocator(AnchoredLocatorBase): + def __init__(self, bbox_to_anchor, x_size, y_size, loc, + borderpad=0.5, bbox_transform=None): + super().__init__( + bbox_to_anchor, None, loc, + borderpad=borderpad, bbox_transform=bbox_transform + ) + + self.x_size = Size.from_any(x_size) + self.y_size = Size.from_any(y_size) + + def get_bbox(self, renderer): + bbox = self.get_bbox_to_anchor() + dpi = renderer.points_to_pixels(72.) + + r, a = self.x_size.get_size(renderer) + width = bbox.width * r + a * dpi + r, a = self.y_size.get_size(renderer) + height = bbox.height * r + a * dpi + + fontsize = renderer.points_to_pixels(self.prop.get_size_in_points()) + pad = self.pad * fontsize + + return Bbox.from_bounds(0, 0, width, height).padded(pad) + + +class AnchoredZoomLocator(AnchoredLocatorBase): + def __init__(self, parent_axes, zoom, loc, + borderpad=0.5, + bbox_to_anchor=None, + bbox_transform=None): + self.parent_axes = parent_axes + self.zoom = zoom + if bbox_to_anchor is None: + bbox_to_anchor = parent_axes.bbox + super().__init__( + bbox_to_anchor, None, loc, borderpad=borderpad, + bbox_transform=bbox_transform) + + def get_bbox(self, renderer): + bb = self.parent_axes.transData.transform_bbox(self.axes.viewLim) + fontsize = renderer.points_to_pixels(self.prop.get_size_in_points()) + pad = self.pad * fontsize + return ( + Bbox.from_bounds( + 0, 0, abs(bb.width * self.zoom), abs(bb.height * self.zoom)) + .padded(pad)) + + +class BboxPatch(Patch): + @_docstring.interpd + def __init__(self, bbox, **kwargs): + """ + Patch showing the shape bounded by a Bbox. + + Parameters + ---------- + bbox : `~matplotlib.transforms.Bbox` + Bbox to use for the extents of this patch. + + **kwargs + Patch properties. Valid arguments include: + + %(Patch:kwdoc)s + """ + if "transform" in kwargs: + raise ValueError("transform should not be set") + + kwargs["transform"] = IdentityTransform() + super().__init__(**kwargs) + self.bbox = bbox + + def get_path(self): + # docstring inherited + x0, y0, x1, y1 = self.bbox.extents + return Path._create_closed([(x0, y0), (x1, y0), (x1, y1), (x0, y1)]) + + +class BboxConnector(Patch): + @staticmethod + def get_bbox_edge_pos(bbox, loc): + """ + Return the ``(x, y)`` coordinates of corner *loc* of *bbox*; parameters + behave as documented for the `.BboxConnector` constructor. + """ + x0, y0, x1, y1 = bbox.extents + if loc == 1: + return x1, y1 + elif loc == 2: + return x0, y1 + elif loc == 3: + return x0, y0 + elif loc == 4: + return x1, y0 + + @staticmethod + def connect_bbox(bbox1, bbox2, loc1, loc2=None): + """ + Construct a `.Path` connecting corner *loc1* of *bbox1* to corner + *loc2* of *bbox2*, where parameters behave as documented as for the + `.BboxConnector` constructor. + """ + if isinstance(bbox1, Rectangle): + bbox1 = TransformedBbox(Bbox.unit(), bbox1.get_transform()) + if isinstance(bbox2, Rectangle): + bbox2 = TransformedBbox(Bbox.unit(), bbox2.get_transform()) + if loc2 is None: + loc2 = loc1 + x1, y1 = BboxConnector.get_bbox_edge_pos(bbox1, loc1) + x2, y2 = BboxConnector.get_bbox_edge_pos(bbox2, loc2) + return Path([[x1, y1], [x2, y2]]) + + @_docstring.interpd + def __init__(self, bbox1, bbox2, loc1, loc2=None, **kwargs): + """ + Connect two bboxes with a straight line. + + Parameters + ---------- + bbox1, bbox2 : `~matplotlib.transforms.Bbox` + Bounding boxes to connect. + + loc1, loc2 : {1, 2, 3, 4} + Corner of *bbox1* and *bbox2* to draw the line. Valid values are:: + + 'upper right' : 1, + 'upper left' : 2, + 'lower left' : 3, + 'lower right' : 4 + + *loc2* is optional and defaults to *loc1*. + + **kwargs + Patch properties for the line drawn. Valid arguments include: + + %(Patch:kwdoc)s + """ + if "transform" in kwargs: + raise ValueError("transform should not be set") + + kwargs["transform"] = IdentityTransform() + kwargs.setdefault( + "fill", bool({'fc', 'facecolor', 'color'}.intersection(kwargs))) + super().__init__(**kwargs) + self.bbox1 = bbox1 + self.bbox2 = bbox2 + self.loc1 = loc1 + self.loc2 = loc2 + + def get_path(self): + # docstring inherited + return self.connect_bbox(self.bbox1, self.bbox2, + self.loc1, self.loc2) + + +class BboxConnectorPatch(BboxConnector): + @_docstring.interpd + def __init__(self, bbox1, bbox2, loc1a, loc2a, loc1b, loc2b, **kwargs): + """ + Connect two bboxes with a quadrilateral. + + The quadrilateral is specified by two lines that start and end at + corners of the bboxes. The four sides of the quadrilateral are defined + by the two lines given, the line between the two corners specified in + *bbox1* and the line between the two corners specified in *bbox2*. + + Parameters + ---------- + bbox1, bbox2 : `~matplotlib.transforms.Bbox` + Bounding boxes to connect. + + loc1a, loc2a, loc1b, loc2b : {1, 2, 3, 4} + The first line connects corners *loc1a* of *bbox1* and *loc2a* of + *bbox2*; the second line connects corners *loc1b* of *bbox1* and + *loc2b* of *bbox2*. Valid values are:: + + 'upper right' : 1, + 'upper left' : 2, + 'lower left' : 3, + 'lower right' : 4 + + **kwargs + Patch properties for the line drawn: + + %(Patch:kwdoc)s + """ + if "transform" in kwargs: + raise ValueError("transform should not be set") + super().__init__(bbox1, bbox2, loc1a, loc2a, **kwargs) + self.loc1b = loc1b + self.loc2b = loc2b + + def get_path(self): + # docstring inherited + path1 = self.connect_bbox(self.bbox1, self.bbox2, self.loc1, self.loc2) + path2 = self.connect_bbox(self.bbox2, self.bbox1, + self.loc2b, self.loc1b) + path_merged = [*path1.vertices, *path2.vertices, path1.vertices[0]] + return Path(path_merged) + + +def _add_inset_axes(parent_axes, axes_class, axes_kwargs, axes_locator): + """Helper function to add an inset axes and disable navigation in it.""" + if axes_class is None: + axes_class = HostAxes + if axes_kwargs is None: + axes_kwargs = {} + fig = parent_axes.get_figure(root=False) + inset_axes = axes_class( + fig, parent_axes.get_position(), + **{"navigate": False, **axes_kwargs, "axes_locator": axes_locator}) + return fig.add_axes(inset_axes) + + +@_docstring.interpd +def inset_axes(parent_axes, width, height, loc='upper right', + bbox_to_anchor=None, bbox_transform=None, + axes_class=None, axes_kwargs=None, + borderpad=0.5): + """ + Create an inset axes with a given width and height. + + Both sizes used can be specified either in inches or percentage. + For example,:: + + inset_axes(parent_axes, width='40%%', height='30%%', loc='lower left') + + creates in inset axes in the lower left corner of *parent_axes* which spans + over 30%% in height and 40%% in width of the *parent_axes*. Since the usage + of `.inset_axes` may become slightly tricky when exceeding such standard + cases, it is recommended to read :doc:`the examples + `. + + Notes + ----- + The meaning of *bbox_to_anchor* and *bbox_to_transform* is interpreted + differently from that of legend. The value of bbox_to_anchor + (or the return value of its get_points method; the default is + *parent_axes.bbox*) is transformed by the bbox_transform (the default + is Identity transform) and then interpreted as points in the pixel + coordinate (which is dpi dependent). + + Thus, following three calls are identical and creates an inset axes + with respect to the *parent_axes*:: + + axins = inset_axes(parent_axes, "30%%", "40%%") + axins = inset_axes(parent_axes, "30%%", "40%%", + bbox_to_anchor=parent_axes.bbox) + axins = inset_axes(parent_axes, "30%%", "40%%", + bbox_to_anchor=(0, 0, 1, 1), + bbox_transform=parent_axes.transAxes) + + Parameters + ---------- + parent_axes : `matplotlib.axes.Axes` + Axes to place the inset axes. + + width, height : float or str + Size of the inset axes to create. If a float is provided, it is + the size in inches, e.g. *width=1.3*. If a string is provided, it is + the size in relative units, e.g. *width='40%%'*. By default, i.e. if + neither *bbox_to_anchor* nor *bbox_transform* are specified, those + are relative to the parent_axes. Otherwise, they are to be understood + relative to the bounding box provided via *bbox_to_anchor*. + + loc : str, default: 'upper right' + Location to place the inset axes. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + + bbox_to_anchor : tuple or `~matplotlib.transforms.BboxBase`, optional + Bbox that the inset axes will be anchored to. If None, + a tuple of (0, 0, 1, 1) is used if *bbox_transform* is set + to *parent_axes.transAxes* or *parent_axes.figure.transFigure*. + Otherwise, *parent_axes.bbox* is used. If a tuple, can be either + [left, bottom, width, height], or [left, bottom]. + If the kwargs *width* and/or *height* are specified in relative units, + the 2-tuple [left, bottom] cannot be used. Note that, + unless *bbox_transform* is set, the units of the bounding box + are interpreted in the pixel coordinate. When using *bbox_to_anchor* + with tuple, it almost always makes sense to also specify + a *bbox_transform*. This might often be the axes transform + *parent_axes.transAxes*. + + bbox_transform : `~matplotlib.transforms.Transform`, optional + Transformation for the bbox that contains the inset axes. + If None, a `.transforms.IdentityTransform` is used. The value + of *bbox_to_anchor* (or the return value of its get_points method) + is transformed by the *bbox_transform* and then interpreted + as points in the pixel coordinate (which is dpi dependent). + You may provide *bbox_to_anchor* in some normalized coordinate, + and give an appropriate transform (e.g., *parent_axes.transAxes*). + + axes_class : `~matplotlib.axes.Axes` type, default: `.HostAxes` + The type of the newly created inset axes. + + axes_kwargs : dict, optional + Keyword arguments to pass to the constructor of the inset axes. + Valid arguments include: + + %(Axes:kwdoc)s + + borderpad : float, default: 0.5 + Padding between inset axes and the bbox_to_anchor. + The units are axes font size, i.e. for a default font size of 10 points + *borderpad = 0.5* is equivalent to a padding of 5 points. + + Returns + ------- + inset_axes : *axes_class* + Inset axes object created. + """ + + if (bbox_transform in [parent_axes.transAxes, + parent_axes.get_figure(root=False).transFigure] + and bbox_to_anchor is None): + _api.warn_external("Using the axes or figure transform requires a " + "bounding box in the respective coordinates. " + "Using bbox_to_anchor=(0, 0, 1, 1) now.") + bbox_to_anchor = (0, 0, 1, 1) + if bbox_to_anchor is None: + bbox_to_anchor = parent_axes.bbox + if (isinstance(bbox_to_anchor, tuple) and + (isinstance(width, str) or isinstance(height, str))): + if len(bbox_to_anchor) != 4: + raise ValueError("Using relative units for width or height " + "requires to provide a 4-tuple or a " + "`Bbox` instance to `bbox_to_anchor.") + return _add_inset_axes( + parent_axes, axes_class, axes_kwargs, + AnchoredSizeLocator( + bbox_to_anchor, width, height, loc=loc, + bbox_transform=bbox_transform, borderpad=borderpad)) + + +@_docstring.interpd +def zoomed_inset_axes(parent_axes, zoom, loc='upper right', + bbox_to_anchor=None, bbox_transform=None, + axes_class=None, axes_kwargs=None, + borderpad=0.5): + """ + Create an anchored inset axes by scaling a parent axes. For usage, also see + :doc:`the examples `. + + Parameters + ---------- + parent_axes : `~matplotlib.axes.Axes` + Axes to place the inset axes. + + zoom : float + Scaling factor of the data axes. *zoom* > 1 will enlarge the + coordinates (i.e., "zoomed in"), while *zoom* < 1 will shrink the + coordinates (i.e., "zoomed out"). + + loc : str, default: 'upper right' + Location to place the inset axes. Valid locations are + 'upper left', 'upper center', 'upper right', + 'center left', 'center', 'center right', + 'lower left', 'lower center', 'lower right'. + For backward compatibility, numeric values are accepted as well. + See the parameter *loc* of `.Legend` for details. + + bbox_to_anchor : tuple or `~matplotlib.transforms.BboxBase`, optional + Bbox that the inset axes will be anchored to. If None, + *parent_axes.bbox* is used. If a tuple, can be either + [left, bottom, width, height], or [left, bottom]. + If the kwargs *width* and/or *height* are specified in relative units, + the 2-tuple [left, bottom] cannot be used. Note that + the units of the bounding box are determined through the transform + in use. When using *bbox_to_anchor* it almost always makes sense to + also specify a *bbox_transform*. This might often be the axes transform + *parent_axes.transAxes*. + + bbox_transform : `~matplotlib.transforms.Transform`, optional + Transformation for the bbox that contains the inset axes. + If None, a `.transforms.IdentityTransform` is used (i.e. pixel + coordinates). This is useful when not providing any argument to + *bbox_to_anchor*. When using *bbox_to_anchor* it almost always makes + sense to also specify a *bbox_transform*. This might often be the + axes transform *parent_axes.transAxes*. Inversely, when specifying + the axes- or figure-transform here, be aware that not specifying + *bbox_to_anchor* will use *parent_axes.bbox*, the units of which are + in display (pixel) coordinates. + + axes_class : `~matplotlib.axes.Axes` type, default: `.HostAxes` + The type of the newly created inset axes. + + axes_kwargs : dict, optional + Keyword arguments to pass to the constructor of the inset axes. + Valid arguments include: + + %(Axes:kwdoc)s + + borderpad : float, default: 0.5 + Padding between inset axes and the bbox_to_anchor. + The units are axes font size, i.e. for a default font size of 10 points + *borderpad = 0.5* is equivalent to a padding of 5 points. + + Returns + ------- + inset_axes : *axes_class* + Inset axes object created. + """ + + return _add_inset_axes( + parent_axes, axes_class, axes_kwargs, + AnchoredZoomLocator( + parent_axes, zoom=zoom, loc=loc, + bbox_to_anchor=bbox_to_anchor, bbox_transform=bbox_transform, + borderpad=borderpad)) + + +class _TransformedBboxWithCallback(TransformedBbox): + """ + Variant of `.TransformBbox` which calls *callback* before returning points. + + Used by `.mark_inset` to unstale the parent axes' viewlim as needed. + """ + + def __init__(self, *args, callback, **kwargs): + super().__init__(*args, **kwargs) + self._callback = callback + + def get_points(self): + self._callback() + return super().get_points() + + +@_docstring.interpd +def mark_inset(parent_axes, inset_axes, loc1, loc2, **kwargs): + """ + Draw a box to mark the location of an area represented by an inset axes. + + This function draws a box in *parent_axes* at the bounding box of + *inset_axes*, and shows a connection with the inset axes by drawing lines + at the corners, giving a "zoomed in" effect. + + Parameters + ---------- + parent_axes : `~matplotlib.axes.Axes` + Axes which contains the area of the inset axes. + + inset_axes : `~matplotlib.axes.Axes` + The inset axes. + + loc1, loc2 : {1, 2, 3, 4} + Corners to use for connecting the inset axes and the area in the + parent axes. + + **kwargs + Patch properties for the lines and box drawn: + + %(Patch:kwdoc)s + + Returns + ------- + pp : `~matplotlib.patches.Patch` + The patch drawn to represent the area of the inset axes. + + p1, p2 : `~matplotlib.patches.Patch` + The patches connecting two corners of the inset axes and its area. + """ + rect = _TransformedBboxWithCallback( + inset_axes.viewLim, parent_axes.transData, + callback=parent_axes._unstale_viewLim) + + kwargs.setdefault("fill", bool({'fc', 'facecolor', 'color'}.intersection(kwargs))) + pp = BboxPatch(rect, **kwargs) + parent_axes.add_patch(pp) + + p1 = BboxConnector(inset_axes.bbox, rect, loc1=loc1, **kwargs) + inset_axes.add_patch(p1) + p1.set_clip_on(False) + p2 = BboxConnector(inset_axes.bbox, rect, loc1=loc2, **kwargs) + inset_axes.add_patch(p2) + p2.set_clip_on(False) + + return pp, p1, p2 diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/mpl_axes.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/mpl_axes.py new file mode 100644 index 0000000000000000000000000000000000000000..51c8748758cb6da3052f0e8b05ceba427d77a3f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/mpl_axes.py @@ -0,0 +1,128 @@ +import matplotlib.axes as maxes +from matplotlib.artist import Artist +from matplotlib.axis import XAxis, YAxis + + +class SimpleChainedObjects: + def __init__(self, objects): + self._objects = objects + + def __getattr__(self, k): + _a = SimpleChainedObjects([getattr(a, k) for a in self._objects]) + return _a + + def __call__(self, *args, **kwargs): + for m in self._objects: + m(*args, **kwargs) + + +class Axes(maxes.Axes): + + class AxisDict(dict): + def __init__(self, axes): + self.axes = axes + super().__init__() + + def __getitem__(self, k): + if isinstance(k, tuple): + r = SimpleChainedObjects( + # super() within a list comprehension needs explicit args. + [super(Axes.AxisDict, self).__getitem__(k1) for k1 in k]) + return r + elif isinstance(k, slice): + if k.start is None and k.stop is None and k.step is None: + return SimpleChainedObjects(list(self.values())) + else: + raise ValueError("Unsupported slice") + else: + return dict.__getitem__(self, k) + + def __call__(self, *v, **kwargs): + return maxes.Axes.axis(self.axes, *v, **kwargs) + + @property + def axis(self): + return self._axislines + + def clear(self): + # docstring inherited + super().clear() + # Init axis artists. + self._axislines = self.AxisDict(self) + self._axislines.update( + bottom=SimpleAxisArtist(self.xaxis, 1, self.spines["bottom"]), + top=SimpleAxisArtist(self.xaxis, 2, self.spines["top"]), + left=SimpleAxisArtist(self.yaxis, 1, self.spines["left"]), + right=SimpleAxisArtist(self.yaxis, 2, self.spines["right"])) + + +class SimpleAxisArtist(Artist): + def __init__(self, axis, axisnum, spine): + self._axis = axis + self._axisnum = axisnum + self.line = spine + + if isinstance(axis, XAxis): + self._axis_direction = ["bottom", "top"][axisnum-1] + elif isinstance(axis, YAxis): + self._axis_direction = ["left", "right"][axisnum-1] + else: + raise ValueError( + f"axis must be instance of XAxis or YAxis, but got {axis}") + super().__init__() + + @property + def major_ticks(self): + tickline = "tick%dline" % self._axisnum + return SimpleChainedObjects([getattr(tick, tickline) + for tick in self._axis.get_major_ticks()]) + + @property + def major_ticklabels(self): + label = "label%d" % self._axisnum + return SimpleChainedObjects([getattr(tick, label) + for tick in self._axis.get_major_ticks()]) + + @property + def label(self): + return self._axis.label + + def set_visible(self, b): + self.toggle(all=b) + self.line.set_visible(b) + self._axis.set_visible(True) + super().set_visible(b) + + def set_label(self, txt): + self._axis.set_label_text(txt) + + def toggle(self, all=None, ticks=None, ticklabels=None, label=None): + + if all: + _ticks, _ticklabels, _label = True, True, True + elif all is not None: + _ticks, _ticklabels, _label = False, False, False + else: + _ticks, _ticklabels, _label = None, None, None + + if ticks is not None: + _ticks = ticks + if ticklabels is not None: + _ticklabels = ticklabels + if label is not None: + _label = label + + if _ticks is not None: + tickparam = {f"tick{self._axisnum}On": _ticks} + self._axis.set_tick_params(**tickparam) + if _ticklabels is not None: + tickparam = {f"label{self._axisnum}On": _ticklabels} + self._axis.set_tick_params(**tickparam) + + if _label is not None: + pos = self._axis.get_label_position() + if (pos == self._axis_direction) and not _label: + self._axis.label.set_visible(False) + elif _label: + self._axis.label.set_visible(True) + self._axis.set_label_position(self._axis_direction) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/parasite_axes.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/parasite_axes.py new file mode 100644 index 0000000000000000000000000000000000000000..f7bc2df6d7e01141254114065316a70df5f35e33 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/parasite_axes.py @@ -0,0 +1,257 @@ +from matplotlib import _api, cbook +import matplotlib.artist as martist +import matplotlib.transforms as mtransforms +from matplotlib.transforms import Bbox +from .mpl_axes import Axes + + +class ParasiteAxesBase: + + def __init__(self, parent_axes, aux_transform=None, + *, viewlim_mode=None, **kwargs): + self._parent_axes = parent_axes + self.transAux = aux_transform + self.set_viewlim_mode(viewlim_mode) + kwargs["frameon"] = False + super().__init__(parent_axes.get_figure(root=False), + parent_axes._position, **kwargs) + + def clear(self): + super().clear() + martist.setp(self.get_children(), visible=False) + self._get_lines = self._parent_axes._get_lines + self._parent_axes.callbacks._connect_picklable( + "xlim_changed", self._sync_lims) + self._parent_axes.callbacks._connect_picklable( + "ylim_changed", self._sync_lims) + + def pick(self, mouseevent): + # This most likely goes to Artist.pick (depending on axes_class given + # to the factory), which only handles pick events registered on the + # axes associated with each child: + super().pick(mouseevent) + # But parasite axes are additionally given pick events from their host + # axes (cf. HostAxesBase.pick), which we handle here: + for a in self.get_children(): + if (hasattr(mouseevent.inaxes, "parasites") + and self in mouseevent.inaxes.parasites): + a.pick(mouseevent) + + # aux_transform support + + def _set_lim_and_transforms(self): + if self.transAux is not None: + self.transAxes = self._parent_axes.transAxes + self.transData = self.transAux + self._parent_axes.transData + self._xaxis_transform = mtransforms.blended_transform_factory( + self.transData, self.transAxes) + self._yaxis_transform = mtransforms.blended_transform_factory( + self.transAxes, self.transData) + else: + super()._set_lim_and_transforms() + + def set_viewlim_mode(self, mode): + _api.check_in_list([None, "equal", "transform"], mode=mode) + self._viewlim_mode = mode + + def get_viewlim_mode(self): + return self._viewlim_mode + + def _sync_lims(self, parent): + viewlim = parent.viewLim.frozen() + mode = self.get_viewlim_mode() + if mode is None: + pass + elif mode == "equal": + self.viewLim.set(viewlim) + elif mode == "transform": + self.viewLim.set(viewlim.transformed(self.transAux.inverted())) + else: + _api.check_in_list([None, "equal", "transform"], mode=mode) + + # end of aux_transform support + + +parasite_axes_class_factory = cbook._make_class_factory( + ParasiteAxesBase, "{}Parasite") +ParasiteAxes = parasite_axes_class_factory(Axes) + + +class HostAxesBase: + def __init__(self, *args, **kwargs): + self.parasites = [] + super().__init__(*args, **kwargs) + + def get_aux_axes( + self, tr=None, viewlim_mode="equal", axes_class=None, **kwargs): + """ + Add a parasite axes to this host. + + Despite this method's name, this should actually be thought of as an + ``add_parasite_axes`` method. + + .. versionchanged:: 3.7 + Defaults to same base axes class as host axes. + + Parameters + ---------- + tr : `~matplotlib.transforms.Transform` or None, default: None + If a `.Transform`, the following relation will hold: + ``parasite.transData = tr + host.transData``. + If None, the parasite's and the host's ``transData`` are unrelated. + viewlim_mode : {"equal", "transform", None}, default: "equal" + How the parasite's view limits are set: directly equal to the + parent axes ("equal"), equal after application of *tr* + ("transform"), or independently (None). + axes_class : subclass type of `~matplotlib.axes.Axes`, optional + The `~.axes.Axes` subclass that is instantiated. If None, the base + class of the host axes is used. + **kwargs + Other parameters are forwarded to the parasite axes constructor. + """ + if axes_class is None: + axes_class = self._base_axes_class + parasite_axes_class = parasite_axes_class_factory(axes_class) + ax2 = parasite_axes_class( + self, tr, viewlim_mode=viewlim_mode, **kwargs) + # note that ax2.transData == tr + ax1.transData + # Anything you draw in ax2 will match the ticks and grids of ax1. + self.parasites.append(ax2) + ax2._remove_method = self.parasites.remove + return ax2 + + def draw(self, renderer): + orig_children_len = len(self._children) + + locator = self.get_axes_locator() + if locator: + pos = locator(self, renderer) + self.set_position(pos, which="active") + self.apply_aspect(pos) + else: + self.apply_aspect() + + rect = self.get_position() + for ax in self.parasites: + ax.apply_aspect(rect) + self._children.extend(ax.get_children()) + + super().draw(renderer) + del self._children[orig_children_len:] + + def clear(self): + super().clear() + for ax in self.parasites: + ax.clear() + + def pick(self, mouseevent): + super().pick(mouseevent) + # Also pass pick events on to parasite axes and, in turn, their + # children (cf. ParasiteAxesBase.pick) + for a in self.parasites: + a.pick(mouseevent) + + def twinx(self, axes_class=None): + """ + Create a twin of Axes with a shared x-axis but independent y-axis. + + The y-axis of self will have ticks on the left and the returned axes + will have ticks on the right. + """ + ax = self._add_twin_axes(axes_class, sharex=self) + self.axis["right"].set_visible(False) + ax.axis["right"].set_visible(True) + ax.axis["left", "top", "bottom"].set_visible(False) + return ax + + def twiny(self, axes_class=None): + """ + Create a twin of Axes with a shared y-axis but independent x-axis. + + The x-axis of self will have ticks on the bottom and the returned axes + will have ticks on the top. + """ + ax = self._add_twin_axes(axes_class, sharey=self) + self.axis["top"].set_visible(False) + ax.axis["top"].set_visible(True) + ax.axis["left", "right", "bottom"].set_visible(False) + return ax + + def twin(self, aux_trans=None, axes_class=None): + """ + Create a twin of Axes with no shared axis. + + While self will have ticks on the left and bottom axis, the returned + axes will have ticks on the top and right axis. + """ + if aux_trans is None: + aux_trans = mtransforms.IdentityTransform() + ax = self._add_twin_axes( + axes_class, aux_transform=aux_trans, viewlim_mode="transform") + self.axis["top", "right"].set_visible(False) + ax.axis["top", "right"].set_visible(True) + ax.axis["left", "bottom"].set_visible(False) + return ax + + def _add_twin_axes(self, axes_class, **kwargs): + """ + Helper for `.twinx`/`.twiny`/`.twin`. + + *kwargs* are forwarded to the parasite axes constructor. + """ + if axes_class is None: + axes_class = self._base_axes_class + ax = parasite_axes_class_factory(axes_class)(self, **kwargs) + self.parasites.append(ax) + ax._remove_method = self._remove_any_twin + return ax + + def _remove_any_twin(self, ax): + self.parasites.remove(ax) + restore = ["top", "right"] + if ax._sharex: + restore.remove("top") + if ax._sharey: + restore.remove("right") + self.axis[tuple(restore)].set_visible(True) + self.axis[tuple(restore)].toggle(ticklabels=False, label=False) + + def get_tightbbox(self, renderer=None, *, call_axes_locator=True, + bbox_extra_artists=None): + bbs = [ + *[ax.get_tightbbox(renderer, call_axes_locator=call_axes_locator) + for ax in self.parasites], + super().get_tightbbox(renderer, + call_axes_locator=call_axes_locator, + bbox_extra_artists=bbox_extra_artists)] + return Bbox.union([b for b in bbs if b.width != 0 or b.height != 0]) + + +host_axes_class_factory = host_subplot_class_factory = \ + cbook._make_class_factory(HostAxesBase, "{}HostAxes", "_base_axes_class") +HostAxes = SubplotHost = host_axes_class_factory(Axes) + + +def host_axes(*args, axes_class=Axes, figure=None, **kwargs): + """ + Create axes that can act as a hosts to parasitic axes. + + Parameters + ---------- + figure : `~matplotlib.figure.Figure` + Figure to which the axes will be added. Defaults to the current figure + `.pyplot.gcf()`. + + *args, **kwargs + Will be passed on to the underlying `~.axes.Axes` object creation. + """ + import matplotlib.pyplot as plt + host_axes_class = host_axes_class_factory(axes_class) + if figure is None: + figure = plt.gcf() + ax = host_axes_class(figure, *args, **kwargs) + figure.add_axes(ax) + return ax + + +host_subplot = host_axes diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__init__.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ea4d8ed16a6a24a8c15ab2956ef678a7f256cd80 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__init__.py @@ -0,0 +1,10 @@ +from pathlib import Path + + +# Check that the test directories exist +if not (Path(__file__).parent / "baseline_images").exists(): + raise OSError( + 'The baseline image directory does not exist. ' + 'This is most likely because the test data is not installed. ' + 'You may need to install matplotlib from source to get the ' + 'test data.') diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8250377ee9651e22f3c67c5159491d79ba709f1b Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/conftest.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/conftest.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13d80330752b7a0276bb2b8962bb77b43e9a5abe Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/conftest.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/test_axes_grid1.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/test_axes_grid1.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fcd69f2d8f70736917dfefa6d5e33b10e7c45e57 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/__pycache__/test_axes_grid1.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/conftest.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..61c2de3e07bac4db323f8704961264d123e01544 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/conftest.py @@ -0,0 +1,2 @@ +from matplotlib.testing.conftest import (mpl_test_settings, # noqa + pytest_configure, pytest_unconfigure) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/test_axes_grid1.py b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/test_axes_grid1.py new file mode 100644 index 0000000000000000000000000000000000000000..b045e8dc87569b0ea5653cf8b0c6307135b69c75 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/axes_grid1/tests/test_axes_grid1.py @@ -0,0 +1,782 @@ +from itertools import product +import io +import platform + +import matplotlib as mpl +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker +from matplotlib import cbook +from matplotlib.backend_bases import MouseEvent +from matplotlib.colors import LogNorm +from matplotlib.patches import Circle, Ellipse +from matplotlib.transforms import Bbox, TransformedBbox +from matplotlib.testing.decorators import ( + check_figures_equal, image_comparison, remove_ticks_and_titles) + +from mpl_toolkits.axes_grid1 import ( + axes_size as Size, + host_subplot, make_axes_locatable, + Grid, AxesGrid, ImageGrid) +from mpl_toolkits.axes_grid1.anchored_artists import ( + AnchoredAuxTransformBox, AnchoredDrawingArea, + AnchoredDirectionArrows, AnchoredSizeBar) +from mpl_toolkits.axes_grid1.axes_divider import ( + Divider, HBoxDivider, make_axes_area_auto_adjustable, SubplotDivider, + VBoxDivider) +from mpl_toolkits.axes_grid1.axes_rgb import RGBAxes +from mpl_toolkits.axes_grid1.inset_locator import ( + zoomed_inset_axes, mark_inset, inset_axes, BboxConnectorPatch) +import mpl_toolkits.axes_grid1.mpl_axes +import pytest + +import numpy as np +from numpy.testing import assert_array_equal, assert_array_almost_equal + + +def test_divider_append_axes(): + fig, ax = plt.subplots() + divider = make_axes_locatable(ax) + axs = { + "main": ax, + "top": divider.append_axes("top", 1.2, pad=0.1, sharex=ax), + "bottom": divider.append_axes("bottom", 1.2, pad=0.1, sharex=ax), + "left": divider.append_axes("left", 1.2, pad=0.1, sharey=ax), + "right": divider.append_axes("right", 1.2, pad=0.1, sharey=ax), + } + fig.canvas.draw() + bboxes = {k: axs[k].get_window_extent() for k in axs} + dpi = fig.dpi + assert bboxes["top"].height == pytest.approx(1.2 * dpi) + assert bboxes["bottom"].height == pytest.approx(1.2 * dpi) + assert bboxes["left"].width == pytest.approx(1.2 * dpi) + assert bboxes["right"].width == pytest.approx(1.2 * dpi) + assert bboxes["top"].y0 - bboxes["main"].y1 == pytest.approx(0.1 * dpi) + assert bboxes["main"].y0 - bboxes["bottom"].y1 == pytest.approx(0.1 * dpi) + assert bboxes["main"].x0 - bboxes["left"].x1 == pytest.approx(0.1 * dpi) + assert bboxes["right"].x0 - bboxes["main"].x1 == pytest.approx(0.1 * dpi) + assert bboxes["left"].y0 == bboxes["main"].y0 == bboxes["right"].y0 + assert bboxes["left"].y1 == bboxes["main"].y1 == bboxes["right"].y1 + assert bboxes["top"].x0 == bboxes["main"].x0 == bboxes["bottom"].x0 + assert bboxes["top"].x1 == bboxes["main"].x1 == bboxes["bottom"].x1 + + +# Update style when regenerating the test image +@image_comparison(['twin_axes_empty_and_removed'], extensions=["png"], tol=1, + style=('classic', '_classic_test_patch')) +def test_twin_axes_empty_and_removed(): + # Purely cosmetic font changes (avoid overlap) + mpl.rcParams.update( + {"font.size": 8, "xtick.labelsize": 8, "ytick.labelsize": 8}) + generators = ["twinx", "twiny", "twin"] + modifiers = ["", "host invisible", "twin removed", "twin invisible", + "twin removed\nhost invisible"] + # Unmodified host subplot at the beginning for reference + h = host_subplot(len(modifiers)+1, len(generators), 2) + h.text(0.5, 0.5, "host_subplot", + horizontalalignment="center", verticalalignment="center") + # Host subplots with various modifications (twin*, visibility) applied + for i, (mod, gen) in enumerate(product(modifiers, generators), + len(generators) + 1): + h = host_subplot(len(modifiers)+1, len(generators), i) + t = getattr(h, gen)() + if "twin invisible" in mod: + t.axis[:].set_visible(False) + if "twin removed" in mod: + t.remove() + if "host invisible" in mod: + h.axis[:].set_visible(False) + h.text(0.5, 0.5, gen + ("\n" + mod if mod else ""), + horizontalalignment="center", verticalalignment="center") + plt.subplots_adjust(wspace=0.5, hspace=1) + + +def test_twin_axes_both_with_units(): + host = host_subplot(111) + with pytest.warns(mpl.MatplotlibDeprecationWarning): + host.plot_date([0, 1, 2], [0, 1, 2], xdate=False, ydate=True) + twin = host.twinx() + twin.plot(["a", "b", "c"]) + assert host.get_yticklabels()[0].get_text() == "00:00:00" + assert twin.get_yticklabels()[0].get_text() == "a" + + +def test_axesgrid_colorbar_log_smoketest(): + fig = plt.figure() + grid = AxesGrid(fig, 111, # modified to be only subplot + nrows_ncols=(1, 1), + ngrids=1, + label_mode="L", + cbar_location="top", + cbar_mode="single", + ) + + Z = 10000 * np.random.rand(10, 10) + im = grid[0].imshow(Z, interpolation="nearest", norm=LogNorm()) + + grid.cbar_axes[0].colorbar(im) + + +def test_inset_colorbar_tight_layout_smoketest(): + fig, ax = plt.subplots(1, 1) + pts = ax.scatter([0, 1], [0, 1], c=[1, 5]) + + cax = inset_axes(ax, width="3%", height="70%") + plt.colorbar(pts, cax=cax) + + with pytest.warns(UserWarning, match="This figure includes Axes"): + # Will warn, but not raise an error + plt.tight_layout() + + +@image_comparison(['inset_locator.png'], style='default', remove_text=True) +def test_inset_locator(): + fig, ax = plt.subplots(figsize=[5, 4]) + + # prepare the demo image + # Z is a 15x15 array + Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy") + extent = (-3, 4, -4, 3) + Z2 = np.zeros((150, 150)) + ny, nx = Z.shape + Z2[30:30+ny, 30:30+nx] = Z + + ax.imshow(Z2, extent=extent, interpolation="nearest", + origin="lower") + + axins = zoomed_inset_axes(ax, zoom=6, loc='upper right') + axins.imshow(Z2, extent=extent, interpolation="nearest", + origin="lower") + axins.yaxis.get_major_locator().set_params(nbins=7) + axins.xaxis.get_major_locator().set_params(nbins=7) + # sub region of the original image + x1, x2, y1, y2 = -1.5, -0.9, -2.5, -1.9 + axins.set_xlim(x1, x2) + axins.set_ylim(y1, y2) + + plt.xticks(visible=False) + plt.yticks(visible=False) + + # draw a bbox of the region of the inset axes in the parent axes and + # connecting lines between the bbox and the inset axes area + mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec="0.5") + + asb = AnchoredSizeBar(ax.transData, + 0.5, + '0.5', + loc='lower center', + pad=0.1, borderpad=0.5, sep=5, + frameon=False) + ax.add_artist(asb) + + +@image_comparison(['inset_axes.png'], style='default', remove_text=True) +def test_inset_axes(): + fig, ax = plt.subplots(figsize=[5, 4]) + + # prepare the demo image + # Z is a 15x15 array + Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy") + extent = (-3, 4, -4, 3) + Z2 = np.zeros((150, 150)) + ny, nx = Z.shape + Z2[30:30+ny, 30:30+nx] = Z + + ax.imshow(Z2, extent=extent, interpolation="nearest", + origin="lower") + + # creating our inset axes with a bbox_transform parameter + axins = inset_axes(ax, width=1., height=1., bbox_to_anchor=(1, 1), + bbox_transform=ax.transAxes) + + axins.imshow(Z2, extent=extent, interpolation="nearest", + origin="lower") + axins.yaxis.get_major_locator().set_params(nbins=7) + axins.xaxis.get_major_locator().set_params(nbins=7) + # sub region of the original image + x1, x2, y1, y2 = -1.5, -0.9, -2.5, -1.9 + axins.set_xlim(x1, x2) + axins.set_ylim(y1, y2) + + plt.xticks(visible=False) + plt.yticks(visible=False) + + # draw a bbox of the region of the inset axes in the parent axes and + # connecting lines between the bbox and the inset axes area + mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec="0.5") + + asb = AnchoredSizeBar(ax.transData, + 0.5, + '0.5', + loc='lower center', + pad=0.1, borderpad=0.5, sep=5, + frameon=False) + ax.add_artist(asb) + + +def test_inset_axes_complete(): + dpi = 100 + figsize = (6, 5) + fig, ax = plt.subplots(figsize=figsize, dpi=dpi) + fig.subplots_adjust(.1, .1, .9, .9) + + ins = inset_axes(ax, width=2., height=2., borderpad=0) + fig.canvas.draw() + assert_array_almost_equal( + ins.get_position().extents, + [(0.9*figsize[0]-2.)/figsize[0], (0.9*figsize[1]-2.)/figsize[1], + 0.9, 0.9]) + + ins = inset_axes(ax, width="40%", height="30%", borderpad=0) + fig.canvas.draw() + assert_array_almost_equal( + ins.get_position().extents, [.9-.8*.4, .9-.8*.3, 0.9, 0.9]) + + ins = inset_axes(ax, width=1., height=1.2, bbox_to_anchor=(200, 100), + loc=3, borderpad=0) + fig.canvas.draw() + assert_array_almost_equal( + ins.get_position().extents, + [200/dpi/figsize[0], 100/dpi/figsize[1], + (200/dpi+1)/figsize[0], (100/dpi+1.2)/figsize[1]]) + + ins1 = inset_axes(ax, width="35%", height="60%", loc=3, borderpad=1) + ins2 = inset_axes(ax, width="100%", height="100%", + bbox_to_anchor=(0, 0, .35, .60), + bbox_transform=ax.transAxes, loc=3, borderpad=1) + fig.canvas.draw() + assert_array_equal(ins1.get_position().extents, + ins2.get_position().extents) + + with pytest.raises(ValueError): + ins = inset_axes(ax, width="40%", height="30%", + bbox_to_anchor=(0.4, 0.5)) + + with pytest.warns(UserWarning): + ins = inset_axes(ax, width="40%", height="30%", + bbox_transform=ax.transAxes) + + +def test_inset_axes_tight(): + # gh-26287 found that inset_axes raised with bbox_inches=tight + fig, ax = plt.subplots() + inset_axes(ax, width=1.3, height=0.9) + + f = io.BytesIO() + fig.savefig(f, bbox_inches="tight") + + +@image_comparison(['fill_facecolor.png'], remove_text=True, style='mpl20') +def test_fill_facecolor(): + fig, ax = plt.subplots(1, 5) + fig.set_size_inches(5, 5) + for i in range(1, 4): + ax[i].yaxis.set_visible(False) + ax[4].yaxis.tick_right() + bbox = Bbox.from_extents(0, 0.4, 1, 0.6) + + # fill with blue by setting 'fc' field + bbox1 = TransformedBbox(bbox, ax[0].transData) + bbox2 = TransformedBbox(bbox, ax[1].transData) + # set color to BboxConnectorPatch + p = BboxConnectorPatch( + bbox1, bbox2, loc1a=1, loc2a=2, loc1b=4, loc2b=3, + ec="r", fc="b") + p.set_clip_on(False) + ax[0].add_patch(p) + # set color to marked area + axins = zoomed_inset_axes(ax[0], 1, loc='upper right') + axins.set_xlim(0, 0.2) + axins.set_ylim(0, 0.2) + plt.gca().axes.xaxis.set_ticks([]) + plt.gca().axes.yaxis.set_ticks([]) + mark_inset(ax[0], axins, loc1=2, loc2=4, fc="b", ec="0.5") + + # fill with yellow by setting 'facecolor' field + bbox3 = TransformedBbox(bbox, ax[1].transData) + bbox4 = TransformedBbox(bbox, ax[2].transData) + # set color to BboxConnectorPatch + p = BboxConnectorPatch( + bbox3, bbox4, loc1a=1, loc2a=2, loc1b=4, loc2b=3, + ec="r", facecolor="y") + p.set_clip_on(False) + ax[1].add_patch(p) + # set color to marked area + axins = zoomed_inset_axes(ax[1], 1, loc='upper right') + axins.set_xlim(0, 0.2) + axins.set_ylim(0, 0.2) + plt.gca().axes.xaxis.set_ticks([]) + plt.gca().axes.yaxis.set_ticks([]) + mark_inset(ax[1], axins, loc1=2, loc2=4, facecolor="y", ec="0.5") + + # fill with green by setting 'color' field + bbox5 = TransformedBbox(bbox, ax[2].transData) + bbox6 = TransformedBbox(bbox, ax[3].transData) + # set color to BboxConnectorPatch + p = BboxConnectorPatch( + bbox5, bbox6, loc1a=1, loc2a=2, loc1b=4, loc2b=3, + ec="r", color="g") + p.set_clip_on(False) + ax[2].add_patch(p) + # set color to marked area + axins = zoomed_inset_axes(ax[2], 1, loc='upper right') + axins.set_xlim(0, 0.2) + axins.set_ylim(0, 0.2) + plt.gca().axes.xaxis.set_ticks([]) + plt.gca().axes.yaxis.set_ticks([]) + mark_inset(ax[2], axins, loc1=2, loc2=4, color="g", ec="0.5") + + # fill with green but color won't show if set fill to False + bbox7 = TransformedBbox(bbox, ax[3].transData) + bbox8 = TransformedBbox(bbox, ax[4].transData) + # BboxConnectorPatch won't show green + p = BboxConnectorPatch( + bbox7, bbox8, loc1a=1, loc2a=2, loc1b=4, loc2b=3, + ec="r", fc="g", fill=False) + p.set_clip_on(False) + ax[3].add_patch(p) + # marked area won't show green + axins = zoomed_inset_axes(ax[3], 1, loc='upper right') + axins.set_xlim(0, 0.2) + axins.set_ylim(0, 0.2) + axins.xaxis.set_ticks([]) + axins.yaxis.set_ticks([]) + mark_inset(ax[3], axins, loc1=2, loc2=4, fc="g", ec="0.5", fill=False) + + +# Update style when regenerating the test image +@image_comparison(['zoomed_axes.png', 'inverted_zoomed_axes.png'], + style=('classic', '_classic_test_patch'), + tol=0 if platform.machine() == 'x86_64' else 0.02) +def test_zooming_with_inverted_axes(): + fig, ax = plt.subplots() + ax.plot([1, 2, 3], [1, 2, 3]) + ax.axis([1, 3, 1, 3]) + inset_ax = zoomed_inset_axes(ax, zoom=2.5, loc='lower right') + inset_ax.axis([1.1, 1.4, 1.1, 1.4]) + + fig, ax = plt.subplots() + ax.plot([1, 2, 3], [1, 2, 3]) + ax.axis([3, 1, 3, 1]) + inset_ax = zoomed_inset_axes(ax, zoom=2.5, loc='lower right') + inset_ax.axis([1.4, 1.1, 1.4, 1.1]) + + +# Update style when regenerating the test image +@image_comparison(['anchored_direction_arrows.png'], + tol=0 if platform.machine() == 'x86_64' else 0.01, + style=('classic', '_classic_test_patch')) +def test_anchored_direction_arrows(): + fig, ax = plt.subplots() + ax.imshow(np.zeros((10, 10)), interpolation='nearest') + + simple_arrow = AnchoredDirectionArrows(ax.transAxes, 'X', 'Y') + ax.add_artist(simple_arrow) + + +# Update style when regenerating the test image +@image_comparison(['anchored_direction_arrows_many_args.png'], + style=('classic', '_classic_test_patch')) +def test_anchored_direction_arrows_many_args(): + fig, ax = plt.subplots() + ax.imshow(np.ones((10, 10))) + + direction_arrows = AnchoredDirectionArrows( + ax.transAxes, 'A', 'B', loc='upper right', color='red', + aspect_ratio=-0.5, pad=0.6, borderpad=2, frameon=True, alpha=0.7, + sep_x=-0.06, sep_y=-0.08, back_length=0.1, head_width=9, + head_length=10, tail_width=5) + ax.add_artist(direction_arrows) + + +def test_axes_locatable_position(): + fig, ax = plt.subplots() + divider = make_axes_locatable(ax) + with mpl.rc_context({"figure.subplot.wspace": 0.02}): + cax = divider.append_axes('right', size='5%') + fig.canvas.draw() + assert np.isclose(cax.get_position(original=False).width, + 0.03621495327102808) + + +@image_comparison(['image_grid_each_left_label_mode_all.png'], style='mpl20', + savefig_kwarg={'bbox_inches': 'tight'}) +def test_image_grid_each_left_label_mode_all(): + imdata = np.arange(100).reshape((10, 10)) + + fig = plt.figure(1, (3, 3)) + grid = ImageGrid(fig, (1, 1, 1), nrows_ncols=(3, 2), axes_pad=(0.5, 0.3), + cbar_mode="each", cbar_location="left", cbar_size="15%", + label_mode="all") + # 3-tuple rect => SubplotDivider + assert isinstance(grid.get_divider(), SubplotDivider) + assert grid.get_axes_pad() == (0.5, 0.3) + assert grid.get_aspect() # True by default for ImageGrid + for ax, cax in zip(grid, grid.cbar_axes): + im = ax.imshow(imdata, interpolation='none') + cax.colorbar(im) + + +@image_comparison(['image_grid_single_bottom_label_mode_1.png'], style='mpl20', + savefig_kwarg={'bbox_inches': 'tight'}) +def test_image_grid_single_bottom(): + imdata = np.arange(100).reshape((10, 10)) + + fig = plt.figure(1, (2.5, 1.5)) + grid = ImageGrid(fig, (0, 0, 1, 1), nrows_ncols=(1, 3), + axes_pad=(0.2, 0.15), cbar_mode="single", cbar_pad=0.3, + cbar_location="bottom", cbar_size="10%", label_mode="1") + # 4-tuple rect => Divider, isinstance will give True for SubplotDivider + assert type(grid.get_divider()) is Divider + for i in range(3): + im = grid[i].imshow(imdata, interpolation='none') + grid.cbar_axes[0].colorbar(im) + + +def test_image_grid_label_mode_invalid(): + fig = plt.figure() + with pytest.raises(ValueError, match="'foo' is not a valid value for mode"): + ImageGrid(fig, (0, 0, 1, 1), (2, 1), label_mode="foo") + + +@image_comparison(['image_grid.png'], + remove_text=True, style='mpl20', + savefig_kwarg={'bbox_inches': 'tight'}) +def test_image_grid(): + # test that image grid works with bbox_inches=tight. + im = np.arange(100).reshape((10, 10)) + + fig = plt.figure(1, (4, 4)) + grid = ImageGrid(fig, 111, nrows_ncols=(2, 2), axes_pad=0.1) + assert grid.get_axes_pad() == (0.1, 0.1) + for i in range(4): + grid[i].imshow(im, interpolation='nearest') + + +def test_gettightbbox(): + fig, ax = plt.subplots(figsize=(8, 6)) + + l, = ax.plot([1, 2, 3], [0, 1, 0]) + + ax_zoom = zoomed_inset_axes(ax, 4) + ax_zoom.plot([1, 2, 3], [0, 1, 0]) + + mark_inset(ax, ax_zoom, loc1=1, loc2=3, fc="none", ec='0.3') + + remove_ticks_and_titles(fig) + bbox = fig.get_tightbbox(fig.canvas.get_renderer()) + np.testing.assert_array_almost_equal(bbox.extents, + [-17.7, -13.9, 7.2, 5.4]) + + +@pytest.mark.parametrize("click_on", ["big", "small"]) +@pytest.mark.parametrize("big_on_axes,small_on_axes", [ + ("gca", "gca"), + ("host", "host"), + ("host", "parasite"), + ("parasite", "host"), + ("parasite", "parasite") +]) +def test_picking_callbacks_overlap(big_on_axes, small_on_axes, click_on): + """Test pick events on normal, host or parasite axes.""" + # Two rectangles are drawn and "clicked on", a small one and a big one + # enclosing the small one. The axis on which they are drawn as well as the + # rectangle that is clicked on are varied. + # In each case we expect that both rectangles are picked if we click on the + # small one and only the big one is picked if we click on the big one. + # Also tests picking on normal axes ("gca") as a control. + big = plt.Rectangle((0.25, 0.25), 0.5, 0.5, picker=5) + small = plt.Rectangle((0.4, 0.4), 0.2, 0.2, facecolor="r", picker=5) + # Machinery for "receiving" events + received_events = [] + def on_pick(event): + received_events.append(event) + plt.gcf().canvas.mpl_connect('pick_event', on_pick) + # Shortcut + rectangles_on_axes = (big_on_axes, small_on_axes) + # Axes setup + axes = {"gca": None, "host": None, "parasite": None} + if "gca" in rectangles_on_axes: + axes["gca"] = plt.gca() + if "host" in rectangles_on_axes or "parasite" in rectangles_on_axes: + axes["host"] = host_subplot(111) + axes["parasite"] = axes["host"].twin() + # Add rectangles to axes + axes[big_on_axes].add_patch(big) + axes[small_on_axes].add_patch(small) + # Simulate picking with click mouse event + if click_on == "big": + click_axes = axes[big_on_axes] + axes_coords = (0.3, 0.3) + else: + click_axes = axes[small_on_axes] + axes_coords = (0.5, 0.5) + # In reality mouse events never happen on parasite axes, only host axes + if click_axes is axes["parasite"]: + click_axes = axes["host"] + (x, y) = click_axes.transAxes.transform(axes_coords) + m = MouseEvent("button_press_event", click_axes.get_figure(root=True).canvas, x, y, + button=1) + click_axes.pick(m) + # Checks + expected_n_events = 2 if click_on == "small" else 1 + assert len(received_events) == expected_n_events + event_rects = [event.artist for event in received_events] + assert big in event_rects + if click_on == "small": + assert small in event_rects + + +@image_comparison(['anchored_artists.png'], remove_text=True, style='mpl20') +def test_anchored_artists(): + fig, ax = plt.subplots(figsize=(3, 3)) + ada = AnchoredDrawingArea(40, 20, 0, 0, loc='upper right', pad=0., + frameon=False) + p1 = Circle((10, 10), 10) + ada.drawing_area.add_artist(p1) + p2 = Circle((30, 10), 5, fc="r") + ada.drawing_area.add_artist(p2) + ax.add_artist(ada) + + box = AnchoredAuxTransformBox(ax.transData, loc='upper left') + el = Ellipse((0, 0), width=0.1, height=0.4, angle=30, color='cyan') + box.drawing_area.add_artist(el) + ax.add_artist(box) + + # This block used to test the AnchoredEllipse class, but that was removed. The block + # remains, though it duplicates the above ellipse, so that the test image doesn't + # need to be regenerated. + box = AnchoredAuxTransformBox(ax.transData, loc='lower left', frameon=True, + pad=0.5, borderpad=0.4) + el = Ellipse((0, 0), width=0.1, height=0.25, angle=-60) + box.drawing_area.add_artist(el) + ax.add_artist(box) + + asb = AnchoredSizeBar(ax.transData, 0.2, r"0.2 units", loc='lower right', + pad=0.3, borderpad=0.4, sep=4, fill_bar=True, + frameon=False, label_top=True, prop={'size': 20}, + size_vertical=0.05, color='green') + ax.add_artist(asb) + + +def test_hbox_divider(): + arr1 = np.arange(20).reshape((4, 5)) + arr2 = np.arange(20).reshape((5, 4)) + + fig, (ax1, ax2) = plt.subplots(1, 2) + ax1.imshow(arr1) + ax2.imshow(arr2) + + pad = 0.5 # inches. + divider = HBoxDivider( + fig, 111, # Position of combined axes. + horizontal=[Size.AxesX(ax1), Size.Fixed(pad), Size.AxesX(ax2)], + vertical=[Size.AxesY(ax1), Size.Scaled(1), Size.AxesY(ax2)]) + ax1.set_axes_locator(divider.new_locator(0)) + ax2.set_axes_locator(divider.new_locator(2)) + + fig.canvas.draw() + p1 = ax1.get_position() + p2 = ax2.get_position() + assert p1.height == p2.height + assert p2.width / p1.width == pytest.approx((4 / 5) ** 2) + + +def test_vbox_divider(): + arr1 = np.arange(20).reshape((4, 5)) + arr2 = np.arange(20).reshape((5, 4)) + + fig, (ax1, ax2) = plt.subplots(1, 2) + ax1.imshow(arr1) + ax2.imshow(arr2) + + pad = 0.5 # inches. + divider = VBoxDivider( + fig, 111, # Position of combined axes. + horizontal=[Size.AxesX(ax1), Size.Scaled(1), Size.AxesX(ax2)], + vertical=[Size.AxesY(ax1), Size.Fixed(pad), Size.AxesY(ax2)]) + ax1.set_axes_locator(divider.new_locator(0)) + ax2.set_axes_locator(divider.new_locator(2)) + + fig.canvas.draw() + p1 = ax1.get_position() + p2 = ax2.get_position() + assert p1.width == p2.width + assert p1.height / p2.height == pytest.approx((4 / 5) ** 2) + + +def test_axes_class_tuple(): + fig = plt.figure() + axes_class = (mpl_toolkits.axes_grid1.mpl_axes.Axes, {}) + gr = AxesGrid(fig, 111, nrows_ncols=(1, 1), axes_class=axes_class) + + +def test_grid_axes_lists(): + """Test Grid axes_all, axes_row and axes_column relationship.""" + fig = plt.figure() + grid = Grid(fig, 111, (2, 3), direction="row") + assert_array_equal(grid, grid.axes_all) + assert_array_equal(grid.axes_row, np.transpose(grid.axes_column)) + assert_array_equal(grid, np.ravel(grid.axes_row), "row") + assert grid.get_geometry() == (2, 3) + grid = Grid(fig, 111, (2, 3), direction="column") + assert_array_equal(grid, np.ravel(grid.axes_column), "column") + + +@pytest.mark.parametrize('direction', ('row', 'column')) +def test_grid_axes_position(direction): + """Test positioning of the axes in Grid.""" + fig = plt.figure() + grid = Grid(fig, 111, (2, 2), direction=direction) + loc = [ax.get_axes_locator() for ax in np.ravel(grid.axes_row)] + # Test nx. + assert loc[1].args[0] > loc[0].args[0] + assert loc[0].args[0] == loc[2].args[0] + assert loc[3].args[0] == loc[1].args[0] + # Test ny. + assert loc[2].args[1] < loc[0].args[1] + assert loc[0].args[1] == loc[1].args[1] + assert loc[3].args[1] == loc[2].args[1] + + +@pytest.mark.parametrize('rect, ngrids, error, message', ( + ((1, 1), None, TypeError, "Incorrect rect format"), + (111, -1, ValueError, "ngrids must be positive"), + (111, 7, ValueError, "ngrids must be positive"), +)) +def test_grid_errors(rect, ngrids, error, message): + fig = plt.figure() + with pytest.raises(error, match=message): + Grid(fig, rect, (2, 3), ngrids=ngrids) + + +@pytest.mark.parametrize('anchor, error, message', ( + (None, TypeError, "anchor must be str"), + ("CC", ValueError, "'CC' is not a valid value for anchor"), + ((1, 1, 1), TypeError, "anchor must be str"), +)) +def test_divider_errors(anchor, error, message): + fig = plt.figure() + with pytest.raises(error, match=message): + Divider(fig, [0, 0, 1, 1], [Size.Fixed(1)], [Size.Fixed(1)], + anchor=anchor) + + +@check_figures_equal(extensions=["png"]) +def test_mark_inset_unstales_viewlim(fig_test, fig_ref): + inset, full = fig_test.subplots(1, 2) + full.plot([0, 5], [0, 5]) + inset.set(xlim=(1, 2), ylim=(1, 2)) + # Check that mark_inset unstales full's viewLim before drawing the marks. + mark_inset(full, inset, 1, 4) + + inset, full = fig_ref.subplots(1, 2) + full.plot([0, 5], [0, 5]) + inset.set(xlim=(1, 2), ylim=(1, 2)) + mark_inset(full, inset, 1, 4) + # Manually unstale the full's viewLim. + fig_ref.canvas.draw() + + +def test_auto_adjustable(): + fig = plt.figure() + ax = fig.add_axes([0, 0, 1, 1]) + pad = 0.1 + make_axes_area_auto_adjustable(ax, pad=pad) + fig.canvas.draw() + tbb = ax.get_tightbbox() + assert tbb.x0 == pytest.approx(pad * fig.dpi) + assert tbb.x1 == pytest.approx(fig.bbox.width - pad * fig.dpi) + assert tbb.y0 == pytest.approx(pad * fig.dpi) + assert tbb.y1 == pytest.approx(fig.bbox.height - pad * fig.dpi) + + +# Update style when regenerating the test image +@image_comparison(['rgb_axes.png'], remove_text=True, + style=('classic', '_classic_test_patch')) +def test_rgb_axes(): + fig = plt.figure() + ax = RGBAxes(fig, (0.1, 0.1, 0.8, 0.8), pad=0.1) + rng = np.random.default_rng(19680801) + r = rng.random((5, 5)) + g = rng.random((5, 5)) + b = rng.random((5, 5)) + ax.imshow_rgb(r, g, b, interpolation='none') + + +# The original version of this test relied on mpl_toolkits's slightly different +# colorbar implementation; moving to matplotlib's own colorbar implementation +# caused the small image comparison error. +@image_comparison(['imagegrid_cbar_mode.png'], + remove_text=True, style='mpl20', tol=0.3) +def test_imagegrid_cbar_mode_edge(): + arr = np.arange(16).reshape((4, 4)) + + fig = plt.figure(figsize=(18, 9)) + + positions = (241, 242, 243, 244, 245, 246, 247, 248) + directions = ['row']*4 + ['column']*4 + cbar_locations = ['left', 'right', 'top', 'bottom']*2 + + for position, direction, location in zip( + positions, directions, cbar_locations): + grid = ImageGrid(fig, position, + nrows_ncols=(2, 2), + direction=direction, + cbar_location=location, + cbar_size='20%', + cbar_mode='edge') + ax1, ax2, ax3, ax4 = grid + + ax1.imshow(arr, cmap='nipy_spectral') + ax2.imshow(arr.T, cmap='hot') + ax3.imshow(np.hypot(arr, arr.T), cmap='jet') + ax4.imshow(np.arctan2(arr, arr.T), cmap='hsv') + + # In each row/column, the "first" colorbars must be overwritten by the + # "second" ones. To achieve this, clear out the axes first. + for ax in grid: + ax.cax.cla() + cb = ax.cax.colorbar(ax.images[0]) + + +def test_imagegrid(): + fig = plt.figure() + grid = ImageGrid(fig, 111, nrows_ncols=(1, 1)) + ax = grid[0] + im = ax.imshow([[1, 2]], norm=mpl.colors.LogNorm()) + cb = ax.cax.colorbar(im) + assert isinstance(cb.locator, mticker.LogLocator) + + +def test_removal(): + import matplotlib.pyplot as plt + import mpl_toolkits.axisartist as AA + fig = plt.figure() + ax = host_subplot(111, axes_class=AA.Axes, figure=fig) + col = ax.fill_between(range(5), 0, range(5)) + fig.canvas.draw() + col.remove() + fig.canvas.draw() + + +@image_comparison(['anchored_locator_base_call.png'], style="mpl20") +def test_anchored_locator_base_call(): + fig = plt.figure(figsize=(3, 3)) + fig1, fig2 = fig.subfigures(nrows=2, ncols=1) + + ax = fig1.subplots() + ax.set(aspect=1, xlim=(-15, 15), ylim=(-20, 5)) + ax.set(xticks=[], yticks=[]) + + Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy") + extent = (-3, 4, -4, 3) + + axins = zoomed_inset_axes(ax, zoom=2, loc="upper left") + axins.set(xticks=[], yticks=[]) + + axins.imshow(Z, extent=extent, origin="lower") + + +def test_grid_with_axes_class_not_overriding_axis(): + Grid(plt.figure(), 111, (2, 2), axes_class=mpl.axes.Axes) + RGBAxes(plt.figure(), 111, axes_class=mpl.axes.Axes) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/axisartist/__pycache__/grid_helper_curvelinear.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/axisartist/__pycache__/grid_helper_curvelinear.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07f52ccd9a4f76e71c9d5b01d6801ce6ad8f9507 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/axisartist/__pycache__/grid_helper_curvelinear.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__init__.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a089fbd6b70ebdafb19f6c902d49e5e3bdc05f43 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__init__.py @@ -0,0 +1,3 @@ +from .axes3d import Axes3D + +__all__ = ['Axes3D'] diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e3d43c744ab533dd44301903b30ab60ca72edb3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/art3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/art3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6356a213c61f2d5588b9f29ff9fb2ec7e543eedd Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/art3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/axis3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/axis3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7cdbba0f16e39be68af2f7f582491cb844c94060 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/axis3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/proj3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/proj3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..919970cb0077d2fa4c8a6c0d3ff8729e22831a8f Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/__pycache__/proj3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/art3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/art3d.py new file mode 100644 index 0000000000000000000000000000000000000000..deb0ca34302c03cda08e57b28a82cb3654e45aaa --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/art3d.py @@ -0,0 +1,1427 @@ +# art3d.py, original mplot3d version by John Porter +# Parts rewritten by Reinier Heeres +# Minor additions by Ben Axelrod + +""" +Module containing 3D artist code and functions to convert 2D +artists into 3D versions which can be added to an Axes3D. +""" + +import math + +import numpy as np + +from contextlib import contextmanager + +from matplotlib import ( + _api, artist, cbook, colors as mcolors, lines, text as mtext, + path as mpath) +from matplotlib.collections import ( + Collection, LineCollection, PolyCollection, PatchCollection, PathCollection) +from matplotlib.colors import Normalize +from matplotlib.patches import Patch +from . import proj3d + + +def _norm_angle(a): + """Return the given angle normalized to -180 < *a* <= 180 degrees.""" + a = (a + 360) % 360 + if a > 180: + a = a - 360 + return a + + +def _norm_text_angle(a): + """Return the given angle normalized to -90 < *a* <= 90 degrees.""" + a = (a + 180) % 180 + if a > 90: + a = a - 180 + return a + + +def get_dir_vector(zdir): + """ + Return a direction vector. + + Parameters + ---------- + zdir : {'x', 'y', 'z', None, 3-tuple} + The direction. Possible values are: + + - 'x': equivalent to (1, 0, 0) + - 'y': equivalent to (0, 1, 0) + - 'z': equivalent to (0, 0, 1) + - *None*: equivalent to (0, 0, 0) + - an iterable (x, y, z) is converted to an array + + Returns + ------- + x, y, z : array + The direction vector. + """ + if zdir == 'x': + return np.array((1, 0, 0)) + elif zdir == 'y': + return np.array((0, 1, 0)) + elif zdir == 'z': + return np.array((0, 0, 1)) + elif zdir is None: + return np.array((0, 0, 0)) + elif np.iterable(zdir) and len(zdir) == 3: + return np.array(zdir) + else: + raise ValueError("'x', 'y', 'z', None or vector of length 3 expected") + + +def _viewlim_mask(xs, ys, zs, axes): + """ + Return original points with points outside the axes view limits masked. + + Parameters + ---------- + xs, ys, zs : array-like + The points to mask. + axes : Axes3D + The axes to use for the view limits. + + Returns + ------- + xs_masked, ys_masked, zs_masked : np.ma.array + The masked points. + """ + mask = np.logical_or.reduce((xs < axes.xy_viewLim.xmin, + xs > axes.xy_viewLim.xmax, + ys < axes.xy_viewLim.ymin, + ys > axes.xy_viewLim.ymax, + zs < axes.zz_viewLim.xmin, + zs > axes.zz_viewLim.xmax)) + xs_masked = np.ma.array(xs, mask=mask) + ys_masked = np.ma.array(ys, mask=mask) + zs_masked = np.ma.array(zs, mask=mask) + return xs_masked, ys_masked, zs_masked + + +class Text3D(mtext.Text): + """ + Text object with 3D position and direction. + + Parameters + ---------- + x, y, z : float + The position of the text. + text : str + The text string to display. + zdir : {'x', 'y', 'z', None, 3-tuple} + The direction of the text. See `.get_dir_vector` for a description of + the values. + axlim_clip : bool, default: False + Whether to hide text outside the axes view limits. + + Other Parameters + ---------------- + **kwargs + All other parameters are passed on to `~matplotlib.text.Text`. + """ + + def __init__(self, x=0, y=0, z=0, text='', zdir='z', axlim_clip=False, + **kwargs): + mtext.Text.__init__(self, x, y, text, **kwargs) + self.set_3d_properties(z, zdir, axlim_clip) + + def get_position_3d(self): + """Return the (x, y, z) position of the text.""" + return self._x, self._y, self._z + + def set_position_3d(self, xyz, zdir=None): + """ + Set the (*x*, *y*, *z*) position of the text. + + Parameters + ---------- + xyz : (float, float, float) + The position in 3D space. + zdir : {'x', 'y', 'z', None, 3-tuple} + The direction of the text. If unspecified, the *zdir* will not be + changed. See `.get_dir_vector` for a description of the values. + """ + super().set_position(xyz[:2]) + self.set_z(xyz[2]) + if zdir is not None: + self._dir_vec = get_dir_vector(zdir) + + def set_z(self, z): + """ + Set the *z* position of the text. + + Parameters + ---------- + z : float + """ + self._z = z + self.stale = True + + def set_3d_properties(self, z=0, zdir='z', axlim_clip=False): + """ + Set the *z* position and direction of the text. + + Parameters + ---------- + z : float + The z-position in 3D space. + zdir : {'x', 'y', 'z', 3-tuple} + The direction of the text. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide text outside the axes view limits. + """ + self._z = z + self._dir_vec = get_dir_vector(zdir) + self._axlim_clip = axlim_clip + self.stale = True + + @artist.allow_rasterization + def draw(self, renderer): + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(self._x, self._y, self._z, self.axes) + position3d = np.ma.row_stack((xs, ys, zs)).ravel().filled(np.nan) + else: + xs, ys, zs = self._x, self._y, self._z + position3d = np.asanyarray([xs, ys, zs]) + + proj = proj3d._proj_trans_points( + [position3d, position3d + self._dir_vec], self.axes.M) + dx = proj[0][1] - proj[0][0] + dy = proj[1][1] - proj[1][0] + angle = math.degrees(math.atan2(dy, dx)) + with cbook._setattr_cm(self, _x=proj[0][0], _y=proj[1][0], + _rotation=_norm_text_angle(angle)): + mtext.Text.draw(self, renderer) + self.stale = False + + def get_tightbbox(self, renderer=None): + # Overwriting the 2d Text behavior which is not valid for 3d. + # For now, just return None to exclude from layout calculation. + return None + + +def text_2d_to_3d(obj, z=0, zdir='z', axlim_clip=False): + """ + Convert a `.Text` to a `.Text3D` object. + + Parameters + ---------- + z : float + The z-position in 3D space. + zdir : {'x', 'y', 'z', 3-tuple} + The direction of the text. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide text outside the axes view limits. + """ + obj.__class__ = Text3D + obj.set_3d_properties(z, zdir, axlim_clip) + + +class Line3D(lines.Line2D): + """ + 3D line object. + + .. note:: Use `get_data_3d` to obtain the data associated with the line. + `~.Line2D.get_data`, `~.Line2D.get_xdata`, and `~.Line2D.get_ydata` return + the x- and y-coordinates of the projected 2D-line, not the x- and y-data of + the 3D-line. Similarly, use `set_data_3d` to set the data, not + `~.Line2D.set_data`, `~.Line2D.set_xdata`, and `~.Line2D.set_ydata`. + """ + + def __init__(self, xs, ys, zs, *args, axlim_clip=False, **kwargs): + """ + + Parameters + ---------- + xs : array-like + The x-data to be plotted. + ys : array-like + The y-data to be plotted. + zs : array-like + The z-data to be plotted. + *args, **kwargs + Additional arguments are passed to `~matplotlib.lines.Line2D`. + """ + super().__init__([], [], *args, **kwargs) + self.set_data_3d(xs, ys, zs) + self._axlim_clip = axlim_clip + + def set_3d_properties(self, zs=0, zdir='z', axlim_clip=False): + """ + Set the *z* position and direction of the line. + + Parameters + ---------- + zs : float or array of floats + The location along the *zdir* axis in 3D space to position the + line. + zdir : {'x', 'y', 'z'} + Plane to plot line orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide lines with an endpoint outside the axes view limits. + """ + xs = self.get_xdata() + ys = self.get_ydata() + zs = cbook._to_unmasked_float_array(zs).ravel() + zs = np.broadcast_to(zs, len(xs)) + self._verts3d = juggle_axes(xs, ys, zs, zdir) + self._axlim_clip = axlim_clip + self.stale = True + + def set_data_3d(self, *args): + """ + Set the x, y and z data + + Parameters + ---------- + x : array-like + The x-data to be plotted. + y : array-like + The y-data to be plotted. + z : array-like + The z-data to be plotted. + + Notes + ----- + Accepts x, y, z arguments or a single array-like (x, y, z) + """ + if len(args) == 1: + args = args[0] + for name, xyz in zip('xyz', args): + if not np.iterable(xyz): + raise RuntimeError(f'{name} must be a sequence') + self._verts3d = args + self.stale = True + + def get_data_3d(self): + """ + Get the current data + + Returns + ------- + verts3d : length-3 tuple or array-like + The current data as a tuple or array-like. + """ + return self._verts3d + + @artist.allow_rasterization + def draw(self, renderer): + if self._axlim_clip: + xs3d, ys3d, zs3d = _viewlim_mask(*self._verts3d, self.axes) + else: + xs3d, ys3d, zs3d = self._verts3d + xs, ys, zs, tis = proj3d._proj_transform_clip(xs3d, ys3d, zs3d, + self.axes.M, + self.axes._focal_length) + self.set_data(xs, ys) + super().draw(renderer) + self.stale = False + + +def line_2d_to_3d(line, zs=0, zdir='z', axlim_clip=False): + """ + Convert a `.Line2D` to a `.Line3D` object. + + Parameters + ---------- + zs : float + The location along the *zdir* axis in 3D space to position the line. + zdir : {'x', 'y', 'z'} + Plane to plot line orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide lines with an endpoint outside the axes view limits. + """ + + line.__class__ = Line3D + line.set_3d_properties(zs, zdir, axlim_clip) + + +def _path_to_3d_segment(path, zs=0, zdir='z'): + """Convert a path to a 3D segment.""" + + zs = np.broadcast_to(zs, len(path)) + pathsegs = path.iter_segments(simplify=False, curves=False) + seg = [(x, y, z) for (((x, y), code), z) in zip(pathsegs, zs)] + seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg] + return seg3d + + +def _paths_to_3d_segments(paths, zs=0, zdir='z'): + """Convert paths from a collection object to 3D segments.""" + + if not np.iterable(zs): + zs = np.broadcast_to(zs, len(paths)) + else: + if len(zs) != len(paths): + raise ValueError('Number of z-coordinates does not match paths.') + + segs = [_path_to_3d_segment(path, pathz, zdir) + for path, pathz in zip(paths, zs)] + return segs + + +def _path_to_3d_segment_with_codes(path, zs=0, zdir='z'): + """Convert a path to a 3D segment with path codes.""" + + zs = np.broadcast_to(zs, len(path)) + pathsegs = path.iter_segments(simplify=False, curves=False) + seg_codes = [((x, y, z), code) for ((x, y), code), z in zip(pathsegs, zs)] + if seg_codes: + seg, codes = zip(*seg_codes) + seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg] + else: + seg3d = [] + codes = [] + return seg3d, list(codes) + + +def _paths_to_3d_segments_with_codes(paths, zs=0, zdir='z'): + """ + Convert paths from a collection object to 3D segments with path codes. + """ + + zs = np.broadcast_to(zs, len(paths)) + segments_codes = [_path_to_3d_segment_with_codes(path, pathz, zdir) + for path, pathz in zip(paths, zs)] + if segments_codes: + segments, codes = zip(*segments_codes) + else: + segments, codes = [], [] + return list(segments), list(codes) + + +class Collection3D(Collection): + """A collection of 3D paths.""" + + def do_3d_projection(self): + """Project the points according to renderer matrix.""" + vs_list = [vs for vs, _ in self._3dverts_codes] + if self._axlim_clip: + vs_list = [np.ma.row_stack(_viewlim_mask(*vs.T, self.axes)).T + for vs in vs_list] + xyzs_list = [proj3d.proj_transform(*vs.T, self.axes.M) for vs in vs_list] + self._paths = [mpath.Path(np.ma.column_stack([xs, ys]), cs) + for (xs, ys, _), (_, cs) in zip(xyzs_list, self._3dverts_codes)] + zs = np.concatenate([zs for _, _, zs in xyzs_list]) + return zs.min() if len(zs) else 1e9 + + +def collection_2d_to_3d(col, zs=0, zdir='z', axlim_clip=False): + """Convert a `.Collection` to a `.Collection3D` object.""" + zs = np.broadcast_to(zs, len(col.get_paths())) + col._3dverts_codes = [ + (np.column_stack(juggle_axes( + *np.column_stack([p.vertices, np.broadcast_to(z, len(p.vertices))]).T, + zdir)), + p.codes) + for p, z in zip(col.get_paths(), zs)] + col.__class__ = cbook._make_class_factory(Collection3D, "{}3D")(type(col)) + col._axlim_clip = axlim_clip + + +class Line3DCollection(LineCollection): + """ + A collection of 3D lines. + """ + def __init__(self, lines, axlim_clip=False, **kwargs): + super().__init__(lines, **kwargs) + self._axlim_clip = axlim_clip + + def set_sort_zpos(self, val): + """Set the position to use for z-sorting.""" + self._sort_zpos = val + self.stale = True + + def set_segments(self, segments): + """ + Set 3D segments. + """ + self._segments3d = segments + super().set_segments([]) + + def do_3d_projection(self): + """ + Project the points according to renderer matrix. + """ + segments = self._segments3d + if self._axlim_clip: + all_points = np.ma.vstack(segments) + masked_points = np.ma.column_stack([*_viewlim_mask(*all_points.T, + self.axes)]) + segment_lengths = [np.shape(segment)[0] for segment in segments] + segments = np.split(masked_points, np.cumsum(segment_lengths[:-1])) + xyslist = [proj3d._proj_trans_points(points, self.axes.M) + for points in segments] + segments_2d = [np.ma.column_stack([xs, ys]) for xs, ys, zs in xyslist] + LineCollection.set_segments(self, segments_2d) + + # FIXME + minz = 1e9 + for xs, ys, zs in xyslist: + minz = min(minz, min(zs)) + return minz + + +def line_collection_2d_to_3d(col, zs=0, zdir='z', axlim_clip=False): + """Convert a `.LineCollection` to a `.Line3DCollection` object.""" + segments3d = _paths_to_3d_segments(col.get_paths(), zs, zdir) + col.__class__ = Line3DCollection + col.set_segments(segments3d) + col._axlim_clip = axlim_clip + + +class Patch3D(Patch): + """ + 3D patch object. + """ + + def __init__(self, *args, zs=(), zdir='z', axlim_clip=False, **kwargs): + """ + Parameters + ---------- + verts : + zs : float + The location along the *zdir* axis in 3D space to position the + patch. + zdir : {'x', 'y', 'z'} + Plane to plot patch orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + """ + super().__init__(*args, **kwargs) + self.set_3d_properties(zs, zdir, axlim_clip) + + def set_3d_properties(self, verts, zs=0, zdir='z', axlim_clip=False): + """ + Set the *z* position and direction of the patch. + + Parameters + ---------- + verts : + zs : float + The location along the *zdir* axis in 3D space to position the + patch. + zdir : {'x', 'y', 'z'} + Plane to plot patch orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + """ + zs = np.broadcast_to(zs, len(verts)) + self._segment3d = [juggle_axes(x, y, z, zdir) + for ((x, y), z) in zip(verts, zs)] + self._axlim_clip = axlim_clip + + def get_path(self): + # docstring inherited + # self._path2d is not initialized until do_3d_projection + if not hasattr(self, '_path2d'): + self.axes.M = self.axes.get_proj() + self.do_3d_projection() + return self._path2d + + def do_3d_projection(self): + s = self._segment3d + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(*zip(*s), self.axes) + else: + xs, ys, zs = zip(*s) + vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs, + self.axes.M, + self.axes._focal_length) + self._path2d = mpath.Path(np.ma.column_stack([vxs, vys])) + return min(vzs) + + +class PathPatch3D(Patch3D): + """ + 3D PathPatch object. + """ + + def __init__(self, path, *, zs=(), zdir='z', axlim_clip=False, **kwargs): + """ + Parameters + ---------- + path : + zs : float + The location along the *zdir* axis in 3D space to position the + path patch. + zdir : {'x', 'y', 'z', 3-tuple} + Plane to plot path patch orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide path patches with a point outside the axes view limits. + """ + # Not super().__init__! + Patch.__init__(self, **kwargs) + self.set_3d_properties(path, zs, zdir, axlim_clip) + + def set_3d_properties(self, path, zs=0, zdir='z', axlim_clip=False): + """ + Set the *z* position and direction of the path patch. + + Parameters + ---------- + path : + zs : float + The location along the *zdir* axis in 3D space to position the + path patch. + zdir : {'x', 'y', 'z', 3-tuple} + Plane to plot path patch orthogonal to. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide path patches with a point outside the axes view limits. + """ + Patch3D.set_3d_properties(self, path.vertices, zs=zs, zdir=zdir, + axlim_clip=axlim_clip) + self._code3d = path.codes + + def do_3d_projection(self): + s = self._segment3d + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(*zip(*s), self.axes) + else: + xs, ys, zs = zip(*s) + vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs, + self.axes.M, + self.axes._focal_length) + self._path2d = mpath.Path(np.ma.column_stack([vxs, vys]), self._code3d) + return min(vzs) + + +def _get_patch_verts(patch): + """Return a list of vertices for the path of a patch.""" + trans = patch.get_patch_transform() + path = patch.get_path() + polygons = path.to_polygons(trans) + return polygons[0] if len(polygons) else np.array([]) + + +def patch_2d_to_3d(patch, z=0, zdir='z', axlim_clip=False): + """Convert a `.Patch` to a `.Patch3D` object.""" + verts = _get_patch_verts(patch) + patch.__class__ = Patch3D + patch.set_3d_properties(verts, z, zdir, axlim_clip) + + +def pathpatch_2d_to_3d(pathpatch, z=0, zdir='z'): + """Convert a `.PathPatch` to a `.PathPatch3D` object.""" + path = pathpatch.get_path() + trans = pathpatch.get_patch_transform() + + mpath = trans.transform_path(path) + pathpatch.__class__ = PathPatch3D + pathpatch.set_3d_properties(mpath, z, zdir) + + +class Patch3DCollection(PatchCollection): + """ + A collection of 3D patches. + """ + + def __init__(self, *args, + zs=0, zdir='z', depthshade=True, axlim_clip=False, **kwargs): + """ + Create a collection of flat 3D patches with its normal vector + pointed in *zdir* direction, and located at *zs* on the *zdir* + axis. 'zs' can be a scalar or an array-like of the same length as + the number of patches in the collection. + + Constructor arguments are the same as for + :class:`~matplotlib.collections.PatchCollection`. In addition, + keywords *zs=0* and *zdir='z'* are available. + + Also, the keyword argument *depthshade* is available to indicate + whether to shade the patches in order to give the appearance of depth + (default is *True*). This is typically desired in scatter plots. + """ + self._depthshade = depthshade + super().__init__(*args, **kwargs) + self.set_3d_properties(zs, zdir, axlim_clip) + + def get_depthshade(self): + return self._depthshade + + def set_depthshade(self, depthshade): + """ + Set whether depth shading is performed on collection members. + + Parameters + ---------- + depthshade : bool + Whether to shade the patches in order to give the appearance of + depth. + """ + self._depthshade = depthshade + self.stale = True + + def set_sort_zpos(self, val): + """Set the position to use for z-sorting.""" + self._sort_zpos = val + self.stale = True + + def set_3d_properties(self, zs, zdir, axlim_clip=False): + """ + Set the *z* positions and direction of the patches. + + Parameters + ---------- + zs : float or array of floats + The location or locations to place the patches in the collection + along the *zdir* axis. + zdir : {'x', 'y', 'z'} + Plane to plot patches orthogonal to. + All patches must have the same direction. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + """ + # Force the collection to initialize the face and edgecolors + # just in case it is a scalarmappable with a colormap. + self.update_scalarmappable() + offsets = self.get_offsets() + if len(offsets) > 0: + xs, ys = offsets.T + else: + xs = [] + ys = [] + self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir) + self._z_markers_idx = slice(-1) + self._vzs = None + self._axlim_clip = axlim_clip + self.stale = True + + def do_3d_projection(self): + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(*self._offsets3d, self.axes) + else: + xs, ys, zs = self._offsets3d + vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs, + self.axes.M, + self.axes._focal_length) + self._vzs = vzs + super().set_offsets(np.ma.column_stack([vxs, vys])) + + if vzs.size > 0: + return min(vzs) + else: + return np.nan + + def _maybe_depth_shade_and_sort_colors(self, color_array): + color_array = ( + _zalpha(color_array, self._vzs) + if self._vzs is not None and self._depthshade + else color_array + ) + if len(color_array) > 1: + color_array = color_array[self._z_markers_idx] + return mcolors.to_rgba_array(color_array, self._alpha) + + def get_facecolor(self): + return self._maybe_depth_shade_and_sort_colors(super().get_facecolor()) + + def get_edgecolor(self): + # We need this check here to make sure we do not double-apply the depth + # based alpha shading when the edge color is "face" which means the + # edge colour should be identical to the face colour. + if cbook._str_equal(self._edgecolors, 'face'): + return self.get_facecolor() + return self._maybe_depth_shade_and_sort_colors(super().get_edgecolor()) + + +class Path3DCollection(PathCollection): + """ + A collection of 3D paths. + """ + + def __init__(self, *args, + zs=0, zdir='z', depthshade=True, axlim_clip=False, **kwargs): + """ + Create a collection of flat 3D paths with its normal vector + pointed in *zdir* direction, and located at *zs* on the *zdir* + axis. 'zs' can be a scalar or an array-like of the same length as + the number of paths in the collection. + + Constructor arguments are the same as for + :class:`~matplotlib.collections.PathCollection`. In addition, + keywords *zs=0* and *zdir='z'* are available. + + Also, the keyword argument *depthshade* is available to indicate + whether to shade the patches in order to give the appearance of depth + (default is *True*). This is typically desired in scatter plots. + """ + self._depthshade = depthshade + self._in_draw = False + super().__init__(*args, **kwargs) + self.set_3d_properties(zs, zdir, axlim_clip) + self._offset_zordered = None + + def draw(self, renderer): + with self._use_zordered_offset(): + with cbook._setattr_cm(self, _in_draw=True): + super().draw(renderer) + + def set_sort_zpos(self, val): + """Set the position to use for z-sorting.""" + self._sort_zpos = val + self.stale = True + + def set_3d_properties(self, zs, zdir, axlim_clip=False): + """ + Set the *z* positions and direction of the paths. + + Parameters + ---------- + zs : float or array of floats + The location or locations to place the paths in the collection + along the *zdir* axis. + zdir : {'x', 'y', 'z'} + Plane to plot paths orthogonal to. + All paths must have the same direction. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide paths with a vertex outside the axes view limits. + """ + # Force the collection to initialize the face and edgecolors + # just in case it is a scalarmappable with a colormap. + self.update_scalarmappable() + offsets = self.get_offsets() + if len(offsets) > 0: + xs, ys = offsets.T + else: + xs = [] + ys = [] + self._zdir = zdir + self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir) + # In the base draw methods we access the attributes directly which + # means we cannot resolve the shuffling in the getter methods like + # we do for the edge and face colors. + # + # This means we need to carry around a cache of the unsorted sizes and + # widths (postfixed with 3d) and in `do_3d_projection` set the + # depth-sorted version of that data into the private state used by the + # base collection class in its draw method. + # + # Grab the current sizes and linewidths to preserve them. + self._sizes3d = self._sizes + self._linewidths3d = np.array(self._linewidths) + xs, ys, zs = self._offsets3d + + # Sort the points based on z coordinates + # Performance optimization: Create a sorted index array and reorder + # points and point properties according to the index array + self._z_markers_idx = slice(-1) + self._vzs = None + + self._axlim_clip = axlim_clip + self.stale = True + + def set_sizes(self, sizes, dpi=72.0): + super().set_sizes(sizes, dpi) + if not self._in_draw: + self._sizes3d = sizes + + def set_linewidth(self, lw): + super().set_linewidth(lw) + if not self._in_draw: + self._linewidths3d = np.array(self._linewidths) + + def get_depthshade(self): + return self._depthshade + + def set_depthshade(self, depthshade): + """ + Set whether depth shading is performed on collection members. + + Parameters + ---------- + depthshade : bool + Whether to shade the patches in order to give the appearance of + depth. + """ + self._depthshade = depthshade + self.stale = True + + def do_3d_projection(self): + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(*self._offsets3d, self.axes) + else: + xs, ys, zs = self._offsets3d + vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs, + self.axes.M, + self.axes._focal_length) + # Sort the points based on z coordinates + # Performance optimization: Create a sorted index array and reorder + # points and point properties according to the index array + z_markers_idx = self._z_markers_idx = np.ma.argsort(vzs)[::-1] + self._vzs = vzs + + # we have to special case the sizes because of code in collections.py + # as the draw method does + # self.set_sizes(self._sizes, self.figure.dpi) + # so we cannot rely on doing the sorting on the way out via get_* + + if len(self._sizes3d) > 1: + self._sizes = self._sizes3d[z_markers_idx] + + if len(self._linewidths3d) > 1: + self._linewidths = self._linewidths3d[z_markers_idx] + + PathCollection.set_offsets(self, np.ma.column_stack((vxs, vys))) + + # Re-order items + vzs = vzs[z_markers_idx] + vxs = vxs[z_markers_idx] + vys = vys[z_markers_idx] + + # Store ordered offset for drawing purpose + self._offset_zordered = np.ma.column_stack((vxs, vys)) + + return np.min(vzs) if vzs.size else np.nan + + @contextmanager + def _use_zordered_offset(self): + if self._offset_zordered is None: + # Do nothing + yield + else: + # Swap offset with z-ordered offset + old_offset = self._offsets + super().set_offsets(self._offset_zordered) + try: + yield + finally: + self._offsets = old_offset + + def _maybe_depth_shade_and_sort_colors(self, color_array): + color_array = ( + _zalpha(color_array, self._vzs) + if self._vzs is not None and self._depthshade + else color_array + ) + if len(color_array) > 1: + color_array = color_array[self._z_markers_idx] + return mcolors.to_rgba_array(color_array, self._alpha) + + def get_facecolor(self): + return self._maybe_depth_shade_and_sort_colors(super().get_facecolor()) + + def get_edgecolor(self): + # We need this check here to make sure we do not double-apply the depth + # based alpha shading when the edge color is "face" which means the + # edge colour should be identical to the face colour. + if cbook._str_equal(self._edgecolors, 'face'): + return self.get_facecolor() + return self._maybe_depth_shade_and_sort_colors(super().get_edgecolor()) + + +def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True, axlim_clip=False): + """ + Convert a `.PatchCollection` into a `.Patch3DCollection` object + (or a `.PathCollection` into a `.Path3DCollection` object). + + Parameters + ---------- + col : `~matplotlib.collections.PatchCollection` or \ +`~matplotlib.collections.PathCollection` + The collection to convert. + zs : float or array of floats + The location or locations to place the patches in the collection along + the *zdir* axis. Default: 0. + zdir : {'x', 'y', 'z'} + The axis in which to place the patches. Default: "z". + See `.get_dir_vector` for a description of the values. + depthshade : bool, default: True + Whether to shade the patches to give a sense of depth. + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + """ + if isinstance(col, PathCollection): + col.__class__ = Path3DCollection + col._offset_zordered = None + elif isinstance(col, PatchCollection): + col.__class__ = Patch3DCollection + col._depthshade = depthshade + col._in_draw = False + col.set_3d_properties(zs, zdir, axlim_clip) + + +class Poly3DCollection(PolyCollection): + """ + A collection of 3D polygons. + + .. note:: + **Filling of 3D polygons** + + There is no simple definition of the enclosed surface of a 3D polygon + unless the polygon is planar. + + In practice, Matplotlib fills the 2D projection of the polygon. This + gives a correct filling appearance only for planar polygons. For all + other polygons, you'll find orientations in which the edges of the + polygon intersect in the projection. This will lead to an incorrect + visualization of the 3D area. + + If you need filled areas, it is recommended to create them via + `~mpl_toolkits.mplot3d.axes3d.Axes3D.plot_trisurf`, which creates a + triangulation and thus generates consistent surfaces. + """ + + def __init__(self, verts, *args, zsort='average', shade=False, + lightsource=None, axlim_clip=False, **kwargs): + """ + Parameters + ---------- + verts : list of (N, 3) array-like + The sequence of polygons [*verts0*, *verts1*, ...] where each + element *verts_i* defines the vertices of polygon *i* as a 2D + array-like of shape (N, 3). + zsort : {'average', 'min', 'max'}, default: 'average' + The calculation method for the z-order. + See `~.Poly3DCollection.set_zsort` for details. + shade : bool, default: False + Whether to shade *facecolors* and *edgecolors*. When activating + *shade*, *facecolors* and/or *edgecolors* must be provided. + + .. versionadded:: 3.7 + + lightsource : `~matplotlib.colors.LightSource`, optional + The lightsource to use when *shade* is True. + + .. versionadded:: 3.7 + + axlim_clip : bool, default: False + Whether to hide polygons with a vertex outside the view limits. + + *args, **kwargs + All other parameters are forwarded to `.PolyCollection`. + + Notes + ----- + Note that this class does a bit of magic with the _facecolors + and _edgecolors properties. + """ + if shade: + normals = _generate_normals(verts) + facecolors = kwargs.get('facecolors', None) + if facecolors is not None: + kwargs['facecolors'] = _shade_colors( + facecolors, normals, lightsource + ) + + edgecolors = kwargs.get('edgecolors', None) + if edgecolors is not None: + kwargs['edgecolors'] = _shade_colors( + edgecolors, normals, lightsource + ) + if facecolors is None and edgecolors is None: + raise ValueError( + "You must provide facecolors, edgecolors, or both for " + "shade to work.") + super().__init__(verts, *args, **kwargs) + if isinstance(verts, np.ndarray): + if verts.ndim != 3: + raise ValueError('verts must be a list of (N, 3) array-like') + else: + if any(len(np.shape(vert)) != 2 for vert in verts): + raise ValueError('verts must be a list of (N, 3) array-like') + self.set_zsort(zsort) + self._codes3d = None + self._axlim_clip = axlim_clip + + _zsort_functions = { + 'average': np.average, + 'min': np.min, + 'max': np.max, + } + + def set_zsort(self, zsort): + """ + Set the calculation method for the z-order. + + Parameters + ---------- + zsort : {'average', 'min', 'max'} + The function applied on the z-coordinates of the vertices in the + viewer's coordinate system, to determine the z-order. + """ + self._zsortfunc = self._zsort_functions[zsort] + self._sort_zpos = None + self.stale = True + + @_api.deprecated("3.10") + def get_vector(self, segments3d): + return self._get_vector(segments3d) + + def _get_vector(self, segments3d): + """Optimize points for projection.""" + if len(segments3d): + xs, ys, zs = np.vstack(segments3d).T + else: # vstack can't stack zero arrays. + xs, ys, zs = [], [], [] + ones = np.ones(len(xs)) + self._vec = np.array([xs, ys, zs, ones]) + + indices = [0, *np.cumsum([len(segment) for segment in segments3d])] + self._segslices = [*map(slice, indices[:-1], indices[1:])] + + def set_verts(self, verts, closed=True): + """ + Set 3D vertices. + + Parameters + ---------- + verts : list of (N, 3) array-like + The sequence of polygons [*verts0*, *verts1*, ...] where each + element *verts_i* defines the vertices of polygon *i* as a 2D + array-like of shape (N, 3). + closed : bool, default: True + Whether the polygon should be closed by adding a CLOSEPOLY + connection at the end. + """ + self._get_vector(verts) + # 2D verts will be updated at draw time + super().set_verts([], False) + self._closed = closed + + def set_verts_and_codes(self, verts, codes): + """Set 3D vertices with path codes.""" + # set vertices with closed=False to prevent PolyCollection from + # setting path codes + self.set_verts(verts, closed=False) + # and set our own codes instead. + self._codes3d = codes + + def set_3d_properties(self, axlim_clip=False): + # Force the collection to initialize the face and edgecolors + # just in case it is a scalarmappable with a colormap. + self.update_scalarmappable() + self._sort_zpos = None + self.set_zsort('average') + self._facecolor3d = PolyCollection.get_facecolor(self) + self._edgecolor3d = PolyCollection.get_edgecolor(self) + self._alpha3d = PolyCollection.get_alpha(self) + self.stale = True + + def set_sort_zpos(self, val): + """Set the position to use for z-sorting.""" + self._sort_zpos = val + self.stale = True + + def do_3d_projection(self): + """ + Perform the 3D projection for this object. + """ + if self._A is not None: + # force update of color mapping because we re-order them + # below. If we do not do this here, the 2D draw will call + # this, but we will never port the color mapped values back + # to the 3D versions. + # + # We hold the 3D versions in a fixed order (the order the user + # passed in) and sort the 2D version by view depth. + self.update_scalarmappable() + if self._face_is_mapped: + self._facecolor3d = self._facecolors + if self._edge_is_mapped: + self._edgecolor3d = self._edgecolors + if self._axlim_clip: + xs, ys, zs = _viewlim_mask(*self._vec[0:3], self.axes) + if self._vec.shape[0] == 4: # Will be 3 (xyz) or 4 (xyzw) + w_masked = np.ma.masked_where(zs.mask, self._vec[3]) + vec = np.ma.array([xs, ys, zs, w_masked]) + else: + vec = np.ma.array([xs, ys, zs]) + else: + vec = self._vec + txs, tys, tzs = proj3d._proj_transform_vec(vec, self.axes.M) + xyzlist = [(txs[sl], tys[sl], tzs[sl]) for sl in self._segslices] + + # This extra fuss is to re-order face / edge colors + cface = self._facecolor3d + cedge = self._edgecolor3d + if len(cface) != len(xyzlist): + cface = cface.repeat(len(xyzlist), axis=0) + if len(cedge) != len(xyzlist): + if len(cedge) == 0: + cedge = cface + else: + cedge = cedge.repeat(len(xyzlist), axis=0) + + if xyzlist: + # sort by depth (furthest drawn first) + z_segments_2d = sorted( + ((self._zsortfunc(zs.data), np.ma.column_stack([xs, ys]), fc, ec, idx) + for idx, ((xs, ys, zs), fc, ec) + in enumerate(zip(xyzlist, cface, cedge))), + key=lambda x: x[0], reverse=True) + + _, segments_2d, self._facecolors2d, self._edgecolors2d, idxs = \ + zip(*z_segments_2d) + else: + segments_2d = [] + self._facecolors2d = np.empty((0, 4)) + self._edgecolors2d = np.empty((0, 4)) + idxs = [] + + if self._codes3d is not None: + codes = [self._codes3d[idx] for idx in idxs] + PolyCollection.set_verts_and_codes(self, segments_2d, codes) + else: + PolyCollection.set_verts(self, segments_2d, self._closed) + + if len(self._edgecolor3d) != len(cface): + self._edgecolors2d = self._edgecolor3d + + # Return zorder value + if self._sort_zpos is not None: + zvec = np.array([[0], [0], [self._sort_zpos], [1]]) + ztrans = proj3d._proj_transform_vec(zvec, self.axes.M) + return ztrans[2][0] + elif tzs.size > 0: + # FIXME: Some results still don't look quite right. + # In particular, examine contourf3d_demo2.py + # with az = -54 and elev = -45. + return np.min(tzs) + else: + return np.nan + + def set_facecolor(self, colors): + # docstring inherited + super().set_facecolor(colors) + self._facecolor3d = PolyCollection.get_facecolor(self) + + def set_edgecolor(self, colors): + # docstring inherited + super().set_edgecolor(colors) + self._edgecolor3d = PolyCollection.get_edgecolor(self) + + def set_alpha(self, alpha): + # docstring inherited + artist.Artist.set_alpha(self, alpha) + try: + self._facecolor3d = mcolors.to_rgba_array( + self._facecolor3d, self._alpha) + except (AttributeError, TypeError, IndexError): + pass + try: + self._edgecolors = mcolors.to_rgba_array( + self._edgecolor3d, self._alpha) + except (AttributeError, TypeError, IndexError): + pass + self.stale = True + + def get_facecolor(self): + # docstring inherited + # self._facecolors2d is not initialized until do_3d_projection + if not hasattr(self, '_facecolors2d'): + self.axes.M = self.axes.get_proj() + self.do_3d_projection() + return np.asarray(self._facecolors2d) + + def get_edgecolor(self): + # docstring inherited + # self._edgecolors2d is not initialized until do_3d_projection + if not hasattr(self, '_edgecolors2d'): + self.axes.M = self.axes.get_proj() + self.do_3d_projection() + return np.asarray(self._edgecolors2d) + + +def poly_collection_2d_to_3d(col, zs=0, zdir='z', axlim_clip=False): + """ + Convert a `.PolyCollection` into a `.Poly3DCollection` object. + + Parameters + ---------- + col : `~matplotlib.collections.PolyCollection` + The collection to convert. + zs : float or array of floats + The location or locations to place the polygons in the collection along + the *zdir* axis. Default: 0. + zdir : {'x', 'y', 'z'} + The axis in which to place the patches. Default: 'z'. + See `.get_dir_vector` for a description of the values. + """ + segments_3d, codes = _paths_to_3d_segments_with_codes( + col.get_paths(), zs, zdir) + col.__class__ = Poly3DCollection + col.set_verts_and_codes(segments_3d, codes) + col.set_3d_properties() + col._axlim_clip = axlim_clip + + +def juggle_axes(xs, ys, zs, zdir): + """ + Reorder coordinates so that 2D *xs*, *ys* can be plotted in the plane + orthogonal to *zdir*. *zdir* is normally 'x', 'y' or 'z'. However, if + *zdir* starts with a '-' it is interpreted as a compensation for + `rotate_axes`. + """ + if zdir == 'x': + return zs, xs, ys + elif zdir == 'y': + return xs, zs, ys + elif zdir[0] == '-': + return rotate_axes(xs, ys, zs, zdir) + else: + return xs, ys, zs + + +def rotate_axes(xs, ys, zs, zdir): + """ + Reorder coordinates so that the axes are rotated with *zdir* along + the original z axis. Prepending the axis with a '-' does the + inverse transform, so *zdir* can be 'x', '-x', 'y', '-y', 'z' or '-z'. + """ + if zdir in ('x', '-y'): + return ys, zs, xs + elif zdir in ('-x', 'y'): + return zs, xs, ys + else: + return xs, ys, zs + + +def _zalpha(colors, zs): + """Modify the alphas of the color list according to depth.""" + # FIXME: This only works well if the points for *zs* are well-spaced + # in all three dimensions. Otherwise, at certain orientations, + # the min and max zs are very close together. + # Should really normalize against the viewing depth. + if len(colors) == 0 or len(zs) == 0: + return np.zeros((0, 4)) + norm = Normalize(min(zs), max(zs)) + sats = 1 - norm(zs) * 0.7 + rgba = np.broadcast_to(mcolors.to_rgba_array(colors), (len(zs), 4)) + return np.column_stack([rgba[:, :3], rgba[:, 3] * sats]) + + +def _all_points_on_plane(xs, ys, zs, atol=1e-8): + """ + Check if all points are on the same plane. Note that NaN values are + ignored. + + Parameters + ---------- + xs, ys, zs : array-like + The x, y, and z coordinates of the points. + atol : float, default: 1e-8 + The tolerance for the equality check. + """ + xs, ys, zs = np.asarray(xs), np.asarray(ys), np.asarray(zs) + points = np.column_stack([xs, ys, zs]) + points = points[~np.isnan(points).any(axis=1)] + # Check for the case where we have less than 3 unique points + points = np.unique(points, axis=0) + if len(points) <= 3: + return True + # Calculate the vectors from the first point to all other points + vs = (points - points[0])[1:] + vs = vs / np.linalg.norm(vs, axis=1)[:, np.newaxis] + # Filter out parallel vectors + vs = np.unique(vs, axis=0) + if len(vs) <= 2: + return True + # Filter out parallel and antiparallel vectors to the first vector + cross_norms = np.linalg.norm(np.cross(vs[0], vs[1:]), axis=1) + zero_cross_norms = np.where(np.isclose(cross_norms, 0, atol=atol))[0] + 1 + vs = np.delete(vs, zero_cross_norms, axis=0) + if len(vs) <= 2: + return True + # Calculate the normal vector from the first three points + n = np.cross(vs[0], vs[1]) + n = n / np.linalg.norm(n) + # If the dot product of the normal vector and all other vectors is zero, + # all points are on the same plane + dots = np.dot(n, vs.transpose()) + return np.allclose(dots, 0, atol=atol) + + +def _generate_normals(polygons): + """ + Compute the normals of a list of polygons, one normal per polygon. + + Normals point towards the viewer for a face with its vertices in + counterclockwise order, following the right hand rule. + + Uses three points equally spaced around the polygon. This method assumes + that the points are in a plane. Otherwise, more than one shade is required, + which is not supported. + + Parameters + ---------- + polygons : list of (M_i, 3) array-like, or (..., M, 3) array-like + A sequence of polygons to compute normals for, which can have + varying numbers of vertices. If the polygons all have the same + number of vertices and array is passed, then the operation will + be vectorized. + + Returns + ------- + normals : (..., 3) array + A normal vector estimated for the polygon. + """ + if isinstance(polygons, np.ndarray): + # optimization: polygons all have the same number of points, so can + # vectorize + n = polygons.shape[-2] + i1, i2, i3 = 0, n//3, 2*n//3 + v1 = polygons[..., i1, :] - polygons[..., i2, :] + v2 = polygons[..., i2, :] - polygons[..., i3, :] + else: + # The subtraction doesn't vectorize because polygons is jagged. + v1 = np.empty((len(polygons), 3)) + v2 = np.empty((len(polygons), 3)) + for poly_i, ps in enumerate(polygons): + n = len(ps) + ps = np.asarray(ps) + i1, i2, i3 = 0, n//3, 2*n//3 + v1[poly_i, :] = ps[i1, :] - ps[i2, :] + v2[poly_i, :] = ps[i2, :] - ps[i3, :] + return np.cross(v1, v2) + + +def _shade_colors(color, normals, lightsource=None): + """ + Shade *color* using normal vectors given by *normals*, + assuming a *lightsource* (using default position if not given). + *color* can also be an array of the same length as *normals*. + """ + if lightsource is None: + # chosen for backwards-compatibility + lightsource = mcolors.LightSource(azdeg=225, altdeg=19.4712) + + with np.errstate(invalid="ignore"): + shade = ((normals / np.linalg.norm(normals, axis=1, keepdims=True)) + @ lightsource.direction) + mask = ~np.isnan(shade) + + if mask.any(): + # convert dot product to allowed shading fractions + in_norm = mcolors.Normalize(-1, 1) + out_norm = mcolors.Normalize(0.3, 1).inverse + + def norm(x): + return out_norm(in_norm(x)) + + shade[~mask] = 0 + + color = mcolors.to_rgba_array(color) + # shape of color should be (M, 4) (where M is number of faces) + # shape of shade should be (M,) + # colors should have final shape of (M, 4) + alpha = color[:, 3] + colors = norm(shade)[:, np.newaxis] * color + colors[:, 3] = alpha + else: + colors = np.asanyarray(color).copy() + + return colors diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axes3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axes3d.py new file mode 100644 index 0000000000000000000000000000000000000000..d0ba360c314bbaa39efd5e64936a07e6cd0ec0f6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axes3d.py @@ -0,0 +1,4162 @@ +""" +axes3d.py, original mplot3d version by John Porter +Created: 23 Sep 2005 + +Parts fixed by Reinier Heeres +Minor additions by Ben Axelrod +Significant updates and revisions by Ben Root + +Module containing Axes3D, an object which can plot 3D objects on a +2D matplotlib figure. +""" + +from collections import defaultdict +import itertools +import math +import textwrap +import warnings + +import numpy as np + +import matplotlib as mpl +from matplotlib import _api, cbook, _docstring, _preprocess_data +import matplotlib.artist as martist +import matplotlib.collections as mcoll +import matplotlib.colors as mcolors +import matplotlib.image as mimage +import matplotlib.lines as mlines +import matplotlib.patches as mpatches +import matplotlib.container as mcontainer +import matplotlib.transforms as mtransforms +from matplotlib.axes import Axes +from matplotlib.axes._base import _axis_method_wrapper, _process_plot_format +from matplotlib.transforms import Bbox +from matplotlib.tri._triangulation import Triangulation + +from . import art3d +from . import proj3d +from . import axis3d + + +@_docstring.interpd +@_api.define_aliases({ + "xlim": ["xlim3d"], "ylim": ["ylim3d"], "zlim": ["zlim3d"]}) +class Axes3D(Axes): + """ + 3D Axes object. + + .. note:: + + As a user, you do not instantiate Axes directly, but use Axes creation + methods instead; e.g. from `.pyplot` or `.Figure`: + `~.pyplot.subplots`, `~.pyplot.subplot_mosaic` or `.Figure.add_axes`. + """ + name = '3d' + + _axis_names = ("x", "y", "z") + Axes._shared_axes["z"] = cbook.Grouper() + Axes._shared_axes["view"] = cbook.Grouper() + + def __init__( + self, fig, rect=None, *args, + elev=30, azim=-60, roll=0, shareview=None, sharez=None, + proj_type='persp', focal_length=None, + box_aspect=None, + computed_zorder=True, + **kwargs, + ): + """ + Parameters + ---------- + fig : Figure + The parent figure. + rect : tuple (left, bottom, width, height), default: None. + The ``(left, bottom, width, height)`` Axes position. + elev : float, default: 30 + The elevation angle in degrees rotates the camera above and below + the x-y plane, with a positive angle corresponding to a location + above the plane. + azim : float, default: -60 + The azimuthal angle in degrees rotates the camera about the z axis, + with a positive angle corresponding to a right-handed rotation. In + other words, a positive azimuth rotates the camera about the origin + from its location along the +x axis towards the +y axis. + roll : float, default: 0 + The roll angle in degrees rotates the camera about the viewing + axis. A positive angle spins the camera clockwise, causing the + scene to rotate counter-clockwise. + shareview : Axes3D, optional + Other Axes to share view angles with. Note that it is not possible + to unshare axes. + sharez : Axes3D, optional + Other Axes to share z-limits with. Note that it is not possible to + unshare axes. + proj_type : {'persp', 'ortho'} + The projection type, default 'persp'. + focal_length : float, default: None + For a projection type of 'persp', the focal length of the virtual + camera. Must be > 0. If None, defaults to 1. + For a projection type of 'ortho', must be set to either None + or infinity (numpy.inf). If None, defaults to infinity. + The focal length can be computed from a desired Field Of View via + the equation: focal_length = 1/tan(FOV/2) + box_aspect : 3-tuple of floats, default: None + Changes the physical dimensions of the Axes3D, such that the ratio + of the axis lengths in display units is x:y:z. + If None, defaults to 4:4:3 + computed_zorder : bool, default: True + If True, the draw order is computed based on the average position + of the `.Artist`\\s along the view direction. + Set to False if you want to manually control the order in which + Artists are drawn on top of each other using their *zorder* + attribute. This can be used for fine-tuning if the automatic order + does not produce the desired result. Note however, that a manual + zorder will only be correct for a limited view angle. If the figure + is rotated by the user, it will look wrong from certain angles. + + **kwargs + Other optional keyword arguments: + + %(Axes3D:kwdoc)s + """ + + if rect is None: + rect = [0.0, 0.0, 1.0, 1.0] + + self.initial_azim = azim + self.initial_elev = elev + self.initial_roll = roll + self.set_proj_type(proj_type, focal_length) + self.computed_zorder = computed_zorder + + self.xy_viewLim = Bbox.unit() + self.zz_viewLim = Bbox.unit() + xymargin = 0.05 * 10/11 # match mpl3.8 appearance + self.xy_dataLim = Bbox([[xymargin, xymargin], + [1 - xymargin, 1 - xymargin]]) + # z-limits are encoded in the x-component of the Bbox, y is un-used + self.zz_dataLim = Bbox.unit() + + # inhibit autoscale_view until the axes are defined + # they can't be defined until Axes.__init__ has been called + self.view_init(self.initial_elev, self.initial_azim, self.initial_roll) + + self._sharez = sharez + if sharez is not None: + self._shared_axes["z"].join(self, sharez) + self._adjustable = 'datalim' + + self._shareview = shareview + if shareview is not None: + self._shared_axes["view"].join(self, shareview) + + if kwargs.pop('auto_add_to_figure', False): + raise AttributeError( + 'auto_add_to_figure is no longer supported for Axes3D. ' + 'Use fig.add_axes(ax) instead.' + ) + + super().__init__( + fig, rect, frameon=True, box_aspect=box_aspect, *args, **kwargs + ) + # Disable drawing of axes by base class + super().set_axis_off() + # Enable drawing of axes by Axes3D class + self.set_axis_on() + self.M = None + self.invM = None + + self._view_margin = 1/48 # default value to match mpl3.8 + self.autoscale_view() + + # func used to format z -- fall back on major formatters + self.fmt_zdata = None + + self.mouse_init() + fig = self.get_figure(root=True) + fig.canvas.callbacks._connect_picklable( + 'motion_notify_event', self._on_move) + fig.canvas.callbacks._connect_picklable( + 'button_press_event', self._button_press) + fig.canvas.callbacks._connect_picklable( + 'button_release_event', self._button_release) + self.set_top_view() + + self.patch.set_linewidth(0) + # Calculate the pseudo-data width and height + pseudo_bbox = self.transLimits.inverted().transform([(0, 0), (1, 1)]) + self._pseudo_w, self._pseudo_h = pseudo_bbox[1] - pseudo_bbox[0] + + # mplot3d currently manages its own spines and needs these turned off + # for bounding box calculations + self.spines[:].set_visible(False) + + def set_axis_off(self): + self._axis3don = False + self.stale = True + + def set_axis_on(self): + self._axis3don = True + self.stale = True + + def convert_zunits(self, z): + """ + For artists in an Axes, if the zaxis has units support, + convert *z* using zaxis unit type + """ + return self.zaxis.convert_units(z) + + def set_top_view(self): + # this happens to be the right view for the viewing coordinates + # moved up and to the left slightly to fit labels and axes + xdwl = 0.95 / self._dist + xdw = 0.9 / self._dist + ydwl = 0.95 / self._dist + ydw = 0.9 / self._dist + # Set the viewing pane. + self.viewLim.intervalx = (-xdwl, xdw) + self.viewLim.intervaly = (-ydwl, ydw) + self.stale = True + + def _init_axis(self): + """Init 3D Axes; overrides creation of regular X/Y Axes.""" + self.xaxis = axis3d.XAxis(self) + self.yaxis = axis3d.YAxis(self) + self.zaxis = axis3d.ZAxis(self) + + def get_zaxis(self): + """Return the ``ZAxis`` (`~.axis3d.Axis`) instance.""" + return self.zaxis + + get_zgridlines = _axis_method_wrapper("zaxis", "get_gridlines") + get_zticklines = _axis_method_wrapper("zaxis", "get_ticklines") + + def _transformed_cube(self, vals): + """Return cube with limits from *vals* transformed by self.M.""" + minx, maxx, miny, maxy, minz, maxz = vals + xyzs = [(minx, miny, minz), + (maxx, miny, minz), + (maxx, maxy, minz), + (minx, maxy, minz), + (minx, miny, maxz), + (maxx, miny, maxz), + (maxx, maxy, maxz), + (minx, maxy, maxz)] + return proj3d._proj_points(xyzs, self.M) + + def set_aspect(self, aspect, adjustable=None, anchor=None, share=False): + """ + Set the aspect ratios. + + Parameters + ---------- + aspect : {'auto', 'equal', 'equalxy', 'equalxz', 'equalyz'} + Possible values: + + ========= ================================================== + value description + ========= ================================================== + 'auto' automatic; fill the position rectangle with data. + 'equal' adapt all the axes to have equal aspect ratios. + 'equalxy' adapt the x and y axes to have equal aspect ratios. + 'equalxz' adapt the x and z axes to have equal aspect ratios. + 'equalyz' adapt the y and z axes to have equal aspect ratios. + ========= ================================================== + + adjustable : None or {'box', 'datalim'}, optional + If not *None*, this defines which parameter will be adjusted to + meet the required aspect. See `.set_adjustable` for further + details. + + anchor : None or str or 2-tuple of float, optional + If not *None*, this defines where the Axes will be drawn if there + is extra space due to aspect constraints. The most common way to + specify the anchor are abbreviations of cardinal directions: + + ===== ===================== + value description + ===== ===================== + 'C' centered + 'SW' lower left corner + 'S' middle of bottom edge + 'SE' lower right corner + etc. + ===== ===================== + + See `~.Axes.set_anchor` for further details. + + share : bool, default: False + If ``True``, apply the settings to all shared Axes. + + See Also + -------- + mpl_toolkits.mplot3d.axes3d.Axes3D.set_box_aspect + """ + _api.check_in_list(('auto', 'equal', 'equalxy', 'equalyz', 'equalxz'), + aspect=aspect) + super().set_aspect( + aspect='auto', adjustable=adjustable, anchor=anchor, share=share) + self._aspect = aspect + + if aspect in ('equal', 'equalxy', 'equalxz', 'equalyz'): + ax_indices = self._equal_aspect_axis_indices(aspect) + + view_intervals = np.array([self.xaxis.get_view_interval(), + self.yaxis.get_view_interval(), + self.zaxis.get_view_interval()]) + ptp = np.ptp(view_intervals, axis=1) + if self._adjustable == 'datalim': + mean = np.mean(view_intervals, axis=1) + scale = max(ptp[ax_indices] / self._box_aspect[ax_indices]) + deltas = scale * self._box_aspect + + for i, set_lim in enumerate((self.set_xlim3d, + self.set_ylim3d, + self.set_zlim3d)): + if i in ax_indices: + set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2., + auto=True, view_margin=None) + else: # 'box' + # Change the box aspect such that the ratio of the length of + # the unmodified axis to the length of the diagonal + # perpendicular to it remains unchanged. + box_aspect = np.array(self._box_aspect) + box_aspect[ax_indices] = ptp[ax_indices] + remaining_ax_indices = {0, 1, 2}.difference(ax_indices) + if remaining_ax_indices: + remaining = remaining_ax_indices.pop() + old_diag = np.linalg.norm(self._box_aspect[ax_indices]) + new_diag = np.linalg.norm(box_aspect[ax_indices]) + box_aspect[remaining] *= new_diag / old_diag + self.set_box_aspect(box_aspect) + + def _equal_aspect_axis_indices(self, aspect): + """ + Get the indices for which of the x, y, z axes are constrained to have + equal aspect ratios. + + Parameters + ---------- + aspect : {'auto', 'equal', 'equalxy', 'equalxz', 'equalyz'} + See descriptions in docstring for `.set_aspect()`. + """ + ax_indices = [] # aspect == 'auto' + if aspect == 'equal': + ax_indices = [0, 1, 2] + elif aspect == 'equalxy': + ax_indices = [0, 1] + elif aspect == 'equalxz': + ax_indices = [0, 2] + elif aspect == 'equalyz': + ax_indices = [1, 2] + return ax_indices + + def set_box_aspect(self, aspect, *, zoom=1): + """ + Set the Axes box aspect. + + The box aspect is the ratio of height to width in display + units for each face of the box when viewed perpendicular to + that face. This is not to be confused with the data aspect (see + `~.Axes3D.set_aspect`). The default ratios are 4:4:3 (x:y:z). + + To simulate having equal aspect in data space, set the box + aspect to match your data range in each dimension. + + *zoom* controls the overall size of the Axes3D in the figure. + + Parameters + ---------- + aspect : 3-tuple of floats or None + Changes the physical dimensions of the Axes3D, such that the ratio + of the axis lengths in display units is x:y:z. + If None, defaults to (4, 4, 3). + + zoom : float, default: 1 + Control overall size of the Axes3D in the figure. Must be > 0. + """ + if zoom <= 0: + raise ValueError(f'Argument zoom = {zoom} must be > 0') + + if aspect is None: + aspect = np.asarray((4, 4, 3), dtype=float) + else: + aspect = np.asarray(aspect, dtype=float) + _api.check_shape((3,), aspect=aspect) + # The scale 1.8294640721620434 is tuned to match the mpl3.2 appearance. + # The 25/24 factor is to compensate for the change in automargin + # behavior in mpl3.9. This comes from the padding of 1/48 on both sides + # of the axes in mpl3.8. + aspect *= 1.8294640721620434 * 25/24 * zoom / np.linalg.norm(aspect) + + self._box_aspect = self._roll_to_vertical(aspect, reverse=True) + self.stale = True + + def apply_aspect(self, position=None): + if position is None: + position = self.get_position(original=True) + + # in the superclass, we would go through and actually deal with axis + # scales and box/datalim. Those are all irrelevant - all we need to do + # is make sure our coordinate system is square. + trans = self.get_figure().transSubfigure + bb = mtransforms.Bbox.unit().transformed(trans) + # this is the physical aspect of the panel (or figure): + fig_aspect = bb.height / bb.width + + box_aspect = 1 + pb = position.frozen() + pb1 = pb.shrunk_to_aspect(box_aspect, pb, fig_aspect) + self._set_position(pb1.anchored(self.get_anchor(), pb), 'active') + + @martist.allow_rasterization + def draw(self, renderer): + if not self.get_visible(): + return + self._unstale_viewLim() + + # draw the background patch + self.patch.draw(renderer) + self._frameon = False + + # first, set the aspect + # this is duplicated from `axes._base._AxesBase.draw` + # but must be called before any of the artist are drawn as + # it adjusts the view limits and the size of the bounding box + # of the Axes + locator = self.get_axes_locator() + self.apply_aspect(locator(self, renderer) if locator else None) + + # add the projection matrix to the renderer + self.M = self.get_proj() + self.invM = np.linalg.inv(self.M) + + collections_and_patches = ( + artist for artist in self._children + if isinstance(artist, (mcoll.Collection, mpatches.Patch)) + and artist.get_visible()) + if self.computed_zorder: + # Calculate projection of collections and patches and zorder + # them. Make sure they are drawn above the grids. + zorder_offset = max(axis.get_zorder() + for axis in self._axis_map.values()) + 1 + collection_zorder = patch_zorder = zorder_offset + + for artist in sorted(collections_and_patches, + key=lambda artist: artist.do_3d_projection(), + reverse=True): + if isinstance(artist, mcoll.Collection): + artist.zorder = collection_zorder + collection_zorder += 1 + elif isinstance(artist, mpatches.Patch): + artist.zorder = patch_zorder + patch_zorder += 1 + else: + for artist in collections_and_patches: + artist.do_3d_projection() + + if self._axis3don: + # Draw panes first + for axis in self._axis_map.values(): + axis.draw_pane(renderer) + # Then gridlines + for axis in self._axis_map.values(): + axis.draw_grid(renderer) + # Then axes, labels, text, and ticks + for axis in self._axis_map.values(): + axis.draw(renderer) + + # Then rest + super().draw(renderer) + + def get_axis_position(self): + tc = self._transformed_cube(self.get_w_lims()) + xhigh = tc[1][2] > tc[2][2] + yhigh = tc[3][2] > tc[2][2] + zhigh = tc[0][2] > tc[2][2] + return xhigh, yhigh, zhigh + + def update_datalim(self, xys, **kwargs): + """ + Not implemented in `~mpl_toolkits.mplot3d.axes3d.Axes3D`. + """ + pass + + get_autoscalez_on = _axis_method_wrapper("zaxis", "_get_autoscale_on") + set_autoscalez_on = _axis_method_wrapper("zaxis", "_set_autoscale_on") + + def get_zmargin(self): + """ + Retrieve autoscaling margin of the z-axis. + + .. versionadded:: 3.9 + + Returns + ------- + zmargin : float + + See Also + -------- + mpl_toolkits.mplot3d.axes3d.Axes3D.set_zmargin + """ + return self._zmargin + + def set_zmargin(self, m): + """ + Set padding of Z data limits prior to autoscaling. + + *m* times the data interval will be added to each end of that interval + before it is used in autoscaling. If *m* is negative, this will clip + the data range instead of expanding it. + + For example, if your data is in the range [0, 2], a margin of 0.1 will + result in a range [-0.2, 2.2]; a margin of -0.1 will result in a range + of [0.2, 1.8]. + + Parameters + ---------- + m : float greater than -0.5 + """ + if m <= -0.5: + raise ValueError("margin must be greater than -0.5") + self._zmargin = m + self._request_autoscale_view("z") + self.stale = True + + def margins(self, *margins, x=None, y=None, z=None, tight=True): + """ + Set or retrieve autoscaling margins. + + See `.Axes.margins` for full documentation. Because this function + applies to 3D Axes, it also takes a *z* argument, and returns + ``(xmargin, ymargin, zmargin)``. + """ + if margins and (x is not None or y is not None or z is not None): + raise TypeError('Cannot pass both positional and keyword ' + 'arguments for x, y, and/or z.') + elif len(margins) == 1: + x = y = z = margins[0] + elif len(margins) == 3: + x, y, z = margins + elif margins: + raise TypeError('Must pass a single positional argument for all ' + 'margins, or one for each margin (x, y, z).') + + if x is None and y is None and z is None: + if tight is not True: + _api.warn_external(f'ignoring tight={tight!r} in get mode') + return self._xmargin, self._ymargin, self._zmargin + + if x is not None: + self.set_xmargin(x) + if y is not None: + self.set_ymargin(y) + if z is not None: + self.set_zmargin(z) + + self.autoscale_view( + tight=tight, scalex=(x is not None), scaley=(y is not None), + scalez=(z is not None) + ) + + def autoscale(self, enable=True, axis='both', tight=None): + """ + Convenience method for simple axis view autoscaling. + + See `.Axes.autoscale` for full documentation. Because this function + applies to 3D Axes, *axis* can also be set to 'z', and setting *axis* + to 'both' autoscales all three axes. + """ + if enable is None: + scalex = True + scaley = True + scalez = True + else: + if axis in ['x', 'both']: + self.set_autoscalex_on(enable) + scalex = self.get_autoscalex_on() + else: + scalex = False + if axis in ['y', 'both']: + self.set_autoscaley_on(enable) + scaley = self.get_autoscaley_on() + else: + scaley = False + if axis in ['z', 'both']: + self.set_autoscalez_on(enable) + scalez = self.get_autoscalez_on() + else: + scalez = False + if scalex: + self._request_autoscale_view("x", tight=tight) + if scaley: + self._request_autoscale_view("y", tight=tight) + if scalez: + self._request_autoscale_view("z", tight=tight) + + def auto_scale_xyz(self, X, Y, Z=None, had_data=None): + # This updates the bounding boxes as to keep a record as to what the + # minimum sized rectangular volume holds the data. + if np.shape(X) == np.shape(Y): + self.xy_dataLim.update_from_data_xy( + np.column_stack([np.ravel(X), np.ravel(Y)]), not had_data) + else: + self.xy_dataLim.update_from_data_x(X, not had_data) + self.xy_dataLim.update_from_data_y(Y, not had_data) + if Z is not None: + self.zz_dataLim.update_from_data_x(Z, not had_data) + # Let autoscale_view figure out how to use this data. + self.autoscale_view() + + def autoscale_view(self, tight=None, + scalex=True, scaley=True, scalez=True): + """ + Autoscale the view limits using the data limits. + + See `.Axes.autoscale_view` for full documentation. Because this + function applies to 3D Axes, it also takes a *scalez* argument. + """ + # This method looks at the rectangular volume (see above) + # of data and decides how to scale the view portal to fit it. + if tight is None: + _tight = self._tight + if not _tight: + # if image data only just use the datalim + for artist in self._children: + if isinstance(artist, mimage.AxesImage): + _tight = True + elif isinstance(artist, (mlines.Line2D, mpatches.Patch)): + _tight = False + break + else: + _tight = self._tight = bool(tight) + + if scalex and self.get_autoscalex_on(): + x0, x1 = self.xy_dataLim.intervalx + xlocator = self.xaxis.get_major_locator() + x0, x1 = xlocator.nonsingular(x0, x1) + if self._xmargin > 0: + delta = (x1 - x0) * self._xmargin + x0 -= delta + x1 += delta + if not _tight: + x0, x1 = xlocator.view_limits(x0, x1) + self.set_xbound(x0, x1, self._view_margin) + + if scaley and self.get_autoscaley_on(): + y0, y1 = self.xy_dataLim.intervaly + ylocator = self.yaxis.get_major_locator() + y0, y1 = ylocator.nonsingular(y0, y1) + if self._ymargin > 0: + delta = (y1 - y0) * self._ymargin + y0 -= delta + y1 += delta + if not _tight: + y0, y1 = ylocator.view_limits(y0, y1) + self.set_ybound(y0, y1, self._view_margin) + + if scalez and self.get_autoscalez_on(): + z0, z1 = self.zz_dataLim.intervalx + zlocator = self.zaxis.get_major_locator() + z0, z1 = zlocator.nonsingular(z0, z1) + if self._zmargin > 0: + delta = (z1 - z0) * self._zmargin + z0 -= delta + z1 += delta + if not _tight: + z0, z1 = zlocator.view_limits(z0, z1) + self.set_zbound(z0, z1, self._view_margin) + + def get_w_lims(self): + """Get 3D world limits.""" + minx, maxx = self.get_xlim3d() + miny, maxy = self.get_ylim3d() + minz, maxz = self.get_zlim3d() + return minx, maxx, miny, maxy, minz, maxz + + def _set_bound3d(self, get_bound, set_lim, axis_inverted, + lower=None, upper=None, view_margin=None): + """ + Set 3D axis bounds. + """ + if upper is None and np.iterable(lower): + lower, upper = lower + + old_lower, old_upper = get_bound() + if lower is None: + lower = old_lower + if upper is None: + upper = old_upper + + set_lim(sorted((lower, upper), reverse=bool(axis_inverted())), + auto=None, view_margin=view_margin) + + def set_xbound(self, lower=None, upper=None, view_margin=None): + """ + Set the lower and upper numerical bounds of the x-axis. + + This method will honor axis inversion regardless of parameter order. + It will not change the autoscaling setting (`.get_autoscalex_on()`). + + Parameters + ---------- + lower, upper : float or None + The lower and upper bounds. If *None*, the respective axis bound + is not modified. + view_margin : float or None + The margin to apply to the bounds. If *None*, the margin is handled + by `.set_xlim`. + + See Also + -------- + get_xbound + get_xlim, set_xlim + invert_xaxis, xaxis_inverted + """ + self._set_bound3d(self.get_xbound, self.set_xlim, self.xaxis_inverted, + lower, upper, view_margin) + + def set_ybound(self, lower=None, upper=None, view_margin=None): + """ + Set the lower and upper numerical bounds of the y-axis. + + This method will honor axis inversion regardless of parameter order. + It will not change the autoscaling setting (`.get_autoscaley_on()`). + + Parameters + ---------- + lower, upper : float or None + The lower and upper bounds. If *None*, the respective axis bound + is not modified. + view_margin : float or None + The margin to apply to the bounds. If *None*, the margin is handled + by `.set_ylim`. + + See Also + -------- + get_ybound + get_ylim, set_ylim + invert_yaxis, yaxis_inverted + """ + self._set_bound3d(self.get_ybound, self.set_ylim, self.yaxis_inverted, + lower, upper, view_margin) + + def set_zbound(self, lower=None, upper=None, view_margin=None): + """ + Set the lower and upper numerical bounds of the z-axis. + This method will honor axis inversion regardless of parameter order. + It will not change the autoscaling setting (`.get_autoscaley_on()`). + + Parameters + ---------- + lower, upper : float or None + The lower and upper bounds. If *None*, the respective axis bound + is not modified. + view_margin : float or None + The margin to apply to the bounds. If *None*, the margin is handled + by `.set_zlim`. + + See Also + -------- + get_zbound + get_zlim, set_zlim + invert_zaxis, zaxis_inverted + """ + self._set_bound3d(self.get_zbound, self.set_zlim, self.zaxis_inverted, + lower, upper, view_margin) + + def _set_lim3d(self, axis, lower=None, upper=None, *, emit=True, + auto=False, view_margin=None, axmin=None, axmax=None): + """ + Set 3D axis limits. + """ + if upper is None: + if np.iterable(lower): + lower, upper = lower + elif axmax is None: + upper = axis.get_view_interval()[1] + if lower is None and axmin is None: + lower = axis.get_view_interval()[0] + if axmin is not None: + if lower is not None: + raise TypeError("Cannot pass both 'lower' and 'min'") + lower = axmin + if axmax is not None: + if upper is not None: + raise TypeError("Cannot pass both 'upper' and 'max'") + upper = axmax + if np.isinf(lower) or np.isinf(upper): + raise ValueError(f"Axis limits {lower}, {upper} cannot be infinite") + if view_margin is None: + if mpl.rcParams['axes3d.automargin']: + view_margin = self._view_margin + else: + view_margin = 0 + delta = (upper - lower) * view_margin + lower -= delta + upper += delta + return axis._set_lim(lower, upper, emit=emit, auto=auto) + + def set_xlim(self, left=None, right=None, *, emit=True, auto=False, + view_margin=None, xmin=None, xmax=None): + """ + Set the 3D x-axis view limits. + + Parameters + ---------- + left : float, optional + The left xlim in data coordinates. Passing *None* leaves the + limit unchanged. + + The left and right xlims may also be passed as the tuple + (*left*, *right*) as the first positional argument (or as + the *left* keyword argument). + + .. ACCEPTS: (left: float, right: float) + + right : float, optional + The right xlim in data coordinates. Passing *None* leaves the + limit unchanged. + + emit : bool, default: True + Whether to notify observers of limit change. + + auto : bool or None, default: False + Whether to turn on autoscaling of the x-axis. *True* turns on, + *False* turns off, *None* leaves unchanged. + + view_margin : float, optional + The additional margin to apply to the limits. + + xmin, xmax : float, optional + They are equivalent to left and right respectively, and it is an + error to pass both *xmin* and *left* or *xmax* and *right*. + + Returns + ------- + left, right : (float, float) + The new x-axis limits in data coordinates. + + See Also + -------- + get_xlim + set_xbound, get_xbound + invert_xaxis, xaxis_inverted + + Notes + ----- + The *left* value may be greater than the *right* value, in which + case the x-axis values will decrease from *left* to *right*. + + Examples + -------- + >>> set_xlim(left, right) + >>> set_xlim((left, right)) + >>> left, right = set_xlim(left, right) + + One limit may be left unchanged. + + >>> set_xlim(right=right_lim) + + Limits may be passed in reverse order to flip the direction of + the x-axis. For example, suppose ``x`` represents depth of the + ocean in m. The x-axis limits might be set like the following + so 5000 m depth is at the left of the plot and the surface, + 0 m, is at the right. + + >>> set_xlim(5000, 0) + """ + return self._set_lim3d(self.xaxis, left, right, emit=emit, auto=auto, + view_margin=view_margin, axmin=xmin, axmax=xmax) + + def set_ylim(self, bottom=None, top=None, *, emit=True, auto=False, + view_margin=None, ymin=None, ymax=None): + """ + Set the 3D y-axis view limits. + + Parameters + ---------- + bottom : float, optional + The bottom ylim in data coordinates. Passing *None* leaves the + limit unchanged. + + The bottom and top ylims may also be passed as the tuple + (*bottom*, *top*) as the first positional argument (or as + the *bottom* keyword argument). + + .. ACCEPTS: (bottom: float, top: float) + + top : float, optional + The top ylim in data coordinates. Passing *None* leaves the + limit unchanged. + + emit : bool, default: True + Whether to notify observers of limit change. + + auto : bool or None, default: False + Whether to turn on autoscaling of the y-axis. *True* turns on, + *False* turns off, *None* leaves unchanged. + + view_margin : float, optional + The additional margin to apply to the limits. + + ymin, ymax : float, optional + They are equivalent to bottom and top respectively, and it is an + error to pass both *ymin* and *bottom* or *ymax* and *top*. + + Returns + ------- + bottom, top : (float, float) + The new y-axis limits in data coordinates. + + See Also + -------- + get_ylim + set_ybound, get_ybound + invert_yaxis, yaxis_inverted + + Notes + ----- + The *bottom* value may be greater than the *top* value, in which + case the y-axis values will decrease from *bottom* to *top*. + + Examples + -------- + >>> set_ylim(bottom, top) + >>> set_ylim((bottom, top)) + >>> bottom, top = set_ylim(bottom, top) + + One limit may be left unchanged. + + >>> set_ylim(top=top_lim) + + Limits may be passed in reverse order to flip the direction of + the y-axis. For example, suppose ``y`` represents depth of the + ocean in m. The y-axis limits might be set like the following + so 5000 m depth is at the bottom of the plot and the surface, + 0 m, is at the top. + + >>> set_ylim(5000, 0) + """ + return self._set_lim3d(self.yaxis, bottom, top, emit=emit, auto=auto, + view_margin=view_margin, axmin=ymin, axmax=ymax) + + def set_zlim(self, bottom=None, top=None, *, emit=True, auto=False, + view_margin=None, zmin=None, zmax=None): + """ + Set the 3D z-axis view limits. + + Parameters + ---------- + bottom : float, optional + The bottom zlim in data coordinates. Passing *None* leaves the + limit unchanged. + + The bottom and top zlims may also be passed as the tuple + (*bottom*, *top*) as the first positional argument (or as + the *bottom* keyword argument). + + .. ACCEPTS: (bottom: float, top: float) + + top : float, optional + The top zlim in data coordinates. Passing *None* leaves the + limit unchanged. + + emit : bool, default: True + Whether to notify observers of limit change. + + auto : bool or None, default: False + Whether to turn on autoscaling of the z-axis. *True* turns on, + *False* turns off, *None* leaves unchanged. + + view_margin : float, optional + The additional margin to apply to the limits. + + zmin, zmax : float, optional + They are equivalent to bottom and top respectively, and it is an + error to pass both *zmin* and *bottom* or *zmax* and *top*. + + Returns + ------- + bottom, top : (float, float) + The new z-axis limits in data coordinates. + + See Also + -------- + get_zlim + set_zbound, get_zbound + invert_zaxis, zaxis_inverted + + Notes + ----- + The *bottom* value may be greater than the *top* value, in which + case the z-axis values will decrease from *bottom* to *top*. + + Examples + -------- + >>> set_zlim(bottom, top) + >>> set_zlim((bottom, top)) + >>> bottom, top = set_zlim(bottom, top) + + One limit may be left unchanged. + + >>> set_zlim(top=top_lim) + + Limits may be passed in reverse order to flip the direction of + the z-axis. For example, suppose ``z`` represents depth of the + ocean in m. The z-axis limits might be set like the following + so 5000 m depth is at the bottom of the plot and the surface, + 0 m, is at the top. + + >>> set_zlim(5000, 0) + """ + return self._set_lim3d(self.zaxis, bottom, top, emit=emit, auto=auto, + view_margin=view_margin, axmin=zmin, axmax=zmax) + + set_xlim3d = set_xlim + set_ylim3d = set_ylim + set_zlim3d = set_zlim + + def get_xlim(self): + # docstring inherited + return tuple(self.xy_viewLim.intervalx) + + def get_ylim(self): + # docstring inherited + return tuple(self.xy_viewLim.intervaly) + + def get_zlim(self): + """ + Return the 3D z-axis view limits. + + Returns + ------- + left, right : (float, float) + The current z-axis limits in data coordinates. + + See Also + -------- + set_zlim + set_zbound, get_zbound + invert_zaxis, zaxis_inverted + + Notes + ----- + The z-axis may be inverted, in which case the *left* value will + be greater than the *right* value. + """ + return tuple(self.zz_viewLim.intervalx) + + get_zscale = _axis_method_wrapper("zaxis", "get_scale") + + # Redefine all three methods to overwrite their docstrings. + set_xscale = _axis_method_wrapper("xaxis", "_set_axes_scale") + set_yscale = _axis_method_wrapper("yaxis", "_set_axes_scale") + set_zscale = _axis_method_wrapper("zaxis", "_set_axes_scale") + set_xscale.__doc__, set_yscale.__doc__, set_zscale.__doc__ = map( + """ + Set the {}-axis scale. + + Parameters + ---------- + value : {{"linear"}} + The axis scale type to apply. 3D Axes currently only support + linear scales; other scales yield nonsensical results. + + **kwargs + Keyword arguments are nominally forwarded to the scale class, but + none of them is applicable for linear scales. + """.format, + ["x", "y", "z"]) + + get_zticks = _axis_method_wrapper("zaxis", "get_ticklocs") + set_zticks = _axis_method_wrapper("zaxis", "set_ticks") + get_zmajorticklabels = _axis_method_wrapper("zaxis", "get_majorticklabels") + get_zminorticklabels = _axis_method_wrapper("zaxis", "get_minorticklabels") + get_zticklabels = _axis_method_wrapper("zaxis", "get_ticklabels") + set_zticklabels = _axis_method_wrapper( + "zaxis", "set_ticklabels", + doc_sub={"Axis.set_ticks": "Axes3D.set_zticks"}) + + zaxis_date = _axis_method_wrapper("zaxis", "axis_date") + if zaxis_date.__doc__: + zaxis_date.__doc__ += textwrap.dedent(""" + + Notes + ----- + This function is merely provided for completeness, but 3D Axes do not + support dates for ticks, and so this may not work as expected. + """) + + def clabel(self, *args, **kwargs): + """Currently not implemented for 3D Axes, and returns *None*.""" + return None + + def view_init(self, elev=None, azim=None, roll=None, vertical_axis="z", + share=False): + """ + Set the elevation and azimuth of the Axes in degrees (not radians). + + This can be used to rotate the Axes programmatically. + + To look normal to the primary planes, the following elevation and + azimuth angles can be used. A roll angle of 0, 90, 180, or 270 deg + will rotate these views while keeping the axes at right angles. + + ========== ==== ==== + view plane elev azim + ========== ==== ==== + XY 90 -90 + XZ 0 -90 + YZ 0 0 + -XY -90 90 + -XZ 0 90 + -YZ 0 180 + ========== ==== ==== + + Parameters + ---------- + elev : float, default: None + The elevation angle in degrees rotates the camera above the plane + pierced by the vertical axis, with a positive angle corresponding + to a location above that plane. For example, with the default + vertical axis of 'z', the elevation defines the angle of the camera + location above the x-y plane. + If None, then the initial value as specified in the `Axes3D` + constructor is used. + azim : float, default: None + The azimuthal angle in degrees rotates the camera about the + vertical axis, with a positive angle corresponding to a + right-handed rotation. For example, with the default vertical axis + of 'z', a positive azimuth rotates the camera about the origin from + its location along the +x axis towards the +y axis. + If None, then the initial value as specified in the `Axes3D` + constructor is used. + roll : float, default: None + The roll angle in degrees rotates the camera about the viewing + axis. A positive angle spins the camera clockwise, causing the + scene to rotate counter-clockwise. + If None, then the initial value as specified in the `Axes3D` + constructor is used. + vertical_axis : {"z", "x", "y"}, default: "z" + The axis to align vertically. *azim* rotates about this axis. + share : bool, default: False + If ``True``, apply the settings to all Axes with shared views. + """ + + self._dist = 10 # The camera distance from origin. Behaves like zoom + + if elev is None: + elev = self.initial_elev + if azim is None: + azim = self.initial_azim + if roll is None: + roll = self.initial_roll + vertical_axis = _api.check_getitem( + {name: idx for idx, name in enumerate(self._axis_names)}, + vertical_axis=vertical_axis, + ) + + if share: + axes = {sibling for sibling + in self._shared_axes['view'].get_siblings(self)} + else: + axes = [self] + + for ax in axes: + ax.elev = elev + ax.azim = azim + ax.roll = roll + ax._vertical_axis = vertical_axis + + def set_proj_type(self, proj_type, focal_length=None): + """ + Set the projection type. + + Parameters + ---------- + proj_type : {'persp', 'ortho'} + The projection type. + focal_length : float, default: None + For a projection type of 'persp', the focal length of the virtual + camera. Must be > 0. If None, defaults to 1. + The focal length can be computed from a desired Field Of View via + the equation: focal_length = 1/tan(FOV/2) + """ + _api.check_in_list(['persp', 'ortho'], proj_type=proj_type) + if proj_type == 'persp': + if focal_length is None: + focal_length = 1 + elif focal_length <= 0: + raise ValueError(f"focal_length = {focal_length} must be " + "greater than 0") + self._focal_length = focal_length + else: # 'ortho': + if focal_length not in (None, np.inf): + raise ValueError(f"focal_length = {focal_length} must be " + f"None for proj_type = {proj_type}") + self._focal_length = np.inf + + def _roll_to_vertical( + self, arr: "np.typing.ArrayLike", reverse: bool = False + ) -> np.ndarray: + """ + Roll arrays to match the different vertical axis. + + Parameters + ---------- + arr : ArrayLike + Array to roll. + reverse : bool, default: False + Reverse the direction of the roll. + """ + if reverse: + return np.roll(arr, (self._vertical_axis - 2) * -1) + else: + return np.roll(arr, (self._vertical_axis - 2)) + + def get_proj(self): + """Create the projection matrix from the current viewing position.""" + + # Transform to uniform world coordinates 0-1, 0-1, 0-1 + box_aspect = self._roll_to_vertical(self._box_aspect) + worldM = proj3d.world_transformation( + *self.get_xlim3d(), + *self.get_ylim3d(), + *self.get_zlim3d(), + pb_aspect=box_aspect, + ) + + # Look into the middle of the world coordinates: + R = 0.5 * box_aspect + + # elev: elevation angle in the z plane. + # azim: azimuth angle in the xy plane. + # Coordinates for a point that rotates around the box of data. + # p0, p1 corresponds to rotating the box only around the vertical axis. + # p2 corresponds to rotating the box only around the horizontal axis. + elev_rad = np.deg2rad(self.elev) + azim_rad = np.deg2rad(self.azim) + p0 = np.cos(elev_rad) * np.cos(azim_rad) + p1 = np.cos(elev_rad) * np.sin(azim_rad) + p2 = np.sin(elev_rad) + + # When changing vertical axis the coordinates changes as well. + # Roll the values to get the same behaviour as the default: + ps = self._roll_to_vertical([p0, p1, p2]) + + # The coordinates for the eye viewing point. The eye is looking + # towards the middle of the box of data from a distance: + eye = R + self._dist * ps + + # Calculate the viewing axes for the eye position + u, v, w = self._calc_view_axes(eye) + self._view_u = u # _view_u is towards the right of the screen + self._view_v = v # _view_v is towards the top of the screen + self._view_w = w # _view_w is out of the screen + + # Generate the view and projection transformation matrices + if self._focal_length == np.inf: + # Orthographic projection + viewM = proj3d._view_transformation_uvw(u, v, w, eye) + projM = proj3d._ortho_transformation(-self._dist, self._dist) + else: + # Perspective projection + # Scale the eye dist to compensate for the focal length zoom effect + eye_focal = R + self._dist * ps * self._focal_length + viewM = proj3d._view_transformation_uvw(u, v, w, eye_focal) + projM = proj3d._persp_transformation(-self._dist, + self._dist, + self._focal_length) + + # Combine all the transformation matrices to get the final projection + M0 = np.dot(viewM, worldM) + M = np.dot(projM, M0) + return M + + def mouse_init(self, rotate_btn=1, pan_btn=2, zoom_btn=3): + """ + Set the mouse buttons for 3D rotation and zooming. + + Parameters + ---------- + rotate_btn : int or list of int, default: 1 + The mouse button or buttons to use for 3D rotation of the Axes. + pan_btn : int or list of int, default: 2 + The mouse button or buttons to use to pan the 3D Axes. + zoom_btn : int or list of int, default: 3 + The mouse button or buttons to use to zoom the 3D Axes. + """ + self.button_pressed = None + # coerce scalars into array-like, then convert into + # a regular list to avoid comparisons against None + # which breaks in recent versions of numpy. + self._rotate_btn = np.atleast_1d(rotate_btn).tolist() + self._pan_btn = np.atleast_1d(pan_btn).tolist() + self._zoom_btn = np.atleast_1d(zoom_btn).tolist() + + def disable_mouse_rotation(self): + """Disable mouse buttons for 3D rotation, panning, and zooming.""" + self.mouse_init(rotate_btn=[], pan_btn=[], zoom_btn=[]) + + def can_zoom(self): + # doc-string inherited + return True + + def can_pan(self): + # doc-string inherited + return True + + def sharez(self, other): + """ + Share the z-axis with *other*. + + This is equivalent to passing ``sharez=other`` when constructing the + Axes, and cannot be used if the z-axis is already being shared with + another Axes. Note that it is not possible to unshare axes. + """ + _api.check_isinstance(Axes3D, other=other) + if self._sharez is not None and other is not self._sharez: + raise ValueError("z-axis is already shared") + self._shared_axes["z"].join(self, other) + self._sharez = other + self.zaxis.major = other.zaxis.major # Ticker instances holding + self.zaxis.minor = other.zaxis.minor # locator and formatter. + z0, z1 = other.get_zlim() + self.set_zlim(z0, z1, emit=False, auto=other.get_autoscalez_on()) + self.zaxis._scale = other.zaxis._scale + + def shareview(self, other): + """ + Share the view angles with *other*. + + This is equivalent to passing ``shareview=other`` when constructing the + Axes, and cannot be used if the view angles are already being shared + with another Axes. Note that it is not possible to unshare axes. + """ + _api.check_isinstance(Axes3D, other=other) + if self._shareview is not None and other is not self._shareview: + raise ValueError("view angles are already shared") + self._shared_axes["view"].join(self, other) + self._shareview = other + vertical_axis = self._axis_names[other._vertical_axis] + self.view_init(elev=other.elev, azim=other.azim, roll=other.roll, + vertical_axis=vertical_axis, share=True) + + def clear(self): + # docstring inherited. + super().clear() + if self._focal_length == np.inf: + self._zmargin = mpl.rcParams['axes.zmargin'] + else: + self._zmargin = 0. + + xymargin = 0.05 * 10/11 # match mpl3.8 appearance + self.xy_dataLim = Bbox([[xymargin, xymargin], + [1 - xymargin, 1 - xymargin]]) + # z-limits are encoded in the x-component of the Bbox, y is un-used + self.zz_dataLim = Bbox.unit() + self._view_margin = 1/48 # default value to match mpl3.8 + self.autoscale_view() + + self.grid(mpl.rcParams['axes3d.grid']) + + def _button_press(self, event): + if event.inaxes == self: + self.button_pressed = event.button + self._sx, self._sy = event.xdata, event.ydata + toolbar = self.get_figure(root=True).canvas.toolbar + if toolbar and toolbar._nav_stack() is None: + toolbar.push_current() + if toolbar: + toolbar.set_message(toolbar._mouse_event_to_message(event)) + + def _button_release(self, event): + self.button_pressed = None + toolbar = self.get_figure(root=True).canvas.toolbar + # backend_bases.release_zoom and backend_bases.release_pan call + # push_current, so check the navigation mode so we don't call it twice + if toolbar and self.get_navigate_mode() is None: + toolbar.push_current() + if toolbar: + toolbar.set_message(toolbar._mouse_event_to_message(event)) + + def _get_view(self): + # docstring inherited + return { + "xlim": self.get_xlim(), "autoscalex_on": self.get_autoscalex_on(), + "ylim": self.get_ylim(), "autoscaley_on": self.get_autoscaley_on(), + "zlim": self.get_zlim(), "autoscalez_on": self.get_autoscalez_on(), + }, (self.elev, self.azim, self.roll) + + def _set_view(self, view): + # docstring inherited + props, (elev, azim, roll) = view + self.set(**props) + self.elev = elev + self.azim = azim + self.roll = roll + + def format_zdata(self, z): + """ + Return *z* string formatted. This function will use the + :attr:`fmt_zdata` attribute if it is callable, else will fall + back on the zaxis major formatter + """ + try: + return self.fmt_zdata(z) + except (AttributeError, TypeError): + func = self.zaxis.get_major_formatter().format_data_short + val = func(z) + return val + + def format_coord(self, xv, yv, renderer=None): + """ + Return a string giving the current view rotation angles, or the x, y, z + coordinates of the point on the nearest axis pane underneath the mouse + cursor, depending on the mouse button pressed. + """ + coords = '' + + if self.button_pressed in self._rotate_btn: + # ignore xv and yv and display angles instead + coords = self._rotation_coords() + + elif self.M is not None: + coords = self._location_coords(xv, yv, renderer) + + return coords + + def _rotation_coords(self): + """ + Return the rotation angles as a string. + """ + norm_elev = art3d._norm_angle(self.elev) + norm_azim = art3d._norm_angle(self.azim) + norm_roll = art3d._norm_angle(self.roll) + coords = (f"elevation={norm_elev:.0f}\N{DEGREE SIGN}, " + f"azimuth={norm_azim:.0f}\N{DEGREE SIGN}, " + f"roll={norm_roll:.0f}\N{DEGREE SIGN}" + ).replace("-", "\N{MINUS SIGN}") + return coords + + def _location_coords(self, xv, yv, renderer): + """ + Return the location on the axis pane underneath the cursor as a string. + """ + p1, pane_idx = self._calc_coord(xv, yv, renderer) + xs = self.format_xdata(p1[0]) + ys = self.format_ydata(p1[1]) + zs = self.format_zdata(p1[2]) + if pane_idx == 0: + coords = f'x pane={xs}, y={ys}, z={zs}' + elif pane_idx == 1: + coords = f'x={xs}, y pane={ys}, z={zs}' + elif pane_idx == 2: + coords = f'x={xs}, y={ys}, z pane={zs}' + return coords + + def _get_camera_loc(self): + """ + Returns the current camera location in data coordinates. + """ + cx, cy, cz, dx, dy, dz = self._get_w_centers_ranges() + c = np.array([cx, cy, cz]) + r = np.array([dx, dy, dz]) + + if self._focal_length == np.inf: # orthographic projection + focal_length = 1e9 # large enough to be effectively infinite + else: # perspective projection + focal_length = self._focal_length + eye = c + self._view_w * self._dist * r / self._box_aspect * focal_length + return eye + + def _calc_coord(self, xv, yv, renderer=None): + """ + Given the 2D view coordinates, find the point on the nearest axis pane + that lies directly below those coordinates. Returns a 3D point in data + coordinates. + """ + if self._focal_length == np.inf: # orthographic projection + zv = 1 + else: # perspective projection + zv = -1 / self._focal_length + + # Convert point on view plane to data coordinates + p1 = np.array(proj3d.inv_transform(xv, yv, zv, self.invM)).ravel() + + # Get the vector from the camera to the point on the view plane + vec = self._get_camera_loc() - p1 + + # Get the pane locations for each of the axes + pane_locs = [] + for axis in self._axis_map.values(): + xys, loc = axis.active_pane() + pane_locs.append(loc) + + # Find the distance to the nearest pane by projecting the view vector + scales = np.zeros(3) + for i in range(3): + if vec[i] == 0: + scales[i] = np.inf + else: + scales[i] = (p1[i] - pane_locs[i]) / vec[i] + pane_idx = np.argmin(abs(scales)) + scale = scales[pane_idx] + + # Calculate the point on the closest pane + p2 = p1 - scale*vec + return p2, pane_idx + + def _arcball(self, x: float, y: float) -> np.ndarray: + """ + Convert a point (x, y) to a point on a virtual trackball. + + This is Ken Shoemake's arcball (a sphere), modified + to soften the abrupt edge (optionally). + See: Ken Shoemake, "ARCBALL: A user interface for specifying + three-dimensional rotation using a mouse." in + Proceedings of Graphics Interface '92, 1992, pp. 151-156, + https://doi.org/10.20380/GI1992.18 + The smoothing of the edge is inspired by Gavin Bell's arcball + (a sphere combined with a hyperbola), but here, the sphere + is combined with a section of a cylinder, so it has finite support. + """ + s = mpl.rcParams['axes3d.trackballsize'] / 2 + b = mpl.rcParams['axes3d.trackballborder'] / s + x /= s + y /= s + r2 = x*x + y*y + r = np.sqrt(r2) + ra = 1 + b + a = b * (1 + b/2) + ri = 2/(ra + 1/ra) + if r < ri: + p = np.array([np.sqrt(1 - r2), x, y]) + elif r < ra: + dr = ra - r + p = np.array([a - np.sqrt((a + dr) * (a - dr)), x, y]) + p /= np.linalg.norm(p) + else: + p = np.array([0, x/r, y/r]) + return p + + def _on_move(self, event): + """ + Mouse moving. + + By default, button-1 rotates, button-2 pans, and button-3 zooms; + these buttons can be modified via `mouse_init`. + """ + + if not self.button_pressed: + return + + if self.get_navigate_mode() is not None: + # we don't want to rotate if we are zooming/panning + # from the toolbar + return + + if self.M is None: + return + + x, y = event.xdata, event.ydata + # In case the mouse is out of bounds. + if x is None or event.inaxes != self: + return + + dx, dy = x - self._sx, y - self._sy + w = self._pseudo_w + h = self._pseudo_h + + # Rotation + if self.button_pressed in self._rotate_btn: + # rotate viewing point + # get the x and y pixel coords + if dx == 0 and dy == 0: + return + + style = mpl.rcParams['axes3d.mouserotationstyle'] + if style == 'azel': + roll = np.deg2rad(self.roll) + delev = -(dy/h)*180*np.cos(roll) + (dx/w)*180*np.sin(roll) + dazim = -(dy/h)*180*np.sin(roll) - (dx/w)*180*np.cos(roll) + elev = self.elev + delev + azim = self.azim + dazim + roll = self.roll + else: + q = _Quaternion.from_cardan_angles( + *np.deg2rad((self.elev, self.azim, self.roll))) + + if style == 'trackball': + k = np.array([0, -dy/h, dx/w]) + nk = np.linalg.norm(k) + th = nk / mpl.rcParams['axes3d.trackballsize'] + dq = _Quaternion(np.cos(th), k*np.sin(th)/nk) + else: # 'sphere', 'arcball' + current_vec = self._arcball(self._sx/w, self._sy/h) + new_vec = self._arcball(x/w, y/h) + if style == 'sphere': + dq = _Quaternion.rotate_from_to(current_vec, new_vec) + else: # 'arcball' + dq = _Quaternion(0, new_vec) * _Quaternion(0, -current_vec) + + q = dq * q + elev, azim, roll = np.rad2deg(q.as_cardan_angles()) + + # update view + vertical_axis = self._axis_names[self._vertical_axis] + self.view_init( + elev=elev, + azim=azim, + roll=roll, + vertical_axis=vertical_axis, + share=True, + ) + self.stale = True + + # Pan + elif self.button_pressed in self._pan_btn: + # Start the pan event with pixel coordinates + px, py = self.transData.transform([self._sx, self._sy]) + self.start_pan(px, py, 2) + # pan view (takes pixel coordinate input) + self.drag_pan(2, None, event.x, event.y) + self.end_pan() + + # Zoom + elif self.button_pressed in self._zoom_btn: + # zoom view (dragging down zooms in) + scale = h/(h - dy) + self._scale_axis_limits(scale, scale, scale) + + # Store the event coordinates for the next time through. + self._sx, self._sy = x, y + # Always request a draw update at the end of interaction + self.get_figure(root=True).canvas.draw_idle() + + def drag_pan(self, button, key, x, y): + # docstring inherited + + # Get the coordinates from the move event + p = self._pan_start + (xdata, ydata), (xdata_start, ydata_start) = p.trans_inverse.transform( + [(x, y), (p.x, p.y)]) + self._sx, self._sy = xdata, ydata + # Calling start_pan() to set the x/y of this event as the starting + # move location for the next event + self.start_pan(x, y, button) + du, dv = xdata - xdata_start, ydata - ydata_start + dw = 0 + if key == 'x': + dv = 0 + elif key == 'y': + du = 0 + if du == 0 and dv == 0: + return + + # Transform the pan from the view axes to the data axes + R = np.array([self._view_u, self._view_v, self._view_w]) + R = -R / self._box_aspect * self._dist + duvw_projected = R.T @ np.array([du, dv, dw]) + + # Calculate pan distance + minx, maxx, miny, maxy, minz, maxz = self.get_w_lims() + dx = (maxx - minx) * duvw_projected[0] + dy = (maxy - miny) * duvw_projected[1] + dz = (maxz - minz) * duvw_projected[2] + + # Set the new axis limits + self.set_xlim3d(minx + dx, maxx + dx, auto=None) + self.set_ylim3d(miny + dy, maxy + dy, auto=None) + self.set_zlim3d(minz + dz, maxz + dz, auto=None) + + def _calc_view_axes(self, eye): + """ + Get the unit vectors for the viewing axes in data coordinates. + `u` is towards the right of the screen + `v` is towards the top of the screen + `w` is out of the screen + """ + elev_rad = np.deg2rad(art3d._norm_angle(self.elev)) + roll_rad = np.deg2rad(art3d._norm_angle(self.roll)) + + # Look into the middle of the world coordinates + R = 0.5 * self._roll_to_vertical(self._box_aspect) + + # Define which axis should be vertical. A negative value + # indicates the plot is upside down and therefore the values + # have been reversed: + V = np.zeros(3) + V[self._vertical_axis] = -1 if abs(elev_rad) > np.pi/2 else 1 + + u, v, w = proj3d._view_axes(eye, R, V, roll_rad) + return u, v, w + + def _set_view_from_bbox(self, bbox, direction='in', + mode=None, twinx=False, twiny=False): + """ + Zoom in or out of the bounding box. + + Will center the view in the center of the bounding box, and zoom by + the ratio of the size of the bounding box to the size of the Axes3D. + """ + (start_x, start_y, stop_x, stop_y) = bbox + if mode == 'x': + start_y = self.bbox.min[1] + stop_y = self.bbox.max[1] + elif mode == 'y': + start_x = self.bbox.min[0] + stop_x = self.bbox.max[0] + + # Clip to bounding box limits + start_x, stop_x = np.clip(sorted([start_x, stop_x]), + self.bbox.min[0], self.bbox.max[0]) + start_y, stop_y = np.clip(sorted([start_y, stop_y]), + self.bbox.min[1], self.bbox.max[1]) + + # Move the center of the view to the center of the bbox + zoom_center_x = (start_x + stop_x)/2 + zoom_center_y = (start_y + stop_y)/2 + + ax_center_x = (self.bbox.max[0] + self.bbox.min[0])/2 + ax_center_y = (self.bbox.max[1] + self.bbox.min[1])/2 + + self.start_pan(zoom_center_x, zoom_center_y, 2) + self.drag_pan(2, None, ax_center_x, ax_center_y) + self.end_pan() + + # Calculate zoom level + dx = abs(start_x - stop_x) + dy = abs(start_y - stop_y) + scale_u = dx / (self.bbox.max[0] - self.bbox.min[0]) + scale_v = dy / (self.bbox.max[1] - self.bbox.min[1]) + + # Keep aspect ratios equal + scale = max(scale_u, scale_v) + + # Zoom out + if direction == 'out': + scale = 1 / scale + + self._zoom_data_limits(scale, scale, scale) + + def _zoom_data_limits(self, scale_u, scale_v, scale_w): + """ + Zoom in or out of a 3D plot. + + Will scale the data limits by the scale factors. These will be + transformed to the x, y, z data axes based on the current view angles. + A scale factor > 1 zooms out and a scale factor < 1 zooms in. + + For an Axes that has had its aspect ratio set to 'equal', 'equalxy', + 'equalyz', or 'equalxz', the relevant axes are constrained to zoom + equally. + + Parameters + ---------- + scale_u : float + Scale factor for the u view axis (view screen horizontal). + scale_v : float + Scale factor for the v view axis (view screen vertical). + scale_w : float + Scale factor for the w view axis (view screen depth). + """ + scale = np.array([scale_u, scale_v, scale_w]) + + # Only perform frame conversion if unequal scale factors + if not np.allclose(scale, scale_u): + # Convert the scale factors from the view frame to the data frame + R = np.array([self._view_u, self._view_v, self._view_w]) + S = scale * np.eye(3) + scale = np.linalg.norm(R.T @ S, axis=1) + + # Set the constrained scale factors to the factor closest to 1 + if self._aspect in ('equal', 'equalxy', 'equalxz', 'equalyz'): + ax_idxs = self._equal_aspect_axis_indices(self._aspect) + min_ax_idxs = np.argmin(np.abs(scale[ax_idxs] - 1)) + scale[ax_idxs] = scale[ax_idxs][min_ax_idxs] + + self._scale_axis_limits(scale[0], scale[1], scale[2]) + + def _scale_axis_limits(self, scale_x, scale_y, scale_z): + """ + Keeping the center of the x, y, and z data axes fixed, scale their + limits by scale factors. A scale factor > 1 zooms out and a scale + factor < 1 zooms in. + + Parameters + ---------- + scale_x : float + Scale factor for the x data axis. + scale_y : float + Scale factor for the y data axis. + scale_z : float + Scale factor for the z data axis. + """ + # Get the axis centers and ranges + cx, cy, cz, dx, dy, dz = self._get_w_centers_ranges() + + # Set the scaled axis limits + self.set_xlim3d(cx - dx*scale_x/2, cx + dx*scale_x/2, auto=None) + self.set_ylim3d(cy - dy*scale_y/2, cy + dy*scale_y/2, auto=None) + self.set_zlim3d(cz - dz*scale_z/2, cz + dz*scale_z/2, auto=None) + + def _get_w_centers_ranges(self): + """Get 3D world centers and axis ranges.""" + # Calculate center of axis limits + minx, maxx, miny, maxy, minz, maxz = self.get_w_lims() + cx = (maxx + minx)/2 + cy = (maxy + miny)/2 + cz = (maxz + minz)/2 + + # Calculate range of axis limits + dx = (maxx - minx) + dy = (maxy - miny) + dz = (maxz - minz) + return cx, cy, cz, dx, dy, dz + + def set_zlabel(self, zlabel, fontdict=None, labelpad=None, **kwargs): + """ + Set zlabel. See doc for `.set_ylabel` for description. + """ + if labelpad is not None: + self.zaxis.labelpad = labelpad + return self.zaxis.set_label_text(zlabel, fontdict, **kwargs) + + def get_zlabel(self): + """ + Get the z-label text string. + """ + label = self.zaxis.label + return label.get_text() + + # Axes rectangle characteristics + + # The frame_on methods are not available for 3D axes. + # Python will raise a TypeError if they are called. + get_frame_on = None + set_frame_on = None + + def grid(self, visible=True, **kwargs): + """ + Set / unset 3D grid. + + .. note:: + + Currently, this function does not behave the same as + `.axes.Axes.grid`, but it is intended to eventually support that + behavior. + """ + # TODO: Operate on each axes separately + if len(kwargs): + visible = True + self._draw_grid = visible + self.stale = True + + def tick_params(self, axis='both', **kwargs): + """ + Convenience method for changing the appearance of ticks and + tick labels. + + See `.Axes.tick_params` for full documentation. Because this function + applies to 3D Axes, *axis* can also be set to 'z', and setting *axis* + to 'both' autoscales all three axes. + + Also, because of how Axes3D objects are drawn very differently + from regular 2D Axes, some of these settings may have + ambiguous meaning. For simplicity, the 'z' axis will + accept settings as if it was like the 'y' axis. + + .. note:: + Axes3D currently ignores some of these settings. + """ + _api.check_in_list(['x', 'y', 'z', 'both'], axis=axis) + if axis in ['x', 'y', 'both']: + super().tick_params(axis, **kwargs) + if axis in ['z', 'both']: + zkw = dict(kwargs) + zkw.pop('top', None) + zkw.pop('bottom', None) + zkw.pop('labeltop', None) + zkw.pop('labelbottom', None) + self.zaxis.set_tick_params(**zkw) + + # data limits, ticks, tick labels, and formatting + + def invert_zaxis(self): + """ + Invert the z-axis. + + See Also + -------- + zaxis_inverted + get_zlim, set_zlim + get_zbound, set_zbound + """ + bottom, top = self.get_zlim() + self.set_zlim(top, bottom, auto=None) + + zaxis_inverted = _axis_method_wrapper("zaxis", "get_inverted") + + def get_zbound(self): + """ + Return the lower and upper z-axis bounds, in increasing order. + + See Also + -------- + set_zbound + get_zlim, set_zlim + invert_zaxis, zaxis_inverted + """ + lower, upper = self.get_zlim() + if lower < upper: + return lower, upper + else: + return upper, lower + + def text(self, x, y, z, s, zdir=None, *, axlim_clip=False, **kwargs): + """ + Add the text *s* to the 3D Axes at location *x*, *y*, *z* in data coordinates. + + Parameters + ---------- + x, y, z : float + The position to place the text. + s : str + The text. + zdir : {'x', 'y', 'z', 3-tuple}, optional + The direction to be used as the z-direction. Default: 'z'. + See `.get_dir_vector` for a description of the values. + axlim_clip : bool, default: False + Whether to hide text that is outside the axes view limits. + + .. versionadded:: 3.10 + **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.text`. + + Returns + ------- + `.Text3D` + The created `.Text3D` instance. + """ + text = super().text(x, y, s, **kwargs) + art3d.text_2d_to_3d(text, z, zdir, axlim_clip) + return text + + text3D = text + text2D = Axes.text + + def plot(self, xs, ys, *args, zdir='z', axlim_clip=False, **kwargs): + """ + Plot 2D or 3D data. + + Parameters + ---------- + xs : 1D array-like + x coordinates of vertices. + ys : 1D array-like + y coordinates of vertices. + zs : float or 1D array-like + z coordinates of vertices; either one for all points or one for + each point. + zdir : {'x', 'y', 'z'}, default: 'z' + When plotting 2D data, the direction to use as z. + axlim_clip : bool, default: False + Whether to hide data that is outside the axes view limits. + + .. versionadded:: 3.10 + **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.plot`. + """ + had_data = self.has_data() + + # `zs` can be passed positionally or as keyword; checking whether + # args[0] is a string matches the behavior of 2D `plot` (via + # `_process_plot_var_args`). + if args and not isinstance(args[0], str): + zs, *args = args + if 'zs' in kwargs: + raise TypeError("plot() for multiple values for argument 'zs'") + else: + zs = kwargs.pop('zs', 0) + + xs, ys, zs = cbook._broadcast_with_masks(xs, ys, zs) + + lines = super().plot(xs, ys, *args, **kwargs) + for line in lines: + art3d.line_2d_to_3d(line, zs=zs, zdir=zdir, axlim_clip=axlim_clip) + + xs, ys, zs = art3d.juggle_axes(xs, ys, zs, zdir) + self.auto_scale_xyz(xs, ys, zs, had_data) + return lines + + plot3D = plot + + def fill_between(self, x1, y1, z1, x2, y2, z2, *, + where=None, mode='auto', facecolors=None, shade=None, + axlim_clip=False, **kwargs): + """ + Fill the area between two 3D curves. + + The curves are defined by the points (*x1*, *y1*, *z1*) and + (*x2*, *y2*, *z2*). This creates one or multiple quadrangle + polygons that are filled. All points must be the same length N, or a + single value to be used for all points. + + Parameters + ---------- + x1, y1, z1 : float or 1D array-like + x, y, and z coordinates of vertices for 1st line. + + x2, y2, z2 : float or 1D array-like + x, y, and z coordinates of vertices for 2nd line. + + where : array of bool (length N), optional + Define *where* to exclude some regions from being filled. The + filled regions are defined by the coordinates ``pts[where]``, + for all x, y, and z pts. More precisely, fill between ``pts[i]`` + and ``pts[i+1]`` if ``where[i] and where[i+1]``. Note that this + definition implies that an isolated *True* value between two + *False* values in *where* will not result in filling. Both sides of + the *True* position remain unfilled due to the adjacent *False* + values. + + mode : {'quad', 'polygon', 'auto'}, default: 'auto' + The fill mode. One of: + + - 'quad': A separate quadrilateral polygon is created for each + pair of subsequent points in the two lines. + - 'polygon': The two lines are connected to form a single polygon. + This is faster and can render more cleanly for simple shapes + (e.g. for filling between two lines that lie within a plane). + - 'auto': If the points all lie on the same 3D plane, 'polygon' is + used. Otherwise, 'quad' is used. + + facecolors : list of :mpltype:`color`, default: None + Colors of each individual patch, or a single color to be used for + all patches. + + shade : bool, default: None + Whether to shade the facecolors. If *None*, then defaults to *True* + for 'quad' mode and *False* for 'polygon' mode. + + axlim_clip : bool, default: False + Whether to hide data that is outside the axes view limits. + + .. versionadded:: 3.10 + + **kwargs + All other keyword arguments are passed on to `.Poly3DCollection`. + + Returns + ------- + `.Poly3DCollection` + A `.Poly3DCollection` containing the plotted polygons. + + """ + _api.check_in_list(['auto', 'quad', 'polygon'], mode=mode) + + had_data = self.has_data() + x1, y1, z1, x2, y2, z2 = cbook._broadcast_with_masks(x1, y1, z1, x2, y2, z2) + + if facecolors is None: + facecolors = [self._get_patches_for_fill.get_next_color()] + facecolors = list(mcolors.to_rgba_array(facecolors)) + + if where is None: + where = True + else: + where = np.asarray(where, dtype=bool) + if where.size != x1.size: + raise ValueError(f"where size ({where.size}) does not match " + f"size ({x1.size})") + where = where & ~np.isnan(x1) # NaNs were broadcast in _broadcast_with_masks + + if mode == 'auto': + if art3d._all_points_on_plane(np.concatenate((x1[where], x2[where])), + np.concatenate((y1[where], y2[where])), + np.concatenate((z1[where], z2[where])), + atol=1e-12): + mode = 'polygon' + else: + mode = 'quad' + + if shade is None: + if mode == 'quad': + shade = True + else: + shade = False + + polys = [] + for idx0, idx1 in cbook.contiguous_regions(where): + x1i = x1[idx0:idx1] + y1i = y1[idx0:idx1] + z1i = z1[idx0:idx1] + x2i = x2[idx0:idx1] + y2i = y2[idx0:idx1] + z2i = z2[idx0:idx1] + + if not len(x1i): + continue + + if mode == 'quad': + # Preallocate the array for the region's vertices, and fill it in + n_polys_i = len(x1i) - 1 + polys_i = np.empty((n_polys_i, 4, 3)) + polys_i[:, 0, :] = np.column_stack((x1i[:-1], y1i[:-1], z1i[:-1])) + polys_i[:, 1, :] = np.column_stack((x1i[1:], y1i[1:], z1i[1:])) + polys_i[:, 2, :] = np.column_stack((x2i[1:], y2i[1:], z2i[1:])) + polys_i[:, 3, :] = np.column_stack((x2i[:-1], y2i[:-1], z2i[:-1])) + polys = polys + [*polys_i] + elif mode == 'polygon': + line1 = np.column_stack((x1i, y1i, z1i)) + line2 = np.column_stack((x2i[::-1], y2i[::-1], z2i[::-1])) + poly = np.concatenate((line1, line2), axis=0) + polys.append(poly) + + polyc = art3d.Poly3DCollection(polys, facecolors=facecolors, shade=shade, + axlim_clip=axlim_clip, **kwargs) + self.add_collection(polyc) + + self.auto_scale_xyz([x1, x2], [y1, y2], [z1, z2], had_data) + return polyc + + def plot_surface(self, X, Y, Z, *, norm=None, vmin=None, + vmax=None, lightsource=None, axlim_clip=False, **kwargs): + """ + Create a surface plot. + + By default, it will be colored in shades of a solid color, but it also + supports colormapping by supplying the *cmap* argument. + + .. note:: + + The *rcount* and *ccount* kwargs, which both default to 50, + determine the maximum number of samples used in each direction. If + the input data is larger, it will be downsampled (by slicing) to + these numbers of points. + + .. note:: + + To maximize rendering speed consider setting *rstride* and *cstride* + to divisors of the number of rows minus 1 and columns minus 1 + respectively. For example, given 51 rows rstride can be any of the + divisors of 50. + + Similarly, a setting of *rstride* and *cstride* equal to 1 (or + *rcount* and *ccount* equal the number of rows and columns) can use + the optimized path. + + Parameters + ---------- + X, Y, Z : 2D arrays + Data values. + + rcount, ccount : int + Maximum number of samples used in each direction. If the input + data is larger, it will be downsampled (by slicing) to these + numbers of points. Defaults to 50. + + rstride, cstride : int + Downsampling stride in each direction. These arguments are + mutually exclusive with *rcount* and *ccount*. If only one of + *rstride* or *cstride* is set, the other defaults to 10. + + 'classic' mode uses a default of ``rstride = cstride = 10`` instead + of the new default of ``rcount = ccount = 50``. + + color : :mpltype:`color` + Color of the surface patches. + + cmap : Colormap, optional + Colormap of the surface patches. + + facecolors : list of :mpltype:`color` + Colors of each individual patch. + + norm : `~matplotlib.colors.Normalize`, optional + Normalization for the colormap. + + vmin, vmax : float, optional + Bounds for the normalization. + + shade : bool, default: True + Whether to shade the facecolors. Shading is always disabled when + *cmap* is specified. + + lightsource : `~matplotlib.colors.LightSource`, optional + The lightsource to use when *shade* is True. + + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + + **kwargs + Other keyword arguments are forwarded to `.Poly3DCollection`. + """ + + had_data = self.has_data() + + if Z.ndim != 2: + raise ValueError("Argument Z must be 2-dimensional.") + + Z = cbook._to_unmasked_float_array(Z) + X, Y, Z = np.broadcast_arrays(X, Y, Z) + rows, cols = Z.shape + + has_stride = 'rstride' in kwargs or 'cstride' in kwargs + has_count = 'rcount' in kwargs or 'ccount' in kwargs + + if has_stride and has_count: + raise ValueError("Cannot specify both stride and count arguments") + + rstride = kwargs.pop('rstride', 10) + cstride = kwargs.pop('cstride', 10) + rcount = kwargs.pop('rcount', 50) + ccount = kwargs.pop('ccount', 50) + + if mpl.rcParams['_internal.classic_mode']: + # Strides have priority over counts in classic mode. + # So, only compute strides from counts + # if counts were explicitly given + compute_strides = has_count + else: + # If the strides are provided then it has priority. + # Otherwise, compute the strides from the counts. + compute_strides = not has_stride + + if compute_strides: + rstride = int(max(np.ceil(rows / rcount), 1)) + cstride = int(max(np.ceil(cols / ccount), 1)) + + fcolors = kwargs.pop('facecolors', None) + + cmap = kwargs.get('cmap', None) + shade = kwargs.pop('shade', cmap is None) + if shade is None: + raise ValueError("shade cannot be None.") + + colset = [] # the sampled facecolor + if (rows - 1) % rstride == 0 and \ + (cols - 1) % cstride == 0 and \ + fcolors is None: + polys = np.stack( + [cbook._array_patch_perimeters(a, rstride, cstride) + for a in (X, Y, Z)], + axis=-1) + else: + # evenly spaced, and including both endpoints + row_inds = list(range(0, rows-1, rstride)) + [rows-1] + col_inds = list(range(0, cols-1, cstride)) + [cols-1] + + polys = [] + for rs, rs_next in itertools.pairwise(row_inds): + for cs, cs_next in itertools.pairwise(col_inds): + ps = [ + # +1 ensures we share edges between polygons + cbook._array_perimeter(a[rs:rs_next+1, cs:cs_next+1]) + for a in (X, Y, Z) + ] + # ps = np.stack(ps, axis=-1) + ps = np.array(ps).T + polys.append(ps) + + if fcolors is not None: + colset.append(fcolors[rs][cs]) + + # In cases where there are non-finite values in the data (possibly NaNs from + # masked arrays), artifacts can be introduced. Here check whether such values + # are present and remove them. + if not isinstance(polys, np.ndarray) or not np.isfinite(polys).all(): + new_polys = [] + new_colset = [] + + # Depending on fcolors, colset is either an empty list or has as + # many elements as polys. In the former case new_colset results in + # a list with None entries, that is discarded later. + for p, col in itertools.zip_longest(polys, colset): + new_poly = np.array(p)[np.isfinite(p).all(axis=1)] + if len(new_poly): + new_polys.append(new_poly) + new_colset.append(col) + + # Replace previous polys and, if fcolors is not None, colset + polys = new_polys + if fcolors is not None: + colset = new_colset + + # note that the striding causes some polygons to have more coordinates + # than others + + if fcolors is not None: + polyc = art3d.Poly3DCollection( + polys, edgecolors=colset, facecolors=colset, shade=shade, + lightsource=lightsource, axlim_clip=axlim_clip, **kwargs) + elif cmap: + polyc = art3d.Poly3DCollection(polys, axlim_clip=axlim_clip, **kwargs) + # can't always vectorize, because polys might be jagged + if isinstance(polys, np.ndarray): + avg_z = polys[..., 2].mean(axis=-1) + else: + avg_z = np.array([ps[:, 2].mean() for ps in polys]) + polyc.set_array(avg_z) + if vmin is not None or vmax is not None: + polyc.set_clim(vmin, vmax) + if norm is not None: + polyc.set_norm(norm) + else: + color = kwargs.pop('color', None) + if color is None: + color = self._get_lines.get_next_color() + color = np.array(mcolors.to_rgba(color)) + + polyc = art3d.Poly3DCollection( + polys, facecolors=color, shade=shade, lightsource=lightsource, + axlim_clip=axlim_clip, **kwargs) + + self.add_collection(polyc) + self.auto_scale_xyz(X, Y, Z, had_data) + + return polyc + + def plot_wireframe(self, X, Y, Z, *, axlim_clip=False, **kwargs): + """ + Plot a 3D wireframe. + + .. note:: + + The *rcount* and *ccount* kwargs, which both default to 50, + determine the maximum number of samples used in each direction. If + the input data is larger, it will be downsampled (by slicing) to + these numbers of points. + + Parameters + ---------- + X, Y, Z : 2D arrays + Data values. + + axlim_clip : bool, default: False + Whether to hide lines and patches with vertices outside the axes + view limits. + + .. versionadded:: 3.10 + + rcount, ccount : int + Maximum number of samples used in each direction. If the input + data is larger, it will be downsampled (by slicing) to these + numbers of points. Setting a count to zero causes the data to be + not sampled in the corresponding direction, producing a 3D line + plot rather than a wireframe plot. Defaults to 50. + + rstride, cstride : int + Downsampling stride in each direction. These arguments are + mutually exclusive with *rcount* and *ccount*. If only one of + *rstride* or *cstride* is set, the other defaults to 1. Setting a + stride to zero causes the data to be not sampled in the + corresponding direction, producing a 3D line plot rather than a + wireframe plot. + + 'classic' mode uses a default of ``rstride = cstride = 1`` instead + of the new default of ``rcount = ccount = 50``. + + **kwargs + Other keyword arguments are forwarded to `.Line3DCollection`. + """ + + had_data = self.has_data() + if Z.ndim != 2: + raise ValueError("Argument Z must be 2-dimensional.") + # FIXME: Support masked arrays + X, Y, Z = np.broadcast_arrays(X, Y, Z) + rows, cols = Z.shape + + has_stride = 'rstride' in kwargs or 'cstride' in kwargs + has_count = 'rcount' in kwargs or 'ccount' in kwargs + + if has_stride and has_count: + raise ValueError("Cannot specify both stride and count arguments") + + rstride = kwargs.pop('rstride', 1) + cstride = kwargs.pop('cstride', 1) + rcount = kwargs.pop('rcount', 50) + ccount = kwargs.pop('ccount', 50) + + if mpl.rcParams['_internal.classic_mode']: + # Strides have priority over counts in classic mode. + # So, only compute strides from counts + # if counts were explicitly given + if has_count: + rstride = int(max(np.ceil(rows / rcount), 1)) if rcount else 0 + cstride = int(max(np.ceil(cols / ccount), 1)) if ccount else 0 + else: + # If the strides are provided then it has priority. + # Otherwise, compute the strides from the counts. + if not has_stride: + rstride = int(max(np.ceil(rows / rcount), 1)) if rcount else 0 + cstride = int(max(np.ceil(cols / ccount), 1)) if ccount else 0 + + # We want two sets of lines, one running along the "rows" of + # Z and another set of lines running along the "columns" of Z. + # This transpose will make it easy to obtain the columns. + tX, tY, tZ = np.transpose(X), np.transpose(Y), np.transpose(Z) + + if rstride: + rii = list(range(0, rows, rstride)) + # Add the last index only if needed + if rows > 0 and rii[-1] != (rows - 1): + rii += [rows-1] + else: + rii = [] + if cstride: + cii = list(range(0, cols, cstride)) + # Add the last index only if needed + if cols > 0 and cii[-1] != (cols - 1): + cii += [cols-1] + else: + cii = [] + + if rstride == 0 and cstride == 0: + raise ValueError("Either rstride or cstride must be non zero") + + # If the inputs were empty, then just + # reset everything. + if Z.size == 0: + rii = [] + cii = [] + + xlines = [X[i] for i in rii] + ylines = [Y[i] for i in rii] + zlines = [Z[i] for i in rii] + + txlines = [tX[i] for i in cii] + tylines = [tY[i] for i in cii] + tzlines = [tZ[i] for i in cii] + + lines = ([list(zip(xl, yl, zl)) + for xl, yl, zl in zip(xlines, ylines, zlines)] + + [list(zip(xl, yl, zl)) + for xl, yl, zl in zip(txlines, tylines, tzlines)]) + + linec = art3d.Line3DCollection(lines, axlim_clip=axlim_clip, **kwargs) + self.add_collection(linec) + self.auto_scale_xyz(X, Y, Z, had_data) + + return linec + + def plot_trisurf(self, *args, color=None, norm=None, vmin=None, vmax=None, + lightsource=None, axlim_clip=False, **kwargs): + """ + Plot a triangulated surface. + + The (optional) triangulation can be specified in one of two ways; + either:: + + plot_trisurf(triangulation, ...) + + where triangulation is a `~matplotlib.tri.Triangulation` object, or:: + + plot_trisurf(X, Y, ...) + plot_trisurf(X, Y, triangles, ...) + plot_trisurf(X, Y, triangles=triangles, ...) + + in which case a Triangulation object will be created. See + `.Triangulation` for an explanation of these possibilities. + + The remaining arguments are:: + + plot_trisurf(..., Z) + + where *Z* is the array of values to contour, one per point + in the triangulation. + + Parameters + ---------- + X, Y, Z : array-like + Data values as 1D arrays. + color + Color of the surface patches. + cmap + A colormap for the surface patches. + norm : `~matplotlib.colors.Normalize`, optional + An instance of Normalize to map values to colors. + vmin, vmax : float, optional + Minimum and maximum value to map. + shade : bool, default: True + Whether to shade the facecolors. Shading is always disabled when + *cmap* is specified. + lightsource : `~matplotlib.colors.LightSource`, optional + The lightsource to use when *shade* is True. + axlim_clip : bool, default: False + Whether to hide patches with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + **kwargs + All other keyword arguments are passed on to + :class:`~mpl_toolkits.mplot3d.art3d.Poly3DCollection` + + Examples + -------- + .. plot:: gallery/mplot3d/trisurf3d.py + .. plot:: gallery/mplot3d/trisurf3d_2.py + """ + + had_data = self.has_data() + + # TODO: Support custom face colours + if color is None: + color = self._get_lines.get_next_color() + color = np.array(mcolors.to_rgba(color)) + + cmap = kwargs.get('cmap', None) + shade = kwargs.pop('shade', cmap is None) + + tri, args, kwargs = \ + Triangulation.get_from_args_and_kwargs(*args, **kwargs) + try: + z = kwargs.pop('Z') + except KeyError: + # We do this so Z doesn't get passed as an arg to PolyCollection + z, *args = args + z = np.asarray(z) + + triangles = tri.get_masked_triangles() + xt = tri.x[triangles] + yt = tri.y[triangles] + zt = z[triangles] + verts = np.stack((xt, yt, zt), axis=-1) + + if cmap: + polyc = art3d.Poly3DCollection(verts, *args, + axlim_clip=axlim_clip, **kwargs) + # average over the three points of each triangle + avg_z = verts[:, :, 2].mean(axis=1) + polyc.set_array(avg_z) + if vmin is not None or vmax is not None: + polyc.set_clim(vmin, vmax) + if norm is not None: + polyc.set_norm(norm) + else: + polyc = art3d.Poly3DCollection( + verts, *args, shade=shade, lightsource=lightsource, + facecolors=color, axlim_clip=axlim_clip, **kwargs) + + self.add_collection(polyc) + self.auto_scale_xyz(tri.x, tri.y, z, had_data) + + return polyc + + def _3d_extend_contour(self, cset, stride=5): + """ + Extend a contour in 3D by creating + """ + + dz = (cset.levels[1] - cset.levels[0]) / 2 + polyverts = [] + colors = [] + for idx, level in enumerate(cset.levels): + path = cset.get_paths()[idx] + subpaths = [*path._iter_connected_components()] + color = cset.get_edgecolor()[idx] + top = art3d._paths_to_3d_segments(subpaths, level - dz) + bot = art3d._paths_to_3d_segments(subpaths, level + dz) + if not len(top[0]): + continue + nsteps = max(round(len(top[0]) / stride), 2) + stepsize = (len(top[0]) - 1) / (nsteps - 1) + polyverts.extend([ + (top[0][round(i * stepsize)], top[0][round((i + 1) * stepsize)], + bot[0][round((i + 1) * stepsize)], bot[0][round(i * stepsize)]) + for i in range(round(nsteps) - 1)]) + colors.extend([color] * (round(nsteps) - 1)) + self.add_collection3d(art3d.Poly3DCollection( + np.array(polyverts), # All polygons have 4 vertices, so vectorize. + facecolors=colors, edgecolors=colors, shade=True)) + cset.remove() + + def add_contour_set( + self, cset, extend3d=False, stride=5, zdir='z', offset=None, + axlim_clip=False): + zdir = '-' + zdir + if extend3d: + self._3d_extend_contour(cset, stride) + else: + art3d.collection_2d_to_3d( + cset, zs=offset if offset is not None else cset.levels, zdir=zdir, + axlim_clip=axlim_clip) + + def add_contourf_set(self, cset, zdir='z', offset=None, *, axlim_clip=False): + self._add_contourf_set(cset, zdir=zdir, offset=offset, + axlim_clip=axlim_clip) + + def _add_contourf_set(self, cset, zdir='z', offset=None, axlim_clip=False): + """ + Returns + ------- + levels : `numpy.ndarray` + Levels at which the filled contours are added. + """ + zdir = '-' + zdir + + midpoints = cset.levels[:-1] + np.diff(cset.levels) / 2 + # Linearly interpolate to get levels for any extensions + if cset._extend_min: + min_level = cset.levels[0] - np.diff(cset.levels[:2]) / 2 + midpoints = np.insert(midpoints, 0, min_level) + if cset._extend_max: + max_level = cset.levels[-1] + np.diff(cset.levels[-2:]) / 2 + midpoints = np.append(midpoints, max_level) + + art3d.collection_2d_to_3d( + cset, zs=offset if offset is not None else midpoints, zdir=zdir, + axlim_clip=axlim_clip) + return midpoints + + @_preprocess_data() + def contour(self, X, Y, Z, *args, + extend3d=False, stride=5, zdir='z', offset=None, axlim_clip=False, + **kwargs): + """ + Create a 3D contour plot. + + Parameters + ---------- + X, Y, Z : array-like, + Input data. See `.Axes.contour` for supported data shapes. + extend3d : bool, default: False + Whether to extend contour in 3D. + stride : int, default: 5 + Step size for extending contour. + zdir : {'x', 'y', 'z'}, default: 'z' + The direction to use. + offset : float, optional + If specified, plot a projection of the contour lines at this + position in a plane normal to *zdir*. + axlim_clip : bool, default: False + Whether to hide lines with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + + *args, **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.contour`. + + Returns + ------- + matplotlib.contour.QuadContourSet + """ + had_data = self.has_data() + + jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir) + cset = super().contour(jX, jY, jZ, *args, **kwargs) + self.add_contour_set(cset, extend3d, stride, zdir, offset, axlim_clip) + + self.auto_scale_xyz(X, Y, Z, had_data) + return cset + + contour3D = contour + + @_preprocess_data() + def tricontour(self, *args, + extend3d=False, stride=5, zdir='z', offset=None, axlim_clip=False, + **kwargs): + """ + Create a 3D contour plot. + + .. note:: + This method currently produces incorrect output due to a + longstanding bug in 3D PolyCollection rendering. + + Parameters + ---------- + X, Y, Z : array-like + Input data. See `.Axes.tricontour` for supported data shapes. + extend3d : bool, default: False + Whether to extend contour in 3D. + stride : int, default: 5 + Step size for extending contour. + zdir : {'x', 'y', 'z'}, default: 'z' + The direction to use. + offset : float, optional + If specified, plot a projection of the contour lines at this + position in a plane normal to *zdir*. + axlim_clip : bool, default: False + Whether to hide lines with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + *args, **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.tricontour`. + + Returns + ------- + matplotlib.tri._tricontour.TriContourSet + """ + had_data = self.has_data() + + tri, args, kwargs = Triangulation.get_from_args_and_kwargs( + *args, **kwargs) + X = tri.x + Y = tri.y + if 'Z' in kwargs: + Z = kwargs.pop('Z') + else: + # We do this so Z doesn't get passed as an arg to Axes.tricontour + Z, *args = args + + jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir) + tri = Triangulation(jX, jY, tri.triangles, tri.mask) + + cset = super().tricontour(tri, jZ, *args, **kwargs) + self.add_contour_set(cset, extend3d, stride, zdir, offset, axlim_clip) + + self.auto_scale_xyz(X, Y, Z, had_data) + return cset + + def _auto_scale_contourf(self, X, Y, Z, zdir, levels, had_data): + # Autoscale in the zdir based on the levels added, which are + # different from data range if any contour extensions are present + dim_vals = {'x': X, 'y': Y, 'z': Z, zdir: levels} + # Input data and levels have different sizes, but auto_scale_xyz + # expected same-size input, so manually take min/max limits + limits = [(np.nanmin(dim_vals[dim]), np.nanmax(dim_vals[dim])) + for dim in ['x', 'y', 'z']] + self.auto_scale_xyz(*limits, had_data) + + @_preprocess_data() + def contourf(self, X, Y, Z, *args, + zdir='z', offset=None, axlim_clip=False, **kwargs): + """ + Create a 3D filled contour plot. + + Parameters + ---------- + X, Y, Z : array-like + Input data. See `.Axes.contourf` for supported data shapes. + zdir : {'x', 'y', 'z'}, default: 'z' + The direction to use. + offset : float, optional + If specified, plot a projection of the contour lines at this + position in a plane normal to *zdir*. + axlim_clip : bool, default: False + Whether to hide lines with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + *args, **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.contourf`. + + Returns + ------- + matplotlib.contour.QuadContourSet + """ + had_data = self.has_data() + + jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir) + cset = super().contourf(jX, jY, jZ, *args, **kwargs) + levels = self._add_contourf_set(cset, zdir, offset, axlim_clip) + + self._auto_scale_contourf(X, Y, Z, zdir, levels, had_data) + return cset + + contourf3D = contourf + + @_preprocess_data() + def tricontourf(self, *args, zdir='z', offset=None, axlim_clip=False, **kwargs): + """ + Create a 3D filled contour plot. + + .. note:: + This method currently produces incorrect output due to a + longstanding bug in 3D PolyCollection rendering. + + Parameters + ---------- + X, Y, Z : array-like + Input data. See `.Axes.tricontourf` for supported data shapes. + zdir : {'x', 'y', 'z'}, default: 'z' + The direction to use. + offset : float, optional + If specified, plot a projection of the contour lines at this + position in a plane normal to zdir. + axlim_clip : bool, default: False + Whether to hide lines with a vertex outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + *args, **kwargs + Other arguments are forwarded to + `matplotlib.axes.Axes.tricontourf`. + + Returns + ------- + matplotlib.tri._tricontour.TriContourSet + """ + had_data = self.has_data() + + tri, args, kwargs = Triangulation.get_from_args_and_kwargs( + *args, **kwargs) + X = tri.x + Y = tri.y + if 'Z' in kwargs: + Z = kwargs.pop('Z') + else: + # We do this so Z doesn't get passed as an arg to Axes.tricontourf + Z, *args = args + + jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir) + tri = Triangulation(jX, jY, tri.triangles, tri.mask) + + cset = super().tricontourf(tri, jZ, *args, **kwargs) + levels = self._add_contourf_set(cset, zdir, offset, axlim_clip) + + self._auto_scale_contourf(X, Y, Z, zdir, levels, had_data) + return cset + + def add_collection3d(self, col, zs=0, zdir='z', autolim=True, *, + axlim_clip=False): + """ + Add a 3D collection object to the plot. + + 2D collection types are converted to a 3D version by + modifying the object and adding z coordinate information, + *zs* and *zdir*. + + Supported 2D collection types are: + + - `.PolyCollection` + - `.LineCollection` + - `.PatchCollection` (currently not supporting *autolim*) + + Parameters + ---------- + col : `.Collection` + A 2D collection object. + zs : float or array-like, default: 0 + The z-positions to be used for the 2D objects. + zdir : {'x', 'y', 'z'}, default: 'z' + The direction to use for the z-positions. + autolim : bool, default: True + Whether to update the data limits. + axlim_clip : bool, default: False + Whether to hide the scatter points outside the axes view limits. + + .. versionadded:: 3.10 + """ + had_data = self.has_data() + + zvals = np.atleast_1d(zs) + zsortval = (np.min(zvals) if zvals.size + else 0) # FIXME: arbitrary default + + # FIXME: use issubclass() (although, then a 3D collection + # object would also pass.) Maybe have a collection3d + # abstract class to test for and exclude? + if type(col) is mcoll.PolyCollection: + art3d.poly_collection_2d_to_3d(col, zs=zs, zdir=zdir, + axlim_clip=axlim_clip) + col.set_sort_zpos(zsortval) + elif type(col) is mcoll.LineCollection: + art3d.line_collection_2d_to_3d(col, zs=zs, zdir=zdir, + axlim_clip=axlim_clip) + col.set_sort_zpos(zsortval) + elif type(col) is mcoll.PatchCollection: + art3d.patch_collection_2d_to_3d(col, zs=zs, zdir=zdir, + axlim_clip=axlim_clip) + col.set_sort_zpos(zsortval) + + if autolim: + if isinstance(col, art3d.Line3DCollection): + self.auto_scale_xyz(*np.array(col._segments3d).transpose(), + had_data=had_data) + elif isinstance(col, art3d.Poly3DCollection): + self.auto_scale_xyz(*col._vec[:-1], had_data=had_data) + elif isinstance(col, art3d.Patch3DCollection): + pass + # FIXME: Implement auto-scaling function for Patch3DCollection + # Currently unable to do so due to issues with Patch3DCollection + # See https://github.com/matplotlib/matplotlib/issues/14298 for details + + collection = super().add_collection(col) + return collection + + @_preprocess_data(replace_names=["xs", "ys", "zs", "s", + "edgecolors", "c", "facecolor", + "facecolors", "color"]) + def scatter(self, xs, ys, + zs=0, zdir='z', s=20, c=None, depthshade=True, *args, + axlim_clip=False, **kwargs): + """ + Create a scatter plot. + + Parameters + ---------- + xs, ys : array-like + The data positions. + zs : float or array-like, default: 0 + The z-positions. Either an array of the same length as *xs* and + *ys* or a single value to place all points in the same plane. + zdir : {'x', 'y', 'z', '-x', '-y', '-z'}, default: 'z' + The axis direction for the *zs*. This is useful when plotting 2D + data on a 3D Axes. The data must be passed as *xs*, *ys*. Setting + *zdir* to 'y' then plots the data to the x-z-plane. + + See also :doc:`/gallery/mplot3d/2dcollections3d`. + + s : float or array-like, default: 20 + The marker size in points**2. Either an array of the same length + as *xs* and *ys* or a single value to make all markers the same + size. + c : :mpltype:`color`, sequence, or sequence of colors, optional + The marker color. Possible values: + + - A single color format string. + - A sequence of colors of length n. + - A sequence of n numbers to be mapped to colors using *cmap* and + *norm*. + - A 2D array in which the rows are RGB or RGBA. + + For more details see the *c* argument of `~.axes.Axes.scatter`. + depthshade : bool, default: True + Whether to shade the scatter markers to give the appearance of + depth. Each call to ``scatter()`` will perform its depthshading + independently. + axlim_clip : bool, default: False + Whether to hide the scatter points outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + **kwargs + All other keyword arguments are passed on to `~.axes.Axes.scatter`. + + Returns + ------- + paths : `~matplotlib.collections.PathCollection` + """ + + had_data = self.has_data() + zs_orig = zs + + xs, ys, zs = cbook._broadcast_with_masks(xs, ys, zs) + s = np.ma.ravel(s) # This doesn't have to match x, y in size. + + xs, ys, zs, s, c, color = cbook.delete_masked_points( + xs, ys, zs, s, c, kwargs.get('color', None) + ) + if kwargs.get("color") is not None: + kwargs['color'] = color + + # For xs and ys, 2D scatter() will do the copying. + if np.may_share_memory(zs_orig, zs): # Avoid unnecessary copies. + zs = zs.copy() + + patches = super().scatter(xs, ys, s=s, c=c, *args, **kwargs) + art3d.patch_collection_2d_to_3d(patches, zs=zs, zdir=zdir, + depthshade=depthshade, + axlim_clip=axlim_clip) + + if self._zmargin < 0.05 and xs.size > 0: + self.set_zmargin(0.05) + + self.auto_scale_xyz(xs, ys, zs, had_data) + + return patches + + scatter3D = scatter + + @_preprocess_data() + def bar(self, left, height, zs=0, zdir='z', *args, + axlim_clip=False, **kwargs): + """ + Add 2D bar(s). + + Parameters + ---------- + left : 1D array-like + The x coordinates of the left sides of the bars. + height : 1D array-like + The height of the bars. + zs : float or 1D array-like, default: 0 + Z coordinate of bars; if a single value is specified, it will be + used for all bars. + zdir : {'x', 'y', 'z'}, default: 'z' + When plotting 2D data, the direction to use as z ('x', 'y' or 'z'). + axlim_clip : bool, default: False + Whether to hide bars with points outside the axes view limits. + + .. versionadded:: 3.10 + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + **kwargs + Other keyword arguments are forwarded to + `matplotlib.axes.Axes.bar`. + + Returns + ------- + mpl_toolkits.mplot3d.art3d.Patch3DCollection + """ + had_data = self.has_data() + + patches = super().bar(left, height, *args, **kwargs) + + zs = np.broadcast_to(zs, len(left), subok=True) + + verts = [] + verts_zs = [] + for p, z in zip(patches, zs): + vs = art3d._get_patch_verts(p) + verts += vs.tolist() + verts_zs += [z] * len(vs) + art3d.patch_2d_to_3d(p, z, zdir, axlim_clip) + if 'alpha' in kwargs: + p.set_alpha(kwargs['alpha']) + + if len(verts) > 0: + # the following has to be skipped if verts is empty + # NOTE: Bugs could still occur if len(verts) > 0, + # but the "2nd dimension" is empty. + xs, ys = zip(*verts) + else: + xs, ys = [], [] + + xs, ys, verts_zs = art3d.juggle_axes(xs, ys, verts_zs, zdir) + self.auto_scale_xyz(xs, ys, verts_zs, had_data) + + return patches + + @_preprocess_data() + def bar3d(self, x, y, z, dx, dy, dz, color=None, + zsort='average', shade=True, lightsource=None, *args, + axlim_clip=False, **kwargs): + """ + Generate a 3D barplot. + + This method creates three-dimensional barplot where the width, + depth, height, and color of the bars can all be uniquely set. + + Parameters + ---------- + x, y, z : array-like + The coordinates of the anchor point of the bars. + + dx, dy, dz : float or array-like + The width, depth, and height of the bars, respectively. + + color : sequence of colors, optional + The color of the bars can be specified globally or + individually. This parameter can be: + + - A single color, to color all bars the same color. + - An array of colors of length N bars, to color each bar + independently. + - An array of colors of length 6, to color the faces of the + bars similarly. + - An array of colors of length 6 * N bars, to color each face + independently. + + When coloring the faces of the boxes specifically, this is + the order of the coloring: + + 1. -Z (bottom of box) + 2. +Z (top of box) + 3. -Y + 4. +Y + 5. -X + 6. +X + + zsort : {'average', 'min', 'max'}, default: 'average' + The z-axis sorting scheme passed onto `~.art3d.Poly3DCollection` + + shade : bool, default: True + When true, this shades the dark sides of the bars (relative + to the plot's source of light). + + lightsource : `~matplotlib.colors.LightSource`, optional + The lightsource to use when *shade* is True. + + axlim_clip : bool, default: False + Whether to hide the bars with points outside the axes view limits. + + .. versionadded:: 3.10 + + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + + **kwargs + Any additional keyword arguments are passed onto + `~.art3d.Poly3DCollection`. + + Returns + ------- + collection : `~.art3d.Poly3DCollection` + A collection of three-dimensional polygons representing the bars. + """ + + had_data = self.has_data() + + x, y, z, dx, dy, dz = np.broadcast_arrays( + np.atleast_1d(x), y, z, dx, dy, dz) + minx = np.min(x) + maxx = np.max(x + dx) + miny = np.min(y) + maxy = np.max(y + dy) + minz = np.min(z) + maxz = np.max(z + dz) + + # shape (6, 4, 3) + # All faces are oriented facing outwards - when viewed from the + # outside, their vertices are in a counterclockwise ordering. + cuboid = np.array([ + # -z + ( + (0, 0, 0), + (0, 1, 0), + (1, 1, 0), + (1, 0, 0), + ), + # +z + ( + (0, 0, 1), + (1, 0, 1), + (1, 1, 1), + (0, 1, 1), + ), + # -y + ( + (0, 0, 0), + (1, 0, 0), + (1, 0, 1), + (0, 0, 1), + ), + # +y + ( + (0, 1, 0), + (0, 1, 1), + (1, 1, 1), + (1, 1, 0), + ), + # -x + ( + (0, 0, 0), + (0, 0, 1), + (0, 1, 1), + (0, 1, 0), + ), + # +x + ( + (1, 0, 0), + (1, 1, 0), + (1, 1, 1), + (1, 0, 1), + ), + ]) + + # indexed by [bar, face, vertex, coord] + polys = np.empty(x.shape + cuboid.shape) + + # handle each coordinate separately + for i, p, dp in [(0, x, dx), (1, y, dy), (2, z, dz)]: + p = p[..., np.newaxis, np.newaxis] + dp = dp[..., np.newaxis, np.newaxis] + polys[..., i] = p + dp * cuboid[..., i] + + # collapse the first two axes + polys = polys.reshape((-1,) + polys.shape[2:]) + + facecolors = [] + if color is None: + color = [self._get_patches_for_fill.get_next_color()] + + color = list(mcolors.to_rgba_array(color)) + + if len(color) == len(x): + # bar colors specified, need to expand to number of faces + for c in color: + facecolors.extend([c] * 6) + else: + # a single color specified, or face colors specified explicitly + facecolors = color + if len(facecolors) < len(x): + facecolors *= (6 * len(x)) + + col = art3d.Poly3DCollection(polys, + zsort=zsort, + facecolors=facecolors, + shade=shade, + lightsource=lightsource, + axlim_clip=axlim_clip, + *args, **kwargs) + self.add_collection(col) + + self.auto_scale_xyz((minx, maxx), (miny, maxy), (minz, maxz), had_data) + + return col + + def set_title(self, label, fontdict=None, loc='center', **kwargs): + # docstring inherited + ret = super().set_title(label, fontdict=fontdict, loc=loc, **kwargs) + (x, y) = self.title.get_position() + self.title.set_y(0.92 * y) + return ret + + @_preprocess_data() + def quiver(self, X, Y, Z, U, V, W, *, + length=1, arrow_length_ratio=.3, pivot='tail', normalize=False, + axlim_clip=False, **kwargs): + """ + Plot a 3D field of arrows. + + The arguments can be array-like or scalars, so long as they can be + broadcast together. The arguments can also be masked arrays. If an + element in any of argument is masked, then that corresponding quiver + element will not be plotted. + + Parameters + ---------- + X, Y, Z : array-like + The x, y and z coordinates of the arrow locations (default is + tail of arrow; see *pivot* kwarg). + + U, V, W : array-like + The x, y and z components of the arrow vectors. + + length : float, default: 1 + The length of each quiver. + + arrow_length_ratio : float, default: 0.3 + The ratio of the arrow head with respect to the quiver. + + pivot : {'tail', 'middle', 'tip'}, default: 'tail' + The part of the arrow that is at the grid point; the arrow + rotates about this point, hence the name *pivot*. + + normalize : bool, default: False + Whether all arrows are normalized to have the same length, or keep + the lengths defined by *u*, *v*, and *w*. + + axlim_clip : bool, default: False + Whether to hide arrows with points outside the axes view limits. + + .. versionadded:: 3.10 + + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + + **kwargs + Any additional keyword arguments are delegated to + :class:`.Line3DCollection` + """ + + def calc_arrows(UVW): + # get unit direction vector perpendicular to (u, v, w) + x = UVW[:, 0] + y = UVW[:, 1] + norm = np.linalg.norm(UVW[:, :2], axis=1) + x_p = np.divide(y, norm, where=norm != 0, out=np.zeros_like(x)) + y_p = np.divide(-x, norm, where=norm != 0, out=np.ones_like(x)) + # compute the two arrowhead direction unit vectors + rangle = math.radians(15) + c = math.cos(rangle) + s = math.sin(rangle) + # construct the rotation matrices of shape (3, 3, n) + r13 = y_p * s + r32 = x_p * s + r12 = x_p * y_p * (1 - c) + Rpos = np.array( + [[c + (x_p ** 2) * (1 - c), r12, r13], + [r12, c + (y_p ** 2) * (1 - c), -r32], + [-r13, r32, np.full_like(x_p, c)]]) + # opposite rotation negates all the sin terms + Rneg = Rpos.copy() + Rneg[[0, 1, 2, 2], [2, 2, 0, 1]] *= -1 + # Batch n (3, 3) x (3) matrix multiplications ((3, 3, n) x (n, 3)). + Rpos_vecs = np.einsum("ij...,...j->...i", Rpos, UVW) + Rneg_vecs = np.einsum("ij...,...j->...i", Rneg, UVW) + # Stack into (n, 2, 3) result. + return np.stack([Rpos_vecs, Rneg_vecs], axis=1) + + had_data = self.has_data() + + input_args = cbook._broadcast_with_masks(X, Y, Z, U, V, W, + compress=True) + + if any(len(v) == 0 for v in input_args): + # No quivers, so just make an empty collection and return early + linec = art3d.Line3DCollection([], **kwargs) + self.add_collection(linec) + return linec + + shaft_dt = np.array([0., length], dtype=float) + arrow_dt = shaft_dt * arrow_length_ratio + + _api.check_in_list(['tail', 'middle', 'tip'], pivot=pivot) + if pivot == 'tail': + shaft_dt -= length + elif pivot == 'middle': + shaft_dt -= length / 2 + + XYZ = np.column_stack(input_args[:3]) + UVW = np.column_stack(input_args[3:]).astype(float) + + # Normalize rows of UVW + if normalize: + norm = np.linalg.norm(UVW, axis=1) + norm[norm == 0] = 1 + UVW = UVW / norm.reshape((-1, 1)) + + if len(XYZ) > 0: + # compute the shaft lines all at once with an outer product + shafts = (XYZ - np.multiply.outer(shaft_dt, UVW)).swapaxes(0, 1) + # compute head direction vectors, n heads x 2 sides x 3 dimensions + head_dirs = calc_arrows(UVW) + # compute all head lines at once, starting from the shaft ends + heads = shafts[:, :1] - np.multiply.outer(arrow_dt, head_dirs) + # stack left and right head lines together + heads = heads.reshape((len(arrow_dt), -1, 3)) + # transpose to get a list of lines + heads = heads.swapaxes(0, 1) + + lines = [*shafts, *heads[::2], *heads[1::2]] + else: + lines = [] + + linec = art3d.Line3DCollection(lines, axlim_clip=axlim_clip, **kwargs) + self.add_collection(linec) + + self.auto_scale_xyz(XYZ[:, 0], XYZ[:, 1], XYZ[:, 2], had_data) + + return linec + + quiver3D = quiver + + def voxels(self, *args, facecolors=None, edgecolors=None, shade=True, + lightsource=None, axlim_clip=False, **kwargs): + """ + ax.voxels([x, y, z,] /, filled, facecolors=None, edgecolors=None, \ +**kwargs) + + Plot a set of filled voxels + + All voxels are plotted as 1x1x1 cubes on the axis, with + ``filled[0, 0, 0]`` placed with its lower corner at the origin. + Occluded faces are not plotted. + + Parameters + ---------- + filled : 3D np.array of bool + A 3D array of values, with truthy values indicating which voxels + to fill + + x, y, z : 3D np.array, optional + The coordinates of the corners of the voxels. This should broadcast + to a shape one larger in every dimension than the shape of + *filled*. These can be used to plot non-cubic voxels. + + If not specified, defaults to increasing integers along each axis, + like those returned by :func:`~numpy.indices`. + As indicated by the ``/`` in the function signature, these + arguments can only be passed positionally. + + facecolors, edgecolors : array-like, optional + The color to draw the faces and edges of the voxels. Can only be + passed as keyword arguments. + These parameters can be: + + - A single color value, to color all voxels the same color. This + can be either a string, or a 1D RGB/RGBA array + - ``None``, the default, to use a single color for the faces, and + the style default for the edges. + - A 3D `~numpy.ndarray` of color names, with each item the color + for the corresponding voxel. The size must match the voxels. + - A 4D `~numpy.ndarray` of RGB/RGBA data, with the components + along the last axis. + + shade : bool, default: True + Whether to shade the facecolors. + + lightsource : `~matplotlib.colors.LightSource`, optional + The lightsource to use when *shade* is True. + + axlim_clip : bool, default: False + Whether to hide voxels with points outside the axes view limits. + + .. versionadded:: 3.10 + + **kwargs + Additional keyword arguments to pass onto + `~mpl_toolkits.mplot3d.art3d.Poly3DCollection`. + + Returns + ------- + faces : dict + A dictionary indexed by coordinate, where ``faces[i, j, k]`` is a + `.Poly3DCollection` of the faces drawn for the voxel + ``filled[i, j, k]``. If no faces were drawn for a given voxel, + either because it was not asked to be drawn, or it is fully + occluded, then ``(i, j, k) not in faces``. + + Examples + -------- + .. plot:: gallery/mplot3d/voxels.py + .. plot:: gallery/mplot3d/voxels_rgb.py + .. plot:: gallery/mplot3d/voxels_torus.py + .. plot:: gallery/mplot3d/voxels_numpy_logo.py + """ + + # work out which signature we should be using, and use it to parse + # the arguments. Name must be voxels for the correct error message + if len(args) >= 3: + # underscores indicate position only + def voxels(__x, __y, __z, filled, **kwargs): + return (__x, __y, __z), filled, kwargs + else: + def voxels(filled, **kwargs): + return None, filled, kwargs + + xyz, filled, kwargs = voxels(*args, **kwargs) + + # check dimensions + if filled.ndim != 3: + raise ValueError("Argument filled must be 3-dimensional") + size = np.array(filled.shape, dtype=np.intp) + + # check xyz coordinates, which are one larger than the filled shape + coord_shape = tuple(size + 1) + if xyz is None: + x, y, z = np.indices(coord_shape) + else: + x, y, z = (np.broadcast_to(c, coord_shape) for c in xyz) + + def _broadcast_color_arg(color, name): + if np.ndim(color) in (0, 1): + # single color, like "red" or [1, 0, 0] + return np.broadcast_to(color, filled.shape + np.shape(color)) + elif np.ndim(color) in (3, 4): + # 3D array of strings, or 4D array with last axis rgb + if np.shape(color)[:3] != filled.shape: + raise ValueError( + f"When multidimensional, {name} must match the shape " + "of filled") + return color + else: + raise ValueError(f"Invalid {name} argument") + + # broadcast and default on facecolors + if facecolors is None: + facecolors = self._get_patches_for_fill.get_next_color() + facecolors = _broadcast_color_arg(facecolors, 'facecolors') + + # broadcast but no default on edgecolors + edgecolors = _broadcast_color_arg(edgecolors, 'edgecolors') + + # scale to the full array, even if the data is only in the center + self.auto_scale_xyz(x, y, z) + + # points lying on corners of a square + square = np.array([ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + ], dtype=np.intp) + + voxel_faces = defaultdict(list) + + def permutation_matrices(n): + """Generate cyclic permutation matrices.""" + mat = np.eye(n, dtype=np.intp) + for i in range(n): + yield mat + mat = np.roll(mat, 1, axis=0) + + # iterate over each of the YZ, ZX, and XY orientations, finding faces + # to render + for permute in permutation_matrices(3): + # find the set of ranges to iterate over + pc, qc, rc = permute.T.dot(size) + pinds = np.arange(pc) + qinds = np.arange(qc) + rinds = np.arange(rc) + + square_rot_pos = square.dot(permute.T) + square_rot_neg = square_rot_pos[::-1] + + # iterate within the current plane + for p in pinds: + for q in qinds: + # iterate perpendicularly to the current plane, handling + # boundaries. We only draw faces between a voxel and an + # empty space, to avoid drawing internal faces. + + # draw lower faces + p0 = permute.dot([p, q, 0]) + i0 = tuple(p0) + if filled[i0]: + voxel_faces[i0].append(p0 + square_rot_neg) + + # draw middle faces + for r1, r2 in itertools.pairwise(rinds): + p1 = permute.dot([p, q, r1]) + p2 = permute.dot([p, q, r2]) + + i1 = tuple(p1) + i2 = tuple(p2) + + if filled[i1] and not filled[i2]: + voxel_faces[i1].append(p2 + square_rot_pos) + elif not filled[i1] and filled[i2]: + voxel_faces[i2].append(p2 + square_rot_neg) + + # draw upper faces + pk = permute.dot([p, q, rc-1]) + pk2 = permute.dot([p, q, rc]) + ik = tuple(pk) + if filled[ik]: + voxel_faces[ik].append(pk2 + square_rot_pos) + + # iterate over the faces, and generate a Poly3DCollection for each + # voxel + polygons = {} + for coord, faces_inds in voxel_faces.items(): + # convert indices into 3D positions + if xyz is None: + faces = faces_inds + else: + faces = [] + for face_inds in faces_inds: + ind = face_inds[:, 0], face_inds[:, 1], face_inds[:, 2] + face = np.empty(face_inds.shape) + face[:, 0] = x[ind] + face[:, 1] = y[ind] + face[:, 2] = z[ind] + faces.append(face) + + # shade the faces + facecolor = facecolors[coord] + edgecolor = edgecolors[coord] + + poly = art3d.Poly3DCollection( + faces, facecolors=facecolor, edgecolors=edgecolor, + shade=shade, lightsource=lightsource, axlim_clip=axlim_clip, + **kwargs) + self.add_collection3d(poly) + polygons[coord] = poly + + return polygons + + @_preprocess_data(replace_names=["x", "y", "z", "xerr", "yerr", "zerr"]) + def errorbar(self, x, y, z, zerr=None, yerr=None, xerr=None, fmt='', + barsabove=False, errorevery=1, ecolor=None, elinewidth=None, + capsize=None, capthick=None, xlolims=False, xuplims=False, + ylolims=False, yuplims=False, zlolims=False, zuplims=False, + axlim_clip=False, + **kwargs): + """ + Plot lines and/or markers with errorbars around them. + + *x*/*y*/*z* define the data locations, and *xerr*/*yerr*/*zerr* define + the errorbar sizes. By default, this draws the data markers/lines as + well the errorbars. Use fmt='none' to draw errorbars only. + + Parameters + ---------- + x, y, z : float or array-like + The data positions. + + xerr, yerr, zerr : float or array-like, shape (N,) or (2, N), optional + The errorbar sizes: + + - scalar: Symmetric +/- values for all data points. + - shape(N,): Symmetric +/-values for each data point. + - shape(2, N): Separate - and + values for each bar. First row + contains the lower errors, the second row contains the upper + errors. + - *None*: No errorbar. + + Note that all error arrays should have *positive* values. + + fmt : str, default: '' + The format for the data points / data lines. See `.plot` for + details. + + Use 'none' (case-insensitive) to plot errorbars without any data + markers. + + ecolor : :mpltype:`color`, default: None + The color of the errorbar lines. If None, use the color of the + line connecting the markers. + + elinewidth : float, default: None + The linewidth of the errorbar lines. If None, the linewidth of + the current style is used. + + capsize : float, default: :rc:`errorbar.capsize` + The length of the error bar caps in points. + + capthick : float, default: None + An alias to the keyword argument *markeredgewidth* (a.k.a. *mew*). + This setting is a more sensible name for the property that + controls the thickness of the error bar cap in points. For + backwards compatibility, if *mew* or *markeredgewidth* are given, + then they will over-ride *capthick*. This may change in future + releases. + + barsabove : bool, default: False + If True, will plot the errorbars above the plot + symbols. Default is below. + + xlolims, ylolims, zlolims : bool, default: False + These arguments can be used to indicate that a value gives only + lower limits. In that case a caret symbol is used to indicate + this. *lims*-arguments may be scalars, or array-likes of the same + length as the errors. To use limits with inverted axes, + `~.set_xlim`, `~.set_ylim`, or `~.set_zlim` must be + called before `errorbar`. Note the tricky parameter names: setting + e.g. *ylolims* to True means that the y-value is a *lower* limit of + the True value, so, only an *upward*-pointing arrow will be drawn! + + xuplims, yuplims, zuplims : bool, default: False + Same as above, but for controlling the upper limits. + + errorevery : int or (int, int), default: 1 + draws error bars on a subset of the data. *errorevery* =N draws + error bars on the points (x[::N], y[::N], z[::N]). + *errorevery* =(start, N) draws error bars on the points + (x[start::N], y[start::N], z[start::N]). e.g. *errorevery* =(6, 3) + adds error bars to the data at (x[6], x[9], x[12], x[15], ...). + Used to avoid overlapping error bars when two series share x-axis + values. + + axlim_clip : bool, default: False + Whether to hide error bars that are outside the axes limits. + + .. versionadded:: 3.10 + + Returns + ------- + errlines : list + List of `~mpl_toolkits.mplot3d.art3d.Line3DCollection` instances + each containing an errorbar line. + caplines : list + List of `~mpl_toolkits.mplot3d.art3d.Line3D` instances each + containing a capline object. + limmarks : list + List of `~mpl_toolkits.mplot3d.art3d.Line3D` instances each + containing a marker with an upper or lower limit. + + Other Parameters + ---------------- + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + + **kwargs + All other keyword arguments for styling errorbar lines are passed + `~mpl_toolkits.mplot3d.art3d.Line3DCollection`. + + Examples + -------- + .. plot:: gallery/mplot3d/errorbar3d.py + """ + had_data = self.has_data() + + kwargs = cbook.normalize_kwargs(kwargs, mlines.Line2D) + # Drop anything that comes in as None to use the default instead. + kwargs = {k: v for k, v in kwargs.items() if v is not None} + kwargs.setdefault('zorder', 2) + + self._process_unit_info([("x", x), ("y", y), ("z", z)], kwargs, + convert=False) + + # make sure all the args are iterable; use lists not arrays to + # preserve units + x = x if np.iterable(x) else [x] + y = y if np.iterable(y) else [y] + z = z if np.iterable(z) else [z] + + if not len(x) == len(y) == len(z): + raise ValueError("'x', 'y', and 'z' must have the same size") + + everymask = self._errorevery_to_mask(x, errorevery) + + label = kwargs.pop("label", None) + kwargs['label'] = '_nolegend_' + + # Create the main line and determine overall kwargs for child artists. + # We avoid calling self.plot() directly, or self._get_lines(), because + # that would call self._process_unit_info again, and do other indirect + # data processing. + (data_line, base_style), = self._get_lines._plot_args( + self, (x, y) if fmt == '' else (x, y, fmt), kwargs, return_kwargs=True) + art3d.line_2d_to_3d(data_line, zs=z, axlim_clip=axlim_clip) + + # Do this after creating `data_line` to avoid modifying `base_style`. + if barsabove: + data_line.set_zorder(kwargs['zorder'] - .1) + else: + data_line.set_zorder(kwargs['zorder'] + .1) + + # Add line to plot, or throw it away and use it to determine kwargs. + if fmt.lower() != 'none': + self.add_line(data_line) + else: + data_line = None + # Remove alpha=0 color that _process_plot_format returns. + base_style.pop('color') + + if 'color' not in base_style: + base_style['color'] = 'C0' + if ecolor is None: + ecolor = base_style['color'] + + # Eject any line-specific information from format string, as it's not + # needed for bars or caps. + for key in ['marker', 'markersize', 'markerfacecolor', + 'markeredgewidth', 'markeredgecolor', 'markevery', + 'linestyle', 'fillstyle', 'drawstyle', 'dash_capstyle', + 'dash_joinstyle', 'solid_capstyle', 'solid_joinstyle']: + base_style.pop(key, None) + + # Make the style dict for the line collections (the bars). + eb_lines_style = {**base_style, 'color': ecolor} + + if elinewidth: + eb_lines_style['linewidth'] = elinewidth + elif 'linewidth' in kwargs: + eb_lines_style['linewidth'] = kwargs['linewidth'] + + for key in ('transform', 'alpha', 'zorder', 'rasterized'): + if key in kwargs: + eb_lines_style[key] = kwargs[key] + + # Make the style dict for caps (the "hats"). + eb_cap_style = {**base_style, 'linestyle': 'None'} + if capsize is None: + capsize = mpl.rcParams["errorbar.capsize"] + if capsize > 0: + eb_cap_style['markersize'] = 2. * capsize + if capthick is not None: + eb_cap_style['markeredgewidth'] = capthick + eb_cap_style['color'] = ecolor + + def _apply_mask(arrays, mask): + # Return, for each array in *arrays*, the elements for which *mask* + # is True, without using fancy indexing. + return [[*itertools.compress(array, mask)] for array in arrays] + + def _extract_errs(err, data, lomask, himask): + # For separate +/- error values we need to unpack err + if len(err.shape) == 2: + low_err, high_err = err + else: + low_err, high_err = err, err + + lows = np.where(lomask | ~everymask, data, data - low_err) + highs = np.where(himask | ~everymask, data, data + high_err) + + return lows, highs + + # collect drawn items while looping over the three coordinates + errlines, caplines, limmarks = [], [], [] + + # list of endpoint coordinates, used for auto-scaling + coorderrs = [] + + # define the markers used for errorbar caps and limits below + # the dictionary key is mapped by the `i_xyz` helper dictionary + capmarker = {0: '|', 1: '|', 2: '_'} + i_xyz = {'x': 0, 'y': 1, 'z': 2} + + # Calculate marker size from points to quiver length. Because these are + # not markers, and 3D Axes do not use the normal transform stack, this + # is a bit involved. Since the quiver arrows will change size as the + # scene is rotated, they are given a standard size based on viewing + # them directly in planar form. + quiversize = eb_cap_style.get('markersize', + mpl.rcParams['lines.markersize']) ** 2 + quiversize *= self.get_figure(root=True).dpi / 72 + quiversize = self.transAxes.inverted().transform([ + (0, 0), (quiversize, quiversize)]) + quiversize = np.mean(np.diff(quiversize, axis=0)) + # quiversize is now in Axes coordinates, and to convert back to data + # coordinates, we need to run it through the inverse 3D transform. For + # consistency, this uses a fixed elevation, azimuth, and roll. + with cbook._setattr_cm(self, elev=0, azim=0, roll=0): + invM = np.linalg.inv(self.get_proj()) + # elev=azim=roll=0 produces the Y-Z plane, so quiversize in 2D 'x' is + # 'y' in 3D, hence the 1 index. + quiversize = np.dot(invM, [quiversize, 0, 0, 0])[1] + # Quivers use a fixed 15-degree arrow head, so scale up the length so + # that the size corresponds to the base. In other words, this constant + # corresponds to the equation tan(15) = (base / 2) / (arrow length). + quiversize *= 1.8660254037844388 + eb_quiver_style = {**eb_cap_style, + 'length': quiversize, 'arrow_length_ratio': 1} + eb_quiver_style.pop('markersize', None) + + # loop over x-, y-, and z-direction and draw relevant elements + for zdir, data, err, lolims, uplims in zip( + ['x', 'y', 'z'], [x, y, z], [xerr, yerr, zerr], + [xlolims, ylolims, zlolims], [xuplims, yuplims, zuplims]): + + dir_vector = art3d.get_dir_vector(zdir) + i_zdir = i_xyz[zdir] + + if err is None: + continue + + if not np.iterable(err): + err = [err] * len(data) + + err = np.atleast_1d(err) + + # arrays fine here, they are booleans and hence not units + lolims = np.broadcast_to(lolims, len(data)).astype(bool) + uplims = np.broadcast_to(uplims, len(data)).astype(bool) + + # a nested list structure that expands to (xl,xh),(yl,yh),(zl,zh), + # where x/y/z and l/h correspond to dimensions and low/high + # positions of errorbars in a dimension we're looping over + coorderr = [ + _extract_errs(err * dir_vector[i], coord, lolims, uplims) + for i, coord in enumerate([x, y, z])] + (xl, xh), (yl, yh), (zl, zh) = coorderr + + # draws capmarkers - flat caps orthogonal to the error bars + nolims = ~(lolims | uplims) + if nolims.any() and capsize > 0: + lo_caps_xyz = _apply_mask([xl, yl, zl], nolims & everymask) + hi_caps_xyz = _apply_mask([xh, yh, zh], nolims & everymask) + + # setting '_' for z-caps and '|' for x- and y-caps; + # these markers will rotate as the viewing angle changes + cap_lo = art3d.Line3D(*lo_caps_xyz, ls='', + marker=capmarker[i_zdir], + axlim_clip=axlim_clip, + **eb_cap_style) + cap_hi = art3d.Line3D(*hi_caps_xyz, ls='', + marker=capmarker[i_zdir], + axlim_clip=axlim_clip, + **eb_cap_style) + self.add_line(cap_lo) + self.add_line(cap_hi) + caplines.append(cap_lo) + caplines.append(cap_hi) + + if lolims.any(): + xh0, yh0, zh0 = _apply_mask([xh, yh, zh], lolims & everymask) + self.quiver(xh0, yh0, zh0, *dir_vector, **eb_quiver_style) + if uplims.any(): + xl0, yl0, zl0 = _apply_mask([xl, yl, zl], uplims & everymask) + self.quiver(xl0, yl0, zl0, *-dir_vector, **eb_quiver_style) + + errline = art3d.Line3DCollection(np.array(coorderr).T, + axlim_clip=axlim_clip, + **eb_lines_style) + self.add_collection(errline) + errlines.append(errline) + coorderrs.append(coorderr) + + coorderrs = np.array(coorderrs) + + def _digout_minmax(err_arr, coord_label): + return (np.nanmin(err_arr[:, i_xyz[coord_label], :, :]), + np.nanmax(err_arr[:, i_xyz[coord_label], :, :])) + + minx, maxx = _digout_minmax(coorderrs, 'x') + miny, maxy = _digout_minmax(coorderrs, 'y') + minz, maxz = _digout_minmax(coorderrs, 'z') + self.auto_scale_xyz((minx, maxx), (miny, maxy), (minz, maxz), had_data) + + # Adapting errorbar containers for 3d case, assuming z-axis points "up" + errorbar_container = mcontainer.ErrorbarContainer( + (data_line, tuple(caplines), tuple(errlines)), + has_xerr=(xerr is not None or yerr is not None), + has_yerr=(zerr is not None), + label=label) + self.containers.append(errorbar_container) + + return errlines, caplines, limmarks + + def get_tightbbox(self, renderer=None, *, call_axes_locator=True, + bbox_extra_artists=None, for_layout_only=False): + ret = super().get_tightbbox(renderer, + call_axes_locator=call_axes_locator, + bbox_extra_artists=bbox_extra_artists, + for_layout_only=for_layout_only) + batch = [ret] + if self._axis3don: + for axis in self._axis_map.values(): + if axis.get_visible(): + axis_bb = martist._get_tightbbox_for_layout_only( + axis, renderer) + if axis_bb: + batch.append(axis_bb) + return mtransforms.Bbox.union(batch) + + @_preprocess_data() + def stem(self, x, y, z, *, linefmt='C0-', markerfmt='C0o', basefmt='C3-', + bottom=0, label=None, orientation='z', axlim_clip=False): + """ + Create a 3D stem plot. + + A stem plot draws lines perpendicular to a baseline, and places markers + at the heads. By default, the baseline is defined by *x* and *y*, and + stems are drawn vertically from *bottom* to *z*. + + Parameters + ---------- + x, y, z : array-like + The positions of the heads of the stems. The stems are drawn along + the *orientation*-direction from the baseline at *bottom* (in the + *orientation*-coordinate) to the heads. By default, the *x* and *y* + positions are used for the baseline and *z* for the head position, + but this can be changed by *orientation*. + + linefmt : str, default: 'C0-' + A string defining the properties of the vertical lines. Usually, + this will be a color or a color and a linestyle: + + ========= ============= + Character Line Style + ========= ============= + ``'-'`` solid line + ``'--'`` dashed line + ``'-.'`` dash-dot line + ``':'`` dotted line + ========= ============= + + Note: While it is technically possible to specify valid formats + other than color or color and linestyle (e.g. 'rx' or '-.'), this + is beyond the intention of the method and will most likely not + result in a reasonable plot. + + markerfmt : str, default: 'C0o' + A string defining the properties of the markers at the stem heads. + + basefmt : str, default: 'C3-' + A format string defining the properties of the baseline. + + bottom : float, default: 0 + The position of the baseline, in *orientation*-coordinates. + + label : str, optional + The label to use for the stems in legends. + + orientation : {'x', 'y', 'z'}, default: 'z' + The direction along which stems are drawn. + + axlim_clip : bool, default: False + Whether to hide stems that are outside the axes limits. + + .. versionadded:: 3.10 + + data : indexable object, optional + DATA_PARAMETER_PLACEHOLDER + + Returns + ------- + `.StemContainer` + The container may be treated like a tuple + (*markerline*, *stemlines*, *baseline*) + + Examples + -------- + .. plot:: gallery/mplot3d/stem3d_demo.py + """ + + from matplotlib.container import StemContainer + + had_data = self.has_data() + + _api.check_in_list(['x', 'y', 'z'], orientation=orientation) + + xlim = (np.min(x), np.max(x)) + ylim = (np.min(y), np.max(y)) + zlim = (np.min(z), np.max(z)) + + # Determine the appropriate plane for the baseline and the direction of + # stemlines based on the value of orientation. + if orientation == 'x': + basex, basexlim = y, ylim + basey, baseylim = z, zlim + lines = [[(bottom, thisy, thisz), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + elif orientation == 'y': + basex, basexlim = x, xlim + basey, baseylim = z, zlim + lines = [[(thisx, bottom, thisz), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + else: + basex, basexlim = x, xlim + basey, baseylim = y, ylim + lines = [[(thisx, thisy, bottom), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + + # Determine style for stem lines. + linestyle, linemarker, linecolor = _process_plot_format(linefmt) + if linestyle is None: + linestyle = mpl.rcParams['lines.linestyle'] + + # Plot everything in required order. + baseline, = self.plot(basex, basey, basefmt, zs=bottom, + zdir=orientation, label='_nolegend_') + stemlines = art3d.Line3DCollection( + lines, linestyles=linestyle, colors=linecolor, label='_nolegend_', + axlim_clip=axlim_clip) + self.add_collection(stemlines) + markerline, = self.plot(x, y, z, markerfmt, label='_nolegend_') + + stem_container = StemContainer((markerline, stemlines, baseline), + label=label) + self.add_container(stem_container) + + jx, jy, jz = art3d.juggle_axes(basexlim, baseylim, [bottom, bottom], + orientation) + self.auto_scale_xyz([*jx, *xlim], [*jy, *ylim], [*jz, *zlim], had_data) + + return stem_container + + stem3D = stem + + +def get_test_data(delta=0.05): + """Return a tuple X, Y, Z with a test data set.""" + x = y = np.arange(-3.0, 3.0, delta) + X, Y = np.meshgrid(x, y) + + Z1 = np.exp(-(X**2 + Y**2) / 2) / (2 * np.pi) + Z2 = (np.exp(-(((X - 1) / 1.5)**2 + ((Y - 1) / 0.5)**2) / 2) / + (2 * np.pi * 0.5 * 1.5)) + Z = Z2 - Z1 + + X = X * 10 + Y = Y * 10 + Z = Z * 500 + return X, Y, Z + + +class _Quaternion: + """ + Quaternions + consisting of scalar, along 1, and vector, with components along i, j, k + """ + + def __init__(self, scalar, vector): + self.scalar = scalar + self.vector = np.array(vector) + + def __neg__(self): + return self.__class__(-self.scalar, -self.vector) + + def __mul__(self, other): + """ + Product of two quaternions + i*i = j*j = k*k = i*j*k = -1 + Quaternion multiplication can be expressed concisely + using scalar and vector parts, + see + """ + return self.__class__( + self.scalar*other.scalar - np.dot(self.vector, other.vector), + self.scalar*other.vector + self.vector*other.scalar + + np.cross(self.vector, other.vector)) + + def conjugate(self): + """The conjugate quaternion -(1/2)*(q+i*q*i+j*q*j+k*q*k)""" + return self.__class__(self.scalar, -self.vector) + + @property + def norm(self): + """The 2-norm, q*q', a scalar""" + return self.scalar*self.scalar + np.dot(self.vector, self.vector) + + def normalize(self): + """Scaling such that norm equals 1""" + n = np.sqrt(self.norm) + return self.__class__(self.scalar/n, self.vector/n) + + def reciprocal(self): + """The reciprocal, 1/q = q'/(q*q') = q' / norm(q)""" + n = self.norm + return self.__class__(self.scalar/n, -self.vector/n) + + def __div__(self, other): + return self*other.reciprocal() + + __truediv__ = __div__ + + def rotate(self, v): + # Rotate the vector v by the quaternion q, i.e., + # calculate (the vector part of) q*v/q + v = self.__class__(0, v) + v = self*v/self + return v.vector + + def __eq__(self, other): + return (self.scalar == other.scalar) and (self.vector == other.vector).all + + def __repr__(self): + return "_Quaternion({}, {})".format(repr(self.scalar), repr(self.vector)) + + @classmethod + def rotate_from_to(cls, r1, r2): + """ + The quaternion for the shortest rotation from vector r1 to vector r2 + i.e., q = sqrt(r2*r1'), normalized. + If r1 and r2 are antiparallel, then the result is ambiguous; + a normal vector will be returned, and a warning will be issued. + """ + k = np.cross(r1, r2) + nk = np.linalg.norm(k) + th = np.arctan2(nk, np.dot(r1, r2)) + th /= 2 + if nk == 0: # r1 and r2 are parallel or anti-parallel + if np.dot(r1, r2) < 0: + warnings.warn("Rotation defined by anti-parallel vectors is ambiguous") + k = np.zeros(3) + k[np.argmin(r1*r1)] = 1 # basis vector most perpendicular to r1-r2 + k = np.cross(r1, k) + k = k / np.linalg.norm(k) # unit vector normal to r1-r2 + q = cls(0, k) + else: + q = cls(1, [0, 0, 0]) # = 1, no rotation + else: + q = cls(np.cos(th), k*np.sin(th)/nk) + return q + + @classmethod + def from_cardan_angles(cls, elev, azim, roll): + """ + Converts the angles to a quaternion + q = exp((roll/2)*e_x)*exp((elev/2)*e_y)*exp((-azim/2)*e_z) + i.e., the angles are a kind of Tait-Bryan angles, -z,y',x". + The angles should be given in radians, not degrees. + """ + ca, sa = np.cos(azim/2), np.sin(azim/2) + ce, se = np.cos(elev/2), np.sin(elev/2) + cr, sr = np.cos(roll/2), np.sin(roll/2) + + qw = ca*ce*cr + sa*se*sr + qx = ca*ce*sr - sa*se*cr + qy = ca*se*cr + sa*ce*sr + qz = ca*se*sr - sa*ce*cr + return cls(qw, [qx, qy, qz]) + + def as_cardan_angles(self): + """ + The inverse of `from_cardan_angles()`. + Note that the angles returned are in radians, not degrees. + The angles are not sensitive to the quaternion's norm(). + """ + qw = self.scalar + qx, qy, qz = self.vector[..., :] + azim = np.arctan2(2*(-qw*qz+qx*qy), qw*qw+qx*qx-qy*qy-qz*qz) + elev = np.arcsin(np.clip(2*(qw*qy+qz*qx)/(qw*qw+qx*qx+qy*qy+qz*qz), -1, 1)) + roll = np.arctan2(2*(qw*qx-qy*qz), qw*qw-qx*qx-qy*qy+qz*qz) + return elev, azim, roll diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axis3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axis3d.py new file mode 100644 index 0000000000000000000000000000000000000000..4da5031b990ca742eb0d303003af5c6a821c4db9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/axis3d.py @@ -0,0 +1,750 @@ +# axis3d.py, original mplot3d version by John Porter +# Created: 23 Sep 2005 +# Parts rewritten by Reinier Heeres + +import inspect + +import numpy as np + +import matplotlib as mpl +from matplotlib import ( + _api, artist, lines as mlines, axis as maxis, patches as mpatches, + transforms as mtransforms, colors as mcolors) +from . import art3d, proj3d + + +def _move_from_center(coord, centers, deltas, axmask=(True, True, True)): + """ + For each coordinate where *axmask* is True, move *coord* away from + *centers* by *deltas*. + """ + coord = np.asarray(coord) + return coord + axmask * np.copysign(1, coord - centers) * deltas + + +def _tick_update_position(tick, tickxs, tickys, labelpos): + """Update tick line and label position and style.""" + + tick.label1.set_position(labelpos) + tick.label2.set_position(labelpos) + tick.tick1line.set_visible(True) + tick.tick2line.set_visible(False) + tick.tick1line.set_linestyle('-') + tick.tick1line.set_marker('') + tick.tick1line.set_data(tickxs, tickys) + tick.gridline.set_data([0], [0]) + + +class Axis(maxis.XAxis): + """An Axis class for the 3D plots.""" + # These points from the unit cube make up the x, y and z-planes + _PLANES = ( + (0, 3, 7, 4), (1, 2, 6, 5), # yz planes + (0, 1, 5, 4), (3, 2, 6, 7), # xz planes + (0, 1, 2, 3), (4, 5, 6, 7), # xy planes + ) + + # Some properties for the axes + _AXINFO = { + 'x': {'i': 0, 'tickdir': 1, 'juggled': (1, 0, 2)}, + 'y': {'i': 1, 'tickdir': 0, 'juggled': (0, 1, 2)}, + 'z': {'i': 2, 'tickdir': 0, 'juggled': (0, 2, 1)}, + } + + def _old_init(self, adir, v_intervalx, d_intervalx, axes, *args, + rotate_label=None, **kwargs): + return locals() + + def _new_init(self, axes, *, rotate_label=None, **kwargs): + return locals() + + def __init__(self, *args, **kwargs): + params = _api.select_matching_signature( + [self._old_init, self._new_init], *args, **kwargs) + if "adir" in params: + _api.warn_deprecated( + "3.6", message=f"The signature of 3D Axis constructors has " + f"changed in %(since)s; the new signature is " + f"{inspect.signature(type(self).__init__)}", pending=True) + if params["adir"] != self.axis_name: + raise ValueError(f"Cannot instantiate {type(self).__name__} " + f"with adir={params['adir']!r}") + axes = params["axes"] + rotate_label = params["rotate_label"] + args = params.get("args", ()) + kwargs = params["kwargs"] + + name = self.axis_name + + self._label_position = 'default' + self._tick_position = 'default' + + # This is a temporary member variable. + # Do not depend on this existing in future releases! + self._axinfo = self._AXINFO[name].copy() + # Common parts + self._axinfo.update({ + 'label': {'va': 'center', 'ha': 'center', + 'rotation_mode': 'anchor'}, + 'color': mpl.rcParams[f'axes3d.{name}axis.panecolor'], + 'tick': { + 'inward_factor': 0.2, + 'outward_factor': 0.1, + }, + }) + + if mpl.rcParams['_internal.classic_mode']: + self._axinfo.update({ + 'axisline': {'linewidth': 0.75, 'color': (0, 0, 0, 1)}, + 'grid': { + 'color': (0.9, 0.9, 0.9, 1), + 'linewidth': 1.0, + 'linestyle': '-', + }, + }) + self._axinfo['tick'].update({ + 'linewidth': { + True: mpl.rcParams['lines.linewidth'], # major + False: mpl.rcParams['lines.linewidth'], # minor + } + }) + else: + self._axinfo.update({ + 'axisline': { + 'linewidth': mpl.rcParams['axes.linewidth'], + 'color': mpl.rcParams['axes.edgecolor'], + }, + 'grid': { + 'color': mpl.rcParams['grid.color'], + 'linewidth': mpl.rcParams['grid.linewidth'], + 'linestyle': mpl.rcParams['grid.linestyle'], + }, + }) + self._axinfo['tick'].update({ + 'linewidth': { + True: ( # major + mpl.rcParams['xtick.major.width'] if name in 'xz' + else mpl.rcParams['ytick.major.width']), + False: ( # minor + mpl.rcParams['xtick.minor.width'] if name in 'xz' + else mpl.rcParams['ytick.minor.width']), + } + }) + + super().__init__(axes, *args, **kwargs) + + # data and viewing intervals for this direction + if "d_intervalx" in params: + self.set_data_interval(*params["d_intervalx"]) + if "v_intervalx" in params: + self.set_view_interval(*params["v_intervalx"]) + self.set_rotate_label(rotate_label) + self._init3d() # Inline after init3d deprecation elapses. + + __init__.__signature__ = inspect.signature(_new_init) + adir = _api.deprecated("3.6", pending=True)( + property(lambda self: self.axis_name)) + + def _init3d(self): + self.line = mlines.Line2D( + xdata=(0, 0), ydata=(0, 0), + linewidth=self._axinfo['axisline']['linewidth'], + color=self._axinfo['axisline']['color'], + antialiased=True) + + # Store dummy data in Polygon object + self.pane = mpatches.Polygon([[0, 0], [0, 1]], closed=False) + self.set_pane_color(self._axinfo['color']) + + self.axes._set_artist_props(self.line) + self.axes._set_artist_props(self.pane) + self.gridlines = art3d.Line3DCollection([]) + self.axes._set_artist_props(self.gridlines) + self.axes._set_artist_props(self.label) + self.axes._set_artist_props(self.offsetText) + # Need to be able to place the label at the correct location + self.label._transform = self.axes.transData + self.offsetText._transform = self.axes.transData + + @_api.deprecated("3.6", pending=True) + def init3d(self): # After deprecation elapses, inline _init3d to __init__. + self._init3d() + + def get_major_ticks(self, numticks=None): + ticks = super().get_major_ticks(numticks) + for t in ticks: + for obj in [ + t.tick1line, t.tick2line, t.gridline, t.label1, t.label2]: + obj.set_transform(self.axes.transData) + return ticks + + def get_minor_ticks(self, numticks=None): + ticks = super().get_minor_ticks(numticks) + for t in ticks: + for obj in [ + t.tick1line, t.tick2line, t.gridline, t.label1, t.label2]: + obj.set_transform(self.axes.transData) + return ticks + + def set_ticks_position(self, position): + """ + Set the ticks position. + + Parameters + ---------- + position : {'lower', 'upper', 'both', 'default', 'none'} + The position of the bolded axis lines, ticks, and tick labels. + """ + _api.check_in_list(['lower', 'upper', 'both', 'default', 'none'], + position=position) + self._tick_position = position + + def get_ticks_position(self): + """ + Get the ticks position. + + Returns + ------- + str : {'lower', 'upper', 'both', 'default', 'none'} + The position of the bolded axis lines, ticks, and tick labels. + """ + return self._tick_position + + def set_label_position(self, position): + """ + Set the label position. + + Parameters + ---------- + position : {'lower', 'upper', 'both', 'default', 'none'} + The position of the axis label. + """ + _api.check_in_list(['lower', 'upper', 'both', 'default', 'none'], + position=position) + self._label_position = position + + def get_label_position(self): + """ + Get the label position. + + Returns + ------- + str : {'lower', 'upper', 'both', 'default', 'none'} + The position of the axis label. + """ + return self._label_position + + def set_pane_color(self, color, alpha=None): + """ + Set pane color. + + Parameters + ---------- + color : :mpltype:`color` + Color for axis pane. + alpha : float, optional + Alpha value for axis pane. If None, base it on *color*. + """ + color = mcolors.to_rgba(color, alpha) + self._axinfo['color'] = color + self.pane.set_edgecolor(color) + self.pane.set_facecolor(color) + self.pane.set_alpha(color[-1]) + self.stale = True + + def set_rotate_label(self, val): + """ + Whether to rotate the axis label: True, False or None. + If set to None the label will be rotated if longer than 4 chars. + """ + self._rotate_label = val + self.stale = True + + def get_rotate_label(self, text): + if self._rotate_label is not None: + return self._rotate_label + else: + return len(text) > 4 + + def _get_coord_info(self): + mins, maxs = np.array([ + self.axes.get_xbound(), + self.axes.get_ybound(), + self.axes.get_zbound(), + ]).T + + # Project the bounds along the current position of the cube: + bounds = mins[0], maxs[0], mins[1], maxs[1], mins[2], maxs[2] + bounds_proj = self.axes._transformed_cube(bounds) + + # Determine which one of the parallel planes are higher up: + means_z0 = np.zeros(3) + means_z1 = np.zeros(3) + for i in range(3): + means_z0[i] = np.mean(bounds_proj[self._PLANES[2 * i], 2]) + means_z1[i] = np.mean(bounds_proj[self._PLANES[2 * i + 1], 2]) + highs = means_z0 < means_z1 + + # Special handling for edge-on views + equals = np.abs(means_z0 - means_z1) <= np.finfo(float).eps + if np.sum(equals) == 2: + vertical = np.where(~equals)[0][0] + if vertical == 2: # looking at XY plane + highs = np.array([True, True, highs[2]]) + elif vertical == 1: # looking at XZ plane + highs = np.array([True, highs[1], False]) + elif vertical == 0: # looking at YZ plane + highs = np.array([highs[0], False, False]) + + return mins, maxs, bounds_proj, highs + + def _calc_centers_deltas(self, maxs, mins): + centers = 0.5 * (maxs + mins) + # In mpl3.8, the scale factor was 1/12. mpl3.9 changes this to + # 1/12 * 24/25 = 0.08 to compensate for the change in automargin + # behavior and keep appearance the same. The 24/25 factor is from the + # 1/48 padding added to each side of the axis in mpl3.8. + scale = 0.08 + deltas = (maxs - mins) * scale + return centers, deltas + + def _get_axis_line_edge_points(self, minmax, maxmin, position=None): + """Get the edge points for the black bolded axis line.""" + # When changing vertical axis some of the axes has to be + # moved to the other plane so it looks the same as if the z-axis + # was the vertical axis. + mb = [minmax, maxmin] # line from origin to nearest corner to camera + mb_rev = mb[::-1] + mm = [[mb, mb_rev, mb_rev], [mb_rev, mb_rev, mb], [mb, mb, mb]] + mm = mm[self.axes._vertical_axis][self._axinfo["i"]] + + juggled = self._axinfo["juggled"] + edge_point_0 = mm[0].copy() # origin point + + if ((position == 'lower' and mm[1][juggled[-1]] < mm[0][juggled[-1]]) or + (position == 'upper' and mm[1][juggled[-1]] > mm[0][juggled[-1]])): + edge_point_0[juggled[-1]] = mm[1][juggled[-1]] + else: + edge_point_0[juggled[0]] = mm[1][juggled[0]] + + edge_point_1 = edge_point_0.copy() + edge_point_1[juggled[1]] = mm[1][juggled[1]] + + return edge_point_0, edge_point_1 + + def _get_all_axis_line_edge_points(self, minmax, maxmin, axis_position=None): + # Determine edge points for the axis lines + edgep1s = [] + edgep2s = [] + position = [] + if axis_position in (None, 'default'): + edgep1, edgep2 = self._get_axis_line_edge_points(minmax, maxmin) + edgep1s = [edgep1] + edgep2s = [edgep2] + position = ['default'] + else: + edgep1_l, edgep2_l = self._get_axis_line_edge_points(minmax, maxmin, + position='lower') + edgep1_u, edgep2_u = self._get_axis_line_edge_points(minmax, maxmin, + position='upper') + if axis_position in ('lower', 'both'): + edgep1s.append(edgep1_l) + edgep2s.append(edgep2_l) + position.append('lower') + if axis_position in ('upper', 'both'): + edgep1s.append(edgep1_u) + edgep2s.append(edgep2_u) + position.append('upper') + return edgep1s, edgep2s, position + + def _get_tickdir(self, position): + """ + Get the direction of the tick. + + Parameters + ---------- + position : str, optional : {'upper', 'lower', 'default'} + The position of the axis. + + Returns + ------- + tickdir : int + Index which indicates which coordinate the tick line will + align with. + """ + _api.check_in_list(('upper', 'lower', 'default'), position=position) + + # TODO: Move somewhere else where it's triggered less: + tickdirs_base = [v["tickdir"] for v in self._AXINFO.values()] # default + elev_mod = np.mod(self.axes.elev + 180, 360) - 180 + azim_mod = np.mod(self.axes.azim, 360) + if position == 'upper': + if elev_mod >= 0: + tickdirs_base = [2, 2, 0] + else: + tickdirs_base = [1, 0, 0] + if 0 <= azim_mod < 180: + tickdirs_base[2] = 1 + elif position == 'lower': + if elev_mod >= 0: + tickdirs_base = [1, 0, 1] + else: + tickdirs_base = [2, 2, 1] + if 0 <= azim_mod < 180: + tickdirs_base[2] = 0 + info_i = [v["i"] for v in self._AXINFO.values()] + + i = self._axinfo["i"] + vert_ax = self.axes._vertical_axis + j = vert_ax - 2 + # default: tickdir = [[1, 2, 1], [2, 2, 0], [1, 0, 0]][vert_ax][i] + tickdir = np.roll(info_i, -j)[np.roll(tickdirs_base, j)][i] + return tickdir + + def active_pane(self): + mins, maxs, tc, highs = self._get_coord_info() + info = self._axinfo + index = info['i'] + if not highs[index]: + loc = mins[index] + plane = self._PLANES[2 * index] + else: + loc = maxs[index] + plane = self._PLANES[2 * index + 1] + xys = np.array([tc[p] for p in plane]) + return xys, loc + + def draw_pane(self, renderer): + """ + Draw pane. + + Parameters + ---------- + renderer : `~matplotlib.backend_bases.RendererBase` subclass + """ + renderer.open_group('pane3d', gid=self.get_gid()) + xys, loc = self.active_pane() + self.pane.xy = xys[:, :2] + self.pane.draw(renderer) + renderer.close_group('pane3d') + + def _axmask(self): + axmask = [True, True, True] + axmask[self._axinfo["i"]] = False + return axmask + + def _draw_ticks(self, renderer, edgep1, centers, deltas, highs, + deltas_per_point, pos): + ticks = self._update_ticks() + info = self._axinfo + index = info["i"] + juggled = info["juggled"] + + mins, maxs, tc, highs = self._get_coord_info() + centers, deltas = self._calc_centers_deltas(maxs, mins) + + # Draw ticks: + tickdir = self._get_tickdir(pos) + tickdelta = deltas[tickdir] if highs[tickdir] else -deltas[tickdir] + + tick_info = info['tick'] + tick_out = tick_info['outward_factor'] * tickdelta + tick_in = tick_info['inward_factor'] * tickdelta + tick_lw = tick_info['linewidth'] + edgep1_tickdir = edgep1[tickdir] + out_tickdir = edgep1_tickdir + tick_out + in_tickdir = edgep1_tickdir - tick_in + + default_label_offset = 8. # A rough estimate + points = deltas_per_point * deltas + for tick in ticks: + # Get tick line positions + pos = edgep1.copy() + pos[index] = tick.get_loc() + pos[tickdir] = out_tickdir + x1, y1, z1 = proj3d.proj_transform(*pos, self.axes.M) + pos[tickdir] = in_tickdir + x2, y2, z2 = proj3d.proj_transform(*pos, self.axes.M) + + # Get position of label + labeldeltas = (tick.get_pad() + default_label_offset) * points + + pos[tickdir] = edgep1_tickdir + pos = _move_from_center(pos, centers, labeldeltas, self._axmask()) + lx, ly, lz = proj3d.proj_transform(*pos, self.axes.M) + + _tick_update_position(tick, (x1, x2), (y1, y2), (lx, ly)) + tick.tick1line.set_linewidth(tick_lw[tick._major]) + tick.draw(renderer) + + def _draw_offset_text(self, renderer, edgep1, edgep2, labeldeltas, centers, + highs, pep, dx, dy): + # Get general axis information: + info = self._axinfo + index = info["i"] + juggled = info["juggled"] + tickdir = info["tickdir"] + + # Which of the two edge points do we want to + # use for locating the offset text? + if juggled[2] == 2: + outeredgep = edgep1 + outerindex = 0 + else: + outeredgep = edgep2 + outerindex = 1 + + pos = _move_from_center(outeredgep, centers, labeldeltas, + self._axmask()) + olx, oly, olz = proj3d.proj_transform(*pos, self.axes.M) + self.offsetText.set_text(self.major.formatter.get_offset()) + self.offsetText.set_position((olx, oly)) + angle = art3d._norm_text_angle(np.rad2deg(np.arctan2(dy, dx))) + self.offsetText.set_rotation(angle) + # Must set rotation mode to "anchor" so that + # the alignment point is used as the "fulcrum" for rotation. + self.offsetText.set_rotation_mode('anchor') + + # ---------------------------------------------------------------------- + # Note: the following statement for determining the proper alignment of + # the offset text. This was determined entirely by trial-and-error + # and should not be in any way considered as "the way". There are + # still some edge cases where alignment is not quite right, but this + # seems to be more of a geometry issue (in other words, I might be + # using the wrong reference points). + # + # (TT, FF, TF, FT) are the shorthand for the tuple of + # (centpt[tickdir] <= pep[tickdir, outerindex], + # centpt[index] <= pep[index, outerindex]) + # + # Three-letters (e.g., TFT, FTT) are short-hand for the array of bools + # from the variable 'highs'. + # --------------------------------------------------------------------- + centpt = proj3d.proj_transform(*centers, self.axes.M) + if centpt[tickdir] > pep[tickdir, outerindex]: + # if FT and if highs has an even number of Trues + if (centpt[index] <= pep[index, outerindex] + and np.count_nonzero(highs) % 2 == 0): + # Usually, this means align right, except for the FTT case, + # in which offset for axis 1 and 2 are aligned left. + if highs.tolist() == [False, True, True] and index in (1, 2): + align = 'left' + else: + align = 'right' + else: + # The FF case + align = 'left' + else: + # if TF and if highs has an even number of Trues + if (centpt[index] > pep[index, outerindex] + and np.count_nonzero(highs) % 2 == 0): + # Usually mean align left, except if it is axis 2 + align = 'right' if index == 2 else 'left' + else: + # The TT case + align = 'right' + + self.offsetText.set_va('center') + self.offsetText.set_ha(align) + self.offsetText.draw(renderer) + + def _draw_labels(self, renderer, edgep1, edgep2, labeldeltas, centers, dx, dy): + label = self._axinfo["label"] + + # Draw labels + lxyz = 0.5 * (edgep1 + edgep2) + lxyz = _move_from_center(lxyz, centers, labeldeltas, self._axmask()) + tlx, tly, tlz = proj3d.proj_transform(*lxyz, self.axes.M) + self.label.set_position((tlx, tly)) + if self.get_rotate_label(self.label.get_text()): + angle = art3d._norm_text_angle(np.rad2deg(np.arctan2(dy, dx))) + self.label.set_rotation(angle) + self.label.set_va(label['va']) + self.label.set_ha(label['ha']) + self.label.set_rotation_mode(label['rotation_mode']) + self.label.draw(renderer) + + @artist.allow_rasterization + def draw(self, renderer): + self.label._transform = self.axes.transData + self.offsetText._transform = self.axes.transData + renderer.open_group("axis3d", gid=self.get_gid()) + + # Get general axis information: + mins, maxs, tc, highs = self._get_coord_info() + centers, deltas = self._calc_centers_deltas(maxs, mins) + + # Calculate offset distances + # A rough estimate; points are ambiguous since 3D plots rotate + reltoinches = self.get_figure(root=False).dpi_scale_trans.inverted() + ax_inches = reltoinches.transform(self.axes.bbox.size) + ax_points_estimate = sum(72. * ax_inches) + deltas_per_point = 48 / ax_points_estimate + default_offset = 21. + labeldeltas = (self.labelpad + default_offset) * deltas_per_point * deltas + + # Determine edge points for the axis lines + minmax = np.where(highs, maxs, mins) # "origin" point + maxmin = np.where(~highs, maxs, mins) # "opposite" corner near camera + + for edgep1, edgep2, pos in zip(*self._get_all_axis_line_edge_points( + minmax, maxmin, self._tick_position)): + # Project the edge points along the current position + pep = proj3d._proj_trans_points([edgep1, edgep2], self.axes.M) + pep = np.asarray(pep) + + # The transAxes transform is used because the Text object + # rotates the text relative to the display coordinate system. + # Therefore, if we want the labels to remain parallel to the + # axis regardless of the aspect ratio, we need to convert the + # edge points of the plane to display coordinates and calculate + # an angle from that. + # TODO: Maybe Text objects should handle this themselves? + dx, dy = (self.axes.transAxes.transform([pep[0:2, 1]]) - + self.axes.transAxes.transform([pep[0:2, 0]]))[0] + + # Draw the lines + self.line.set_data(pep[0], pep[1]) + self.line.draw(renderer) + + # Draw ticks + self._draw_ticks(renderer, edgep1, centers, deltas, highs, + deltas_per_point, pos) + + # Draw Offset text + self._draw_offset_text(renderer, edgep1, edgep2, labeldeltas, + centers, highs, pep, dx, dy) + + for edgep1, edgep2, pos in zip(*self._get_all_axis_line_edge_points( + minmax, maxmin, self._label_position)): + # See comments above + pep = proj3d._proj_trans_points([edgep1, edgep2], self.axes.M) + pep = np.asarray(pep) + dx, dy = (self.axes.transAxes.transform([pep[0:2, 1]]) - + self.axes.transAxes.transform([pep[0:2, 0]]))[0] + + # Draw labels + self._draw_labels(renderer, edgep1, edgep2, labeldeltas, centers, dx, dy) + + renderer.close_group('axis3d') + self.stale = False + + @artist.allow_rasterization + def draw_grid(self, renderer): + if not self.axes._draw_grid: + return + + renderer.open_group("grid3d", gid=self.get_gid()) + + ticks = self._update_ticks() + if len(ticks): + # Get general axis information: + info = self._axinfo + index = info["i"] + + mins, maxs, tc, highs = self._get_coord_info() + + minmax = np.where(highs, maxs, mins) + maxmin = np.where(~highs, maxs, mins) + + # Grid points where the planes meet + xyz0 = np.tile(minmax, (len(ticks), 1)) + xyz0[:, index] = [tick.get_loc() for tick in ticks] + + # Grid lines go from the end of one plane through the plane + # intersection (at xyz0) to the end of the other plane. The first + # point (0) differs along dimension index-2 and the last (2) along + # dimension index-1. + lines = np.stack([xyz0, xyz0, xyz0], axis=1) + lines[:, 0, index - 2] = maxmin[index - 2] + lines[:, 2, index - 1] = maxmin[index - 1] + self.gridlines.set_segments(lines) + gridinfo = info['grid'] + self.gridlines.set_color(gridinfo['color']) + self.gridlines.set_linewidth(gridinfo['linewidth']) + self.gridlines.set_linestyle(gridinfo['linestyle']) + self.gridlines.do_3d_projection() + self.gridlines.draw(renderer) + + renderer.close_group('grid3d') + + # TODO: Get this to work (more) properly when mplot3d supports the + # transforms framework. + def get_tightbbox(self, renderer=None, *, for_layout_only=False): + # docstring inherited + if not self.get_visible(): + return + # We have to directly access the internal data structures + # (and hope they are up to date) because at draw time we + # shift the ticks and their labels around in (x, y) space + # based on the projection, the current view port, and their + # position in 3D space. If we extend the transforms framework + # into 3D we would not need to do this different book keeping + # than we do in the normal axis + major_locs = self.get_majorticklocs() + minor_locs = self.get_minorticklocs() + + ticks = [*self.get_minor_ticks(len(minor_locs)), + *self.get_major_ticks(len(major_locs))] + view_low, view_high = self.get_view_interval() + if view_low > view_high: + view_low, view_high = view_high, view_low + interval_t = self.get_transform().transform([view_low, view_high]) + + ticks_to_draw = [] + for tick in ticks: + try: + loc_t = self.get_transform().transform(tick.get_loc()) + except AssertionError: + # Transform.transform doesn't allow masked values but + # some scales might make them, so we need this try/except. + pass + else: + if mtransforms._interval_contains_close(interval_t, loc_t): + ticks_to_draw.append(tick) + + ticks = ticks_to_draw + + bb_1, bb_2 = self._get_ticklabel_bboxes(ticks, renderer) + other = [] + + if self.line.get_visible(): + other.append(self.line.get_window_extent(renderer)) + if (self.label.get_visible() and not for_layout_only and + self.label.get_text()): + other.append(self.label.get_window_extent(renderer)) + + return mtransforms.Bbox.union([*bb_1, *bb_2, *other]) + + d_interval = _api.deprecated( + "3.6", alternative="get_data_interval", pending=True)( + property(lambda self: self.get_data_interval(), + lambda self, minmax: self.set_data_interval(*minmax))) + v_interval = _api.deprecated( + "3.6", alternative="get_view_interval", pending=True)( + property(lambda self: self.get_view_interval(), + lambda self, minmax: self.set_view_interval(*minmax))) + + +class XAxis(Axis): + axis_name = "x" + get_view_interval, set_view_interval = maxis._make_getset_interval( + "view", "xy_viewLim", "intervalx") + get_data_interval, set_data_interval = maxis._make_getset_interval( + "data", "xy_dataLim", "intervalx") + + +class YAxis(Axis): + axis_name = "y" + get_view_interval, set_view_interval = maxis._make_getset_interval( + "view", "xy_viewLim", "intervaly") + get_data_interval, set_data_interval = maxis._make_getset_interval( + "data", "xy_dataLim", "intervaly") + + +class ZAxis(Axis): + axis_name = "z" + get_view_interval, set_view_interval = maxis._make_getset_interval( + "view", "zz_viewLim", "intervalx") + get_data_interval, set_data_interval = maxis._make_getset_interval( + "data", "zz_dataLim", "intervalx") diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/proj3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/proj3d.py new file mode 100644 index 0000000000000000000000000000000000000000..923bd32c9ce032ad27379f8d8c6436564ec16d75 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/proj3d.py @@ -0,0 +1,219 @@ +""" +Various transforms used for by the 3D code +""" + +import numpy as np + +from matplotlib import _api + + +def world_transformation(xmin, xmax, + ymin, ymax, + zmin, zmax, pb_aspect=None): + """ + Produce a matrix that scales homogeneous coords in the specified ranges + to [0, 1], or [0, pb_aspect[i]] if the plotbox aspect ratio is specified. + """ + dx = xmax - xmin + dy = ymax - ymin + dz = zmax - zmin + if pb_aspect is not None: + ax, ay, az = pb_aspect + dx /= ax + dy /= ay + dz /= az + + return np.array([[1/dx, 0, 0, -xmin/dx], + [ 0, 1/dy, 0, -ymin/dy], + [ 0, 0, 1/dz, -zmin/dz], + [ 0, 0, 0, 1]]) + + +def _rotation_about_vector(v, angle): + """ + Produce a rotation matrix for an angle in radians about a vector. + """ + vx, vy, vz = v / np.linalg.norm(v) + s = np.sin(angle) + c = np.cos(angle) + t = 2*np.sin(angle/2)**2 # more numerically stable than t = 1-c + + R = np.array([ + [t*vx*vx + c, t*vx*vy - vz*s, t*vx*vz + vy*s], + [t*vy*vx + vz*s, t*vy*vy + c, t*vy*vz - vx*s], + [t*vz*vx - vy*s, t*vz*vy + vx*s, t*vz*vz + c]]) + + return R + + +def _view_axes(E, R, V, roll): + """ + Get the unit viewing axes in data coordinates. + + Parameters + ---------- + E : 3-element numpy array + The coordinates of the eye/camera. + R : 3-element numpy array + The coordinates of the center of the view box. + V : 3-element numpy array + Unit vector in the direction of the vertical axis. + roll : float + The roll angle in radians. + + Returns + ------- + u : 3-element numpy array + Unit vector pointing towards the right of the screen. + v : 3-element numpy array + Unit vector pointing towards the top of the screen. + w : 3-element numpy array + Unit vector pointing out of the screen. + """ + w = (E - R) + w = w/np.linalg.norm(w) + u = np.cross(V, w) + u = u/np.linalg.norm(u) + v = np.cross(w, u) # Will be a unit vector + + # Save some computation for the default roll=0 + if roll != 0: + # A positive rotation of the camera is a negative rotation of the world + Rroll = _rotation_about_vector(w, -roll) + u = np.dot(Rroll, u) + v = np.dot(Rroll, v) + return u, v, w + + +def _view_transformation_uvw(u, v, w, E): + """ + Return the view transformation matrix. + + Parameters + ---------- + u : 3-element numpy array + Unit vector pointing towards the right of the screen. + v : 3-element numpy array + Unit vector pointing towards the top of the screen. + w : 3-element numpy array + Unit vector pointing out of the screen. + E : 3-element numpy array + The coordinates of the eye/camera. + """ + Mr = np.eye(4) + Mt = np.eye(4) + Mr[:3, :3] = [u, v, w] + Mt[:3, -1] = -E + M = np.dot(Mr, Mt) + return M + + +def _persp_transformation(zfront, zback, focal_length): + e = focal_length + a = 1 # aspect ratio + b = (zfront+zback)/(zfront-zback) + c = -2*(zfront*zback)/(zfront-zback) + proj_matrix = np.array([[e, 0, 0, 0], + [0, e/a, 0, 0], + [0, 0, b, c], + [0, 0, -1, 0]]) + return proj_matrix + + +def _ortho_transformation(zfront, zback): + # note: w component in the resulting vector will be (zback-zfront), not 1 + a = -(zfront + zback) + b = -(zfront - zback) + proj_matrix = np.array([[2, 0, 0, 0], + [0, 2, 0, 0], + [0, 0, -2, 0], + [0, 0, a, b]]) + return proj_matrix + + +def _proj_transform_vec(vec, M): + vecw = np.dot(M, vec.data) + w = vecw[3] + txs, tys, tzs = vecw[0]/w, vecw[1]/w, vecw[2]/w + if np.ma.isMA(vec[0]): # we check each to protect for scalars + txs = np.ma.array(txs, mask=vec[0].mask) + if np.ma.isMA(vec[1]): + tys = np.ma.array(tys, mask=vec[1].mask) + if np.ma.isMA(vec[2]): + tzs = np.ma.array(tzs, mask=vec[2].mask) + return txs, tys, tzs + + +def _proj_transform_vec_clip(vec, M, focal_length): + vecw = np.dot(M, vec.data) + w = vecw[3] + txs, tys, tzs = vecw[0] / w, vecw[1] / w, vecw[2] / w + if np.isinf(focal_length): # don't clip orthographic projection + tis = np.ones(txs.shape, dtype=bool) + else: + tis = (-1 <= txs) & (txs <= 1) & (-1 <= tys) & (tys <= 1) & (tzs <= 0) + if np.ma.isMA(vec[0]): + tis = tis & ~vec[0].mask + if np.ma.isMA(vec[1]): + tis = tis & ~vec[1].mask + if np.ma.isMA(vec[2]): + tis = tis & ~vec[2].mask + + txs = np.ma.masked_array(txs, ~tis) + tys = np.ma.masked_array(tys, ~tis) + tzs = np.ma.masked_array(tzs, ~tis) + return txs, tys, tzs, tis + + +def inv_transform(xs, ys, zs, invM): + """ + Transform the points by the inverse of the projection matrix, *invM*. + """ + vec = _vec_pad_ones(xs, ys, zs) + vecr = np.dot(invM, vec) + if vecr.shape == (4,): + vecr = vecr.reshape((4, 1)) + for i in range(vecr.shape[1]): + if vecr[3][i] != 0: + vecr[:, i] = vecr[:, i] / vecr[3][i] + return vecr[0], vecr[1], vecr[2] + + +def _vec_pad_ones(xs, ys, zs): + if np.ma.isMA(xs) or np.ma.isMA(ys) or np.ma.isMA(zs): + return np.ma.array([xs, ys, zs, np.ones_like(xs)]) + else: + return np.array([xs, ys, zs, np.ones_like(xs)]) + + +def proj_transform(xs, ys, zs, M): + """ + Transform the points by the projection matrix *M*. + """ + vec = _vec_pad_ones(xs, ys, zs) + return _proj_transform_vec(vec, M) + + +@_api.deprecated("3.10") +def proj_transform_clip(xs, ys, zs, M): + return _proj_transform_clip(xs, ys, zs, M, focal_length=np.inf) + + +def _proj_transform_clip(xs, ys, zs, M, focal_length): + """ + Transform the points by the projection matrix + and return the clipping result + returns txs, tys, tzs, tis + """ + vec = _vec_pad_ones(xs, ys, zs) + return _proj_transform_vec_clip(vec, M, focal_length) + + +def _proj_points(points, M): + return np.column_stack(_proj_trans_points(points, M)) + + +def _proj_trans_points(points, M): + points = np.asanyarray(points) + xs, ys, zs = points[:, 0], points[:, 1], points[:, 2] + return proj_transform(xs, ys, zs, M) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__init__.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ea4d8ed16a6a24a8c15ab2956ef678a7f256cd80 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__init__.py @@ -0,0 +1,10 @@ +from pathlib import Path + + +# Check that the test directories exist +if not (Path(__file__).parent / "baseline_images").exists(): + raise OSError( + 'The baseline image directory does not exist. ' + 'This is most likely because the test data is not installed. ' + 'You may need to install matplotlib from source to get the ' + 'test data.') diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b3b5e64564b5341e3b0f0546402e14bfecf28e33 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/conftest.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/conftest.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..243be8e765c2ade60b557422282b1a2ee6e0536f Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/conftest.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_art3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_art3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72ad9ec6d06bbaa9470032c5408e7f2fe5742948 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_art3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_axes3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_axes3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff07a88ca678723df24c69e5c8b1965ae9ac9e38 Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_axes3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_legend3d.cpython-310.pyc b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_legend3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8c78f217af76867e0f77067d849f767a782dc2b Binary files /dev/null and b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/__pycache__/test_legend3d.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/conftest.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..61c2de3e07bac4db323f8704961264d123e01544 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/conftest.py @@ -0,0 +1,2 @@ +from matplotlib.testing.conftest import (mpl_test_settings, # noqa + pytest_configure, pytest_unconfigure) diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_art3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_art3d.py new file mode 100644 index 0000000000000000000000000000000000000000..174c12608ae933c1b918f45e57abfe6a2fb552bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_art3d.py @@ -0,0 +1,102 @@ +import numpy as np + +import matplotlib.pyplot as plt + +from matplotlib.backend_bases import MouseEvent +from mpl_toolkits.mplot3d.art3d import ( + Line3DCollection, + Poly3DCollection, + _all_points_on_plane, +) + + +def test_scatter_3d_projection_conservation(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + # fix axes3d projection + ax.roll = 0 + ax.elev = 0 + ax.azim = -45 + ax.stale = True + + x = [0, 1, 2, 3, 4] + scatter_collection = ax.scatter(x, x, x) + fig.canvas.draw_idle() + + # Get scatter location on canvas and freeze the data + scatter_offset = scatter_collection.get_offsets() + scatter_location = ax.transData.transform(scatter_offset) + + # Yaw -44 and -46 are enough to produce two set of scatter + # with opposite z-order without moving points too far + for azim in (-44, -46): + ax.azim = azim + ax.stale = True + fig.canvas.draw_idle() + + for i in range(5): + # Create a mouse event used to locate and to get index + # from each dots + event = MouseEvent("button_press_event", fig.canvas, + *scatter_location[i, :]) + contains, ind = scatter_collection.contains(event) + assert contains is True + assert len(ind["ind"]) == 1 + assert ind["ind"][0] == i + + +def test_zordered_error(): + # Smoke test for https://github.com/matplotlib/matplotlib/issues/26497 + lc = [(np.fromiter([0.0, 0.0, 0.0], dtype="float"), + np.fromiter([1.0, 1.0, 1.0], dtype="float"))] + pc = [np.fromiter([0.0, 0.0], dtype="float"), + np.fromiter([0.0, 1.0], dtype="float"), + np.fromiter([1.0, 1.0], dtype="float")] + + fig = plt.figure() + ax = fig.add_subplot(projection="3d") + ax.add_collection(Line3DCollection(lc)) + ax.scatter(*pc, visible=False) + plt.draw() + + +def test_all_points_on_plane(): + # Non-coplanar points + points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]) + assert not _all_points_on_plane(*points.T) + + # Duplicate points + points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 0]]) + assert _all_points_on_plane(*points.T) + + # NaN values + points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, np.nan]]) + assert _all_points_on_plane(*points.T) + + # Less than 3 unique points + points = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]]) + assert _all_points_on_plane(*points.T) + + # All points lie on a line + points = np.array([[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0]]) + assert _all_points_on_plane(*points.T) + + # All points lie on two lines, with antiparallel vectors + points = np.array([[-2, 2, 0], [-1, 1, 0], [1, -1, 0], + [0, 0, 0], [2, 0, 0], [1, 0, 0]]) + assert _all_points_on_plane(*points.T) + + # All points lie on a plane + points = np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0], [1, 2, 0]]) + assert _all_points_on_plane(*points.T) + + +def test_generate_normals(): + # Smoke test for https://github.com/matplotlib/matplotlib/issues/29156 + vertices = ((0, 0, 0), (0, 5, 0), (5, 5, 0), (5, 0, 0)) + shape = Poly3DCollection([vertices], edgecolors='r', shade=True) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.add_collection3d(shape) + plt.draw() diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_axes3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_axes3d.py new file mode 100644 index 0000000000000000000000000000000000000000..b30c2556b105e70e3e08fa3877c712f78f09d95c --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_axes3d.py @@ -0,0 +1,2688 @@ +import functools +import itertools +import platform +import sys + +import pytest + +from mpl_toolkits.mplot3d import Axes3D, axes3d, proj3d, art3d +from mpl_toolkits.mplot3d.axes3d import _Quaternion as Quaternion +import matplotlib as mpl +from matplotlib.backend_bases import (MouseButton, MouseEvent, + NavigationToolbar2) +from matplotlib import cm +from matplotlib import colors as mcolors, patches as mpatch +from matplotlib.testing.decorators import image_comparison, check_figures_equal +from matplotlib.testing.widgets import mock_event +from matplotlib.collections import LineCollection, PolyCollection +from matplotlib.patches import Circle, PathPatch +from matplotlib.path import Path +from matplotlib.text import Text + +import matplotlib.pyplot as plt +import numpy as np + + +mpl3d_image_comparison = functools.partial( + image_comparison, remove_text=True, style='default') + + +def plot_cuboid(ax, scale): + # plot a rectangular cuboid with side lengths given by scale (x, y, z) + r = [0, 1] + pts = itertools.combinations(np.array(list(itertools.product(r, r, r))), 2) + for start, end in pts: + if np.sum(np.abs(start - end)) == r[1] - r[0]: + ax.plot3D(*zip(start*np.array(scale), end*np.array(scale))) + + +@check_figures_equal(extensions=["png"]) +def test_invisible_axes(fig_test, fig_ref): + ax = fig_test.subplots(subplot_kw=dict(projection='3d')) + ax.set_visible(False) + + +@mpl3d_image_comparison(['grid_off.png'], style='mpl20') +def test_grid_off(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.grid(False) + + +@mpl3d_image_comparison(['invisible_ticks_axis.png'], style='mpl20') +def test_invisible_ticks_axis(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_zticks([]) + for axis in [ax.xaxis, ax.yaxis, ax.zaxis]: + axis.line.set_visible(False) + + +@mpl3d_image_comparison(['axis_positions.png'], remove_text=False, style='mpl20') +def test_axis_positions(): + positions = ['upper', 'lower', 'both', 'none'] + fig, axs = plt.subplots(2, 2, subplot_kw={'projection': '3d'}) + for ax, pos in zip(axs.flatten(), positions): + for axis in ax.xaxis, ax.yaxis, ax.zaxis: + axis.set_label_position(pos) + axis.set_ticks_position(pos) + title = f'{pos}' + ax.set(xlabel='x', ylabel='y', zlabel='z', title=title) + + +@mpl3d_image_comparison(['aspects.png'], remove_text=False, style='mpl20') +def test_aspects(): + aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz', 'equal') + _, axs = plt.subplots(2, 3, subplot_kw={'projection': '3d'}) + + for ax in axs.flatten()[0:-1]: + plot_cuboid(ax, scale=[1, 1, 5]) + # plot a cube as well to cover github #25443 + plot_cuboid(axs[1][2], scale=[1, 1, 1]) + + for i, ax in enumerate(axs.flatten()): + ax.set_title(aspects[i]) + ax.set_box_aspect((3, 4, 5)) + ax.set_aspect(aspects[i], adjustable='datalim') + axs[1][2].set_title('equal (cube)') + + +@mpl3d_image_comparison(['aspects_adjust_box.png'], + remove_text=False, style='mpl20') +def test_aspects_adjust_box(): + aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz') + fig, axs = plt.subplots(1, len(aspects), subplot_kw={'projection': '3d'}, + figsize=(11, 3)) + + for i, ax in enumerate(axs): + plot_cuboid(ax, scale=[4, 3, 5]) + ax.set_title(aspects[i]) + ax.set_aspect(aspects[i], adjustable='box') + + +def test_axes3d_repr(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.set_label('label') + ax.set_title('title') + ax.set_xlabel('x') + ax.set_ylabel('y') + ax.set_zlabel('z') + assert repr(ax) == ( + "") + + +@mpl3d_image_comparison(['axes3d_primary_views.png'], style='mpl20', + tol=0.05 if sys.platform == "darwin" else 0) +def test_axes3d_primary_views(): + # (elev, azim, roll) + views = [(90, -90, 0), # XY + (0, -90, 0), # XZ + (0, 0, 0), # YZ + (-90, 90, 0), # -XY + (0, 90, 0), # -XZ + (0, 180, 0)] # -YZ + # When viewing primary planes, draw the two visible axes so they intersect + # at their low values + fig, axs = plt.subplots(2, 3, subplot_kw={'projection': '3d'}) + for i, ax in enumerate(axs.flat): + ax.set_xlabel('x') + ax.set_ylabel('y') + ax.set_zlabel('z') + ax.set_proj_type('ortho') + ax.view_init(elev=views[i][0], azim=views[i][1], roll=views[i][2]) + plt.tight_layout() + + +@mpl3d_image_comparison(['bar3d.png'], style='mpl20') +def test_bar3d(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + for c, z in zip(['r', 'g', 'b', 'y'], [30, 20, 10, 0]): + xs = np.arange(20) + ys = np.arange(20) + cs = [c] * len(xs) + cs[0] = 'c' + ax.bar(xs, ys, zs=z, zdir='y', align='edge', color=cs, alpha=0.8) + + +def test_bar3d_colors(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + for c in ['red', 'green', 'blue', 'yellow']: + xs = np.arange(len(c)) + ys = np.zeros_like(xs) + zs = np.zeros_like(ys) + # Color names with same length as xs/ys/zs should not be split into + # individual letters. + ax.bar3d(xs, ys, zs, 1, 1, 1, color=c) + + +@mpl3d_image_comparison(['bar3d_shaded.png'], style='mpl20') +def test_bar3d_shaded(): + x = np.arange(4) + y = np.arange(5) + x2d, y2d = np.meshgrid(x, y) + x2d, y2d = x2d.ravel(), y2d.ravel() + z = x2d + y2d + 1 # Avoid triggering bug with zero-depth boxes. + + views = [(30, -60, 0), (30, 30, 30), (-30, 30, -90), (300, -30, 0)] + fig = plt.figure(figsize=plt.figaspect(1 / len(views))) + axs = fig.subplots( + 1, len(views), + subplot_kw=dict(projection='3d') + ) + for ax, (elev, azim, roll) in zip(axs, views): + ax.bar3d(x2d, y2d, x2d * 0, 1, 1, z, shade=True) + ax.view_init(elev=elev, azim=azim, roll=roll) + fig.canvas.draw() + + +@mpl3d_image_comparison(['bar3d_notshaded.png'], style='mpl20') +def test_bar3d_notshaded(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + x = np.arange(4) + y = np.arange(5) + x2d, y2d = np.meshgrid(x, y) + x2d, y2d = x2d.ravel(), y2d.ravel() + z = x2d + y2d + ax.bar3d(x2d, y2d, x2d * 0, 1, 1, z, shade=False) + fig.canvas.draw() + + +def test_bar3d_lightsource(): + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection="3d") + + ls = mcolors.LightSource(azdeg=0, altdeg=90) + + length, width = 3, 4 + area = length * width + + x, y = np.meshgrid(np.arange(length), np.arange(width)) + x = x.ravel() + y = y.ravel() + dz = x + y + + color = [cm.coolwarm(i/area) for i in range(area)] + + collection = ax.bar3d(x=x, y=y, z=0, + dx=1, dy=1, dz=dz, + color=color, shade=True, lightsource=ls) + + # Testing that the custom 90° lightsource produces different shading on + # the top facecolors compared to the default, and that those colors are + # precisely (within floating point rounding errors of 4 ULP) the colors + # from the colormap, due to the illumination parallel to the z-axis. + np.testing.assert_array_max_ulp(color, collection._facecolor3d[1::6], 4) + + +@mpl3d_image_comparison(['contour3d.png'], style='mpl20', + tol=0 if platform.machine() == 'x86_64' else 0.002) +def test_contour3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.contour(X, Y, Z, zdir='z', offset=-100, cmap=cm.coolwarm) + ax.contour(X, Y, Z, zdir='x', offset=-40, cmap=cm.coolwarm) + ax.contour(X, Y, Z, zdir='y', offset=40, cmap=cm.coolwarm) + ax.axis(xmin=-40, xmax=40, ymin=-40, ymax=40, zmin=-100, zmax=100) + + +@mpl3d_image_comparison(['contour3d_extend3d.png'], style='mpl20') +def test_contour3d_extend3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.contour(X, Y, Z, zdir='z', offset=-100, cmap=cm.coolwarm, extend3d=True) + ax.set_xlim(-30, 30) + ax.set_ylim(-20, 40) + ax.set_zlim(-80, 80) + + +@mpl3d_image_comparison(['contourf3d.png'], style='mpl20') +def test_contourf3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.contourf(X, Y, Z, zdir='z', offset=-100, cmap=cm.coolwarm) + ax.contourf(X, Y, Z, zdir='x', offset=-40, cmap=cm.coolwarm) + ax.contourf(X, Y, Z, zdir='y', offset=40, cmap=cm.coolwarm) + ax.set_xlim(-40, 40) + ax.set_ylim(-40, 40) + ax.set_zlim(-100, 100) + + +@mpl3d_image_comparison(['contourf3d_fill.png'], style='mpl20') +def test_contourf3d_fill(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y = np.meshgrid(np.arange(-2, 2, 0.25), np.arange(-2, 2, 0.25)) + Z = X.clip(0, 0) + # This produces holes in the z=0 surface that causes rendering errors if + # the Poly3DCollection is not aware of path code information (issue #4784) + Z[::5, ::5] = 0.1 + ax.contourf(X, Y, Z, offset=0, levels=[-0.1, 0], cmap=cm.coolwarm) + ax.set_xlim(-2, 2) + ax.set_ylim(-2, 2) + ax.set_zlim(-1, 1) + + +@pytest.mark.parametrize('extend, levels', [['both', [2, 4, 6]], + ['min', [2, 4, 6, 8]], + ['max', [0, 2, 4, 6]]]) +@check_figures_equal(extensions=["png"]) +def test_contourf3d_extend(fig_test, fig_ref, extend, levels): + X, Y = np.meshgrid(np.arange(-2, 2, 0.25), np.arange(-2, 2, 0.25)) + # Z is in the range [0, 8] + Z = X**2 + Y**2 + + # Manually set the over/under colors to be the end of the colormap + cmap = mpl.colormaps['viridis'].copy() + cmap.set_under(cmap(0)) + cmap.set_over(cmap(255)) + # Set vmin/max to be the min/max values plotted on the reference image + kwargs = {'vmin': 1, 'vmax': 7, 'cmap': cmap} + + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.contourf(X, Y, Z, levels=[0, 2, 4, 6, 8], **kwargs) + + ax_test = fig_test.add_subplot(projection='3d') + ax_test.contourf(X, Y, Z, levels, extend=extend, **kwargs) + + for ax in [ax_ref, ax_test]: + ax.set_xlim(-2, 2) + ax.set_ylim(-2, 2) + ax.set_zlim(-10, 10) + + +@mpl3d_image_comparison(['tricontour.png'], tol=0.02, style='mpl20') +def test_tricontour(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + + np.random.seed(19680801) + x = np.random.rand(1000) - 0.5 + y = np.random.rand(1000) - 0.5 + z = -(x**2 + y**2) + + ax = fig.add_subplot(1, 2, 1, projection='3d') + ax.tricontour(x, y, z) + ax = fig.add_subplot(1, 2, 2, projection='3d') + ax.tricontourf(x, y, z) + + +def test_contour3d_1d_input(): + # Check that 1D sequences of different length for {x, y} doesn't error + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + nx, ny = 30, 20 + x = np.linspace(-10, 10, nx) + y = np.linspace(-10, 10, ny) + z = np.random.randint(0, 2, [ny, nx]) + ax.contour(x, y, z, [0.5]) + + +@mpl3d_image_comparison(['lines3d.png'], style='mpl20') +def test_lines3d(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + theta = np.linspace(-4 * np.pi, 4 * np.pi, 100) + z = np.linspace(-2, 2, 100) + r = z ** 2 + 1 + x = r * np.sin(theta) + y = r * np.cos(theta) + ax.plot(x, y, z) + + +@check_figures_equal(extensions=["png"]) +def test_plot_scalar(fig_test, fig_ref): + ax1 = fig_test.add_subplot(projection='3d') + ax1.plot([1], [1], "o") + ax2 = fig_ref.add_subplot(projection='3d') + ax2.plot(1, 1, "o") + + +def test_invalid_line_data(): + with pytest.raises(RuntimeError, match='x must be'): + art3d.Line3D(0, [], []) + with pytest.raises(RuntimeError, match='y must be'): + art3d.Line3D([], 0, []) + with pytest.raises(RuntimeError, match='z must be'): + art3d.Line3D([], [], 0) + + line = art3d.Line3D([], [], []) + with pytest.raises(RuntimeError, match='x must be'): + line.set_data_3d(0, [], []) + with pytest.raises(RuntimeError, match='y must be'): + line.set_data_3d([], 0, []) + with pytest.raises(RuntimeError, match='z must be'): + line.set_data_3d([], [], 0) + + +@mpl3d_image_comparison(['mixedsubplot.png'], style='mpl20') +def test_mixedsubplots(): + def f(t): + return np.cos(2*np.pi*t) * np.exp(-t) + + t1 = np.arange(0.0, 5.0, 0.1) + t2 = np.arange(0.0, 5.0, 0.02) + + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure(figsize=plt.figaspect(2.)) + ax = fig.add_subplot(2, 1, 1) + ax.plot(t1, f(t1), 'bo', t2, f(t2), 'k--', markerfacecolor='green') + ax.grid(True) + + ax = fig.add_subplot(2, 1, 2, projection='3d') + X, Y = np.meshgrid(np.arange(-5, 5, 0.25), np.arange(-5, 5, 0.25)) + R = np.hypot(X, Y) + Z = np.sin(R) + + ax.plot_surface(X, Y, Z, rcount=40, ccount=40, + linewidth=0, antialiased=False) + + ax.set_zlim3d(-1, 1) + + +@check_figures_equal(extensions=['png']) +def test_tight_layout_text(fig_test, fig_ref): + # text is currently ignored in tight layout. So the order of text() and + # tight_layout() calls should not influence the result. + ax1 = fig_test.add_subplot(projection='3d') + ax1.text(.5, .5, .5, s='some string') + fig_test.tight_layout() + + ax2 = fig_ref.add_subplot(projection='3d') + fig_ref.tight_layout() + ax2.text(.5, .5, .5, s='some string') + + +@mpl3d_image_comparison(['scatter3d.png'], style='mpl20') +def test_scatter3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.scatter(np.arange(10), np.arange(10), np.arange(10), + c='r', marker='o') + x = y = z = np.arange(10, 20) + ax.scatter(x, y, z, c='b', marker='^') + z[-1] = 0 # Check that scatter() copies the data. + # Ensure empty scatters do not break. + ax.scatter([], [], [], c='r', marker='X') + + +@mpl3d_image_comparison(['scatter3d_color.png'], style='mpl20') +def test_scatter3d_color(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + # Check that 'none' color works; these two should overlay to produce the + # same as setting just `color`. + ax.scatter(np.arange(10), np.arange(10), np.arange(10), + facecolor='r', edgecolor='none', marker='o') + ax.scatter(np.arange(10), np.arange(10), np.arange(10), + facecolor='none', edgecolor='r', marker='o') + + ax.scatter(np.arange(10, 20), np.arange(10, 20), np.arange(10, 20), + color='b', marker='s') + + +@mpl3d_image_comparison(['scatter3d_linewidth.png'], style='mpl20') +def test_scatter3d_linewidth(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + # Check that array-like linewidth can be set + ax.scatter(np.arange(10), np.arange(10), np.arange(10), + marker='o', linewidth=np.arange(10)) + + +@check_figures_equal(extensions=['png']) +def test_scatter3d_linewidth_modification(fig_ref, fig_test): + # Changing Path3DCollection linewidths with array-like post-creation + # should work correctly. + ax_test = fig_test.add_subplot(projection='3d') + c = ax_test.scatter(np.arange(10), np.arange(10), np.arange(10), + marker='o') + c.set_linewidths(np.arange(10)) + + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.scatter(np.arange(10), np.arange(10), np.arange(10), marker='o', + linewidths=np.arange(10)) + + +@check_figures_equal(extensions=['png']) +def test_scatter3d_modification(fig_ref, fig_test): + # Changing Path3DCollection properties post-creation should work correctly. + ax_test = fig_test.add_subplot(projection='3d') + c = ax_test.scatter(np.arange(10), np.arange(10), np.arange(10), + marker='o') + c.set_facecolor('C1') + c.set_edgecolor('C2') + c.set_alpha([0.3, 0.7] * 5) + assert c.get_depthshade() + c.set_depthshade(False) + assert not c.get_depthshade() + c.set_sizes(np.full(10, 75)) + c.set_linewidths(3) + + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.scatter(np.arange(10), np.arange(10), np.arange(10), marker='o', + facecolor='C1', edgecolor='C2', alpha=[0.3, 0.7] * 5, + depthshade=False, s=75, linewidths=3) + + +@pytest.mark.parametrize('depthshade', [True, False]) +@check_figures_equal(extensions=['png']) +def test_scatter3d_sorting(fig_ref, fig_test, depthshade): + """Test that marker properties are correctly sorted.""" + + y, x = np.mgrid[:10, :10] + z = np.arange(x.size).reshape(x.shape) + + sizes = np.full(z.shape, 25) + sizes[0::2, 0::2] = 100 + sizes[1::2, 1::2] = 100 + + facecolors = np.full(z.shape, 'C0') + facecolors[:5, :5] = 'C1' + facecolors[6:, :4] = 'C2' + facecolors[6:, 6:] = 'C3' + + edgecolors = np.full(z.shape, 'C4') + edgecolors[1:5, 1:5] = 'C5' + edgecolors[5:9, 1:5] = 'C6' + edgecolors[5:9, 5:9] = 'C7' + + linewidths = np.full(z.shape, 2) + linewidths[0::2, 0::2] = 5 + linewidths[1::2, 1::2] = 5 + + x, y, z, sizes, facecolors, edgecolors, linewidths = ( + a.flatten() + for a in [x, y, z, sizes, facecolors, edgecolors, linewidths] + ) + + ax_ref = fig_ref.add_subplot(projection='3d') + sets = (np.unique(a) for a in [sizes, facecolors, edgecolors, linewidths]) + for s, fc, ec, lw in itertools.product(*sets): + subset = ( + (sizes != s) | + (facecolors != fc) | + (edgecolors != ec) | + (linewidths != lw) + ) + subset = np.ma.masked_array(z, subset, dtype=float) + + # When depth shading is disabled, the colors are passed through as + # single-item lists; this triggers single path optimization. The + # following reshaping is a hack to disable that, since the optimization + # would not occur for the full scatter which has multiple colors. + fc = np.repeat(fc, sum(~subset.mask)) + + ax_ref.scatter(x, y, subset, s=s, fc=fc, ec=ec, lw=lw, alpha=1, + depthshade=depthshade) + + ax_test = fig_test.add_subplot(projection='3d') + ax_test.scatter(x, y, z, s=sizes, fc=facecolors, ec=edgecolors, + lw=linewidths, alpha=1, depthshade=depthshade) + + +@pytest.mark.parametrize('azim', [-50, 130]) # yellow first, blue first +@check_figures_equal(extensions=['png']) +def test_marker_draw_order_data_reversed(fig_test, fig_ref, azim): + """ + Test that the draw order does not depend on the data point order. + + For the given viewing angle at azim=-50, the yellow marker should be in + front. For azim=130, the blue marker should be in front. + """ + x = [-1, 1] + y = [1, -1] + z = [0, 0] + color = ['b', 'y'] + ax = fig_test.add_subplot(projection='3d') + ax.scatter(x, y, z, s=3500, c=color) + ax.view_init(elev=0, azim=azim, roll=0) + ax = fig_ref.add_subplot(projection='3d') + ax.scatter(x[::-1], y[::-1], z[::-1], s=3500, c=color[::-1]) + ax.view_init(elev=0, azim=azim, roll=0) + + +@check_figures_equal(extensions=['png']) +def test_marker_draw_order_view_rotated(fig_test, fig_ref): + """ + Test that the draw order changes with the direction. + + If we rotate *azim* by 180 degrees and exchange the colors, the plot + plot should look the same again. + """ + azim = 130 + x = [-1, 1] + y = [1, -1] + z = [0, 0] + color = ['b', 'y'] + ax = fig_test.add_subplot(projection='3d') + # axis are not exactly invariant under 180 degree rotation -> deactivate + ax.set_axis_off() + ax.scatter(x, y, z, s=3500, c=color) + ax.view_init(elev=0, azim=azim, roll=0) + ax = fig_ref.add_subplot(projection='3d') + ax.set_axis_off() + ax.scatter(x, y, z, s=3500, c=color[::-1]) # color reversed + ax.view_init(elev=0, azim=azim - 180, roll=0) # view rotated by 180 deg + + +@mpl3d_image_comparison(['plot_3d_from_2d.png'], tol=0.019, style='mpl20') +def test_plot_3d_from_2d(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + xs = np.arange(0, 5) + ys = np.arange(5, 10) + ax.plot(xs, ys, zs=0, zdir='x') + ax.plot(xs, ys, zs=0, zdir='y') + + +@mpl3d_image_comparison(['fill_between_quad.png'], style='mpl20') +def test_fill_between_quad(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + theta = np.linspace(0, 2*np.pi, 50) + + x1 = np.cos(theta) + y1 = np.sin(theta) + z1 = 0.1 * np.sin(6 * theta) + + x2 = 0.6 * np.cos(theta) + y2 = 0.6 * np.sin(theta) + z2 = 2 + + where = (theta < np.pi/2) | (theta > 3*np.pi/2) + + # Since none of x1 == x2, y1 == y2, or z1 == z2 is True, the fill_between + # mode will map to 'quad' + ax.fill_between(x1, y1, z1, x2, y2, z2, + where=where, mode='auto', alpha=0.5, edgecolor='k') + + +@mpl3d_image_comparison(['fill_between_polygon.png'], style='mpl20') +def test_fill_between_polygon(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + theta = np.linspace(0, 2*np.pi, 50) + + x1 = x2 = theta + y1 = y2 = 0 + z1 = np.cos(theta) + z2 = z1 + 1 + + where = (theta < np.pi/2) | (theta > 3*np.pi/2) + + # Since x1 == x2 and y1 == y2, the fill_between mode will be 'polygon' + ax.fill_between(x1, y1, z1, x2, y2, z2, + where=where, mode='auto', edgecolor='k') + + +@mpl3d_image_comparison(['surface3d.png'], style='mpl20') +def test_surface3d(): + # Remove this line when this test image is regenerated. + plt.rcParams['pcolormesh.snap'] = False + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X = np.arange(-5, 5, 0.25) + Y = np.arange(-5, 5, 0.25) + X, Y = np.meshgrid(X, Y) + R = np.hypot(X, Y) + Z = np.sin(R) + surf = ax.plot_surface(X, Y, Z, rcount=40, ccount=40, cmap=cm.coolwarm, + lw=0, antialiased=False) + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_zlim(-1.01, 1.01) + fig.colorbar(surf, shrink=0.5, aspect=5) + + +@image_comparison(['surface3d_label_offset_tick_position.png'], style='mpl20') +def test_surface3d_label_offset_tick_position(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax = plt.figure().add_subplot(projection="3d") + + x, y = np.mgrid[0:6 * np.pi:0.25, 0:4 * np.pi:0.25] + z = np.sqrt(np.abs(np.cos(x) + np.cos(y))) + + ax.plot_surface(x * 1e5, y * 1e6, z * 1e8, cmap='autumn', cstride=2, rstride=2) + ax.set_xlabel("X label") + ax.set_ylabel("Y label") + ax.set_zlabel("Z label") + + +@mpl3d_image_comparison(['surface3d_shaded.png'], style='mpl20') +def test_surface3d_shaded(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X = np.arange(-5, 5, 0.25) + Y = np.arange(-5, 5, 0.25) + X, Y = np.meshgrid(X, Y) + R = np.sqrt(X ** 2 + Y ** 2) + Z = np.sin(R) + ax.plot_surface(X, Y, Z, rstride=5, cstride=5, + color=[0.25, 1, 0.25], lw=1, antialiased=False) + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_zlim(-1.01, 1.01) + + +@mpl3d_image_comparison(['surface3d_masked.png'], style='mpl20') +def test_surface3d_masked(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] + y = [1, 2, 3, 4, 5, 6, 7, 8] + + x, y = np.meshgrid(x, y) + matrix = np.array( + [ + [-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [-1, 1, 2, 3, 4, 4, 4, 3, 2, 1, 1], + [-1, -1., 4, 5, 6, 8, 6, 5, 4, 3, -1.], + [-1, -1., 7, 8, 11, 12, 11, 8, 7, -1., -1.], + [-1, -1., 8, 9, 10, 16, 10, 9, 10, 7, -1.], + [-1, -1., -1., 12, 16, 20, 16, 12, 11, -1., -1.], + [-1, -1., -1., -1., 22, 24, 22, 20, 18, -1., -1.], + [-1, -1., -1., -1., -1., 28, 26, 25, -1., -1., -1.], + ] + ) + z = np.ma.masked_less(matrix, 0) + norm = mcolors.Normalize(vmax=z.max(), vmin=z.min()) + colors = mpl.colormaps["plasma"](norm(z)) + ax.plot_surface(x, y, z, facecolors=colors) + ax.view_init(30, -80, 0) + + +@check_figures_equal(extensions=["png"]) +def test_plot_scatter_masks(fig_test, fig_ref): + x = np.linspace(0, 10, 100) + y = np.linspace(0, 10, 100) + z = np.sin(x) * np.cos(y) + mask = z > 0 + + z_masked = np.ma.array(z, mask=mask) + ax_test = fig_test.add_subplot(projection='3d') + ax_test.scatter(x, y, z_masked) + ax_test.plot(x, y, z_masked) + + x[mask] = y[mask] = z[mask] = np.nan + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.scatter(x, y, z) + ax_ref.plot(x, y, z) + + +@check_figures_equal(extensions=["png"]) +def test_plot_surface_None_arg(fig_test, fig_ref): + x, y = np.meshgrid(np.arange(5), np.arange(5)) + z = x + y + ax_test = fig_test.add_subplot(projection='3d') + ax_test.plot_surface(x, y, z, facecolors=None) + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.plot_surface(x, y, z) + + +@mpl3d_image_comparison(['surface3d_masked_strides.png'], style='mpl20') +def test_surface3d_masked_strides(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + x, y = np.mgrid[-6:6.1:1, -6:6.1:1] + z = np.ma.masked_less(x * y, 2) + + ax.plot_surface(x, y, z, rstride=4, cstride=4) + ax.view_init(60, -45, 0) + + +@mpl3d_image_comparison(['text3d.png'], remove_text=False, style='mpl20') +def test_text3d(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + zdirs = (None, 'x', 'y', 'z', (1, 1, 0), (1, 1, 1)) + xs = (2, 6, 4, 9, 7, 2) + ys = (6, 4, 8, 7, 2, 2) + zs = (4, 2, 5, 6, 1, 7) + + for zdir, x, y, z in zip(zdirs, xs, ys, zs): + label = '(%d, %d, %d), dir=%s' % (x, y, z, zdir) + ax.text(x, y, z, label, zdir) + + ax.text(1, 1, 1, "red", color='red') + ax.text2D(0.05, 0.95, "2D Text", transform=ax.transAxes) + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_xlim3d(0, 10) + ax.set_ylim3d(0, 10) + ax.set_zlim3d(0, 10) + ax.set_xlabel('X axis') + ax.set_ylabel('Y axis') + ax.set_zlabel('Z axis') + + +@check_figures_equal(extensions=['png']) +def test_text3d_modification(fig_ref, fig_test): + # Modifying the Text position after the fact should work the same as + # setting it directly. + zdirs = (None, 'x', 'y', 'z', (1, 1, 0), (1, 1, 1)) + xs = (2, 6, 4, 9, 7, 2) + ys = (6, 4, 8, 7, 2, 2) + zs = (4, 2, 5, 6, 1, 7) + + ax_test = fig_test.add_subplot(projection='3d') + ax_test.set_xlim3d(0, 10) + ax_test.set_ylim3d(0, 10) + ax_test.set_zlim3d(0, 10) + for zdir, x, y, z in zip(zdirs, xs, ys, zs): + t = ax_test.text(0, 0, 0, f'({x}, {y}, {z}), dir={zdir}') + t.set_position_3d((x, y, z), zdir=zdir) + + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.set_xlim3d(0, 10) + ax_ref.set_ylim3d(0, 10) + ax_ref.set_zlim3d(0, 10) + for zdir, x, y, z in zip(zdirs, xs, ys, zs): + ax_ref.text(x, y, z, f'({x}, {y}, {z}), dir={zdir}', zdir=zdir) + + +@mpl3d_image_comparison(['trisurf3d.png'], tol=0.061, style='mpl20') +def test_trisurf3d(): + n_angles = 36 + n_radii = 8 + radii = np.linspace(0.125, 1.0, n_radii) + angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False) + angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) + angles[:, 1::2] += np.pi/n_angles + + x = np.append(0, (radii*np.cos(angles)).flatten()) + y = np.append(0, (radii*np.sin(angles)).flatten()) + z = np.sin(-x*y) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.plot_trisurf(x, y, z, cmap=cm.jet, linewidth=0.2) + + +@mpl3d_image_comparison(['trisurf3d_shaded.png'], tol=0.03, style='mpl20') +def test_trisurf3d_shaded(): + n_angles = 36 + n_radii = 8 + radii = np.linspace(0.125, 1.0, n_radii) + angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False) + angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) + angles[:, 1::2] += np.pi/n_angles + + x = np.append(0, (radii*np.cos(angles)).flatten()) + y = np.append(0, (radii*np.sin(angles)).flatten()) + z = np.sin(-x*y) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.plot_trisurf(x, y, z, color=[1, 0.5, 0], linewidth=0.2) + + +@mpl3d_image_comparison(['wireframe3d.png'], style='mpl20') +def test_wireframe3d(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.plot_wireframe(X, Y, Z, rcount=13, ccount=13) + + +@mpl3d_image_comparison(['wireframe3dzerocstride.png'], style='mpl20') +def test_wireframe3dzerocstride(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.plot_wireframe(X, Y, Z, rcount=13, ccount=0) + + +@mpl3d_image_comparison(['wireframe3dzerorstride.png'], style='mpl20') +def test_wireframe3dzerorstride(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + ax.plot_wireframe(X, Y, Z, rstride=0, cstride=10) + + +def test_wireframe3dzerostrideraises(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + with pytest.raises(ValueError): + ax.plot_wireframe(X, Y, Z, rstride=0, cstride=0) + + +def test_mixedsamplesraises(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + X, Y, Z = axes3d.get_test_data(0.05) + with pytest.raises(ValueError): + ax.plot_wireframe(X, Y, Z, rstride=10, ccount=50) + with pytest.raises(ValueError): + ax.plot_surface(X, Y, Z, cstride=50, rcount=10) + + +# remove tolerance when regenerating the test image +@mpl3d_image_comparison(['quiver3d.png'], style='mpl20', tol=0.003) +def test_quiver3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + pivots = ['tip', 'middle', 'tail'] + colors = ['tab:blue', 'tab:orange', 'tab:green'] + for i, (pivot, color) in enumerate(zip(pivots, colors)): + x, y, z = np.meshgrid([-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]) + u = -x + v = -y + w = -z + # Offset each set in z direction + z += 2 * i + ax.quiver(x, y, z, u, v, w, length=1, pivot=pivot, color=color) + ax.scatter(x, y, z, color=color) + + ax.set_xlim(-3, 3) + ax.set_ylim(-3, 3) + ax.set_zlim(-1, 5) + + +@check_figures_equal(extensions=["png"]) +def test_quiver3d_empty(fig_test, fig_ref): + fig_ref.add_subplot(projection='3d') + x = y = z = u = v = w = [] + ax = fig_test.add_subplot(projection='3d') + ax.quiver(x, y, z, u, v, w, length=0.1, pivot='tip', normalize=True) + + +@mpl3d_image_comparison(['quiver3d_masked.png'], style='mpl20') +def test_quiver3d_masked(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + # Using mgrid here instead of ogrid because masked_where doesn't + # seem to like broadcasting very much... + x, y, z = np.mgrid[-1:0.8:10j, -1:0.8:10j, -1:0.6:3j] + + u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z) + v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z) + w = (2/3)**0.5 * np.cos(np.pi * x) * np.cos(np.pi * y) * np.sin(np.pi * z) + u = np.ma.masked_where((-0.4 < x) & (x < 0.1), u, copy=False) + v = np.ma.masked_where((0.1 < y) & (y < 0.7), v, copy=False) + + ax.quiver(x, y, z, u, v, w, length=0.1, pivot='tip', normalize=True) + + +@mpl3d_image_comparison(['quiver3d_colorcoded.png'], style='mpl20') +def test_quiver3d_colorcoded(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + x = y = dx = dz = np.zeros(10) + z = dy = np.arange(10.) + + color = plt.cm.Reds(dy/dy.max()) + ax.quiver(x, y, z, dx, dy, dz, colors=color) + ax.set_ylim(0, 10) + + +def test_patch_modification(): + fig = plt.figure() + ax = fig.add_subplot(projection="3d") + circle = Circle((0, 0)) + ax.add_patch(circle) + art3d.patch_2d_to_3d(circle) + circle.set_facecolor((1.0, 0.0, 0.0, 1)) + + assert mcolors.same_color(circle.get_facecolor(), (1, 0, 0, 1)) + fig.canvas.draw() + assert mcolors.same_color(circle.get_facecolor(), (1, 0, 0, 1)) + + +@check_figures_equal(extensions=['png']) +def test_patch_collection_modification(fig_test, fig_ref): + # Test that modifying Patch3DCollection properties after creation works. + patch1 = Circle((0, 0), 0.05) + patch2 = Circle((0.1, 0.1), 0.03) + facecolors = np.array([[0., 0.5, 0., 1.], [0.5, 0., 0., 0.5]]) + c = art3d.Patch3DCollection([patch1, patch2], linewidths=3) + + ax_test = fig_test.add_subplot(projection='3d') + ax_test.add_collection3d(c) + c.set_edgecolor('C2') + c.set_facecolor(facecolors) + c.set_alpha(0.7) + assert c.get_depthshade() + c.set_depthshade(False) + assert not c.get_depthshade() + + patch1 = Circle((0, 0), 0.05) + patch2 = Circle((0.1, 0.1), 0.03) + facecolors = np.array([[0., 0.5, 0., 1.], [0.5, 0., 0., 0.5]]) + c = art3d.Patch3DCollection([patch1, patch2], linewidths=3, + edgecolor='C2', facecolor=facecolors, + alpha=0.7, depthshade=False) + + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.add_collection3d(c) + + +def test_poly3dcollection_verts_validation(): + poly = [[0, 0, 1], [0, 1, 1], [0, 1, 0], [0, 0, 0]] + with pytest.raises(ValueError, match=r'list of \(N, 3\) array-like'): + art3d.Poly3DCollection(poly) # should be Poly3DCollection([poly]) + + poly = np.array(poly, dtype=float) + with pytest.raises(ValueError, match=r'list of \(N, 3\) array-like'): + art3d.Poly3DCollection(poly) # should be Poly3DCollection([poly]) + + +@mpl3d_image_comparison(['poly3dcollection_closed.png'], style='mpl20') +def test_poly3dcollection_closed(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + poly1 = np.array([[0, 0, 1], [0, 1, 1], [0, 0, 0]], float) + poly2 = np.array([[0, 1, 1], [1, 1, 1], [1, 1, 0]], float) + c1 = art3d.Poly3DCollection([poly1], linewidths=3, edgecolor='k', + facecolor=(0.5, 0.5, 1, 0.5), closed=True) + c2 = art3d.Poly3DCollection([poly2], linewidths=3, edgecolor='k', + facecolor=(1, 0.5, 0.5, 0.5), closed=False) + ax.add_collection3d(c1, autolim=False) + ax.add_collection3d(c2, autolim=False) + + +def test_poly_collection_2d_to_3d_empty(): + poly = PolyCollection([]) + art3d.poly_collection_2d_to_3d(poly) + assert isinstance(poly, art3d.Poly3DCollection) + assert poly.get_paths() == [] + + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + ax.add_artist(poly) + minz = poly.do_3d_projection() + assert np.isnan(minz) + + # Ensure drawing actually works. + fig.canvas.draw() + + +@mpl3d_image_comparison(['poly3dcollection_alpha.png'], style='mpl20') +def test_poly3dcollection_alpha(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + poly1 = np.array([[0, 0, 1], [0, 1, 1], [0, 0, 0]], float) + poly2 = np.array([[0, 1, 1], [1, 1, 1], [1, 1, 0]], float) + c1 = art3d.Poly3DCollection([poly1], linewidths=3, edgecolor='k', + facecolor=(0.5, 0.5, 1), closed=True) + c1.set_alpha(0.5) + c2 = art3d.Poly3DCollection([poly2], linewidths=3, closed=False) + # Post-creation modification should work. + c2.set_facecolor((1, 0.5, 0.5)) + c2.set_edgecolor('k') + c2.set_alpha(0.5) + ax.add_collection3d(c1, autolim=False) + ax.add_collection3d(c2, autolim=False) + + +@mpl3d_image_comparison(['add_collection3d_zs_array.png'], style='mpl20') +def test_add_collection3d_zs_array(): + theta = np.linspace(-4 * np.pi, 4 * np.pi, 100) + z = np.linspace(-2, 2, 100) + r = z**2 + 1 + x = r * np.sin(theta) + y = r * np.cos(theta) + + points = np.column_stack([x, y, z]).reshape(-1, 1, 3) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + norm = plt.Normalize(0, 2*np.pi) + # 2D LineCollection from x & y values + lc = LineCollection(segments[:, :, :2], cmap='twilight', norm=norm) + lc.set_array(np.mod(theta, 2*np.pi)) + # Add 2D collection at z values to ax + line = ax.add_collection3d(lc, zs=segments[:, :, 2]) + + assert line is not None + + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_xlim(-5, 5) + ax.set_ylim(-4, 6) + ax.set_zlim(-2, 2) + + +@mpl3d_image_comparison(['add_collection3d_zs_scalar.png'], style='mpl20') +def test_add_collection3d_zs_scalar(): + theta = np.linspace(0, 2 * np.pi, 100) + z = 1 + r = z**2 + 1 + x = r * np.sin(theta) + y = r * np.cos(theta) + + points = np.column_stack([x, y]).reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + norm = plt.Normalize(0, 2*np.pi) + lc = LineCollection(segments, cmap='twilight', norm=norm) + lc.set_array(theta) + line = ax.add_collection3d(lc, zs=z) + + assert line is not None + + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_xlim(-5, 5) + ax.set_ylim(-4, 6) + ax.set_zlim(0, 2) + + +def test_line3dCollection_autoscaling(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + lines = [[(0, 0, 0), (1, 4, 2)], + [(1, 1, 3), (2, 0, 2)], + [(1, 0, 4), (1, 4, 5)]] + + lc = art3d.Line3DCollection(lines) + ax.add_collection3d(lc) + assert np.allclose(ax.get_xlim3d(), (-0.041666666666666664, 2.0416666666666665)) + assert np.allclose(ax.get_ylim3d(), (-0.08333333333333333, 4.083333333333333)) + assert np.allclose(ax.get_zlim3d(), (-0.10416666666666666, 5.104166666666667)) + + +def test_poly3dCollection_autoscaling(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + poly = np.array([[0, 0, 0], [1, 1, 3], [1, 0, 4]]) + col = art3d.Poly3DCollection([poly]) + ax.add_collection3d(col) + assert np.allclose(ax.get_xlim3d(), (-0.020833333333333332, 1.0208333333333333)) + assert np.allclose(ax.get_ylim3d(), (-0.020833333333333332, 1.0208333333333333)) + assert np.allclose(ax.get_zlim3d(), (-0.0833333333333333, 4.083333333333333)) + + +@mpl3d_image_comparison(['axes3d_labelpad.png'], + remove_text=False, style='mpl20') +def test_axes3d_labelpad(): + fig = plt.figure() + ax = fig.add_axes(Axes3D(fig)) + # labelpad respects rcParams + assert ax.xaxis.labelpad == mpl.rcParams['axes.labelpad'] + # labelpad can be set in set_label + ax.set_xlabel('X LABEL', labelpad=10) + assert ax.xaxis.labelpad == 10 + ax.set_ylabel('Y LABEL') + ax.set_zlabel('Z LABEL', labelpad=20) + assert ax.zaxis.labelpad == 20 + assert ax.get_zlabel() == 'Z LABEL' + # or manually + ax.yaxis.labelpad = 20 + ax.zaxis.labelpad = -40 + + # Tick labels also respect tick.pad (also from rcParams) + for i, tick in enumerate(ax.yaxis.get_major_ticks()): + tick.set_pad(tick.get_pad() + 5 - i * 5) + + +@mpl3d_image_comparison(['axes3d_cla.png'], remove_text=False, style='mpl20') +def test_axes3d_cla(): + # fixed in pull request 4553 + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection='3d') + ax.set_axis_off() + ax.cla() # make sure the axis displayed is 3D (not 2D) + + +@mpl3d_image_comparison(['axes3d_rotated.png'], + remove_text=False, style='mpl20') +def test_axes3d_rotated(): + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection='3d') + ax.view_init(90, 45, 0) # look down, rotated. Should be square + + +def test_plotsurface_1d_raises(): + x = np.linspace(0.5, 10, num=100) + y = np.linspace(0.5, 10, num=100) + X, Y = np.meshgrid(x, y) + z = np.random.randn(100) + + fig = plt.figure(figsize=(14, 6)) + ax = fig.add_subplot(1, 2, 1, projection='3d') + with pytest.raises(ValueError): + ax.plot_surface(X, Y, z) + + +def _test_proj_make_M(): + # eye point + E = np.array([1000, -1000, 2000]) + R = np.array([100, 100, 100]) + V = np.array([0, 0, 1]) + roll = 0 + u, v, w = proj3d._view_axes(E, R, V, roll) + viewM = proj3d._view_transformation_uvw(u, v, w, E) + perspM = proj3d._persp_transformation(100, -100, 1) + M = np.dot(perspM, viewM) + return M + + +def test_proj_transform(): + M = _test_proj_make_M() + invM = np.linalg.inv(M) + + xs = np.array([0, 1, 1, 0, 0, 0, 1, 1, 0, 0]) * 300.0 + ys = np.array([0, 0, 1, 1, 0, 0, 0, 1, 1, 0]) * 300.0 + zs = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) * 300.0 + + txs, tys, tzs = proj3d.proj_transform(xs, ys, zs, M) + ixs, iys, izs = proj3d.inv_transform(txs, tys, tzs, invM) + + np.testing.assert_almost_equal(ixs, xs) + np.testing.assert_almost_equal(iys, ys) + np.testing.assert_almost_equal(izs, zs) + + +def _test_proj_draw_axes(M, s=1, *args, **kwargs): + xs = [0, s, 0, 0] + ys = [0, 0, s, 0] + zs = [0, 0, 0, s] + txs, tys, tzs = proj3d.proj_transform(xs, ys, zs, M) + o, ax, ay, az = zip(txs, tys) + lines = [(o, ax), (o, ay), (o, az)] + + fig, ax = plt.subplots(*args, **kwargs) + linec = LineCollection(lines) + ax.add_collection(linec) + for x, y, t in zip(txs, tys, ['o', 'x', 'y', 'z']): + ax.text(x, y, t) + + return fig, ax + + +@mpl3d_image_comparison(['proj3d_axes_cube.png'], style='mpl20') +def test_proj_axes_cube(): + M = _test_proj_make_M() + + ts = '0 1 2 3 0 4 5 6 7 4'.split() + xs = np.array([0, 1, 1, 0, 0, 0, 1, 1, 0, 0]) * 300.0 + ys = np.array([0, 0, 1, 1, 0, 0, 0, 1, 1, 0]) * 300.0 + zs = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) * 300.0 + + txs, tys, tzs = proj3d.proj_transform(xs, ys, zs, M) + + fig, ax = _test_proj_draw_axes(M, s=400) + + ax.scatter(txs, tys, c=tzs) + ax.plot(txs, tys, c='r') + for x, y, t in zip(txs, tys, ts): + ax.text(x, y, t) + + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_xlim(-0.2, 0.2) + ax.set_ylim(-0.2, 0.2) + + +@mpl3d_image_comparison(['proj3d_axes_cube_ortho.png'], style='mpl20') +def test_proj_axes_cube_ortho(): + E = np.array([200, 100, 100]) + R = np.array([0, 0, 0]) + V = np.array([0, 0, 1]) + roll = 0 + u, v, w = proj3d._view_axes(E, R, V, roll) + viewM = proj3d._view_transformation_uvw(u, v, w, E) + orthoM = proj3d._ortho_transformation(-1, 1) + M = np.dot(orthoM, viewM) + + ts = '0 1 2 3 0 4 5 6 7 4'.split() + xs = np.array([0, 1, 1, 0, 0, 0, 1, 1, 0, 0]) * 100 + ys = np.array([0, 0, 1, 1, 0, 0, 0, 1, 1, 0]) * 100 + zs = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) * 100 + + txs, tys, tzs = proj3d.proj_transform(xs, ys, zs, M) + + fig, ax = _test_proj_draw_axes(M, s=150) + + ax.scatter(txs, tys, s=300-tzs) + ax.plot(txs, tys, c='r') + for x, y, t in zip(txs, tys, ts): + ax.text(x, y, t) + + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + ax.set_xlim(-200, 200) + ax.set_ylim(-200, 200) + + +def test_world(): + xmin, xmax = 100, 120 + ymin, ymax = -100, 100 + zmin, zmax = 0.1, 0.2 + M = proj3d.world_transformation(xmin, xmax, ymin, ymax, zmin, zmax) + np.testing.assert_allclose(M, + [[5e-2, 0, 0, -5], + [0, 5e-3, 0, 5e-1], + [0, 0, 1e1, -1], + [0, 0, 0, 1]]) + + +def test_autoscale(): + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + assert ax.get_zscale() == 'linear' + ax._view_margin = 0 + ax.margins(x=0, y=.1, z=.2) + ax.plot([0, 1], [0, 1], [0, 1]) + assert ax.get_w_lims() == (0, 1, -.1, 1.1, -.2, 1.2) + ax.autoscale(False) + ax.set_autoscalez_on(True) + ax.plot([0, 2], [0, 2], [0, 2]) + assert ax.get_w_lims() == (0, 1, -.1, 1.1, -.4, 2.4) + ax.autoscale(axis='x') + ax.plot([0, 2], [0, 2], [0, 2]) + assert ax.get_w_lims() == (0, 2, -.1, 1.1, -.4, 2.4) + + +@pytest.mark.parametrize('axis', ('x', 'y', 'z')) +@pytest.mark.parametrize('auto', (True, False, None)) +def test_unautoscale(axis, auto): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + x = np.arange(100) + y = np.linspace(-0.1, 0.1, 100) + ax.scatter(x, y) + + get_autoscale_on = getattr(ax, f'get_autoscale{axis}_on') + set_lim = getattr(ax, f'set_{axis}lim') + get_lim = getattr(ax, f'get_{axis}lim') + + post_auto = get_autoscale_on() if auto is None else auto + + set_lim((-0.5, 0.5), auto=auto) + assert post_auto == get_autoscale_on() + fig.canvas.draw() + np.testing.assert_array_equal(get_lim(), (-0.5, 0.5)) + + +@check_figures_equal(extensions=["png"]) +def test_culling(fig_test, fig_ref): + xmins = (-100, -50) + for fig, xmin in zip((fig_test, fig_ref), xmins): + ax = fig.add_subplot(projection='3d') + n = abs(xmin) + 1 + xs = np.linspace(0, xmin, n) + ys = np.ones(n) + zs = np.zeros(n) + ax.plot(xs, ys, zs, 'k') + + ax.set(xlim=(-5, 5), ylim=(-5, 5), zlim=(-5, 5)) + ax.view_init(5, 180, 0) + + +def test_axes3d_focal_length_checks(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + with pytest.raises(ValueError): + ax.set_proj_type('persp', focal_length=0) + with pytest.raises(ValueError): + ax.set_proj_type('ortho', focal_length=1) + + +@mpl3d_image_comparison(['axes3d_focal_length.png'], + remove_text=False, style='mpl20') +def test_axes3d_focal_length(): + fig, axs = plt.subplots(1, 2, subplot_kw={'projection': '3d'}) + axs[0].set_proj_type('persp', focal_length=np.inf) + axs[1].set_proj_type('persp', focal_length=0.15) + + +@mpl3d_image_comparison(['axes3d_ortho.png'], remove_text=False, style='mpl20') +def test_axes3d_ortho(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.set_proj_type('ortho') + + +@mpl3d_image_comparison(['axes3d_isometric.png'], style='mpl20') +def test_axes3d_isometric(): + from itertools import combinations, product + fig, ax = plt.subplots(subplot_kw=dict( + projection='3d', + proj_type='ortho', + box_aspect=(4, 4, 4) + )) + r = (-1, 1) # stackoverflow.com/a/11156353 + for s, e in combinations(np.array(list(product(r, r, r))), 2): + if abs(s - e).sum() == r[1] - r[0]: + ax.plot3D(*zip(s, e), c='k') + ax.view_init(elev=np.degrees(np.arctan(1. / np.sqrt(2))), azim=-45, roll=0) + ax.grid(True) + + +@check_figures_equal(extensions=["png"]) +def test_axlim_clip(fig_test, fig_ref): + # With axlim clipping + ax = fig_test.add_subplot(projection="3d") + x = np.linspace(0, 1, 11) + y = np.linspace(0, 1, 11) + X, Y = np.meshgrid(x, y) + Z = X + Y + ax.plot_surface(X, Y, Z, facecolor='C1', edgecolors=None, + rcount=50, ccount=50, axlim_clip=True) + # This ax.plot is to cover the extra surface edge which is not clipped out + ax.plot([0.5, 0.5], [0, 1], [0.5, 1.5], + color='k', linewidth=3, zorder=5, axlim_clip=True) + ax.scatter(X.ravel(), Y.ravel(), Z.ravel() + 1, axlim_clip=True) + ax.quiver(X.ravel(), Y.ravel(), Z.ravel() + 2, + 0*X.ravel(), 0*Y.ravel(), 0*Z.ravel() + 1, + arrow_length_ratio=0, axlim_clip=True) + ax.plot(X[0], Y[0], Z[0] + 3, color='C2', axlim_clip=True) + ax.text(1.1, 0.5, 4, 'test', axlim_clip=True) # won't be visible + ax.set(xlim=(0, 0.5), ylim=(0, 1), zlim=(0, 5)) + + # With manual clipping + ax = fig_ref.add_subplot(projection="3d") + idx = (X <= 0.5) + X = X[idx].reshape(11, 6) + Y = Y[idx].reshape(11, 6) + Z = Z[idx].reshape(11, 6) + ax.plot_surface(X, Y, Z, facecolor='C1', edgecolors=None, + rcount=50, ccount=50, axlim_clip=False) + ax.plot([0.5, 0.5], [0, 1], [0.5, 1.5], + color='k', linewidth=3, zorder=5, axlim_clip=False) + ax.scatter(X.ravel(), Y.ravel(), Z.ravel() + 1, axlim_clip=False) + ax.quiver(X.ravel(), Y.ravel(), Z.ravel() + 2, + 0*X.ravel(), 0*Y.ravel(), 0*Z.ravel() + 1, + arrow_length_ratio=0, axlim_clip=False) + ax.plot(X[0], Y[0], Z[0] + 3, color='C2', axlim_clip=False) + ax.set(xlim=(0, 0.5), ylim=(0, 1), zlim=(0, 5)) + + +@pytest.mark.parametrize('value', [np.inf, np.nan]) +@pytest.mark.parametrize(('setter', 'side'), [ + ('set_xlim3d', 'left'), + ('set_xlim3d', 'right'), + ('set_ylim3d', 'bottom'), + ('set_ylim3d', 'top'), + ('set_zlim3d', 'bottom'), + ('set_zlim3d', 'top'), +]) +def test_invalid_axes_limits(setter, side, value): + limit = {side: value} + fig = plt.figure() + obj = fig.add_subplot(projection='3d') + with pytest.raises(ValueError): + getattr(obj, setter)(**limit) + + +class TestVoxels: + @mpl3d_image_comparison(['voxels-simple.png'], style='mpl20') + def test_simple(self): + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + x, y, z = np.indices((5, 4, 3)) + voxels = (x == y) | (y == z) + ax.voxels(voxels) + + @mpl3d_image_comparison(['voxels-edge-style.png'], style='mpl20') + def test_edge_style(self): + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + x, y, z = np.indices((5, 5, 4)) + voxels = ((x - 2)**2 + (y - 2)**2 + (z-1.5)**2) < 2.2**2 + v = ax.voxels(voxels, linewidths=3, edgecolor='C1') + + # change the edge color of one voxel + v[max(v.keys())].set_edgecolor('C2') + + @mpl3d_image_comparison(['voxels-named-colors.png'], style='mpl20') + def test_named_colors(self): + """Test with colors set to a 3D object array of strings.""" + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + x, y, z = np.indices((10, 10, 10)) + voxels = (x == y) | (y == z) + voxels = voxels & ~(x * y * z < 1) + colors = np.full((10, 10, 10), 'C0', dtype=np.object_) + colors[(x < 5) & (y < 5)] = '0.25' + colors[(x + z) < 10] = 'cyan' + ax.voxels(voxels, facecolors=colors) + + @mpl3d_image_comparison(['voxels-rgb-data.png'], style='mpl20') + def test_rgb_data(self): + """Test with colors set to a 4d float array of rgb data.""" + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + x, y, z = np.indices((10, 10, 10)) + voxels = (x == y) | (y == z) + colors = np.zeros((10, 10, 10, 3)) + colors[..., 0] = x / 9 + colors[..., 1] = y / 9 + colors[..., 2] = z / 9 + ax.voxels(voxels, facecolors=colors) + + @mpl3d_image_comparison(['voxels-alpha.png'], style='mpl20') + def test_alpha(self): + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + x, y, z = np.indices((10, 10, 10)) + v1 = x == y + v2 = np.abs(x - y) < 2 + voxels = v1 | v2 + colors = np.zeros((10, 10, 10, 4)) + colors[v2] = [1, 0, 0, 0.5] + colors[v1] = [0, 1, 0, 0.5] + v = ax.voxels(voxels, facecolors=colors) + + assert type(v) is dict + for coord, poly in v.items(): + assert voxels[coord], "faces returned for absent voxel" + assert isinstance(poly, art3d.Poly3DCollection) + + @mpl3d_image_comparison(['voxels-xyz.png'], + tol=0.01, remove_text=False, style='mpl20') + def test_xyz(self): + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + def midpoints(x): + sl = () + for i in range(x.ndim): + x = (x[sl + np.index_exp[:-1]] + + x[sl + np.index_exp[1:]]) / 2.0 + sl += np.index_exp[:] + return x + + # prepare some coordinates, and attach rgb values to each + r, g, b = np.indices((17, 17, 17)) / 16.0 + rc = midpoints(r) + gc = midpoints(g) + bc = midpoints(b) + + # define a sphere about [0.5, 0.5, 0.5] + sphere = (rc - 0.5)**2 + (gc - 0.5)**2 + (bc - 0.5)**2 < 0.5**2 + + # combine the color components + colors = np.zeros(sphere.shape + (3,)) + colors[..., 0] = rc + colors[..., 1] = gc + colors[..., 2] = bc + + # and plot everything + ax.voxels(r, g, b, sphere, + facecolors=colors, + edgecolors=np.clip(2*colors - 0.5, 0, 1), # brighter + linewidth=0.5) + + def test_calling_conventions(self): + x, y, z = np.indices((3, 4, 5)) + filled = np.ones((2, 3, 4)) + + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + + # all the valid calling conventions + for kw in (dict(), dict(edgecolor='k')): + ax.voxels(filled, **kw) + ax.voxels(filled=filled, **kw) + ax.voxels(x, y, z, filled, **kw) + ax.voxels(x, y, z, filled=filled, **kw) + + # duplicate argument + with pytest.raises(TypeError, match='voxels'): + ax.voxels(x, y, z, filled, filled=filled) + # missing arguments + with pytest.raises(TypeError, match='voxels'): + ax.voxels(x, y) + # x, y, z are positional only - this passes them on as attributes of + # Poly3DCollection + with pytest.raises(AttributeError, match="keyword argument 'x'") as exec_info: + ax.voxels(filled=filled, x=x, y=y, z=z) + assert exec_info.value.name == 'x' + + +def test_line3d_set_get_data_3d(): + x, y, z = [0, 1], [2, 3], [4, 5] + x2, y2, z2 = [6, 7], [8, 9], [10, 11] + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + lines = ax.plot(x, y, z) + line = lines[0] + np.testing.assert_array_equal((x, y, z), line.get_data_3d()) + line.set_data_3d(x2, y2, z2) + np.testing.assert_array_equal((x2, y2, z2), line.get_data_3d()) + line.set_xdata(x) + line.set_ydata(y) + line.set_3d_properties(zs=z, zdir='z') + np.testing.assert_array_equal((x, y, z), line.get_data_3d()) + line.set_3d_properties(zs=0, zdir='z') + np.testing.assert_array_equal((x, y, np.zeros_like(z)), line.get_data_3d()) + + +@check_figures_equal(extensions=["png"]) +def test_inverted(fig_test, fig_ref): + # Plot then invert. + ax = fig_test.add_subplot(projection="3d") + ax.plot([1, 1, 10, 10], [1, 10, 10, 10], [1, 1, 1, 10]) + ax.invert_yaxis() + # Invert then plot. + ax = fig_ref.add_subplot(projection="3d") + ax.invert_yaxis() + ax.plot([1, 1, 10, 10], [1, 10, 10, 10], [1, 1, 1, 10]) + + +def test_inverted_cla(): + # GitHub PR #5450. Setting autoscale should reset + # axes to be non-inverted. + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + # 1. test that a new axis is not inverted per default + assert not ax.xaxis_inverted() + assert not ax.yaxis_inverted() + assert not ax.zaxis_inverted() + ax.set_xlim(1, 0) + ax.set_ylim(1, 0) + ax.set_zlim(1, 0) + assert ax.xaxis_inverted() + assert ax.yaxis_inverted() + assert ax.zaxis_inverted() + ax.cla() + assert not ax.xaxis_inverted() + assert not ax.yaxis_inverted() + assert not ax.zaxis_inverted() + + +def test_ax3d_tickcolour(): + fig = plt.figure() + ax = Axes3D(fig) + + ax.tick_params(axis='x', colors='red') + ax.tick_params(axis='y', colors='red') + ax.tick_params(axis='z', colors='red') + fig.canvas.draw() + + for tick in ax.xaxis.get_major_ticks(): + assert tick.tick1line._color == 'red' + for tick in ax.yaxis.get_major_ticks(): + assert tick.tick1line._color == 'red' + for tick in ax.zaxis.get_major_ticks(): + assert tick.tick1line._color == 'red' + + +@check_figures_equal(extensions=["png"]) +def test_ticklabel_format(fig_test, fig_ref): + axs = fig_test.subplots(4, 5, subplot_kw={"projection": "3d"}) + for ax in axs.flat: + ax.set_xlim(1e7, 1e7 + 10) + for row, name in zip(axs, ["x", "y", "z", "both"]): + row[0].ticklabel_format( + axis=name, style="plain") + row[1].ticklabel_format( + axis=name, scilimits=(-2, 2)) + row[2].ticklabel_format( + axis=name, useOffset=not mpl.rcParams["axes.formatter.useoffset"]) + row[3].ticklabel_format( + axis=name, useLocale=not mpl.rcParams["axes.formatter.use_locale"]) + row[4].ticklabel_format( + axis=name, + useMathText=not mpl.rcParams["axes.formatter.use_mathtext"]) + + def get_formatters(ax, names): + return [getattr(ax, name).get_major_formatter() for name in names] + + axs = fig_ref.subplots(4, 5, subplot_kw={"projection": "3d"}) + for ax in axs.flat: + ax.set_xlim(1e7, 1e7 + 10) + for row, names in zip( + axs, [["xaxis"], ["yaxis"], ["zaxis"], ["xaxis", "yaxis", "zaxis"]] + ): + for fmt in get_formatters(row[0], names): + fmt.set_scientific(False) + for fmt in get_formatters(row[1], names): + fmt.set_powerlimits((-2, 2)) + for fmt in get_formatters(row[2], names): + fmt.set_useOffset(not mpl.rcParams["axes.formatter.useoffset"]) + for fmt in get_formatters(row[3], names): + fmt.set_useLocale(not mpl.rcParams["axes.formatter.use_locale"]) + for fmt in get_formatters(row[4], names): + fmt.set_useMathText( + not mpl.rcParams["axes.formatter.use_mathtext"]) + + +@check_figures_equal(extensions=["png"]) +def test_quiver3D_smoke(fig_test, fig_ref): + pivot = "middle" + # Make the grid + x, y, z = np.meshgrid( + np.arange(-0.8, 1, 0.2), + np.arange(-0.8, 1, 0.2), + np.arange(-0.8, 1, 0.8) + ) + u = v = w = np.ones_like(x) + + for fig, length in zip((fig_ref, fig_test), (1, 1.0)): + ax = fig.add_subplot(projection="3d") + ax.quiver(x, y, z, u, v, w, length=length, pivot=pivot) + + +@image_comparison(["minor_ticks.png"], style="mpl20") +def test_minor_ticks(): + ax = plt.figure().add_subplot(projection="3d") + ax.set_xticks([0.25], minor=True) + ax.set_xticklabels(["quarter"], minor=True) + ax.set_yticks([0.33], minor=True) + ax.set_yticklabels(["third"], minor=True) + ax.set_zticks([0.50], minor=True) + ax.set_zticklabels(["half"], minor=True) + + +# remove tolerance when regenerating the test image +@mpl3d_image_comparison(['errorbar3d_errorevery.png'], style='mpl20', tol=0.003) +def test_errorbar3d_errorevery(): + """Tests errorevery functionality for 3D errorbars.""" + t = np.arange(0, 2*np.pi+.1, 0.01) + x, y, z = np.sin(t), np.cos(3*t), np.sin(5*t) + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + estep = 15 + i = np.arange(t.size) + zuplims = (i % estep == 0) & (i // estep % 3 == 0) + zlolims = (i % estep == 0) & (i // estep % 3 == 2) + + ax.errorbar(x, y, z, 0.2, zuplims=zuplims, zlolims=zlolims, + errorevery=estep) + + +@mpl3d_image_comparison(['errorbar3d.png'], style='mpl20', + tol=0 if platform.machine() == 'x86_64' else 0.02) +def test_errorbar3d(): + """Tests limits, color styling, and legend for 3D errorbars.""" + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + d = [1, 2, 3, 4, 5] + e = [.5, .5, .5, .5, .5] + ax.errorbar(x=d, y=d, z=d, xerr=e, yerr=e, zerr=e, capsize=3, + zuplims=[False, True, False, True, True], + zlolims=[True, False, False, True, False], + yuplims=True, + ecolor='purple', label='Error lines') + ax.legend() + + +@image_comparison(['stem3d.png'], style='mpl20', tol=0.009) +def test_stem3d(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig, axs = plt.subplots(2, 3, figsize=(8, 6), + constrained_layout=True, + subplot_kw={'projection': '3d'}) + + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta - np.pi/2) + y = np.sin(theta - np.pi/2) + z = theta + + for ax, zdir in zip(axs[0], ['x', 'y', 'z']): + ax.stem(x, y, z, orientation=zdir) + ax.set_title(f'orientation={zdir}') + + x = np.linspace(-np.pi/2, np.pi/2, 20) + y = np.ones_like(x) + z = np.cos(x) + + for ax, zdir in zip(axs[1], ['x', 'y', 'z']): + markerline, stemlines, baseline = ax.stem( + x, y, z, + linefmt='C4-.', markerfmt='C1D', basefmt='C2', + orientation=zdir) + ax.set_title(f'orientation={zdir}') + markerline.set(markerfacecolor='none', markeredgewidth=2) + baseline.set_linewidth(3) + + +@image_comparison(["equal_box_aspect.png"], style="mpl20") +def test_equal_box_aspect(): + from itertools import product, combinations + + fig = plt.figure() + ax = fig.add_subplot(projection="3d") + + # Make data + u = np.linspace(0, 2 * np.pi, 100) + v = np.linspace(0, np.pi, 100) + x = np.outer(np.cos(u), np.sin(v)) + y = np.outer(np.sin(u), np.sin(v)) + z = np.outer(np.ones_like(u), np.cos(v)) + + # Plot the surface + ax.plot_surface(x, y, z) + + # draw cube + r = [-1, 1] + for s, e in combinations(np.array(list(product(r, r, r))), 2): + if np.sum(np.abs(s - e)) == r[1] - r[0]: + ax.plot3D(*zip(s, e), color="b") + + # Make axes limits + xyzlim = np.column_stack( + [ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()] + ) + XYZlim = [min(xyzlim[0]), max(xyzlim[1])] + ax.set_xlim3d(XYZlim) + ax.set_ylim3d(XYZlim) + ax.set_zlim3d(XYZlim) + ax.axis('off') + ax.set_box_aspect((1, 1, 1)) + + with pytest.raises(ValueError, match="Argument zoom ="): + ax.set_box_aspect((1, 1, 1), zoom=-1) + + +def test_colorbar_pos(): + num_plots = 2 + fig, axs = plt.subplots(1, num_plots, figsize=(4, 5), + constrained_layout=True, + subplot_kw={'projection': '3d'}) + for ax in axs: + p_tri = ax.plot_trisurf(np.random.randn(5), np.random.randn(5), + np.random.randn(5)) + + cbar = plt.colorbar(p_tri, ax=axs, orientation='horizontal') + + fig.canvas.draw() + # check that actually on the bottom + assert cbar.ax.get_position().extents[1] < 0.2 + + +def test_inverted_zaxis(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.set_zlim(0, 1) + assert not ax.zaxis_inverted() + assert ax.get_zlim() == (0, 1) + assert ax.get_zbound() == (0, 1) + + # Change bound + ax.set_zbound((0, 2)) + assert not ax.zaxis_inverted() + assert ax.get_zlim() == (0, 2) + assert ax.get_zbound() == (0, 2) + + # Change invert + ax.invert_zaxis() + assert ax.zaxis_inverted() + assert ax.get_zlim() == (2, 0) + assert ax.get_zbound() == (0, 2) + + # Set upper bound + ax.set_zbound(upper=1) + assert ax.zaxis_inverted() + assert ax.get_zlim() == (1, 0) + assert ax.get_zbound() == (0, 1) + + # Set lower bound + ax.set_zbound(lower=2) + assert ax.zaxis_inverted() + assert ax.get_zlim() == (2, 1) + assert ax.get_zbound() == (1, 2) + + +def test_set_zlim(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + assert np.allclose(ax.get_zlim(), (-1/48, 49/48)) + ax.set_zlim(zmax=2) + assert np.allclose(ax.get_zlim(), (-1/48, 2)) + ax.set_zlim(zmin=1) + assert ax.get_zlim() == (1, 2) + + with pytest.raises( + TypeError, match="Cannot pass both 'lower' and 'min'"): + ax.set_zlim(bottom=0, zmin=1) + with pytest.raises( + TypeError, match="Cannot pass both 'upper' and 'max'"): + ax.set_zlim(top=0, zmax=1) + + +@check_figures_equal(extensions=["png"]) +def test_shared_view(fig_test, fig_ref): + elev, azim, roll = 5, 20, 30 + ax1 = fig_test.add_subplot(131, projection="3d") + ax2 = fig_test.add_subplot(132, projection="3d", shareview=ax1) + ax3 = fig_test.add_subplot(133, projection="3d") + ax3.shareview(ax1) + ax2.view_init(elev=elev, azim=azim, roll=roll, share=True) + + for subplot_num in (131, 132, 133): + ax = fig_ref.add_subplot(subplot_num, projection="3d") + ax.view_init(elev=elev, azim=azim, roll=roll) + + +def test_shared_axes_retick(): + fig = plt.figure() + ax1 = fig.add_subplot(211, projection="3d") + ax2 = fig.add_subplot(212, projection="3d", sharez=ax1) + ax1.plot([0, 1], [0, 1], [0, 2]) + ax2.plot([0, 1], [0, 1], [0, 2]) + ax1.set_zticks([-0.5, 0, 2, 2.5]) + # check that setting ticks on a shared axis is synchronized + assert ax1.get_zlim() == (-0.5, 2.5) + assert ax2.get_zlim() == (-0.5, 2.5) + + +def test_quaternion(): + # 1: + q1 = Quaternion(1, [0, 0, 0]) + assert q1.scalar == 1 + assert (q1.vector == [0, 0, 0]).all + # __neg__: + assert (-q1).scalar == -1 + assert ((-q1).vector == [0, 0, 0]).all + # i, j, k: + qi = Quaternion(0, [1, 0, 0]) + assert qi.scalar == 0 + assert (qi.vector == [1, 0, 0]).all + qj = Quaternion(0, [0, 1, 0]) + assert qj.scalar == 0 + assert (qj.vector == [0, 1, 0]).all + qk = Quaternion(0, [0, 0, 1]) + assert qk.scalar == 0 + assert (qk.vector == [0, 0, 1]).all + # i^2 = j^2 = k^2 = -1: + assert qi*qi == -q1 + assert qj*qj == -q1 + assert qk*qk == -q1 + # identity: + assert q1*qi == qi + assert q1*qj == qj + assert q1*qk == qk + # i*j=k, j*k=i, k*i=j: + assert qi*qj == qk + assert qj*qk == qi + assert qk*qi == qj + assert qj*qi == -qk + assert qk*qj == -qi + assert qi*qk == -qj + # __mul__: + assert (Quaternion(2, [3, 4, 5]) * Quaternion(6, [7, 8, 9]) + == Quaternion(-86, [28, 48, 44])) + # conjugate(): + for q in [q1, qi, qj, qk]: + assert q.conjugate().scalar == q.scalar + assert (q.conjugate().vector == -q.vector).all + assert q.conjugate().conjugate() == q + assert ((q*q.conjugate()).vector == 0).all + # norm: + q0 = Quaternion(0, [0, 0, 0]) + assert q0.norm == 0 + assert q1.norm == 1 + assert qi.norm == 1 + assert qj.norm == 1 + assert qk.norm == 1 + for q in [q0, q1, qi, qj, qk]: + assert q.norm == (q*q.conjugate()).scalar + # normalize(): + for q in [ + Quaternion(2, [0, 0, 0]), + Quaternion(0, [3, 0, 0]), + Quaternion(0, [0, 4, 0]), + Quaternion(0, [0, 0, 5]), + Quaternion(6, [7, 8, 9]) + ]: + assert q.normalize().norm == 1 + # reciprocal(): + for q in [q1, qi, qj, qk]: + assert q*q.reciprocal() == q1 + assert q.reciprocal()*q == q1 + # rotate(): + assert (qi.rotate([1, 2, 3]) == np.array([1, -2, -3])).all + # rotate_from_to(): + for r1, r2, q in [ + ([1, 0, 0], [0, 1, 0], Quaternion(np.sqrt(1/2), [0, 0, np.sqrt(1/2)])), + ([1, 0, 0], [0, 0, 1], Quaternion(np.sqrt(1/2), [0, -np.sqrt(1/2), 0])), + ([1, 0, 0], [1, 0, 0], Quaternion(1, [0, 0, 0])) + ]: + assert Quaternion.rotate_from_to(r1, r2) == q + # rotate_from_to(), special case: + for r1 in [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1]]: + r1 = np.array(r1) + with pytest.warns(UserWarning): + q = Quaternion.rotate_from_to(r1, -r1) + assert np.isclose(q.norm, 1) + assert np.dot(q.vector, r1) == 0 + # from_cardan_angles(), as_cardan_angles(): + for elev, azim, roll in [(0, 0, 0), + (90, 0, 0), (0, 90, 0), (0, 0, 90), + (0, 30, 30), (30, 0, 30), (30, 30, 0), + (47, 11, -24)]: + for mag in [1, 2]: + q = Quaternion.from_cardan_angles( + np.deg2rad(elev), np.deg2rad(azim), np.deg2rad(roll)) + assert np.isclose(q.norm, 1) + q = Quaternion(mag * q.scalar, mag * q.vector) + np.testing.assert_allclose(np.rad2deg(Quaternion.as_cardan_angles(q)), + (elev, azim, roll), atol=1e-6) + + +@pytest.mark.parametrize('style', + ('azel', 'trackball', 'sphere', 'arcball')) +def test_rotate(style): + """Test rotating using the left mouse button.""" + if style == 'azel': + s = 0.5 + else: + s = mpl.rcParams['axes3d.trackballsize'] / 2 + s *= 0.5 + mpl.rcParams['axes3d.trackballborder'] = 0 + with mpl.rc_context({'axes3d.mouserotationstyle': style}): + for roll, dx, dy in [ + [0, 1, 0], + [30, 1, 0], + [0, 0, 1], + [30, 0, 1], + [0, 0.5, np.sqrt(3)/2], + [30, 0.5, np.sqrt(3)/2], + [0, 2, 0]]: + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection='3d') + ax.view_init(0, 0, roll) + ax.figure.canvas.draw() + + # drag mouse to change orientation + ax._button_press( + mock_event(ax, button=MouseButton.LEFT, xdata=0, ydata=0)) + ax._on_move( + mock_event(ax, button=MouseButton.LEFT, + xdata=s*dx*ax._pseudo_w, ydata=s*dy*ax._pseudo_h)) + ax.figure.canvas.draw() + + c = np.sqrt(3)/2 + expectations = { + ('azel', 0, 1, 0): (0, -45, 0), + ('azel', 0, 0, 1): (-45, 0, 0), + ('azel', 0, 0.5, c): (-38.971143, -22.5, 0), + ('azel', 0, 2, 0): (0, -90, 0), + ('azel', 30, 1, 0): (22.5, -38.971143, 30), + ('azel', 30, 0, 1): (-38.971143, -22.5, 30), + ('azel', 30, 0.5, c): (-22.5, -38.971143, 30), + + ('trackball', 0, 1, 0): (0, -28.64789, 0), + ('trackball', 0, 0, 1): (-28.64789, 0, 0), + ('trackball', 0, 0.5, c): (-24.531578, -15.277726, 3.340403), + ('trackball', 0, 2, 0): (0, -180/np.pi, 0), + ('trackball', 30, 1, 0): (13.869588, -25.319385, 26.87008), + ('trackball', 30, 0, 1): (-24.531578, -15.277726, 33.340403), + ('trackball', 30, 0.5, c): (-13.869588, -25.319385, 33.129920), + + ('sphere', 0, 1, 0): (0, -30, 0), + ('sphere', 0, 0, 1): (-30, 0, 0), + ('sphere', 0, 0.5, c): (-25.658906, -16.102114, 3.690068), + ('sphere', 0, 2, 0): (0, -90, 0), + ('sphere', 30, 1, 0): (14.477512, -26.565051, 26.565051), + ('sphere', 30, 0, 1): (-25.658906, -16.102114, 33.690068), + ('sphere', 30, 0.5, c): (-14.477512, -26.565051, 33.434949), + + ('arcball', 0, 1, 0): (0, -60, 0), + ('arcball', 0, 0, 1): (-60, 0, 0), + ('arcball', 0, 0.5, c): (-48.590378, -40.893395, 19.106605), + ('arcball', 0, 2, 0): (0, 180, 0), + ('arcball', 30, 1, 0): (25.658906, -56.309932, 16.102114), + ('arcball', 30, 0, 1): (-48.590378, -40.893395, 49.106605), + ('arcball', 30, 0.5, c): (-25.658906, -56.309932, 43.897886)} + new_elev, new_azim, new_roll = expectations[(style, roll, dx, dy)] + np.testing.assert_allclose((ax.elev, ax.azim, ax.roll), + (new_elev, new_azim, new_roll), atol=1e-6) + + +def test_pan(): + """Test mouse panning using the middle mouse button.""" + + def convert_lim(dmin, dmax): + """Convert min/max limits to center and range.""" + center = (dmin + dmax) / 2 + range_ = dmax - dmin + return center, range_ + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.scatter(0, 0, 0) + fig.canvas.draw() + + x_center0, x_range0 = convert_lim(*ax.get_xlim3d()) + y_center0, y_range0 = convert_lim(*ax.get_ylim3d()) + z_center0, z_range0 = convert_lim(*ax.get_zlim3d()) + + # move mouse diagonally to pan along all axis. + ax._button_press( + mock_event(ax, button=MouseButton.MIDDLE, xdata=0, ydata=0)) + ax._on_move( + mock_event(ax, button=MouseButton.MIDDLE, xdata=1, ydata=1)) + + x_center, x_range = convert_lim(*ax.get_xlim3d()) + y_center, y_range = convert_lim(*ax.get_ylim3d()) + z_center, z_range = convert_lim(*ax.get_zlim3d()) + + # Ranges have not changed + assert x_range == pytest.approx(x_range0) + assert y_range == pytest.approx(y_range0) + assert z_range == pytest.approx(z_range0) + + # But center positions have + assert x_center != pytest.approx(x_center0) + assert y_center != pytest.approx(y_center0) + assert z_center != pytest.approx(z_center0) + + +@pytest.mark.parametrize("tool,button,key,expected", + [("zoom", MouseButton.LEFT, None, # zoom in + ((0.00, 0.06), (0.01, 0.07), (0.02, 0.08))), + ("zoom", MouseButton.LEFT, 'x', # zoom in + ((-0.01, 0.10), (-0.03, 0.08), (-0.06, 0.06))), + ("zoom", MouseButton.LEFT, 'y', # zoom in + ((-0.07, 0.05), (-0.04, 0.08), (0.00, 0.12))), + ("zoom", MouseButton.RIGHT, None, # zoom out + ((-0.09, 0.15), (-0.08, 0.17), (-0.07, 0.18))), + ("pan", MouseButton.LEFT, None, + ((-0.70, -0.58), (-1.04, -0.91), (-1.27, -1.15))), + ("pan", MouseButton.LEFT, 'x', + ((-0.97, -0.84), (-0.58, -0.46), (-0.06, 0.06))), + ("pan", MouseButton.LEFT, 'y', + ((0.20, 0.32), (-0.51, -0.39), (-1.27, -1.15)))]) +def test_toolbar_zoom_pan(tool, button, key, expected): + # NOTE: The expected zoom values are rough ballparks of moving in the view + # to make sure we are getting the right direction of motion. + # The specific values can and should change if the zoom movement + # scaling factor gets updated. + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.scatter(0, 0, 0) + fig.canvas.draw() + xlim0, ylim0, zlim0 = ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d() + + # Mouse from (0, 0) to (1, 1) + d0 = (0, 0) + d1 = (1, 1) + # Convert to screen coordinates ("s"). Events are defined only with pixel + # precision, so round the pixel values, and below, check against the + # corresponding xdata/ydata, which are close but not equal to d0/d1. + s0 = ax.transData.transform(d0).astype(int) + s1 = ax.transData.transform(d1).astype(int) + + # Set up the mouse movements + start_event = MouseEvent( + "button_press_event", fig.canvas, *s0, button, key=key) + stop_event = MouseEvent( + "button_release_event", fig.canvas, *s1, button, key=key) + + tb = NavigationToolbar2(fig.canvas) + if tool == "zoom": + tb.zoom() + tb.press_zoom(start_event) + tb.drag_zoom(stop_event) + tb.release_zoom(stop_event) + else: + tb.pan() + tb.press_pan(start_event) + tb.drag_pan(stop_event) + tb.release_pan(stop_event) + + # Should be close, but won't be exact due to screen integer resolution + xlim, ylim, zlim = expected + assert ax.get_xlim3d() == pytest.approx(xlim, abs=0.01) + assert ax.get_ylim3d() == pytest.approx(ylim, abs=0.01) + assert ax.get_zlim3d() == pytest.approx(zlim, abs=0.01) + + # Ensure that back, forward, and home buttons work + tb.back() + assert ax.get_xlim3d() == pytest.approx(xlim0) + assert ax.get_ylim3d() == pytest.approx(ylim0) + assert ax.get_zlim3d() == pytest.approx(zlim0) + + tb.forward() + assert ax.get_xlim3d() == pytest.approx(xlim, abs=0.01) + assert ax.get_ylim3d() == pytest.approx(ylim, abs=0.01) + assert ax.get_zlim3d() == pytest.approx(zlim, abs=0.01) + + tb.home() + assert ax.get_xlim3d() == pytest.approx(xlim0) + assert ax.get_ylim3d() == pytest.approx(ylim0) + assert ax.get_zlim3d() == pytest.approx(zlim0) + + +@mpl.style.context('default') +@check_figures_equal(extensions=["png"]) +def test_scalarmap_update(fig_test, fig_ref): + + x, y, z = np.array(list(itertools.product(*[np.arange(0, 5, 1), + np.arange(0, 5, 1), + np.arange(0, 5, 1)]))).T + c = x + y + + # test + ax_test = fig_test.add_subplot(111, projection='3d') + sc_test = ax_test.scatter(x, y, z, c=c, s=40, cmap='viridis') + # force a draw + fig_test.canvas.draw() + # mark it as "stale" + sc_test.changed() + + # ref + ax_ref = fig_ref.add_subplot(111, projection='3d') + sc_ref = ax_ref.scatter(x, y, z, c=c, s=40, cmap='viridis') + + +def test_subfigure_simple(): + # smoketest that subfigures can work... + fig = plt.figure() + sf = fig.subfigures(1, 2) + ax = sf[0].add_subplot(1, 1, 1, projection='3d') + ax = sf[1].add_subplot(1, 1, 1, projection='3d', label='other') + + +# Update style when regenerating the test image +@image_comparison(baseline_images=['computed_zorder'], remove_text=True, + extensions=['png'], style=('mpl20')) +def test_computed_zorder(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax1 = fig.add_subplot(221, projection='3d') + ax2 = fig.add_subplot(222, projection='3d') + ax2.computed_zorder = False + + # create a horizontal plane + corners = ((0, 0, 0), (0, 5, 0), (5, 5, 0), (5, 0, 0)) + for ax in (ax1, ax2): + tri = art3d.Poly3DCollection([corners], + facecolors='white', + edgecolors='black', + zorder=1) + ax.add_collection3d(tri) + + # plot a vector + ax.plot((2, 2), (2, 2), (0, 4), c='red', zorder=2) + + # plot some points + ax.scatter((3, 3), (1, 3), (1, 3), c='red', zorder=10) + + ax.set_xlim((0, 5.0)) + ax.set_ylim((0, 5.0)) + ax.set_zlim((0, 2.5)) + + ax3 = fig.add_subplot(223, projection='3d') + ax4 = fig.add_subplot(224, projection='3d') + ax4.computed_zorder = False + + dim = 10 + X, Y = np.meshgrid((-dim, dim), (-dim, dim)) + Z = np.zeros((2, 2)) + + angle = 0.5 + X2, Y2 = np.meshgrid((-dim, dim), (0, dim)) + Z2 = Y2 * angle + X3, Y3 = np.meshgrid((-dim, dim), (-dim, 0)) + Z3 = Y3 * angle + + r = 7 + M = 1000 + th = np.linspace(0, 2 * np.pi, M) + x, y, z = r * np.cos(th), r * np.sin(th), angle * r * np.sin(th) + for ax in (ax3, ax4): + ax.plot_surface(X2, Y3, Z3, + color='blue', + alpha=0.5, + linewidth=0, + zorder=-1) + ax.plot(x[y < 0], y[y < 0], z[y < 0], + lw=5, + linestyle='--', + color='green', + zorder=0) + + ax.plot_surface(X, Y, Z, + color='red', + alpha=0.5, + linewidth=0, + zorder=1) + + ax.plot(r * np.sin(th), r * np.cos(th), np.zeros(M), + lw=5, + linestyle='--', + color='black', + zorder=2) + + ax.plot_surface(X2, Y2, Z2, + color='blue', + alpha=0.5, + linewidth=0, + zorder=3) + + ax.plot(x[y > 0], y[y > 0], z[y > 0], lw=5, + linestyle='--', + color='green', + zorder=4) + ax.view_init(elev=20, azim=-20, roll=0) + ax.axis('off') + + +def test_format_coord(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + x = np.arange(10) + ax.plot(x, np.sin(x)) + xv = 0.1 + yv = 0.1 + fig.canvas.draw() + assert ax.format_coord(xv, yv) == 'x=10.5227, y pane=1.0417, z=0.1444' + + # Modify parameters + ax.view_init(roll=30, vertical_axis="y") + fig.canvas.draw() + assert ax.format_coord(xv, yv) == 'x pane=9.1875, y=0.9761, z=0.1291' + + # Reset parameters + ax.view_init() + fig.canvas.draw() + assert ax.format_coord(xv, yv) == 'x=10.5227, y pane=1.0417, z=0.1444' + + # Check orthographic projection + ax.set_proj_type('ortho') + fig.canvas.draw() + assert ax.format_coord(xv, yv) == 'x=10.8869, y pane=1.0417, z=0.1528' + + # Check non-default perspective projection + ax.set_proj_type('persp', focal_length=0.1) + fig.canvas.draw() + assert ax.format_coord(xv, yv) == 'x=9.0620, y pane=1.0417, z=0.1110' + + +def test_get_axis_position(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + x = np.arange(10) + ax.plot(x, np.sin(x)) + fig.canvas.draw() + assert ax.get_axis_position() == (False, True, False) + + +def test_margins(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.margins(0.2) + assert ax.margins() == (0.2, 0.2, 0.2) + ax.margins(0.1, 0.2, 0.3) + assert ax.margins() == (0.1, 0.2, 0.3) + ax.margins(x=0) + assert ax.margins() == (0, 0.2, 0.3) + ax.margins(y=0.1) + assert ax.margins() == (0, 0.1, 0.3) + ax.margins(z=0) + assert ax.margins() == (0, 0.1, 0) + + +def test_margin_getters(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.margins(0.1, 0.2, 0.3) + assert ax.get_xmargin() == 0.1 + assert ax.get_ymargin() == 0.2 + assert ax.get_zmargin() == 0.3 + + +@pytest.mark.parametrize('err, args, kwargs, match', ( + (ValueError, (-1,), {}, r'margin must be greater than -0\.5'), + (ValueError, (1, -1, 1), {}, r'margin must be greater than -0\.5'), + (ValueError, (1, 1, -1), {}, r'margin must be greater than -0\.5'), + (ValueError, tuple(), {'x': -1}, r'margin must be greater than -0\.5'), + (ValueError, tuple(), {'y': -1}, r'margin must be greater than -0\.5'), + (ValueError, tuple(), {'z': -1}, r'margin must be greater than -0\.5'), + (TypeError, (1, ), {'x': 1}, + 'Cannot pass both positional and keyword'), + (TypeError, (1, ), {'x': 1, 'y': 1, 'z': 1}, + 'Cannot pass both positional and keyword'), + (TypeError, (1, ), {'x': 1, 'y': 1}, + 'Cannot pass both positional and keyword'), + (TypeError, (1, 1), {}, 'Must pass a single positional argument for'), +)) +def test_margins_errors(err, args, kwargs, match): + with pytest.raises(err, match=match): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.margins(*args, **kwargs) + + +@check_figures_equal(extensions=["png"]) +def test_text_3d(fig_test, fig_ref): + ax = fig_ref.add_subplot(projection="3d") + txt = Text(0.5, 0.5, r'Foo bar $\int$') + art3d.text_2d_to_3d(txt, z=1) + ax.add_artist(txt) + assert txt.get_position_3d() == (0.5, 0.5, 1) + + ax = fig_test.add_subplot(projection="3d") + t3d = art3d.Text3D(0.5, 0.5, 1, r'Foo bar $\int$') + ax.add_artist(t3d) + assert t3d.get_position_3d() == (0.5, 0.5, 1) + + +def test_draw_single_lines_from_Nx1(): + # Smoke test for GH#23459 + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + ax.plot([[0], [1]], [[0], [1]], [[0], [1]]) + + +@check_figures_equal(extensions=["png"]) +def test_pathpatch_3d(fig_test, fig_ref): + ax = fig_ref.add_subplot(projection="3d") + path = Path.unit_rectangle() + patch = PathPatch(path) + art3d.pathpatch_2d_to_3d(patch, z=(0, 0.5, 0.7, 1, 0), zdir='y') + ax.add_artist(patch) + + ax = fig_test.add_subplot(projection="3d") + pp3d = art3d.PathPatch3D(path, zs=(0, 0.5, 0.7, 1, 0), zdir='y') + ax.add_artist(pp3d) + + +@image_comparison(baseline_images=['scatter_spiral.png'], + remove_text=True, + style='mpl20') +def test_scatter_spiral(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + th = np.linspace(0, 2 * np.pi * 6, 256) + sc = ax.scatter(np.sin(th), np.cos(th), th, s=(1 + th * 5), c=th ** 2) + + # force at least 1 draw! + fig.canvas.draw() + + +def test_Poly3DCollection_get_path(): + # Smoke test to see that get_path does not raise + # See GH#27361 + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + p = Circle((0, 0), 1.0) + ax.add_patch(p) + art3d.pathpatch_2d_to_3d(p) + p.get_path() + + +def test_Poly3DCollection_get_facecolor(): + # Smoke test to see that get_facecolor does not raise + # See GH#4067 + y, x = np.ogrid[1:10:100j, 1:10:100j] + z2 = np.cos(x) ** 3 - np.sin(y) ** 2 + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + r = ax.plot_surface(x, y, z2, cmap='hot') + r.get_facecolor() + + +def test_Poly3DCollection_get_edgecolor(): + # Smoke test to see that get_edgecolor does not raise + # See GH#4067 + y, x = np.ogrid[1:10:100j, 1:10:100j] + z2 = np.cos(x) ** 3 - np.sin(y) ** 2 + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + r = ax.plot_surface(x, y, z2, cmap='hot') + r.get_edgecolor() + + +@pytest.mark.parametrize( + "vertical_axis, proj_expected, axis_lines_expected, tickdirs_expected", + [ + ( + "z", + [ + [0.0, 1.142857, 0.0, -0.571429], + [0.0, 0.0, 0.857143, -0.428571], + [0.0, 0.0, 0.0, -10.0], + [-1.142857, 0.0, 0.0, 10.571429], + ], + [ + ([0.05617978, 0.06329114], [-0.04213483, -0.04746835]), + ([-0.06329114, 0.06329114], [-0.04746835, -0.04746835]), + ([-0.06329114, -0.06329114], [-0.04746835, 0.04746835]), + ], + [1, 0, 0], + ), + ( + "y", + [ + [1.142857, 0.0, 0.0, -0.571429], + [0.0, 0.857143, 0.0, -0.428571], + [0.0, 0.0, 0.0, -10.0], + [0.0, 0.0, -1.142857, 10.571429], + ], + [ + ([-0.06329114, 0.06329114], [0.04746835, 0.04746835]), + ([0.06329114, 0.06329114], [-0.04746835, 0.04746835]), + ([-0.05617978, -0.06329114], [0.04213483, 0.04746835]), + ], + [2, 2, 0], + ), + ( + "x", + [ + [0.0, 0.0, 1.142857, -0.571429], + [0.857143, 0.0, 0.0, -0.428571], + [0.0, 0.0, 0.0, -10.0], + [0.0, -1.142857, 0.0, 10.571429], + ], + [ + ([-0.06329114, -0.06329114], [0.04746835, -0.04746835]), + ([0.06329114, 0.05617978], [0.04746835, 0.04213483]), + ([0.06329114, -0.06329114], [0.04746835, 0.04746835]), + ], + [1, 2, 1], + ), + ], +) +def test_view_init_vertical_axis( + vertical_axis, proj_expected, axis_lines_expected, tickdirs_expected +): + """ + Test the actual projection, axis lines and ticks matches expected values. + + Parameters + ---------- + vertical_axis : str + Axis to align vertically. + proj_expected : ndarray + Expected values from ax.get_proj(). + axis_lines_expected : tuple of arrays + Edgepoints of the axis line. Expected values retrieved according + to ``ax.get_[xyz]axis().line.get_data()``. + tickdirs_expected : list of int + indexes indicating which axis to create a tick line along. + """ + rtol = 2e-06 + ax = plt.subplot(1, 1, 1, projection="3d") + ax.view_init(elev=0, azim=0, roll=0, vertical_axis=vertical_axis) + ax.get_figure().canvas.draw() + + # Assert the projection matrix: + proj_actual = ax.get_proj() + np.testing.assert_allclose(proj_expected, proj_actual, rtol=rtol) + + for i, axis in enumerate([ax.get_xaxis(), ax.get_yaxis(), ax.get_zaxis()]): + # Assert black lines are correctly aligned: + axis_line_expected = axis_lines_expected[i] + axis_line_actual = axis.line.get_data() + np.testing.assert_allclose(axis_line_expected, axis_line_actual, + rtol=rtol) + + # Assert ticks are correctly aligned: + tickdir_expected = tickdirs_expected[i] + tickdir_actual = axis._get_tickdir('default') + np.testing.assert_array_equal(tickdir_expected, tickdir_actual) + + +@pytest.mark.parametrize("vertical_axis", ["x", "y", "z"]) +def test_on_move_vertical_axis(vertical_axis: str) -> None: + """ + Test vertical axis is respected when rotating the plot interactively. + """ + ax = plt.subplot(1, 1, 1, projection="3d") + ax.view_init(elev=0, azim=0, roll=0, vertical_axis=vertical_axis) + ax.get_figure().canvas.draw() + + proj_before = ax.get_proj() + event_click = mock_event(ax, button=MouseButton.LEFT, xdata=0, ydata=1) + ax._button_press(event_click) + + event_move = mock_event(ax, button=MouseButton.LEFT, xdata=0.5, ydata=0.8) + ax._on_move(event_move) + + assert ax._axis_names.index(vertical_axis) == ax._vertical_axis + + # Make sure plot has actually moved: + proj_after = ax.get_proj() + np.testing.assert_raises( + AssertionError, np.testing.assert_allclose, proj_before, proj_after + ) + + +@pytest.mark.parametrize( + "vertical_axis, aspect_expected", + [ + ("x", [1.190476, 0.892857, 1.190476]), + ("y", [0.892857, 1.190476, 1.190476]), + ("z", [1.190476, 1.190476, 0.892857]), + ], +) +def test_set_box_aspect_vertical_axis(vertical_axis, aspect_expected): + ax = plt.subplot(1, 1, 1, projection="3d") + ax.view_init(elev=0, azim=0, roll=0, vertical_axis=vertical_axis) + ax.get_figure().canvas.draw() + + ax.set_box_aspect(None) + + np.testing.assert_allclose(aspect_expected, ax._box_aspect, rtol=1e-6) + + +@image_comparison(baseline_images=['arc_pathpatch.png'], + remove_text=True, + style='mpl20') +def test_arc_pathpatch(): + ax = plt.subplot(1, 1, 1, projection="3d") + a = mpatch.Arc((0.5, 0.5), width=0.5, height=0.9, + angle=20, theta1=10, theta2=130) + ax.add_patch(a) + art3d.pathpatch_2d_to_3d(a, z=0, zdir='z') + + +@image_comparison(baseline_images=['panecolor_rcparams.png'], + remove_text=True, + style='mpl20') +def test_panecolor_rcparams(): + with plt.rc_context({'axes3d.xaxis.panecolor': 'r', + 'axes3d.yaxis.panecolor': 'g', + 'axes3d.zaxis.panecolor': 'b'}): + fig = plt.figure(figsize=(1, 1)) + fig.add_subplot(projection='3d') + + +@check_figures_equal(extensions=["png"]) +def test_mutating_input_arrays_y_and_z(fig_test, fig_ref): + """ + Test to see if the `z` axis does not get mutated + after a call to `Axes3D.plot` + + test cases came from GH#8990 + """ + ax1 = fig_test.add_subplot(111, projection='3d') + x = [1, 2, 3] + y = [0.0, 0.0, 0.0] + z = [0.0, 0.0, 0.0] + ax1.plot(x, y, z, 'o-') + + # mutate y,z to get a nontrivial line + y[:] = [1, 2, 3] + z[:] = [1, 2, 3] + + # draw the same plot without mutating x and y + ax2 = fig_ref.add_subplot(111, projection='3d') + x = [1, 2, 3] + y = [0.0, 0.0, 0.0] + z = [0.0, 0.0, 0.0] + ax2.plot(x, y, z, 'o-') + + +def test_scatter_masked_color(): + """ + Test color parameter usage with non-finite coordinate arrays. + + GH#26236 + """ + + x = [np.nan, 1, 2, 1] + y = [0, np.inf, 2, 1] + z = [0, 1, -np.inf, 1] + colors = [ + [0.0, 0.0, 0.0, 1], + [0.0, 0.0, 0.0, 1], + [0.0, 0.0, 0.0, 1], + [0.0, 0.0, 0.0, 1] + ] + + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + path3d = ax.scatter(x, y, z, color=colors) + + # Assert sizes' equality + assert len(path3d.get_offsets()) ==\ + len(super(type(path3d), path3d).get_facecolors()) + + +@mpl3d_image_comparison(['surface3d_zsort_inf.png'], style='mpl20') +def test_surface3d_zsort_inf(): + plt.rcParams['axes3d.automargin'] = True # Remove when image is regenerated + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + + x, y = np.mgrid[-2:2:0.1, -2:2:0.1] + z = np.sin(x)**2 + np.cos(y)**2 + z[x.shape[0] // 2:, x.shape[1] // 2:] = np.inf + + ax.plot_surface(x, y, z, cmap='jet') + ax.view_init(elev=45, azim=145) + + +def test_Poly3DCollection_init_value_error(): + # smoke test to ensure the input check works + # GH#26420 + with pytest.raises(ValueError, + match='You must provide facecolors, edgecolors, ' + 'or both for shade to work.'): + poly = np.array([[0, 0, 1], [0, 1, 1], [0, 0, 0]], float) + c = art3d.Poly3DCollection([poly], shade=True) + + +def test_ndarray_color_kwargs_value_error(): + # smoke test + # ensures ndarray can be passed to color in kwargs for 3d projection plot + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + ax.scatter(1, 0, 0, color=np.array([0, 0, 0, 1])) + fig.canvas.draw() diff --git a/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_legend3d.py b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_legend3d.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd676df1e31591f343a4a1bd27c27b98d564949 --- /dev/null +++ b/venv/lib/python3.10/site-packages/mpl_toolkits/mplot3d/tests/test_legend3d.py @@ -0,0 +1,117 @@ +import platform + +import numpy as np + +import matplotlib as mpl +from matplotlib.colors import same_color +from matplotlib.testing.decorators import image_comparison +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import art3d + + +@image_comparison(['legend_plot.png'], remove_text=True, style='mpl20') +def test_legend_plot(): + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + x = np.arange(10) + ax.plot(x, 5 - x, 'o', zdir='y', label='z=1') + ax.plot(x, x - 5, 'o', zdir='y', label='z=-1') + ax.legend() + + +@image_comparison(['legend_bar.png'], remove_text=True, style='mpl20') +def test_legend_bar(): + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + x = np.arange(10) + b1 = ax.bar(x, x, zdir='y', align='edge', color='m') + b2 = ax.bar(x, x[::-1], zdir='x', align='edge', color='g') + ax.legend([b1[0], b2[0]], ['up', 'down']) + + +@image_comparison(['fancy.png'], remove_text=True, style='mpl20', + tol=0 if platform.machine() == 'x86_64' else 0.011) +def test_fancy(): + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + ax.plot(np.arange(10), np.full(10, 5), np.full(10, 5), 'o--', label='line') + ax.scatter(np.arange(10), np.arange(10, 0, -1), label='scatter') + ax.errorbar(np.full(10, 5), np.arange(10), np.full(10, 10), + xerr=0.5, zerr=0.5, label='errorbar') + ax.legend(loc='lower left', ncols=2, title='My legend', numpoints=1) + + +def test_linecollection_scaled_dashes(): + lines1 = [[(0, .5), (.5, 1)], [(.3, .6), (.2, .2)]] + lines2 = [[[0.7, .2], [.8, .4]], [[.5, .7], [.6, .1]]] + lines3 = [[[0.6, .2], [.8, .4]], [[.5, .7], [.1, .1]]] + lc1 = art3d.Line3DCollection(lines1, linestyles="--", lw=3) + lc2 = art3d.Line3DCollection(lines2, linestyles="-.") + lc3 = art3d.Line3DCollection(lines3, linestyles=":", lw=.5) + + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + ax.add_collection(lc1) + ax.add_collection(lc2) + ax.add_collection(lc3) + + leg = ax.legend([lc1, lc2, lc3], ['line1', 'line2', 'line 3']) + h1, h2, h3 = leg.legend_handles + + for oh, lh in zip((lc1, lc2, lc3), (h1, h2, h3)): + assert oh.get_linestyles()[0] == lh._dash_pattern + + +def test_handlerline3d(): + # Test marker consistency for monolithic Line3D legend handler. + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + ax.scatter([0, 1], [0, 1], marker="v") + handles = [art3d.Line3D([0], [0], [0], marker="v")] + leg = ax.legend(handles, ["Aardvark"], numpoints=1) + assert handles[0].get_marker() == leg.legend_handles[0].get_marker() + + +def test_contour_legend_elements(): + x, y = np.mgrid[1:10, 1:10] + h = x * y + colors = ['blue', '#00FF00', 'red'] + + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + cs = ax.contour(x, y, h, levels=[10, 30, 50], colors=colors, extend='both') + + artists, labels = cs.legend_elements() + assert labels == ['$x = 10.0$', '$x = 30.0$', '$x = 50.0$'] + assert all(isinstance(a, mpl.lines.Line2D) for a in artists) + assert all(same_color(a.get_color(), c) + for a, c in zip(artists, colors)) + + +def test_contourf_legend_elements(): + x, y = np.mgrid[1:10, 1:10] + h = x * y + + fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) + cs = ax.contourf(x, y, h, levels=[10, 30, 50], + colors=['#FFFF00', '#FF00FF', '#00FFFF'], + extend='both') + cs.cmap.set_over('red') + cs.cmap.set_under('blue') + cs.changed() + artists, labels = cs.legend_elements() + assert labels == ['$x \\leq -1e+250s$', + '$10.0 < x \\leq 30.0$', + '$30.0 < x \\leq 50.0$', + '$x > 1e+250s$'] + expected_colors = ('blue', '#FFFF00', '#FF00FF', 'red') + assert all(isinstance(a, mpl.patches.Rectangle) for a in artists) + assert all(same_color(a.get_facecolor(), c) + for a, c in zip(artists, expected_colors)) + + +def test_legend_Poly3dCollection(): + + verts = np.asarray([[0, 0, 0], [0, 1, 1], [1, 0, 1]]) + mesh = art3d.Poly3DCollection([verts], label="surface") + + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + mesh.set_edgecolor('k') + handle = ax.add_collection3d(mesh) + leg = ax.legend() + assert (leg.legend_handles[0].get_facecolor() + == handle.get_facecolor()).all() diff --git a/venv/lib/python3.10/site-packages/multidict/__init__.py b/venv/lib/python3.10/site-packages/multidict/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a688932e52f4e017972878f3307641435fb928bf --- /dev/null +++ b/venv/lib/python3.10/site-packages/multidict/__init__.py @@ -0,0 +1,60 @@ +""" +Multidict implementation. + +HTTP Headers and URL query string require specific data structure: +multidict. It behaves mostly like a dict but it can have +several values for the same key. +""" + +from typing import TYPE_CHECKING + +from ._abc import MultiMapping, MutableMultiMapping +from ._compat import USE_EXTENSIONS + +__all__ = ( + "CIMultiDict", + "CIMultiDictProxy", + "MultiDict", + "MultiDictProxy", + "MultiMapping", + "MutableMultiMapping", + "getversion", + "istr", + "upstr", +) + +__version__ = "6.7.0" + + +if TYPE_CHECKING or not USE_EXTENSIONS: + from ._multidict_py import ( + CIMultiDict, + CIMultiDictProxy, + MultiDict, + MultiDictProxy, + getversion, + istr, + ) +else: + from collections.abc import ItemsView, KeysView, ValuesView + + from ._multidict import ( + CIMultiDict, + CIMultiDictProxy, + MultiDict, + MultiDictProxy, + _ItemsView, + _KeysView, + _ValuesView, + getversion, + istr, + ) + + MultiMapping.register(MultiDictProxy) + MutableMultiMapping.register(MultiDict) + KeysView.register(_KeysView) + ItemsView.register(_ItemsView) + ValuesView.register(_ValuesView) + + +upstr = istr diff --git a/venv/lib/python3.10/site-packages/multidict/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/multidict/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80312c045e4671ccad750370a66f4f86cf9b5c18 Binary files /dev/null and b/venv/lib/python3.10/site-packages/multidict/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/multidict/__pycache__/_abc.cpython-310.pyc b/venv/lib/python3.10/site-packages/multidict/__pycache__/_abc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4abf1f37ad37907f5bc9b4e5d2622f6bae3fcd54 Binary files /dev/null and b/venv/lib/python3.10/site-packages/multidict/__pycache__/_abc.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/multidict/__pycache__/_compat.cpython-310.pyc b/venv/lib/python3.10/site-packages/multidict/__pycache__/_compat.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de0cfa4534005e3e6af8fa1487bd80c53db7a540 Binary files /dev/null and b/venv/lib/python3.10/site-packages/multidict/__pycache__/_compat.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/multidict/__pycache__/_multidict_py.cpython-310.pyc b/venv/lib/python3.10/site-packages/multidict/__pycache__/_multidict_py.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6eeb3089670e4ba2f196674cb5c16f42e8299760 Binary files /dev/null and b/venv/lib/python3.10/site-packages/multidict/__pycache__/_multidict_py.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/multidict/_abc.py b/venv/lib/python3.10/site-packages/multidict/_abc.py new file mode 100644 index 0000000000000000000000000000000000000000..54253e9e779915aa8741313563359a8d9d87ec16 --- /dev/null +++ b/venv/lib/python3.10/site-packages/multidict/_abc.py @@ -0,0 +1,73 @@ +import abc +from collections.abc import Iterable, Mapping, MutableMapping +from typing import TYPE_CHECKING, Protocol, TypeVar, Union, overload + +if TYPE_CHECKING: + from ._multidict_py import istr +else: + istr = str + +_V = TypeVar("_V") +_V_co = TypeVar("_V_co", covariant=True) +_T = TypeVar("_T") + + +class SupportsKeys(Protocol[_V_co]): + def keys(self) -> Iterable[str]: ... + def __getitem__(self, key: str, /) -> _V_co: ... + + +class SupportsIKeys(Protocol[_V_co]): + def keys(self) -> Iterable[istr]: ... + def __getitem__(self, key: istr, /) -> _V_co: ... + + +MDArg = Union[SupportsKeys[_V], SupportsIKeys[_V], Iterable[tuple[str, _V]], None] + + +class MultiMapping(Mapping[str, _V_co]): + @overload + def getall(self, key: str) -> list[_V_co]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V_co], _T]: ... + @abc.abstractmethod + def getall(self, key: str, default: _T = ...) -> Union[list[_V_co], _T]: + """Return all values for key.""" + + @overload + def getone(self, key: str) -> _V_co: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V_co, _T]: ... + @abc.abstractmethod + def getone(self, key: str, default: _T = ...) -> Union[_V_co, _T]: + """Return first value for key.""" + + +class MutableMultiMapping(MultiMapping[_V], MutableMapping[str, _V]): + @abc.abstractmethod + def add(self, key: str, value: _V) -> None: + """Add value to list.""" + + @abc.abstractmethod + def extend(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Add everything from arg and kwargs to the mapping.""" + + @abc.abstractmethod + def merge(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Merge into the mapping, adding non-existing keys.""" + + @overload + def popone(self, key: str) -> _V: ... + @overload + def popone(self, key: str, default: _T) -> Union[_V, _T]: ... + @abc.abstractmethod + def popone(self, key: str, default: _T = ...) -> Union[_V, _T]: + """Remove specified key and return the corresponding value.""" + + @overload + def popall(self, key: str) -> list[_V]: ... + @overload + def popall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + @abc.abstractmethod + def popall(self, key: str, default: _T = ...) -> Union[list[_V], _T]: + """Remove all occurrences of key and return the list of corresponding values.""" diff --git a/venv/lib/python3.10/site-packages/multidict/_compat.py b/venv/lib/python3.10/site-packages/multidict/_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..264d327e6a920503a72cc9e6174b26b63677faf5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/multidict/_compat.py @@ -0,0 +1,15 @@ +import os +import platform + +NO_EXTENSIONS = bool(os.environ.get("MULTIDICT_NO_EXTENSIONS")) + +PYPY = platform.python_implementation() == "PyPy" + +USE_EXTENSIONS = not NO_EXTENSIONS and not PYPY + +if USE_EXTENSIONS: + try: + from . import _multidict # type: ignore[attr-defined] # noqa: F401 + except ImportError: # pragma: no cover + # FIXME: Refactor for coverage. See #837. + USE_EXTENSIONS = False diff --git a/venv/lib/python3.10/site-packages/multidict/_multidict_py.py b/venv/lib/python3.10/site-packages/multidict/_multidict_py.py new file mode 100644 index 0000000000000000000000000000000000000000..6b68d52eeff981cc8082e3426ba037ae1aedb2c1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/multidict/_multidict_py.py @@ -0,0 +1,1242 @@ +import enum +import functools +import reprlib +import sys +from array import array +from collections.abc import ( + ItemsView, + Iterable, + Iterator, + KeysView, + Mapping, + ValuesView, +) +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Generic, + NoReturn, + Optional, + TypeVar, + Union, + cast, + overload, +) + +from ._abc import MDArg, MultiMapping, MutableMultiMapping, SupportsKeys + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + + +class istr(str): + """Case insensitive str.""" + + __is_istr__ = True + __istr_identity__: Optional[str] = None + + +_V = TypeVar("_V") +_T = TypeVar("_T") + +_SENTINEL = enum.Enum("_SENTINEL", "sentinel") +sentinel = _SENTINEL.sentinel + +_version = array("Q", [0]) + + +class _Iter(Generic[_T]): + __slots__ = ("_size", "_iter") + + def __init__(self, size: int, iterator: Iterator[_T]): + self._size = size + self._iter = iterator + + def __iter__(self) -> Self: + return self + + def __next__(self) -> _T: + return next(self._iter) + + def __length_hint__(self) -> int: + return self._size + + +class _ViewBase(Generic[_V]): + def __init__( + self, + md: "MultiDict[_V]", + ): + self._md = md + + def __len__(self) -> int: + return len(self._md) + + +class _ItemsView(_ViewBase[_V], ItemsView[str, _V]): + def __contains__(self, item: object) -> bool: + if not isinstance(item, (tuple, list)) or len(item) != 2: + return False + key, value = item + try: + identity = self._md._identity(key) + except TypeError: + return False + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and value == e.value: + return True + return False + + def __iter__(self) -> _Iter[tuple[str, _V]]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[tuple[str, _V]]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield self._md._key(e.key), e.value + + @reprlib.recursive_repr() + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(f"'{e.key}': {e.value!r}") + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + def _parse_item( + self, arg: Union[tuple[str, _V], _T] + ) -> Optional[tuple[int, str, str, _V]]: + if not isinstance(arg, tuple): + return None + if len(arg) != 2: + return None + try: + identity = self._md._identity(arg[0]) + return (hash(identity), identity, arg[0], arg[1]) + except TypeError: + return None + + def _tmp_set(self, it: Iterable[_T]) -> set[tuple[str, _V]]: + tmp = set() + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + else: + tmp.add((item[1], item[3])) + return tmp + + def __and__(self, other: Iterable[Any]) -> set[tuple[str, _V]]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + e.hash = -1 + if e.identity == identity and e.value == value: + ret.add((e.key, e.value)) + self._md._keys.restore_hash(hash_) + return ret + + def __rand__(self, other: Iterable[_T]) -> set[_T]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: + ret.add(arg) + break + return ret + + def __or__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + ret: set[Union[tuple[str, _V], _T]] = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item: Optional[tuple[int, str, str, _V]] = self._parse_item(arg) + if item is None: + ret.add(arg) + continue + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + break + else: + ret.add(arg) + return ret + + def __ror__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + try: + ret: set[Union[tuple[str, _V], _T]] = set(other) + except TypeError: + return NotImplemented + tmp = self._tmp_set(ret) + + for e in self._md._keys.iter_entries(): + if (e.identity, e.value) not in tmp: + ret.add((e.key, e.value)) + return ret + + def __sub__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + ret: set[Union[tuple[str, _V], _T]] = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + tmp = self._tmp_set(it) + + for e in self._md._keys.iter_entries(): + if (e.identity, e.value) not in tmp: + ret.add((e.key, e.value)) + + return ret + + def __rsub__(self, other: Iterable[_T]) -> set[_T]: + ret: set[_T] = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for arg in it: + item = self._parse_item(arg) + if item is None: + ret.add(arg) + continue + + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + break + else: + ret.add(arg) + return ret + + def __xor__(self, other: Iterable[_T]) -> set[Union[tuple[str, _V], _T]]: + try: + rgt = set(other) + except TypeError: + return NotImplemented + ret: set[Union[tuple[str, _V], _T]] = self - rgt + ret |= rgt - self + return ret + + __rxor__ = __xor__ + + def isdisjoint(self, other: Iterable[tuple[str, _V]]) -> bool: + for arg in other: + item = self._parse_item(arg) + if item is None: + continue + + hash_, identity, key, value = item + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity and e.value == value: # pragma: no branch + return False + return True + + +class _ValuesView(_ViewBase[_V], ValuesView[_V]): + def __contains__(self, value: object) -> bool: + for e in self._md._keys.iter_entries(): + if e.value == value: + return True + return False + + def __iter__(self) -> _Iter[_V]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[_V]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield e.value + + @reprlib.recursive_repr() + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(repr(e.value)) + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + +class _KeysView(_ViewBase[_V], KeysView[str]): + def __contains__(self, key: object) -> bool: + if not isinstance(key, str): + return False + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return True + return False + + def __iter__(self) -> _Iter[str]: + return _Iter(len(self), self._iter(self._md._version)) + + def _iter(self, version: int) -> Iterator[str]: + for e in self._md._keys.iter_entries(): + if version != self._md._version: + raise RuntimeError("Dictionary changed during iteration") + yield self._md._key(e.key) + + def __repr__(self) -> str: + lst = [] + for e in self._md._keys.iter_entries(): + lst.append(f"'{e.key}'") + body = ", ".join(lst) + return f"<{self.__class__.__name__}({body})>" + + def __and__(self, other: Iterable[object]) -> set[str]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + ret.add(e.key) + break + return ret + + def __rand__(self, other: Iterable[_T]) -> set[_T]: + ret = set() + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + if key in self._md: + ret.add(key) + return cast(set[_T], ret) + + def __or__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + ret: set[Union[str, _T]] = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + ret.add(key) + continue + if key not in self._md: + ret.add(key) + return ret + + def __ror__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + try: + ret: set[Union[str, _T]] = set(other) + except TypeError: + return NotImplemented + + tmp = set() + for key in ret: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + tmp.add(identity) + + for e in self._md._keys.iter_entries(): + if e.identity not in tmp: + ret.add(e.key) + return ret + + def __sub__(self, other: Iterable[object]) -> set[str]: + ret = set(self) + try: + it = iter(other) + except TypeError: + return NotImplemented + for key in it: + if not isinstance(key, str): + continue + identity = self._md._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._md._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + ret.discard(e.key) + break + return ret + + def __rsub__(self, other: Iterable[_T]) -> set[_T]: + try: + ret: set[_T] = set(other) + except TypeError: + return NotImplemented + for key in other: + if not isinstance(key, str): + continue + if key in self._md: + ret.discard(key) # type: ignore[arg-type] + return ret + + def __xor__(self, other: Iterable[_T]) -> set[Union[str, _T]]: + try: + rgt = set(other) + except TypeError: + return NotImplemented + ret: set[Union[str, _T]] = self - rgt # type: ignore[assignment] + ret |= rgt - self + return ret + + __rxor__ = __xor__ + + def isdisjoint(self, other: Iterable[object]) -> bool: + for key in other: + if not isinstance(key, str): + continue + if key in self._md: + return False + return True + + +class _CSMixin: + _ci: ClassVar[bool] = False + + def _key(self, key: str) -> str: + return key + + def _identity(self, key: str) -> str: + if isinstance(key, str): + return key + else: + raise TypeError("MultiDict keys should be either str or subclasses of str") + + +class _CIMixin: + _ci: ClassVar[bool] = True + + def _key(self, key: str) -> str: + if type(key) is istr: + return key + else: + return istr(key) + + def _identity(self, key: str) -> str: + if isinstance(key, istr): + ret = key.__istr_identity__ + if ret is None: + ret = key.lower() + key.__istr_identity__ = ret + return ret + if isinstance(key, str): + return key.lower() + else: + raise TypeError("MultiDict keys should be either str or subclasses of str") + + +def estimate_log2_keysize(n: int) -> int: + # 7 == HT_MINSIZE - 1 + return (((n * 3 + 1) // 2) | 7).bit_length() + + +@dataclass +class _Entry(Generic[_V]): + hash: int + identity: str + key: str + value: _V + + +@dataclass +class _HtKeys(Generic[_V]): # type: ignore[misc] + LOG_MINSIZE: ClassVar[int] = 3 + MINSIZE: ClassVar[int] = 8 + PREALLOCATED_INDICES: ClassVar[dict[int, array]] = { # type: ignore[type-arg] + log2_size: array( + "b" if log2_size < 8 else "h", (-1 for i in range(1 << log2_size)) + ) + for log2_size in range(3, 10) + } + + log2_size: int + usable: int + + indices: array # type: ignore[type-arg] # in py3.9 array is not generic + entries: list[Optional[_Entry[_V]]] + + @functools.cached_property + def nslots(self) -> int: + return 1 << self.log2_size + + @functools.cached_property + def mask(self) -> int: + return self.nslots - 1 + + if sys.implementation.name != "pypy": + + def __sizeof__(self) -> int: + return ( + object.__sizeof__(self) + + sys.getsizeof(self.indices) + + sys.getsizeof(self.entries) + ) + + @classmethod + def new(cls, log2_size: int, entries: list[Optional[_Entry[_V]]]) -> Self: + size = 1 << log2_size + usable = (size << 1) // 3 + if log2_size < 10: + indices = cls.PREALLOCATED_INDICES[log2_size].__copy__() + elif log2_size < 16: + indices = array("h", (-1 for i in range(size))) + elif log2_size < 32: + indices = array("l", (-1 for i in range(size))) + else: # pragma: no cover # don't test huge multidicts + indices = array("q", (-1 for i in range(size))) + ret = cls( + log2_size=log2_size, + usable=usable, + indices=indices, + entries=entries, + ) + return ret + + def clone(self) -> "_HtKeys[_V]": + entries = [ + _Entry(e.hash, e.identity, e.key, e.value) if e is not None else None + for e in self.entries + ] + + return _HtKeys( + log2_size=self.log2_size, + usable=self.usable, + indices=self.indices.__copy__(), + entries=entries, + ) + + def build_indices(self, update: bool) -> None: + mask = self.mask + indices = self.indices + for idx, e in enumerate(self.entries): + assert e is not None + hash_ = e.hash + if update: + if hash_ == -1: + hash_ = hash(e.identity) + else: + assert hash_ != -1 + i = hash_ & mask + perturb = hash_ & sys.maxsize + while indices[i] != -1: + perturb >>= 5 + i = mask & (i * 5 + perturb + 1) + indices[i] = idx + + def find_empty_slot(self, hash_: int) -> int: + mask = self.mask + indices = self.indices + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + return i + + def iter_hash(self, hash_: int) -> Iterator[tuple[int, int, _Entry[_V]]]: + mask = self.mask + indices = self.indices + entries = self.entries + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + if ix != -2: + e = entries[ix] + if e.hash == hash_: + yield i, ix, e + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + + def del_idx(self, hash_: int, idx: int) -> None: + mask = self.mask + indices = self.indices + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != idx: + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + indices[i] = -2 + + def iter_entries(self) -> Iterator[_Entry[_V]]: + return filter(None, self.entries) + + def restore_hash(self, hash_: int) -> None: + mask = self.mask + indices = self.indices + entries = self.entries + i = hash_ & mask + perturb = hash_ & sys.maxsize + ix = indices[i] + while ix != -1: + if ix != -2: + entry = entries[ix] + if entry.hash == -1: + entry.hash = hash_ + perturb >>= 5 + i = (i * 5 + perturb + 1) & mask + ix = indices[i] + + +class MultiDict(_CSMixin, MutableMultiMapping[_V]): + """Dictionary with the support for duplicate keys.""" + + __slots__ = ("_keys", "_used", "_version") + + def __init__(self, arg: MDArg[_V] = None, /, **kwargs: _V): + self._used = 0 + v = _version + v[0] += 1 + self._version = v[0] + if not kwargs: + md = None + if isinstance(arg, MultiDictProxy): + md = arg._md + elif isinstance(arg, MultiDict): + md = arg + if md is not None and md._ci is self._ci: + self._from_md(md) + return + + it = self._parse_args(arg, kwargs) + log2_size = estimate_log2_keysize(cast(int, next(it))) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in init + log2_size = 17 + self._keys: _HtKeys[_V] = _HtKeys.new(log2_size, []) + self._extend_items(cast(Iterator[_Entry[_V]], it)) + + def _from_md(self, md: "MultiDict[_V]") -> None: + # Copy everything as-is without compacting the new multidict, + # otherwise it requires reindexing + self._keys = md._keys.clone() + self._used = md._used + + @overload + def getall(self, key: str) -> list[_V]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def getall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Return a list of all values matching the key.""" + identity = self._identity(key) + hash_ = hash(identity) + res = [] + restore = [] + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + res.append(e.value) + e.hash = -1 + restore.append(idx) + + if res: + entries = self._keys.entries + for idx in restore: + entries[idx].hash = hash_ # type: ignore[union-attr] + return res + if not res and default is not sentinel: + return default + raise KeyError("Key not found: %r" % key) + + @overload + def getone(self, key: str) -> _V: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V, _T]: ... + def getone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Get first value matching the key. + + Raises KeyError if the key is not found and no default is provided. + """ + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return e.value + if default is not sentinel: + return default + raise KeyError("Key not found: %r" % key) + + # Mapping interface # + + def __getitem__(self, key: str) -> _V: + return self.getone(key) + + @overload + def get(self, key: str, /) -> Union[_V, None]: ... + @overload + def get(self, key: str, /, default: _T) -> Union[_V, _T]: ... + def get(self, key: str, default: Union[_T, None] = None) -> Union[_V, _T, None]: + """Get first value matching the key. + + If the key is not found, returns the default (or None if no default is provided) + """ + return self.getone(key, default) + + def __iter__(self) -> Iterator[str]: + return iter(self.keys()) + + def __len__(self) -> int: + return self._used + + def keys(self) -> KeysView[str]: + """Return a new view of the dictionary's keys.""" + return _KeysView(self) + + def items(self) -> ItemsView[str, _V]: + """Return a new view of the dictionary's items *(key, value) pairs).""" + return _ItemsView(self) + + def values(self) -> _ValuesView[_V]: + """Return a new view of the dictionary's values.""" + return _ValuesView(self) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Mapping): + return NotImplemented + if isinstance(other, MultiDictProxy): + return self == other._md + if isinstance(other, MultiDict): + lft = self._keys + rht = other._keys + if self._used != other._used: + return False + for e1, e2 in zip(lft.iter_entries(), rht.iter_entries()): + if e1.identity != e2.identity or e1.value != e2.value: + return False + return True + if self._used != len(other): + return False + for k, v in self.items(): + nv = other.get(k, sentinel) + if v != nv: + return False + return True + + def __contains__(self, key: object) -> bool: + if not isinstance(key, str): + return False + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return True + return False + + @reprlib.recursive_repr() + def __repr__(self) -> str: + body = ", ".join(f"'{e.key}': {e.value!r}" for e in self._keys.iter_entries()) + return f"<{self.__class__.__name__}({body})>" + + if sys.implementation.name != "pypy": + + def __sizeof__(self) -> int: + return object.__sizeof__(self) + sys.getsizeof(self._keys) + + def __reduce__(self) -> tuple[type[Self], tuple[list[tuple[str, _V]]]]: + return (self.__class__, (list(self.items()),)) + + def add(self, key: str, value: _V) -> None: + identity = self._identity(key) + hash_ = hash(identity) + self._add_with_hash(_Entry(hash_, identity, key, value)) + self._incr_version() + + def copy(self) -> Self: + """Return a copy of itself.""" + cls = self.__class__ + return cls(self) + + __copy__ = copy + + def extend(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Extend current MultiDict with more values. + + This method must be used instead of update. + """ + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + self._resize(estimate_log2_keysize(newsize), False) + self._extend_items(cast(Iterator[_Entry[_V]], it)) + + def _parse_args( + self, + arg: MDArg[_V], + kwargs: Mapping[str, _V], + ) -> Iterator[Union[int, _Entry[_V]]]: + identity_func = self._identity + if arg: + if isinstance(arg, MultiDictProxy): + arg = arg._md + if isinstance(arg, MultiDict): + yield len(arg) + len(kwargs) + if self._ci is not arg._ci: + for e in arg._keys.iter_entries(): + identity = identity_func(e.key) + yield _Entry(hash(identity), identity, e.key, e.value) + else: + for e in arg._keys.iter_entries(): + yield _Entry(e.hash, e.identity, e.key, e.value) + if kwargs: + for key, value in kwargs.items(): + identity = identity_func(key) + yield _Entry(hash(identity), identity, key, value) + else: + if hasattr(arg, "keys"): + arg = cast(SupportsKeys[_V], arg) + arg = [(k, arg[k]) for k in arg.keys()] + if kwargs: + arg = list(arg) + arg.extend(list(kwargs.items())) + try: + yield len(arg) + len(kwargs) # type: ignore[arg-type] + except TypeError: + yield 0 + for pos, item in enumerate(arg): + if not len(item) == 2: + raise ValueError( + f"multidict update sequence element #{pos}" + f"has length {len(item)}; 2 is required" + ) + identity = identity_func(item[0]) + yield _Entry(hash(identity), identity, item[0], item[1]) + else: + yield len(kwargs) + for key, value in kwargs.items(): + identity = identity_func(key) + yield _Entry(hash(identity), identity, key, value) + + def _extend_items(self, items: Iterable[_Entry[_V]]) -> None: + for e in items: + self._add_with_hash(e) + self._incr_version() + + def clear(self) -> None: + """Remove all items from MultiDict.""" + self._used = 0 + self._keys = _HtKeys.new(_HtKeys.LOG_MINSIZE, []) + self._incr_version() + + # Mapping interface # + + def __setitem__(self, key: str, value: _V) -> None: + identity = self._identity(key) + hash_ = hash(identity) + found = False + + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + if not found: + e.key = key + e.value = value + e.hash = -1 + found = True + self._incr_version() + elif e.hash != -1: # pragma: no branch + self._del_at(slot, idx) + + if not found: + self._add_with_hash(_Entry(hash_, identity, key, value)) + else: + self._keys.restore_hash(hash_) + + def __delitem__(self, key: str) -> None: + found = False + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + self._del_at(slot, idx) + found = True + if not found: + raise KeyError(key) + else: + self._incr_version() + + @overload + def setdefault( + self: "MultiDict[Union[_T, None]]", key: str, default: None = None + ) -> Union[_T, None]: ... + @overload + def setdefault(self, key: str, default: _V) -> _V: ... + def setdefault(self, key: str, default: Union[_V, None] = None) -> Union[_V, None]: # type: ignore[misc] + """Return value for key, set value to default if key is not present.""" + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + return e.value + self.add(key, default) # type: ignore[arg-type] + return default + + @overload + def popone(self, key: str) -> _V: ... + @overload + def popone(self, key: str, default: _T) -> Union[_V, _T]: ... + def popone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Remove specified key and return the corresponding value. + + If key is not found, d is returned if given, otherwise + KeyError is raised. + + """ + identity = self._identity(key) + hash_ = hash(identity) + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + value = e.value + self._del_at(slot, idx) + self._incr_version() + return value + if default is sentinel: + raise KeyError(key) + else: + return default + + # Type checking will inherit signature for pop() if we don't confuse it here. + if not TYPE_CHECKING: + pop = popone + + @overload + def popall(self, key: str) -> list[_V]: ... + @overload + def popall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def popall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Remove all occurrences of key and return the list of corresponding + values. + + If key is not found, default is returned if given, otherwise + KeyError is raised. + + """ + found = False + identity = self._identity(key) + hash_ = hash(identity) + ret = [] + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + found = True + ret.append(e.value) + self._del_at(slot, idx) + self._incr_version() + + if not found: + if default is sentinel: + raise KeyError(key) + else: + return default + else: + return ret + + def popitem(self) -> tuple[str, _V]: + """Remove and return an arbitrary (key, value) pair.""" + if self._used <= 0: + raise KeyError("empty multidict") + + pos = len(self._keys.entries) - 1 + entry = self._keys.entries.pop() + + while entry is None: + pos -= 1 + entry = self._keys.entries.pop() + + ret = self._key(entry.key), entry.value + self._keys.del_idx(entry.hash, pos) + self._used -= 1 + self._incr_version() + return ret + + def update(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Update the dictionary, overwriting existing keys.""" + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + log2_size = estimate_log2_keysize(newsize) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in update, + # duplicate keys could reduce the resulting anount of entries + log2_size = 17 + if log2_size > self._keys.log2_size: + self._resize(log2_size, False) + try: + self._update_items(cast(Iterator[_Entry[_V]], it)) + finally: + self._post_update() + + def _update_items(self, items: Iterator[_Entry[_V]]) -> None: + for entry in items: + found = False + hash_ = entry.hash + identity = entry.identity + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + if not found: + found = True + e.key = entry.key + e.value = entry.value + e.hash = -1 + else: + self._del_at_for_upd(e) + if not found: + self._add_with_hash_for_upd(entry) + + def _post_update(self) -> None: + keys = self._keys + indices = keys.indices + entries = keys.entries + for slot in range(keys.nslots): + idx = indices[slot] + if idx >= 0: + e2 = entries[idx] + assert e2 is not None + if e2.key is None: + entries[idx] = None + indices[slot] = -2 + self._used -= 1 + if e2.hash == -1: + e2.hash = hash(e2.identity) + + self._incr_version() + + def merge(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None: + """Merge into the dictionary, adding non-existing keys.""" + it = self._parse_args(arg, kwargs) + newsize = self._used + cast(int, next(it)) + log2_size = estimate_log2_keysize(newsize) + if log2_size > 17: # pragma: no cover + # Don't overallocate really huge keys space in update, + # duplicate keys could reduce the resulting anount of entries + log2_size = 17 + if log2_size > self._keys.log2_size: + self._resize(log2_size, False) + try: + self._merge_items(cast(Iterator[_Entry[_V]], it)) + finally: + self._post_update() + + def _merge_items(self, items: Iterator[_Entry[_V]]) -> None: + for entry in items: + hash_ = entry.hash + identity = entry.identity + for slot, idx, e in self._keys.iter_hash(hash_): + if e.identity == identity: # pragma: no branch + break + else: + self._add_with_hash_for_upd(entry) + + def _incr_version(self) -> None: + v = _version + v[0] += 1 + self._version = v[0] + + def _resize(self, log2_newsize: int, update: bool) -> None: + oldkeys = self._keys + newentries = self._used + + if len(oldkeys.entries) == newentries: + entries = oldkeys.entries + else: + entries = [e for e in oldkeys.entries if e is not None] + newkeys: _HtKeys[_V] = _HtKeys.new(log2_newsize, entries) + newkeys.usable -= newentries + newkeys.build_indices(update) + self._keys = newkeys + + def _add_with_hash(self, entry: _Entry[_V]) -> None: + if self._keys.usable <= 0: + self._resize((self._used * 3 | _HtKeys.MINSIZE - 1).bit_length(), False) + keys = self._keys + slot = keys.find_empty_slot(entry.hash) + keys.indices[slot] = len(keys.entries) + keys.entries.append(entry) + self._incr_version() + self._used += 1 + keys.usable -= 1 + + def _add_with_hash_for_upd(self, entry: _Entry[_V]) -> None: + if self._keys.usable <= 0: + self._resize((self._used * 3 | _HtKeys.MINSIZE - 1).bit_length(), True) + keys = self._keys + slot = keys.find_empty_slot(entry.hash) + keys.indices[slot] = len(keys.entries) + entry.hash = -1 + keys.entries.append(entry) + self._incr_version() + self._used += 1 + keys.usable -= 1 + + def _del_at(self, slot: int, idx: int) -> None: + self._keys.entries[idx] = None + self._keys.indices[slot] = -2 + self._used -= 1 + + def _del_at_for_upd(self, entry: _Entry[_V]) -> None: + entry.key = None # type: ignore[assignment] + entry.value = None # type: ignore[assignment] + + +class CIMultiDict(_CIMixin, MultiDict[_V]): + """Dictionary with the support for duplicate case-insensitive keys.""" + + +class MultiDictProxy(_CSMixin, MultiMapping[_V]): + """Read-only proxy for MultiDict instance.""" + + __slots__ = ("_md",) + + _md: MultiDict[_V] + + def __init__(self, arg: Union[MultiDict[_V], "MultiDictProxy[_V]"]): + if not isinstance(arg, (MultiDict, MultiDictProxy)): + raise TypeError( + f"ctor requires MultiDict or MultiDictProxy instance, not {type(arg)}" + ) + if isinstance(arg, MultiDictProxy): + self._md = arg._md + else: + self._md = arg + + def __reduce__(self) -> NoReturn: + raise TypeError(f"can't pickle {self.__class__.__name__} objects") + + @overload + def getall(self, key: str) -> list[_V]: ... + @overload + def getall(self, key: str, default: _T) -> Union[list[_V], _T]: ... + def getall( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[list[_V], _T]: + """Return a list of all values matching the key.""" + if default is not sentinel: + return self._md.getall(key, default) + else: + return self._md.getall(key) + + @overload + def getone(self, key: str) -> _V: ... + @overload + def getone(self, key: str, default: _T) -> Union[_V, _T]: ... + def getone( + self, key: str, default: Union[_T, _SENTINEL] = sentinel + ) -> Union[_V, _T]: + """Get first value matching the key. + + Raises KeyError if the key is not found and no default is provided. + """ + if default is not sentinel: + return self._md.getone(key, default) + else: + return self._md.getone(key) + + # Mapping interface # + + def __getitem__(self, key: str) -> _V: + return self.getone(key) + + @overload + def get(self, key: str, /) -> Union[_V, None]: ... + @overload + def get(self, key: str, /, default: _T) -> Union[_V, _T]: ... + def get(self, key: str, default: Union[_T, None] = None) -> Union[_V, _T, None]: + """Get first value matching the key. + + If the key is not found, returns the default (or None if no default is provided) + """ + return self._md.getone(key, default) + + def __iter__(self) -> Iterator[str]: + return iter(self._md.keys()) + + def __len__(self) -> int: + return len(self._md) + + def keys(self) -> KeysView[str]: + """Return a new view of the dictionary's keys.""" + return self._md.keys() + + def items(self) -> ItemsView[str, _V]: + """Return a new view of the dictionary's items *(key, value) pairs).""" + return self._md.items() + + def values(self) -> _ValuesView[_V]: + """Return a new view of the dictionary's values.""" + return self._md.values() + + def __eq__(self, other: object) -> bool: + return self._md == other + + def __contains__(self, key: object) -> bool: + return key in self._md + + @reprlib.recursive_repr() + def __repr__(self) -> str: + body = ", ".join(f"'{k}': {v!r}" for k, v in self.items()) + return f"<{self.__class__.__name__}({body})>" + + def copy(self) -> MultiDict[_V]: + """Return a copy of itself.""" + return MultiDict(self._md) + + +class CIMultiDictProxy(_CIMixin, MultiDictProxy[_V]): + """Read-only proxy for CIMultiDict instance.""" + + def __init__(self, arg: Union[MultiDict[_V], MultiDictProxy[_V]]): + if not isinstance(arg, (CIMultiDict, CIMultiDictProxy)): + raise TypeError( + "ctor requires CIMultiDict or CIMultiDictProxy instance" + f", not {type(arg)}" + ) + + super().__init__(arg) + + def copy(self) -> CIMultiDict[_V]: + """Return a copy of itself.""" + return CIMultiDict(self._md) + + +def getversion(md: Union[MultiDict[object], MultiDictProxy[object]]) -> int: + if isinstance(md, MultiDictProxy): + md = md._md + elif not isinstance(md, MultiDict): + raise TypeError("Parameter should be multidict or proxy") + return md._version diff --git a/venv/lib/python3.10/site-packages/multidict/py.typed b/venv/lib/python3.10/site-packages/multidict/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..dfe8cc048e7100a97025b954fffa31e4ff859a7d --- /dev/null +++ b/venv/lib/python3.10/site-packages/multidict/py.typed @@ -0,0 +1 @@ +PEP-561 marker. \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/__config__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/__config__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2be3f45242c19c44746d0ecc9ec51cb727abbe8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/__config__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf88b3f4f824f9bc8a9e0253fec9772519c58d13 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_array_api_info.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_array_api_info.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60ec2f00fbfbb1911894d466e39fd6b3f770165f Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_array_api_info.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_configtool.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_configtool.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..918aa27a90f8bf90aed88c848671b61e235ad7ab Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_configtool.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_distributor_init.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_distributor_init.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f76986893e23e9e598b89af567080bf308036dc4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_distributor_init.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_expired_attrs_2_0.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_expired_attrs_2_0.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c2a28327dca6267b9808ee7037eddb21694dbfd Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_expired_attrs_2_0.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_globals.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_globals.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13bcdf43d3fb5b9ee619f8384dca195430567c37 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_globals.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/_pytesttester.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/_pytesttester.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76db8254dd5be345a9634e91f43a01c6a75e1900 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/_pytesttester.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/conftest.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/conftest.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a441b1063282a016143181daf6c16fc24cbb166c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/conftest.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/ctypeslib.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/ctypeslib.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92362ae0b99e7da9122938db1bd5e0dced2fe679 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/ctypeslib.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/dtypes.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/dtypes.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06cba2041b1da1a7468098b73bc19baed2c1d743 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/dtypes.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/exceptions.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/exceptions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c1def268159e9c00175fffc04e0745a0a6b3639 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/exceptions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/matlib.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/matlib.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68375533fe80dcd596e7e1a51a2199ac04cba97f Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/matlib.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/__pycache__/version.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/__pycache__/version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c40a44541e073b7f1bdde526e5e79bccd3ca169a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/__pycache__/version.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/__pycache__/_ufunc_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/_core/__pycache__/_ufunc_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3eba6542b58be5c8ce1b2c7fc9189704dd575b0d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/_core/__pycache__/_ufunc_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.c b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.c new file mode 100644 index 0000000000000000000000000000000000000000..c1a152704fa18c79f928b8fa6644f362579383f1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.c @@ -0,0 +1,376 @@ + +/* These pointers will be stored in the C-object for use in other + extension modules +*/ + +void *PyArray_API[] = { + (void *) PyArray_GetNDArrayCVersion, + NULL, + (void *) &PyArray_Type, + (void *) &PyArrayDescr_Type, + NULL, + (void *) &PyArrayIter_Type, + (void *) &PyArrayMultiIter_Type, + (int *) &NPY_NUMUSERTYPES, + (void *) &PyBoolArrType_Type, + (void *) &_PyArrayScalar_BoolValues, + (void *) &PyGenericArrType_Type, + (void *) &PyNumberArrType_Type, + (void *) &PyIntegerArrType_Type, + (void *) &PySignedIntegerArrType_Type, + (void *) &PyUnsignedIntegerArrType_Type, + (void *) &PyInexactArrType_Type, + (void *) &PyFloatingArrType_Type, + (void *) &PyComplexFloatingArrType_Type, + (void *) &PyFlexibleArrType_Type, + (void *) &PyCharacterArrType_Type, + (void *) &PyByteArrType_Type, + (void *) &PyShortArrType_Type, + (void *) &PyIntArrType_Type, + (void *) &PyLongArrType_Type, + (void *) &PyLongLongArrType_Type, + (void *) &PyUByteArrType_Type, + (void *) &PyUShortArrType_Type, + (void *) &PyUIntArrType_Type, + (void *) &PyULongArrType_Type, + (void *) &PyULongLongArrType_Type, + (void *) &PyFloatArrType_Type, + (void *) &PyDoubleArrType_Type, + (void *) &PyLongDoubleArrType_Type, + (void *) &PyCFloatArrType_Type, + (void *) &PyCDoubleArrType_Type, + (void *) &PyCLongDoubleArrType_Type, + (void *) &PyObjectArrType_Type, + (void *) &PyStringArrType_Type, + (void *) &PyUnicodeArrType_Type, + (void *) &PyVoidArrType_Type, + NULL, + NULL, + (void *) PyArray_INCREF, + (void *) PyArray_XDECREF, + (void *) PyArray_SetStringFunction, + (void *) PyArray_DescrFromType, + (void *) PyArray_TypeObjectFromType, + (void *) PyArray_Zero, + (void *) PyArray_One, + (void *) PyArray_CastToType, + (void *) PyArray_CopyInto, + (void *) PyArray_CopyAnyInto, + (void *) PyArray_CanCastSafely, + (void *) PyArray_CanCastTo, + (void *) PyArray_ObjectType, + (void *) PyArray_DescrFromObject, + (void *) PyArray_ConvertToCommonType, + (void *) PyArray_DescrFromScalar, + (void *) PyArray_DescrFromTypeObject, + (void *) PyArray_Size, + (void *) PyArray_Scalar, + (void *) PyArray_FromScalar, + (void *) PyArray_ScalarAsCtype, + (void *) PyArray_CastScalarToCtype, + (void *) PyArray_CastScalarDirect, + (void *) PyArray_Pack, + NULL, + NULL, + NULL, + (void *) PyArray_FromAny, + (void *) PyArray_EnsureArray, + (void *) PyArray_EnsureAnyArray, + (void *) PyArray_FromFile, + (void *) PyArray_FromString, + (void *) PyArray_FromBuffer, + (void *) PyArray_FromIter, + (void *) PyArray_Return, + (void *) PyArray_GetField, + (void *) PyArray_SetField, + (void *) PyArray_Byteswap, + (void *) PyArray_Resize, + NULL, + NULL, + NULL, + (void *) PyArray_CopyObject, + (void *) PyArray_NewCopy, + (void *) PyArray_ToList, + (void *) PyArray_ToString, + (void *) PyArray_ToFile, + (void *) PyArray_Dump, + (void *) PyArray_Dumps, + (void *) PyArray_ValidType, + (void *) PyArray_UpdateFlags, + (void *) PyArray_New, + (void *) PyArray_NewFromDescr, + (void *) PyArray_DescrNew, + (void *) PyArray_DescrNewFromType, + (void *) PyArray_GetPriority, + (void *) PyArray_IterNew, + (void *) PyArray_MultiIterNew, + (void *) PyArray_PyIntAsInt, + (void *) PyArray_PyIntAsIntp, + (void *) PyArray_Broadcast, + NULL, + (void *) PyArray_FillWithScalar, + (void *) PyArray_CheckStrides, + (void *) PyArray_DescrNewByteorder, + (void *) PyArray_IterAllButAxis, + (void *) PyArray_CheckFromAny, + (void *) PyArray_FromArray, + (void *) PyArray_FromInterface, + (void *) PyArray_FromStructInterface, + (void *) PyArray_FromArrayAttr, + (void *) PyArray_ScalarKind, + (void *) PyArray_CanCoerceScalar, + NULL, + (void *) PyArray_CanCastScalar, + NULL, + (void *) PyArray_RemoveSmallest, + (void *) PyArray_ElementStrides, + (void *) PyArray_Item_INCREF, + (void *) PyArray_Item_XDECREF, + NULL, + (void *) PyArray_Transpose, + (void *) PyArray_TakeFrom, + (void *) PyArray_PutTo, + (void *) PyArray_PutMask, + (void *) PyArray_Repeat, + (void *) PyArray_Choose, + (void *) PyArray_Sort, + (void *) PyArray_ArgSort, + (void *) PyArray_SearchSorted, + (void *) PyArray_ArgMax, + (void *) PyArray_ArgMin, + (void *) PyArray_Reshape, + (void *) PyArray_Newshape, + (void *) PyArray_Squeeze, + (void *) PyArray_View, + (void *) PyArray_SwapAxes, + (void *) PyArray_Max, + (void *) PyArray_Min, + (void *) PyArray_Ptp, + (void *) PyArray_Mean, + (void *) PyArray_Trace, + (void *) PyArray_Diagonal, + (void *) PyArray_Clip, + (void *) PyArray_Conjugate, + (void *) PyArray_Nonzero, + (void *) PyArray_Std, + (void *) PyArray_Sum, + (void *) PyArray_CumSum, + (void *) PyArray_Prod, + (void *) PyArray_CumProd, + (void *) PyArray_All, + (void *) PyArray_Any, + (void *) PyArray_Compress, + (void *) PyArray_Flatten, + (void *) PyArray_Ravel, + (void *) PyArray_MultiplyList, + (void *) PyArray_MultiplyIntList, + (void *) PyArray_GetPtr, + (void *) PyArray_CompareLists, + (void *) PyArray_AsCArray, + NULL, + NULL, + (void *) PyArray_Free, + (void *) PyArray_Converter, + (void *) PyArray_IntpFromSequence, + (void *) PyArray_Concatenate, + (void *) PyArray_InnerProduct, + (void *) PyArray_MatrixProduct, + NULL, + (void *) PyArray_Correlate, + NULL, + (void *) PyArray_DescrConverter, + (void *) PyArray_DescrConverter2, + (void *) PyArray_IntpConverter, + (void *) PyArray_BufferConverter, + (void *) PyArray_AxisConverter, + (void *) PyArray_BoolConverter, + (void *) PyArray_ByteorderConverter, + (void *) PyArray_OrderConverter, + (void *) PyArray_EquivTypes, + (void *) PyArray_Zeros, + (void *) PyArray_Empty, + (void *) PyArray_Where, + (void *) PyArray_Arange, + (void *) PyArray_ArangeObj, + (void *) PyArray_SortkindConverter, + (void *) PyArray_LexSort, + (void *) PyArray_Round, + (void *) PyArray_EquivTypenums, + (void *) PyArray_RegisterDataType, + (void *) PyArray_RegisterCastFunc, + (void *) PyArray_RegisterCanCast, + (void *) PyArray_InitArrFuncs, + (void *) PyArray_IntTupleFromIntp, + NULL, + (void *) PyArray_ClipmodeConverter, + (void *) PyArray_OutputConverter, + (void *) PyArray_BroadcastToShape, + NULL, + NULL, + (void *) PyArray_DescrAlignConverter, + (void *) PyArray_DescrAlignConverter2, + (void *) PyArray_SearchsideConverter, + (void *) PyArray_CheckAxis, + (void *) PyArray_OverflowMultiplyList, + NULL, + (void *) PyArray_MultiIterFromObjects, + (void *) PyArray_GetEndianness, + (void *) PyArray_GetNDArrayCFeatureVersion, + (void *) PyArray_Correlate2, + (void *) PyArray_NeighborhoodIterNew, + (void *) &PyTimeIntegerArrType_Type, + (void *) &PyDatetimeArrType_Type, + (void *) &PyTimedeltaArrType_Type, + (void *) &PyHalfArrType_Type, + (void *) &NpyIter_Type, + NULL, + NULL, + NULL, + NULL, + NULL, + (void *) NpyIter_New, + (void *) NpyIter_MultiNew, + (void *) NpyIter_AdvancedNew, + (void *) NpyIter_Copy, + (void *) NpyIter_Deallocate, + (void *) NpyIter_HasDelayedBufAlloc, + (void *) NpyIter_HasExternalLoop, + (void *) NpyIter_EnableExternalLoop, + (void *) NpyIter_GetInnerStrideArray, + (void *) NpyIter_GetInnerLoopSizePtr, + (void *) NpyIter_Reset, + (void *) NpyIter_ResetBasePointers, + (void *) NpyIter_ResetToIterIndexRange, + (void *) NpyIter_GetNDim, + (void *) NpyIter_GetNOp, + (void *) NpyIter_GetIterNext, + (void *) NpyIter_GetIterSize, + (void *) NpyIter_GetIterIndexRange, + (void *) NpyIter_GetIterIndex, + (void *) NpyIter_GotoIterIndex, + (void *) NpyIter_HasMultiIndex, + (void *) NpyIter_GetShape, + (void *) NpyIter_GetGetMultiIndex, + (void *) NpyIter_GotoMultiIndex, + (void *) NpyIter_RemoveMultiIndex, + (void *) NpyIter_HasIndex, + (void *) NpyIter_IsBuffered, + (void *) NpyIter_IsGrowInner, + (void *) NpyIter_GetBufferSize, + (void *) NpyIter_GetIndexPtr, + (void *) NpyIter_GotoIndex, + (void *) NpyIter_GetDataPtrArray, + (void *) NpyIter_GetDescrArray, + (void *) NpyIter_GetOperandArray, + (void *) NpyIter_GetIterView, + (void *) NpyIter_GetReadFlags, + (void *) NpyIter_GetWriteFlags, + (void *) NpyIter_DebugPrint, + (void *) NpyIter_IterationNeedsAPI, + (void *) NpyIter_GetInnerFixedStrideArray, + (void *) NpyIter_RemoveAxis, + (void *) NpyIter_GetAxisStrideArray, + (void *) NpyIter_RequiresBuffering, + (void *) NpyIter_GetInitialDataPtrArray, + (void *) NpyIter_CreateCompatibleStrides, + (void *) PyArray_CastingConverter, + (void *) PyArray_CountNonzero, + (void *) PyArray_PromoteTypes, + (void *) PyArray_MinScalarType, + (void *) PyArray_ResultType, + (void *) PyArray_CanCastArrayTo, + (void *) PyArray_CanCastTypeTo, + (void *) PyArray_EinsteinSum, + (void *) PyArray_NewLikeArray, + NULL, + (void *) PyArray_ConvertClipmodeSequence, + (void *) PyArray_MatrixProduct2, + (void *) NpyIter_IsFirstVisit, + (void *) PyArray_SetBaseObject, + (void *) PyArray_CreateSortedStridePerm, + (void *) PyArray_RemoveAxesInPlace, + (void *) PyArray_DebugPrint, + (void *) PyArray_FailUnlessWriteable, + (void *) PyArray_SetUpdateIfCopyBase, + (void *) PyDataMem_NEW, + (void *) PyDataMem_FREE, + (void *) PyDataMem_RENEW, + NULL, + (NPY_CASTING *) &NPY_DEFAULT_ASSIGN_CASTING, + NULL, + NULL, + NULL, + (void *) PyArray_Partition, + (void *) PyArray_ArgPartition, + (void *) PyArray_SelectkindConverter, + (void *) PyDataMem_NEW_ZEROED, + (void *) PyArray_CheckAnyScalarExact, + NULL, + (void *) PyArray_ResolveWritebackIfCopy, + (void *) PyArray_SetWritebackIfCopyBase, + (void *) PyDataMem_SetHandler, + (void *) PyDataMem_GetHandler, + (PyObject* *) &PyDataMem_DefaultHandler, + (void *) NpyDatetime_ConvertDatetime64ToDatetimeStruct, + (void *) NpyDatetime_ConvertDatetimeStructToDatetime64, + (void *) NpyDatetime_ConvertPyDateTimeToDatetimeStruct, + (void *) NpyDatetime_GetDatetimeISO8601StrLen, + (void *) NpyDatetime_MakeISO8601Datetime, + (void *) NpyDatetime_ParseISO8601Datetime, + (void *) NpyString_load, + (void *) NpyString_pack, + (void *) NpyString_pack_null, + (void *) NpyString_acquire_allocator, + (void *) NpyString_acquire_allocators, + (void *) NpyString_release_allocator, + (void *) NpyString_release_allocators, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, + (void *) PyArray_GetDefaultDescr, + (void *) PyArrayInitDTypeMeta_FromSpec, + (void *) PyArray_CommonDType, + (void *) PyArray_PromoteDTypeSequence, + (void *) _PyDataType_GetArrFuncs, + NULL, + NULL, + NULL +}; diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.h new file mode 100644 index 0000000000000000000000000000000000000000..cfc3628aa53e41f22d81877df31d40020519bebf --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__multiarray_api.h @@ -0,0 +1,1613 @@ + +#if defined(_MULTIARRAYMODULE) || defined(WITH_CPYCHECKER_STEALS_REFERENCE_TO_ARG_ATTRIBUTE) + +typedef struct { + PyObject_HEAD + npy_bool obval; +} PyBoolScalarObject; + +extern NPY_NO_EXPORT PyTypeObject PyArrayNeighborhoodIter_Type; +extern NPY_NO_EXPORT PyBoolScalarObject _PyArrayScalar_BoolValues[2]; + +NPY_NO_EXPORT unsigned int PyArray_GetNDArrayCVersion \ + (void); +extern NPY_NO_EXPORT PyTypeObject PyArray_Type; + +extern NPY_NO_EXPORT PyArray_DTypeMeta PyArrayDescr_TypeFull; +#define PyArrayDescr_Type (*(PyTypeObject *)(&PyArrayDescr_TypeFull)) + +extern NPY_NO_EXPORT PyTypeObject PyArrayIter_Type; + +extern NPY_NO_EXPORT PyTypeObject PyArrayMultiIter_Type; + +extern NPY_NO_EXPORT int NPY_NUMUSERTYPES; + +extern NPY_NO_EXPORT PyTypeObject PyBoolArrType_Type; + +extern NPY_NO_EXPORT PyBoolScalarObject _PyArrayScalar_BoolValues[2]; + +extern NPY_NO_EXPORT PyTypeObject PyGenericArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyNumberArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyIntegerArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PySignedIntegerArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUnsignedIntegerArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyInexactArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyFloatingArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyComplexFloatingArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyFlexibleArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyCharacterArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyByteArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyShortArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyIntArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyLongArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyLongLongArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUByteArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUShortArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUIntArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyULongArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyULongLongArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyFloatArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyDoubleArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyLongDoubleArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyCFloatArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyCDoubleArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyCLongDoubleArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyObjectArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyStringArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUnicodeArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyVoidArrType_Type; + +NPY_NO_EXPORT int PyArray_INCREF \ + (PyArrayObject *); +NPY_NO_EXPORT int PyArray_XDECREF \ + (PyArrayObject *); +NPY_NO_EXPORT void PyArray_SetStringFunction \ + (PyObject *, int); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrFromType \ + (int); +NPY_NO_EXPORT PyObject * PyArray_TypeObjectFromType \ + (int); +NPY_NO_EXPORT char * PyArray_Zero \ + (PyArrayObject *); +NPY_NO_EXPORT char * PyArray_One \ + (PyArrayObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_CastToType \ + (PyArrayObject *, PyArray_Descr *, int); +NPY_NO_EXPORT int PyArray_CopyInto \ + (PyArrayObject *, PyArrayObject *); +NPY_NO_EXPORT int PyArray_CopyAnyInto \ + (PyArrayObject *, PyArrayObject *); +NPY_NO_EXPORT int PyArray_CanCastSafely \ + (int, int); +NPY_NO_EXPORT npy_bool PyArray_CanCastTo \ + (PyArray_Descr *, PyArray_Descr *); +NPY_NO_EXPORT int PyArray_ObjectType \ + (PyObject *, int); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrFromObject \ + (PyObject *, PyArray_Descr *); +NPY_NO_EXPORT PyArrayObject ** PyArray_ConvertToCommonType \ + (PyObject *, int *); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrFromScalar \ + (PyObject *); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrFromTypeObject \ + (PyObject *); +NPY_NO_EXPORT npy_intp PyArray_Size \ + (PyObject *); +NPY_NO_EXPORT PyObject * PyArray_Scalar \ + (void *, PyArray_Descr *, PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_FromScalar \ + (PyObject *, PyArray_Descr *); +NPY_NO_EXPORT void PyArray_ScalarAsCtype \ + (PyObject *, void *); +NPY_NO_EXPORT int PyArray_CastScalarToCtype \ + (PyObject *, void *, PyArray_Descr *); +NPY_NO_EXPORT int PyArray_CastScalarDirect \ + (PyObject *, PyArray_Descr *, void *, int); +NPY_NO_EXPORT int PyArray_Pack \ + (PyArray_Descr *, void *, PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_FromAny \ + (PyObject *, PyArray_Descr *, int, int, int, PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(1) PyObject * PyArray_EnsureArray \ + (PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(1) PyObject * PyArray_EnsureAnyArray \ + (PyObject *); +NPY_NO_EXPORT PyObject * PyArray_FromFile \ + (FILE *, PyArray_Descr *, npy_intp, char *); +NPY_NO_EXPORT PyObject * PyArray_FromString \ + (char *, npy_intp, PyArray_Descr *, npy_intp, char *); +NPY_NO_EXPORT PyObject * PyArray_FromBuffer \ + (PyObject *, PyArray_Descr *, npy_intp, npy_intp); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_FromIter \ + (PyObject *, PyArray_Descr *, npy_intp); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(1) PyObject * PyArray_Return \ + (PyArrayObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_GetField \ + (PyArrayObject *, PyArray_Descr *, int); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) int PyArray_SetField \ + (PyArrayObject *, PyArray_Descr *, int, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_Byteswap \ + (PyArrayObject *, npy_bool); +NPY_NO_EXPORT PyObject * PyArray_Resize \ + (PyArrayObject *, PyArray_Dims *, int, NPY_ORDER NPY_UNUSED(order)); +NPY_NO_EXPORT int PyArray_CopyObject \ + (PyArrayObject *, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_NewCopy \ + (PyArrayObject *, NPY_ORDER); +NPY_NO_EXPORT PyObject * PyArray_ToList \ + (PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_ToString \ + (PyArrayObject *, NPY_ORDER); +NPY_NO_EXPORT int PyArray_ToFile \ + (PyArrayObject *, FILE *, char *, char *); +NPY_NO_EXPORT int PyArray_Dump \ + (PyObject *, PyObject *, int); +NPY_NO_EXPORT PyObject * PyArray_Dumps \ + (PyObject *, int); +NPY_NO_EXPORT int PyArray_ValidType \ + (int); +NPY_NO_EXPORT void PyArray_UpdateFlags \ + (PyArrayObject *, int); +NPY_NO_EXPORT PyObject * PyArray_New \ + (PyTypeObject *, int, npy_intp const *, int, npy_intp const *, void *, int, int, PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_NewFromDescr \ + (PyTypeObject *, PyArray_Descr *, int, npy_intp const *, npy_intp const *, void *, int, PyObject *); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrNew \ + (PyArray_Descr *); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrNewFromType \ + (int); +NPY_NO_EXPORT double PyArray_GetPriority \ + (PyObject *, double); +NPY_NO_EXPORT PyObject * PyArray_IterNew \ + (PyObject *); +NPY_NO_EXPORT PyObject* PyArray_MultiIterNew \ + (int, ...); +NPY_NO_EXPORT int PyArray_PyIntAsInt \ + (PyObject *); +NPY_NO_EXPORT npy_intp PyArray_PyIntAsIntp \ + (PyObject *); +NPY_NO_EXPORT int PyArray_Broadcast \ + (PyArrayMultiIterObject *); +NPY_NO_EXPORT int PyArray_FillWithScalar \ + (PyArrayObject *, PyObject *); +NPY_NO_EXPORT npy_bool PyArray_CheckStrides \ + (int, int, npy_intp, npy_intp, npy_intp const *, npy_intp const *); +NPY_NO_EXPORT PyArray_Descr * PyArray_DescrNewByteorder \ + (PyArray_Descr *, char); +NPY_NO_EXPORT PyObject * PyArray_IterAllButAxis \ + (PyObject *, int *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_CheckFromAny \ + (PyObject *, PyArray_Descr *, int, int, int, PyObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_FromArray \ + (PyArrayObject *, PyArray_Descr *, int); +NPY_NO_EXPORT PyObject * PyArray_FromInterface \ + (PyObject *); +NPY_NO_EXPORT PyObject * PyArray_FromStructInterface \ + (PyObject *); +NPY_NO_EXPORT PyObject * PyArray_FromArrayAttr \ + (PyObject *, PyArray_Descr *, PyObject *); +NPY_NO_EXPORT NPY_SCALARKIND PyArray_ScalarKind \ + (int, PyArrayObject **); +NPY_NO_EXPORT int PyArray_CanCoerceScalar \ + (int, int, NPY_SCALARKIND); +NPY_NO_EXPORT npy_bool PyArray_CanCastScalar \ + (PyTypeObject *, PyTypeObject *); +NPY_NO_EXPORT int PyArray_RemoveSmallest \ + (PyArrayMultiIterObject *); +NPY_NO_EXPORT int PyArray_ElementStrides \ + (PyObject *); +NPY_NO_EXPORT void PyArray_Item_INCREF \ + (char *, PyArray_Descr *); +NPY_NO_EXPORT void PyArray_Item_XDECREF \ + (char *, PyArray_Descr *); +NPY_NO_EXPORT PyObject * PyArray_Transpose \ + (PyArrayObject *, PyArray_Dims *); +NPY_NO_EXPORT PyObject * PyArray_TakeFrom \ + (PyArrayObject *, PyObject *, int, PyArrayObject *, NPY_CLIPMODE); +NPY_NO_EXPORT PyObject * PyArray_PutTo \ + (PyArrayObject *, PyObject*, PyObject *, NPY_CLIPMODE); +NPY_NO_EXPORT PyObject * PyArray_PutMask \ + (PyArrayObject *, PyObject*, PyObject*); +NPY_NO_EXPORT PyObject * PyArray_Repeat \ + (PyArrayObject *, PyObject *, int); +NPY_NO_EXPORT PyObject * PyArray_Choose \ + (PyArrayObject *, PyObject *, PyArrayObject *, NPY_CLIPMODE); +NPY_NO_EXPORT int PyArray_Sort \ + (PyArrayObject *, int, NPY_SORTKIND); +NPY_NO_EXPORT PyObject * PyArray_ArgSort \ + (PyArrayObject *, int, NPY_SORTKIND); +NPY_NO_EXPORT PyObject * PyArray_SearchSorted \ + (PyArrayObject *, PyObject *, NPY_SEARCHSIDE, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_ArgMax \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_ArgMin \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Reshape \ + (PyArrayObject *, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_Newshape \ + (PyArrayObject *, PyArray_Dims *, NPY_ORDER); +NPY_NO_EXPORT PyObject * PyArray_Squeeze \ + (PyArrayObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) PyObject * PyArray_View \ + (PyArrayObject *, PyArray_Descr *, PyTypeObject *); +NPY_NO_EXPORT PyObject * PyArray_SwapAxes \ + (PyArrayObject *, int, int); +NPY_NO_EXPORT PyObject * PyArray_Max \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Min \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Ptp \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Mean \ + (PyArrayObject *, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Trace \ + (PyArrayObject *, int, int, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Diagonal \ + (PyArrayObject *, int, int, int); +NPY_NO_EXPORT PyObject * PyArray_Clip \ + (PyArrayObject *, PyObject *, PyObject *, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Conjugate \ + (PyArrayObject *, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Nonzero \ + (PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Std \ + (PyArrayObject *, int, int, PyArrayObject *, int); +NPY_NO_EXPORT PyObject * PyArray_Sum \ + (PyArrayObject *, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_CumSum \ + (PyArrayObject *, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Prod \ + (PyArrayObject *, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_CumProd \ + (PyArrayObject *, int, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_All \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Any \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Compress \ + (PyArrayObject *, PyObject *, int, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyArray_Flatten \ + (PyArrayObject *, NPY_ORDER); +NPY_NO_EXPORT PyObject * PyArray_Ravel \ + (PyArrayObject *, NPY_ORDER); +NPY_NO_EXPORT npy_intp PyArray_MultiplyList \ + (npy_intp const *, int); +NPY_NO_EXPORT int PyArray_MultiplyIntList \ + (int const *, int); +NPY_NO_EXPORT void * PyArray_GetPtr \ + (PyArrayObject *, npy_intp const*); +NPY_NO_EXPORT int PyArray_CompareLists \ + (npy_intp const *, npy_intp const *, int); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(5) int PyArray_AsCArray \ + (PyObject **, void *, npy_intp *, int, PyArray_Descr*); +NPY_NO_EXPORT int PyArray_Free \ + (PyObject *, void *); +NPY_NO_EXPORT int PyArray_Converter \ + (PyObject *, PyObject **); +NPY_NO_EXPORT int PyArray_IntpFromSequence \ + (PyObject *, npy_intp *, int); +NPY_NO_EXPORT PyObject * PyArray_Concatenate \ + (PyObject *, int); +NPY_NO_EXPORT PyObject * PyArray_InnerProduct \ + (PyObject *, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_MatrixProduct \ + (PyObject *, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_Correlate \ + (PyObject *, PyObject *, int); +NPY_NO_EXPORT int PyArray_DescrConverter \ + (PyObject *, PyArray_Descr **); +NPY_NO_EXPORT int PyArray_DescrConverter2 \ + (PyObject *, PyArray_Descr **); +NPY_NO_EXPORT int PyArray_IntpConverter \ + (PyObject *, PyArray_Dims *); +NPY_NO_EXPORT int PyArray_BufferConverter \ + (PyObject *, PyArray_Chunk *); +NPY_NO_EXPORT int PyArray_AxisConverter \ + (PyObject *, int *); +NPY_NO_EXPORT int PyArray_BoolConverter \ + (PyObject *, npy_bool *); +NPY_NO_EXPORT int PyArray_ByteorderConverter \ + (PyObject *, char *); +NPY_NO_EXPORT int PyArray_OrderConverter \ + (PyObject *, NPY_ORDER *); +NPY_NO_EXPORT unsigned char PyArray_EquivTypes \ + (PyArray_Descr *, PyArray_Descr *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(3) PyObject * PyArray_Zeros \ + (int, npy_intp const *, PyArray_Descr *, int); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(3) PyObject * PyArray_Empty \ + (int, npy_intp const *, PyArray_Descr *, int); +NPY_NO_EXPORT PyObject * PyArray_Where \ + (PyObject *, PyObject *, PyObject *); +NPY_NO_EXPORT PyObject * PyArray_Arange \ + (double, double, double, int); +NPY_NO_EXPORT PyObject * PyArray_ArangeObj \ + (PyObject *, PyObject *, PyObject *, PyArray_Descr *); +NPY_NO_EXPORT int PyArray_SortkindConverter \ + (PyObject *, NPY_SORTKIND *); +NPY_NO_EXPORT PyObject * PyArray_LexSort \ + (PyObject *, int); +NPY_NO_EXPORT PyObject * PyArray_Round \ + (PyArrayObject *, int, PyArrayObject *); +NPY_NO_EXPORT unsigned char PyArray_EquivTypenums \ + (int, int); +NPY_NO_EXPORT int PyArray_RegisterDataType \ + (PyArray_DescrProto *); +NPY_NO_EXPORT int PyArray_RegisterCastFunc \ + (PyArray_Descr *, int, PyArray_VectorUnaryFunc *); +NPY_NO_EXPORT int PyArray_RegisterCanCast \ + (PyArray_Descr *, int, NPY_SCALARKIND); +NPY_NO_EXPORT void PyArray_InitArrFuncs \ + (PyArray_ArrFuncs *); +NPY_NO_EXPORT PyObject * PyArray_IntTupleFromIntp \ + (int, npy_intp const *); +NPY_NO_EXPORT int PyArray_ClipmodeConverter \ + (PyObject *, NPY_CLIPMODE *); +NPY_NO_EXPORT int PyArray_OutputConverter \ + (PyObject *, PyArrayObject **); +NPY_NO_EXPORT PyObject * PyArray_BroadcastToShape \ + (PyObject *, npy_intp *, int); +NPY_NO_EXPORT int PyArray_DescrAlignConverter \ + (PyObject *, PyArray_Descr **); +NPY_NO_EXPORT int PyArray_DescrAlignConverter2 \ + (PyObject *, PyArray_Descr **); +NPY_NO_EXPORT int PyArray_SearchsideConverter \ + (PyObject *, void *); +NPY_NO_EXPORT PyObject * PyArray_CheckAxis \ + (PyArrayObject *, int *, int); +NPY_NO_EXPORT npy_intp PyArray_OverflowMultiplyList \ + (npy_intp const *, int); +NPY_NO_EXPORT PyObject* PyArray_MultiIterFromObjects \ + (PyObject **, int, int, ...); +NPY_NO_EXPORT int PyArray_GetEndianness \ + (void); +NPY_NO_EXPORT unsigned int PyArray_GetNDArrayCFeatureVersion \ + (void); +NPY_NO_EXPORT PyObject * PyArray_Correlate2 \ + (PyObject *, PyObject *, int); +NPY_NO_EXPORT PyObject* PyArray_NeighborhoodIterNew \ + (PyArrayIterObject *, const npy_intp *, int, PyArrayObject*); +extern NPY_NO_EXPORT PyTypeObject PyTimeIntegerArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyDatetimeArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyTimedeltaArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject PyHalfArrType_Type; + +extern NPY_NO_EXPORT PyTypeObject NpyIter_Type; + +NPY_NO_EXPORT NpyIter * NpyIter_New \ + (PyArrayObject *, npy_uint32, NPY_ORDER, NPY_CASTING, PyArray_Descr*); +NPY_NO_EXPORT NpyIter * NpyIter_MultiNew \ + (int, PyArrayObject **, npy_uint32, NPY_ORDER, NPY_CASTING, npy_uint32 *, PyArray_Descr **); +NPY_NO_EXPORT NpyIter * NpyIter_AdvancedNew \ + (int, PyArrayObject **, npy_uint32, NPY_ORDER, NPY_CASTING, npy_uint32 *, PyArray_Descr **, int, int **, npy_intp *, npy_intp); +NPY_NO_EXPORT NpyIter * NpyIter_Copy \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_Deallocate \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_HasDelayedBufAlloc \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_HasExternalLoop \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_EnableExternalLoop \ + (NpyIter *); +NPY_NO_EXPORT npy_intp * NpyIter_GetInnerStrideArray \ + (NpyIter *); +NPY_NO_EXPORT npy_intp * NpyIter_GetInnerLoopSizePtr \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_Reset \ + (NpyIter *, char **); +NPY_NO_EXPORT int NpyIter_ResetBasePointers \ + (NpyIter *, char **, char **); +NPY_NO_EXPORT int NpyIter_ResetToIterIndexRange \ + (NpyIter *, npy_intp, npy_intp, char **); +NPY_NO_EXPORT int NpyIter_GetNDim \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_GetNOp \ + (NpyIter *); +NPY_NO_EXPORT NpyIter_IterNextFunc * NpyIter_GetIterNext \ + (NpyIter *, char **); +NPY_NO_EXPORT npy_intp NpyIter_GetIterSize \ + (NpyIter *); +NPY_NO_EXPORT void NpyIter_GetIterIndexRange \ + (NpyIter *, npy_intp *, npy_intp *); +NPY_NO_EXPORT npy_intp NpyIter_GetIterIndex \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_GotoIterIndex \ + (NpyIter *, npy_intp); +NPY_NO_EXPORT npy_bool NpyIter_HasMultiIndex \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_GetShape \ + (NpyIter *, npy_intp *); +NPY_NO_EXPORT NpyIter_GetMultiIndexFunc * NpyIter_GetGetMultiIndex \ + (NpyIter *, char **); +NPY_NO_EXPORT int NpyIter_GotoMultiIndex \ + (NpyIter *, npy_intp const *); +NPY_NO_EXPORT int NpyIter_RemoveMultiIndex \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_HasIndex \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_IsBuffered \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_IsGrowInner \ + (NpyIter *); +NPY_NO_EXPORT npy_intp NpyIter_GetBufferSize \ + (NpyIter *); +NPY_NO_EXPORT npy_intp * NpyIter_GetIndexPtr \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_GotoIndex \ + (NpyIter *, npy_intp); +NPY_NO_EXPORT char ** NpyIter_GetDataPtrArray \ + (NpyIter *); +NPY_NO_EXPORT PyArray_Descr ** NpyIter_GetDescrArray \ + (NpyIter *); +NPY_NO_EXPORT PyArrayObject ** NpyIter_GetOperandArray \ + (NpyIter *); +NPY_NO_EXPORT PyArrayObject * NpyIter_GetIterView \ + (NpyIter *, npy_intp); +NPY_NO_EXPORT void NpyIter_GetReadFlags \ + (NpyIter *, char *); +NPY_NO_EXPORT void NpyIter_GetWriteFlags \ + (NpyIter *, char *); +NPY_NO_EXPORT void NpyIter_DebugPrint \ + (NpyIter *); +NPY_NO_EXPORT npy_bool NpyIter_IterationNeedsAPI \ + (NpyIter *); +NPY_NO_EXPORT void NpyIter_GetInnerFixedStrideArray \ + (NpyIter *, npy_intp *); +NPY_NO_EXPORT int NpyIter_RemoveAxis \ + (NpyIter *, int); +NPY_NO_EXPORT npy_intp * NpyIter_GetAxisStrideArray \ + (NpyIter *, int); +NPY_NO_EXPORT npy_bool NpyIter_RequiresBuffering \ + (NpyIter *); +NPY_NO_EXPORT char ** NpyIter_GetInitialDataPtrArray \ + (NpyIter *); +NPY_NO_EXPORT int NpyIter_CreateCompatibleStrides \ + (NpyIter *, npy_intp, npy_intp *); +NPY_NO_EXPORT int PyArray_CastingConverter \ + (PyObject *, NPY_CASTING *); +NPY_NO_EXPORT npy_intp PyArray_CountNonzero \ + (PyArrayObject *); +NPY_NO_EXPORT PyArray_Descr * PyArray_PromoteTypes \ + (PyArray_Descr *, PyArray_Descr *); +NPY_NO_EXPORT PyArray_Descr * PyArray_MinScalarType \ + (PyArrayObject *); +NPY_NO_EXPORT PyArray_Descr * PyArray_ResultType \ + (npy_intp, PyArrayObject *arrs[], npy_intp, PyArray_Descr *descrs[]); +NPY_NO_EXPORT npy_bool PyArray_CanCastArrayTo \ + (PyArrayObject *, PyArray_Descr *, NPY_CASTING); +NPY_NO_EXPORT npy_bool PyArray_CanCastTypeTo \ + (PyArray_Descr *, PyArray_Descr *, NPY_CASTING); +NPY_NO_EXPORT PyArrayObject * PyArray_EinsteinSum \ + (char *, npy_intp, PyArrayObject **, PyArray_Descr *, NPY_ORDER, NPY_CASTING, PyArrayObject *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(3) PyObject * PyArray_NewLikeArray \ + (PyArrayObject *, NPY_ORDER, PyArray_Descr *, int); +NPY_NO_EXPORT int PyArray_ConvertClipmodeSequence \ + (PyObject *, NPY_CLIPMODE *, int); +NPY_NO_EXPORT PyObject * PyArray_MatrixProduct2 \ + (PyObject *, PyObject *, PyArrayObject*); +NPY_NO_EXPORT npy_bool NpyIter_IsFirstVisit \ + (NpyIter *, int); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) int PyArray_SetBaseObject \ + (PyArrayObject *, PyObject *); +NPY_NO_EXPORT void PyArray_CreateSortedStridePerm \ + (int, npy_intp const *, npy_stride_sort_item *); +NPY_NO_EXPORT void PyArray_RemoveAxesInPlace \ + (PyArrayObject *, const npy_bool *); +NPY_NO_EXPORT void PyArray_DebugPrint \ + (PyArrayObject *); +NPY_NO_EXPORT int PyArray_FailUnlessWriteable \ + (PyArrayObject *, const char *); +NPY_NO_EXPORT NPY_STEALS_REF_TO_ARG(2) int PyArray_SetUpdateIfCopyBase \ + (PyArrayObject *, PyArrayObject *); +NPY_NO_EXPORT void * PyDataMem_NEW \ + (size_t); +NPY_NO_EXPORT void PyDataMem_FREE \ + (void *); +NPY_NO_EXPORT void * PyDataMem_RENEW \ + (void *, size_t); +extern NPY_NO_EXPORT NPY_CASTING NPY_DEFAULT_ASSIGN_CASTING; + +NPY_NO_EXPORT int PyArray_Partition \ + (PyArrayObject *, PyArrayObject *, int, NPY_SELECTKIND); +NPY_NO_EXPORT PyObject * PyArray_ArgPartition \ + (PyArrayObject *, PyArrayObject *, int, NPY_SELECTKIND); +NPY_NO_EXPORT int PyArray_SelectkindConverter \ + (PyObject *, NPY_SELECTKIND *); +NPY_NO_EXPORT void * PyDataMem_NEW_ZEROED \ + (size_t, size_t); +NPY_NO_EXPORT int PyArray_CheckAnyScalarExact \ + (PyObject *); +NPY_NO_EXPORT int PyArray_ResolveWritebackIfCopy \ + (PyArrayObject *); +NPY_NO_EXPORT int PyArray_SetWritebackIfCopyBase \ + (PyArrayObject *, PyArrayObject *); +NPY_NO_EXPORT PyObject * PyDataMem_SetHandler \ + (PyObject *); +NPY_NO_EXPORT PyObject * PyDataMem_GetHandler \ + (void); +extern NPY_NO_EXPORT PyObject* PyDataMem_DefaultHandler; + +NPY_NO_EXPORT int NpyDatetime_ConvertDatetime64ToDatetimeStruct \ + (PyArray_DatetimeMetaData *, npy_datetime, npy_datetimestruct *); +NPY_NO_EXPORT int NpyDatetime_ConvertDatetimeStructToDatetime64 \ + (PyArray_DatetimeMetaData *, const npy_datetimestruct *, npy_datetime *); +NPY_NO_EXPORT int NpyDatetime_ConvertPyDateTimeToDatetimeStruct \ + (PyObject *, npy_datetimestruct *, NPY_DATETIMEUNIT *, int); +NPY_NO_EXPORT int NpyDatetime_GetDatetimeISO8601StrLen \ + (int, NPY_DATETIMEUNIT); +NPY_NO_EXPORT int NpyDatetime_MakeISO8601Datetime \ + (npy_datetimestruct *, char *, npy_intp, int, int, NPY_DATETIMEUNIT, int, NPY_CASTING); +NPY_NO_EXPORT int NpyDatetime_ParseISO8601Datetime \ + (char const *, Py_ssize_t, NPY_DATETIMEUNIT, NPY_CASTING, npy_datetimestruct *, NPY_DATETIMEUNIT *, npy_bool *); +NPY_NO_EXPORT int NpyString_load \ + (npy_string_allocator *, const npy_packed_static_string *, npy_static_string *); +NPY_NO_EXPORT int NpyString_pack \ + (npy_string_allocator *, npy_packed_static_string *, const char *, size_t); +NPY_NO_EXPORT int NpyString_pack_null \ + (npy_string_allocator *, npy_packed_static_string *); +NPY_NO_EXPORT npy_string_allocator * NpyString_acquire_allocator \ + (const PyArray_StringDTypeObject *); +NPY_NO_EXPORT void NpyString_acquire_allocators \ + (size_t, PyArray_Descr *const descrs[], npy_string_allocator *allocators[]); +NPY_NO_EXPORT void NpyString_release_allocator \ + (npy_string_allocator *); +NPY_NO_EXPORT void NpyString_release_allocators \ + (size_t, npy_string_allocator *allocators[]); +NPY_NO_EXPORT PyArray_Descr * PyArray_GetDefaultDescr \ + (PyArray_DTypeMeta *); +NPY_NO_EXPORT int PyArrayInitDTypeMeta_FromSpec \ + (PyArray_DTypeMeta *, PyArrayDTypeMeta_Spec *); +NPY_NO_EXPORT PyArray_DTypeMeta * PyArray_CommonDType \ + (PyArray_DTypeMeta *, PyArray_DTypeMeta *); +NPY_NO_EXPORT PyArray_DTypeMeta * PyArray_PromoteDTypeSequence \ + (npy_intp, PyArray_DTypeMeta **); +NPY_NO_EXPORT PyArray_ArrFuncs * _PyDataType_GetArrFuncs \ + (const PyArray_Descr *); + +#else + +#if defined(PY_ARRAY_UNIQUE_SYMBOL) + #define PyArray_API PY_ARRAY_UNIQUE_SYMBOL + #define _NPY_VERSION_CONCAT_HELPER2(x, y) x ## y + #define _NPY_VERSION_CONCAT_HELPER(arg) \ + _NPY_VERSION_CONCAT_HELPER2(arg, PyArray_RUNTIME_VERSION) + #define PyArray_RUNTIME_VERSION \ + _NPY_VERSION_CONCAT_HELPER(PY_ARRAY_UNIQUE_SYMBOL) +#endif + +/* By default do not export API in an .so (was never the case on windows) */ +#ifndef NPY_API_SYMBOL_ATTRIBUTE + #define NPY_API_SYMBOL_ATTRIBUTE NPY_VISIBILITY_HIDDEN +#endif + +#if defined(NO_IMPORT) || defined(NO_IMPORT_ARRAY) +extern NPY_API_SYMBOL_ATTRIBUTE void **PyArray_API; +extern NPY_API_SYMBOL_ATTRIBUTE int PyArray_RUNTIME_VERSION; +#else +#if defined(PY_ARRAY_UNIQUE_SYMBOL) +NPY_API_SYMBOL_ATTRIBUTE void **PyArray_API; +NPY_API_SYMBOL_ATTRIBUTE int PyArray_RUNTIME_VERSION; +#else +static void **PyArray_API = NULL; +static int PyArray_RUNTIME_VERSION = 0; +#endif +#endif + +#define PyArray_GetNDArrayCVersion \ + (*(unsigned int (*)(void)) \ + PyArray_API[0]) +#define PyArray_Type (*(PyTypeObject *)PyArray_API[2]) +#define PyArrayDescr_Type (*(PyTypeObject *)PyArray_API[3]) +#define PyArrayIter_Type (*(PyTypeObject *)PyArray_API[5]) +#define PyArrayMultiIter_Type (*(PyTypeObject *)PyArray_API[6]) +#define NPY_NUMUSERTYPES (*(int *)PyArray_API[7]) +#define PyBoolArrType_Type (*(PyTypeObject *)PyArray_API[8]) +#define _PyArrayScalar_BoolValues ((PyBoolScalarObject *)PyArray_API[9]) +#define PyGenericArrType_Type (*(PyTypeObject *)PyArray_API[10]) +#define PyNumberArrType_Type (*(PyTypeObject *)PyArray_API[11]) +#define PyIntegerArrType_Type (*(PyTypeObject *)PyArray_API[12]) +#define PySignedIntegerArrType_Type (*(PyTypeObject *)PyArray_API[13]) +#define PyUnsignedIntegerArrType_Type (*(PyTypeObject *)PyArray_API[14]) +#define PyInexactArrType_Type (*(PyTypeObject *)PyArray_API[15]) +#define PyFloatingArrType_Type (*(PyTypeObject *)PyArray_API[16]) +#define PyComplexFloatingArrType_Type (*(PyTypeObject *)PyArray_API[17]) +#define PyFlexibleArrType_Type (*(PyTypeObject *)PyArray_API[18]) +#define PyCharacterArrType_Type (*(PyTypeObject *)PyArray_API[19]) +#define PyByteArrType_Type (*(PyTypeObject *)PyArray_API[20]) +#define PyShortArrType_Type (*(PyTypeObject *)PyArray_API[21]) +#define PyIntArrType_Type (*(PyTypeObject *)PyArray_API[22]) +#define PyLongArrType_Type (*(PyTypeObject *)PyArray_API[23]) +#define PyLongLongArrType_Type (*(PyTypeObject *)PyArray_API[24]) +#define PyUByteArrType_Type (*(PyTypeObject *)PyArray_API[25]) +#define PyUShortArrType_Type (*(PyTypeObject *)PyArray_API[26]) +#define PyUIntArrType_Type (*(PyTypeObject *)PyArray_API[27]) +#define PyULongArrType_Type (*(PyTypeObject *)PyArray_API[28]) +#define PyULongLongArrType_Type (*(PyTypeObject *)PyArray_API[29]) +#define PyFloatArrType_Type (*(PyTypeObject *)PyArray_API[30]) +#define PyDoubleArrType_Type (*(PyTypeObject *)PyArray_API[31]) +#define PyLongDoubleArrType_Type (*(PyTypeObject *)PyArray_API[32]) +#define PyCFloatArrType_Type (*(PyTypeObject *)PyArray_API[33]) +#define PyCDoubleArrType_Type (*(PyTypeObject *)PyArray_API[34]) +#define PyCLongDoubleArrType_Type (*(PyTypeObject *)PyArray_API[35]) +#define PyObjectArrType_Type (*(PyTypeObject *)PyArray_API[36]) +#define PyStringArrType_Type (*(PyTypeObject *)PyArray_API[37]) +#define PyUnicodeArrType_Type (*(PyTypeObject *)PyArray_API[38]) +#define PyVoidArrType_Type (*(PyTypeObject *)PyArray_API[39]) +#define PyArray_INCREF \ + (*(int (*)(PyArrayObject *)) \ + PyArray_API[42]) +#define PyArray_XDECREF \ + (*(int (*)(PyArrayObject *)) \ + PyArray_API[43]) +#define PyArray_SetStringFunction \ + (*(void (*)(PyObject *, int)) \ + PyArray_API[44]) +#define PyArray_DescrFromType \ + (*(PyArray_Descr * (*)(int)) \ + PyArray_API[45]) +#define PyArray_TypeObjectFromType \ + (*(PyObject * (*)(int)) \ + PyArray_API[46]) +#define PyArray_Zero \ + (*(char * (*)(PyArrayObject *)) \ + PyArray_API[47]) +#define PyArray_One \ + (*(char * (*)(PyArrayObject *)) \ + PyArray_API[48]) +#define PyArray_CastToType \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Descr *, int)) \ + PyArray_API[49]) +#define PyArray_CopyInto \ + (*(int (*)(PyArrayObject *, PyArrayObject *)) \ + PyArray_API[50]) +#define PyArray_CopyAnyInto \ + (*(int (*)(PyArrayObject *, PyArrayObject *)) \ + PyArray_API[51]) +#define PyArray_CanCastSafely \ + (*(int (*)(int, int)) \ + PyArray_API[52]) +#define PyArray_CanCastTo \ + (*(npy_bool (*)(PyArray_Descr *, PyArray_Descr *)) \ + PyArray_API[53]) +#define PyArray_ObjectType \ + (*(int (*)(PyObject *, int)) \ + PyArray_API[54]) +#define PyArray_DescrFromObject \ + (*(PyArray_Descr * (*)(PyObject *, PyArray_Descr *)) \ + PyArray_API[55]) +#define PyArray_ConvertToCommonType \ + (*(PyArrayObject ** (*)(PyObject *, int *)) \ + PyArray_API[56]) +#define PyArray_DescrFromScalar \ + (*(PyArray_Descr * (*)(PyObject *)) \ + PyArray_API[57]) +#define PyArray_DescrFromTypeObject \ + (*(PyArray_Descr * (*)(PyObject *)) \ + PyArray_API[58]) +#define PyArray_Size \ + (*(npy_intp (*)(PyObject *)) \ + PyArray_API[59]) +#define PyArray_Scalar \ + (*(PyObject * (*)(void *, PyArray_Descr *, PyObject *)) \ + PyArray_API[60]) +#define PyArray_FromScalar \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *)) \ + PyArray_API[61]) +#define PyArray_ScalarAsCtype \ + (*(void (*)(PyObject *, void *)) \ + PyArray_API[62]) +#define PyArray_CastScalarToCtype \ + (*(int (*)(PyObject *, void *, PyArray_Descr *)) \ + PyArray_API[63]) +#define PyArray_CastScalarDirect \ + (*(int (*)(PyObject *, PyArray_Descr *, void *, int)) \ + PyArray_API[64]) + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyArray_Pack \ + (*(int (*)(PyArray_Descr *, void *, PyObject *)) \ + PyArray_API[65]) +#endif +#define PyArray_FromAny \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *, int, int, int, PyObject *)) \ + PyArray_API[69]) +#define PyArray_EnsureArray \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[70]) +#define PyArray_EnsureAnyArray \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[71]) +#define PyArray_FromFile \ + (*(PyObject * (*)(FILE *, PyArray_Descr *, npy_intp, char *)) \ + PyArray_API[72]) +#define PyArray_FromString \ + (*(PyObject * (*)(char *, npy_intp, PyArray_Descr *, npy_intp, char *)) \ + PyArray_API[73]) +#define PyArray_FromBuffer \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *, npy_intp, npy_intp)) \ + PyArray_API[74]) +#define PyArray_FromIter \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *, npy_intp)) \ + PyArray_API[75]) +#define PyArray_Return \ + (*(PyObject * (*)(PyArrayObject *)) \ + PyArray_API[76]) +#define PyArray_GetField \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Descr *, int)) \ + PyArray_API[77]) +#define PyArray_SetField \ + (*(int (*)(PyArrayObject *, PyArray_Descr *, int, PyObject *)) \ + PyArray_API[78]) +#define PyArray_Byteswap \ + (*(PyObject * (*)(PyArrayObject *, npy_bool)) \ + PyArray_API[79]) +#define PyArray_Resize \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Dims *, int, NPY_ORDER NPY_UNUSED(order))) \ + PyArray_API[80]) +#define PyArray_CopyObject \ + (*(int (*)(PyArrayObject *, PyObject *)) \ + PyArray_API[84]) +#define PyArray_NewCopy \ + (*(PyObject * (*)(PyArrayObject *, NPY_ORDER)) \ + PyArray_API[85]) +#define PyArray_ToList \ + (*(PyObject * (*)(PyArrayObject *)) \ + PyArray_API[86]) +#define PyArray_ToString \ + (*(PyObject * (*)(PyArrayObject *, NPY_ORDER)) \ + PyArray_API[87]) +#define PyArray_ToFile \ + (*(int (*)(PyArrayObject *, FILE *, char *, char *)) \ + PyArray_API[88]) +#define PyArray_Dump \ + (*(int (*)(PyObject *, PyObject *, int)) \ + PyArray_API[89]) +#define PyArray_Dumps \ + (*(PyObject * (*)(PyObject *, int)) \ + PyArray_API[90]) +#define PyArray_ValidType \ + (*(int (*)(int)) \ + PyArray_API[91]) +#define PyArray_UpdateFlags \ + (*(void (*)(PyArrayObject *, int)) \ + PyArray_API[92]) +#define PyArray_New \ + (*(PyObject * (*)(PyTypeObject *, int, npy_intp const *, int, npy_intp const *, void *, int, int, PyObject *)) \ + PyArray_API[93]) +#define PyArray_NewFromDescr \ + (*(PyObject * (*)(PyTypeObject *, PyArray_Descr *, int, npy_intp const *, npy_intp const *, void *, int, PyObject *)) \ + PyArray_API[94]) +#define PyArray_DescrNew \ + (*(PyArray_Descr * (*)(PyArray_Descr *)) \ + PyArray_API[95]) +#define PyArray_DescrNewFromType \ + (*(PyArray_Descr * (*)(int)) \ + PyArray_API[96]) +#define PyArray_GetPriority \ + (*(double (*)(PyObject *, double)) \ + PyArray_API[97]) +#define PyArray_IterNew \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[98]) +#define PyArray_MultiIterNew \ + (*(PyObject* (*)(int, ...)) \ + PyArray_API[99]) +#define PyArray_PyIntAsInt \ + (*(int (*)(PyObject *)) \ + PyArray_API[100]) +#define PyArray_PyIntAsIntp \ + (*(npy_intp (*)(PyObject *)) \ + PyArray_API[101]) +#define PyArray_Broadcast \ + (*(int (*)(PyArrayMultiIterObject *)) \ + PyArray_API[102]) +#define PyArray_FillWithScalar \ + (*(int (*)(PyArrayObject *, PyObject *)) \ + PyArray_API[104]) +#define PyArray_CheckStrides \ + (*(npy_bool (*)(int, int, npy_intp, npy_intp, npy_intp const *, npy_intp const *)) \ + PyArray_API[105]) +#define PyArray_DescrNewByteorder \ + (*(PyArray_Descr * (*)(PyArray_Descr *, char)) \ + PyArray_API[106]) +#define PyArray_IterAllButAxis \ + (*(PyObject * (*)(PyObject *, int *)) \ + PyArray_API[107]) +#define PyArray_CheckFromAny \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *, int, int, int, PyObject *)) \ + PyArray_API[108]) +#define PyArray_FromArray \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Descr *, int)) \ + PyArray_API[109]) +#define PyArray_FromInterface \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[110]) +#define PyArray_FromStructInterface \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[111]) +#define PyArray_FromArrayAttr \ + (*(PyObject * (*)(PyObject *, PyArray_Descr *, PyObject *)) \ + PyArray_API[112]) +#define PyArray_ScalarKind \ + (*(NPY_SCALARKIND (*)(int, PyArrayObject **)) \ + PyArray_API[113]) +#define PyArray_CanCoerceScalar \ + (*(int (*)(int, int, NPY_SCALARKIND)) \ + PyArray_API[114]) +#define PyArray_CanCastScalar \ + (*(npy_bool (*)(PyTypeObject *, PyTypeObject *)) \ + PyArray_API[116]) +#define PyArray_RemoveSmallest \ + (*(int (*)(PyArrayMultiIterObject *)) \ + PyArray_API[118]) +#define PyArray_ElementStrides \ + (*(int (*)(PyObject *)) \ + PyArray_API[119]) +#define PyArray_Item_INCREF \ + (*(void (*)(char *, PyArray_Descr *)) \ + PyArray_API[120]) +#define PyArray_Item_XDECREF \ + (*(void (*)(char *, PyArray_Descr *)) \ + PyArray_API[121]) +#define PyArray_Transpose \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Dims *)) \ + PyArray_API[123]) +#define PyArray_TakeFrom \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, int, PyArrayObject *, NPY_CLIPMODE)) \ + PyArray_API[124]) +#define PyArray_PutTo \ + (*(PyObject * (*)(PyArrayObject *, PyObject*, PyObject *, NPY_CLIPMODE)) \ + PyArray_API[125]) +#define PyArray_PutMask \ + (*(PyObject * (*)(PyArrayObject *, PyObject*, PyObject*)) \ + PyArray_API[126]) +#define PyArray_Repeat \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, int)) \ + PyArray_API[127]) +#define PyArray_Choose \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, PyArrayObject *, NPY_CLIPMODE)) \ + PyArray_API[128]) +#define PyArray_Sort \ + (*(int (*)(PyArrayObject *, int, NPY_SORTKIND)) \ + PyArray_API[129]) +#define PyArray_ArgSort \ + (*(PyObject * (*)(PyArrayObject *, int, NPY_SORTKIND)) \ + PyArray_API[130]) +#define PyArray_SearchSorted \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, NPY_SEARCHSIDE, PyObject *)) \ + PyArray_API[131]) +#define PyArray_ArgMax \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[132]) +#define PyArray_ArgMin \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[133]) +#define PyArray_Reshape \ + (*(PyObject * (*)(PyArrayObject *, PyObject *)) \ + PyArray_API[134]) +#define PyArray_Newshape \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Dims *, NPY_ORDER)) \ + PyArray_API[135]) +#define PyArray_Squeeze \ + (*(PyObject * (*)(PyArrayObject *)) \ + PyArray_API[136]) +#define PyArray_View \ + (*(PyObject * (*)(PyArrayObject *, PyArray_Descr *, PyTypeObject *)) \ + PyArray_API[137]) +#define PyArray_SwapAxes \ + (*(PyObject * (*)(PyArrayObject *, int, int)) \ + PyArray_API[138]) +#define PyArray_Max \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[139]) +#define PyArray_Min \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[140]) +#define PyArray_Ptp \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[141]) +#define PyArray_Mean \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *)) \ + PyArray_API[142]) +#define PyArray_Trace \ + (*(PyObject * (*)(PyArrayObject *, int, int, int, int, PyArrayObject *)) \ + PyArray_API[143]) +#define PyArray_Diagonal \ + (*(PyObject * (*)(PyArrayObject *, int, int, int)) \ + PyArray_API[144]) +#define PyArray_Clip \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, PyObject *, PyArrayObject *)) \ + PyArray_API[145]) +#define PyArray_Conjugate \ + (*(PyObject * (*)(PyArrayObject *, PyArrayObject *)) \ + PyArray_API[146]) +#define PyArray_Nonzero \ + (*(PyObject * (*)(PyArrayObject *)) \ + PyArray_API[147]) +#define PyArray_Std \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *, int)) \ + PyArray_API[148]) +#define PyArray_Sum \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *)) \ + PyArray_API[149]) +#define PyArray_CumSum \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *)) \ + PyArray_API[150]) +#define PyArray_Prod \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *)) \ + PyArray_API[151]) +#define PyArray_CumProd \ + (*(PyObject * (*)(PyArrayObject *, int, int, PyArrayObject *)) \ + PyArray_API[152]) +#define PyArray_All \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[153]) +#define PyArray_Any \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[154]) +#define PyArray_Compress \ + (*(PyObject * (*)(PyArrayObject *, PyObject *, int, PyArrayObject *)) \ + PyArray_API[155]) +#define PyArray_Flatten \ + (*(PyObject * (*)(PyArrayObject *, NPY_ORDER)) \ + PyArray_API[156]) +#define PyArray_Ravel \ + (*(PyObject * (*)(PyArrayObject *, NPY_ORDER)) \ + PyArray_API[157]) +#define PyArray_MultiplyList \ + (*(npy_intp (*)(npy_intp const *, int)) \ + PyArray_API[158]) +#define PyArray_MultiplyIntList \ + (*(int (*)(int const *, int)) \ + PyArray_API[159]) +#define PyArray_GetPtr \ + (*(void * (*)(PyArrayObject *, npy_intp const*)) \ + PyArray_API[160]) +#define PyArray_CompareLists \ + (*(int (*)(npy_intp const *, npy_intp const *, int)) \ + PyArray_API[161]) +#define PyArray_AsCArray \ + (*(int (*)(PyObject **, void *, npy_intp *, int, PyArray_Descr*)) \ + PyArray_API[162]) +#define PyArray_Free \ + (*(int (*)(PyObject *, void *)) \ + PyArray_API[165]) +#define PyArray_Converter \ + (*(int (*)(PyObject *, PyObject **)) \ + PyArray_API[166]) +#define PyArray_IntpFromSequence \ + (*(int (*)(PyObject *, npy_intp *, int)) \ + PyArray_API[167]) +#define PyArray_Concatenate \ + (*(PyObject * (*)(PyObject *, int)) \ + PyArray_API[168]) +#define PyArray_InnerProduct \ + (*(PyObject * (*)(PyObject *, PyObject *)) \ + PyArray_API[169]) +#define PyArray_MatrixProduct \ + (*(PyObject * (*)(PyObject *, PyObject *)) \ + PyArray_API[170]) +#define PyArray_Correlate \ + (*(PyObject * (*)(PyObject *, PyObject *, int)) \ + PyArray_API[172]) +#define PyArray_DescrConverter \ + (*(int (*)(PyObject *, PyArray_Descr **)) \ + PyArray_API[174]) +#define PyArray_DescrConverter2 \ + (*(int (*)(PyObject *, PyArray_Descr **)) \ + PyArray_API[175]) +#define PyArray_IntpConverter \ + (*(int (*)(PyObject *, PyArray_Dims *)) \ + PyArray_API[176]) +#define PyArray_BufferConverter \ + (*(int (*)(PyObject *, PyArray_Chunk *)) \ + PyArray_API[177]) +#define PyArray_AxisConverter \ + (*(int (*)(PyObject *, int *)) \ + PyArray_API[178]) +#define PyArray_BoolConverter \ + (*(int (*)(PyObject *, npy_bool *)) \ + PyArray_API[179]) +#define PyArray_ByteorderConverter \ + (*(int (*)(PyObject *, char *)) \ + PyArray_API[180]) +#define PyArray_OrderConverter \ + (*(int (*)(PyObject *, NPY_ORDER *)) \ + PyArray_API[181]) +#define PyArray_EquivTypes \ + (*(unsigned char (*)(PyArray_Descr *, PyArray_Descr *)) \ + PyArray_API[182]) +#define PyArray_Zeros \ + (*(PyObject * (*)(int, npy_intp const *, PyArray_Descr *, int)) \ + PyArray_API[183]) +#define PyArray_Empty \ + (*(PyObject * (*)(int, npy_intp const *, PyArray_Descr *, int)) \ + PyArray_API[184]) +#define PyArray_Where \ + (*(PyObject * (*)(PyObject *, PyObject *, PyObject *)) \ + PyArray_API[185]) +#define PyArray_Arange \ + (*(PyObject * (*)(double, double, double, int)) \ + PyArray_API[186]) +#define PyArray_ArangeObj \ + (*(PyObject * (*)(PyObject *, PyObject *, PyObject *, PyArray_Descr *)) \ + PyArray_API[187]) +#define PyArray_SortkindConverter \ + (*(int (*)(PyObject *, NPY_SORTKIND *)) \ + PyArray_API[188]) +#define PyArray_LexSort \ + (*(PyObject * (*)(PyObject *, int)) \ + PyArray_API[189]) +#define PyArray_Round \ + (*(PyObject * (*)(PyArrayObject *, int, PyArrayObject *)) \ + PyArray_API[190]) +#define PyArray_EquivTypenums \ + (*(unsigned char (*)(int, int)) \ + PyArray_API[191]) +#define PyArray_RegisterDataType \ + (*(int (*)(PyArray_DescrProto *)) \ + PyArray_API[192]) +#define PyArray_RegisterCastFunc \ + (*(int (*)(PyArray_Descr *, int, PyArray_VectorUnaryFunc *)) \ + PyArray_API[193]) +#define PyArray_RegisterCanCast \ + (*(int (*)(PyArray_Descr *, int, NPY_SCALARKIND)) \ + PyArray_API[194]) +#define PyArray_InitArrFuncs \ + (*(void (*)(PyArray_ArrFuncs *)) \ + PyArray_API[195]) +#define PyArray_IntTupleFromIntp \ + (*(PyObject * (*)(int, npy_intp const *)) \ + PyArray_API[196]) +#define PyArray_ClipmodeConverter \ + (*(int (*)(PyObject *, NPY_CLIPMODE *)) \ + PyArray_API[198]) +#define PyArray_OutputConverter \ + (*(int (*)(PyObject *, PyArrayObject **)) \ + PyArray_API[199]) +#define PyArray_BroadcastToShape \ + (*(PyObject * (*)(PyObject *, npy_intp *, int)) \ + PyArray_API[200]) +#define PyArray_DescrAlignConverter \ + (*(int (*)(PyObject *, PyArray_Descr **)) \ + PyArray_API[203]) +#define PyArray_DescrAlignConverter2 \ + (*(int (*)(PyObject *, PyArray_Descr **)) \ + PyArray_API[204]) +#define PyArray_SearchsideConverter \ + (*(int (*)(PyObject *, void *)) \ + PyArray_API[205]) +#define PyArray_CheckAxis \ + (*(PyObject * (*)(PyArrayObject *, int *, int)) \ + PyArray_API[206]) +#define PyArray_OverflowMultiplyList \ + (*(npy_intp (*)(npy_intp const *, int)) \ + PyArray_API[207]) +#define PyArray_MultiIterFromObjects \ + (*(PyObject* (*)(PyObject **, int, int, ...)) \ + PyArray_API[209]) +#define PyArray_GetEndianness \ + (*(int (*)(void)) \ + PyArray_API[210]) +#define PyArray_GetNDArrayCFeatureVersion \ + (*(unsigned int (*)(void)) \ + PyArray_API[211]) +#define PyArray_Correlate2 \ + (*(PyObject * (*)(PyObject *, PyObject *, int)) \ + PyArray_API[212]) +#define PyArray_NeighborhoodIterNew \ + (*(PyObject* (*)(PyArrayIterObject *, const npy_intp *, int, PyArrayObject*)) \ + PyArray_API[213]) +#define PyTimeIntegerArrType_Type (*(PyTypeObject *)PyArray_API[214]) +#define PyDatetimeArrType_Type (*(PyTypeObject *)PyArray_API[215]) +#define PyTimedeltaArrType_Type (*(PyTypeObject *)PyArray_API[216]) +#define PyHalfArrType_Type (*(PyTypeObject *)PyArray_API[217]) +#define NpyIter_Type (*(PyTypeObject *)PyArray_API[218]) +#define NpyIter_New \ + (*(NpyIter * (*)(PyArrayObject *, npy_uint32, NPY_ORDER, NPY_CASTING, PyArray_Descr*)) \ + PyArray_API[224]) +#define NpyIter_MultiNew \ + (*(NpyIter * (*)(int, PyArrayObject **, npy_uint32, NPY_ORDER, NPY_CASTING, npy_uint32 *, PyArray_Descr **)) \ + PyArray_API[225]) +#define NpyIter_AdvancedNew \ + (*(NpyIter * (*)(int, PyArrayObject **, npy_uint32, NPY_ORDER, NPY_CASTING, npy_uint32 *, PyArray_Descr **, int, int **, npy_intp *, npy_intp)) \ + PyArray_API[226]) +#define NpyIter_Copy \ + (*(NpyIter * (*)(NpyIter *)) \ + PyArray_API[227]) +#define NpyIter_Deallocate \ + (*(int (*)(NpyIter *)) \ + PyArray_API[228]) +#define NpyIter_HasDelayedBufAlloc \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[229]) +#define NpyIter_HasExternalLoop \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[230]) +#define NpyIter_EnableExternalLoop \ + (*(int (*)(NpyIter *)) \ + PyArray_API[231]) +#define NpyIter_GetInnerStrideArray \ + (*(npy_intp * (*)(NpyIter *)) \ + PyArray_API[232]) +#define NpyIter_GetInnerLoopSizePtr \ + (*(npy_intp * (*)(NpyIter *)) \ + PyArray_API[233]) +#define NpyIter_Reset \ + (*(int (*)(NpyIter *, char **)) \ + PyArray_API[234]) +#define NpyIter_ResetBasePointers \ + (*(int (*)(NpyIter *, char **, char **)) \ + PyArray_API[235]) +#define NpyIter_ResetToIterIndexRange \ + (*(int (*)(NpyIter *, npy_intp, npy_intp, char **)) \ + PyArray_API[236]) +#define NpyIter_GetNDim \ + (*(int (*)(NpyIter *)) \ + PyArray_API[237]) +#define NpyIter_GetNOp \ + (*(int (*)(NpyIter *)) \ + PyArray_API[238]) +#define NpyIter_GetIterNext \ + (*(NpyIter_IterNextFunc * (*)(NpyIter *, char **)) \ + PyArray_API[239]) +#define NpyIter_GetIterSize \ + (*(npy_intp (*)(NpyIter *)) \ + PyArray_API[240]) +#define NpyIter_GetIterIndexRange \ + (*(void (*)(NpyIter *, npy_intp *, npy_intp *)) \ + PyArray_API[241]) +#define NpyIter_GetIterIndex \ + (*(npy_intp (*)(NpyIter *)) \ + PyArray_API[242]) +#define NpyIter_GotoIterIndex \ + (*(int (*)(NpyIter *, npy_intp)) \ + PyArray_API[243]) +#define NpyIter_HasMultiIndex \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[244]) +#define NpyIter_GetShape \ + (*(int (*)(NpyIter *, npy_intp *)) \ + PyArray_API[245]) +#define NpyIter_GetGetMultiIndex \ + (*(NpyIter_GetMultiIndexFunc * (*)(NpyIter *, char **)) \ + PyArray_API[246]) +#define NpyIter_GotoMultiIndex \ + (*(int (*)(NpyIter *, npy_intp const *)) \ + PyArray_API[247]) +#define NpyIter_RemoveMultiIndex \ + (*(int (*)(NpyIter *)) \ + PyArray_API[248]) +#define NpyIter_HasIndex \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[249]) +#define NpyIter_IsBuffered \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[250]) +#define NpyIter_IsGrowInner \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[251]) +#define NpyIter_GetBufferSize \ + (*(npy_intp (*)(NpyIter *)) \ + PyArray_API[252]) +#define NpyIter_GetIndexPtr \ + (*(npy_intp * (*)(NpyIter *)) \ + PyArray_API[253]) +#define NpyIter_GotoIndex \ + (*(int (*)(NpyIter *, npy_intp)) \ + PyArray_API[254]) +#define NpyIter_GetDataPtrArray \ + (*(char ** (*)(NpyIter *)) \ + PyArray_API[255]) +#define NpyIter_GetDescrArray \ + (*(PyArray_Descr ** (*)(NpyIter *)) \ + PyArray_API[256]) +#define NpyIter_GetOperandArray \ + (*(PyArrayObject ** (*)(NpyIter *)) \ + PyArray_API[257]) +#define NpyIter_GetIterView \ + (*(PyArrayObject * (*)(NpyIter *, npy_intp)) \ + PyArray_API[258]) +#define NpyIter_GetReadFlags \ + (*(void (*)(NpyIter *, char *)) \ + PyArray_API[259]) +#define NpyIter_GetWriteFlags \ + (*(void (*)(NpyIter *, char *)) \ + PyArray_API[260]) +#define NpyIter_DebugPrint \ + (*(void (*)(NpyIter *)) \ + PyArray_API[261]) +#define NpyIter_IterationNeedsAPI \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[262]) +#define NpyIter_GetInnerFixedStrideArray \ + (*(void (*)(NpyIter *, npy_intp *)) \ + PyArray_API[263]) +#define NpyIter_RemoveAxis \ + (*(int (*)(NpyIter *, int)) \ + PyArray_API[264]) +#define NpyIter_GetAxisStrideArray \ + (*(npy_intp * (*)(NpyIter *, int)) \ + PyArray_API[265]) +#define NpyIter_RequiresBuffering \ + (*(npy_bool (*)(NpyIter *)) \ + PyArray_API[266]) +#define NpyIter_GetInitialDataPtrArray \ + (*(char ** (*)(NpyIter *)) \ + PyArray_API[267]) +#define NpyIter_CreateCompatibleStrides \ + (*(int (*)(NpyIter *, npy_intp, npy_intp *)) \ + PyArray_API[268]) +#define PyArray_CastingConverter \ + (*(int (*)(PyObject *, NPY_CASTING *)) \ + PyArray_API[269]) +#define PyArray_CountNonzero \ + (*(npy_intp (*)(PyArrayObject *)) \ + PyArray_API[270]) +#define PyArray_PromoteTypes \ + (*(PyArray_Descr * (*)(PyArray_Descr *, PyArray_Descr *)) \ + PyArray_API[271]) +#define PyArray_MinScalarType \ + (*(PyArray_Descr * (*)(PyArrayObject *)) \ + PyArray_API[272]) +#define PyArray_ResultType \ + (*(PyArray_Descr * (*)(npy_intp, PyArrayObject *arrs[], npy_intp, PyArray_Descr *descrs[])) \ + PyArray_API[273]) +#define PyArray_CanCastArrayTo \ + (*(npy_bool (*)(PyArrayObject *, PyArray_Descr *, NPY_CASTING)) \ + PyArray_API[274]) +#define PyArray_CanCastTypeTo \ + (*(npy_bool (*)(PyArray_Descr *, PyArray_Descr *, NPY_CASTING)) \ + PyArray_API[275]) +#define PyArray_EinsteinSum \ + (*(PyArrayObject * (*)(char *, npy_intp, PyArrayObject **, PyArray_Descr *, NPY_ORDER, NPY_CASTING, PyArrayObject *)) \ + PyArray_API[276]) +#define PyArray_NewLikeArray \ + (*(PyObject * (*)(PyArrayObject *, NPY_ORDER, PyArray_Descr *, int)) \ + PyArray_API[277]) +#define PyArray_ConvertClipmodeSequence \ + (*(int (*)(PyObject *, NPY_CLIPMODE *, int)) \ + PyArray_API[279]) +#define PyArray_MatrixProduct2 \ + (*(PyObject * (*)(PyObject *, PyObject *, PyArrayObject*)) \ + PyArray_API[280]) +#define NpyIter_IsFirstVisit \ + (*(npy_bool (*)(NpyIter *, int)) \ + PyArray_API[281]) +#define PyArray_SetBaseObject \ + (*(int (*)(PyArrayObject *, PyObject *)) \ + PyArray_API[282]) +#define PyArray_CreateSortedStridePerm \ + (*(void (*)(int, npy_intp const *, npy_stride_sort_item *)) \ + PyArray_API[283]) +#define PyArray_RemoveAxesInPlace \ + (*(void (*)(PyArrayObject *, const npy_bool *)) \ + PyArray_API[284]) +#define PyArray_DebugPrint \ + (*(void (*)(PyArrayObject *)) \ + PyArray_API[285]) +#define PyArray_FailUnlessWriteable \ + (*(int (*)(PyArrayObject *, const char *)) \ + PyArray_API[286]) +#define PyArray_SetUpdateIfCopyBase \ + (*(int (*)(PyArrayObject *, PyArrayObject *)) \ + PyArray_API[287]) +#define PyDataMem_NEW \ + (*(void * (*)(size_t)) \ + PyArray_API[288]) +#define PyDataMem_FREE \ + (*(void (*)(void *)) \ + PyArray_API[289]) +#define PyDataMem_RENEW \ + (*(void * (*)(void *, size_t)) \ + PyArray_API[290]) +#define NPY_DEFAULT_ASSIGN_CASTING (*(NPY_CASTING *)PyArray_API[292]) +#define PyArray_Partition \ + (*(int (*)(PyArrayObject *, PyArrayObject *, int, NPY_SELECTKIND)) \ + PyArray_API[296]) +#define PyArray_ArgPartition \ + (*(PyObject * (*)(PyArrayObject *, PyArrayObject *, int, NPY_SELECTKIND)) \ + PyArray_API[297]) +#define PyArray_SelectkindConverter \ + (*(int (*)(PyObject *, NPY_SELECTKIND *)) \ + PyArray_API[298]) +#define PyDataMem_NEW_ZEROED \ + (*(void * (*)(size_t, size_t)) \ + PyArray_API[299]) +#define PyArray_CheckAnyScalarExact \ + (*(int (*)(PyObject *)) \ + PyArray_API[300]) +#define PyArray_ResolveWritebackIfCopy \ + (*(int (*)(PyArrayObject *)) \ + PyArray_API[302]) +#define PyArray_SetWritebackIfCopyBase \ + (*(int (*)(PyArrayObject *, PyArrayObject *)) \ + PyArray_API[303]) + +#if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION +#define PyDataMem_SetHandler \ + (*(PyObject * (*)(PyObject *)) \ + PyArray_API[304]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION +#define PyDataMem_GetHandler \ + (*(PyObject * (*)(void)) \ + PyArray_API[305]) +#endif +#define PyDataMem_DefaultHandler (*(PyObject* *)PyArray_API[306]) + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_ConvertDatetime64ToDatetimeStruct \ + (*(int (*)(PyArray_DatetimeMetaData *, npy_datetime, npy_datetimestruct *)) \ + PyArray_API[307]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_ConvertDatetimeStructToDatetime64 \ + (*(int (*)(PyArray_DatetimeMetaData *, const npy_datetimestruct *, npy_datetime *)) \ + PyArray_API[308]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_ConvertPyDateTimeToDatetimeStruct \ + (*(int (*)(PyObject *, npy_datetimestruct *, NPY_DATETIMEUNIT *, int)) \ + PyArray_API[309]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_GetDatetimeISO8601StrLen \ + (*(int (*)(int, NPY_DATETIMEUNIT)) \ + PyArray_API[310]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_MakeISO8601Datetime \ + (*(int (*)(npy_datetimestruct *, char *, npy_intp, int, int, NPY_DATETIMEUNIT, int, NPY_CASTING)) \ + PyArray_API[311]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyDatetime_ParseISO8601Datetime \ + (*(int (*)(char const *, Py_ssize_t, NPY_DATETIMEUNIT, NPY_CASTING, npy_datetimestruct *, NPY_DATETIMEUNIT *, npy_bool *)) \ + PyArray_API[312]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_load \ + (*(int (*)(npy_string_allocator *, const npy_packed_static_string *, npy_static_string *)) \ + PyArray_API[313]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_pack \ + (*(int (*)(npy_string_allocator *, npy_packed_static_string *, const char *, size_t)) \ + PyArray_API[314]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_pack_null \ + (*(int (*)(npy_string_allocator *, npy_packed_static_string *)) \ + PyArray_API[315]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_acquire_allocator \ + (*(npy_string_allocator * (*)(const PyArray_StringDTypeObject *)) \ + PyArray_API[316]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_acquire_allocators \ + (*(void (*)(size_t, PyArray_Descr *const descrs[], npy_string_allocator *allocators[])) \ + PyArray_API[317]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_release_allocator \ + (*(void (*)(npy_string_allocator *)) \ + PyArray_API[318]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define NpyString_release_allocators \ + (*(void (*)(size_t, npy_string_allocator *allocators[])) \ + PyArray_API[319]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyArray_GetDefaultDescr \ + (*(PyArray_Descr * (*)(PyArray_DTypeMeta *)) \ + PyArray_API[361]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyArrayInitDTypeMeta_FromSpec \ + (*(int (*)(PyArray_DTypeMeta *, PyArrayDTypeMeta_Spec *)) \ + PyArray_API[362]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyArray_CommonDType \ + (*(PyArray_DTypeMeta * (*)(PyArray_DTypeMeta *, PyArray_DTypeMeta *)) \ + PyArray_API[363]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyArray_PromoteDTypeSequence \ + (*(PyArray_DTypeMeta * (*)(npy_intp, PyArray_DTypeMeta **)) \ + PyArray_API[364]) +#endif +#define _PyDataType_GetArrFuncs \ + (*(PyArray_ArrFuncs * (*)(const PyArray_Descr *)) \ + PyArray_API[365]) + +/* + * The DType classes are inconvenient for the Python generation so exposed + * manually in the header below (may be moved). + */ +#include "numpy/_public_dtype_api_table.h" + +#if !defined(NO_IMPORT_ARRAY) && !defined(NO_IMPORT) +static int +_import_array(void) +{ + int st; + PyObject *numpy = PyImport_ImportModule("numpy._core._multiarray_umath"); + if (numpy == NULL && PyErr_ExceptionMatches(PyExc_ModuleNotFoundError)) { + PyErr_Clear(); + numpy = PyImport_ImportModule("numpy.core._multiarray_umath"); + } + + if (numpy == NULL) { + return -1; + } + + PyObject *c_api = PyObject_GetAttrString(numpy, "_ARRAY_API"); + Py_DECREF(numpy); + if (c_api == NULL) { + return -1; + } + + if (!PyCapsule_CheckExact(c_api)) { + PyErr_SetString(PyExc_RuntimeError, "_ARRAY_API is not PyCapsule object"); + Py_DECREF(c_api); + return -1; + } + PyArray_API = (void **)PyCapsule_GetPointer(c_api, NULL); + Py_DECREF(c_api); + if (PyArray_API == NULL) { + PyErr_SetString(PyExc_RuntimeError, "_ARRAY_API is NULL pointer"); + return -1; + } + + /* + * On exceedingly few platforms these sizes may not match, in which case + * We do not support older NumPy versions at all. + */ + if (sizeof(Py_ssize_t) != sizeof(Py_intptr_t) && + PyArray_RUNTIME_VERSION < NPY_2_0_API_VERSION) { + PyErr_Format(PyExc_RuntimeError, + "module compiled against NumPy 2.0 but running on NumPy 1.x. " + "Unfortunately, this is not supported on niche platforms where " + "`sizeof(size_t) != sizeof(inptr_t)`."); + } + /* + * Perform runtime check of C API version. As of now NumPy 2.0 is ABI + * backwards compatible (in the exposed feature subset!) for all practical + * purposes. + */ + if (NPY_VERSION < PyArray_GetNDArrayCVersion()) { + PyErr_Format(PyExc_RuntimeError, "module compiled against "\ + "ABI version 0x%x but this version of numpy is 0x%x", \ + (int) NPY_VERSION, (int) PyArray_GetNDArrayCVersion()); + return -1; + } + PyArray_RUNTIME_VERSION = (int)PyArray_GetNDArrayCFeatureVersion(); + if (NPY_FEATURE_VERSION > PyArray_RUNTIME_VERSION) { + PyErr_Format(PyExc_RuntimeError, + "module was compiled against NumPy C-API version 0x%x " + "(NumPy " NPY_FEATURE_VERSION_STRING ") " + "but the running NumPy has C-API version 0x%x. " + "Check the section C-API incompatibility at the " + "Troubleshooting ImportError section at " + "https://numpy.org/devdocs/user/troubleshooting-importerror.html" + "#c-api-incompatibility " + "for indications on how to solve this problem.", + (int)NPY_FEATURE_VERSION, PyArray_RUNTIME_VERSION); + return -1; + } + + /* + * Perform runtime check of endianness and check it matches the one set by + * the headers (npy_endian.h) as a safeguard + */ + st = PyArray_GetEndianness(); + if (st == NPY_CPU_UNKNOWN_ENDIAN) { + PyErr_SetString(PyExc_RuntimeError, + "FATAL: module compiled as unknown endian"); + return -1; + } +#if NPY_BYTE_ORDER == NPY_BIG_ENDIAN + if (st != NPY_CPU_BIG) { + PyErr_SetString(PyExc_RuntimeError, + "FATAL: module compiled as big endian, but " + "detected different endianness at runtime"); + return -1; + } +#elif NPY_BYTE_ORDER == NPY_LITTLE_ENDIAN + if (st != NPY_CPU_LITTLE) { + PyErr_SetString(PyExc_RuntimeError, + "FATAL: module compiled as little endian, but " + "detected different endianness at runtime"); + return -1; + } +#endif + + return 0; +} + +#define import_array() { \ + if (_import_array() < 0) { \ + PyErr_Print(); \ + PyErr_SetString( \ + PyExc_ImportError, \ + "numpy._core.multiarray failed to import" \ + ); \ + return NULL; \ + } \ +} + +#define import_array1(ret) { \ + if (_import_array() < 0) { \ + PyErr_Print(); \ + PyErr_SetString( \ + PyExc_ImportError, \ + "numpy._core.multiarray failed to import" \ + ); \ + return ret; \ + } \ +} + +#define import_array2(msg, ret) { \ + if (_import_array() < 0) { \ + PyErr_Print(); \ + PyErr_SetString(PyExc_ImportError, msg); \ + return ret; \ + } \ +} + +#endif + +#endif diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.c b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.c new file mode 100644 index 0000000000000000000000000000000000000000..10fcbc4553989057667a90ce9d587deefc13f5f1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.c @@ -0,0 +1,54 @@ + +/* These pointers will be stored in the C-object for use in other + extension modules +*/ + +void *PyUFunc_API[] = { + (void *) &PyUFunc_Type, + (void *) PyUFunc_FromFuncAndData, + (void *) PyUFunc_RegisterLoopForType, + NULL, + (void *) PyUFunc_f_f_As_d_d, + (void *) PyUFunc_d_d, + (void *) PyUFunc_f_f, + (void *) PyUFunc_g_g, + (void *) PyUFunc_F_F_As_D_D, + (void *) PyUFunc_F_F, + (void *) PyUFunc_D_D, + (void *) PyUFunc_G_G, + (void *) PyUFunc_O_O, + (void *) PyUFunc_ff_f_As_dd_d, + (void *) PyUFunc_ff_f, + (void *) PyUFunc_dd_d, + (void *) PyUFunc_gg_g, + (void *) PyUFunc_FF_F_As_DD_D, + (void *) PyUFunc_DD_D, + (void *) PyUFunc_FF_F, + (void *) PyUFunc_GG_G, + (void *) PyUFunc_OO_O, + (void *) PyUFunc_O_O_method, + (void *) PyUFunc_OO_O_method, + (void *) PyUFunc_On_Om, + NULL, + NULL, + (void *) PyUFunc_clearfperr, + (void *) PyUFunc_getfperr, + NULL, + (void *) PyUFunc_ReplaceLoopBySignature, + (void *) PyUFunc_FromFuncAndDataAndSignature, + NULL, + (void *) PyUFunc_e_e, + (void *) PyUFunc_e_e_As_f_f, + (void *) PyUFunc_e_e_As_d_d, + (void *) PyUFunc_ee_e, + (void *) PyUFunc_ee_e_As_ff_f, + (void *) PyUFunc_ee_e_As_dd_d, + (void *) PyUFunc_DefaultTypeResolver, + (void *) PyUFunc_ValidateCasting, + (void *) PyUFunc_RegisterLoopForDescr, + (void *) PyUFunc_FromFuncAndDataAndSignatureAndIdentity, + (void *) PyUFunc_AddLoopFromSpec, + (void *) PyUFunc_AddPromoter, + (void *) PyUFunc_AddWrappingLoop, + (void *) PyUFunc_GiveFloatingpointErrors +}; diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.h new file mode 100644 index 0000000000000000000000000000000000000000..df7ded10b548d495338dc368d727cf77ff70740a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/__ufunc_api.h @@ -0,0 +1,340 @@ + +#ifdef _UMATHMODULE + +extern NPY_NO_EXPORT PyTypeObject PyUFunc_Type; + +extern NPY_NO_EXPORT PyTypeObject PyUFunc_Type; + +NPY_NO_EXPORT PyObject * PyUFunc_FromFuncAndData \ + (PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, int); +NPY_NO_EXPORT int PyUFunc_RegisterLoopForType \ + (PyUFuncObject *, int, PyUFuncGenericFunction, const int *, void *); +NPY_NO_EXPORT void PyUFunc_f_f_As_d_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_d_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_f_f \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_g_g \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_F_F_As_D_D \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_F_F \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_D_D \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_G_G \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_O_O \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_ff_f_As_dd_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_ff_f \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_dd_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_gg_g \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_FF_F_As_DD_D \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_DD_D \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_FF_F \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_GG_G \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_OO_O \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_O_O_method \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_OO_O_method \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_On_Om \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_clearfperr \ + (void); +NPY_NO_EXPORT int PyUFunc_getfperr \ + (void); +NPY_NO_EXPORT int PyUFunc_ReplaceLoopBySignature \ + (PyUFuncObject *, PyUFuncGenericFunction, const int *, PyUFuncGenericFunction *); +NPY_NO_EXPORT PyObject * PyUFunc_FromFuncAndDataAndSignature \ + (PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, int, const char *); +NPY_NO_EXPORT void PyUFunc_e_e \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_e_e_As_f_f \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_e_e_As_d_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_ee_e \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_ee_e_As_ff_f \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT void PyUFunc_ee_e_As_dd_d \ + (char **, npy_intp const *, npy_intp const *, void *); +NPY_NO_EXPORT int PyUFunc_DefaultTypeResolver \ + (PyUFuncObject *, NPY_CASTING, PyArrayObject **, PyObject *, PyArray_Descr **); +NPY_NO_EXPORT int PyUFunc_ValidateCasting \ + (PyUFuncObject *, NPY_CASTING, PyArrayObject **, PyArray_Descr *const *); +NPY_NO_EXPORT int PyUFunc_RegisterLoopForDescr \ + (PyUFuncObject *, PyArray_Descr *, PyUFuncGenericFunction, PyArray_Descr **, void *); +NPY_NO_EXPORT PyObject * PyUFunc_FromFuncAndDataAndSignatureAndIdentity \ + (PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, const int, const char *, PyObject *); +NPY_NO_EXPORT int PyUFunc_AddLoopFromSpec \ + (PyObject *, PyArrayMethod_Spec *); +NPY_NO_EXPORT int PyUFunc_AddPromoter \ + (PyObject *, PyObject *, PyObject *); +NPY_NO_EXPORT int PyUFunc_AddWrappingLoop \ + (PyObject *, PyArray_DTypeMeta *new_dtypes[], PyArray_DTypeMeta *wrapped_dtypes[], PyArrayMethod_TranslateGivenDescriptors *, PyArrayMethod_TranslateLoopDescriptors *); +NPY_NO_EXPORT int PyUFunc_GiveFloatingpointErrors \ + (const char *, int); + +#else + +#if defined(PY_UFUNC_UNIQUE_SYMBOL) +#define PyUFunc_API PY_UFUNC_UNIQUE_SYMBOL +#endif + +/* By default do not export API in an .so (was never the case on windows) */ +#ifndef NPY_API_SYMBOL_ATTRIBUTE + #define NPY_API_SYMBOL_ATTRIBUTE NPY_VISIBILITY_HIDDEN +#endif + +#if defined(NO_IMPORT) || defined(NO_IMPORT_UFUNC) +extern NPY_API_SYMBOL_ATTRIBUTE void **PyUFunc_API; +#else +#if defined(PY_UFUNC_UNIQUE_SYMBOL) +NPY_API_SYMBOL_ATTRIBUTE void **PyUFunc_API; +#else +static void **PyUFunc_API=NULL; +#endif +#endif + +#define PyUFunc_Type (*(PyTypeObject *)PyUFunc_API[0]) +#define PyUFunc_FromFuncAndData \ + (*(PyObject * (*)(PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, int)) \ + PyUFunc_API[1]) +#define PyUFunc_RegisterLoopForType \ + (*(int (*)(PyUFuncObject *, int, PyUFuncGenericFunction, const int *, void *)) \ + PyUFunc_API[2]) +#define PyUFunc_f_f_As_d_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[4]) +#define PyUFunc_d_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[5]) +#define PyUFunc_f_f \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[6]) +#define PyUFunc_g_g \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[7]) +#define PyUFunc_F_F_As_D_D \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[8]) +#define PyUFunc_F_F \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[9]) +#define PyUFunc_D_D \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[10]) +#define PyUFunc_G_G \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[11]) +#define PyUFunc_O_O \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[12]) +#define PyUFunc_ff_f_As_dd_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[13]) +#define PyUFunc_ff_f \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[14]) +#define PyUFunc_dd_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[15]) +#define PyUFunc_gg_g \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[16]) +#define PyUFunc_FF_F_As_DD_D \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[17]) +#define PyUFunc_DD_D \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[18]) +#define PyUFunc_FF_F \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[19]) +#define PyUFunc_GG_G \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[20]) +#define PyUFunc_OO_O \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[21]) +#define PyUFunc_O_O_method \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[22]) +#define PyUFunc_OO_O_method \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[23]) +#define PyUFunc_On_Om \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[24]) +#define PyUFunc_clearfperr \ + (*(void (*)(void)) \ + PyUFunc_API[27]) +#define PyUFunc_getfperr \ + (*(int (*)(void)) \ + PyUFunc_API[28]) +#define PyUFunc_ReplaceLoopBySignature \ + (*(int (*)(PyUFuncObject *, PyUFuncGenericFunction, const int *, PyUFuncGenericFunction *)) \ + PyUFunc_API[30]) +#define PyUFunc_FromFuncAndDataAndSignature \ + (*(PyObject * (*)(PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, int, const char *)) \ + PyUFunc_API[31]) +#define PyUFunc_e_e \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[33]) +#define PyUFunc_e_e_As_f_f \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[34]) +#define PyUFunc_e_e_As_d_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[35]) +#define PyUFunc_ee_e \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[36]) +#define PyUFunc_ee_e_As_ff_f \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[37]) +#define PyUFunc_ee_e_As_dd_d \ + (*(void (*)(char **, npy_intp const *, npy_intp const *, void *)) \ + PyUFunc_API[38]) +#define PyUFunc_DefaultTypeResolver \ + (*(int (*)(PyUFuncObject *, NPY_CASTING, PyArrayObject **, PyObject *, PyArray_Descr **)) \ + PyUFunc_API[39]) +#define PyUFunc_ValidateCasting \ + (*(int (*)(PyUFuncObject *, NPY_CASTING, PyArrayObject **, PyArray_Descr *const *)) \ + PyUFunc_API[40]) +#define PyUFunc_RegisterLoopForDescr \ + (*(int (*)(PyUFuncObject *, PyArray_Descr *, PyUFuncGenericFunction, PyArray_Descr **, void *)) \ + PyUFunc_API[41]) + +#if NPY_FEATURE_VERSION >= NPY_1_16_API_VERSION +#define PyUFunc_FromFuncAndDataAndSignatureAndIdentity \ + (*(PyObject * (*)(PyUFuncGenericFunction *, void *const *, const char *, int, int, int, int, const char *, const char *, const int, const char *, PyObject *)) \ + PyUFunc_API[42]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyUFunc_AddLoopFromSpec \ + (*(int (*)(PyObject *, PyArrayMethod_Spec *)) \ + PyUFunc_API[43]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyUFunc_AddPromoter \ + (*(int (*)(PyObject *, PyObject *, PyObject *)) \ + PyUFunc_API[44]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyUFunc_AddWrappingLoop \ + (*(int (*)(PyObject *, PyArray_DTypeMeta *new_dtypes[], PyArray_DTypeMeta *wrapped_dtypes[], PyArrayMethod_TranslateGivenDescriptors *, PyArrayMethod_TranslateLoopDescriptors *)) \ + PyUFunc_API[45]) +#endif + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +#define PyUFunc_GiveFloatingpointErrors \ + (*(int (*)(const char *, int)) \ + PyUFunc_API[46]) +#endif + +static inline int +_import_umath(void) +{ + PyObject *numpy = PyImport_ImportModule("numpy._core._multiarray_umath"); + if (numpy == NULL && PyErr_ExceptionMatches(PyExc_ModuleNotFoundError)) { + PyErr_Clear(); + numpy = PyImport_ImportModule("numpy.core._multiarray_umath"); + } + + if (numpy == NULL) { + PyErr_SetString(PyExc_ImportError, + "_multiarray_umath failed to import"); + return -1; + } + + PyObject *c_api = PyObject_GetAttrString(numpy, "_UFUNC_API"); + Py_DECREF(numpy); + if (c_api == NULL) { + PyErr_SetString(PyExc_AttributeError, "_UFUNC_API not found"); + return -1; + } + + if (!PyCapsule_CheckExact(c_api)) { + PyErr_SetString(PyExc_RuntimeError, "_UFUNC_API is not PyCapsule object"); + Py_DECREF(c_api); + return -1; + } + PyUFunc_API = (void **)PyCapsule_GetPointer(c_api, NULL); + Py_DECREF(c_api); + if (PyUFunc_API == NULL) { + PyErr_SetString(PyExc_RuntimeError, "_UFUNC_API is NULL pointer"); + return -1; + } + return 0; +} + +#define import_umath() \ + do {\ + UFUNC_NOFPE\ + if (_import_umath() < 0) {\ + PyErr_Print();\ + PyErr_SetString(PyExc_ImportError,\ + "numpy._core.umath failed to import");\ + return NULL;\ + }\ + } while(0) + +#define import_umath1(ret) \ + do {\ + UFUNC_NOFPE\ + if (_import_umath() < 0) {\ + PyErr_Print();\ + PyErr_SetString(PyExc_ImportError,\ + "numpy._core.umath failed to import");\ + return ret;\ + }\ + } while(0) + +#define import_umath2(ret, msg) \ + do {\ + UFUNC_NOFPE\ + if (_import_umath() < 0) {\ + PyErr_Print();\ + PyErr_SetString(PyExc_ImportError, msg);\ + return ret;\ + }\ + } while(0) + +#define import_ufunc() \ + do {\ + UFUNC_NOFPE\ + if (_import_umath() < 0) {\ + PyErr_Print();\ + PyErr_SetString(PyExc_ImportError,\ + "numpy._core.umath failed to import");\ + }\ + } while(0) + + +static inline int +PyUFunc_ImportUFuncAPI() +{ + if (NPY_UNLIKELY(PyUFunc_API == NULL)) { + import_umath1(-1); + } + return 0; +} + +#endif diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/_numpyconfig.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/_numpyconfig.h new file mode 100644 index 0000000000000000000000000000000000000000..6e50e16505bbc9ad5f859ff63dd16978c7cd3ecb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/_numpyconfig.h @@ -0,0 +1,33 @@ +#define NPY_HAVE_ENDIAN_H 1 + +#define NPY_SIZEOF_SHORT 2 +#define NPY_SIZEOF_INT 4 +#define NPY_SIZEOF_LONG 8 +#define NPY_SIZEOF_FLOAT 4 +#define NPY_SIZEOF_COMPLEX_FLOAT 8 +#define NPY_SIZEOF_DOUBLE 8 +#define NPY_SIZEOF_COMPLEX_DOUBLE 16 +#define NPY_SIZEOF_LONGDOUBLE 16 +#define NPY_SIZEOF_COMPLEX_LONGDOUBLE 32 +#define NPY_SIZEOF_PY_INTPTR_T 8 +#define NPY_SIZEOF_INTP 8 +#define NPY_SIZEOF_UINTP 8 +#define NPY_SIZEOF_WCHAR_T 4 +#define NPY_SIZEOF_OFF_T 8 +#define NPY_SIZEOF_PY_LONG_LONG 8 +#define NPY_SIZEOF_LONGLONG 8 + +/* + * Defined to 1 or 0. Note that Pyodide hardcodes NPY_NO_SMP (and other defines + * in this header) for better cross-compilation, so don't rename them without a + * good reason. + */ +#define NPY_NO_SMP 0 + +#define NPY_VISIBILITY_HIDDEN __attribute__((visibility("hidden"))) +#define NPY_ABI_VERSION 0x02000000 +#define NPY_API_VERSION 0x00000013 + +#ifndef __STDC_FORMAT_MACROS +#define __STDC_FORMAT_MACROS 1 +#endif diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayobject.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayobject.h new file mode 100644 index 0000000000000000000000000000000000000000..97d93590e401fc878b72281d938170623f01da97 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayobject.h @@ -0,0 +1,7 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ +#define Py_ARRAYOBJECT_H + +#include "ndarrayobject.h" + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_ARRAYOBJECT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayscalars.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayscalars.h new file mode 100644 index 0000000000000000000000000000000000000000..ff048061f70abda097606b6247e32234d4eacf29 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/arrayscalars.h @@ -0,0 +1,196 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ + +#ifndef _MULTIARRAYMODULE +typedef struct { + PyObject_HEAD + npy_bool obval; +} PyBoolScalarObject; +#endif + + +typedef struct { + PyObject_HEAD + signed char obval; +} PyByteScalarObject; + + +typedef struct { + PyObject_HEAD + short obval; +} PyShortScalarObject; + + +typedef struct { + PyObject_HEAD + int obval; +} PyIntScalarObject; + + +typedef struct { + PyObject_HEAD + long obval; +} PyLongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_longlong obval; +} PyLongLongScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned char obval; +} PyUByteScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned short obval; +} PyUShortScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned int obval; +} PyUIntScalarObject; + + +typedef struct { + PyObject_HEAD + unsigned long obval; +} PyULongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_ulonglong obval; +} PyULongLongScalarObject; + + +typedef struct { + PyObject_HEAD + npy_half obval; +} PyHalfScalarObject; + + +typedef struct { + PyObject_HEAD + float obval; +} PyFloatScalarObject; + + +typedef struct { + PyObject_HEAD + double obval; +} PyDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_longdouble obval; +} PyLongDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_cfloat obval; +} PyCFloatScalarObject; + + +typedef struct { + PyObject_HEAD + npy_cdouble obval; +} PyCDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + npy_clongdouble obval; +} PyCLongDoubleScalarObject; + + +typedef struct { + PyObject_HEAD + PyObject * obval; +} PyObjectScalarObject; + +typedef struct { + PyObject_HEAD + npy_datetime obval; + PyArray_DatetimeMetaData obmeta; +} PyDatetimeScalarObject; + +typedef struct { + PyObject_HEAD + npy_timedelta obval; + PyArray_DatetimeMetaData obmeta; +} PyTimedeltaScalarObject; + + +typedef struct { + PyObject_HEAD + char obval; +} PyScalarObject; + +#define PyStringScalarObject PyBytesObject +#ifndef Py_LIMITED_API +typedef struct { + /* note that the PyObject_HEAD macro lives right here */ + PyUnicodeObject base; + Py_UCS4 *obval; + #if NPY_FEATURE_VERSION >= NPY_1_20_API_VERSION + char *buffer_fmt; + #endif +} PyUnicodeScalarObject; +#endif + + +typedef struct { + PyObject_VAR_HEAD + char *obval; +#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD + /* Internally use the subclass to allow accessing names/fields */ + _PyArray_LegacyDescr *descr; +#else + PyArray_Descr *descr; +#endif + int flags; + PyObject *base; + #if NPY_FEATURE_VERSION >= NPY_1_20_API_VERSION + void *_buffer_info; /* private buffer info, tagged to allow warning */ + #endif +} PyVoidScalarObject; + +/* Macros + PyScalarObject + PyArrType_Type + are defined in ndarrayobject.h +*/ + +#define PyArrayScalar_False ((PyObject *)(&(_PyArrayScalar_BoolValues[0]))) +#define PyArrayScalar_True ((PyObject *)(&(_PyArrayScalar_BoolValues[1]))) +#define PyArrayScalar_FromLong(i) \ + ((PyObject *)(&(_PyArrayScalar_BoolValues[((i)!=0)]))) +#define PyArrayScalar_RETURN_BOOL_FROM_LONG(i) \ + return Py_INCREF(PyArrayScalar_FromLong(i)), \ + PyArrayScalar_FromLong(i) +#define PyArrayScalar_RETURN_FALSE \ + return Py_INCREF(PyArrayScalar_False), \ + PyArrayScalar_False +#define PyArrayScalar_RETURN_TRUE \ + return Py_INCREF(PyArrayScalar_True), \ + PyArrayScalar_True + +#define PyArrayScalar_New(cls) \ + Py##cls##ArrType_Type.tp_alloc(&Py##cls##ArrType_Type, 0) +#ifndef Py_LIMITED_API +/* For the limited API, use PyArray_ScalarAsCtype instead */ +#define PyArrayScalar_VAL(obj, cls) \ + ((Py##cls##ScalarObject *)obj)->obval +#define PyArrayScalar_ASSIGN(obj, cls, val) \ + PyArrayScalar_VAL(obj, cls) = val +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_ARRAYSCALARS_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/dtype_api.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/dtype_api.h new file mode 100644 index 0000000000000000000000000000000000000000..b37c9fbb68213e8ff7817093bc4c1f495d84b79d --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/dtype_api.h @@ -0,0 +1,480 @@ +/* + * The public DType API + */ + +#ifndef NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ +#define NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ + +struct PyArrayMethodObject_tag; + +/* + * Largely opaque struct for DType classes (i.e. metaclass instances). + * The internal definition is currently in `ndarraytypes.h` (export is a bit + * more complex because `PyArray_Descr` is a DTypeMeta internally but not + * externally). + */ +#if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) + +#ifndef Py_LIMITED_API + + typedef struct PyArray_DTypeMeta_tag { + PyHeapTypeObject super; + + /* + * Most DTypes will have a singleton default instance, for the + * parametric legacy DTypes (bytes, string, void, datetime) this + * may be a pointer to the *prototype* instance? + */ + PyArray_Descr *singleton; + /* Copy of the legacy DTypes type number, usually invalid. */ + int type_num; + + /* The type object of the scalar instances (may be NULL?) */ + PyTypeObject *scalar_type; + /* + * DType flags to signal legacy, parametric, or + * abstract. But plenty of space for additional information/flags. + */ + npy_uint64 flags; + + /* + * Use indirection in order to allow a fixed size for this struct. + * A stable ABI size makes creating a static DType less painful + * while also ensuring flexibility for all opaque API (with one + * indirection due the pointer lookup). + */ + void *dt_slots; + /* Allow growing (at the moment also beyond this) */ + void *reserved[3]; + } PyArray_DTypeMeta; + +#else + +typedef PyTypeObject PyArray_DTypeMeta; + +#endif /* Py_LIMITED_API */ + +#endif /* not internal build */ + +/* + * ****************************************************** + * ArrayMethod API (Casting and UFuncs) + * ****************************************************** + */ + + +typedef enum { + /* Flag for whether the GIL is required */ + NPY_METH_REQUIRES_PYAPI = 1 << 0, + /* + * Some functions cannot set floating point error flags, this flag + * gives us the option (not requirement) to skip floating point error + * setup/check. No function should set error flags and ignore them + * since it would interfere with chaining operations (e.g. casting). + */ + NPY_METH_NO_FLOATINGPOINT_ERRORS = 1 << 1, + /* Whether the method supports unaligned access (not runtime) */ + NPY_METH_SUPPORTS_UNALIGNED = 1 << 2, + /* + * Used for reductions to allow reordering the operation. At this point + * assume that if set, it also applies to normal operations though! + */ + NPY_METH_IS_REORDERABLE = 1 << 3, + /* + * Private flag for now for *logic* functions. The logical functions + * `logical_or` and `logical_and` can always cast the inputs to booleans + * "safely" (because that is how the cast to bool is defined). + * @seberg: I am not sure this is the best way to handle this, so its + * private for now (also it is very limited anyway). + * There is one "exception". NA aware dtypes cannot cast to bool + * (hopefully), so the `??->?` loop should error even with this flag. + * But a second NA fallback loop will be necessary. + */ + _NPY_METH_FORCE_CAST_INPUTS = 1 << 17, + + /* All flags which can change at runtime */ + NPY_METH_RUNTIME_FLAGS = ( + NPY_METH_REQUIRES_PYAPI | + NPY_METH_NO_FLOATINGPOINT_ERRORS), +} NPY_ARRAYMETHOD_FLAGS; + + +typedef struct PyArrayMethod_Context_tag { + /* The caller, which is typically the original ufunc. May be NULL */ + PyObject *caller; + /* The method "self". Currently an opaque object. */ + struct PyArrayMethodObject_tag *method; + + /* Operand descriptors, filled in by resolve_descriptors */ + PyArray_Descr *const *descriptors; + /* Structure may grow (this is harmless for DType authors) */ +} PyArrayMethod_Context; + + +/* + * The main object for creating a new ArrayMethod. We use the typical `slots` + * mechanism used by the Python limited API (see below for the slot defs). + */ +typedef struct { + const char *name; + int nin, nout; + NPY_CASTING casting; + NPY_ARRAYMETHOD_FLAGS flags; + PyArray_DTypeMeta **dtypes; + PyType_Slot *slots; +} PyArrayMethod_Spec; + + +/* + * ArrayMethod slots + * ----------------- + * + * SLOTS IDs For the ArrayMethod creation, once fully public, IDs are fixed + * but can be deprecated and arbitrarily extended. + */ +#define _NPY_METH_resolve_descriptors_with_scalars 1 +#define NPY_METH_resolve_descriptors 2 +#define NPY_METH_get_loop 3 +#define NPY_METH_get_reduction_initial 4 +/* specific loops for constructions/default get_loop: */ +#define NPY_METH_strided_loop 5 +#define NPY_METH_contiguous_loop 6 +#define NPY_METH_unaligned_strided_loop 7 +#define NPY_METH_unaligned_contiguous_loop 8 +#define NPY_METH_contiguous_indexed_loop 9 +#define _NPY_METH_static_data 10 + + +/* + * The resolve descriptors function, must be able to handle NULL values for + * all output (but not input) `given_descrs` and fill `loop_descrs`. + * Return -1 on error or 0 if the operation is not possible without an error + * set. (This may still be in flux.) + * Otherwise must return the "casting safety", for normal functions, this is + * almost always "safe" (or even "equivalent"?). + * + * `resolve_descriptors` is optional if all output DTypes are non-parametric. + */ +typedef NPY_CASTING (PyArrayMethod_ResolveDescriptors)( + /* "method" is currently opaque (necessary e.g. to wrap Python) */ + struct PyArrayMethodObject_tag *method, + /* DTypes the method was created for */ + PyArray_DTypeMeta *const *dtypes, + /* Input descriptors (instances). Outputs may be NULL. */ + PyArray_Descr *const *given_descrs, + /* Exact loop descriptors to use, must not hold references on error */ + PyArray_Descr **loop_descrs, + npy_intp *view_offset); + + +/* + * Rarely needed, slightly more powerful version of `resolve_descriptors`. + * See also `PyArrayMethod_ResolveDescriptors` for details on shared arguments. + * + * NOTE: This function is private now as it is unclear how and what to pass + * exactly as additional information to allow dealing with the scalars. + * See also gh-24915. + */ +typedef NPY_CASTING (PyArrayMethod_ResolveDescriptorsWithScalar)( + struct PyArrayMethodObject_tag *method, + PyArray_DTypeMeta *const *dtypes, + /* Unlike above, these can have any DType and we may allow NULL. */ + PyArray_Descr *const *given_descrs, + /* + * Input scalars or NULL. Only ever passed for python scalars. + * WARNING: In some cases, a loop may be explicitly selected and the + * value passed is not available (NULL) or does not have the + * expected type. + */ + PyObject *const *input_scalars, + PyArray_Descr **loop_descrs, + npy_intp *view_offset); + + + +typedef int (PyArrayMethod_StridedLoop)(PyArrayMethod_Context *context, + char *const *data, const npy_intp *dimensions, const npy_intp *strides, + NpyAuxData *transferdata); + + +typedef int (PyArrayMethod_GetLoop)( + PyArrayMethod_Context *context, + int aligned, int move_references, + const npy_intp *strides, + PyArrayMethod_StridedLoop **out_loop, + NpyAuxData **out_transferdata, + NPY_ARRAYMETHOD_FLAGS *flags); + +/** + * Query an ArrayMethod for the initial value for use in reduction. + * + * @param context The arraymethod context, mainly to access the descriptors. + * @param reduction_is_empty Whether the reduction is empty. When it is, the + * value returned may differ. In this case it is a "default" value that + * may differ from the "identity" value normally used. For example: + * - `0.0` is the default for `sum([])`. But `-0.0` is the correct + * identity otherwise as it preserves the sign for `sum([-0.0])`. + * - We use no identity for object, but return the default of `0` and `1` + * for the empty `sum([], dtype=object)` and `prod([], dtype=object)`. + * This allows `np.sum(np.array(["a", "b"], dtype=object))` to work. + * - `-inf` or `INT_MIN` for `max` is an identity, but at least `INT_MIN` + * not a good *default* when there are no items. + * @param initial Pointer to initial data to be filled (if possible) + * + * @returns -1, 0, or 1 indicating error, no initial value, and initial being + * successfully filled. Errors must not be given where 0 is correct, NumPy + * may call this even when not strictly necessary. + */ +typedef int (PyArrayMethod_GetReductionInitial)( + PyArrayMethod_Context *context, npy_bool reduction_is_empty, + void *initial); + +/* + * The following functions are only used by the wrapping array method defined + * in umath/wrapping_array_method.c + */ + + +/* + * The function to convert the given descriptors (passed in to + * `resolve_descriptors`) and translates them for the wrapped loop. + * The new descriptors MUST be viewable with the old ones, `NULL` must be + * supported (for outputs) and should normally be forwarded. + * + * The function must clean up on error. + * + * NOTE: We currently assume that this translation gives "viewable" results. + * I.e. there is no additional casting related to the wrapping process. + * In principle that could be supported, but not sure it is useful. + * This currently also means that e.g. alignment must apply identically + * to the new dtypes. + * + * TODO: Due to the fact that `resolve_descriptors` is also used for `can_cast` + * there is no way to "pass out" the result of this function. This means + * it will be called twice for every ufunc call. + * (I am considering including `auxdata` as an "optional" parameter to + * `resolve_descriptors`, so that it can be filled there if not NULL.) + */ +typedef int (PyArrayMethod_TranslateGivenDescriptors)(int nin, int nout, + PyArray_DTypeMeta *const wrapped_dtypes[], + PyArray_Descr *const given_descrs[], PyArray_Descr *new_descrs[]); + +/** + * The function to convert the actual loop descriptors (as returned by the + * original `resolve_descriptors` function) to the ones the output array + * should use. + * This function must return "viewable" types, it must not mutate them in any + * form that would break the inner-loop logic. Does not need to support NULL. + * + * The function must clean up on error. + * + * @param nin Number of input arguments + * @param nout Number of output arguments + * @param new_dtypes The DTypes of the output (usually probably not needed) + * @param given_descrs Original given_descrs to the resolver, necessary to + * fetch any information related to the new dtypes from the original. + * @param original_descrs The `loop_descrs` returned by the wrapped loop. + * @param loop_descrs The output descriptors, compatible to `original_descrs`. + * + * @returns 0 on success, -1 on failure. + */ +typedef int (PyArrayMethod_TranslateLoopDescriptors)(int nin, int nout, + PyArray_DTypeMeta *const new_dtypes[], PyArray_Descr *const given_descrs[], + PyArray_Descr *original_descrs[], PyArray_Descr *loop_descrs[]); + + + +/* + * A traverse loop working on a single array. This is similar to the general + * strided-loop function. This is designed for loops that need to visit every + * element of a single array. + * + * Currently this is used for array clearing, via the NPY_DT_get_clear_loop + * API hook, and zero-filling, via the NPY_DT_get_fill_zero_loop API hook. + * These are most useful for handling arrays storing embedded references to + * python objects or heap-allocated data. + * + * The `void *traverse_context` is passed in because we may need to pass in + * Interpreter state or similar in the future, but we don't want to pass in + * a full context (with pointers to dtypes, method, caller which all make + * no sense for a traverse function). + * + * We assume for now that this context can be just passed through in the + * the future (for structured dtypes). + * + */ +typedef int (PyArrayMethod_TraverseLoop)( + void *traverse_context, const PyArray_Descr *descr, char *data, + npy_intp size, npy_intp stride, NpyAuxData *auxdata); + + +/* + * Simplified get_loop function specific to dtype traversal + * + * It should set the flags needed for the traversal loop and set out_loop to the + * loop function, which must be a valid PyArrayMethod_TraverseLoop + * pointer. Currently this is used for zero-filling and clearing arrays storing + * embedded references. + * + */ +typedef int (PyArrayMethod_GetTraverseLoop)( + void *traverse_context, const PyArray_Descr *descr, + int aligned, npy_intp fixed_stride, + PyArrayMethod_TraverseLoop **out_loop, NpyAuxData **out_auxdata, + NPY_ARRAYMETHOD_FLAGS *flags); + + +/* + * Type of the C promoter function, which must be wrapped into a + * PyCapsule with name "numpy._ufunc_promoter". + * + * Note that currently the output dtypes are always NULL unless they are + * also part of the signature. This is an implementation detail and could + * change in the future. However, in general promoters should not have a + * need for output dtypes. + * (There are potential use-cases, these are currently unsupported.) + */ +typedef int (PyArrayMethod_PromoterFunction)(PyObject *ufunc, + PyArray_DTypeMeta *const op_dtypes[], PyArray_DTypeMeta *const signature[], + PyArray_DTypeMeta *new_op_dtypes[]); + +/* + * **************************** + * DTYPE API + * **************************** + */ + +#define NPY_DT_ABSTRACT 1 << 1 +#define NPY_DT_PARAMETRIC 1 << 2 +#define NPY_DT_NUMERIC 1 << 3 + +/* + * These correspond to slots in the NPY_DType_Slots struct and must + * be in the same order as the members of that struct. If new slots + * get added or old slots get removed NPY_NUM_DTYPE_SLOTS must also + * be updated + */ + +#define NPY_DT_discover_descr_from_pyobject 1 +// this slot is considered private because its API hasn't been decided +#define _NPY_DT_is_known_scalar_type 2 +#define NPY_DT_default_descr 3 +#define NPY_DT_common_dtype 4 +#define NPY_DT_common_instance 5 +#define NPY_DT_ensure_canonical 6 +#define NPY_DT_setitem 7 +#define NPY_DT_getitem 8 +#define NPY_DT_get_clear_loop 9 +#define NPY_DT_get_fill_zero_loop 10 +#define NPY_DT_finalize_descr 11 + +// These PyArray_ArrFunc slots will be deprecated and replaced eventually +// getitem and setitem can be defined as a performance optimization; +// by default the user dtypes call `legacy_getitem_using_DType` and +// `legacy_setitem_using_DType`, respectively. This functionality is +// only supported for basic NumPy DTypes. + + +// used to separate dtype slots from arrfuncs slots +// intended only for internal use but defined here for clarity +#define _NPY_DT_ARRFUNCS_OFFSET (1 << 10) + +// Cast is disabled +// #define NPY_DT_PyArray_ArrFuncs_cast 0 + _NPY_DT_ARRFUNCS_OFFSET + +#define NPY_DT_PyArray_ArrFuncs_getitem 1 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_setitem 2 + _NPY_DT_ARRFUNCS_OFFSET + +// Copyswap is disabled +// #define NPY_DT_PyArray_ArrFuncs_copyswapn 3 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_copyswap 4 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_compare 5 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argmax 6 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_dotfunc 7 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_scanfunc 8 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fromstr 9 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_nonzero 10 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fill 11 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_fillwithscalar 12 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_sort 13 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argsort 14 + _NPY_DT_ARRFUNCS_OFFSET + +// Casting related slots are disabled. See +// https://github.com/numpy/numpy/pull/23173#discussion_r1101098163 +// #define NPY_DT_PyArray_ArrFuncs_castdict 15 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_scalarkind 16 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_cancastscalarkindto 17 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_cancastto 18 + _NPY_DT_ARRFUNCS_OFFSET + +// These are deprecated in NumPy 1.19, so are disabled here. +// #define NPY_DT_PyArray_ArrFuncs_fastclip 19 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_fastputmask 20 + _NPY_DT_ARRFUNCS_OFFSET +// #define NPY_DT_PyArray_ArrFuncs_fasttake 21 + _NPY_DT_ARRFUNCS_OFFSET +#define NPY_DT_PyArray_ArrFuncs_argmin 22 + _NPY_DT_ARRFUNCS_OFFSET + + +// TODO: These slots probably still need some thought, and/or a way to "grow"? +typedef struct { + PyTypeObject *typeobj; /* type of python scalar or NULL */ + int flags; /* flags, including parametric and abstract */ + /* NULL terminated cast definitions. Use NULL for the newly created DType */ + PyArrayMethod_Spec **casts; + PyType_Slot *slots; + /* Baseclass or NULL (will always subclass `np.dtype`) */ + PyTypeObject *baseclass; +} PyArrayDTypeMeta_Spec; + + +typedef PyArray_Descr *(PyArrayDTypeMeta_DiscoverDescrFromPyobject)( + PyArray_DTypeMeta *cls, PyObject *obj); + +/* + * Before making this public, we should decide whether it should pass + * the type, or allow looking at the object. A possible use-case: + * `np.array(np.array([0]), dtype=np.ndarray)` + * Could consider arrays that are not `dtype=ndarray` "scalars". + */ +typedef int (PyArrayDTypeMeta_IsKnownScalarType)( + PyArray_DTypeMeta *cls, PyTypeObject *obj); + +typedef PyArray_Descr *(PyArrayDTypeMeta_DefaultDescriptor)(PyArray_DTypeMeta *cls); +typedef PyArray_DTypeMeta *(PyArrayDTypeMeta_CommonDType)( + PyArray_DTypeMeta *dtype1, PyArray_DTypeMeta *dtype2); + + +/* + * Convenience utility for getting a reference to the DType metaclass associated + * with a dtype instance. + */ +#define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr)) + +static inline PyArray_DTypeMeta * +NPY_DT_NewRef(PyArray_DTypeMeta *o) { + Py_INCREF((PyObject *)o); + return o; +} + + +typedef PyArray_Descr *(PyArrayDTypeMeta_CommonInstance)( + PyArray_Descr *dtype1, PyArray_Descr *dtype2); +typedef PyArray_Descr *(PyArrayDTypeMeta_EnsureCanonical)(PyArray_Descr *dtype); +/* + * Returns either a new reference to *dtype* or a new descriptor instance + * initialized with the same parameters as *dtype*. The caller cannot know + * which choice a dtype will make. This function is called just before the + * array buffer is created for a newly created array, it is not called for + * views and the descriptor returned by this function is attached to the array. + */ +typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtype); + +/* + * TODO: These two functions are currently only used for experimental DType + * API support. Their relation should be "reversed": NumPy should + * always use them internally. + * There are open points about "casting safety" though, e.g. setting + * elements is currently always unsafe. + */ +typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *); +typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *); + +#endif /* NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/halffloat.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/halffloat.h new file mode 100644 index 0000000000000000000000000000000000000000..950401664e101d17ce5461ecb21c7d3bdf824f42 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/halffloat.h @@ -0,0 +1,70 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Half-precision routines + */ + +/* Conversions */ +float npy_half_to_float(npy_half h); +double npy_half_to_double(npy_half h); +npy_half npy_float_to_half(float f); +npy_half npy_double_to_half(double d); +/* Comparisons */ +int npy_half_eq(npy_half h1, npy_half h2); +int npy_half_ne(npy_half h1, npy_half h2); +int npy_half_le(npy_half h1, npy_half h2); +int npy_half_lt(npy_half h1, npy_half h2); +int npy_half_ge(npy_half h1, npy_half h2); +int npy_half_gt(npy_half h1, npy_half h2); +/* faster *_nonan variants for when you know h1 and h2 are not NaN */ +int npy_half_eq_nonan(npy_half h1, npy_half h2); +int npy_half_lt_nonan(npy_half h1, npy_half h2); +int npy_half_le_nonan(npy_half h1, npy_half h2); +/* Miscellaneous functions */ +int npy_half_iszero(npy_half h); +int npy_half_isnan(npy_half h); +int npy_half_isinf(npy_half h); +int npy_half_isfinite(npy_half h); +int npy_half_signbit(npy_half h); +npy_half npy_half_copysign(npy_half x, npy_half y); +npy_half npy_half_spacing(npy_half h); +npy_half npy_half_nextafter(npy_half x, npy_half y); +npy_half npy_half_divmod(npy_half x, npy_half y, npy_half *modulus); + +/* + * Half-precision constants + */ + +#define NPY_HALF_ZERO (0x0000u) +#define NPY_HALF_PZERO (0x0000u) +#define NPY_HALF_NZERO (0x8000u) +#define NPY_HALF_ONE (0x3c00u) +#define NPY_HALF_NEGONE (0xbc00u) +#define NPY_HALF_PINF (0x7c00u) +#define NPY_HALF_NINF (0xfc00u) +#define NPY_HALF_NAN (0x7e00u) + +#define NPY_MAX_HALF (0x7bffu) + +/* + * Bit-level conversions + */ + +npy_uint16 npy_floatbits_to_halfbits(npy_uint32 f); +npy_uint16 npy_doublebits_to_halfbits(npy_uint64 d); +npy_uint32 npy_halfbits_to_floatbits(npy_uint16 h); +npy_uint64 npy_halfbits_to_doublebits(npy_uint16 h); + +#ifdef __cplusplus +} +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_HALFFLOAT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarrayobject.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarrayobject.h new file mode 100644 index 0000000000000000000000000000000000000000..f06bafe5b52a4ab930c2aad0c45d95d0b6d8104b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarrayobject.h @@ -0,0 +1,304 @@ +/* + * DON'T INCLUDE THIS DIRECTLY. + */ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ + +#ifdef __cplusplus +extern "C" { +#endif + +#include +#include "ndarraytypes.h" +#include "dtype_api.h" + +/* Includes the "function" C-API -- these are all stored in a + list of pointers --- one for each file + The two lists are concatenated into one in multiarray. + + They are available as import_array() +*/ + +#include "__multiarray_api.h" + +/* + * Include any definitions which are defined differently for 1.x and 2.x + * (Symbols only available on 2.x are not there, but rather guarded.) + */ +#include "npy_2_compat.h" + +/* C-API that requires previous API to be defined */ + +#define PyArray_DescrCheck(op) PyObject_TypeCheck(op, &PyArrayDescr_Type) + +#define PyArray_Check(op) PyObject_TypeCheck(op, &PyArray_Type) +#define PyArray_CheckExact(op) (((PyObject*)(op))->ob_type == &PyArray_Type) + +#define PyArray_HasArrayInterfaceType(op, type, context, out) \ + ((((out)=PyArray_FromStructInterface(op)) != Py_NotImplemented) || \ + (((out)=PyArray_FromInterface(op)) != Py_NotImplemented) || \ + (((out)=PyArray_FromArrayAttr(op, type, context)) != \ + Py_NotImplemented)) + +#define PyArray_HasArrayInterface(op, out) \ + PyArray_HasArrayInterfaceType(op, NULL, NULL, out) + +#define PyArray_IsZeroDim(op) (PyArray_Check(op) && \ + (PyArray_NDIM((PyArrayObject *)op) == 0)) + +#define PyArray_IsScalar(obj, cls) \ + (PyObject_TypeCheck(obj, &Py##cls##ArrType_Type)) + +#define PyArray_CheckScalar(m) (PyArray_IsScalar(m, Generic) || \ + PyArray_IsZeroDim(m)) +#define PyArray_IsPythonNumber(obj) \ + (PyFloat_Check(obj) || PyComplex_Check(obj) || \ + PyLong_Check(obj) || PyBool_Check(obj)) +#define PyArray_IsIntegerScalar(obj) (PyLong_Check(obj) \ + || PyArray_IsScalar((obj), Integer)) +#define PyArray_IsPythonScalar(obj) \ + (PyArray_IsPythonNumber(obj) || PyBytes_Check(obj) || \ + PyUnicode_Check(obj)) + +#define PyArray_IsAnyScalar(obj) \ + (PyArray_IsScalar(obj, Generic) || PyArray_IsPythonScalar(obj)) + +#define PyArray_CheckAnyScalar(obj) (PyArray_IsPythonScalar(obj) || \ + PyArray_CheckScalar(obj)) + + +#define PyArray_GETCONTIGUOUS(m) (PyArray_ISCONTIGUOUS(m) ? \ + Py_INCREF(m), (m) : \ + (PyArrayObject *)(PyArray_Copy(m))) + +#define PyArray_SAMESHAPE(a1,a2) ((PyArray_NDIM(a1) == PyArray_NDIM(a2)) && \ + PyArray_CompareLists(PyArray_DIMS(a1), \ + PyArray_DIMS(a2), \ + PyArray_NDIM(a1))) + +#define PyArray_SIZE(m) PyArray_MultiplyList(PyArray_DIMS(m), PyArray_NDIM(m)) +#define PyArray_NBYTES(m) (PyArray_ITEMSIZE(m) * PyArray_SIZE(m)) +#define PyArray_FROM_O(m) PyArray_FromAny(m, NULL, 0, 0, 0, NULL) + +#define PyArray_FROM_OF(m,flags) PyArray_CheckFromAny(m, NULL, 0, 0, flags, \ + NULL) + +#define PyArray_FROM_OT(m,type) PyArray_FromAny(m, \ + PyArray_DescrFromType(type), 0, 0, 0, NULL) + +#define PyArray_FROM_OTF(m, type, flags) \ + PyArray_FromAny(m, PyArray_DescrFromType(type), 0, 0, \ + (((flags) & NPY_ARRAY_ENSURECOPY) ? \ + ((flags) | NPY_ARRAY_DEFAULT) : (flags)), NULL) + +#define PyArray_FROMANY(m, type, min, max, flags) \ + PyArray_FromAny(m, PyArray_DescrFromType(type), min, max, \ + (((flags) & NPY_ARRAY_ENSURECOPY) ? \ + (flags) | NPY_ARRAY_DEFAULT : (flags)), NULL) + +#define PyArray_ZEROS(m, dims, type, is_f_order) \ + PyArray_Zeros(m, dims, PyArray_DescrFromType(type), is_f_order) + +#define PyArray_EMPTY(m, dims, type, is_f_order) \ + PyArray_Empty(m, dims, PyArray_DescrFromType(type), is_f_order) + +#define PyArray_FILLWBYTE(obj, val) memset(PyArray_DATA(obj), val, \ + PyArray_NBYTES(obj)) + +#define PyArray_ContiguousFromAny(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_DEFAULT, NULL) + +#define PyArray_EquivArrTypes(a1, a2) \ + PyArray_EquivTypes(PyArray_DESCR(a1), PyArray_DESCR(a2)) + +#define PyArray_EquivByteorders(b1, b2) \ + (((b1) == (b2)) || (PyArray_ISNBO(b1) == PyArray_ISNBO(b2))) + +#define PyArray_SimpleNew(nd, dims, typenum) \ + PyArray_New(&PyArray_Type, nd, dims, typenum, NULL, NULL, 0, 0, NULL) + +#define PyArray_SimpleNewFromData(nd, dims, typenum, data) \ + PyArray_New(&PyArray_Type, nd, dims, typenum, NULL, \ + data, 0, NPY_ARRAY_CARRAY, NULL) + +#define PyArray_SimpleNewFromDescr(nd, dims, descr) \ + PyArray_NewFromDescr(&PyArray_Type, descr, nd, dims, \ + NULL, NULL, 0, NULL) + +#define PyArray_ToScalar(data, arr) \ + PyArray_Scalar(data, PyArray_DESCR(arr), (PyObject *)arr) + + +/* These might be faster without the dereferencing of obj + going on inside -- of course an optimizing compiler should + inline the constants inside a for loop making it a moot point +*/ + +#define PyArray_GETPTR1(obj, i) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0])) + +#define PyArray_GETPTR2(obj, i, j) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1])) + +#define PyArray_GETPTR3(obj, i, j, k) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1] + \ + (k)*PyArray_STRIDES(obj)[2])) + +#define PyArray_GETPTR4(obj, i, j, k, l) ((void *)(PyArray_BYTES(obj) + \ + (i)*PyArray_STRIDES(obj)[0] + \ + (j)*PyArray_STRIDES(obj)[1] + \ + (k)*PyArray_STRIDES(obj)[2] + \ + (l)*PyArray_STRIDES(obj)[3])) + +static inline void +PyArray_DiscardWritebackIfCopy(PyArrayObject *arr) +{ + PyArrayObject_fields *fa = (PyArrayObject_fields *)arr; + if (fa && fa->base) { + if (fa->flags & NPY_ARRAY_WRITEBACKIFCOPY) { + PyArray_ENABLEFLAGS((PyArrayObject*)fa->base, NPY_ARRAY_WRITEABLE); + Py_DECREF(fa->base); + fa->base = NULL; + PyArray_CLEARFLAGS(arr, NPY_ARRAY_WRITEBACKIFCOPY); + } + } +} + +#define PyArray_DESCR_REPLACE(descr) do { \ + PyArray_Descr *_new_; \ + _new_ = PyArray_DescrNew(descr); \ + Py_XDECREF(descr); \ + descr = _new_; \ + } while(0) + +/* Copy should always return contiguous array */ +#define PyArray_Copy(obj) PyArray_NewCopy(obj, NPY_CORDER) + +#define PyArray_FromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_BEHAVED | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_ContiguousFromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_DEFAULT | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_CopyFromObject(op, type, min_depth, max_depth) \ + PyArray_FromAny(op, PyArray_DescrFromType(type), min_depth, \ + max_depth, NPY_ARRAY_ENSURECOPY | \ + NPY_ARRAY_DEFAULT | \ + NPY_ARRAY_ENSUREARRAY, NULL) + +#define PyArray_Cast(mp, type_num) \ + PyArray_CastToType(mp, PyArray_DescrFromType(type_num), 0) + +#define PyArray_Take(ap, items, axis) \ + PyArray_TakeFrom(ap, items, axis, NULL, NPY_RAISE) + +#define PyArray_Put(ap, items, values) \ + PyArray_PutTo(ap, items, values, NPY_RAISE) + + +/* + Check to see if this key in the dictionary is the "title" + entry of the tuple (i.e. a duplicate dictionary entry in the fields + dict). +*/ + +static inline int +NPY_TITLE_KEY_check(PyObject *key, PyObject *value) +{ + PyObject *title; + if (PyTuple_Size(value) != 3) { + return 0; + } + title = PyTuple_GetItem(value, 2); + if (key == title) { + return 1; + } +#ifdef PYPY_VERSION + /* + * On PyPy, dictionary keys do not always preserve object identity. + * Fall back to comparison by value. + */ + if (PyUnicode_Check(title) && PyUnicode_Check(key)) { + return PyUnicode_Compare(title, key) == 0 ? 1 : 0; + } +#endif + return 0; +} + +/* Macro, for backward compat with "if NPY_TITLE_KEY(key, value) { ..." */ +#define NPY_TITLE_KEY(key, value) (NPY_TITLE_KEY_check((key), (value))) + +#define DEPRECATE(msg) PyErr_WarnEx(PyExc_DeprecationWarning,msg,1) +#define DEPRECATE_FUTUREWARNING(msg) PyErr_WarnEx(PyExc_FutureWarning,msg,1) + + +/* + * These macros and functions unfortunately require runtime version checks + * that are only defined in `npy_2_compat.h`. For that reasons they cannot be + * part of `ndarraytypes.h` which tries to be self contained. + */ + +static inline npy_intp +PyArray_ITEMSIZE(const PyArrayObject *arr) +{ + return PyDataType_ELSIZE(((PyArrayObject_fields *)arr)->descr); +} + +#define PyDataType_HASFIELDS(obj) (PyDataType_ISLEGACY((PyArray_Descr*)(obj)) && PyDataType_NAMES((PyArray_Descr*)(obj)) != NULL) +#define PyDataType_HASSUBARRAY(dtype) (PyDataType_ISLEGACY(dtype) && PyDataType_SUBARRAY(dtype) != NULL) +#define PyDataType_ISUNSIZED(dtype) ((dtype)->elsize == 0 && \ + !PyDataType_HASFIELDS(dtype)) + +#define PyDataType_FLAGCHK(dtype, flag) \ + ((PyDataType_FLAGS(dtype) & (flag)) == (flag)) + +#define PyDataType_REFCHK(dtype) \ + PyDataType_FLAGCHK(dtype, NPY_ITEM_REFCOUNT) + +#define NPY_BEGIN_THREADS_DESCR(dtype) \ + do {if (!(PyDataType_FLAGCHK((dtype), NPY_NEEDS_PYAPI))) \ + NPY_BEGIN_THREADS;} while (0); + +#define NPY_END_THREADS_DESCR(dtype) \ + do {if (!(PyDataType_FLAGCHK((dtype), NPY_NEEDS_PYAPI))) \ + NPY_END_THREADS; } while (0); + +#if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) +/* The internal copy of this is now defined in `dtypemeta.h` */ +/* + * `PyArray_Scalar` is the same as this function but converts will convert + * most NumPy types to Python scalars. + */ +static inline PyObject * +PyArray_GETITEM(const PyArrayObject *arr, const char *itemptr) +{ + return PyDataType_GetArrFuncs(((PyArrayObject_fields *)arr)->descr)->getitem( + (void *)itemptr, (PyArrayObject *)arr); +} + +/* + * SETITEM should only be used if it is known that the value is a scalar + * and of a type understood by the arrays dtype. + * Use `PyArray_Pack` if the value may be of a different dtype. + */ +static inline int +PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v) +{ + return PyDataType_GetArrFuncs(((PyArrayObject_fields *)arr)->descr)->setitem(v, itemptr, arr); +} +#endif /* not internal */ + + +#ifdef __cplusplus +} +#endif + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NDARRAYOBJECT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarraytypes.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarraytypes.h new file mode 100644 index 0000000000000000000000000000000000000000..37788a74557f4361f78d611470ffe9318f3c8752 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ndarraytypes.h @@ -0,0 +1,1933 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NDARRAYTYPES_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NDARRAYTYPES_H_ + +#include "npy_common.h" +#include "npy_endian.h" +#include "npy_cpu.h" +#include "utils.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define NPY_NO_EXPORT NPY_VISIBILITY_HIDDEN + +/* Always allow threading unless it was explicitly disabled at build time */ +#if !NPY_NO_SMP + #define NPY_ALLOW_THREADS 1 +#else + #define NPY_ALLOW_THREADS 0 +#endif + +#ifndef __has_extension +#define __has_extension(x) 0 +#endif + +/* + * There are several places in the code where an array of dimensions + * is allocated statically. This is the size of that static + * allocation. + * + * The array creation itself could have arbitrary dimensions but all + * the places where static allocation is used would need to be changed + * to dynamic (including inside of several structures) + * + * As of NumPy 2.0, we strongly discourage the downstream use of NPY_MAXDIMS, + * but since auditing everything seems a big ask, define it as 64. + * A future version could: + * - Increase or remove the limit and require recompilation (like 2.0 did) + * - Deprecate or remove the macro but keep the limit (at basically any time) + */ +#define NPY_MAXDIMS 64 +/* We cannot change this as it would break ABI: */ +#define NPY_MAXDIMS_LEGACY_ITERS 32 +/* NPY_MAXARGS is version dependent and defined in npy_2_compat.h */ + +/* Used for Converter Functions "O&" code in ParseTuple */ +#define NPY_FAIL 0 +#define NPY_SUCCEED 1 + + +enum NPY_TYPES { NPY_BOOL=0, + 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=17, + NPY_STRING, NPY_UNICODE, + NPY_VOID, + /* + * New 1.6 types appended, may be integrated + * into the above in 2.0. + */ + NPY_DATETIME, NPY_TIMEDELTA, NPY_HALF, + + NPY_CHAR, /* Deprecated, will raise if used */ + + /* The number of *legacy* dtypes */ + NPY_NTYPES_LEGACY=24, + + /* assign a high value to avoid changing this in the + future when new dtypes are added */ + NPY_NOTYPE=25, + + NPY_USERDEF=256, /* leave room for characters */ + + /* The number of types not including the new 1.6 types */ + NPY_NTYPES_ABI_COMPATIBLE=21, + + /* + * New DTypes which do not share the legacy layout + * (added after NumPy 2.0). VSTRING is the first of these + * we may open up a block for user-defined dtypes in the + * future. + */ + NPY_VSTRING=2056, +}; + + +/* basetype array priority */ +#define NPY_PRIORITY 0.0 + +/* default subtype priority */ +#define NPY_SUBTYPE_PRIORITY 1.0 + +/* default scalar priority */ +#define NPY_SCALAR_PRIORITY -1000000.0 + +/* How many floating point types are there (excluding half) */ +#define NPY_NUM_FLOATTYPE 3 + +/* + * These characters correspond to the array type and the struct + * module + */ + +enum NPY_TYPECHAR { + NPY_BOOLLTR = '?', + NPY_BYTELTR = 'b', + NPY_UBYTELTR = 'B', + NPY_SHORTLTR = 'h', + NPY_USHORTLTR = 'H', + NPY_INTLTR = 'i', + NPY_UINTLTR = 'I', + NPY_LONGLTR = 'l', + NPY_ULONGLTR = 'L', + NPY_LONGLONGLTR = 'q', + NPY_ULONGLONGLTR = 'Q', + NPY_HALFLTR = 'e', + NPY_FLOATLTR = 'f', + NPY_DOUBLELTR = 'd', + NPY_LONGDOUBLELTR = 'g', + NPY_CFLOATLTR = 'F', + NPY_CDOUBLELTR = 'D', + NPY_CLONGDOUBLELTR = 'G', + NPY_OBJECTLTR = 'O', + NPY_STRINGLTR = 'S', + NPY_DEPRECATED_STRINGLTR2 = 'a', + NPY_UNICODELTR = 'U', + NPY_VOIDLTR = 'V', + NPY_DATETIMELTR = 'M', + NPY_TIMEDELTALTR = 'm', + NPY_CHARLTR = 'c', + + /* + * New non-legacy DTypes + */ + NPY_VSTRINGLTR = 'T', + + /* + * Note, we removed `NPY_INTPLTR` due to changing its definition + * to 'n', rather than 'p'. On any typical platform this is the + * same integer. 'n' should be used for the `np.intp` with the same + * size as `size_t` while 'p' remains pointer sized. + * + * 'p', 'P', 'n', and 'N' are valid and defined explicitly + * in `arraytypes.c.src`. + */ + + /* + * These are for dtype 'kinds', not dtype 'typecodes' + * as the above are for. + */ + NPY_GENBOOLLTR ='b', + NPY_SIGNEDLTR = 'i', + NPY_UNSIGNEDLTR = 'u', + NPY_FLOATINGLTR = 'f', + NPY_COMPLEXLTR = 'c', + +}; + +/* + * Changing this may break Numpy API compatibility + * due to changing offsets in PyArray_ArrFuncs, so be + * careful. Here we have reused the mergesort slot for + * any kind of stable sort, the actual implementation will + * depend on the data type. + */ +typedef enum { + _NPY_SORT_UNDEFINED=-1, + NPY_QUICKSORT=0, + NPY_HEAPSORT=1, + NPY_MERGESORT=2, + NPY_STABLESORT=2, +} NPY_SORTKIND; +#define NPY_NSORTS (NPY_STABLESORT + 1) + + +typedef enum { + NPY_INTROSELECT=0 +} NPY_SELECTKIND; +#define NPY_NSELECTS (NPY_INTROSELECT + 1) + + +typedef enum { + NPY_SEARCHLEFT=0, + NPY_SEARCHRIGHT=1 +} NPY_SEARCHSIDE; +#define NPY_NSEARCHSIDES (NPY_SEARCHRIGHT + 1) + + +typedef enum { + NPY_NOSCALAR=-1, + NPY_BOOL_SCALAR, + NPY_INTPOS_SCALAR, + NPY_INTNEG_SCALAR, + NPY_FLOAT_SCALAR, + NPY_COMPLEX_SCALAR, + NPY_OBJECT_SCALAR +} NPY_SCALARKIND; +#define NPY_NSCALARKINDS (NPY_OBJECT_SCALAR + 1) + +/* For specifying array memory layout or iteration order */ +typedef enum { + /* Fortran order if inputs are all Fortran, C otherwise */ + NPY_ANYORDER=-1, + /* C order */ + NPY_CORDER=0, + /* Fortran order */ + NPY_FORTRANORDER=1, + /* An order as close to the inputs as possible */ + NPY_KEEPORDER=2 +} NPY_ORDER; + +/* For specifying allowed casting in operations which support it */ +typedef enum { + _NPY_ERROR_OCCURRED_IN_CAST = -1, + /* Only allow identical types */ + NPY_NO_CASTING=0, + /* Allow identical and byte swapped types */ + NPY_EQUIV_CASTING=1, + /* Only allow safe casts */ + NPY_SAFE_CASTING=2, + /* Allow safe casts or casts within the same kind */ + NPY_SAME_KIND_CASTING=3, + /* Allow any casts */ + NPY_UNSAFE_CASTING=4, +} NPY_CASTING; + +typedef enum { + NPY_CLIP=0, + NPY_WRAP=1, + NPY_RAISE=2 +} NPY_CLIPMODE; + +typedef enum { + NPY_VALID=0, + NPY_SAME=1, + NPY_FULL=2 +} NPY_CORRELATEMODE; + +/* The special not-a-time (NaT) value */ +#define NPY_DATETIME_NAT NPY_MIN_INT64 + +/* + * Upper bound on the length of a DATETIME ISO 8601 string + * YEAR: 21 (64-bit year) + * MONTH: 3 + * DAY: 3 + * HOURS: 3 + * MINUTES: 3 + * SECONDS: 3 + * ATTOSECONDS: 1 + 3*6 + * TIMEZONE: 5 + * NULL TERMINATOR: 1 + */ +#define NPY_DATETIME_MAX_ISO8601_STRLEN (21 + 3*5 + 1 + 3*6 + 6 + 1) + +/* The FR in the unit names stands for frequency */ +typedef enum { + /* Force signed enum type, must be -1 for code compatibility */ + NPY_FR_ERROR = -1, /* error or undetermined */ + + /* Start of valid units */ + NPY_FR_Y = 0, /* Years */ + NPY_FR_M = 1, /* Months */ + NPY_FR_W = 2, /* Weeks */ + /* Gap where 1.6 NPY_FR_B (value 3) was */ + NPY_FR_D = 4, /* Days */ + NPY_FR_h = 5, /* hours */ + NPY_FR_m = 6, /* minutes */ + NPY_FR_s = 7, /* seconds */ + NPY_FR_ms = 8, /* milliseconds */ + NPY_FR_us = 9, /* microseconds */ + NPY_FR_ns = 10, /* nanoseconds */ + NPY_FR_ps = 11, /* picoseconds */ + NPY_FR_fs = 12, /* femtoseconds */ + NPY_FR_as = 13, /* attoseconds */ + NPY_FR_GENERIC = 14 /* unbound units, can convert to anything */ +} NPY_DATETIMEUNIT; + +/* + * NOTE: With the NPY_FR_B gap for 1.6 ABI compatibility, NPY_DATETIME_NUMUNITS + * is technically one more than the actual number of units. + */ +#define NPY_DATETIME_NUMUNITS (NPY_FR_GENERIC + 1) +#define NPY_DATETIME_DEFAULTUNIT NPY_FR_GENERIC + +/* + * Business day conventions for mapping invalid business + * days to valid business days. + */ +typedef enum { + /* Go forward in time to the following business day. */ + NPY_BUSDAY_FORWARD, + NPY_BUSDAY_FOLLOWING = NPY_BUSDAY_FORWARD, + /* Go backward in time to the preceding business day. */ + NPY_BUSDAY_BACKWARD, + NPY_BUSDAY_PRECEDING = NPY_BUSDAY_BACKWARD, + /* + * Go forward in time to the following business day, unless it + * crosses a month boundary, in which case go backward + */ + NPY_BUSDAY_MODIFIEDFOLLOWING, + /* + * Go backward in time to the preceding business day, unless it + * crosses a month boundary, in which case go forward. + */ + NPY_BUSDAY_MODIFIEDPRECEDING, + /* Produce a NaT for non-business days. */ + NPY_BUSDAY_NAT, + /* Raise an exception for non-business days. */ + NPY_BUSDAY_RAISE +} NPY_BUSDAY_ROLL; + + +/************************************************************ + * NumPy Auxiliary Data for inner loops, sort functions, etc. + ************************************************************/ + +/* + * When creating an auxiliary data struct, this should always appear + * as the first member, like this: + * + * typedef struct { + * NpyAuxData base; + * double constant; + * } constant_multiplier_aux_data; + */ +typedef struct NpyAuxData_tag NpyAuxData; + +/* Function pointers for freeing or cloning auxiliary data */ +typedef void (NpyAuxData_FreeFunc) (NpyAuxData *); +typedef NpyAuxData *(NpyAuxData_CloneFunc) (NpyAuxData *); + +struct NpyAuxData_tag { + NpyAuxData_FreeFunc *free; + NpyAuxData_CloneFunc *clone; + /* To allow for a bit of expansion without breaking the ABI */ + void *reserved[2]; +}; + +/* Macros to use for freeing and cloning auxiliary data */ +#define NPY_AUXDATA_FREE(auxdata) \ + do { \ + if ((auxdata) != NULL) { \ + (auxdata)->free(auxdata); \ + } \ + } while(0) +#define NPY_AUXDATA_CLONE(auxdata) \ + ((auxdata)->clone(auxdata)) + +#define NPY_ERR(str) fprintf(stderr, #str); fflush(stderr); +#define NPY_ERR2(str) fprintf(stderr, str); fflush(stderr); + +/* +* Macros to define how array, and dimension/strides data is +* allocated. These should be made private +*/ + +#define NPY_USE_PYMEM 1 + + +#if NPY_USE_PYMEM == 1 +/* use the Raw versions which are safe to call with the GIL released */ +#define PyArray_malloc PyMem_RawMalloc +#define PyArray_free PyMem_RawFree +#define PyArray_realloc PyMem_RawRealloc +#else +#define PyArray_malloc malloc +#define PyArray_free free +#define PyArray_realloc realloc +#endif + +/* Dimensions and strides */ +#define PyDimMem_NEW(size) \ + ((npy_intp *)PyArray_malloc(size*sizeof(npy_intp))) + +#define PyDimMem_FREE(ptr) PyArray_free(ptr) + +#define PyDimMem_RENEW(ptr,size) \ + ((npy_intp *)PyArray_realloc(ptr,size*sizeof(npy_intp))) + +/* forward declaration */ +struct _PyArray_Descr; + +/* These must deal with unaligned and swapped data if necessary */ +typedef PyObject * (PyArray_GetItemFunc) (void *, void *); +typedef int (PyArray_SetItemFunc)(PyObject *, void *, void *); + +typedef void (PyArray_CopySwapNFunc)(void *, npy_intp, void *, npy_intp, + npy_intp, int, void *); + +typedef void (PyArray_CopySwapFunc)(void *, void *, int, void *); +typedef npy_bool (PyArray_NonzeroFunc)(void *, void *); + + +/* + * These assume aligned and notswapped data -- a buffer will be used + * before or contiguous data will be obtained + */ + +typedef int (PyArray_CompareFunc)(const void *, const void *, void *); +typedef int (PyArray_ArgFunc)(void*, npy_intp, npy_intp*, void *); + +typedef void (PyArray_DotFunc)(void *, npy_intp, void *, npy_intp, void *, + npy_intp, void *); + +typedef void (PyArray_VectorUnaryFunc)(void *, void *, npy_intp, void *, + void *); + +/* + * XXX the ignore argument should be removed next time the API version + * is bumped. It used to be the separator. + */ +typedef int (PyArray_ScanFunc)(FILE *fp, void *dptr, + char *ignore, struct _PyArray_Descr *); +typedef int (PyArray_FromStrFunc)(char *s, void *dptr, char **endptr, + struct _PyArray_Descr *); + +typedef int (PyArray_FillFunc)(void *, npy_intp, void *); + +typedef int (PyArray_SortFunc)(void *, npy_intp, void *); +typedef int (PyArray_ArgSortFunc)(void *, npy_intp *, npy_intp, void *); + +typedef int (PyArray_FillWithScalarFunc)(void *, npy_intp, void *, void *); + +typedef int (PyArray_ScalarKindFunc)(void *); + +typedef struct { + npy_intp *ptr; + int len; +} PyArray_Dims; + +typedef struct { + /* + * Functions to cast to most other standard types + * Can have some NULL entries. The types + * DATETIME, TIMEDELTA, and HALF go into the castdict + * even though they are built-in. + */ + PyArray_VectorUnaryFunc *cast[NPY_NTYPES_ABI_COMPATIBLE]; + + /* The next four functions *cannot* be NULL */ + + /* + * Functions to get and set items with standard Python types + * -- not array scalars + */ + PyArray_GetItemFunc *getitem; + PyArray_SetItemFunc *setitem; + + /* + * Copy and/or swap data. Memory areas may not overlap + * Use memmove first if they might + */ + PyArray_CopySwapNFunc *copyswapn; + PyArray_CopySwapFunc *copyswap; + + /* + * Function to compare items + * Can be NULL + */ + PyArray_CompareFunc *compare; + + /* + * Function to select largest + * Can be NULL + */ + PyArray_ArgFunc *argmax; + + /* + * Function to compute dot product + * Can be NULL + */ + PyArray_DotFunc *dotfunc; + + /* + * Function to scan an ASCII file and + * place a single value plus possible separator + * Can be NULL + */ + PyArray_ScanFunc *scanfunc; + + /* + * Function to read a single value from a string + * and adjust the pointer; Can be NULL + */ + PyArray_FromStrFunc *fromstr; + + /* + * Function to determine if data is zero or not + * If NULL a default version is + * used at Registration time. + */ + PyArray_NonzeroFunc *nonzero; + + /* + * Used for arange. Should return 0 on success + * and -1 on failure. + * Can be NULL. + */ + PyArray_FillFunc *fill; + + /* + * Function to fill arrays with scalar values + * Can be NULL + */ + PyArray_FillWithScalarFunc *fillwithscalar; + + /* + * Sorting functions + * Can be NULL + */ + PyArray_SortFunc *sort[NPY_NSORTS]; + PyArray_ArgSortFunc *argsort[NPY_NSORTS]; + + /* + * Dictionary of additional casting functions + * PyArray_VectorUnaryFuncs + * which can be populated to support casting + * to other registered types. Can be NULL + */ + PyObject *castdict; + + /* + * Functions useful for generalizing + * the casting rules. + * Can be NULL; + */ + PyArray_ScalarKindFunc *scalarkind; + int **cancastscalarkindto; + int *cancastto; + + void *_unused1; + void *_unused2; + void *_unused3; + + /* + * Function to select smallest + * Can be NULL + */ + PyArray_ArgFunc *argmin; + +} PyArray_ArrFuncs; + + +/* The item must be reference counted when it is inserted or extracted. */ +#define NPY_ITEM_REFCOUNT 0x01 +/* Same as needing REFCOUNT */ +#define NPY_ITEM_HASOBJECT 0x01 +/* Convert to list for pickling */ +#define NPY_LIST_PICKLE 0x02 +/* The item is a POINTER */ +#define NPY_ITEM_IS_POINTER 0x04 +/* memory needs to be initialized for this data-type */ +#define NPY_NEEDS_INIT 0x08 +/* operations need Python C-API so don't give-up thread. */ +#define NPY_NEEDS_PYAPI 0x10 +/* Use f.getitem when extracting elements of this data-type */ +#define NPY_USE_GETITEM 0x20 +/* Use f.setitem when setting creating 0-d array from this data-type.*/ +#define NPY_USE_SETITEM 0x40 +/* A sticky flag specifically for structured arrays */ +#define NPY_ALIGNED_STRUCT 0x80 + +/* + *These are inherited for global data-type if any data-types in the + * field have them + */ +#define NPY_FROM_FIELDS (NPY_NEEDS_INIT | NPY_LIST_PICKLE | \ + NPY_ITEM_REFCOUNT | NPY_NEEDS_PYAPI) + +#define NPY_OBJECT_DTYPE_FLAGS (NPY_LIST_PICKLE | NPY_USE_GETITEM | \ + NPY_ITEM_IS_POINTER | NPY_ITEM_REFCOUNT | \ + NPY_NEEDS_INIT | NPY_NEEDS_PYAPI) + +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION +/* + * Public version of the Descriptor struct as of 2.x + */ +typedef struct _PyArray_Descr { + PyObject_HEAD + /* + * the type object representing an + * instance of this type -- should not + * be two type_numbers with the same type + * object. + */ + PyTypeObject *typeobj; + /* kind for this type */ + char kind; + /* unique-character representing this type */ + char type; + /* + * '>' (big), '<' (little), '|' + * (not-applicable), or '=' (native). + */ + char byteorder; + /* Former flags flags space (unused) to ensure type_num is stable. */ + char _former_flags; + /* number representing this type */ + int type_num; + /* Space for dtype instance specific flags. */ + npy_uint64 flags; + /* element size (itemsize) for this type */ + npy_intp elsize; + /* alignment needed for this type */ + npy_intp alignment; + /* metadata dict or NULL */ + PyObject *metadata; + /* Cached hash value (-1 if not yet computed). */ + npy_hash_t hash; + /* Unused slot (must be initialized to NULL) for future use */ + void *reserved_null[2]; +} PyArray_Descr; + +#else /* 1.x and 2.x compatible version (only shared fields): */ + +typedef struct _PyArray_Descr { + PyObject_HEAD + PyTypeObject *typeobj; + char kind; + char type; + char byteorder; + char _former_flags; + int type_num; +} PyArray_Descr; + +/* To access modified fields, define the full 2.0 struct: */ +typedef struct { + PyObject_HEAD + PyTypeObject *typeobj; + char kind; + char type; + char byteorder; + char _former_flags; + int type_num; + npy_uint64 flags; + npy_intp elsize; + npy_intp alignment; + PyObject *metadata; + npy_hash_t hash; + void *reserved_null[2]; +} _PyArray_DescrNumPy2; + +#endif /* 1.x and 2.x compatible version */ + +/* + * Semi-private struct with additional field of legacy descriptors (must + * check NPY_DT_is_legacy before casting/accessing). The struct is also not + * valid when running on 1.x (i.e. in public API use). + */ +typedef struct { + PyObject_HEAD + PyTypeObject *typeobj; + char kind; + char type; + char byteorder; + char _former_flags; + int type_num; + npy_uint64 flags; + npy_intp elsize; + npy_intp alignment; + PyObject *metadata; + npy_hash_t hash; + void *reserved_null[2]; + struct _arr_descr *subarray; + PyObject *fields; + PyObject *names; + NpyAuxData *c_metadata; +} _PyArray_LegacyDescr; + + +/* + * Umodified PyArray_Descr struct identical to NumPy 1.x. This struct is + * used as a prototype for registering a new legacy DType. + * It is also used to access the fields in user code running on 1.x. + */ +typedef struct { + PyObject_HEAD + PyTypeObject *typeobj; + char kind; + char type; + char byteorder; + char flags; + int type_num; + int elsize; + int alignment; + struct _arr_descr *subarray; + PyObject *fields; + PyObject *names; + PyArray_ArrFuncs *f; + PyObject *metadata; + NpyAuxData *c_metadata; + npy_hash_t hash; +} PyArray_DescrProto; + + +typedef struct _arr_descr { + PyArray_Descr *base; + PyObject *shape; /* a tuple */ +} PyArray_ArrayDescr; + +/* + * Memory handler structure for array data. + */ +/* The declaration of free differs from PyMemAllocatorEx */ +typedef struct { + void *ctx; + void* (*malloc) (void *ctx, size_t size); + void* (*calloc) (void *ctx, size_t nelem, size_t elsize); + void* (*realloc) (void *ctx, void *ptr, size_t new_size); + void (*free) (void *ctx, void *ptr, size_t size); + /* + * This is the end of the version=1 struct. Only add new fields after + * this line + */ +} PyDataMemAllocator; + +typedef struct { + char name[127]; /* multiple of 64 to keep the struct aligned */ + uint8_t version; /* currently 1 */ + PyDataMemAllocator allocator; +} PyDataMem_Handler; + + +/* + * The main array object structure. + * + * It has been recommended to use the inline functions defined below + * (PyArray_DATA and friends) to access fields here for a number of + * releases. Direct access to the members themselves is deprecated. + * To ensure that your code does not use deprecated access, + * #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION + * (or NPY_1_8_API_VERSION or higher as required). + */ +/* This struct will be moved to a private header in a future release */ +typedef struct tagPyArrayObject_fields { + PyObject_HEAD + /* Pointer to the raw data buffer */ + char *data; + /* The number of dimensions, also called 'ndim' */ + int nd; + /* The size in each dimension, also called 'shape' */ + npy_intp *dimensions; + /* + * Number of bytes to jump to get to the + * next element in each dimension + */ + npy_intp *strides; + /* + * This object is decref'd upon + * deletion of array. Except in the + * case of WRITEBACKIFCOPY which has + * special handling. + * + * For views it points to the original + * array, collapsed so no chains of + * views occur. + * + * For creation from buffer object it + * points to an object that should be + * decref'd on deletion + * + * For WRITEBACKIFCOPY flag this is an + * array to-be-updated upon calling + * PyArray_ResolveWritebackIfCopy + */ + PyObject *base; + /* Pointer to type structure */ + PyArray_Descr *descr; + /* Flags describing array -- see below */ + int flags; + /* For weak references */ + PyObject *weakreflist; +#if NPY_FEATURE_VERSION >= NPY_1_20_API_VERSION + void *_buffer_info; /* private buffer info, tagged to allow warning */ +#endif + /* + * For malloc/calloc/realloc/free per object + */ +#if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION + PyObject *mem_handler; +#endif +} PyArrayObject_fields; + +/* + * To hide the implementation details, we only expose + * the Python struct HEAD. + */ +#if !defined(NPY_NO_DEPRECATED_API) || \ + (NPY_NO_DEPRECATED_API < NPY_1_7_API_VERSION) +/* + * Can't put this in npy_deprecated_api.h like the others. + * PyArrayObject field access is deprecated as of NumPy 1.7. + */ +typedef PyArrayObject_fields PyArrayObject; +#else +typedef struct tagPyArrayObject { + PyObject_HEAD +} PyArrayObject; +#endif + +/* + * Removed 2020-Nov-25, NumPy 1.20 + * #define NPY_SIZEOF_PYARRAYOBJECT (sizeof(PyArrayObject_fields)) + * + * The above macro was removed as it gave a false sense of a stable ABI + * with respect to the structures size. If you require a runtime constant, + * you can use `PyArray_Type.tp_basicsize` instead. Otherwise, please + * see the PyArrayObject documentation or ask the NumPy developers for + * information on how to correctly replace the macro in a way that is + * compatible with multiple NumPy versions. + */ + +/* Mirrors buffer object to ptr */ + +typedef struct { + PyObject_HEAD + PyObject *base; + void *ptr; + npy_intp len; + int flags; +} PyArray_Chunk; + +typedef struct { + NPY_DATETIMEUNIT base; + int num; +} PyArray_DatetimeMetaData; + +typedef struct { + NpyAuxData base; + PyArray_DatetimeMetaData meta; +} PyArray_DatetimeDTypeMetaData; + +/* + * This structure contains an exploded view of a date-time value. + * NaT is represented by year == NPY_DATETIME_NAT. + */ +typedef struct { + npy_int64 year; + npy_int32 month, day, hour, min, sec, us, ps, as; +} npy_datetimestruct; + +/* This structure contains an exploded view of a timedelta value */ +typedef struct { + npy_int64 day; + npy_int32 sec, us, ps, as; +} npy_timedeltastruct; + +typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *); + +/* + * Means c-style contiguous (last index varies the fastest). The data + * elements right after each other. + * + * This flag may be requested in constructor functions. + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_C_CONTIGUOUS 0x0001 + +/* + * Set if array is a contiguous Fortran array: the first index varies + * the fastest in memory (strides array is reverse of C-contiguous + * array) + * + * This flag may be requested in constructor functions. + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_F_CONTIGUOUS 0x0002 + +/* + * Note: all 0-d arrays are C_CONTIGUOUS and F_CONTIGUOUS. If a + * 1-d array is C_CONTIGUOUS it is also F_CONTIGUOUS. Arrays with + * more then one dimension can be C_CONTIGUOUS and F_CONTIGUOUS + * at the same time if they have either zero or one element. + * A higher dimensional array always has the same contiguity flags as + * `array.squeeze()`; dimensions with `array.shape[dimension] == 1` are + * effectively ignored when checking for contiguity. + */ + +/* + * If set, the array owns the data: it will be free'd when the array + * is deleted. + * + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_OWNDATA 0x0004 + +/* + * An array never has the next four set; they're only used as parameter + * flags to the various FromAny functions + * + * This flag may be requested in constructor functions. + */ + +/* Cause a cast to occur regardless of whether or not it is safe. */ +#define NPY_ARRAY_FORCECAST 0x0010 + +/* + * Always copy the array. Returned arrays are always CONTIGUOUS, + * ALIGNED, and WRITEABLE. See also: NPY_ARRAY_ENSURENOCOPY = 0x4000. + * + * This flag may be requested in constructor functions. + */ +#define NPY_ARRAY_ENSURECOPY 0x0020 + +/* + * Make sure the returned array is a base-class ndarray + * + * This flag may be requested in constructor functions. + */ +#define NPY_ARRAY_ENSUREARRAY 0x0040 + +/* + * Make sure that the strides are in units of the element size Needed + * for some operations with record-arrays. + * + * This flag may be requested in constructor functions. + */ +#define NPY_ARRAY_ELEMENTSTRIDES 0x0080 + +/* + * Array data is aligned on the appropriate memory address for the type + * stored according to how the compiler would align things (e.g., an + * array of integers (4 bytes each) starts on a memory address that's + * a multiple of 4) + * + * This flag may be requested in constructor functions. + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_ALIGNED 0x0100 + +/* + * Array data has the native endianness + * + * This flag may be requested in constructor functions. + */ +#define NPY_ARRAY_NOTSWAPPED 0x0200 + +/* + * Array data is writeable + * + * This flag may be requested in constructor functions. + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_WRITEABLE 0x0400 + +/* + * If this flag is set, then base contains a pointer to an array of + * the same size that should be updated with the current contents of + * this array when PyArray_ResolveWritebackIfCopy is called. + * + * This flag may be requested in constructor functions. + * This flag may be tested for in PyArray_FLAGS(arr). + */ +#define NPY_ARRAY_WRITEBACKIFCOPY 0x2000 + +/* + * No copy may be made while converting from an object/array (result is a view) + * + * This flag may be requested in constructor functions. + */ +#define NPY_ARRAY_ENSURENOCOPY 0x4000 + +/* + * NOTE: there are also internal flags defined in multiarray/arrayobject.h, + * which start at bit 31 and work down. + */ + +#define NPY_ARRAY_BEHAVED (NPY_ARRAY_ALIGNED | \ + NPY_ARRAY_WRITEABLE) +#define NPY_ARRAY_BEHAVED_NS (NPY_ARRAY_ALIGNED | \ + NPY_ARRAY_WRITEABLE | \ + NPY_ARRAY_NOTSWAPPED) +#define NPY_ARRAY_CARRAY (NPY_ARRAY_C_CONTIGUOUS | \ + NPY_ARRAY_BEHAVED) +#define NPY_ARRAY_CARRAY_RO (NPY_ARRAY_C_CONTIGUOUS | \ + NPY_ARRAY_ALIGNED) +#define NPY_ARRAY_FARRAY (NPY_ARRAY_F_CONTIGUOUS | \ + NPY_ARRAY_BEHAVED) +#define NPY_ARRAY_FARRAY_RO (NPY_ARRAY_F_CONTIGUOUS | \ + NPY_ARRAY_ALIGNED) +#define NPY_ARRAY_DEFAULT (NPY_ARRAY_CARRAY) +#define NPY_ARRAY_IN_ARRAY (NPY_ARRAY_CARRAY_RO) +#define NPY_ARRAY_OUT_ARRAY (NPY_ARRAY_CARRAY) +#define NPY_ARRAY_INOUT_ARRAY (NPY_ARRAY_CARRAY) +#define NPY_ARRAY_INOUT_ARRAY2 (NPY_ARRAY_CARRAY | \ + NPY_ARRAY_WRITEBACKIFCOPY) +#define NPY_ARRAY_IN_FARRAY (NPY_ARRAY_FARRAY_RO) +#define NPY_ARRAY_OUT_FARRAY (NPY_ARRAY_FARRAY) +#define NPY_ARRAY_INOUT_FARRAY (NPY_ARRAY_FARRAY) +#define NPY_ARRAY_INOUT_FARRAY2 (NPY_ARRAY_FARRAY | \ + NPY_ARRAY_WRITEBACKIFCOPY) + +#define NPY_ARRAY_UPDATE_ALL (NPY_ARRAY_C_CONTIGUOUS | \ + NPY_ARRAY_F_CONTIGUOUS | \ + NPY_ARRAY_ALIGNED) + +/* This flag is for the array interface, not PyArrayObject */ +#define NPY_ARR_HAS_DESCR 0x0800 + + + + +/* + * Size of internal buffers used for alignment Make BUFSIZE a multiple + * of sizeof(npy_cdouble) -- usually 16 so that ufunc buffers are aligned + */ +#define NPY_MIN_BUFSIZE ((int)sizeof(npy_cdouble)) +#define NPY_MAX_BUFSIZE (((int)sizeof(npy_cdouble))*1000000) +#define NPY_BUFSIZE 8192 +/* buffer stress test size: */ +/*#define NPY_BUFSIZE 17*/ + +/* + * C API: consists of Macros and functions. The MACROS are defined + * here. + */ + + +#define PyArray_ISCONTIGUOUS(m) PyArray_CHKFLAGS((m), NPY_ARRAY_C_CONTIGUOUS) +#define PyArray_ISWRITEABLE(m) PyArray_CHKFLAGS((m), NPY_ARRAY_WRITEABLE) +#define PyArray_ISALIGNED(m) PyArray_CHKFLAGS((m), NPY_ARRAY_ALIGNED) + +#define PyArray_IS_C_CONTIGUOUS(m) PyArray_CHKFLAGS((m), NPY_ARRAY_C_CONTIGUOUS) +#define PyArray_IS_F_CONTIGUOUS(m) PyArray_CHKFLAGS((m), NPY_ARRAY_F_CONTIGUOUS) + +/* the variable is used in some places, so always define it */ +#define NPY_BEGIN_THREADS_DEF PyThreadState *_save=NULL; +#if NPY_ALLOW_THREADS +#define NPY_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS +#define NPY_END_ALLOW_THREADS Py_END_ALLOW_THREADS +#define NPY_BEGIN_THREADS do {_save = PyEval_SaveThread();} while (0); +#define NPY_END_THREADS do { if (_save) \ + { PyEval_RestoreThread(_save); _save = NULL;} } while (0); +#define NPY_BEGIN_THREADS_THRESHOLDED(loop_size) do { if ((loop_size) > 500) \ + { _save = PyEval_SaveThread();} } while (0); + + +#define NPY_ALLOW_C_API_DEF PyGILState_STATE __save__; +#define NPY_ALLOW_C_API do {__save__ = PyGILState_Ensure();} while (0); +#define NPY_DISABLE_C_API do {PyGILState_Release(__save__);} while (0); +#else +#define NPY_BEGIN_ALLOW_THREADS +#define NPY_END_ALLOW_THREADS +#define NPY_BEGIN_THREADS +#define NPY_END_THREADS +#define NPY_BEGIN_THREADS_THRESHOLDED(loop_size) +#define NPY_BEGIN_THREADS_DESCR(dtype) +#define NPY_END_THREADS_DESCR(dtype) +#define NPY_ALLOW_C_API_DEF +#define NPY_ALLOW_C_API +#define NPY_DISABLE_C_API +#endif + +/********************************** + * The nditer object, added in 1.6 + **********************************/ + +/* The actual structure of the iterator is an internal detail */ +typedef struct NpyIter_InternalOnly NpyIter; + +/* Iterator function pointers that may be specialized */ +typedef int (NpyIter_IterNextFunc)(NpyIter *iter); +typedef void (NpyIter_GetMultiIndexFunc)(NpyIter *iter, + npy_intp *outcoords); + +/*** Global flags that may be passed to the iterator constructors ***/ + +/* Track an index representing C order */ +#define NPY_ITER_C_INDEX 0x00000001 +/* Track an index representing Fortran order */ +#define NPY_ITER_F_INDEX 0x00000002 +/* Track a multi-index */ +#define NPY_ITER_MULTI_INDEX 0x00000004 +/* User code external to the iterator does the 1-dimensional innermost loop */ +#define NPY_ITER_EXTERNAL_LOOP 0x00000008 +/* Convert all the operands to a common data type */ +#define NPY_ITER_COMMON_DTYPE 0x00000010 +/* Operands may hold references, requiring API access during iteration */ +#define NPY_ITER_REFS_OK 0x00000020 +/* Zero-sized operands should be permitted, iteration checks IterSize for 0 */ +#define NPY_ITER_ZEROSIZE_OK 0x00000040 +/* Permits reductions (size-0 stride with dimension size > 1) */ +#define NPY_ITER_REDUCE_OK 0x00000080 +/* Enables sub-range iteration */ +#define NPY_ITER_RANGED 0x00000100 +/* Enables buffering */ +#define NPY_ITER_BUFFERED 0x00000200 +/* When buffering is enabled, grows the inner loop if possible */ +#define NPY_ITER_GROWINNER 0x00000400 +/* Delay allocation of buffers until first Reset* call */ +#define NPY_ITER_DELAY_BUFALLOC 0x00000800 +/* When NPY_KEEPORDER is specified, disable reversing negative-stride axes */ +#define NPY_ITER_DONT_NEGATE_STRIDES 0x00001000 +/* + * If output operands overlap with other operands (based on heuristics that + * has false positives but no false negatives), make temporary copies to + * eliminate overlap. + */ +#define NPY_ITER_COPY_IF_OVERLAP 0x00002000 + +/*** Per-operand flags that may be passed to the iterator constructors ***/ + +/* The operand will be read from and written to */ +#define NPY_ITER_READWRITE 0x00010000 +/* The operand will only be read from */ +#define NPY_ITER_READONLY 0x00020000 +/* The operand will only be written to */ +#define NPY_ITER_WRITEONLY 0x00040000 +/* The operand's data must be in native byte order */ +#define NPY_ITER_NBO 0x00080000 +/* The operand's data must be aligned */ +#define NPY_ITER_ALIGNED 0x00100000 +/* The operand's data must be contiguous (within the inner loop) */ +#define NPY_ITER_CONTIG 0x00200000 +/* The operand may be copied to satisfy requirements */ +#define NPY_ITER_COPY 0x00400000 +/* The operand may be copied with WRITEBACKIFCOPY to satisfy requirements */ +#define NPY_ITER_UPDATEIFCOPY 0x00800000 +/* Allocate the operand if it is NULL */ +#define NPY_ITER_ALLOCATE 0x01000000 +/* If an operand is allocated, don't use any subtype */ +#define NPY_ITER_NO_SUBTYPE 0x02000000 +/* This is a virtual array slot, operand is NULL but temporary data is there */ +#define NPY_ITER_VIRTUAL 0x04000000 +/* Require that the dimension match the iterator dimensions exactly */ +#define NPY_ITER_NO_BROADCAST 0x08000000 +/* A mask is being used on this array, affects buffer -> array copy */ +#define NPY_ITER_WRITEMASKED 0x10000000 +/* This array is the mask for all WRITEMASKED operands */ +#define NPY_ITER_ARRAYMASK 0x20000000 +/* Assume iterator order data access for COPY_IF_OVERLAP */ +#define NPY_ITER_OVERLAP_ASSUME_ELEMENTWISE 0x40000000 + +#define NPY_ITER_GLOBAL_FLAGS 0x0000ffff +#define NPY_ITER_PER_OP_FLAGS 0xffff0000 + + +/***************************** + * Basic iterator object + *****************************/ + +/* FWD declaration */ +typedef struct PyArrayIterObject_tag PyArrayIterObject; + +/* + * type of the function which translates a set of coordinates to a + * pointer to the data + */ +typedef char* (*npy_iter_get_dataptr_t)( + PyArrayIterObject* iter, const npy_intp*); + +struct PyArrayIterObject_tag { + PyObject_HEAD + int nd_m1; /* number of dimensions - 1 */ + npy_intp index, size; + npy_intp coordinates[NPY_MAXDIMS_LEGACY_ITERS];/* N-dimensional loop */ + npy_intp dims_m1[NPY_MAXDIMS_LEGACY_ITERS]; /* ao->dimensions - 1 */ + npy_intp strides[NPY_MAXDIMS_LEGACY_ITERS]; /* ao->strides or fake */ + npy_intp backstrides[NPY_MAXDIMS_LEGACY_ITERS];/* how far to jump back */ + npy_intp factors[NPY_MAXDIMS_LEGACY_ITERS]; /* shape factors */ + PyArrayObject *ao; + char *dataptr; /* pointer to current item*/ + npy_bool contiguous; + + npy_intp bounds[NPY_MAXDIMS_LEGACY_ITERS][2]; + npy_intp limits[NPY_MAXDIMS_LEGACY_ITERS][2]; + npy_intp limits_sizes[NPY_MAXDIMS_LEGACY_ITERS]; + npy_iter_get_dataptr_t translate; +} ; + + +/* Iterator API */ +#define PyArrayIter_Check(op) PyObject_TypeCheck((op), &PyArrayIter_Type) + +#define _PyAIT(it) ((PyArrayIterObject *)(it)) +#define PyArray_ITER_RESET(it) do { \ + _PyAIT(it)->index = 0; \ + _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao); \ + memset(_PyAIT(it)->coordinates, 0, \ + (_PyAIT(it)->nd_m1+1)*sizeof(npy_intp)); \ +} while (0) + +#define _PyArray_ITER_NEXT1(it) do { \ + (it)->dataptr += _PyAIT(it)->strides[0]; \ + (it)->coordinates[0]++; \ +} while (0) + +#define _PyArray_ITER_NEXT2(it) do { \ + if ((it)->coordinates[1] < (it)->dims_m1[1]) { \ + (it)->coordinates[1]++; \ + (it)->dataptr += (it)->strides[1]; \ + } \ + else { \ + (it)->coordinates[1] = 0; \ + (it)->coordinates[0]++; \ + (it)->dataptr += (it)->strides[0] - \ + (it)->backstrides[1]; \ + } \ +} while (0) + +#define PyArray_ITER_NEXT(it) do { \ + _PyAIT(it)->index++; \ + if (_PyAIT(it)->nd_m1 == 0) { \ + _PyArray_ITER_NEXT1(_PyAIT(it)); \ + } \ + else if (_PyAIT(it)->contiguous) \ + _PyAIT(it)->dataptr += PyArray_ITEMSIZE(_PyAIT(it)->ao); \ + else if (_PyAIT(it)->nd_m1 == 1) { \ + _PyArray_ITER_NEXT2(_PyAIT(it)); \ + } \ + else { \ + int __npy_i; \ + for (__npy_i=_PyAIT(it)->nd_m1; __npy_i >= 0; __npy_i--) { \ + if (_PyAIT(it)->coordinates[__npy_i] < \ + _PyAIT(it)->dims_m1[__npy_i]) { \ + _PyAIT(it)->coordinates[__npy_i]++; \ + _PyAIT(it)->dataptr += \ + _PyAIT(it)->strides[__npy_i]; \ + break; \ + } \ + else { \ + _PyAIT(it)->coordinates[__npy_i] = 0; \ + _PyAIT(it)->dataptr -= \ + _PyAIT(it)->backstrides[__npy_i]; \ + } \ + } \ + } \ +} while (0) + +#define PyArray_ITER_GOTO(it, destination) do { \ + int __npy_i; \ + _PyAIT(it)->index = 0; \ + _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao); \ + for (__npy_i = _PyAIT(it)->nd_m1; __npy_i>=0; __npy_i--) { \ + if (destination[__npy_i] < 0) { \ + destination[__npy_i] += \ + _PyAIT(it)->dims_m1[__npy_i]+1; \ + } \ + _PyAIT(it)->dataptr += destination[__npy_i] * \ + _PyAIT(it)->strides[__npy_i]; \ + _PyAIT(it)->coordinates[__npy_i] = \ + destination[__npy_i]; \ + _PyAIT(it)->index += destination[__npy_i] * \ + ( __npy_i==_PyAIT(it)->nd_m1 ? 1 : \ + _PyAIT(it)->dims_m1[__npy_i+1]+1) ; \ + } \ +} while (0) + +#define PyArray_ITER_GOTO1D(it, ind) do { \ + int __npy_i; \ + npy_intp __npy_ind = (npy_intp)(ind); \ + if (__npy_ind < 0) __npy_ind += _PyAIT(it)->size; \ + _PyAIT(it)->index = __npy_ind; \ + if (_PyAIT(it)->nd_m1 == 0) { \ + _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao) + \ + __npy_ind * _PyAIT(it)->strides[0]; \ + } \ + else if (_PyAIT(it)->contiguous) \ + _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao) + \ + __npy_ind * PyArray_ITEMSIZE(_PyAIT(it)->ao); \ + else { \ + _PyAIT(it)->dataptr = PyArray_BYTES(_PyAIT(it)->ao); \ + for (__npy_i = 0; __npy_i<=_PyAIT(it)->nd_m1; \ + __npy_i++) { \ + _PyAIT(it)->coordinates[__npy_i] = \ + (__npy_ind / _PyAIT(it)->factors[__npy_i]); \ + _PyAIT(it)->dataptr += \ + (__npy_ind / _PyAIT(it)->factors[__npy_i]) \ + * _PyAIT(it)->strides[__npy_i]; \ + __npy_ind %= _PyAIT(it)->factors[__npy_i]; \ + } \ + } \ +} while (0) + +#define PyArray_ITER_DATA(it) ((void *)(_PyAIT(it)->dataptr)) + +#define PyArray_ITER_NOTDONE(it) (_PyAIT(it)->index < _PyAIT(it)->size) + + +/* + * Any object passed to PyArray_Broadcast must be binary compatible + * with this structure. + */ + +typedef struct { + PyObject_HEAD + int numiter; /* number of iters */ + npy_intp size; /* broadcasted size */ + npy_intp index; /* current index */ + int nd; /* number of dims */ + npy_intp dimensions[NPY_MAXDIMS_LEGACY_ITERS]; /* dimensions */ + /* + * Space for the individual iterators, do not specify size publicly + * to allow changing it more easily. + * One reason is that Cython uses this for checks and only allows + * growing structs (as of Cython 3.0.6). It also allows NPY_MAXARGS + * to be runtime dependent. + */ +#if (defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) + PyArrayIterObject *iters[64]; +#elif defined(__cplusplus) + /* + * C++ doesn't strictly support flexible members and gives compilers + * warnings (pedantic only), so we lie. We can't make it 64 because + * then Cython is unhappy (larger struct at runtime is OK smaller not). + */ + PyArrayIterObject *iters[32]; +#else + PyArrayIterObject *iters[]; +#endif +} PyArrayMultiIterObject; + +#define _PyMIT(m) ((PyArrayMultiIterObject *)(m)) +#define PyArray_MultiIter_RESET(multi) do { \ + int __npy_mi; \ + _PyMIT(multi)->index = 0; \ + for (__npy_mi=0; __npy_mi < _PyMIT(multi)->numiter; __npy_mi++) { \ + PyArray_ITER_RESET(_PyMIT(multi)->iters[__npy_mi]); \ + } \ +} while (0) + +#define PyArray_MultiIter_NEXT(multi) do { \ + int __npy_mi; \ + _PyMIT(multi)->index++; \ + for (__npy_mi=0; __npy_mi < _PyMIT(multi)->numiter; __npy_mi++) { \ + PyArray_ITER_NEXT(_PyMIT(multi)->iters[__npy_mi]); \ + } \ +} while (0) + +#define PyArray_MultiIter_GOTO(multi, dest) do { \ + int __npy_mi; \ + for (__npy_mi=0; __npy_mi < _PyMIT(multi)->numiter; __npy_mi++) { \ + PyArray_ITER_GOTO(_PyMIT(multi)->iters[__npy_mi], dest); \ + } \ + _PyMIT(multi)->index = _PyMIT(multi)->iters[0]->index; \ +} while (0) + +#define PyArray_MultiIter_GOTO1D(multi, ind) do { \ + int __npy_mi; \ + for (__npy_mi=0; __npy_mi < _PyMIT(multi)->numiter; __npy_mi++) { \ + PyArray_ITER_GOTO1D(_PyMIT(multi)->iters[__npy_mi], ind); \ + } \ + _PyMIT(multi)->index = _PyMIT(multi)->iters[0]->index; \ +} while (0) + +#define PyArray_MultiIter_DATA(multi, i) \ + ((void *)(_PyMIT(multi)->iters[i]->dataptr)) + +#define PyArray_MultiIter_NEXTi(multi, i) \ + PyArray_ITER_NEXT(_PyMIT(multi)->iters[i]) + +#define PyArray_MultiIter_NOTDONE(multi) \ + (_PyMIT(multi)->index < _PyMIT(multi)->size) + + +static NPY_INLINE int +PyArray_MultiIter_NUMITER(PyArrayMultiIterObject *multi) +{ + return multi->numiter; +} + + +static NPY_INLINE npy_intp +PyArray_MultiIter_SIZE(PyArrayMultiIterObject *multi) +{ + return multi->size; +} + + +static NPY_INLINE npy_intp +PyArray_MultiIter_INDEX(PyArrayMultiIterObject *multi) +{ + return multi->index; +} + + +static NPY_INLINE int +PyArray_MultiIter_NDIM(PyArrayMultiIterObject *multi) +{ + return multi->nd; +} + + +static NPY_INLINE npy_intp * +PyArray_MultiIter_DIMS(PyArrayMultiIterObject *multi) +{ + return multi->dimensions; +} + + +static NPY_INLINE void ** +PyArray_MultiIter_ITERS(PyArrayMultiIterObject *multi) +{ + return (void**)multi->iters; +} + + +enum { + NPY_NEIGHBORHOOD_ITER_ZERO_PADDING, + NPY_NEIGHBORHOOD_ITER_ONE_PADDING, + NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING, + NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING, + NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING +}; + +typedef struct { + PyObject_HEAD + + /* + * PyArrayIterObject part: keep this in this exact order + */ + int nd_m1; /* number of dimensions - 1 */ + npy_intp index, size; + npy_intp coordinates[NPY_MAXDIMS_LEGACY_ITERS];/* N-dimensional loop */ + npy_intp dims_m1[NPY_MAXDIMS_LEGACY_ITERS]; /* ao->dimensions - 1 */ + npy_intp strides[NPY_MAXDIMS_LEGACY_ITERS]; /* ao->strides or fake */ + npy_intp backstrides[NPY_MAXDIMS_LEGACY_ITERS];/* how far to jump back */ + npy_intp factors[NPY_MAXDIMS_LEGACY_ITERS]; /* shape factors */ + PyArrayObject *ao; + char *dataptr; /* pointer to current item*/ + npy_bool contiguous; + + npy_intp bounds[NPY_MAXDIMS_LEGACY_ITERS][2]; + npy_intp limits[NPY_MAXDIMS_LEGACY_ITERS][2]; + npy_intp limits_sizes[NPY_MAXDIMS_LEGACY_ITERS]; + npy_iter_get_dataptr_t translate; + + /* + * New members + */ + npy_intp nd; + + /* Dimensions is the dimension of the array */ + npy_intp dimensions[NPY_MAXDIMS_LEGACY_ITERS]; + + /* + * Neighborhood points coordinates are computed relatively to the + * point pointed by _internal_iter + */ + PyArrayIterObject* _internal_iter; + /* + * To keep a reference to the representation of the constant value + * for constant padding + */ + char* constant; + + int mode; +} PyArrayNeighborhoodIterObject; + +/* + * Neighborhood iterator API + */ + +/* General: those work for any mode */ +static inline int +PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter); +static inline int +PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter); +#if 0 +static inline int +PyArrayNeighborhoodIter_Next2D(PyArrayNeighborhoodIterObject* iter); +#endif + +/* + * Include inline implementations - functions defined there are not + * considered public API + */ +#define NUMPY_CORE_INCLUDE_NUMPY__NEIGHBORHOOD_IMP_H_ +#include "_neighborhood_iterator_imp.h" +#undef NUMPY_CORE_INCLUDE_NUMPY__NEIGHBORHOOD_IMP_H_ + + + +/* The default array type */ +#define NPY_DEFAULT_TYPE NPY_DOUBLE +/* default integer type defined in npy_2_compat header */ + +/* + * All sorts of useful ways to look into a PyArrayObject. It is recommended + * to use PyArrayObject * objects instead of always casting from PyObject *, + * for improved type checking. + * + * In many cases here the macro versions of the accessors are deprecated, + * but can't be immediately changed to inline functions because the + * preexisting macros accept PyObject * and do automatic casts. Inline + * functions accepting PyArrayObject * provides for some compile-time + * checking of correctness when working with these objects in C. + */ + +#define PyArray_ISONESEGMENT(m) (PyArray_CHKFLAGS(m, NPY_ARRAY_C_CONTIGUOUS) || \ + PyArray_CHKFLAGS(m, NPY_ARRAY_F_CONTIGUOUS)) + +#define PyArray_ISFORTRAN(m) (PyArray_CHKFLAGS(m, NPY_ARRAY_F_CONTIGUOUS) && \ + (!PyArray_CHKFLAGS(m, NPY_ARRAY_C_CONTIGUOUS))) + +#define PyArray_FORTRAN_IF(m) ((PyArray_CHKFLAGS(m, NPY_ARRAY_F_CONTIGUOUS) ? \ + NPY_ARRAY_F_CONTIGUOUS : 0)) + +static inline int +PyArray_NDIM(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->nd; +} + +static inline void * +PyArray_DATA(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->data; +} + +static inline char * +PyArray_BYTES(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->data; +} + +static inline npy_intp * +PyArray_DIMS(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->dimensions; +} + +static inline npy_intp * +PyArray_STRIDES(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->strides; +} + +static inline npy_intp +PyArray_DIM(const PyArrayObject *arr, int idim) +{ + return ((PyArrayObject_fields *)arr)->dimensions[idim]; +} + +static inline npy_intp +PyArray_STRIDE(const PyArrayObject *arr, int istride) +{ + return ((PyArrayObject_fields *)arr)->strides[istride]; +} + +static inline NPY_RETURNS_BORROWED_REF PyObject * +PyArray_BASE(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->base; +} + +static inline NPY_RETURNS_BORROWED_REF PyArray_Descr * +PyArray_DESCR(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->descr; +} + +static inline int +PyArray_FLAGS(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->flags; +} + + +static inline int +PyArray_TYPE(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->descr->type_num; +} + +static inline int +PyArray_CHKFLAGS(const PyArrayObject *arr, int flags) +{ + return (PyArray_FLAGS(arr) & flags) == flags; +} + +static inline PyArray_Descr * +PyArray_DTYPE(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->descr; +} + +static inline npy_intp * +PyArray_SHAPE(const PyArrayObject *arr) +{ + return ((PyArrayObject_fields *)arr)->dimensions; +} + +/* + * Enables the specified array flags. Does no checking, + * assumes you know what you're doing. + */ +static inline void +PyArray_ENABLEFLAGS(PyArrayObject *arr, int flags) +{ + ((PyArrayObject_fields *)arr)->flags |= flags; +} + +/* + * Clears the specified array flags. Does no checking, + * assumes you know what you're doing. + */ +static inline void +PyArray_CLEARFLAGS(PyArrayObject *arr, int flags) +{ + ((PyArrayObject_fields *)arr)->flags &= ~flags; +} + +#if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION + static inline NPY_RETURNS_BORROWED_REF PyObject * + PyArray_HANDLER(PyArrayObject *arr) + { + return ((PyArrayObject_fields *)arr)->mem_handler; + } +#endif + +#define PyTypeNum_ISBOOL(type) ((type) == NPY_BOOL) + +#define PyTypeNum_ISUNSIGNED(type) (((type) == NPY_UBYTE) || \ + ((type) == NPY_USHORT) || \ + ((type) == NPY_UINT) || \ + ((type) == NPY_ULONG) || \ + ((type) == NPY_ULONGLONG)) + +#define PyTypeNum_ISSIGNED(type) (((type) == NPY_BYTE) || \ + ((type) == NPY_SHORT) || \ + ((type) == NPY_INT) || \ + ((type) == NPY_LONG) || \ + ((type) == NPY_LONGLONG)) + +#define PyTypeNum_ISINTEGER(type) (((type) >= NPY_BYTE) && \ + ((type) <= NPY_ULONGLONG)) + +#define PyTypeNum_ISFLOAT(type) ((((type) >= NPY_FLOAT) && \ + ((type) <= NPY_LONGDOUBLE)) || \ + ((type) == NPY_HALF)) + +#define PyTypeNum_ISNUMBER(type) (((type) <= NPY_CLONGDOUBLE) || \ + ((type) == NPY_HALF)) + +#define PyTypeNum_ISSTRING(type) (((type) == NPY_STRING) || \ + ((type) == NPY_UNICODE)) + +#define PyTypeNum_ISCOMPLEX(type) (((type) >= NPY_CFLOAT) && \ + ((type) <= NPY_CLONGDOUBLE)) + +#define PyTypeNum_ISFLEXIBLE(type) (((type) >=NPY_STRING) && \ + ((type) <=NPY_VOID)) + +#define PyTypeNum_ISDATETIME(type) (((type) >=NPY_DATETIME) && \ + ((type) <=NPY_TIMEDELTA)) + +#define PyTypeNum_ISUSERDEF(type) (((type) >= NPY_USERDEF) && \ + ((type) < NPY_USERDEF+ \ + NPY_NUMUSERTYPES)) + +#define PyTypeNum_ISEXTENDED(type) (PyTypeNum_ISFLEXIBLE(type) || \ + PyTypeNum_ISUSERDEF(type)) + +#define PyTypeNum_ISOBJECT(type) ((type) == NPY_OBJECT) + + +#define PyDataType_ISLEGACY(dtype) ((dtype)->type_num < NPY_VSTRING && ((dtype)->type_num >= 0)) +#define PyDataType_ISBOOL(obj) PyTypeNum_ISBOOL(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISUNSIGNED(obj) PyTypeNum_ISUNSIGNED(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISSIGNED(obj) PyTypeNum_ISSIGNED(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISINTEGER(obj) PyTypeNum_ISINTEGER(((PyArray_Descr*)(obj))->type_num ) +#define PyDataType_ISFLOAT(obj) PyTypeNum_ISFLOAT(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISNUMBER(obj) PyTypeNum_ISNUMBER(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISSTRING(obj) PyTypeNum_ISSTRING(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISCOMPLEX(obj) PyTypeNum_ISCOMPLEX(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISFLEXIBLE(obj) PyTypeNum_ISFLEXIBLE(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISDATETIME(obj) PyTypeNum_ISDATETIME(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISUSERDEF(obj) PyTypeNum_ISUSERDEF(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISEXTENDED(obj) PyTypeNum_ISEXTENDED(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_ISOBJECT(obj) PyTypeNum_ISOBJECT(((PyArray_Descr*)(obj))->type_num) +#define PyDataType_MAKEUNSIZED(dtype) ((dtype)->elsize = 0) +/* + * PyDataType_* FLAGS, FLACHK, REFCHK, HASFIELDS, HASSUBARRAY, UNSIZED, + * SUBARRAY, NAMES, FIELDS, C_METADATA, and METADATA require version specific + * lookup and are defined in npy_2_compat.h. + */ + + +#define PyArray_ISBOOL(obj) PyTypeNum_ISBOOL(PyArray_TYPE(obj)) +#define PyArray_ISUNSIGNED(obj) PyTypeNum_ISUNSIGNED(PyArray_TYPE(obj)) +#define PyArray_ISSIGNED(obj) PyTypeNum_ISSIGNED(PyArray_TYPE(obj)) +#define PyArray_ISINTEGER(obj) PyTypeNum_ISINTEGER(PyArray_TYPE(obj)) +#define PyArray_ISFLOAT(obj) PyTypeNum_ISFLOAT(PyArray_TYPE(obj)) +#define PyArray_ISNUMBER(obj) PyTypeNum_ISNUMBER(PyArray_TYPE(obj)) +#define PyArray_ISSTRING(obj) PyTypeNum_ISSTRING(PyArray_TYPE(obj)) +#define PyArray_ISCOMPLEX(obj) PyTypeNum_ISCOMPLEX(PyArray_TYPE(obj)) +#define PyArray_ISFLEXIBLE(obj) PyTypeNum_ISFLEXIBLE(PyArray_TYPE(obj)) +#define PyArray_ISDATETIME(obj) PyTypeNum_ISDATETIME(PyArray_TYPE(obj)) +#define PyArray_ISUSERDEF(obj) PyTypeNum_ISUSERDEF(PyArray_TYPE(obj)) +#define PyArray_ISEXTENDED(obj) PyTypeNum_ISEXTENDED(PyArray_TYPE(obj)) +#define PyArray_ISOBJECT(obj) PyTypeNum_ISOBJECT(PyArray_TYPE(obj)) +#define PyArray_HASFIELDS(obj) PyDataType_HASFIELDS(PyArray_DESCR(obj)) + + /* + * FIXME: This should check for a flag on the data-type that + * states whether or not it is variable length. Because the + * ISFLEXIBLE check is hard-coded to the built-in data-types. + */ +#define PyArray_ISVARIABLE(obj) PyTypeNum_ISFLEXIBLE(PyArray_TYPE(obj)) + +#define PyArray_SAFEALIGNEDCOPY(obj) (PyArray_ISALIGNED(obj) && !PyArray_ISVARIABLE(obj)) + + +#define NPY_LITTLE '<' +#define NPY_BIG '>' +#define NPY_NATIVE '=' +#define NPY_SWAP 's' +#define NPY_IGNORE '|' + +#if NPY_BYTE_ORDER == NPY_BIG_ENDIAN +#define NPY_NATBYTE NPY_BIG +#define NPY_OPPBYTE NPY_LITTLE +#else +#define NPY_NATBYTE NPY_LITTLE +#define NPY_OPPBYTE NPY_BIG +#endif + +#define PyArray_ISNBO(arg) ((arg) != NPY_OPPBYTE) +#define PyArray_IsNativeByteOrder PyArray_ISNBO +#define PyArray_ISNOTSWAPPED(m) PyArray_ISNBO(PyArray_DESCR(m)->byteorder) +#define PyArray_ISBYTESWAPPED(m) (!PyArray_ISNOTSWAPPED(m)) + +#define PyArray_FLAGSWAP(m, flags) (PyArray_CHKFLAGS(m, flags) && \ + PyArray_ISNOTSWAPPED(m)) + +#define PyArray_ISCARRAY(m) PyArray_FLAGSWAP(m, NPY_ARRAY_CARRAY) +#define PyArray_ISCARRAY_RO(m) PyArray_FLAGSWAP(m, NPY_ARRAY_CARRAY_RO) +#define PyArray_ISFARRAY(m) PyArray_FLAGSWAP(m, NPY_ARRAY_FARRAY) +#define PyArray_ISFARRAY_RO(m) PyArray_FLAGSWAP(m, NPY_ARRAY_FARRAY_RO) +#define PyArray_ISBEHAVED(m) PyArray_FLAGSWAP(m, NPY_ARRAY_BEHAVED) +#define PyArray_ISBEHAVED_RO(m) PyArray_FLAGSWAP(m, NPY_ARRAY_ALIGNED) + + +#define PyDataType_ISNOTSWAPPED(d) PyArray_ISNBO(((PyArray_Descr *)(d))->byteorder) +#define PyDataType_ISBYTESWAPPED(d) (!PyDataType_ISNOTSWAPPED(d)) + +/************************************************************ + * A struct used by PyArray_CreateSortedStridePerm, new in 1.7. + ************************************************************/ + +typedef struct { + npy_intp perm, stride; +} npy_stride_sort_item; + +/************************************************************ + * This is the form of the struct that's stored in the + * PyCapsule returned by an array's __array_struct__ attribute. See + * https://docs.scipy.org/doc/numpy/reference/arrays.interface.html for the full + * documentation. + ************************************************************/ +typedef struct { + int two; /* + * contains the integer 2 as a sanity + * check + */ + + int nd; /* number of dimensions */ + + char typekind; /* + * kind in array --- character code of + * typestr + */ + + int itemsize; /* size of each element */ + + int flags; /* + * how should be data interpreted. Valid + * flags are CONTIGUOUS (1), F_CONTIGUOUS (2), + * ALIGNED (0x100), NOTSWAPPED (0x200), and + * WRITEABLE (0x400). ARR_HAS_DESCR (0x800) + * states that arrdescr field is present in + * structure + */ + + npy_intp *shape; /* + * A length-nd array of shape + * information + */ + + npy_intp *strides; /* A length-nd array of stride information */ + + void *data; /* A pointer to the first element of the array */ + + PyObject *descr; /* + * A list of fields or NULL (ignored if flags + * does not have ARR_HAS_DESCR flag set) + */ +} PyArrayInterface; + + +/**************************************** + * NpyString + * + * Types used by the NpyString API. + ****************************************/ + +/* + * A "packed" encoded string. The string data must be accessed by first unpacking the string. + */ +typedef struct npy_packed_static_string npy_packed_static_string; + +/* + * An unpacked read-only view onto the data in a packed string + */ +typedef struct npy_unpacked_static_string { + size_t size; + const char *buf; +} npy_static_string; + +/* + * Handles heap allocations for static strings. + */ +typedef struct npy_string_allocator npy_string_allocator; + +typedef struct { + PyArray_Descr base; + // The object representing a null value + PyObject *na_object; + // Flag indicating whether or not to coerce arbitrary objects to strings + char coerce; + // Flag indicating the na object is NaN-like + char has_nan_na; + // Flag indicating the na object is a string + char has_string_na; + // If nonzero, indicates that this instance is owned by an array already + char array_owned; + // The string data to use when a default string is needed + npy_static_string default_string; + // The name of the missing data object, if any + npy_static_string na_name; + // the allocator should only be directly accessed after + // acquiring the allocator_lock and the lock should + // be released immediately after the allocator is + // no longer needed + npy_string_allocator *allocator; +} PyArray_StringDTypeObject; + +/* + * PyArray_DTypeMeta related definitions. + * + * As of now, this API is preliminary and will be extended as necessary. + */ +#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD + /* + * The Structures defined in this block are currently considered + * private API and may change without warning! + * Part of this (at least the size) is expected to be public API without + * further modifications. + */ + /* TODO: Make this definition public in the API, as soon as its settled */ + NPY_NO_EXPORT extern PyTypeObject PyArrayDTypeMeta_Type; + + /* + * While NumPy DTypes would not need to be heap types the plan is to + * make DTypes available in Python at which point they will be heap types. + * Since we also wish to add fields to the DType class, this looks like + * a typical instance definition, but with PyHeapTypeObject instead of + * only the PyObject_HEAD. + * This must only be exposed very extremely careful consideration, since + * it is a fairly complex construct which may be better to allow + * refactoring of. + */ + typedef struct { + PyHeapTypeObject super; + + /* + * Most DTypes will have a singleton default instance, for the + * parametric legacy DTypes (bytes, string, void, datetime) this + * may be a pointer to the *prototype* instance? + */ + PyArray_Descr *singleton; + /* Copy of the legacy DTypes type number, usually invalid. */ + int type_num; + + /* The type object of the scalar instances (may be NULL?) */ + PyTypeObject *scalar_type; + /* + * DType flags to signal legacy, parametric, or + * abstract. But plenty of space for additional information/flags. + */ + npy_uint64 flags; + + /* + * Use indirection in order to allow a fixed size for this struct. + * A stable ABI size makes creating a static DType less painful + * while also ensuring flexibility for all opaque API (with one + * indirection due the pointer lookup). + */ + void *dt_slots; + void *reserved[3]; + } PyArray_DTypeMeta; + +#endif /* NPY_INTERNAL_BUILD */ + + +/* + * Use the keyword NPY_DEPRECATED_INCLUDES to ensure that the header files + * npy_*_*_deprecated_api.h are only included from here and nowhere else. + */ +#ifdef NPY_DEPRECATED_INCLUDES +#error "Do not use the reserved keyword NPY_DEPRECATED_INCLUDES." +#endif +#define NPY_DEPRECATED_INCLUDES +#if !defined(NPY_NO_DEPRECATED_API) || \ + (NPY_NO_DEPRECATED_API < NPY_1_7_API_VERSION) +#include "npy_1_7_deprecated_api.h" +#endif +/* + * There is no file npy_1_8_deprecated_api.h since there are no additional + * deprecated API features in NumPy 1.8. + * + * Note to maintainers: insert code like the following in future NumPy + * versions. + * + * #if !defined(NPY_NO_DEPRECATED_API) || \ + * (NPY_NO_DEPRECATED_API < NPY_1_9_API_VERSION) + * #include "npy_1_9_deprecated_api.h" + * #endif + */ +#undef NPY_DEPRECATED_INCLUDES + +#ifdef __cplusplus +} +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NDARRAYTYPES_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_1_7_deprecated_api.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_1_7_deprecated_api.h new file mode 100644 index 0000000000000000000000000000000000000000..be53cded488dd373072ab798fd7c25a5041e20ad --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_1_7_deprecated_api.h @@ -0,0 +1,112 @@ +#ifndef NPY_DEPRECATED_INCLUDES +#error "Should never include npy_*_*_deprecated_api directly." +#endif + +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_1_7_DEPRECATED_API_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_1_7_DEPRECATED_API_H_ + +/* Emit a warning if the user did not specifically request the old API */ +#ifndef NPY_NO_DEPRECATED_API +#if defined(_WIN32) +#define _WARN___STR2__(x) #x +#define _WARN___STR1__(x) _WARN___STR2__(x) +#define _WARN___LOC__ __FILE__ "(" _WARN___STR1__(__LINE__) ") : Warning Msg: " +#pragma message(_WARN___LOC__"Using deprecated NumPy API, disable it with " \ + "#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION") +#else +#warning "Using deprecated NumPy API, disable it with " \ + "#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION" +#endif +#endif + +/* + * This header exists to collect all dangerous/deprecated NumPy API + * as of NumPy 1.7. + * + * This is an attempt to remove bad API, the proliferation of macros, + * and namespace pollution currently produced by the NumPy headers. + */ + +/* These array flags are deprecated as of NumPy 1.7 */ +#define NPY_CONTIGUOUS NPY_ARRAY_C_CONTIGUOUS +#define NPY_FORTRAN NPY_ARRAY_F_CONTIGUOUS + +/* + * The consistent NPY_ARRAY_* names which don't pollute the NPY_* + * namespace were added in NumPy 1.7. + * + * These versions of the carray flags are deprecated, but + * probably should only be removed after two releases instead of one. + */ +#define NPY_C_CONTIGUOUS NPY_ARRAY_C_CONTIGUOUS +#define NPY_F_CONTIGUOUS NPY_ARRAY_F_CONTIGUOUS +#define NPY_OWNDATA NPY_ARRAY_OWNDATA +#define NPY_FORCECAST NPY_ARRAY_FORCECAST +#define NPY_ENSURECOPY NPY_ARRAY_ENSURECOPY +#define NPY_ENSUREARRAY NPY_ARRAY_ENSUREARRAY +#define NPY_ELEMENTSTRIDES NPY_ARRAY_ELEMENTSTRIDES +#define NPY_ALIGNED NPY_ARRAY_ALIGNED +#define NPY_NOTSWAPPED NPY_ARRAY_NOTSWAPPED +#define NPY_WRITEABLE NPY_ARRAY_WRITEABLE +#define NPY_BEHAVED NPY_ARRAY_BEHAVED +#define NPY_BEHAVED_NS NPY_ARRAY_BEHAVED_NS +#define NPY_CARRAY NPY_ARRAY_CARRAY +#define NPY_CARRAY_RO NPY_ARRAY_CARRAY_RO +#define NPY_FARRAY NPY_ARRAY_FARRAY +#define NPY_FARRAY_RO NPY_ARRAY_FARRAY_RO +#define NPY_DEFAULT NPY_ARRAY_DEFAULT +#define NPY_IN_ARRAY NPY_ARRAY_IN_ARRAY +#define NPY_OUT_ARRAY NPY_ARRAY_OUT_ARRAY +#define NPY_INOUT_ARRAY NPY_ARRAY_INOUT_ARRAY +#define NPY_IN_FARRAY NPY_ARRAY_IN_FARRAY +#define NPY_OUT_FARRAY NPY_ARRAY_OUT_FARRAY +#define NPY_INOUT_FARRAY NPY_ARRAY_INOUT_FARRAY +#define NPY_UPDATE_ALL NPY_ARRAY_UPDATE_ALL + +/* This way of accessing the default type is deprecated as of NumPy 1.7 */ +#define PyArray_DEFAULT NPY_DEFAULT_TYPE + +/* + * Deprecated as of NumPy 1.7, this kind of shortcut doesn't + * belong in the public API. + */ +#define NPY_AO PyArrayObject + +/* + * Deprecated as of NumPy 1.7, an all-lowercase macro doesn't + * belong in the public API. + */ +#define fortran fortran_ + +/* + * Deprecated as of NumPy 1.7, as it is a namespace-polluting + * macro. + */ +#define FORTRAN_IF PyArray_FORTRAN_IF + +/* Deprecated as of NumPy 1.7, datetime64 uses c_metadata instead */ +#define NPY_METADATA_DTSTR "__timeunit__" + +/* + * Deprecated as of NumPy 1.7. + * The reasoning: + * - These are for datetime, but there's no datetime "namespace". + * - They just turn NPY_STR_ into "", which is just + * making something simple be indirected. + */ +#define NPY_STR_Y "Y" +#define NPY_STR_M "M" +#define NPY_STR_W "W" +#define NPY_STR_D "D" +#define NPY_STR_h "h" +#define NPY_STR_m "m" +#define NPY_STR_s "s" +#define NPY_STR_ms "ms" +#define NPY_STR_us "us" +#define NPY_STR_ns "ns" +#define NPY_STR_ps "ps" +#define NPY_STR_fs "fs" +#define NPY_STR_as "as" + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_1_7_DEPRECATED_API_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_compat.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_compat.h new file mode 100644 index 0000000000000000000000000000000000000000..e39e65aedea7dc95a1130ea94518bb778d1238cd --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_compat.h @@ -0,0 +1,249 @@ +/* + * This header file defines relevant features which: + * - Require runtime inspection depending on the NumPy version. + * - May be needed when compiling with an older version of NumPy to allow + * a smooth transition. + * + * As such, it is shipped with NumPy 2.0, but designed to be vendored in full + * or parts by downstream projects. + * + * It must be included after any other includes. `import_array()` must have + * been called in the scope or version dependency will misbehave, even when + * only `PyUFunc_` API is used. + * + * If required complicated defs (with inline functions) should be written as: + * + * #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION + * Simple definition when NumPy 2.0 API is guaranteed. + * #else + * static inline definition of a 1.x compatibility shim + * #if NPY_ABI_VERSION < 0x02000000 + * Make 1.x compatibility shim the public API (1.x only branch) + * #else + * Runtime dispatched version (1.x or 2.x) + * #endif + * #endif + * + * An internal build always passes NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION + */ + +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_ + +/* + * New macros for accessing real and complex part of a complex number can be + * found in "npy_2_complexcompat.h". + */ + + +/* + * This header is meant to be included by downstream directly for 1.x compat. + * In that case we need to ensure that users first included the full headers + * and not just `ndarraytypes.h`. + */ + +#ifndef NPY_FEATURE_VERSION + #error "The NumPy 2 compat header requires `import_array()` for which " \ + "the `ndarraytypes.h` header include is not sufficient. Please " \ + "include it after `numpy/ndarrayobject.h` or similar.\n" \ + "To simplify inclusion, you may use `PyArray_ImportNumPy()` " \ + "which is defined in the compat header and is lightweight (can be)." +#endif + +#if NPY_ABI_VERSION < 0x02000000 + /* + * Define 2.0 feature version as it is needed below to decide whether we + * compile for both 1.x and 2.x (defining it guarantees 1.x only). + */ + #define NPY_2_0_API_VERSION 0x00000012 + /* + * If we are compiling with NumPy 1.x, PyArray_RUNTIME_VERSION so we + * pretend the `PyArray_RUNTIME_VERSION` is `NPY_FEATURE_VERSION`. + * This allows downstream to use `PyArray_RUNTIME_VERSION` if they need to. + */ + #define PyArray_RUNTIME_VERSION NPY_FEATURE_VERSION + /* Compiling on NumPy 1.x where these are the same: */ + #define PyArray_DescrProto PyArray_Descr +#endif + + +/* + * Define a better way to call `_import_array()` to simplify backporting as + * we now require imports more often (necessary to make ABI flexible). + */ +#ifdef import_array1 + +static inline int +PyArray_ImportNumPyAPI(void) +{ + if (NPY_UNLIKELY(PyArray_API == NULL)) { + import_array1(-1); + } + return 0; +} + +#endif /* import_array1 */ + + +/* + * NPY_DEFAULT_INT + * + * The default integer has changed, `NPY_DEFAULT_INT` is available at runtime + * for use as type number, e.g. `PyArray_DescrFromType(NPY_DEFAULT_INT)`. + * + * NPY_RAVEL_AXIS + * + * This was introduced in NumPy 2.0 to allow indicating that an axis should be + * raveled in an operation. Before NumPy 2.0, NPY_MAXDIMS was used for this purpose. + * + * NPY_MAXDIMS + * + * A constant indicating the maximum number dimensions allowed when creating + * an ndarray. + * + * NPY_NTYPES_LEGACY + * + * The number of built-in NumPy dtypes. + */ +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION + #define NPY_DEFAULT_INT NPY_INTP + #define NPY_RAVEL_AXIS NPY_MIN_INT + #define NPY_MAXARGS 64 + +#elif NPY_ABI_VERSION < 0x02000000 + #define NPY_DEFAULT_INT NPY_LONG + #define NPY_RAVEL_AXIS 32 + #define NPY_MAXARGS 32 + + /* Aliases of 2.x names to 1.x only equivalent names */ + #define NPY_NTYPES NPY_NTYPES_LEGACY + #define PyArray_DescrProto PyArray_Descr + #define _PyArray_LegacyDescr PyArray_Descr + /* NumPy 2 definition always works, but add it for 1.x only */ + #define PyDataType_ISLEGACY(dtype) (1) +#else + #define NPY_DEFAULT_INT \ + (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? NPY_INTP : NPY_LONG) + #define NPY_RAVEL_AXIS \ + (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? NPY_MIN_INT : 32) + #define NPY_MAXARGS \ + (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? 64 : 32) +#endif + + +/* + * Access inline functions for descriptor fields. Except for the first + * few fields, these needed to be moved (elsize, alignment) for + * additional space. Or they are descriptor specific and are not generally + * available anymore (metadata, c_metadata, subarray, names, fields). + * + * Most of these are defined via the `DESCR_ACCESSOR` macro helper. + */ +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION || NPY_ABI_VERSION < 0x02000000 + /* Compiling for 1.x or 2.x only, direct field access is OK: */ + + static inline void + PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size) + { + dtype->elsize = size; + } + + static inline npy_uint64 + PyDataType_FLAGS(const PyArray_Descr *dtype) + { + #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION + return dtype->flags; + #else + return (unsigned char)dtype->flags; /* Need unsigned cast on 1.x */ + #endif + } + + #define DESCR_ACCESSOR(FIELD, field, type, legacy_only) \ + static inline type \ + PyDataType_##FIELD(const PyArray_Descr *dtype) { \ + if (legacy_only && !PyDataType_ISLEGACY(dtype)) { \ + return (type)0; \ + } \ + return ((_PyArray_LegacyDescr *)dtype)->field; \ + } +#else /* compiling for both 1.x and 2.x */ + + static inline void + PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size) + { + if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) { + ((_PyArray_DescrNumPy2 *)dtype)->elsize = size; + } + else { + ((PyArray_DescrProto *)dtype)->elsize = (int)size; + } + } + + static inline npy_uint64 + PyDataType_FLAGS(const PyArray_Descr *dtype) + { + if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) { + return ((_PyArray_DescrNumPy2 *)dtype)->flags; + } + else { + return (unsigned char)((PyArray_DescrProto *)dtype)->flags; + } + } + + /* Cast to LegacyDescr always fine but needed when `legacy_only` */ + #define DESCR_ACCESSOR(FIELD, field, type, legacy_only) \ + static inline type \ + PyDataType_##FIELD(const PyArray_Descr *dtype) { \ + if (legacy_only && !PyDataType_ISLEGACY(dtype)) { \ + return (type)0; \ + } \ + if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) { \ + return ((_PyArray_LegacyDescr *)dtype)->field; \ + } \ + else { \ + return ((PyArray_DescrProto *)dtype)->field; \ + } \ + } +#endif + +DESCR_ACCESSOR(ELSIZE, elsize, npy_intp, 0) +DESCR_ACCESSOR(ALIGNMENT, alignment, npy_intp, 0) +DESCR_ACCESSOR(METADATA, metadata, PyObject *, 1) +DESCR_ACCESSOR(SUBARRAY, subarray, PyArray_ArrayDescr *, 1) +DESCR_ACCESSOR(NAMES, names, PyObject *, 1) +DESCR_ACCESSOR(FIELDS, fields, PyObject *, 1) +DESCR_ACCESSOR(C_METADATA, c_metadata, NpyAuxData *, 1) + +#undef DESCR_ACCESSOR + + +#if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD) +#if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION + static inline PyArray_ArrFuncs * + PyDataType_GetArrFuncs(const PyArray_Descr *descr) + { + return _PyDataType_GetArrFuncs(descr); + } +#elif NPY_ABI_VERSION < 0x02000000 + static inline PyArray_ArrFuncs * + PyDataType_GetArrFuncs(const PyArray_Descr *descr) + { + return descr->f; + } +#else + static inline PyArray_ArrFuncs * + PyDataType_GetArrFuncs(const PyArray_Descr *descr) + { + if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) { + return _PyDataType_GetArrFuncs(descr); + } + else { + return ((PyArray_DescrProto *)descr)->f; + } + } +#endif + + +#endif /* not internal build */ + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_complexcompat.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_complexcompat.h new file mode 100644 index 0000000000000000000000000000000000000000..0b509011b2803b978bd913191816a9bbd6635d35 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_2_complexcompat.h @@ -0,0 +1,28 @@ +/* This header is designed to be copy-pasted into downstream packages, since it provides + a compatibility layer between the old C struct complex types and the new native C99 + complex types. The new macros are in numpy/npy_math.h, which is why it is included here. */ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPLEXCOMPAT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPLEXCOMPAT_H_ + +#include + +#ifndef NPY_CSETREALF +#define NPY_CSETREALF(c, r) (c)->real = (r) +#endif +#ifndef NPY_CSETIMAGF +#define NPY_CSETIMAGF(c, i) (c)->imag = (i) +#endif +#ifndef NPY_CSETREAL +#define NPY_CSETREAL(c, r) (c)->real = (r) +#endif +#ifndef NPY_CSETIMAG +#define NPY_CSETIMAG(c, i) (c)->imag = (i) +#endif +#ifndef NPY_CSETREALL +#define NPY_CSETREALL(c, r) (c)->real = (r) +#endif +#ifndef NPY_CSETIMAGL +#define NPY_CSETIMAGL(c, i) (c)->imag = (i) +#endif + +#endif diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_3kcompat.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_3kcompat.h new file mode 100644 index 0000000000000000000000000000000000000000..c2bf74faf09d9df0df07f23e7d8b8550fe79ba5f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_3kcompat.h @@ -0,0 +1,374 @@ +/* + * This is a convenience header file providing compatibility utilities + * for supporting different minor versions of Python 3. + * It was originally used to support the transition from Python 2, + * hence the "3k" naming. + * + * If you want to use this for your own projects, it's recommended to make a + * copy of it. We don't provide backwards compatibility guarantees. + */ + +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_3KCOMPAT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_3KCOMPAT_H_ + +#include +#include + +#include "npy_common.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/* Python13 removes _PyLong_AsInt */ +static inline int +Npy__PyLong_AsInt(PyObject *obj) +{ + int overflow; + long result = PyLong_AsLongAndOverflow(obj, &overflow); + + /* INT_MAX and INT_MIN are defined in Python.h */ + if (overflow || result > INT_MAX || result < INT_MIN) { + /* XXX: could be cute and give a different + message for overflow == -1 */ + PyErr_SetString(PyExc_OverflowError, + "Python int too large to convert to C int"); + return -1; + } + return (int)result; +} + +#if defined _MSC_VER && _MSC_VER >= 1900 + +#include + +/* + * Macros to protect CRT calls against instant termination when passed an + * invalid parameter (https://bugs.python.org/issue23524). + */ +extern _invalid_parameter_handler _Py_silent_invalid_parameter_handler; +#define NPY_BEGIN_SUPPRESS_IPH { _invalid_parameter_handler _Py_old_handler = \ + _set_thread_local_invalid_parameter_handler(_Py_silent_invalid_parameter_handler); +#define NPY_END_SUPPRESS_IPH _set_thread_local_invalid_parameter_handler(_Py_old_handler); } + +#else + +#define NPY_BEGIN_SUPPRESS_IPH +#define NPY_END_SUPPRESS_IPH + +#endif /* _MSC_VER >= 1900 */ + +/* + * PyFile_* compatibility + */ + +/* + * Get a FILE* handle to the file represented by the Python object + */ +static inline FILE* +npy_PyFile_Dup2(PyObject *file, char *mode, npy_off_t *orig_pos) +{ + int fd, fd2, unbuf; + Py_ssize_t fd2_tmp; + PyObject *ret, *os, *io, *io_raw; + npy_off_t pos; + FILE *handle; + + /* Flush first to ensure things end up in the file in the correct order */ + ret = PyObject_CallMethod(file, "flush", ""); + if (ret == NULL) { + return NULL; + } + Py_DECREF(ret); + fd = PyObject_AsFileDescriptor(file); + if (fd == -1) { + return NULL; + } + + /* + * The handle needs to be dup'd because we have to call fclose + * at the end + */ + os = PyImport_ImportModule("os"); + if (os == NULL) { + return NULL; + } + ret = PyObject_CallMethod(os, "dup", "i", fd); + Py_DECREF(os); + if (ret == NULL) { + return NULL; + } + fd2_tmp = PyNumber_AsSsize_t(ret, PyExc_IOError); + Py_DECREF(ret); + if (fd2_tmp == -1 && PyErr_Occurred()) { + return NULL; + } + if (fd2_tmp < INT_MIN || fd2_tmp > INT_MAX) { + PyErr_SetString(PyExc_IOError, + "Getting an 'int' from os.dup() failed"); + return NULL; + } + fd2 = (int)fd2_tmp; + + /* Convert to FILE* handle */ +#ifdef _WIN32 + NPY_BEGIN_SUPPRESS_IPH + handle = _fdopen(fd2, mode); + NPY_END_SUPPRESS_IPH +#else + handle = fdopen(fd2, mode); +#endif + if (handle == NULL) { + PyErr_SetString(PyExc_IOError, + "Getting a FILE* from a Python file object via " + "_fdopen failed. If you built NumPy, you probably " + "linked with the wrong debug/release runtime"); + return NULL; + } + + /* Record the original raw file handle position */ + *orig_pos = npy_ftell(handle); + if (*orig_pos == -1) { + /* The io module is needed to determine if buffering is used */ + io = PyImport_ImportModule("io"); + if (io == NULL) { + fclose(handle); + return NULL; + } + /* File object instances of RawIOBase are unbuffered */ + io_raw = PyObject_GetAttrString(io, "RawIOBase"); + Py_DECREF(io); + if (io_raw == NULL) { + fclose(handle); + return NULL; + } + unbuf = PyObject_IsInstance(file, io_raw); + Py_DECREF(io_raw); + if (unbuf == 1) { + /* Succeed if the IO is unbuffered */ + return handle; + } + else { + PyErr_SetString(PyExc_IOError, "obtaining file position failed"); + fclose(handle); + return NULL; + } + } + + /* Seek raw handle to the Python-side position */ + ret = PyObject_CallMethod(file, "tell", ""); + if (ret == NULL) { + fclose(handle); + return NULL; + } + pos = PyLong_AsLongLong(ret); + Py_DECREF(ret); + if (PyErr_Occurred()) { + fclose(handle); + return NULL; + } + if (npy_fseek(handle, pos, SEEK_SET) == -1) { + PyErr_SetString(PyExc_IOError, "seeking file failed"); + fclose(handle); + return NULL; + } + return handle; +} + +/* + * Close the dup-ed file handle, and seek the Python one to the current position + */ +static inline int +npy_PyFile_DupClose2(PyObject *file, FILE* handle, npy_off_t orig_pos) +{ + int fd, unbuf; + PyObject *ret, *io, *io_raw; + npy_off_t position; + + position = npy_ftell(handle); + + /* Close the FILE* handle */ + fclose(handle); + + /* + * Restore original file handle position, in order to not confuse + * Python-side data structures + */ + fd = PyObject_AsFileDescriptor(file); + if (fd == -1) { + return -1; + } + + if (npy_lseek(fd, orig_pos, SEEK_SET) == -1) { + + /* The io module is needed to determine if buffering is used */ + io = PyImport_ImportModule("io"); + if (io == NULL) { + return -1; + } + /* File object instances of RawIOBase are unbuffered */ + io_raw = PyObject_GetAttrString(io, "RawIOBase"); + Py_DECREF(io); + if (io_raw == NULL) { + return -1; + } + unbuf = PyObject_IsInstance(file, io_raw); + Py_DECREF(io_raw); + if (unbuf == 1) { + /* Succeed if the IO is unbuffered */ + return 0; + } + else { + PyErr_SetString(PyExc_IOError, "seeking file failed"); + return -1; + } + } + + if (position == -1) { + PyErr_SetString(PyExc_IOError, "obtaining file position failed"); + return -1; + } + + /* Seek Python-side handle to the FILE* handle position */ + ret = PyObject_CallMethod(file, "seek", NPY_OFF_T_PYFMT "i", position, 0); + if (ret == NULL) { + return -1; + } + Py_DECREF(ret); + return 0; +} + +static inline PyObject* +npy_PyFile_OpenFile(PyObject *filename, const char *mode) +{ + PyObject *open; + open = PyDict_GetItemString(PyEval_GetBuiltins(), "open"); + if (open == NULL) { + return NULL; + } + return PyObject_CallFunction(open, "Os", filename, mode); +} + +static inline int +npy_PyFile_CloseFile(PyObject *file) +{ + PyObject *ret; + + ret = PyObject_CallMethod(file, "close", NULL); + if (ret == NULL) { + return -1; + } + Py_DECREF(ret); + return 0; +} + +/* This is a copy of _PyErr_ChainExceptions, which + * is no longer exported from Python3.12 + */ +static inline void +npy_PyErr_ChainExceptions(PyObject *exc, PyObject *val, PyObject *tb) +{ + if (exc == NULL) + return; + + if (PyErr_Occurred()) { + PyObject *exc2, *val2, *tb2; + PyErr_Fetch(&exc2, &val2, &tb2); + PyErr_NormalizeException(&exc, &val, &tb); + if (tb != NULL) { + PyException_SetTraceback(val, tb); + Py_DECREF(tb); + } + Py_DECREF(exc); + PyErr_NormalizeException(&exc2, &val2, &tb2); + PyException_SetContext(val2, val); + PyErr_Restore(exc2, val2, tb2); + } + else { + PyErr_Restore(exc, val, tb); + } +} + +/* This is a copy of _PyErr_ChainExceptions, with: + * __cause__ used instead of __context__ + */ +static inline void +npy_PyErr_ChainExceptionsCause(PyObject *exc, PyObject *val, PyObject *tb) +{ + if (exc == NULL) + return; + + if (PyErr_Occurred()) { + PyObject *exc2, *val2, *tb2; + PyErr_Fetch(&exc2, &val2, &tb2); + PyErr_NormalizeException(&exc, &val, &tb); + if (tb != NULL) { + PyException_SetTraceback(val, tb); + Py_DECREF(tb); + } + Py_DECREF(exc); + PyErr_NormalizeException(&exc2, &val2, &tb2); + PyException_SetCause(val2, val); + PyErr_Restore(exc2, val2, tb2); + } + else { + PyErr_Restore(exc, val, tb); + } +} + +/* + * PyCObject functions adapted to PyCapsules. + * + * The main job here is to get rid of the improved error handling + * of PyCapsules. It's a shame... + */ +static inline PyObject * +NpyCapsule_FromVoidPtr(void *ptr, void (*dtor)(PyObject *)) +{ + PyObject *ret = PyCapsule_New(ptr, NULL, dtor); + if (ret == NULL) { + PyErr_Clear(); + } + return ret; +} + +static inline PyObject * +NpyCapsule_FromVoidPtrAndDesc(void *ptr, void* context, void (*dtor)(PyObject *)) +{ + PyObject *ret = NpyCapsule_FromVoidPtr(ptr, dtor); + if (ret != NULL && PyCapsule_SetContext(ret, context) != 0) { + PyErr_Clear(); + Py_DECREF(ret); + ret = NULL; + } + return ret; +} + +static inline void * +NpyCapsule_AsVoidPtr(PyObject *obj) +{ + void *ret = PyCapsule_GetPointer(obj, NULL); + if (ret == NULL) { + PyErr_Clear(); + } + return ret; +} + +static inline void * +NpyCapsule_GetDesc(PyObject *obj) +{ + return PyCapsule_GetContext(obj); +} + +static inline int +NpyCapsule_Check(PyObject *ptr) +{ + return PyCapsule_CheckExact(ptr); +} + +#ifdef __cplusplus +} +#endif + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_3KCOMPAT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_common.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_common.h new file mode 100644 index 0000000000000000000000000000000000000000..79ad8ad78cb2bc68867b3002b091e5d7e7b6bc9a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_common.h @@ -0,0 +1,1070 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_COMMON_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_COMMON_H_ + +/* need Python.h for npy_intp, npy_uintp */ +#include + +/* numpconfig.h is auto-generated */ +#include "numpyconfig.h" +#ifdef HAVE_NPY_CONFIG_H +#include +#endif + +/* + * using static inline modifiers when defining npy_math functions + * allows the compiler to make optimizations when possible + */ +#ifndef NPY_INLINE_MATH +#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD + #define NPY_INLINE_MATH 1 +#else + #define NPY_INLINE_MATH 0 +#endif +#endif + +/* + * gcc does not unroll even with -O3 + * use with care, unrolling on modern cpus rarely speeds things up + */ +#ifdef HAVE_ATTRIBUTE_OPTIMIZE_UNROLL_LOOPS +#define NPY_GCC_UNROLL_LOOPS \ + __attribute__((optimize("unroll-loops"))) +#else +#define NPY_GCC_UNROLL_LOOPS +#endif + +/* highest gcc optimization level, enabled autovectorizer */ +#ifdef HAVE_ATTRIBUTE_OPTIMIZE_OPT_3 +#define NPY_GCC_OPT_3 __attribute__((optimize("O3"))) +#else +#define NPY_GCC_OPT_3 +#endif + +/* + * mark an argument (starting from 1) that must not be NULL and is not checked + * DO NOT USE IF FUNCTION CHECKS FOR NULL!! the compiler will remove the check + */ +#ifdef HAVE_ATTRIBUTE_NONNULL +#define NPY_GCC_NONNULL(n) __attribute__((nonnull(n))) +#else +#define NPY_GCC_NONNULL(n) +#endif + +/* + * give a hint to the compiler which branch is more likely or unlikely + * to occur, e.g. rare error cases: + * + * if (NPY_UNLIKELY(failure == 0)) + * return NULL; + * + * the double !! is to cast the expression (e.g. NULL) to a boolean required by + * the intrinsic + */ +#ifdef HAVE___BUILTIN_EXPECT +#define NPY_LIKELY(x) __builtin_expect(!!(x), 1) +#define NPY_UNLIKELY(x) __builtin_expect(!!(x), 0) +#else +#define NPY_LIKELY(x) (x) +#define NPY_UNLIKELY(x) (x) +#endif + +#ifdef HAVE___BUILTIN_PREFETCH +/* unlike _mm_prefetch also works on non-x86 */ +#define NPY_PREFETCH(x, rw, loc) __builtin_prefetch((x), (rw), (loc)) +#else +#ifdef NPY_HAVE_SSE +/* _MM_HINT_ET[01] (rw = 1) unsupported, only available in gcc >= 4.9 */ +#define NPY_PREFETCH(x, rw, loc) _mm_prefetch((x), loc == 0 ? _MM_HINT_NTA : \ + (loc == 1 ? _MM_HINT_T2 : \ + (loc == 2 ? _MM_HINT_T1 : \ + (loc == 3 ? _MM_HINT_T0 : -1)))) +#else +#define NPY_PREFETCH(x, rw,loc) +#endif +#endif + +/* `NPY_INLINE` kept for backwards compatibility; use `inline` instead */ +#if defined(_MSC_VER) && !defined(__clang__) + #define NPY_INLINE __inline +/* clang included here to handle clang-cl on Windows */ +#elif defined(__GNUC__) || defined(__clang__) + #if defined(__STRICT_ANSI__) + #define NPY_INLINE __inline__ + #else + #define NPY_INLINE inline + #endif +#else + #define NPY_INLINE +#endif + +#ifdef _MSC_VER + #define NPY_FINLINE static __forceinline +#elif defined(__GNUC__) + #define NPY_FINLINE static inline __attribute__((always_inline)) +#else + #define NPY_FINLINE static +#endif + +#if defined(_MSC_VER) + #define NPY_NOINLINE static __declspec(noinline) +#elif defined(__GNUC__) || defined(__clang__) + #define NPY_NOINLINE static __attribute__((noinline)) +#else + #define NPY_NOINLINE static +#endif + +#ifdef __cplusplus + #define NPY_TLS thread_local +#elif defined(HAVE_THREAD_LOCAL) + #define NPY_TLS thread_local +#elif defined(HAVE__THREAD_LOCAL) + #define NPY_TLS _Thread_local +#elif defined(HAVE___THREAD) + #define NPY_TLS __thread +#elif defined(HAVE___DECLSPEC_THREAD_) + #define NPY_TLS __declspec(thread) +#else + #define NPY_TLS +#endif + +#ifdef WITH_CPYCHECKER_RETURNS_BORROWED_REF_ATTRIBUTE + #define NPY_RETURNS_BORROWED_REF \ + __attribute__((cpychecker_returns_borrowed_ref)) +#else + #define NPY_RETURNS_BORROWED_REF +#endif + +#ifdef WITH_CPYCHECKER_STEALS_REFERENCE_TO_ARG_ATTRIBUTE + #define NPY_STEALS_REF_TO_ARG(n) \ + __attribute__((cpychecker_steals_reference_to_arg(n))) +#else + #define NPY_STEALS_REF_TO_ARG(n) +#endif + +/* 64 bit file position support, also on win-amd64. Issue gh-2256 */ +#if defined(_MSC_VER) && defined(_WIN64) && (_MSC_VER > 1400) || \ + defined(__MINGW32__) || defined(__MINGW64__) + #include + + #define npy_fseek _fseeki64 + #define npy_ftell _ftelli64 + #define npy_lseek _lseeki64 + #define npy_off_t npy_int64 + + #if NPY_SIZEOF_INT == 8 + #define NPY_OFF_T_PYFMT "i" + #elif NPY_SIZEOF_LONG == 8 + #define NPY_OFF_T_PYFMT "l" + #elif NPY_SIZEOF_LONGLONG == 8 + #define NPY_OFF_T_PYFMT "L" + #else + #error Unsupported size for type off_t + #endif +#else +#ifdef HAVE_FSEEKO + #define npy_fseek fseeko +#else + #define npy_fseek fseek +#endif +#ifdef HAVE_FTELLO + #define npy_ftell ftello +#else + #define npy_ftell ftell +#endif + #include + #ifndef _WIN32 + #include + #endif + #define npy_lseek lseek + #define npy_off_t off_t + + #if NPY_SIZEOF_OFF_T == NPY_SIZEOF_SHORT + #define NPY_OFF_T_PYFMT "h" + #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_INT + #define NPY_OFF_T_PYFMT "i" + #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_LONG + #define NPY_OFF_T_PYFMT "l" + #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_LONGLONG + #define NPY_OFF_T_PYFMT "L" + #else + #error Unsupported size for type off_t + #endif +#endif + +/* enums for detected endianness */ +enum { + NPY_CPU_UNKNOWN_ENDIAN, + NPY_CPU_LITTLE, + NPY_CPU_BIG +}; + +/* + * This is to typedef npy_intp to the appropriate size for Py_ssize_t. + * (Before NumPy 2.0 we used Py_intptr_t and Py_uintptr_t from `pyport.h`.) + */ +typedef Py_ssize_t npy_intp; +typedef size_t npy_uintp; + +/* + * Define sizes that were not defined in numpyconfig.h. + */ +#define NPY_SIZEOF_CHAR 1 +#define NPY_SIZEOF_BYTE 1 +#define NPY_SIZEOF_DATETIME 8 +#define NPY_SIZEOF_TIMEDELTA 8 +#define NPY_SIZEOF_HALF 2 +#define NPY_SIZEOF_CFLOAT NPY_SIZEOF_COMPLEX_FLOAT +#define NPY_SIZEOF_CDOUBLE NPY_SIZEOF_COMPLEX_DOUBLE +#define NPY_SIZEOF_CLONGDOUBLE NPY_SIZEOF_COMPLEX_LONGDOUBLE + +#ifdef constchar +#undef constchar +#endif + +#define NPY_SSIZE_T_PYFMT "n" +#define constchar char + +/* NPY_INTP_FMT Note: + * Unlike the other NPY_*_FMT macros, which are used with PyOS_snprintf, + * NPY_INTP_FMT is used with PyErr_Format and PyUnicode_FromFormat. Those + * functions use different formatting codes that are portably specified + * according to the Python documentation. See issue gh-2388. + */ +#if NPY_SIZEOF_INTP == NPY_SIZEOF_LONG + #define NPY_INTP NPY_LONG + #define NPY_UINTP NPY_ULONG + #define PyIntpArrType_Type PyLongArrType_Type + #define PyUIntpArrType_Type PyULongArrType_Type + #define NPY_MAX_INTP NPY_MAX_LONG + #define NPY_MIN_INTP NPY_MIN_LONG + #define NPY_MAX_UINTP NPY_MAX_ULONG + #define NPY_INTP_FMT "ld" +#elif NPY_SIZEOF_INTP == NPY_SIZEOF_INT + #define NPY_INTP NPY_INT + #define NPY_UINTP NPY_UINT + #define PyIntpArrType_Type PyIntArrType_Type + #define PyUIntpArrType_Type PyUIntArrType_Type + #define NPY_MAX_INTP NPY_MAX_INT + #define NPY_MIN_INTP NPY_MIN_INT + #define NPY_MAX_UINTP NPY_MAX_UINT + #define NPY_INTP_FMT "d" +#elif defined(PY_LONG_LONG) && (NPY_SIZEOF_INTP == NPY_SIZEOF_LONGLONG) + #define NPY_INTP NPY_LONGLONG + #define NPY_UINTP NPY_ULONGLONG + #define PyIntpArrType_Type PyLongLongArrType_Type + #define PyUIntpArrType_Type PyULongLongArrType_Type + #define NPY_MAX_INTP NPY_MAX_LONGLONG + #define NPY_MIN_INTP NPY_MIN_LONGLONG + #define NPY_MAX_UINTP NPY_MAX_ULONGLONG + #define NPY_INTP_FMT "lld" +#else + #error "Failed to correctly define NPY_INTP and NPY_UINTP" +#endif + + +/* + * Some platforms don't define bool, long long, or long double. + * Handle that here. + */ +#define NPY_BYTE_FMT "hhd" +#define NPY_UBYTE_FMT "hhu" +#define NPY_SHORT_FMT "hd" +#define NPY_USHORT_FMT "hu" +#define NPY_INT_FMT "d" +#define NPY_UINT_FMT "u" +#define NPY_LONG_FMT "ld" +#define NPY_ULONG_FMT "lu" +#define NPY_HALF_FMT "g" +#define NPY_FLOAT_FMT "g" +#define NPY_DOUBLE_FMT "g" + + +#ifdef PY_LONG_LONG +typedef PY_LONG_LONG npy_longlong; +typedef unsigned PY_LONG_LONG npy_ulonglong; +# ifdef _MSC_VER +# define NPY_LONGLONG_FMT "I64d" +# define NPY_ULONGLONG_FMT "I64u" +# else +# define NPY_LONGLONG_FMT "lld" +# define NPY_ULONGLONG_FMT "llu" +# endif +# ifdef _MSC_VER +# define NPY_LONGLONG_SUFFIX(x) (x##i64) +# define NPY_ULONGLONG_SUFFIX(x) (x##Ui64) +# else +# define NPY_LONGLONG_SUFFIX(x) (x##LL) +# define NPY_ULONGLONG_SUFFIX(x) (x##ULL) +# endif +#else +typedef long npy_longlong; +typedef unsigned long npy_ulonglong; +# define NPY_LONGLONG_SUFFIX(x) (x##L) +# define NPY_ULONGLONG_SUFFIX(x) (x##UL) +#endif + + +typedef unsigned char npy_bool; +#define NPY_FALSE 0 +#define NPY_TRUE 1 +/* + * `NPY_SIZEOF_LONGDOUBLE` isn't usually equal to sizeof(long double). + * In some certain cases, it may forced to be equal to sizeof(double) + * even against the compiler implementation and the same goes for + * `complex long double`. + * + * Therefore, avoid `long double`, use `npy_longdouble` instead, + * and when it comes to standard math functions make sure of using + * the double version when `NPY_SIZEOF_LONGDOUBLE` == `NPY_SIZEOF_DOUBLE`. + * For example: + * npy_longdouble *ptr, x; + * #if NPY_SIZEOF_LONGDOUBLE == NPY_SIZEOF_DOUBLE + * npy_longdouble r = modf(x, ptr); + * #else + * npy_longdouble r = modfl(x, ptr); + * #endif + * + * See https://github.com/numpy/numpy/issues/20348 + */ +#if NPY_SIZEOF_LONGDOUBLE == NPY_SIZEOF_DOUBLE + #define NPY_LONGDOUBLE_FMT "g" + #define longdouble_t double + typedef double npy_longdouble; +#else + #define NPY_LONGDOUBLE_FMT "Lg" + #define longdouble_t long double + typedef long double npy_longdouble; +#endif + +#ifndef Py_USING_UNICODE +#error Must use Python with unicode enabled. +#endif + + +typedef signed char npy_byte; +typedef unsigned char npy_ubyte; +typedef unsigned short npy_ushort; +typedef unsigned int npy_uint; +typedef unsigned long npy_ulong; + +/* These are for completeness */ +typedef char npy_char; +typedef short npy_short; +typedef int npy_int; +typedef long npy_long; +typedef float npy_float; +typedef double npy_double; + +typedef Py_hash_t npy_hash_t; +#define NPY_SIZEOF_HASH_T NPY_SIZEOF_INTP + +#if defined(__cplusplus) + +typedef struct +{ + double _Val[2]; +} npy_cdouble; + +typedef struct +{ + float _Val[2]; +} npy_cfloat; + +typedef struct +{ + long double _Val[2]; +} npy_clongdouble; + +#else + +#include + + +#if defined(_MSC_VER) && !defined(__INTEL_COMPILER) +typedef _Dcomplex npy_cdouble; +typedef _Fcomplex npy_cfloat; +typedef _Lcomplex npy_clongdouble; +#else /* !defined(_MSC_VER) || defined(__INTEL_COMPILER) */ +typedef double _Complex npy_cdouble; +typedef float _Complex npy_cfloat; +typedef longdouble_t _Complex npy_clongdouble; +#endif + +#endif + +/* + * numarray-style bit-width typedefs + */ +#define NPY_MAX_INT8 127 +#define NPY_MIN_INT8 -128 +#define NPY_MAX_UINT8 255 +#define NPY_MAX_INT16 32767 +#define NPY_MIN_INT16 -32768 +#define NPY_MAX_UINT16 65535 +#define NPY_MAX_INT32 2147483647 +#define NPY_MIN_INT32 (-NPY_MAX_INT32 - 1) +#define NPY_MAX_UINT32 4294967295U +#define NPY_MAX_INT64 NPY_LONGLONG_SUFFIX(9223372036854775807) +#define NPY_MIN_INT64 (-NPY_MAX_INT64 - NPY_LONGLONG_SUFFIX(1)) +#define NPY_MAX_UINT64 NPY_ULONGLONG_SUFFIX(18446744073709551615) +#define NPY_MAX_INT128 NPY_LONGLONG_SUFFIX(85070591730234615865843651857942052864) +#define NPY_MIN_INT128 (-NPY_MAX_INT128 - NPY_LONGLONG_SUFFIX(1)) +#define NPY_MAX_UINT128 NPY_ULONGLONG_SUFFIX(170141183460469231731687303715884105728) +#define NPY_MAX_INT256 NPY_LONGLONG_SUFFIX(57896044618658097711785492504343953926634992332820282019728792003956564819967) +#define NPY_MIN_INT256 (-NPY_MAX_INT256 - NPY_LONGLONG_SUFFIX(1)) +#define NPY_MAX_UINT256 NPY_ULONGLONG_SUFFIX(115792089237316195423570985008687907853269984665640564039457584007913129639935) +#define NPY_MIN_DATETIME NPY_MIN_INT64 +#define NPY_MAX_DATETIME NPY_MAX_INT64 +#define NPY_MIN_TIMEDELTA NPY_MIN_INT64 +#define NPY_MAX_TIMEDELTA NPY_MAX_INT64 + + /* Need to find the number of bits for each type and + make definitions accordingly. + + C states that sizeof(char) == 1 by definition + + So, just using the sizeof keyword won't help. + + It also looks like Python itself uses sizeof(char) quite a + bit, which by definition should be 1 all the time. + + Idea: Make Use of CHAR_BIT which should tell us how many + BITS per CHARACTER + */ + + /* Include platform definitions -- These are in the C89/90 standard */ +#include +#define NPY_MAX_BYTE SCHAR_MAX +#define NPY_MIN_BYTE SCHAR_MIN +#define NPY_MAX_UBYTE UCHAR_MAX +#define NPY_MAX_SHORT SHRT_MAX +#define NPY_MIN_SHORT SHRT_MIN +#define NPY_MAX_USHORT USHRT_MAX +#define NPY_MAX_INT INT_MAX +#ifndef INT_MIN +#define INT_MIN (-INT_MAX - 1) +#endif +#define NPY_MIN_INT INT_MIN +#define NPY_MAX_UINT UINT_MAX +#define NPY_MAX_LONG LONG_MAX +#define NPY_MIN_LONG LONG_MIN +#define NPY_MAX_ULONG ULONG_MAX + +#define NPY_BITSOF_BOOL (sizeof(npy_bool) * CHAR_BIT) +#define NPY_BITSOF_CHAR CHAR_BIT +#define NPY_BITSOF_BYTE (NPY_SIZEOF_BYTE * CHAR_BIT) +#define NPY_BITSOF_SHORT (NPY_SIZEOF_SHORT * CHAR_BIT) +#define NPY_BITSOF_INT (NPY_SIZEOF_INT * CHAR_BIT) +#define NPY_BITSOF_LONG (NPY_SIZEOF_LONG * CHAR_BIT) +#define NPY_BITSOF_LONGLONG (NPY_SIZEOF_LONGLONG * CHAR_BIT) +#define NPY_BITSOF_INTP (NPY_SIZEOF_INTP * CHAR_BIT) +#define NPY_BITSOF_HALF (NPY_SIZEOF_HALF * CHAR_BIT) +#define NPY_BITSOF_FLOAT (NPY_SIZEOF_FLOAT * CHAR_BIT) +#define NPY_BITSOF_DOUBLE (NPY_SIZEOF_DOUBLE * CHAR_BIT) +#define NPY_BITSOF_LONGDOUBLE (NPY_SIZEOF_LONGDOUBLE * CHAR_BIT) +#define NPY_BITSOF_CFLOAT (NPY_SIZEOF_CFLOAT * CHAR_BIT) +#define NPY_BITSOF_CDOUBLE (NPY_SIZEOF_CDOUBLE * CHAR_BIT) +#define NPY_BITSOF_CLONGDOUBLE (NPY_SIZEOF_CLONGDOUBLE * CHAR_BIT) +#define NPY_BITSOF_DATETIME (NPY_SIZEOF_DATETIME * CHAR_BIT) +#define NPY_BITSOF_TIMEDELTA (NPY_SIZEOF_TIMEDELTA * CHAR_BIT) + +#if NPY_BITSOF_LONG == 8 +#define NPY_INT8 NPY_LONG +#define NPY_UINT8 NPY_ULONG + typedef long npy_int8; + typedef unsigned long npy_uint8; +#define PyInt8ScalarObject PyLongScalarObject +#define PyInt8ArrType_Type PyLongArrType_Type +#define PyUInt8ScalarObject PyULongScalarObject +#define PyUInt8ArrType_Type PyULongArrType_Type +#define NPY_INT8_FMT NPY_LONG_FMT +#define NPY_UINT8_FMT NPY_ULONG_FMT +#elif NPY_BITSOF_LONG == 16 +#define NPY_INT16 NPY_LONG +#define NPY_UINT16 NPY_ULONG + typedef long npy_int16; + typedef unsigned long npy_uint16; +#define PyInt16ScalarObject PyLongScalarObject +#define PyInt16ArrType_Type PyLongArrType_Type +#define PyUInt16ScalarObject PyULongScalarObject +#define PyUInt16ArrType_Type PyULongArrType_Type +#define NPY_INT16_FMT NPY_LONG_FMT +#define NPY_UINT16_FMT NPY_ULONG_FMT +#elif NPY_BITSOF_LONG == 32 +#define NPY_INT32 NPY_LONG +#define NPY_UINT32 NPY_ULONG + typedef long npy_int32; + typedef unsigned long npy_uint32; + typedef unsigned long npy_ucs4; +#define PyInt32ScalarObject PyLongScalarObject +#define PyInt32ArrType_Type PyLongArrType_Type +#define PyUInt32ScalarObject PyULongScalarObject +#define PyUInt32ArrType_Type PyULongArrType_Type +#define NPY_INT32_FMT NPY_LONG_FMT +#define NPY_UINT32_FMT NPY_ULONG_FMT +#elif NPY_BITSOF_LONG == 64 +#define NPY_INT64 NPY_LONG +#define NPY_UINT64 NPY_ULONG + typedef long npy_int64; + typedef unsigned long npy_uint64; +#define PyInt64ScalarObject PyLongScalarObject +#define PyInt64ArrType_Type PyLongArrType_Type +#define PyUInt64ScalarObject PyULongScalarObject +#define PyUInt64ArrType_Type PyULongArrType_Type +#define NPY_INT64_FMT NPY_LONG_FMT +#define NPY_UINT64_FMT NPY_ULONG_FMT +#define MyPyLong_FromInt64 PyLong_FromLong +#define MyPyLong_AsInt64 PyLong_AsLong +#elif NPY_BITSOF_LONG == 128 +#define NPY_INT128 NPY_LONG +#define NPY_UINT128 NPY_ULONG + typedef long npy_int128; + typedef unsigned long npy_uint128; +#define PyInt128ScalarObject PyLongScalarObject +#define PyInt128ArrType_Type PyLongArrType_Type +#define PyUInt128ScalarObject PyULongScalarObject +#define PyUInt128ArrType_Type PyULongArrType_Type +#define NPY_INT128_FMT NPY_LONG_FMT +#define NPY_UINT128_FMT NPY_ULONG_FMT +#endif + +#if NPY_BITSOF_LONGLONG == 8 +# ifndef NPY_INT8 +# define NPY_INT8 NPY_LONGLONG +# define NPY_UINT8 NPY_ULONGLONG + typedef npy_longlong npy_int8; + typedef npy_ulonglong npy_uint8; +# define PyInt8ScalarObject PyLongLongScalarObject +# define PyInt8ArrType_Type PyLongLongArrType_Type +# define PyUInt8ScalarObject PyULongLongScalarObject +# define PyUInt8ArrType_Type PyULongLongArrType_Type +#define NPY_INT8_FMT NPY_LONGLONG_FMT +#define NPY_UINT8_FMT NPY_ULONGLONG_FMT +# endif +# define NPY_MAX_LONGLONG NPY_MAX_INT8 +# define NPY_MIN_LONGLONG NPY_MIN_INT8 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT8 +#elif NPY_BITSOF_LONGLONG == 16 +# ifndef NPY_INT16 +# define NPY_INT16 NPY_LONGLONG +# define NPY_UINT16 NPY_ULONGLONG + typedef npy_longlong npy_int16; + typedef npy_ulonglong npy_uint16; +# define PyInt16ScalarObject PyLongLongScalarObject +# define PyInt16ArrType_Type PyLongLongArrType_Type +# define PyUInt16ScalarObject PyULongLongScalarObject +# define PyUInt16ArrType_Type PyULongLongArrType_Type +#define NPY_INT16_FMT NPY_LONGLONG_FMT +#define NPY_UINT16_FMT NPY_ULONGLONG_FMT +# endif +# define NPY_MAX_LONGLONG NPY_MAX_INT16 +# define NPY_MIN_LONGLONG NPY_MIN_INT16 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT16 +#elif NPY_BITSOF_LONGLONG == 32 +# ifndef NPY_INT32 +# define NPY_INT32 NPY_LONGLONG +# define NPY_UINT32 NPY_ULONGLONG + typedef npy_longlong npy_int32; + typedef npy_ulonglong npy_uint32; + typedef npy_ulonglong npy_ucs4; +# define PyInt32ScalarObject PyLongLongScalarObject +# define PyInt32ArrType_Type PyLongLongArrType_Type +# define PyUInt32ScalarObject PyULongLongScalarObject +# define PyUInt32ArrType_Type PyULongLongArrType_Type +#define NPY_INT32_FMT NPY_LONGLONG_FMT +#define NPY_UINT32_FMT NPY_ULONGLONG_FMT +# endif +# define NPY_MAX_LONGLONG NPY_MAX_INT32 +# define NPY_MIN_LONGLONG NPY_MIN_INT32 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT32 +#elif NPY_BITSOF_LONGLONG == 64 +# ifndef NPY_INT64 +# define NPY_INT64 NPY_LONGLONG +# define NPY_UINT64 NPY_ULONGLONG + typedef npy_longlong npy_int64; + typedef npy_ulonglong npy_uint64; +# define PyInt64ScalarObject PyLongLongScalarObject +# define PyInt64ArrType_Type PyLongLongArrType_Type +# define PyUInt64ScalarObject PyULongLongScalarObject +# define PyUInt64ArrType_Type PyULongLongArrType_Type +#define NPY_INT64_FMT NPY_LONGLONG_FMT +#define NPY_UINT64_FMT NPY_ULONGLONG_FMT +# define MyPyLong_FromInt64 PyLong_FromLongLong +# define MyPyLong_AsInt64 PyLong_AsLongLong +# endif +# define NPY_MAX_LONGLONG NPY_MAX_INT64 +# define NPY_MIN_LONGLONG NPY_MIN_INT64 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT64 +#elif NPY_BITSOF_LONGLONG == 128 +# ifndef NPY_INT128 +# define NPY_INT128 NPY_LONGLONG +# define NPY_UINT128 NPY_ULONGLONG + typedef npy_longlong npy_int128; + typedef npy_ulonglong npy_uint128; +# define PyInt128ScalarObject PyLongLongScalarObject +# define PyInt128ArrType_Type PyLongLongArrType_Type +# define PyUInt128ScalarObject PyULongLongScalarObject +# define PyUInt128ArrType_Type PyULongLongArrType_Type +#define NPY_INT128_FMT NPY_LONGLONG_FMT +#define NPY_UINT128_FMT NPY_ULONGLONG_FMT +# endif +# define NPY_MAX_LONGLONG NPY_MAX_INT128 +# define NPY_MIN_LONGLONG NPY_MIN_INT128 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT128 +#elif NPY_BITSOF_LONGLONG == 256 +# define NPY_INT256 NPY_LONGLONG +# define NPY_UINT256 NPY_ULONGLONG + typedef npy_longlong npy_int256; + typedef npy_ulonglong npy_uint256; +# define PyInt256ScalarObject PyLongLongScalarObject +# define PyInt256ArrType_Type PyLongLongArrType_Type +# define PyUInt256ScalarObject PyULongLongScalarObject +# define PyUInt256ArrType_Type PyULongLongArrType_Type +#define NPY_INT256_FMT NPY_LONGLONG_FMT +#define NPY_UINT256_FMT NPY_ULONGLONG_FMT +# define NPY_MAX_LONGLONG NPY_MAX_INT256 +# define NPY_MIN_LONGLONG NPY_MIN_INT256 +# define NPY_MAX_ULONGLONG NPY_MAX_UINT256 +#endif + +#if NPY_BITSOF_INT == 8 +#ifndef NPY_INT8 +#define NPY_INT8 NPY_INT +#define NPY_UINT8 NPY_UINT + typedef int npy_int8; + typedef unsigned int npy_uint8; +# define PyInt8ScalarObject PyIntScalarObject +# define PyInt8ArrType_Type PyIntArrType_Type +# define PyUInt8ScalarObject PyUIntScalarObject +# define PyUInt8ArrType_Type PyUIntArrType_Type +#define NPY_INT8_FMT NPY_INT_FMT +#define NPY_UINT8_FMT NPY_UINT_FMT +#endif +#elif NPY_BITSOF_INT == 16 +#ifndef NPY_INT16 +#define NPY_INT16 NPY_INT +#define NPY_UINT16 NPY_UINT + typedef int npy_int16; + typedef unsigned int npy_uint16; +# define PyInt16ScalarObject PyIntScalarObject +# define PyInt16ArrType_Type PyIntArrType_Type +# define PyUInt16ScalarObject PyIntUScalarObject +# define PyUInt16ArrType_Type PyIntUArrType_Type +#define NPY_INT16_FMT NPY_INT_FMT +#define NPY_UINT16_FMT NPY_UINT_FMT +#endif +#elif NPY_BITSOF_INT == 32 +#ifndef NPY_INT32 +#define NPY_INT32 NPY_INT +#define NPY_UINT32 NPY_UINT + typedef int npy_int32; + typedef unsigned int npy_uint32; + typedef unsigned int npy_ucs4; +# define PyInt32ScalarObject PyIntScalarObject +# define PyInt32ArrType_Type PyIntArrType_Type +# define PyUInt32ScalarObject PyUIntScalarObject +# define PyUInt32ArrType_Type PyUIntArrType_Type +#define NPY_INT32_FMT NPY_INT_FMT +#define NPY_UINT32_FMT NPY_UINT_FMT +#endif +#elif NPY_BITSOF_INT == 64 +#ifndef NPY_INT64 +#define NPY_INT64 NPY_INT +#define NPY_UINT64 NPY_UINT + typedef int npy_int64; + typedef unsigned int npy_uint64; +# define PyInt64ScalarObject PyIntScalarObject +# define PyInt64ArrType_Type PyIntArrType_Type +# define PyUInt64ScalarObject PyUIntScalarObject +# define PyUInt64ArrType_Type PyUIntArrType_Type +#define NPY_INT64_FMT NPY_INT_FMT +#define NPY_UINT64_FMT NPY_UINT_FMT +# define MyPyLong_FromInt64 PyLong_FromLong +# define MyPyLong_AsInt64 PyLong_AsLong +#endif +#elif NPY_BITSOF_INT == 128 +#ifndef NPY_INT128 +#define NPY_INT128 NPY_INT +#define NPY_UINT128 NPY_UINT + typedef int npy_int128; + typedef unsigned int npy_uint128; +# define PyInt128ScalarObject PyIntScalarObject +# define PyInt128ArrType_Type PyIntArrType_Type +# define PyUInt128ScalarObject PyUIntScalarObject +# define PyUInt128ArrType_Type PyUIntArrType_Type +#define NPY_INT128_FMT NPY_INT_FMT +#define NPY_UINT128_FMT NPY_UINT_FMT +#endif +#endif + +#if NPY_BITSOF_SHORT == 8 +#ifndef NPY_INT8 +#define NPY_INT8 NPY_SHORT +#define NPY_UINT8 NPY_USHORT + typedef short npy_int8; + typedef unsigned short npy_uint8; +# define PyInt8ScalarObject PyShortScalarObject +# define PyInt8ArrType_Type PyShortArrType_Type +# define PyUInt8ScalarObject PyUShortScalarObject +# define PyUInt8ArrType_Type PyUShortArrType_Type +#define NPY_INT8_FMT NPY_SHORT_FMT +#define NPY_UINT8_FMT NPY_USHORT_FMT +#endif +#elif NPY_BITSOF_SHORT == 16 +#ifndef NPY_INT16 +#define NPY_INT16 NPY_SHORT +#define NPY_UINT16 NPY_USHORT + typedef short npy_int16; + typedef unsigned short npy_uint16; +# define PyInt16ScalarObject PyShortScalarObject +# define PyInt16ArrType_Type PyShortArrType_Type +# define PyUInt16ScalarObject PyUShortScalarObject +# define PyUInt16ArrType_Type PyUShortArrType_Type +#define NPY_INT16_FMT NPY_SHORT_FMT +#define NPY_UINT16_FMT NPY_USHORT_FMT +#endif +#elif NPY_BITSOF_SHORT == 32 +#ifndef NPY_INT32 +#define NPY_INT32 NPY_SHORT +#define NPY_UINT32 NPY_USHORT + typedef short npy_int32; + typedef unsigned short npy_uint32; + typedef unsigned short npy_ucs4; +# define PyInt32ScalarObject PyShortScalarObject +# define PyInt32ArrType_Type PyShortArrType_Type +# define PyUInt32ScalarObject PyUShortScalarObject +# define PyUInt32ArrType_Type PyUShortArrType_Type +#define NPY_INT32_FMT NPY_SHORT_FMT +#define NPY_UINT32_FMT NPY_USHORT_FMT +#endif +#elif NPY_BITSOF_SHORT == 64 +#ifndef NPY_INT64 +#define NPY_INT64 NPY_SHORT +#define NPY_UINT64 NPY_USHORT + typedef short npy_int64; + typedef unsigned short npy_uint64; +# define PyInt64ScalarObject PyShortScalarObject +# define PyInt64ArrType_Type PyShortArrType_Type +# define PyUInt64ScalarObject PyUShortScalarObject +# define PyUInt64ArrType_Type PyUShortArrType_Type +#define NPY_INT64_FMT NPY_SHORT_FMT +#define NPY_UINT64_FMT NPY_USHORT_FMT +# define MyPyLong_FromInt64 PyLong_FromLong +# define MyPyLong_AsInt64 PyLong_AsLong +#endif +#elif NPY_BITSOF_SHORT == 128 +#ifndef NPY_INT128 +#define NPY_INT128 NPY_SHORT +#define NPY_UINT128 NPY_USHORT + typedef short npy_int128; + typedef unsigned short npy_uint128; +# define PyInt128ScalarObject PyShortScalarObject +# define PyInt128ArrType_Type PyShortArrType_Type +# define PyUInt128ScalarObject PyUShortScalarObject +# define PyUInt128ArrType_Type PyUShortArrType_Type +#define NPY_INT128_FMT NPY_SHORT_FMT +#define NPY_UINT128_FMT NPY_USHORT_FMT +#endif +#endif + + +#if NPY_BITSOF_CHAR == 8 +#ifndef NPY_INT8 +#define NPY_INT8 NPY_BYTE +#define NPY_UINT8 NPY_UBYTE + typedef signed char npy_int8; + typedef unsigned char npy_uint8; +# define PyInt8ScalarObject PyByteScalarObject +# define PyInt8ArrType_Type PyByteArrType_Type +# define PyUInt8ScalarObject PyUByteScalarObject +# define PyUInt8ArrType_Type PyUByteArrType_Type +#define NPY_INT8_FMT NPY_BYTE_FMT +#define NPY_UINT8_FMT NPY_UBYTE_FMT +#endif +#elif NPY_BITSOF_CHAR == 16 +#ifndef NPY_INT16 +#define NPY_INT16 NPY_BYTE +#define NPY_UINT16 NPY_UBYTE + typedef signed char npy_int16; + typedef unsigned char npy_uint16; +# define PyInt16ScalarObject PyByteScalarObject +# define PyInt16ArrType_Type PyByteArrType_Type +# define PyUInt16ScalarObject PyUByteScalarObject +# define PyUInt16ArrType_Type PyUByteArrType_Type +#define NPY_INT16_FMT NPY_BYTE_FMT +#define NPY_UINT16_FMT NPY_UBYTE_FMT +#endif +#elif NPY_BITSOF_CHAR == 32 +#ifndef NPY_INT32 +#define NPY_INT32 NPY_BYTE +#define NPY_UINT32 NPY_UBYTE + typedef signed char npy_int32; + typedef unsigned char npy_uint32; + typedef unsigned char npy_ucs4; +# define PyInt32ScalarObject PyByteScalarObject +# define PyInt32ArrType_Type PyByteArrType_Type +# define PyUInt32ScalarObject PyUByteScalarObject +# define PyUInt32ArrType_Type PyUByteArrType_Type +#define NPY_INT32_FMT NPY_BYTE_FMT +#define NPY_UINT32_FMT NPY_UBYTE_FMT +#endif +#elif NPY_BITSOF_CHAR == 64 +#ifndef NPY_INT64 +#define NPY_INT64 NPY_BYTE +#define NPY_UINT64 NPY_UBYTE + typedef signed char npy_int64; + typedef unsigned char npy_uint64; +# define PyInt64ScalarObject PyByteScalarObject +# define PyInt64ArrType_Type PyByteArrType_Type +# define PyUInt64ScalarObject PyUByteScalarObject +# define PyUInt64ArrType_Type PyUByteArrType_Type +#define NPY_INT64_FMT NPY_BYTE_FMT +#define NPY_UINT64_FMT NPY_UBYTE_FMT +# define MyPyLong_FromInt64 PyLong_FromLong +# define MyPyLong_AsInt64 PyLong_AsLong +#endif +#elif NPY_BITSOF_CHAR == 128 +#ifndef NPY_INT128 +#define NPY_INT128 NPY_BYTE +#define NPY_UINT128 NPY_UBYTE + typedef signed char npy_int128; + typedef unsigned char npy_uint128; +# define PyInt128ScalarObject PyByteScalarObject +# define PyInt128ArrType_Type PyByteArrType_Type +# define PyUInt128ScalarObject PyUByteScalarObject +# define PyUInt128ArrType_Type PyUByteArrType_Type +#define NPY_INT128_FMT NPY_BYTE_FMT +#define NPY_UINT128_FMT NPY_UBYTE_FMT +#endif +#endif + + + +#if NPY_BITSOF_DOUBLE == 32 +#ifndef NPY_FLOAT32 +#define NPY_FLOAT32 NPY_DOUBLE +#define NPY_COMPLEX64 NPY_CDOUBLE + typedef double npy_float32; + typedef npy_cdouble npy_complex64; +# define PyFloat32ScalarObject PyDoubleScalarObject +# define PyComplex64ScalarObject PyCDoubleScalarObject +# define PyFloat32ArrType_Type PyDoubleArrType_Type +# define PyComplex64ArrType_Type PyCDoubleArrType_Type +#define NPY_FLOAT32_FMT NPY_DOUBLE_FMT +#define NPY_COMPLEX64_FMT NPY_CDOUBLE_FMT +#endif +#elif NPY_BITSOF_DOUBLE == 64 +#ifndef NPY_FLOAT64 +#define NPY_FLOAT64 NPY_DOUBLE +#define NPY_COMPLEX128 NPY_CDOUBLE + typedef double npy_float64; + typedef npy_cdouble npy_complex128; +# define PyFloat64ScalarObject PyDoubleScalarObject +# define PyComplex128ScalarObject PyCDoubleScalarObject +# define PyFloat64ArrType_Type PyDoubleArrType_Type +# define PyComplex128ArrType_Type PyCDoubleArrType_Type +#define NPY_FLOAT64_FMT NPY_DOUBLE_FMT +#define NPY_COMPLEX128_FMT NPY_CDOUBLE_FMT +#endif +#elif NPY_BITSOF_DOUBLE == 80 +#ifndef NPY_FLOAT80 +#define NPY_FLOAT80 NPY_DOUBLE +#define NPY_COMPLEX160 NPY_CDOUBLE + typedef double npy_float80; + typedef npy_cdouble npy_complex160; +# define PyFloat80ScalarObject PyDoubleScalarObject +# define PyComplex160ScalarObject PyCDoubleScalarObject +# define PyFloat80ArrType_Type PyDoubleArrType_Type +# define PyComplex160ArrType_Type PyCDoubleArrType_Type +#define NPY_FLOAT80_FMT NPY_DOUBLE_FMT +#define NPY_COMPLEX160_FMT NPY_CDOUBLE_FMT +#endif +#elif NPY_BITSOF_DOUBLE == 96 +#ifndef NPY_FLOAT96 +#define NPY_FLOAT96 NPY_DOUBLE +#define NPY_COMPLEX192 NPY_CDOUBLE + typedef double npy_float96; + typedef npy_cdouble npy_complex192; +# define PyFloat96ScalarObject PyDoubleScalarObject +# define PyComplex192ScalarObject PyCDoubleScalarObject +# define PyFloat96ArrType_Type PyDoubleArrType_Type +# define PyComplex192ArrType_Type PyCDoubleArrType_Type +#define NPY_FLOAT96_FMT NPY_DOUBLE_FMT +#define NPY_COMPLEX192_FMT NPY_CDOUBLE_FMT +#endif +#elif NPY_BITSOF_DOUBLE == 128 +#ifndef NPY_FLOAT128 +#define NPY_FLOAT128 NPY_DOUBLE +#define NPY_COMPLEX256 NPY_CDOUBLE + typedef double npy_float128; + typedef npy_cdouble npy_complex256; +# define PyFloat128ScalarObject PyDoubleScalarObject +# define PyComplex256ScalarObject PyCDoubleScalarObject +# define PyFloat128ArrType_Type PyDoubleArrType_Type +# define PyComplex256ArrType_Type PyCDoubleArrType_Type +#define NPY_FLOAT128_FMT NPY_DOUBLE_FMT +#define NPY_COMPLEX256_FMT NPY_CDOUBLE_FMT +#endif +#endif + + + +#if NPY_BITSOF_FLOAT == 32 +#ifndef NPY_FLOAT32 +#define NPY_FLOAT32 NPY_FLOAT +#define NPY_COMPLEX64 NPY_CFLOAT + typedef float npy_float32; + typedef npy_cfloat npy_complex64; +# define PyFloat32ScalarObject PyFloatScalarObject +# define PyComplex64ScalarObject PyCFloatScalarObject +# define PyFloat32ArrType_Type PyFloatArrType_Type +# define PyComplex64ArrType_Type PyCFloatArrType_Type +#define NPY_FLOAT32_FMT NPY_FLOAT_FMT +#define NPY_COMPLEX64_FMT NPY_CFLOAT_FMT +#endif +#elif NPY_BITSOF_FLOAT == 64 +#ifndef NPY_FLOAT64 +#define NPY_FLOAT64 NPY_FLOAT +#define NPY_COMPLEX128 NPY_CFLOAT + typedef float npy_float64; + typedef npy_cfloat npy_complex128; +# define PyFloat64ScalarObject PyFloatScalarObject +# define PyComplex128ScalarObject PyCFloatScalarObject +# define PyFloat64ArrType_Type PyFloatArrType_Type +# define PyComplex128ArrType_Type PyCFloatArrType_Type +#define NPY_FLOAT64_FMT NPY_FLOAT_FMT +#define NPY_COMPLEX128_FMT NPY_CFLOAT_FMT +#endif +#elif NPY_BITSOF_FLOAT == 80 +#ifndef NPY_FLOAT80 +#define NPY_FLOAT80 NPY_FLOAT +#define NPY_COMPLEX160 NPY_CFLOAT + typedef float npy_float80; + typedef npy_cfloat npy_complex160; +# define PyFloat80ScalarObject PyFloatScalarObject +# define PyComplex160ScalarObject PyCFloatScalarObject +# define PyFloat80ArrType_Type PyFloatArrType_Type +# define PyComplex160ArrType_Type PyCFloatArrType_Type +#define NPY_FLOAT80_FMT NPY_FLOAT_FMT +#define NPY_COMPLEX160_FMT NPY_CFLOAT_FMT +#endif +#elif NPY_BITSOF_FLOAT == 96 +#ifndef NPY_FLOAT96 +#define NPY_FLOAT96 NPY_FLOAT +#define NPY_COMPLEX192 NPY_CFLOAT + typedef float npy_float96; + typedef npy_cfloat npy_complex192; +# define PyFloat96ScalarObject PyFloatScalarObject +# define PyComplex192ScalarObject PyCFloatScalarObject +# define PyFloat96ArrType_Type PyFloatArrType_Type +# define PyComplex192ArrType_Type PyCFloatArrType_Type +#define NPY_FLOAT96_FMT NPY_FLOAT_FMT +#define NPY_COMPLEX192_FMT NPY_CFLOAT_FMT +#endif +#elif NPY_BITSOF_FLOAT == 128 +#ifndef NPY_FLOAT128 +#define NPY_FLOAT128 NPY_FLOAT +#define NPY_COMPLEX256 NPY_CFLOAT + typedef float npy_float128; + typedef npy_cfloat npy_complex256; +# define PyFloat128ScalarObject PyFloatScalarObject +# define PyComplex256ScalarObject PyCFloatScalarObject +# define PyFloat128ArrType_Type PyFloatArrType_Type +# define PyComplex256ArrType_Type PyCFloatArrType_Type +#define NPY_FLOAT128_FMT NPY_FLOAT_FMT +#define NPY_COMPLEX256_FMT NPY_CFLOAT_FMT +#endif +#endif + +/* half/float16 isn't a floating-point type in C */ +#define NPY_FLOAT16 NPY_HALF +typedef npy_uint16 npy_half; +typedef npy_half npy_float16; + +#if NPY_BITSOF_LONGDOUBLE == 32 +#ifndef NPY_FLOAT32 +#define NPY_FLOAT32 NPY_LONGDOUBLE +#define NPY_COMPLEX64 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float32; + typedef npy_clongdouble npy_complex64; +# define PyFloat32ScalarObject PyLongDoubleScalarObject +# define PyComplex64ScalarObject PyCLongDoubleScalarObject +# define PyFloat32ArrType_Type PyLongDoubleArrType_Type +# define PyComplex64ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT32_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX64_FMT NPY_CLONGDOUBLE_FMT +#endif +#elif NPY_BITSOF_LONGDOUBLE == 64 +#ifndef NPY_FLOAT64 +#define NPY_FLOAT64 NPY_LONGDOUBLE +#define NPY_COMPLEX128 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float64; + typedef npy_clongdouble npy_complex128; +# define PyFloat64ScalarObject PyLongDoubleScalarObject +# define PyComplex128ScalarObject PyCLongDoubleScalarObject +# define PyFloat64ArrType_Type PyLongDoubleArrType_Type +# define PyComplex128ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT64_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX128_FMT NPY_CLONGDOUBLE_FMT +#endif +#elif NPY_BITSOF_LONGDOUBLE == 80 +#ifndef NPY_FLOAT80 +#define NPY_FLOAT80 NPY_LONGDOUBLE +#define NPY_COMPLEX160 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float80; + typedef npy_clongdouble npy_complex160; +# define PyFloat80ScalarObject PyLongDoubleScalarObject +# define PyComplex160ScalarObject PyCLongDoubleScalarObject +# define PyFloat80ArrType_Type PyLongDoubleArrType_Type +# define PyComplex160ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT80_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX160_FMT NPY_CLONGDOUBLE_FMT +#endif +#elif NPY_BITSOF_LONGDOUBLE == 96 +#ifndef NPY_FLOAT96 +#define NPY_FLOAT96 NPY_LONGDOUBLE +#define NPY_COMPLEX192 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float96; + typedef npy_clongdouble npy_complex192; +# define PyFloat96ScalarObject PyLongDoubleScalarObject +# define PyComplex192ScalarObject PyCLongDoubleScalarObject +# define PyFloat96ArrType_Type PyLongDoubleArrType_Type +# define PyComplex192ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT96_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX192_FMT NPY_CLONGDOUBLE_FMT +#endif +#elif NPY_BITSOF_LONGDOUBLE == 128 +#ifndef NPY_FLOAT128 +#define NPY_FLOAT128 NPY_LONGDOUBLE +#define NPY_COMPLEX256 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float128; + typedef npy_clongdouble npy_complex256; +# define PyFloat128ScalarObject PyLongDoubleScalarObject +# define PyComplex256ScalarObject PyCLongDoubleScalarObject +# define PyFloat128ArrType_Type PyLongDoubleArrType_Type +# define PyComplex256ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT128_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX256_FMT NPY_CLONGDOUBLE_FMT +#endif +#elif NPY_BITSOF_LONGDOUBLE == 256 +#define NPY_FLOAT256 NPY_LONGDOUBLE +#define NPY_COMPLEX512 NPY_CLONGDOUBLE + typedef npy_longdouble npy_float256; + typedef npy_clongdouble npy_complex512; +# define PyFloat256ScalarObject PyLongDoubleScalarObject +# define PyComplex512ScalarObject PyCLongDoubleScalarObject +# define PyFloat256ArrType_Type PyLongDoubleArrType_Type +# define PyComplex512ArrType_Type PyCLongDoubleArrType_Type +#define NPY_FLOAT256_FMT NPY_LONGDOUBLE_FMT +#define NPY_COMPLEX512_FMT NPY_CLONGDOUBLE_FMT +#endif + +/* datetime typedefs */ +typedef npy_int64 npy_timedelta; +typedef npy_int64 npy_datetime; +#define NPY_DATETIME_FMT NPY_INT64_FMT +#define NPY_TIMEDELTA_FMT NPY_INT64_FMT + +/* End of typedefs for numarray style bit-width names */ + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_COMMON_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_endian.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_endian.h new file mode 100644 index 0000000000000000000000000000000000000000..09262120bf82c5065d1f914a5825ce2de0063d4d --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_endian.h @@ -0,0 +1,78 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_ENDIAN_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_ENDIAN_H_ + +/* + * NPY_BYTE_ORDER is set to the same value as BYTE_ORDER set by glibc in + * endian.h + */ + +#if defined(NPY_HAVE_ENDIAN_H) || defined(NPY_HAVE_SYS_ENDIAN_H) + /* Use endian.h if available */ + + #if defined(NPY_HAVE_ENDIAN_H) + #include + #elif defined(NPY_HAVE_SYS_ENDIAN_H) + #include + #endif + + #if defined(BYTE_ORDER) && defined(BIG_ENDIAN) && defined(LITTLE_ENDIAN) + #define NPY_BYTE_ORDER BYTE_ORDER + #define NPY_LITTLE_ENDIAN LITTLE_ENDIAN + #define NPY_BIG_ENDIAN BIG_ENDIAN + #elif defined(_BYTE_ORDER) && defined(_BIG_ENDIAN) && defined(_LITTLE_ENDIAN) + #define NPY_BYTE_ORDER _BYTE_ORDER + #define NPY_LITTLE_ENDIAN _LITTLE_ENDIAN + #define NPY_BIG_ENDIAN _BIG_ENDIAN + #elif defined(__BYTE_ORDER) && defined(__BIG_ENDIAN) && defined(__LITTLE_ENDIAN) + #define NPY_BYTE_ORDER __BYTE_ORDER + #define NPY_LITTLE_ENDIAN __LITTLE_ENDIAN + #define NPY_BIG_ENDIAN __BIG_ENDIAN + #endif +#endif + +#ifndef NPY_BYTE_ORDER + /* Set endianness info using target CPU */ + #include "npy_cpu.h" + + #define NPY_LITTLE_ENDIAN 1234 + #define NPY_BIG_ENDIAN 4321 + + #if defined(NPY_CPU_X86) \ + || defined(NPY_CPU_AMD64) \ + || defined(NPY_CPU_IA64) \ + || defined(NPY_CPU_ALPHA) \ + || defined(NPY_CPU_ARMEL) \ + || defined(NPY_CPU_ARMEL_AARCH32) \ + || defined(NPY_CPU_ARMEL_AARCH64) \ + || defined(NPY_CPU_SH_LE) \ + || defined(NPY_CPU_MIPSEL) \ + || defined(NPY_CPU_PPC64LE) \ + || defined(NPY_CPU_ARCEL) \ + || defined(NPY_CPU_RISCV64) \ + || defined(NPY_CPU_RISCV32) \ + || defined(NPY_CPU_LOONGARCH) \ + || defined(NPY_CPU_WASM) + #define NPY_BYTE_ORDER NPY_LITTLE_ENDIAN + + #elif defined(NPY_CPU_PPC) \ + || defined(NPY_CPU_SPARC) \ + || defined(NPY_CPU_S390) \ + || defined(NPY_CPU_HPPA) \ + || defined(NPY_CPU_PPC64) \ + || defined(NPY_CPU_ARMEB) \ + || defined(NPY_CPU_ARMEB_AARCH32) \ + || defined(NPY_CPU_ARMEB_AARCH64) \ + || defined(NPY_CPU_SH_BE) \ + || defined(NPY_CPU_MIPSEB) \ + || defined(NPY_CPU_OR1K) \ + || defined(NPY_CPU_M68K) \ + || defined(NPY_CPU_ARCEB) + #define NPY_BYTE_ORDER NPY_BIG_ENDIAN + + #else + #error Unknown CPU: can not set endianness + #endif + +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_ENDIAN_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_math.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_math.h new file mode 100644 index 0000000000000000000000000000000000000000..abc784bc686c18d162a4730a7c60393622023aae --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_math.h @@ -0,0 +1,602 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_MATH_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_MATH_H_ + +#include + +#include + +/* By adding static inline specifiers to npy_math function definitions when + appropriate, compiler is given the opportunity to optimize */ +#if NPY_INLINE_MATH +#define NPY_INPLACE static inline +#else +#define NPY_INPLACE +#endif + + +#ifdef __cplusplus +extern "C" { +#endif + +#define PyArray_MAX(a,b) (((a)>(b))?(a):(b)) +#define PyArray_MIN(a,b) (((a)<(b))?(a):(b)) + +/* + * NAN and INFINITY like macros (same behavior as glibc for NAN, same as C99 + * for INFINITY) + * + * XXX: I should test whether INFINITY and NAN are available on the platform + */ +static inline float __npy_inff(void) +{ + const union { npy_uint32 __i; float __f;} __bint = {0x7f800000UL}; + return __bint.__f; +} + +static inline float __npy_nanf(void) +{ + const union { npy_uint32 __i; float __f;} __bint = {0x7fc00000UL}; + return __bint.__f; +} + +static inline float __npy_pzerof(void) +{ + const union { npy_uint32 __i; float __f;} __bint = {0x00000000UL}; + return __bint.__f; +} + +static inline float __npy_nzerof(void) +{ + const union { npy_uint32 __i; float __f;} __bint = {0x80000000UL}; + return __bint.__f; +} + +#define NPY_INFINITYF __npy_inff() +#define NPY_NANF __npy_nanf() +#define NPY_PZEROF __npy_pzerof() +#define NPY_NZEROF __npy_nzerof() + +#define NPY_INFINITY ((npy_double)NPY_INFINITYF) +#define NPY_NAN ((npy_double)NPY_NANF) +#define NPY_PZERO ((npy_double)NPY_PZEROF) +#define NPY_NZERO ((npy_double)NPY_NZEROF) + +#define NPY_INFINITYL ((npy_longdouble)NPY_INFINITYF) +#define NPY_NANL ((npy_longdouble)NPY_NANF) +#define NPY_PZEROL ((npy_longdouble)NPY_PZEROF) +#define NPY_NZEROL ((npy_longdouble)NPY_NZEROF) + +/* + * Useful constants + */ +#define NPY_E 2.718281828459045235360287471352662498 /* e */ +#define NPY_LOG2E 1.442695040888963407359924681001892137 /* log_2 e */ +#define NPY_LOG10E 0.434294481903251827651128918916605082 /* log_10 e */ +#define NPY_LOGE2 0.693147180559945309417232121458176568 /* log_e 2 */ +#define NPY_LOGE10 2.302585092994045684017991454684364208 /* log_e 10 */ +#define NPY_PI 3.141592653589793238462643383279502884 /* pi */ +#define NPY_PI_2 1.570796326794896619231321691639751442 /* pi/2 */ +#define NPY_PI_4 0.785398163397448309615660845819875721 /* pi/4 */ +#define NPY_1_PI 0.318309886183790671537767526745028724 /* 1/pi */ +#define NPY_2_PI 0.636619772367581343075535053490057448 /* 2/pi */ +#define NPY_EULER 0.577215664901532860606512090082402431 /* Euler constant */ +#define NPY_SQRT2 1.414213562373095048801688724209698079 /* sqrt(2) */ +#define NPY_SQRT1_2 0.707106781186547524400844362104849039 /* 1/sqrt(2) */ + +#define NPY_Ef 2.718281828459045235360287471352662498F /* e */ +#define NPY_LOG2Ef 1.442695040888963407359924681001892137F /* log_2 e */ +#define NPY_LOG10Ef 0.434294481903251827651128918916605082F /* log_10 e */ +#define NPY_LOGE2f 0.693147180559945309417232121458176568F /* log_e 2 */ +#define NPY_LOGE10f 2.302585092994045684017991454684364208F /* log_e 10 */ +#define NPY_PIf 3.141592653589793238462643383279502884F /* pi */ +#define NPY_PI_2f 1.570796326794896619231321691639751442F /* pi/2 */ +#define NPY_PI_4f 0.785398163397448309615660845819875721F /* pi/4 */ +#define NPY_1_PIf 0.318309886183790671537767526745028724F /* 1/pi */ +#define NPY_2_PIf 0.636619772367581343075535053490057448F /* 2/pi */ +#define NPY_EULERf 0.577215664901532860606512090082402431F /* Euler constant */ +#define NPY_SQRT2f 1.414213562373095048801688724209698079F /* sqrt(2) */ +#define NPY_SQRT1_2f 0.707106781186547524400844362104849039F /* 1/sqrt(2) */ + +#define NPY_El 2.718281828459045235360287471352662498L /* e */ +#define NPY_LOG2El 1.442695040888963407359924681001892137L /* log_2 e */ +#define NPY_LOG10El 0.434294481903251827651128918916605082L /* log_10 e */ +#define NPY_LOGE2l 0.693147180559945309417232121458176568L /* log_e 2 */ +#define NPY_LOGE10l 2.302585092994045684017991454684364208L /* log_e 10 */ +#define NPY_PIl 3.141592653589793238462643383279502884L /* pi */ +#define NPY_PI_2l 1.570796326794896619231321691639751442L /* pi/2 */ +#define NPY_PI_4l 0.785398163397448309615660845819875721L /* pi/4 */ +#define NPY_1_PIl 0.318309886183790671537767526745028724L /* 1/pi */ +#define NPY_2_PIl 0.636619772367581343075535053490057448L /* 2/pi */ +#define NPY_EULERl 0.577215664901532860606512090082402431L /* Euler constant */ +#define NPY_SQRT2l 1.414213562373095048801688724209698079L /* sqrt(2) */ +#define NPY_SQRT1_2l 0.707106781186547524400844362104849039L /* 1/sqrt(2) */ + +/* + * Integer functions. + */ +NPY_INPLACE npy_uint npy_gcdu(npy_uint a, npy_uint b); +NPY_INPLACE npy_uint npy_lcmu(npy_uint a, npy_uint b); +NPY_INPLACE npy_ulong npy_gcdul(npy_ulong a, npy_ulong b); +NPY_INPLACE npy_ulong npy_lcmul(npy_ulong a, npy_ulong b); +NPY_INPLACE npy_ulonglong npy_gcdull(npy_ulonglong a, npy_ulonglong b); +NPY_INPLACE npy_ulonglong npy_lcmull(npy_ulonglong a, npy_ulonglong b); + +NPY_INPLACE npy_int npy_gcd(npy_int a, npy_int b); +NPY_INPLACE npy_int npy_lcm(npy_int a, npy_int b); +NPY_INPLACE npy_long npy_gcdl(npy_long a, npy_long b); +NPY_INPLACE npy_long npy_lcml(npy_long a, npy_long b); +NPY_INPLACE npy_longlong npy_gcdll(npy_longlong a, npy_longlong b); +NPY_INPLACE npy_longlong npy_lcmll(npy_longlong a, npy_longlong b); + +NPY_INPLACE npy_ubyte npy_rshiftuhh(npy_ubyte a, npy_ubyte b); +NPY_INPLACE npy_ubyte npy_lshiftuhh(npy_ubyte a, npy_ubyte b); +NPY_INPLACE npy_ushort npy_rshiftuh(npy_ushort a, npy_ushort b); +NPY_INPLACE npy_ushort npy_lshiftuh(npy_ushort a, npy_ushort b); +NPY_INPLACE npy_uint npy_rshiftu(npy_uint a, npy_uint b); +NPY_INPLACE npy_uint npy_lshiftu(npy_uint a, npy_uint b); +NPY_INPLACE npy_ulong npy_rshiftul(npy_ulong a, npy_ulong b); +NPY_INPLACE npy_ulong npy_lshiftul(npy_ulong a, npy_ulong b); +NPY_INPLACE npy_ulonglong npy_rshiftull(npy_ulonglong a, npy_ulonglong b); +NPY_INPLACE npy_ulonglong npy_lshiftull(npy_ulonglong a, npy_ulonglong b); + +NPY_INPLACE npy_byte npy_rshifthh(npy_byte a, npy_byte b); +NPY_INPLACE npy_byte npy_lshifthh(npy_byte a, npy_byte b); +NPY_INPLACE npy_short npy_rshifth(npy_short a, npy_short b); +NPY_INPLACE npy_short npy_lshifth(npy_short a, npy_short b); +NPY_INPLACE npy_int npy_rshift(npy_int a, npy_int b); +NPY_INPLACE npy_int npy_lshift(npy_int a, npy_int b); +NPY_INPLACE npy_long npy_rshiftl(npy_long a, npy_long b); +NPY_INPLACE npy_long npy_lshiftl(npy_long a, npy_long b); +NPY_INPLACE npy_longlong npy_rshiftll(npy_longlong a, npy_longlong b); +NPY_INPLACE npy_longlong npy_lshiftll(npy_longlong a, npy_longlong b); + +NPY_INPLACE uint8_t npy_popcountuhh(npy_ubyte a); +NPY_INPLACE uint8_t npy_popcountuh(npy_ushort a); +NPY_INPLACE uint8_t npy_popcountu(npy_uint a); +NPY_INPLACE uint8_t npy_popcountul(npy_ulong a); +NPY_INPLACE uint8_t npy_popcountull(npy_ulonglong a); +NPY_INPLACE uint8_t npy_popcounthh(npy_byte a); +NPY_INPLACE uint8_t npy_popcounth(npy_short a); +NPY_INPLACE uint8_t npy_popcount(npy_int a); +NPY_INPLACE uint8_t npy_popcountl(npy_long a); +NPY_INPLACE uint8_t npy_popcountll(npy_longlong a); + +/* + * C99 double math funcs that need fixups or are blocklist-able + */ +NPY_INPLACE double npy_sin(double x); +NPY_INPLACE double npy_cos(double x); +NPY_INPLACE double npy_tan(double x); +NPY_INPLACE double npy_hypot(double x, double y); +NPY_INPLACE double npy_log2(double x); +NPY_INPLACE double npy_atan2(double x, double y); + +/* Mandatory C99 double math funcs, no blocklisting or fixups */ +/* defined for legacy reasons, should be deprecated at some point */ +#define npy_sinh sinh +#define npy_cosh cosh +#define npy_tanh tanh +#define npy_asin asin +#define npy_acos acos +#define npy_atan atan +#define npy_log log +#define npy_log10 log10 +#define npy_cbrt cbrt +#define npy_fabs fabs +#define npy_ceil ceil +#define npy_fmod fmod +#define npy_floor floor +#define npy_expm1 expm1 +#define npy_log1p log1p +#define npy_acosh acosh +#define npy_asinh asinh +#define npy_atanh atanh +#define npy_rint rint +#define npy_trunc trunc +#define npy_exp2 exp2 +#define npy_frexp frexp +#define npy_ldexp ldexp +#define npy_copysign copysign +#define npy_exp exp +#define npy_sqrt sqrt +#define npy_pow pow +#define npy_modf modf +#define npy_nextafter nextafter + +double npy_spacing(double x); + +/* + * IEEE 754 fpu handling + */ + +/* use builtins to avoid function calls in tight loops + * only available if npy_config.h is available (= numpys own build) */ +#ifdef HAVE___BUILTIN_ISNAN + #define npy_isnan(x) __builtin_isnan(x) +#else + #define npy_isnan(x) isnan(x) +#endif + + +/* only available if npy_config.h is available (= numpys own build) */ +#ifdef HAVE___BUILTIN_ISFINITE + #define npy_isfinite(x) __builtin_isfinite(x) +#else + #define npy_isfinite(x) isfinite((x)) +#endif + +/* only available if npy_config.h is available (= numpys own build) */ +#ifdef HAVE___BUILTIN_ISINF + #define npy_isinf(x) __builtin_isinf(x) +#else + #define npy_isinf(x) isinf((x)) +#endif + +#define npy_signbit(x) signbit((x)) + +/* + * float C99 math funcs that need fixups or are blocklist-able + */ +NPY_INPLACE float npy_sinf(float x); +NPY_INPLACE float npy_cosf(float x); +NPY_INPLACE float npy_tanf(float x); +NPY_INPLACE float npy_expf(float x); +NPY_INPLACE float npy_sqrtf(float x); +NPY_INPLACE float npy_hypotf(float x, float y); +NPY_INPLACE float npy_log2f(float x); +NPY_INPLACE float npy_atan2f(float x, float y); +NPY_INPLACE float npy_powf(float x, float y); +NPY_INPLACE float npy_modff(float x, float* y); + +/* Mandatory C99 float math funcs, no blocklisting or fixups */ +/* defined for legacy reasons, should be deprecated at some point */ + +#define npy_sinhf sinhf +#define npy_coshf coshf +#define npy_tanhf tanhf +#define npy_asinf asinf +#define npy_acosf acosf +#define npy_atanf atanf +#define npy_logf logf +#define npy_log10f log10f +#define npy_cbrtf cbrtf +#define npy_fabsf fabsf +#define npy_ceilf ceilf +#define npy_fmodf fmodf +#define npy_floorf floorf +#define npy_expm1f expm1f +#define npy_log1pf log1pf +#define npy_asinhf asinhf +#define npy_acoshf acoshf +#define npy_atanhf atanhf +#define npy_rintf rintf +#define npy_truncf truncf +#define npy_exp2f exp2f +#define npy_frexpf frexpf +#define npy_ldexpf ldexpf +#define npy_copysignf copysignf +#define npy_nextafterf nextafterf + +float npy_spacingf(float x); + +/* + * long double C99 double math funcs that need fixups or are blocklist-able + */ +NPY_INPLACE npy_longdouble npy_sinl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_cosl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_tanl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_expl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_sqrtl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_hypotl(npy_longdouble x, npy_longdouble y); +NPY_INPLACE npy_longdouble npy_log2l(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_atan2l(npy_longdouble x, npy_longdouble y); +NPY_INPLACE npy_longdouble npy_powl(npy_longdouble x, npy_longdouble y); +NPY_INPLACE npy_longdouble npy_modfl(npy_longdouble x, npy_longdouble* y); + +/* Mandatory C99 double math funcs, no blocklisting or fixups */ +/* defined for legacy reasons, should be deprecated at some point */ +#define npy_sinhl sinhl +#define npy_coshl coshl +#define npy_tanhl tanhl +#define npy_fabsl fabsl +#define npy_floorl floorl +#define npy_ceill ceill +#define npy_rintl rintl +#define npy_truncl truncl +#define npy_cbrtl cbrtl +#define npy_log10l log10l +#define npy_logl logl +#define npy_expm1l expm1l +#define npy_asinl asinl +#define npy_acosl acosl +#define npy_atanl atanl +#define npy_asinhl asinhl +#define npy_acoshl acoshl +#define npy_atanhl atanhl +#define npy_log1pl log1pl +#define npy_exp2l exp2l +#define npy_fmodl fmodl +#define npy_frexpl frexpl +#define npy_ldexpl ldexpl +#define npy_copysignl copysignl +#define npy_nextafterl nextafterl + +npy_longdouble npy_spacingl(npy_longdouble x); + +/* + * Non standard functions + */ +NPY_INPLACE double npy_deg2rad(double x); +NPY_INPLACE double npy_rad2deg(double x); +NPY_INPLACE double npy_logaddexp(double x, double y); +NPY_INPLACE double npy_logaddexp2(double x, double y); +NPY_INPLACE double npy_divmod(double x, double y, double *modulus); +NPY_INPLACE double npy_heaviside(double x, double h0); + +NPY_INPLACE float npy_deg2radf(float x); +NPY_INPLACE float npy_rad2degf(float x); +NPY_INPLACE float npy_logaddexpf(float x, float y); +NPY_INPLACE float npy_logaddexp2f(float x, float y); +NPY_INPLACE float npy_divmodf(float x, float y, float *modulus); +NPY_INPLACE float npy_heavisidef(float x, float h0); + +NPY_INPLACE npy_longdouble npy_deg2radl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_rad2degl(npy_longdouble x); +NPY_INPLACE npy_longdouble npy_logaddexpl(npy_longdouble x, npy_longdouble y); +NPY_INPLACE npy_longdouble npy_logaddexp2l(npy_longdouble x, npy_longdouble y); +NPY_INPLACE npy_longdouble npy_divmodl(npy_longdouble x, npy_longdouble y, + npy_longdouble *modulus); +NPY_INPLACE npy_longdouble npy_heavisidel(npy_longdouble x, npy_longdouble h0); + +#define npy_degrees npy_rad2deg +#define npy_degreesf npy_rad2degf +#define npy_degreesl npy_rad2degl + +#define npy_radians npy_deg2rad +#define npy_radiansf npy_deg2radf +#define npy_radiansl npy_deg2radl + +/* + * Complex declarations + */ + +static inline double npy_creal(const npy_cdouble z) +{ +#if defined(__cplusplus) + return z._Val[0]; +#else + return creal(z); +#endif +} + +static inline void npy_csetreal(npy_cdouble *z, const double r) +{ + ((double *) z)[0] = r; +} + +static inline double npy_cimag(const npy_cdouble z) +{ +#if defined(__cplusplus) + return z._Val[1]; +#else + return cimag(z); +#endif +} + +static inline void npy_csetimag(npy_cdouble *z, const double i) +{ + ((double *) z)[1] = i; +} + +static inline float npy_crealf(const npy_cfloat z) +{ +#if defined(__cplusplus) + return z._Val[0]; +#else + return crealf(z); +#endif +} + +static inline void npy_csetrealf(npy_cfloat *z, const float r) +{ + ((float *) z)[0] = r; +} + +static inline float npy_cimagf(const npy_cfloat z) +{ +#if defined(__cplusplus) + return z._Val[1]; +#else + return cimagf(z); +#endif +} + +static inline void npy_csetimagf(npy_cfloat *z, const float i) +{ + ((float *) z)[1] = i; +} + +static inline npy_longdouble npy_creall(const npy_clongdouble z) +{ +#if defined(__cplusplus) + return (npy_longdouble)z._Val[0]; +#else + return creall(z); +#endif +} + +static inline void npy_csetreall(npy_clongdouble *z, const longdouble_t r) +{ + ((longdouble_t *) z)[0] = r; +} + +static inline npy_longdouble npy_cimagl(const npy_clongdouble z) +{ +#if defined(__cplusplus) + return (npy_longdouble)z._Val[1]; +#else + return cimagl(z); +#endif +} + +static inline void npy_csetimagl(npy_clongdouble *z, const longdouble_t i) +{ + ((longdouble_t *) z)[1] = i; +} + +#define NPY_CSETREAL(z, r) npy_csetreal(z, r) +#define NPY_CSETIMAG(z, i) npy_csetimag(z, i) +#define NPY_CSETREALF(z, r) npy_csetrealf(z, r) +#define NPY_CSETIMAGF(z, i) npy_csetimagf(z, i) +#define NPY_CSETREALL(z, r) npy_csetreall(z, r) +#define NPY_CSETIMAGL(z, i) npy_csetimagl(z, i) + +static inline npy_cdouble npy_cpack(double x, double y) +{ + npy_cdouble z; + npy_csetreal(&z, x); + npy_csetimag(&z, y); + return z; +} + +static inline npy_cfloat npy_cpackf(float x, float y) +{ + npy_cfloat z; + npy_csetrealf(&z, x); + npy_csetimagf(&z, y); + return z; +} + +static inline npy_clongdouble npy_cpackl(npy_longdouble x, npy_longdouble y) +{ + npy_clongdouble z; + npy_csetreall(&z, x); + npy_csetimagl(&z, y); + return z; +} + +/* + * Double precision complex functions + */ +double npy_cabs(npy_cdouble z); +double npy_carg(npy_cdouble z); + +npy_cdouble npy_cexp(npy_cdouble z); +npy_cdouble npy_clog(npy_cdouble z); +npy_cdouble npy_cpow(npy_cdouble x, npy_cdouble y); + +npy_cdouble npy_csqrt(npy_cdouble z); + +npy_cdouble npy_ccos(npy_cdouble z); +npy_cdouble npy_csin(npy_cdouble z); +npy_cdouble npy_ctan(npy_cdouble z); + +npy_cdouble npy_ccosh(npy_cdouble z); +npy_cdouble npy_csinh(npy_cdouble z); +npy_cdouble npy_ctanh(npy_cdouble z); + +npy_cdouble npy_cacos(npy_cdouble z); +npy_cdouble npy_casin(npy_cdouble z); +npy_cdouble npy_catan(npy_cdouble z); + +npy_cdouble npy_cacosh(npy_cdouble z); +npy_cdouble npy_casinh(npy_cdouble z); +npy_cdouble npy_catanh(npy_cdouble z); + +/* + * Single precision complex functions + */ +float npy_cabsf(npy_cfloat z); +float npy_cargf(npy_cfloat z); + +npy_cfloat npy_cexpf(npy_cfloat z); +npy_cfloat npy_clogf(npy_cfloat z); +npy_cfloat npy_cpowf(npy_cfloat x, npy_cfloat y); + +npy_cfloat npy_csqrtf(npy_cfloat z); + +npy_cfloat npy_ccosf(npy_cfloat z); +npy_cfloat npy_csinf(npy_cfloat z); +npy_cfloat npy_ctanf(npy_cfloat z); + +npy_cfloat npy_ccoshf(npy_cfloat z); +npy_cfloat npy_csinhf(npy_cfloat z); +npy_cfloat npy_ctanhf(npy_cfloat z); + +npy_cfloat npy_cacosf(npy_cfloat z); +npy_cfloat npy_casinf(npy_cfloat z); +npy_cfloat npy_catanf(npy_cfloat z); + +npy_cfloat npy_cacoshf(npy_cfloat z); +npy_cfloat npy_casinhf(npy_cfloat z); +npy_cfloat npy_catanhf(npy_cfloat z); + + +/* + * Extended precision complex functions + */ +npy_longdouble npy_cabsl(npy_clongdouble z); +npy_longdouble npy_cargl(npy_clongdouble z); + +npy_clongdouble npy_cexpl(npy_clongdouble z); +npy_clongdouble npy_clogl(npy_clongdouble z); +npy_clongdouble npy_cpowl(npy_clongdouble x, npy_clongdouble y); + +npy_clongdouble npy_csqrtl(npy_clongdouble z); + +npy_clongdouble npy_ccosl(npy_clongdouble z); +npy_clongdouble npy_csinl(npy_clongdouble z); +npy_clongdouble npy_ctanl(npy_clongdouble z); + +npy_clongdouble npy_ccoshl(npy_clongdouble z); +npy_clongdouble npy_csinhl(npy_clongdouble z); +npy_clongdouble npy_ctanhl(npy_clongdouble z); + +npy_clongdouble npy_cacosl(npy_clongdouble z); +npy_clongdouble npy_casinl(npy_clongdouble z); +npy_clongdouble npy_catanl(npy_clongdouble z); + +npy_clongdouble npy_cacoshl(npy_clongdouble z); +npy_clongdouble npy_casinhl(npy_clongdouble z); +npy_clongdouble npy_catanhl(npy_clongdouble z); + + +/* + * Functions that set the floating point error + * status word. + */ + +/* + * platform-dependent code translates floating point + * status to an integer sum of these values + */ +#define NPY_FPE_DIVIDEBYZERO 1 +#define NPY_FPE_OVERFLOW 2 +#define NPY_FPE_UNDERFLOW 4 +#define NPY_FPE_INVALID 8 + +int npy_clear_floatstatus_barrier(char*); +int npy_get_floatstatus_barrier(char*); +/* + * use caution with these - clang and gcc8.1 are known to reorder calls + * to this form of the function which can defeat the check. The _barrier + * form of the call is preferable, where the argument is + * (char*)&local_variable + */ +int npy_clear_floatstatus(void); +int npy_get_floatstatus(void); + +void npy_set_floatstatus_divbyzero(void); +void npy_set_floatstatus_overflow(void); +void npy_set_floatstatus_underflow(void); +void npy_set_floatstatus_invalid(void); + +#ifdef __cplusplus +} +#endif + +#if NPY_INLINE_MATH +#include "npy_math_internal.h" +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_MATH_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_no_deprecated_api.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_no_deprecated_api.h new file mode 100644 index 0000000000000000000000000000000000000000..39658c0bd2d61aacd25f75439e81ea16c3e33db8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_no_deprecated_api.h @@ -0,0 +1,20 @@ +/* + * This include file is provided for inclusion in Cython *.pyd files where + * one would like to define the NPY_NO_DEPRECATED_API macro. It can be + * included by + * + * cdef extern from "npy_no_deprecated_api.h": pass + * + */ +#ifndef NPY_NO_DEPRECATED_API + +/* put this check here since there may be multiple includes in C extensions. */ +#if defined(NUMPY_CORE_INCLUDE_NUMPY_NDARRAYTYPES_H_) || \ + defined(NUMPY_CORE_INCLUDE_NUMPY_NPY_DEPRECATED_API_H) || \ + defined(NUMPY_CORE_INCLUDE_NUMPY_OLD_DEFINES_H_) +#error "npy_no_deprecated_api.h" must be first among numpy includes. +#else +#define NPY_NO_DEPRECATED_API NPY_API_VERSION +#endif + +#endif /* NPY_NO_DEPRECATED_API */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_os.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_os.h new file mode 100644 index 0000000000000000000000000000000000000000..0ce5d78b42c0e53c660654e297446d7811901aa2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/npy_os.h @@ -0,0 +1,42 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_OS_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_OS_H_ + +#if defined(linux) || defined(__linux) || defined(__linux__) + #define NPY_OS_LINUX +#elif defined(__FreeBSD__) || defined(__NetBSD__) || \ + defined(__OpenBSD__) || defined(__DragonFly__) + #define NPY_OS_BSD + #ifdef __FreeBSD__ + #define NPY_OS_FREEBSD + #elif defined(__NetBSD__) + #define NPY_OS_NETBSD + #elif defined(__OpenBSD__) + #define NPY_OS_OPENBSD + #elif defined(__DragonFly__) + #define NPY_OS_DRAGONFLY + #endif +#elif defined(sun) || defined(__sun) + #define NPY_OS_SOLARIS +#elif defined(__CYGWIN__) + #define NPY_OS_CYGWIN +/* We are on Windows.*/ +#elif defined(_WIN32) + /* We are using MinGW (64-bit or 32-bit)*/ + #if defined(__MINGW32__) || defined(__MINGW64__) + #define NPY_OS_MINGW + /* Otherwise, if _WIN64 is defined, we are targeting 64-bit Windows*/ + #elif defined(_WIN64) + #define NPY_OS_WIN64 + /* Otherwise assume we are targeting 32-bit Windows*/ + #else + #define NPY_OS_WIN32 + #endif +#elif defined(__APPLE__) + #define NPY_OS_DARWIN +#elif defined(__HAIKU__) + #define NPY_OS_HAIKU +#else + #define NPY_OS_UNKNOWN +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_OS_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/numpyconfig.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/numpyconfig.h new file mode 100644 index 0000000000000000000000000000000000000000..46ecade41ada08fad4c1ae389cc4e9b6e902de24 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/numpyconfig.h @@ -0,0 +1,178 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_NUMPYCONFIG_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_NPY_NUMPYCONFIG_H_ + +#include "_numpyconfig.h" + +/* + * On Mac OS X, because there is only one configuration stage for all the archs + * in universal builds, any macro which depends on the arch needs to be + * hardcoded. + * + * Note that distutils/pip will attempt a universal2 build when Python itself + * is built as universal2, hence this hardcoding is needed even if we do not + * support universal2 wheels anymore (see gh-22796). + * This code block can be removed after we have dropped the setup.py based + * build completely. + */ +#ifdef __APPLE__ + #undef NPY_SIZEOF_LONG + + #ifdef __LP64__ + #define NPY_SIZEOF_LONG 8 + #else + #define NPY_SIZEOF_LONG 4 + #endif + + #undef NPY_SIZEOF_LONGDOUBLE + #undef NPY_SIZEOF_COMPLEX_LONGDOUBLE + #ifdef HAVE_LDOUBLE_IEEE_DOUBLE_LE + #undef HAVE_LDOUBLE_IEEE_DOUBLE_LE + #endif + #ifdef HAVE_LDOUBLE_INTEL_EXTENDED_16_BYTES_LE + #undef HAVE_LDOUBLE_INTEL_EXTENDED_16_BYTES_LE + #endif + + #if defined(__arm64__) + #define NPY_SIZEOF_LONGDOUBLE 8 + #define NPY_SIZEOF_COMPLEX_LONGDOUBLE 16 + #define HAVE_LDOUBLE_IEEE_DOUBLE_LE 1 + #elif defined(__x86_64) + #define NPY_SIZEOF_LONGDOUBLE 16 + #define NPY_SIZEOF_COMPLEX_LONGDOUBLE 32 + #define HAVE_LDOUBLE_INTEL_EXTENDED_16_BYTES_LE 1 + #elif defined (__i386) + #define NPY_SIZEOF_LONGDOUBLE 12 + #define NPY_SIZEOF_COMPLEX_LONGDOUBLE 24 + #elif defined(__ppc__) || defined (__ppc64__) + #define NPY_SIZEOF_LONGDOUBLE 16 + #define NPY_SIZEOF_COMPLEX_LONGDOUBLE 32 + #else + #error "unknown architecture" + #endif +#endif + + +/** + * To help with both NPY_TARGET_VERSION and the NPY_NO_DEPRECATED_API macro, + * we include API version numbers for specific versions of NumPy. + * To exclude all API that was deprecated as of 1.7, add the following before + * #including any NumPy headers: + * #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION + * The same is true for NPY_TARGET_VERSION, although NumPy will default to + * a backwards compatible build anyway. + */ +#define NPY_1_7_API_VERSION 0x00000007 +#define NPY_1_8_API_VERSION 0x00000008 +#define NPY_1_9_API_VERSION 0x00000009 +#define NPY_1_10_API_VERSION 0x0000000a +#define NPY_1_11_API_VERSION 0x0000000a +#define NPY_1_12_API_VERSION 0x0000000a +#define NPY_1_13_API_VERSION 0x0000000b +#define NPY_1_14_API_VERSION 0x0000000c +#define NPY_1_15_API_VERSION 0x0000000c +#define NPY_1_16_API_VERSION 0x0000000d +#define NPY_1_17_API_VERSION 0x0000000d +#define NPY_1_18_API_VERSION 0x0000000d +#define NPY_1_19_API_VERSION 0x0000000d +#define NPY_1_20_API_VERSION 0x0000000e +#define NPY_1_21_API_VERSION 0x0000000e +#define NPY_1_22_API_VERSION 0x0000000f +#define NPY_1_23_API_VERSION 0x00000010 +#define NPY_1_24_API_VERSION 0x00000010 +#define NPY_1_25_API_VERSION 0x00000011 +#define NPY_2_0_API_VERSION 0x00000012 +#define NPY_2_1_API_VERSION 0x00000013 + + +/* + * Binary compatibility version number. This number is increased + * whenever the C-API is changed such that binary compatibility is + * broken, i.e. whenever a recompile of extension modules is needed. + */ +#define NPY_VERSION NPY_ABI_VERSION + +/* + * Minor API version we are compiling to be compatible with. The version + * Number is always increased when the API changes via: `NPY_API_VERSION` + * (and should maybe just track the NumPy version). + * + * If we have an internal build, we always target the current version of + * course. + * + * For downstream users, we default to an older version to provide them with + * maximum compatibility by default. Downstream can choose to extend that + * default, or narrow it down if they wish to use newer API. If you adjust + * this, consider the Python version support (example for 1.25.x): + * + * NumPy 1.25.x supports Python: 3.9 3.10 3.11 (3.12) + * NumPy 1.19.x supports Python: 3.6 3.7 3.8 3.9 + * NumPy 1.17.x supports Python: 3.5 3.6 3.7 3.8 + * NumPy 1.15.x supports Python: ... 3.6 3.7 + * + * Users of the stable ABI may wish to target the last Python that is not + * end of life. This would be 3.8 at NumPy 1.25 release time. + * 1.17 as default was the choice of oldest-support-numpy at the time and + * has in practice no limit (compared to 1.19). Even earlier becomes legacy. + */ +#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD + /* NumPy internal build, always use current version. */ + #define NPY_FEATURE_VERSION NPY_API_VERSION +#elif defined(NPY_TARGET_VERSION) && NPY_TARGET_VERSION + /* user provided a target version, use it */ + #define NPY_FEATURE_VERSION NPY_TARGET_VERSION +#else + /* Use the default (increase when dropping Python 3.10 support) */ + #define NPY_FEATURE_VERSION NPY_1_21_API_VERSION +#endif + +/* Sanity check the (requested) feature version */ +#if NPY_FEATURE_VERSION > NPY_API_VERSION + #error "NPY_TARGET_VERSION higher than NumPy headers!" +#elif NPY_FEATURE_VERSION < NPY_1_15_API_VERSION + /* No support for irrelevant old targets, no need for error, but warn. */ + #ifndef _MSC_VER + #warning "Requested NumPy target lower than supported NumPy 1.15." + #else + #define _WARN___STR2__(x) #x + #define _WARN___STR1__(x) _WARN___STR2__(x) + #define _WARN___LOC__ __FILE__ "(" _WARN___STR1__(__LINE__) ") : Warning Msg: " + #pragma message(_WARN___LOC__"Requested NumPy target lower than supported NumPy 1.15.") + #endif +#endif + +/* + * We define a human readable translation to the Python version of NumPy + * for error messages (and also to allow grepping the binaries for conda). + */ +#if NPY_FEATURE_VERSION == NPY_1_7_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.7" +#elif NPY_FEATURE_VERSION == NPY_1_8_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.8" +#elif NPY_FEATURE_VERSION == NPY_1_9_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.9" +#elif NPY_FEATURE_VERSION == NPY_1_10_API_VERSION /* also 1.11, 1.12 */ + #define NPY_FEATURE_VERSION_STRING "1.10" +#elif NPY_FEATURE_VERSION == NPY_1_13_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.13" +#elif NPY_FEATURE_VERSION == NPY_1_14_API_VERSION /* also 1.15 */ + #define NPY_FEATURE_VERSION_STRING "1.14" +#elif NPY_FEATURE_VERSION == NPY_1_16_API_VERSION /* also 1.17, 1.18, 1.19 */ + #define NPY_FEATURE_VERSION_STRING "1.16" +#elif NPY_FEATURE_VERSION == NPY_1_20_API_VERSION /* also 1.21 */ + #define NPY_FEATURE_VERSION_STRING "1.20" +#elif NPY_FEATURE_VERSION == NPY_1_22_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.22" +#elif NPY_FEATURE_VERSION == NPY_1_23_API_VERSION /* also 1.24 */ + #define NPY_FEATURE_VERSION_STRING "1.23" +#elif NPY_FEATURE_VERSION == NPY_1_25_API_VERSION + #define NPY_FEATURE_VERSION_STRING "1.25" +#elif NPY_FEATURE_VERSION == NPY_2_0_API_VERSION + #define NPY_FEATURE_VERSION_STRING "2.0" +#elif NPY_FEATURE_VERSION == NPY_2_1_API_VERSION + #define NPY_FEATURE_VERSION_STRING "2.1" +#else + #error "Missing version string define for new NumPy version." +#endif + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_NPY_NUMPYCONFIG_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/LICENSE.txt b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..d72a7c388d406191f2b3113efa5f89916a39d9b2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/LICENSE.txt @@ -0,0 +1,21 @@ + zlib License + ------------ + + Copyright (C) 2010 - 2019 ridiculous_fish, + Copyright (C) 2016 - 2019 Kim Walisch, + + This software is provided 'as-is', without any express or implied + warranty. In no event will the authors be held liable for any damages + arising from the use of this software. + + Permission is granted to anyone to use this software for any purpose, + including commercial applications, and to alter it and redistribute it + freely, subject to the following restrictions: + + 1. The origin of this software must not be misrepresented; you must not + claim that you wrote the original software. If you use this software + in a product, an acknowledgment in the product documentation would be + appreciated but is not required. + 2. Altered source versions must be plainly marked as such, and must not be + misrepresented as being the original software. + 3. This notice may not be removed or altered from any source distribution. diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/bitgen.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/bitgen.h new file mode 100644 index 0000000000000000000000000000000000000000..162dd5c5753079eb1d76efa7fc8a3847c2ad6602 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/bitgen.h @@ -0,0 +1,20 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_RANDOM_BITGEN_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_RANDOM_BITGEN_H_ + +#pragma once +#include +#include +#include + +/* Must match the declaration in numpy/random/.pxd */ + +typedef struct bitgen { + void *state; + uint64_t (*next_uint64)(void *st); + uint32_t (*next_uint32)(void *st); + double (*next_double)(void *st); + uint64_t (*next_raw)(void *st); +} bitgen_t; + + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_RANDOM_BITGEN_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/distributions.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/distributions.h new file mode 100644 index 0000000000000000000000000000000000000000..e7fa4bd00d43430eb1da23bd577688d8733bb6e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/distributions.h @@ -0,0 +1,209 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_RANDOM_DISTRIBUTIONS_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_RANDOM_DISTRIBUTIONS_H_ + +#ifdef __cplusplus +extern "C" { +#endif + +#include +#include "numpy/npy_common.h" +#include +#include +#include + +#include "numpy/npy_math.h" +#include "numpy/random/bitgen.h" + +/* + * RAND_INT_TYPE is used to share integer generators with RandomState which + * used long in place of int64_t. If changing a distribution that uses + * RAND_INT_TYPE, then the original unmodified copy must be retained for + * use in RandomState by copying to the legacy distributions source file. + */ +#ifdef NP_RANDOM_LEGACY +#define RAND_INT_TYPE long +#define RAND_INT_MAX LONG_MAX +#else +#define RAND_INT_TYPE int64_t +#define RAND_INT_MAX INT64_MAX +#endif + +#ifdef _MSC_VER +#define DECLDIR __declspec(dllexport) +#else +#define DECLDIR extern +#endif + +#ifndef MIN +#define MIN(x, y) (((x) < (y)) ? x : y) +#define MAX(x, y) (((x) > (y)) ? x : y) +#endif + +#ifndef M_PI +#define M_PI 3.14159265358979323846264338328 +#endif + +typedef struct s_binomial_t { + int has_binomial; /* !=0: following parameters initialized for binomial */ + double psave; + RAND_INT_TYPE nsave; + double r; + double q; + double fm; + RAND_INT_TYPE m; + double p1; + double xm; + double xl; + double xr; + double c; + double laml; + double lamr; + double p2; + double p3; + double p4; +} binomial_t; + +DECLDIR float random_standard_uniform_f(bitgen_t *bitgen_state); +DECLDIR double random_standard_uniform(bitgen_t *bitgen_state); +DECLDIR void random_standard_uniform_fill(bitgen_t *, npy_intp, double *); +DECLDIR void random_standard_uniform_fill_f(bitgen_t *, npy_intp, float *); + +DECLDIR int64_t random_positive_int64(bitgen_t *bitgen_state); +DECLDIR int32_t random_positive_int32(bitgen_t *bitgen_state); +DECLDIR int64_t random_positive_int(bitgen_t *bitgen_state); +DECLDIR uint64_t random_uint(bitgen_t *bitgen_state); + +DECLDIR double random_standard_exponential(bitgen_t *bitgen_state); +DECLDIR float random_standard_exponential_f(bitgen_t *bitgen_state); +DECLDIR void random_standard_exponential_fill(bitgen_t *, npy_intp, double *); +DECLDIR void random_standard_exponential_fill_f(bitgen_t *, npy_intp, float *); +DECLDIR void random_standard_exponential_inv_fill(bitgen_t *, npy_intp, double *); +DECLDIR void random_standard_exponential_inv_fill_f(bitgen_t *, npy_intp, float *); + +DECLDIR double random_standard_normal(bitgen_t *bitgen_state); +DECLDIR float random_standard_normal_f(bitgen_t *bitgen_state); +DECLDIR void random_standard_normal_fill(bitgen_t *, npy_intp, double *); +DECLDIR void random_standard_normal_fill_f(bitgen_t *, npy_intp, float *); +DECLDIR double random_standard_gamma(bitgen_t *bitgen_state, double shape); +DECLDIR float random_standard_gamma_f(bitgen_t *bitgen_state, float shape); + +DECLDIR double random_normal(bitgen_t *bitgen_state, double loc, double scale); + +DECLDIR double random_gamma(bitgen_t *bitgen_state, double shape, double scale); +DECLDIR float random_gamma_f(bitgen_t *bitgen_state, float shape, float scale); + +DECLDIR double random_exponential(bitgen_t *bitgen_state, double scale); +DECLDIR double random_uniform(bitgen_t *bitgen_state, double lower, double range); +DECLDIR double random_beta(bitgen_t *bitgen_state, double a, double b); +DECLDIR double random_chisquare(bitgen_t *bitgen_state, double df); +DECLDIR double random_f(bitgen_t *bitgen_state, double dfnum, double dfden); +DECLDIR double random_standard_cauchy(bitgen_t *bitgen_state); +DECLDIR double random_pareto(bitgen_t *bitgen_state, double a); +DECLDIR double random_weibull(bitgen_t *bitgen_state, double a); +DECLDIR double random_power(bitgen_t *bitgen_state, double a); +DECLDIR double random_laplace(bitgen_t *bitgen_state, double loc, double scale); +DECLDIR double random_gumbel(bitgen_t *bitgen_state, double loc, double scale); +DECLDIR double random_logistic(bitgen_t *bitgen_state, double loc, double scale); +DECLDIR double random_lognormal(bitgen_t *bitgen_state, double mean, double sigma); +DECLDIR double random_rayleigh(bitgen_t *bitgen_state, double mode); +DECLDIR double random_standard_t(bitgen_t *bitgen_state, double df); +DECLDIR double random_noncentral_chisquare(bitgen_t *bitgen_state, double df, + double nonc); +DECLDIR double random_noncentral_f(bitgen_t *bitgen_state, double dfnum, + double dfden, double nonc); +DECLDIR double random_wald(bitgen_t *bitgen_state, double mean, double scale); +DECLDIR double random_vonmises(bitgen_t *bitgen_state, double mu, double kappa); +DECLDIR double random_triangular(bitgen_t *bitgen_state, double left, double mode, + double right); + +DECLDIR RAND_INT_TYPE random_poisson(bitgen_t *bitgen_state, double lam); +DECLDIR RAND_INT_TYPE random_negative_binomial(bitgen_t *bitgen_state, double n, + double p); + +DECLDIR int64_t random_binomial(bitgen_t *bitgen_state, double p, + int64_t n, binomial_t *binomial); + +DECLDIR int64_t random_logseries(bitgen_t *bitgen_state, double p); +DECLDIR int64_t random_geometric(bitgen_t *bitgen_state, double p); +DECLDIR RAND_INT_TYPE random_geometric_search(bitgen_t *bitgen_state, double p); +DECLDIR RAND_INT_TYPE random_zipf(bitgen_t *bitgen_state, double a); +DECLDIR int64_t random_hypergeometric(bitgen_t *bitgen_state, + int64_t good, int64_t bad, int64_t sample); +DECLDIR uint64_t random_interval(bitgen_t *bitgen_state, uint64_t max); + +/* Generate random uint64 numbers in closed interval [off, off + rng]. */ +DECLDIR uint64_t random_bounded_uint64(bitgen_t *bitgen_state, uint64_t off, + uint64_t rng, uint64_t mask, + bool use_masked); + +/* Generate random uint32 numbers in closed interval [off, off + rng]. */ +DECLDIR uint32_t random_buffered_bounded_uint32(bitgen_t *bitgen_state, + uint32_t off, uint32_t rng, + uint32_t mask, bool use_masked, + int *bcnt, uint32_t *buf); +DECLDIR uint16_t random_buffered_bounded_uint16(bitgen_t *bitgen_state, + uint16_t off, uint16_t rng, + uint16_t mask, bool use_masked, + int *bcnt, uint32_t *buf); +DECLDIR uint8_t random_buffered_bounded_uint8(bitgen_t *bitgen_state, uint8_t off, + uint8_t rng, uint8_t mask, + bool use_masked, int *bcnt, + uint32_t *buf); +DECLDIR npy_bool random_buffered_bounded_bool(bitgen_t *bitgen_state, npy_bool off, + npy_bool rng, npy_bool mask, + bool use_masked, int *bcnt, + uint32_t *buf); + +DECLDIR void random_bounded_uint64_fill(bitgen_t *bitgen_state, uint64_t off, + uint64_t rng, npy_intp cnt, + bool use_masked, uint64_t *out); +DECLDIR void random_bounded_uint32_fill(bitgen_t *bitgen_state, uint32_t off, + uint32_t rng, npy_intp cnt, + bool use_masked, uint32_t *out); +DECLDIR void random_bounded_uint16_fill(bitgen_t *bitgen_state, uint16_t off, + uint16_t rng, npy_intp cnt, + bool use_masked, uint16_t *out); +DECLDIR void random_bounded_uint8_fill(bitgen_t *bitgen_state, uint8_t off, + uint8_t rng, npy_intp cnt, + bool use_masked, uint8_t *out); +DECLDIR void random_bounded_bool_fill(bitgen_t *bitgen_state, npy_bool off, + npy_bool rng, npy_intp cnt, + bool use_masked, npy_bool *out); + +DECLDIR void random_multinomial(bitgen_t *bitgen_state, RAND_INT_TYPE n, RAND_INT_TYPE *mnix, + double *pix, npy_intp d, binomial_t *binomial); + +/* multivariate hypergeometric, "count" method */ +DECLDIR int random_multivariate_hypergeometric_count(bitgen_t *bitgen_state, + int64_t total, + size_t num_colors, int64_t *colors, + int64_t nsample, + size_t num_variates, int64_t *variates); + +/* multivariate hypergeometric, "marginals" method */ +DECLDIR void random_multivariate_hypergeometric_marginals(bitgen_t *bitgen_state, + int64_t total, + size_t num_colors, int64_t *colors, + int64_t nsample, + size_t num_variates, int64_t *variates); + +/* Common to legacy-distributions.c and distributions.c but not exported */ + +RAND_INT_TYPE random_binomial_btpe(bitgen_t *bitgen_state, + RAND_INT_TYPE n, + double p, + binomial_t *binomial); +RAND_INT_TYPE random_binomial_inversion(bitgen_t *bitgen_state, + RAND_INT_TYPE n, + double p, + binomial_t *binomial); +double random_loggam(double x); +static inline double next_double(bitgen_t *bitgen_state) { + return bitgen_state->next_double(bitgen_state->state); +} + +#ifdef __cplusplus +} +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_RANDOM_DISTRIBUTIONS_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/libdivide.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/libdivide.h new file mode 100644 index 0000000000000000000000000000000000000000..f4eb8039b50c21c0977a38ee5f47c8f89307d482 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/random/libdivide.h @@ -0,0 +1,2079 @@ +// libdivide.h - Optimized integer division +// https://libdivide.com +// +// Copyright (C) 2010 - 2019 ridiculous_fish, +// Copyright (C) 2016 - 2019 Kim Walisch, +// +// libdivide is dual-licensed under the Boost or zlib licenses. +// You may use libdivide under the terms of either of these. +// See LICENSE.txt for more details. + +#ifndef NUMPY_CORE_INCLUDE_NUMPY_LIBDIVIDE_LIBDIVIDE_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_LIBDIVIDE_LIBDIVIDE_H_ + +#define LIBDIVIDE_VERSION "3.0" +#define LIBDIVIDE_VERSION_MAJOR 3 +#define LIBDIVIDE_VERSION_MINOR 0 + +#include + +#if defined(__cplusplus) + #include + #include + #include +#else + #include + #include +#endif + +#if defined(LIBDIVIDE_AVX512) + #include +#elif defined(LIBDIVIDE_AVX2) + #include +#elif defined(LIBDIVIDE_SSE2) + #include +#endif + +#if defined(_MSC_VER) + #include + // disable warning C4146: unary minus operator applied + // to unsigned type, result still unsigned + #pragma warning(disable: 4146) + #define LIBDIVIDE_VC +#endif + +#if !defined(__has_builtin) + #define __has_builtin(x) 0 +#endif + +#if defined(__SIZEOF_INT128__) + #define HAS_INT128_T + // clang-cl on Windows does not yet support 128-bit division + #if !(defined(__clang__) && defined(LIBDIVIDE_VC)) + #define HAS_INT128_DIV + #endif +#endif + +#if defined(__x86_64__) || defined(_M_X64) + #define LIBDIVIDE_X86_64 +#endif + +#if defined(__i386__) + #define LIBDIVIDE_i386 +#endif + +#if defined(__GNUC__) || defined(__clang__) + #define LIBDIVIDE_GCC_STYLE_ASM +#endif + +#if defined(__cplusplus) || defined(LIBDIVIDE_VC) + #define LIBDIVIDE_FUNCTION __FUNCTION__ +#else + #define LIBDIVIDE_FUNCTION __func__ +#endif + +#define LIBDIVIDE_ERROR(msg) \ + do { \ + fprintf(stderr, "libdivide.h:%d: %s(): Error: %s\n", \ + __LINE__, LIBDIVIDE_FUNCTION, msg); \ + abort(); \ + } while (0) + +#if defined(LIBDIVIDE_ASSERTIONS_ON) + #define LIBDIVIDE_ASSERT(x) \ + do { \ + if (!(x)) { \ + fprintf(stderr, "libdivide.h:%d: %s(): Assertion failed: %s\n", \ + __LINE__, LIBDIVIDE_FUNCTION, #x); \ + abort(); \ + } \ + } while (0) +#else + #define LIBDIVIDE_ASSERT(x) +#endif + +#ifdef __cplusplus +namespace libdivide { +#endif + +// pack divider structs to prevent compilers from padding. +// This reduces memory usage by up to 43% when using a large +// array of libdivide dividers and improves performance +// by up to 10% because of reduced memory bandwidth. +#pragma pack(push, 1) + +struct libdivide_u32_t { + uint32_t magic; + uint8_t more; +}; + +struct libdivide_s32_t { + int32_t magic; + uint8_t more; +}; + +struct libdivide_u64_t { + uint64_t magic; + uint8_t more; +}; + +struct libdivide_s64_t { + int64_t magic; + uint8_t more; +}; + +struct libdivide_u32_branchfree_t { + uint32_t magic; + uint8_t more; +}; + +struct libdivide_s32_branchfree_t { + int32_t magic; + uint8_t more; +}; + +struct libdivide_u64_branchfree_t { + uint64_t magic; + uint8_t more; +}; + +struct libdivide_s64_branchfree_t { + int64_t magic; + uint8_t more; +}; + +#pragma pack(pop) + +// Explanation of the "more" field: +// +// * Bits 0-5 is the shift value (for shift path or mult path). +// * Bit 6 is the add indicator for mult path. +// * Bit 7 is set if the divisor is negative. We use bit 7 as the negative +// divisor indicator so that we can efficiently use sign extension to +// create a bitmask with all bits set to 1 (if the divisor is negative) +// or 0 (if the divisor is positive). +// +// u32: [0-4] shift value +// [5] ignored +// [6] add indicator +// magic number of 0 indicates shift path +// +// s32: [0-4] shift value +// [5] ignored +// [6] add indicator +// [7] indicates negative divisor +// magic number of 0 indicates shift path +// +// u64: [0-5] shift value +// [6] add indicator +// magic number of 0 indicates shift path +// +// s64: [0-5] shift value +// [6] add indicator +// [7] indicates negative divisor +// magic number of 0 indicates shift path +// +// In s32 and s64 branchfree modes, the magic number is negated according to +// whether the divisor is negated. In branchfree strategy, it is not negated. + +enum { + LIBDIVIDE_32_SHIFT_MASK = 0x1F, + LIBDIVIDE_64_SHIFT_MASK = 0x3F, + LIBDIVIDE_ADD_MARKER = 0x40, + LIBDIVIDE_NEGATIVE_DIVISOR = 0x80 +}; + +static inline struct libdivide_s32_t libdivide_s32_gen(int32_t d); +static inline struct libdivide_u32_t libdivide_u32_gen(uint32_t d); +static inline struct libdivide_s64_t libdivide_s64_gen(int64_t d); +static inline struct libdivide_u64_t libdivide_u64_gen(uint64_t d); + +static inline struct libdivide_s32_branchfree_t libdivide_s32_branchfree_gen(int32_t d); +static inline struct libdivide_u32_branchfree_t libdivide_u32_branchfree_gen(uint32_t d); +static inline struct libdivide_s64_branchfree_t libdivide_s64_branchfree_gen(int64_t d); +static inline struct libdivide_u64_branchfree_t libdivide_u64_branchfree_gen(uint64_t d); + +static inline int32_t libdivide_s32_do(int32_t numer, const struct libdivide_s32_t *denom); +static inline uint32_t libdivide_u32_do(uint32_t numer, const struct libdivide_u32_t *denom); +static inline int64_t libdivide_s64_do(int64_t numer, const struct libdivide_s64_t *denom); +static inline uint64_t libdivide_u64_do(uint64_t numer, const struct libdivide_u64_t *denom); + +static inline int32_t libdivide_s32_branchfree_do(int32_t numer, const struct libdivide_s32_branchfree_t *denom); +static inline uint32_t libdivide_u32_branchfree_do(uint32_t numer, const struct libdivide_u32_branchfree_t *denom); +static inline int64_t libdivide_s64_branchfree_do(int64_t numer, const struct libdivide_s64_branchfree_t *denom); +static inline uint64_t libdivide_u64_branchfree_do(uint64_t numer, const struct libdivide_u64_branchfree_t *denom); + +static inline int32_t libdivide_s32_recover(const struct libdivide_s32_t *denom); +static inline uint32_t libdivide_u32_recover(const struct libdivide_u32_t *denom); +static inline int64_t libdivide_s64_recover(const struct libdivide_s64_t *denom); +static inline uint64_t libdivide_u64_recover(const struct libdivide_u64_t *denom); + +static inline int32_t libdivide_s32_branchfree_recover(const struct libdivide_s32_branchfree_t *denom); +static inline uint32_t libdivide_u32_branchfree_recover(const struct libdivide_u32_branchfree_t *denom); +static inline int64_t libdivide_s64_branchfree_recover(const struct libdivide_s64_branchfree_t *denom); +static inline uint64_t libdivide_u64_branchfree_recover(const struct libdivide_u64_branchfree_t *denom); + +//////// Internal Utility Functions + +static inline uint32_t libdivide_mullhi_u32(uint32_t x, uint32_t y) { + uint64_t xl = x, yl = y; + uint64_t rl = xl * yl; + return (uint32_t)(rl >> 32); +} + +static inline int32_t libdivide_mullhi_s32(int32_t x, int32_t y) { + int64_t xl = x, yl = y; + int64_t rl = xl * yl; + // needs to be arithmetic shift + return (int32_t)(rl >> 32); +} + +static inline uint64_t libdivide_mullhi_u64(uint64_t x, uint64_t y) { +#if defined(LIBDIVIDE_VC) && \ + defined(LIBDIVIDE_X86_64) + return __umulh(x, y); +#elif defined(HAS_INT128_T) + __uint128_t xl = x, yl = y; + __uint128_t rl = xl * yl; + return (uint64_t)(rl >> 64); +#else + // full 128 bits are x0 * y0 + (x0 * y1 << 32) + (x1 * y0 << 32) + (x1 * y1 << 64) + uint32_t mask = 0xFFFFFFFF; + uint32_t x0 = (uint32_t)(x & mask); + uint32_t x1 = (uint32_t)(x >> 32); + uint32_t y0 = (uint32_t)(y & mask); + uint32_t y1 = (uint32_t)(y >> 32); + uint32_t x0y0_hi = libdivide_mullhi_u32(x0, y0); + uint64_t x0y1 = x0 * (uint64_t)y1; + uint64_t x1y0 = x1 * (uint64_t)y0; + uint64_t x1y1 = x1 * (uint64_t)y1; + uint64_t temp = x1y0 + x0y0_hi; + uint64_t temp_lo = temp & mask; + uint64_t temp_hi = temp >> 32; + + return x1y1 + temp_hi + ((temp_lo + x0y1) >> 32); +#endif +} + +static inline int64_t libdivide_mullhi_s64(int64_t x, int64_t y) { +#if defined(LIBDIVIDE_VC) && \ + defined(LIBDIVIDE_X86_64) + return __mulh(x, y); +#elif defined(HAS_INT128_T) + __int128_t xl = x, yl = y; + __int128_t rl = xl * yl; + return (int64_t)(rl >> 64); +#else + // full 128 bits are x0 * y0 + (x0 * y1 << 32) + (x1 * y0 << 32) + (x1 * y1 << 64) + uint32_t mask = 0xFFFFFFFF; + uint32_t x0 = (uint32_t)(x & mask); + uint32_t y0 = (uint32_t)(y & mask); + int32_t x1 = (int32_t)(x >> 32); + int32_t y1 = (int32_t)(y >> 32); + uint32_t x0y0_hi = libdivide_mullhi_u32(x0, y0); + int64_t t = x1 * (int64_t)y0 + x0y0_hi; + int64_t w1 = x0 * (int64_t)y1 + (t & mask); + + return x1 * (int64_t)y1 + (t >> 32) + (w1 >> 32); +#endif +} + +static inline int32_t libdivide_count_leading_zeros32(uint32_t val) { +#if defined(__GNUC__) || \ + __has_builtin(__builtin_clz) + // Fast way to count leading zeros + return __builtin_clz(val); +#elif defined(LIBDIVIDE_VC) + unsigned long result; + if (_BitScanReverse(&result, val)) { + return 31 - result; + } + return 0; +#else + if (val == 0) + return 32; + int32_t result = 8; + uint32_t hi = 0xFFU << 24; + while ((val & hi) == 0) { + hi >>= 8; + result += 8; + } + while (val & hi) { + result -= 1; + hi <<= 1; + } + return result; +#endif +} + +static inline int32_t libdivide_count_leading_zeros64(uint64_t val) { +#if defined(__GNUC__) || \ + __has_builtin(__builtin_clzll) + // Fast way to count leading zeros + return __builtin_clzll(val); +#elif defined(LIBDIVIDE_VC) && defined(_WIN64) + unsigned long result; + if (_BitScanReverse64(&result, val)) { + return 63 - result; + } + return 0; +#else + uint32_t hi = val >> 32; + uint32_t lo = val & 0xFFFFFFFF; + if (hi != 0) return libdivide_count_leading_zeros32(hi); + return 32 + libdivide_count_leading_zeros32(lo); +#endif +} + +// libdivide_64_div_32_to_32: divides a 64-bit uint {u1, u0} by a 32-bit +// uint {v}. The result must fit in 32 bits. +// Returns the quotient directly and the remainder in *r +static inline uint32_t libdivide_64_div_32_to_32(uint32_t u1, uint32_t u0, uint32_t v, uint32_t *r) { +#if (defined(LIBDIVIDE_i386) || defined(LIBDIVIDE_X86_64)) && \ + defined(LIBDIVIDE_GCC_STYLE_ASM) + uint32_t result; + __asm__("divl %[v]" + : "=a"(result), "=d"(*r) + : [v] "r"(v), "a"(u0), "d"(u1) + ); + return result; +#else + uint64_t n = ((uint64_t)u1 << 32) | u0; + uint32_t result = (uint32_t)(n / v); + *r = (uint32_t)(n - result * (uint64_t)v); + return result; +#endif +} + +// libdivide_128_div_64_to_64: divides a 128-bit uint {u1, u0} by a 64-bit +// uint {v}. The result must fit in 64 bits. +// Returns the quotient directly and the remainder in *r +static uint64_t libdivide_128_div_64_to_64(uint64_t u1, uint64_t u0, uint64_t v, uint64_t *r) { +#if defined(LIBDIVIDE_X86_64) && \ + defined(LIBDIVIDE_GCC_STYLE_ASM) + uint64_t result; + __asm__("divq %[v]" + : "=a"(result), "=d"(*r) + : [v] "r"(v), "a"(u0), "d"(u1) + ); + return result; +#elif defined(HAS_INT128_T) && \ + defined(HAS_INT128_DIV) + __uint128_t n = ((__uint128_t)u1 << 64) | u0; + uint64_t result = (uint64_t)(n / v); + *r = (uint64_t)(n - result * (__uint128_t)v); + return result; +#else + // Code taken from Hacker's Delight: + // http://www.hackersdelight.org/HDcode/divlu.c. + // License permits inclusion here per: + // http://www.hackersdelight.org/permissions.htm + + const uint64_t b = (1ULL << 32); // Number base (32 bits) + uint64_t un1, un0; // Norm. dividend LSD's + uint64_t vn1, vn0; // Norm. divisor digits + uint64_t q1, q0; // Quotient digits + uint64_t un64, un21, un10; // Dividend digit pairs + uint64_t rhat; // A remainder + int32_t s; // Shift amount for norm + + // If overflow, set rem. to an impossible value, + // and return the largest possible quotient + if (u1 >= v) { + *r = (uint64_t) -1; + return (uint64_t) -1; + } + + // count leading zeros + s = libdivide_count_leading_zeros64(v); + if (s > 0) { + // Normalize divisor + v = v << s; + un64 = (u1 << s) | (u0 >> (64 - s)); + un10 = u0 << s; // Shift dividend left + } else { + // Avoid undefined behavior of (u0 >> 64). + // The behavior is undefined if the right operand is + // negative, or greater than or equal to the length + // in bits of the promoted left operand. + un64 = u1; + un10 = u0; + } + + // Break divisor up into two 32-bit digits + vn1 = v >> 32; + vn0 = v & 0xFFFFFFFF; + + // Break right half of dividend into two digits + un1 = un10 >> 32; + un0 = un10 & 0xFFFFFFFF; + + // Compute the first quotient digit, q1 + q1 = un64 / vn1; + rhat = un64 - q1 * vn1; + + while (q1 >= b || q1 * vn0 > b * rhat + un1) { + q1 = q1 - 1; + rhat = rhat + vn1; + if (rhat >= b) + break; + } + + // Multiply and subtract + un21 = un64 * b + un1 - q1 * v; + + // Compute the second quotient digit + q0 = un21 / vn1; + rhat = un21 - q0 * vn1; + + while (q0 >= b || q0 * vn0 > b * rhat + un0) { + q0 = q0 - 1; + rhat = rhat + vn1; + if (rhat >= b) + break; + } + + *r = (un21 * b + un0 - q0 * v) >> s; + return q1 * b + q0; +#endif +} + +// Bitshift a u128 in place, left (signed_shift > 0) or right (signed_shift < 0) +static inline void libdivide_u128_shift(uint64_t *u1, uint64_t *u0, int32_t signed_shift) { + if (signed_shift > 0) { + uint32_t shift = signed_shift; + *u1 <<= shift; + *u1 |= *u0 >> (64 - shift); + *u0 <<= shift; + } + else if (signed_shift < 0) { + uint32_t shift = -signed_shift; + *u0 >>= shift; + *u0 |= *u1 << (64 - shift); + *u1 >>= shift; + } +} + +// Computes a 128 / 128 -> 64 bit division, with a 128 bit remainder. +static uint64_t libdivide_128_div_128_to_64(uint64_t u_hi, uint64_t u_lo, uint64_t v_hi, uint64_t v_lo, uint64_t *r_hi, uint64_t *r_lo) { +#if defined(HAS_INT128_T) && \ + defined(HAS_INT128_DIV) + __uint128_t ufull = u_hi; + __uint128_t vfull = v_hi; + ufull = (ufull << 64) | u_lo; + vfull = (vfull << 64) | v_lo; + uint64_t res = (uint64_t)(ufull / vfull); + __uint128_t remainder = ufull - (vfull * res); + *r_lo = (uint64_t)remainder; + *r_hi = (uint64_t)(remainder >> 64); + return res; +#else + // Adapted from "Unsigned Doubleword Division" in Hacker's Delight + // We want to compute u / v + typedef struct { uint64_t hi; uint64_t lo; } u128_t; + u128_t u = {u_hi, u_lo}; + u128_t v = {v_hi, v_lo}; + + if (v.hi == 0) { + // divisor v is a 64 bit value, so we just need one 128/64 division + // Note that we are simpler than Hacker's Delight here, because we know + // the quotient fits in 64 bits whereas Hacker's Delight demands a full + // 128 bit quotient + *r_hi = 0; + return libdivide_128_div_64_to_64(u.hi, u.lo, v.lo, r_lo); + } + // Here v >= 2**64 + // We know that v.hi != 0, so count leading zeros is OK + // We have 0 <= n <= 63 + uint32_t n = libdivide_count_leading_zeros64(v.hi); + + // Normalize the divisor so its MSB is 1 + u128_t v1t = v; + libdivide_u128_shift(&v1t.hi, &v1t.lo, n); + uint64_t v1 = v1t.hi; // i.e. v1 = v1t >> 64 + + // To ensure no overflow + u128_t u1 = u; + libdivide_u128_shift(&u1.hi, &u1.lo, -1); + + // Get quotient from divide unsigned insn. + uint64_t rem_ignored; + uint64_t q1 = libdivide_128_div_64_to_64(u1.hi, u1.lo, v1, &rem_ignored); + + // Undo normalization and division of u by 2. + u128_t q0 = {0, q1}; + libdivide_u128_shift(&q0.hi, &q0.lo, n); + libdivide_u128_shift(&q0.hi, &q0.lo, -63); + + // Make q0 correct or too small by 1 + // Equivalent to `if (q0 != 0) q0 = q0 - 1;` + if (q0.hi != 0 || q0.lo != 0) { + q0.hi -= (q0.lo == 0); // borrow + q0.lo -= 1; + } + + // Now q0 is correct. + // Compute q0 * v as q0v + // = (q0.hi << 64 + q0.lo) * (v.hi << 64 + v.lo) + // = (q0.hi * v.hi << 128) + (q0.hi * v.lo << 64) + + // (q0.lo * v.hi << 64) + q0.lo * v.lo) + // Each term is 128 bit + // High half of full product (upper 128 bits!) are dropped + u128_t q0v = {0, 0}; + q0v.hi = q0.hi*v.lo + q0.lo*v.hi + libdivide_mullhi_u64(q0.lo, v.lo); + q0v.lo = q0.lo*v.lo; + + // Compute u - q0v as u_q0v + // This is the remainder + u128_t u_q0v = u; + u_q0v.hi -= q0v.hi + (u.lo < q0v.lo); // second term is borrow + u_q0v.lo -= q0v.lo; + + // Check if u_q0v >= v + // This checks if our remainder is larger than the divisor + if ((u_q0v.hi > v.hi) || + (u_q0v.hi == v.hi && u_q0v.lo >= v.lo)) { + // Increment q0 + q0.lo += 1; + q0.hi += (q0.lo == 0); // carry + + // Subtract v from remainder + u_q0v.hi -= v.hi + (u_q0v.lo < v.lo); + u_q0v.lo -= v.lo; + } + + *r_hi = u_q0v.hi; + *r_lo = u_q0v.lo; + + LIBDIVIDE_ASSERT(q0.hi == 0); + return q0.lo; +#endif +} + +////////// UINT32 + +static inline struct libdivide_u32_t libdivide_internal_u32_gen(uint32_t d, int branchfree) { + if (d == 0) { + LIBDIVIDE_ERROR("divider must be != 0"); + } + + struct libdivide_u32_t result; + uint32_t floor_log_2_d = 31 - libdivide_count_leading_zeros32(d); + + // Power of 2 + if ((d & (d - 1)) == 0) { + // We need to subtract 1 from the shift value in case of an unsigned + // branchfree divider because there is a hardcoded right shift by 1 + // in its division algorithm. Because of this we also need to add back + // 1 in its recovery algorithm. + result.magic = 0; + result.more = (uint8_t)(floor_log_2_d - (branchfree != 0)); + } else { + uint8_t more; + uint32_t rem, proposed_m; + proposed_m = libdivide_64_div_32_to_32(1U << floor_log_2_d, 0, d, &rem); + + LIBDIVIDE_ASSERT(rem > 0 && rem < d); + const uint32_t e = d - rem; + + // This power works if e < 2**floor_log_2_d. + if (!branchfree && (e < (1U << floor_log_2_d))) { + // This power works + more = floor_log_2_d; + } else { + // We have to use the general 33-bit algorithm. We need to compute + // (2**power) / d. However, we already have (2**(power-1))/d and + // its remainder. By doubling both, and then correcting the + // remainder, we can compute the larger division. + // don't care about overflow here - in fact, we expect it + proposed_m += proposed_m; + const uint32_t twice_rem = rem + rem; + if (twice_rem >= d || twice_rem < rem) proposed_m += 1; + more = floor_log_2_d | LIBDIVIDE_ADD_MARKER; + } + result.magic = 1 + proposed_m; + result.more = more; + // result.more's shift should in general be ceil_log_2_d. But if we + // used the smaller power, we subtract one from the shift because we're + // using the smaller power. If we're using the larger power, we + // subtract one from the shift because it's taken care of by the add + // indicator. So floor_log_2_d happens to be correct in both cases. + } + return result; +} + +struct libdivide_u32_t libdivide_u32_gen(uint32_t d) { + return libdivide_internal_u32_gen(d, 0); +} + +struct libdivide_u32_branchfree_t libdivide_u32_branchfree_gen(uint32_t d) { + if (d == 1) { + LIBDIVIDE_ERROR("branchfree divider must be != 1"); + } + struct libdivide_u32_t tmp = libdivide_internal_u32_gen(d, 1); + struct libdivide_u32_branchfree_t ret = {tmp.magic, (uint8_t)(tmp.more & LIBDIVIDE_32_SHIFT_MASK)}; + return ret; +} + +uint32_t libdivide_u32_do(uint32_t numer, const struct libdivide_u32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return numer >> more; + } + else { + uint32_t q = libdivide_mullhi_u32(denom->magic, numer); + if (more & LIBDIVIDE_ADD_MARKER) { + uint32_t t = ((numer - q) >> 1) + q; + return t >> (more & LIBDIVIDE_32_SHIFT_MASK); + } + else { + // All upper bits are 0, + // don't need to mask them off. + return q >> more; + } + } +} + +uint32_t libdivide_u32_branchfree_do(uint32_t numer, const struct libdivide_u32_branchfree_t *denom) { + uint32_t q = libdivide_mullhi_u32(denom->magic, numer); + uint32_t t = ((numer - q) >> 1) + q; + return t >> denom->more; +} + +uint32_t libdivide_u32_recover(const struct libdivide_u32_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + + if (!denom->magic) { + return 1U << shift; + } else if (!(more & LIBDIVIDE_ADD_MARKER)) { + // We compute q = n/d = n*m / 2^(32 + shift) + // Therefore we have d = 2^(32 + shift) / m + // We need to ceil it. + // We know d is not a power of 2, so m is not a power of 2, + // so we can just add 1 to the floor + uint32_t hi_dividend = 1U << shift; + uint32_t rem_ignored; + return 1 + libdivide_64_div_32_to_32(hi_dividend, 0, denom->magic, &rem_ignored); + } else { + // Here we wish to compute d = 2^(32+shift+1)/(m+2^32). + // Notice (m + 2^32) is a 33 bit number. Use 64 bit division for now + // Also note that shift may be as high as 31, so shift + 1 will + // overflow. So we have to compute it as 2^(32+shift)/(m+2^32), and + // then double the quotient and remainder. + uint64_t half_n = 1ULL << (32 + shift); + uint64_t d = (1ULL << 32) | denom->magic; + // Note that the quotient is guaranteed <= 32 bits, but the remainder + // may need 33! + uint32_t half_q = (uint32_t)(half_n / d); + uint64_t rem = half_n % d; + // We computed 2^(32+shift)/(m+2^32) + // Need to double it, and then add 1 to the quotient if doubling th + // remainder would increase the quotient. + // Note that rem<<1 cannot overflow, since rem < d and d is 33 bits + uint32_t full_q = half_q + half_q + ((rem<<1) >= d); + + // We rounded down in gen (hence +1) + return full_q + 1; + } +} + +uint32_t libdivide_u32_branchfree_recover(const struct libdivide_u32_branchfree_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + + if (!denom->magic) { + return 1U << (shift + 1); + } else { + // Here we wish to compute d = 2^(32+shift+1)/(m+2^32). + // Notice (m + 2^32) is a 33 bit number. Use 64 bit division for now + // Also note that shift may be as high as 31, so shift + 1 will + // overflow. So we have to compute it as 2^(32+shift)/(m+2^32), and + // then double the quotient and remainder. + uint64_t half_n = 1ULL << (32 + shift); + uint64_t d = (1ULL << 32) | denom->magic; + // Note that the quotient is guaranteed <= 32 bits, but the remainder + // may need 33! + uint32_t half_q = (uint32_t)(half_n / d); + uint64_t rem = half_n % d; + // We computed 2^(32+shift)/(m+2^32) + // Need to double it, and then add 1 to the quotient if doubling th + // remainder would increase the quotient. + // Note that rem<<1 cannot overflow, since rem < d and d is 33 bits + uint32_t full_q = half_q + half_q + ((rem<<1) >= d); + + // We rounded down in gen (hence +1) + return full_q + 1; + } +} + +/////////// UINT64 + +static inline struct libdivide_u64_t libdivide_internal_u64_gen(uint64_t d, int branchfree) { + if (d == 0) { + LIBDIVIDE_ERROR("divider must be != 0"); + } + + struct libdivide_u64_t result; + uint32_t floor_log_2_d = 63 - libdivide_count_leading_zeros64(d); + + // Power of 2 + if ((d & (d - 1)) == 0) { + // We need to subtract 1 from the shift value in case of an unsigned + // branchfree divider because there is a hardcoded right shift by 1 + // in its division algorithm. Because of this we also need to add back + // 1 in its recovery algorithm. + result.magic = 0; + result.more = (uint8_t)(floor_log_2_d - (branchfree != 0)); + } else { + uint64_t proposed_m, rem; + uint8_t more; + // (1 << (64 + floor_log_2_d)) / d + proposed_m = libdivide_128_div_64_to_64(1ULL << floor_log_2_d, 0, d, &rem); + + LIBDIVIDE_ASSERT(rem > 0 && rem < d); + const uint64_t e = d - rem; + + // This power works if e < 2**floor_log_2_d. + if (!branchfree && e < (1ULL << floor_log_2_d)) { + // This power works + more = floor_log_2_d; + } else { + // We have to use the general 65-bit algorithm. We need to compute + // (2**power) / d. However, we already have (2**(power-1))/d and + // its remainder. By doubling both, and then correcting the + // remainder, we can compute the larger division. + // don't care about overflow here - in fact, we expect it + proposed_m += proposed_m; + const uint64_t twice_rem = rem + rem; + if (twice_rem >= d || twice_rem < rem) proposed_m += 1; + more = floor_log_2_d | LIBDIVIDE_ADD_MARKER; + } + result.magic = 1 + proposed_m; + result.more = more; + // result.more's shift should in general be ceil_log_2_d. But if we + // used the smaller power, we subtract one from the shift because we're + // using the smaller power. If we're using the larger power, we + // subtract one from the shift because it's taken care of by the add + // indicator. So floor_log_2_d happens to be correct in both cases, + // which is why we do it outside of the if statement. + } + return result; +} + +struct libdivide_u64_t libdivide_u64_gen(uint64_t d) { + return libdivide_internal_u64_gen(d, 0); +} + +struct libdivide_u64_branchfree_t libdivide_u64_branchfree_gen(uint64_t d) { + if (d == 1) { + LIBDIVIDE_ERROR("branchfree divider must be != 1"); + } + struct libdivide_u64_t tmp = libdivide_internal_u64_gen(d, 1); + struct libdivide_u64_branchfree_t ret = {tmp.magic, (uint8_t)(tmp.more & LIBDIVIDE_64_SHIFT_MASK)}; + return ret; +} + +uint64_t libdivide_u64_do(uint64_t numer, const struct libdivide_u64_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return numer >> more; + } + else { + uint64_t q = libdivide_mullhi_u64(denom->magic, numer); + if (more & LIBDIVIDE_ADD_MARKER) { + uint64_t t = ((numer - q) >> 1) + q; + return t >> (more & LIBDIVIDE_64_SHIFT_MASK); + } + else { + // All upper bits are 0, + // don't need to mask them off. + return q >> more; + } + } +} + +uint64_t libdivide_u64_branchfree_do(uint64_t numer, const struct libdivide_u64_branchfree_t *denom) { + uint64_t q = libdivide_mullhi_u64(denom->magic, numer); + uint64_t t = ((numer - q) >> 1) + q; + return t >> denom->more; +} + +uint64_t libdivide_u64_recover(const struct libdivide_u64_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + + if (!denom->magic) { + return 1ULL << shift; + } else if (!(more & LIBDIVIDE_ADD_MARKER)) { + // We compute q = n/d = n*m / 2^(64 + shift) + // Therefore we have d = 2^(64 + shift) / m + // We need to ceil it. + // We know d is not a power of 2, so m is not a power of 2, + // so we can just add 1 to the floor + uint64_t hi_dividend = 1ULL << shift; + uint64_t rem_ignored; + return 1 + libdivide_128_div_64_to_64(hi_dividend, 0, denom->magic, &rem_ignored); + } else { + // Here we wish to compute d = 2^(64+shift+1)/(m+2^64). + // Notice (m + 2^64) is a 65 bit number. This gets hairy. See + // libdivide_u32_recover for more on what we do here. + // TODO: do something better than 128 bit math + + // Full n is a (potentially) 129 bit value + // half_n is a 128 bit value + // Compute the hi half of half_n. Low half is 0. + uint64_t half_n_hi = 1ULL << shift, half_n_lo = 0; + // d is a 65 bit value. The high bit is always set to 1. + const uint64_t d_hi = 1, d_lo = denom->magic; + // Note that the quotient is guaranteed <= 64 bits, + // but the remainder may need 65! + uint64_t r_hi, r_lo; + uint64_t half_q = libdivide_128_div_128_to_64(half_n_hi, half_n_lo, d_hi, d_lo, &r_hi, &r_lo); + // We computed 2^(64+shift)/(m+2^64) + // Double the remainder ('dr') and check if that is larger than d + // Note that d is a 65 bit value, so r1 is small and so r1 + r1 + // cannot overflow + uint64_t dr_lo = r_lo + r_lo; + uint64_t dr_hi = r_hi + r_hi + (dr_lo < r_lo); // last term is carry + int dr_exceeds_d = (dr_hi > d_hi) || (dr_hi == d_hi && dr_lo >= d_lo); + uint64_t full_q = half_q + half_q + (dr_exceeds_d ? 1 : 0); + return full_q + 1; + } +} + +uint64_t libdivide_u64_branchfree_recover(const struct libdivide_u64_branchfree_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + + if (!denom->magic) { + return 1ULL << (shift + 1); + } else { + // Here we wish to compute d = 2^(64+shift+1)/(m+2^64). + // Notice (m + 2^64) is a 65 bit number. This gets hairy. See + // libdivide_u32_recover for more on what we do here. + // TODO: do something better than 128 bit math + + // Full n is a (potentially) 129 bit value + // half_n is a 128 bit value + // Compute the hi half of half_n. Low half is 0. + uint64_t half_n_hi = 1ULL << shift, half_n_lo = 0; + // d is a 65 bit value. The high bit is always set to 1. + const uint64_t d_hi = 1, d_lo = denom->magic; + // Note that the quotient is guaranteed <= 64 bits, + // but the remainder may need 65! + uint64_t r_hi, r_lo; + uint64_t half_q = libdivide_128_div_128_to_64(half_n_hi, half_n_lo, d_hi, d_lo, &r_hi, &r_lo); + // We computed 2^(64+shift)/(m+2^64) + // Double the remainder ('dr') and check if that is larger than d + // Note that d is a 65 bit value, so r1 is small and so r1 + r1 + // cannot overflow + uint64_t dr_lo = r_lo + r_lo; + uint64_t dr_hi = r_hi + r_hi + (dr_lo < r_lo); // last term is carry + int dr_exceeds_d = (dr_hi > d_hi) || (dr_hi == d_hi && dr_lo >= d_lo); + uint64_t full_q = half_q + half_q + (dr_exceeds_d ? 1 : 0); + return full_q + 1; + } +} + +/////////// SINT32 + +static inline struct libdivide_s32_t libdivide_internal_s32_gen(int32_t d, int branchfree) { + if (d == 0) { + LIBDIVIDE_ERROR("divider must be != 0"); + } + + struct libdivide_s32_t result; + + // If d is a power of 2, or negative a power of 2, we have to use a shift. + // This is especially important because the magic algorithm fails for -1. + // To check if d is a power of 2 or its inverse, it suffices to check + // whether its absolute value has exactly one bit set. This works even for + // INT_MIN, because abs(INT_MIN) == INT_MIN, and INT_MIN has one bit set + // and is a power of 2. + uint32_t ud = (uint32_t)d; + uint32_t absD = (d < 0) ? -ud : ud; + uint32_t floor_log_2_d = 31 - libdivide_count_leading_zeros32(absD); + // check if exactly one bit is set, + // don't care if absD is 0 since that's divide by zero + if ((absD & (absD - 1)) == 0) { + // Branchfree and normal paths are exactly the same + result.magic = 0; + result.more = floor_log_2_d | (d < 0 ? LIBDIVIDE_NEGATIVE_DIVISOR : 0); + } else { + LIBDIVIDE_ASSERT(floor_log_2_d >= 1); + + uint8_t more; + // the dividend here is 2**(floor_log_2_d + 31), so the low 32 bit word + // is 0 and the high word is floor_log_2_d - 1 + uint32_t rem, proposed_m; + proposed_m = libdivide_64_div_32_to_32(1U << (floor_log_2_d - 1), 0, absD, &rem); + const uint32_t e = absD - rem; + + // We are going to start with a power of floor_log_2_d - 1. + // This works if works if e < 2**floor_log_2_d. + if (!branchfree && e < (1U << floor_log_2_d)) { + // This power works + more = floor_log_2_d - 1; + } else { + // We need to go one higher. This should not make proposed_m + // overflow, but it will make it negative when interpreted as an + // int32_t. + proposed_m += proposed_m; + const uint32_t twice_rem = rem + rem; + if (twice_rem >= absD || twice_rem < rem) proposed_m += 1; + more = floor_log_2_d | LIBDIVIDE_ADD_MARKER; + } + + proposed_m += 1; + int32_t magic = (int32_t)proposed_m; + + // Mark if we are negative. Note we only negate the magic number in the + // branchfull case. + if (d < 0) { + more |= LIBDIVIDE_NEGATIVE_DIVISOR; + if (!branchfree) { + magic = -magic; + } + } + + result.more = more; + result.magic = magic; + } + return result; +} + +struct libdivide_s32_t libdivide_s32_gen(int32_t d) { + return libdivide_internal_s32_gen(d, 0); +} + +struct libdivide_s32_branchfree_t libdivide_s32_branchfree_gen(int32_t d) { + struct libdivide_s32_t tmp = libdivide_internal_s32_gen(d, 1); + struct libdivide_s32_branchfree_t result = {tmp.magic, tmp.more}; + return result; +} + +int32_t libdivide_s32_do(int32_t numer, const struct libdivide_s32_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + + if (!denom->magic) { + uint32_t sign = (int8_t)more >> 7; + uint32_t mask = (1U << shift) - 1; + uint32_t uq = numer + ((numer >> 31) & mask); + int32_t q = (int32_t)uq; + q >>= shift; + q = (q ^ sign) - sign; + return q; + } else { + uint32_t uq = (uint32_t)libdivide_mullhi_s32(denom->magic, numer); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift and then sign extend + int32_t sign = (int8_t)more >> 7; + // q += (more < 0 ? -numer : numer) + // cast required to avoid UB + uq += ((uint32_t)numer ^ sign) - sign; + } + int32_t q = (int32_t)uq; + q >>= shift; + q += (q < 0); + return q; + } +} + +int32_t libdivide_s32_branchfree_do(int32_t numer, const struct libdivide_s32_branchfree_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + // must be arithmetic shift and then sign extend + int32_t sign = (int8_t)more >> 7; + int32_t magic = denom->magic; + int32_t q = libdivide_mullhi_s32(magic, numer); + q += numer; + + // If q is non-negative, we have nothing to do + // If q is negative, we want to add either (2**shift)-1 if d is a power of + // 2, or (2**shift) if it is not a power of 2 + uint32_t is_power_of_2 = (magic == 0); + uint32_t q_sign = (uint32_t)(q >> 31); + q += q_sign & ((1U << shift) - is_power_of_2); + + // Now arithmetic right shift + q >>= shift; + // Negate if needed + q = (q ^ sign) - sign; + + return q; +} + +int32_t libdivide_s32_recover(const struct libdivide_s32_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + if (!denom->magic) { + uint32_t absD = 1U << shift; + if (more & LIBDIVIDE_NEGATIVE_DIVISOR) { + absD = -absD; + } + return (int32_t)absD; + } else { + // Unsigned math is much easier + // We negate the magic number only in the branchfull case, and we don't + // know which case we're in. However we have enough information to + // determine the correct sign of the magic number. The divisor was + // negative if LIBDIVIDE_NEGATIVE_DIVISOR is set. If ADD_MARKER is set, + // the magic number's sign is opposite that of the divisor. + // We want to compute the positive magic number. + int negative_divisor = (more & LIBDIVIDE_NEGATIVE_DIVISOR); + int magic_was_negated = (more & LIBDIVIDE_ADD_MARKER) + ? denom->magic > 0 : denom->magic < 0; + + // Handle the power of 2 case (including branchfree) + if (denom->magic == 0) { + int32_t result = 1U << shift; + return negative_divisor ? -result : result; + } + + uint32_t d = (uint32_t)(magic_was_negated ? -denom->magic : denom->magic); + uint64_t n = 1ULL << (32 + shift); // this shift cannot exceed 30 + uint32_t q = (uint32_t)(n / d); + int32_t result = (int32_t)q; + result += 1; + return negative_divisor ? -result : result; + } +} + +int32_t libdivide_s32_branchfree_recover(const struct libdivide_s32_branchfree_t *denom) { + return libdivide_s32_recover((const struct libdivide_s32_t *)denom); +} + +///////////// SINT64 + +static inline struct libdivide_s64_t libdivide_internal_s64_gen(int64_t d, int branchfree) { + if (d == 0) { + LIBDIVIDE_ERROR("divider must be != 0"); + } + + struct libdivide_s64_t result; + + // If d is a power of 2, or negative a power of 2, we have to use a shift. + // This is especially important because the magic algorithm fails for -1. + // To check if d is a power of 2 or its inverse, it suffices to check + // whether its absolute value has exactly one bit set. This works even for + // INT_MIN, because abs(INT_MIN) == INT_MIN, and INT_MIN has one bit set + // and is a power of 2. + uint64_t ud = (uint64_t)d; + uint64_t absD = (d < 0) ? -ud : ud; + uint32_t floor_log_2_d = 63 - libdivide_count_leading_zeros64(absD); + // check if exactly one bit is set, + // don't care if absD is 0 since that's divide by zero + if ((absD & (absD - 1)) == 0) { + // Branchfree and non-branchfree cases are the same + result.magic = 0; + result.more = floor_log_2_d | (d < 0 ? LIBDIVIDE_NEGATIVE_DIVISOR : 0); + } else { + // the dividend here is 2**(floor_log_2_d + 63), so the low 64 bit word + // is 0 and the high word is floor_log_2_d - 1 + uint8_t more; + uint64_t rem, proposed_m; + proposed_m = libdivide_128_div_64_to_64(1ULL << (floor_log_2_d - 1), 0, absD, &rem); + const uint64_t e = absD - rem; + + // We are going to start with a power of floor_log_2_d - 1. + // This works if works if e < 2**floor_log_2_d. + if (!branchfree && e < (1ULL << floor_log_2_d)) { + // This power works + more = floor_log_2_d - 1; + } else { + // We need to go one higher. This should not make proposed_m + // overflow, but it will make it negative when interpreted as an + // int32_t. + proposed_m += proposed_m; + const uint64_t twice_rem = rem + rem; + if (twice_rem >= absD || twice_rem < rem) proposed_m += 1; + // note that we only set the LIBDIVIDE_NEGATIVE_DIVISOR bit if we + // also set ADD_MARKER this is an annoying optimization that + // enables algorithm #4 to avoid the mask. However we always set it + // in the branchfree case + more = floor_log_2_d | LIBDIVIDE_ADD_MARKER; + } + proposed_m += 1; + int64_t magic = (int64_t)proposed_m; + + // Mark if we are negative + if (d < 0) { + more |= LIBDIVIDE_NEGATIVE_DIVISOR; + if (!branchfree) { + magic = -magic; + } + } + + result.more = more; + result.magic = magic; + } + return result; +} + +struct libdivide_s64_t libdivide_s64_gen(int64_t d) { + return libdivide_internal_s64_gen(d, 0); +} + +struct libdivide_s64_branchfree_t libdivide_s64_branchfree_gen(int64_t d) { + struct libdivide_s64_t tmp = libdivide_internal_s64_gen(d, 1); + struct libdivide_s64_branchfree_t ret = {tmp.magic, tmp.more}; + return ret; +} + +int64_t libdivide_s64_do(int64_t numer, const struct libdivide_s64_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + + if (!denom->magic) { // shift path + uint64_t mask = (1ULL << shift) - 1; + uint64_t uq = numer + ((numer >> 63) & mask); + int64_t q = (int64_t)uq; + q >>= shift; + // must be arithmetic shift and then sign-extend + int64_t sign = (int8_t)more >> 7; + q = (q ^ sign) - sign; + return q; + } else { + uint64_t uq = (uint64_t)libdivide_mullhi_s64(denom->magic, numer); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift and then sign extend + int64_t sign = (int8_t)more >> 7; + // q += (more < 0 ? -numer : numer) + // cast required to avoid UB + uq += ((uint64_t)numer ^ sign) - sign; + } + int64_t q = (int64_t)uq; + q >>= shift; + q += (q < 0); + return q; + } +} + +int64_t libdivide_s64_branchfree_do(int64_t numer, const struct libdivide_s64_branchfree_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + // must be arithmetic shift and then sign extend + int64_t sign = (int8_t)more >> 7; + int64_t magic = denom->magic; + int64_t q = libdivide_mullhi_s64(magic, numer); + q += numer; + + // If q is non-negative, we have nothing to do. + // If q is negative, we want to add either (2**shift)-1 if d is a power of + // 2, or (2**shift) if it is not a power of 2. + uint64_t is_power_of_2 = (magic == 0); + uint64_t q_sign = (uint64_t)(q >> 63); + q += q_sign & ((1ULL << shift) - is_power_of_2); + + // Arithmetic right shift + q >>= shift; + // Negate if needed + q = (q ^ sign) - sign; + + return q; +} + +int64_t libdivide_s64_recover(const struct libdivide_s64_t *denom) { + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + if (denom->magic == 0) { // shift path + uint64_t absD = 1ULL << shift; + if (more & LIBDIVIDE_NEGATIVE_DIVISOR) { + absD = -absD; + } + return (int64_t)absD; + } else { + // Unsigned math is much easier + int negative_divisor = (more & LIBDIVIDE_NEGATIVE_DIVISOR); + int magic_was_negated = (more & LIBDIVIDE_ADD_MARKER) + ? denom->magic > 0 : denom->magic < 0; + + uint64_t d = (uint64_t)(magic_was_negated ? -denom->magic : denom->magic); + uint64_t n_hi = 1ULL << shift, n_lo = 0; + uint64_t rem_ignored; + uint64_t q = libdivide_128_div_64_to_64(n_hi, n_lo, d, &rem_ignored); + int64_t result = (int64_t)(q + 1); + if (negative_divisor) { + result = -result; + } + return result; + } +} + +int64_t libdivide_s64_branchfree_recover(const struct libdivide_s64_branchfree_t *denom) { + return libdivide_s64_recover((const struct libdivide_s64_t *)denom); +} + +#if defined(LIBDIVIDE_AVX512) + +static inline __m512i libdivide_u32_do_vector(__m512i numers, const struct libdivide_u32_t *denom); +static inline __m512i libdivide_s32_do_vector(__m512i numers, const struct libdivide_s32_t *denom); +static inline __m512i libdivide_u64_do_vector(__m512i numers, const struct libdivide_u64_t *denom); +static inline __m512i libdivide_s64_do_vector(__m512i numers, const struct libdivide_s64_t *denom); + +static inline __m512i libdivide_u32_branchfree_do_vector(__m512i numers, const struct libdivide_u32_branchfree_t *denom); +static inline __m512i libdivide_s32_branchfree_do_vector(__m512i numers, const struct libdivide_s32_branchfree_t *denom); +static inline __m512i libdivide_u64_branchfree_do_vector(__m512i numers, const struct libdivide_u64_branchfree_t *denom); +static inline __m512i libdivide_s64_branchfree_do_vector(__m512i numers, const struct libdivide_s64_branchfree_t *denom); + +//////// Internal Utility Functions + +static inline __m512i libdivide_s64_signbits(__m512i v) {; + return _mm512_srai_epi64(v, 63); +} + +static inline __m512i libdivide_s64_shift_right_vector(__m512i v, int amt) { + return _mm512_srai_epi64(v, amt); +} + +// Here, b is assumed to contain one 32-bit value repeated. +static inline __m512i libdivide_mullhi_u32_vector(__m512i a, __m512i b) { + __m512i hi_product_0Z2Z = _mm512_srli_epi64(_mm512_mul_epu32(a, b), 32); + __m512i a1X3X = _mm512_srli_epi64(a, 32); + __m512i mask = _mm512_set_epi32(-1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0); + __m512i hi_product_Z1Z3 = _mm512_and_si512(_mm512_mul_epu32(a1X3X, b), mask); + return _mm512_or_si512(hi_product_0Z2Z, hi_product_Z1Z3); +} + +// b is one 32-bit value repeated. +static inline __m512i libdivide_mullhi_s32_vector(__m512i a, __m512i b) { + __m512i hi_product_0Z2Z = _mm512_srli_epi64(_mm512_mul_epi32(a, b), 32); + __m512i a1X3X = _mm512_srli_epi64(a, 32); + __m512i mask = _mm512_set_epi32(-1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0, -1, 0); + __m512i hi_product_Z1Z3 = _mm512_and_si512(_mm512_mul_epi32(a1X3X, b), mask); + return _mm512_or_si512(hi_product_0Z2Z, hi_product_Z1Z3); +} + +// Here, y is assumed to contain one 64-bit value repeated. +// https://stackoverflow.com/a/28827013 +static inline __m512i libdivide_mullhi_u64_vector(__m512i x, __m512i y) { + __m512i lomask = _mm512_set1_epi64(0xffffffff); + __m512i xh = _mm512_shuffle_epi32(x, (_MM_PERM_ENUM) 0xB1); + __m512i yh = _mm512_shuffle_epi32(y, (_MM_PERM_ENUM) 0xB1); + __m512i w0 = _mm512_mul_epu32(x, y); + __m512i w1 = _mm512_mul_epu32(x, yh); + __m512i w2 = _mm512_mul_epu32(xh, y); + __m512i w3 = _mm512_mul_epu32(xh, yh); + __m512i w0h = _mm512_srli_epi64(w0, 32); + __m512i s1 = _mm512_add_epi64(w1, w0h); + __m512i s1l = _mm512_and_si512(s1, lomask); + __m512i s1h = _mm512_srli_epi64(s1, 32); + __m512i s2 = _mm512_add_epi64(w2, s1l); + __m512i s2h = _mm512_srli_epi64(s2, 32); + __m512i hi = _mm512_add_epi64(w3, s1h); + hi = _mm512_add_epi64(hi, s2h); + + return hi; +} + +// y is one 64-bit value repeated. +static inline __m512i libdivide_mullhi_s64_vector(__m512i x, __m512i y) { + __m512i p = libdivide_mullhi_u64_vector(x, y); + __m512i t1 = _mm512_and_si512(libdivide_s64_signbits(x), y); + __m512i t2 = _mm512_and_si512(libdivide_s64_signbits(y), x); + p = _mm512_sub_epi64(p, t1); + p = _mm512_sub_epi64(p, t2); + return p; +} + +////////// UINT32 + +__m512i libdivide_u32_do_vector(__m512i numers, const struct libdivide_u32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm512_srli_epi32(numers, more); + } + else { + __m512i q = libdivide_mullhi_u32_vector(numers, _mm512_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + __m512i t = _mm512_add_epi32(_mm512_srli_epi32(_mm512_sub_epi32(numers, q), 1), q); + return _mm512_srli_epi32(t, shift); + } + else { + return _mm512_srli_epi32(q, more); + } + } +} + +__m512i libdivide_u32_branchfree_do_vector(__m512i numers, const struct libdivide_u32_branchfree_t *denom) { + __m512i q = libdivide_mullhi_u32_vector(numers, _mm512_set1_epi32(denom->magic)); + __m512i t = _mm512_add_epi32(_mm512_srli_epi32(_mm512_sub_epi32(numers, q), 1), q); + return _mm512_srli_epi32(t, denom->more); +} + +////////// UINT64 + +__m512i libdivide_u64_do_vector(__m512i numers, const struct libdivide_u64_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm512_srli_epi64(numers, more); + } + else { + __m512i q = libdivide_mullhi_u64_vector(numers, _mm512_set1_epi64(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + __m512i t = _mm512_add_epi64(_mm512_srli_epi64(_mm512_sub_epi64(numers, q), 1), q); + return _mm512_srli_epi64(t, shift); + } + else { + return _mm512_srli_epi64(q, more); + } + } +} + +__m512i libdivide_u64_branchfree_do_vector(__m512i numers, const struct libdivide_u64_branchfree_t *denom) { + __m512i q = libdivide_mullhi_u64_vector(numers, _mm512_set1_epi64(denom->magic)); + __m512i t = _mm512_add_epi64(_mm512_srli_epi64(_mm512_sub_epi64(numers, q), 1), q); + return _mm512_srli_epi64(t, denom->more); +} + +////////// SINT32 + +__m512i libdivide_s32_do_vector(__m512i numers, const struct libdivide_s32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + uint32_t mask = (1U << shift) - 1; + __m512i roundToZeroTweak = _mm512_set1_epi32(mask); + // q = numer + ((numer >> 31) & roundToZeroTweak); + __m512i q = _mm512_add_epi32(numers, _mm512_and_si512(_mm512_srai_epi32(numers, 31), roundToZeroTweak)); + q = _mm512_srai_epi32(q, shift); + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm512_sub_epi32(_mm512_xor_si512(q, sign), sign); + return q; + } + else { + __m512i q = libdivide_mullhi_s32_vector(numers, _mm512_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm512_add_epi32(q, _mm512_sub_epi32(_mm512_xor_si512(numers, sign), sign)); + } + // q >>= shift + q = _mm512_srai_epi32(q, more & LIBDIVIDE_32_SHIFT_MASK); + q = _mm512_add_epi32(q, _mm512_srli_epi32(q, 31)); // q += (q < 0) + return q; + } +} + +__m512i libdivide_s32_branchfree_do_vector(__m512i numers, const struct libdivide_s32_branchfree_t *denom) { + int32_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + // must be arithmetic shift + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + __m512i q = libdivide_mullhi_s32_vector(numers, _mm512_set1_epi32(magic)); + q = _mm512_add_epi32(q, numers); // q += numers + + // If q is non-negative, we have nothing to do + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2 + uint32_t is_power_of_2 = (magic == 0); + __m512i q_sign = _mm512_srai_epi32(q, 31); // q_sign = q >> 31 + __m512i mask = _mm512_set1_epi32((1U << shift) - is_power_of_2); + q = _mm512_add_epi32(q, _mm512_and_si512(q_sign, mask)); // q = q + (q_sign & mask) + q = _mm512_srai_epi32(q, shift); // q >>= shift + q = _mm512_sub_epi32(_mm512_xor_si512(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +////////// SINT64 + +__m512i libdivide_s64_do_vector(__m512i numers, const struct libdivide_s64_t *denom) { + uint8_t more = denom->more; + int64_t magic = denom->magic; + if (magic == 0) { // shift path + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + uint64_t mask = (1ULL << shift) - 1; + __m512i roundToZeroTweak = _mm512_set1_epi64(mask); + // q = numer + ((numer >> 63) & roundToZeroTweak); + __m512i q = _mm512_add_epi64(numers, _mm512_and_si512(libdivide_s64_signbits(numers), roundToZeroTweak)); + q = libdivide_s64_shift_right_vector(q, shift); + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm512_sub_epi64(_mm512_xor_si512(q, sign), sign); + return q; + } + else { + __m512i q = libdivide_mullhi_s64_vector(numers, _mm512_set1_epi64(magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm512_add_epi64(q, _mm512_sub_epi64(_mm512_xor_si512(numers, sign), sign)); + } + // q >>= denom->mult_path.shift + q = libdivide_s64_shift_right_vector(q, more & LIBDIVIDE_64_SHIFT_MASK); + q = _mm512_add_epi64(q, _mm512_srli_epi64(q, 63)); // q += (q < 0) + return q; + } +} + +__m512i libdivide_s64_branchfree_do_vector(__m512i numers, const struct libdivide_s64_branchfree_t *denom) { + int64_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + // must be arithmetic shift + __m512i sign = _mm512_set1_epi32((int8_t)more >> 7); + + // libdivide_mullhi_s64(numers, magic); + __m512i q = libdivide_mullhi_s64_vector(numers, _mm512_set1_epi64(magic)); + q = _mm512_add_epi64(q, numers); // q += numers + + // If q is non-negative, we have nothing to do. + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2. + uint32_t is_power_of_2 = (magic == 0); + __m512i q_sign = libdivide_s64_signbits(q); // q_sign = q >> 63 + __m512i mask = _mm512_set1_epi64((1ULL << shift) - is_power_of_2); + q = _mm512_add_epi64(q, _mm512_and_si512(q_sign, mask)); // q = q + (q_sign & mask) + q = libdivide_s64_shift_right_vector(q, shift); // q >>= shift + q = _mm512_sub_epi64(_mm512_xor_si512(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +#elif defined(LIBDIVIDE_AVX2) + +static inline __m256i libdivide_u32_do_vector(__m256i numers, const struct libdivide_u32_t *denom); +static inline __m256i libdivide_s32_do_vector(__m256i numers, const struct libdivide_s32_t *denom); +static inline __m256i libdivide_u64_do_vector(__m256i numers, const struct libdivide_u64_t *denom); +static inline __m256i libdivide_s64_do_vector(__m256i numers, const struct libdivide_s64_t *denom); + +static inline __m256i libdivide_u32_branchfree_do_vector(__m256i numers, const struct libdivide_u32_branchfree_t *denom); +static inline __m256i libdivide_s32_branchfree_do_vector(__m256i numers, const struct libdivide_s32_branchfree_t *denom); +static inline __m256i libdivide_u64_branchfree_do_vector(__m256i numers, const struct libdivide_u64_branchfree_t *denom); +static inline __m256i libdivide_s64_branchfree_do_vector(__m256i numers, const struct libdivide_s64_branchfree_t *denom); + +//////// Internal Utility Functions + +// Implementation of _mm256_srai_epi64(v, 63) (from AVX512). +static inline __m256i libdivide_s64_signbits(__m256i v) { + __m256i hiBitsDuped = _mm256_shuffle_epi32(v, _MM_SHUFFLE(3, 3, 1, 1)); + __m256i signBits = _mm256_srai_epi32(hiBitsDuped, 31); + return signBits; +} + +// Implementation of _mm256_srai_epi64 (from AVX512). +static inline __m256i libdivide_s64_shift_right_vector(__m256i v, int amt) { + const int b = 64 - amt; + __m256i m = _mm256_set1_epi64x(1ULL << (b - 1)); + __m256i x = _mm256_srli_epi64(v, amt); + __m256i result = _mm256_sub_epi64(_mm256_xor_si256(x, m), m); + return result; +} + +// Here, b is assumed to contain one 32-bit value repeated. +static inline __m256i libdivide_mullhi_u32_vector(__m256i a, __m256i b) { + __m256i hi_product_0Z2Z = _mm256_srli_epi64(_mm256_mul_epu32(a, b), 32); + __m256i a1X3X = _mm256_srli_epi64(a, 32); + __m256i mask = _mm256_set_epi32(-1, 0, -1, 0, -1, 0, -1, 0); + __m256i hi_product_Z1Z3 = _mm256_and_si256(_mm256_mul_epu32(a1X3X, b), mask); + return _mm256_or_si256(hi_product_0Z2Z, hi_product_Z1Z3); +} + +// b is one 32-bit value repeated. +static inline __m256i libdivide_mullhi_s32_vector(__m256i a, __m256i b) { + __m256i hi_product_0Z2Z = _mm256_srli_epi64(_mm256_mul_epi32(a, b), 32); + __m256i a1X3X = _mm256_srli_epi64(a, 32); + __m256i mask = _mm256_set_epi32(-1, 0, -1, 0, -1, 0, -1, 0); + __m256i hi_product_Z1Z3 = _mm256_and_si256(_mm256_mul_epi32(a1X3X, b), mask); + return _mm256_or_si256(hi_product_0Z2Z, hi_product_Z1Z3); +} + +// Here, y is assumed to contain one 64-bit value repeated. +// https://stackoverflow.com/a/28827013 +static inline __m256i libdivide_mullhi_u64_vector(__m256i x, __m256i y) { + __m256i lomask = _mm256_set1_epi64x(0xffffffff); + __m256i xh = _mm256_shuffle_epi32(x, 0xB1); // x0l, x0h, x1l, x1h + __m256i yh = _mm256_shuffle_epi32(y, 0xB1); // y0l, y0h, y1l, y1h + __m256i w0 = _mm256_mul_epu32(x, y); // x0l*y0l, x1l*y1l + __m256i w1 = _mm256_mul_epu32(x, yh); // x0l*y0h, x1l*y1h + __m256i w2 = _mm256_mul_epu32(xh, y); // x0h*y0l, x1h*y0l + __m256i w3 = _mm256_mul_epu32(xh, yh); // x0h*y0h, x1h*y1h + __m256i w0h = _mm256_srli_epi64(w0, 32); + __m256i s1 = _mm256_add_epi64(w1, w0h); + __m256i s1l = _mm256_and_si256(s1, lomask); + __m256i s1h = _mm256_srli_epi64(s1, 32); + __m256i s2 = _mm256_add_epi64(w2, s1l); + __m256i s2h = _mm256_srli_epi64(s2, 32); + __m256i hi = _mm256_add_epi64(w3, s1h); + hi = _mm256_add_epi64(hi, s2h); + + return hi; +} + +// y is one 64-bit value repeated. +static inline __m256i libdivide_mullhi_s64_vector(__m256i x, __m256i y) { + __m256i p = libdivide_mullhi_u64_vector(x, y); + __m256i t1 = _mm256_and_si256(libdivide_s64_signbits(x), y); + __m256i t2 = _mm256_and_si256(libdivide_s64_signbits(y), x); + p = _mm256_sub_epi64(p, t1); + p = _mm256_sub_epi64(p, t2); + return p; +} + +////////// UINT32 + +__m256i libdivide_u32_do_vector(__m256i numers, const struct libdivide_u32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm256_srli_epi32(numers, more); + } + else { + __m256i q = libdivide_mullhi_u32_vector(numers, _mm256_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + __m256i t = _mm256_add_epi32(_mm256_srli_epi32(_mm256_sub_epi32(numers, q), 1), q); + return _mm256_srli_epi32(t, shift); + } + else { + return _mm256_srli_epi32(q, more); + } + } +} + +__m256i libdivide_u32_branchfree_do_vector(__m256i numers, const struct libdivide_u32_branchfree_t *denom) { + __m256i q = libdivide_mullhi_u32_vector(numers, _mm256_set1_epi32(denom->magic)); + __m256i t = _mm256_add_epi32(_mm256_srli_epi32(_mm256_sub_epi32(numers, q), 1), q); + return _mm256_srli_epi32(t, denom->more); +} + +////////// UINT64 + +__m256i libdivide_u64_do_vector(__m256i numers, const struct libdivide_u64_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm256_srli_epi64(numers, more); + } + else { + __m256i q = libdivide_mullhi_u64_vector(numers, _mm256_set1_epi64x(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + __m256i t = _mm256_add_epi64(_mm256_srli_epi64(_mm256_sub_epi64(numers, q), 1), q); + return _mm256_srli_epi64(t, shift); + } + else { + return _mm256_srli_epi64(q, more); + } + } +} + +__m256i libdivide_u64_branchfree_do_vector(__m256i numers, const struct libdivide_u64_branchfree_t *denom) { + __m256i q = libdivide_mullhi_u64_vector(numers, _mm256_set1_epi64x(denom->magic)); + __m256i t = _mm256_add_epi64(_mm256_srli_epi64(_mm256_sub_epi64(numers, q), 1), q); + return _mm256_srli_epi64(t, denom->more); +} + +////////// SINT32 + +__m256i libdivide_s32_do_vector(__m256i numers, const struct libdivide_s32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + uint32_t mask = (1U << shift) - 1; + __m256i roundToZeroTweak = _mm256_set1_epi32(mask); + // q = numer + ((numer >> 31) & roundToZeroTweak); + __m256i q = _mm256_add_epi32(numers, _mm256_and_si256(_mm256_srai_epi32(numers, 31), roundToZeroTweak)); + q = _mm256_srai_epi32(q, shift); + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm256_sub_epi32(_mm256_xor_si256(q, sign), sign); + return q; + } + else { + __m256i q = libdivide_mullhi_s32_vector(numers, _mm256_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm256_add_epi32(q, _mm256_sub_epi32(_mm256_xor_si256(numers, sign), sign)); + } + // q >>= shift + q = _mm256_srai_epi32(q, more & LIBDIVIDE_32_SHIFT_MASK); + q = _mm256_add_epi32(q, _mm256_srli_epi32(q, 31)); // q += (q < 0) + return q; + } +} + +__m256i libdivide_s32_branchfree_do_vector(__m256i numers, const struct libdivide_s32_branchfree_t *denom) { + int32_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + // must be arithmetic shift + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + __m256i q = libdivide_mullhi_s32_vector(numers, _mm256_set1_epi32(magic)); + q = _mm256_add_epi32(q, numers); // q += numers + + // If q is non-negative, we have nothing to do + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2 + uint32_t is_power_of_2 = (magic == 0); + __m256i q_sign = _mm256_srai_epi32(q, 31); // q_sign = q >> 31 + __m256i mask = _mm256_set1_epi32((1U << shift) - is_power_of_2); + q = _mm256_add_epi32(q, _mm256_and_si256(q_sign, mask)); // q = q + (q_sign & mask) + q = _mm256_srai_epi32(q, shift); // q >>= shift + q = _mm256_sub_epi32(_mm256_xor_si256(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +////////// SINT64 + +__m256i libdivide_s64_do_vector(__m256i numers, const struct libdivide_s64_t *denom) { + uint8_t more = denom->more; + int64_t magic = denom->magic; + if (magic == 0) { // shift path + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + uint64_t mask = (1ULL << shift) - 1; + __m256i roundToZeroTweak = _mm256_set1_epi64x(mask); + // q = numer + ((numer >> 63) & roundToZeroTweak); + __m256i q = _mm256_add_epi64(numers, _mm256_and_si256(libdivide_s64_signbits(numers), roundToZeroTweak)); + q = libdivide_s64_shift_right_vector(q, shift); + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm256_sub_epi64(_mm256_xor_si256(q, sign), sign); + return q; + } + else { + __m256i q = libdivide_mullhi_s64_vector(numers, _mm256_set1_epi64x(magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm256_add_epi64(q, _mm256_sub_epi64(_mm256_xor_si256(numers, sign), sign)); + } + // q >>= denom->mult_path.shift + q = libdivide_s64_shift_right_vector(q, more & LIBDIVIDE_64_SHIFT_MASK); + q = _mm256_add_epi64(q, _mm256_srli_epi64(q, 63)); // q += (q < 0) + return q; + } +} + +__m256i libdivide_s64_branchfree_do_vector(__m256i numers, const struct libdivide_s64_branchfree_t *denom) { + int64_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + // must be arithmetic shift + __m256i sign = _mm256_set1_epi32((int8_t)more >> 7); + + // libdivide_mullhi_s64(numers, magic); + __m256i q = libdivide_mullhi_s64_vector(numers, _mm256_set1_epi64x(magic)); + q = _mm256_add_epi64(q, numers); // q += numers + + // If q is non-negative, we have nothing to do. + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2. + uint32_t is_power_of_2 = (magic == 0); + __m256i q_sign = libdivide_s64_signbits(q); // q_sign = q >> 63 + __m256i mask = _mm256_set1_epi64x((1ULL << shift) - is_power_of_2); + q = _mm256_add_epi64(q, _mm256_and_si256(q_sign, mask)); // q = q + (q_sign & mask) + q = libdivide_s64_shift_right_vector(q, shift); // q >>= shift + q = _mm256_sub_epi64(_mm256_xor_si256(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +#elif defined(LIBDIVIDE_SSE2) + +static inline __m128i libdivide_u32_do_vector(__m128i numers, const struct libdivide_u32_t *denom); +static inline __m128i libdivide_s32_do_vector(__m128i numers, const struct libdivide_s32_t *denom); +static inline __m128i libdivide_u64_do_vector(__m128i numers, const struct libdivide_u64_t *denom); +static inline __m128i libdivide_s64_do_vector(__m128i numers, const struct libdivide_s64_t *denom); + +static inline __m128i libdivide_u32_branchfree_do_vector(__m128i numers, const struct libdivide_u32_branchfree_t *denom); +static inline __m128i libdivide_s32_branchfree_do_vector(__m128i numers, const struct libdivide_s32_branchfree_t *denom); +static inline __m128i libdivide_u64_branchfree_do_vector(__m128i numers, const struct libdivide_u64_branchfree_t *denom); +static inline __m128i libdivide_s64_branchfree_do_vector(__m128i numers, const struct libdivide_s64_branchfree_t *denom); + +//////// Internal Utility Functions + +// Implementation of _mm_srai_epi64(v, 63) (from AVX512). +static inline __m128i libdivide_s64_signbits(__m128i v) { + __m128i hiBitsDuped = _mm_shuffle_epi32(v, _MM_SHUFFLE(3, 3, 1, 1)); + __m128i signBits = _mm_srai_epi32(hiBitsDuped, 31); + return signBits; +} + +// Implementation of _mm_srai_epi64 (from AVX512). +static inline __m128i libdivide_s64_shift_right_vector(__m128i v, int amt) { + const int b = 64 - amt; + __m128i m = _mm_set1_epi64x(1ULL << (b - 1)); + __m128i x = _mm_srli_epi64(v, amt); + __m128i result = _mm_sub_epi64(_mm_xor_si128(x, m), m); + return result; +} + +// Here, b is assumed to contain one 32-bit value repeated. +static inline __m128i libdivide_mullhi_u32_vector(__m128i a, __m128i b) { + __m128i hi_product_0Z2Z = _mm_srli_epi64(_mm_mul_epu32(a, b), 32); + __m128i a1X3X = _mm_srli_epi64(a, 32); + __m128i mask = _mm_set_epi32(-1, 0, -1, 0); + __m128i hi_product_Z1Z3 = _mm_and_si128(_mm_mul_epu32(a1X3X, b), mask); + return _mm_or_si128(hi_product_0Z2Z, hi_product_Z1Z3); +} + +// SSE2 does not have a signed multiplication instruction, but we can convert +// unsigned to signed pretty efficiently. Again, b is just a 32 bit value +// repeated four times. +static inline __m128i libdivide_mullhi_s32_vector(__m128i a, __m128i b) { + __m128i p = libdivide_mullhi_u32_vector(a, b); + // t1 = (a >> 31) & y, arithmetic shift + __m128i t1 = _mm_and_si128(_mm_srai_epi32(a, 31), b); + __m128i t2 = _mm_and_si128(_mm_srai_epi32(b, 31), a); + p = _mm_sub_epi32(p, t1); + p = _mm_sub_epi32(p, t2); + return p; +} + +// Here, y is assumed to contain one 64-bit value repeated. +// https://stackoverflow.com/a/28827013 +static inline __m128i libdivide_mullhi_u64_vector(__m128i x, __m128i y) { + __m128i lomask = _mm_set1_epi64x(0xffffffff); + __m128i xh = _mm_shuffle_epi32(x, 0xB1); // x0l, x0h, x1l, x1h + __m128i yh = _mm_shuffle_epi32(y, 0xB1); // y0l, y0h, y1l, y1h + __m128i w0 = _mm_mul_epu32(x, y); // x0l*y0l, x1l*y1l + __m128i w1 = _mm_mul_epu32(x, yh); // x0l*y0h, x1l*y1h + __m128i w2 = _mm_mul_epu32(xh, y); // x0h*y0l, x1h*y0l + __m128i w3 = _mm_mul_epu32(xh, yh); // x0h*y0h, x1h*y1h + __m128i w0h = _mm_srli_epi64(w0, 32); + __m128i s1 = _mm_add_epi64(w1, w0h); + __m128i s1l = _mm_and_si128(s1, lomask); + __m128i s1h = _mm_srli_epi64(s1, 32); + __m128i s2 = _mm_add_epi64(w2, s1l); + __m128i s2h = _mm_srli_epi64(s2, 32); + __m128i hi = _mm_add_epi64(w3, s1h); + hi = _mm_add_epi64(hi, s2h); + + return hi; +} + +// y is one 64-bit value repeated. +static inline __m128i libdivide_mullhi_s64_vector(__m128i x, __m128i y) { + __m128i p = libdivide_mullhi_u64_vector(x, y); + __m128i t1 = _mm_and_si128(libdivide_s64_signbits(x), y); + __m128i t2 = _mm_and_si128(libdivide_s64_signbits(y), x); + p = _mm_sub_epi64(p, t1); + p = _mm_sub_epi64(p, t2); + return p; +} + +////////// UINT32 + +__m128i libdivide_u32_do_vector(__m128i numers, const struct libdivide_u32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm_srli_epi32(numers, more); + } + else { + __m128i q = libdivide_mullhi_u32_vector(numers, _mm_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + __m128i t = _mm_add_epi32(_mm_srli_epi32(_mm_sub_epi32(numers, q), 1), q); + return _mm_srli_epi32(t, shift); + } + else { + return _mm_srli_epi32(q, more); + } + } +} + +__m128i libdivide_u32_branchfree_do_vector(__m128i numers, const struct libdivide_u32_branchfree_t *denom) { + __m128i q = libdivide_mullhi_u32_vector(numers, _mm_set1_epi32(denom->magic)); + __m128i t = _mm_add_epi32(_mm_srli_epi32(_mm_sub_epi32(numers, q), 1), q); + return _mm_srli_epi32(t, denom->more); +} + +////////// UINT64 + +__m128i libdivide_u64_do_vector(__m128i numers, const struct libdivide_u64_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + return _mm_srli_epi64(numers, more); + } + else { + __m128i q = libdivide_mullhi_u64_vector(numers, _mm_set1_epi64x(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // uint32_t t = ((numer - q) >> 1) + q; + // return t >> denom->shift; + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + __m128i t = _mm_add_epi64(_mm_srli_epi64(_mm_sub_epi64(numers, q), 1), q); + return _mm_srli_epi64(t, shift); + } + else { + return _mm_srli_epi64(q, more); + } + } +} + +__m128i libdivide_u64_branchfree_do_vector(__m128i numers, const struct libdivide_u64_branchfree_t *denom) { + __m128i q = libdivide_mullhi_u64_vector(numers, _mm_set1_epi64x(denom->magic)); + __m128i t = _mm_add_epi64(_mm_srli_epi64(_mm_sub_epi64(numers, q), 1), q); + return _mm_srli_epi64(t, denom->more); +} + +////////// SINT32 + +__m128i libdivide_s32_do_vector(__m128i numers, const struct libdivide_s32_t *denom) { + uint8_t more = denom->more; + if (!denom->magic) { + uint32_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + uint32_t mask = (1U << shift) - 1; + __m128i roundToZeroTweak = _mm_set1_epi32(mask); + // q = numer + ((numer >> 31) & roundToZeroTweak); + __m128i q = _mm_add_epi32(numers, _mm_and_si128(_mm_srai_epi32(numers, 31), roundToZeroTweak)); + q = _mm_srai_epi32(q, shift); + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm_sub_epi32(_mm_xor_si128(q, sign), sign); + return q; + } + else { + __m128i q = libdivide_mullhi_s32_vector(numers, _mm_set1_epi32(denom->magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm_add_epi32(q, _mm_sub_epi32(_mm_xor_si128(numers, sign), sign)); + } + // q >>= shift + q = _mm_srai_epi32(q, more & LIBDIVIDE_32_SHIFT_MASK); + q = _mm_add_epi32(q, _mm_srli_epi32(q, 31)); // q += (q < 0) + return q; + } +} + +__m128i libdivide_s32_branchfree_do_vector(__m128i numers, const struct libdivide_s32_branchfree_t *denom) { + int32_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_32_SHIFT_MASK; + // must be arithmetic shift + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + __m128i q = libdivide_mullhi_s32_vector(numers, _mm_set1_epi32(magic)); + q = _mm_add_epi32(q, numers); // q += numers + + // If q is non-negative, we have nothing to do + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2 + uint32_t is_power_of_2 = (magic == 0); + __m128i q_sign = _mm_srai_epi32(q, 31); // q_sign = q >> 31 + __m128i mask = _mm_set1_epi32((1U << shift) - is_power_of_2); + q = _mm_add_epi32(q, _mm_and_si128(q_sign, mask)); // q = q + (q_sign & mask) + q = _mm_srai_epi32(q, shift); // q >>= shift + q = _mm_sub_epi32(_mm_xor_si128(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +////////// SINT64 + +__m128i libdivide_s64_do_vector(__m128i numers, const struct libdivide_s64_t *denom) { + uint8_t more = denom->more; + int64_t magic = denom->magic; + if (magic == 0) { // shift path + uint32_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + uint64_t mask = (1ULL << shift) - 1; + __m128i roundToZeroTweak = _mm_set1_epi64x(mask); + // q = numer + ((numer >> 63) & roundToZeroTweak); + __m128i q = _mm_add_epi64(numers, _mm_and_si128(libdivide_s64_signbits(numers), roundToZeroTweak)); + q = libdivide_s64_shift_right_vector(q, shift); + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + // q = (q ^ sign) - sign; + q = _mm_sub_epi64(_mm_xor_si128(q, sign), sign); + return q; + } + else { + __m128i q = libdivide_mullhi_s64_vector(numers, _mm_set1_epi64x(magic)); + if (more & LIBDIVIDE_ADD_MARKER) { + // must be arithmetic shift + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + // q += ((numer ^ sign) - sign); + q = _mm_add_epi64(q, _mm_sub_epi64(_mm_xor_si128(numers, sign), sign)); + } + // q >>= denom->mult_path.shift + q = libdivide_s64_shift_right_vector(q, more & LIBDIVIDE_64_SHIFT_MASK); + q = _mm_add_epi64(q, _mm_srli_epi64(q, 63)); // q += (q < 0) + return q; + } +} + +__m128i libdivide_s64_branchfree_do_vector(__m128i numers, const struct libdivide_s64_branchfree_t *denom) { + int64_t magic = denom->magic; + uint8_t more = denom->more; + uint8_t shift = more & LIBDIVIDE_64_SHIFT_MASK; + // must be arithmetic shift + __m128i sign = _mm_set1_epi32((int8_t)more >> 7); + + // libdivide_mullhi_s64(numers, magic); + __m128i q = libdivide_mullhi_s64_vector(numers, _mm_set1_epi64x(magic)); + q = _mm_add_epi64(q, numers); // q += numers + + // If q is non-negative, we have nothing to do. + // If q is negative, we want to add either (2**shift)-1 if d is + // a power of 2, or (2**shift) if it is not a power of 2. + uint32_t is_power_of_2 = (magic == 0); + __m128i q_sign = libdivide_s64_signbits(q); // q_sign = q >> 63 + __m128i mask = _mm_set1_epi64x((1ULL << shift) - is_power_of_2); + q = _mm_add_epi64(q, _mm_and_si128(q_sign, mask)); // q = q + (q_sign & mask) + q = libdivide_s64_shift_right_vector(q, shift); // q >>= shift + q = _mm_sub_epi64(_mm_xor_si128(q, sign), sign); // q = (q ^ sign) - sign + return q; +} + +#endif + +/////////// C++ stuff + +#ifdef __cplusplus + +// The C++ divider class is templated on both an integer type +// (like uint64_t) and an algorithm type. +// * BRANCHFULL is the default algorithm type. +// * BRANCHFREE is the branchfree algorithm type. +enum { + BRANCHFULL, + BRANCHFREE +}; + +#if defined(LIBDIVIDE_AVX512) + #define LIBDIVIDE_VECTOR_TYPE __m512i +#elif defined(LIBDIVIDE_AVX2) + #define LIBDIVIDE_VECTOR_TYPE __m256i +#elif defined(LIBDIVIDE_SSE2) + #define LIBDIVIDE_VECTOR_TYPE __m128i +#endif + +#if !defined(LIBDIVIDE_VECTOR_TYPE) + #define LIBDIVIDE_DIVIDE_VECTOR(ALGO) +#else + #define LIBDIVIDE_DIVIDE_VECTOR(ALGO) \ + LIBDIVIDE_VECTOR_TYPE divide(LIBDIVIDE_VECTOR_TYPE n) const { \ + return libdivide_##ALGO##_do_vector(n, &denom); \ + } +#endif + +// The DISPATCHER_GEN() macro generates C++ methods (for the given integer +// and algorithm types) that redirect to libdivide's C API. +#define DISPATCHER_GEN(T, ALGO) \ + libdivide_##ALGO##_t denom; \ + dispatcher() { } \ + dispatcher(T d) \ + : denom(libdivide_##ALGO##_gen(d)) \ + { } \ + T divide(T n) const { \ + return libdivide_##ALGO##_do(n, &denom); \ + } \ + LIBDIVIDE_DIVIDE_VECTOR(ALGO) \ + T recover() const { \ + return libdivide_##ALGO##_recover(&denom); \ + } + +// The dispatcher selects a specific division algorithm for a given +// type and ALGO using partial template specialization. +template struct dispatcher { }; + +template<> struct dispatcher { DISPATCHER_GEN(int32_t, s32) }; +template<> struct dispatcher { DISPATCHER_GEN(int32_t, s32_branchfree) }; +template<> struct dispatcher { DISPATCHER_GEN(uint32_t, u32) }; +template<> struct dispatcher { DISPATCHER_GEN(uint32_t, u32_branchfree) }; +template<> struct dispatcher { DISPATCHER_GEN(int64_t, s64) }; +template<> struct dispatcher { DISPATCHER_GEN(int64_t, s64_branchfree) }; +template<> struct dispatcher { DISPATCHER_GEN(uint64_t, u64) }; +template<> struct dispatcher { DISPATCHER_GEN(uint64_t, u64_branchfree) }; + +// This is the main divider class for use by the user (C++ API). +// The actual division algorithm is selected using the dispatcher struct +// based on the integer and algorithm template parameters. +template +class divider { +public: + // We leave the default constructor empty so that creating + // an array of dividers and then initializing them + // later doesn't slow us down. + divider() { } + + // Constructor that takes the divisor as a parameter + divider(T d) : div(d) { } + + // Divides n by the divisor + T divide(T n) const { + return div.divide(n); + } + + // Recovers the divisor, returns the value that was + // used to initialize this divider object. + T recover() const { + return div.recover(); + } + + bool operator==(const divider& other) const { + return div.denom.magic == other.denom.magic && + div.denom.more == other.denom.more; + } + + bool operator!=(const divider& other) const { + return !(*this == other); + } + +#if defined(LIBDIVIDE_VECTOR_TYPE) + // Treats the vector as packed integer values with the same type as + // the divider (e.g. s32, u32, s64, u64) and divides each of + // them by the divider, returning the packed quotients. + LIBDIVIDE_VECTOR_TYPE divide(LIBDIVIDE_VECTOR_TYPE n) const { + return div.divide(n); + } +#endif + +private: + // Storage for the actual divisor + dispatcher::value, + std::is_signed::value, sizeof(T), ALGO> div; +}; + +// Overload of operator / for scalar division +template +T operator/(T n, const divider& div) { + return div.divide(n); +} + +// Overload of operator /= for scalar division +template +T& operator/=(T& n, const divider& div) { + n = div.divide(n); + return n; +} + +#if defined(LIBDIVIDE_VECTOR_TYPE) + // Overload of operator / for vector division + template + LIBDIVIDE_VECTOR_TYPE operator/(LIBDIVIDE_VECTOR_TYPE n, const divider& div) { + return div.divide(n); + } + // Overload of operator /= for vector division + template + LIBDIVIDE_VECTOR_TYPE& operator/=(LIBDIVIDE_VECTOR_TYPE& n, const divider& div) { + n = div.divide(n); + return n; + } +#endif + +// libdivdie::branchfree_divider +template +using branchfree_divider = divider; + +} // namespace libdivide + +#endif // __cplusplus + +#endif // NUMPY_CORE_INCLUDE_NUMPY_LIBDIVIDE_LIBDIVIDE_H_ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ufuncobject.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ufuncobject.h new file mode 100644 index 0000000000000000000000000000000000000000..169a93eb5597777b67733cb5806bca7a2d38cc56 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/ufuncobject.h @@ -0,0 +1,347 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_UFUNCOBJECT_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_UFUNCOBJECT_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * The legacy generic inner loop for a standard element-wise or + * generalized ufunc. + */ +typedef void (*PyUFuncGenericFunction) + (char **args, + npy_intp const *dimensions, + npy_intp const *strides, + void *innerloopdata); + +/* + * The most generic one-dimensional inner loop for + * a masked standard element-wise ufunc. "Masked" here means that it skips + * doing calculations on any items for which the maskptr array has a true + * value. + */ +typedef void (PyUFunc_MaskedStridedInnerLoopFunc)( + char **dataptrs, npy_intp *strides, + char *maskptr, npy_intp mask_stride, + npy_intp count, + NpyAuxData *innerloopdata); + +/* Forward declaration for the type resolver and loop selector typedefs */ +struct _tagPyUFuncObject; + +/* + * Given the operands for calling a ufunc, should determine the + * calculation input and output data types and return an inner loop function. + * This function should validate that the casting rule is being followed, + * and fail if it is not. + * + * For backwards compatibility, the regular type resolution function does not + * support auxiliary data with object semantics. The type resolution call + * which returns a masked generic function returns a standard NpyAuxData + * object, for which the NPY_AUXDATA_FREE and NPY_AUXDATA_CLONE macros + * work. + * + * ufunc: The ufunc object. + * casting: The 'casting' parameter provided to the ufunc. + * operands: An array of length (ufunc->nin + ufunc->nout), + * with the output parameters possibly NULL. + * type_tup: Either NULL, or the type_tup passed to the ufunc. + * out_dtypes: An array which should be populated with new + * references to (ufunc->nin + ufunc->nout) new + * dtypes, one for each input and output. These + * dtypes should all be in native-endian format. + * + * Should return 0 on success, -1 on failure (with exception set), + * or -2 if Py_NotImplemented should be returned. + */ +typedef int (PyUFunc_TypeResolutionFunc)( + struct _tagPyUFuncObject *ufunc, + NPY_CASTING casting, + PyArrayObject **operands, + PyObject *type_tup, + PyArray_Descr **out_dtypes); + +/* + * This is the signature for the functions that may be assigned to the + * `process_core_dims_func` field of the PyUFuncObject structure. + * Implementation of this function is optional. This function is only used + * by generalized ufuncs (i.e. those with the field `core_enabled` set to 1). + * The function is called by the ufunc during the processing of the arguments + * of a call of the ufunc. The function can check the core dimensions of the + * input and output arrays and return -1 with an exception set if any + * requirements are not satisfied. If the caller of the ufunc didn't provide + * output arrays, the core dimensions associated with the output arrays (i.e. + * those that are not also used in input arrays) will have the value -1 in + * `core_dim_sizes`. This function can replace any output core dimensions + * that are -1 with a value that is appropriate for the ufunc. + * + * Parameter Description + * --------------- ------------------------------------------------------ + * ufunc The ufunc object + * core_dim_sizes An array with length `ufunc->core_num_dim_ix`. + * The core dimensions of the arrays passed to the ufunc + * will have been set. If the caller of the ufunc didn't + * provide the output array(s), the output-only core + * dimensions will have the value -1. + * + * The function must not change any element in `core_dim_sizes` that is + * not -1 on input. Doing so will result in incorrect output from the + * ufunc, and could result in a crash of the Python interpreter. + * + * The function must return 0 on success, -1 on failure (with an exception + * set). + */ +typedef int (PyUFunc_ProcessCoreDimsFunc)( + struct _tagPyUFuncObject *ufunc, + npy_intp *core_dim_sizes); + +typedef struct _tagPyUFuncObject { + PyObject_HEAD + /* + * nin: Number of inputs + * nout: Number of outputs + * nargs: Always nin + nout (Why is it stored?) + */ + int nin, nout, nargs; + + /* + * Identity for reduction, any of PyUFunc_One, PyUFunc_Zero + * PyUFunc_MinusOne, PyUFunc_None, PyUFunc_ReorderableNone, + * PyUFunc_IdentityValue. + */ + int identity; + + /* Array of one-dimensional core loops */ + PyUFuncGenericFunction *functions; + /* Array of funcdata that gets passed into the functions */ + void *const *data; + /* The number of elements in 'functions' and 'data' */ + int ntypes; + + /* Used to be unused field 'check_return' */ + int reserved1; + + /* The name of the ufunc */ + const char *name; + + /* Array of type numbers, of size ('nargs' * 'ntypes') */ + const char *types; + + /* Documentation string */ + const char *doc; + + void *ptr; + PyObject *obj; + PyObject *userloops; + + /* generalized ufunc parameters */ + + /* 0 for scalar ufunc; 1 for generalized ufunc */ + int core_enabled; + /* number of distinct dimension names in signature */ + int core_num_dim_ix; + + /* + * dimension indices of input/output argument k are stored in + * core_dim_ixs[core_offsets[k]..core_offsets[k]+core_num_dims[k]-1] + */ + + /* numbers of core dimensions of each argument */ + int *core_num_dims; + /* + * dimension indices in a flatted form; indices + * are in the range of [0,core_num_dim_ix) + */ + int *core_dim_ixs; + /* + * positions of 1st core dimensions of each + * argument in core_dim_ixs, equivalent to cumsum(core_num_dims) + */ + int *core_offsets; + /* signature string for printing purpose */ + char *core_signature; + + /* + * A function which resolves the types and fills an array + * with the dtypes for the inputs and outputs. + */ + PyUFunc_TypeResolutionFunc *type_resolver; + + /* A dictionary to monkeypatch ufuncs */ + PyObject *dict; + + /* + * This was blocked off to be the "new" inner loop selector in 1.7, + * but this was never implemented. (This is also why the above + * selector is called the "legacy" selector.) + */ + #ifndef Py_LIMITED_API + vectorcallfunc vectorcall; + #else + void *vectorcall; + #endif + + /* Was previously the `PyUFunc_MaskedInnerLoopSelectionFunc` */ + void *reserved3; + + /* + * List of flags for each operand when ufunc is called by nditer object. + * These flags will be used in addition to the default flags for each + * operand set by nditer object. + */ + npy_uint32 *op_flags; + + /* + * List of global flags used when ufunc is called by nditer object. + * These flags will be used in addition to the default global flags + * set by nditer object. + */ + npy_uint32 iter_flags; + + /* New in NPY_API_VERSION 0x0000000D and above */ + #if NPY_FEATURE_VERSION >= NPY_1_16_API_VERSION + /* + * for each core_num_dim_ix distinct dimension names, + * the possible "frozen" size (-1 if not frozen). + */ + npy_intp *core_dim_sizes; + + /* + * for each distinct core dimension, a set of UFUNC_CORE_DIM* flags + */ + npy_uint32 *core_dim_flags; + + /* Identity for reduction, when identity == PyUFunc_IdentityValue */ + PyObject *identity_value; + #endif /* NPY_FEATURE_VERSION >= NPY_1_16_API_VERSION */ + + /* New in NPY_API_VERSION 0x0000000F and above */ + #if NPY_FEATURE_VERSION >= NPY_1_22_API_VERSION + /* New private fields related to dispatching */ + void *_dispatch_cache; + /* A PyListObject of `(tuple of DTypes, ArrayMethod/Promoter)` */ + PyObject *_loops; + #endif + #if NPY_FEATURE_VERSION >= NPY_2_1_API_VERSION + /* + * Optional function to process core dimensions of a gufunc. + */ + PyUFunc_ProcessCoreDimsFunc *process_core_dims_func; + #endif +} PyUFuncObject; + +#include "arrayobject.h" +/* Generalized ufunc; 0x0001 reserved for possible use as CORE_ENABLED */ +/* the core dimension's size will be determined by the operands. */ +#define UFUNC_CORE_DIM_SIZE_INFERRED 0x0002 +/* the core dimension may be absent */ +#define UFUNC_CORE_DIM_CAN_IGNORE 0x0004 +/* flags inferred during execution */ +#define UFUNC_CORE_DIM_MISSING 0x00040000 + + +#define UFUNC_OBJ_ISOBJECT 1 +#define UFUNC_OBJ_NEEDS_API 2 + + +#if NPY_ALLOW_THREADS +#define NPY_LOOP_BEGIN_THREADS do {if (!(loop->obj & UFUNC_OBJ_NEEDS_API)) _save = PyEval_SaveThread();} while (0); +#define NPY_LOOP_END_THREADS do {if (!(loop->obj & UFUNC_OBJ_NEEDS_API)) PyEval_RestoreThread(_save);} while (0); +#else +#define NPY_LOOP_BEGIN_THREADS +#define NPY_LOOP_END_THREADS +#endif + +/* + * UFunc has unit of 0, and the order of operations can be reordered + * This case allows reduction with multiple axes at once. + */ +#define PyUFunc_Zero 0 +/* + * UFunc has unit of 1, and the order of operations can be reordered + * This case allows reduction with multiple axes at once. + */ +#define PyUFunc_One 1 +/* + * UFunc has unit of -1, and the order of operations can be reordered + * This case allows reduction with multiple axes at once. Intended for + * bitwise_and reduction. + */ +#define PyUFunc_MinusOne 2 +/* + * UFunc has no unit, and the order of operations cannot be reordered. + * This case does not allow reduction with multiple axes at once. + */ +#define PyUFunc_None -1 +/* + * UFunc has no unit, and the order of operations can be reordered + * This case allows reduction with multiple axes at once. + */ +#define PyUFunc_ReorderableNone -2 +/* + * UFunc unit is an identity_value, and the order of operations can be reordered + * This case allows reduction with multiple axes at once. + */ +#define PyUFunc_IdentityValue -3 + + +#define UFUNC_REDUCE 0 +#define UFUNC_ACCUMULATE 1 +#define UFUNC_REDUCEAT 2 +#define UFUNC_OUTER 3 + + +typedef struct { + int nin; + int nout; + PyObject *callable; +} PyUFunc_PyFuncData; + +/* A linked-list of function information for + user-defined 1-d loops. + */ +typedef struct _loop1d_info { + PyUFuncGenericFunction func; + void *data; + int *arg_types; + struct _loop1d_info *next; + int nargs; + PyArray_Descr **arg_dtypes; +} PyUFunc_Loop1d; + + +#define UFUNC_PYVALS_NAME "UFUNC_PYVALS" + +/* + * THESE MACROS ARE DEPRECATED. + * Use npy_set_floatstatus_* in the npymath library. + */ +#define UFUNC_FPE_DIVIDEBYZERO NPY_FPE_DIVIDEBYZERO +#define UFUNC_FPE_OVERFLOW NPY_FPE_OVERFLOW +#define UFUNC_FPE_UNDERFLOW NPY_FPE_UNDERFLOW +#define UFUNC_FPE_INVALID NPY_FPE_INVALID + +#define generate_divbyzero_error() npy_set_floatstatus_divbyzero() +#define generate_overflow_error() npy_set_floatstatus_overflow() + + /* Make sure it gets defined if it isn't already */ +#ifndef UFUNC_NOFPE +/* Clear the floating point exception default of Borland C++ */ +#if defined(__BORLANDC__) +#define UFUNC_NOFPE _control87(MCW_EM, MCW_EM); +#else +#define UFUNC_NOFPE +#endif +#endif + +#include "__ufunc_api.h" + +#ifdef __cplusplus +} +#endif + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_UFUNCOBJECT_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/utils.h b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/utils.h new file mode 100644 index 0000000000000000000000000000000000000000..97f06092e54050baf3c2fc4372429cbd110429e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/_core/include/numpy/utils.h @@ -0,0 +1,37 @@ +#ifndef NUMPY_CORE_INCLUDE_NUMPY_UTILS_H_ +#define NUMPY_CORE_INCLUDE_NUMPY_UTILS_H_ + +#ifndef __COMP_NPY_UNUSED + #if defined(__GNUC__) + #define __COMP_NPY_UNUSED __attribute__ ((__unused__)) + #elif defined(__ICC) + #define __COMP_NPY_UNUSED __attribute__ ((__unused__)) + #elif defined(__clang__) + #define __COMP_NPY_UNUSED __attribute__ ((unused)) + #else + #define __COMP_NPY_UNUSED + #endif +#endif + +#if defined(__GNUC__) || defined(__ICC) || defined(__clang__) + #define NPY_DECL_ALIGNED(x) __attribute__ ((aligned (x))) +#elif defined(_MSC_VER) + #define NPY_DECL_ALIGNED(x) __declspec(align(x)) +#else + #define NPY_DECL_ALIGNED(x) +#endif + +/* Use this to tag a variable as not used. It will remove unused variable + * warning on support platforms (see __COM_NPY_UNUSED) and mangle the variable + * to avoid accidental use */ +#define NPY_UNUSED(x) __NPY_UNUSED_TAGGED ## x __COMP_NPY_UNUSED +#define NPY_EXPAND(x) x + +#define NPY_STRINGIFY(x) #x +#define NPY_TOSTRING(x) NPY_STRINGIFY(x) + +#define NPY_CAT__(a, b) a ## b +#define NPY_CAT_(a, b) NPY_CAT__(a, b) +#define NPY_CAT(a, b) NPY_CAT_(a, b) + +#endif /* NUMPY_CORE_INCLUDE_NUMPY_UTILS_H_ */ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__init__.py b/venv/lib/python3.10/site-packages/numpy/f2py/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8bf1d637ec0c394d249f4e24d2c778915306a244 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/__init__.py @@ -0,0 +1,87 @@ +"""Fortran to Python Interface Generator. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the terms +of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +__all__ = ['run_main', 'get_include'] + +import sys +import subprocess +import os +import warnings + +from numpy.exceptions import VisibleDeprecationWarning +from . import f2py2e +from . import diagnose + +run_main = f2py2e.run_main +main = f2py2e.main + + +def get_include(): + """ + Return the directory that contains the ``fortranobject.c`` and ``.h`` files. + + .. note:: + + This function is not needed when building an extension with + `numpy.distutils` directly from ``.f`` and/or ``.pyf`` files + in one go. + + Python extension modules built with f2py-generated code need to use + ``fortranobject.c`` as a source file, and include the ``fortranobject.h`` + header. This function can be used to obtain the directory containing + both of these files. + + Returns + ------- + include_path : str + Absolute path to the directory containing ``fortranobject.c`` and + ``fortranobject.h``. + + Notes + ----- + .. versionadded:: 1.21.1 + + Unless the build system you are using has specific support for f2py, + building a Python extension using a ``.pyf`` signature file is a two-step + process. For a module ``mymod``: + + * Step 1: run ``python -m numpy.f2py mymod.pyf --quiet``. This + generates ``mymodmodule.c`` and (if needed) + ``mymod-f2pywrappers.f`` files next to ``mymod.pyf``. + * Step 2: build your Python extension module. This requires the + following source files: + + * ``mymodmodule.c`` + * ``mymod-f2pywrappers.f`` (if it was generated in Step 1) + * ``fortranobject.c`` + + See Also + -------- + numpy.get_include : function that returns the numpy include directory + + """ + return os.path.join(os.path.dirname(__file__), 'src') + + +def __getattr__(attr): + + # Avoid importing things that aren't needed for building + # which might import the main numpy module + if attr == "test": + from numpy._pytesttester import PytestTester + test = PytestTester(__name__) + return test + + else: + raise AttributeError("module {!r} has no attribute " + "{!r}".format(__name__, attr)) + + +def __dir__(): + return list(globals().keys() | {"test"}) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__init__.pyi b/venv/lib/python3.10/site-packages/numpy/f2py/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9cf1247f77976664992c0403a98bf3d9c23e639b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/__init__.pyi @@ -0,0 +1,42 @@ +from _typeshed import StrOrBytesPath +import subprocess +from collections.abc import Iterable +from typing import Literal as L, overload, TypedDict, type_check_only + +__all__ = ["run_main", "get_include"] + +@type_check_only +class _F2PyDictBase(TypedDict): + csrc: list[str] + h: list[str] + +@type_check_only +class _F2PyDict(_F2PyDictBase, total=False): + fsrc: list[str] + ltx: list[str] + +def run_main(comline_list: Iterable[str]) -> dict[str, _F2PyDict]: ... + +@overload +def compile( + source: str | bytes, + modulename: str = ..., + extra_args: str | list[str] = ..., + verbose: bool = ..., + source_fn: StrOrBytesPath | None = ..., + extension: L[".f", ".f90"] = ..., + full_output: L[False] = ..., +) -> int: ... +@overload +def compile( + source: str | bytes, + modulename: str = ..., + extra_args: str | list[str] = ..., + verbose: bool = ..., + source_fn: StrOrBytesPath | None = ..., + extension: L[".f", ".f90"] = ..., + *, + full_output: L[True], +) -> subprocess.CompletedProcess[bytes]: ... + +def get_include() -> str: ... diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__main__.py b/venv/lib/python3.10/site-packages/numpy/f2py/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..936a753a2796896667aa782277be41b40af061d3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/__main__.py @@ -0,0 +1,5 @@ +# See: +# https://web.archive.org/web/20140822061353/http://cens.ioc.ee/projects/f2py2e +from numpy.f2py.f2py2e import main + +main() diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e01d9f840021c3bc25db2a5e2afe5077c9c1f86d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__main__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__main__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6254c3702f9017ca7412a38b8f480defd39e018b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__main__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__version__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__version__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..99f686a63a50e55a23bd4fbd0eb45f3dcdb53380 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/__version__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_isocbind.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_isocbind.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb604bc990529d098ee3f425d26965af88956a57 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_isocbind.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_src_pyf.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_src_pyf.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32edab1d02d7d64f214acba5aec2b4037ede4967 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/_src_pyf.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/auxfuncs.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/auxfuncs.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..249499661e31cd6909b401d8e4405311b533e7a8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/auxfuncs.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/capi_maps.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/capi_maps.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6d5d8b74ff6dcdfdcb53bc8a339aa295a50b3d4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/capi_maps.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cb_rules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cb_rules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67628c265e56a9acd7f3135ee1a9e111f7c9187c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cb_rules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cfuncs.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cfuncs.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e0565f3f74cf9193f9a5edb3c39fc2757dce6c4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/cfuncs.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/common_rules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/common_rules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ca89d5c7935be6323ca2f2600002dbfe3b8986a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/common_rules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/crackfortran.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/crackfortran.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5278cff00f22aac29b4d14c09f630b08d766a05 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/crackfortran.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/diagnose.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/diagnose.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffd33b613c7111943cb0bb818aca762cb8991697 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/diagnose.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f2py2e.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f2py2e.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c44dfb486584518dddb810f5825d9ec9301f9d3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f2py2e.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f90mod_rules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f90mod_rules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28f92b820f12c37b9b8f83fe01221d076a6bfdc0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/f90mod_rules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/func2subr.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/func2subr.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d77ab9a5be5984e9653541325548c5733c3ee1d5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/func2subr.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/rules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/rules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c53e8d59cd9d2050283f911725a0d990a4116bc2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/rules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/symbolic.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/symbolic.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70febf6fcde56efddec84ee2393cb9bdb36f078b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/symbolic.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/use_rules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/use_rules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e03faba6c4ab690bb178bca1e6ce30737692ba9d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/__pycache__/use_rules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/__version__.py b/venv/lib/python3.10/site-packages/numpy/f2py/__version__.py new file mode 100644 index 0000000000000000000000000000000000000000..e20d7c1dbb38807d248ff886e30425e7ff597299 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/__version__.py @@ -0,0 +1 @@ +from numpy.version import version diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__init__.py b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e91393c14be39b20d5e94e262e91a05052681318 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__init__.py @@ -0,0 +1,9 @@ +def f2py_build_generator(name): + if name == "meson": + from ._meson import MesonBackend + return MesonBackend + elif name == "distutils": + from ._distutils import DistutilsBackend + return DistutilsBackend + else: + raise ValueError(f"Unknown backend: {name}") diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c3b907fbbe08fdcd85ea0212c62faf614bbc6db Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_backend.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_backend.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4c18365be1986f34f026e120c92963b8db6ff7c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_backend.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_distutils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_distutils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df4185444fd7e54f06dfc605988bdceecd2904be Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_distutils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_meson.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_meson.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc7563eb7e03daa09aa2e7c0ad82f969a7b81ab0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/__pycache__/_meson.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_backend.py b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..a7d43d2587b2f4886372f44c9bac7f5b840d7612 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_backend.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class Backend(ABC): + def __init__( + self, + modulename, + sources, + extra_objects, + build_dir, + include_dirs, + library_dirs, + libraries, + define_macros, + undef_macros, + f2py_flags, + sysinfo_flags, + fc_flags, + flib_flags, + setup_flags, + remove_build_dir, + extra_dat, + ): + self.modulename = modulename + self.sources = sources + self.extra_objects = extra_objects + self.build_dir = build_dir + self.include_dirs = include_dirs + self.library_dirs = library_dirs + self.libraries = libraries + self.define_macros = define_macros + self.undef_macros = undef_macros + self.f2py_flags = f2py_flags + self.sysinfo_flags = sysinfo_flags + self.fc_flags = fc_flags + self.flib_flags = flib_flags + self.setup_flags = setup_flags + self.remove_build_dir = remove_build_dir + self.extra_dat = extra_dat + + @abstractmethod + def compile(self) -> None: + """Compile the wrapper.""" + pass diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_distutils.py b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_distutils.py new file mode 100644 index 0000000000000000000000000000000000000000..aa7680a07ff9f3cd96226989fc66762f12d4e92e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_distutils.py @@ -0,0 +1,75 @@ +from ._backend import Backend + +from numpy.distutils.core import setup, Extension +from numpy.distutils.system_info import get_info +from numpy.distutils.misc_util import dict_append +from numpy.exceptions import VisibleDeprecationWarning +import os +import sys +import shutil +import warnings + + +class DistutilsBackend(Backend): + def __init__(sef, *args, **kwargs): + warnings.warn( + "\ndistutils has been deprecated since NumPy 1.26.x\n" + "Use the Meson backend instead, or generate wrappers" + " without -c and use a custom build script", + VisibleDeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) + + def compile(self): + num_info = {} + if num_info: + self.include_dirs.extend(num_info.get("include_dirs", [])) + ext_args = { + "name": self.modulename, + "sources": self.sources, + "include_dirs": self.include_dirs, + "library_dirs": self.library_dirs, + "libraries": self.libraries, + "define_macros": self.define_macros, + "undef_macros": self.undef_macros, + "extra_objects": self.extra_objects, + "f2py_options": self.f2py_flags, + } + + if self.sysinfo_flags: + for n in self.sysinfo_flags: + i = get_info(n) + if not i: + print( + f"No {n!r} resources found" + "in system (try `f2py --help-link`)" + ) + dict_append(ext_args, **i) + + ext = Extension(**ext_args) + + sys.argv = [sys.argv[0]] + self.setup_flags + sys.argv.extend( + [ + "build", + "--build-temp", + self.build_dir, + "--build-base", + self.build_dir, + "--build-platlib", + ".", + "--disable-optimization", + ] + ) + + if self.fc_flags: + sys.argv.extend(["config_fc"] + self.fc_flags) + if self.flib_flags: + sys.argv.extend(["build_ext"] + self.flib_flags) + + setup(ext_modules=[ext]) + + if self.remove_build_dir and os.path.exists(self.build_dir): + print(f"Removing build directory {self.build_dir}") + shutil.rmtree(self.build_dir) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_meson.py b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_meson.py new file mode 100644 index 0000000000000000000000000000000000000000..9195e51f02fd8769775f5fb07caa67bb82583758 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/_meson.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import os +import errno +import shutil +import subprocess +import sys +import re +from pathlib import Path + +from ._backend import Backend +from string import Template +from itertools import chain + + + +class MesonTemplate: + """Template meson build file generation class.""" + + def __init__( + self, + modulename: str, + sources: list[Path], + deps: list[str], + libraries: list[str], + library_dirs: list[Path], + include_dirs: list[Path], + object_files: list[Path], + linker_args: list[str], + fortran_args: list[str], + build_type: str, + python_exe: str, + ): + self.modulename = modulename + self.build_template_path = ( + Path(__file__).parent.absolute() / "meson.build.template" + ) + self.sources = sources + self.deps = deps + self.libraries = libraries + self.library_dirs = library_dirs + if include_dirs is not None: + self.include_dirs = include_dirs + else: + self.include_dirs = [] + self.substitutions = {} + self.objects = object_files + # Convert args to '' wrapped variant for meson + self.fortran_args = [ + f"'{x}'" if not (x.startswith("'") and x.endswith("'")) else x + for x in fortran_args + ] + self.pipeline = [ + self.initialize_template, + self.sources_substitution, + self.deps_substitution, + self.include_substitution, + self.libraries_substitution, + self.fortran_args_substitution, + ] + self.build_type = build_type + self.python_exe = python_exe + self.indent = " " * 21 + + def meson_build_template(self) -> str: + if not self.build_template_path.is_file(): + raise FileNotFoundError( + errno.ENOENT, + "Meson build template" + f" {self.build_template_path.absolute()}" + " does not exist.", + ) + return self.build_template_path.read_text() + + def initialize_template(self) -> None: + self.substitutions["modulename"] = self.modulename + self.substitutions["buildtype"] = self.build_type + self.substitutions["python"] = self.python_exe + + def sources_substitution(self) -> None: + self.substitutions["source_list"] = ",\n".join( + [f"{self.indent}'''{source}'''," for source in self.sources] + ) + + def deps_substitution(self) -> None: + self.substitutions["dep_list"] = f",\n{self.indent}".join( + [f"{self.indent}dependency('{dep}')," for dep in self.deps] + ) + + def libraries_substitution(self) -> None: + self.substitutions["lib_dir_declarations"] = "\n".join( + [ + f"lib_dir_{i} = declare_dependency(link_args : ['''-L{lib_dir}'''])" + for i, lib_dir in enumerate(self.library_dirs) + ] + ) + + self.substitutions["lib_declarations"] = "\n".join( + [ + f"{lib.replace('.','_')} = declare_dependency(link_args : ['-l{lib}'])" + for lib in self.libraries + ] + ) + + self.substitutions["lib_list"] = f"\n{self.indent}".join( + [f"{self.indent}{lib.replace('.','_')}," for lib in self.libraries] + ) + self.substitutions["lib_dir_list"] = f"\n{self.indent}".join( + [f"{self.indent}lib_dir_{i}," for i in range(len(self.library_dirs))] + ) + + def include_substitution(self) -> None: + self.substitutions["inc_list"] = f",\n{self.indent}".join( + [f"{self.indent}'''{inc}'''," for inc in self.include_dirs] + ) + + def fortran_args_substitution(self) -> None: + if self.fortran_args: + self.substitutions["fortran_args"] = ( + f"{self.indent}fortran_args: [{', '.join(list(self.fortran_args))}]," + ) + else: + self.substitutions["fortran_args"] = "" + + def generate_meson_build(self): + for node in self.pipeline: + node() + template = Template(self.meson_build_template()) + meson_build = template.substitute(self.substitutions) + meson_build = re.sub(r",,", ",", meson_build) + return meson_build + + +class MesonBackend(Backend): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.dependencies = self.extra_dat.get("dependencies", []) + self.meson_build_dir = "bbdir" + self.build_type = ( + "debug" if any("debug" in flag for flag in self.fc_flags) else "release" + ) + self.fc_flags = _get_flags(self.fc_flags) + + def _move_exec_to_root(self, build_dir: Path): + walk_dir = Path(build_dir) / self.meson_build_dir + path_objects = chain( + walk_dir.glob(f"{self.modulename}*.so"), + walk_dir.glob(f"{self.modulename}*.pyd"), + ) + # Same behavior as distutils + # https://github.com/numpy/numpy/issues/24874#issuecomment-1835632293 + for path_object in path_objects: + dest_path = Path.cwd() / path_object.name + if dest_path.exists(): + dest_path.unlink() + shutil.copy2(path_object, dest_path) + os.remove(path_object) + + def write_meson_build(self, build_dir: Path) -> None: + """Writes the meson build file at specified location""" + meson_template = MesonTemplate( + self.modulename, + self.sources, + self.dependencies, + self.libraries, + self.library_dirs, + self.include_dirs, + self.extra_objects, + self.flib_flags, + self.fc_flags, + self.build_type, + sys.executable, + ) + src = meson_template.generate_meson_build() + Path(build_dir).mkdir(parents=True, exist_ok=True) + meson_build_file = Path(build_dir) / "meson.build" + meson_build_file.write_text(src) + return meson_build_file + + def _run_subprocess_command(self, command, cwd): + subprocess.run(command, cwd=cwd, check=True) + + def run_meson(self, build_dir: Path): + setup_command = ["meson", "setup", self.meson_build_dir] + self._run_subprocess_command(setup_command, build_dir) + compile_command = ["meson", "compile", "-C", self.meson_build_dir] + self._run_subprocess_command(compile_command, build_dir) + + def compile(self) -> None: + self.sources = _prepare_sources(self.modulename, self.sources, self.build_dir) + self.write_meson_build(self.build_dir) + self.run_meson(self.build_dir) + self._move_exec_to_root(self.build_dir) + + +def _prepare_sources(mname, sources, bdir): + extended_sources = sources.copy() + Path(bdir).mkdir(parents=True, exist_ok=True) + # Copy sources + for source in sources: + if Path(source).exists() and Path(source).is_file(): + shutil.copy(source, bdir) + generated_sources = [ + Path(f"{mname}module.c"), + Path(f"{mname}-f2pywrappers2.f90"), + Path(f"{mname}-f2pywrappers.f"), + ] + bdir = Path(bdir) + for generated_source in generated_sources: + if generated_source.exists(): + shutil.copy(generated_source, bdir / generated_source.name) + extended_sources.append(generated_source.name) + generated_source.unlink() + extended_sources = [ + Path(source).name + for source in extended_sources + if not Path(source).suffix == ".pyf" + ] + return extended_sources + + +def _get_flags(fc_flags): + flag_values = [] + flag_pattern = re.compile(r"--f(77|90)flags=(.*)") + for flag in fc_flags: + match_result = flag_pattern.match(flag) + if match_result: + values = match_result.group(2).strip().split() + values = [val.strip("'\"") for val in values] + flag_values.extend(values) + # Hacky way to preserve order of flags + unique_flags = list(dict.fromkeys(flag_values)) + return unique_flags diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_backends/meson.build.template b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/meson.build.template new file mode 100644 index 0000000000000000000000000000000000000000..fdcc1b17ce2118543266526c129d3a0a718eae63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_backends/meson.build.template @@ -0,0 +1,55 @@ +project('${modulename}', + ['c', 'fortran'], + version : '0.1', + meson_version: '>= 1.1.0', + default_options : [ + 'warning_level=1', + 'buildtype=${buildtype}' + ]) +fc = meson.get_compiler('fortran') + +py = import('python').find_installation('''${python}''', pure: false) +py_dep = py.dependency() + +incdir_numpy = run_command(py, + ['-c', 'import os; os.chdir(".."); import numpy; print(numpy.get_include())'], + check : true +).stdout().strip() + +incdir_f2py = run_command(py, + ['-c', 'import os; os.chdir(".."); import numpy.f2py; print(numpy.f2py.get_include())'], + check : true +).stdout().strip() + +inc_np = include_directories(incdir_numpy) +np_dep = declare_dependency(include_directories: inc_np) + +incdir_f2py = incdir_numpy / '..' / '..' / 'f2py' / 'src' +inc_f2py = include_directories(incdir_f2py) +fortranobject_c = incdir_f2py / 'fortranobject.c' + +inc_np = include_directories(incdir_numpy, incdir_f2py) +# gh-25000 +quadmath_dep = fc.find_library('quadmath', required: false) + +${lib_declarations} +${lib_dir_declarations} + +py.extension_module('${modulename}', + [ +${source_list}, + fortranobject_c + ], + include_directories: [ + inc_np, +${inc_list} + ], + dependencies : [ + py_dep, + quadmath_dep, +${dep_list} +${lib_list} +${lib_dir_list} + ], +${fortran_args} + install : true) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_isocbind.py b/venv/lib/python3.10/site-packages/numpy/f2py/_isocbind.py new file mode 100644 index 0000000000000000000000000000000000000000..3043c5d9163f7101d165ca08e33adf0970547612 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_isocbind.py @@ -0,0 +1,62 @@ +""" +ISO_C_BINDING maps for f2py2e. +Only required declarations/macros/functions will be used. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +# These map to keys in c2py_map, via forced casting for now, see gh-25229 +iso_c_binding_map = { + 'integer': { + 'c_int': 'int', + 'c_short': 'short', # 'short' <=> 'int' for now + 'c_long': 'long', # 'long' <=> 'int' for now + 'c_long_long': 'long_long', + 'c_signed_char': 'signed_char', + 'c_size_t': 'unsigned', # size_t <=> 'unsigned' for now + 'c_int8_t': 'signed_char', # int8_t <=> 'signed_char' for now + 'c_int16_t': 'short', # int16_t <=> 'short' for now + 'c_int32_t': 'int', # int32_t <=> 'int' for now + 'c_int64_t': 'long_long', + 'c_int_least8_t': 'signed_char', # int_least8_t <=> 'signed_char' for now + 'c_int_least16_t': 'short', # int_least16_t <=> 'short' for now + 'c_int_least32_t': 'int', # int_least32_t <=> 'int' for now + 'c_int_least64_t': 'long_long', + 'c_int_fast8_t': 'signed_char', # int_fast8_t <=> 'signed_char' for now + 'c_int_fast16_t': 'short', # int_fast16_t <=> 'short' for now + 'c_int_fast32_t': 'int', # int_fast32_t <=> 'int' for now + 'c_int_fast64_t': 'long_long', + 'c_intmax_t': 'long_long', # intmax_t <=> 'long_long' for now + 'c_intptr_t': 'long', # intptr_t <=> 'long' for now + 'c_ptrdiff_t': 'long', # ptrdiff_t <=> 'long' for now + }, + 'real': { + 'c_float': 'float', + 'c_double': 'double', + 'c_long_double': 'long_double' + }, + 'complex': { + 'c_float_complex': 'complex_float', + 'c_double_complex': 'complex_double', + 'c_long_double_complex': 'complex_long_double' + }, + 'logical': { + 'c_bool': 'unsigned_char' # _Bool <=> 'unsigned_char' for now + }, + 'character': { + 'c_char': 'char' + } +} + +# TODO: See gh-25229 +isoc_c2pycode_map = {} +iso_c2py_map = {} + +isoc_kindmap = {} +for fortran_type, c_type_dict in iso_c_binding_map.items(): + for c_type in c_type_dict.keys(): + isoc_kindmap[c_type] = fortran_type diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/_src_pyf.py b/venv/lib/python3.10/site-packages/numpy/f2py/_src_pyf.py new file mode 100644 index 0000000000000000000000000000000000000000..ce59a35fed3d5fd1b704e7d826a63bbc5ee76a0e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/_src_pyf.py @@ -0,0 +1,240 @@ +import os +import re + +# START OF CODE VENDORED FROM `numpy.distutils.from_template` +############################################################# +""" +process_file(filename) + + takes templated file .xxx.src and produces .xxx file where .xxx + is .pyf .f90 or .f using the following template rules: + + '<..>' denotes a template. + + All function and subroutine blocks in a source file with names that + contain '<..>' will be replicated according to the rules in '<..>'. + + The number of comma-separated words in '<..>' will determine the number of + replicates. + + '<..>' may have two different forms, named and short. For example, + + named: + where anywhere inside a block '

' will be replaced with + 'd', 's', 'z', and 'c' for each replicate of the block. + + <_c> is already defined: <_c=s,d,c,z> + <_t> is already defined: <_t=real,double precision,complex,double complex> + + short: + , a short form of the named, useful when no

appears inside + a block. + + In general, '<..>' contains a comma separated list of arbitrary + expressions. If these expression must contain a comma|leftarrow|rightarrow, + then prepend the comma|leftarrow|rightarrow with a backslash. + + If an expression matches '\\' then it will be replaced + by -th expression. + + Note that all '<..>' forms in a block must have the same number of + comma-separated entries. + + Predefined named template rules: + + + + + +""" + +routine_start_re = re.compile(r'(\n|\A)(( (\$|\*))|)\s*(subroutine|function)\b', re.I) +routine_end_re = re.compile(r'\n\s*end\s*(subroutine|function)\b.*(\n|\Z)', re.I) +function_start_re = re.compile(r'\n (\$|\*)\s*function\b', re.I) + +def parse_structure(astr): + """ Return a list of tuples for each function or subroutine each + tuple is the start and end of a subroutine or function to be + expanded. + """ + + spanlist = [] + ind = 0 + while True: + m = routine_start_re.search(astr, ind) + if m is None: + break + start = m.start() + if function_start_re.match(astr, start, m.end()): + while True: + i = astr.rfind('\n', ind, start) + if i==-1: + break + start = i + if astr[i:i+7]!='\n $': + break + start += 1 + m = routine_end_re.search(astr, m.end()) + ind = end = m and m.end()-1 or len(astr) + spanlist.append((start, end)) + return spanlist + +template_re = re.compile(r"<\s*(\w[\w\d]*)\s*>") +named_re = re.compile(r"<\s*(\w[\w\d]*)\s*=\s*(.*?)\s*>") +list_re = re.compile(r"<\s*((.*?))\s*>") + +def find_repl_patterns(astr): + reps = named_re.findall(astr) + names = {} + for rep in reps: + name = rep[0].strip() or unique_key(names) + repl = rep[1].replace(r'\,', '@comma@') + thelist = conv(repl) + names[name] = thelist + return names + +def find_and_remove_repl_patterns(astr): + names = find_repl_patterns(astr) + astr = re.subn(named_re, '', astr)[0] + return astr, names + +item_re = re.compile(r"\A\\(?P\d+)\Z") +def conv(astr): + b = astr.split(',') + l = [x.strip() for x in b] + for i in range(len(l)): + m = item_re.match(l[i]) + if m: + j = int(m.group('index')) + l[i] = l[j] + return ','.join(l) + +def unique_key(adict): + """ Obtain a unique key given a dictionary.""" + allkeys = list(adict.keys()) + done = False + n = 1 + while not done: + newkey = '__l%s' % (n) + if newkey in allkeys: + n += 1 + else: + done = True + return newkey + + +template_name_re = re.compile(r'\A\s*(\w[\w\d]*)\s*\Z') +def expand_sub(substr, names): + substr = substr.replace(r'\>', '@rightarrow@') + substr = substr.replace(r'\<', '@leftarrow@') + lnames = find_repl_patterns(substr) + substr = named_re.sub(r"<\1>", substr) # get rid of definition templates + + def listrepl(mobj): + thelist = conv(mobj.group(1).replace(r'\,', '@comma@')) + if template_name_re.match(thelist): + return "<%s>" % (thelist) + name = None + for key in lnames.keys(): # see if list is already in dictionary + if lnames[key] == thelist: + name = key + if name is None: # this list is not in the dictionary yet + name = unique_key(lnames) + lnames[name] = thelist + return "<%s>" % name + + substr = list_re.sub(listrepl, substr) # convert all lists to named templates + # newnames are constructed as needed + + numsubs = None + base_rule = None + rules = {} + for r in template_re.findall(substr): + if r not in rules: + thelist = lnames.get(r, names.get(r, None)) + if thelist is None: + raise ValueError('No replicates found for <%s>' % (r)) + if r not in names and not thelist.startswith('_'): + names[r] = thelist + rule = [i.replace('@comma@', ',') for i in thelist.split(',')] + num = len(rule) + + if numsubs is None: + numsubs = num + rules[r] = rule + base_rule = r + elif num == numsubs: + rules[r] = rule + else: + print("Mismatch in number of replacements (base <{}={}>) " + "for <{}={}>. Ignoring.".format(base_rule, ','.join(rules[base_rule]), r, thelist)) + if not rules: + return substr + + def namerepl(mobj): + name = mobj.group(1) + return rules.get(name, (k+1)*[name])[k] + + newstr = '' + for k in range(numsubs): + newstr += template_re.sub(namerepl, substr) + '\n\n' + + newstr = newstr.replace('@rightarrow@', '>') + newstr = newstr.replace('@leftarrow@', '<') + return newstr + +def process_str(allstr): + newstr = allstr + writestr = '' + + struct = parse_structure(newstr) + + oldend = 0 + names = {} + names.update(_special_names) + for sub in struct: + cleanedstr, defs = find_and_remove_repl_patterns(newstr[oldend:sub[0]]) + writestr += cleanedstr + names.update(defs) + writestr += expand_sub(newstr[sub[0]:sub[1]], names) + oldend = sub[1] + writestr += newstr[oldend:] + + return writestr + +include_src_re = re.compile(r"(\n|\A)\s*include\s*['\"](?P[\w\d./\\]+\.src)['\"]", re.I) + +def resolve_includes(source): + d = os.path.dirname(source) + with open(source) as fid: + lines = [] + for line in fid: + m = include_src_re.match(line) + if m: + fn = m.group('name') + if not os.path.isabs(fn): + fn = os.path.join(d, fn) + if os.path.isfile(fn): + lines.extend(resolve_includes(fn)) + else: + lines.append(line) + else: + lines.append(line) + return lines + +def process_file(source): + lines = resolve_includes(source) + return process_str(''.join(lines)) + +_special_names = find_repl_patterns(''' +<_c=s,d,c,z> +<_t=real,double precision,complex,double complex> + + + + + +''') + +# END OF CODE VENDORED FROM `numpy.distutils.from_template` +########################################################### diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/auxfuncs.py b/venv/lib/python3.10/site-packages/numpy/f2py/auxfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..e926a52d1b51fdfd62a7d3cbe8311c63d7f9e22f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/auxfuncs.py @@ -0,0 +1,1000 @@ +""" +Auxiliary functions for f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy (BSD style) LICENSE. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +import pprint +import sys +import re +import types +from functools import reduce + +from . import __version__ +from . import cfuncs +from .cfuncs import errmess + +__all__ = [ + 'applyrules', 'debugcapi', 'dictappend', 'errmess', 'gentitle', + 'getargs2', 'getcallprotoargument', 'getcallstatement', + 'getfortranname', 'getpymethoddef', 'getrestdoc', 'getusercode', + 'getusercode1', 'getdimension', 'hasbody', 'hascallstatement', 'hascommon', + 'hasexternals', 'hasinitvalue', 'hasnote', 'hasresultnote', + 'isallocatable', 'isarray', 'isarrayofstrings', + 'ischaracter', 'ischaracterarray', 'ischaracter_or_characterarray', + 'iscomplex', 'iscstyledirective', + 'iscomplexarray', 'iscomplexfunction', 'iscomplexfunction_warn', + 'isdouble', 'isdummyroutine', 'isexternal', 'isfunction', + 'isfunction_wrap', 'isint1', 'isint1array', 'isinteger', 'isintent_aux', + 'isintent_c', 'isintent_callback', 'isintent_copy', 'isintent_dict', + 'isintent_hide', 'isintent_in', 'isintent_inout', 'isintent_inplace', + 'isintent_nothide', 'isintent_out', 'isintent_overwrite', 'islogical', + 'islogicalfunction', 'islong_complex', 'islong_double', + 'islong_doublefunction', 'islong_long', 'islong_longfunction', + 'ismodule', 'ismoduleroutine', 'isoptional', 'isprivate', 'isvariable', + 'isrequired', 'isroutine', 'isscalar', 'issigned_long_longarray', + 'isstring', 'isstringarray', 'isstring_or_stringarray', 'isstringfunction', + 'issubroutine', 'get_f2py_modulename', 'issubroutine_wrap', 'isthreadsafe', + 'isunsigned', 'isunsigned_char', 'isunsigned_chararray', + 'isunsigned_long_long', 'isunsigned_long_longarray', 'isunsigned_short', + 'isunsigned_shortarray', 'l_and', 'l_not', 'l_or', 'outmess', 'replace', + 'show', 'stripcomma', 'throw_error', 'isattr_value', 'getuseblocks', + 'process_f2cmap_dict', 'containscommon' +] + + +f2py_version = __version__.version + + +show = pprint.pprint + +options = {} +debugoptions = [] +wrapfuncs = 1 + + +def outmess(t): + if options.get('verbose', 1): + sys.stdout.write(t) + + +def debugcapi(var): + return 'capi' in debugoptions + + +def _ischaracter(var): + return 'typespec' in var and var['typespec'] == 'character' and \ + not isexternal(var) + + +def _isstring(var): + return 'typespec' in var and var['typespec'] == 'character' and \ + not isexternal(var) + + +def ischaracter_or_characterarray(var): + return _ischaracter(var) and 'charselector' not in var + + +def ischaracter(var): + return ischaracter_or_characterarray(var) and not isarray(var) + + +def ischaracterarray(var): + return ischaracter_or_characterarray(var) and isarray(var) + + +def isstring_or_stringarray(var): + return _ischaracter(var) and 'charselector' in var + + +def isstring(var): + return isstring_or_stringarray(var) and not isarray(var) + + +def isstringarray(var): + return isstring_or_stringarray(var) and isarray(var) + + +def isarrayofstrings(var): # obsolete? + # leaving out '*' for now so that `character*(*) a(m)` and `character + # a(m,*)` are treated differently. Luckily `character**` is illegal. + return isstringarray(var) and var['dimension'][-1] == '(*)' + + +def isarray(var): + return 'dimension' in var and not isexternal(var) + + +def isscalar(var): + return not (isarray(var) or isstring(var) or isexternal(var)) + + +def iscomplex(var): + return isscalar(var) and \ + var.get('typespec') in ['complex', 'double complex'] + + +def islogical(var): + return isscalar(var) and var.get('typespec') == 'logical' + + +def isinteger(var): + return isscalar(var) and var.get('typespec') == 'integer' + + +def isreal(var): + return isscalar(var) and var.get('typespec') == 'real' + + +def get_kind(var): + try: + return var['kindselector']['*'] + except KeyError: + try: + return var['kindselector']['kind'] + except KeyError: + pass + + +def isint1(var): + return var.get('typespec') == 'integer' \ + and get_kind(var) == '1' and not isarray(var) + + +def islong_long(var): + if not isscalar(var): + return 0 + if var.get('typespec') not in ['integer', 'logical']: + return 0 + return get_kind(var) == '8' + + +def isunsigned_char(var): + if not isscalar(var): + return 0 + if var.get('typespec') != 'integer': + return 0 + return get_kind(var) == '-1' + + +def isunsigned_short(var): + if not isscalar(var): + return 0 + if var.get('typespec') != 'integer': + return 0 + return get_kind(var) == '-2' + + +def isunsigned(var): + if not isscalar(var): + return 0 + if var.get('typespec') != 'integer': + return 0 + return get_kind(var) == '-4' + + +def isunsigned_long_long(var): + if not isscalar(var): + return 0 + if var.get('typespec') != 'integer': + return 0 + return get_kind(var) == '-8' + + +def isdouble(var): + if not isscalar(var): + return 0 + if not var.get('typespec') == 'real': + return 0 + return get_kind(var) == '8' + + +def islong_double(var): + if not isscalar(var): + return 0 + if not var.get('typespec') == 'real': + return 0 + return get_kind(var) == '16' + + +def islong_complex(var): + if not iscomplex(var): + return 0 + return get_kind(var) == '32' + + +def iscomplexarray(var): + return isarray(var) and \ + var.get('typespec') in ['complex', 'double complex'] + + +def isint1array(var): + return isarray(var) and var.get('typespec') == 'integer' \ + and get_kind(var) == '1' + + +def isunsigned_chararray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '-1' + + +def isunsigned_shortarray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '-2' + + +def isunsignedarray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '-4' + + +def isunsigned_long_longarray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '-8' + + +def issigned_chararray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '1' + + +def issigned_shortarray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '2' + + +def issigned_array(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '4' + + +def issigned_long_longarray(var): + return isarray(var) and var.get('typespec') in ['integer', 'logical']\ + and get_kind(var) == '8' + + +def isallocatable(var): + return 'attrspec' in var and 'allocatable' in var['attrspec'] + + +def ismutable(var): + return not ('dimension' not in var or isstring(var)) + + +def ismoduleroutine(rout): + return 'modulename' in rout + + +def ismodule(rout): + return 'block' in rout and 'module' == rout['block'] + + +def isfunction(rout): + return 'block' in rout and 'function' == rout['block'] + + +def isfunction_wrap(rout): + if isintent_c(rout): + return 0 + return wrapfuncs and isfunction(rout) and (not isexternal(rout)) + + +def issubroutine(rout): + return 'block' in rout and 'subroutine' == rout['block'] + + +def issubroutine_wrap(rout): + if isintent_c(rout): + return 0 + return issubroutine(rout) and hasassumedshape(rout) + +def isattr_value(var): + return 'value' in var.get('attrspec', []) + + +def hasassumedshape(rout): + if rout.get('hasassumedshape'): + return True + for a in rout['args']: + for d in rout['vars'].get(a, {}).get('dimension', []): + if d == ':': + rout['hasassumedshape'] = True + return True + return False + + +def requiresf90wrapper(rout): + return ismoduleroutine(rout) or hasassumedshape(rout) + + +def isroutine(rout): + return isfunction(rout) or issubroutine(rout) + + +def islogicalfunction(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return islogical(rout['vars'][a]) + return 0 + + +def islong_longfunction(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return islong_long(rout['vars'][a]) + return 0 + + +def islong_doublefunction(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return islong_double(rout['vars'][a]) + return 0 + + +def iscomplexfunction(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return iscomplex(rout['vars'][a]) + return 0 + + +def iscomplexfunction_warn(rout): + if iscomplexfunction(rout): + outmess("""\ + ************************************************************** + Warning: code with a function returning complex value + may not work correctly with your Fortran compiler. + When using GNU gcc/g77 compilers, codes should work + correctly for callbacks with: + f2py -c -DF2PY_CB_RETURNCOMPLEX + **************************************************************\n""") + return 1 + return 0 + + +def isstringfunction(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return isstring(rout['vars'][a]) + return 0 + + +def hasexternals(rout): + return 'externals' in rout and rout['externals'] + + +def isthreadsafe(rout): + return 'f2pyenhancements' in rout and \ + 'threadsafe' in rout['f2pyenhancements'] + + +def hasvariables(rout): + return 'vars' in rout and rout['vars'] + + +def isoptional(var): + return ('attrspec' in var and 'optional' in var['attrspec'] and + 'required' not in var['attrspec']) and isintent_nothide(var) + + +def isexternal(var): + return 'attrspec' in var and 'external' in var['attrspec'] + + +def getdimension(var): + dimpattern = r"\((.*?)\)" + if 'attrspec' in var.keys(): + if any('dimension' in s for s in var['attrspec']): + return [re.findall(dimpattern, v) for v in var['attrspec']][0] + + +def isrequired(var): + return not isoptional(var) and isintent_nothide(var) + + +def iscstyledirective(f2py_line): + directives = {"callstatement", "callprotoargument", "pymethoddef"} + return any(directive in f2py_line.lower() for directive in directives) + + +def isintent_in(var): + if 'intent' not in var: + return 1 + if 'hide' in var['intent']: + return 0 + if 'inplace' in var['intent']: + return 0 + if 'in' in var['intent']: + return 1 + if 'out' in var['intent']: + return 0 + if 'inout' in var['intent']: + return 0 + if 'outin' in var['intent']: + return 0 + return 1 + + +def isintent_inout(var): + return ('intent' in var and ('inout' in var['intent'] or + 'outin' in var['intent']) and 'in' not in var['intent'] and + 'hide' not in var['intent'] and 'inplace' not in var['intent']) + + +def isintent_out(var): + return 'out' in var.get('intent', []) + + +def isintent_hide(var): + return ('intent' in var and ('hide' in var['intent'] or + ('out' in var['intent'] and 'in' not in var['intent'] and + (not l_or(isintent_inout, isintent_inplace)(var))))) + + +def isintent_nothide(var): + return not isintent_hide(var) + + +def isintent_c(var): + return 'c' in var.get('intent', []) + + +def isintent_cache(var): + return 'cache' in var.get('intent', []) + + +def isintent_copy(var): + return 'copy' in var.get('intent', []) + + +def isintent_overwrite(var): + return 'overwrite' in var.get('intent', []) + + +def isintent_callback(var): + return 'callback' in var.get('intent', []) + + +def isintent_inplace(var): + return 'inplace' in var.get('intent', []) + + +def isintent_aux(var): + return 'aux' in var.get('intent', []) + + +def isintent_aligned4(var): + return 'aligned4' in var.get('intent', []) + + +def isintent_aligned8(var): + return 'aligned8' in var.get('intent', []) + + +def isintent_aligned16(var): + return 'aligned16' in var.get('intent', []) + + +isintent_dict = {isintent_in: 'INTENT_IN', isintent_inout: 'INTENT_INOUT', + isintent_out: 'INTENT_OUT', isintent_hide: 'INTENT_HIDE', + isintent_cache: 'INTENT_CACHE', + isintent_c: 'INTENT_C', isoptional: 'OPTIONAL', + isintent_inplace: 'INTENT_INPLACE', + isintent_aligned4: 'INTENT_ALIGNED4', + isintent_aligned8: 'INTENT_ALIGNED8', + isintent_aligned16: 'INTENT_ALIGNED16', + } + + +def isprivate(var): + return 'attrspec' in var and 'private' in var['attrspec'] + + +def isvariable(var): + # heuristic to find public/private declarations of filtered subroutines + if len(var) == 1 and 'attrspec' in var and \ + var['attrspec'][0] in ('public', 'private'): + is_var = False + else: + is_var = True + return is_var + +def hasinitvalue(var): + return '=' in var + + +def hasinitvalueasstring(var): + if not hasinitvalue(var): + return 0 + return var['='][0] in ['"', "'"] + + +def hasnote(var): + return 'note' in var + + +def hasresultnote(rout): + if not isfunction(rout): + return 0 + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if a in rout['vars']: + return hasnote(rout['vars'][a]) + return 0 + + +def hascommon(rout): + return 'common' in rout + + +def containscommon(rout): + if hascommon(rout): + return 1 + if hasbody(rout): + for b in rout['body']: + if containscommon(b): + return 1 + return 0 + + +def containsmodule(block): + if ismodule(block): + return 1 + if not hasbody(block): + return 0 + for b in block['body']: + if containsmodule(b): + return 1 + return 0 + + +def hasbody(rout): + return 'body' in rout + + +def hascallstatement(rout): + return getcallstatement(rout) is not None + + +def istrue(var): + return 1 + + +def isfalse(var): + return 0 + + +class F2PYError(Exception): + pass + + +class throw_error: + + def __init__(self, mess): + self.mess = mess + + def __call__(self, var): + mess = '\n\n var = %s\n Message: %s\n' % (var, self.mess) + raise F2PYError(mess) + + +def l_and(*f): + l1, l2 = 'lambda v', [] + for i in range(len(f)): + l1 = '%s,f%d=f[%d]' % (l1, i, i) + l2.append('f%d(v)' % (i)) + return eval('%s:%s' % (l1, ' and '.join(l2))) + + +def l_or(*f): + l1, l2 = 'lambda v', [] + for i in range(len(f)): + l1 = '%s,f%d=f[%d]' % (l1, i, i) + l2.append('f%d(v)' % (i)) + return eval('%s:%s' % (l1, ' or '.join(l2))) + + +def l_not(f): + return eval('lambda v,f=f:not f(v)') + + +def isdummyroutine(rout): + try: + return rout['f2pyenhancements']['fortranname'] == '' + except KeyError: + return 0 + + +def getfortranname(rout): + try: + name = rout['f2pyenhancements']['fortranname'] + if name == '': + raise KeyError + if not name: + errmess('Failed to use fortranname from %s\n' % + (rout['f2pyenhancements'])) + raise KeyError + except KeyError: + name = rout['name'] + return name + + +def getmultilineblock(rout, blockname, comment=1, counter=0): + try: + r = rout['f2pyenhancements'].get(blockname) + except KeyError: + return + if not r: + return + if counter > 0 and isinstance(r, str): + return + if isinstance(r, list): + if counter >= len(r): + return + r = r[counter] + if r[:3] == "'''": + if comment: + r = '\t/* start ' + blockname + \ + ' multiline (' + repr(counter) + ') */\n' + r[3:] + else: + r = r[3:] + if r[-3:] == "'''": + if comment: + r = r[:-3] + '\n\t/* end multiline (' + repr(counter) + ')*/' + else: + r = r[:-3] + else: + errmess("%s multiline block should end with `'''`: %s\n" + % (blockname, repr(r))) + return r + + +def getcallstatement(rout): + return getmultilineblock(rout, 'callstatement') + + +def getcallprotoargument(rout, cb_map={}): + r = getmultilineblock(rout, 'callprotoargument', comment=0) + if r: + return r + if hascallstatement(rout): + outmess( + 'warning: callstatement is defined without callprotoargument\n') + return + from .capi_maps import getctype + arg_types, arg_types2 = [], [] + if l_and(isstringfunction, l_not(isfunction_wrap))(rout): + arg_types.extend(['char*', 'size_t']) + for n in rout['args']: + var = rout['vars'][n] + if isintent_callback(var): + continue + if n in cb_map: + ctype = cb_map[n] + '_typedef' + else: + ctype = getctype(var) + if l_and(isintent_c, l_or(isscalar, iscomplex))(var): + pass + elif isstring(var): + pass + else: + if not isattr_value(var): + ctype = ctype + '*' + if (isstring(var) + or isarrayofstrings(var) # obsolete? + or isstringarray(var)): + arg_types2.append('size_t') + arg_types.append(ctype) + + proto_args = ','.join(arg_types + arg_types2) + if not proto_args: + proto_args = 'void' + return proto_args + + +def getusercode(rout): + return getmultilineblock(rout, 'usercode') + + +def getusercode1(rout): + return getmultilineblock(rout, 'usercode', counter=1) + + +def getpymethoddef(rout): + return getmultilineblock(rout, 'pymethoddef') + + +def getargs(rout): + sortargs, args = [], [] + if 'args' in rout: + args = rout['args'] + if 'sortvars' in rout: + for a in rout['sortvars']: + if a in args: + sortargs.append(a) + for a in args: + if a not in sortargs: + sortargs.append(a) + else: + sortargs = rout['args'] + return args, sortargs + + +def getargs2(rout): + sortargs, args = [], rout.get('args', []) + auxvars = [a for a in rout['vars'].keys() if isintent_aux(rout['vars'][a]) + and a not in args] + args = auxvars + args + if 'sortvars' in rout: + for a in rout['sortvars']: + if a in args: + sortargs.append(a) + for a in args: + if a not in sortargs: + sortargs.append(a) + else: + sortargs = auxvars + rout['args'] + return args, sortargs + + +def getrestdoc(rout): + if 'f2pymultilines' not in rout: + return None + k = None + if rout['block'] == 'python module': + k = rout['block'], rout['name'] + return rout['f2pymultilines'].get(k, None) + + +def gentitle(name): + ln = (80 - len(name) - 6) // 2 + return '/*%s %s %s*/' % (ln * '*', name, ln * '*') + + +def flatlist(lst): + if isinstance(lst, list): + return reduce(lambda x, y, f=flatlist: x + f(y), lst, []) + return [lst] + + +def stripcomma(s): + if s and s[-1] == ',': + return s[:-1] + return s + + +def replace(str, d, defaultsep=''): + if isinstance(d, list): + return [replace(str, _m, defaultsep) for _m in d] + if isinstance(str, list): + return [replace(_m, d, defaultsep) for _m in str] + for k in 2 * list(d.keys()): + if k == 'separatorsfor': + continue + if 'separatorsfor' in d and k in d['separatorsfor']: + sep = d['separatorsfor'][k] + else: + sep = defaultsep + if isinstance(d[k], list): + str = str.replace('#%s#' % (k), sep.join(flatlist(d[k]))) + else: + str = str.replace('#%s#' % (k), d[k]) + return str + + +def dictappend(rd, ar): + if isinstance(ar, list): + for a in ar: + rd = dictappend(rd, a) + return rd + for k in ar.keys(): + if k[0] == '_': + continue + if k in rd: + if isinstance(rd[k], str): + rd[k] = [rd[k]] + if isinstance(rd[k], list): + if isinstance(ar[k], list): + rd[k] = rd[k] + ar[k] + else: + rd[k].append(ar[k]) + elif isinstance(rd[k], dict): + if isinstance(ar[k], dict): + if k == 'separatorsfor': + for k1 in ar[k].keys(): + if k1 not in rd[k]: + rd[k][k1] = ar[k][k1] + else: + rd[k] = dictappend(rd[k], ar[k]) + else: + rd[k] = ar[k] + return rd + + +def applyrules(rules, d, var={}): + ret = {} + if isinstance(rules, list): + for r in rules: + rr = applyrules(r, d, var) + ret = dictappend(ret, rr) + if '_break' in rr: + break + return ret + if '_check' in rules and (not rules['_check'](var)): + return ret + if 'need' in rules: + res = applyrules({'needs': rules['need']}, d, var) + if 'needs' in res: + cfuncs.append_needs(res['needs']) + + for k in rules.keys(): + if k == 'separatorsfor': + ret[k] = rules[k] + continue + if isinstance(rules[k], str): + ret[k] = replace(rules[k], d) + elif isinstance(rules[k], list): + ret[k] = [] + for i in rules[k]: + ar = applyrules({k: i}, d, var) + if k in ar: + ret[k].append(ar[k]) + elif k[0] == '_': + continue + elif isinstance(rules[k], dict): + ret[k] = [] + for k1 in rules[k].keys(): + if isinstance(k1, types.FunctionType) and k1(var): + if isinstance(rules[k][k1], list): + for i in rules[k][k1]: + if isinstance(i, dict): + res = applyrules({'supertext': i}, d, var) + if 'supertext' in res: + i = res['supertext'] + else: + i = '' + ret[k].append(replace(i, d)) + else: + i = rules[k][k1] + if isinstance(i, dict): + res = applyrules({'supertext': i}, d) + if 'supertext' in res: + i = res['supertext'] + else: + i = '' + ret[k].append(replace(i, d)) + else: + errmess('applyrules: ignoring rule %s.\n' % repr(rules[k])) + if isinstance(ret[k], list): + if len(ret[k]) == 1: + ret[k] = ret[k][0] + if ret[k] == []: + del ret[k] + return ret + +_f2py_module_name_match = re.compile(r'\s*python\s*module\s*(?P[\w_]+)', + re.I).match +_f2py_user_module_name_match = re.compile(r'\s*python\s*module\s*(?P[\w_]*?' + r'__user__[\w_]*)', re.I).match + +def get_f2py_modulename(source): + name = None + with open(source) as f: + for line in f: + m = _f2py_module_name_match(line) + if m: + if _f2py_user_module_name_match(line): # skip *__user__* names + continue + name = m.group('name') + break + return name + +def getuseblocks(pymod): + all_uses = [] + for inner in pymod['body']: + for modblock in inner['body']: + if modblock.get('use'): + all_uses.extend([x for x in modblock.get("use").keys() if "__" not in x]) + return all_uses + +def process_f2cmap_dict(f2cmap_all, new_map, c2py_map, verbose = False): + """ + Update the Fortran-to-C type mapping dictionary with new mappings and + return a list of successfully mapped C types. + + This function integrates a new mapping dictionary into an existing + Fortran-to-C type mapping dictionary. It ensures that all keys are in + lowercase and validates new entries against a given C-to-Python mapping + dictionary. Redefinitions and invalid entries are reported with a warning. + + Parameters + ---------- + f2cmap_all : dict + The existing Fortran-to-C type mapping dictionary that will be updated. + It should be a dictionary of dictionaries where the main keys represent + Fortran types and the nested dictionaries map Fortran type specifiers + to corresponding C types. + + new_map : dict + A dictionary containing new type mappings to be added to `f2cmap_all`. + The structure should be similar to `f2cmap_all`, with keys representing + Fortran types and values being dictionaries of type specifiers and their + C type equivalents. + + c2py_map : dict + A dictionary used for validating the C types in `new_map`. It maps C + types to corresponding Python types and is used to ensure that the C + types specified in `new_map` are valid. + + verbose : boolean + A flag used to provide information about the types mapped + + Returns + ------- + tuple of (dict, list) + The updated Fortran-to-C type mapping dictionary and a list of + successfully mapped C types. + """ + f2cmap_mapped = [] + + new_map_lower = {} + for k, d1 in new_map.items(): + d1_lower = {k1.lower(): v1 for k1, v1 in d1.items()} + new_map_lower[k.lower()] = d1_lower + + for k, d1 in new_map_lower.items(): + if k not in f2cmap_all: + f2cmap_all[k] = {} + + for k1, v1 in d1.items(): + if v1 in c2py_map: + if k1 in f2cmap_all[k]: + outmess( + "\tWarning: redefinition of {'%s':{'%s':'%s'->'%s'}}\n" + % (k, k1, f2cmap_all[k][k1], v1) + ) + f2cmap_all[k][k1] = v1 + if verbose: + outmess('\tMapping "%s(kind=%s)" to "%s"\n' % (k, k1, v1)) + f2cmap_mapped.append(v1) + else: + if verbose: + errmess( + "\tIgnoring map {'%s':{'%s':'%s'}}: '%s' must be in %s\n" + % (k, k1, v1, v1, list(c2py_map.keys())) + ) + + return f2cmap_all, f2cmap_mapped diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/capi_maps.py b/venv/lib/python3.10/site-packages/numpy/f2py/capi_maps.py new file mode 100644 index 0000000000000000000000000000000000000000..83e5b1ba945a794f21e4c8e92caa8645ea42294f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/capi_maps.py @@ -0,0 +1,821 @@ +""" +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +from . import __version__ +f2py_version = __version__.version + +import copy +import re +import os +from .crackfortran import markoutercomma +from . import cb_rules +from ._isocbind import iso_c_binding_map, isoc_c2pycode_map, iso_c2py_map + +# The environment provided by auxfuncs.py is needed for some calls to eval. +# As the needed functions cannot be determined by static inspection of the +# code, it is safest to use import * pending a major refactoring of f2py. +from .auxfuncs import * + +__all__ = [ + 'getctype', 'getstrlength', 'getarrdims', 'getpydocsign', + 'getarrdocsign', 'getinit', 'sign2map', 'routsign2map', 'modsign2map', + 'cb_sign2map', 'cb_routsign2map', 'common_sign2map', 'process_f2cmap_dict' +] + + +depargs = [] +lcb_map = {} +lcb2_map = {} +# forced casting: mainly caused by the fact that Python or Numeric +# C/APIs do not support the corresponding C types. +c2py_map = {'double': 'float', + 'float': 'float', # forced casting + 'long_double': 'float', # forced casting + 'char': 'int', # forced casting + 'signed_char': 'int', # forced casting + 'unsigned_char': 'int', # forced casting + 'short': 'int', # forced casting + 'unsigned_short': 'int', # forced casting + 'int': 'int', # forced casting + 'long': 'int', + 'long_long': 'long', + 'unsigned': 'int', # forced casting + 'complex_float': 'complex', # forced casting + 'complex_double': 'complex', + 'complex_long_double': 'complex', # forced casting + 'string': 'string', + 'character': 'bytes', + } + +c2capi_map = {'double': 'NPY_DOUBLE', + 'float': 'NPY_FLOAT', + 'long_double': 'NPY_LONGDOUBLE', + 'char': 'NPY_BYTE', + 'unsigned_char': 'NPY_UBYTE', + 'signed_char': 'NPY_BYTE', + 'short': 'NPY_SHORT', + 'unsigned_short': 'NPY_USHORT', + 'int': 'NPY_INT', + 'unsigned': 'NPY_UINT', + 'long': 'NPY_LONG', + 'unsigned_long': 'NPY_ULONG', + 'long_long': 'NPY_LONGLONG', + 'unsigned_long_long': 'NPY_ULONGLONG', + 'complex_float': 'NPY_CFLOAT', + 'complex_double': 'NPY_CDOUBLE', + 'complex_long_double': 'NPY_CDOUBLE', + 'string': 'NPY_STRING', + 'character': 'NPY_STRING'} + +c2pycode_map = {'double': 'd', + 'float': 'f', + 'long_double': 'g', + 'char': 'b', + 'unsigned_char': 'B', + 'signed_char': 'b', + 'short': 'h', + 'unsigned_short': 'H', + 'int': 'i', + 'unsigned': 'I', + 'long': 'l', + 'unsigned_long': 'L', + 'long_long': 'q', + 'unsigned_long_long': 'Q', + 'complex_float': 'F', + 'complex_double': 'D', + 'complex_long_double': 'G', + 'string': 'S', + 'character': 'c'} + +# https://docs.python.org/3/c-api/arg.html#building-values +c2buildvalue_map = {'double': 'd', + 'float': 'f', + 'char': 'b', + 'signed_char': 'b', + 'short': 'h', + 'int': 'i', + 'long': 'l', + 'long_long': 'L', + 'complex_float': 'N', + 'complex_double': 'N', + 'complex_long_double': 'N', + 'string': 'y', + 'character': 'c'} + +f2cmap_all = {'real': {'': 'float', '4': 'float', '8': 'double', + '12': 'long_double', '16': 'long_double'}, + 'integer': {'': 'int', '1': 'signed_char', '2': 'short', + '4': 'int', '8': 'long_long', + '-1': 'unsigned_char', '-2': 'unsigned_short', + '-4': 'unsigned', '-8': 'unsigned_long_long'}, + 'complex': {'': 'complex_float', '8': 'complex_float', + '16': 'complex_double', '24': 'complex_long_double', + '32': 'complex_long_double'}, + 'complexkind': {'': 'complex_float', '4': 'complex_float', + '8': 'complex_double', '12': 'complex_long_double', + '16': 'complex_long_double'}, + 'logical': {'': 'int', '1': 'char', '2': 'short', '4': 'int', + '8': 'long_long'}, + 'double complex': {'': 'complex_double'}, + 'double precision': {'': 'double'}, + 'byte': {'': 'char'}, + } + +# Add ISO_C handling +c2pycode_map.update(isoc_c2pycode_map) +c2py_map.update(iso_c2py_map) +f2cmap_all, _ = process_f2cmap_dict(f2cmap_all, iso_c_binding_map, c2py_map) +# End ISO_C handling +f2cmap_default = copy.deepcopy(f2cmap_all) + +f2cmap_mapped = [] + +def load_f2cmap_file(f2cmap_file): + global f2cmap_all, f2cmap_mapped + + f2cmap_all = copy.deepcopy(f2cmap_default) + + if f2cmap_file is None: + # Default value + f2cmap_file = '.f2py_f2cmap' + if not os.path.isfile(f2cmap_file): + return + + # User defined additions to f2cmap_all. + # f2cmap_file must contain a dictionary of dictionaries, only. For + # example, {'real':{'low':'float'}} means that Fortran 'real(low)' is + # interpreted as C 'float'. This feature is useful for F90/95 users if + # they use PARAMETERS in type specifications. + try: + outmess('Reading f2cmap from {!r} ...\n'.format(f2cmap_file)) + with open(f2cmap_file) as f: + d = eval(f.read().lower(), {}, {}) + f2cmap_all, f2cmap_mapped = process_f2cmap_dict(f2cmap_all, d, c2py_map, True) + outmess('Successfully applied user defined f2cmap changes\n') + except Exception as msg: + errmess('Failed to apply user defined f2cmap changes: %s. Skipping.\n' % (msg)) + + +cformat_map = {'double': '%g', + 'float': '%g', + 'long_double': '%Lg', + 'char': '%d', + 'signed_char': '%d', + 'unsigned_char': '%hhu', + 'short': '%hd', + 'unsigned_short': '%hu', + 'int': '%d', + 'unsigned': '%u', + 'long': '%ld', + 'unsigned_long': '%lu', + 'long_long': '%ld', + 'complex_float': '(%g,%g)', + 'complex_double': '(%g,%g)', + 'complex_long_double': '(%Lg,%Lg)', + 'string': '\\"%s\\"', + 'character': "'%c'", + } + +# Auxiliary functions + + +def getctype(var): + """ + Determines C type + """ + ctype = 'void' + if isfunction(var): + if 'result' in var: + a = var['result'] + else: + a = var['name'] + if a in var['vars']: + return getctype(var['vars'][a]) + else: + errmess('getctype: function %s has no return value?!\n' % a) + elif issubroutine(var): + return ctype + elif ischaracter_or_characterarray(var): + return 'character' + elif isstring_or_stringarray(var): + return 'string' + elif 'typespec' in var and var['typespec'].lower() in f2cmap_all: + typespec = var['typespec'].lower() + f2cmap = f2cmap_all[typespec] + ctype = f2cmap[''] # default type + if 'kindselector' in var: + if '*' in var['kindselector']: + try: + ctype = f2cmap[var['kindselector']['*']] + except KeyError: + errmess('getctype: "%s %s %s" not supported.\n' % + (var['typespec'], '*', var['kindselector']['*'])) + elif 'kind' in var['kindselector']: + if typespec + 'kind' in f2cmap_all: + f2cmap = f2cmap_all[typespec + 'kind'] + try: + ctype = f2cmap[var['kindselector']['kind']] + except KeyError: + if typespec in f2cmap_all: + f2cmap = f2cmap_all[typespec] + try: + ctype = f2cmap[str(var['kindselector']['kind'])] + except KeyError: + errmess('getctype: "%s(kind=%s)" is mapped to C "%s" (to override define dict(%s = dict(%s="")) in %s/.f2py_f2cmap file).\n' + % (typespec, var['kindselector']['kind'], ctype, + typespec, var['kindselector']['kind'], os.getcwd())) + else: + if not isexternal(var): + errmess('getctype: No C-type found in "%s", assuming void.\n' % var) + return ctype + + +def f2cexpr(expr): + """Rewrite Fortran expression as f2py supported C expression. + + Due to the lack of a proper expression parser in f2py, this + function uses a heuristic approach that assumes that Fortran + arithmetic expressions are valid C arithmetic expressions when + mapping Fortran function calls to the corresponding C function/CPP + macros calls. + + """ + # TODO: support Fortran `len` function with optional kind parameter + expr = re.sub(r'\blen\b', 'f2py_slen', expr) + return expr + + +def getstrlength(var): + if isstringfunction(var): + if 'result' in var: + a = var['result'] + else: + a = var['name'] + if a in var['vars']: + return getstrlength(var['vars'][a]) + else: + errmess('getstrlength: function %s has no return value?!\n' % a) + if not isstring(var): + errmess( + 'getstrlength: expected a signature of a string but got: %s\n' % (repr(var))) + len = '1' + if 'charselector' in var: + a = var['charselector'] + if '*' in a: + len = a['*'] + elif 'len' in a: + len = f2cexpr(a['len']) + if re.match(r'\(\s*(\*|:)\s*\)', len) or re.match(r'(\*|:)', len): + if isintent_hide(var): + errmess('getstrlength:intent(hide): expected a string with defined length but got: %s\n' % ( + repr(var))) + len = '-1' + return len + + +def getarrdims(a, var, verbose=0): + ret = {} + if isstring(var) and not isarray(var): + ret['size'] = getstrlength(var) + ret['rank'] = '0' + ret['dims'] = '' + elif isscalar(var): + ret['size'] = '1' + ret['rank'] = '0' + ret['dims'] = '' + elif isarray(var): + dim = copy.copy(var['dimension']) + ret['size'] = '*'.join(dim) + try: + ret['size'] = repr(eval(ret['size'])) + except Exception: + pass + ret['dims'] = ','.join(dim) + ret['rank'] = repr(len(dim)) + ret['rank*[-1]'] = repr(len(dim) * [-1])[1:-1] + for i in range(len(dim)): # solve dim for dependencies + v = [] + if dim[i] in depargs: + v = [dim[i]] + else: + for va in depargs: + if re.match(r'.*?\b%s\b.*' % va, dim[i]): + v.append(va) + for va in v: + if depargs.index(va) > depargs.index(a): + dim[i] = '*' + break + ret['setdims'], i = '', -1 + for d in dim: + i = i + 1 + if d not in ['*', ':', '(*)', '(:)']: + ret['setdims'] = '%s#varname#_Dims[%d]=%s,' % ( + ret['setdims'], i, d) + if ret['setdims']: + ret['setdims'] = ret['setdims'][:-1] + ret['cbsetdims'], i = '', -1 + for d in var['dimension']: + i = i + 1 + if d not in ['*', ':', '(*)', '(:)']: + ret['cbsetdims'] = '%s#varname#_Dims[%d]=%s,' % ( + ret['cbsetdims'], i, d) + elif isintent_in(var): + outmess('getarrdims:warning: assumed shape array, using 0 instead of %r\n' + % (d)) + ret['cbsetdims'] = '%s#varname#_Dims[%d]=%s,' % ( + ret['cbsetdims'], i, 0) + elif verbose: + errmess( + 'getarrdims: If in call-back function: array argument %s must have bounded dimensions: got %s\n' % (repr(a), repr(d))) + if ret['cbsetdims']: + ret['cbsetdims'] = ret['cbsetdims'][:-1] +# if not isintent_c(var): +# var['dimension'].reverse() + return ret + + +def getpydocsign(a, var): + global lcb_map + if isfunction(var): + if 'result' in var: + af = var['result'] + else: + af = var['name'] + if af in var['vars']: + return getpydocsign(af, var['vars'][af]) + else: + errmess('getctype: function %s has no return value?!\n' % af) + return '', '' + sig, sigout = a, a + opt = '' + if isintent_in(var): + opt = 'input' + elif isintent_inout(var): + opt = 'in/output' + out_a = a + if isintent_out(var): + for k in var['intent']: + if k[:4] == 'out=': + out_a = k[4:] + break + init = '' + ctype = getctype(var) + + if hasinitvalue(var): + init, showinit = getinit(a, var) + init = ', optional\\n Default: %s' % showinit + if isscalar(var): + if isintent_inout(var): + sig = '%s : %s rank-0 array(%s,\'%s\')%s' % (a, opt, c2py_map[ctype], + c2pycode_map[ctype], init) + else: + sig = '%s : %s %s%s' % (a, opt, c2py_map[ctype], init) + sigout = '%s : %s' % (out_a, c2py_map[ctype]) + elif isstring(var): + if isintent_inout(var): + sig = '%s : %s rank-0 array(string(len=%s),\'c\')%s' % ( + a, opt, getstrlength(var), init) + else: + sig = '%s : %s string(len=%s)%s' % ( + a, opt, getstrlength(var), init) + sigout = '%s : string(len=%s)' % (out_a, getstrlength(var)) + elif isarray(var): + dim = var['dimension'] + rank = repr(len(dim)) + sig = '%s : %s rank-%s array(\'%s\') with bounds (%s)%s' % (a, opt, rank, + c2pycode_map[ + ctype], + ','.join(dim), init) + if a == out_a: + sigout = '%s : rank-%s array(\'%s\') with bounds (%s)'\ + % (a, rank, c2pycode_map[ctype], ','.join(dim)) + else: + sigout = '%s : rank-%s array(\'%s\') with bounds (%s) and %s storage'\ + % (out_a, rank, c2pycode_map[ctype], ','.join(dim), a) + elif isexternal(var): + ua = '' + if a in lcb_map and lcb_map[a] in lcb2_map and 'argname' in lcb2_map[lcb_map[a]]: + ua = lcb2_map[lcb_map[a]]['argname'] + if not ua == a: + ua = ' => %s' % ua + else: + ua = '' + sig = '%s : call-back function%s' % (a, ua) + sigout = sig + else: + errmess( + 'getpydocsign: Could not resolve docsignature for "%s".\n' % a) + return sig, sigout + + +def getarrdocsign(a, var): + ctype = getctype(var) + if isstring(var) and (not isarray(var)): + sig = '%s : rank-0 array(string(len=%s),\'c\')' % (a, + getstrlength(var)) + elif isscalar(var): + sig = '%s : rank-0 array(%s,\'%s\')' % (a, c2py_map[ctype], + c2pycode_map[ctype],) + elif isarray(var): + dim = var['dimension'] + rank = repr(len(dim)) + sig = '%s : rank-%s array(\'%s\') with bounds (%s)' % (a, rank, + c2pycode_map[ + ctype], + ','.join(dim)) + return sig + + +def getinit(a, var): + if isstring(var): + init, showinit = '""', "''" + else: + init, showinit = '', '' + if hasinitvalue(var): + init = var['='] + showinit = init + if iscomplex(var) or iscomplexarray(var): + ret = {} + + try: + v = var["="] + if ',' in v: + ret['init.r'], ret['init.i'] = markoutercomma( + v[1:-1]).split('@,@') + else: + v = eval(v, {}, {}) + ret['init.r'], ret['init.i'] = str(v.real), str(v.imag) + except Exception: + raise ValueError( + 'getinit: expected complex number `(r,i)\' but got `%s\' as initial value of %r.' % (init, a)) + if isarray(var): + init = '(capi_c.r=%s,capi_c.i=%s,capi_c)' % ( + ret['init.r'], ret['init.i']) + elif isstring(var): + if not init: + init, showinit = '""', "''" + if init[0] == "'": + init = '"%s"' % (init[1:-1].replace('"', '\\"')) + if init[0] == '"': + showinit = "'%s'" % (init[1:-1]) + return init, showinit + + +def get_elsize(var): + if isstring(var) or isstringarray(var): + elsize = getstrlength(var) + # override with user-specified length when available: + elsize = var['charselector'].get('f2py_len', elsize) + return elsize + if ischaracter(var) or ischaracterarray(var): + return '1' + # for numerical types, PyArray_New* functions ignore specified + # elsize, so we just return 1 and let elsize be determined at + # runtime, see fortranobject.c + return '1' + + +def sign2map(a, var): + """ + varname,ctype,atype + init,init.r,init.i,pytype + vardebuginfo,vardebugshowvalue,varshowvalue + varrformat + + intent + """ + out_a = a + if isintent_out(var): + for k in var['intent']: + if k[:4] == 'out=': + out_a = k[4:] + break + ret = {'varname': a, 'outvarname': out_a, 'ctype': getctype(var)} + intent_flags = [] + for f, s in isintent_dict.items(): + if f(var): + intent_flags.append('F2PY_%s' % s) + if intent_flags: + # TODO: Evaluate intent_flags here. + ret['intent'] = '|'.join(intent_flags) + else: + ret['intent'] = 'F2PY_INTENT_IN' + if isarray(var): + ret['varrformat'] = 'N' + elif ret['ctype'] in c2buildvalue_map: + ret['varrformat'] = c2buildvalue_map[ret['ctype']] + else: + ret['varrformat'] = 'O' + ret['init'], ret['showinit'] = getinit(a, var) + if hasinitvalue(var) and iscomplex(var) and not isarray(var): + ret['init.r'], ret['init.i'] = markoutercomma( + ret['init'][1:-1]).split('@,@') + if isexternal(var): + ret['cbnamekey'] = a + if a in lcb_map: + ret['cbname'] = lcb_map[a] + ret['maxnofargs'] = lcb2_map[lcb_map[a]]['maxnofargs'] + ret['nofoptargs'] = lcb2_map[lcb_map[a]]['nofoptargs'] + ret['cbdocstr'] = lcb2_map[lcb_map[a]]['docstr'] + ret['cblatexdocstr'] = lcb2_map[lcb_map[a]]['latexdocstr'] + else: + ret['cbname'] = a + errmess('sign2map: Confused: external %s is not in lcb_map%s.\n' % ( + a, list(lcb_map.keys()))) + if isstring(var): + ret['length'] = getstrlength(var) + if isarray(var): + ret = dictappend(ret, getarrdims(a, var)) + dim = copy.copy(var['dimension']) + if ret['ctype'] in c2capi_map: + ret['atype'] = c2capi_map[ret['ctype']] + ret['elsize'] = get_elsize(var) + # Debug info + if debugcapi(var): + il = [isintent_in, 'input', isintent_out, 'output', + isintent_inout, 'inoutput', isrequired, 'required', + isoptional, 'optional', isintent_hide, 'hidden', + iscomplex, 'complex scalar', + l_and(isscalar, l_not(iscomplex)), 'scalar', + isstring, 'string', isarray, 'array', + iscomplexarray, 'complex array', isstringarray, 'string array', + iscomplexfunction, 'complex function', + l_and(isfunction, l_not(iscomplexfunction)), 'function', + isexternal, 'callback', + isintent_callback, 'callback', + isintent_aux, 'auxiliary', + ] + rl = [] + for i in range(0, len(il), 2): + if il[i](var): + rl.append(il[i + 1]) + if isstring(var): + rl.append('slen(%s)=%s' % (a, ret['length'])) + if isarray(var): + ddim = ','.join( + map(lambda x, y: '%s|%s' % (x, y), var['dimension'], dim)) + rl.append('dims(%s)' % ddim) + if isexternal(var): + ret['vardebuginfo'] = 'debug-capi:%s=>%s:%s' % ( + a, ret['cbname'], ','.join(rl)) + else: + ret['vardebuginfo'] = 'debug-capi:%s %s=%s:%s' % ( + ret['ctype'], a, ret['showinit'], ','.join(rl)) + if isscalar(var): + if ret['ctype'] in cformat_map: + ret['vardebugshowvalue'] = 'debug-capi:%s=%s' % ( + a, cformat_map[ret['ctype']]) + if isstring(var): + ret['vardebugshowvalue'] = 'debug-capi:slen(%s)=%%d %s=\\"%%s\\"' % ( + a, a) + if isexternal(var): + ret['vardebugshowvalue'] = 'debug-capi:%s=%%p' % (a) + if ret['ctype'] in cformat_map: + ret['varshowvalue'] = '#name#:%s=%s' % (a, cformat_map[ret['ctype']]) + ret['showvalueformat'] = '%s' % (cformat_map[ret['ctype']]) + if isstring(var): + ret['varshowvalue'] = '#name#:slen(%s)=%%d %s=\\"%%s\\"' % (a, a) + ret['pydocsign'], ret['pydocsignout'] = getpydocsign(a, var) + if hasnote(var): + ret['note'] = var['note'] + return ret + + +def routsign2map(rout): + """ + name,NAME,begintitle,endtitle + rname,ctype,rformat + routdebugshowvalue + """ + global lcb_map + name = rout['name'] + fname = getfortranname(rout) + ret = {'name': name, + 'texname': name.replace('_', '\\_'), + 'name_lower': name.lower(), + 'NAME': name.upper(), + 'begintitle': gentitle(name), + 'endtitle': gentitle('end of %s' % name), + 'fortranname': fname, + 'FORTRANNAME': fname.upper(), + 'callstatement': getcallstatement(rout) or '', + 'usercode': getusercode(rout) or '', + 'usercode1': getusercode1(rout) or '', + } + if '_' in fname: + ret['F_FUNC'] = 'F_FUNC_US' + else: + ret['F_FUNC'] = 'F_FUNC' + if '_' in name: + ret['F_WRAPPEDFUNC'] = 'F_WRAPPEDFUNC_US' + else: + ret['F_WRAPPEDFUNC'] = 'F_WRAPPEDFUNC' + lcb_map = {} + if 'use' in rout: + for u in rout['use'].keys(): + if u in cb_rules.cb_map: + for un in cb_rules.cb_map[u]: + ln = un[0] + if 'map' in rout['use'][u]: + for k in rout['use'][u]['map'].keys(): + if rout['use'][u]['map'][k] == un[0]: + ln = k + break + lcb_map[ln] = un[1] + elif rout.get('externals'): + errmess('routsign2map: Confused: function %s has externals %s but no "use" statement.\n' % ( + ret['name'], repr(rout['externals']))) + ret['callprotoargument'] = getcallprotoargument(rout, lcb_map) or '' + if isfunction(rout): + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + ret['rname'] = a + ret['pydocsign'], ret['pydocsignout'] = getpydocsign(a, rout) + ret['ctype'] = getctype(rout['vars'][a]) + if hasresultnote(rout): + ret['resultnote'] = rout['vars'][a]['note'] + rout['vars'][a]['note'] = ['See elsewhere.'] + if ret['ctype'] in c2buildvalue_map: + ret['rformat'] = c2buildvalue_map[ret['ctype']] + else: + ret['rformat'] = 'O' + errmess('routsign2map: no c2buildvalue key for type %s\n' % + (repr(ret['ctype']))) + if debugcapi(rout): + if ret['ctype'] in cformat_map: + ret['routdebugshowvalue'] = 'debug-capi:%s=%s' % ( + a, cformat_map[ret['ctype']]) + if isstringfunction(rout): + ret['routdebugshowvalue'] = 'debug-capi:slen(%s)=%%d %s=\\"%%s\\"' % ( + a, a) + if isstringfunction(rout): + ret['rlength'] = getstrlength(rout['vars'][a]) + if ret['rlength'] == '-1': + errmess('routsign2map: expected explicit specification of the length of the string returned by the fortran function %s; taking 10.\n' % ( + repr(rout['name']))) + ret['rlength'] = '10' + if hasnote(rout): + ret['note'] = rout['note'] + rout['note'] = ['See elsewhere.'] + return ret + + +def modsign2map(m): + """ + modulename + """ + if ismodule(m): + ret = {'f90modulename': m['name'], + 'F90MODULENAME': m['name'].upper(), + 'texf90modulename': m['name'].replace('_', '\\_')} + else: + ret = {'modulename': m['name'], + 'MODULENAME': m['name'].upper(), + 'texmodulename': m['name'].replace('_', '\\_')} + ret['restdoc'] = getrestdoc(m) or [] + if hasnote(m): + ret['note'] = m['note'] + ret['usercode'] = getusercode(m) or '' + ret['usercode1'] = getusercode1(m) or '' + if m['body']: + ret['interface_usercode'] = getusercode(m['body'][0]) or '' + else: + ret['interface_usercode'] = '' + ret['pymethoddef'] = getpymethoddef(m) or '' + if 'gil_used' in m: + ret['gil_used'] = m['gil_used'] + if 'coutput' in m: + ret['coutput'] = m['coutput'] + if 'f2py_wrapper_output' in m: + ret['f2py_wrapper_output'] = m['f2py_wrapper_output'] + return ret + + +def cb_sign2map(a, var, index=None): + ret = {'varname': a} + ret['varname_i'] = ret['varname'] + ret['ctype'] = getctype(var) + if ret['ctype'] in c2capi_map: + ret['atype'] = c2capi_map[ret['ctype']] + ret['elsize'] = get_elsize(var) + if ret['ctype'] in cformat_map: + ret['showvalueformat'] = '%s' % (cformat_map[ret['ctype']]) + if isarray(var): + ret = dictappend(ret, getarrdims(a, var)) + ret['pydocsign'], ret['pydocsignout'] = getpydocsign(a, var) + if hasnote(var): + ret['note'] = var['note'] + var['note'] = ['See elsewhere.'] + return ret + + +def cb_routsign2map(rout, um): + """ + name,begintitle,endtitle,argname + ctype,rctype,maxnofargs,nofoptargs,returncptr + """ + ret = {'name': 'cb_%s_in_%s' % (rout['name'], um), + 'returncptr': ''} + if isintent_callback(rout): + if '_' in rout['name']: + F_FUNC = 'F_FUNC_US' + else: + F_FUNC = 'F_FUNC' + ret['callbackname'] = '%s(%s,%s)' \ + % (F_FUNC, + rout['name'].lower(), + rout['name'].upper(), + ) + ret['static'] = 'extern' + else: + ret['callbackname'] = ret['name'] + ret['static'] = 'static' + ret['argname'] = rout['name'] + ret['begintitle'] = gentitle(ret['name']) + ret['endtitle'] = gentitle('end of %s' % ret['name']) + ret['ctype'] = getctype(rout) + ret['rctype'] = 'void' + if ret['ctype'] == 'string': + ret['rctype'] = 'void' + else: + ret['rctype'] = ret['ctype'] + if ret['rctype'] != 'void': + if iscomplexfunction(rout): + ret['returncptr'] = """ +#ifdef F2PY_CB_RETURNCOMPLEX +return_value= +#endif +""" + else: + ret['returncptr'] = 'return_value=' + if ret['ctype'] in cformat_map: + ret['showvalueformat'] = '%s' % (cformat_map[ret['ctype']]) + if isstringfunction(rout): + ret['strlength'] = getstrlength(rout) + if isfunction(rout): + if 'result' in rout: + a = rout['result'] + else: + a = rout['name'] + if hasnote(rout['vars'][a]): + ret['note'] = rout['vars'][a]['note'] + rout['vars'][a]['note'] = ['See elsewhere.'] + ret['rname'] = a + ret['pydocsign'], ret['pydocsignout'] = getpydocsign(a, rout) + if iscomplexfunction(rout): + ret['rctype'] = """ +#ifdef F2PY_CB_RETURNCOMPLEX +#ctype# +#else +void +#endif +""" + else: + if hasnote(rout): + ret['note'] = rout['note'] + rout['note'] = ['See elsewhere.'] + nofargs = 0 + nofoptargs = 0 + if 'args' in rout and 'vars' in rout: + for a in rout['args']: + var = rout['vars'][a] + if l_or(isintent_in, isintent_inout)(var): + nofargs = nofargs + 1 + if isoptional(var): + nofoptargs = nofoptargs + 1 + ret['maxnofargs'] = repr(nofargs) + ret['nofoptargs'] = repr(nofoptargs) + if hasnote(rout) and isfunction(rout) and 'result' in rout: + ret['routnote'] = rout['note'] + rout['note'] = ['See elsewhere.'] + return ret + + +def common_sign2map(a, var): # obsolete + ret = {'varname': a, 'ctype': getctype(var)} + if isstringarray(var): + ret['ctype'] = 'char' + if ret['ctype'] in c2capi_map: + ret['atype'] = c2capi_map[ret['ctype']] + ret['elsize'] = get_elsize(var) + if ret['ctype'] in cformat_map: + ret['showvalueformat'] = '%s' % (cformat_map[ret['ctype']]) + if isarray(var): + ret = dictappend(ret, getarrdims(a, var)) + elif isstring(var): + ret['size'] = getstrlength(var) + ret['rank'] = '1' + ret['pydocsign'], ret['pydocsignout'] = getpydocsign(a, var) + if hasnote(var): + ret['note'] = var['note'] + var['note'] = ['See elsewhere.'] + # for strings this returns 0-rank but actually is 1-rank + ret['arrdocstr'] = getarrdocsign(a, var) + return ret diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/cb_rules.py b/venv/lib/python3.10/site-packages/numpy/f2py/cb_rules.py new file mode 100644 index 0000000000000000000000000000000000000000..faf8dd4013018a3fd1cfc1afdf62104ed98b16b1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/cb_rules.py @@ -0,0 +1,644 @@ +""" +Build call-back mechanism for f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +from . import __version__ +from .auxfuncs import ( + applyrules, debugcapi, dictappend, errmess, getargs, hasnote, isarray, + iscomplex, iscomplexarray, iscomplexfunction, isfunction, isintent_c, + isintent_hide, isintent_in, isintent_inout, isintent_nothide, + isintent_out, isoptional, isrequired, isscalar, isstring, + isstringfunction, issubroutine, l_and, l_not, l_or, outmess, replace, + stripcomma, throw_error +) +from . import cfuncs + +f2py_version = __version__.version + + +################## Rules for callback function ############## + +cb_routine_rules = { + 'cbtypedefs': 'typedef #rctype#(*#name#_typedef)(#optargs_td##args_td##strarglens_td##noargs#);', + 'body': """ +#begintitle# +typedef struct { + PyObject *capi; + PyTupleObject *args_capi; + int nofargs; + jmp_buf jmpbuf; +} #name#_t; + +#if defined(F2PY_THREAD_LOCAL_DECL) && !defined(F2PY_USE_PYTHON_TLS) + +static F2PY_THREAD_LOCAL_DECL #name#_t *_active_#name# = NULL; + +static #name#_t *swap_active_#name#(#name#_t *ptr) { + #name#_t *prev = _active_#name#; + _active_#name# = ptr; + return prev; +} + +static #name#_t *get_active_#name#(void) { + return _active_#name#; +} + +#else + +static #name#_t *swap_active_#name#(#name#_t *ptr) { + char *key = "__f2py_cb_#name#"; + return (#name#_t *)F2PySwapThreadLocalCallbackPtr(key, ptr); +} + +static #name#_t *get_active_#name#(void) { + char *key = "__f2py_cb_#name#"; + return (#name#_t *)F2PyGetThreadLocalCallbackPtr(key); +} + +#endif + +/*typedef #rctype#(*#name#_typedef)(#optargs_td##args_td##strarglens_td##noargs#);*/ +#static# #rctype# #callbackname# (#optargs##args##strarglens##noargs#) { + #name#_t cb_local = { NULL, NULL, 0 }; + #name#_t *cb = NULL; + PyTupleObject *capi_arglist = NULL; + PyObject *capi_return = NULL; + PyObject *capi_tmp = NULL; + PyObject *capi_arglist_list = NULL; + int capi_j,capi_i = 0; + int capi_longjmp_ok = 1; +#decl# +#ifdef F2PY_REPORT_ATEXIT +f2py_cb_start_clock(); +#endif + cb = get_active_#name#(); + if (cb == NULL) { + capi_longjmp_ok = 0; + cb = &cb_local; + } + capi_arglist = cb->args_capi; + CFUNCSMESS(\"cb:Call-back function #name# (maxnofargs=#maxnofargs#(-#nofoptargs#))\\n\"); + CFUNCSMESSPY(\"cb:#name#_capi=\",cb->capi); + if (cb->capi==NULL) { + capi_longjmp_ok = 0; + cb->capi = PyObject_GetAttrString(#modulename#_module,\"#argname#\"); + CFUNCSMESSPY(\"cb:#name#_capi=\",cb->capi); + } + if (cb->capi==NULL) { + PyErr_SetString(#modulename#_error,\"cb: Callback #argname# not defined (as an argument or module #modulename# attribute).\\n\"); + goto capi_fail; + } + if (F2PyCapsule_Check(cb->capi)) { + #name#_typedef #name#_cptr; + #name#_cptr = F2PyCapsule_AsVoidPtr(cb->capi); + #returncptr#(*#name#_cptr)(#optargs_nm##args_nm##strarglens_nm#); + #return# + } + if (capi_arglist==NULL) { + capi_longjmp_ok = 0; + capi_tmp = PyObject_GetAttrString(#modulename#_module,\"#argname#_extra_args\"); + if (capi_tmp) { + capi_arglist = (PyTupleObject *)PySequence_Tuple(capi_tmp); + Py_DECREF(capi_tmp); + if (capi_arglist==NULL) { + PyErr_SetString(#modulename#_error,\"Failed to convert #modulename#.#argname#_extra_args to tuple.\\n\"); + goto capi_fail; + } + } else { + PyErr_Clear(); + capi_arglist = (PyTupleObject *)Py_BuildValue(\"()\"); + } + } + if (capi_arglist == NULL) { + PyErr_SetString(#modulename#_error,\"Callback #argname# argument list is not set.\\n\"); + goto capi_fail; + } +#setdims# +#ifdef PYPY_VERSION +#define CAPI_ARGLIST_SETITEM(idx, value) PyList_SetItem((PyObject *)capi_arglist_list, idx, value) + capi_arglist_list = PySequence_List((PyObject *)capi_arglist); + if (capi_arglist_list == NULL) goto capi_fail; +#else +#define CAPI_ARGLIST_SETITEM(idx, value) PyTuple_SetItem((PyObject *)capi_arglist, idx, value) +#endif +#pyobjfrom# +#undef CAPI_ARGLIST_SETITEM +#ifdef PYPY_VERSION + CFUNCSMESSPY(\"cb:capi_arglist=\",capi_arglist_list); +#else + CFUNCSMESSPY(\"cb:capi_arglist=\",capi_arglist); +#endif + CFUNCSMESS(\"cb:Call-back calling Python function #argname#.\\n\"); +#ifdef F2PY_REPORT_ATEXIT +f2py_cb_start_call_clock(); +#endif +#ifdef PYPY_VERSION + capi_return = PyObject_CallObject(cb->capi,(PyObject *)capi_arglist_list); + Py_DECREF(capi_arglist_list); + capi_arglist_list = NULL; +#else + capi_return = PyObject_CallObject(cb->capi,(PyObject *)capi_arglist); +#endif +#ifdef F2PY_REPORT_ATEXIT +f2py_cb_stop_call_clock(); +#endif + CFUNCSMESSPY(\"cb:capi_return=\",capi_return); + if (capi_return == NULL) { + fprintf(stderr,\"capi_return is NULL\\n\"); + goto capi_fail; + } + if (capi_return == Py_None) { + Py_DECREF(capi_return); + capi_return = Py_BuildValue(\"()\"); + } + else if (!PyTuple_Check(capi_return)) { + capi_return = Py_BuildValue(\"(N)\",capi_return); + } + capi_j = PyTuple_Size(capi_return); + capi_i = 0; +#frompyobj# + CFUNCSMESS(\"cb:#name#:successful\\n\"); + Py_DECREF(capi_return); +#ifdef F2PY_REPORT_ATEXIT +f2py_cb_stop_clock(); +#endif + goto capi_return_pt; +capi_fail: + fprintf(stderr,\"Call-back #name# failed.\\n\"); + Py_XDECREF(capi_return); + Py_XDECREF(capi_arglist_list); + if (capi_longjmp_ok) { + longjmp(cb->jmpbuf,-1); + } +capi_return_pt: + ; +#return# +} +#endtitle# +""", + 'need': ['setjmp.h', 'CFUNCSMESS', 'F2PY_THREAD_LOCAL_DECL'], + 'maxnofargs': '#maxnofargs#', + 'nofoptargs': '#nofoptargs#', + 'docstr': """\ + def #argname#(#docsignature#): return #docreturn#\\n\\ +#docstrsigns#""", + 'latexdocstr': """ +{{}\\verb@def #argname#(#latexdocsignature#): return #docreturn#@{}} +#routnote# + +#latexdocstrsigns#""", + 'docstrshort': 'def #argname#(#docsignature#): return #docreturn#' +} +cb_rout_rules = [ + { # Init + 'separatorsfor': {'decl': '\n', + 'args': ',', 'optargs': '', 'pyobjfrom': '\n', 'freemem': '\n', + 'args_td': ',', 'optargs_td': '', + 'args_nm': ',', 'optargs_nm': '', + 'frompyobj': '\n', 'setdims': '\n', + 'docstrsigns': '\\n"\n"', + 'latexdocstrsigns': '\n', + 'latexdocstrreq': '\n', 'latexdocstropt': '\n', + 'latexdocstrout': '\n', 'latexdocstrcbs': '\n', + }, + 'decl': '/*decl*/', 'pyobjfrom': '/*pyobjfrom*/', 'frompyobj': '/*frompyobj*/', + 'args': [], 'optargs': '', 'return': '', 'strarglens': '', 'freemem': '/*freemem*/', + 'args_td': [], 'optargs_td': '', 'strarglens_td': '', + 'args_nm': [], 'optargs_nm': '', 'strarglens_nm': '', + 'noargs': '', + 'setdims': '/*setdims*/', + 'docstrsigns': '', 'latexdocstrsigns': '', + 'docstrreq': ' Required arguments:', + 'docstropt': ' Optional arguments:', + 'docstrout': ' Return objects:', + 'docstrcbs': ' Call-back functions:', + 'docreturn': '', 'docsign': '', 'docsignopt': '', + 'latexdocstrreq': '\\noindent Required arguments:', + 'latexdocstropt': '\\noindent Optional arguments:', + 'latexdocstrout': '\\noindent Return objects:', + 'latexdocstrcbs': '\\noindent Call-back functions:', + 'routnote': {hasnote: '--- #note#', l_not(hasnote): ''}, + }, { # Function + 'decl': ' #ctype# return_value = 0;', + 'frompyobj': [ + {debugcapi: ' CFUNCSMESS("cb:Getting return_value->");'}, + '''\ + if (capi_j>capi_i) { + GETSCALARFROMPYTUPLE(capi_return,capi_i++,&return_value,#ctype#, + "#ctype#_from_pyobj failed in converting return_value of" + " call-back function #name# to C #ctype#\\n"); + } else { + fprintf(stderr,"Warning: call-back function #name# did not provide" + " return value (index=%d, type=#ctype#)\\n",capi_i); + }''', + {debugcapi: + ' fprintf(stderr,"#showvalueformat#.\\n",return_value);'} + ], + 'need': ['#ctype#_from_pyobj', {debugcapi: 'CFUNCSMESS'}, 'GETSCALARFROMPYTUPLE'], + 'return': ' return return_value;', + '_check': l_and(isfunction, l_not(isstringfunction), l_not(iscomplexfunction)) + }, + { # String function + 'pyobjfrom': {debugcapi: ' fprintf(stderr,"debug-capi:cb:#name#:%d:\\n",return_value_len);'}, + 'args': '#ctype# return_value,int return_value_len', + 'args_nm': 'return_value,&return_value_len', + 'args_td': '#ctype# ,int', + 'frompyobj': [ + {debugcapi: ' CFUNCSMESS("cb:Getting return_value->\\"");'}, + """\ + if (capi_j>capi_i) { + GETSTRFROMPYTUPLE(capi_return,capi_i++,return_value,return_value_len); + } else { + fprintf(stderr,"Warning: call-back function #name# did not provide" + " return value (index=%d, type=#ctype#)\\n",capi_i); + }""", + {debugcapi: + ' fprintf(stderr,"#showvalueformat#\\".\\n",return_value);'} + ], + 'need': ['#ctype#_from_pyobj', {debugcapi: 'CFUNCSMESS'}, + 'string.h', 'GETSTRFROMPYTUPLE'], + 'return': 'return;', + '_check': isstringfunction + }, + { # Complex function + 'optargs': """ +#ifndef F2PY_CB_RETURNCOMPLEX +#ctype# *return_value +#endif +""", + 'optargs_nm': """ +#ifndef F2PY_CB_RETURNCOMPLEX +return_value +#endif +""", + 'optargs_td': """ +#ifndef F2PY_CB_RETURNCOMPLEX +#ctype# * +#endif +""", + 'decl': """ +#ifdef F2PY_CB_RETURNCOMPLEX + #ctype# return_value = {0, 0}; +#endif +""", + 'frompyobj': [ + {debugcapi: ' CFUNCSMESS("cb:Getting return_value->");'}, + """\ + if (capi_j>capi_i) { +#ifdef F2PY_CB_RETURNCOMPLEX + GETSCALARFROMPYTUPLE(capi_return,capi_i++,&return_value,#ctype#, + \"#ctype#_from_pyobj failed in converting return_value of call-back\" + \" function #name# to C #ctype#\\n\"); +#else + GETSCALARFROMPYTUPLE(capi_return,capi_i++,return_value,#ctype#, + \"#ctype#_from_pyobj failed in converting return_value of call-back\" + \" function #name# to C #ctype#\\n\"); +#endif + } else { + fprintf(stderr, + \"Warning: call-back function #name# did not provide\" + \" return value (index=%d, type=#ctype#)\\n\",capi_i); + }""", + {debugcapi: """\ +#ifdef F2PY_CB_RETURNCOMPLEX + fprintf(stderr,\"#showvalueformat#.\\n\",(return_value).r,(return_value).i); +#else + fprintf(stderr,\"#showvalueformat#.\\n\",(*return_value).r,(*return_value).i); +#endif +"""} + ], + 'return': """ +#ifdef F2PY_CB_RETURNCOMPLEX + return return_value; +#else + return; +#endif +""", + 'need': ['#ctype#_from_pyobj', {debugcapi: 'CFUNCSMESS'}, + 'string.h', 'GETSCALARFROMPYTUPLE', '#ctype#'], + '_check': iscomplexfunction + }, + {'docstrout': ' #pydocsignout#', + 'latexdocstrout': ['\\item[]{{}\\verb@#pydocsignout#@{}}', + {hasnote: '--- #note#'}], + 'docreturn': '#rname#,', + '_check': isfunction}, + {'_check': issubroutine, 'return': 'return;'} +] + +cb_arg_rules = [ + { # Doc + 'docstropt': {l_and(isoptional, isintent_nothide): ' #pydocsign#'}, + 'docstrreq': {l_and(isrequired, isintent_nothide): ' #pydocsign#'}, + 'docstrout': {isintent_out: ' #pydocsignout#'}, + 'latexdocstropt': {l_and(isoptional, isintent_nothide): ['\\item[]{{}\\verb@#pydocsign#@{}}', + {hasnote: '--- #note#'}]}, + 'latexdocstrreq': {l_and(isrequired, isintent_nothide): ['\\item[]{{}\\verb@#pydocsign#@{}}', + {hasnote: '--- #note#'}]}, + 'latexdocstrout': {isintent_out: ['\\item[]{{}\\verb@#pydocsignout#@{}}', + {l_and(hasnote, isintent_hide): '--- #note#', + l_and(hasnote, isintent_nothide): '--- See above.'}]}, + 'docsign': {l_and(isrequired, isintent_nothide): '#varname#,'}, + 'docsignopt': {l_and(isoptional, isintent_nothide): '#varname#,'}, + 'depend': '' + }, + { + 'args': { + l_and(isscalar, isintent_c): '#ctype# #varname_i#', + l_and(isscalar, l_not(isintent_c)): '#ctype# *#varname_i#_cb_capi', + isarray: '#ctype# *#varname_i#', + isstring: '#ctype# #varname_i#' + }, + 'args_nm': { + l_and(isscalar, isintent_c): '#varname_i#', + l_and(isscalar, l_not(isintent_c)): '#varname_i#_cb_capi', + isarray: '#varname_i#', + isstring: '#varname_i#' + }, + 'args_td': { + l_and(isscalar, isintent_c): '#ctype#', + l_and(isscalar, l_not(isintent_c)): '#ctype# *', + isarray: '#ctype# *', + isstring: '#ctype#' + }, + 'need': {l_or(isscalar, isarray, isstring): '#ctype#'}, + # untested with multiple args + 'strarglens': {isstring: ',int #varname_i#_cb_len'}, + 'strarglens_td': {isstring: ',int'}, # untested with multiple args + # untested with multiple args + 'strarglens_nm': {isstring: ',#varname_i#_cb_len'}, + }, + { # Scalars + 'decl': {l_not(isintent_c): ' #ctype# #varname_i#=(*#varname_i#_cb_capi);'}, + 'error': {l_and(isintent_c, isintent_out, + throw_error('intent(c,out) is forbidden for callback scalar arguments')): + ''}, + 'frompyobj': [{debugcapi: ' CFUNCSMESS("cb:Getting #varname#->");'}, + {isintent_out: + ' if (capi_j>capi_i)\n GETSCALARFROMPYTUPLE(capi_return,capi_i++,#varname_i#_cb_capi,#ctype#,"#ctype#_from_pyobj failed in converting argument #varname# of call-back function #name# to C #ctype#\\n");'}, + {l_and(debugcapi, l_and(l_not(iscomplex), isintent_c)): + ' fprintf(stderr,"#showvalueformat#.\\n",#varname_i#);'}, + {l_and(debugcapi, l_and(l_not(iscomplex), l_not( isintent_c))): + ' fprintf(stderr,"#showvalueformat#.\\n",*#varname_i#_cb_capi);'}, + {l_and(debugcapi, l_and(iscomplex, isintent_c)): + ' fprintf(stderr,"#showvalueformat#.\\n",(#varname_i#).r,(#varname_i#).i);'}, + {l_and(debugcapi, l_and(iscomplex, l_not( isintent_c))): + ' fprintf(stderr,"#showvalueformat#.\\n",(*#varname_i#_cb_capi).r,(*#varname_i#_cb_capi).i);'}, + ], + 'need': [{isintent_out: ['#ctype#_from_pyobj', 'GETSCALARFROMPYTUPLE']}, + {debugcapi: 'CFUNCSMESS'}], + '_check': isscalar + }, { + 'pyobjfrom': [{isintent_in: """\ + if (cb->nofargs>capi_i) + if (CAPI_ARGLIST_SETITEM(capi_i++,pyobj_from_#ctype#1(#varname_i#))) + goto capi_fail;"""}, + {isintent_inout: """\ + if (cb->nofargs>capi_i) + if (CAPI_ARGLIST_SETITEM(capi_i++,pyarr_from_p_#ctype#1(#varname_i#_cb_capi))) + goto capi_fail;"""}], + 'need': [{isintent_in: 'pyobj_from_#ctype#1'}, + {isintent_inout: 'pyarr_from_p_#ctype#1'}, + {iscomplex: '#ctype#'}], + '_check': l_and(isscalar, isintent_nothide), + '_optional': '' + }, { # String + 'frompyobj': [{debugcapi: ' CFUNCSMESS("cb:Getting #varname#->\\"");'}, + """ if (capi_j>capi_i) + GETSTRFROMPYTUPLE(capi_return,capi_i++,#varname_i#,#varname_i#_cb_len);""", + {debugcapi: + ' fprintf(stderr,"#showvalueformat#\\":%d:.\\n",#varname_i#,#varname_i#_cb_len);'}, + ], + 'need': ['#ctype#', 'GETSTRFROMPYTUPLE', + {debugcapi: 'CFUNCSMESS'}, 'string.h'], + '_check': l_and(isstring, isintent_out) + }, { + 'pyobjfrom': [ + {debugcapi: + (' fprintf(stderr,"debug-capi:cb:#varname#=#showvalueformat#:' + '%d:\\n",#varname_i#,#varname_i#_cb_len);')}, + {isintent_in: """\ + if (cb->nofargs>capi_i) + if (CAPI_ARGLIST_SETITEM(capi_i++,pyobj_from_#ctype#1size(#varname_i#,#varname_i#_cb_len))) + goto capi_fail;"""}, + {isintent_inout: """\ + if (cb->nofargs>capi_i) { + int #varname_i#_cb_dims[] = {#varname_i#_cb_len}; + if (CAPI_ARGLIST_SETITEM(capi_i++,pyarr_from_p_#ctype#1(#varname_i#,#varname_i#_cb_dims))) + goto capi_fail; + }"""}], + 'need': [{isintent_in: 'pyobj_from_#ctype#1size'}, + {isintent_inout: 'pyarr_from_p_#ctype#1'}], + '_check': l_and(isstring, isintent_nothide), + '_optional': '' + }, + # Array ... + { + 'decl': ' npy_intp #varname_i#_Dims[#rank#] = {#rank*[-1]#};', + 'setdims': ' #cbsetdims#;', + '_check': isarray, + '_depend': '' + }, + { + 'pyobjfrom': [{debugcapi: ' fprintf(stderr,"debug-capi:cb:#varname#\\n");'}, + {isintent_c: """\ + if (cb->nofargs>capi_i) { + /* tmp_arr will be inserted to capi_arglist_list that will be + destroyed when leaving callback function wrapper together + with tmp_arr. */ + PyArrayObject *tmp_arr = (PyArrayObject *)PyArray_New(&PyArray_Type, + #rank#,#varname_i#_Dims,#atype#,NULL,(char*)#varname_i#,#elsize#, + NPY_ARRAY_CARRAY,NULL); +""", + l_not(isintent_c): """\ + if (cb->nofargs>capi_i) { + /* tmp_arr will be inserted to capi_arglist_list that will be + destroyed when leaving callback function wrapper together + with tmp_arr. */ + PyArrayObject *tmp_arr = (PyArrayObject *)PyArray_New(&PyArray_Type, + #rank#,#varname_i#_Dims,#atype#,NULL,(char*)#varname_i#,#elsize#, + NPY_ARRAY_FARRAY,NULL); +""", + }, + """ + if (tmp_arr==NULL) + goto capi_fail; + if (CAPI_ARGLIST_SETITEM(capi_i++,(PyObject *)tmp_arr)) + goto capi_fail; +}"""], + '_check': l_and(isarray, isintent_nothide, l_or(isintent_in, isintent_inout)), + '_optional': '', + }, { + 'frompyobj': [{debugcapi: ' CFUNCSMESS("cb:Getting #varname#->");'}, + """ if (capi_j>capi_i) { + PyArrayObject *rv_cb_arr = NULL; + if ((capi_tmp = PyTuple_GetItem(capi_return,capi_i++))==NULL) goto capi_fail; + rv_cb_arr = array_from_pyobj(#atype#,#varname_i#_Dims,#rank#,F2PY_INTENT_IN""", + {isintent_c: '|F2PY_INTENT_C'}, + """,capi_tmp); + if (rv_cb_arr == NULL) { + fprintf(stderr,\"rv_cb_arr is NULL\\n\"); + goto capi_fail; + } + MEMCOPY(#varname_i#,PyArray_DATA(rv_cb_arr),PyArray_NBYTES(rv_cb_arr)); + if (capi_tmp != (PyObject *)rv_cb_arr) { + Py_DECREF(rv_cb_arr); + } + }""", + {debugcapi: ' fprintf(stderr,"<-.\\n");'}, + ], + 'need': ['MEMCOPY', {iscomplexarray: '#ctype#'}], + '_check': l_and(isarray, isintent_out) + }, { + 'docreturn': '#varname#,', + '_check': isintent_out + } +] + +################## Build call-back module ############# +cb_map = {} + + +def buildcallbacks(m): + cb_map[m['name']] = [] + for bi in m['body']: + if bi['block'] == 'interface': + for b in bi['body']: + if b: + buildcallback(b, m['name']) + else: + errmess('warning: empty body for %s\n' % (m['name'])) + + +def buildcallback(rout, um): + from . import capi_maps + + outmess(' Constructing call-back function "cb_%s_in_%s"\n' % + (rout['name'], um)) + args, depargs = getargs(rout) + capi_maps.depargs = depargs + var = rout['vars'] + vrd = capi_maps.cb_routsign2map(rout, um) + rd = dictappend({}, vrd) + cb_map[um].append([rout['name'], rd['name']]) + for r in cb_rout_rules: + if ('_check' in r and r['_check'](rout)) or ('_check' not in r): + ar = applyrules(r, vrd, rout) + rd = dictappend(rd, ar) + savevrd = {} + for i, a in enumerate(args): + vrd = capi_maps.cb_sign2map(a, var[a], index=i) + savevrd[a] = vrd + for r in cb_arg_rules: + if '_depend' in r: + continue + if '_optional' in r and isoptional(var[a]): + continue + if ('_check' in r and r['_check'](var[a])) or ('_check' not in r): + ar = applyrules(r, vrd, var[a]) + rd = dictappend(rd, ar) + if '_break' in r: + break + for a in args: + vrd = savevrd[a] + for r in cb_arg_rules: + if '_depend' in r: + continue + if ('_optional' not in r) or ('_optional' in r and isrequired(var[a])): + continue + if ('_check' in r and r['_check'](var[a])) or ('_check' not in r): + ar = applyrules(r, vrd, var[a]) + rd = dictappend(rd, ar) + if '_break' in r: + break + for a in depargs: + vrd = savevrd[a] + for r in cb_arg_rules: + if '_depend' not in r: + continue + if '_optional' in r: + continue + if ('_check' in r and r['_check'](var[a])) or ('_check' not in r): + ar = applyrules(r, vrd, var[a]) + rd = dictappend(rd, ar) + if '_break' in r: + break + if 'args' in rd and 'optargs' in rd: + if isinstance(rd['optargs'], list): + rd['optargs'] = rd['optargs'] + [""" +#ifndef F2PY_CB_RETURNCOMPLEX +, +#endif +"""] + rd['optargs_nm'] = rd['optargs_nm'] + [""" +#ifndef F2PY_CB_RETURNCOMPLEX +, +#endif +"""] + rd['optargs_td'] = rd['optargs_td'] + [""" +#ifndef F2PY_CB_RETURNCOMPLEX +, +#endif +"""] + if isinstance(rd['docreturn'], list): + rd['docreturn'] = stripcomma( + replace('#docreturn#', {'docreturn': rd['docreturn']})) + optargs = stripcomma(replace('#docsignopt#', + {'docsignopt': rd['docsignopt']} + )) + if optargs == '': + rd['docsignature'] = stripcomma( + replace('#docsign#', {'docsign': rd['docsign']})) + else: + rd['docsignature'] = replace('#docsign#[#docsignopt#]', + {'docsign': rd['docsign'], + 'docsignopt': optargs, + }) + rd['latexdocsignature'] = rd['docsignature'].replace('_', '\\_') + rd['latexdocsignature'] = rd['latexdocsignature'].replace(',', ', ') + rd['docstrsigns'] = [] + rd['latexdocstrsigns'] = [] + for k in ['docstrreq', 'docstropt', 'docstrout', 'docstrcbs']: + if k in rd and isinstance(rd[k], list): + rd['docstrsigns'] = rd['docstrsigns'] + rd[k] + k = 'latex' + k + if k in rd and isinstance(rd[k], list): + rd['latexdocstrsigns'] = rd['latexdocstrsigns'] + rd[k][0:1] +\ + ['\\begin{description}'] + rd[k][1:] +\ + ['\\end{description}'] + if 'args' not in rd: + rd['args'] = '' + rd['args_td'] = '' + rd['args_nm'] = '' + if not (rd.get('args') or rd.get('optargs') or rd.get('strarglens')): + rd['noargs'] = 'void' + + ar = applyrules(cb_routine_rules, rd) + cfuncs.callbacks[rd['name']] = ar['body'] + if isinstance(ar['need'], str): + ar['need'] = [ar['need']] + + if 'need' in rd: + for t in cfuncs.typedefs.keys(): + if t in rd['need']: + ar['need'].append(t) + + cfuncs.typedefs_generated[rd['name'] + '_typedef'] = ar['cbtypedefs'] + ar['need'].append(rd['name'] + '_typedef') + cfuncs.needs[rd['name']] = ar['need'] + + capi_maps.lcb2_map[rd['name']] = {'maxnofargs': ar['maxnofargs'], + 'nofoptargs': ar['nofoptargs'], + 'docstr': ar['docstr'], + 'latexdocstr': ar['latexdocstr'], + 'argname': rd['argname'] + } + outmess(' %s\n' % (ar['docstrshort'])) + return +################## Build call-back function ############# diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/cfuncs.py b/venv/lib/python3.10/site-packages/numpy/f2py/cfuncs.py new file mode 100644 index 0000000000000000000000000000000000000000..6856416fd04ab2c87cbec959020e8ddc1713564b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/cfuncs.py @@ -0,0 +1,1552 @@ +""" +C declarations, CPP macros, and C functions for f2py2e. +Only required declarations/macros/functions will be used. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +import sys +import copy + +from . import __version__ + +f2py_version = __version__.version + + +def errmess(s: str) -> None: + """ + Write an error message to stderr. + + This indirection is needed because sys.stderr might not always be available (see #26862). + """ + if sys.stderr is not None: + sys.stderr.write(s) + +##################### Definitions ################## + +outneeds = {'includes0': [], 'includes': [], 'typedefs': [], 'typedefs_generated': [], + 'userincludes': [], + 'cppmacros': [], 'cfuncs': [], 'callbacks': [], 'f90modhooks': [], + 'commonhooks': []} +needs = {} +includes0 = {'includes0': '/*need_includes0*/'} +includes = {'includes': '/*need_includes*/'} +userincludes = {'userincludes': '/*need_userincludes*/'} +typedefs = {'typedefs': '/*need_typedefs*/'} +typedefs_generated = {'typedefs_generated': '/*need_typedefs_generated*/'} +cppmacros = {'cppmacros': '/*need_cppmacros*/'} +cfuncs = {'cfuncs': '/*need_cfuncs*/'} +callbacks = {'callbacks': '/*need_callbacks*/'} +f90modhooks = {'f90modhooks': '/*need_f90modhooks*/', + 'initf90modhooksstatic': '/*initf90modhooksstatic*/', + 'initf90modhooksdynamic': '/*initf90modhooksdynamic*/', + } +commonhooks = {'commonhooks': '/*need_commonhooks*/', + 'initcommonhooks': '/*need_initcommonhooks*/', + } + +############ Includes ################### + +includes0['math.h'] = '#include ' +includes0['string.h'] = '#include ' +includes0['setjmp.h'] = '#include ' + +includes['arrayobject.h'] = '''#define PY_ARRAY_UNIQUE_SYMBOL PyArray_API +#include "arrayobject.h"''' +includes['npy_math.h'] = '#include "numpy/npy_math.h"' + +includes['arrayobject.h'] = '#include "fortranobject.h"' +includes['stdarg.h'] = '#include ' + +############# Type definitions ############### + +typedefs['unsigned_char'] = 'typedef unsigned char unsigned_char;' +typedefs['unsigned_short'] = 'typedef unsigned short unsigned_short;' +typedefs['unsigned_long'] = 'typedef unsigned long unsigned_long;' +typedefs['signed_char'] = 'typedef signed char signed_char;' +typedefs['long_long'] = """ +#if defined(NPY_OS_WIN32) +typedef __int64 long_long; +#else +typedef long long long_long; +typedef unsigned long long unsigned_long_long; +#endif +""" +typedefs['unsigned_long_long'] = """ +#if defined(NPY_OS_WIN32) +typedef __uint64 long_long; +#else +typedef unsigned long long unsigned_long_long; +#endif +""" +typedefs['long_double'] = """ +#ifndef _LONG_DOUBLE +typedef long double long_double; +#endif +""" +typedefs[ + 'complex_long_double'] = 'typedef struct {long double r,i;} complex_long_double;' +typedefs['complex_float'] = 'typedef struct {float r,i;} complex_float;' +typedefs['complex_double'] = 'typedef struct {double r,i;} complex_double;' +typedefs['string'] = """typedef char * string;""" +typedefs['character'] = """typedef char character;""" + + +############### CPP macros #################### +cppmacros['CFUNCSMESS'] = """ +#ifdef DEBUGCFUNCS +#define CFUNCSMESS(mess) fprintf(stderr,\"debug-capi:\"mess); +#define CFUNCSMESSPY(mess,obj) CFUNCSMESS(mess) \\ + PyObject_Print((PyObject *)obj,stderr,Py_PRINT_RAW);\\ + fprintf(stderr,\"\\n\"); +#else +#define CFUNCSMESS(mess) +#define CFUNCSMESSPY(mess,obj) +#endif +""" +cppmacros['F_FUNC'] = """ +#if defined(PREPEND_FORTRAN) +#if defined(NO_APPEND_FORTRAN) +#if defined(UPPERCASE_FORTRAN) +#define F_FUNC(f,F) _##F +#else +#define F_FUNC(f,F) _##f +#endif +#else +#if defined(UPPERCASE_FORTRAN) +#define F_FUNC(f,F) _##F##_ +#else +#define F_FUNC(f,F) _##f##_ +#endif +#endif +#else +#if defined(NO_APPEND_FORTRAN) +#if defined(UPPERCASE_FORTRAN) +#define F_FUNC(f,F) F +#else +#define F_FUNC(f,F) f +#endif +#else +#if defined(UPPERCASE_FORTRAN) +#define F_FUNC(f,F) F##_ +#else +#define F_FUNC(f,F) f##_ +#endif +#endif +#endif +#if defined(UNDERSCORE_G77) +#define F_FUNC_US(f,F) F_FUNC(f##_,F##_) +#else +#define F_FUNC_US(f,F) F_FUNC(f,F) +#endif +""" +cppmacros['F_WRAPPEDFUNC'] = """ +#if defined(PREPEND_FORTRAN) +#if defined(NO_APPEND_FORTRAN) +#if defined(UPPERCASE_FORTRAN) +#define F_WRAPPEDFUNC(f,F) _F2PYWRAP##F +#else +#define F_WRAPPEDFUNC(f,F) _f2pywrap##f +#endif +#else +#if defined(UPPERCASE_FORTRAN) +#define F_WRAPPEDFUNC(f,F) _F2PYWRAP##F##_ +#else +#define F_WRAPPEDFUNC(f,F) _f2pywrap##f##_ +#endif +#endif +#else +#if defined(NO_APPEND_FORTRAN) +#if defined(UPPERCASE_FORTRAN) +#define F_WRAPPEDFUNC(f,F) F2PYWRAP##F +#else +#define F_WRAPPEDFUNC(f,F) f2pywrap##f +#endif +#else +#if defined(UPPERCASE_FORTRAN) +#define F_WRAPPEDFUNC(f,F) F2PYWRAP##F##_ +#else +#define F_WRAPPEDFUNC(f,F) f2pywrap##f##_ +#endif +#endif +#endif +#if defined(UNDERSCORE_G77) +#define F_WRAPPEDFUNC_US(f,F) F_WRAPPEDFUNC(f##_,F##_) +#else +#define F_WRAPPEDFUNC_US(f,F) F_WRAPPEDFUNC(f,F) +#endif +""" +cppmacros['F_MODFUNC'] = """ +#if defined(F90MOD2CCONV1) /*E.g. Compaq Fortran */ +#if defined(NO_APPEND_FORTRAN) +#define F_MODFUNCNAME(m,f) $ ## m ## $ ## f +#else +#define F_MODFUNCNAME(m,f) $ ## m ## $ ## f ## _ +#endif +#endif + +#if defined(F90MOD2CCONV2) /*E.g. IBM XL Fortran, not tested though */ +#if defined(NO_APPEND_FORTRAN) +#define F_MODFUNCNAME(m,f) __ ## m ## _MOD_ ## f +#else +#define F_MODFUNCNAME(m,f) __ ## m ## _MOD_ ## f ## _ +#endif +#endif + +#if defined(F90MOD2CCONV3) /*E.g. MIPSPro Compilers */ +#if defined(NO_APPEND_FORTRAN) +#define F_MODFUNCNAME(m,f) f ## .in. ## m +#else +#define F_MODFUNCNAME(m,f) f ## .in. ## m ## _ +#endif +#endif +/* +#if defined(UPPERCASE_FORTRAN) +#define F_MODFUNC(m,M,f,F) F_MODFUNCNAME(M,F) +#else +#define F_MODFUNC(m,M,f,F) F_MODFUNCNAME(m,f) +#endif +*/ + +#define F_MODFUNC(m,f) (*(f2pymodstruct##m##.##f)) +""" +cppmacros['SWAPUNSAFE'] = """ +#define SWAP(a,b) (size_t)(a) = ((size_t)(a) ^ (size_t)(b));\\ + (size_t)(b) = ((size_t)(a) ^ (size_t)(b));\\ + (size_t)(a) = ((size_t)(a) ^ (size_t)(b)) +""" +cppmacros['SWAP'] = """ +#define SWAP(a,b,t) {\\ + t *c;\\ + c = a;\\ + a = b;\\ + b = c;} +""" +# cppmacros['ISCONTIGUOUS']='#define ISCONTIGUOUS(m) (PyArray_FLAGS(m) & +# NPY_ARRAY_C_CONTIGUOUS)' +cppmacros['PRINTPYOBJERR'] = """ +#define PRINTPYOBJERR(obj)\\ + fprintf(stderr,\"#modulename#.error is related to \");\\ + PyObject_Print((PyObject *)obj,stderr,Py_PRINT_RAW);\\ + fprintf(stderr,\"\\n\"); +""" +cppmacros['MINMAX'] = """ +#ifndef max +#define max(a,b) ((a > b) ? (a) : (b)) +#endif +#ifndef min +#define min(a,b) ((a < b) ? (a) : (b)) +#endif +#ifndef MAX +#define MAX(a,b) ((a > b) ? (a) : (b)) +#endif +#ifndef MIN +#define MIN(a,b) ((a < b) ? (a) : (b)) +#endif +""" +cppmacros['len..'] = """ +/* See fortranobject.h for definitions. The macros here are provided for BC. */ +#define rank f2py_rank +#define shape f2py_shape +#define fshape f2py_shape +#define len f2py_len +#define flen f2py_flen +#define slen f2py_slen +#define size f2py_size +""" +cppmacros['pyobj_from_char1'] = r""" +#define pyobj_from_char1(v) (PyLong_FromLong(v)) +""" +cppmacros['pyobj_from_short1'] = r""" +#define pyobj_from_short1(v) (PyLong_FromLong(v)) +""" +needs['pyobj_from_int1'] = ['signed_char'] +cppmacros['pyobj_from_int1'] = r""" +#define pyobj_from_int1(v) (PyLong_FromLong(v)) +""" +cppmacros['pyobj_from_long1'] = r""" +#define pyobj_from_long1(v) (PyLong_FromLong(v)) +""" +needs['pyobj_from_long_long1'] = ['long_long'] +cppmacros['pyobj_from_long_long1'] = """ +#ifdef HAVE_LONG_LONG +#define pyobj_from_long_long1(v) (PyLong_FromLongLong(v)) +#else +#warning HAVE_LONG_LONG is not available. Redefining pyobj_from_long_long. +#define pyobj_from_long_long1(v) (PyLong_FromLong(v)) +#endif +""" +needs['pyobj_from_long_double1'] = ['long_double'] +cppmacros['pyobj_from_long_double1'] = """ +#define pyobj_from_long_double1(v) (PyFloat_FromDouble(v))""" +cppmacros['pyobj_from_double1'] = """ +#define pyobj_from_double1(v) (PyFloat_FromDouble(v))""" +cppmacros['pyobj_from_float1'] = """ +#define pyobj_from_float1(v) (PyFloat_FromDouble(v))""" +needs['pyobj_from_complex_long_double1'] = ['complex_long_double'] +cppmacros['pyobj_from_complex_long_double1'] = """ +#define pyobj_from_complex_long_double1(v) (PyComplex_FromDoubles(v.r,v.i))""" +needs['pyobj_from_complex_double1'] = ['complex_double'] +cppmacros['pyobj_from_complex_double1'] = """ +#define pyobj_from_complex_double1(v) (PyComplex_FromDoubles(v.r,v.i))""" +needs['pyobj_from_complex_float1'] = ['complex_float'] +cppmacros['pyobj_from_complex_float1'] = """ +#define pyobj_from_complex_float1(v) (PyComplex_FromDoubles(v.r,v.i))""" +needs['pyobj_from_string1'] = ['string'] +cppmacros['pyobj_from_string1'] = """ +#define pyobj_from_string1(v) (PyUnicode_FromString((char *)v))""" +needs['pyobj_from_string1size'] = ['string'] +cppmacros['pyobj_from_string1size'] = """ +#define pyobj_from_string1size(v,len) (PyUnicode_FromStringAndSize((char *)v, len))""" +needs['TRYPYARRAYTEMPLATE'] = ['PRINTPYOBJERR'] +cppmacros['TRYPYARRAYTEMPLATE'] = """ +/* New SciPy */ +#define TRYPYARRAYTEMPLATECHAR case NPY_STRING: *(char *)(PyArray_DATA(arr))=*v; break; +#define TRYPYARRAYTEMPLATELONG case NPY_LONG: *(long *)(PyArray_DATA(arr))=*v; break; +#define TRYPYARRAYTEMPLATEOBJECT case NPY_OBJECT: PyArray_SETITEM(arr,PyArray_DATA(arr),pyobj_from_ ## ctype ## 1(*v)); break; + +#define TRYPYARRAYTEMPLATE(ctype,typecode) \\ + PyArrayObject *arr = NULL;\\ + if (!obj) return -2;\\ + if (!PyArray_Check(obj)) return -1;\\ + if (!(arr=(PyArrayObject *)obj)) {fprintf(stderr,\"TRYPYARRAYTEMPLATE:\");PRINTPYOBJERR(obj);return 0;}\\ + if (PyArray_DESCR(arr)->type==typecode) {*(ctype *)(PyArray_DATA(arr))=*v; return 1;}\\ + switch (PyArray_TYPE(arr)) {\\ + case NPY_DOUBLE: *(npy_double *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_INT: *(npy_int *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_LONG: *(npy_long *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_FLOAT: *(npy_float *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_CDOUBLE: *(npy_double *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_CFLOAT: *(npy_float *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_BOOL: *(npy_bool *)(PyArray_DATA(arr))=(*v!=0); break;\\ + case NPY_UBYTE: *(npy_ubyte *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_BYTE: *(npy_byte *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_SHORT: *(npy_short *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_USHORT: *(npy_ushort *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_UINT: *(npy_uint *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_ULONG: *(npy_ulong *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_LONGLONG: *(npy_longlong *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_ULONGLONG: *(npy_ulonglong *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_LONGDOUBLE: *(npy_longdouble *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_CLONGDOUBLE: *(npy_longdouble *)(PyArray_DATA(arr))=*v; break;\\ + case NPY_OBJECT: PyArray_SETITEM(arr, PyArray_DATA(arr), pyobj_from_ ## ctype ## 1(*v)); break;\\ + default: return -2;\\ + };\\ + return 1 +""" + +needs['TRYCOMPLEXPYARRAYTEMPLATE'] = ['PRINTPYOBJERR'] +cppmacros['TRYCOMPLEXPYARRAYTEMPLATE'] = """ +#define TRYCOMPLEXPYARRAYTEMPLATEOBJECT case NPY_OBJECT: PyArray_SETITEM(arr, PyArray_DATA(arr), pyobj_from_complex_ ## ctype ## 1((*v))); break; +#define TRYCOMPLEXPYARRAYTEMPLATE(ctype,typecode)\\ + PyArrayObject *arr = NULL;\\ + if (!obj) return -2;\\ + if (!PyArray_Check(obj)) return -1;\\ + if (!(arr=(PyArrayObject *)obj)) {fprintf(stderr,\"TRYCOMPLEXPYARRAYTEMPLATE:\");PRINTPYOBJERR(obj);return 0;}\\ + if (PyArray_DESCR(arr)->type==typecode) {\\ + *(ctype *)(PyArray_DATA(arr))=(*v).r;\\ + *(ctype *)(PyArray_DATA(arr)+sizeof(ctype))=(*v).i;\\ + return 1;\\ + }\\ + switch (PyArray_TYPE(arr)) {\\ + case NPY_CDOUBLE: *(npy_double *)(PyArray_DATA(arr))=(*v).r;\\ + *(npy_double *)(PyArray_DATA(arr)+sizeof(npy_double))=(*v).i;\\ + break;\\ + case NPY_CFLOAT: *(npy_float *)(PyArray_DATA(arr))=(*v).r;\\ + *(npy_float *)(PyArray_DATA(arr)+sizeof(npy_float))=(*v).i;\\ + break;\\ + case NPY_DOUBLE: *(npy_double *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_LONG: *(npy_long *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_FLOAT: *(npy_float *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_INT: *(npy_int *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_SHORT: *(npy_short *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_UBYTE: *(npy_ubyte *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_BYTE: *(npy_byte *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_BOOL: *(npy_bool *)(PyArray_DATA(arr))=((*v).r!=0 && (*v).i!=0); break;\\ + case NPY_USHORT: *(npy_ushort *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_UINT: *(npy_uint *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_ULONG: *(npy_ulong *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_LONGLONG: *(npy_longlong *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_ULONGLONG: *(npy_ulonglong *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_LONGDOUBLE: *(npy_longdouble *)(PyArray_DATA(arr))=(*v).r; break;\\ + case NPY_CLONGDOUBLE: *(npy_longdouble *)(PyArray_DATA(arr))=(*v).r;\\ + *(npy_longdouble *)(PyArray_DATA(arr)+sizeof(npy_longdouble))=(*v).i;\\ + break;\\ + case NPY_OBJECT: PyArray_SETITEM(arr, PyArray_DATA(arr), pyobj_from_complex_ ## ctype ## 1((*v))); break;\\ + default: return -2;\\ + };\\ + return -1; +""" +# cppmacros['NUMFROMARROBJ']=""" +# define NUMFROMARROBJ(typenum,ctype) \\ +# if (PyArray_Check(obj)) arr = (PyArrayObject *)obj;\\ +# else arr = (PyArrayObject *)PyArray_ContiguousFromObject(obj,typenum,0,0);\\ +# if (arr) {\\ +# if (PyArray_TYPE(arr)==NPY_OBJECT) {\\ +# if (!ctype ## _from_pyobj(v,(PyArray_DESCR(arr)->getitem)(PyArray_DATA(arr)),\"\"))\\ +# goto capi_fail;\\ +# } else {\\ +# (PyArray_DESCR(arr)->cast[typenum])(PyArray_DATA(arr),1,(char*)v,1,1);\\ +# }\\ +# if ((PyObject *)arr != obj) { Py_DECREF(arr); }\\ +# return 1;\\ +# } +# """ +# XXX: Note that CNUMFROMARROBJ is identical with NUMFROMARROBJ +# cppmacros['CNUMFROMARROBJ']=""" +# define CNUMFROMARROBJ(typenum,ctype) \\ +# if (PyArray_Check(obj)) arr = (PyArrayObject *)obj;\\ +# else arr = (PyArrayObject *)PyArray_ContiguousFromObject(obj,typenum,0,0);\\ +# if (arr) {\\ +# if (PyArray_TYPE(arr)==NPY_OBJECT) {\\ +# if (!ctype ## _from_pyobj(v,(PyArray_DESCR(arr)->getitem)(PyArray_DATA(arr)),\"\"))\\ +# goto capi_fail;\\ +# } else {\\ +# (PyArray_DESCR(arr)->cast[typenum])((void *)(PyArray_DATA(arr)),1,(void *)(v),1,1);\\ +# }\\ +# if ((PyObject *)arr != obj) { Py_DECREF(arr); }\\ +# return 1;\\ +# } +# """ + + +needs['GETSTRFROMPYTUPLE'] = ['STRINGCOPYN', 'PRINTPYOBJERR'] +cppmacros['GETSTRFROMPYTUPLE'] = """ +#define GETSTRFROMPYTUPLE(tuple,index,str,len) {\\ + PyObject *rv_cb_str = PyTuple_GetItem((tuple),(index));\\ + if (rv_cb_str == NULL)\\ + goto capi_fail;\\ + if (PyBytes_Check(rv_cb_str)) {\\ + str[len-1]='\\0';\\ + STRINGCOPYN((str),PyBytes_AS_STRING((PyBytesObject*)rv_cb_str),(len));\\ + } else {\\ + PRINTPYOBJERR(rv_cb_str);\\ + PyErr_SetString(#modulename#_error,\"string object expected\");\\ + goto capi_fail;\\ + }\\ + } +""" +cppmacros['GETSCALARFROMPYTUPLE'] = """ +#define GETSCALARFROMPYTUPLE(tuple,index,var,ctype,mess) {\\ + if ((capi_tmp = PyTuple_GetItem((tuple),(index)))==NULL) goto capi_fail;\\ + if (!(ctype ## _from_pyobj((var),capi_tmp,mess)))\\ + goto capi_fail;\\ + } +""" + +cppmacros['FAILNULL'] = """\ +#define FAILNULL(p) do { \\ + if ((p) == NULL) { \\ + PyErr_SetString(PyExc_MemoryError, "NULL pointer found"); \\ + goto capi_fail; \\ + } \\ +} while (0) +""" +needs['MEMCOPY'] = ['string.h', 'FAILNULL'] +cppmacros['MEMCOPY'] = """ +#define MEMCOPY(to,from,n)\\ + do { FAILNULL(to); FAILNULL(from); (void)memcpy(to,from,n); } while (0) +""" +cppmacros['STRINGMALLOC'] = """ +#define STRINGMALLOC(str,len)\\ + if ((str = (string)malloc(len+1)) == NULL) {\\ + PyErr_SetString(PyExc_MemoryError, \"out of memory\");\\ + goto capi_fail;\\ + } else {\\ + (str)[len] = '\\0';\\ + } +""" +cppmacros['STRINGFREE'] = """ +#define STRINGFREE(str) do {if (!(str == NULL)) free(str);} while (0) +""" +needs['STRINGPADN'] = ['string.h'] +cppmacros['STRINGPADN'] = """ +/* +STRINGPADN replaces null values with padding values from the right. + +`to` must have size of at least N bytes. + +If the `to[N-1]` has null value, then replace it and all the +preceding, nulls with the given padding. + +STRINGPADN(to, N, PADDING, NULLVALUE) is an inverse operation. +*/ +#define STRINGPADN(to, N, NULLVALUE, PADDING) \\ + do { \\ + int _m = (N); \\ + char *_to = (to); \\ + for (_m -= 1; _m >= 0 && _to[_m] == NULLVALUE; _m--) { \\ + _to[_m] = PADDING; \\ + } \\ + } while (0) +""" +needs['STRINGCOPYN'] = ['string.h', 'FAILNULL'] +cppmacros['STRINGCOPYN'] = """ +/* +STRINGCOPYN copies N bytes. + +`to` and `from` buffers must have sizes of at least N bytes. +*/ +#define STRINGCOPYN(to,from,N) \\ + do { \\ + int _m = (N); \\ + char *_to = (to); \\ + char *_from = (from); \\ + FAILNULL(_to); FAILNULL(_from); \\ + (void)strncpy(_to, _from, _m); \\ + } while (0) +""" +needs['STRINGCOPY'] = ['string.h', 'FAILNULL'] +cppmacros['STRINGCOPY'] = """ +#define STRINGCOPY(to,from)\\ + do { FAILNULL(to); FAILNULL(from); (void)strcpy(to,from); } while (0) +""" +cppmacros['CHECKGENERIC'] = """ +#define CHECKGENERIC(check,tcheck,name) \\ + if (!(check)) {\\ + PyErr_SetString(#modulename#_error,\"(\"tcheck\") failed for \"name);\\ + /*goto capi_fail;*/\\ + } else """ +cppmacros['CHECKARRAY'] = """ +#define CHECKARRAY(check,tcheck,name) \\ + if (!(check)) {\\ + PyErr_SetString(#modulename#_error,\"(\"tcheck\") failed for \"name);\\ + /*goto capi_fail;*/\\ + } else """ +cppmacros['CHECKSTRING'] = """ +#define CHECKSTRING(check,tcheck,name,show,var)\\ + if (!(check)) {\\ + char errstring[256];\\ + sprintf(errstring, \"%s: \"show, \"(\"tcheck\") failed for \"name, slen(var), var);\\ + PyErr_SetString(#modulename#_error, errstring);\\ + /*goto capi_fail;*/\\ + } else """ +cppmacros['CHECKSCALAR'] = """ +#define CHECKSCALAR(check,tcheck,name,show,var)\\ + if (!(check)) {\\ + char errstring[256];\\ + sprintf(errstring, \"%s: \"show, \"(\"tcheck\") failed for \"name, var);\\ + PyErr_SetString(#modulename#_error,errstring);\\ + /*goto capi_fail;*/\\ + } else """ +# cppmacros['CHECKDIMS']=""" +# define CHECKDIMS(dims,rank) \\ +# for (int i=0;i<(rank);i++)\\ +# if (dims[i]<0) {\\ +# fprintf(stderr,\"Unspecified array argument requires a complete dimension specification.\\n\");\\ +# goto capi_fail;\\ +# } +# """ +cppmacros[ + 'ARRSIZE'] = '#define ARRSIZE(dims,rank) (_PyArray_multiply_list(dims,rank))' +cppmacros['OLDPYNUM'] = """ +#ifdef OLDPYNUM +#error You need to install NumPy version 0.13 or higher. See https://scipy.org/install.html +#endif +""" + +# Defining the correct value to indicate thread-local storage in C without +# running a compile-time check (which we have no control over in generated +# code used outside of NumPy) is hard. Therefore we support overriding this +# via an external define - the f2py-using package can then use the same +# compile-time checks as we use for `NPY_TLS` when building NumPy (see +# scipy#21860 for an example of that). +# +# __STDC_NO_THREADS__ should not be coupled to the availability of _Thread_local. +# In case we get a bug report, guard it with __STDC_NO_THREADS__ after all. +# +# `thread_local` has become a keyword in C23, but don't try to use that yet +# (too new, doing so while C23 support is preliminary will likely cause more +# problems than it solves). +# +# Note: do not try to use `threads.h`, its availability is very low +# *and* threads.h isn't actually used where `F2PY_THREAD_LOCAL_DECL` is +# in the generated code. See gh-27718 for more details. +cppmacros["F2PY_THREAD_LOCAL_DECL"] = """ +#ifndef F2PY_THREAD_LOCAL_DECL +#if defined(_MSC_VER) +#define F2PY_THREAD_LOCAL_DECL __declspec(thread) +#elif defined(NPY_OS_MINGW) +#define F2PY_THREAD_LOCAL_DECL __thread +#elif defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) +#define F2PY_THREAD_LOCAL_DECL _Thread_local +#elif defined(__GNUC__) \\ + && (__GNUC__ > 4 || (__GNUC__ == 4 && (__GNUC_MINOR__ >= 4))) +#define F2PY_THREAD_LOCAL_DECL __thread +#endif +#endif +""" +################# C functions ############### + +cfuncs['calcarrindex'] = """ +static int calcarrindex(int *i,PyArrayObject *arr) { + int k,ii = i[0]; + for (k=1; k < PyArray_NDIM(arr); k++) + ii += (ii*(PyArray_DIM(arr,k) - 1)+i[k]); /* assuming contiguous arr */ + return ii; +}""" +cfuncs['calcarrindextr'] = """ +static int calcarrindextr(int *i,PyArrayObject *arr) { + int k,ii = i[PyArray_NDIM(arr)-1]; + for (k=1; k < PyArray_NDIM(arr); k++) + ii += (ii*(PyArray_DIM(arr,PyArray_NDIM(arr)-k-1) - 1)+i[PyArray_NDIM(arr)-k-1]); /* assuming contiguous arr */ + return ii; +}""" +cfuncs['forcomb'] = """ +static struct { int nd;npy_intp *d;int *i,*i_tr,tr; } forcombcache; +static int initforcomb(npy_intp *dims,int nd,int tr) { + int k; + if (dims==NULL) return 0; + if (nd<0) return 0; + forcombcache.nd = nd; + forcombcache.d = dims; + forcombcache.tr = tr; + if ((forcombcache.i = (int *)malloc(sizeof(int)*nd))==NULL) return 0; + if ((forcombcache.i_tr = (int *)malloc(sizeof(int)*nd))==NULL) return 0; + for (k=1;k PyArray_NBYTES(arr)) { + n = PyArray_NBYTES(arr); + } + STRINGCOPYN(buf, str, n); + return 1; + } +capi_fail: + PRINTPYOBJERR(obj); + PyErr_SetString(#modulename#_error, \"try_pyarr_from_string failed\"); + return 0; +} +""" +needs['string_from_pyobj'] = ['string', 'STRINGMALLOC', 'STRINGCOPYN'] +cfuncs['string_from_pyobj'] = """ +/* + Create a new string buffer `str` of at most length `len` from a + Python string-like object `obj`. + + The string buffer has given size (len) or the size of inistr when len==-1. + + The string buffer is padded with blanks: in Fortran, trailing blanks + are insignificant contrary to C nulls. + */ +static int +string_from_pyobj(string *str, int *len, const string inistr, PyObject *obj, + const char *errmess) +{ + PyObject *tmp = NULL; + string buf = NULL; + npy_intp n = -1; +#ifdef DEBUGCFUNCS +fprintf(stderr,\"string_from_pyobj(str='%s',len=%d,inistr='%s',obj=%p)\\n\", + (char*)str, *len, (char *)inistr, obj); +#endif + if (obj == Py_None) { + n = strlen(inistr); + buf = inistr; + } + else if (PyArray_Check(obj)) { + PyArrayObject *arr = (PyArrayObject *)obj; + if (!ISCONTIGUOUS(arr)) { + PyErr_SetString(PyExc_ValueError, + \"array object is non-contiguous.\"); + goto capi_fail; + } + n = PyArray_NBYTES(arr); + buf = PyArray_DATA(arr); + n = strnlen(buf, n); + } + else { + if (PyBytes_Check(obj)) { + tmp = obj; + Py_INCREF(tmp); + } + else if (PyUnicode_Check(obj)) { + tmp = PyUnicode_AsASCIIString(obj); + } + else { + PyObject *tmp2; + tmp2 = PyObject_Str(obj); + if (tmp2) { + tmp = PyUnicode_AsASCIIString(tmp2); + Py_DECREF(tmp2); + } + else { + tmp = NULL; + } + } + if (tmp == NULL) goto capi_fail; + n = PyBytes_GET_SIZE(tmp); + buf = PyBytes_AS_STRING(tmp); + } + if (*len == -1) { + /* TODO: change the type of `len` so that we can remove this */ + if (n > NPY_MAX_INT) { + PyErr_SetString(PyExc_OverflowError, + "object too large for a 32-bit int"); + goto capi_fail; + } + *len = n; + } + else if (*len < n) { + /* discard the last (len-n) bytes of input buf */ + n = *len; + } + if (n < 0 || *len < 0 || buf == NULL) { + goto capi_fail; + } + STRINGMALLOC(*str, *len); // *str is allocated with size (*len + 1) + if (n < *len) { + /* + Pad fixed-width string with nulls. The caller will replace + nulls with blanks when the corresponding argument is not + intent(c). + */ + memset(*str + n, '\\0', *len - n); + } + STRINGCOPYN(*str, buf, n); + Py_XDECREF(tmp); + return 1; +capi_fail: + Py_XDECREF(tmp); + { + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = #modulename#_error; + } + PyErr_SetString(err, errmess); + } + return 0; +} +""" + +cfuncs['character_from_pyobj'] = """ +static int +character_from_pyobj(character* v, PyObject *obj, const char *errmess) { + if (PyBytes_Check(obj)) { + /* empty bytes has trailing null, so dereferencing is always safe */ + *v = PyBytes_AS_STRING(obj)[0]; + return 1; + } else if (PyUnicode_Check(obj)) { + PyObject* tmp = PyUnicode_AsASCIIString(obj); + if (tmp != NULL) { + *v = PyBytes_AS_STRING(tmp)[0]; + Py_DECREF(tmp); + return 1; + } + } else if (PyArray_Check(obj)) { + PyArrayObject* arr = (PyArrayObject*)obj; + if (F2PY_ARRAY_IS_CHARACTER_COMPATIBLE(arr)) { + *v = PyArray_BYTES(arr)[0]; + return 1; + } else if (F2PY_IS_UNICODE_ARRAY(arr)) { + // TODO: update when numpy will support 1-byte and + // 2-byte unicode dtypes + PyObject* tmp = PyUnicode_FromKindAndData( + PyUnicode_4BYTE_KIND, + PyArray_BYTES(arr), + (PyArray_NBYTES(arr)>0?1:0)); + if (tmp != NULL) { + if (character_from_pyobj(v, tmp, errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + } + } else if (PySequence_Check(obj)) { + PyObject* tmp = PySequence_GetItem(obj,0); + if (tmp != NULL) { + if (character_from_pyobj(v, tmp, errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + } + { + /* TODO: This error (and most other) error handling needs cleaning. */ + char mess[F2PY_MESSAGE_BUFFER_SIZE]; + strcpy(mess, errmess); + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = PyExc_TypeError; + Py_INCREF(err); + } + else { + Py_INCREF(err); + PyErr_Clear(); + } + sprintf(mess + strlen(mess), + " -- expected str|bytes|sequence-of-str-or-bytes, got "); + f2py_describe(obj, mess + strlen(mess)); + PyErr_SetString(err, mess); + Py_DECREF(err); + } + return 0; +} +""" + +# TODO: These should be dynamically generated, too many mapped to int things, +# see note in _isocbind.py +needs['char_from_pyobj'] = ['int_from_pyobj'] +cfuncs['char_from_pyobj'] = """ +static int +char_from_pyobj(char* v, PyObject *obj, const char *errmess) { + int i = 0; + if (int_from_pyobj(&i, obj, errmess)) { + *v = (char)i; + return 1; + } + return 0; +} +""" + + +needs['signed_char_from_pyobj'] = ['int_from_pyobj', 'signed_char'] +cfuncs['signed_char_from_pyobj'] = """ +static int +signed_char_from_pyobj(signed_char* v, PyObject *obj, const char *errmess) { + int i = 0; + if (int_from_pyobj(&i, obj, errmess)) { + *v = (signed_char)i; + return 1; + } + return 0; +} +""" + + +needs['short_from_pyobj'] = ['int_from_pyobj'] +cfuncs['short_from_pyobj'] = """ +static int +short_from_pyobj(short* v, PyObject *obj, const char *errmess) { + int i = 0; + if (int_from_pyobj(&i, obj, errmess)) { + *v = (short)i; + return 1; + } + return 0; +} +""" + + +cfuncs['int_from_pyobj'] = """ +static int +int_from_pyobj(int* v, PyObject *obj, const char *errmess) +{ + PyObject* tmp = NULL; + + if (PyLong_Check(obj)) { + *v = Npy__PyLong_AsInt(obj); + return !(*v == -1 && PyErr_Occurred()); + } + + tmp = PyNumber_Long(obj); + if (tmp) { + *v = Npy__PyLong_AsInt(tmp); + Py_DECREF(tmp); + return !(*v == -1 && PyErr_Occurred()); + } + + if (PyComplex_Check(obj)) { + PyErr_Clear(); + tmp = PyObject_GetAttrString(obj,\"real\"); + } + else if (PyBytes_Check(obj) || PyUnicode_Check(obj)) { + /*pass*/; + } + else if (PySequence_Check(obj)) { + PyErr_Clear(); + tmp = PySequence_GetItem(obj, 0); + } + + if (tmp) { + if (int_from_pyobj(v, tmp, errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + + { + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = #modulename#_error; + } + PyErr_SetString(err, errmess); + } + return 0; +} +""" + + +cfuncs['long_from_pyobj'] = """ +static int +long_from_pyobj(long* v, PyObject *obj, const char *errmess) { + PyObject* tmp = NULL; + + if (PyLong_Check(obj)) { + *v = PyLong_AsLong(obj); + return !(*v == -1 && PyErr_Occurred()); + } + + tmp = PyNumber_Long(obj); + if (tmp) { + *v = PyLong_AsLong(tmp); + Py_DECREF(tmp); + return !(*v == -1 && PyErr_Occurred()); + } + + if (PyComplex_Check(obj)) { + PyErr_Clear(); + tmp = PyObject_GetAttrString(obj,\"real\"); + } + else if (PyBytes_Check(obj) || PyUnicode_Check(obj)) { + /*pass*/; + } + else if (PySequence_Check(obj)) { + PyErr_Clear(); + tmp = PySequence_GetItem(obj, 0); + } + + if (tmp) { + if (long_from_pyobj(v, tmp, errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + { + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = #modulename#_error; + } + PyErr_SetString(err, errmess); + } + return 0; +} +""" + + +needs['long_long_from_pyobj'] = ['long_long'] +cfuncs['long_long_from_pyobj'] = """ +static int +long_long_from_pyobj(long_long* v, PyObject *obj, const char *errmess) +{ + PyObject* tmp = NULL; + + if (PyLong_Check(obj)) { + *v = PyLong_AsLongLong(obj); + return !(*v == -1 && PyErr_Occurred()); + } + + tmp = PyNumber_Long(obj); + if (tmp) { + *v = PyLong_AsLongLong(tmp); + Py_DECREF(tmp); + return !(*v == -1 && PyErr_Occurred()); + } + + if (PyComplex_Check(obj)) { + PyErr_Clear(); + tmp = PyObject_GetAttrString(obj,\"real\"); + } + else if (PyBytes_Check(obj) || PyUnicode_Check(obj)) { + /*pass*/; + } + else if (PySequence_Check(obj)) { + PyErr_Clear(); + tmp = PySequence_GetItem(obj, 0); + } + + if (tmp) { + if (long_long_from_pyobj(v, tmp, errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + { + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = #modulename#_error; + } + PyErr_SetString(err,errmess); + } + return 0; +} +""" + + +needs['long_double_from_pyobj'] = ['double_from_pyobj', 'long_double'] +cfuncs['long_double_from_pyobj'] = """ +static int +long_double_from_pyobj(long_double* v, PyObject *obj, const char *errmess) +{ + double d=0; + if (PyArray_CheckScalar(obj)){ + if PyArray_IsScalar(obj, LongDouble) { + PyArray_ScalarAsCtype(obj, v); + return 1; + } + else if (PyArray_Check(obj) && PyArray_TYPE(obj) == NPY_LONGDOUBLE) { + (*v) = *((npy_longdouble *)PyArray_DATA(obj)); + return 1; + } + } + if (double_from_pyobj(&d, obj, errmess)) { + *v = (long_double)d; + return 1; + } + return 0; +} +""" + + +cfuncs['double_from_pyobj'] = """ +static int +double_from_pyobj(double* v, PyObject *obj, const char *errmess) +{ + PyObject* tmp = NULL; + if (PyFloat_Check(obj)) { + *v = PyFloat_AsDouble(obj); + return !(*v == -1.0 && PyErr_Occurred()); + } + + tmp = PyNumber_Float(obj); + if (tmp) { + *v = PyFloat_AsDouble(tmp); + Py_DECREF(tmp); + return !(*v == -1.0 && PyErr_Occurred()); + } + + if (PyComplex_Check(obj)) { + PyErr_Clear(); + tmp = PyObject_GetAttrString(obj,\"real\"); + } + else if (PyBytes_Check(obj) || PyUnicode_Check(obj)) { + /*pass*/; + } + else if (PySequence_Check(obj)) { + PyErr_Clear(); + tmp = PySequence_GetItem(obj, 0); + } + + if (tmp) { + if (double_from_pyobj(v,tmp,errmess)) {Py_DECREF(tmp); return 1;} + Py_DECREF(tmp); + } + { + PyObject* err = PyErr_Occurred(); + if (err==NULL) err = #modulename#_error; + PyErr_SetString(err,errmess); + } + return 0; +} +""" + + +needs['float_from_pyobj'] = ['double_from_pyobj'] +cfuncs['float_from_pyobj'] = """ +static int +float_from_pyobj(float* v, PyObject *obj, const char *errmess) +{ + double d=0.0; + if (double_from_pyobj(&d,obj,errmess)) { + *v = (float)d; + return 1; + } + return 0; +} +""" + + +needs['complex_long_double_from_pyobj'] = ['complex_long_double', 'long_double', + 'complex_double_from_pyobj', 'npy_math.h'] +cfuncs['complex_long_double_from_pyobj'] = """ +static int +complex_long_double_from_pyobj(complex_long_double* v, PyObject *obj, const char *errmess) +{ + complex_double cd = {0.0,0.0}; + if (PyArray_CheckScalar(obj)){ + if PyArray_IsScalar(obj, CLongDouble) { + PyArray_ScalarAsCtype(obj, v); + return 1; + } + else if (PyArray_Check(obj) && PyArray_TYPE(obj)==NPY_CLONGDOUBLE) { + (*v).r = npy_creall(*(((npy_clongdouble *)PyArray_DATA(obj)))); + (*v).i = npy_cimagl(*(((npy_clongdouble *)PyArray_DATA(obj)))); + return 1; + } + } + if (complex_double_from_pyobj(&cd,obj,errmess)) { + (*v).r = (long_double)cd.r; + (*v).i = (long_double)cd.i; + return 1; + } + return 0; +} +""" + + +needs['complex_double_from_pyobj'] = ['complex_double', 'npy_math.h'] +cfuncs['complex_double_from_pyobj'] = """ +static int +complex_double_from_pyobj(complex_double* v, PyObject *obj, const char *errmess) { + Py_complex c; + if (PyComplex_Check(obj)) { + c = PyComplex_AsCComplex(obj); + (*v).r = c.real; + (*v).i = c.imag; + return 1; + } + if (PyArray_IsScalar(obj, ComplexFloating)) { + if (PyArray_IsScalar(obj, CFloat)) { + npy_cfloat new; + PyArray_ScalarAsCtype(obj, &new); + (*v).r = (double)npy_crealf(new); + (*v).i = (double)npy_cimagf(new); + } + else if (PyArray_IsScalar(obj, CLongDouble)) { + npy_clongdouble new; + PyArray_ScalarAsCtype(obj, &new); + (*v).r = (double)npy_creall(new); + (*v).i = (double)npy_cimagl(new); + } + else { /* if (PyArray_IsScalar(obj, CDouble)) */ + PyArray_ScalarAsCtype(obj, v); + } + return 1; + } + if (PyArray_CheckScalar(obj)) { /* 0-dim array or still array scalar */ + PyArrayObject *arr; + if (PyArray_Check(obj)) { + arr = (PyArrayObject *)PyArray_Cast((PyArrayObject *)obj, NPY_CDOUBLE); + } + else { + arr = (PyArrayObject *)PyArray_FromScalar(obj, PyArray_DescrFromType(NPY_CDOUBLE)); + } + if (arr == NULL) { + return 0; + } + (*v).r = npy_creal(*(((npy_cdouble *)PyArray_DATA(arr)))); + (*v).i = npy_cimag(*(((npy_cdouble *)PyArray_DATA(arr)))); + Py_DECREF(arr); + return 1; + } + /* Python does not provide PyNumber_Complex function :-( */ + (*v).i = 0.0; + if (PyFloat_Check(obj)) { + (*v).r = PyFloat_AsDouble(obj); + return !((*v).r == -1.0 && PyErr_Occurred()); + } + if (PyLong_Check(obj)) { + (*v).r = PyLong_AsDouble(obj); + return !((*v).r == -1.0 && PyErr_Occurred()); + } + if (PySequence_Check(obj) && !(PyBytes_Check(obj) || PyUnicode_Check(obj))) { + PyObject *tmp = PySequence_GetItem(obj,0); + if (tmp) { + if (complex_double_from_pyobj(v,tmp,errmess)) { + Py_DECREF(tmp); + return 1; + } + Py_DECREF(tmp); + } + } + { + PyObject* err = PyErr_Occurred(); + if (err==NULL) + err = PyExc_TypeError; + PyErr_SetString(err,errmess); + } + return 0; +} +""" + + +needs['complex_float_from_pyobj'] = [ + 'complex_float', 'complex_double_from_pyobj'] +cfuncs['complex_float_from_pyobj'] = """ +static int +complex_float_from_pyobj(complex_float* v,PyObject *obj,const char *errmess) +{ + complex_double cd={0.0,0.0}; + if (complex_double_from_pyobj(&cd,obj,errmess)) { + (*v).r = (float)cd.r; + (*v).i = (float)cd.i; + return 1; + } + return 0; +} +""" + + +cfuncs['try_pyarr_from_character'] = """ +static int try_pyarr_from_character(PyObject* obj, character* v) { + PyArrayObject *arr = (PyArrayObject*)obj; + if (!obj) return -2; + if (PyArray_Check(obj)) { + if (F2PY_ARRAY_IS_CHARACTER_COMPATIBLE(arr)) { + *(character *)(PyArray_DATA(arr)) = *v; + return 1; + } + } + { + char mess[F2PY_MESSAGE_BUFFER_SIZE]; + PyObject* err = PyErr_Occurred(); + if (err == NULL) { + err = PyExc_ValueError; + strcpy(mess, "try_pyarr_from_character failed" + " -- expected bytes array-scalar|array, got "); + f2py_describe(obj, mess + strlen(mess)); + PyErr_SetString(err, mess); + } + } + return 0; +} +""" + +needs['try_pyarr_from_char'] = ['pyobj_from_char1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_char'] = 'static int try_pyarr_from_char(PyObject* obj,char* v) {\n TRYPYARRAYTEMPLATE(char,\'c\');\n}\n' +needs['try_pyarr_from_signed_char'] = ['TRYPYARRAYTEMPLATE', 'unsigned_char'] +cfuncs[ + 'try_pyarr_from_unsigned_char'] = 'static int try_pyarr_from_unsigned_char(PyObject* obj,unsigned_char* v) {\n TRYPYARRAYTEMPLATE(unsigned_char,\'b\');\n}\n' +needs['try_pyarr_from_signed_char'] = ['TRYPYARRAYTEMPLATE', 'signed_char'] +cfuncs[ + 'try_pyarr_from_signed_char'] = 'static int try_pyarr_from_signed_char(PyObject* obj,signed_char* v) {\n TRYPYARRAYTEMPLATE(signed_char,\'1\');\n}\n' +needs['try_pyarr_from_short'] = ['pyobj_from_short1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_short'] = 'static int try_pyarr_from_short(PyObject* obj,short* v) {\n TRYPYARRAYTEMPLATE(short,\'s\');\n}\n' +needs['try_pyarr_from_int'] = ['pyobj_from_int1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_int'] = 'static int try_pyarr_from_int(PyObject* obj,int* v) {\n TRYPYARRAYTEMPLATE(int,\'i\');\n}\n' +needs['try_pyarr_from_long'] = ['pyobj_from_long1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_long'] = 'static int try_pyarr_from_long(PyObject* obj,long* v) {\n TRYPYARRAYTEMPLATE(long,\'l\');\n}\n' +needs['try_pyarr_from_long_long'] = [ + 'pyobj_from_long_long1', 'TRYPYARRAYTEMPLATE', 'long_long'] +cfuncs[ + 'try_pyarr_from_long_long'] = 'static int try_pyarr_from_long_long(PyObject* obj,long_long* v) {\n TRYPYARRAYTEMPLATE(long_long,\'L\');\n}\n' +needs['try_pyarr_from_float'] = ['pyobj_from_float1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_float'] = 'static int try_pyarr_from_float(PyObject* obj,float* v) {\n TRYPYARRAYTEMPLATE(float,\'f\');\n}\n' +needs['try_pyarr_from_double'] = ['pyobj_from_double1', 'TRYPYARRAYTEMPLATE'] +cfuncs[ + 'try_pyarr_from_double'] = 'static int try_pyarr_from_double(PyObject* obj,double* v) {\n TRYPYARRAYTEMPLATE(double,\'d\');\n}\n' +needs['try_pyarr_from_complex_float'] = [ + 'pyobj_from_complex_float1', 'TRYCOMPLEXPYARRAYTEMPLATE', 'complex_float'] +cfuncs[ + 'try_pyarr_from_complex_float'] = 'static int try_pyarr_from_complex_float(PyObject* obj,complex_float* v) {\n TRYCOMPLEXPYARRAYTEMPLATE(float,\'F\');\n}\n' +needs['try_pyarr_from_complex_double'] = [ + 'pyobj_from_complex_double1', 'TRYCOMPLEXPYARRAYTEMPLATE', 'complex_double'] +cfuncs[ + 'try_pyarr_from_complex_double'] = 'static int try_pyarr_from_complex_double(PyObject* obj,complex_double* v) {\n TRYCOMPLEXPYARRAYTEMPLATE(double,\'D\');\n}\n' + + +needs['create_cb_arglist'] = ['CFUNCSMESS', 'PRINTPYOBJERR', 'MINMAX'] +# create the list of arguments to be used when calling back to python +cfuncs['create_cb_arglist'] = """ +static int +create_cb_arglist(PyObject* fun, PyTupleObject* xa , const int maxnofargs, + const int nofoptargs, int *nofargs, PyTupleObject **args, + const char *errmess) +{ + PyObject *tmp = NULL; + PyObject *tmp_fun = NULL; + Py_ssize_t tot, opt, ext, siz, i, di = 0; + CFUNCSMESS(\"create_cb_arglist\\n\"); + tot=opt=ext=siz=0; + /* Get the total number of arguments */ + if (PyFunction_Check(fun)) { + tmp_fun = fun; + Py_INCREF(tmp_fun); + } + else { + di = 1; + if (PyObject_HasAttrString(fun,\"im_func\")) { + tmp_fun = PyObject_GetAttrString(fun,\"im_func\"); + } + else if (PyObject_HasAttrString(fun,\"__call__\")) { + tmp = PyObject_GetAttrString(fun,\"__call__\"); + if (PyObject_HasAttrString(tmp,\"im_func\")) + tmp_fun = PyObject_GetAttrString(tmp,\"im_func\"); + else { + tmp_fun = fun; /* built-in function */ + Py_INCREF(tmp_fun); + tot = maxnofargs; + if (PyCFunction_Check(fun)) { + /* In case the function has a co_argcount (like on PyPy) */ + di = 0; + } + if (xa != NULL) + tot += PyTuple_Size((PyObject *)xa); + } + Py_XDECREF(tmp); + } + else if (PyFortran_Check(fun) || PyFortran_Check1(fun)) { + tot = maxnofargs; + if (xa != NULL) + tot += PyTuple_Size((PyObject *)xa); + tmp_fun = fun; + Py_INCREF(tmp_fun); + } + else if (F2PyCapsule_Check(fun)) { + tot = maxnofargs; + if (xa != NULL) + ext = PyTuple_Size((PyObject *)xa); + if(ext>0) { + fprintf(stderr,\"extra arguments tuple cannot be used with PyCapsule call-back\\n\"); + goto capi_fail; + } + tmp_fun = fun; + Py_INCREF(tmp_fun); + } + } + + if (tmp_fun == NULL) { + fprintf(stderr, + \"Call-back argument must be function|instance|instance.__call__|f2py-function \" + \"but got %s.\\n\", + ((fun == NULL) ? \"NULL\" : Py_TYPE(fun)->tp_name)); + goto capi_fail; + } + + if (PyObject_HasAttrString(tmp_fun,\"__code__\")) { + if (PyObject_HasAttrString(tmp = PyObject_GetAttrString(tmp_fun,\"__code__\"),\"co_argcount\")) { + PyObject *tmp_argcount = PyObject_GetAttrString(tmp,\"co_argcount\"); + Py_DECREF(tmp); + if (tmp_argcount == NULL) { + goto capi_fail; + } + tot = PyLong_AsSsize_t(tmp_argcount) - di; + Py_DECREF(tmp_argcount); + } + } + /* Get the number of optional arguments */ + if (PyObject_HasAttrString(tmp_fun,\"__defaults__\")) { + if (PyTuple_Check(tmp = PyObject_GetAttrString(tmp_fun,\"__defaults__\"))) + opt = PyTuple_Size(tmp); + Py_XDECREF(tmp); + } + /* Get the number of extra arguments */ + if (xa != NULL) + ext = PyTuple_Size((PyObject *)xa); + /* Calculate the size of call-backs argument list */ + siz = MIN(maxnofargs+ext,tot); + *nofargs = MAX(0,siz-ext); + +#ifdef DEBUGCFUNCS + fprintf(stderr, + \"debug-capi:create_cb_arglist:maxnofargs(-nofoptargs),\" + \"tot,opt,ext,siz,nofargs = %d(-%d), %zd, %zd, %zd, %zd, %d\\n\", + maxnofargs, nofoptargs, tot, opt, ext, siz, *nofargs); +#endif + + if (siz < tot-opt) { + fprintf(stderr, + \"create_cb_arglist: Failed to build argument list \" + \"(siz) with enough arguments (tot-opt) required by \" + \"user-supplied function (siz,tot,opt=%zd, %zd, %zd).\\n\", + siz, tot, opt); + goto capi_fail; + } + + /* Initialize argument list */ + *args = (PyTupleObject *)PyTuple_New(siz); + for (i=0;i<*nofargs;i++) { + Py_INCREF(Py_None); + PyTuple_SET_ITEM((PyObject *)(*args),i,Py_None); + } + if (xa != NULL) + for (i=(*nofargs);i 0: + if outneeds[n][0] not in needs: + out.append(outneeds[n][0]) + del outneeds[n][0] + else: + flag = 0 + for k in outneeds[n][1:]: + if k in needs[outneeds[n][0]]: + flag = 1 + break + if flag: + outneeds[n] = outneeds[n][1:] + [outneeds[n][0]] + else: + out.append(outneeds[n][0]) + del outneeds[n][0] + if saveout and (0 not in map(lambda x, y: x == y, saveout, outneeds[n])) \ + and outneeds[n] != []: + print(n, saveout) + errmess( + 'get_needs: no progress in sorting needs, probably circular dependence, skipping.\n') + out = out + saveout + break + saveout = copy.copy(outneeds[n]) + if out == []: + out = [n] + res[n] = out + return res diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/common_rules.py b/venv/lib/python3.10/site-packages/numpy/f2py/common_rules.py new file mode 100644 index 0000000000000000000000000000000000000000..64347b737454fe1bae544b6630de2729157d7f71 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/common_rules.py @@ -0,0 +1,146 @@ +""" +Build common block mechanism for f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +from . import __version__ +f2py_version = __version__.version + +from .auxfuncs import ( + hasbody, hascommon, hasnote, isintent_hide, outmess, getuseblocks +) +from . import capi_maps +from . import func2subr +from .crackfortran import rmbadname + + +def findcommonblocks(block, top=1): + ret = [] + if hascommon(block): + for key, value in block['common'].items(): + vars_ = {v: block['vars'][v] for v in value} + ret.append((key, value, vars_)) + elif hasbody(block): + for b in block['body']: + ret = ret + findcommonblocks(b, 0) + if top: + tret = [] + names = [] + for t in ret: + if t[0] not in names: + names.append(t[0]) + tret.append(t) + return tret + return ret + + +def buildhooks(m): + ret = {'commonhooks': [], 'initcommonhooks': [], + 'docs': ['"COMMON blocks:\\n"']} + fwrap = [''] + + def fadd(line, s=fwrap): + s[0] = '%s\n %s' % (s[0], line) + chooks = [''] + + def cadd(line, s=chooks): + s[0] = '%s\n%s' % (s[0], line) + ihooks = [''] + + def iadd(line, s=ihooks): + s[0] = '%s\n%s' % (s[0], line) + doc = [''] + + def dadd(line, s=doc): + s[0] = '%s\n%s' % (s[0], line) + for (name, vnames, vars) in findcommonblocks(m): + lower_name = name.lower() + hnames, inames = [], [] + for n in vnames: + if isintent_hide(vars[n]): + hnames.append(n) + else: + inames.append(n) + if hnames: + outmess('\t\tConstructing COMMON block support for "%s"...\n\t\t %s\n\t\t Hidden: %s\n' % ( + name, ','.join(inames), ','.join(hnames))) + else: + outmess('\t\tConstructing COMMON block support for "%s"...\n\t\t %s\n' % ( + name, ','.join(inames))) + fadd('subroutine f2pyinit%s(setupfunc)' % name) + for usename in getuseblocks(m): + fadd(f'use {usename}') + fadd('external setupfunc') + for n in vnames: + fadd(func2subr.var2fixfortran(vars, n)) + if name == '_BLNK_': + fadd('common %s' % (','.join(vnames))) + else: + fadd('common /%s/ %s' % (name, ','.join(vnames))) + fadd('call setupfunc(%s)' % (','.join(inames))) + fadd('end\n') + cadd('static FortranDataDef f2py_%s_def[] = {' % (name)) + idims = [] + for n in inames: + ct = capi_maps.getctype(vars[n]) + elsize = capi_maps.get_elsize(vars[n]) + at = capi_maps.c2capi_map[ct] + dm = capi_maps.getarrdims(n, vars[n]) + if dm['dims']: + idims.append('(%s)' % (dm['dims'])) + else: + idims.append('') + dms = dm['dims'].strip() + if not dms: + dms = '-1' + cadd('\t{\"%s\",%s,{{%s}},%s, %s},' + % (n, dm['rank'], dms, at, elsize)) + cadd('\t{NULL}\n};') + inames1 = rmbadname(inames) + inames1_tps = ','.join(['char *' + s for s in inames1]) + cadd('static void f2py_setup_%s(%s) {' % (name, inames1_tps)) + cadd('\tint i_f2py=0;') + for n in inames1: + cadd('\tf2py_%s_def[i_f2py++].data = %s;' % (name, n)) + cadd('}') + if '_' in lower_name: + F_FUNC = 'F_FUNC_US' + else: + F_FUNC = 'F_FUNC' + cadd('extern void %s(f2pyinit%s,F2PYINIT%s)(void(*)(%s));' + % (F_FUNC, lower_name, name.upper(), + ','.join(['char*'] * len(inames1)))) + cadd('static void f2py_init_%s(void) {' % name) + cadd('\t%s(f2pyinit%s,F2PYINIT%s)(f2py_setup_%s);' + % (F_FUNC, lower_name, name.upper(), name)) + cadd('}\n') + iadd('\ttmp = PyFortranObject_New(f2py_%s_def,f2py_init_%s);' % (name, name)) + iadd('\tif (tmp == NULL) return NULL;') + iadd('\tif (F2PyDict_SetItemString(d, \"%s\", tmp) == -1) return NULL;' + % name) + iadd('\tPy_DECREF(tmp);') + tname = name.replace('_', '\\_') + dadd('\\subsection{Common block \\texttt{%s}}\n' % (tname)) + dadd('\\begin{description}') + for n in inames: + dadd('\\item[]{{}\\verb@%s@{}}' % + (capi_maps.getarrdocsign(n, vars[n]))) + if hasnote(vars[n]): + note = vars[n]['note'] + if isinstance(note, list): + note = '\n'.join(note) + dadd('--- %s' % (note)) + dadd('\\end{description}') + ret['docs'].append( + '"\t/%s/ %s\\n"' % (name, ','.join(map(lambda v, d: v + d, inames, idims)))) + ret['commonhooks'] = chooks + ret['initcommonhooks'] = ihooks + ret['latexdoc'] = doc[0] + if len(ret['docs']) <= 1: + ret['docs'] = '' + return ret, fwrap[0] diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/crackfortran.py b/venv/lib/python3.10/site-packages/numpy/f2py/crackfortran.py new file mode 100644 index 0000000000000000000000000000000000000000..3ea1888df113686fb36f2c0cafed0786f12411a6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/crackfortran.py @@ -0,0 +1,3746 @@ +""" +crackfortran --- read fortran (77,90) code and extract declaration information. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. + + +Usage of crackfortran: +====================== +Command line keys: -quiet,-verbose,-fix,-f77,-f90,-show,-h + -m ,--ignore-contains +Functions: crackfortran, crack2fortran +The following Fortran statements/constructions are supported +(or will be if needed): + block data,byte,call,character,common,complex,contains,data, + dimension,double complex,double precision,end,external,function, + implicit,integer,intent,interface,intrinsic, + logical,module,optional,parameter,private,public, + program,real,(sequence?),subroutine,type,use,virtual, + include,pythonmodule +Note: 'virtual' is mapped to 'dimension'. +Note: 'implicit integer (z) static (z)' is 'implicit static (z)' (this is minor bug). +Note: code after 'contains' will be ignored until its scope ends. +Note: 'common' statement is extended: dimensions are moved to variable definitions +Note: f2py directive: f2py is read as +Note: pythonmodule is introduced to represent Python module + +Usage: + `postlist=crackfortran(files)` + `postlist` contains declaration information read from the list of files `files`. + `crack2fortran(postlist)` returns a fortran code to be saved to pyf-file + + `postlist` has the following structure: + *** it is a list of dictionaries containing `blocks': + B = {'block','body','vars','parent_block'[,'name','prefix','args','result', + 'implicit','externals','interfaced','common','sortvars', + 'commonvars','note']} + B['block'] = 'interface' | 'function' | 'subroutine' | 'module' | + 'program' | 'block data' | 'type' | 'pythonmodule' | + 'abstract interface' + B['body'] --- list containing `subblocks' with the same structure as `blocks' + B['parent_block'] --- dictionary of a parent block: + C['body'][]['parent_block'] is C + B['vars'] --- dictionary of variable definitions + B['sortvars'] --- dictionary of variable definitions sorted by dependence (independent first) + B['name'] --- name of the block (not if B['block']=='interface') + B['prefix'] --- prefix string (only if B['block']=='function') + B['args'] --- list of argument names if B['block']== 'function' | 'subroutine' + B['result'] --- name of the return value (only if B['block']=='function') + B['implicit'] --- dictionary {'a':,'b':...} | None + B['externals'] --- list of variables being external + B['interfaced'] --- list of variables being external and defined + B['common'] --- dictionary of common blocks (list of objects) + B['commonvars'] --- list of variables used in common blocks (dimensions are moved to variable definitions) + B['from'] --- string showing the 'parents' of the current block + B['use'] --- dictionary of modules used in current block: + {:{['only':<0|1>],['map':{:,...}]}} + B['note'] --- list of LaTeX comments on the block + B['f2pyenhancements'] --- optional dictionary + {'threadsafe':'','fortranname':, + 'callstatement':|, + 'callprotoargument':, + 'usercode':|, + 'pymethoddef:' + } + B['entry'] --- dictionary {entryname:argslist,..} + B['varnames'] --- list of variable names given in the order of reading the + Fortran code, useful for derived types. + B['saved_interface'] --- a string of scanned routine signature, defines explicit interface + *** Variable definition is a dictionary + D = B['vars'][] = + {'typespec'[,'attrspec','kindselector','charselector','=','typename']} + D['typespec'] = 'byte' | 'character' | 'complex' | 'double complex' | + 'double precision' | 'integer' | 'logical' | 'real' | 'type' + D['attrspec'] --- list of attributes (e.g. 'dimension()', + 'external','intent(in|out|inout|hide|c|callback|cache|aligned4|aligned8|aligned16)', + 'optional','required', etc) + K = D['kindselector'] = {['*','kind']} (only if D['typespec'] = + 'complex' | 'integer' | 'logical' | 'real' ) + C = D['charselector'] = {['*','len','kind','f2py_len']} + (only if D['typespec']=='character') + D['='] --- initialization expression string + D['typename'] --- name of the type if D['typespec']=='type' + D['dimension'] --- list of dimension bounds + D['intent'] --- list of intent specifications + D['depend'] --- list of variable names on which current variable depends on + D['check'] --- list of C-expressions; if C-expr returns zero, exception is raised + D['note'] --- list of LaTeX comments on the variable + *** Meaning of kind/char selectors (few examples): + D['typespec>']*K['*'] + D['typespec'](kind=K['kind']) + character*C['*'] + character(len=C['len'],kind=C['kind'], f2py_len=C['f2py_len']) + (see also fortran type declaration statement formats below) + +Fortran 90 type declaration statement format (F77 is subset of F90) +==================================================================== +(Main source: IBM XL Fortran 5.1 Language Reference Manual) +type declaration = [[]::] + = byte | + character[] | + complex[] | + double complex | + double precision | + integer[] | + logical[] | + real[] | + type() + = * | + ([len=][,[kind=]]) | + (kind=[,len=]) + = * | + ([kind=]) + = comma separated list of attributes. + Only the following attributes are used in + building up the interface: + external + (parameter --- affects '=' key) + optional + intent + Other attributes are ignored. + = in | out | inout + = comma separated list of dimension bounds. + = [[*][()] | [()]*] + [// | =] [,] + +In addition, the following attributes are used: check,depend,note + +TODO: + * Apply 'parameter' attribute (e.g. 'integer parameter :: i=2' 'real x(i)' + -> 'real x(2)') + The above may be solved by creating appropriate preprocessor program, for example. + +""" +import sys +import string +import fileinput +import re +import os +import copy +import platform +import codecs +from pathlib import Path +try: + import charset_normalizer +except ImportError: + charset_normalizer = None + +from . import __version__ + +# The environment provided by auxfuncs.py is needed for some calls to eval. +# As the needed functions cannot be determined by static inspection of the +# code, it is safest to use import * pending a major refactoring of f2py. +from .auxfuncs import * +from . import symbolic + +f2py_version = __version__.version + +# Global flags: +strictf77 = 1 # Ignore `!' comments unless line[0]=='!' +sourcecodeform = 'fix' # 'fix','free' +quiet = 0 # Be verbose if 0 (Obsolete: not used any more) +verbose = 1 # Be quiet if 0, extra verbose if > 1. +tabchar = 4 * ' ' +pyffilename = '' +f77modulename = '' +skipemptyends = 0 # for old F77 programs without 'program' statement +ignorecontains = 1 +dolowercase = 1 +debug = [] + +# Global variables +beginpattern = '' +currentfilename = '' +expectbegin = 1 +f90modulevars = {} +filepositiontext = '' +gotnextfile = 1 +groupcache = None +groupcounter = 0 +grouplist = {groupcounter: []} +groupname = '' +include_paths = [] +neededmodule = -1 +onlyfuncs = [] +previous_context = None +skipblocksuntil = -1 +skipfuncs = [] +skipfunctions = [] +usermodules = [] + + +def reset_global_f2py_vars(): + global groupcounter, grouplist, neededmodule, expectbegin + global skipblocksuntil, usermodules, f90modulevars, gotnextfile + global filepositiontext, currentfilename, skipfunctions, skipfuncs + global onlyfuncs, include_paths, previous_context + global strictf77, sourcecodeform, quiet, verbose, tabchar, pyffilename + global f77modulename, skipemptyends, ignorecontains, dolowercase, debug + + # flags + strictf77 = 1 + sourcecodeform = 'fix' + quiet = 0 + verbose = 1 + tabchar = 4 * ' ' + pyffilename = '' + f77modulename = '' + skipemptyends = 0 + ignorecontains = 1 + dolowercase = 1 + debug = [] + # variables + groupcounter = 0 + grouplist = {groupcounter: []} + neededmodule = -1 + expectbegin = 1 + skipblocksuntil = -1 + usermodules = [] + f90modulevars = {} + gotnextfile = 1 + filepositiontext = '' + currentfilename = '' + skipfunctions = [] + skipfuncs = [] + onlyfuncs = [] + include_paths = [] + previous_context = None + + +def outmess(line, flag=1): + global filepositiontext + + if not verbose: + return + if not quiet: + if flag: + sys.stdout.write(filepositiontext) + sys.stdout.write(line) + +re._MAXCACHE = 50 +defaultimplicitrules = {} +for c in "abcdefghopqrstuvwxyz$_": + defaultimplicitrules[c] = {'typespec': 'real'} +for c in "ijklmn": + defaultimplicitrules[c] = {'typespec': 'integer'} +badnames = {} +invbadnames = {} +for n in ['int', 'double', 'float', 'char', 'short', 'long', 'void', 'case', 'while', + 'return', 'signed', 'unsigned', 'if', 'for', 'typedef', 'sizeof', 'union', + 'struct', 'static', 'register', 'new', 'break', 'do', 'goto', 'switch', + 'continue', 'else', 'inline', 'extern', 'delete', 'const', 'auto', + 'len', 'rank', 'shape', 'index', 'slen', 'size', '_i', + 'max', 'min', + 'flen', 'fshape', + 'string', 'complex_double', 'float_double', 'stdin', 'stderr', 'stdout', + 'type', 'default']: + badnames[n] = n + '_bn' + invbadnames[n + '_bn'] = n + + +def rmbadname1(name): + if name in badnames: + errmess('rmbadname1: Replacing "%s" with "%s".\n' % + (name, badnames[name])) + return badnames[name] + return name + + +def rmbadname(names): + return [rmbadname1(_m) for _m in names] + + +def undo_rmbadname1(name): + if name in invbadnames: + errmess('undo_rmbadname1: Replacing "%s" with "%s".\n' + % (name, invbadnames[name])) + return invbadnames[name] + return name + + +def undo_rmbadname(names): + return [undo_rmbadname1(_m) for _m in names] + + +_has_f_header = re.compile(r'-\*-\s*fortran\s*-\*-', re.I).search +_has_f90_header = re.compile(r'-\*-\s*f90\s*-\*-', re.I).search +_has_fix_header = re.compile(r'-\*-\s*fix\s*-\*-', re.I).search +_free_f90_start = re.compile(r'[^c*]\s*[^\s\d\t]', re.I).match + +# Extensions +COMMON_FREE_EXTENSIONS = ['.f90', '.f95', '.f03', '.f08'] +COMMON_FIXED_EXTENSIONS = ['.for', '.ftn', '.f77', '.f'] + + +def openhook(filename, mode): + """Ensures that filename is opened with correct encoding parameter. + + This function uses charset_normalizer package, when available, for + determining the encoding of the file to be opened. When charset_normalizer + is not available, the function detects only UTF encodings, otherwise, ASCII + encoding is used as fallback. + """ + # Reads in the entire file. Robust detection of encoding. + # Correctly handles comments or late stage unicode characters + # gh-22871 + if charset_normalizer is not None: + encoding = charset_normalizer.from_path(filename).best().encoding + else: + # hint: install charset_normalizer for correct encoding handling + # No need to read the whole file for trying with startswith + nbytes = min(32, os.path.getsize(filename)) + with open(filename, 'rb') as fhandle: + raw = fhandle.read(nbytes) + if raw.startswith(codecs.BOM_UTF8): + encoding = 'UTF-8-SIG' + elif raw.startswith((codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE)): + encoding = 'UTF-32' + elif raw.startswith((codecs.BOM_LE, codecs.BOM_BE)): + encoding = 'UTF-16' + else: + # Fallback, without charset_normalizer + encoding = 'ascii' + return open(filename, mode, encoding=encoding) + + +def is_free_format(fname): + """Check if file is in free format Fortran.""" + # f90 allows both fixed and free format, assuming fixed unless + # signs of free format are detected. + result = False + if Path(fname).suffix.lower() in COMMON_FREE_EXTENSIONS: + result = True + with openhook(fname, 'r') as fhandle: + line = fhandle.readline() + n = 15 # the number of non-comment lines to scan for hints + if _has_f_header(line): + n = 0 + elif _has_f90_header(line): + n = 0 + result = True + while n > 0 and line: + if line[0] != '!' and line.strip(): + n -= 1 + if (line[0] != '\t' and _free_f90_start(line[:5])) or line[-2:-1] == '&': + result = True + break + line = fhandle.readline() + return result + + +# Read fortran (77,90) code +def readfortrancode(ffile, dowithline=show, istop=1): + """ + Read fortran codes from files and + 1) Get rid of comments, line continuations, and empty lines; lower cases. + 2) Call dowithline(line) on every line. + 3) Recursively call itself when statement \"include ''\" is met. + """ + global gotnextfile, filepositiontext, currentfilename, sourcecodeform, strictf77 + global beginpattern, quiet, verbose, dolowercase, include_paths + + if not istop: + saveglobals = gotnextfile, filepositiontext, currentfilename, sourcecodeform, strictf77,\ + beginpattern, quiet, verbose, dolowercase + if ffile == []: + return + localdolowercase = dolowercase + # cont: set to True when the content of the last line read + # indicates statement continuation + cont = False + finalline = '' + ll = '' + includeline = re.compile( + r'\s*include\s*(\'|")(?P[^\'"]*)(\'|")', re.I) + cont1 = re.compile(r'(?P.*)&\s*\Z') + cont2 = re.compile(r'(\s*&|)(?P.*)') + mline_mark = re.compile(r".*?'''") + if istop: + dowithline('', -1) + ll, l1 = '', '' + spacedigits = [' '] + [str(_m) for _m in range(10)] + filepositiontext = '' + fin = fileinput.FileInput(ffile, openhook=openhook) + while True: + try: + l = fin.readline() + except UnicodeDecodeError as msg: + raise Exception( + f'readfortrancode: reading {fin.filename()}#{fin.lineno()}' + f' failed with\n{msg}.\nIt is likely that installing charset_normalizer' + ' package will help f2py determine the input file encoding' + ' correctly.') + if not l: + break + if fin.isfirstline(): + filepositiontext = '' + currentfilename = fin.filename() + gotnextfile = 1 + l1 = l + strictf77 = 0 + sourcecodeform = 'fix' + ext = os.path.splitext(currentfilename)[1] + if Path(currentfilename).suffix.lower() in COMMON_FIXED_EXTENSIONS and \ + not (_has_f90_header(l) or _has_fix_header(l)): + strictf77 = 1 + elif is_free_format(currentfilename) and not _has_fix_header(l): + sourcecodeform = 'free' + if strictf77: + beginpattern = beginpattern77 + else: + beginpattern = beginpattern90 + outmess('\tReading file %s (format:%s%s)\n' + % (repr(currentfilename), sourcecodeform, + strictf77 and ',strict' or '')) + + l = l.expandtabs().replace('\xa0', ' ') + # Get rid of newline characters + while not l == '': + if l[-1] not in "\n\r\f": + break + l = l[:-1] + # Do not lower for directives, gh-2547, gh-27697, gh-26681 + is_f2py_directive = False + # Unconditionally remove comments + (l, rl) = split_by_unquoted(l, '!') + l += ' ' + if rl[:5].lower() == '!f2py': # f2py directive + l, _ = split_by_unquoted(l + 4 * ' ' + rl[5:], '!') + is_f2py_directive = True + if l.strip() == '': # Skip empty line + if sourcecodeform == 'free': + # In free form, a statement continues in the next line + # that is not a comment line [3.3.2.4^1], lines with + # blanks are comment lines [3.3.2.3^1]. Hence, the + # line continuation flag must retain its state. + pass + else: + # In fixed form, statement continuation is determined + # by a non-blank character at the 6-th position. Empty + # line indicates a start of a new statement + # [3.3.3.3^1]. Hence, the line continuation flag must + # be reset. + cont = False + continue + if sourcecodeform == 'fix': + if l[0] in ['*', 'c', '!', 'C', '#']: + if l[1:5].lower() == 'f2py': # f2py directive + l = ' ' + l[5:] + is_f2py_directive = True + else: # Skip comment line + cont = False + is_f2py_directive = False + continue + elif strictf77: + if len(l) > 72: + l = l[:72] + if l[0] not in spacedigits: + raise Exception('readfortrancode: Found non-(space,digit) char ' + 'in the first column.\n\tAre you sure that ' + 'this code is in fix form?\n\tline=%s' % repr(l)) + + if (not cont or strictf77) and (len(l) > 5 and not l[5] == ' '): + # Continuation of a previous line + ll = ll + l[6:] + finalline = '' + origfinalline = '' + else: + r = cont1.match(l) + if r: + l = r.group('line') # Continuation follows .. + if cont: + ll = ll + cont2.match(l).group('line') + finalline = '' + origfinalline = '' + else: + # clean up line beginning from possible digits. + l = ' ' + l[5:] + # f2py directives are already stripped by this point + if localdolowercase: + finalline = ll.lower() + else: + finalline = ll + origfinalline = ll + ll = l + + elif sourcecodeform == 'free': + if not cont and ext == '.pyf' and mline_mark.match(l): + l = l + '\n' + while True: + lc = fin.readline() + if not lc: + errmess( + 'Unexpected end of file when reading multiline\n') + break + l = l + lc + if mline_mark.match(lc): + break + l = l.rstrip() + r = cont1.match(l) + if r: + l = r.group('line') # Continuation follows .. + if cont: + ll = ll + cont2.match(l).group('line') + finalline = '' + origfinalline = '' + else: + if localdolowercase: + # only skip lowering for C style constructs + # gh-2547, gh-27697, gh-26681, gh-28014 + finalline = ll.lower() if not (is_f2py_directive and iscstyledirective(ll)) else ll + else: + finalline = ll + origfinalline = ll + ll = l + cont = (r is not None) + else: + raise ValueError( + "Flag sourcecodeform must be either 'fix' or 'free': %s" % repr(sourcecodeform)) + filepositiontext = 'Line #%d in %s:"%s"\n\t' % ( + fin.filelineno() - 1, currentfilename, l1) + m = includeline.match(origfinalline) + if m: + fn = m.group('name') + if os.path.isfile(fn): + readfortrancode(fn, dowithline=dowithline, istop=0) + else: + include_dirs = [ + os.path.dirname(currentfilename)] + include_paths + foundfile = 0 + for inc_dir in include_dirs: + fn1 = os.path.join(inc_dir, fn) + if os.path.isfile(fn1): + foundfile = 1 + readfortrancode(fn1, dowithline=dowithline, istop=0) + break + if not foundfile: + outmess('readfortrancode: could not find include file %s in %s. Ignoring.\n' % ( + repr(fn), os.pathsep.join(include_dirs))) + else: + dowithline(finalline) + l1 = ll + # Last line should never have an f2py directive anyway + if localdolowercase: + finalline = ll.lower() + else: + finalline = ll + origfinalline = ll + filepositiontext = 'Line #%d in %s:"%s"\n\t' % ( + fin.filelineno() - 1, currentfilename, l1) + m = includeline.match(origfinalline) + if m: + fn = m.group('name') + if os.path.isfile(fn): + readfortrancode(fn, dowithline=dowithline, istop=0) + else: + include_dirs = [os.path.dirname(currentfilename)] + include_paths + foundfile = 0 + for inc_dir in include_dirs: + fn1 = os.path.join(inc_dir, fn) + if os.path.isfile(fn1): + foundfile = 1 + readfortrancode(fn1, dowithline=dowithline, istop=0) + break + if not foundfile: + outmess('readfortrancode: could not find include file %s in %s. Ignoring.\n' % ( + repr(fn), os.pathsep.join(include_dirs))) + else: + dowithline(finalline) + filepositiontext = '' + fin.close() + if istop: + dowithline('', 1) + else: + gotnextfile, filepositiontext, currentfilename, sourcecodeform, strictf77,\ + beginpattern, quiet, verbose, dolowercase = saveglobals + +# Crack line +beforethisafter = r'\s*(?P%s(?=\s*(\b(%s)\b)))' + \ + r'\s*(?P(\b(%s)\b))' + \ + r'\s*(?P%s)\s*\Z' +## +fortrantypes = r'character|logical|integer|real|complex|double\s*(precision\s*(complex|)|complex)|type(?=\s*\([\w\s,=(*)]*\))|byte' +typespattern = re.compile( + beforethisafter % ('', fortrantypes, fortrantypes, '.*'), re.I), 'type' +typespattern4implicit = re.compile(beforethisafter % ( + '', fortrantypes + '|static|automatic|undefined', fortrantypes + '|static|automatic|undefined', '.*'), re.I) +# +functionpattern = re.compile(beforethisafter % ( + r'([a-z]+[\w\s(=*+-/)]*?|)', 'function', 'function', '.*'), re.I), 'begin' +subroutinepattern = re.compile(beforethisafter % ( + r'[a-z\s]*?', 'subroutine', 'subroutine', '.*'), re.I), 'begin' +# modulepattern=re.compile(beforethisafter%('[a-z\s]*?','module','module','.*'),re.I),'begin' +# +groupbegins77 = r'program|block\s*data' +beginpattern77 = re.compile( + beforethisafter % ('', groupbegins77, groupbegins77, '.*'), re.I), 'begin' +groupbegins90 = groupbegins77 + \ + r'|module(?!\s*procedure)|python\s*module|(abstract|)\s*interface|' + \ + r'type(?!\s*\()' +beginpattern90 = re.compile( + beforethisafter % ('', groupbegins90, groupbegins90, '.*'), re.I), 'begin' +groupends = (r'end|endprogram|endblockdata|endmodule|endpythonmodule|' + r'endinterface|endsubroutine|endfunction') +endpattern = re.compile( + beforethisafter % ('', groupends, groupends, '.*'), re.I), 'end' +# block, the Fortran 2008 construct needs special handling in the rest of the file +endifs = r'end\s*(if|do|where|select|while|forall|associate|' + \ + r'critical|enum|team)' +endifpattern = re.compile( + beforethisafter % (r'[\w]*?', endifs, endifs, '.*'), re.I), 'endif' +# +moduleprocedures = r'module\s*procedure' +moduleprocedurepattern = re.compile( + beforethisafter % ('', moduleprocedures, moduleprocedures, '.*'), re.I), \ + 'moduleprocedure' +implicitpattern = re.compile( + beforethisafter % ('', 'implicit', 'implicit', '.*'), re.I), 'implicit' +dimensionpattern = re.compile(beforethisafter % ( + '', 'dimension|virtual', 'dimension|virtual', '.*'), re.I), 'dimension' +externalpattern = re.compile( + beforethisafter % ('', 'external', 'external', '.*'), re.I), 'external' +optionalpattern = re.compile( + beforethisafter % ('', 'optional', 'optional', '.*'), re.I), 'optional' +requiredpattern = re.compile( + beforethisafter % ('', 'required', 'required', '.*'), re.I), 'required' +publicpattern = re.compile( + beforethisafter % ('', 'public', 'public', '.*'), re.I), 'public' +privatepattern = re.compile( + beforethisafter % ('', 'private', 'private', '.*'), re.I), 'private' +intrinsicpattern = re.compile( + beforethisafter % ('', 'intrinsic', 'intrinsic', '.*'), re.I), 'intrinsic' +intentpattern = re.compile(beforethisafter % ( + '', 'intent|depend|note|check', 'intent|depend|note|check', r'\s*\(.*?\).*'), re.I), 'intent' +parameterpattern = re.compile( + beforethisafter % ('', 'parameter', 'parameter', r'\s*\(.*'), re.I), 'parameter' +datapattern = re.compile( + beforethisafter % ('', 'data', 'data', '.*'), re.I), 'data' +callpattern = re.compile( + beforethisafter % ('', 'call', 'call', '.*'), re.I), 'call' +entrypattern = re.compile( + beforethisafter % ('', 'entry', 'entry', '.*'), re.I), 'entry' +callfunpattern = re.compile( + beforethisafter % ('', 'callfun', 'callfun', '.*'), re.I), 'callfun' +commonpattern = re.compile( + beforethisafter % ('', 'common', 'common', '.*'), re.I), 'common' +usepattern = re.compile( + beforethisafter % ('', 'use', 'use', '.*'), re.I), 'use' +containspattern = re.compile( + beforethisafter % ('', 'contains', 'contains', ''), re.I), 'contains' +formatpattern = re.compile( + beforethisafter % ('', 'format', 'format', '.*'), re.I), 'format' +# Non-fortran and f2py-specific statements +f2pyenhancementspattern = re.compile(beforethisafter % ('', 'threadsafe|fortranname|callstatement|callprotoargument|usercode|pymethoddef', + 'threadsafe|fortranname|callstatement|callprotoargument|usercode|pymethoddef', '.*'), re.I | re.S), 'f2pyenhancements' +multilinepattern = re.compile( + r"\s*(?P''')(?P.*?)(?P''')\s*\Z", re.S), 'multiline' +## + +def split_by_unquoted(line, characters): + """ + Splits the line into (line[:i], line[i:]), + where i is the index of first occurrence of one of the characters + not within quotes, or len(line) if no such index exists + """ + assert not (set('"\'') & set(characters)), "cannot split by unquoted quotes" + r = re.compile( + r"\A(?P({single_quoted}|{double_quoted}|{not_quoted})*)" + r"(?P{char}.*)\Z".format( + not_quoted="[^\"'{}]".format(re.escape(characters)), + char="[{}]".format(re.escape(characters)), + single_quoted=r"('([^'\\]|(\\.))*')", + double_quoted=r'("([^"\\]|(\\.))*")')) + m = r.match(line) + if m: + d = m.groupdict() + return (d["before"], d["after"]) + return (line, "") + +def _simplifyargs(argsline): + a = [] + for n in markoutercomma(argsline).split('@,@'): + for r in '(),': + n = n.replace(r, '_') + a.append(n) + return ','.join(a) + +crackline_re_1 = re.compile(r'\s*(?P\b[a-z]+\w*\b)\s*=.*', re.I) +crackline_bind_1 = re.compile(r'\s*(?P\b[a-z]+\w*\b)\s*=.*', re.I) +crackline_bindlang = re.compile(r'\s*bind\(\s*(?P[^,]+)\s*,\s*name\s*=\s*"(?P[^"]+)"\s*\)', re.I) + +def crackline(line, reset=0): + """ + reset=-1 --- initialize + reset=0 --- crack the line + reset=1 --- final check if mismatch of blocks occurred + + Cracked data is saved in grouplist[0]. + """ + global beginpattern, groupcounter, groupname, groupcache, grouplist + global filepositiontext, currentfilename, neededmodule, expectbegin + global skipblocksuntil, skipemptyends, previous_context, gotnextfile + + _, has_semicolon = split_by_unquoted(line, ";") + if has_semicolon and not (f2pyenhancementspattern[0].match(line) or + multilinepattern[0].match(line)): + # XXX: non-zero reset values need testing + assert reset == 0, repr(reset) + # split line on unquoted semicolons + line, semicolon_line = split_by_unquoted(line, ";") + while semicolon_line: + crackline(line, reset) + line, semicolon_line = split_by_unquoted(semicolon_line[1:], ";") + crackline(line, reset) + return + if reset < 0: + groupcounter = 0 + groupname = {groupcounter: ''} + groupcache = {groupcounter: {}} + grouplist = {groupcounter: []} + groupcache[groupcounter]['body'] = [] + groupcache[groupcounter]['vars'] = {} + groupcache[groupcounter]['block'] = '' + groupcache[groupcounter]['name'] = '' + neededmodule = -1 + skipblocksuntil = -1 + return + if reset > 0: + fl = 0 + if f77modulename and neededmodule == groupcounter: + fl = 2 + while groupcounter > fl: + outmess('crackline: groupcounter=%s groupname=%s\n' % + (repr(groupcounter), repr(groupname))) + outmess( + 'crackline: Mismatch of blocks encountered. Trying to fix it by assuming "end" statement.\n') + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 + if f77modulename and neededmodule == groupcounter: + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 # end interface + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 # end module + neededmodule = -1 + return + if line == '': + return + flag = 0 + for pat in [dimensionpattern, externalpattern, intentpattern, optionalpattern, + requiredpattern, + parameterpattern, datapattern, publicpattern, privatepattern, + intrinsicpattern, + endifpattern, endpattern, + formatpattern, + beginpattern, functionpattern, subroutinepattern, + implicitpattern, typespattern, commonpattern, + callpattern, usepattern, containspattern, + entrypattern, + f2pyenhancementspattern, + multilinepattern, + moduleprocedurepattern + ]: + m = pat[0].match(line) + if m: + break + flag = flag + 1 + if not m: + re_1 = crackline_re_1 + if 0 <= skipblocksuntil <= groupcounter: + return + if 'externals' in groupcache[groupcounter]: + for name in groupcache[groupcounter]['externals']: + if name in invbadnames: + name = invbadnames[name] + if 'interfaced' in groupcache[groupcounter] and name in groupcache[groupcounter]['interfaced']: + continue + m1 = re.match( + r'(?P[^"]*)\b%s\b\s*@\(@(?P[^@]*)@\)@.*\Z' % name, markouterparen(line), re.I) + if m1: + m2 = re_1.match(m1.group('before')) + a = _simplifyargs(m1.group('args')) + if m2: + line = 'callfun %s(%s) result (%s)' % ( + name, a, m2.group('result')) + else: + line = 'callfun %s(%s)' % (name, a) + m = callfunpattern[0].match(line) + if not m: + outmess( + 'crackline: could not resolve function call for line=%s.\n' % repr(line)) + return + analyzeline(m, 'callfun', line) + return + if verbose > 1 or (verbose == 1 and currentfilename.lower().endswith('.pyf')): + previous_context = None + outmess('crackline:%d: No pattern for line\n' % (groupcounter)) + return + elif pat[1] == 'end': + if 0 <= skipblocksuntil < groupcounter: + groupcounter = groupcounter - 1 + if skipblocksuntil <= groupcounter: + return + if groupcounter <= 0: + raise Exception('crackline: groupcounter(=%s) is nonpositive. ' + 'Check the blocks.' + % (groupcounter)) + m1 = beginpattern[0].match(line) + if (m1) and (not m1.group('this') == groupname[groupcounter]): + raise Exception('crackline: End group %s does not match with ' + 'previous Begin group %s\n\t%s' % + (repr(m1.group('this')), repr(groupname[groupcounter]), + filepositiontext) + ) + if skipblocksuntil == groupcounter: + skipblocksuntil = -1 + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 + if not skipemptyends: + expectbegin = 1 + elif pat[1] == 'begin': + if 0 <= skipblocksuntil <= groupcounter: + groupcounter = groupcounter + 1 + return + gotnextfile = 0 + analyzeline(m, pat[1], line) + expectbegin = 0 + elif pat[1] == 'endif': + pass + elif pat[1] == 'moduleprocedure': + analyzeline(m, pat[1], line) + elif pat[1] == 'contains': + if ignorecontains: + return + if 0 <= skipblocksuntil <= groupcounter: + return + skipblocksuntil = groupcounter + else: + if 0 <= skipblocksuntil <= groupcounter: + return + analyzeline(m, pat[1], line) + + +def markouterparen(line): + l = '' + f = 0 + for c in line: + if c == '(': + f = f + 1 + if f == 1: + l = l + '@(@' + continue + elif c == ')': + f = f - 1 + if f == 0: + l = l + '@)@' + continue + l = l + c + return l + + +def markoutercomma(line, comma=','): + l = '' + f = 0 + before, after = split_by_unquoted(line, comma + '()') + l += before + while after: + if (after[0] == comma) and (f == 0): + l += '@' + comma + '@' + else: + l += after[0] + if after[0] == '(': + f += 1 + elif after[0] == ')': + f -= 1 + before, after = split_by_unquoted(after[1:], comma + '()') + l += before + assert not f, repr((f, line, l)) + return l + +def unmarkouterparen(line): + r = line.replace('@(@', '(').replace('@)@', ')') + return r + + +def appenddecl(decl, decl2, force=1): + if not decl: + decl = {} + if not decl2: + return decl + if decl is decl2: + return decl + for k in list(decl2.keys()): + if k == 'typespec': + if force or k not in decl: + decl[k] = decl2[k] + elif k == 'attrspec': + for l in decl2[k]: + decl = setattrspec(decl, l, force) + elif k == 'kindselector': + decl = setkindselector(decl, decl2[k], force) + elif k == 'charselector': + decl = setcharselector(decl, decl2[k], force) + elif k in ['=', 'typename']: + if force or k not in decl: + decl[k] = decl2[k] + elif k == 'note': + pass + elif k in ['intent', 'check', 'dimension', 'optional', + 'required', 'depend']: + errmess('appenddecl: "%s" not implemented.\n' % k) + else: + raise Exception('appenddecl: Unknown variable definition key: ' + + str(k)) + return decl + +selectpattern = re.compile( + r'\s*(?P(@\(@.*?@\)@|\*[\d*]+|\*\s*@\(@.*?@\)@|))(?P.*)\Z', re.I) +typedefpattern = re.compile( + r'(?:,(?P[\w(),]+))?(::)?(?P\b[a-z$_][\w$]*\b)' + r'(?:\((?P[\w,]*)\))?\Z', re.I) +nameargspattern = re.compile( + r'\s*(?P\b[\w$]+\b)\s*(@\(@\s*(?P[\w\s,]*)\s*@\)@|)\s*((result(\s*@\(@\s*(?P\b[\w$]+\b)\s*@\)@|))|(bind\s*@\(@\s*(?P(?:(?!@\)@).)*)\s*@\)@))*\s*\Z', re.I) +operatorpattern = re.compile( + r'\s*(?P(operator|assignment))' + r'@\(@\s*(?P[^)]+)\s*@\)@\s*\Z', re.I) +callnameargspattern = re.compile( + r'\s*(?P\b[\w$]+\b)\s*@\(@\s*(?P.*)\s*@\)@\s*\Z', re.I) +real16pattern = re.compile( + r'([-+]?(?:\d+(?:\.\d*)?|\d*\.\d+))[dD]((?:[-+]?\d+)?)') +real8pattern = re.compile( + r'([-+]?((?:\d+(?:\.\d*)?|\d*\.\d+))[eE]((?:[-+]?\d+)?)|(\d+\.\d*))') + +_intentcallbackpattern = re.compile(r'intent\s*\(.*?\bcallback\b', re.I) + + +def _is_intent_callback(vdecl): + for a in vdecl.get('attrspec', []): + if _intentcallbackpattern.match(a): + return 1 + return 0 + + +def _resolvetypedefpattern(line): + line = ''.join(line.split()) # removes whitespace + m1 = typedefpattern.match(line) + print(line, m1) + if m1: + attrs = m1.group('attributes') + attrs = [a.lower() for a in attrs.split(',')] if attrs else [] + return m1.group('name'), attrs, m1.group('params') + return None, [], None + +def parse_name_for_bind(line): + pattern = re.compile(r'bind\(\s*(?P[^,]+)(?:\s*,\s*name\s*=\s*["\'](?P[^"\']+)["\']\s*)?\)', re.I) + match = pattern.search(line) + bind_statement = None + if match: + bind_statement = match.group(0) + # Remove the 'bind' construct from the line. + line = line[:match.start()] + line[match.end():] + return line, bind_statement + +def _resolvenameargspattern(line): + line, bind_cname = parse_name_for_bind(line) + line = markouterparen(line) + m1 = nameargspattern.match(line) + if m1: + return m1.group('name'), m1.group('args'), m1.group('result'), bind_cname + m1 = operatorpattern.match(line) + if m1: + name = m1.group('scheme') + '(' + m1.group('name') + ')' + return name, [], None, None + m1 = callnameargspattern.match(line) + if m1: + return m1.group('name'), m1.group('args'), None, None + return None, [], None, None + + +def analyzeline(m, case, line): + """ + Reads each line in the input file in sequence and updates global vars. + + Effectively reads and collects information from the input file to the + global variable groupcache, a dictionary containing info about each part + of the fortran module. + + At the end of analyzeline, information is filtered into the correct dict + keys, but parameter values and dimensions are not yet interpreted. + """ + global groupcounter, groupname, groupcache, grouplist, filepositiontext + global currentfilename, f77modulename, neededinterface, neededmodule + global expectbegin, gotnextfile, previous_context + + block = m.group('this') + if case != 'multiline': + previous_context = None + if expectbegin and case not in ['begin', 'call', 'callfun', 'type'] \ + and not skipemptyends and groupcounter < 1: + newname = os.path.basename(currentfilename).split('.')[0] + outmess( + 'analyzeline: no group yet. Creating program group with name "%s".\n' % newname) + gotnextfile = 0 + groupcounter = groupcounter + 1 + groupname[groupcounter] = 'program' + groupcache[groupcounter] = {} + grouplist[groupcounter] = [] + groupcache[groupcounter]['body'] = [] + groupcache[groupcounter]['vars'] = {} + groupcache[groupcounter]['block'] = 'program' + groupcache[groupcounter]['name'] = newname + groupcache[groupcounter]['from'] = 'fromsky' + expectbegin = 0 + if case in ['begin', 'call', 'callfun']: + # Crack line => block,name,args,result + block = block.lower() + if re.match(r'block\s*data', block, re.I): + block = 'block data' + elif re.match(r'python\s*module', block, re.I): + block = 'python module' + elif re.match(r'abstract\s*interface', block, re.I): + block = 'abstract interface' + if block == 'type': + name, attrs, _ = _resolvetypedefpattern(m.group('after')) + groupcache[groupcounter]['vars'][name] = dict(attrspec = attrs) + args = [] + result = None + else: + name, args, result, bindcline = _resolvenameargspattern(m.group('after')) + if name is None: + if block == 'block data': + name = '_BLOCK_DATA_' + else: + name = '' + if block not in ['interface', 'block data', 'abstract interface']: + outmess('analyzeline: No name/args pattern found for line.\n') + + previous_context = (block, name, groupcounter) + if args: + args = rmbadname([x.strip() + for x in markoutercomma(args).split('@,@')]) + else: + args = [] + if '' in args: + while '' in args: + args.remove('') + outmess( + 'analyzeline: argument list is malformed (missing argument).\n') + + # end of crack line => block,name,args,result + needmodule = 0 + needinterface = 0 + + if case in ['call', 'callfun']: + needinterface = 1 + if 'args' not in groupcache[groupcounter]: + return + if name not in groupcache[groupcounter]['args']: + return + for it in grouplist[groupcounter]: + if it['name'] == name: + return + if name in groupcache[groupcounter]['interfaced']: + return + block = {'call': 'subroutine', 'callfun': 'function'}[case] + if f77modulename and neededmodule == -1 and groupcounter <= 1: + neededmodule = groupcounter + 2 + needmodule = 1 + if block not in ['interface', 'abstract interface']: + needinterface = 1 + # Create new block(s) + groupcounter = groupcounter + 1 + groupcache[groupcounter] = {} + grouplist[groupcounter] = [] + if needmodule: + if verbose > 1: + outmess('analyzeline: Creating module block %s\n' % + repr(f77modulename), 0) + groupname[groupcounter] = 'module' + groupcache[groupcounter]['block'] = 'python module' + groupcache[groupcounter]['name'] = f77modulename + groupcache[groupcounter]['from'] = '' + groupcache[groupcounter]['body'] = [] + groupcache[groupcounter]['externals'] = [] + groupcache[groupcounter]['interfaced'] = [] + groupcache[groupcounter]['vars'] = {} + groupcounter = groupcounter + 1 + groupcache[groupcounter] = {} + grouplist[groupcounter] = [] + if needinterface: + if verbose > 1: + outmess('analyzeline: Creating additional interface block (groupcounter=%s).\n' % ( + groupcounter), 0) + groupname[groupcounter] = 'interface' + groupcache[groupcounter]['block'] = 'interface' + groupcache[groupcounter]['name'] = 'unknown_interface' + groupcache[groupcounter]['from'] = '%s:%s' % ( + groupcache[groupcounter - 1]['from'], groupcache[groupcounter - 1]['name']) + groupcache[groupcounter]['body'] = [] + groupcache[groupcounter]['externals'] = [] + groupcache[groupcounter]['interfaced'] = [] + groupcache[groupcounter]['vars'] = {} + groupcounter = groupcounter + 1 + groupcache[groupcounter] = {} + grouplist[groupcounter] = [] + groupname[groupcounter] = block + groupcache[groupcounter]['block'] = block + if not name: + name = 'unknown_' + block.replace(' ', '_') + groupcache[groupcounter]['prefix'] = m.group('before') + groupcache[groupcounter]['name'] = rmbadname1(name) + groupcache[groupcounter]['result'] = result + if groupcounter == 1: + groupcache[groupcounter]['from'] = currentfilename + else: + if f77modulename and groupcounter == 3: + groupcache[groupcounter]['from'] = '%s:%s' % ( + groupcache[groupcounter - 1]['from'], currentfilename) + else: + groupcache[groupcounter]['from'] = '%s:%s' % ( + groupcache[groupcounter - 1]['from'], groupcache[groupcounter - 1]['name']) + for k in list(groupcache[groupcounter].keys()): + if not groupcache[groupcounter][k]: + del groupcache[groupcounter][k] + + groupcache[groupcounter]['args'] = args + groupcache[groupcounter]['body'] = [] + groupcache[groupcounter]['externals'] = [] + groupcache[groupcounter]['interfaced'] = [] + groupcache[groupcounter]['vars'] = {} + groupcache[groupcounter]['entry'] = {} + # end of creation + if block == 'type': + groupcache[groupcounter]['varnames'] = [] + + if case in ['call', 'callfun']: # set parents variables + if name not in groupcache[groupcounter - 2]['externals']: + groupcache[groupcounter - 2]['externals'].append(name) + groupcache[groupcounter]['vars'] = copy.deepcopy( + groupcache[groupcounter - 2]['vars']) + try: + del groupcache[groupcounter]['vars'][name][ + groupcache[groupcounter]['vars'][name]['attrspec'].index('external')] + except Exception: + pass + if block in ['function', 'subroutine']: # set global attributes + # name is fortran name + if bindcline: + bindcdat = re.search(crackline_bindlang, bindcline) + if bindcdat: + groupcache[groupcounter]['bindlang'] = {name : {}} + groupcache[groupcounter]['bindlang'][name]["lang"] = bindcdat.group('lang') + if bindcdat.group('lang_name'): + groupcache[groupcounter]['bindlang'][name]["name"] = bindcdat.group('lang_name') + try: + groupcache[groupcounter]['vars'][name] = appenddecl( + groupcache[groupcounter]['vars'][name], groupcache[groupcounter - 2]['vars']['']) + except Exception: + pass + if case == 'callfun': # return type + if result and result in groupcache[groupcounter]['vars']: + if not name == result: + groupcache[groupcounter]['vars'][name] = appenddecl( + groupcache[groupcounter]['vars'][name], groupcache[groupcounter]['vars'][result]) + # if groupcounter>1: # name is interfaced + try: + groupcache[groupcounter - 2]['interfaced'].append(name) + except Exception: + pass + if block == 'function': + t = typespattern[0].match(m.group('before') + ' ' + name) + if t: + typespec, selector, attr, edecl = cracktypespec0( + t.group('this'), t.group('after')) + updatevars(typespec, selector, attr, edecl) + + if case in ['call', 'callfun']: + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 # end routine + grouplist[groupcounter - 1].append(groupcache[groupcounter]) + grouplist[groupcounter - 1][-1]['body'] = grouplist[groupcounter] + del grouplist[groupcounter] + groupcounter = groupcounter - 1 # end interface + + elif case == 'entry': + name, args, result, _= _resolvenameargspattern(m.group('after')) + if name is not None: + if args: + args = rmbadname([x.strip() + for x in markoutercomma(args).split('@,@')]) + else: + args = [] + assert result is None, repr(result) + groupcache[groupcounter]['entry'][name] = args + previous_context = ('entry', name, groupcounter) + elif case == 'type': + typespec, selector, attr, edecl = cracktypespec0( + block, m.group('after')) + last_name = updatevars(typespec, selector, attr, edecl) + if last_name is not None: + previous_context = ('variable', last_name, groupcounter) + elif case in ['dimension', 'intent', 'optional', 'required', 'external', 'public', 'private', 'intrinsic']: + edecl = groupcache[groupcounter]['vars'] + ll = m.group('after').strip() + i = ll.find('::') + if i < 0 and case == 'intent': + i = markouterparen(ll).find('@)@') - 2 + ll = ll[:i + 1] + '::' + ll[i + 1:] + i = ll.find('::') + if ll[i:] == '::' and 'args' in groupcache[groupcounter]: + outmess('All arguments will have attribute %s%s\n' % + (m.group('this'), ll[:i])) + ll = ll + ','.join(groupcache[groupcounter]['args']) + if i < 0: + i = 0 + pl = '' + else: + pl = ll[:i].strip() + ll = ll[i + 2:] + ch = markoutercomma(pl).split('@,@') + if len(ch) > 1: + pl = ch[0] + outmess('analyzeline: cannot handle multiple attributes without type specification. Ignoring %r.\n' % ( + ','.join(ch[1:]))) + last_name = None + + for e in [x.strip() for x in markoutercomma(ll).split('@,@')]: + m1 = namepattern.match(e) + if not m1: + if case in ['public', 'private']: + k = '' + else: + print(m.groupdict()) + outmess('analyzeline: no name pattern found in %s statement for %s. Skipping.\n' % ( + case, repr(e))) + continue + else: + k = rmbadname1(m1.group('name')) + if case in ['public', 'private'] and \ + (k == 'operator' or k == 'assignment'): + k += m1.group('after') + if k not in edecl: + edecl[k] = {} + if case == 'dimension': + ap = case + m1.group('after') + if case == 'intent': + ap = m.group('this') + pl + if _intentcallbackpattern.match(ap): + if k not in groupcache[groupcounter]['args']: + if groupcounter > 1: + if '__user__' not in groupcache[groupcounter - 2]['name']: + outmess( + 'analyzeline: missing __user__ module (could be nothing)\n') + # fixes ticket 1693 + if k != groupcache[groupcounter]['name']: + outmess('analyzeline: appending intent(callback) %s' + ' to %s arguments\n' % (k, groupcache[groupcounter]['name'])) + groupcache[groupcounter]['args'].append(k) + else: + errmess( + 'analyzeline: intent(callback) %s is ignored\n' % (k)) + else: + errmess('analyzeline: intent(callback) %s is already' + ' in argument list\n' % (k)) + if case in ['optional', 'required', 'public', 'external', 'private', 'intrinsic']: + ap = case + if 'attrspec' in edecl[k]: + edecl[k]['attrspec'].append(ap) + else: + edecl[k]['attrspec'] = [ap] + if case == 'external': + if groupcache[groupcounter]['block'] == 'program': + outmess('analyzeline: ignoring program arguments\n') + continue + if k not in groupcache[groupcounter]['args']: + continue + if 'externals' not in groupcache[groupcounter]: + groupcache[groupcounter]['externals'] = [] + groupcache[groupcounter]['externals'].append(k) + last_name = k + groupcache[groupcounter]['vars'] = edecl + if last_name is not None: + previous_context = ('variable', last_name, groupcounter) + elif case == 'moduleprocedure': + groupcache[groupcounter]['implementedby'] = \ + [x.strip() for x in m.group('after').split(',')] + elif case == 'parameter': + edecl = groupcache[groupcounter]['vars'] + ll = m.group('after').strip()[1:-1] + last_name = None + for e in markoutercomma(ll).split('@,@'): + try: + k, initexpr = [x.strip() for x in e.split('=')] + except Exception: + outmess( + 'analyzeline: could not extract name,expr in parameter statement "%s" of "%s"\n' % (e, ll)) + continue + params = get_parameters(edecl) + k = rmbadname1(k) + if k not in edecl: + edecl[k] = {} + if '=' in edecl[k] and (not edecl[k]['='] == initexpr): + outmess('analyzeline: Overwriting the value of parameter "%s" ("%s") with "%s".\n' % ( + k, edecl[k]['='], initexpr)) + t = determineexprtype(initexpr, params) + if t: + if t.get('typespec') == 'real': + tt = list(initexpr) + for m in real16pattern.finditer(initexpr): + tt[m.start():m.end()] = list( + initexpr[m.start():m.end()].lower().replace('d', 'e')) + initexpr = ''.join(tt) + elif t.get('typespec') == 'complex': + initexpr = initexpr[1:].lower().replace('d', 'e').\ + replace(',', '+1j*(') + try: + v = eval(initexpr, {}, params) + except (SyntaxError, NameError, TypeError) as msg: + errmess('analyzeline: Failed to evaluate %r. Ignoring: %s\n' + % (initexpr, msg)) + continue + edecl[k]['='] = repr(v) + if 'attrspec' in edecl[k]: + edecl[k]['attrspec'].append('parameter') + else: + edecl[k]['attrspec'] = ['parameter'] + last_name = k + groupcache[groupcounter]['vars'] = edecl + if last_name is not None: + previous_context = ('variable', last_name, groupcounter) + elif case == 'implicit': + if m.group('after').strip().lower() == 'none': + groupcache[groupcounter]['implicit'] = None + elif m.group('after'): + if 'implicit' in groupcache[groupcounter]: + impl = groupcache[groupcounter]['implicit'] + else: + impl = {} + if impl is None: + outmess( + 'analyzeline: Overwriting earlier "implicit none" statement.\n') + impl = {} + for e in markoutercomma(m.group('after')).split('@,@'): + decl = {} + m1 = re.match( + r'\s*(?P.*?)\s*(\(\s*(?P[a-z-, ]+)\s*\)\s*|)\Z', e, re.I) + if not m1: + outmess( + 'analyzeline: could not extract info of implicit statement part "%s"\n' % (e)) + continue + m2 = typespattern4implicit.match(m1.group('this')) + if not m2: + outmess( + 'analyzeline: could not extract types pattern of implicit statement part "%s"\n' % (e)) + continue + typespec, selector, attr, edecl = cracktypespec0( + m2.group('this'), m2.group('after')) + kindselect, charselect, typename = cracktypespec( + typespec, selector) + decl['typespec'] = typespec + decl['kindselector'] = kindselect + decl['charselector'] = charselect + decl['typename'] = typename + for k in list(decl.keys()): + if not decl[k]: + del decl[k] + for r in markoutercomma(m1.group('after')).split('@,@'): + if '-' in r: + try: + begc, endc = [x.strip() for x in r.split('-')] + except Exception: + outmess( + 'analyzeline: expected "-" instead of "%s" in range list of implicit statement\n' % r) + continue + else: + begc = endc = r.strip() + if not len(begc) == len(endc) == 1: + outmess( + 'analyzeline: expected "-" instead of "%s" in range list of implicit statement (2)\n' % r) + continue + for o in range(ord(begc), ord(endc) + 1): + impl[chr(o)] = decl + groupcache[groupcounter]['implicit'] = impl + elif case == 'data': + ll = [] + dl = '' + il = '' + f = 0 + fc = 1 + inp = 0 + for c in m.group('after'): + if not inp: + if c == "'": + fc = not fc + if c == '/' and fc: + f = f + 1 + continue + if c == '(': + inp = inp + 1 + elif c == ')': + inp = inp - 1 + if f == 0: + dl = dl + c + elif f == 1: + il = il + c + elif f == 2: + dl = dl.strip() + if dl.startswith(','): + dl = dl[1:].strip() + ll.append([dl, il]) + dl = c + il = '' + f = 0 + if f == 2: + dl = dl.strip() + if dl.startswith(','): + dl = dl[1:].strip() + ll.append([dl, il]) + vars = groupcache[groupcounter].get('vars', {}) + last_name = None + for l in ll: + l[0], l[1] = l[0].strip(), l[1].strip() + if l[0].startswith(','): + l[0] = l[0][1:] + if l[0].startswith('('): + outmess('analyzeline: implied-DO list "%s" is not supported. Skipping.\n' % l[0]) + continue + for idx, v in enumerate(rmbadname([x.strip() for x in markoutercomma(l[0]).split('@,@')])): + if v.startswith('('): + outmess('analyzeline: implied-DO list "%s" is not supported. Skipping.\n' % v) + # XXX: subsequent init expressions may get wrong values. + # Ignoring since data statements are irrelevant for + # wrapping. + continue + if '!' in l[1]: + # Fixes gh-24746 pyf generation + # XXX: This essentially ignores the value for generating the pyf which is fine: + # integer dimension(3) :: mytab + # common /mycom/ mytab + # Since in any case it is initialized in the Fortran code + outmess('Comment line in declaration "%s" is not supported. Skipping.\n' % l[1]) + continue + vars.setdefault(v, {}) + vtype = vars[v].get('typespec') + vdim = getdimension(vars[v]) + matches = re.findall(r"\(.*?\)", l[1]) if vtype == 'complex' else l[1].split(',') + try: + new_val = "(/{}/)".format(", ".join(matches)) if vdim else matches[idx] + except IndexError: + # gh-24746 + # Runs only if above code fails. Fixes the line + # DATA IVAR1, IVAR2, IVAR3, IVAR4, EVAR5 /4*0,0.0D0/ + # by expanding to ['0', '0', '0', '0', '0.0d0'] + if any("*" in m for m in matches): + expanded_list = [] + for match in matches: + if "*" in match: + try: + multiplier, value = match.split("*") + expanded_list.extend([value.strip()] * int(multiplier)) + except ValueError: # if int(multiplier) fails + expanded_list.append(match.strip()) + else: + expanded_list.append(match.strip()) + matches = expanded_list + new_val = "(/{}/)".format(", ".join(matches)) if vdim else matches[idx] + current_val = vars[v].get('=') + if current_val and (current_val != new_val): + outmess('analyzeline: changing init expression of "%s" ("%s") to "%s"\n' % (v, current_val, new_val)) + vars[v]['='] = new_val + last_name = v + groupcache[groupcounter]['vars'] = vars + if last_name: + previous_context = ('variable', last_name, groupcounter) + elif case == 'common': + line = m.group('after').strip() + if not line[0] == '/': + line = '//' + line + + cl = [] + [_, bn, ol] = re.split('/', line, maxsplit=2) + bn = bn.strip() + if not bn: + bn = '_BLNK_' + cl.append([bn, ol]) + commonkey = {} + if 'common' in groupcache[groupcounter]: + commonkey = groupcache[groupcounter]['common'] + for c in cl: + if c[0] not in commonkey: + commonkey[c[0]] = [] + for i in [x.strip() for x in markoutercomma(c[1]).split('@,@')]: + if i: + commonkey[c[0]].append(i) + groupcache[groupcounter]['common'] = commonkey + previous_context = ('common', bn, groupcounter) + elif case == 'use': + m1 = re.match( + r'\A\s*(?P\b\w+\b)\s*((,(\s*\bonly\b\s*:|(?P))\s*(?P.*))|)\s*\Z', m.group('after'), re.I) + if m1: + mm = m1.groupdict() + if 'use' not in groupcache[groupcounter]: + groupcache[groupcounter]['use'] = {} + name = m1.group('name') + groupcache[groupcounter]['use'][name] = {} + isonly = 0 + if 'list' in mm and mm['list'] is not None: + if 'notonly' in mm and mm['notonly'] is None: + isonly = 1 + groupcache[groupcounter]['use'][name]['only'] = isonly + ll = [x.strip() for x in mm['list'].split(',')] + rl = {} + for l in ll: + if '=' in l: + m2 = re.match( + r'\A\s*(?P\b\w+\b)\s*=\s*>\s*(?P\b\w+\b)\s*\Z', l, re.I) + if m2: + rl[m2.group('local').strip()] = m2.group( + 'use').strip() + else: + outmess( + 'analyzeline: Not local=>use pattern found in %s\n' % repr(l)) + else: + rl[l] = l + groupcache[groupcounter]['use'][name]['map'] = rl + else: + pass + else: + print(m.groupdict()) + outmess('analyzeline: Could not crack the use statement.\n') + elif case in ['f2pyenhancements']: + if 'f2pyenhancements' not in groupcache[groupcounter]: + groupcache[groupcounter]['f2pyenhancements'] = {} + d = groupcache[groupcounter]['f2pyenhancements'] + if m.group('this') == 'usercode' and 'usercode' in d: + if isinstance(d['usercode'], str): + d['usercode'] = [d['usercode']] + d['usercode'].append(m.group('after')) + else: + d[m.group('this')] = m.group('after') + elif case == 'multiline': + if previous_context is None: + if verbose: + outmess('analyzeline: No context for multiline block.\n') + return + gc = groupcounter + appendmultiline(groupcache[gc], + previous_context[:2], + m.group('this')) + else: + if verbose > 1: + print(m.groupdict()) + outmess('analyzeline: No code implemented for line.\n') + + +def appendmultiline(group, context_name, ml): + if 'f2pymultilines' not in group: + group['f2pymultilines'] = {} + d = group['f2pymultilines'] + if context_name not in d: + d[context_name] = [] + d[context_name].append(ml) + return + + +def cracktypespec0(typespec, ll): + selector = None + attr = None + if re.match(r'double\s*complex', typespec, re.I): + typespec = 'double complex' + elif re.match(r'double\s*precision', typespec, re.I): + typespec = 'double precision' + else: + typespec = typespec.strip().lower() + m1 = selectpattern.match(markouterparen(ll)) + if not m1: + outmess( + 'cracktypespec0: no kind/char_selector pattern found for line.\n') + return + d = m1.groupdict() + for k in list(d.keys()): + d[k] = unmarkouterparen(d[k]) + if typespec in ['complex', 'integer', 'logical', 'real', 'character', 'type']: + selector = d['this'] + ll = d['after'] + i = ll.find('::') + if i >= 0: + attr = ll[:i].strip() + ll = ll[i + 2:] + return typespec, selector, attr, ll +##### +namepattern = re.compile(r'\s*(?P\b\w+\b)\s*(?P.*)\s*\Z', re.I) +kindselector = re.compile( + r'\s*(\(\s*(kind\s*=)?\s*(?P.*)\s*\)|\*\s*(?P.*?))\s*\Z', re.I) +charselector = re.compile( + r'\s*(\((?P.*)\)|\*\s*(?P.*))\s*\Z', re.I) +lenkindpattern = re.compile( + r'\s*(kind\s*=\s*(?P.*?)\s*(@,@\s*len\s*=\s*(?P.*)|)' + r'|(len\s*=\s*|)(?P.*?)\s*(@,@\s*(kind\s*=\s*|)(?P.*)' + r'|(f2py_len\s*=\s*(?P.*))|))\s*\Z', re.I) +lenarraypattern = re.compile( + r'\s*(@\(@\s*(?!/)\s*(?P.*?)\s*@\)@\s*\*\s*(?P.*?)|(\*\s*(?P.*?)|)\s*(@\(@\s*(?!/)\s*(?P.*?)\s*@\)@|))\s*(=\s*(?P.*?)|(@\(@|)/\s*(?P.*?)\s*/(@\)@|)|)\s*\Z', re.I) + + +def removespaces(expr): + expr = expr.strip() + if len(expr) <= 1: + return expr + expr2 = expr[0] + for i in range(1, len(expr) - 1): + if (expr[i] == ' ' and + ((expr[i + 1] in "()[]{}=+-/* ") or + (expr[i - 1] in "()[]{}=+-/* "))): + continue + expr2 = expr2 + expr[i] + expr2 = expr2 + expr[-1] + return expr2 + + +def markinnerspaces(line): + """ + The function replace all spaces in the input variable line which are + surrounded with quotation marks, with the triplet "@_@". + + For instance, for the input "a 'b c'" the function returns "a 'b@_@c'" + + Parameters + ---------- + line : str + + Returns + ------- + str + + """ + fragment = '' + inside = False + current_quote = None + escaped = '' + for c in line: + if escaped == '\\' and c in ['\\', '\'', '"']: + fragment += c + escaped = c + continue + if not inside and c in ['\'', '"']: + current_quote = c + if c == current_quote: + inside = not inside + elif c == ' ' and inside: + fragment += '@_@' + continue + fragment += c + escaped = c # reset to non-backslash + return fragment + + +def updatevars(typespec, selector, attrspec, entitydecl): + """ + Returns last_name, the variable name without special chars, parenthesis + or dimension specifiers. + + Alters groupcache to add the name, typespec, attrspec (and possibly value) + of current variable. + """ + global groupcache, groupcounter + + last_name = None + kindselect, charselect, typename = cracktypespec(typespec, selector) + # Clean up outer commas, whitespace and undesired chars from attrspec + if attrspec: + attrspec = [x.strip() for x in markoutercomma(attrspec).split('@,@')] + l = [] + c = re.compile(r'(?P[a-zA-Z]+)') + for a in attrspec: + if not a: + continue + m = c.match(a) + if m: + s = m.group('start').lower() + a = s + a[len(s):] + l.append(a) + attrspec = l + el = [x.strip() for x in markoutercomma(entitydecl).split('@,@')] + el1 = [] + for e in el: + for e1 in [x.strip() for x in markoutercomma(removespaces(markinnerspaces(e)), comma=' ').split('@ @')]: + if e1: + el1.append(e1.replace('@_@', ' ')) + for e in el1: + m = namepattern.match(e) + if not m: + outmess( + 'updatevars: no name pattern found for entity=%s. Skipping.\n' % (repr(e))) + continue + ename = rmbadname1(m.group('name')) + edecl = {} + if ename in groupcache[groupcounter]['vars']: + edecl = groupcache[groupcounter]['vars'][ename].copy() + not_has_typespec = 'typespec' not in edecl + if not_has_typespec: + edecl['typespec'] = typespec + elif typespec and (not typespec == edecl['typespec']): + outmess('updatevars: attempt to change the type of "%s" ("%s") to "%s". Ignoring.\n' % ( + ename, edecl['typespec'], typespec)) + if 'kindselector' not in edecl: + edecl['kindselector'] = copy.copy(kindselect) + elif kindselect: + for k in list(kindselect.keys()): + if k in edecl['kindselector'] and (not kindselect[k] == edecl['kindselector'][k]): + outmess('updatevars: attempt to change the kindselector "%s" of "%s" ("%s") to "%s". Ignoring.\n' % ( + k, ename, edecl['kindselector'][k], kindselect[k])) + else: + edecl['kindselector'][k] = copy.copy(kindselect[k]) + if 'charselector' not in edecl and charselect: + if not_has_typespec: + edecl['charselector'] = charselect + else: + errmess('updatevars:%s: attempt to change empty charselector to %r. Ignoring.\n' + % (ename, charselect)) + elif charselect: + for k in list(charselect.keys()): + if k in edecl['charselector'] and (not charselect[k] == edecl['charselector'][k]): + outmess('updatevars: attempt to change the charselector "%s" of "%s" ("%s") to "%s". Ignoring.\n' % ( + k, ename, edecl['charselector'][k], charselect[k])) + else: + edecl['charselector'][k] = copy.copy(charselect[k]) + if 'typename' not in edecl: + edecl['typename'] = typename + elif typename and (not edecl['typename'] == typename): + outmess('updatevars: attempt to change the typename of "%s" ("%s") to "%s". Ignoring.\n' % ( + ename, edecl['typename'], typename)) + if 'attrspec' not in edecl: + edecl['attrspec'] = copy.copy(attrspec) + elif attrspec: + for a in attrspec: + if a not in edecl['attrspec']: + edecl['attrspec'].append(a) + else: + edecl['typespec'] = copy.copy(typespec) + edecl['kindselector'] = copy.copy(kindselect) + edecl['charselector'] = copy.copy(charselect) + edecl['typename'] = typename + edecl['attrspec'] = copy.copy(attrspec) + if 'external' in (edecl.get('attrspec') or []) and e in groupcache[groupcounter]['args']: + if 'externals' not in groupcache[groupcounter]: + groupcache[groupcounter]['externals'] = [] + groupcache[groupcounter]['externals'].append(e) + if m.group('after'): + m1 = lenarraypattern.match(markouterparen(m.group('after'))) + if m1: + d1 = m1.groupdict() + for lk in ['len', 'array', 'init']: + if d1[lk + '2'] is not None: + d1[lk] = d1[lk + '2'] + del d1[lk + '2'] + for k in list(d1.keys()): + if d1[k] is not None: + d1[k] = unmarkouterparen(d1[k]) + else: + del d1[k] + + if 'len' in d1 and 'array' in d1: + if d1['len'] == '': + d1['len'] = d1['array'] + del d1['array'] + elif typespec == 'character': + if ('charselector' not in edecl) or (not edecl['charselector']): + edecl['charselector'] = {} + if 'len' in edecl['charselector']: + del edecl['charselector']['len'] + edecl['charselector']['*'] = d1['len'] + del d1['len'] + else: + d1['array'] = d1['array'] + ',' + d1['len'] + del d1['len'] + errmess('updatevars: "%s %s" is mapped to "%s %s(%s)"\n' % ( + typespec, e, typespec, ename, d1['array'])) + + if 'len' in d1: + if typespec in ['complex', 'integer', 'logical', 'real']: + if ('kindselector' not in edecl) or (not edecl['kindselector']): + edecl['kindselector'] = {} + edecl['kindselector']['*'] = d1['len'] + del d1['len'] + elif typespec == 'character': + if ('charselector' not in edecl) or (not edecl['charselector']): + edecl['charselector'] = {} + if 'len' in edecl['charselector']: + del edecl['charselector']['len'] + edecl['charselector']['*'] = d1['len'] + del d1['len'] + + if 'init' in d1: + if '=' in edecl and (not edecl['='] == d1['init']): + outmess('updatevars: attempt to change the init expression of "%s" ("%s") to "%s". Ignoring.\n' % ( + ename, edecl['='], d1['init'])) + else: + edecl['='] = d1['init'] + + if 'array' in d1: + dm = 'dimension(%s)' % d1['array'] + if 'attrspec' not in edecl or (not edecl['attrspec']): + edecl['attrspec'] = [dm] + else: + edecl['attrspec'].append(dm) + for dm1 in edecl['attrspec']: + if dm1[:9] == 'dimension' and dm1 != dm: + del edecl['attrspec'][-1] + errmess('updatevars:%s: attempt to change %r to %r. Ignoring.\n' + % (ename, dm1, dm)) + break + + else: + outmess('updatevars: could not crack entity declaration "%s". Ignoring.\n' % ( + ename + m.group('after'))) + for k in list(edecl.keys()): + if not edecl[k]: + del edecl[k] + groupcache[groupcounter]['vars'][ename] = edecl + if 'varnames' in groupcache[groupcounter]: + groupcache[groupcounter]['varnames'].append(ename) + last_name = ename + return last_name + + +def cracktypespec(typespec, selector): + kindselect = None + charselect = None + typename = None + if selector: + if typespec in ['complex', 'integer', 'logical', 'real']: + kindselect = kindselector.match(selector) + if not kindselect: + outmess( + 'cracktypespec: no kindselector pattern found for %s\n' % (repr(selector))) + return + kindselect = kindselect.groupdict() + kindselect['*'] = kindselect['kind2'] + del kindselect['kind2'] + for k in list(kindselect.keys()): + if not kindselect[k]: + del kindselect[k] + for k, i in list(kindselect.items()): + kindselect[k] = rmbadname1(i) + elif typespec == 'character': + charselect = charselector.match(selector) + if not charselect: + outmess( + 'cracktypespec: no charselector pattern found for %s\n' % (repr(selector))) + return + charselect = charselect.groupdict() + charselect['*'] = charselect['charlen'] + del charselect['charlen'] + if charselect['lenkind']: + lenkind = lenkindpattern.match( + markoutercomma(charselect['lenkind'])) + lenkind = lenkind.groupdict() + for lk in ['len', 'kind']: + if lenkind[lk + '2']: + lenkind[lk] = lenkind[lk + '2'] + charselect[lk] = lenkind[lk] + del lenkind[lk + '2'] + if lenkind['f2py_len'] is not None: + # used to specify the length of assumed length strings + charselect['f2py_len'] = lenkind['f2py_len'] + del charselect['lenkind'] + for k in list(charselect.keys()): + if not charselect[k]: + del charselect[k] + for k, i in list(charselect.items()): + charselect[k] = rmbadname1(i) + elif typespec == 'type': + typename = re.match(r'\s*\(\s*(?P\w+)\s*\)', selector, re.I) + if typename: + typename = typename.group('name') + else: + outmess('cracktypespec: no typename found in %s\n' % + (repr(typespec + selector))) + else: + outmess('cracktypespec: no selector used for %s\n' % + (repr(selector))) + return kindselect, charselect, typename +###### + + +def setattrspec(decl, attr, force=0): + if not decl: + decl = {} + if not attr: + return decl + if 'attrspec' not in decl: + decl['attrspec'] = [attr] + return decl + if force: + decl['attrspec'].append(attr) + if attr in decl['attrspec']: + return decl + if attr == 'static' and 'automatic' not in decl['attrspec']: + decl['attrspec'].append(attr) + elif attr == 'automatic' and 'static' not in decl['attrspec']: + decl['attrspec'].append(attr) + elif attr == 'public': + if 'private' not in decl['attrspec']: + decl['attrspec'].append(attr) + elif attr == 'private': + if 'public' not in decl['attrspec']: + decl['attrspec'].append(attr) + else: + decl['attrspec'].append(attr) + return decl + + +def setkindselector(decl, sel, force=0): + if not decl: + decl = {} + if not sel: + return decl + if 'kindselector' not in decl: + decl['kindselector'] = sel + return decl + for k in list(sel.keys()): + if force or k not in decl['kindselector']: + decl['kindselector'][k] = sel[k] + return decl + + +def setcharselector(decl, sel, force=0): + if not decl: + decl = {} + if not sel: + return decl + if 'charselector' not in decl: + decl['charselector'] = sel + return decl + + for k in list(sel.keys()): + if force or k not in decl['charselector']: + decl['charselector'][k] = sel[k] + return decl + + +def getblockname(block, unknown='unknown'): + if 'name' in block: + return block['name'] + return unknown + +# post processing + + +def setmesstext(block): + global filepositiontext + + try: + filepositiontext = 'In: %s:%s\n' % (block['from'], block['name']) + except Exception: + pass + + +def get_usedict(block): + usedict = {} + if 'parent_block' in block: + usedict = get_usedict(block['parent_block']) + if 'use' in block: + usedict.update(block['use']) + return usedict + + +def get_useparameters(block, param_map=None): + global f90modulevars + + if param_map is None: + param_map = {} + usedict = get_usedict(block) + if not usedict: + return param_map + for usename, mapping in list(usedict.items()): + usename = usename.lower() + if usename not in f90modulevars: + outmess('get_useparameters: no module %s info used by %s\n' % + (usename, block.get('name'))) + continue + mvars = f90modulevars[usename] + params = get_parameters(mvars) + if not params: + continue + # XXX: apply mapping + if mapping: + errmess('get_useparameters: mapping for %s not impl.\n' % (mapping)) + for k, v in list(params.items()): + if k in param_map: + outmess('get_useparameters: overriding parameter %s with' + ' value from module %s\n' % (repr(k), repr(usename))) + param_map[k] = v + + return param_map + + +def postcrack2(block, tab='', param_map=None): + global f90modulevars + + if not f90modulevars: + return block + if isinstance(block, list): + ret = [postcrack2(g, tab=tab + '\t', param_map=param_map) + for g in block] + return ret + setmesstext(block) + outmess('%sBlock: %s\n' % (tab, block['name']), 0) + + if param_map is None: + param_map = get_useparameters(block) + + if param_map is not None and 'vars' in block: + vars = block['vars'] + for n in list(vars.keys()): + var = vars[n] + if 'kindselector' in var: + kind = var['kindselector'] + if 'kind' in kind: + val = kind['kind'] + if val in param_map: + kind['kind'] = param_map[val] + new_body = [postcrack2(b, tab=tab + '\t', param_map=param_map) + for b in block['body']] + block['body'] = new_body + + return block + + +def postcrack(block, args=None, tab=''): + """ + TODO: + function return values + determine expression types if in argument list + """ + global usermodules, onlyfunctions + + if isinstance(block, list): + gret = [] + uret = [] + for g in block: + setmesstext(g) + g = postcrack(g, tab=tab + '\t') + # sort user routines to appear first + if 'name' in g and '__user__' in g['name']: + uret.append(g) + else: + gret.append(g) + return uret + gret + setmesstext(block) + if not isinstance(block, dict) and 'block' not in block: + raise Exception('postcrack: Expected block dictionary instead of ' + + str(block)) + if 'name' in block and not block['name'] == 'unknown_interface': + outmess('%sBlock: %s\n' % (tab, block['name']), 0) + block = analyzeargs(block) + block = analyzecommon(block) + block['vars'] = analyzevars(block) + block['sortvars'] = sortvarnames(block['vars']) + if block.get('args'): + args = block['args'] + block['body'] = analyzebody(block, args, tab=tab) + + userisdefined = [] + if 'use' in block: + useblock = block['use'] + for k in list(useblock.keys()): + if '__user__' in k: + userisdefined.append(k) + else: + useblock = {} + name = '' + if 'name' in block: + name = block['name'] + # and not userisdefined: # Build a __user__ module + if block.get('externals'): + interfaced = [] + if 'interfaced' in block: + interfaced = block['interfaced'] + mvars = copy.copy(block['vars']) + if name: + mname = name + '__user__routines' + else: + mname = 'unknown__user__routines' + if mname in userisdefined: + i = 1 + while '%s_%i' % (mname, i) in userisdefined: + i = i + 1 + mname = '%s_%i' % (mname, i) + interface = {'block': 'interface', 'body': [], + 'vars': {}, 'name': name + '_user_interface'} + for e in block['externals']: + if e in interfaced: + edef = [] + j = -1 + for b in block['body']: + j = j + 1 + if b['block'] == 'interface': + i = -1 + for bb in b['body']: + i = i + 1 + if 'name' in bb and bb['name'] == e: + edef = copy.copy(bb) + del b['body'][i] + break + if edef: + if not b['body']: + del block['body'][j] + del interfaced[interfaced.index(e)] + break + interface['body'].append(edef) + else: + if e in mvars and not isexternal(mvars[e]): + interface['vars'][e] = mvars[e] + if interface['vars'] or interface['body']: + block['interfaced'] = interfaced + mblock = {'block': 'python module', 'body': [ + interface], 'vars': {}, 'name': mname, 'interfaced': block['externals']} + useblock[mname] = {} + usermodules.append(mblock) + if useblock: + block['use'] = useblock + return block + + +def sortvarnames(vars): + indep = [] + dep = [] + for v in list(vars.keys()): + if 'depend' in vars[v] and vars[v]['depend']: + dep.append(v) + else: + indep.append(v) + n = len(dep) + i = 0 + while dep: # XXX: How to catch dependence cycles correctly? + v = dep[0] + fl = 0 + for w in dep[1:]: + if w in vars[v]['depend']: + fl = 1 + break + if fl: + dep = dep[1:] + [v] + i = i + 1 + if i > n: + errmess('sortvarnames: failed to compute dependencies because' + ' of cyclic dependencies between ' + + ', '.join(dep) + '\n') + indep = indep + dep + break + else: + indep.append(v) + dep = dep[1:] + n = len(dep) + i = 0 + return indep + + +def analyzecommon(block): + if not hascommon(block): + return block + commonvars = [] + for k in list(block['common'].keys()): + comvars = [] + for e in block['common'][k]: + m = re.match( + r'\A\s*\b(?P.*?)\b\s*(\((?P.*?)\)|)\s*\Z', e, re.I) + if m: + dims = [] + if m.group('dims'): + dims = [x.strip() + for x in markoutercomma(m.group('dims')).split('@,@')] + n = rmbadname1(m.group('name').strip()) + if n in block['vars']: + if 'attrspec' in block['vars'][n]: + block['vars'][n]['attrspec'].append( + 'dimension(%s)' % (','.join(dims))) + else: + block['vars'][n]['attrspec'] = [ + 'dimension(%s)' % (','.join(dims))] + else: + if dims: + block['vars'][n] = { + 'attrspec': ['dimension(%s)' % (','.join(dims))]} + else: + block['vars'][n] = {} + if n not in commonvars: + commonvars.append(n) + else: + n = e + errmess( + 'analyzecommon: failed to extract "[()]" from "%s" in common /%s/.\n' % (e, k)) + comvars.append(n) + block['common'][k] = comvars + if 'commonvars' not in block: + block['commonvars'] = commonvars + else: + block['commonvars'] = block['commonvars'] + commonvars + return block + + +def analyzebody(block, args, tab=''): + global usermodules, skipfuncs, onlyfuncs, f90modulevars + + setmesstext(block) + + maybe_private = { + key: value + for key, value in block['vars'].items() + if 'attrspec' not in value or 'public' not in value['attrspec'] + } + + body = [] + for b in block['body']: + b['parent_block'] = block + if b['block'] in ['function', 'subroutine']: + if args is not None and b['name'] not in args: + continue + else: + as_ = b['args'] + # Add private members to skipfuncs for gh-23879 + if b['name'] in maybe_private.keys(): + skipfuncs.append(b['name']) + if b['name'] in skipfuncs: + continue + if onlyfuncs and b['name'] not in onlyfuncs: + continue + b['saved_interface'] = crack2fortrangen( + b, '\n' + ' ' * 6, as_interface=True) + + else: + as_ = args + b = postcrack(b, as_, tab=tab + '\t') + if b['block'] in ['interface', 'abstract interface'] and \ + not b['body'] and not b.get('implementedby'): + if 'f2pyenhancements' not in b: + continue + if b['block'].replace(' ', '') == 'pythonmodule': + usermodules.append(b) + else: + if b['block'] == 'module': + f90modulevars[b['name']] = b['vars'] + body.append(b) + return body + + +def buildimplicitrules(block): + setmesstext(block) + implicitrules = defaultimplicitrules + attrrules = {} + if 'implicit' in block: + if block['implicit'] is None: + implicitrules = None + if verbose > 1: + outmess( + 'buildimplicitrules: no implicit rules for routine %s.\n' % repr(block['name'])) + else: + for k in list(block['implicit'].keys()): + if block['implicit'][k].get('typespec') not in ['static', 'automatic']: + implicitrules[k] = block['implicit'][k] + else: + attrrules[k] = block['implicit'][k]['typespec'] + return implicitrules, attrrules + + +def myeval(e, g=None, l=None): + """ Like `eval` but returns only integers and floats """ + r = eval(e, g, l) + if type(r) in [int, float]: + return r + raise ValueError('r=%r' % (r)) + +getlincoef_re_1 = re.compile(r'\A\b\w+\b\Z', re.I) + + +def getlincoef(e, xset): # e = a*x+b ; x in xset + """ + Obtain ``a`` and ``b`` when ``e == "a*x+b"``, where ``x`` is a symbol in + xset. + + >>> getlincoef('2*x + 1', {'x'}) + (2, 1, 'x') + >>> getlincoef('3*x + x*2 + 2 + 1', {'x'}) + (5, 3, 'x') + >>> getlincoef('0', {'x'}) + (0, 0, None) + >>> getlincoef('0*x', {'x'}) + (0, 0, 'x') + >>> getlincoef('x*x', {'x'}) + (None, None, None) + + This can be tricked by sufficiently complex expressions + + >>> getlincoef('(x - 0.5)*(x - 1.5)*(x - 1)*x + 2*x + 3', {'x'}) + (2.0, 3.0, 'x') + """ + try: + c = int(myeval(e, {}, {})) + return 0, c, None + except Exception: + pass + if getlincoef_re_1.match(e): + return 1, 0, e + len_e = len(e) + for x in xset: + if len(x) > len_e: + continue + if re.search(r'\w\s*\([^)]*\b' + x + r'\b', e): + # skip function calls having x as an argument, e.g max(1, x) + continue + re_1 = re.compile(r'(?P.*?)\b' + x + r'\b(?P.*)', re.I) + m = re_1.match(e) + if m: + try: + m1 = re_1.match(e) + while m1: + ee = '%s(%s)%s' % ( + m1.group('before'), 0, m1.group('after')) + m1 = re_1.match(ee) + b = myeval(ee, {}, {}) + m1 = re_1.match(e) + while m1: + ee = '%s(%s)%s' % ( + m1.group('before'), 1, m1.group('after')) + m1 = re_1.match(ee) + a = myeval(ee, {}, {}) - b + m1 = re_1.match(e) + while m1: + ee = '%s(%s)%s' % ( + m1.group('before'), 0.5, m1.group('after')) + m1 = re_1.match(ee) + c = myeval(ee, {}, {}) + # computing another point to be sure that expression is linear + m1 = re_1.match(e) + while m1: + ee = '%s(%s)%s' % ( + m1.group('before'), 1.5, m1.group('after')) + m1 = re_1.match(ee) + c2 = myeval(ee, {}, {}) + if (a * 0.5 + b == c and a * 1.5 + b == c2): + return a, b, x + except Exception: + pass + break + return None, None, None + + +word_pattern = re.compile(r'\b[a-z][\w$]*\b', re.I) + + +def _get_depend_dict(name, vars, deps): + if name in vars: + words = vars[name].get('depend', []) + + if '=' in vars[name] and not isstring(vars[name]): + for word in word_pattern.findall(vars[name]['=']): + # The word_pattern may return values that are not + # only variables, they can be string content for instance + if word not in words and word in vars and word != name: + words.append(word) + for word in words[:]: + for w in deps.get(word, []) \ + or _get_depend_dict(word, vars, deps): + if w not in words: + words.append(w) + else: + outmess('_get_depend_dict: no dependence info for %s\n' % (repr(name))) + words = [] + deps[name] = words + return words + + +def _calc_depend_dict(vars): + names = list(vars.keys()) + depend_dict = {} + for n in names: + _get_depend_dict(n, vars, depend_dict) + return depend_dict + + +def get_sorted_names(vars): + depend_dict = _calc_depend_dict(vars) + names = [] + for name in list(depend_dict.keys()): + if not depend_dict[name]: + names.append(name) + del depend_dict[name] + while depend_dict: + for name, lst in list(depend_dict.items()): + new_lst = [n for n in lst if n in depend_dict] + if not new_lst: + names.append(name) + del depend_dict[name] + else: + depend_dict[name] = new_lst + return [name for name in names if name in vars] + + +def _kind_func(string): + # XXX: return something sensible. + if string[0] in "'\"": + string = string[1:-1] + if real16pattern.match(string): + return 8 + elif real8pattern.match(string): + return 4 + return 'kind(' + string + ')' + + +def _selected_int_kind_func(r): + # XXX: This should be processor dependent + m = 10 ** r + if m <= 2 ** 8: + return 1 + if m <= 2 ** 16: + return 2 + if m <= 2 ** 32: + return 4 + if m <= 2 ** 63: + return 8 + if m <= 2 ** 128: + return 16 + return -1 + + +def _selected_real_kind_func(p, r=0, radix=0): + # XXX: This should be processor dependent + # This is only verified for 0 <= p <= 20, possibly good for p <= 33 and above + if p < 7: + return 4 + if p < 16: + return 8 + machine = platform.machine().lower() + if machine.startswith(('aarch64', 'alpha', 'arm64', 'loongarch', 'mips', 'power', 'ppc', 'riscv', 's390x', 'sparc')): + if p <= 33: + return 16 + else: + if p < 19: + return 10 + elif p <= 33: + return 16 + return -1 + + +def get_parameters(vars, global_params={}): + params = copy.copy(global_params) + g_params = copy.copy(global_params) + for name, func in [('kind', _kind_func), + ('selected_int_kind', _selected_int_kind_func), + ('selected_real_kind', _selected_real_kind_func), ]: + if name not in g_params: + g_params[name] = func + param_names = [] + for n in get_sorted_names(vars): + if 'attrspec' in vars[n] and 'parameter' in vars[n]['attrspec']: + param_names.append(n) + kind_re = re.compile(r'\bkind\s*\(\s*(?P.*)\s*\)', re.I) + selected_int_kind_re = re.compile( + r'\bselected_int_kind\s*\(\s*(?P.*)\s*\)', re.I) + selected_kind_re = re.compile( + r'\bselected_(int|real)_kind\s*\(\s*(?P.*)\s*\)', re.I) + for n in param_names: + if '=' in vars[n]: + v = vars[n]['='] + if islogical(vars[n]): + v = v.lower() + for repl in [ + ('.false.', 'False'), + ('.true.', 'True'), + # TODO: test .eq., .neq., etc replacements. + ]: + v = v.replace(*repl) + + v = kind_re.sub(r'kind("\1")', v) + v = selected_int_kind_re.sub(r'selected_int_kind(\1)', v) + + # We need to act according to the data. + # The easy case is if the data has a kind-specifier, + # then we may easily remove those specifiers. + # However, it may be that the user uses other specifiers...(!) + is_replaced = False + + if 'kindselector' in vars[n]: + # Remove kind specifier (including those defined + # by parameters) + if 'kind' in vars[n]['kindselector']: + orig_v_len = len(v) + v = v.replace('_' + vars[n]['kindselector']['kind'], '') + # Again, this will be true if even a single specifier + # has been replaced, see comment above. + is_replaced = len(v) < orig_v_len + + if not is_replaced: + if not selected_kind_re.match(v): + v_ = v.split('_') + # In case there are additive parameters + if len(v_) > 1: + v = ''.join(v_[:-1]).lower().replace(v_[-1].lower(), '') + + # Currently this will not work for complex numbers. + # There is missing code for extracting a complex number, + # which may be defined in either of these: + # a) (Re, Im) + # b) cmplx(Re, Im) + # c) dcmplx(Re, Im) + # d) cmplx(Re, Im, ) + + if isdouble(vars[n]): + tt = list(v) + for m in real16pattern.finditer(v): + tt[m.start():m.end()] = list( + v[m.start():m.end()].lower().replace('d', 'e')) + v = ''.join(tt) + + elif iscomplex(vars[n]): + outmess(f'get_parameters[TODO]: ' + f'implement evaluation of complex expression {v}\n') + + dimspec = ([s.removeprefix('dimension').strip() + for s in vars[n]['attrspec'] + if s.startswith('dimension')] or [None])[0] + + # Handle _dp for gh-6624 + # Also fixes gh-20460 + if real16pattern.search(v): + v = 8 + elif real8pattern.search(v): + v = 4 + try: + params[n] = param_eval(v, g_params, params, dimspec=dimspec) + except Exception as msg: + params[n] = v + outmess(f'get_parameters: got "{msg}" on {n!r}\n') + + if isstring(vars[n]) and isinstance(params[n], int): + params[n] = chr(params[n]) + nl = n.lower() + if nl != n: + params[nl] = params[n] + else: + print(vars[n]) + outmess(f'get_parameters:parameter {n!r} does not have value?!\n') + return params + + +def _eval_length(length, params): + if length in ['(:)', '(*)', '*']: + return '(*)' + return _eval_scalar(length, params) + + +_is_kind_number = re.compile(r'\d+_').match + + +def _eval_scalar(value, params): + if _is_kind_number(value): + value = value.split('_')[0] + try: + # TODO: use symbolic from PR #19805 + value = eval(value, {}, params) + value = (repr if isinstance(value, str) else str)(value) + except (NameError, SyntaxError, TypeError): + return value + except Exception as msg: + errmess('"%s" in evaluating %r ' + '(available names: %s)\n' + % (msg, value, list(params.keys()))) + return value + + +def analyzevars(block): + """ + Sets correct dimension information for each variable/parameter + """ + + global f90modulevars + + setmesstext(block) + implicitrules, attrrules = buildimplicitrules(block) + vars = copy.copy(block['vars']) + if block['block'] == 'function' and block['name'] not in vars: + vars[block['name']] = {} + if '' in block['vars']: + del vars[''] + if 'attrspec' in block['vars']['']: + gen = block['vars']['']['attrspec'] + for n in set(vars) | set(b['name'] for b in block['body']): + for k in ['public', 'private']: + if k in gen: + vars[n] = setattrspec(vars.get(n, {}), k) + svars = [] + args = block['args'] + for a in args: + try: + vars[a] + svars.append(a) + except KeyError: + pass + for n in list(vars.keys()): + if n not in args: + svars.append(n) + + params = get_parameters(vars, get_useparameters(block)) + # At this point, params are read and interpreted, but + # the params used to define vars are not yet parsed + dep_matches = {} + name_match = re.compile(r'[A-Za-z][\w$]*').match + for v in list(vars.keys()): + m = name_match(v) + if m: + n = v[m.start():m.end()] + try: + dep_matches[n] + except KeyError: + dep_matches[n] = re.compile(r'.*\b%s\b' % (v), re.I).match + for n in svars: + if n[0] in list(attrrules.keys()): + vars[n] = setattrspec(vars[n], attrrules[n[0]]) + if 'typespec' not in vars[n]: + if not('attrspec' in vars[n] and 'external' in vars[n]['attrspec']): + if implicitrules: + ln0 = n[0].lower() + for k in list(implicitrules[ln0].keys()): + if k == 'typespec' and implicitrules[ln0][k] == 'undefined': + continue + if k not in vars[n]: + vars[n][k] = implicitrules[ln0][k] + elif k == 'attrspec': + for l in implicitrules[ln0][k]: + vars[n] = setattrspec(vars[n], l) + elif n in block['args']: + outmess('analyzevars: typespec of variable %s is not defined in routine %s.\n' % ( + repr(n), block['name'])) + if 'charselector' in vars[n]: + if 'len' in vars[n]['charselector']: + l = vars[n]['charselector']['len'] + try: + l = str(eval(l, {}, params)) + except Exception: + pass + vars[n]['charselector']['len'] = l + + if 'kindselector' in vars[n]: + if 'kind' in vars[n]['kindselector']: + l = vars[n]['kindselector']['kind'] + try: + l = str(eval(l, {}, params)) + except Exception: + pass + vars[n]['kindselector']['kind'] = l + + dimension_exprs = {} + if 'attrspec' in vars[n]: + attr = vars[n]['attrspec'] + attr.reverse() + vars[n]['attrspec'] = [] + dim, intent, depend, check, note = None, None, None, None, None + for a in attr: + if a[:9] == 'dimension': + dim = (a[9:].strip())[1:-1] + elif a[:6] == 'intent': + intent = (a[6:].strip())[1:-1] + elif a[:6] == 'depend': + depend = (a[6:].strip())[1:-1] + elif a[:5] == 'check': + check = (a[5:].strip())[1:-1] + elif a[:4] == 'note': + note = (a[4:].strip())[1:-1] + else: + vars[n] = setattrspec(vars[n], a) + if intent: + if 'intent' not in vars[n]: + vars[n]['intent'] = [] + for c in [x.strip() for x in markoutercomma(intent).split('@,@')]: + # Remove spaces so that 'in out' becomes 'inout' + tmp = c.replace(' ', '') + if tmp not in vars[n]['intent']: + vars[n]['intent'].append(tmp) + intent = None + if note: + note = note.replace('\\n\\n', '\n\n') + note = note.replace('\\n ', '\n') + if 'note' not in vars[n]: + vars[n]['note'] = [note] + else: + vars[n]['note'].append(note) + note = None + if depend is not None: + if 'depend' not in vars[n]: + vars[n]['depend'] = [] + for c in rmbadname([x.strip() for x in markoutercomma(depend).split('@,@')]): + if c not in vars[n]['depend']: + vars[n]['depend'].append(c) + depend = None + if check is not None: + if 'check' not in vars[n]: + vars[n]['check'] = [] + for c in [x.strip() for x in markoutercomma(check).split('@,@')]: + if c not in vars[n]['check']: + vars[n]['check'].append(c) + check = None + if dim and 'dimension' not in vars[n]: + vars[n]['dimension'] = [] + for d in rmbadname( + [x.strip() for x in markoutercomma(dim).split('@,@')] + ): + # d is the expression inside the dimension declaration + # Evaluate `d` with respect to params + try: + # the dimension for this variable depends on a + # previously defined parameter + d = param_parse(d, params) + except (ValueError, IndexError, KeyError): + outmess( + 'analyzevars: could not parse dimension for ' + f'variable {d!r}\n' + ) + + dim_char = ':' if d == ':' else '*' + if d == dim_char: + dl = [dim_char] + else: + dl = markoutercomma(d, ':').split('@:@') + if len(dl) == 2 and '*' in dl: # e.g. dimension(5:*) + dl = ['*'] + d = '*' + if len(dl) == 1 and dl[0] != dim_char: + dl = ['1', dl[0]] + if len(dl) == 2: + d1, d2 = map(symbolic.Expr.parse, dl) + dsize = d2 - d1 + 1 + d = dsize.tostring(language=symbolic.Language.C) + # find variables v that define d as a linear + # function, `d == a * v + b`, and store + # coefficients a and b for further analysis. + solver_and_deps = {} + for v in block['vars']: + s = symbolic.as_symbol(v) + if dsize.contains(s): + try: + a, b = dsize.linear_solve(s) + + def solve_v(s, a=a, b=b): + return (s - b) / a + + all_symbols = set(a.symbols()) + all_symbols.update(b.symbols()) + except RuntimeError as msg: + # d is not a linear function of v, + # however, if v can be determined + # from d using other means, + # implement the corresponding + # solve_v function here. + solve_v = None + all_symbols = set(dsize.symbols()) + v_deps = set( + s.data for s in all_symbols + if s.data in vars) + solver_and_deps[v] = solve_v, list(v_deps) + # Note that dsize may contain symbols that are + # not defined in block['vars']. Here we assume + # these correspond to Fortran/C intrinsic + # functions or that are defined by other + # means. We'll let the compiler validate the + # definiteness of such symbols. + dimension_exprs[d] = solver_and_deps + vars[n]['dimension'].append(d) + + if 'check' not in vars[n] and 'args' in block and n in block['args']: + # n is an argument that has no checks defined. Here we + # generate some consistency checks for n, and when n is an + # array, generate checks for its dimensions and construct + # initialization expressions. + n_deps = vars[n].get('depend', []) + n_checks = [] + n_is_input = l_or(isintent_in, isintent_inout, + isintent_inplace)(vars[n]) + if isarray(vars[n]): # n is array + for i, d in enumerate(vars[n]['dimension']): + coeffs_and_deps = dimension_exprs.get(d) + if coeffs_and_deps is None: + # d is `:` or `*` or a constant expression + pass + elif n_is_input: + # n is an input array argument and its shape + # may define variables used in dimension + # specifications. + for v, (solver, deps) in coeffs_and_deps.items(): + def compute_deps(v, deps): + for v1 in coeffs_and_deps.get(v, [None, []])[1]: + if v1 not in deps: + deps.add(v1) + compute_deps(v1, deps) + all_deps = set() + compute_deps(v, all_deps) + if (v in n_deps + or '=' in vars[v] + or 'depend' in vars[v]): + # Skip a variable that + # - n depends on + # - has user-defined initialization expression + # - has user-defined dependencies + continue + if solver is not None and v not in all_deps: + # v can be solved from d, hence, we + # make it an optional argument with + # initialization expression: + is_required = False + init = solver(symbolic.as_symbol( + f'shape({n}, {i})')) + init = init.tostring( + language=symbolic.Language.C) + vars[v]['='] = init + # n needs to be initialized before v. So, + # making v dependent on n and on any + # variables in solver or d. + vars[v]['depend'] = [n] + deps + if 'check' not in vars[v]: + # add check only when no + # user-specified checks exist + vars[v]['check'] = [ + f'shape({n}, {i}) == {d}'] + else: + # d is a non-linear function on v, + # hence, v must be a required input + # argument that n will depend on + is_required = True + if 'intent' not in vars[v]: + vars[v]['intent'] = [] + if 'in' not in vars[v]['intent']: + vars[v]['intent'].append('in') + # v needs to be initialized before n + n_deps.append(v) + n_checks.append( + f'shape({n}, {i}) == {d}') + v_attr = vars[v].get('attrspec', []) + if not ('optional' in v_attr + or 'required' in v_attr): + v_attr.append( + 'required' if is_required else 'optional') + if v_attr: + vars[v]['attrspec'] = v_attr + if coeffs_and_deps is not None: + # extend v dependencies with ones specified in attrspec + for v, (solver, deps) in coeffs_and_deps.items(): + v_deps = vars[v].get('depend', []) + for aa in vars[v].get('attrspec', []): + if aa.startswith('depend'): + aa = ''.join(aa.split()) + v_deps.extend(aa[7:-1].split(',')) + if v_deps: + vars[v]['depend'] = list(set(v_deps)) + if n not in v_deps: + n_deps.append(v) + elif isstring(vars[n]): + if 'charselector' in vars[n]: + if '*' in vars[n]['charselector']: + length = _eval_length(vars[n]['charselector']['*'], + params) + vars[n]['charselector']['*'] = length + elif 'len' in vars[n]['charselector']: + length = _eval_length(vars[n]['charselector']['len'], + params) + del vars[n]['charselector']['len'] + vars[n]['charselector']['*'] = length + if n_checks: + vars[n]['check'] = n_checks + if n_deps: + vars[n]['depend'] = list(set(n_deps)) + + if '=' in vars[n]: + if 'attrspec' not in vars[n]: + vars[n]['attrspec'] = [] + if ('optional' not in vars[n]['attrspec']) and \ + ('required' not in vars[n]['attrspec']): + vars[n]['attrspec'].append('optional') + if 'depend' not in vars[n]: + vars[n]['depend'] = [] + for v, m in list(dep_matches.items()): + if m(vars[n]['=']): + vars[n]['depend'].append(v) + if not vars[n]['depend']: + del vars[n]['depend'] + if isscalar(vars[n]): + vars[n]['='] = _eval_scalar(vars[n]['='], params) + + for n in list(vars.keys()): + if n == block['name']: # n is block name + if 'note' in vars[n]: + block['note'] = vars[n]['note'] + if block['block'] == 'function': + if 'result' in block and block['result'] in vars: + vars[n] = appenddecl(vars[n], vars[block['result']]) + if 'prefix' in block: + pr = block['prefix'] + pr1 = pr.replace('pure', '') + ispure = (not pr == pr1) + pr = pr1.replace('recursive', '') + isrec = (not pr == pr1) + m = typespattern[0].match(pr) + if m: + typespec, selector, attr, edecl = cracktypespec0( + m.group('this'), m.group('after')) + kindselect, charselect, typename = cracktypespec( + typespec, selector) + vars[n]['typespec'] = typespec + try: + if block['result']: + vars[block['result']]['typespec'] = typespec + except Exception: + pass + if kindselect: + if 'kind' in kindselect: + try: + kindselect['kind'] = eval( + kindselect['kind'], {}, params) + except Exception: + pass + vars[n]['kindselector'] = kindselect + if charselect: + vars[n]['charselector'] = charselect + if typename: + vars[n]['typename'] = typename + if ispure: + vars[n] = setattrspec(vars[n], 'pure') + if isrec: + vars[n] = setattrspec(vars[n], 'recursive') + else: + outmess( + 'analyzevars: prefix (%s) were not used\n' % repr(block['prefix'])) + if block['block'] not in ['module', 'pythonmodule', 'python module', 'block data']: + if 'commonvars' in block: + neededvars = copy.copy(block['args'] + block['commonvars']) + else: + neededvars = copy.copy(block['args']) + for n in list(vars.keys()): + if l_or(isintent_callback, isintent_aux)(vars[n]): + neededvars.append(n) + if 'entry' in block: + neededvars.extend(list(block['entry'].keys())) + for k in list(block['entry'].keys()): + for n in block['entry'][k]: + if n not in neededvars: + neededvars.append(n) + if block['block'] == 'function': + if 'result' in block: + neededvars.append(block['result']) + else: + neededvars.append(block['name']) + if block['block'] in ['subroutine', 'function']: + name = block['name'] + if name in vars and 'intent' in vars[name]: + block['intent'] = vars[name]['intent'] + if block['block'] == 'type': + neededvars.extend(list(vars.keys())) + for n in list(vars.keys()): + if n not in neededvars: + del vars[n] + return vars + + +analyzeargs_re_1 = re.compile(r'\A[a-z]+[\w$]*\Z', re.I) + + +def param_eval(v, g_params, params, dimspec=None): + """ + Creates a dictionary of indices and values for each parameter in a + parameter array to be evaluated later. + + WARNING: It is not possible to initialize multidimensional array + parameters e.g. dimension(-3:1, 4, 3:5) at this point. This is because in + Fortran initialization through array constructor requires the RESHAPE + intrinsic function. Since the right-hand side of the parameter declaration + is not executed in f2py, but rather at the compiled c/fortran extension, + later, it is not possible to execute a reshape of a parameter array. + One issue remains: if the user wants to access the array parameter from + python, we should either + 1) allow them to access the parameter array using python standard indexing + (which is often incompatible with the original fortran indexing) + 2) allow the parameter array to be accessed in python as a dictionary with + fortran indices as keys + We are choosing 2 for now. + """ + if dimspec is None: + try: + p = eval(v, g_params, params) + except Exception as msg: + p = v + outmess(f'param_eval: got "{msg}" on {v!r}\n') + return p + + # This is an array parameter. + # First, we parse the dimension information + if len(dimspec) < 2 or dimspec[::len(dimspec)-1] != "()": + raise ValueError(f'param_eval: dimension {dimspec} can\'t be parsed') + dimrange = dimspec[1:-1].split(',') + if len(dimrange) == 1: + # e.g. dimension(2) or dimension(-1:1) + dimrange = dimrange[0].split(':') + # now, dimrange is a list of 1 or 2 elements + if len(dimrange) == 1: + bound = param_parse(dimrange[0], params) + dimrange = range(1, int(bound)+1) + else: + lbound = param_parse(dimrange[0], params) + ubound = param_parse(dimrange[1], params) + dimrange = range(int(lbound), int(ubound)+1) + else: + raise ValueError('param_eval: multidimensional array parameters ' + f'{dimspec} not supported') + + # Parse parameter value + v = (v[2:-2] if v.startswith('(/') else v).split(',') + v_eval = [] + for item in v: + try: + item = eval(item, g_params, params) + except Exception as msg: + outmess(f'param_eval: got "{msg}" on {item!r}\n') + v_eval.append(item) + + p = dict(zip(dimrange, v_eval)) + + return p + + +def param_parse(d, params): + """Recursively parse array dimensions. + + Parses the declaration of an array variable or parameter + `dimension` keyword, and is called recursively if the + dimension for this array is a previously defined parameter + (found in `params`). + + Parameters + ---------- + d : str + Fortran expression describing the dimension of an array. + params : dict + Previously parsed parameters declared in the Fortran source file. + + Returns + ------- + out : str + Parsed dimension expression. + + Examples + -------- + + * If the line being analyzed is + + `integer, parameter, dimension(2) :: pa = (/ 3, 5 /)` + + then `d = 2` and we return immediately, with + + >>> d = '2' + >>> param_parse(d, params) + 2 + + * If the line being analyzed is + + `integer, parameter, dimension(pa) :: pb = (/1, 2, 3/)` + + then `d = 'pa'`; since `pa` is a previously parsed parameter, + and `pa = 3`, we call `param_parse` recursively, to obtain + + >>> d = 'pa' + >>> params = {'pa': 3} + >>> param_parse(d, params) + 3 + + * If the line being analyzed is + + `integer, parameter, dimension(pa(1)) :: pb = (/1, 2, 3/)` + + then `d = 'pa(1)'`; since `pa` is a previously parsed parameter, + and `pa(1) = 3`, we call `param_parse` recursively, to obtain + + >>> d = 'pa(1)' + >>> params = dict(pa={1: 3, 2: 5}) + >>> param_parse(d, params) + 3 + """ + if "(" in d: + # this dimension expression is an array + dname = d[:d.find("(")] + ddims = d[d.find("(")+1:d.rfind(")")] + # this dimension expression is also a parameter; + # parse it recursively + index = int(param_parse(ddims, params)) + return str(params[dname][index]) + elif d in params: + return str(params[d]) + else: + for p in params: + re_1 = re.compile( + r'(?P.*?)\b' + p + r'\b(?P.*)', re.I + ) + m = re_1.match(d) + while m: + d = m.group('before') + \ + str(params[p]) + m.group('after') + m = re_1.match(d) + return d + + +def expr2name(a, block, args=[]): + orig_a = a + a_is_expr = not analyzeargs_re_1.match(a) + if a_is_expr: # `a` is an expression + implicitrules, attrrules = buildimplicitrules(block) + at = determineexprtype(a, block['vars'], implicitrules) + na = 'e_' + for c in a: + c = c.lower() + if c not in string.ascii_lowercase + string.digits: + c = '_' + na = na + c + if na[-1] == '_': + na = na + 'e' + else: + na = na + '_e' + a = na + while a in block['vars'] or a in block['args']: + a = a + 'r' + if a in args: + k = 1 + while a + str(k) in args: + k = k + 1 + a = a + str(k) + if a_is_expr: + block['vars'][a] = at + else: + if a not in block['vars']: + if orig_a in block['vars']: + block['vars'][a] = block['vars'][orig_a] + else: + block['vars'][a] = {} + if 'externals' in block and orig_a in block['externals'] + block['interfaced']: + block['vars'][a] = setattrspec(block['vars'][a], 'external') + return a + + +def analyzeargs(block): + setmesstext(block) + implicitrules, _ = buildimplicitrules(block) + if 'args' not in block: + block['args'] = [] + args = [] + for a in block['args']: + a = expr2name(a, block, args) + args.append(a) + block['args'] = args + if 'entry' in block: + for k, args1 in list(block['entry'].items()): + for a in args1: + if a not in block['vars']: + block['vars'][a] = {} + + for b in block['body']: + if b['name'] in args: + if 'externals' not in block: + block['externals'] = [] + if b['name'] not in block['externals']: + block['externals'].append(b['name']) + if 'result' in block and block['result'] not in block['vars']: + block['vars'][block['result']] = {} + return block + +determineexprtype_re_1 = re.compile(r'\A\(.+?,.+?\)\Z', re.I) +determineexprtype_re_2 = re.compile(r'\A[+-]?\d+(_(?P\w+)|)\Z', re.I) +determineexprtype_re_3 = re.compile( + r'\A[+-]?[\d.]+[-\d+de.]*(_(?P\w+)|)\Z', re.I) +determineexprtype_re_4 = re.compile(r'\A\(.*\)\Z', re.I) +determineexprtype_re_5 = re.compile(r'\A(?P\w+)\s*\(.*?\)\s*\Z', re.I) + + +def _ensure_exprdict(r): + if isinstance(r, int): + return {'typespec': 'integer'} + if isinstance(r, float): + return {'typespec': 'real'} + if isinstance(r, complex): + return {'typespec': 'complex'} + if isinstance(r, dict): + return r + raise AssertionError(repr(r)) + + +def determineexprtype(expr, vars, rules={}): + if expr in vars: + return _ensure_exprdict(vars[expr]) + expr = expr.strip() + if determineexprtype_re_1.match(expr): + return {'typespec': 'complex'} + m = determineexprtype_re_2.match(expr) + if m: + if 'name' in m.groupdict() and m.group('name'): + outmess( + 'determineexprtype: selected kind types not supported (%s)\n' % repr(expr)) + return {'typespec': 'integer'} + m = determineexprtype_re_3.match(expr) + if m: + if 'name' in m.groupdict() and m.group('name'): + outmess( + 'determineexprtype: selected kind types not supported (%s)\n' % repr(expr)) + return {'typespec': 'real'} + for op in ['+', '-', '*', '/']: + for e in [x.strip() for x in markoutercomma(expr, comma=op).split('@' + op + '@')]: + if e in vars: + return _ensure_exprdict(vars[e]) + t = {} + if determineexprtype_re_4.match(expr): # in parenthesis + t = determineexprtype(expr[1:-1], vars, rules) + else: + m = determineexprtype_re_5.match(expr) + if m: + rn = m.group('name') + t = determineexprtype(m.group('name'), vars, rules) + if t and 'attrspec' in t: + del t['attrspec'] + if not t: + if rn[0] in rules: + return _ensure_exprdict(rules[rn[0]]) + if expr[0] in '\'"': + return {'typespec': 'character', 'charselector': {'*': '*'}} + if not t: + outmess( + 'determineexprtype: could not determine expressions (%s) type.\n' % (repr(expr))) + return t + +###### + + +def crack2fortrangen(block, tab='\n', as_interface=False): + global skipfuncs, onlyfuncs + + setmesstext(block) + ret = '' + if isinstance(block, list): + for g in block: + if g and g['block'] in ['function', 'subroutine']: + if g['name'] in skipfuncs: + continue + if onlyfuncs and g['name'] not in onlyfuncs: + continue + ret = ret + crack2fortrangen(g, tab, as_interface=as_interface) + return ret + prefix = '' + name = '' + args = '' + blocktype = block['block'] + if blocktype == 'program': + return '' + argsl = [] + if 'name' in block: + name = block['name'] + if 'args' in block: + vars = block['vars'] + for a in block['args']: + a = expr2name(a, block, argsl) + if not isintent_callback(vars[a]): + argsl.append(a) + if block['block'] == 'function' or argsl: + args = '(%s)' % ','.join(argsl) + f2pyenhancements = '' + if 'f2pyenhancements' in block: + for k in list(block['f2pyenhancements'].keys()): + f2pyenhancements = '%s%s%s %s' % ( + f2pyenhancements, tab + tabchar, k, block['f2pyenhancements'][k]) + intent_lst = block.get('intent', [])[:] + if blocktype == 'function' and 'callback' in intent_lst: + intent_lst.remove('callback') + if intent_lst: + f2pyenhancements = '%s%sintent(%s) %s' %\ + (f2pyenhancements, tab + tabchar, + ','.join(intent_lst), name) + use = '' + if 'use' in block: + use = use2fortran(block['use'], tab + tabchar) + common = '' + if 'common' in block: + common = common2fortran(block['common'], tab + tabchar) + if name == 'unknown_interface': + name = '' + result = '' + if 'result' in block: + result = ' result (%s)' % block['result'] + if block['result'] not in argsl: + argsl.append(block['result']) + body = crack2fortrangen(block['body'], tab + tabchar, as_interface=as_interface) + vars = vars2fortran( + block, block['vars'], argsl, tab + tabchar, as_interface=as_interface) + mess = '' + if 'from' in block and not as_interface: + mess = '! in %s' % block['from'] + if 'entry' in block: + entry_stmts = '' + for k, i in list(block['entry'].items()): + entry_stmts = '%s%sentry %s(%s)' \ + % (entry_stmts, tab + tabchar, k, ','.join(i)) + body = body + entry_stmts + if blocktype == 'block data' and name == '_BLOCK_DATA_': + name = '' + ret = '%s%s%s %s%s%s %s%s%s%s%s%s%send %s %s' % ( + tab, prefix, blocktype, name, args, result, mess, f2pyenhancements, use, vars, common, body, tab, blocktype, name) + return ret + + +def common2fortran(common, tab=''): + ret = '' + for k in list(common.keys()): + if k == '_BLNK_': + ret = '%s%scommon %s' % (ret, tab, ','.join(common[k])) + else: + ret = '%s%scommon /%s/ %s' % (ret, tab, k, ','.join(common[k])) + return ret + + +def use2fortran(use, tab=''): + ret = '' + for m in list(use.keys()): + ret = '%s%suse %s,' % (ret, tab, m) + if use[m] == {}: + if ret and ret[-1] == ',': + ret = ret[:-1] + continue + if 'only' in use[m] and use[m]['only']: + ret = '%s only:' % (ret) + if 'map' in use[m] and use[m]['map']: + c = ' ' + for k in list(use[m]['map'].keys()): + if k == use[m]['map'][k]: + ret = '%s%s%s' % (ret, c, k) + c = ',' + else: + ret = '%s%s%s=>%s' % (ret, c, k, use[m]['map'][k]) + c = ',' + if ret and ret[-1] == ',': + ret = ret[:-1] + return ret + + +def true_intent_list(var): + lst = var['intent'] + ret = [] + for intent in lst: + try: + f = globals()['isintent_%s' % intent] + except KeyError: + pass + else: + if f(var): + ret.append(intent) + return ret + + +def vars2fortran(block, vars, args, tab='', as_interface=False): + setmesstext(block) + ret = '' + nout = [] + for a in args: + if a in block['vars']: + nout.append(a) + if 'commonvars' in block: + for a in block['commonvars']: + if a in vars: + if a not in nout: + nout.append(a) + else: + errmess( + 'vars2fortran: Confused?!: "%s" is not defined in vars.\n' % a) + if 'varnames' in block: + nout.extend(block['varnames']) + if not as_interface: + for a in list(vars.keys()): + if a not in nout: + nout.append(a) + for a in nout: + if 'depend' in vars[a]: + for d in vars[a]['depend']: + if d in vars and 'depend' in vars[d] and a in vars[d]['depend']: + errmess( + 'vars2fortran: Warning: cross-dependence between variables "%s" and "%s"\n' % (a, d)) + if 'externals' in block and a in block['externals']: + if isintent_callback(vars[a]): + ret = '%s%sintent(callback) %s' % (ret, tab, a) + ret = '%s%sexternal %s' % (ret, tab, a) + if isoptional(vars[a]): + ret = '%s%soptional %s' % (ret, tab, a) + if a in vars and 'typespec' not in vars[a]: + continue + cont = 1 + for b in block['body']: + if a == b['name'] and b['block'] == 'function': + cont = 0 + break + if cont: + continue + if a not in vars: + show(vars) + outmess('vars2fortran: No definition for argument "%s".\n' % a) + continue + if a == block['name']: + if block['block'] != 'function' or block.get('result'): + # 1) skip declaring a variable that name matches with + # subroutine name + # 2) skip declaring function when its type is + # declared via `result` construction + continue + if 'typespec' not in vars[a]: + if 'attrspec' in vars[a] and 'external' in vars[a]['attrspec']: + if a in args: + ret = '%s%sexternal %s' % (ret, tab, a) + continue + show(vars[a]) + outmess('vars2fortran: No typespec for argument "%s".\n' % a) + continue + vardef = vars[a]['typespec'] + if vardef == 'type' and 'typename' in vars[a]: + vardef = '%s(%s)' % (vardef, vars[a]['typename']) + selector = {} + if 'kindselector' in vars[a]: + selector = vars[a]['kindselector'] + elif 'charselector' in vars[a]: + selector = vars[a]['charselector'] + if '*' in selector: + if selector['*'] in ['*', ':']: + vardef = '%s*(%s)' % (vardef, selector['*']) + else: + vardef = '%s*%s' % (vardef, selector['*']) + else: + if 'len' in selector: + vardef = '%s(len=%s' % (vardef, selector['len']) + if 'kind' in selector: + vardef = '%s,kind=%s)' % (vardef, selector['kind']) + else: + vardef = '%s)' % (vardef) + elif 'kind' in selector: + vardef = '%s(kind=%s)' % (vardef, selector['kind']) + c = ' ' + if 'attrspec' in vars[a]: + attr = [l for l in vars[a]['attrspec'] + if l not in ['external']] + if as_interface and 'intent(in)' in attr and 'intent(out)' in attr: + # In Fortran, intent(in, out) are conflicting while + # intent(in, out) can be specified only via + # `!f2py intent(out) ..`. + # So, for the Fortran interface, we'll drop + # intent(out) to resolve the conflict. + attr.remove('intent(out)') + if attr: + vardef = '%s, %s' % (vardef, ','.join(attr)) + c = ',' + if 'dimension' in vars[a]: + vardef = '%s%sdimension(%s)' % ( + vardef, c, ','.join(vars[a]['dimension'])) + c = ',' + if 'intent' in vars[a]: + lst = true_intent_list(vars[a]) + if lst: + vardef = '%s%sintent(%s)' % (vardef, c, ','.join(lst)) + c = ',' + if 'check' in vars[a]: + vardef = '%s%scheck(%s)' % (vardef, c, ','.join(vars[a]['check'])) + c = ',' + if 'depend' in vars[a]: + vardef = '%s%sdepend(%s)' % ( + vardef, c, ','.join(vars[a]['depend'])) + c = ',' + if '=' in vars[a]: + v = vars[a]['='] + if vars[a]['typespec'] in ['complex', 'double complex']: + try: + v = eval(v) + v = '(%s,%s)' % (v.real, v.imag) + except Exception: + pass + vardef = '%s :: %s=%s' % (vardef, a, v) + else: + vardef = '%s :: %s' % (vardef, a) + ret = '%s%s%s' % (ret, tab, vardef) + return ret +###### + + +# We expose post_processing_hooks as global variable so that +# user-libraries could register their own hooks to f2py. +post_processing_hooks = [] + + +def crackfortran(files): + global usermodules, post_processing_hooks + + outmess('Reading fortran codes...\n', 0) + readfortrancode(files, crackline) + outmess('Post-processing...\n', 0) + usermodules = [] + postlist = postcrack(grouplist[0]) + outmess('Applying post-processing hooks...\n', 0) + for hook in post_processing_hooks: + outmess(f' {hook.__name__}\n', 0) + postlist = traverse(postlist, hook) + outmess('Post-processing (stage 2)...\n', 0) + postlist = postcrack2(postlist) + return usermodules + postlist + + +def crack2fortran(block): + global f2py_version + + pyf = crack2fortrangen(block) + '\n' + header = """! -*- f90 -*- +! Note: the context of this file is case sensitive. +""" + footer = """ +! This file was auto-generated with f2py (version:%s). +! See: +! https://web.archive.org/web/20140822061353/http://cens.ioc.ee/projects/f2py2e +""" % (f2py_version) + return header + pyf + footer + + +def _is_visit_pair(obj): + return (isinstance(obj, tuple) + and len(obj) == 2 + and isinstance(obj[0], (int, str))) + + +def traverse(obj, visit, parents=[], result=None, *args, **kwargs): + '''Traverse f2py data structure with the following visit function: + + def visit(item, parents, result, *args, **kwargs): + """ + + parents is a list of key-"f2py data structure" pairs from which + items are taken from. + + result is a f2py data structure that is filled with the + return value of the visit function. + + item is 2-tuple (index, value) if parents[-1][1] is a list + item is 2-tuple (key, value) if parents[-1][1] is a dict + + The return value of visit must be None, or of the same kind as + item, that is, if parents[-1] is a list, the return value must + be 2-tuple (new_index, new_value), or if parents[-1] is a + dict, the return value must be 2-tuple (new_key, new_value). + + If new_index or new_value is None, the return value of visit + is ignored, that is, it will not be added to the result. + + If the return value is None, the content of obj will be + traversed, otherwise not. + """ + ''' + + if _is_visit_pair(obj): + if obj[0] == 'parent_block': + # avoid infinite recursion + return obj + new_result = visit(obj, parents, result, *args, **kwargs) + if new_result is not None: + assert _is_visit_pair(new_result) + return new_result + parent = obj + result_key, obj = obj + else: + parent = (None, obj) + result_key = None + + if isinstance(obj, list): + new_result = [] + for index, value in enumerate(obj): + new_index, new_item = traverse((index, value), visit, + parents=parents + [parent], + result=result, *args, **kwargs) + if new_index is not None: + new_result.append(new_item) + elif isinstance(obj, dict): + new_result = dict() + for key, value in obj.items(): + new_key, new_value = traverse((key, value), visit, + parents=parents + [parent], + result=result, *args, **kwargs) + if new_key is not None: + new_result[new_key] = new_value + else: + new_result = obj + + if result_key is None: + return new_result + return result_key, new_result + + +def character_backward_compatibility_hook(item, parents, result, + *args, **kwargs): + """Previously, Fortran character was incorrectly treated as + character*1. This hook fixes the usage of the corresponding + variables in `check`, `dimension`, `=`, and `callstatement` + expressions. + + The usage of `char*` in `callprotoargument` expression can be left + unchanged because C `character` is C typedef of `char`, although, + new implementations should use `character*` in the corresponding + expressions. + + See https://github.com/numpy/numpy/pull/19388 for more information. + + """ + parent_key, parent_value = parents[-1] + key, value = item + + def fix_usage(varname, value): + value = re.sub(r'[*]\s*\b' + varname + r'\b', varname, value) + value = re.sub(r'\b' + varname + r'\b\s*[\[]\s*0\s*[\]]', + varname, value) + return value + + if parent_key in ['dimension', 'check']: + assert parents[-3][0] == 'vars' + vars_dict = parents[-3][1] + elif key == '=': + assert parents[-2][0] == 'vars' + vars_dict = parents[-2][1] + else: + vars_dict = None + + new_value = None + if vars_dict is not None: + new_value = value + for varname, vd in vars_dict.items(): + if ischaracter(vd): + new_value = fix_usage(varname, new_value) + elif key == 'callstatement': + vars_dict = parents[-2][1]['vars'] + new_value = value + for varname, vd in vars_dict.items(): + if ischaracter(vd): + # replace all occurrences of `` with + # `&` in argument passing + new_value = re.sub( + r'(? `{new_value}`\n', 1) + return (key, new_value) + + +post_processing_hooks.append(character_backward_compatibility_hook) + + +if __name__ == "__main__": + files = [] + funcs = [] + f = 1 + f2 = 0 + f3 = 0 + showblocklist = 0 + for l in sys.argv[1:]: + if l == '': + pass + elif l[0] == ':': + f = 0 + elif l == '-quiet': + quiet = 1 + verbose = 0 + elif l == '-verbose': + verbose = 2 + quiet = 0 + elif l == '-fix': + if strictf77: + outmess( + 'Use option -f90 before -fix if Fortran 90 code is in fix form.\n', 0) + skipemptyends = 1 + sourcecodeform = 'fix' + elif l == '-skipemptyends': + skipemptyends = 1 + elif l == '--ignore-contains': + ignorecontains = 1 + elif l == '-f77': + strictf77 = 1 + sourcecodeform = 'fix' + elif l == '-f90': + strictf77 = 0 + sourcecodeform = 'free' + skipemptyends = 1 + elif l == '-h': + f2 = 1 + elif l == '-show': + showblocklist = 1 + elif l == '-m': + f3 = 1 + elif l[0] == '-': + errmess('Unknown option %s\n' % repr(l)) + elif f2: + f2 = 0 + pyffilename = l + elif f3: + f3 = 0 + f77modulename = l + elif f: + try: + open(l).close() + files.append(l) + except OSError as detail: + errmess(f'OSError: {detail!s}\n') + else: + funcs.append(l) + if not strictf77 and f77modulename and not skipemptyends: + outmess("""\ + Warning: You have specified module name for non Fortran 77 code that + should not need one (expect if you are scanning F90 code for non + module blocks but then you should use flag -skipemptyends and also + be sure that the files do not contain programs without program + statement). +""", 0) + + postlist = crackfortran(files) + if pyffilename: + outmess('Writing fortran code to file %s\n' % repr(pyffilename), 0) + pyf = crack2fortran(postlist) + with open(pyffilename, 'w') as f: + f.write(pyf) + if showblocklist: + show(postlist) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/diagnose.py b/venv/lib/python3.10/site-packages/numpy/f2py/diagnose.py new file mode 100644 index 0000000000000000000000000000000000000000..523c2c679d9edf78b7d82bf77904348f0f99a3e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/diagnose.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +import os +import sys +import tempfile + + +def run_command(cmd): + print('Running %r:' % (cmd)) + os.system(cmd) + print('------') + + +def run(): + _path = os.getcwd() + os.chdir(tempfile.gettempdir()) + print('------') + print('os.name=%r' % (os.name)) + print('------') + print('sys.platform=%r' % (sys.platform)) + print('------') + print('sys.version:') + print(sys.version) + print('------') + print('sys.prefix:') + print(sys.prefix) + print('------') + print('sys.path=%r' % (':'.join(sys.path))) + print('------') + + try: + import numpy + has_newnumpy = 1 + except ImportError as e: + print('Failed to import new numpy:', e) + has_newnumpy = 0 + + try: + from numpy.f2py import f2py2e + has_f2py2e = 1 + except ImportError as e: + print('Failed to import f2py2e:', e) + has_f2py2e = 0 + + try: + import numpy.distutils + has_numpy_distutils = 2 + except ImportError: + try: + import numpy_distutils + has_numpy_distutils = 1 + except ImportError as e: + print('Failed to import numpy_distutils:', e) + has_numpy_distutils = 0 + + if has_newnumpy: + try: + print('Found new numpy version %r in %s' % + (numpy.__version__, numpy.__file__)) + except Exception as msg: + print('error:', msg) + print('------') + + if has_f2py2e: + try: + print('Found f2py2e version %r in %s' % + (f2py2e.__version__.version, f2py2e.__file__)) + except Exception as msg: + print('error:', msg) + print('------') + + if has_numpy_distutils: + try: + if has_numpy_distutils == 2: + print('Found numpy.distutils version %r in %r' % ( + numpy.distutils.__version__, + numpy.distutils.__file__)) + else: + print('Found numpy_distutils version %r in %r' % ( + numpy_distutils.numpy_distutils_version.numpy_distutils_version, + numpy_distutils.__file__)) + print('------') + except Exception as msg: + print('error:', msg) + print('------') + try: + if has_numpy_distutils == 1: + print( + 'Importing numpy_distutils.command.build_flib ...', end=' ') + import numpy_distutils.command.build_flib as build_flib + print('ok') + print('------') + try: + print( + 'Checking availability of supported Fortran compilers:') + for compiler_class in build_flib.all_compilers: + compiler_class(verbose=1).is_available() + print('------') + except Exception as msg: + print('error:', msg) + print('------') + except Exception as msg: + print( + 'error:', msg, '(ignore it, build_flib is obsolete for numpy.distutils 0.2.2 and up)') + print('------') + try: + if has_numpy_distutils == 2: + print('Importing numpy.distutils.fcompiler ...', end=' ') + import numpy.distutils.fcompiler as fcompiler + else: + print('Importing numpy_distutils.fcompiler ...', end=' ') + import numpy_distutils.fcompiler as fcompiler + print('ok') + print('------') + try: + print('Checking availability of supported Fortran compilers:') + fcompiler.show_fcompilers() + print('------') + except Exception as msg: + print('error:', msg) + print('------') + except Exception as msg: + print('error:', msg) + print('------') + try: + if has_numpy_distutils == 2: + print('Importing numpy.distutils.cpuinfo ...', end=' ') + from numpy.distutils.cpuinfo import cpuinfo + print('ok') + print('------') + else: + try: + print( + 'Importing numpy_distutils.command.cpuinfo ...', end=' ') + from numpy_distutils.command.cpuinfo import cpuinfo + print('ok') + print('------') + except Exception as msg: + print('error:', msg, '(ignore it)') + print('Importing numpy_distutils.cpuinfo ...', end=' ') + from numpy_distutils.cpuinfo import cpuinfo + print('ok') + print('------') + cpu = cpuinfo() + print('CPU information:', end=' ') + for name in dir(cpuinfo): + if name[0] == '_' and name[1] != '_' and getattr(cpu, name[1:])(): + print(name[1:], end=' ') + print('------') + except Exception as msg: + print('error:', msg) + print('------') + os.chdir(_path) +if __name__ == "__main__": + run() diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/f2py2e.py b/venv/lib/python3.10/site-packages/numpy/f2py/f2py2e.py new file mode 100644 index 0000000000000000000000000000000000000000..c0f801e06c7fc0f9d9634c695028edd333f6502b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/f2py2e.py @@ -0,0 +1,783 @@ +""" + +f2py2e - Fortran to Python C/API generator. 2nd Edition. + See __usage__ below. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +import sys +import os +import pprint +import re +import argparse + +from . import crackfortran +from . import rules +from . import cb_rules +from . import auxfuncs +from . import cfuncs +from . import f90mod_rules +from . import __version__ +from . import capi_maps +from .cfuncs import errmess +from numpy.f2py._backends import f2py_build_generator + +f2py_version = __version__.version +numpy_version = __version__.version + +# outmess=sys.stdout.write +show = pprint.pprint +outmess = auxfuncs.outmess +MESON_ONLY_VER = (sys.version_info >= (3, 12)) + +__usage__ =\ +f"""Usage: + +1) To construct extension module sources: + + f2py [] [[[only:]||[skip:]] \\ + ] \\ + [: ...] + +2) To compile fortran files and build extension modules: + + f2py -c [, , ] + +3) To generate signature files: + + f2py -h ...< same options as in (1) > + +Description: This program generates a Python C/API file (module.c) + that contains wrappers for given fortran functions so that they + can be called from Python. With the -c option the corresponding + extension modules are built. + +Options: + + -h Write signatures of the fortran routines to file + and exit. You can then edit and use it instead + of . If ==stdout then the + signatures are printed to stdout. + Names of fortran routines for which Python C/API + functions will be generated. Default is all that are found + in . + Paths to fortran/signature files that will be scanned for + in order to determine their signatures. + skip: Ignore fortran functions that follow until `:'. + only: Use only fortran functions that follow until `:'. + : Get back to mode. + + -m Name of the module; f2py generates a Python/C API + file module.c or extension module . + Default is 'untitled'. + + '-include

' Writes additional headers in the C wrapper, can be passed + multiple times, generates #include
each time. + + --[no-]lower Do [not] lower the cases in . By default, + --lower is assumed with -h key, and --no-lower without -h key. + + --build-dir All f2py generated files are created in . + Default is tempfile.mkdtemp(). + + --overwrite-signature Overwrite existing signature file. + + --[no-]latex-doc Create (or not) module.tex. + Default is --no-latex-doc. + --short-latex Create 'incomplete' LaTeX document (without commands + \\documentclass, \\tableofcontents, and \\begin{{document}}, + \\end{{document}}). + + --[no-]rest-doc Create (or not) module.rst. + Default is --no-rest-doc. + + --debug-capi Create C/API code that reports the state of the wrappers + during runtime. Useful for debugging. + + --[no-]wrap-functions Create Fortran subroutine wrappers to Fortran 77 + functions. --wrap-functions is default because it ensures + maximum portability/compiler independence. + + --[no-]freethreading-compatible Create a module that declares it does or + doesn't require the GIL. The default is + --freethreading-compatible for backward + compatibility. Inspect the Fortran code you are wrapping for + thread safety issues before passing + --no-freethreading-compatible, as f2py does not analyze + fortran code for thread safety issues. + + --include-paths ::... Search include files from the given + directories. + + --help-link [..] List system resources found by system_info.py. See also + --link- switch below. [..] is optional list + of resources names. E.g. try 'f2py --help-link lapack_opt'. + + --f2cmap Load Fortran-to-Python KIND specification from the given + file. Default: .f2py_f2cmap in current directory. + + --quiet Run quietly. + --verbose Run with extra verbosity. + --skip-empty-wrappers Only generate wrapper files when needed. + -v Print f2py version ID and exit. + + +build backend options (only effective with -c) +[NO_MESON] is used to indicate an option not meant to be used +with the meson backend or above Python 3.12: + + --fcompiler= Specify Fortran compiler type by vendor [NO_MESON] + --compiler= Specify distutils C compiler type [NO_MESON] + + --help-fcompiler List available Fortran compilers and exit [NO_MESON] + --f77exec= Specify the path to F77 compiler [NO_MESON] + --f90exec= Specify the path to F90 compiler [NO_MESON] + --f77flags= Specify F77 compiler flags + --f90flags= Specify F90 compiler flags + --opt= Specify optimization flags [NO_MESON] + --arch= Specify architecture specific optimization flags [NO_MESON] + --noopt Compile without optimization [NO_MESON] + --noarch Compile without arch-dependent optimization [NO_MESON] + --debug Compile with debugging information + + --dep + Specify a meson dependency for the module. This may + be passed multiple times for multiple dependencies. + Dependencies are stored in a list for further processing. + + Example: --dep lapack --dep scalapack + This will identify "lapack" and "scalapack" as dependencies + and remove them from argv, leaving a dependencies list + containing ["lapack", "scalapack"]. + + --backend + Specify the build backend for the compilation process. + The supported backends are 'meson' and 'distutils'. + If not specified, defaults to 'distutils'. On + Python 3.12 or higher, the default is 'meson'. + +Extra options (only effective with -c): + + --link- Link extension module with as defined + by numpy.distutils/system_info.py. E.g. to link + with optimized LAPACK libraries (vecLib on MacOSX, + ATLAS elsewhere), use --link-lapack_opt. + See also --help-link switch. [NO_MESON] + + -L/path/to/lib/ -l + -D -U + -I/path/to/include/ + .o .so .a + + Using the following macros may be required with non-gcc Fortran + compilers: + -DPREPEND_FORTRAN -DNO_APPEND_FORTRAN -DUPPERCASE_FORTRAN + + When using -DF2PY_REPORT_ATEXIT, a performance report of F2PY + interface is printed out at exit (platforms: Linux). + + When using -DF2PY_REPORT_ON_ARRAY_COPY=, a message is + sent to stderr whenever F2PY interface makes a copy of an + array. Integer sets the threshold for array sizes when + a message should be shown. + +Version: {f2py_version} +numpy Version: {numpy_version} +License: NumPy license (see LICENSE.txt in the NumPy source code) +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +https://numpy.org/doc/stable/f2py/index.html\n""" + + +def scaninputline(inputline): + files, skipfuncs, onlyfuncs, debug = [], [], [], [] + f, f2, f3, f5, f6, f8, f9, f10 = 1, 0, 0, 0, 0, 0, 0, 0 + verbose = 1 + emptygen = True + dolc = -1 + dolatexdoc = 0 + dorestdoc = 0 + wrapfuncs = 1 + buildpath = '.' + include_paths, freethreading_compatible, inputline = get_newer_options(inputline) + signsfile, modulename = None, None + options = {'buildpath': buildpath, + 'coutput': None, + 'f2py_wrapper_output': None} + for l in inputline: + if l == '': + pass + elif l == 'only:': + f = 0 + elif l == 'skip:': + f = -1 + elif l == ':': + f = 1 + elif l[:8] == '--debug-': + debug.append(l[8:]) + elif l == '--lower': + dolc = 1 + elif l == '--build-dir': + f6 = 1 + elif l == '--no-lower': + dolc = 0 + elif l == '--quiet': + verbose = 0 + elif l == '--verbose': + verbose += 1 + elif l == '--latex-doc': + dolatexdoc = 1 + elif l == '--no-latex-doc': + dolatexdoc = 0 + elif l == '--rest-doc': + dorestdoc = 1 + elif l == '--no-rest-doc': + dorestdoc = 0 + elif l == '--wrap-functions': + wrapfuncs = 1 + elif l == '--no-wrap-functions': + wrapfuncs = 0 + elif l == '--short-latex': + options['shortlatex'] = 1 + elif l == '--coutput': + f8 = 1 + elif l == '--f2py-wrapper-output': + f9 = 1 + elif l == '--f2cmap': + f10 = 1 + elif l == '--overwrite-signature': + options['h-overwrite'] = 1 + elif l == '-h': + f2 = 1 + elif l == '-m': + f3 = 1 + elif l[:2] == '-v': + print(f2py_version) + sys.exit() + elif l == '--show-compilers': + f5 = 1 + elif l[:8] == '-include': + cfuncs.outneeds['userincludes'].append(l[9:-1]) + cfuncs.userincludes[l[9:-1]] = '#include ' + l[8:] + elif l == '--skip-empty-wrappers': + emptygen = False + elif l[0] == '-': + errmess('Unknown option %s\n' % repr(l)) + sys.exit() + elif f2: + f2 = 0 + signsfile = l + elif f3: + f3 = 0 + modulename = l + elif f6: + f6 = 0 + buildpath = l + elif f8: + f8 = 0 + options["coutput"] = l + elif f9: + f9 = 0 + options["f2py_wrapper_output"] = l + elif f10: + f10 = 0 + options["f2cmap_file"] = l + elif f == 1: + try: + with open(l): + pass + files.append(l) + except OSError as detail: + errmess(f'OSError: {detail!s}. Skipping file "{l!s}".\n') + elif f == -1: + skipfuncs.append(l) + elif f == 0: + onlyfuncs.append(l) + if not f5 and not files and not modulename: + print(__usage__) + sys.exit() + if not os.path.isdir(buildpath): + if not verbose: + outmess('Creating build directory %s\n' % (buildpath)) + os.mkdir(buildpath) + if signsfile: + signsfile = os.path.join(buildpath, signsfile) + if signsfile and os.path.isfile(signsfile) and 'h-overwrite' not in options: + errmess( + 'Signature file "%s" exists!!! Use --overwrite-signature to overwrite.\n' % (signsfile)) + sys.exit() + + options['emptygen'] = emptygen + options['debug'] = debug + options['verbose'] = verbose + if dolc == -1 and not signsfile: + options['do-lower'] = 0 + else: + options['do-lower'] = dolc + if modulename: + options['module'] = modulename + if signsfile: + options['signsfile'] = signsfile + if onlyfuncs: + options['onlyfuncs'] = onlyfuncs + if skipfuncs: + options['skipfuncs'] = skipfuncs + options['dolatexdoc'] = dolatexdoc + options['dorestdoc'] = dorestdoc + options['wrapfuncs'] = wrapfuncs + options['buildpath'] = buildpath + options['include_paths'] = include_paths + options['requires_gil'] = not freethreading_compatible + options.setdefault('f2cmap_file', None) + return files, options + + +def callcrackfortran(files, options): + rules.options = options + crackfortran.debug = options['debug'] + crackfortran.verbose = options['verbose'] + if 'module' in options: + crackfortran.f77modulename = options['module'] + if 'skipfuncs' in options: + crackfortran.skipfuncs = options['skipfuncs'] + if 'onlyfuncs' in options: + crackfortran.onlyfuncs = options['onlyfuncs'] + crackfortran.include_paths[:] = options['include_paths'] + crackfortran.dolowercase = options['do-lower'] + postlist = crackfortran.crackfortran(files) + if 'signsfile' in options: + outmess('Saving signatures to file "%s"\n' % (options['signsfile'])) + pyf = crackfortran.crack2fortran(postlist) + if options['signsfile'][-6:] == 'stdout': + sys.stdout.write(pyf) + else: + with open(options['signsfile'], 'w') as f: + f.write(pyf) + if options["coutput"] is None: + for mod in postlist: + mod["coutput"] = "%smodule.c" % mod["name"] + else: + for mod in postlist: + mod["coutput"] = options["coutput"] + if options["f2py_wrapper_output"] is None: + for mod in postlist: + mod["f2py_wrapper_output"] = "%s-f2pywrappers.f" % mod["name"] + else: + for mod in postlist: + mod["f2py_wrapper_output"] = options["f2py_wrapper_output"] + for mod in postlist: + if options["requires_gil"]: + mod['gil_used'] = 'Py_MOD_GIL_USED' + else: + mod['gil_used'] = 'Py_MOD_GIL_NOT_USED' + return postlist + + +def buildmodules(lst): + cfuncs.buildcfuncs() + outmess('Building modules...\n') + modules, mnames, isusedby = [], [], {} + for item in lst: + if '__user__' in item['name']: + cb_rules.buildcallbacks(item) + else: + if 'use' in item: + for u in item['use'].keys(): + if u not in isusedby: + isusedby[u] = [] + isusedby[u].append(item['name']) + modules.append(item) + mnames.append(item['name']) + ret = {} + for module, name in zip(modules, mnames): + if name in isusedby: + outmess('\tSkipping module "%s" which is used by %s.\n' % ( + name, ','.join('"%s"' % s for s in isusedby[name]))) + else: + um = [] + if 'use' in module: + for u in module['use'].keys(): + if u in isusedby and u in mnames: + um.append(modules[mnames.index(u)]) + else: + outmess( + f'\tModule "{name}" uses nonexisting "{u}" ' + 'which will be ignored.\n') + ret[name] = {} + dict_append(ret[name], rules.buildmodule(module, um)) + return ret + + +def dict_append(d_out, d_in): + for (k, v) in d_in.items(): + if k not in d_out: + d_out[k] = [] + if isinstance(v, list): + d_out[k] = d_out[k] + v + else: + d_out[k].append(v) + + +def run_main(comline_list): + """ + Equivalent to running:: + + f2py + + where ``=string.join(,' ')``, but in Python. Unless + ``-h`` is used, this function returns a dictionary containing + information on generated modules and their dependencies on source + files. + + You cannot build extension modules with this function, that is, + using ``-c`` is not allowed. Use the ``compile`` command instead. + + Examples + -------- + The command ``f2py -m scalar scalar.f`` can be executed from Python as + follows. + + .. literalinclude:: ../../source/f2py/code/results/run_main_session.dat + :language: python + + """ + crackfortran.reset_global_f2py_vars() + f2pydir = os.path.dirname(os.path.abspath(cfuncs.__file__)) + fobjhsrc = os.path.join(f2pydir, 'src', 'fortranobject.h') + fobjcsrc = os.path.join(f2pydir, 'src', 'fortranobject.c') + # gh-22819 -- begin + parser = make_f2py_compile_parser() + args, comline_list = parser.parse_known_args(comline_list) + pyf_files, _ = filter_files("", "[.]pyf([.]src|)", comline_list) + # Checks that no existing modulename is defined in a pyf file + # TODO: Remove all this when scaninputline is replaced + if args.module_name: + if "-h" in comline_list: + modname = ( + args.module_name + ) # Directly use from args when -h is present + else: + modname = validate_modulename( + pyf_files, args.module_name + ) # Validate modname when -h is not present + comline_list += ['-m', modname] # needed for the rest of scaninputline + # gh-22819 -- end + files, options = scaninputline(comline_list) + auxfuncs.options = options + capi_maps.load_f2cmap_file(options['f2cmap_file']) + postlist = callcrackfortran(files, options) + isusedby = {} + for plist in postlist: + if 'use' in plist: + for u in plist['use'].keys(): + if u not in isusedby: + isusedby[u] = [] + isusedby[u].append(plist['name']) + for plist in postlist: + if plist['block'] == 'python module' and '__user__' in plist['name']: + if plist['name'] in isusedby: + # if not quiet: + outmess( + f'Skipping Makefile build for module "{plist["name"]}" ' + 'which is used by {}\n'.format( + ','.join(f'"{s}"' for s in isusedby[plist['name']]))) + if 'signsfile' in options: + if options['verbose'] > 1: + outmess( + 'Stopping. Edit the signature file and then run f2py on the signature file: ') + outmess('%s %s\n' % + (os.path.basename(sys.argv[0]), options['signsfile'])) + return + for plist in postlist: + if plist['block'] != 'python module': + if 'python module' not in options: + errmess( + 'Tip: If your original code is Fortran source then you must use -m option.\n') + raise TypeError('All blocks must be python module blocks but got %s' % ( + repr(plist['block']))) + auxfuncs.debugoptions = options['debug'] + f90mod_rules.options = options + auxfuncs.wrapfuncs = options['wrapfuncs'] + + ret = buildmodules(postlist) + + for mn in ret.keys(): + dict_append(ret[mn], {'csrc': fobjcsrc, 'h': fobjhsrc}) + return ret + + +def filter_files(prefix, suffix, files, remove_prefix=None): + """ + Filter files by prefix and suffix. + """ + filtered, rest = [], [] + match = re.compile(prefix + r'.*' + suffix + r'\Z').match + if remove_prefix: + ind = len(prefix) + else: + ind = 0 + for file in [x.strip() for x in files]: + if match(file): + filtered.append(file[ind:]) + else: + rest.append(file) + return filtered, rest + + +def get_prefix(module): + p = os.path.dirname(os.path.dirname(module.__file__)) + return p + + +class CombineIncludePaths(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + include_paths_set = set(getattr(namespace, 'include_paths', []) or []) + if option_string == "--include_paths": + outmess("Use --include-paths or -I instead of --include_paths which will be removed") + if option_string == "--include-paths" or option_string == "--include_paths": + include_paths_set.update(values.split(':')) + else: + include_paths_set.add(values) + namespace.include_paths = list(include_paths_set) + +def f2py_parser(): + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("-I", dest="include_paths", action=CombineIncludePaths) + parser.add_argument("--include-paths", dest="include_paths", action=CombineIncludePaths) + parser.add_argument("--include_paths", dest="include_paths", action=CombineIncludePaths) + parser.add_argument("--freethreading-compatible", dest="ftcompat", action=argparse.BooleanOptionalAction) + return parser + +def get_newer_options(iline): + iline = (' '.join(iline)).split() + parser = f2py_parser() + args, remain = parser.parse_known_args(iline) + ipaths = args.include_paths + if args.include_paths is None: + ipaths = [] + return ipaths, args.ftcompat, remain + +def make_f2py_compile_parser(): + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--dep", action="append", dest="dependencies") + parser.add_argument("--backend", choices=['meson', 'distutils'], default='distutils') + parser.add_argument("-m", dest="module_name") + return parser + +def preparse_sysargv(): + # To keep backwards bug compatibility, newer flags are handled by argparse, + # and `sys.argv` is passed to the rest of `f2py` as is. + parser = make_f2py_compile_parser() + + args, remaining_argv = parser.parse_known_args() + sys.argv = [sys.argv[0]] + remaining_argv + + backend_key = args.backend + if MESON_ONLY_VER and backend_key == 'distutils': + outmess("Cannot use distutils backend with Python>=3.12," + " using meson backend instead.\n") + backend_key = "meson" + + return { + "dependencies": args.dependencies or [], + "backend": backend_key, + "modulename": args.module_name, + } + +def run_compile(): + """ + Do it all in one call! + """ + import tempfile + + # Collect dependency flags, preprocess sys.argv + argy = preparse_sysargv() + modulename = argy["modulename"] + if modulename is None: + modulename = 'untitled' + dependencies = argy["dependencies"] + backend_key = argy["backend"] + build_backend = f2py_build_generator(backend_key) + + i = sys.argv.index('-c') + del sys.argv[i] + + remove_build_dir = 0 + try: + i = sys.argv.index('--build-dir') + except ValueError: + i = None + if i is not None: + build_dir = sys.argv[i + 1] + del sys.argv[i + 1] + del sys.argv[i] + else: + remove_build_dir = 1 + build_dir = tempfile.mkdtemp() + + _reg1 = re.compile(r'--link-') + sysinfo_flags = [_m for _m in sys.argv[1:] if _reg1.match(_m)] + sys.argv = [_m for _m in sys.argv if _m not in sysinfo_flags] + if sysinfo_flags: + sysinfo_flags = [f[7:] for f in sysinfo_flags] + + _reg2 = re.compile( + r'--((no-|)(wrap-functions|lower|freethreading-compatible)|debug-capi|quiet|skip-empty-wrappers)|-include') + f2py_flags = [_m for _m in sys.argv[1:] if _reg2.match(_m)] + sys.argv = [_m for _m in sys.argv if _m not in f2py_flags] + f2py_flags2 = [] + fl = 0 + for a in sys.argv[1:]: + if a in ['only:', 'skip:']: + fl = 1 + elif a == ':': + fl = 0 + if fl or a == ':': + f2py_flags2.append(a) + if f2py_flags2 and f2py_flags2[-1] != ':': + f2py_flags2.append(':') + f2py_flags.extend(f2py_flags2) + sys.argv = [_m for _m in sys.argv if _m not in f2py_flags2] + _reg3 = re.compile( + r'--((f(90)?compiler(-exec|)|compiler)=|help-compiler)') + flib_flags = [_m for _m in sys.argv[1:] if _reg3.match(_m)] + sys.argv = [_m for _m in sys.argv if _m not in flib_flags] + # TODO: Once distutils is dropped completely, i.e. min_ver >= 3.12, unify into --fflags + reg_f77_f90_flags = re.compile(r'--f(77|90)flags=') + reg_distutils_flags = re.compile(r'--((f(77|90)exec|opt|arch)=|(debug|noopt|noarch|help-fcompiler))') + fc_flags = [_m for _m in sys.argv[1:] if reg_f77_f90_flags.match(_m)] + distutils_flags = [_m for _m in sys.argv[1:] if reg_distutils_flags.match(_m)] + if not (MESON_ONLY_VER or backend_key == 'meson'): + fc_flags.extend(distutils_flags) + sys.argv = [_m for _m in sys.argv if _m not in (fc_flags + distutils_flags)] + + del_list = [] + for s in flib_flags: + v = '--fcompiler=' + if s[:len(v)] == v: + if MESON_ONLY_VER or backend_key == 'meson': + outmess( + "--fcompiler cannot be used with meson," + "set compiler with the FC environment variable\n" + ) + else: + from numpy.distutils import fcompiler + fcompiler.load_all_fcompiler_classes() + allowed_keys = list(fcompiler.fcompiler_class.keys()) + nv = ov = s[len(v):].lower() + if ov not in allowed_keys: + vmap = {} # XXX + try: + nv = vmap[ov] + except KeyError: + if ov not in vmap.values(): + print('Unknown vendor: "%s"' % (s[len(v):])) + nv = ov + i = flib_flags.index(s) + flib_flags[i] = '--fcompiler=' + nv + continue + for s in del_list: + i = flib_flags.index(s) + del flib_flags[i] + assert len(flib_flags) <= 2, repr(flib_flags) + + _reg5 = re.compile(r'--(verbose)') + setup_flags = [_m for _m in sys.argv[1:] if _reg5.match(_m)] + sys.argv = [_m for _m in sys.argv if _m not in setup_flags] + + if '--quiet' in f2py_flags: + setup_flags.append('--quiet') + + # Ugly filter to remove everything but sources + sources = sys.argv[1:] + f2cmapopt = '--f2cmap' + if f2cmapopt in sys.argv: + i = sys.argv.index(f2cmapopt) + f2py_flags.extend(sys.argv[i:i + 2]) + del sys.argv[i + 1], sys.argv[i] + sources = sys.argv[1:] + + pyf_files, _sources = filter_files("", "[.]pyf([.]src|)", sources) + sources = pyf_files + _sources + modulename = validate_modulename(pyf_files, modulename) + extra_objects, sources = filter_files('', '[.](o|a|so|dylib)', sources) + library_dirs, sources = filter_files('-L', '', sources, remove_prefix=1) + libraries, sources = filter_files('-l', '', sources, remove_prefix=1) + undef_macros, sources = filter_files('-U', '', sources, remove_prefix=1) + define_macros, sources = filter_files('-D', '', sources, remove_prefix=1) + for i in range(len(define_macros)): + name_value = define_macros[i].split('=', 1) + if len(name_value) == 1: + name_value.append(None) + if len(name_value) == 2: + define_macros[i] = tuple(name_value) + else: + print('Invalid use of -D:', name_value) + + # Construct wrappers / signatures / things + if backend_key == 'meson': + if not pyf_files: + outmess('Using meson backend\nWill pass --lower to f2py\nSee https://numpy.org/doc/stable/f2py/buildtools/meson.html\n') + f2py_flags.append('--lower') + run_main(f" {' '.join(f2py_flags)} -m {modulename} {' '.join(sources)}".split()) + else: + run_main(f" {' '.join(f2py_flags)} {' '.join(pyf_files)}".split()) + + # Order matters here, includes are needed for run_main above + include_dirs, _, sources = get_newer_options(sources) + # Now use the builder + builder = build_backend( + modulename, + sources, + extra_objects, + build_dir, + include_dirs, + library_dirs, + libraries, + define_macros, + undef_macros, + f2py_flags, + sysinfo_flags, + fc_flags, + flib_flags, + setup_flags, + remove_build_dir, + {"dependencies": dependencies}, + ) + + builder.compile() + + +def validate_modulename(pyf_files, modulename='untitled'): + if len(pyf_files) > 1: + raise ValueError("Only one .pyf file per call") + if pyf_files: + pyff = pyf_files[0] + pyf_modname = auxfuncs.get_f2py_modulename(pyff) + if modulename != pyf_modname: + outmess( + f"Ignoring -m {modulename}.\n" + f"{pyff} defines {pyf_modname} to be the modulename.\n" + ) + modulename = pyf_modname + return modulename + +def main(): + if '--help-link' in sys.argv[1:]: + sys.argv.remove('--help-link') + if MESON_ONLY_VER: + outmess("Use --dep for meson builds\n") + else: + from numpy.distutils.system_info import show_all + show_all() + return + + if '-c' in sys.argv[1:]: + run_compile() + else: + run_main(sys.argv[1:]) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/f90mod_rules.py b/venv/lib/python3.10/site-packages/numpy/f2py/f90mod_rules.py new file mode 100644 index 0000000000000000000000000000000000000000..b1cd1532065744822991d7d1f4c9a29a216dede9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/f90mod_rules.py @@ -0,0 +1,270 @@ +""" +Build F90 module support for f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +__version__ = "$Revision: 1.27 $"[10:-1] + +f2py_version = 'See `f2py -v`' + +import numpy as np + +from . import capi_maps +from . import func2subr +from .crackfortran import undo_rmbadname, undo_rmbadname1 + +# The environment provided by auxfuncs.py is needed for some calls to eval. +# As the needed functions cannot be determined by static inspection of the +# code, it is safest to use import * pending a major refactoring of f2py. +from .auxfuncs import * + +options = {} + + +def findf90modules(m): + if ismodule(m): + return [m] + if not hasbody(m): + return [] + ret = [] + for b in m['body']: + if ismodule(b): + ret.append(b) + else: + ret = ret + findf90modules(b) + return ret + +fgetdims1 = """\ + external f2pysetdata + logical ns + integer r,i + integer(%d) s(*) + ns = .FALSE. + if (allocated(d)) then + do i=1,r + if ((size(d,i).ne.s(i)).and.(s(i).ge.0)) then + ns = .TRUE. + end if + end do + if (ns) then + deallocate(d) + end if + end if + if ((.not.allocated(d)).and.(s(1).ge.1)) then""" % np.intp().itemsize + +fgetdims2 = """\ + end if + if (allocated(d)) then + do i=1,r + s(i) = size(d,i) + end do + end if + flag = 1 + call f2pysetdata(d,allocated(d))""" + +fgetdims2_sa = """\ + end if + if (allocated(d)) then + do i=1,r + s(i) = size(d,i) + end do + !s(r) must be equal to len(d(1)) + end if + flag = 2 + call f2pysetdata(d,allocated(d))""" + + +def buildhooks(pymod): + from . import rules + ret = {'f90modhooks': [], 'initf90modhooks': [], 'body': [], + 'need': ['F_FUNC', 'arrayobject.h'], + 'separatorsfor': {'includes0': '\n', 'includes': '\n'}, + 'docs': ['"Fortran 90/95 modules:\\n"'], + 'latexdoc': []} + fhooks = [''] + + def fadd(line, s=fhooks): + s[0] = '%s\n %s' % (s[0], line) + doc = [''] + + def dadd(line, s=doc): + s[0] = '%s\n%s' % (s[0], line) + + usenames = getuseblocks(pymod) + for m in findf90modules(pymod): + sargs, fargs, efargs, modobjs, notvars, onlyvars = [], [], [], [], [ + m['name']], [] + sargsp = [] + ifargs = [] + mfargs = [] + if hasbody(m): + for b in m['body']: + notvars.append(b['name']) + for n in m['vars'].keys(): + var = m['vars'][n] + + if (n not in notvars and isvariable(var)) and (not l_or(isintent_hide, isprivate)(var)): + onlyvars.append(n) + mfargs.append(n) + outmess('\t\tConstructing F90 module support for "%s"...\n' % + (m['name'])) + if len(onlyvars) == 0 and len(notvars) == 1 and m['name'] in notvars: + outmess(f"\t\t\tSkipping {m['name']} since there are no public vars/func in this module...\n") + continue + + # gh-25186 + if m['name'] in usenames and containscommon(m): + outmess(f"\t\t\tSkipping {m['name']} since it is in 'use' and contains a common block...\n") + continue + if onlyvars: + outmess('\t\t Variables: %s\n' % (' '.join(onlyvars))) + chooks = [''] + + def cadd(line, s=chooks): + s[0] = '%s\n%s' % (s[0], line) + ihooks = [''] + + def iadd(line, s=ihooks): + s[0] = '%s\n%s' % (s[0], line) + + vrd = capi_maps.modsign2map(m) + cadd('static FortranDataDef f2py_%s_def[] = {' % (m['name'])) + dadd('\\subsection{Fortran 90/95 module \\texttt{%s}}\n' % (m['name'])) + if hasnote(m): + note = m['note'] + if isinstance(note, list): + note = '\n'.join(note) + dadd(note) + if onlyvars: + dadd('\\begin{description}') + for n in onlyvars: + var = m['vars'][n] + modobjs.append(n) + ct = capi_maps.getctype(var) + at = capi_maps.c2capi_map[ct] + dm = capi_maps.getarrdims(n, var) + dms = dm['dims'].replace('*', '-1').strip() + dms = dms.replace(':', '-1').strip() + if not dms: + dms = '-1' + use_fgetdims2 = fgetdims2 + cadd('\t{"%s",%s,{{%s}},%s, %s},' % + (undo_rmbadname1(n), dm['rank'], dms, at, + capi_maps.get_elsize(var))) + dadd('\\item[]{{}\\verb@%s@{}}' % + (capi_maps.getarrdocsign(n, var))) + if hasnote(var): + note = var['note'] + if isinstance(note, list): + note = '\n'.join(note) + dadd('--- %s' % (note)) + if isallocatable(var): + fargs.append('f2py_%s_getdims_%s' % (m['name'], n)) + efargs.append(fargs[-1]) + sargs.append( + 'void (*%s)(int*,npy_intp*,void(*)(char*,npy_intp*),int*)' % (n)) + sargsp.append('void (*)(int*,npy_intp*,void(*)(char*,npy_intp*),int*)') + iadd('\tf2py_%s_def[i_f2py++].func = %s;' % (m['name'], n)) + fadd('subroutine %s(r,s,f2pysetdata,flag)' % (fargs[-1])) + fadd('use %s, only: d => %s\n' % + (m['name'], undo_rmbadname1(n))) + fadd('integer flag\n') + fhooks[0] = fhooks[0] + fgetdims1 + dms = range(1, int(dm['rank']) + 1) + fadd(' allocate(d(%s))\n' % + (','.join(['s(%s)' % i for i in dms]))) + fhooks[0] = fhooks[0] + use_fgetdims2 + fadd('end subroutine %s' % (fargs[-1])) + else: + fargs.append(n) + sargs.append('char *%s' % (n)) + sargsp.append('char*') + iadd('\tf2py_%s_def[i_f2py++].data = %s;' % (m['name'], n)) + if onlyvars: + dadd('\\end{description}') + if hasbody(m): + for b in m['body']: + if not isroutine(b): + outmess("f90mod_rules.buildhooks:" + f" skipping {b['block']} {b['name']}\n") + continue + modobjs.append('%s()' % (b['name'])) + b['modulename'] = m['name'] + api, wrap = rules.buildapi(b) + if isfunction(b): + fhooks[0] = fhooks[0] + wrap + fargs.append('f2pywrap_%s_%s' % (m['name'], b['name'])) + ifargs.append(func2subr.createfuncwrapper(b, signature=1)) + else: + if wrap: + fhooks[0] = fhooks[0] + wrap + fargs.append('f2pywrap_%s_%s' % (m['name'], b['name'])) + ifargs.append( + func2subr.createsubrwrapper(b, signature=1)) + else: + fargs.append(b['name']) + mfargs.append(fargs[-1]) + api['externroutines'] = [] + ar = applyrules(api, vrd) + ar['docs'] = [] + ar['docshort'] = [] + ret = dictappend(ret, ar) + cadd(('\t{"%s",-1,{{-1}},0,0,NULL,(void *)' + 'f2py_rout_#modulename#_%s_%s,' + 'doc_f2py_rout_#modulename#_%s_%s},') + % (b['name'], m['name'], b['name'], m['name'], b['name'])) + sargs.append('char *%s' % (b['name'])) + sargsp.append('char *') + iadd('\tf2py_%s_def[i_f2py++].data = %s;' % + (m['name'], b['name'])) + cadd('\t{NULL}\n};\n') + iadd('}') + ihooks[0] = 'static void f2py_setup_%s(%s) {\n\tint i_f2py=0;%s' % ( + m['name'], ','.join(sargs), ihooks[0]) + if '_' in m['name']: + F_FUNC = 'F_FUNC_US' + else: + F_FUNC = 'F_FUNC' + iadd('extern void %s(f2pyinit%s,F2PYINIT%s)(void (*)(%s));' + % (F_FUNC, m['name'], m['name'].upper(), ','.join(sargsp))) + iadd('static void f2py_init_%s(void) {' % (m['name'])) + iadd('\t%s(f2pyinit%s,F2PYINIT%s)(f2py_setup_%s);' + % (F_FUNC, m['name'], m['name'].upper(), m['name'])) + iadd('}\n') + ret['f90modhooks'] = ret['f90modhooks'] + chooks + ihooks + ret['initf90modhooks'] = ['\tPyDict_SetItemString(d, "%s", PyFortranObject_New(f2py_%s_def,f2py_init_%s));' % ( + m['name'], m['name'], m['name'])] + ret['initf90modhooks'] + fadd('') + fadd('subroutine f2pyinit%s(f2pysetupfunc)' % (m['name'])) + if mfargs: + for a in undo_rmbadname(mfargs): + fadd('use %s, only : %s' % (m['name'], a)) + if ifargs: + fadd(' '.join(['interface'] + ifargs)) + fadd('end interface') + fadd('external f2pysetupfunc') + if efargs: + for a in undo_rmbadname(efargs): + fadd('external %s' % (a)) + fadd('call f2pysetupfunc(%s)' % (','.join(undo_rmbadname(fargs)))) + fadd('end subroutine f2pyinit%s\n' % (m['name'])) + + dadd('\n'.join(ret['latexdoc']).replace( + r'\subsection{', r'\subsubsection{')) + + ret['latexdoc'] = [] + ret['docs'].append('"\t%s --- %s"' % (m['name'], + ','.join(undo_rmbadname(modobjs)))) + + ret['routine_defs'] = '' + ret['doc'] = [] + ret['docshort'] = [] + ret['latexdoc'] = doc[0] + if len(ret['docs']) <= 1: + ret['docs'] = '' + return ret, fhooks[0] diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/func2subr.py b/venv/lib/python3.10/site-packages/numpy/f2py/func2subr.py new file mode 100644 index 0000000000000000000000000000000000000000..b9aa9fc007cb8efdfdd13138671f0412d45d63a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/func2subr.py @@ -0,0 +1,323 @@ +""" + +Rules for building C/API module with f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +import copy + +from .auxfuncs import ( + getfortranname, isexternal, isfunction, isfunction_wrap, isintent_in, + isintent_out, islogicalfunction, ismoduleroutine, isscalar, + issubroutine, issubroutine_wrap, outmess, show +) + +from ._isocbind import isoc_kindmap + +def var2fixfortran(vars, a, fa=None, f90mode=None): + if fa is None: + fa = a + if a not in vars: + show(vars) + outmess('var2fixfortran: No definition for argument "%s".\n' % a) + return '' + if 'typespec' not in vars[a]: + show(vars[a]) + outmess('var2fixfortran: No typespec for argument "%s".\n' % a) + return '' + vardef = vars[a]['typespec'] + if vardef == 'type' and 'typename' in vars[a]: + vardef = '%s(%s)' % (vardef, vars[a]['typename']) + selector = {} + lk = '' + if 'kindselector' in vars[a]: + selector = vars[a]['kindselector'] + lk = 'kind' + elif 'charselector' in vars[a]: + selector = vars[a]['charselector'] + lk = 'len' + if '*' in selector: + if f90mode: + if selector['*'] in ['*', ':', '(*)']: + vardef = '%s(len=*)' % (vardef) + else: + vardef = '%s(%s=%s)' % (vardef, lk, selector['*']) + else: + if selector['*'] in ['*', ':']: + vardef = '%s*(%s)' % (vardef, selector['*']) + else: + vardef = '%s*%s' % (vardef, selector['*']) + else: + if 'len' in selector: + vardef = '%s(len=%s' % (vardef, selector['len']) + if 'kind' in selector: + vardef = '%s,kind=%s)' % (vardef, selector['kind']) + else: + vardef = '%s)' % (vardef) + elif 'kind' in selector: + vardef = '%s(kind=%s)' % (vardef, selector['kind']) + + vardef = '%s %s' % (vardef, fa) + if 'dimension' in vars[a]: + vardef = '%s(%s)' % (vardef, ','.join(vars[a]['dimension'])) + return vardef + +def useiso_c_binding(rout): + useisoc = False + for key, value in rout['vars'].items(): + kind_value = value.get('kindselector', {}).get('kind') + if kind_value in isoc_kindmap: + return True + return useisoc + +def createfuncwrapper(rout, signature=0): + assert isfunction(rout) + + extra_args = [] + vars = rout['vars'] + for a in rout['args']: + v = rout['vars'][a] + for i, d in enumerate(v.get('dimension', [])): + if d == ':': + dn = 'f2py_%s_d%s' % (a, i) + dv = dict(typespec='integer', intent=['hide']) + dv['='] = 'shape(%s, %s)' % (a, i) + extra_args.append(dn) + vars[dn] = dv + v['dimension'][i] = dn + rout['args'].extend(extra_args) + need_interface = bool(extra_args) + + ret = [''] + + def add(line, ret=ret): + ret[0] = '%s\n %s' % (ret[0], line) + name = rout['name'] + fortranname = getfortranname(rout) + f90mode = ismoduleroutine(rout) + newname = '%sf2pywrap' % (name) + + if newname not in vars: + vars[newname] = vars[name] + args = [newname] + rout['args'][1:] + else: + args = [newname] + rout['args'] + + l_tmpl = var2fixfortran(vars, name, '@@@NAME@@@', f90mode) + if l_tmpl[:13] == 'character*(*)': + if f90mode: + l_tmpl = 'character(len=10)' + l_tmpl[13:] + else: + l_tmpl = 'character*10' + l_tmpl[13:] + charselect = vars[name]['charselector'] + if charselect.get('*', '') == '(*)': + charselect['*'] = '10' + + l1 = l_tmpl.replace('@@@NAME@@@', newname) + rl = None + + useisoc = useiso_c_binding(rout) + sargs = ', '.join(args) + if f90mode: + # gh-23598 fix warning + # Essentially, this gets called again with modules where the name of the + # function is added to the arguments, which is not required, and removed + sargs = sargs.replace(f"{name}, ", '') + args = [arg for arg in args if arg != name] + rout['args'] = args + add('subroutine f2pywrap_%s_%s (%s)' % + (rout['modulename'], name, sargs)) + if not signature: + add('use %s, only : %s' % (rout['modulename'], fortranname)) + if useisoc: + add('use iso_c_binding') + else: + add('subroutine f2pywrap%s (%s)' % (name, sargs)) + if useisoc: + add('use iso_c_binding') + if not need_interface: + add('external %s' % (fortranname)) + rl = l_tmpl.replace('@@@NAME@@@', '') + ' ' + fortranname + + if need_interface: + for line in rout['saved_interface'].split('\n'): + if line.lstrip().startswith('use ') and '__user__' not in line: + add(line) + + args = args[1:] + dumped_args = [] + for a in args: + if isexternal(vars[a]): + add('external %s' % (a)) + dumped_args.append(a) + for a in args: + if a in dumped_args: + continue + if isscalar(vars[a]): + add(var2fixfortran(vars, a, f90mode=f90mode)) + dumped_args.append(a) + for a in args: + if a in dumped_args: + continue + if isintent_in(vars[a]): + add(var2fixfortran(vars, a, f90mode=f90mode)) + dumped_args.append(a) + for a in args: + if a in dumped_args: + continue + add(var2fixfortran(vars, a, f90mode=f90mode)) + + add(l1) + if rl is not None: + add(rl) + + if need_interface: + if f90mode: + # f90 module already defines needed interface + pass + else: + add('interface') + add(rout['saved_interface'].lstrip()) + add('end interface') + + sargs = ', '.join([a for a in args if a not in extra_args]) + + if not signature: + if islogicalfunction(rout): + add('%s = .not.(.not.%s(%s))' % (newname, fortranname, sargs)) + else: + add('%s = %s(%s)' % (newname, fortranname, sargs)) + if f90mode: + add('end subroutine f2pywrap_%s_%s' % (rout['modulename'], name)) + else: + add('end') + return ret[0] + + +def createsubrwrapper(rout, signature=0): + assert issubroutine(rout) + + extra_args = [] + vars = rout['vars'] + for a in rout['args']: + v = rout['vars'][a] + for i, d in enumerate(v.get('dimension', [])): + if d == ':': + dn = 'f2py_%s_d%s' % (a, i) + dv = dict(typespec='integer', intent=['hide']) + dv['='] = 'shape(%s, %s)' % (a, i) + extra_args.append(dn) + vars[dn] = dv + v['dimension'][i] = dn + rout['args'].extend(extra_args) + need_interface = bool(extra_args) + + ret = [''] + + def add(line, ret=ret): + ret[0] = '%s\n %s' % (ret[0], line) + name = rout['name'] + fortranname = getfortranname(rout) + f90mode = ismoduleroutine(rout) + + args = rout['args'] + + useisoc = useiso_c_binding(rout) + sargs = ', '.join(args) + if f90mode: + add('subroutine f2pywrap_%s_%s (%s)' % + (rout['modulename'], name, sargs)) + if useisoc: + add('use iso_c_binding') + if not signature: + add('use %s, only : %s' % (rout['modulename'], fortranname)) + else: + add('subroutine f2pywrap%s (%s)' % (name, sargs)) + if useisoc: + add('use iso_c_binding') + if not need_interface: + add('external %s' % (fortranname)) + + if need_interface: + for line in rout['saved_interface'].split('\n'): + if line.lstrip().startswith('use ') and '__user__' not in line: + add(line) + + dumped_args = [] + for a in args: + if isexternal(vars[a]): + add('external %s' % (a)) + dumped_args.append(a) + for a in args: + if a in dumped_args: + continue + if isscalar(vars[a]): + add(var2fixfortran(vars, a, f90mode=f90mode)) + dumped_args.append(a) + for a in args: + if a in dumped_args: + continue + add(var2fixfortran(vars, a, f90mode=f90mode)) + + if need_interface: + if f90mode: + # f90 module already defines needed interface + pass + else: + add('interface') + for line in rout['saved_interface'].split('\n'): + if line.lstrip().startswith('use ') and '__user__' in line: + continue + add(line) + add('end interface') + + sargs = ', '.join([a for a in args if a not in extra_args]) + + if not signature: + add('call %s(%s)' % (fortranname, sargs)) + if f90mode: + add('end subroutine f2pywrap_%s_%s' % (rout['modulename'], name)) + else: + add('end') + return ret[0] + + +def assubr(rout): + if isfunction_wrap(rout): + fortranname = getfortranname(rout) + name = rout['name'] + outmess('\t\tCreating wrapper for Fortran function "%s"("%s")...\n' % ( + name, fortranname)) + rout = copy.copy(rout) + fname = name + rname = fname + if 'result' in rout: + rname = rout['result'] + rout['vars'][fname] = rout['vars'][rname] + fvar = rout['vars'][fname] + if not isintent_out(fvar): + if 'intent' not in fvar: + fvar['intent'] = [] + fvar['intent'].append('out') + flag = 1 + for i in fvar['intent']: + if i.startswith('out='): + flag = 0 + break + if flag: + fvar['intent'].append('out=%s' % (rname)) + rout['args'][:] = [fname] + rout['args'] + return rout, createfuncwrapper(rout) + if issubroutine_wrap(rout): + fortranname = getfortranname(rout) + name = rout['name'] + outmess('\t\tCreating wrapper for Fortran subroutine "%s"("%s")...\n' + % (name, fortranname)) + rout = copy.copy(rout) + return rout, createsubrwrapper(rout) + return rout, '' diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/rules.py b/venv/lib/python3.10/site-packages/numpy/f2py/rules.py new file mode 100644 index 0000000000000000000000000000000000000000..84137811a4462b03e8e4e72d09564fbbc4086ecb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/rules.py @@ -0,0 +1,1578 @@ +""" + +Rules for building C/API module with f2py2e. + +Here is a skeleton of a new wrapper function (13Dec2001): + +wrapper_function(args) + declarations + get_python_arguments, say, `a' and `b' + + get_a_from_python + if (successful) { + + get_b_from_python + if (successful) { + + callfortran + if (successful) { + + put_a_to_python + if (successful) { + + put_b_to_python + if (successful) { + + buildvalue = ... + + } + + } + + } + + } + cleanup_b + + } + cleanup_a + + return buildvalue + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +import os +import sys +import time +import copy +from pathlib import Path + +# __version__.version is now the same as the NumPy version +from . import __version__ + +from .auxfuncs import ( + applyrules, debugcapi, dictappend, errmess, gentitle, getargs2, + hascallstatement, hasexternals, hasinitvalue, hasnote, + hasresultnote, isarray, isarrayofstrings, ischaracter, + ischaracterarray, ischaracter_or_characterarray, iscomplex, + iscomplexarray, iscomplexfunction, iscomplexfunction_warn, + isdummyroutine, isexternal, isfunction, isfunction_wrap, isint1, + isint1array, isintent_aux, isintent_c, isintent_callback, + isintent_copy, isintent_hide, isintent_inout, isintent_nothide, + isintent_out, isintent_overwrite, islogical, islong_complex, + islong_double, islong_doublefunction, islong_long, + islong_longfunction, ismoduleroutine, isoptional, isrequired, + isscalar, issigned_long_longarray, isstring, isstringarray, + isstringfunction, issubroutine, isattr_value, + issubroutine_wrap, isthreadsafe, isunsigned, isunsigned_char, + isunsigned_chararray, isunsigned_long_long, + isunsigned_long_longarray, isunsigned_short, isunsigned_shortarray, + l_and, l_not, l_or, outmess, replace, stripcomma, requiresf90wrapper +) + +from . import capi_maps +from . import cfuncs +from . import common_rules +from . import use_rules +from . import f90mod_rules +from . import func2subr + +f2py_version = __version__.version +numpy_version = __version__.version + +options = {} +sepdict = {} +# for k in ['need_cfuncs']: sepdict[k]=',' +for k in ['decl', + 'frompyobj', + 'cleanupfrompyobj', + 'topyarr', 'method', + 'pyobjfrom', 'closepyobjfrom', + 'freemem', + 'userincludes', + 'includes0', 'includes', 'typedefs', 'typedefs_generated', + 'cppmacros', 'cfuncs', 'callbacks', + 'latexdoc', + 'restdoc', + 'routine_defs', 'externroutines', + 'initf2pywraphooks', + 'commonhooks', 'initcommonhooks', + 'f90modhooks', 'initf90modhooks']: + sepdict[k] = '\n' + +#################### Rules for C/API module ################# + +generationtime = int(os.environ.get('SOURCE_DATE_EPOCH', time.time())) +module_rules = { + 'modulebody': """\ +/* File: #modulename#module.c + * This file is auto-generated with f2py (version:#f2py_version#). + * f2py is a Fortran to Python Interface Generator (FPIG), Second Edition, + * written by Pearu Peterson . + * Generation date: """ + time.asctime(time.gmtime(generationtime)) + """ + * Do not edit this file directly unless you know what you are doing!!! + */ + +#ifdef __cplusplus +extern \"C\" { +#endif + +#ifndef PY_SSIZE_T_CLEAN +#define PY_SSIZE_T_CLEAN +#endif /* PY_SSIZE_T_CLEAN */ + +/* Unconditionally included */ +#include +#include + +""" + gentitle("See f2py2e/cfuncs.py: includes") + """ +#includes# +#includes0# + +""" + gentitle("See f2py2e/rules.py: mod_rules['modulebody']") + """ +static PyObject *#modulename#_error; +static PyObject *#modulename#_module; + +""" + gentitle("See f2py2e/cfuncs.py: typedefs") + """ +#typedefs# + +""" + gentitle("See f2py2e/cfuncs.py: typedefs_generated") + """ +#typedefs_generated# + +""" + gentitle("See f2py2e/cfuncs.py: cppmacros") + """ +#cppmacros# + +""" + gentitle("See f2py2e/cfuncs.py: cfuncs") + """ +#cfuncs# + +""" + gentitle("See f2py2e/cfuncs.py: userincludes") + """ +#userincludes# + +""" + gentitle("See f2py2e/capi_rules.py: usercode") + """ +#usercode# + +/* See f2py2e/rules.py */ +#externroutines# + +""" + gentitle("See f2py2e/capi_rules.py: usercode1") + """ +#usercode1# + +""" + gentitle("See f2py2e/cb_rules.py: buildcallback") + """ +#callbacks# + +""" + gentitle("See f2py2e/rules.py: buildapi") + """ +#body# + +""" + gentitle("See f2py2e/f90mod_rules.py: buildhooks") + """ +#f90modhooks# + +""" + gentitle("See f2py2e/rules.py: module_rules['modulebody']") + """ + +""" + gentitle("See f2py2e/common_rules.py: buildhooks") + """ +#commonhooks# + +""" + gentitle("See f2py2e/rules.py") + """ + +static FortranDataDef f2py_routine_defs[] = { +#routine_defs# + {NULL} +}; + +static PyMethodDef f2py_module_methods[] = { +#pymethoddef# + {NULL,NULL} +}; + +static struct PyModuleDef moduledef = { + PyModuleDef_HEAD_INIT, + "#modulename#", + NULL, + -1, + f2py_module_methods, + NULL, + NULL, + NULL, + NULL +}; + +PyMODINIT_FUNC PyInit_#modulename#(void) { + int i; + PyObject *m,*d, *s, *tmp; + m = #modulename#_module = PyModule_Create(&moduledef); + Py_SET_TYPE(&PyFortran_Type, &PyType_Type); + import_array(); + if (PyErr_Occurred()) + {PyErr_SetString(PyExc_ImportError, \"can't initialize module #modulename# (failed to import numpy)\"); return m;} + d = PyModule_GetDict(m); + s = PyUnicode_FromString(\"#f2py_version#\"); + PyDict_SetItemString(d, \"__version__\", s); + Py_DECREF(s); + s = PyUnicode_FromString( + \"This module '#modulename#' is auto-generated with f2py (version:#f2py_version#).\\nFunctions:\\n\"\n#docs#\".\"); + PyDict_SetItemString(d, \"__doc__\", s); + Py_DECREF(s); + s = PyUnicode_FromString(\"""" + numpy_version + """\"); + PyDict_SetItemString(d, \"__f2py_numpy_version__\", s); + Py_DECREF(s); + #modulename#_error = PyErr_NewException (\"#modulename#.error\", NULL, NULL); + /* + * Store the error object inside the dict, so that it could get deallocated. + * (in practice, this is a module, so it likely will not and cannot.) + */ + PyDict_SetItemString(d, \"_#modulename#_error\", #modulename#_error); + Py_DECREF(#modulename#_error); + for(i=0;f2py_routine_defs[i].name!=NULL;i++) { + tmp = PyFortranObject_NewAsAttr(&f2py_routine_defs[i]); + PyDict_SetItemString(d, f2py_routine_defs[i].name, tmp); + Py_DECREF(tmp); + } +#initf2pywraphooks# +#initf90modhooks# +#initcommonhooks# +#interface_usercode# + +#if Py_GIL_DISABLED + // signal whether this module supports running with the GIL disabled + PyUnstable_Module_SetGIL(m , #gil_used#); +#endif + +#ifdef F2PY_REPORT_ATEXIT + if (! PyErr_Occurred()) + on_exit(f2py_report_on_exit,(void*)\"#modulename#\"); +#endif + + if (PyType_Ready(&PyFortran_Type) < 0) { + return NULL; + } + + return m; +} +#ifdef __cplusplus +} +#endif +""", + 'separatorsfor': {'latexdoc': '\n\n', + 'restdoc': '\n\n'}, + 'latexdoc': ['\\section{Module \\texttt{#texmodulename#}}\n', + '#modnote#\n', + '#latexdoc#'], + 'restdoc': ['Module #modulename#\n' + '=' * 80, + '\n#restdoc#'] +} + +defmod_rules = [ + {'body': '/*eof body*/', + 'method': '/*eof method*/', + 'externroutines': '/*eof externroutines*/', + 'routine_defs': '/*eof routine_defs*/', + 'initf90modhooks': '/*eof initf90modhooks*/', + 'initf2pywraphooks': '/*eof initf2pywraphooks*/', + 'initcommonhooks': '/*eof initcommonhooks*/', + 'latexdoc': '', + 'restdoc': '', + 'modnote': {hasnote: '#note#', l_not(hasnote): ''}, + } +] + +routine_rules = { + 'separatorsfor': sepdict, + 'body': """ +#begintitle# +static char doc_#apiname#[] = \"\\\n#docreturn##name#(#docsignatureshort#)\\n\\nWrapper for ``#name#``.\\\n\\n#docstrsigns#\"; +/* #declfortranroutine# */ +static PyObject *#apiname#(const PyObject *capi_self, + PyObject *capi_args, + PyObject *capi_keywds, + #functype# (*f2py_func)(#callprotoargument#)) { + PyObject * volatile capi_buildvalue = NULL; + volatile int f2py_success = 1; +#decl# + static char *capi_kwlist[] = {#kwlist##kwlistopt##kwlistxa#NULL}; +#usercode# +#routdebugenter# +#ifdef F2PY_REPORT_ATEXIT +f2py_start_clock(); +#endif + if (!PyArg_ParseTupleAndKeywords(capi_args,capi_keywds,\\ + \"#argformat#|#keyformat##xaformat#:#pyname#\",\\ + capi_kwlist#args_capi##keys_capi##keys_xa#))\n return NULL; +#frompyobj# +/*end of frompyobj*/ +#ifdef F2PY_REPORT_ATEXIT +f2py_start_call_clock(); +#endif +#callfortranroutine# +if (PyErr_Occurred()) + f2py_success = 0; +#ifdef F2PY_REPORT_ATEXIT +f2py_stop_call_clock(); +#endif +/*end of callfortranroutine*/ + if (f2py_success) { +#pyobjfrom# +/*end of pyobjfrom*/ + CFUNCSMESS(\"Building return value.\\n\"); + capi_buildvalue = Py_BuildValue(\"#returnformat#\"#return#); +/*closepyobjfrom*/ +#closepyobjfrom# + } /*if (f2py_success) after callfortranroutine*/ +/*cleanupfrompyobj*/ +#cleanupfrompyobj# + if (capi_buildvalue == NULL) { +#routdebugfailure# + } else { +#routdebugleave# + } + CFUNCSMESS(\"Freeing memory.\\n\"); +#freemem# +#ifdef F2PY_REPORT_ATEXIT +f2py_stop_clock(); +#endif + return capi_buildvalue; +} +#endtitle# +""", + 'routine_defs': '#routine_def#', + 'initf2pywraphooks': '#initf2pywraphook#', + 'externroutines': '#declfortranroutine#', + 'doc': '#docreturn##name#(#docsignature#)', + 'docshort': '#docreturn##name#(#docsignatureshort#)', + 'docs': '" #docreturn##name#(#docsignature#)\\n"\n', + 'need': ['arrayobject.h', 'CFUNCSMESS', 'MINMAX'], + 'cppmacros': {debugcapi: '#define DEBUGCFUNCS'}, + 'latexdoc': ['\\subsection{Wrapper function \\texttt{#texname#}}\n', + """ +\\noindent{{}\\verb@#docreturn##name#@{}}\\texttt{(#latexdocsignatureshort#)} +#routnote# + +#latexdocstrsigns# +"""], + 'restdoc': ['Wrapped function ``#name#``\n' + '-' * 80, + + ] +} + +################## Rules for C/API function ############## + +rout_rules = [ + { # Init + 'separatorsfor': {'callfortranroutine': '\n', 'routdebugenter': '\n', 'decl': '\n', + 'routdebugleave': '\n', 'routdebugfailure': '\n', + 'setjmpbuf': ' || ', + 'docstrreq': '\n', 'docstropt': '\n', 'docstrout': '\n', + 'docstrcbs': '\n', 'docstrsigns': '\\n"\n"', + 'latexdocstrsigns': '\n', + 'latexdocstrreq': '\n', 'latexdocstropt': '\n', + 'latexdocstrout': '\n', 'latexdocstrcbs': '\n', + }, + 'kwlist': '', 'kwlistopt': '', 'callfortran': '', 'callfortranappend': '', + 'docsign': '', 'docsignopt': '', 'decl': '/*decl*/', + 'freemem': '/*freemem*/', + 'docsignshort': '', 'docsignoptshort': '', + 'docstrsigns': '', 'latexdocstrsigns': '', + 'docstrreq': '\\nParameters\\n----------', + 'docstropt': '\\nOther Parameters\\n----------------', + 'docstrout': '\\nReturns\\n-------', + 'docstrcbs': '\\nNotes\\n-----\\nCall-back functions::\\n', + 'latexdocstrreq': '\\noindent Required arguments:', + 'latexdocstropt': '\\noindent Optional arguments:', + 'latexdocstrout': '\\noindent Return objects:', + 'latexdocstrcbs': '\\noindent Call-back functions:', + 'args_capi': '', 'keys_capi': '', 'functype': '', + 'frompyobj': '/*frompyobj*/', + # this list will be reversed + 'cleanupfrompyobj': ['/*end of cleanupfrompyobj*/'], + 'pyobjfrom': '/*pyobjfrom*/', + # this list will be reversed + 'closepyobjfrom': ['/*end of closepyobjfrom*/'], + 'topyarr': '/*topyarr*/', 'routdebugleave': '/*routdebugleave*/', + 'routdebugenter': '/*routdebugenter*/', + 'routdebugfailure': '/*routdebugfailure*/', + 'callfortranroutine': '/*callfortranroutine*/', + 'argformat': '', 'keyformat': '', 'need_cfuncs': '', + 'docreturn': '', 'return': '', 'returnformat': '', 'rformat': '', + 'kwlistxa': '', 'keys_xa': '', 'xaformat': '', 'docsignxa': '', 'docsignxashort': '', + 'initf2pywraphook': '', + 'routnote': {hasnote: '--- #note#', l_not(hasnote): ''}, + }, { + 'apiname': 'f2py_rout_#modulename#_#name#', + 'pyname': '#modulename#.#name#', + 'decl': '', + '_check': l_not(ismoduleroutine) + }, { + 'apiname': 'f2py_rout_#modulename#_#f90modulename#_#name#', + 'pyname': '#modulename#.#f90modulename#.#name#', + 'decl': '', + '_check': ismoduleroutine + }, { # Subroutine + 'functype': 'void', + 'declfortranroutine': {l_and(l_not(l_or(ismoduleroutine, isintent_c)), l_not(isdummyroutine)): 'extern void #F_FUNC#(#fortranname#,#FORTRANNAME#)(#callprotoargument#);', + l_and(l_not(ismoduleroutine), isintent_c, l_not(isdummyroutine)): 'extern void #fortranname#(#callprotoargument#);', + ismoduleroutine: '', + isdummyroutine: '' + }, + 'routine_def': { + l_not(l_or(ismoduleroutine, isintent_c, isdummyroutine)): + ' {\"#name#\",-1,{{-1}},0,0,(char *)' + ' #F_FUNC#(#fortranname#,#FORTRANNAME#),' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + l_and(l_not(ismoduleroutine), isintent_c, l_not(isdummyroutine)): + ' {\"#name#\",-1,{{-1}},0,0,(char *)#fortranname#,' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + l_and(l_not(ismoduleroutine), isdummyroutine): + ' {\"#name#\",-1,{{-1}},0,0,NULL,' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + }, + 'need': {l_and(l_not(l_or(ismoduleroutine, isintent_c)), l_not(isdummyroutine)): 'F_FUNC'}, + 'callfortranroutine': [ + {debugcapi: [ + """ fprintf(stderr,\"debug-capi:Fortran subroutine `#fortranname#(#callfortran#)\'\\n\");"""]}, + {hasexternals: """\ + if (#setjmpbuf#) { + f2py_success = 0; + } else {"""}, + {isthreadsafe: ' Py_BEGIN_ALLOW_THREADS'}, + {hascallstatement: ''' #callstatement#; + /*(*f2py_func)(#callfortran#);*/'''}, + {l_not(l_or(hascallstatement, isdummyroutine)) + : ' (*f2py_func)(#callfortran#);'}, + {isthreadsafe: ' Py_END_ALLOW_THREADS'}, + {hasexternals: """ }"""} + ], + '_check': l_and(issubroutine, l_not(issubroutine_wrap)), + }, { # Wrapped function + 'functype': 'void', + 'declfortranroutine': {l_not(l_or(ismoduleroutine, isdummyroutine)): 'extern void #F_WRAPPEDFUNC#(#name_lower#,#NAME#)(#callprotoargument#);', + isdummyroutine: '', + }, + + 'routine_def': { + l_not(l_or(ismoduleroutine, isdummyroutine)): + ' {\"#name#\",-1,{{-1}},0,0,(char *)' + ' #F_WRAPPEDFUNC#(#name_lower#,#NAME#),' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + isdummyroutine: + ' {\"#name#\",-1,{{-1}},0,0,NULL,' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + }, + 'initf2pywraphook': {l_not(l_or(ismoduleroutine, isdummyroutine)): ''' + { + extern #ctype# #F_FUNC#(#name_lower#,#NAME#)(void); + PyObject* o = PyDict_GetItemString(d,"#name#"); + tmp = F2PyCapsule_FromVoidPtr((void*)#F_WRAPPEDFUNC#(#name_lower#,#NAME#),NULL); + PyObject_SetAttrString(o,"_cpointer", tmp); + Py_DECREF(tmp); + s = PyUnicode_FromString("#name#"); + PyObject_SetAttrString(o,"__name__", s); + Py_DECREF(s); + } + '''}, + 'need': {l_not(l_or(ismoduleroutine, isdummyroutine)): ['F_WRAPPEDFUNC', 'F_FUNC']}, + 'callfortranroutine': [ + {debugcapi: [ + """ fprintf(stderr,\"debug-capi:Fortran subroutine `f2pywrap#name_lower#(#callfortran#)\'\\n\");"""]}, + {hasexternals: """\ + if (#setjmpbuf#) { + f2py_success = 0; + } else {"""}, + {isthreadsafe: ' Py_BEGIN_ALLOW_THREADS'}, + {l_not(l_or(hascallstatement, isdummyroutine)) + : ' (*f2py_func)(#callfortran#);'}, + {hascallstatement: + ' #callstatement#;\n /*(*f2py_func)(#callfortran#);*/'}, + {isthreadsafe: ' Py_END_ALLOW_THREADS'}, + {hasexternals: ' }'} + ], + '_check': isfunction_wrap, + }, { # Wrapped subroutine + 'functype': 'void', + 'declfortranroutine': {l_not(l_or(ismoduleroutine, isdummyroutine)): 'extern void #F_WRAPPEDFUNC#(#name_lower#,#NAME#)(#callprotoargument#);', + isdummyroutine: '', + }, + + 'routine_def': { + l_not(l_or(ismoduleroutine, isdummyroutine)): + ' {\"#name#\",-1,{{-1}},0,0,(char *)' + ' #F_WRAPPEDFUNC#(#name_lower#,#NAME#),' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + isdummyroutine: + ' {\"#name#\",-1,{{-1}},0,0,NULL,' + ' (f2py_init_func)#apiname#,doc_#apiname#},', + }, + 'initf2pywraphook': {l_not(l_or(ismoduleroutine, isdummyroutine)): ''' + { + extern void #F_FUNC#(#name_lower#,#NAME#)(void); + PyObject* o = PyDict_GetItemString(d,"#name#"); + tmp = F2PyCapsule_FromVoidPtr((void*)#F_FUNC#(#name_lower#,#NAME#),NULL); + PyObject_SetAttrString(o,"_cpointer", tmp); + Py_DECREF(tmp); + s = PyUnicode_FromString("#name#"); + PyObject_SetAttrString(o,"__name__", s); + Py_DECREF(s); + } + '''}, + 'need': {l_not(l_or(ismoduleroutine, isdummyroutine)): ['F_WRAPPEDFUNC', 'F_FUNC']}, + 'callfortranroutine': [ + {debugcapi: [ + """ fprintf(stderr,\"debug-capi:Fortran subroutine `f2pywrap#name_lower#(#callfortran#)\'\\n\");"""]}, + {hasexternals: """\ + if (#setjmpbuf#) { + f2py_success = 0; + } else {"""}, + {isthreadsafe: ' Py_BEGIN_ALLOW_THREADS'}, + {l_not(l_or(hascallstatement, isdummyroutine)) + : ' (*f2py_func)(#callfortran#);'}, + {hascallstatement: + ' #callstatement#;\n /*(*f2py_func)(#callfortran#);*/'}, + {isthreadsafe: ' Py_END_ALLOW_THREADS'}, + {hasexternals: ' }'} + ], + '_check': issubroutine_wrap, + }, { # Function + 'functype': '#ctype#', + 'docreturn': {l_not(isintent_hide): '#rname#,'}, + 'docstrout': '#pydocsignout#', + 'latexdocstrout': ['\\item[]{{}\\verb@#pydocsignout#@{}}', + {hasresultnote: '--- #resultnote#'}], + 'callfortranroutine': [{l_and(debugcapi, isstringfunction): """\ +#ifdef USESCOMPAQFORTRAN + fprintf(stderr,\"debug-capi:Fortran function #ctype# #fortranname#(#callcompaqfortran#)\\n\"); +#else + fprintf(stderr,\"debug-capi:Fortran function #ctype# #fortranname#(#callfortran#)\\n\"); +#endif +"""}, + {l_and(debugcapi, l_not(isstringfunction)): """\ + fprintf(stderr,\"debug-capi:Fortran function #ctype# #fortranname#(#callfortran#)\\n\"); +"""} + ], + '_check': l_and(isfunction, l_not(isfunction_wrap)) + }, { # Scalar function + 'declfortranroutine': {l_and(l_not(l_or(ismoduleroutine, isintent_c)), l_not(isdummyroutine)): 'extern #ctype# #F_FUNC#(#fortranname#,#FORTRANNAME#)(#callprotoargument#);', + l_and(l_not(ismoduleroutine), isintent_c, l_not(isdummyroutine)): 'extern #ctype# #fortranname#(#callprotoargument#);', + isdummyroutine: '' + }, + 'routine_def': { + l_and(l_not(l_or(ismoduleroutine, isintent_c)), + l_not(isdummyroutine)): + (' {\"#name#\",-1,{{-1}},0,0,(char *)' + ' #F_FUNC#(#fortranname#,#FORTRANNAME#),' + ' (f2py_init_func)#apiname#,doc_#apiname#},'), + l_and(l_not(ismoduleroutine), isintent_c, l_not(isdummyroutine)): + (' {\"#name#\",-1,{{-1}},0,0,(char *)#fortranname#,' + ' (f2py_init_func)#apiname#,doc_#apiname#},'), + isdummyroutine: + ' {\"#name#\",-1,{{-1}},0,0,NULL,' + '(f2py_init_func)#apiname#,doc_#apiname#},', + }, + 'decl': [{iscomplexfunction_warn: ' #ctype# #name#_return_value={0,0};', + l_not(iscomplexfunction): ' #ctype# #name#_return_value=0;'}, + {iscomplexfunction: + ' PyObject *#name#_return_value_capi = Py_None;'} + ], + 'callfortranroutine': [ + {hasexternals: """\ + if (#setjmpbuf#) { + f2py_success = 0; + } else {"""}, + {isthreadsafe: ' Py_BEGIN_ALLOW_THREADS'}, + {hascallstatement: ''' #callstatement#; +/* #name#_return_value = (*f2py_func)(#callfortran#);*/ +'''}, + {l_not(l_or(hascallstatement, isdummyroutine)) + : ' #name#_return_value = (*f2py_func)(#callfortran#);'}, + {isthreadsafe: ' Py_END_ALLOW_THREADS'}, + {hasexternals: ' }'}, + {l_and(debugcapi, iscomplexfunction) + : ' fprintf(stderr,"#routdebugshowvalue#\\n",#name#_return_value.r,#name#_return_value.i);'}, + {l_and(debugcapi, l_not(iscomplexfunction)): ' fprintf(stderr,"#routdebugshowvalue#\\n",#name#_return_value);'}], + 'pyobjfrom': {iscomplexfunction: ' #name#_return_value_capi = pyobj_from_#ctype#1(#name#_return_value);'}, + 'need': [{l_not(isdummyroutine): 'F_FUNC'}, + {iscomplexfunction: 'pyobj_from_#ctype#1'}, + {islong_longfunction: 'long_long'}, + {islong_doublefunction: 'long_double'}], + 'returnformat': {l_not(isintent_hide): '#rformat#'}, + 'return': {iscomplexfunction: ',#name#_return_value_capi', + l_not(l_or(iscomplexfunction, isintent_hide)): ',#name#_return_value'}, + '_check': l_and(isfunction, l_not(isstringfunction), l_not(isfunction_wrap)) + }, { # String function # in use for --no-wrap + 'declfortranroutine': 'extern void #F_FUNC#(#fortranname#,#FORTRANNAME#)(#callprotoargument#);', + 'routine_def': {l_not(l_or(ismoduleroutine, isintent_c)): + ' {\"#name#\",-1,{{-1}},0,0,(char *)#F_FUNC#(#fortranname#,#FORTRANNAME#),(f2py_init_func)#apiname#,doc_#apiname#},', + l_and(l_not(ismoduleroutine), isintent_c): + ' {\"#name#\",-1,{{-1}},0,0,(char *)#fortranname#,(f2py_init_func)#apiname#,doc_#apiname#},' + }, + 'decl': [' #ctype# #name#_return_value = NULL;', + ' int #name#_return_value_len = 0;'], + 'callfortran':'#name#_return_value,#name#_return_value_len,', + 'callfortranroutine':[' #name#_return_value_len = #rlength#;', + ' if ((#name#_return_value = (string)malloc(' + + '#name#_return_value_len+1) == NULL) {', + ' PyErr_SetString(PyExc_MemoryError, \"out of memory\");', + ' f2py_success = 0;', + ' } else {', + " (#name#_return_value)[#name#_return_value_len] = '\\0';", + ' }', + ' if (f2py_success) {', + {hasexternals: """\ + if (#setjmpbuf#) { + f2py_success = 0; + } else {"""}, + {isthreadsafe: ' Py_BEGIN_ALLOW_THREADS'}, + """\ +#ifdef USESCOMPAQFORTRAN + (*f2py_func)(#callcompaqfortran#); +#else + (*f2py_func)(#callfortran#); +#endif +""", + {isthreadsafe: ' Py_END_ALLOW_THREADS'}, + {hasexternals: ' }'}, + {debugcapi: + ' fprintf(stderr,"#routdebugshowvalue#\\n",#name#_return_value_len,#name#_return_value);'}, + ' } /* if (f2py_success) after (string)malloc */', + ], + 'returnformat': '#rformat#', + 'return': ',#name#_return_value', + 'freemem': ' STRINGFREE(#name#_return_value);', + 'need': ['F_FUNC', '#ctype#', 'STRINGFREE'], + '_check':l_and(isstringfunction, l_not(isfunction_wrap)) # ???obsolete + }, + { # Debugging + 'routdebugenter': ' fprintf(stderr,"debug-capi:Python C/API function #modulename#.#name#(#docsignature#)\\n");', + 'routdebugleave': ' fprintf(stderr,"debug-capi:Python C/API function #modulename#.#name#: successful.\\n");', + 'routdebugfailure': ' fprintf(stderr,"debug-capi:Python C/API function #modulename#.#name#: failure.\\n");', + '_check': debugcapi + } +] + +################ Rules for arguments ################## + +typedef_need_dict = {islong_long: 'long_long', + islong_double: 'long_double', + islong_complex: 'complex_long_double', + isunsigned_char: 'unsigned_char', + isunsigned_short: 'unsigned_short', + isunsigned: 'unsigned', + isunsigned_long_long: 'unsigned_long_long', + isunsigned_chararray: 'unsigned_char', + isunsigned_shortarray: 'unsigned_short', + isunsigned_long_longarray: 'unsigned_long_long', + issigned_long_longarray: 'long_long', + isint1: 'signed_char', + ischaracter_or_characterarray: 'character', + } + +aux_rules = [ + { + 'separatorsfor': sepdict + }, + { # Common + 'frompyobj': [' /* Processing auxiliary variable #varname# */', + {debugcapi: ' fprintf(stderr,"#vardebuginfo#\\n");'}, ], + 'cleanupfrompyobj': ' /* End of cleaning variable #varname# */', + 'need': typedef_need_dict, + }, + # Scalars (not complex) + { # Common + 'decl': ' #ctype# #varname# = 0;', + 'need': {hasinitvalue: 'math.h'}, + 'frompyobj': {hasinitvalue: ' #varname# = #init#;'}, + '_check': l_and(isscalar, l_not(iscomplex)), + }, + { + 'return': ',#varname#', + 'docstrout': '#pydocsignout#', + 'docreturn': '#outvarname#,', + 'returnformat': '#varrformat#', + '_check': l_and(isscalar, l_not(iscomplex), isintent_out), + }, + # Complex scalars + { # Common + 'decl': ' #ctype# #varname#;', + 'frompyobj': {hasinitvalue: ' #varname#.r = #init.r#, #varname#.i = #init.i#;'}, + '_check': iscomplex + }, + # String + { # Common + 'decl': [' #ctype# #varname# = NULL;', + ' int slen(#varname#);', + ], + 'need':['len..'], + '_check':isstring + }, + # Array + { # Common + 'decl': [' #ctype# *#varname# = NULL;', + ' npy_intp #varname#_Dims[#rank#] = {#rank*[-1]#};', + ' const int #varname#_Rank = #rank#;', + ], + 'need':['len..', {hasinitvalue: 'forcomb'}, {hasinitvalue: 'CFUNCSMESS'}], + '_check': isarray + }, + # Scalararray + { # Common + '_check': l_and(isarray, l_not(iscomplexarray)) + }, { # Not hidden + '_check': l_and(isarray, l_not(iscomplexarray), isintent_nothide) + }, + # Integer*1 array + {'need': '#ctype#', + '_check': isint1array, + '_depend': '' + }, + # Integer*-1 array + {'need': '#ctype#', + '_check': l_or(isunsigned_chararray, isunsigned_char), + '_depend': '' + }, + # Integer*-2 array + {'need': '#ctype#', + '_check': isunsigned_shortarray, + '_depend': '' + }, + # Integer*-8 array + {'need': '#ctype#', + '_check': isunsigned_long_longarray, + '_depend': '' + }, + # Complexarray + {'need': '#ctype#', + '_check': iscomplexarray, + '_depend': '' + }, + # Stringarray + { + 'callfortranappend': {isarrayofstrings: 'flen(#varname#),'}, + 'need': 'string', + '_check': isstringarray + } +] + +arg_rules = [ + { + 'separatorsfor': sepdict + }, + { # Common + 'frompyobj': [' /* Processing variable #varname# */', + {debugcapi: ' fprintf(stderr,"#vardebuginfo#\\n");'}, ], + 'cleanupfrompyobj': ' /* End of cleaning variable #varname# */', + '_depend': '', + 'need': typedef_need_dict, + }, + # Doc signatures + { + 'docstropt': {l_and(isoptional, isintent_nothide): '#pydocsign#'}, + 'docstrreq': {l_and(isrequired, isintent_nothide): '#pydocsign#'}, + 'docstrout': {isintent_out: '#pydocsignout#'}, + 'latexdocstropt': {l_and(isoptional, isintent_nothide): ['\\item[]{{}\\verb@#pydocsign#@{}}', + {hasnote: '--- #note#'}]}, + 'latexdocstrreq': {l_and(isrequired, isintent_nothide): ['\\item[]{{}\\verb@#pydocsign#@{}}', + {hasnote: '--- #note#'}]}, + 'latexdocstrout': {isintent_out: ['\\item[]{{}\\verb@#pydocsignout#@{}}', + {l_and(hasnote, isintent_hide): '--- #note#', + l_and(hasnote, isintent_nothide): '--- See above.'}]}, + 'depend': '' + }, + # Required/Optional arguments + { + 'kwlist': '"#varname#",', + 'docsign': '#varname#,', + '_check': l_and(isintent_nothide, l_not(isoptional)) + }, + { + 'kwlistopt': '"#varname#",', + 'docsignopt': '#varname#=#showinit#,', + 'docsignoptshort': '#varname#,', + '_check': l_and(isintent_nothide, isoptional) + }, + # Docstring/BuildValue + { + 'docreturn': '#outvarname#,', + 'returnformat': '#varrformat#', + '_check': isintent_out + }, + # Externals (call-back functions) + { # Common + 'docsignxa': {isintent_nothide: '#varname#_extra_args=(),'}, + 'docsignxashort': {isintent_nothide: '#varname#_extra_args,'}, + 'docstropt': {isintent_nothide: '#varname#_extra_args : input tuple, optional\\n Default: ()'}, + 'docstrcbs': '#cbdocstr#', + 'latexdocstrcbs': '\\item[] #cblatexdocstr#', + 'latexdocstropt': {isintent_nothide: '\\item[]{{}\\verb@#varname#_extra_args := () input tuple@{}} --- Extra arguments for call-back function {{}\\verb@#varname#@{}}.'}, + 'decl': [' #cbname#_t #varname#_cb = { Py_None, NULL, 0 };', + ' #cbname#_t *#varname#_cb_ptr = &#varname#_cb;', + ' PyTupleObject *#varname#_xa_capi = NULL;', + {l_not(isintent_callback): + ' #cbname#_typedef #varname#_cptr;'} + ], + 'kwlistxa': {isintent_nothide: '"#varname#_extra_args",'}, + 'argformat': {isrequired: 'O'}, + 'keyformat': {isoptional: 'O'}, + 'xaformat': {isintent_nothide: 'O!'}, + 'args_capi': {isrequired: ',&#varname#_cb.capi'}, + 'keys_capi': {isoptional: ',&#varname#_cb.capi'}, + 'keys_xa': ',&PyTuple_Type,&#varname#_xa_capi', + 'setjmpbuf': '(setjmp(#varname#_cb.jmpbuf))', + 'callfortran': {l_not(isintent_callback): '#varname#_cptr,'}, + 'need': ['#cbname#', 'setjmp.h'], + '_check':isexternal + }, + { + 'frompyobj': [{l_not(isintent_callback): """\ +if(F2PyCapsule_Check(#varname#_cb.capi)) { + #varname#_cptr = F2PyCapsule_AsVoidPtr(#varname#_cb.capi); +} else { + #varname#_cptr = #cbname#; +} +"""}, {isintent_callback: """\ +if (#varname#_cb.capi==Py_None) { + #varname#_cb.capi = PyObject_GetAttrString(#modulename#_module,\"#varname#\"); + if (#varname#_cb.capi) { + if (#varname#_xa_capi==NULL) { + if (PyObject_HasAttrString(#modulename#_module,\"#varname#_extra_args\")) { + PyObject* capi_tmp = PyObject_GetAttrString(#modulename#_module,\"#varname#_extra_args\"); + if (capi_tmp) { + #varname#_xa_capi = (PyTupleObject *)PySequence_Tuple(capi_tmp); + Py_DECREF(capi_tmp); + } + else { + #varname#_xa_capi = (PyTupleObject *)Py_BuildValue(\"()\"); + } + if (#varname#_xa_capi==NULL) { + PyErr_SetString(#modulename#_error,\"Failed to convert #modulename#.#varname#_extra_args to tuple.\\n\"); + return NULL; + } + } + } + } + if (#varname#_cb.capi==NULL) { + PyErr_SetString(#modulename#_error,\"Callback #varname# not defined (as an argument or module #modulename# attribute).\\n\"); + return NULL; + } +} +"""}, + """\ + if (create_cb_arglist(#varname#_cb.capi,#varname#_xa_capi,#maxnofargs#,#nofoptargs#,&#varname#_cb.nofargs,&#varname#_cb.args_capi,\"failed in processing argument list for call-back #varname#.\")) { +""", + {debugcapi: ["""\ + fprintf(stderr,\"debug-capi:Assuming %d arguments; at most #maxnofargs#(-#nofoptargs#) is expected.\\n\",#varname#_cb.nofargs); + CFUNCSMESSPY(\"for #varname#=\",#varname#_cb.capi);""", + {l_not(isintent_callback): """ fprintf(stderr,\"#vardebugshowvalue# (call-back in C).\\n\",#cbname#);"""}]}, + """\ + CFUNCSMESS(\"Saving callback variables for `#varname#`.\\n\"); + #varname#_cb_ptr = swap_active_#cbname#(#varname#_cb_ptr);""", + ], + 'cleanupfrompyobj': + """\ + CFUNCSMESS(\"Restoring callback variables for `#varname#`.\\n\"); + #varname#_cb_ptr = swap_active_#cbname#(#varname#_cb_ptr); + Py_DECREF(#varname#_cb.args_capi); + }""", + 'need': ['SWAP', 'create_cb_arglist'], + '_check':isexternal, + '_depend':'' + }, + # Scalars (not complex) + { # Common + 'decl': ' #ctype# #varname# = 0;', + 'pyobjfrom': {debugcapi: ' fprintf(stderr,"#vardebugshowvalue#\\n",#varname#);'}, + 'callfortran': {l_or(isintent_c, isattr_value): '#varname#,', l_not(l_or(isintent_c, isattr_value)): '&#varname#,'}, + 'return': {isintent_out: ',#varname#'}, + '_check': l_and(isscalar, l_not(iscomplex)) + }, { + 'need': {hasinitvalue: 'math.h'}, + '_check': l_and(isscalar, l_not(iscomplex)), + }, { # Not hidden + 'decl': ' PyObject *#varname#_capi = Py_None;', + 'argformat': {isrequired: 'O'}, + 'keyformat': {isoptional: 'O'}, + 'args_capi': {isrequired: ',&#varname#_capi'}, + 'keys_capi': {isoptional: ',&#varname#_capi'}, + 'pyobjfrom': {isintent_inout: """\ + f2py_success = try_pyarr_from_#ctype#(#varname#_capi,&#varname#); + if (f2py_success) {"""}, + 'closepyobjfrom': {isintent_inout: " } /*if (f2py_success) of #varname# pyobjfrom*/"}, + 'need': {isintent_inout: 'try_pyarr_from_#ctype#'}, + '_check': l_and(isscalar, l_not(iscomplex), l_not(isstring), + isintent_nothide) + }, { + 'frompyobj': [ + # hasinitvalue... + # if pyobj is None: + # varname = init + # else + # from_pyobj(varname) + # + # isoptional and noinitvalue... + # if pyobj is not None: + # from_pyobj(varname) + # else: + # varname is uninitialized + # + # ... + # from_pyobj(varname) + # + {hasinitvalue: ' if (#varname#_capi == Py_None) #varname# = #init#; else', + '_depend': ''}, + {l_and(isoptional, l_not(hasinitvalue)): ' if (#varname#_capi != Py_None)', + '_depend': ''}, + {l_not(islogical): '''\ + f2py_success = #ctype#_from_pyobj(&#varname#,#varname#_capi,"#pyname#() #nth# (#varname#) can\'t be converted to #ctype#"); + if (f2py_success) {'''}, + {islogical: '''\ + #varname# = (#ctype#)PyObject_IsTrue(#varname#_capi); + f2py_success = 1; + if (f2py_success) {'''}, + ], + 'cleanupfrompyobj': ' } /*if (f2py_success) of #varname#*/', + 'need': {l_not(islogical): '#ctype#_from_pyobj'}, + '_check': l_and(isscalar, l_not(iscomplex), isintent_nothide), + '_depend': '' + }, { # Hidden + 'frompyobj': {hasinitvalue: ' #varname# = #init#;'}, + 'need': typedef_need_dict, + '_check': l_and(isscalar, l_not(iscomplex), isintent_hide), + '_depend': '' + }, { # Common + 'frompyobj': {debugcapi: ' fprintf(stderr,"#vardebugshowvalue#\\n",#varname#);'}, + '_check': l_and(isscalar, l_not(iscomplex)), + '_depend': '' + }, + # Complex scalars + { # Common + 'decl': ' #ctype# #varname#;', + 'callfortran': {isintent_c: '#varname#,', l_not(isintent_c): '&#varname#,'}, + 'pyobjfrom': {debugcapi: ' fprintf(stderr,"#vardebugshowvalue#\\n",#varname#.r,#varname#.i);'}, + 'return': {isintent_out: ',#varname#_capi'}, + '_check': iscomplex + }, { # Not hidden + 'decl': ' PyObject *#varname#_capi = Py_None;', + 'argformat': {isrequired: 'O'}, + 'keyformat': {isoptional: 'O'}, + 'args_capi': {isrequired: ',&#varname#_capi'}, + 'keys_capi': {isoptional: ',&#varname#_capi'}, + 'need': {isintent_inout: 'try_pyarr_from_#ctype#'}, + 'pyobjfrom': {isintent_inout: """\ + f2py_success = try_pyarr_from_#ctype#(#varname#_capi,&#varname#); + if (f2py_success) {"""}, + 'closepyobjfrom': {isintent_inout: " } /*if (f2py_success) of #varname# pyobjfrom*/"}, + '_check': l_and(iscomplex, isintent_nothide) + }, { + 'frompyobj': [{hasinitvalue: ' if (#varname#_capi==Py_None) {#varname#.r = #init.r#, #varname#.i = #init.i#;} else'}, + {l_and(isoptional, l_not(hasinitvalue)) + : ' if (#varname#_capi != Py_None)'}, + ' f2py_success = #ctype#_from_pyobj(&#varname#,#varname#_capi,"#pyname#() #nth# (#varname#) can\'t be converted to #ctype#");' + '\n if (f2py_success) {'], + 'cleanupfrompyobj': ' } /*if (f2py_success) of #varname# frompyobj*/', + 'need': ['#ctype#_from_pyobj'], + '_check': l_and(iscomplex, isintent_nothide), + '_depend': '' + }, { # Hidden + 'decl': {isintent_out: ' PyObject *#varname#_capi = Py_None;'}, + '_check': l_and(iscomplex, isintent_hide) + }, { + 'frompyobj': {hasinitvalue: ' #varname#.r = #init.r#, #varname#.i = #init.i#;'}, + '_check': l_and(iscomplex, isintent_hide), + '_depend': '' + }, { # Common + 'pyobjfrom': {isintent_out: ' #varname#_capi = pyobj_from_#ctype#1(#varname#);'}, + 'need': ['pyobj_from_#ctype#1'], + '_check': iscomplex + }, { + 'frompyobj': {debugcapi: ' fprintf(stderr,"#vardebugshowvalue#\\n",#varname#.r,#varname#.i);'}, + '_check': iscomplex, + '_depend': '' + }, + # String + { # Common + 'decl': [' #ctype# #varname# = NULL;', + ' int slen(#varname#);', + ' PyObject *#varname#_capi = Py_None;'], + 'callfortran':'#varname#,', + 'callfortranappend':'slen(#varname#),', + 'pyobjfrom':[ + {debugcapi: + ' fprintf(stderr,' + '"#vardebugshowvalue#\\n",slen(#varname#),#varname#);'}, + # The trailing null value for Fortran is blank. + {l_and(isintent_out, l_not(isintent_c)): + " STRINGPADN(#varname#, slen(#varname#), ' ', '\\0');"}, + ], + 'return': {isintent_out: ',#varname#'}, + 'need': ['len..', + {l_and(isintent_out, l_not(isintent_c)): 'STRINGPADN'}], + '_check': isstring + }, { # Common + 'frompyobj': [ + """\ + slen(#varname#) = #elsize#; + f2py_success = #ctype#_from_pyobj(&#varname#,&slen(#varname#),#init#,""" +"""#varname#_capi,\"#ctype#_from_pyobj failed in converting #nth#""" +"""`#varname#\' of #pyname# to C #ctype#\"); + if (f2py_success) {""", + # The trailing null value for Fortran is blank. + {l_not(isintent_c): + " STRINGPADN(#varname#, slen(#varname#), '\\0', ' ');"}, + ], + 'cleanupfrompyobj': """\ + STRINGFREE(#varname#); + } /*if (f2py_success) of #varname#*/""", + 'need': ['#ctype#_from_pyobj', 'len..', 'STRINGFREE', + {l_not(isintent_c): 'STRINGPADN'}], + '_check':isstring, + '_depend':'' + }, { # Not hidden + 'argformat': {isrequired: 'O'}, + 'keyformat': {isoptional: 'O'}, + 'args_capi': {isrequired: ',&#varname#_capi'}, + 'keys_capi': {isoptional: ',&#varname#_capi'}, + 'pyobjfrom': [ + {l_and(isintent_inout, l_not(isintent_c)): + " STRINGPADN(#varname#, slen(#varname#), ' ', '\\0');"}, + {isintent_inout: '''\ + f2py_success = try_pyarr_from_#ctype#(#varname#_capi, #varname#, + slen(#varname#)); + if (f2py_success) {'''}], + 'closepyobjfrom': {isintent_inout: ' } /*if (f2py_success) of #varname# pyobjfrom*/'}, + 'need': {isintent_inout: 'try_pyarr_from_#ctype#', + l_and(isintent_inout, l_not(isintent_c)): 'STRINGPADN'}, + '_check': l_and(isstring, isintent_nothide) + }, { # Hidden + '_check': l_and(isstring, isintent_hide) + }, { + 'frompyobj': {debugcapi: ' fprintf(stderr,"#vardebugshowvalue#\\n",slen(#varname#),#varname#);'}, + '_check': isstring, + '_depend': '' + }, + # Array + { # Common + 'decl': [' #ctype# *#varname# = NULL;', + ' npy_intp #varname#_Dims[#rank#] = {#rank*[-1]#};', + ' const int #varname#_Rank = #rank#;', + ' PyArrayObject *capi_#varname#_as_array = NULL;', + ' int capi_#varname#_intent = 0;', + {isstringarray: ' int slen(#varname#) = 0;'}, + ], + 'callfortran':'#varname#,', + 'callfortranappend': {isstringarray: 'slen(#varname#),'}, + 'return': {isintent_out: ',capi_#varname#_as_array'}, + 'need': 'len..', + '_check': isarray + }, { # intent(overwrite) array + 'decl': ' int capi_overwrite_#varname# = 1;', + 'kwlistxa': '"overwrite_#varname#",', + 'xaformat': 'i', + 'keys_xa': ',&capi_overwrite_#varname#', + 'docsignxa': 'overwrite_#varname#=1,', + 'docsignxashort': 'overwrite_#varname#,', + 'docstropt': 'overwrite_#varname# : input int, optional\\n Default: 1', + '_check': l_and(isarray, isintent_overwrite), + }, { + 'frompyobj': ' capi_#varname#_intent |= (capi_overwrite_#varname#?0:F2PY_INTENT_COPY);', + '_check': l_and(isarray, isintent_overwrite), + '_depend': '', + }, + { # intent(copy) array + 'decl': ' int capi_overwrite_#varname# = 0;', + 'kwlistxa': '"overwrite_#varname#",', + 'xaformat': 'i', + 'keys_xa': ',&capi_overwrite_#varname#', + 'docsignxa': 'overwrite_#varname#=0,', + 'docsignxashort': 'overwrite_#varname#,', + 'docstropt': 'overwrite_#varname# : input int, optional\\n Default: 0', + '_check': l_and(isarray, isintent_copy), + }, { + 'frompyobj': ' capi_#varname#_intent |= (capi_overwrite_#varname#?0:F2PY_INTENT_COPY);', + '_check': l_and(isarray, isintent_copy), + '_depend': '', + }, { + 'need': [{hasinitvalue: 'forcomb'}, {hasinitvalue: 'CFUNCSMESS'}], + '_check': isarray, + '_depend': '' + }, { # Not hidden + 'decl': ' PyObject *#varname#_capi = Py_None;', + 'argformat': {isrequired: 'O'}, + 'keyformat': {isoptional: 'O'}, + 'args_capi': {isrequired: ',&#varname#_capi'}, + 'keys_capi': {isoptional: ',&#varname#_capi'}, + '_check': l_and(isarray, isintent_nothide) + }, { + 'frompyobj': [ + ' #setdims#;', + ' capi_#varname#_intent |= #intent#;', + (' const char * capi_errmess = "#modulename#.#pyname#:' + ' failed to create array from the #nth# `#varname#`";'), + {isintent_hide: + ' capi_#varname#_as_array = ndarray_from_pyobj(' + ' #atype#,#elsize#,#varname#_Dims,#varname#_Rank,' + ' capi_#varname#_intent,Py_None,capi_errmess);'}, + {isintent_nothide: + ' capi_#varname#_as_array = ndarray_from_pyobj(' + ' #atype#,#elsize#,#varname#_Dims,#varname#_Rank,' + ' capi_#varname#_intent,#varname#_capi,capi_errmess);'}, + """\ + if (capi_#varname#_as_array == NULL) { + PyObject* capi_err = PyErr_Occurred(); + if (capi_err == NULL) { + capi_err = #modulename#_error; + PyErr_SetString(capi_err, capi_errmess); + } + } else { + #varname# = (#ctype# *)(PyArray_DATA(capi_#varname#_as_array)); +""", + {isstringarray: + ' slen(#varname#) = f2py_itemsize(#varname#);'}, + {hasinitvalue: [ + {isintent_nothide: + ' if (#varname#_capi == Py_None) {'}, + {isintent_hide: ' {'}, + {iscomplexarray: ' #ctype# capi_c;'}, + """\ + int *_i,capi_i=0; + CFUNCSMESS(\"#name#: Initializing #varname#=#init#\\n\"); + if (initforcomb(PyArray_DIMS(capi_#varname#_as_array), + PyArray_NDIM(capi_#varname#_as_array),1)) { + while ((_i = nextforcomb())) + #varname#[capi_i++] = #init#; /* fortran way */ + } else { + PyObject *exc, *val, *tb; + PyErr_Fetch(&exc, &val, &tb); + PyErr_SetString(exc ? exc : #modulename#_error, + \"Initialization of #nth# #varname# failed (initforcomb).\"); + npy_PyErr_ChainExceptionsCause(exc, val, tb); + f2py_success = 0; + } + } + if (f2py_success) {"""]}, + ], + 'cleanupfrompyobj': [ # note that this list will be reversed + ' } ' + '/* if (capi_#varname#_as_array == NULL) ... else of #varname# */', + {l_not(l_or(isintent_out, isintent_hide)): """\ + if((PyObject *)capi_#varname#_as_array!=#varname#_capi) { + Py_XDECREF(capi_#varname#_as_array); }"""}, + {l_and(isintent_hide, l_not(isintent_out)) + : """ Py_XDECREF(capi_#varname#_as_array);"""}, + {hasinitvalue: ' } /*if (f2py_success) of #varname# init*/'}, + ], + '_check': isarray, + '_depend': '' + }, + # Scalararray + { # Common + '_check': l_and(isarray, l_not(iscomplexarray)) + }, { # Not hidden + '_check': l_and(isarray, l_not(iscomplexarray), isintent_nothide) + }, + # Integer*1 array + {'need': '#ctype#', + '_check': isint1array, + '_depend': '' + }, + # Integer*-1 array + {'need': '#ctype#', + '_check': isunsigned_chararray, + '_depend': '' + }, + # Integer*-2 array + {'need': '#ctype#', + '_check': isunsigned_shortarray, + '_depend': '' + }, + # Integer*-8 array + {'need': '#ctype#', + '_check': isunsigned_long_longarray, + '_depend': '' + }, + # Complexarray + {'need': '#ctype#', + '_check': iscomplexarray, + '_depend': '' + }, + # Character + { + 'need': 'string', + '_check': ischaracter, + }, + # Character array + { + 'need': 'string', + '_check': ischaracterarray, + }, + # Stringarray + { + 'callfortranappend': {isarrayofstrings: 'flen(#varname#),'}, + 'need': 'string', + '_check': isstringarray + } +] + +################# Rules for checking ############### + +check_rules = [ + { + 'frompyobj': {debugcapi: ' fprintf(stderr,\"debug-capi:Checking `#check#\'\\n\");'}, + 'need': 'len..' + }, { + 'frompyobj': ' CHECKSCALAR(#check#,\"#check#\",\"#nth# #varname#\",\"#varshowvalue#\",#varname#) {', + 'cleanupfrompyobj': ' } /*CHECKSCALAR(#check#)*/', + 'need': 'CHECKSCALAR', + '_check': l_and(isscalar, l_not(iscomplex)), + '_break': '' + }, { + 'frompyobj': ' CHECKSTRING(#check#,\"#check#\",\"#nth# #varname#\",\"#varshowvalue#\",#varname#) {', + 'cleanupfrompyobj': ' } /*CHECKSTRING(#check#)*/', + 'need': 'CHECKSTRING', + '_check': isstring, + '_break': '' + }, { + 'need': 'CHECKARRAY', + 'frompyobj': ' CHECKARRAY(#check#,\"#check#\",\"#nth# #varname#\") {', + 'cleanupfrompyobj': ' } /*CHECKARRAY(#check#)*/', + '_check': isarray, + '_break': '' + }, { + 'need': 'CHECKGENERIC', + 'frompyobj': ' CHECKGENERIC(#check#,\"#check#\",\"#nth# #varname#\") {', + 'cleanupfrompyobj': ' } /*CHECKGENERIC(#check#)*/', + } +] + +########## Applying the rules. No need to modify what follows ############# + +#################### Build C/API module ####################### + + +def buildmodule(m, um): + """ + Return + """ + outmess(' Building module "%s"...\n' % (m['name'])) + ret = {} + mod_rules = defmod_rules[:] + vrd = capi_maps.modsign2map(m) + rd = dictappend({'f2py_version': f2py_version}, vrd) + funcwrappers = [] + funcwrappers2 = [] # F90 codes + for n in m['interfaced']: + nb = None + for bi in m['body']: + if bi['block'] not in ['interface', 'abstract interface']: + errmess('buildmodule: Expected interface block. Skipping.\n') + continue + for b in bi['body']: + if b['name'] == n: + nb = b + break + + if not nb: + print( + 'buildmodule: Could not find the body of interfaced routine "%s". Skipping.\n' % (n), file=sys.stderr) + continue + nb_list = [nb] + if 'entry' in nb: + for k, a in nb['entry'].items(): + nb1 = copy.deepcopy(nb) + del nb1['entry'] + nb1['name'] = k + nb1['args'] = a + nb_list.append(nb1) + for nb in nb_list: + # requiresf90wrapper must be called before buildapi as it + # rewrites assumed shape arrays as automatic arrays. + isf90 = requiresf90wrapper(nb) + # options is in scope here + if options['emptygen']: + b_path = options['buildpath'] + m_name = vrd['modulename'] + outmess(' Generating possibly empty wrappers"\n') + Path(f"{b_path}/{vrd['coutput']}").touch() + if isf90: + # f77 + f90 wrappers + outmess(f' Maybe empty "{m_name}-f2pywrappers2.f90"\n') + Path(f'{b_path}/{m_name}-f2pywrappers2.f90').touch() + outmess(f' Maybe empty "{m_name}-f2pywrappers.f"\n') + Path(f'{b_path}/{m_name}-f2pywrappers.f').touch() + else: + # only f77 wrappers + outmess(f' Maybe empty "{m_name}-f2pywrappers.f"\n') + Path(f'{b_path}/{m_name}-f2pywrappers.f').touch() + api, wrap = buildapi(nb) + if wrap: + if isf90: + funcwrappers2.append(wrap) + else: + funcwrappers.append(wrap) + ar = applyrules(api, vrd) + rd = dictappend(rd, ar) + + # Construct COMMON block support + cr, wrap = common_rules.buildhooks(m) + if wrap: + funcwrappers.append(wrap) + ar = applyrules(cr, vrd) + rd = dictappend(rd, ar) + + # Construct F90 module support + mr, wrap = f90mod_rules.buildhooks(m) + if wrap: + funcwrappers2.append(wrap) + ar = applyrules(mr, vrd) + rd = dictappend(rd, ar) + + for u in um: + ar = use_rules.buildusevars(u, m['use'][u['name']]) + rd = dictappend(rd, ar) + + needs = cfuncs.get_needs() + # Add mapped definitions + needs['typedefs'] += [cvar for cvar in capi_maps.f2cmap_mapped # + if cvar in typedef_need_dict.values()] + code = {} + for n in needs.keys(): + code[n] = [] + for k in needs[n]: + c = '' + if k in cfuncs.includes0: + c = cfuncs.includes0[k] + elif k in cfuncs.includes: + c = cfuncs.includes[k] + elif k in cfuncs.userincludes: + c = cfuncs.userincludes[k] + elif k in cfuncs.typedefs: + c = cfuncs.typedefs[k] + elif k in cfuncs.typedefs_generated: + c = cfuncs.typedefs_generated[k] + elif k in cfuncs.cppmacros: + c = cfuncs.cppmacros[k] + elif k in cfuncs.cfuncs: + c = cfuncs.cfuncs[k] + elif k in cfuncs.callbacks: + c = cfuncs.callbacks[k] + elif k in cfuncs.f90modhooks: + c = cfuncs.f90modhooks[k] + elif k in cfuncs.commonhooks: + c = cfuncs.commonhooks[k] + else: + errmess('buildmodule: unknown need %s.\n' % (repr(k))) + continue + code[n].append(c) + mod_rules.append(code) + for r in mod_rules: + if ('_check' in r and r['_check'](m)) or ('_check' not in r): + ar = applyrules(r, vrd, m) + rd = dictappend(rd, ar) + ar = applyrules(module_rules, rd) + + fn = os.path.join(options['buildpath'], vrd['coutput']) + ret['csrc'] = fn + with open(fn, 'w') as f: + f.write(ar['modulebody'].replace('\t', 2 * ' ')) + outmess(' Wrote C/API module "%s" to file "%s"\n' % (m['name'], fn)) + + if options['dorestdoc']: + fn = os.path.join( + options['buildpath'], vrd['modulename'] + 'module.rest') + with open(fn, 'w') as f: + f.write('.. -*- rest -*-\n') + f.write('\n'.join(ar['restdoc'])) + outmess(' ReST Documentation is saved to file "%s/%smodule.rest"\n' % + (options['buildpath'], vrd['modulename'])) + if options['dolatexdoc']: + fn = os.path.join( + options['buildpath'], vrd['modulename'] + 'module.tex') + ret['ltx'] = fn + with open(fn, 'w') as f: + f.write( + '%% This file is auto-generated with f2py (version:%s)\n' % (f2py_version)) + if 'shortlatex' not in options: + f.write( + '\\documentclass{article}\n\\usepackage{a4wide}\n\\begin{document}\n\\tableofcontents\n\n') + f.write('\n'.join(ar['latexdoc'])) + if 'shortlatex' not in options: + f.write('\\end{document}') + outmess(' Documentation is saved to file "%s/%smodule.tex"\n' % + (options['buildpath'], vrd['modulename'])) + if funcwrappers: + wn = os.path.join(options['buildpath'], vrd['f2py_wrapper_output']) + ret['fsrc'] = wn + with open(wn, 'w') as f: + f.write('C -*- fortran -*-\n') + f.write( + 'C This file is autogenerated with f2py (version:%s)\n' % (f2py_version)) + f.write( + 'C It contains Fortran 77 wrappers to fortran functions.\n') + lines = [] + for l in ('\n\n'.join(funcwrappers) + '\n').split('\n'): + if 0 <= l.find('!') < 66: + # don't split comment lines + lines.append(l + '\n') + elif l and l[0] == ' ': + while len(l) >= 66: + lines.append(l[:66] + '\n &') + l = l[66:] + lines.append(l + '\n') + else: + lines.append(l + '\n') + lines = ''.join(lines).replace('\n &\n', '\n') + f.write(lines) + outmess(' Fortran 77 wrappers are saved to "%s"\n' % (wn)) + if funcwrappers2: + wn = os.path.join( + options['buildpath'], '%s-f2pywrappers2.f90' % (vrd['modulename'])) + ret['fsrc'] = wn + with open(wn, 'w') as f: + f.write('! -*- f90 -*-\n') + f.write( + '! This file is autogenerated with f2py (version:%s)\n' % (f2py_version)) + f.write( + '! It contains Fortran 90 wrappers to fortran functions.\n') + lines = [] + for l in ('\n\n'.join(funcwrappers2) + '\n').split('\n'): + if 0 <= l.find('!') < 72: + # don't split comment lines + lines.append(l + '\n') + elif len(l) > 72 and l[0] == ' ': + lines.append(l[:72] + '&\n &') + l = l[72:] + while len(l) > 66: + lines.append(l[:66] + '&\n &') + l = l[66:] + lines.append(l + '\n') + else: + lines.append(l + '\n') + lines = ''.join(lines).replace('\n &\n', '\n') + f.write(lines) + outmess(' Fortran 90 wrappers are saved to "%s"\n' % (wn)) + return ret + +################## Build C/API function ############# + +stnd = {1: 'st', 2: 'nd', 3: 'rd', 4: 'th', 5: 'th', + 6: 'th', 7: 'th', 8: 'th', 9: 'th', 0: 'th'} + + +def buildapi(rout): + rout, wrap = func2subr.assubr(rout) + args, depargs = getargs2(rout) + capi_maps.depargs = depargs + var = rout['vars'] + + if ismoduleroutine(rout): + outmess(' Constructing wrapper function "%s.%s"...\n' % + (rout['modulename'], rout['name'])) + else: + outmess(' Constructing wrapper function "%s"...\n' % (rout['name'])) + # Routine + vrd = capi_maps.routsign2map(rout) + rd = dictappend({}, vrd) + for r in rout_rules: + if ('_check' in r and r['_check'](rout)) or ('_check' not in r): + ar = applyrules(r, vrd, rout) + rd = dictappend(rd, ar) + + # Args + nth, nthk = 0, 0 + savevrd = {} + for a in args: + vrd = capi_maps.sign2map(a, var[a]) + if isintent_aux(var[a]): + _rules = aux_rules + else: + _rules = arg_rules + if not isintent_hide(var[a]): + if not isoptional(var[a]): + nth = nth + 1 + vrd['nth'] = repr(nth) + stnd[nth % 10] + ' argument' + else: + nthk = nthk + 1 + vrd['nth'] = repr(nthk) + stnd[nthk % 10] + ' keyword' + else: + vrd['nth'] = 'hidden' + savevrd[a] = vrd + for r in _rules: + if '_depend' in r: + continue + if ('_check' in r and r['_check'](var[a])) or ('_check' not in r): + ar = applyrules(r, vrd, var[a]) + rd = dictappend(rd, ar) + if '_break' in r: + break + for a in depargs: + if isintent_aux(var[a]): + _rules = aux_rules + else: + _rules = arg_rules + vrd = savevrd[a] + for r in _rules: + if '_depend' not in r: + continue + if ('_check' in r and r['_check'](var[a])) or ('_check' not in r): + ar = applyrules(r, vrd, var[a]) + rd = dictappend(rd, ar) + if '_break' in r: + break + if 'check' in var[a]: + for c in var[a]['check']: + vrd['check'] = c + ar = applyrules(check_rules, vrd, var[a]) + rd = dictappend(rd, ar) + if isinstance(rd['cleanupfrompyobj'], list): + rd['cleanupfrompyobj'].reverse() + if isinstance(rd['closepyobjfrom'], list): + rd['closepyobjfrom'].reverse() + rd['docsignature'] = stripcomma(replace('#docsign##docsignopt##docsignxa#', + {'docsign': rd['docsign'], + 'docsignopt': rd['docsignopt'], + 'docsignxa': rd['docsignxa']})) + optargs = stripcomma(replace('#docsignopt##docsignxa#', + {'docsignxa': rd['docsignxashort'], + 'docsignopt': rd['docsignoptshort']} + )) + if optargs == '': + rd['docsignatureshort'] = stripcomma( + replace('#docsign#', {'docsign': rd['docsign']})) + else: + rd['docsignatureshort'] = replace('#docsign#[#docsignopt#]', + {'docsign': rd['docsign'], + 'docsignopt': optargs, + }) + rd['latexdocsignatureshort'] = rd['docsignatureshort'].replace('_', '\\_') + rd['latexdocsignatureshort'] = rd[ + 'latexdocsignatureshort'].replace(',', ', ') + cfs = stripcomma(replace('#callfortran##callfortranappend#', { + 'callfortran': rd['callfortran'], 'callfortranappend': rd['callfortranappend']})) + if len(rd['callfortranappend']) > 1: + rd['callcompaqfortran'] = stripcomma(replace('#callfortran# 0,#callfortranappend#', { + 'callfortran': rd['callfortran'], 'callfortranappend': rd['callfortranappend']})) + else: + rd['callcompaqfortran'] = cfs + rd['callfortran'] = cfs + if isinstance(rd['docreturn'], list): + rd['docreturn'] = stripcomma( + replace('#docreturn#', {'docreturn': rd['docreturn']})) + ' = ' + rd['docstrsigns'] = [] + rd['latexdocstrsigns'] = [] + for k in ['docstrreq', 'docstropt', 'docstrout', 'docstrcbs']: + if k in rd and isinstance(rd[k], list): + rd['docstrsigns'] = rd['docstrsigns'] + rd[k] + k = 'latex' + k + if k in rd and isinstance(rd[k], list): + rd['latexdocstrsigns'] = rd['latexdocstrsigns'] + rd[k][0:1] +\ + ['\\begin{description}'] + rd[k][1:] +\ + ['\\end{description}'] + + ar = applyrules(routine_rules, rd) + if ismoduleroutine(rout): + outmess(' %s\n' % (ar['docshort'])) + else: + outmess(' %s\n' % (ar['docshort'])) + return ar, wrap + + +#################### EOF rules.py ####################### diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/setup.cfg b/venv/lib/python3.10/site-packages/numpy/f2py/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..14669544cc9ec345373bf5f719e321348fc96a40 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/setup.cfg @@ -0,0 +1,3 @@ +[bdist_rpm] +doc_files = docs/ + tests/ \ No newline at end of file diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.c b/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.c new file mode 100644 index 0000000000000000000000000000000000000000..4e2aa370b643e62db1955673b9f8922da6721ebe --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.c @@ -0,0 +1,1423 @@ +#define FORTRANOBJECT_C +#include "fortranobject.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#include +#include +#include + +/* + This file implements: FortranObject, array_from_pyobj, copy_ND_array + + Author: Pearu Peterson + $Revision: 1.52 $ + $Date: 2005/07/11 07:44:20 $ +*/ + +int +F2PyDict_SetItemString(PyObject *dict, char *name, PyObject *obj) +{ + if (obj == NULL) { + fprintf(stderr, "Error loading %s\n", name); + if (PyErr_Occurred()) { + PyErr_Print(); + PyErr_Clear(); + } + return -1; + } + return PyDict_SetItemString(dict, name, obj); +} + +/* + * Python-only fallback for thread-local callback pointers + */ +void * +F2PySwapThreadLocalCallbackPtr(char *key, void *ptr) +{ + PyObject *local_dict, *value; + void *prev; + + local_dict = PyThreadState_GetDict(); + if (local_dict == NULL) { + Py_FatalError( + "F2PySwapThreadLocalCallbackPtr: PyThreadState_GetDict " + "failed"); + } + + value = PyDict_GetItemString(local_dict, key); + if (value != NULL) { + prev = PyLong_AsVoidPtr(value); + if (PyErr_Occurred()) { + Py_FatalError( + "F2PySwapThreadLocalCallbackPtr: PyLong_AsVoidPtr failed"); + } + } + else { + prev = NULL; + } + + value = PyLong_FromVoidPtr((void *)ptr); + if (value == NULL) { + Py_FatalError( + "F2PySwapThreadLocalCallbackPtr: PyLong_FromVoidPtr failed"); + } + + if (PyDict_SetItemString(local_dict, key, value) != 0) { + Py_FatalError( + "F2PySwapThreadLocalCallbackPtr: PyDict_SetItemString failed"); + } + + Py_DECREF(value); + + return prev; +} + +void * +F2PyGetThreadLocalCallbackPtr(char *key) +{ + PyObject *local_dict, *value; + void *prev; + + local_dict = PyThreadState_GetDict(); + if (local_dict == NULL) { + Py_FatalError( + "F2PyGetThreadLocalCallbackPtr: PyThreadState_GetDict failed"); + } + + value = PyDict_GetItemString(local_dict, key); + if (value != NULL) { + prev = PyLong_AsVoidPtr(value); + if (PyErr_Occurred()) { + Py_FatalError( + "F2PyGetThreadLocalCallbackPtr: PyLong_AsVoidPtr failed"); + } + } + else { + prev = NULL; + } + + return prev; +} + +static PyArray_Descr * +get_descr_from_type_and_elsize(const int type_num, const int elsize) { + PyArray_Descr * descr = PyArray_DescrFromType(type_num); + if (type_num == NPY_STRING) { + // PyArray_DescrFromType returns descr with elsize = 0. + PyArray_DESCR_REPLACE(descr); + if (descr == NULL) { + return NULL; + } + PyDataType_SET_ELSIZE(descr, elsize); + } + return descr; +} + +/************************* FortranObject *******************************/ + +typedef PyObject *(*fortranfunc)(PyObject *, PyObject *, PyObject *, void *); + +PyObject * +PyFortranObject_New(FortranDataDef *defs, f2py_void_func init) +{ + int i; + PyFortranObject *fp = NULL; + PyObject *v = NULL; + if (init != NULL) { /* Initialize F90 module objects */ + (*(init))(); + } + fp = PyObject_New(PyFortranObject, &PyFortran_Type); + if (fp == NULL) { + return NULL; + } + if ((fp->dict = PyDict_New()) == NULL) { + Py_DECREF(fp); + return NULL; + } + fp->len = 0; + while (defs[fp->len].name != NULL) { + fp->len++; + } + if (fp->len == 0) { + goto fail; + } + fp->defs = defs; + for (i = 0; i < fp->len; i++) { + if (fp->defs[i].rank == -1) { /* Is Fortran routine */ + v = PyFortranObject_NewAsAttr(&(fp->defs[i])); + if (v == NULL) { + goto fail; + } + PyDict_SetItemString(fp->dict, fp->defs[i].name, v); + Py_XDECREF(v); + } + else if ((fp->defs[i].data) != + NULL) { /* Is Fortran variable or array (not allocatable) */ + PyArray_Descr * + descr = get_descr_from_type_and_elsize(fp->defs[i].type, + fp->defs[i].elsize); + if (descr == NULL) { + goto fail; + } + v = PyArray_NewFromDescr(&PyArray_Type, descr, fp->defs[i].rank, + fp->defs[i].dims.d, NULL, fp->defs[i].data, + NPY_ARRAY_FARRAY, NULL); + if (v == NULL) { + Py_DECREF(descr); + goto fail; + } + PyDict_SetItemString(fp->dict, fp->defs[i].name, v); + Py_XDECREF(v); + } + } + return (PyObject *)fp; +fail: + Py_XDECREF(fp); + return NULL; +} + +PyObject * +PyFortranObject_NewAsAttr(FortranDataDef *defs) +{ /* used for calling F90 module routines */ + PyFortranObject *fp = NULL; + fp = PyObject_New(PyFortranObject, &PyFortran_Type); + if (fp == NULL) + return NULL; + if ((fp->dict = PyDict_New()) == NULL) { + PyObject_Del(fp); + return NULL; + } + fp->len = 1; + fp->defs = defs; + if (defs->rank == -1) { + PyDict_SetItemString(fp->dict, "__name__", PyUnicode_FromFormat("function %s", defs->name)); + } else if (defs->rank == 0) { + PyDict_SetItemString(fp->dict, "__name__", PyUnicode_FromFormat("scalar %s", defs->name)); + } else { + PyDict_SetItemString(fp->dict, "__name__", PyUnicode_FromFormat("array %s", defs->name)); + } + return (PyObject *)fp; +} + +/* Fortran methods */ + +static void +fortran_dealloc(PyFortranObject *fp) +{ + Py_XDECREF(fp->dict); + PyObject_Del(fp); +} + +/* Returns number of bytes consumed from buf, or -1 on error. */ +static Py_ssize_t +format_def(char *buf, Py_ssize_t size, FortranDataDef def) +{ + char *p = buf; + int i; + npy_intp n; + + n = PyOS_snprintf(p, size, "array(%" NPY_INTP_FMT, def.dims.d[0]); + if (n < 0 || n >= size) { + return -1; + } + p += n; + size -= n; + + for (i = 1; i < def.rank; i++) { + n = PyOS_snprintf(p, size, ",%" NPY_INTP_FMT, def.dims.d[i]); + if (n < 0 || n >= size) { + return -1; + } + p += n; + size -= n; + } + + if (size <= 0) { + return -1; + } + + *p++ = ')'; + size--; + + if (def.data == NULL) { + static const char notalloc[] = ", not allocated"; + if ((size_t)size < sizeof(notalloc)) { + return -1; + } + memcpy(p, notalloc, sizeof(notalloc)); + p += sizeof(notalloc); + size -= sizeof(notalloc); + } + + return p - buf; +} + +static PyObject * +fortran_doc(FortranDataDef def) +{ + char *buf, *p; + PyObject *s = NULL; + Py_ssize_t n, origsize, size = 100; + + if (def.doc != NULL) { + size += strlen(def.doc); + } + origsize = size; + buf = p = (char *)PyMem_Malloc(size); + if (buf == NULL) { + return PyErr_NoMemory(); + } + + if (def.rank == -1) { + if (def.doc) { + n = strlen(def.doc); + if (n > size) { + goto fail; + } + memcpy(p, def.doc, n); + p += n; + size -= n; + } + else { + n = PyOS_snprintf(p, size, "%s - no docs available", def.name); + if (n < 0 || n >= size) { + goto fail; + } + p += n; + size -= n; + } + } + else { + PyArray_Descr *d = PyArray_DescrFromType(def.type); + n = PyOS_snprintf(p, size, "%s : '%c'-", def.name, d->type); + Py_DECREF(d); + if (n < 0 || n >= size) { + goto fail; + } + p += n; + size -= n; + + if (def.data == NULL) { + n = format_def(p, size, def); + if (n < 0) { + goto fail; + } + p += n; + size -= n; + } + else if (def.rank > 0) { + n = format_def(p, size, def); + if (n < 0) { + goto fail; + } + p += n; + size -= n; + } + else { + n = strlen("scalar"); + if (size < n) { + goto fail; + } + memcpy(p, "scalar", n); + p += n; + size -= n; + } + } + if (size <= 1) { + goto fail; + } + *p++ = '\n'; + size--; + + /* p now points one beyond the last character of the string in buf */ + s = PyUnicode_FromStringAndSize(buf, p - buf); + + PyMem_Free(buf); + return s; + +fail: + fprintf(stderr, + "fortranobject.c: fortran_doc: len(p)=%zd>%zd=size:" + " too long docstring required, increase size\n", + p - buf, origsize); + PyMem_Free(buf); + return NULL; +} + +static FortranDataDef *save_def; /* save pointer of an allocatable array */ +static void +set_data(char *d, npy_intp *f) +{ /* callback from Fortran */ + if (*f) /* In fortran f=allocated(d) */ + save_def->data = d; + else + save_def->data = NULL; + /* printf("set_data: d=%p,f=%d\n",d,*f); */ +} + +static PyObject * +fortran_getattr(PyFortranObject *fp, char *name) +{ + int i, j, k, flag; + if (fp->dict != NULL) { + PyObject *v = _PyDict_GetItemStringWithError(fp->dict, name); + if (v == NULL && PyErr_Occurred()) { + return NULL; + } + else if (v != NULL) { + Py_INCREF(v); + return v; + } + } + for (i = 0, j = 1; i < fp->len && (j = strcmp(name, fp->defs[i].name)); + i++) + ; + if (j == 0) + if (fp->defs[i].rank != -1) { /* F90 allocatable array */ + if (fp->defs[i].func == NULL) + return NULL; + for (k = 0; k < fp->defs[i].rank; ++k) fp->defs[i].dims.d[k] = -1; + save_def = &fp->defs[i]; + (*(fp->defs[i].func))(&fp->defs[i].rank, fp->defs[i].dims.d, + set_data, &flag); + if (flag == 2) + k = fp->defs[i].rank + 1; + else + k = fp->defs[i].rank; + if (fp->defs[i].data != NULL) { /* array is allocated */ + PyObject *v = PyArray_New( + &PyArray_Type, k, fp->defs[i].dims.d, fp->defs[i].type, + NULL, fp->defs[i].data, 0, NPY_ARRAY_FARRAY, NULL); + if (v == NULL) + return NULL; + /* Py_INCREF(v); */ + return v; + } + else { /* array is not allocated */ + Py_RETURN_NONE; + } + } + if (strcmp(name, "__dict__") == 0) { + Py_INCREF(fp->dict); + return fp->dict; + } + if (strcmp(name, "__doc__") == 0) { + PyObject *s = PyUnicode_FromString(""), *s2, *s3; + for (i = 0; i < fp->len; i++) { + s2 = fortran_doc(fp->defs[i]); + s3 = PyUnicode_Concat(s, s2); + Py_DECREF(s2); + Py_DECREF(s); + s = s3; + } + if (PyDict_SetItemString(fp->dict, name, s)) + return NULL; + return s; + } + if ((strcmp(name, "_cpointer") == 0) && (fp->len == 1)) { + PyObject *cobj = + F2PyCapsule_FromVoidPtr((void *)(fp->defs[0].data), NULL); + if (PyDict_SetItemString(fp->dict, name, cobj)) + return NULL; + return cobj; + } + PyObject *str, *ret; + str = PyUnicode_FromString(name); + ret = PyObject_GenericGetAttr((PyObject *)fp, str); + Py_DECREF(str); + return ret; +} + +static int +fortran_setattr(PyFortranObject *fp, char *name, PyObject *v) +{ + int i, j, flag; + PyArrayObject *arr = NULL; + for (i = 0, j = 1; i < fp->len && (j = strcmp(name, fp->defs[i].name)); + i++) + ; + if (j == 0) { + if (fp->defs[i].rank == -1) { + PyErr_SetString(PyExc_AttributeError, + "over-writing fortran routine"); + return -1; + } + if (fp->defs[i].func != NULL) { /* is allocatable array */ + npy_intp dims[F2PY_MAX_DIMS]; + int k; + save_def = &fp->defs[i]; + if (v != Py_None) { /* set new value (reallocate if needed -- + see f2py generated code for more + details ) */ + for (k = 0; k < fp->defs[i].rank; k++) dims[k] = -1; + if ((arr = array_from_pyobj(fp->defs[i].type, dims, + fp->defs[i].rank, F2PY_INTENT_IN, + v)) == NULL) + return -1; + (*(fp->defs[i].func))(&fp->defs[i].rank, PyArray_DIMS(arr), + set_data, &flag); + } + else { /* deallocate */ + for (k = 0; k < fp->defs[i].rank; k++) dims[k] = 0; + (*(fp->defs[i].func))(&fp->defs[i].rank, dims, set_data, + &flag); + for (k = 0; k < fp->defs[i].rank; k++) dims[k] = -1; + } + memcpy(fp->defs[i].dims.d, dims, + fp->defs[i].rank * sizeof(npy_intp)); + } + else { /* not allocatable array */ + if ((arr = array_from_pyobj(fp->defs[i].type, fp->defs[i].dims.d, + fp->defs[i].rank, F2PY_INTENT_IN, + v)) == NULL) + return -1; + } + if (fp->defs[i].data != + NULL) { /* copy Python object to Fortran array */ + npy_intp s = PyArray_MultiplyList(fp->defs[i].dims.d, + PyArray_NDIM(arr)); + if (s == -1) + s = PyArray_MultiplyList(PyArray_DIMS(arr), PyArray_NDIM(arr)); + if (s < 0 || (memcpy(fp->defs[i].data, PyArray_DATA(arr), + s * PyArray_ITEMSIZE(arr))) == NULL) { + if ((PyObject *)arr != v) { + Py_DECREF(arr); + } + return -1; + } + if ((PyObject *)arr != v) { + Py_DECREF(arr); + } + } + else + return (fp->defs[i].func == NULL ? -1 : 0); + return 0; /* successful */ + } + if (fp->dict == NULL) { + fp->dict = PyDict_New(); + if (fp->dict == NULL) + return -1; + } + if (v == NULL) { + int rv = PyDict_DelItemString(fp->dict, name); + if (rv < 0) + PyErr_SetString(PyExc_AttributeError, + "delete non-existing fortran attribute"); + return rv; + } + else + return PyDict_SetItemString(fp->dict, name, v); +} + +static PyObject * +fortran_call(PyFortranObject *fp, PyObject *arg, PyObject *kw) +{ + int i = 0; + /* printf("fortran call + name=%s,func=%p,data=%p,%p\n",fp->defs[i].name, + fp->defs[i].func,fp->defs[i].data,&fp->defs[i].data); */ + if (fp->defs[i].rank == -1) { /* is Fortran routine */ + if (fp->defs[i].func == NULL) { + PyErr_Format(PyExc_RuntimeError, "no function to call"); + return NULL; + } + else if (fp->defs[i].data == NULL) + /* dummy routine */ + return (*((fortranfunc)(fp->defs[i].func)))((PyObject *)fp, arg, + kw, NULL); + else + return (*((fortranfunc)(fp->defs[i].func)))( + (PyObject *)fp, arg, kw, (void *)fp->defs[i].data); + } + PyErr_Format(PyExc_TypeError, "this fortran object is not callable"); + return NULL; +} + +static PyObject * +fortran_repr(PyFortranObject *fp) +{ + PyObject *name = NULL, *repr = NULL; + name = PyObject_GetAttrString((PyObject *)fp, "__name__"); + PyErr_Clear(); + if (name != NULL && PyUnicode_Check(name)) { + repr = PyUnicode_FromFormat("", name); + } + else { + repr = PyUnicode_FromString(""); + } + Py_XDECREF(name); + return repr; +} + +PyTypeObject PyFortran_Type = { + PyVarObject_HEAD_INIT(NULL, 0).tp_name = "fortran", + .tp_basicsize = sizeof(PyFortranObject), + .tp_dealloc = (destructor)fortran_dealloc, + .tp_getattr = (getattrfunc)fortran_getattr, + .tp_setattr = (setattrfunc)fortran_setattr, + .tp_repr = (reprfunc)fortran_repr, + .tp_call = (ternaryfunc)fortran_call, +}; + +/************************* f2py_report_atexit *******************************/ + +#ifdef F2PY_REPORT_ATEXIT +static int passed_time = 0; +static int passed_counter = 0; +static int passed_call_time = 0; +static struct timeb start_time; +static struct timeb stop_time; +static struct timeb start_call_time; +static struct timeb stop_call_time; +static int cb_passed_time = 0; +static int cb_passed_counter = 0; +static int cb_passed_call_time = 0; +static struct timeb cb_start_time; +static struct timeb cb_stop_time; +static struct timeb cb_start_call_time; +static struct timeb cb_stop_call_time; + +extern void +f2py_start_clock(void) +{ + ftime(&start_time); +} +extern void +f2py_start_call_clock(void) +{ + f2py_stop_clock(); + ftime(&start_call_time); +} +extern void +f2py_stop_clock(void) +{ + ftime(&stop_time); + passed_time += 1000 * (stop_time.time - start_time.time); + passed_time += stop_time.millitm - start_time.millitm; +} +extern void +f2py_stop_call_clock(void) +{ + ftime(&stop_call_time); + passed_call_time += 1000 * (stop_call_time.time - start_call_time.time); + passed_call_time += stop_call_time.millitm - start_call_time.millitm; + passed_counter += 1; + f2py_start_clock(); +} + +extern void +f2py_cb_start_clock(void) +{ + ftime(&cb_start_time); +} +extern void +f2py_cb_start_call_clock(void) +{ + f2py_cb_stop_clock(); + ftime(&cb_start_call_time); +} +extern void +f2py_cb_stop_clock(void) +{ + ftime(&cb_stop_time); + cb_passed_time += 1000 * (cb_stop_time.time - cb_start_time.time); + cb_passed_time += cb_stop_time.millitm - cb_start_time.millitm; +} +extern void +f2py_cb_stop_call_clock(void) +{ + ftime(&cb_stop_call_time); + cb_passed_call_time += + 1000 * (cb_stop_call_time.time - cb_start_call_time.time); + cb_passed_call_time += + cb_stop_call_time.millitm - cb_start_call_time.millitm; + cb_passed_counter += 1; + f2py_cb_start_clock(); +} + +static int f2py_report_on_exit_been_here = 0; +extern void +f2py_report_on_exit(int exit_flag, void *name) +{ + if (f2py_report_on_exit_been_here) { + fprintf(stderr, " %s\n", (char *)name); + return; + } + f2py_report_on_exit_been_here = 1; + fprintf(stderr, " /-----------------------\\\n"); + fprintf(stderr, " < F2PY performance report >\n"); + fprintf(stderr, " \\-----------------------/\n"); + fprintf(stderr, "Overall time spent in ...\n"); + fprintf(stderr, "(a) wrapped (Fortran/C) functions : %8d msec\n", + passed_call_time); + fprintf(stderr, "(b) f2py interface, %6d calls : %8d msec\n", + passed_counter, passed_time); + fprintf(stderr, "(c) call-back (Python) functions : %8d msec\n", + cb_passed_call_time); + fprintf(stderr, "(d) f2py call-back interface, %6d calls : %8d msec\n", + cb_passed_counter, cb_passed_time); + + fprintf(stderr, + "(e) wrapped (Fortran/C) functions (actual) : %8d msec\n\n", + passed_call_time - cb_passed_call_time - cb_passed_time); + fprintf(stderr, + "Use -DF2PY_REPORT_ATEXIT_DISABLE to disable this message.\n"); + fprintf(stderr, "Exit status: %d\n", exit_flag); + fprintf(stderr, "Modules : %s\n", (char *)name); +} +#endif + +/********************** report on array copy ****************************/ + +#ifdef F2PY_REPORT_ON_ARRAY_COPY +static void +f2py_report_on_array_copy(PyArrayObject *arr) +{ + const npy_intp arr_size = PyArray_Size((PyObject *)arr); + if (arr_size > F2PY_REPORT_ON_ARRAY_COPY) { + fprintf(stderr, + "copied an array: size=%ld, elsize=%" NPY_INTP_FMT "\n", + arr_size, (npy_intp)PyArray_ITEMSIZE(arr)); + } +} +static void +f2py_report_on_array_copy_fromany(void) +{ + fprintf(stderr, "created an array from object\n"); +} + +#define F2PY_REPORT_ON_ARRAY_COPY_FROMARR \ + f2py_report_on_array_copy((PyArrayObject *)arr) +#define F2PY_REPORT_ON_ARRAY_COPY_FROMANY f2py_report_on_array_copy_fromany() +#else +#define F2PY_REPORT_ON_ARRAY_COPY_FROMARR +#define F2PY_REPORT_ON_ARRAY_COPY_FROMANY +#endif + +/************************* array_from_obj *******************************/ + +/* + * File: array_from_pyobj.c + * + * Description: + * ------------ + * Provides array_from_pyobj function that returns a contiguous array + * object with the given dimensions and required storage order, either + * in row-major (C) or column-major (Fortran) order. The function + * array_from_pyobj is very flexible about its Python object argument + * that can be any number, list, tuple, or array. + * + * array_from_pyobj is used in f2py generated Python extension + * modules. + * + * Author: Pearu Peterson + * Created: 13-16 January 2002 + * $Id: fortranobject.c,v 1.52 2005/07/11 07:44:20 pearu Exp $ + */ + +static int check_and_fix_dimensions(const PyArrayObject* arr, + const int rank, + npy_intp *dims, + const char *errmess); + +static int +find_first_negative_dimension(const int rank, const npy_intp *dims) +{ + int i; + for (i = 0; i < rank; ++i) { + if (dims[i] < 0) { + return i; + } + } + return -1; +} + +#ifdef DEBUG_COPY_ND_ARRAY +void +dump_dims(int rank, npy_intp const *dims) +{ + int i; + printf("["); + for (i = 0; i < rank; ++i) { + printf("%3" NPY_INTP_FMT, dims[i]); + } + printf("]\n"); +} +void +dump_attrs(const PyArrayObject *obj) +{ + const PyArrayObject_fields *arr = (const PyArrayObject_fields *)obj; + int rank = PyArray_NDIM(arr); + npy_intp size = PyArray_Size((PyObject *)arr); + printf("\trank = %d, flags = %d, size = %" NPY_INTP_FMT "\n", rank, + arr->flags, size); + printf("\tstrides = "); + dump_dims(rank, arr->strides); + printf("\tdimensions = "); + dump_dims(rank, arr->dimensions); +} +#endif + +#define SWAPTYPE(a, b, t) \ + { \ + t c; \ + c = (a); \ + (a) = (b); \ + (b) = c; \ + } + +static int +swap_arrays(PyArrayObject *obj1, PyArrayObject *obj2) +{ + PyArrayObject_fields *arr1 = (PyArrayObject_fields *)obj1, + *arr2 = (PyArrayObject_fields *)obj2; + SWAPTYPE(arr1->data, arr2->data, char *); + SWAPTYPE(arr1->nd, arr2->nd, int); + SWAPTYPE(arr1->dimensions, arr2->dimensions, npy_intp *); + SWAPTYPE(arr1->strides, arr2->strides, npy_intp *); + SWAPTYPE(arr1->base, arr2->base, PyObject *); + SWAPTYPE(arr1->descr, arr2->descr, PyArray_Descr *); + SWAPTYPE(arr1->flags, arr2->flags, int); + /* SWAPTYPE(arr1->weakreflist,arr2->weakreflist,PyObject*); */ + return 0; +} + +#define ARRAY_ISCOMPATIBLE(arr,type_num) \ + ((PyArray_ISINTEGER(arr) && PyTypeNum_ISINTEGER(type_num)) || \ + (PyArray_ISFLOAT(arr) && PyTypeNum_ISFLOAT(type_num)) || \ + (PyArray_ISCOMPLEX(arr) && PyTypeNum_ISCOMPLEX(type_num)) || \ + (PyArray_ISBOOL(arr) && PyTypeNum_ISBOOL(type_num)) || \ + (PyArray_ISSTRING(arr) && PyTypeNum_ISSTRING(type_num))) + +static int +get_elsize(PyObject *obj) { + /* + get_elsize determines array itemsize from a Python object. Returns + elsize if successful, -1 otherwise. + + Supported types of the input are: numpy.ndarray, bytes, str, tuple, + list. + */ + + if (PyArray_Check(obj)) { + return PyArray_ITEMSIZE((PyArrayObject *)obj); + } else if (PyBytes_Check(obj)) { + return PyBytes_GET_SIZE(obj); + } else if (PyUnicode_Check(obj)) { + return PyUnicode_GET_LENGTH(obj); + } else if (PySequence_Check(obj)) { + PyObject* fast = PySequence_Fast(obj, "f2py:fortranobject.c:get_elsize"); + if (fast != NULL) { + Py_ssize_t i, n = PySequence_Fast_GET_SIZE(fast); + int sz, elsize = 0; + for (i=0; i elsize) { + elsize = sz; + } + } + Py_DECREF(fast); + return elsize; + } + } + return -1; +} + +extern PyArrayObject * +ndarray_from_pyobj(const int type_num, + const int elsize_, + npy_intp *dims, + const int rank, + const int intent, + PyObject *obj, + const char *errmess) { + /* + * Return an array with given element type and shape from a Python + * object while taking into account the usage intent of the array. + * + * - element type is defined by type_num and elsize + * - shape is defined by dims and rank + * + * ndarray_from_pyobj is used to convert Python object arguments + * to numpy ndarrays with given type and shape that data is passed + * to interfaced Fortran or C functions. + * + * errmess (if not NULL), contains a prefix of an error message + * for an exception to be triggered within this function. + * + * Negative elsize value means that elsize is to be determined + * from the Python object in runtime. + * + * Note on strings + * --------------- + * + * String type (type_num == NPY_STRING) does not have fixed + * element size and, by default, the type object sets it to + * 0. Therefore, for string types, one has to use elsize + * argument. For other types, elsize value is ignored. + * + * NumPy defines the type of a fixed-width string as + * dtype('S'). In addition, there is also dtype('c'), that + * appears as dtype('S1') (these have the same type_num value), + * but is actually different (.char attribute is either 'S' or + * 'c', respectively). + * + * In Fortran, character arrays and strings are different + * concepts. The relation between Fortran types, NumPy dtypes, + * and type_num-elsize pairs, is defined as follows: + * + * character*5 foo | dtype('S5') | elsize=5, shape=() + * character(5) foo | dtype('S1') | elsize=1, shape=(5) + * character*5 foo(n) | dtype('S5') | elsize=5, shape=(n,) + * character(5) foo(n) | dtype('S1') | elsize=1, shape=(5, n) + * character*(*) foo | dtype('S') | elsize=-1, shape=() + * + * Note about reference counting + * ----------------------------- + * + * If the caller returns the array to Python, it must be done with + * Py_BuildValue("N",arr). Otherwise, if obj!=arr then the caller + * must call Py_DECREF(arr). + * + * Note on intent(cache,out,..) + * ---------------------------- + * Don't expect correct data when returning intent(cache) array. + * + */ + char mess[F2PY_MESSAGE_BUFFER_SIZE]; + PyArrayObject *arr = NULL; + int elsize = (elsize_ < 0 ? get_elsize(obj) : elsize_); + if (elsize < 0) { + if (errmess != NULL) { + strcpy(mess, errmess); + } + sprintf(mess + strlen(mess), + " -- failed to determine element size from %s", + Py_TYPE(obj)->tp_name); + PyErr_SetString(PyExc_SystemError, mess); + return NULL; + } + PyArray_Descr * descr = get_descr_from_type_and_elsize(type_num, elsize); // new reference + if (descr == NULL) { + return NULL; + } + elsize = PyDataType_ELSIZE(descr); + if ((intent & F2PY_INTENT_HIDE) + || ((intent & F2PY_INTENT_CACHE) && (obj == Py_None)) + || ((intent & F2PY_OPTIONAL) && (obj == Py_None)) + ) { + /* intent(cache), optional, intent(hide) */ + int ineg = find_first_negative_dimension(rank, dims); + if (ineg >= 0) { + int i; + strcpy(mess, "failed to create intent(cache|hide)|optional array" + "-- must have defined dimensions but got ("); + for(i = 0; i < rank; ++i) + sprintf(mess + strlen(mess), "%" NPY_INTP_FMT ",", dims[i]); + strcat(mess, ")"); + PyErr_SetString(PyExc_ValueError, mess); + Py_DECREF(descr); + return NULL; + } + arr = (PyArrayObject *) \ + PyArray_NewFromDescr(&PyArray_Type, descr, rank, dims, + NULL, NULL, !(intent & F2PY_INTENT_C), NULL); + if (arr == NULL) { + Py_DECREF(descr); + return NULL; + } + if (PyArray_ITEMSIZE(arr) != elsize) { + strcpy(mess, "failed to create intent(cache|hide)|optional array"); + sprintf(mess+strlen(mess)," -- expected elsize=%d got %" NPY_INTP_FMT, elsize, (npy_intp)PyArray_ITEMSIZE(arr)); + PyErr_SetString(PyExc_ValueError,mess); + Py_DECREF(arr); + return NULL; + } + if (!(intent & F2PY_INTENT_CACHE)) { + PyArray_FILLWBYTE(arr, 0); + } + return arr; + } + + if (PyArray_Check(obj)) { + arr = (PyArrayObject *)obj; + if (intent & F2PY_INTENT_CACHE) { + /* intent(cache) */ + if (PyArray_ISONESEGMENT(arr) + && PyArray_ITEMSIZE(arr) >= elsize) { + if (check_and_fix_dimensions(arr, rank, dims, errmess)) { + Py_DECREF(descr); + return NULL; + } + if (intent & F2PY_INTENT_OUT) + Py_INCREF(arr); + Py_DECREF(descr); + return arr; + } + strcpy(mess, "failed to initialize intent(cache) array"); + if (!PyArray_ISONESEGMENT(arr)) + strcat(mess, " -- input must be in one segment"); + if (PyArray_ITEMSIZE(arr) < elsize) + sprintf(mess + strlen(mess), + " -- expected at least elsize=%d but got " + "%" NPY_INTP_FMT, + elsize, (npy_intp)PyArray_ITEMSIZE(arr)); + PyErr_SetString(PyExc_ValueError, mess); + Py_DECREF(descr); + return NULL; + } + + /* here we have always intent(in) or intent(inout) or intent(inplace) + */ + + if (check_and_fix_dimensions(arr, rank, dims, errmess)) { + Py_DECREF(descr); + return NULL; + } + /* + printf("intent alignment=%d\n", F2PY_GET_ALIGNMENT(intent)); + printf("alignment check=%d\n", F2PY_CHECK_ALIGNMENT(arr, intent)); + int i; + for (i=1;i<=16;i++) + printf("i=%d isaligned=%d\n", i, ARRAY_ISALIGNED(arr, i)); + */ + if ((! (intent & F2PY_INTENT_COPY)) && + PyArray_ITEMSIZE(arr) == elsize && + ARRAY_ISCOMPATIBLE(arr,type_num) && + F2PY_CHECK_ALIGNMENT(arr, intent)) { + if ((intent & F2PY_INTENT_INOUT || intent & F2PY_INTENT_INPLACE) + ? ((intent & F2PY_INTENT_C) ? PyArray_ISCARRAY(arr) : PyArray_ISFARRAY(arr)) + : ((intent & F2PY_INTENT_C) ? PyArray_ISCARRAY_RO(arr) : PyArray_ISFARRAY_RO(arr))) { + if ((intent & F2PY_INTENT_OUT)) { + Py_INCREF(arr); + } + /* Returning input array */ + Py_DECREF(descr); + return arr; + } + } + if (intent & F2PY_INTENT_INOUT) { + strcpy(mess, "failed to initialize intent(inout) array"); + /* Must use PyArray_IS*ARRAY because intent(inout) requires + * writable input */ + if ((intent & F2PY_INTENT_C) && !PyArray_ISCARRAY(arr)) + strcat(mess, " -- input not contiguous"); + if (!(intent & F2PY_INTENT_C) && !PyArray_ISFARRAY(arr)) + strcat(mess, " -- input not fortran contiguous"); + if (PyArray_ITEMSIZE(arr) != elsize) + sprintf(mess + strlen(mess), + " -- expected elsize=%d but got %" NPY_INTP_FMT, + elsize, + (npy_intp)PyArray_ITEMSIZE(arr) + ); + if (!(ARRAY_ISCOMPATIBLE(arr, type_num))) { + sprintf(mess + strlen(mess), + " -- input '%c' not compatible to '%c'", + PyArray_DESCR(arr)->type, descr->type); + } + if (!(F2PY_CHECK_ALIGNMENT(arr, intent))) + sprintf(mess + strlen(mess), " -- input not %d-aligned", + F2PY_GET_ALIGNMENT(intent)); + PyErr_SetString(PyExc_ValueError, mess); + Py_DECREF(descr); + return NULL; + } + + /* here we have always intent(in) or intent(inplace) */ + + { + PyArrayObject * retarr = (PyArrayObject *) \ + PyArray_NewFromDescr(&PyArray_Type, descr, PyArray_NDIM(arr), PyArray_DIMS(arr), + NULL, NULL, !(intent & F2PY_INTENT_C), NULL); + if (retarr==NULL) { + Py_DECREF(descr); + return NULL; + } + F2PY_REPORT_ON_ARRAY_COPY_FROMARR; + if (PyArray_CopyInto(retarr, arr)) { + Py_DECREF(retarr); + return NULL; + } + if (intent & F2PY_INTENT_INPLACE) { + if (swap_arrays(arr,retarr)) { + Py_DECREF(retarr); + return NULL; /* XXX: set exception */ + } + Py_XDECREF(retarr); + if (intent & F2PY_INTENT_OUT) + Py_INCREF(arr); + } else { + arr = retarr; + } + } + return arr; + } + + if ((intent & F2PY_INTENT_INOUT) || (intent & F2PY_INTENT_INPLACE) || + (intent & F2PY_INTENT_CACHE)) { + PyErr_Format(PyExc_TypeError, + "failed to initialize intent(inout|inplace|cache) " + "array, input '%s' object is not an array", + Py_TYPE(obj)->tp_name); + Py_DECREF(descr); + return NULL; + } + + { + F2PY_REPORT_ON_ARRAY_COPY_FROMANY; + arr = (PyArrayObject *)PyArray_FromAny( + obj, descr, 0, 0, + ((intent & F2PY_INTENT_C) ? NPY_ARRAY_CARRAY + : NPY_ARRAY_FARRAY) | + NPY_ARRAY_FORCECAST, + NULL); + // Warning: in the case of NPY_STRING, PyArray_FromAny may + // reset descr->elsize, e.g. dtype('S0') becomes dtype('S1'). + if (arr == NULL) { + Py_DECREF(descr); + return NULL; + } + if (type_num != NPY_STRING && PyArray_ITEMSIZE(arr) != elsize) { + // This is internal sanity tests: elsize has been set to + // descr->elsize in the beginning of this function. + strcpy(mess, "failed to initialize intent(in) array"); + sprintf(mess + strlen(mess), + " -- expected elsize=%d got %" NPY_INTP_FMT, elsize, + (npy_intp)PyArray_ITEMSIZE(arr)); + PyErr_SetString(PyExc_ValueError, mess); + Py_DECREF(arr); + return NULL; + } + if (check_and_fix_dimensions(arr, rank, dims, errmess)) { + Py_DECREF(arr); + return NULL; + } + return arr; + } +} + +extern PyArrayObject * +array_from_pyobj(const int type_num, + npy_intp *dims, + const int rank, + const int intent, + PyObject *obj) { + /* + Same as ndarray_from_pyobj but with elsize determined from type, + if possible. Provided for backward compatibility. + */ + PyArray_Descr* descr = PyArray_DescrFromType(type_num); + int elsize = PyDataType_ELSIZE(descr); + Py_DECREF(descr); + return ndarray_from_pyobj(type_num, elsize, dims, rank, intent, obj, NULL); +} + +/*****************************************/ +/* Helper functions for array_from_pyobj */ +/*****************************************/ + +static int +check_and_fix_dimensions(const PyArrayObject* arr, const int rank, + npy_intp *dims, const char *errmess) +{ + /* + * This function fills in blanks (that are -1's) in dims list using + * the dimensions from arr. It also checks that non-blank dims will + * match with the corresponding values in arr dimensions. + * + * Returns 0 if the function is successful. + * + * If an error condition is detected, an exception is set and 1 is + * returned. + */ + char mess[F2PY_MESSAGE_BUFFER_SIZE]; + const npy_intp arr_size = + (PyArray_NDIM(arr)) ? PyArray_Size((PyObject *)arr) : 1; +#ifdef DEBUG_COPY_ND_ARRAY + dump_attrs(arr); + printf("check_and_fix_dimensions:init: dims="); + dump_dims(rank, dims); +#endif + if (rank > PyArray_NDIM(arr)) { /* [1,2] -> [[1],[2]]; 1 -> [[1]] */ + npy_intp new_size = 1; + int free_axe = -1; + int i; + npy_intp d; + /* Fill dims where -1 or 0; check dimensions; calc new_size; */ + for (i = 0; i < PyArray_NDIM(arr); ++i) { + d = PyArray_DIM(arr, i); + if (dims[i] >= 0) { + if (d > 1 && dims[i] != d) { + PyErr_Format( + PyExc_ValueError, + "%d-th dimension must be fixed to %" NPY_INTP_FMT + " but got %" NPY_INTP_FMT "\n", + i, dims[i], d); + return 1; + } + if (!dims[i]) + dims[i] = 1; + } + else { + dims[i] = d ? d : 1; + } + new_size *= dims[i]; + } + for (i = PyArray_NDIM(arr); i < rank; ++i) + if (dims[i] > 1) { + PyErr_Format(PyExc_ValueError, + "%d-th dimension must be %" NPY_INTP_FMT + " but got 0 (not defined).\n", + i, dims[i]); + return 1; + } + else if (free_axe < 0) + free_axe = i; + else + dims[i] = 1; + if (free_axe >= 0) { + dims[free_axe] = arr_size / new_size; + new_size *= dims[free_axe]; + } + if (new_size != arr_size) { + PyErr_Format(PyExc_ValueError, + "unexpected array size: new_size=%" NPY_INTP_FMT + ", got array with arr_size=%" NPY_INTP_FMT + " (maybe too many free indices)\n", + new_size, arr_size); + return 1; + } + } + else if (rank == PyArray_NDIM(arr)) { + npy_intp new_size = 1; + int i; + npy_intp d; + for (i = 0; i < rank; ++i) { + d = PyArray_DIM(arr, i); + if (dims[i] >= 0) { + if (d > 1 && d != dims[i]) { + if (errmess != NULL) { + strcpy(mess, errmess); + } + sprintf(mess + strlen(mess), + " -- %d-th dimension must be fixed to %" + NPY_INTP_FMT " but got %" NPY_INTP_FMT, + i, dims[i], d); + PyErr_SetString(PyExc_ValueError, mess); + return 1; + } + if (!dims[i]) + dims[i] = 1; + } + else + dims[i] = d; + new_size *= dims[i]; + } + if (new_size != arr_size) { + PyErr_Format(PyExc_ValueError, + "unexpected array size: new_size=%" NPY_INTP_FMT + ", got array with arr_size=%" NPY_INTP_FMT "\n", + new_size, arr_size); + return 1; + } + } + else { /* [[1,2]] -> [[1],[2]] */ + int i, j; + npy_intp d; + int effrank; + npy_intp size; + for (i = 0, effrank = 0; i < PyArray_NDIM(arr); ++i) + if (PyArray_DIM(arr, i) > 1) + ++effrank; + if (dims[rank - 1] >= 0) + if (effrank > rank) { + PyErr_Format(PyExc_ValueError, + "too many axes: %d (effrank=%d), " + "expected rank=%d\n", + PyArray_NDIM(arr), effrank, rank); + return 1; + } + + for (i = 0, j = 0; i < rank; ++i) { + while (j < PyArray_NDIM(arr) && PyArray_DIM(arr, j) < 2) ++j; + if (j >= PyArray_NDIM(arr)) + d = 1; + else + d = PyArray_DIM(arr, j++); + if (dims[i] >= 0) { + if (d > 1 && d != dims[i]) { + if (errmess != NULL) { + strcpy(mess, errmess); + } + sprintf(mess + strlen(mess), + " -- %d-th dimension must be fixed to %" + NPY_INTP_FMT " but got %" NPY_INTP_FMT + " (real index=%d)\n", + i, dims[i], d, j-1); + PyErr_SetString(PyExc_ValueError, mess); + return 1; + } + if (!dims[i]) + dims[i] = 1; + } + else + dims[i] = d; + } + + for (i = rank; i < PyArray_NDIM(arr); + ++i) { /* [[1,2],[3,4]] -> [1,2,3,4] */ + while (j < PyArray_NDIM(arr) && PyArray_DIM(arr, j) < 2) ++j; + if (j >= PyArray_NDIM(arr)) + d = 1; + else + d = PyArray_DIM(arr, j++); + dims[rank - 1] *= d; + } + for (i = 0, size = 1; i < rank; ++i) size *= dims[i]; + if (size != arr_size) { + char msg[200]; + int len; + snprintf(msg, sizeof(msg), + "unexpected array size: size=%" NPY_INTP_FMT + ", arr_size=%" NPY_INTP_FMT + ", rank=%d, effrank=%d, arr.nd=%d, dims=[", + size, arr_size, rank, effrank, PyArray_NDIM(arr)); + for (i = 0; i < rank; ++i) { + len = strlen(msg); + snprintf(msg + len, sizeof(msg) - len, " %" NPY_INTP_FMT, + dims[i]); + } + len = strlen(msg); + snprintf(msg + len, sizeof(msg) - len, " ], arr.dims=["); + for (i = 0; i < PyArray_NDIM(arr); ++i) { + len = strlen(msg); + snprintf(msg + len, sizeof(msg) - len, " %" NPY_INTP_FMT, + PyArray_DIM(arr, i)); + } + len = strlen(msg); + snprintf(msg + len, sizeof(msg) - len, " ]\n"); + PyErr_SetString(PyExc_ValueError, msg); + return 1; + } + } +#ifdef DEBUG_COPY_ND_ARRAY + printf("check_and_fix_dimensions:end: dims="); + dump_dims(rank, dims); +#endif + return 0; +} + +/* End of file: array_from_pyobj.c */ + +/************************* copy_ND_array *******************************/ + +extern int +copy_ND_array(const PyArrayObject *arr, PyArrayObject *out) +{ + F2PY_REPORT_ON_ARRAY_COPY_FROMARR; + return PyArray_CopyInto(out, (PyArrayObject *)arr); +} + +/********************* Various utility functions ***********************/ + +extern int +f2py_describe(PyObject *obj, char *buf) { + /* + Write the description of a Python object to buf. The caller must + provide buffer with size sufficient to write the description. + + Return 1 on success. + */ + char localbuf[F2PY_MESSAGE_BUFFER_SIZE]; + if (PyBytes_Check(obj)) { + sprintf(localbuf, "%d-%s", (npy_int)PyBytes_GET_SIZE(obj), Py_TYPE(obj)->tp_name); + } else if (PyUnicode_Check(obj)) { + sprintf(localbuf, "%d-%s", (npy_int)PyUnicode_GET_LENGTH(obj), Py_TYPE(obj)->tp_name); + } else if (PyArray_CheckScalar(obj)) { + PyArrayObject* arr = (PyArrayObject*)obj; + sprintf(localbuf, "%c%" NPY_INTP_FMT "-%s-scalar", PyArray_DESCR(arr)->kind, PyArray_ITEMSIZE(arr), Py_TYPE(obj)->tp_name); + } else if (PyArray_Check(obj)) { + int i; + PyArrayObject* arr = (PyArrayObject*)obj; + strcpy(localbuf, "("); + for (i=0; ikind, PyArray_ITEMSIZE(arr), Py_TYPE(obj)->tp_name); + } else if (PySequence_Check(obj)) { + sprintf(localbuf, "%d-%s", (npy_int)PySequence_Length(obj), Py_TYPE(obj)->tp_name); + } else { + sprintf(localbuf, "%s instance", Py_TYPE(obj)->tp_name); + } + // TODO: detect the size of buf and make sure that size(buf) >= size(localbuf). + strcpy(buf, localbuf); + return 1; +} + +extern npy_intp +f2py_size_impl(PyArrayObject* var, ...) +{ + npy_intp sz = 0; + npy_intp dim; + npy_intp rank; + va_list argp; + va_start(argp, var); + dim = va_arg(argp, npy_int); + if (dim==-1) + { + sz = PyArray_SIZE(var); + } + else + { + rank = PyArray_NDIM(var); + if (dim>=1 && dim<=rank) + sz = PyArray_DIM(var, dim-1); + else + fprintf(stderr, "f2py_size: 2nd argument value=%" NPY_INTP_FMT + " fails to satisfy 1<=value<=%" NPY_INTP_FMT + ". Result will be 0.\n", dim, rank); + } + va_end(argp); + return sz; +} + +/*********************************************/ +/* Compatibility functions for Python >= 3.0 */ +/*********************************************/ + +PyObject * +F2PyCapsule_FromVoidPtr(void *ptr, void (*dtor)(PyObject *)) +{ + PyObject *ret = PyCapsule_New(ptr, NULL, dtor); + if (ret == NULL) { + PyErr_Clear(); + } + return ret; +} + +void * +F2PyCapsule_AsVoidPtr(PyObject *obj) +{ + void *ret = PyCapsule_GetPointer(obj, NULL); + if (ret == NULL) { + PyErr_Clear(); + } + return ret; +} + +int +F2PyCapsule_Check(PyObject *ptr) +{ + return PyCapsule_CheckExact(ptr); +} + +#ifdef __cplusplus +} +#endif +/************************* EOF fortranobject.c *******************************/ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.h b/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.h new file mode 100644 index 0000000000000000000000000000000000000000..4aed2f60891b9e9ea2c16373b5959b0a346e7470 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/src/fortranobject.h @@ -0,0 +1,173 @@ +#ifndef Py_FORTRANOBJECT_H +#define Py_FORTRANOBJECT_H +#ifdef __cplusplus +extern "C" { +#endif + +#include + +#ifndef NPY_NO_DEPRECATED_API +#define NPY_NO_DEPRECATED_API NPY_API_VERSION +#endif +#ifdef FORTRANOBJECT_C +#define NO_IMPORT_ARRAY +#endif +#define PY_ARRAY_UNIQUE_SYMBOL _npy_f2py_ARRAY_API +#include "numpy/arrayobject.h" +#include "numpy/npy_3kcompat.h" + +#ifdef F2PY_REPORT_ATEXIT +#include +// clang-format off +extern void f2py_start_clock(void); +extern void f2py_stop_clock(void); +extern void f2py_start_call_clock(void); +extern void f2py_stop_call_clock(void); +extern void f2py_cb_start_clock(void); +extern void f2py_cb_stop_clock(void); +extern void f2py_cb_start_call_clock(void); +extern void f2py_cb_stop_call_clock(void); +extern void f2py_report_on_exit(int, void *); +// clang-format on +#endif + +#ifdef DMALLOC +#include "dmalloc.h" +#endif + +/* Fortran object interface */ + +/* +123456789-123456789-123456789-123456789-123456789-123456789-123456789-12 + +PyFortranObject represents various Fortran objects: +Fortran (module) routines, COMMON blocks, module data. + +Author: Pearu Peterson +*/ + +#define F2PY_MAX_DIMS 40 +#define F2PY_MESSAGE_BUFFER_SIZE 300 // Increase on "stack smashing detected" + +typedef void (*f2py_set_data_func)(char *, npy_intp *); +typedef void (*f2py_void_func)(void); +typedef void (*f2py_init_func)(int *, npy_intp *, f2py_set_data_func, int *); + +/*typedef void* (*f2py_c_func)(void*,...);*/ + +typedef void *(*f2pycfunc)(void); + +typedef struct { + char *name; /* attribute (array||routine) name */ + int rank; /* array rank, 0 for scalar, max is F2PY_MAX_DIMS, + || rank=-1 for Fortran routine */ + struct { + npy_intp d[F2PY_MAX_DIMS]; + } dims; /* dimensions of the array, || not used */ + int type; /* PyArray_ || not used */ + int elsize; /* Element size || not used */ + char *data; /* pointer to array || Fortran routine */ + f2py_init_func func; /* initialization function for + allocatable arrays: + func(&rank,dims,set_ptr_func,name,len(name)) + || C/API wrapper for Fortran routine */ + char *doc; /* documentation string; only recommended + for routines. */ +} FortranDataDef; + +typedef struct { + PyObject_HEAD + int len; /* Number of attributes */ + FortranDataDef *defs; /* An array of FortranDataDef's */ + PyObject *dict; /* Fortran object attribute dictionary */ +} PyFortranObject; + +#define PyFortran_Check(op) (Py_TYPE(op) == &PyFortran_Type) +#define PyFortran_Check1(op) (0 == strcmp(Py_TYPE(op)->tp_name, "fortran")) + +extern PyTypeObject PyFortran_Type; +extern int +F2PyDict_SetItemString(PyObject *dict, char *name, PyObject *obj); +extern PyObject * +PyFortranObject_New(FortranDataDef *defs, f2py_void_func init); +extern PyObject * +PyFortranObject_NewAsAttr(FortranDataDef *defs); + +PyObject * +F2PyCapsule_FromVoidPtr(void *ptr, void (*dtor)(PyObject *)); +void * +F2PyCapsule_AsVoidPtr(PyObject *obj); +int +F2PyCapsule_Check(PyObject *ptr); + +extern void * +F2PySwapThreadLocalCallbackPtr(char *key, void *ptr); +extern void * +F2PyGetThreadLocalCallbackPtr(char *key); + +#define ISCONTIGUOUS(m) (PyArray_FLAGS(m) & NPY_ARRAY_C_CONTIGUOUS) +#define F2PY_INTENT_IN 1 +#define F2PY_INTENT_INOUT 2 +#define F2PY_INTENT_OUT 4 +#define F2PY_INTENT_HIDE 8 +#define F2PY_INTENT_CACHE 16 +#define F2PY_INTENT_COPY 32 +#define F2PY_INTENT_C 64 +#define F2PY_OPTIONAL 128 +#define F2PY_INTENT_INPLACE 256 +#define F2PY_INTENT_ALIGNED4 512 +#define F2PY_INTENT_ALIGNED8 1024 +#define F2PY_INTENT_ALIGNED16 2048 + +#define ARRAY_ISALIGNED(ARR, SIZE) ((size_t)(PyArray_DATA(ARR)) % (SIZE) == 0) +#define F2PY_ALIGN4(intent) (intent & F2PY_INTENT_ALIGNED4) +#define F2PY_ALIGN8(intent) (intent & F2PY_INTENT_ALIGNED8) +#define F2PY_ALIGN16(intent) (intent & F2PY_INTENT_ALIGNED16) + +#define F2PY_GET_ALIGNMENT(intent) \ + (F2PY_ALIGN4(intent) \ + ? 4 \ + : (F2PY_ALIGN8(intent) ? 8 : (F2PY_ALIGN16(intent) ? 16 : 1))) +#define F2PY_CHECK_ALIGNMENT(arr, intent) \ + ARRAY_ISALIGNED(arr, F2PY_GET_ALIGNMENT(intent)) +#define F2PY_ARRAY_IS_CHARACTER_COMPATIBLE(arr) ((PyArray_DESCR(arr)->type_num == NPY_STRING && PyArray_ITEMSIZE(arr) >= 1) \ + || PyArray_DESCR(arr)->type_num == NPY_UINT8) +#define F2PY_IS_UNICODE_ARRAY(arr) (PyArray_DESCR(arr)->type_num == NPY_UNICODE) + +extern PyArrayObject * +ndarray_from_pyobj(const int type_num, const int elsize_, npy_intp *dims, + const int rank, const int intent, PyObject *obj, + const char *errmess); + +extern PyArrayObject * +array_from_pyobj(const int type_num, npy_intp *dims, const int rank, + const int intent, PyObject *obj); +extern int +copy_ND_array(const PyArrayObject *in, PyArrayObject *out); + +#ifdef DEBUG_COPY_ND_ARRAY +extern void +dump_attrs(const PyArrayObject *arr); +#endif + + extern int f2py_describe(PyObject *obj, char *buf); + + /* Utility CPP macros and functions that can be used in signature file + expressions. See signature-file.rst for documentation. + */ + +#define f2py_itemsize(var) (PyArray_ITEMSIZE(capi_ ## var ## _as_array)) +#define f2py_size(var, ...) f2py_size_impl((PyArrayObject *)(capi_ ## var ## _as_array), ## __VA_ARGS__, -1) +#define f2py_rank(var) var ## _Rank +#define f2py_shape(var,dim) var ## _Dims[dim] +#define f2py_len(var) f2py_shape(var,0) +#define f2py_fshape(var,dim) f2py_shape(var,rank(var)-dim-1) +#define f2py_flen(var) f2py_fshape(var,0) +#define f2py_slen(var) capi_ ## var ## _len + + extern npy_intp f2py_size_impl(PyArrayObject* var, ...); + +#ifdef __cplusplus +} +#endif +#endif /* !Py_FORTRANOBJECT_H */ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/symbolic.py b/venv/lib/python3.10/site-packages/numpy/f2py/symbolic.py new file mode 100644 index 0000000000000000000000000000000000000000..63d277d9b01d487175bbb078ce7a4bfadd2ad917 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/symbolic.py @@ -0,0 +1,1517 @@ +"""Fortran/C symbolic expressions + +References: +- J3/21-007: Draft Fortran 202x. https://j3-fortran.org/doc/year/21/21-007.pdf + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" + +# To analyze Fortran expressions to solve dimensions specifications, +# for instances, we implement a minimal symbolic engine for parsing +# expressions into a tree of expression instances. As a first +# instance, we care only about arithmetic expressions involving +# integers and operations like addition (+), subtraction (-), +# multiplication (*), division (Fortran / is Python //, Fortran // is +# concatenate), and exponentiation (**). In addition, .pyf files may +# contain C expressions that support here is implemented as well. +# +# TODO: support logical constants (Op.BOOLEAN) +# TODO: support logical operators (.AND., ...) +# TODO: support defined operators (.MYOP., ...) +# +__all__ = ['Expr'] + + +import re +import warnings +from enum import Enum +from math import gcd + + +class Language(Enum): + """ + Used as Expr.tostring language argument. + """ + Python = 0 + Fortran = 1 + C = 2 + + +class Op(Enum): + """ + Used as Expr op attribute. + """ + INTEGER = 10 + REAL = 12 + COMPLEX = 15 + STRING = 20 + ARRAY = 30 + SYMBOL = 40 + TERNARY = 100 + APPLY = 200 + INDEXING = 210 + CONCAT = 220 + RELATIONAL = 300 + TERMS = 1000 + FACTORS = 2000 + REF = 3000 + DEREF = 3001 + + +class RelOp(Enum): + """ + Used in Op.RELATIONAL expression to specify the function part. + """ + EQ = 1 + NE = 2 + LT = 3 + LE = 4 + GT = 5 + GE = 6 + + @classmethod + def fromstring(cls, s, language=Language.C): + if language is Language.Fortran: + return {'.eq.': RelOp.EQ, '.ne.': RelOp.NE, + '.lt.': RelOp.LT, '.le.': RelOp.LE, + '.gt.': RelOp.GT, '.ge.': RelOp.GE}[s.lower()] + return {'==': RelOp.EQ, '!=': RelOp.NE, '<': RelOp.LT, + '<=': RelOp.LE, '>': RelOp.GT, '>=': RelOp.GE}[s] + + def tostring(self, language=Language.C): + if language is Language.Fortran: + return {RelOp.EQ: '.eq.', RelOp.NE: '.ne.', + RelOp.LT: '.lt.', RelOp.LE: '.le.', + RelOp.GT: '.gt.', RelOp.GE: '.ge.'}[self] + return {RelOp.EQ: '==', RelOp.NE: '!=', + RelOp.LT: '<', RelOp.LE: '<=', + RelOp.GT: '>', RelOp.GE: '>='}[self] + + +class ArithOp(Enum): + """ + Used in Op.APPLY expression to specify the function part. + """ + POS = 1 + NEG = 2 + ADD = 3 + SUB = 4 + MUL = 5 + DIV = 6 + POW = 7 + + +class OpError(Exception): + pass + + +class Precedence(Enum): + """ + Used as Expr.tostring precedence argument. + """ + ATOM = 0 + POWER = 1 + UNARY = 2 + PRODUCT = 3 + SUM = 4 + LT = 6 + EQ = 7 + LAND = 11 + LOR = 12 + TERNARY = 13 + ASSIGN = 14 + TUPLE = 15 + NONE = 100 + + +integer_types = (int,) +number_types = (int, float) + + +def _pairs_add(d, k, v): + # Internal utility method for updating terms and factors data. + c = d.get(k) + if c is None: + d[k] = v + else: + c = c + v + if c: + d[k] = c + else: + del d[k] + + +class ExprWarning(UserWarning): + pass + + +def ewarn(message): + warnings.warn(message, ExprWarning, stacklevel=2) + + +class Expr: + """Represents a Fortran expression as a op-data pair. + + Expr instances are hashable and sortable. + """ + + @staticmethod + def parse(s, language=Language.C): + """Parse a Fortran expression to a Expr. + """ + return fromstring(s, language=language) + + def __init__(self, op, data): + assert isinstance(op, Op) + + # sanity checks + if op is Op.INTEGER: + # data is a 2-tuple of numeric object and a kind value + # (default is 4) + assert isinstance(data, tuple) and len(data) == 2 + assert isinstance(data[0], int) + assert isinstance(data[1], (int, str)), data + elif op is Op.REAL: + # data is a 2-tuple of numeric object and a kind value + # (default is 4) + assert isinstance(data, tuple) and len(data) == 2 + assert isinstance(data[0], float) + assert isinstance(data[1], (int, str)), data + elif op is Op.COMPLEX: + # data is a 2-tuple of constant expressions + assert isinstance(data, tuple) and len(data) == 2 + elif op is Op.STRING: + # data is a 2-tuple of quoted string and a kind value + # (default is 1) + assert isinstance(data, tuple) and len(data) == 2 + assert (isinstance(data[0], str) + and data[0][::len(data[0])-1] in ('""', "''", '@@')) + assert isinstance(data[1], (int, str)), data + elif op is Op.SYMBOL: + # data is any hashable object + assert hash(data) is not None + elif op in (Op.ARRAY, Op.CONCAT): + # data is a tuple of expressions + assert isinstance(data, tuple) + assert all(isinstance(item, Expr) for item in data), data + elif op in (Op.TERMS, Op.FACTORS): + # data is {:} where dict values + # are nonzero Python integers + assert isinstance(data, dict) + elif op is Op.APPLY: + # data is (, , ) where + # operands are Expr instances + assert isinstance(data, tuple) and len(data) == 3 + # function is any hashable object + assert hash(data[0]) is not None + assert isinstance(data[1], tuple) + assert isinstance(data[2], dict) + elif op is Op.INDEXING: + # data is (, ) + assert isinstance(data, tuple) and len(data) == 2 + # function is any hashable object + assert hash(data[0]) is not None + elif op is Op.TERNARY: + # data is (, , ) + assert isinstance(data, tuple) and len(data) == 3 + elif op in (Op.REF, Op.DEREF): + # data is Expr instance + assert isinstance(data, Expr) + elif op is Op.RELATIONAL: + # data is (, , ) + assert isinstance(data, tuple) and len(data) == 3 + else: + raise NotImplementedError( + f'unknown op or missing sanity check: {op}') + + self.op = op + self.data = data + + def __eq__(self, other): + return (isinstance(other, Expr) + and self.op is other.op + and self.data == other.data) + + def __hash__(self): + if self.op in (Op.TERMS, Op.FACTORS): + data = tuple(sorted(self.data.items())) + elif self.op is Op.APPLY: + data = self.data[:2] + tuple(sorted(self.data[2].items())) + else: + data = self.data + return hash((self.op, data)) + + def __lt__(self, other): + if isinstance(other, Expr): + if self.op is not other.op: + return self.op.value < other.op.value + if self.op in (Op.TERMS, Op.FACTORS): + return (tuple(sorted(self.data.items())) + < tuple(sorted(other.data.items()))) + if self.op is Op.APPLY: + if self.data[:2] != other.data[:2]: + return self.data[:2] < other.data[:2] + return tuple(sorted(self.data[2].items())) < tuple( + sorted(other.data[2].items())) + return self.data < other.data + return NotImplemented + + def __le__(self, other): return self == other or self < other + + def __gt__(self, other): return not (self <= other) + + def __ge__(self, other): return not (self < other) + + def __repr__(self): + return f'{type(self).__name__}({self.op}, {self.data!r})' + + def __str__(self): + return self.tostring() + + def tostring(self, parent_precedence=Precedence.NONE, + language=Language.Fortran): + """Return a string representation of Expr. + """ + if self.op in (Op.INTEGER, Op.REAL): + precedence = (Precedence.SUM if self.data[0] < 0 + else Precedence.ATOM) + r = str(self.data[0]) + (f'_{self.data[1]}' + if self.data[1] != 4 else '') + elif self.op is Op.COMPLEX: + r = ', '.join(item.tostring(Precedence.TUPLE, language=language) + for item in self.data) + r = '(' + r + ')' + precedence = Precedence.ATOM + elif self.op is Op.SYMBOL: + precedence = Precedence.ATOM + r = str(self.data) + elif self.op is Op.STRING: + r = self.data[0] + if self.data[1] != 1: + r = self.data[1] + '_' + r + precedence = Precedence.ATOM + elif self.op is Op.ARRAY: + r = ', '.join(item.tostring(Precedence.TUPLE, language=language) + for item in self.data) + r = '[' + r + ']' + precedence = Precedence.ATOM + elif self.op is Op.TERMS: + terms = [] + for term, coeff in sorted(self.data.items()): + if coeff < 0: + op = ' - ' + coeff = -coeff + else: + op = ' + ' + if coeff == 1: + term = term.tostring(Precedence.SUM, language=language) + else: + if term == as_number(1): + term = str(coeff) + else: + term = f'{coeff} * ' + term.tostring( + Precedence.PRODUCT, language=language) + if terms: + terms.append(op) + elif op == ' - ': + terms.append('-') + terms.append(term) + r = ''.join(terms) or '0' + precedence = Precedence.SUM if terms else Precedence.ATOM + elif self.op is Op.FACTORS: + factors = [] + tail = [] + for base, exp in sorted(self.data.items()): + op = ' * ' + if exp == 1: + factor = base.tostring(Precedence.PRODUCT, + language=language) + elif language is Language.C: + if exp in range(2, 10): + factor = base.tostring(Precedence.PRODUCT, + language=language) + factor = ' * '.join([factor] * exp) + elif exp in range(-10, 0): + factor = base.tostring(Precedence.PRODUCT, + language=language) + tail += [factor] * -exp + continue + else: + factor = base.tostring(Precedence.TUPLE, + language=language) + factor = f'pow({factor}, {exp})' + else: + factor = base.tostring(Precedence.POWER, + language=language) + f' ** {exp}' + if factors: + factors.append(op) + factors.append(factor) + if tail: + if not factors: + factors += ['1'] + factors += ['/', '(', ' * '.join(tail), ')'] + r = ''.join(factors) or '1' + precedence = Precedence.PRODUCT if factors else Precedence.ATOM + elif self.op is Op.APPLY: + name, args, kwargs = self.data + if name is ArithOp.DIV and language is Language.C: + numer, denom = [arg.tostring(Precedence.PRODUCT, + language=language) + for arg in args] + r = f'{numer} / {denom}' + precedence = Precedence.PRODUCT + else: + args = [arg.tostring(Precedence.TUPLE, language=language) + for arg in args] + args += [k + '=' + v.tostring(Precedence.NONE) + for k, v in kwargs.items()] + r = f'{name}({", ".join(args)})' + precedence = Precedence.ATOM + elif self.op is Op.INDEXING: + name = self.data[0] + args = [arg.tostring(Precedence.TUPLE, language=language) + for arg in self.data[1:]] + r = f'{name}[{", ".join(args)}]' + precedence = Precedence.ATOM + elif self.op is Op.CONCAT: + args = [arg.tostring(Precedence.PRODUCT, language=language) + for arg in self.data] + r = " // ".join(args) + precedence = Precedence.PRODUCT + elif self.op is Op.TERNARY: + cond, expr1, expr2 = [a.tostring(Precedence.TUPLE, + language=language) + for a in self.data] + if language is Language.C: + r = f'({cond}?{expr1}:{expr2})' + elif language is Language.Python: + r = f'({expr1} if {cond} else {expr2})' + elif language is Language.Fortran: + r = f'merge({expr1}, {expr2}, {cond})' + else: + raise NotImplementedError( + f'tostring for {self.op} and {language}') + precedence = Precedence.ATOM + elif self.op is Op.REF: + r = '&' + self.data.tostring(Precedence.UNARY, language=language) + precedence = Precedence.UNARY + elif self.op is Op.DEREF: + r = '*' + self.data.tostring(Precedence.UNARY, language=language) + precedence = Precedence.UNARY + elif self.op is Op.RELATIONAL: + rop, left, right = self.data + precedence = (Precedence.EQ if rop in (RelOp.EQ, RelOp.NE) + else Precedence.LT) + left = left.tostring(precedence, language=language) + right = right.tostring(precedence, language=language) + rop = rop.tostring(language=language) + r = f'{left} {rop} {right}' + else: + raise NotImplementedError(f'tostring for op {self.op}') + if parent_precedence.value < precedence.value: + # If parent precedence is higher than operand precedence, + # operand will be enclosed in parenthesis. + return '(' + r + ')' + return r + + def __pos__(self): + return self + + def __neg__(self): + return self * -1 + + def __add__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + if self.op is other.op: + if self.op in (Op.INTEGER, Op.REAL): + return as_number( + self.data[0] + other.data[0], + max(self.data[1], other.data[1])) + if self.op is Op.COMPLEX: + r1, i1 = self.data + r2, i2 = other.data + return as_complex(r1 + r2, i1 + i2) + if self.op is Op.TERMS: + r = Expr(self.op, dict(self.data)) + for k, v in other.data.items(): + _pairs_add(r.data, k, v) + return normalize(r) + if self.op is Op.COMPLEX and other.op in (Op.INTEGER, Op.REAL): + return self + as_complex(other) + elif self.op in (Op.INTEGER, Op.REAL) and other.op is Op.COMPLEX: + return as_complex(self) + other + elif self.op is Op.REAL and other.op is Op.INTEGER: + return self + as_real(other, kind=self.data[1]) + elif self.op is Op.INTEGER and other.op is Op.REAL: + return as_real(self, kind=other.data[1]) + other + return as_terms(self) + as_terms(other) + return NotImplemented + + def __radd__(self, other): + if isinstance(other, number_types): + return as_number(other) + self + return NotImplemented + + def __sub__(self, other): + return self + (-other) + + def __rsub__(self, other): + if isinstance(other, number_types): + return as_number(other) - self + return NotImplemented + + def __mul__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + if self.op is other.op: + if self.op in (Op.INTEGER, Op.REAL): + return as_number(self.data[0] * other.data[0], + max(self.data[1], other.data[1])) + elif self.op is Op.COMPLEX: + r1, i1 = self.data + r2, i2 = other.data + return as_complex(r1 * r2 - i1 * i2, r1 * i2 + r2 * i1) + + if self.op is Op.FACTORS: + r = Expr(self.op, dict(self.data)) + for k, v in other.data.items(): + _pairs_add(r.data, k, v) + return normalize(r) + elif self.op is Op.TERMS: + r = Expr(self.op, {}) + for t1, c1 in self.data.items(): + for t2, c2 in other.data.items(): + _pairs_add(r.data, t1 * t2, c1 * c2) + return normalize(r) + + if self.op is Op.COMPLEX and other.op in (Op.INTEGER, Op.REAL): + return self * as_complex(other) + elif other.op is Op.COMPLEX and self.op in (Op.INTEGER, Op.REAL): + return as_complex(self) * other + elif self.op is Op.REAL and other.op is Op.INTEGER: + return self * as_real(other, kind=self.data[1]) + elif self.op is Op.INTEGER and other.op is Op.REAL: + return as_real(self, kind=other.data[1]) * other + + if self.op is Op.TERMS: + return self * as_terms(other) + elif other.op is Op.TERMS: + return as_terms(self) * other + + return as_factors(self) * as_factors(other) + return NotImplemented + + def __rmul__(self, other): + if isinstance(other, number_types): + return as_number(other) * self + return NotImplemented + + def __pow__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + if other.op is Op.INTEGER: + exponent = other.data[0] + # TODO: other kind not used + if exponent == 0: + return as_number(1) + if exponent == 1: + return self + if exponent > 0: + if self.op is Op.FACTORS: + r = Expr(self.op, {}) + for k, v in self.data.items(): + r.data[k] = v * exponent + return normalize(r) + return self * (self ** (exponent - 1)) + elif exponent != -1: + return (self ** (-exponent)) ** -1 + return Expr(Op.FACTORS, {self: exponent}) + return as_apply(ArithOp.POW, self, other) + return NotImplemented + + def __truediv__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + # Fortran / is different from Python /: + # - `/` is a truncate operation for integer operands + return normalize(as_apply(ArithOp.DIV, self, other)) + return NotImplemented + + def __rtruediv__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + return other / self + return NotImplemented + + def __floordiv__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + # Fortran // is different from Python //: + # - `//` is a concatenate operation for string operands + return normalize(Expr(Op.CONCAT, (self, other))) + return NotImplemented + + def __rfloordiv__(self, other): + other = as_expr(other) + if isinstance(other, Expr): + return other // self + return NotImplemented + + def __call__(self, *args, **kwargs): + # In Fortran, parenthesis () are use for both function call as + # well as indexing operations. + # + # TODO: implement a method for deciding when __call__ should + # return an INDEXING expression. + return as_apply(self, *map(as_expr, args), + **dict((k, as_expr(v)) for k, v in kwargs.items())) + + def __getitem__(self, index): + # Provided to support C indexing operations that .pyf files + # may contain. + index = as_expr(index) + if not isinstance(index, tuple): + index = index, + if len(index) > 1: + ewarn(f'C-index should be a single expression but got `{index}`') + return Expr(Op.INDEXING, (self,) + index) + + def substitute(self, symbols_map): + """Recursively substitute symbols with values in symbols map. + + Symbols map is a dictionary of symbol-expression pairs. + """ + if self.op is Op.SYMBOL: + value = symbols_map.get(self) + if value is None: + return self + m = re.match(r'\A(@__f2py_PARENTHESIS_(\w+)_\d+@)\Z', self.data) + if m: + # complement to fromstring method + items, paren = m.groups() + if paren in ['ROUNDDIV', 'SQUARE']: + return as_array(value) + assert paren == 'ROUND', (paren, value) + return value + if self.op in (Op.INTEGER, Op.REAL, Op.STRING): + return self + if self.op in (Op.ARRAY, Op.COMPLEX): + return Expr(self.op, tuple(item.substitute(symbols_map) + for item in self.data)) + if self.op is Op.CONCAT: + return normalize(Expr(self.op, tuple(item.substitute(symbols_map) + for item in self.data))) + if self.op is Op.TERMS: + r = None + for term, coeff in self.data.items(): + if r is None: + r = term.substitute(symbols_map) * coeff + else: + r += term.substitute(symbols_map) * coeff + if r is None: + ewarn('substitute: empty TERMS expression interpreted as' + ' int-literal 0') + return as_number(0) + return r + if self.op is Op.FACTORS: + r = None + for base, exponent in self.data.items(): + if r is None: + r = base.substitute(symbols_map) ** exponent + else: + r *= base.substitute(symbols_map) ** exponent + if r is None: + ewarn('substitute: empty FACTORS expression interpreted' + ' as int-literal 1') + return as_number(1) + return r + if self.op is Op.APPLY: + target, args, kwargs = self.data + if isinstance(target, Expr): + target = target.substitute(symbols_map) + args = tuple(a.substitute(symbols_map) for a in args) + kwargs = dict((k, v.substitute(symbols_map)) + for k, v in kwargs.items()) + return normalize(Expr(self.op, (target, args, kwargs))) + if self.op is Op.INDEXING: + func = self.data[0] + if isinstance(func, Expr): + func = func.substitute(symbols_map) + args = tuple(a.substitute(symbols_map) for a in self.data[1:]) + return normalize(Expr(self.op, (func,) + args)) + if self.op is Op.TERNARY: + operands = tuple(a.substitute(symbols_map) for a in self.data) + return normalize(Expr(self.op, operands)) + if self.op in (Op.REF, Op.DEREF): + return normalize(Expr(self.op, self.data.substitute(symbols_map))) + if self.op is Op.RELATIONAL: + rop, left, right = self.data + left = left.substitute(symbols_map) + right = right.substitute(symbols_map) + return normalize(Expr(self.op, (rop, left, right))) + raise NotImplementedError(f'substitute method for {self.op}: {self!r}') + + def traverse(self, visit, *args, **kwargs): + """Traverse expression tree with visit function. + + The visit function is applied to an expression with given args + and kwargs. + + Traverse call returns an expression returned by visit when not + None, otherwise return a new normalized expression with + traverse-visit sub-expressions. + """ + result = visit(self, *args, **kwargs) + if result is not None: + return result + + if self.op in (Op.INTEGER, Op.REAL, Op.STRING, Op.SYMBOL): + return self + elif self.op in (Op.COMPLEX, Op.ARRAY, Op.CONCAT, Op.TERNARY): + return normalize(Expr(self.op, tuple( + item.traverse(visit, *args, **kwargs) + for item in self.data))) + elif self.op in (Op.TERMS, Op.FACTORS): + data = {} + for k, v in self.data.items(): + k = k.traverse(visit, *args, **kwargs) + v = (v.traverse(visit, *args, **kwargs) + if isinstance(v, Expr) else v) + if k in data: + v = data[k] + v + data[k] = v + return normalize(Expr(self.op, data)) + elif self.op is Op.APPLY: + obj = self.data[0] + func = (obj.traverse(visit, *args, **kwargs) + if isinstance(obj, Expr) else obj) + operands = tuple(operand.traverse(visit, *args, **kwargs) + for operand in self.data[1]) + kwoperands = dict((k, v.traverse(visit, *args, **kwargs)) + for k, v in self.data[2].items()) + return normalize(Expr(self.op, (func, operands, kwoperands))) + elif self.op is Op.INDEXING: + obj = self.data[0] + obj = (obj.traverse(visit, *args, **kwargs) + if isinstance(obj, Expr) else obj) + indices = tuple(index.traverse(visit, *args, **kwargs) + for index in self.data[1:]) + return normalize(Expr(self.op, (obj,) + indices)) + elif self.op in (Op.REF, Op.DEREF): + return normalize(Expr(self.op, + self.data.traverse(visit, *args, **kwargs))) + elif self.op is Op.RELATIONAL: + rop, left, right = self.data + left = left.traverse(visit, *args, **kwargs) + right = right.traverse(visit, *args, **kwargs) + return normalize(Expr(self.op, (rop, left, right))) + raise NotImplementedError(f'traverse method for {self.op}') + + def contains(self, other): + """Check if self contains other. + """ + found = [] + + def visit(expr, found=found): + if found: + return expr + elif expr == other: + found.append(1) + return expr + + self.traverse(visit) + + return len(found) != 0 + + def symbols(self): + """Return a set of symbols contained in self. + """ + found = set() + + def visit(expr, found=found): + if expr.op is Op.SYMBOL: + found.add(expr) + + self.traverse(visit) + + return found + + def polynomial_atoms(self): + """Return a set of expressions used as atoms in polynomial self. + """ + found = set() + + def visit(expr, found=found): + if expr.op is Op.FACTORS: + for b in expr.data: + b.traverse(visit) + return expr + if expr.op in (Op.TERMS, Op.COMPLEX): + return + if expr.op is Op.APPLY and isinstance(expr.data[0], ArithOp): + if expr.data[0] is ArithOp.POW: + expr.data[1][0].traverse(visit) + return expr + return + if expr.op in (Op.INTEGER, Op.REAL): + return expr + + found.add(expr) + + if expr.op in (Op.INDEXING, Op.APPLY): + return expr + + self.traverse(visit) + + return found + + def linear_solve(self, symbol): + """Return a, b such that a * symbol + b == self. + + If self is not linear with respect to symbol, raise RuntimeError. + """ + b = self.substitute({symbol: as_number(0)}) + ax = self - b + a = ax.substitute({symbol: as_number(1)}) + + zero, _ = as_numer_denom(a * symbol - ax) + + if zero != as_number(0): + raise RuntimeError(f'not a {symbol}-linear equation:' + f' {a} * {symbol} + {b} == {self}') + return a, b + + +def normalize(obj): + """Normalize Expr and apply basic evaluation methods. + """ + if not isinstance(obj, Expr): + return obj + + if obj.op is Op.TERMS: + d = {} + for t, c in obj.data.items(): + if c == 0: + continue + if t.op is Op.COMPLEX and c != 1: + t = t * c + c = 1 + if t.op is Op.TERMS: + for t1, c1 in t.data.items(): + _pairs_add(d, t1, c1 * c) + else: + _pairs_add(d, t, c) + if len(d) == 0: + # TODO: determine correct kind + return as_number(0) + elif len(d) == 1: + (t, c), = d.items() + if c == 1: + return t + return Expr(Op.TERMS, d) + + if obj.op is Op.FACTORS: + coeff = 1 + d = {} + for b, e in obj.data.items(): + if e == 0: + continue + if b.op is Op.TERMS and isinstance(e, integer_types) and e > 1: + # expand integer powers of sums + b = b * (b ** (e - 1)) + e = 1 + + if b.op in (Op.INTEGER, Op.REAL): + if e == 1: + coeff *= b.data[0] + elif e > 0: + coeff *= b.data[0] ** e + else: + _pairs_add(d, b, e) + elif b.op is Op.FACTORS: + if e > 0 and isinstance(e, integer_types): + for b1, e1 in b.data.items(): + _pairs_add(d, b1, e1 * e) + else: + _pairs_add(d, b, e) + else: + _pairs_add(d, b, e) + if len(d) == 0 or coeff == 0: + # TODO: determine correct kind + assert isinstance(coeff, number_types) + return as_number(coeff) + elif len(d) == 1: + (b, e), = d.items() + if e == 1: + t = b + else: + t = Expr(Op.FACTORS, d) + if coeff == 1: + return t + return Expr(Op.TERMS, {t: coeff}) + elif coeff == 1: + return Expr(Op.FACTORS, d) + else: + return Expr(Op.TERMS, {Expr(Op.FACTORS, d): coeff}) + + if obj.op is Op.APPLY and obj.data[0] is ArithOp.DIV: + dividend, divisor = obj.data[1] + t1, c1 = as_term_coeff(dividend) + t2, c2 = as_term_coeff(divisor) + if isinstance(c1, integer_types) and isinstance(c2, integer_types): + g = gcd(c1, c2) + c1, c2 = c1//g, c2//g + else: + c1, c2 = c1/c2, 1 + + if t1.op is Op.APPLY and t1.data[0] is ArithOp.DIV: + numer = t1.data[1][0] * c1 + denom = t1.data[1][1] * t2 * c2 + return as_apply(ArithOp.DIV, numer, denom) + + if t2.op is Op.APPLY and t2.data[0] is ArithOp.DIV: + numer = t2.data[1][1] * t1 * c1 + denom = t2.data[1][0] * c2 + return as_apply(ArithOp.DIV, numer, denom) + + d = dict(as_factors(t1).data) + for b, e in as_factors(t2).data.items(): + _pairs_add(d, b, -e) + numer, denom = {}, {} + for b, e in d.items(): + if e > 0: + numer[b] = e + else: + denom[b] = -e + numer = normalize(Expr(Op.FACTORS, numer)) * c1 + denom = normalize(Expr(Op.FACTORS, denom)) * c2 + + if denom.op in (Op.INTEGER, Op.REAL) and denom.data[0] == 1: + # TODO: denom kind not used + return numer + return as_apply(ArithOp.DIV, numer, denom) + + if obj.op is Op.CONCAT: + lst = [obj.data[0]] + for s in obj.data[1:]: + last = lst[-1] + if ( + last.op is Op.STRING + and s.op is Op.STRING + and last.data[0][0] in '"\'' + and s.data[0][0] == last.data[0][-1] + ): + new_last = as_string(last.data[0][:-1] + s.data[0][1:], + max(last.data[1], s.data[1])) + lst[-1] = new_last + else: + lst.append(s) + if len(lst) == 1: + return lst[0] + return Expr(Op.CONCAT, tuple(lst)) + + if obj.op is Op.TERNARY: + cond, expr1, expr2 = map(normalize, obj.data) + if cond.op is Op.INTEGER: + return expr1 if cond.data[0] else expr2 + return Expr(Op.TERNARY, (cond, expr1, expr2)) + + return obj + + +def as_expr(obj): + """Convert non-Expr objects to Expr objects. + """ + if isinstance(obj, complex): + return as_complex(obj.real, obj.imag) + if isinstance(obj, number_types): + return as_number(obj) + if isinstance(obj, str): + # STRING expression holds string with boundary quotes, hence + # applying repr: + return as_string(repr(obj)) + if isinstance(obj, tuple): + return tuple(map(as_expr, obj)) + return obj + + +def as_symbol(obj): + """Return object as SYMBOL expression (variable or unparsed expression). + """ + return Expr(Op.SYMBOL, obj) + + +def as_number(obj, kind=4): + """Return object as INTEGER or REAL constant. + """ + if isinstance(obj, int): + return Expr(Op.INTEGER, (obj, kind)) + if isinstance(obj, float): + return Expr(Op.REAL, (obj, kind)) + if isinstance(obj, Expr): + if obj.op in (Op.INTEGER, Op.REAL): + return obj + raise OpError(f'cannot convert {obj} to INTEGER or REAL constant') + + +def as_integer(obj, kind=4): + """Return object as INTEGER constant. + """ + if isinstance(obj, int): + return Expr(Op.INTEGER, (obj, kind)) + if isinstance(obj, Expr): + if obj.op is Op.INTEGER: + return obj + raise OpError(f'cannot convert {obj} to INTEGER constant') + + +def as_real(obj, kind=4): + """Return object as REAL constant. + """ + if isinstance(obj, int): + return Expr(Op.REAL, (float(obj), kind)) + if isinstance(obj, float): + return Expr(Op.REAL, (obj, kind)) + if isinstance(obj, Expr): + if obj.op is Op.REAL: + return obj + elif obj.op is Op.INTEGER: + return Expr(Op.REAL, (float(obj.data[0]), kind)) + raise OpError(f'cannot convert {obj} to REAL constant') + + +def as_string(obj, kind=1): + """Return object as STRING expression (string literal constant). + """ + return Expr(Op.STRING, (obj, kind)) + + +def as_array(obj): + """Return object as ARRAY expression (array constant). + """ + if isinstance(obj, Expr): + obj = obj, + return Expr(Op.ARRAY, obj) + + +def as_complex(real, imag=0): + """Return object as COMPLEX expression (complex literal constant). + """ + return Expr(Op.COMPLEX, (as_expr(real), as_expr(imag))) + + +def as_apply(func, *args, **kwargs): + """Return object as APPLY expression (function call, constructor, etc.) + """ + return Expr(Op.APPLY, + (func, tuple(map(as_expr, args)), + dict((k, as_expr(v)) for k, v in kwargs.items()))) + + +def as_ternary(cond, expr1, expr2): + """Return object as TERNARY expression (cond?expr1:expr2). + """ + return Expr(Op.TERNARY, (cond, expr1, expr2)) + + +def as_ref(expr): + """Return object as referencing expression. + """ + return Expr(Op.REF, expr) + + +def as_deref(expr): + """Return object as dereferencing expression. + """ + return Expr(Op.DEREF, expr) + + +def as_eq(left, right): + return Expr(Op.RELATIONAL, (RelOp.EQ, left, right)) + + +def as_ne(left, right): + return Expr(Op.RELATIONAL, (RelOp.NE, left, right)) + + +def as_lt(left, right): + return Expr(Op.RELATIONAL, (RelOp.LT, left, right)) + + +def as_le(left, right): + return Expr(Op.RELATIONAL, (RelOp.LE, left, right)) + + +def as_gt(left, right): + return Expr(Op.RELATIONAL, (RelOp.GT, left, right)) + + +def as_ge(left, right): + return Expr(Op.RELATIONAL, (RelOp.GE, left, right)) + + +def as_terms(obj): + """Return expression as TERMS expression. + """ + if isinstance(obj, Expr): + obj = normalize(obj) + if obj.op is Op.TERMS: + return obj + if obj.op is Op.INTEGER: + return Expr(Op.TERMS, {as_integer(1, obj.data[1]): obj.data[0]}) + if obj.op is Op.REAL: + return Expr(Op.TERMS, {as_real(1, obj.data[1]): obj.data[0]}) + return Expr(Op.TERMS, {obj: 1}) + raise OpError(f'cannot convert {type(obj)} to terms Expr') + + +def as_factors(obj): + """Return expression as FACTORS expression. + """ + if isinstance(obj, Expr): + obj = normalize(obj) + if obj.op is Op.FACTORS: + return obj + if obj.op is Op.TERMS: + if len(obj.data) == 1: + (term, coeff), = obj.data.items() + if coeff == 1: + return Expr(Op.FACTORS, {term: 1}) + return Expr(Op.FACTORS, {term: 1, Expr.number(coeff): 1}) + if (obj.op is Op.APPLY + and obj.data[0] is ArithOp.DIV + and not obj.data[2]): + return Expr(Op.FACTORS, {obj.data[1][0]: 1, obj.data[1][1]: -1}) + return Expr(Op.FACTORS, {obj: 1}) + raise OpError(f'cannot convert {type(obj)} to terms Expr') + + +def as_term_coeff(obj): + """Return expression as term-coefficient pair. + """ + if isinstance(obj, Expr): + obj = normalize(obj) + if obj.op is Op.INTEGER: + return as_integer(1, obj.data[1]), obj.data[0] + if obj.op is Op.REAL: + return as_real(1, obj.data[1]), obj.data[0] + if obj.op is Op.TERMS: + if len(obj.data) == 1: + (term, coeff), = obj.data.items() + return term, coeff + # TODO: find common divisor of coefficients + if obj.op is Op.APPLY and obj.data[0] is ArithOp.DIV: + t, c = as_term_coeff(obj.data[1][0]) + return as_apply(ArithOp.DIV, t, obj.data[1][1]), c + return obj, 1 + raise OpError(f'cannot convert {type(obj)} to term and coeff') + + +def as_numer_denom(obj): + """Return expression as numer-denom pair. + """ + if isinstance(obj, Expr): + obj = normalize(obj) + if obj.op in (Op.INTEGER, Op.REAL, Op.COMPLEX, Op.SYMBOL, + Op.INDEXING, Op.TERNARY): + return obj, as_number(1) + elif obj.op is Op.APPLY: + if obj.data[0] is ArithOp.DIV and not obj.data[2]: + numers, denoms = map(as_numer_denom, obj.data[1]) + return numers[0] * denoms[1], numers[1] * denoms[0] + return obj, as_number(1) + elif obj.op is Op.TERMS: + numers, denoms = [], [] + for term, coeff in obj.data.items(): + n, d = as_numer_denom(term) + n = n * coeff + numers.append(n) + denoms.append(d) + numer, denom = as_number(0), as_number(1) + for i in range(len(numers)): + n = numers[i] + for j in range(len(numers)): + if i != j: + n *= denoms[j] + numer += n + denom *= denoms[i] + if denom.op in (Op.INTEGER, Op.REAL) and denom.data[0] < 0: + numer, denom = -numer, -denom + return numer, denom + elif obj.op is Op.FACTORS: + numer, denom = as_number(1), as_number(1) + for b, e in obj.data.items(): + bnumer, bdenom = as_numer_denom(b) + if e > 0: + numer *= bnumer ** e + denom *= bdenom ** e + elif e < 0: + numer *= bdenom ** (-e) + denom *= bnumer ** (-e) + return numer, denom + raise OpError(f'cannot convert {type(obj)} to numer and denom') + + +def _counter(): + # Used internally to generate unique dummy symbols + counter = 0 + while True: + counter += 1 + yield counter + + +COUNTER = _counter() + + +def eliminate_quotes(s): + """Replace quoted substrings of input string. + + Return a new string and a mapping of replacements. + """ + d = {} + + def repl(m): + kind, value = m.groups()[:2] + if kind: + # remove trailing underscore + kind = kind[:-1] + p = {"'": "SINGLE", '"': "DOUBLE"}[value[0]] + k = f'{kind}@__f2py_QUOTES_{p}_{COUNTER.__next__()}@' + d[k] = value + return k + + new_s = re.sub(r'({kind}_|)({single_quoted}|{double_quoted})'.format( + kind=r'\w[\w\d_]*', + single_quoted=r"('([^'\\]|(\\.))*')", + double_quoted=r'("([^"\\]|(\\.))*")'), + repl, s) + + assert '"' not in new_s + assert "'" not in new_s + + return new_s, d + + +def insert_quotes(s, d): + """Inverse of eliminate_quotes. + """ + for k, v in d.items(): + kind = k[:k.find('@')] + if kind: + kind += '_' + s = s.replace(k, kind + v) + return s + + +def replace_parenthesis(s): + """Replace substrings of input that are enclosed in parenthesis. + + Return a new string and a mapping of replacements. + """ + # Find a parenthesis pair that appears first. + + # Fortran deliminator are `(`, `)`, `[`, `]`, `(/', '/)`, `/`. + # We don't handle `/` deliminator because it is not a part of an + # expression. + left, right = None, None + mn_i = len(s) + for left_, right_ in (('(/', '/)'), + '()', + '{}', # to support C literal structs + '[]'): + i = s.find(left_) + if i == -1: + continue + if i < mn_i: + mn_i = i + left, right = left_, right_ + + if left is None: + return s, {} + + i = mn_i + j = s.find(right, i) + + while s.count(left, i + 1, j) != s.count(right, i + 1, j): + j = s.find(right, j + 1) + if j == -1: + raise ValueError(f'Mismatch of {left+right} parenthesis in {s!r}') + + p = {'(': 'ROUND', '[': 'SQUARE', '{': 'CURLY', '(/': 'ROUNDDIV'}[left] + + k = f'@__f2py_PARENTHESIS_{p}_{COUNTER.__next__()}@' + v = s[i+len(left):j] + r, d = replace_parenthesis(s[j+len(right):]) + d[k] = v + return s[:i] + k + r, d + + +def _get_parenthesis_kind(s): + assert s.startswith('@__f2py_PARENTHESIS_'), s + return s.split('_')[4] + + +def unreplace_parenthesis(s, d): + """Inverse of replace_parenthesis. + """ + for k, v in d.items(): + p = _get_parenthesis_kind(k) + left = dict(ROUND='(', SQUARE='[', CURLY='{', ROUNDDIV='(/')[p] + right = dict(ROUND=')', SQUARE=']', CURLY='}', ROUNDDIV='/)')[p] + s = s.replace(k, left + v + right) + return s + + +def fromstring(s, language=Language.C): + """Create an expression from a string. + + This is a "lazy" parser, that is, only arithmetic operations are + resolved, non-arithmetic operations are treated as symbols. + """ + r = _FromStringWorker(language=language).parse(s) + if isinstance(r, Expr): + return r + raise ValueError(f'failed to parse `{s}` to Expr instance: got `{r}`') + + +class _Pair: + # Internal class to represent a pair of expressions + + def __init__(self, left, right): + self.left = left + self.right = right + + def substitute(self, symbols_map): + left, right = self.left, self.right + if isinstance(left, Expr): + left = left.substitute(symbols_map) + if isinstance(right, Expr): + right = right.substitute(symbols_map) + return _Pair(left, right) + + def __repr__(self): + return f'{type(self).__name__}({self.left}, {self.right})' + + +class _FromStringWorker: + + def __init__(self, language=Language.C): + self.original = None + self.quotes_map = None + self.language = language + + def finalize_string(self, s): + return insert_quotes(s, self.quotes_map) + + def parse(self, inp): + self.original = inp + unquoted, self.quotes_map = eliminate_quotes(inp) + return self.process(unquoted) + + def process(self, s, context='expr'): + """Parse string within the given context. + + The context may define the result in case of ambiguous + expressions. For instance, consider expressions `f(x, y)` and + `(x, y) + (a, b)` where `f` is a function and pair `(x, y)` + denotes complex number. Specifying context as "args" or + "expr", the subexpression `(x, y)` will be parse to an + argument list or to a complex number, respectively. + """ + if isinstance(s, (list, tuple)): + return type(s)(self.process(s_, context) for s_ in s) + + assert isinstance(s, str), (type(s), s) + + # replace subexpressions in parenthesis with f2py @-names + r, raw_symbols_map = replace_parenthesis(s) + r = r.strip() + + def restore(r): + # restores subexpressions marked with f2py @-names + if isinstance(r, (list, tuple)): + return type(r)(map(restore, r)) + return unreplace_parenthesis(r, raw_symbols_map) + + # comma-separated tuple + if ',' in r: + operands = restore(r.split(',')) + if context == 'args': + return tuple(self.process(operands)) + if context == 'expr': + if len(operands) == 2: + # complex number literal + return as_complex(*self.process(operands)) + raise NotImplementedError( + f'parsing comma-separated list (context={context}): {r}') + + # ternary operation + m = re.match(r'\A([^?]+)[?]([^:]+)[:](.+)\Z', r) + if m: + assert context == 'expr', context + oper, expr1, expr2 = restore(m.groups()) + oper = self.process(oper) + expr1 = self.process(expr1) + expr2 = self.process(expr2) + return as_ternary(oper, expr1, expr2) + + # relational expression + if self.language is Language.Fortran: + m = re.match( + r'\A(.+)\s*[.](eq|ne|lt|le|gt|ge)[.]\s*(.+)\Z', r, re.I) + else: + m = re.match( + r'\A(.+)\s*([=][=]|[!][=]|[<][=]|[<]|[>][=]|[>])\s*(.+)\Z', r) + if m: + left, rop, right = m.groups() + if self.language is Language.Fortran: + rop = '.' + rop + '.' + left, right = self.process(restore((left, right))) + rop = RelOp.fromstring(rop, language=self.language) + return Expr(Op.RELATIONAL, (rop, left, right)) + + # keyword argument + m = re.match(r'\A(\w[\w\d_]*)\s*[=](.*)\Z', r) + if m: + keyname, value = m.groups() + value = restore(value) + return _Pair(keyname, self.process(value)) + + # addition/subtraction operations + operands = re.split(r'((? 1: + result = self.process(restore(operands[0] or '0')) + for op, operand in zip(operands[1::2], operands[2::2]): + operand = self.process(restore(operand)) + op = op.strip() + if op == '+': + result += operand + else: + assert op == '-' + result -= operand + return result + + # string concatenate operation + if self.language is Language.Fortran and '//' in r: + operands = restore(r.split('//')) + return Expr(Op.CONCAT, + tuple(self.process(operands))) + + # multiplication/division operations + operands = re.split(r'(?<=[@\w\d_])\s*([*]|/)', + (r if self.language is Language.C + else r.replace('**', '@__f2py_DOUBLE_STAR@'))) + if len(operands) > 1: + operands = restore(operands) + if self.language is not Language.C: + operands = [operand.replace('@__f2py_DOUBLE_STAR@', '**') + for operand in operands] + # Expression is an arithmetic product + result = self.process(operands[0]) + for op, operand in zip(operands[1::2], operands[2::2]): + operand = self.process(operand) + op = op.strip() + if op == '*': + result *= operand + else: + assert op == '/' + result /= operand + return result + + # referencing/dereferencing + if r.startswith(('*', '&')): + op = {'*': Op.DEREF, '&': Op.REF}[r[0]] + operand = self.process(restore(r[1:])) + return Expr(op, operand) + + # exponentiation operations + if self.language is not Language.C and '**' in r: + operands = list(reversed(restore(r.split('**')))) + result = self.process(operands[0]) + for operand in operands[1:]: + operand = self.process(operand) + result = operand ** result + return result + + # int-literal-constant + m = re.match(r'\A({digit_string})({kind}|)\Z'.format( + digit_string=r'\d+', + kind=r'_(\d+|\w[\w\d_]*)'), r) + if m: + value, _, kind = m.groups() + if kind and kind.isdigit(): + kind = int(kind) + return as_integer(int(value), kind or 4) + + # real-literal-constant + m = re.match(r'\A({significant}({exponent}|)|\d+{exponent})({kind}|)\Z' + .format( + significant=r'[.]\d+|\d+[.]\d*', + exponent=r'[edED][+-]?\d+', + kind=r'_(\d+|\w[\w\d_]*)'), r) + if m: + value, _, _, kind = m.groups() + if kind and kind.isdigit(): + kind = int(kind) + value = value.lower() + if 'd' in value: + return as_real(float(value.replace('d', 'e')), kind or 8) + return as_real(float(value), kind or 4) + + # string-literal-constant with kind parameter specification + if r in self.quotes_map: + kind = r[:r.find('@')] + return as_string(self.quotes_map[r], kind or 1) + + # array constructor or literal complex constant or + # parenthesized expression + if r in raw_symbols_map: + paren = _get_parenthesis_kind(r) + items = self.process(restore(raw_symbols_map[r]), + 'expr' if paren == 'ROUND' else 'args') + if paren == 'ROUND': + if isinstance(items, Expr): + return items + if paren in ['ROUNDDIV', 'SQUARE']: + # Expression is a array constructor + if isinstance(items, Expr): + items = (items,) + return as_array(items) + + # function call/indexing + m = re.match(r'\A(.+)\s*(@__f2py_PARENTHESIS_(ROUND|SQUARE)_\d+@)\Z', + r) + if m: + target, args, paren = m.groups() + target = self.process(restore(target)) + args = self.process(restore(args)[1:-1], 'args') + if not isinstance(args, tuple): + args = args, + if paren == 'ROUND': + kwargs = dict((a.left, a.right) for a in args + if isinstance(a, _Pair)) + args = tuple(a for a in args if not isinstance(a, _Pair)) + # Warning: this could also be Fortran indexing operation.. + return as_apply(target, *args, **kwargs) + else: + # Expression is a C/Python indexing operation + # (e.g. used in .pyf files) + assert paren == 'SQUARE' + return target[args] + + # Fortran standard conforming identifier + m = re.match(r'\A\w[\w\d_]*\Z', r) + if m: + return as_symbol(r) + + # fall-back to symbol + r = self.finalize_string(restore(r)) + ewarn( + f'fromstring: treating {r!r} as symbol (original={self.original})') + return as_symbol(r) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__init__.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ecb68077b943bf401a3ef268656a67c094078ea --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__init__.py @@ -0,0 +1,15 @@ +from numpy.testing import IS_WASM, IS_EDITABLE +import pytest + +if IS_WASM: + pytest.skip( + "WASM/Pyodide does not use or support Fortran", + allow_module_level=True + ) + + +if IS_EDITABLE: + pytest.skip( + "Editable install doesn't support tests with a compile step", + allow_module_level=True + ) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..670924da632e736bc67a062c22e6c47fdb22410b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_abstract_interface.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_abstract_interface.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aff35948fd113115616c87e47172d23d633ac254 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_abstract_interface.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_array_from_pyobj.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_array_from_pyobj.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e99f345cb5e1b9afe139ef9da16d6c7d8dce7ae Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_array_from_pyobj.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_assumed_shape.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_assumed_shape.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1eba69afd920242b4e611f5fdec689e8e5fa313a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_assumed_shape.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_block_docstring.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_block_docstring.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..228de200de278bac601ac375c452ba843199eac7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_block_docstring.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_callback.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_callback.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da77ca12ae449926e5500455268067afd2f4ed75 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_callback.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_character.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_character.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..742f9305ef61dbd2fb6995bfea0e03030ea31eb2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_character.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_common.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_common.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf44a263079dc006369ecbb98e6bc238334785ce Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_common.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_crackfortran.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_crackfortran.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be1f9a339120e56f2f79d312ad86e1185f591d08 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_crackfortran.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_data.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_data.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd15e8a45b672785c0271eb0fe4ec104ff4bc8a1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_data.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_docs.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_docs.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e38a991d51d5d193206be0df25bb3757de94a6d0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_docs.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2cmap.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2cmap.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0da11804b0460114010e5ebd547d78607f034372 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2cmap.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2py2e.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2py2e.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de01f89230713a40fc4970b6964f3bfed784e4d7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_f2py2e.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_isoc.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_isoc.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef8dd9dcd7ecb308b1267a052306ec427b1213ed Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_isoc.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_kind.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_kind.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6894fd800f8aaede6fe7fa63c9fc6f80451b447 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_kind.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_mixed.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_mixed.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2982bfd6a0a9ea3800254d6de9c21ee16620b7e8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_mixed.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_modules.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_modules.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b5e037529f0543071276574afb11b476dc4453a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_modules.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_parameter.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_parameter.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b70874e409cb20502466b88cdfe67c2bb262929a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_parameter.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_pyf_src.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_pyf_src.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e3dc415a85f10ff9495ade4966e7eab7e1422650 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_pyf_src.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_quoted_character.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_quoted_character.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9813738376d0a74699c4cf607e3a022e66cd6ec8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_quoted_character.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_regression.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_regression.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca42a683505728a5f4641f87a65a545ec25faa85 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_regression.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_character.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_character.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbece4b55c92709fab2c33866673da5b9ab2a885 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_character.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_complex.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_complex.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0213da49bbee2a8064b2baf36d6a9e3c768205d7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_complex.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_integer.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_integer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4dee98c326e7a7c42581b38b02f84f61d20426bf Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_integer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_logical.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_logical.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71875b92bb3d0bb7e790ba53df1bc8160ab15f44 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_logical.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_real.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_real.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d52a3f8f498cca4f23d59793bd63f68460dbdaa Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_return_real.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_routines.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_routines.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53ae0d245e823914881e7644dfd6555db49c431b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_routines.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_semicolon_split.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_semicolon_split.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..542605df1991425ee80aece438cc704dc5b20e02 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_semicolon_split.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_size.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_size.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36271aa0820655683df26580faab2f4fb401c60a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_size.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_string.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_string.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16e20e81df62d505dc73cee213a36f243ba2d3e1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_string.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_symbolic.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_symbolic.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64b6b7e43b1c1cb4ae2a24df94470de91be1d66d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_symbolic.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_value_attrspec.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_value_attrspec.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b796863db6bca91863fa5001028c958c10a2c6f5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/test_value_attrspec.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/util.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/util.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a0cb08ce84b990147b54d66986535cc078223e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/f2py/tests/__pycache__/util.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/foo.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/foo.f90 new file mode 100644 index 0000000000000000000000000000000000000000..76d16aae2b57160228f41c00128ac0067eaf5249 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/foo.f90 @@ -0,0 +1,34 @@ +module ops_module + + abstract interface + subroutine op(x, y, z) + integer, intent(in) :: x, y + integer, intent(out) :: z + end subroutine + end interface + +contains + + subroutine foo(x, y, r1, r2) + integer, intent(in) :: x, y + integer, intent(out) :: r1, r2 + procedure (op) add1, add2 + procedure (op), pointer::p + p=>add1 + call p(x, y, r1) + p=>add2 + call p(x, y, r2) + end subroutine +end module + +subroutine add1(x, y, z) + integer, intent(in) :: x, y + integer, intent(out) :: z + z = x + y +end subroutine + +subroutine add2(x, y, z) + integer, intent(in) :: x, y + integer, intent(out) :: z + z = x + 2 * y +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/gh18403_mod.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/gh18403_mod.f90 new file mode 100644 index 0000000000000000000000000000000000000000..36791e469f5aee1d5fe15b121abeb9c62a45fadf --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/abstract_interface/gh18403_mod.f90 @@ -0,0 +1,6 @@ +module test + abstract interface + subroutine foo() + end subroutine + end interface +end module test diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c new file mode 100644 index 0000000000000000000000000000000000000000..b66672a43e21dd0641bb29085db716181f5e94ce --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/array_from_pyobj/wrapmodule.c @@ -0,0 +1,235 @@ +/* + * This file was auto-generated with f2py (version:2_1330) and hand edited by + * Pearu for testing purposes. Do not edit this file unless you know what you + * are doing!!! + */ + +#ifdef __cplusplus +extern "C" { +#endif + +/*********************** See f2py2e/cfuncs.py: includes ***********************/ + +#define PY_SSIZE_T_CLEAN +#include +#include "fortranobject.h" +#include + +static PyObject *wrap_error; +static PyObject *wrap_module; + +/************************************ call ************************************/ +static char doc_f2py_rout_wrap_call[] = "\ +Function signature:\n\ + arr = call(type_num,dims,intent,obj)\n\ +Required arguments:\n" +" type_num : input int\n" +" dims : input int-sequence\n" +" intent : input int\n" +" obj : input python object\n" +"Return objects:\n" +" arr : array"; +static PyObject *f2py_rout_wrap_call(PyObject *capi_self, + PyObject *capi_args) { + PyObject * volatile capi_buildvalue = NULL; + int type_num = 0; + int elsize = 0; + npy_intp *dims = NULL; + PyObject *dims_capi = Py_None; + int rank = 0; + int intent = 0; + PyArrayObject *capi_arr_tmp = NULL; + PyObject *arr_capi = Py_None; + int i; + + if (!PyArg_ParseTuple(capi_args,"iiOiO|:wrap.call",\ + &type_num,&elsize,&dims_capi,&intent,&arr_capi)) + return NULL; + rank = PySequence_Length(dims_capi); + dims = malloc(rank*sizeof(npy_intp)); + for (i=0;ikind, + PyArray_DESCR(arr)->type, + PyArray_TYPE(arr), + PyArray_ITEMSIZE(arr), + PyDataType_ALIGNMENT(PyArray_DESCR(arr)), + PyArray_FLAGS(arr), + PyArray_ITEMSIZE(arr)); +} + +static PyMethodDef f2py_module_methods[] = { + + {"call",f2py_rout_wrap_call,METH_VARARGS,doc_f2py_rout_wrap_call}, + {"array_attrs",f2py_rout_wrap_attrs,METH_VARARGS,doc_f2py_rout_wrap_attrs}, + {NULL,NULL} +}; + +static struct PyModuleDef moduledef = { + PyModuleDef_HEAD_INIT, + "test_array_from_pyobj_ext", + NULL, + -1, + f2py_module_methods, + NULL, + NULL, + NULL, + NULL +}; + +PyMODINIT_FUNC PyInit_test_array_from_pyobj_ext(void) { + PyObject *m,*d, *s; + m = wrap_module = PyModule_Create(&moduledef); + Py_SET_TYPE(&PyFortran_Type, &PyType_Type); + import_array(); + if (PyErr_Occurred()) + Py_FatalError("can't initialize module wrap (failed to import numpy)"); + d = PyModule_GetDict(m); + s = PyUnicode_FromString("This module 'wrap' is auto-generated with f2py (version:2_1330).\nFunctions:\n" + " arr = call(type_num,dims,intent,obj)\n" + "."); + PyDict_SetItemString(d, "__doc__", s); + wrap_error = PyErr_NewException ("wrap.error", NULL, NULL); + Py_DECREF(s); + +#define ADDCONST(NAME, CONST) \ + s = PyLong_FromLong(CONST); \ + PyDict_SetItemString(d, NAME, s); \ + Py_DECREF(s) + + ADDCONST("F2PY_INTENT_IN", F2PY_INTENT_IN); + ADDCONST("F2PY_INTENT_INOUT", F2PY_INTENT_INOUT); + ADDCONST("F2PY_INTENT_OUT", F2PY_INTENT_OUT); + ADDCONST("F2PY_INTENT_HIDE", F2PY_INTENT_HIDE); + ADDCONST("F2PY_INTENT_CACHE", F2PY_INTENT_CACHE); + ADDCONST("F2PY_INTENT_COPY", F2PY_INTENT_COPY); + ADDCONST("F2PY_INTENT_C", F2PY_INTENT_C); + ADDCONST("F2PY_OPTIONAL", F2PY_OPTIONAL); + ADDCONST("F2PY_INTENT_INPLACE", F2PY_INTENT_INPLACE); + ADDCONST("NPY_BOOL", NPY_BOOL); + ADDCONST("NPY_BYTE", NPY_BYTE); + ADDCONST("NPY_UBYTE", NPY_UBYTE); + ADDCONST("NPY_SHORT", NPY_SHORT); + ADDCONST("NPY_USHORT", NPY_USHORT); + ADDCONST("NPY_INT", NPY_INT); + ADDCONST("NPY_UINT", NPY_UINT); + ADDCONST("NPY_INTP", NPY_INTP); + ADDCONST("NPY_UINTP", NPY_UINTP); + ADDCONST("NPY_LONG", NPY_LONG); + ADDCONST("NPY_ULONG", NPY_ULONG); + ADDCONST("NPY_LONGLONG", NPY_LONGLONG); + ADDCONST("NPY_ULONGLONG", NPY_ULONGLONG); + ADDCONST("NPY_FLOAT", NPY_FLOAT); + ADDCONST("NPY_DOUBLE", NPY_DOUBLE); + ADDCONST("NPY_LONGDOUBLE", NPY_LONGDOUBLE); + ADDCONST("NPY_CFLOAT", NPY_CFLOAT); + ADDCONST("NPY_CDOUBLE", NPY_CDOUBLE); + ADDCONST("NPY_CLONGDOUBLE", NPY_CLONGDOUBLE); + ADDCONST("NPY_OBJECT", NPY_OBJECT); + ADDCONST("NPY_STRING", NPY_STRING); + ADDCONST("NPY_UNICODE", NPY_UNICODE); + ADDCONST("NPY_VOID", NPY_VOID); + ADDCONST("NPY_NTYPES_LEGACY", NPY_NTYPES_LEGACY); + ADDCONST("NPY_NOTYPE", NPY_NOTYPE); + ADDCONST("NPY_USERDEF", NPY_USERDEF); + + ADDCONST("CONTIGUOUS", NPY_ARRAY_C_CONTIGUOUS); + ADDCONST("FORTRAN", NPY_ARRAY_F_CONTIGUOUS); + ADDCONST("OWNDATA", NPY_ARRAY_OWNDATA); + ADDCONST("FORCECAST", NPY_ARRAY_FORCECAST); + ADDCONST("ENSURECOPY", NPY_ARRAY_ENSURECOPY); + ADDCONST("ENSUREARRAY", NPY_ARRAY_ENSUREARRAY); + ADDCONST("ALIGNED", NPY_ARRAY_ALIGNED); + ADDCONST("WRITEABLE", NPY_ARRAY_WRITEABLE); + ADDCONST("WRITEBACKIFCOPY", NPY_ARRAY_WRITEBACKIFCOPY); + + ADDCONST("BEHAVED", NPY_ARRAY_BEHAVED); + ADDCONST("BEHAVED_NS", NPY_ARRAY_BEHAVED_NS); + ADDCONST("CARRAY", NPY_ARRAY_CARRAY); + ADDCONST("FARRAY", NPY_ARRAY_FARRAY); + ADDCONST("CARRAY_RO", NPY_ARRAY_CARRAY_RO); + ADDCONST("FARRAY_RO", NPY_ARRAY_FARRAY_RO); + ADDCONST("DEFAULT", NPY_ARRAY_DEFAULT); + ADDCONST("UPDATE_ALL", NPY_ARRAY_UPDATE_ALL); + +#undef ADDCONST + + if (PyErr_Occurred()) + Py_FatalError("can't initialize module wrap"); + +#ifdef F2PY_REPORT_ATEXIT + on_exit(f2py_report_on_exit,(void*)"array_from_pyobj.wrap.call"); +#endif + +#if Py_GIL_DISABLED + // signal whether this module supports running with the GIL disabled + PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED); +#endif + + return m; +} +#ifdef __cplusplus +} +#endif diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/.f2py_f2cmap b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/.f2py_f2cmap new file mode 100644 index 0000000000000000000000000000000000000000..2665f89b52d2f17ce7b0a857bea73ec5fab2df88 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/.f2py_f2cmap @@ -0,0 +1 @@ +dict(real=dict(rk="double")) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_free.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_free.f90 new file mode 100644 index 0000000000000000000000000000000000000000..b301710f5dda005e67e40cc21a5e0d62d0ec116a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_free.f90 @@ -0,0 +1,34 @@ + +subroutine sum(x, res) + implicit none + real, intent(in) :: x(:) + real, intent(out) :: res + + integer :: i + + !print *, "sum: size(x) = ", size(x) + + res = 0.0 + + do i = 1, size(x) + res = res + x(i) + enddo + +end subroutine sum + +function fsum(x) result (res) + implicit none + real, intent(in) :: x(:) + real :: res + + integer :: i + + !print *, "fsum: size(x) = ", size(x) + + res = 0.0 + + do i = 1, size(x) + res = res + x(i) + enddo + +end function fsum diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_mod.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_mod.f90 new file mode 100644 index 0000000000000000000000000000000000000000..cbe6317ed8f39f8a633b058a4ab64fe1797ea7b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_mod.f90 @@ -0,0 +1,41 @@ + +module mod + +contains + +subroutine sum(x, res) + implicit none + real, intent(in) :: x(:) + real, intent(out) :: res + + integer :: i + + !print *, "sum: size(x) = ", size(x) + + res = 0.0 + + do i = 1, size(x) + res = res + x(i) + enddo + +end subroutine sum + +function fsum(x) result (res) + implicit none + real, intent(in) :: x(:) + real :: res + + integer :: i + + !print *, "fsum: size(x) = ", size(x) + + res = 0.0 + + do i = 1, size(x) + res = res + x(i) + enddo + +end function fsum + + +end module mod diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_use.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_use.f90 new file mode 100644 index 0000000000000000000000000000000000000000..337465ac540440fc8e8e10d23757af202e8a52a4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/foo_use.f90 @@ -0,0 +1,19 @@ +subroutine sum_with_use(x, res) + use precision + + implicit none + + real(kind=rk), intent(in) :: x(:) + real(kind=rk), intent(out) :: res + + integer :: i + + !print *, "size(x) = ", size(x) + + res = 0.0 + + do i = 1, size(x) + res = res + x(i) + enddo + + end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/precision.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/precision.f90 new file mode 100644 index 0000000000000000000000000000000000000000..ed6c70cbbe7dadfd50b616a8cc1570939ef5afd8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/assumed_shape/precision.f90 @@ -0,0 +1,4 @@ +module precision + integer, parameter :: rk = selected_real_kind(8) + integer, parameter :: ik = selected_real_kind(4) +end module diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/block_docstring/foo.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/block_docstring/foo.f new file mode 100644 index 0000000000000000000000000000000000000000..c8315f12ce0f5cf3dbc4c965ad8843d0c10441cd --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/block_docstring/foo.f @@ -0,0 +1,6 @@ + SUBROUTINE FOO() + INTEGER BAR(2, 3) + + COMMON /BLOCK/ BAR + RETURN + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/foo.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/foo.f new file mode 100644 index 0000000000000000000000000000000000000000..ba397bb38133faa8d502807368074e6b296749b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/foo.f @@ -0,0 +1,62 @@ + subroutine t(fun,a) + integer a +cf2py intent(out) a + external fun + call fun(a) + end + + subroutine func(a) +cf2py intent(in,out) a + integer a + a = a + 11 + end + + subroutine func0(a) +cf2py intent(out) a + integer a + a = 11 + end + + subroutine t2(a) +cf2py intent(callback) fun + integer a +cf2py intent(out) a + external fun + call fun(a) + end + + subroutine string_callback(callback, a) + external callback + double precision callback + double precision a + character*1 r +cf2py intent(out) a + r = 'r' + a = callback(r) + end + + subroutine string_callback_array(callback, cu, lencu, a) + external callback + integer callback + integer lencu + character*8 cu(lencu) + integer a +cf2py intent(out) a + + a = callback(cu, lencu) + end + + subroutine hidden_callback(a, r) + external global_f +cf2py intent(callback, hide) global_f + integer a, r, global_f +cf2py intent(out) r + r = global_f(a) + end + + subroutine hidden_callback2(a, r) + external global_f + integer a, r, global_f +cf2py intent(out) r + r = global_f(a) + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh17797.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh17797.f90 new file mode 100644 index 0000000000000000000000000000000000000000..49853afd766a90e521104081bf77236a252d3c70 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh17797.f90 @@ -0,0 +1,7 @@ +function gh17797(f, y) result(r) + external f + integer(8) :: r, f + integer(8), dimension(:) :: y + r = f(0) + r = r + sum(y) +end function gh17797 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh18335.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh18335.f90 new file mode 100644 index 0000000000000000000000000000000000000000..92b6d7540c827d20c7d2169c56f14653954d7ff9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh18335.f90 @@ -0,0 +1,17 @@ + ! When gh18335_workaround is defined as an extension, + ! the issue cannot be reproduced. + !subroutine gh18335_workaround(f, y) + ! implicit none + ! external f + ! integer(kind=1) :: y(1) + ! call f(y) + !end subroutine gh18335_workaround + + function gh18335(f) result (r) + implicit none + external f + integer(kind=1) :: y(1), r + y(1) = 123 + call f(y) + r = y(1) + end function gh18335 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.f new file mode 100644 index 0000000000000000000000000000000000000000..ba727a10a07ebec77845f8a67746cf5d5bb3d32a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.f @@ -0,0 +1,10 @@ + SUBROUTINE FOO(FUN,R) + EXTERNAL FUN + INTEGER I + REAL*8 R, FUN +Cf2py intent(out) r + R = 0D0 + DO I=-5,5 + R = R + FUN(I) + ENDDO + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.pyf new file mode 100644 index 0000000000000000000000000000000000000000..f12011153370b022a2686222655a12245a1eb14e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh25211.pyf @@ -0,0 +1,18 @@ +python module __user__routines + interface + function fun(i) result (r) + integer :: i + real*8 :: r + end function fun + end interface +end python module __user__routines + +python module callback2 + interface + subroutine foo(f,r) + use __user__routines, f=>fun + external f + real*8 intent(out) :: r + end subroutine foo + end interface +end python module callback2 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh26681.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh26681.f90 new file mode 100644 index 0000000000000000000000000000000000000000..00c0ec93df059b0d1952471fae440016abbac3e5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/callback/gh26681.f90 @@ -0,0 +1,18 @@ +module utils + implicit none + contains + subroutine my_abort(message) + implicit none + character(len=*), intent(in) :: message + !f2py callstatement PyErr_SetString(PyExc_ValueError, message);f2py_success = 0; + !f2py callprotoargument char* + write(0,*) "THIS SHOULD NOT APPEAR" + stop 1 + end subroutine my_abort + + subroutine do_something(message) + !f2py intent(callback, hide) mypy_abort + character(len=*), intent(in) :: message + call mypy_abort(message) + end subroutine do_something +end module utils diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/gh_22819.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/gh_22819.pyf new file mode 100644 index 0000000000000000000000000000000000000000..8eb5bb106a366ec214944c19e53d9788c0596e55 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/gh_22819.pyf @@ -0,0 +1,6 @@ +python module test_22819 + interface + subroutine hello() + end subroutine hello + end interface +end python module test_22819 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hi77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hi77.f new file mode 100644 index 0000000000000000000000000000000000000000..8b916ebe0459eb812baad694aa671011a1381f8a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hi77.f @@ -0,0 +1,3 @@ + SUBROUTINE HI + PRINT*, "HELLO WORLD" + END SUBROUTINE diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hiworld.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hiworld.f90 new file mode 100644 index 0000000000000000000000000000000000000000..981f877547a4caec513a15dea1401bbc98ce3f23 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/cli/hiworld.f90 @@ -0,0 +1,3 @@ +function hi() + print*, "Hello World" +end function diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/block.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/block.f new file mode 100644 index 0000000000000000000000000000000000000000..7ea7968fe935182bd17a697b316569546937b715 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/block.f @@ -0,0 +1,11 @@ + SUBROUTINE INITCB + DOUBLE PRECISION LONG + CHARACTER STRING + INTEGER OK + + COMMON /BLOCK/ LONG, STRING, OK + LONG = 1.0 + STRING = '2' + OK = 3 + RETURN + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/gh19161.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/gh19161.f90 new file mode 100644 index 0000000000000000000000000000000000000000..a2f40735ad66a3cb70cfc10a3938882c77ff54ea --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/common/gh19161.f90 @@ -0,0 +1,10 @@ +module typedefmod + use iso_fortran_env, only: real32 +end module typedefmod + +module data + use typedefmod, only: real32 + implicit none + real(kind=real32) :: x + common/test/x +end module data diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/accesstype.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/accesstype.f90 new file mode 100644 index 0000000000000000000000000000000000000000..e2cbd445daf57f21e2d727f42a3891ec28725175 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/accesstype.f90 @@ -0,0 +1,13 @@ +module foo + public + type, private, bind(c) :: a + integer :: i + end type a + type, bind(c) :: b_ + integer :: j + end type b_ + public :: b_ + type :: c + integer :: k + end type c +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/common_with_division.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/common_with_division.f new file mode 100644 index 0000000000000000000000000000000000000000..4aa12cf6dcee9b4936d10f1c5d55f8d59762062a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/common_with_division.f @@ -0,0 +1,17 @@ + subroutine common_with_division + integer lmu,lb,lub,lpmin + parameter (lmu=1) + parameter (lb=20) +c crackfortran fails to parse this +c parameter (lub=(lb-1)*lmu+1) +c crackfortran can successfully parse this though + parameter (lub=lb*lmu-lmu+1) + parameter (lpmin=2) + +c crackfortran fails to parse this correctly +c common /mortmp/ ctmp((lub*(lub+1)*(lub+1))/lpmin+1) + + common /mortmp/ ctmp(lub/lpmin+1) + + return + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_common.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_common.f new file mode 100644 index 0000000000000000000000000000000000000000..5ffd865c837997f8aae2d8faebfd519df61d8cd2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_common.f @@ -0,0 +1,8 @@ + BLOCK DATA PARAM_INI + COMMON /MYCOM/ MYDATA + DATA MYDATA /0/ + END + SUBROUTINE SUB1 + COMMON /MYCOM/ MYDATA + MYDATA = MYDATA + 1 + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_multiplier.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_multiplier.f new file mode 100644 index 0000000000000000000000000000000000000000..19ff8a83e97b7a1fa9ef82a2f4d5241ec422cb01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_multiplier.f @@ -0,0 +1,5 @@ + BLOCK DATA MYBLK + IMPLICIT DOUBLE PRECISION (A-H,O-Z) + COMMON /MYCOM/ IVAR1, IVAR2, IVAR3, IVAR4, EVAR5 + DATA IVAR1, IVAR2, IVAR3, IVAR4, EVAR5 /2*3,2*2,0.0D0/ + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_stmts.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_stmts.f90 new file mode 100644 index 0000000000000000000000000000000000000000..576c5e485baf209aea79f566fc09cb20138a0a25 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_stmts.f90 @@ -0,0 +1,20 @@ +! gh-23276 +module cmplxdat + implicit none + integer :: i, j + real :: x, y + real, dimension(2) :: z + real(kind=8) :: pi + complex(kind=8), target :: medium_ref_index + complex(kind=8), target :: ref_index_one, ref_index_two + complex(kind=8), dimension(2) :: my_array + real(kind=8), dimension(3) :: my_real_array = (/1.0d0, 2.0d0, 3.0d0/) + + data i, j / 2, 3 / + data x, y / 1.5, 2.0 / + data z / 3.5, 7.0 / + data medium_ref_index / (1.d0, 0.d0) / + data ref_index_one, ref_index_two / (13.0d0, 21.0d0), (-30.0d0, 43.0d0) / + data my_array / (1.0d0, 2.0d0), (-3.0d0, 4.0d0) / + data pi / 3.1415926535897932384626433832795028841971693993751058209749445923078164062d0 / +end module cmplxdat diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_with_comments.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_with_comments.f new file mode 100644 index 0000000000000000000000000000000000000000..4128f004e840087ab8e08a06c76995b249a561b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/data_with_comments.f @@ -0,0 +1,8 @@ + BLOCK DATA PARAM_INI + COMMON /MYCOM/ MYTAB + INTEGER MYTAB(3) + DATA MYTAB/ + * 0, ! 1 and more commenty stuff + * 4, ! 2 + * 0 / + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/foo_deps.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/foo_deps.f90 new file mode 100644 index 0000000000000000000000000000000000000000..e327b25c81986b2191fc740991ca9e907b5b0fb6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/foo_deps.f90 @@ -0,0 +1,6 @@ +module foo + type bar + character(len = 4) :: text + end type bar + type(bar), parameter :: abar = bar('abar') +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh15035.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh15035.f new file mode 100644 index 0000000000000000000000000000000000000000..1bb2e6745952cb10067116e9ae3337c8314061ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh15035.f @@ -0,0 +1,16 @@ + subroutine subb(k) + real(8), intent(inout) :: k(:) + k=k+1 + endsubroutine + + subroutine subc(w,k) + real(8), intent(in) :: w(:) + real(8), intent(out) :: k(size(w)) + k=w+1 + endsubroutine + + function t0(value) + character value + character t0 + t0 = value + endfunction diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh17859.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh17859.f new file mode 100644 index 0000000000000000000000000000000000000000..995953845c5eb1b4fa2bdf70c18e0296d38e5252 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh17859.f @@ -0,0 +1,12 @@ + integer(8) function external_as_statement(fcn) + implicit none + external fcn + integer(8) :: fcn + external_as_statement = fcn(0) + end + + integer(8) function external_as_attribute(fcn) + implicit none + integer(8), external :: fcn + external_as_attribute = fcn(0) + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh22648.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh22648.pyf new file mode 100644 index 0000000000000000000000000000000000000000..b3454f18635fc8fe2b8ea5d15f85a9d77af9a22b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh22648.pyf @@ -0,0 +1,7 @@ +python module iri16py ! in + interface ! in :iri16py + block data ! in :iri16py:iridreg_modified.for + COMMON /fircom/ eden,tabhe,tabla,tabmo,tabza,tabfl + end block data + end interface +end python module iri16py diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23533.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23533.f new file mode 100644 index 0000000000000000000000000000000000000000..db522afa7d2fdd09e26f2d02a649a659d9ed7d60 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23533.f @@ -0,0 +1,5 @@ + SUBROUTINE EXAMPLE( ) + IF( .TRUE. ) THEN + CALL DO_SOMETHING() + END IF ! ** .TRUE. ** + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598.f90 new file mode 100644 index 0000000000000000000000000000000000000000..e0dffb5ef29e3d5ba853ff4dfeda57b2bed6a9dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598.f90 @@ -0,0 +1,4 @@ +integer function intproduct(a, b) result(res) + integer, intent(in) :: a, b + res = a*b +end function diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598Warn.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598Warn.f90 new file mode 100644 index 0000000000000000000000000000000000000000..3b44efc5ef16e9f7e1105229371ae48ecc069ee5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23598Warn.f90 @@ -0,0 +1,11 @@ +module test_bug + implicit none + private + public :: intproduct + +contains + integer function intproduct(a, b) result(res) + integer, intent(in) :: a, b + res = a*b + end function +end module diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23879.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23879.f90 new file mode 100644 index 0000000000000000000000000000000000000000..fac262d53c9d3f0f3a5ba1138594f5b694b95717 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh23879.f90 @@ -0,0 +1,20 @@ +module gh23879 + implicit none + private + public :: foo + + contains + + subroutine foo(a, b) + integer, intent(in) :: a + integer, intent(out) :: b + b = a + call bar(b) + end subroutine + + subroutine bar(x) + integer, intent(inout) :: x + x = 2*x + end subroutine + + end module gh23879 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh27697.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh27697.f90 new file mode 100644 index 0000000000000000000000000000000000000000..a5eae4e79b25cd086df4ea2aa26d39ae48a8ca88 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh27697.f90 @@ -0,0 +1,12 @@ +module utils + implicit none + contains + subroutine my_abort(message) + implicit none + character(len=*), intent(in) :: message + !f2py callstatement PyErr_SetString(PyExc_ValueError, message);f2py_success = 0; + !f2py callprotoargument char* + write(0,*) "THIS SHOULD NOT APPEAR" + stop 1 + end subroutine my_abort +end module utils diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh2848.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh2848.f90 new file mode 100644 index 0000000000000000000000000000000000000000..31ea9327a4d9134011cfc668cc88961968756d77 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/gh2848.f90 @@ -0,0 +1,13 @@ + subroutine gh2848( & + ! first 2 parameters + par1, par2,& + ! last 2 parameters + par3, par4) + + integer, intent(in) :: par1, par2 + integer, intent(out) :: par3, par4 + + par3 = par1 + par4 = par2 + + end subroutine gh2848 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/operators.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/operators.f90 new file mode 100644 index 0000000000000000000000000000000000000000..1d060a3d2bd5abd12732e6003cec53f36baeba7c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/operators.f90 @@ -0,0 +1,49 @@ +module foo + type bar + character(len = 32) :: item + end type bar + interface operator(.item.) + module procedure item_int, item_real + end interface operator(.item.) + interface operator(==) + module procedure items_are_equal + end interface operator(==) + interface assignment(=) + module procedure get_int, get_real + end interface assignment(=) +contains + function item_int(val) result(elem) + integer, intent(in) :: val + type(bar) :: elem + + write(elem%item, "(I32)") val + end function item_int + + function item_real(val) result(elem) + real, intent(in) :: val + type(bar) :: elem + + write(elem%item, "(1PE32.12)") val + end function item_real + + function items_are_equal(val1, val2) result(equal) + type(bar), intent(in) :: val1, val2 + logical :: equal + + equal = (val1%item == val2%item) + end function items_are_equal + + subroutine get_real(rval, item) + real, intent(out) :: rval + type(bar), intent(in) :: item + + read(item%item, *) rval + end subroutine get_real + + subroutine get_int(rval, item) + integer, intent(out) :: rval + type(bar), intent(in) :: item + + read(item%item, *) rval + end subroutine get_int +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/privatemod.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/privatemod.f90 new file mode 100644 index 0000000000000000000000000000000000000000..2674c214767b33663e51ee1d32ad7a1792c92680 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/privatemod.f90 @@ -0,0 +1,11 @@ +module foo + private + integer :: a + public :: setA + integer :: b +contains + subroutine setA(v) + integer, intent(in) :: v + a = v + end subroutine setA +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/publicmod.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/publicmod.f90 new file mode 100644 index 0000000000000000000000000000000000000000..1db76e3fe06828ba0d1b640720ec70422cde6872 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/publicmod.f90 @@ -0,0 +1,10 @@ +module foo + public + integer, private :: a + public :: setA +contains + subroutine setA(v) + integer, intent(in) :: v + a = v + end subroutine setA +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/pubprivmod.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/pubprivmod.f90 new file mode 100644 index 0000000000000000000000000000000000000000..46bef7cb91122281ddac7d0f0474c2c01b2a5e6f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/pubprivmod.f90 @@ -0,0 +1,10 @@ +module foo + public + integer, private :: a + integer :: b +contains + subroutine setA(v) + integer, intent(in) :: v + a = v + end subroutine setA +end module foo diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/unicode_comment.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/unicode_comment.f90 new file mode 100644 index 0000000000000000000000000000000000000000..13515ce98c50e88a03004161fb135e8502005a82 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/crackfortran/unicode_comment.f90 @@ -0,0 +1,4 @@ +subroutine foo(x) + real(8), intent(in) :: x + ! Écrit à l'écran la valeur de x +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/.f2py_f2cmap b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/.f2py_f2cmap new file mode 100644 index 0000000000000000000000000000000000000000..a4425f8876f5b7ec9c72a11862a8cd4574d33ea8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/.f2py_f2cmap @@ -0,0 +1 @@ +dict(real=dict(real32='float', real64='double'), integer=dict(int64='long_long')) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/isoFortranEnvMap.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/isoFortranEnvMap.f90 new file mode 100644 index 0000000000000000000000000000000000000000..1e1dc1d4054b36d2b2d9104e8d6ab708361bfbe8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/f2cmap/isoFortranEnvMap.f90 @@ -0,0 +1,9 @@ + subroutine func1(n, x, res) + use, intrinsic :: iso_fortran_env, only: int64, real64 + implicit none + integer(int64), intent(in) :: n + real(real64), intent(in) :: x(n) + real(real64), intent(out) :: res +!f2py intent(hide) :: n + res = sum(x) + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/isocintrin/isoCtests.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/isocintrin/isoCtests.f90 new file mode 100644 index 0000000000000000000000000000000000000000..765f7c1ce6608a0c8588b6c20edd052e2d3e07bf --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/isocintrin/isoCtests.f90 @@ -0,0 +1,34 @@ + module coddity + use iso_c_binding, only: c_double, c_int, c_int64_t + implicit none + contains + subroutine c_add(a, b, c) bind(c, name="c_add") + real(c_double), intent(in) :: a, b + real(c_double), intent(out) :: c + c = a + b + end subroutine c_add + ! gh-9693 + function wat(x, y) result(z) bind(c) + integer(c_int), intent(in) :: x, y + integer(c_int) :: z + + z = x + 7 + end function wat + ! gh-25207 + subroutine c_add_int64(a, b, c) bind(c) + integer(c_int64_t), intent(in) :: a, b + integer(c_int64_t), intent(out) :: c + c = a + b + end subroutine c_add_int64 + ! gh-25207 + subroutine add_arr(A, B, C) + integer(c_int64_t), intent(in) :: A(3) + integer(c_int64_t), intent(in) :: B(3) + integer(c_int64_t), intent(out) :: C(3) + integer :: j + + do j = 1, 3 + C(j) = A(j)+B(j) + end do + end subroutine + end module coddity diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/kind/foo.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/kind/foo.f90 new file mode 100644 index 0000000000000000000000000000000000000000..d3d15cfb20a15004ed86e45dc91792d1c089033a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/kind/foo.f90 @@ -0,0 +1,20 @@ + + +subroutine selectedrealkind(p, r, res) + implicit none + + integer, intent(in) :: p, r + !f2py integer :: r=0 + integer, intent(out) :: res + res = selected_real_kind(p, r) + +end subroutine + +subroutine selectedintkind(p, res) + implicit none + + integer, intent(in) :: p + integer, intent(out) :: res + res = selected_int_kind(p) + +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo.f new file mode 100644 index 0000000000000000000000000000000000000000..c34742578f8551729fdc3474d86e92c87e2868d2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo.f @@ -0,0 +1,5 @@ + subroutine bar11(a) +cf2py intent(out) a + integer a + a = 11 + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_fixed.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_fixed.f90 new file mode 100644 index 0000000000000000000000000000000000000000..7543a6acb7375872388cb9f2ced109db5faa17b0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_fixed.f90 @@ -0,0 +1,8 @@ + module foo_fixed + contains + subroutine bar12(a) +!f2py intent(out) a + integer a + a = 12 + end subroutine bar12 + end module foo_fixed diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_free.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_free.f90 new file mode 100644 index 0000000000000000000000000000000000000000..c1b641f13ec2943b9dd23ba85beecda738b51825 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/mixed/foo_free.f90 @@ -0,0 +1,8 @@ +module foo_free +contains + subroutine bar13(a) + !f2py intent(out) a + integer a + a = 13 + end subroutine bar13 +end module foo_free diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/data.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/data.f90 new file mode 100644 index 0000000000000000000000000000000000000000..483d13ceb95c08bf38b74d8218932fc109792b09 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/data.f90 @@ -0,0 +1,8 @@ +module data + real(8) :: shift +contains + subroutine set_shift(in_shift) + real(8), intent(in) :: in_shift + shift = in_shift + end subroutine set_shift +end module data diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/use_data.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/use_data.f90 new file mode 100644 index 0000000000000000000000000000000000000000..b3fae8b875d03d75199f4cf06d544edb4aab1a89 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh25337/use_data.f90 @@ -0,0 +1,6 @@ +subroutine shift_a(dim_a, a) + use data, only: shift + integer, intent(in) :: dim_a + real(8), intent(inout), dimension(dim_a) :: a + a = a + shift +end subroutine shift_a diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_no_public_entities.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_no_public_entities.f90 new file mode 100644 index 0000000000000000000000000000000000000000..07adce591f35756ca2fe2f3e2dd38fcaf01d2fad --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_no_public_entities.f90 @@ -0,0 +1,21 @@ + module mod2 + implicit none + private mod2_func1 + contains + + subroutine mod2_func1() + print*, "mod2_func1" + end subroutine mod2_func1 + + end module mod2 + + module mod1 + implicit none + private :: mod1_func1 + contains + + subroutine mod1_func1() + print*, "mod1_func1" + end subroutine mod1_func1 + + end module mod1 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_one_public_routine.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_one_public_routine.f90 new file mode 100644 index 0000000000000000000000000000000000000000..b7fb95b010a6657df3abeed4c71466b82dcfdab3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/gh26920/two_mods_with_one_public_routine.f90 @@ -0,0 +1,21 @@ + module mod2 + implicit none + PUBLIC :: mod2_func1 + contains + + subroutine mod2_func1() + print*, "mod2_func1" + end subroutine mod2_func1 + + end module mod2 + + module mod1 + implicit none + PUBLIC :: mod1_func1 + contains + + subroutine mod1_func1() + print*, "mod1_func1" + end subroutine mod1_func1 + + end module mod1 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/module_data_docstring.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/module_data_docstring.f90 new file mode 100644 index 0000000000000000000000000000000000000000..4505e0cbc31e50a75df94b30cd53cf923659379d --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/module_data_docstring.f90 @@ -0,0 +1,12 @@ +module mod + integer :: i + integer :: x(4) + real, dimension(2,3) :: a + real, allocatable, dimension(:,:) :: b +contains + subroutine foo + integer :: k + k = 1 + a(1,2) = a(1,2)+3 + end subroutine foo +end module mod diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/use_modules.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/use_modules.f90 new file mode 100644 index 0000000000000000000000000000000000000000..aa40c86ca39d1dbcfacca9dcb2addbc6ede73140 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/modules/use_modules.f90 @@ -0,0 +1,20 @@ +module mathops + implicit none +contains + function add(a, b) result(c) + integer, intent(in) :: a, b + integer :: c + c = a + b + end function add +end module mathops + +module useops + use mathops, only: add + implicit none +contains + function sum_and_double(a, b) result(d) + integer, intent(in) :: a, b + integer :: d + d = 2 * add(a, b) + end function sum_and_double +end module useops diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/negative_bounds/issue_20853.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/negative_bounds/issue_20853.f90 new file mode 100644 index 0000000000000000000000000000000000000000..bf1fa92853316cc31f825c024855088f42577a1c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/negative_bounds/issue_20853.f90 @@ -0,0 +1,7 @@ +subroutine foo(is_, ie_, arr, tout) + implicit none + integer :: is_,ie_ + real, intent(in) :: arr(is_:ie_) + real, intent(out) :: tout(is_:ie_) + tout = arr +end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_array.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_array.f90 new file mode 100644 index 0000000000000000000000000000000000000000..9a6bf81610d41cf5f480d6e3c25fd7ab4cf5bfff --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_array.f90 @@ -0,0 +1,45 @@ +! Check that parameter arrays are correctly intercepted. +subroutine foo_array(x, y, z) + implicit none + integer, parameter :: dp = selected_real_kind(15) + integer, parameter :: pa = 2 + integer, parameter :: intparamarray(2) = (/ 3, 5 /) + integer, dimension(pa), parameter :: pb = (/ 2, 10 /) + integer, parameter, dimension(intparamarray(1)) :: pc = (/ 2, 10, 20 /) + real(dp), parameter :: doubleparamarray(3) = (/ 3.14_dp, 4._dp, 6.44_dp /) + real(dp), intent(inout) :: x(intparamarray(1)) + real(dp), intent(inout) :: y(intparamarray(2)) + real(dp), intent(out) :: z + + x = x/pb(2) + y = y*pc(2) + z = doubleparamarray(1)*doubleparamarray(2) + doubleparamarray(3) + + return +end subroutine + +subroutine foo_array_any_index(x, y) + implicit none + integer, parameter :: dp = selected_real_kind(15) + integer, parameter, dimension(-1:1) :: myparamarray = (/ 6, 3, 1 /) + integer, parameter, dimension(2) :: nested = (/ 2, 0 /) + integer, parameter :: dim = 2 + real(dp), intent(in) :: x(myparamarray(-1)) + real(dp), intent(out) :: y(nested(1), myparamarray(nested(dim))) + + y = reshape(x, (/nested(1), myparamarray(nested(2))/)) + + return +end subroutine + +subroutine foo_array_delims(x) + implicit none + integer, parameter :: dp = selected_real_kind(15) + integer, parameter, dimension(2) :: myparamarray = (/ (6), 1 /) + integer, parameter, dimension(3) :: test = (/2, 1, (3)/) + real(dp), intent(out) :: x + + x = myparamarray(1)+test(3) + + return +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_both.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_both.f90 new file mode 100644 index 0000000000000000000000000000000000000000..ac90cedc525a6172a9b72f3bc76f57b79d641b6c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_both.f90 @@ -0,0 +1,57 @@ +! Check that parameters are correct intercepted. +! Constants with comma separations are commonly +! used, for instance Pi = 3._dp +subroutine foo(x) + implicit none + integer, parameter :: sp = selected_real_kind(6) + integer, parameter :: dp = selected_real_kind(15) + integer, parameter :: ii = selected_int_kind(9) + integer, parameter :: il = selected_int_kind(18) + real(dp), intent(inout) :: x + dimension x(3) + real(sp), parameter :: three_s = 3._sp + real(dp), parameter :: three_d = 3._dp + integer(ii), parameter :: three_i = 3_ii + integer(il), parameter :: three_l = 3_il + x(1) = x(1) + x(2) * three_s * three_i + x(3) * three_d * three_l + x(2) = x(2) * three_s + x(3) = x(3) * three_l + return +end subroutine + + +subroutine foo_no(x) + implicit none + integer, parameter :: sp = selected_real_kind(6) + integer, parameter :: dp = selected_real_kind(15) + integer, parameter :: ii = selected_int_kind(9) + integer, parameter :: il = selected_int_kind(18) + real(dp), intent(inout) :: x + dimension x(3) + real(sp), parameter :: three_s = 3. + real(dp), parameter :: three_d = 3. + integer(ii), parameter :: three_i = 3 + integer(il), parameter :: three_l = 3 + x(1) = x(1) + x(2) * three_s * three_i + x(3) * three_d * three_l + x(2) = x(2) * three_s + x(3) = x(3) * three_l + return +end subroutine + +subroutine foo_sum(x) + implicit none + integer, parameter :: sp = selected_real_kind(6) + integer, parameter :: dp = selected_real_kind(15) + integer, parameter :: ii = selected_int_kind(9) + integer, parameter :: il = selected_int_kind(18) + real(dp), intent(inout) :: x + dimension x(3) + real(sp), parameter :: three_s = 2._sp + 1._sp + real(dp), parameter :: three_d = 1._dp + 2._dp + integer(ii), parameter :: three_i = 2_ii + 1_ii + integer(il), parameter :: three_l = 1_il + 2_il + x(1) = x(1) + x(2) * three_s * three_i + x(3) * three_d * three_l + x(2) = x(2) * three_s + x(3) = x(3) * three_l + return +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_compound.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_compound.f90 new file mode 100644 index 0000000000000000000000000000000000000000..e51f5e9b2fb166a6b7d9cba57af03617024b7f2a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_compound.f90 @@ -0,0 +1,15 @@ +! Check that parameters are correct intercepted. +! Constants with comma separations are commonly +! used, for instance Pi = 3._dp +subroutine foo_compound_int(x) + implicit none + integer, parameter :: ii = selected_int_kind(9) + integer(ii), intent(inout) :: x + dimension x(3) + integer(ii), parameter :: three = 3_ii + integer(ii), parameter :: two = 2_ii + integer(ii), parameter :: six = three * 1_ii * two + + x(1) = x(1) + x(2) + x(3) * six + return +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_integer.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_integer.f90 new file mode 100644 index 0000000000000000000000000000000000000000..aaa83d2eb241274231130b6243ca2c970b5664e0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_integer.f90 @@ -0,0 +1,22 @@ +! Check that parameters are correct intercepted. +! Constants with comma separations are commonly +! used, for instance Pi = 3._dp +subroutine foo_int(x) + implicit none + integer, parameter :: ii = selected_int_kind(9) + integer(ii), intent(inout) :: x + dimension x(3) + integer(ii), parameter :: three = 3_ii + x(1) = x(1) + x(2) + x(3) * three + return +end subroutine + +subroutine foo_long(x) + implicit none + integer, parameter :: ii = selected_int_kind(18) + integer(ii), intent(inout) :: x + dimension x(3) + integer(ii), parameter :: three = 3_ii + x(1) = x(1) + x(2) + x(3) * three + return +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_non_compound.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_non_compound.f90 new file mode 100644 index 0000000000000000000000000000000000000000..62c9a5b943cb768c9270a04d1dbf36d526a4e6e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_non_compound.f90 @@ -0,0 +1,23 @@ +! Check that parameters are correct intercepted. +! Specifically that types of constants without +! compound kind specs are correctly inferred +! adapted Gibbs iteration code from pymc +! for this test case +subroutine foo_non_compound_int(x) + implicit none + integer, parameter :: ii = selected_int_kind(9) + + integer(ii) maxiterates + parameter (maxiterates=2) + + integer(ii) maxseries + parameter (maxseries=2) + + integer(ii) wasize + parameter (wasize=maxiterates*maxseries) + integer(ii), intent(inout) :: x + dimension x(wasize) + + x(1) = x(1) + x(2) + x(3) + x(4) * wasize + return +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_real.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_real.f90 new file mode 100644 index 0000000000000000000000000000000000000000..02ac9dd993b39dbb69a233ed1f0d031f15f84639 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/parameter/constant_real.f90 @@ -0,0 +1,23 @@ +! Check that parameters are correct intercepted. +! Constants with comma separations are commonly +! used, for instance Pi = 3._dp +subroutine foo_single(x) + implicit none + integer, parameter :: rp = selected_real_kind(6) + real(rp), intent(inout) :: x + dimension x(3) + real(rp), parameter :: three = 3._rp + x(1) = x(1) + x(2) + x(3) * three + return +end subroutine + +subroutine foo_double(x) + implicit none + integer, parameter :: rp = selected_real_kind(15) + real(rp), intent(inout) :: x + dimension x(3) + real(rp), parameter :: three = 3._rp + x(1) = x(1) + x(2) + x(3) * three + return +end subroutine + diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/quoted_character/foo.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/quoted_character/foo.f new file mode 100644 index 0000000000000000000000000000000000000000..9dc1cfa4446d8c05c0fc0bb1c69360a687d003c3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/quoted_character/foo.f @@ -0,0 +1,14 @@ + SUBROUTINE FOO(OUT1, OUT2, OUT3, OUT4, OUT5, OUT6) + CHARACTER SINGLE, DOUBLE, SEMICOL, EXCLA, OPENPAR, CLOSEPAR + PARAMETER (SINGLE="'", DOUBLE='"', SEMICOL=';', EXCLA="!", + 1 OPENPAR="(", CLOSEPAR=")") + CHARACTER OUT1, OUT2, OUT3, OUT4, OUT5, OUT6 +Cf2py intent(out) OUT1, OUT2, OUT3, OUT4, OUT5, OUT6 + OUT1 = SINGLE + OUT2 = DOUBLE + OUT3 = SEMICOL + OUT4 = EXCLA + OUT5 = OPENPAR + OUT6 = CLOSEPAR + RETURN + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/AB.inc b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/AB.inc new file mode 100644 index 0000000000000000000000000000000000000000..8a02f631f43a15f17f65280d651aa46e9cdb7510 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/AB.inc @@ -0,0 +1 @@ +real(8) b, n, m diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/assignOnlyModule.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/assignOnlyModule.f90 new file mode 100644 index 0000000000000000000000000000000000000000..479ac7980c226ac19a8d97dd2e43f8833d95e66e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/assignOnlyModule.f90 @@ -0,0 +1,25 @@ + MODULE MOD_TYPES + INTEGER, PARAMETER :: SP = SELECTED_REAL_KIND(6, 37) + INTEGER, PARAMETER :: DP = SELECTED_REAL_KIND(15, 307) + END MODULE +! + MODULE F_GLOBALS + USE MOD_TYPES + IMPLICIT NONE + INTEGER, PARAMETER :: N_MAX = 16 + INTEGER, PARAMETER :: I_MAX = 18 + INTEGER, PARAMETER :: J_MAX = 72 + REAL(SP) :: XREF + END MODULE F_GLOBALS +! + SUBROUTINE DUMMY () +! + USE F_GLOBALS + USE MOD_TYPES + IMPLICIT NONE +! + REAL(SP) :: MINIMAL + MINIMAL = 0.01*XREF + RETURN +! + END SUBROUTINE DUMMY diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/datonly.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/datonly.f90 new file mode 100644 index 0000000000000000000000000000000000000000..67fc4aca82e3fd4aa9ebb4f725d1b13761dbade9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/datonly.f90 @@ -0,0 +1,17 @@ +module datonly + implicit none + integer, parameter :: max_value = 100 + real, dimension(:), allocatable :: data_array +end module datonly + +module dat + implicit none + integer, parameter :: max_= 1009 +end module dat + +subroutine simple_subroutine(ain, aout) + use dat, only: max_ + integer, intent(in) :: ain + integer, intent(out) :: aout + aout = ain + max_ +end subroutine simple_subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77comments.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77comments.f new file mode 100644 index 0000000000000000000000000000000000000000..452a01a14439b1e7a9a731a31d7f720992826698 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77comments.f @@ -0,0 +1,26 @@ + SUBROUTINE TESTSUB( + & INPUT1, INPUT2, !Input + & OUTPUT1, OUTPUT2) !Output + + IMPLICIT NONE + INTEGER, INTENT(IN) :: INPUT1, INPUT2 + INTEGER, INTENT(OUT) :: OUTPUT1, OUTPUT2 + + OUTPUT1 = INPUT1 + INPUT2 + OUTPUT2 = INPUT1 * INPUT2 + + RETURN + END SUBROUTINE TESTSUB + + SUBROUTINE TESTSUB2(OUTPUT) + IMPLICIT NONE + INTEGER, PARAMETER :: N = 10 ! Array dimension + REAL, INTENT(OUT) :: OUTPUT(N) + INTEGER :: I + + DO I = 1, N + OUTPUT(I) = I * 2.0 + END DO + + RETURN + END diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77fixedform.f95 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77fixedform.f95 new file mode 100644 index 0000000000000000000000000000000000000000..e47a13f7e851969f30aa38344377c5ec90708839 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f77fixedform.f95 @@ -0,0 +1,5 @@ +C This is an invalid file, but it does compile with -ffixed-form + subroutine mwe( + & x) + real x + end subroutine mwe diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f90continuation.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f90continuation.f90 new file mode 100644 index 0000000000000000000000000000000000000000..879e716bbec69e94a28745ecc0d8f9c7a8ea02eb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/f90continuation.f90 @@ -0,0 +1,9 @@ +SUBROUTINE TESTSUB(INPUT1, & ! Hello +! commenty +INPUT2, OUTPUT1, OUTPUT2) ! more comments + INTEGER, INTENT(IN) :: INPUT1, INPUT2 + INTEGER, INTENT(OUT) :: OUTPUT1, OUTPUT2 + OUTPUT1 = INPUT1 + & + INPUT2 + OUTPUT2 = INPUT1 * INPUT2 +END SUBROUTINE TESTSUB diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/incfile.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/incfile.f90 new file mode 100644 index 0000000000000000000000000000000000000000..276ef3a67352b29d37985529ad550b45f68456f0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/incfile.f90 @@ -0,0 +1,5 @@ +function add(n,m) result(b) + implicit none + include 'AB.inc' + b = n + m +end function add diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/inout.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/inout.f90 new file mode 100644 index 0000000000000000000000000000000000000000..80cdad90cec56de2226979fa0c9b0f9dfa39c7c9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/inout.f90 @@ -0,0 +1,9 @@ +! Check that intent(in out) translates as intent(inout). +! The separation seems to be a common usage. + subroutine foo(x) + implicit none + real(4), intent(in out) :: x + dimension x(3) + x(1) = x(1) + x(2) + x(3) + return + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/lower_f2py_fortran.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/lower_f2py_fortran.f90 new file mode 100644 index 0000000000000000000000000000000000000000..1c4b8c192b1b2783a4f5c2275731106c2a6109f2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/regression/lower_f2py_fortran.f90 @@ -0,0 +1,5 @@ +subroutine inquire_next(IU) + IMPLICIT NONE + integer :: IU + !f2py intent(in) IU +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo77.f new file mode 100644 index 0000000000000000000000000000000000000000..facae1016a39010cca10929837d0a95c44376e21 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo77.f @@ -0,0 +1,45 @@ + function t0(value) + character value + character t0 + t0 = value + end + function t1(value) + character*1 value + character*1 t1 + t1 = value + end + function t5(value) + character*5 value + character*5 t5 + t5 = value + end + function ts(value) + character*(*) value + character*(*) ts + ts = value + end + + subroutine s0(t0,value) + character value + character t0 +cf2py intent(out) t0 + t0 = value + end + subroutine s1(t1,value) + character*1 value + character*1 t1 +cf2py intent(out) t1 + t1 = value + end + subroutine s5(t5,value) + character*5 value + character*5 t5 +cf2py intent(out) t5 + t5 = value + end + subroutine ss(ts,value) + character*(*) value + character*10 ts +cf2py intent(out) ts + ts = value + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo90.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo90.f90 new file mode 100644 index 0000000000000000000000000000000000000000..36182bcf2dd71649130f5afe7ef665ac80d64af9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_character/foo90.f90 @@ -0,0 +1,48 @@ +module f90_return_char + contains + function t0(value) + character :: value + character :: t0 + t0 = value + end function t0 + function t1(value) + character(len=1) :: value + character(len=1) :: t1 + t1 = value + end function t1 + function t5(value) + character(len=5) :: value + character(len=5) :: t5 + t5 = value + end function t5 + function ts(value) + character(len=*) :: value + character(len=10) :: ts + ts = value + end function ts + + subroutine s0(t0,value) + character :: value + character :: t0 +!f2py intent(out) t0 + t0 = value + end subroutine s0 + subroutine s1(t1,value) + character(len=1) :: value + character(len=1) :: t1 +!f2py intent(out) t1 + t1 = value + end subroutine s1 + subroutine s5(t5,value) + character(len=5) :: value + character(len=5) :: t5 +!f2py intent(out) t5 + t5 = value + end subroutine s5 + subroutine ss(ts,value) + character(len=*) :: value + character(len=10) :: ts +!f2py intent(out) ts + ts = value + end subroutine ss +end module f90_return_char diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo77.f new file mode 100644 index 0000000000000000000000000000000000000000..37a1ec845ecacc7fbc228f1ee5f628ec73075c12 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo77.f @@ -0,0 +1,45 @@ + function t0(value) + complex value + complex t0 + t0 = value + end + function t8(value) + complex*8 value + complex*8 t8 + t8 = value + end + function t16(value) + complex*16 value + complex*16 t16 + t16 = value + end + function td(value) + double complex value + double complex td + td = value + end + + subroutine s0(t0,value) + complex value + complex t0 +cf2py intent(out) t0 + t0 = value + end + subroutine s8(t8,value) + complex*8 value + complex*8 t8 +cf2py intent(out) t8 + t8 = value + end + subroutine s16(t16,value) + complex*16 value + complex*16 t16 +cf2py intent(out) t16 + t16 = value + end + subroutine sd(td,value) + double complex value + double complex td +cf2py intent(out) td + td = value + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo90.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo90.f90 new file mode 100644 index 0000000000000000000000000000000000000000..adc27b470538bc663416fb512a29c4b2bbe8d3dd --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_complex/foo90.f90 @@ -0,0 +1,48 @@ +module f90_return_complex + contains + function t0(value) + complex :: value + complex :: t0 + t0 = value + end function t0 + function t8(value) + complex(kind=4) :: value + complex(kind=4) :: t8 + t8 = value + end function t8 + function t16(value) + complex(kind=8) :: value + complex(kind=8) :: t16 + t16 = value + end function t16 + function td(value) + double complex :: value + double complex :: td + td = value + end function td + + subroutine s0(t0,value) + complex :: value + complex :: t0 +!f2py intent(out) t0 + t0 = value + end subroutine s0 + subroutine s8(t8,value) + complex(kind=4) :: value + complex(kind=4) :: t8 +!f2py intent(out) t8 + t8 = value + end subroutine s8 + subroutine s16(t16,value) + complex(kind=8) :: value + complex(kind=8) :: t16 +!f2py intent(out) t16 + t16 = value + end subroutine s16 + subroutine sd(td,value) + double complex :: value + double complex :: td +!f2py intent(out) td + td = value + end subroutine sd +end module f90_return_complex diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo77.f new file mode 100644 index 0000000000000000000000000000000000000000..1ab895b9ac340ca91c5d3a4080334bab9f033a55 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo77.f @@ -0,0 +1,56 @@ + function t0(value) + integer value + integer t0 + t0 = value + end + function t1(value) + integer*1 value + integer*1 t1 + t1 = value + end + function t2(value) + integer*2 value + integer*2 t2 + t2 = value + end + function t4(value) + integer*4 value + integer*4 t4 + t4 = value + end + function t8(value) + integer*8 value + integer*8 t8 + t8 = value + end + + subroutine s0(t0,value) + integer value + integer t0 +cf2py intent(out) t0 + t0 = value + end + subroutine s1(t1,value) + integer*1 value + integer*1 t1 +cf2py intent(out) t1 + t1 = value + end + subroutine s2(t2,value) + integer*2 value + integer*2 t2 +cf2py intent(out) t2 + t2 = value + end + subroutine s4(t4,value) + integer*4 value + integer*4 t4 +cf2py intent(out) t4 + t4 = value + end + subroutine s8(t8,value) + integer*8 value + integer*8 t8 +cf2py intent(out) t8 + t8 = value + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo90.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo90.f90 new file mode 100644 index 0000000000000000000000000000000000000000..ba9249aa20f941dbf00f060ad5d7e8820745b0f4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_integer/foo90.f90 @@ -0,0 +1,59 @@ +module f90_return_integer + contains + function t0(value) + integer :: value + integer :: t0 + t0 = value + end function t0 + function t1(value) + integer(kind=1) :: value + integer(kind=1) :: t1 + t1 = value + end function t1 + function t2(value) + integer(kind=2) :: value + integer(kind=2) :: t2 + t2 = value + end function t2 + function t4(value) + integer(kind=4) :: value + integer(kind=4) :: t4 + t4 = value + end function t4 + function t8(value) + integer(kind=8) :: value + integer(kind=8) :: t8 + t8 = value + end function t8 + + subroutine s0(t0,value) + integer :: value + integer :: t0 +!f2py intent(out) t0 + t0 = value + end subroutine s0 + subroutine s1(t1,value) + integer(kind=1) :: value + integer(kind=1) :: t1 +!f2py intent(out) t1 + t1 = value + end subroutine s1 + subroutine s2(t2,value) + integer(kind=2) :: value + integer(kind=2) :: t2 +!f2py intent(out) t2 + t2 = value + end subroutine s2 + subroutine s4(t4,value) + integer(kind=4) :: value + integer(kind=4) :: t4 +!f2py intent(out) t4 + t4 = value + end subroutine s4 + subroutine s8(t8,value) + integer(kind=8) :: value + integer(kind=8) :: t8 +!f2py intent(out) t8 + t8 = value + end subroutine s8 +end module f90_return_integer diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo77.f new file mode 100644 index 0000000000000000000000000000000000000000..ef530145fedf8b5cf3a05bdf0a46a4e22150007b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo77.f @@ -0,0 +1,56 @@ + function t0(value) + logical value + logical t0 + t0 = value + end + function t1(value) + logical*1 value + logical*1 t1 + t1 = value + end + function t2(value) + logical*2 value + logical*2 t2 + t2 = value + end + function t4(value) + logical*4 value + logical*4 t4 + t4 = value + end +c function t8(value) +c logical*8 value +c logical*8 t8 +c t8 = value +c end + + subroutine s0(t0,value) + logical value + logical t0 +cf2py intent(out) t0 + t0 = value + end + subroutine s1(t1,value) + logical*1 value + logical*1 t1 +cf2py intent(out) t1 + t1 = value + end + subroutine s2(t2,value) + logical*2 value + logical*2 t2 +cf2py intent(out) t2 + t2 = value + end + subroutine s4(t4,value) + logical*4 value + logical*4 t4 +cf2py intent(out) t4 + t4 = value + end +c subroutine s8(t8,value) +c logical*8 value +c logical*8 t8 +cf2py intent(out) t8 +c t8 = value +c end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo90.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo90.f90 new file mode 100644 index 0000000000000000000000000000000000000000..a4526468e3719140f0ed7d50a5f3a31d78d1d2de --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_logical/foo90.f90 @@ -0,0 +1,59 @@ +module f90_return_logical + contains + function t0(value) + logical :: value + logical :: t0 + t0 = value + end function t0 + function t1(value) + logical(kind=1) :: value + logical(kind=1) :: t1 + t1 = value + end function t1 + function t2(value) + logical(kind=2) :: value + logical(kind=2) :: t2 + t2 = value + end function t2 + function t4(value) + logical(kind=4) :: value + logical(kind=4) :: t4 + t4 = value + end function t4 + function t8(value) + logical(kind=8) :: value + logical(kind=8) :: t8 + t8 = value + end function t8 + + subroutine s0(t0,value) + logical :: value + logical :: t0 +!f2py intent(out) t0 + t0 = value + end subroutine s0 + subroutine s1(t1,value) + logical(kind=1) :: value + logical(kind=1) :: t1 +!f2py intent(out) t1 + t1 = value + end subroutine s1 + subroutine s2(t2,value) + logical(kind=2) :: value + logical(kind=2) :: t2 +!f2py intent(out) t2 + t2 = value + end subroutine s2 + subroutine s4(t4,value) + logical(kind=4) :: value + logical(kind=4) :: t4 +!f2py intent(out) t4 + t4 = value + end subroutine s4 + subroutine s8(t8,value) + logical(kind=8) :: value + logical(kind=8) :: t8 +!f2py intent(out) t8 + t8 = value + end subroutine s8 +end module f90_return_logical diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo77.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo77.f new file mode 100644 index 0000000000000000000000000000000000000000..bf43dbf11773d8282f3b9a7d7c4ba9da23ee6f27 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo77.f @@ -0,0 +1,45 @@ + function t0(value) + real value + real t0 + t0 = value + end + function t4(value) + real*4 value + real*4 t4 + t4 = value + end + function t8(value) + real*8 value + real*8 t8 + t8 = value + end + function td(value) + double precision value + double precision td + td = value + end + + subroutine s0(t0,value) + real value + real t0 +cf2py intent(out) t0 + t0 = value + end + subroutine s4(t4,value) + real*4 value + real*4 t4 +cf2py intent(out) t4 + t4 = value + end + subroutine s8(t8,value) + real*8 value + real*8 t8 +cf2py intent(out) t8 + t8 = value + end + subroutine sd(td,value) + double precision value + double precision td +cf2py intent(out) td + td = value + end diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo90.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo90.f90 new file mode 100644 index 0000000000000000000000000000000000000000..df9719980f2861678d5c1e4b0529a9eb0e375021 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/return_real/foo90.f90 @@ -0,0 +1,48 @@ +module f90_return_real + contains + function t0(value) + real :: value + real :: t0 + t0 = value + end function t0 + function t4(value) + real(kind=4) :: value + real(kind=4) :: t4 + t4 = value + end function t4 + function t8(value) + real(kind=8) :: value + real(kind=8) :: t8 + t8 = value + end function t8 + function td(value) + double precision :: value + double precision :: td + td = value + end function td + + subroutine s0(t0,value) + real :: value + real :: t0 +!f2py intent(out) t0 + t0 = value + end subroutine s0 + subroutine s4(t4,value) + real(kind=4) :: value + real(kind=4) :: t4 +!f2py intent(out) t4 + t4 = value + end subroutine s4 + subroutine s8(t8,value) + real(kind=8) :: value + real(kind=8) :: t8 +!f2py intent(out) t8 + t8 = value + end subroutine s8 + subroutine sd(td,value) + double precision :: value + double precision :: td +!f2py intent(out) td + td = value + end subroutine sd +end module f90_return_real diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.f new file mode 100644 index 0000000000000000000000000000000000000000..89be972d341966170c9f1a36af1f098e62dd4174 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.f @@ -0,0 +1,5 @@ + REAL*8 FUNCTION FUNCFORTRANNAME(A,B) + REAL*8 A, B + FUNCFORTRANNAME = A + B + RETURN + END FUNCTION diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.pyf new file mode 100644 index 0000000000000000000000000000000000000000..8730ca6a67edd4dea8df456369e5f8ef4cb5f831 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/funcfortranname.pyf @@ -0,0 +1,11 @@ +python module funcfortranname ! in + interface ! in :funcfortranname + function funcfortranname_default(a,b) ! in :funcfortranname:funcfortranname.f + fortranname funcfortranname + real*8 :: a + real*8 :: b + real*8 :: funcfortranname_default + real*8, intent(out) :: funcfortranname + end function funcfortranname_default + end interface +end python module funcfortranname diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.f new file mode 100644 index 0000000000000000000000000000000000000000..1d1eeaeb5a4588141e15c19b3b8e70fef6425f1a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.f @@ -0,0 +1,4 @@ + SUBROUTINE SUBROUT(A,B,C) + REAL*8 A, B, C + C = A + B + END SUBROUTINE diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.pyf new file mode 100644 index 0000000000000000000000000000000000000000..e27cbe1c7455d47ea312966c997bbcf5f690be74 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/routines/subrout.pyf @@ -0,0 +1,10 @@ +python module subrout ! in + interface ! in :subrout + subroutine subrout_default(a,b,c) ! in :subrout:subrout.f + fortranname subrout + real*8 :: a + real*8 :: b + real*8, intent(out) :: c + end subroutine subrout_default + end interface +end python module subrout diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/size/foo.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/size/foo.f90 new file mode 100644 index 0000000000000000000000000000000000000000..5b66f8c430d79a8438ad062466a97cf8c00dfb16 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/size/foo.f90 @@ -0,0 +1,44 @@ + +subroutine foo(a, n, m, b) + implicit none + + real, intent(in) :: a(n, m) + integer, intent(in) :: n, m + real, intent(out) :: b(size(a, 1)) + + integer :: i + + do i = 1, size(b) + b(i) = sum(a(i,:)) + enddo +end subroutine + +subroutine trans(x,y) + implicit none + real, intent(in), dimension(:,:) :: x + real, intent(out), dimension( size(x,2), size(x,1) ) :: y + integer :: N, M, i, j + N = size(x,1) + M = size(x,2) + DO i=1,N + do j=1,M + y(j,i) = x(i,j) + END DO + END DO +end subroutine trans + +subroutine flatten(x,y) + implicit none + real, intent(in), dimension(:,:) :: x + real, intent(out), dimension( size(x) ) :: y + integer :: N, M, i, j, k + N = size(x,1) + M = size(x,2) + k = 1 + DO i=1,N + do j=1,M + y(k) = x(i,j) + k = k + 1 + END DO + END DO +end subroutine flatten diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/char.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/char.f90 new file mode 100644 index 0000000000000000000000000000000000000000..bb7985ce50f2aa252aaca96aba6ef5d5f5d51844 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/char.f90 @@ -0,0 +1,29 @@ +MODULE char_test + +CONTAINS + +SUBROUTINE change_strings(strings, n_strs, out_strings) + IMPLICIT NONE + + ! Inputs + INTEGER, INTENT(IN) :: n_strs + CHARACTER, INTENT(IN), DIMENSION(2,n_strs) :: strings + CHARACTER, INTENT(OUT), DIMENSION(2,n_strs) :: out_strings + +!f2py INTEGER, INTENT(IN) :: n_strs +!f2py CHARACTER, INTENT(IN), DIMENSION(2,n_strs) :: strings +!f2py CHARACTER, INTENT(OUT), DIMENSION(2,n_strs) :: strings + + ! Misc. + INTEGER*4 :: j + + + DO j=1, n_strs + out_strings(1,j) = strings(1,j) + out_strings(2,j) = 'A' + END DO + +END SUBROUTINE change_strings + +END MODULE char_test + diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/fixed_string.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/fixed_string.f90 new file mode 100644 index 0000000000000000000000000000000000000000..7fd1585430fb05f84fb850ef4656d94e8a0804e9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/fixed_string.f90 @@ -0,0 +1,34 @@ +function sint(s) result(i) + implicit none + character(len=*) :: s + integer :: j, i + i = 0 + do j=len(s), 1, -1 + if (.not.((i.eq.0).and.(s(j:j).eq.' '))) then + i = i + ichar(s(j:j)) * 10 ** (j - 1) + endif + end do + return + end function sint + + function test_in_bytes4(a) result (i) + implicit none + integer :: sint + character(len=4) :: a + integer :: i + i = sint(a) + a(1:1) = 'A' + return + end function test_in_bytes4 + + function test_inout_bytes4(a) result (i) + implicit none + integer :: sint + character(len=4), intent(inout) :: a + integer :: i + if (a(1:1).ne.' ') then + a(1:1) = 'E' + endif + i = sint(a) + return + end function test_inout_bytes4 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24008.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24008.f new file mode 100644 index 0000000000000000000000000000000000000000..ab64cf771f68bbcecc8ac2d5d649545fc357e15b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24008.f @@ -0,0 +1,8 @@ + SUBROUTINE GREET(NAME, GREETING) + CHARACTER NAME*(*), GREETING*(*) + CHARACTER*(50) MESSAGE + + MESSAGE = 'Hello, ' // NAME // ', ' // GREETING +c$$$ PRINT *, MESSAGE + + END SUBROUTINE GREET diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24662.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24662.f90 new file mode 100644 index 0000000000000000000000000000000000000000..ca53413cc9b48f1c8d476d329eb4b773695dd08c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh24662.f90 @@ -0,0 +1,7 @@ +subroutine string_inout_optional(output) + implicit none + character*(32), optional, intent(inout) :: output + if (present(output)) then + output="output string" + endif +end subroutine diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.f90 new file mode 100644 index 0000000000000000000000000000000000000000..db1c7100d2ab812de5d212c1bbd255cf2ae60be3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.f90 @@ -0,0 +1,14 @@ +subroutine charint(trans, info) + character, intent(in) :: trans + integer, intent(out) :: info + if (trans == 'N') then + info = 1 + else if (trans == 'T') then + info = 2 + else if (trans == 'C') then + info = 3 + else + info = -1 + end if + +end subroutine charint diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.pyf new file mode 100644 index 0000000000000000000000000000000000000000..7b9609071bce3e775703b12c430f411af09e6eee --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286.pyf @@ -0,0 +1,12 @@ +python module _char_handling_test + interface + subroutine charint(trans, info) + callstatement (*f2py_func)(&trans, &info) + callprotoargument char*, int* + + character, intent(in), check(trans=='N'||trans=='T'||trans=='C') :: trans = 'N' + integer intent(out) :: info + + end subroutine charint + end interface +end python module _char_handling_test diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286_bc.pyf b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286_bc.pyf new file mode 100644 index 0000000000000000000000000000000000000000..e7b10fa9215e88e56794e9c73d0b13872cbd953c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/gh25286_bc.pyf @@ -0,0 +1,12 @@ +python module _char_handling_test + interface + subroutine charint(trans, info) + callstatement (*f2py_func)(&trans, &info) + callprotoargument char*, int* + + character, intent(in), check(*trans=='N'||*trans=='T'||*trans=='C') :: trans = 'N' + integer intent(out) :: info + + end subroutine charint + end interface +end python module _char_handling_test diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/scalar_string.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/scalar_string.f90 new file mode 100644 index 0000000000000000000000000000000000000000..f8f076172ab48ca4834d631b362f47ca374db5e4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/scalar_string.f90 @@ -0,0 +1,9 @@ +MODULE string_test + + character(len=8) :: string + character string77 * 8 + + character(len=12), dimension(5,7) :: strarr + character strarr77(5,7) * 12 + +END MODULE string_test diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/string.f b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/string.f new file mode 100644 index 0000000000000000000000000000000000000000..5210ca4dc054de60488e45baa12b6c8bc89fc9eb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/string/string.f @@ -0,0 +1,12 @@ +C FILE: STRING.F + SUBROUTINE FOO(A,B,C,D) + CHARACTER*5 A, B + CHARACTER*(*) C,D +Cf2py intent(in) a,c +Cf2py intent(inout) b,d + A(1:1) = 'A' + B(1:1) = 'B' + C(1:1) = 'C' + D(1:1) = 'D' + END +C END OF FILE STRING.F diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/value_attrspec/gh21665.f90 b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/value_attrspec/gh21665.f90 new file mode 100644 index 0000000000000000000000000000000000000000..7d9dc0fd4acbc081f55edfafb5dea981dcf279d5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/src/value_attrspec/gh21665.f90 @@ -0,0 +1,9 @@ +module fortfuncs + implicit none +contains + subroutine square(x,y) + integer, intent(in), value :: x + integer, intent(out) :: y + y = x*x + end subroutine square +end module fortfuncs diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_abstract_interface.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_abstract_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..0bc38b51f95d71dcd4a4b9c723211e1c4398c966 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_abstract_interface.py @@ -0,0 +1,24 @@ +import pytest +from . import util +from numpy.f2py import crackfortran +from numpy.testing import IS_WASM + + +@pytest.mark.skipif(IS_WASM, reason="Cannot start subprocess") +@pytest.mark.slow +class TestAbstractInterface(util.F2PyTest): + sources = [util.getpath("tests", "src", "abstract_interface", "foo.f90")] + + skip = ["add1", "add2"] + + def test_abstract_interface(self): + assert self.module.ops_module.foo(3, 5) == (8, 13) + + def test_parse_abstract_interface(self): + # Test gh18403 + fpath = util.getpath("tests", "src", "abstract_interface", + "gh18403_mod.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + assert len(mod[0]["body"]) == 1 + assert mod[0]["body"][0]["block"] == "abstract interface" diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_array_from_pyobj.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_array_from_pyobj.py new file mode 100644 index 0000000000000000000000000000000000000000..41ed2c7a0dfe7e40fcb99553757c41c0dfde7349 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_array_from_pyobj.py @@ -0,0 +1,677 @@ +import sys +import copy +import platform +import pytest +from pathlib import Path + +import numpy as np + +from numpy._core._type_aliases import c_names_dict as _c_names_dict +from . import util + +wrap = None + +# Extend core typeinfo with CHARACTER to test dtype('c') +c_names_dict = dict( + CHARACTER=np.dtype("c"), + **_c_names_dict +) + + +def get_testdir(): + testroot = Path(__file__).resolve().parent / "src" + return testroot / "array_from_pyobj" + +def setup_module(): + """ + Build the required testing extension module + + """ + global wrap + + if wrap is None: + src = [ + get_testdir() / "wrapmodule.c", + ] + wrap = util.build_meson(src, module_name = "test_array_from_pyobj_ext") + + +def flags_info(arr): + flags = wrap.array_attrs(arr)[6] + return flags2names(flags) + + +def flags2names(flags): + info = [] + for flagname in [ + "CONTIGUOUS", + "FORTRAN", + "OWNDATA", + "ENSURECOPY", + "ENSUREARRAY", + "ALIGNED", + "NOTSWAPPED", + "WRITEABLE", + "WRITEBACKIFCOPY", + "UPDATEIFCOPY", + "BEHAVED", + "BEHAVED_RO", + "CARRAY", + "FARRAY", + ]: + if abs(flags) & getattr(wrap, flagname, 0): + info.append(flagname) + return info + + +class Intent: + def __init__(self, intent_list=[]): + self.intent_list = intent_list[:] + flags = 0 + for i in intent_list: + if i == "optional": + flags |= wrap.F2PY_OPTIONAL + else: + flags |= getattr(wrap, "F2PY_INTENT_" + i.upper()) + self.flags = flags + + def __getattr__(self, name): + name = name.lower() + if name == "in_": + name = "in" + return self.__class__(self.intent_list + [name]) + + def __str__(self): + return "intent(%s)" % (",".join(self.intent_list)) + + def __repr__(self): + return "Intent(%r)" % (self.intent_list) + + def is_intent(self, *names): + return all(name in self.intent_list for name in names) + + def is_intent_exact(self, *names): + return len(self.intent_list) == len(names) and self.is_intent(*names) + + +intent = Intent() + +_type_names = [ + "BOOL", + "BYTE", + "UBYTE", + "SHORT", + "USHORT", + "INT", + "UINT", + "LONG", + "ULONG", + "LONGLONG", + "ULONGLONG", + "FLOAT", + "DOUBLE", + "CFLOAT", + "STRING1", + "STRING5", + "CHARACTER", +] + +_cast_dict = {"BOOL": ["BOOL"]} +_cast_dict["BYTE"] = _cast_dict["BOOL"] + ["BYTE"] +_cast_dict["UBYTE"] = _cast_dict["BOOL"] + ["UBYTE"] +_cast_dict["BYTE"] = ["BYTE"] +_cast_dict["UBYTE"] = ["UBYTE"] +_cast_dict["SHORT"] = _cast_dict["BYTE"] + ["UBYTE", "SHORT"] +_cast_dict["USHORT"] = _cast_dict["UBYTE"] + ["BYTE", "USHORT"] +_cast_dict["INT"] = _cast_dict["SHORT"] + ["USHORT", "INT"] +_cast_dict["UINT"] = _cast_dict["USHORT"] + ["SHORT", "UINT"] + +_cast_dict["LONG"] = _cast_dict["INT"] + ["LONG"] +_cast_dict["ULONG"] = _cast_dict["UINT"] + ["ULONG"] + +_cast_dict["LONGLONG"] = _cast_dict["LONG"] + ["LONGLONG"] +_cast_dict["ULONGLONG"] = _cast_dict["ULONG"] + ["ULONGLONG"] + +_cast_dict["FLOAT"] = _cast_dict["SHORT"] + ["USHORT", "FLOAT"] +_cast_dict["DOUBLE"] = _cast_dict["INT"] + ["UINT", "FLOAT", "DOUBLE"] + +_cast_dict["CFLOAT"] = _cast_dict["FLOAT"] + ["CFLOAT"] + +_cast_dict['STRING1'] = ['STRING1'] +_cast_dict['STRING5'] = ['STRING5'] +_cast_dict['CHARACTER'] = ['CHARACTER'] + +# 32 bit system malloc typically does not provide the alignment required by +# 16 byte long double types this means the inout intent cannot be satisfied +# and several tests fail as the alignment flag can be randomly true or false +# when numpy gains an aligned allocator the tests could be enabled again +# +# Furthermore, on macOS ARM64, LONGDOUBLE is an alias for DOUBLE. +if ((np.intp().dtype.itemsize != 4 or np.clongdouble().dtype.alignment <= 8) + and sys.platform != "win32" + and (platform.system(), platform.processor()) != ("Darwin", "arm")): + _type_names.extend(["LONGDOUBLE", "CDOUBLE", "CLONGDOUBLE"]) + _cast_dict["LONGDOUBLE"] = _cast_dict["LONG"] + [ + "ULONG", + "FLOAT", + "DOUBLE", + "LONGDOUBLE", + ] + _cast_dict["CLONGDOUBLE"] = _cast_dict["LONGDOUBLE"] + [ + "CFLOAT", + "CDOUBLE", + "CLONGDOUBLE", + ] + _cast_dict["CDOUBLE"] = _cast_dict["DOUBLE"] + ["CFLOAT", "CDOUBLE"] + + +class Type: + _type_cache = {} + + def __new__(cls, name): + if isinstance(name, np.dtype): + dtype0 = name + name = None + for n, i in c_names_dict.items(): + if not isinstance(i, type) and dtype0.type is i.type: + name = n + break + obj = cls._type_cache.get(name.upper(), None) + if obj is not None: + return obj + obj = object.__new__(cls) + obj._init(name) + cls._type_cache[name.upper()] = obj + return obj + + def _init(self, name): + self.NAME = name.upper() + + if self.NAME == 'CHARACTER': + info = c_names_dict[self.NAME] + self.type_num = wrap.NPY_STRING + self.elsize = 1 + self.dtype = np.dtype('c') + elif self.NAME.startswith('STRING'): + info = c_names_dict[self.NAME[:6]] + self.type_num = wrap.NPY_STRING + self.elsize = int(self.NAME[6:] or 0) + self.dtype = np.dtype(f'S{self.elsize}') + else: + info = c_names_dict[self.NAME] + self.type_num = getattr(wrap, 'NPY_' + self.NAME) + self.elsize = info.itemsize + self.dtype = np.dtype(info.type) + + assert self.type_num == info.num + self.type = info.type + self.dtypechar = info.char + + def __repr__(self): + return (f"Type({self.NAME})|type_num={self.type_num}," + f" dtype={self.dtype}," + f" type={self.type}, elsize={self.elsize}," + f" dtypechar={self.dtypechar}") + + def cast_types(self): + return [self.__class__(_m) for _m in _cast_dict[self.NAME]] + + def all_types(self): + return [self.__class__(_m) for _m in _type_names] + + def smaller_types(self): + bits = c_names_dict[self.NAME].alignment + types = [] + for name in _type_names: + if c_names_dict[name].alignment < bits: + types.append(Type(name)) + return types + + def equal_types(self): + bits = c_names_dict[self.NAME].alignment + types = [] + for name in _type_names: + if name == self.NAME: + continue + if c_names_dict[name].alignment == bits: + types.append(Type(name)) + return types + + def larger_types(self): + bits = c_names_dict[self.NAME].alignment + types = [] + for name in _type_names: + if c_names_dict[name].alignment > bits: + types.append(Type(name)) + return types + + +class Array: + + def __repr__(self): + return (f'Array({self.type}, {self.dims}, {self.intent},' + f' {self.obj})|arr={self.arr}') + + def __init__(self, typ, dims, intent, obj): + self.type = typ + self.dims = dims + self.intent = intent + self.obj_copy = copy.deepcopy(obj) + self.obj = obj + + # arr.dtypechar may be different from typ.dtypechar + self.arr = wrap.call(typ.type_num, + typ.elsize, + dims, intent.flags, obj) + + assert isinstance(self.arr, np.ndarray) + + self.arr_attr = wrap.array_attrs(self.arr) + + if len(dims) > 1: + if self.intent.is_intent("c"): + assert (intent.flags & wrap.F2PY_INTENT_C) + assert not self.arr.flags["FORTRAN"] + assert self.arr.flags["CONTIGUOUS"] + assert (not self.arr_attr[6] & wrap.FORTRAN) + else: + assert (not intent.flags & wrap.F2PY_INTENT_C) + assert self.arr.flags["FORTRAN"] + assert not self.arr.flags["CONTIGUOUS"] + assert (self.arr_attr[6] & wrap.FORTRAN) + + if obj is None: + self.pyarr = None + self.pyarr_attr = None + return + + if intent.is_intent("cache"): + assert isinstance(obj, np.ndarray), repr(type(obj)) + self.pyarr = np.array(obj).reshape(*dims).copy() + else: + self.pyarr = np.array( + np.array(obj, dtype=typ.dtypechar).reshape(*dims), + order=self.intent.is_intent("c") and "C" or "F", + ) + assert self.pyarr.dtype == typ + self.pyarr.setflags(write=self.arr.flags["WRITEABLE"]) + assert self.pyarr.flags["OWNDATA"], (obj, intent) + self.pyarr_attr = wrap.array_attrs(self.pyarr) + + if len(dims) > 1: + if self.intent.is_intent("c"): + assert not self.pyarr.flags["FORTRAN"] + assert self.pyarr.flags["CONTIGUOUS"] + assert (not self.pyarr_attr[6] & wrap.FORTRAN) + else: + assert self.pyarr.flags["FORTRAN"] + assert not self.pyarr.flags["CONTIGUOUS"] + assert (self.pyarr_attr[6] & wrap.FORTRAN) + + assert self.arr_attr[1] == self.pyarr_attr[1] # nd + assert self.arr_attr[2] == self.pyarr_attr[2] # dimensions + if self.arr_attr[1] <= 1: + assert self.arr_attr[3] == self.pyarr_attr[3], repr(( + self.arr_attr[3], + self.pyarr_attr[3], + self.arr.tobytes(), + self.pyarr.tobytes(), + )) # strides + assert self.arr_attr[5][-2:] == self.pyarr_attr[5][-2:], repr(( + self.arr_attr[5], self.pyarr_attr[5] + )) # descr + assert self.arr_attr[6] == self.pyarr_attr[6], repr(( + self.arr_attr[6], + self.pyarr_attr[6], + flags2names(0 * self.arr_attr[6] - self.pyarr_attr[6]), + flags2names(self.arr_attr[6]), + intent, + )) # flags + + if intent.is_intent("cache"): + assert self.arr_attr[5][3] >= self.type.elsize + else: + assert self.arr_attr[5][3] == self.type.elsize + assert (self.arr_equal(self.pyarr, self.arr)) + + if isinstance(self.obj, np.ndarray): + if typ.elsize == Type(obj.dtype).elsize: + if not intent.is_intent("copy") and self.arr_attr[1] <= 1: + assert self.has_shared_memory() + + def arr_equal(self, arr1, arr2): + if arr1.shape != arr2.shape: + return False + return (arr1 == arr2).all() + + def __str__(self): + return str(self.arr) + + def has_shared_memory(self): + """Check that created array shares data with input array.""" + if self.obj is self.arr: + return True + if not isinstance(self.obj, np.ndarray): + return False + obj_attr = wrap.array_attrs(self.obj) + return obj_attr[0] == self.arr_attr[0] + + +class TestIntent: + def test_in_out(self): + assert str(intent.in_.out) == "intent(in,out)" + assert intent.in_.c.is_intent("c") + assert not intent.in_.c.is_intent_exact("c") + assert intent.in_.c.is_intent_exact("c", "in") + assert intent.in_.c.is_intent_exact("in", "c") + assert not intent.in_.is_intent("c") + + +class TestSharedMemory: + + @pytest.fixture(autouse=True, scope="class", params=_type_names) + def setup_type(self, request): + request.cls.type = Type(request.param) + request.cls.array = lambda self, dims, intent, obj: Array( + Type(request.param), dims, intent, obj) + + @property + def num2seq(self): + if self.type.NAME.startswith('STRING'): + elsize = self.type.elsize + return ['1' * elsize, '2' * elsize] + return [1, 2] + + @property + def num23seq(self): + if self.type.NAME.startswith('STRING'): + elsize = self.type.elsize + return [['1' * elsize, '2' * elsize, '3' * elsize], + ['4' * elsize, '5' * elsize, '6' * elsize]] + return [[1, 2, 3], [4, 5, 6]] + + def test_in_from_2seq(self): + a = self.array([2], intent.in_, self.num2seq) + assert not a.has_shared_memory() + + def test_in_from_2casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num2seq, dtype=t.dtype) + a = self.array([len(self.num2seq)], intent.in_, obj) + if t.elsize == self.type.elsize: + assert a.has_shared_memory(), repr((self.type.dtype, t.dtype)) + else: + assert not a.has_shared_memory() + + @pytest.mark.parametrize("write", ["w", "ro"]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("inp", ["2seq", "23seq"]) + def test_in_nocopy(self, write, order, inp): + """Test if intent(in) array can be passed without copies""" + seq = getattr(self, "num" + inp) + obj = np.array(seq, dtype=self.type.dtype, order=order) + obj.setflags(write=(write == 'w')) + a = self.array(obj.shape, + ((order == 'C' and intent.in_.c) or intent.in_), obj) + assert a.has_shared_memory() + + def test_inout_2seq(self): + obj = np.array(self.num2seq, dtype=self.type.dtype) + a = self.array([len(self.num2seq)], intent.inout, obj) + assert a.has_shared_memory() + + try: + a = self.array([2], intent.in_.inout, self.num2seq) + except TypeError as msg: + if not str(msg).startswith( + "failed to initialize intent(inout|inplace|cache) array"): + raise + else: + raise SystemError("intent(inout) should have failed on sequence") + + def test_f_inout_23seq(self): + obj = np.array(self.num23seq, dtype=self.type.dtype, order="F") + shape = (len(self.num23seq), len(self.num23seq[0])) + a = self.array(shape, intent.in_.inout, obj) + assert a.has_shared_memory() + + obj = np.array(self.num23seq, dtype=self.type.dtype, order="C") + shape = (len(self.num23seq), len(self.num23seq[0])) + try: + a = self.array(shape, intent.in_.inout, obj) + except ValueError as msg: + if not str(msg).startswith( + "failed to initialize intent(inout) array"): + raise + else: + raise SystemError( + "intent(inout) should have failed on improper array") + + def test_c_inout_23seq(self): + obj = np.array(self.num23seq, dtype=self.type.dtype) + shape = (len(self.num23seq), len(self.num23seq[0])) + a = self.array(shape, intent.in_.c.inout, obj) + assert a.has_shared_memory() + + def test_in_copy_from_2casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num2seq, dtype=t.dtype) + a = self.array([len(self.num2seq)], intent.in_.copy, obj) + assert not a.has_shared_memory() + + def test_c_in_from_23seq(self): + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_, + self.num23seq) + assert not a.has_shared_memory() + + def test_in_from_23casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num23seq, dtype=t.dtype) + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_, obj) + assert not a.has_shared_memory() + + def test_f_in_from_23casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num23seq, dtype=t.dtype, order="F") + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_, obj) + if t.elsize == self.type.elsize: + assert a.has_shared_memory() + else: + assert not a.has_shared_memory() + + def test_c_in_from_23casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num23seq, dtype=t.dtype) + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_.c, obj) + if t.elsize == self.type.elsize: + assert a.has_shared_memory() + else: + assert not a.has_shared_memory() + + def test_f_copy_in_from_23casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num23seq, dtype=t.dtype, order="F") + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_.copy, + obj) + assert not a.has_shared_memory() + + def test_c_copy_in_from_23casttype(self): + for t in self.type.cast_types(): + obj = np.array(self.num23seq, dtype=t.dtype) + a = self.array( + [len(self.num23seq), len(self.num23seq[0])], intent.in_.c.copy, + obj) + assert not a.has_shared_memory() + + def test_in_cache_from_2casttype(self): + for t in self.type.all_types(): + if t.elsize != self.type.elsize: + continue + obj = np.array(self.num2seq, dtype=t.dtype) + shape = (len(self.num2seq), ) + a = self.array(shape, intent.in_.c.cache, obj) + assert a.has_shared_memory() + + a = self.array(shape, intent.in_.cache, obj) + assert a.has_shared_memory() + + obj = np.array(self.num2seq, dtype=t.dtype, order="F") + a = self.array(shape, intent.in_.c.cache, obj) + assert a.has_shared_memory() + + a = self.array(shape, intent.in_.cache, obj) + assert a.has_shared_memory(), repr(t.dtype) + + try: + a = self.array(shape, intent.in_.cache, obj[::-1]) + except ValueError as msg: + if not str(msg).startswith( + "failed to initialize intent(cache) array"): + raise + else: + raise SystemError( + "intent(cache) should have failed on multisegmented array") + + def test_in_cache_from_2casttype_failure(self): + for t in self.type.all_types(): + if t.NAME == 'STRING': + # string elsize is 0, so skipping the test + continue + if t.elsize >= self.type.elsize: + continue + is_int = np.issubdtype(t.dtype, np.integer) + if is_int and int(self.num2seq[0]) > np.iinfo(t.dtype).max: + # skip test if num2seq would trigger an overflow error + continue + obj = np.array(self.num2seq, dtype=t.dtype) + shape = (len(self.num2seq), ) + try: + self.array(shape, intent.in_.cache, obj) # Should succeed + except ValueError as msg: + if not str(msg).startswith( + "failed to initialize intent(cache) array"): + raise + else: + raise SystemError( + "intent(cache) should have failed on smaller array") + + def test_cache_hidden(self): + shape = (2, ) + a = self.array(shape, intent.cache.hide, None) + assert a.arr.shape == shape + + shape = (2, 3) + a = self.array(shape, intent.cache.hide, None) + assert a.arr.shape == shape + + shape = (-1, 3) + try: + a = self.array(shape, intent.cache.hide, None) + except ValueError as msg: + if not str(msg).startswith( + "failed to create intent(cache|hide)|optional array"): + raise + else: + raise SystemError( + "intent(cache) should have failed on undefined dimensions") + + def test_hidden(self): + shape = (2, ) + a = self.array(shape, intent.hide, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + + shape = (2, 3) + a = self.array(shape, intent.hide, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + assert a.arr.flags["FORTRAN"] and not a.arr.flags["CONTIGUOUS"] + + shape = (2, 3) + a = self.array(shape, intent.c.hide, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + assert not a.arr.flags["FORTRAN"] and a.arr.flags["CONTIGUOUS"] + + shape = (-1, 3) + try: + a = self.array(shape, intent.hide, None) + except ValueError as msg: + if not str(msg).startswith( + "failed to create intent(cache|hide)|optional array"): + raise + else: + raise SystemError( + "intent(hide) should have failed on undefined dimensions") + + def test_optional_none(self): + shape = (2, ) + a = self.array(shape, intent.optional, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + + shape = (2, 3) + a = self.array(shape, intent.optional, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + assert a.arr.flags["FORTRAN"] and not a.arr.flags["CONTIGUOUS"] + + shape = (2, 3) + a = self.array(shape, intent.c.optional, None) + assert a.arr.shape == shape + assert a.arr_equal(a.arr, np.zeros(shape, dtype=self.type.dtype)) + assert not a.arr.flags["FORTRAN"] and a.arr.flags["CONTIGUOUS"] + + def test_optional_from_2seq(self): + obj = self.num2seq + shape = (len(obj), ) + a = self.array(shape, intent.optional, obj) + assert a.arr.shape == shape + assert not a.has_shared_memory() + + def test_optional_from_23seq(self): + obj = self.num23seq + shape = (len(obj), len(obj[0])) + a = self.array(shape, intent.optional, obj) + assert a.arr.shape == shape + assert not a.has_shared_memory() + + a = self.array(shape, intent.optional.c, obj) + assert a.arr.shape == shape + assert not a.has_shared_memory() + + def test_inplace(self): + obj = np.array(self.num23seq, dtype=self.type.dtype) + assert not obj.flags["FORTRAN"] and obj.flags["CONTIGUOUS"] + shape = obj.shape + a = self.array(shape, intent.inplace, obj) + assert obj[1][2] == a.arr[1][2], repr((obj, a.arr)) + a.arr[1][2] = 54 + assert obj[1][2] == a.arr[1][2] == np.array(54, dtype=self.type.dtype) + assert a.arr is obj + assert obj.flags["FORTRAN"] # obj attributes are changed inplace! + assert not obj.flags["CONTIGUOUS"] + + def test_inplace_from_casttype(self): + for t in self.type.cast_types(): + if t is self.type: + continue + obj = np.array(self.num23seq, dtype=t.dtype) + assert obj.dtype.type == t.type + assert obj.dtype.type is not self.type.type + assert not obj.flags["FORTRAN"] and obj.flags["CONTIGUOUS"] + shape = obj.shape + a = self.array(shape, intent.inplace, obj) + assert obj[1][2] == a.arr[1][2], repr((obj, a.arr)) + a.arr[1][2] = 54 + assert obj[1][2] == a.arr[1][2] == np.array(54, + dtype=self.type.dtype) + assert a.arr is obj + assert obj.flags["FORTRAN"] # obj attributes changed inplace! + assert not obj.flags["CONTIGUOUS"] + assert obj.dtype.type is self.type.type # obj changed inplace! diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_assumed_shape.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_assumed_shape.py new file mode 100644 index 0000000000000000000000000000000000000000..d4664cf88cbe9701105a5d428332e3aa0d623930 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_assumed_shape.py @@ -0,0 +1,49 @@ +import os +import pytest +import tempfile + +from . import util + + +class TestAssumedShapeSumExample(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "assumed_shape", "foo_free.f90"), + util.getpath("tests", "src", "assumed_shape", "foo_use.f90"), + util.getpath("tests", "src", "assumed_shape", "precision.f90"), + util.getpath("tests", "src", "assumed_shape", "foo_mod.f90"), + util.getpath("tests", "src", "assumed_shape", ".f2py_f2cmap"), + ] + + @pytest.mark.slow + def test_all(self): + r = self.module.fsum([1, 2]) + assert r == 3 + r = self.module.sum([1, 2]) + assert r == 3 + r = self.module.sum_with_use([1, 2]) + assert r == 3 + + r = self.module.mod.sum([1, 2]) + assert r == 3 + r = self.module.mod.fsum([1, 2]) + assert r == 3 + + +class TestF2cmapOption(TestAssumedShapeSumExample): + def setup_method(self): + # Use a custom file name for .f2py_f2cmap + self.sources = list(self.sources) + f2cmap_src = self.sources.pop(-1) + + self.f2cmap_file = tempfile.NamedTemporaryFile(delete=False) + with open(f2cmap_src, "rb") as f: + self.f2cmap_file.write(f.read()) + self.f2cmap_file.close() + + self.sources.append(self.f2cmap_file.name) + self.options = ["--f2cmap", self.f2cmap_file.name] + + super().setup_method() + + def teardown_method(self): + os.unlink(self.f2cmap_file.name) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_block_docstring.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_block_docstring.py new file mode 100644 index 0000000000000000000000000000000000000000..16b5559e8e42da3fdcf0890fc08fb0a248bd69c1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_block_docstring.py @@ -0,0 +1,18 @@ +import sys +import pytest +from . import util + +from numpy.testing import IS_PYPY + + +@pytest.mark.slow +class TestBlockDocString(util.F2PyTest): + sources = [util.getpath("tests", "src", "block_docstring", "foo.f")] + + @pytest.mark.skipif(sys.platform == "win32", + reason="Fails with MinGW64 Gfortran (Issue #9673)") + @pytest.mark.xfail(IS_PYPY, + reason="PyPy cannot modify tp_doc after PyType_Ready") + def test_block_docstring(self): + expected = "bar : 'i'-array(2,3)\n" + assert self.module.block.__doc__ == expected diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_callback.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..4a9ed484a4a4a370f0c2d82d6a3d5b6f7ab5582e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_callback.py @@ -0,0 +1,261 @@ +import math +import textwrap +import sys +import pytest +import threading +import traceback +import time +import platform + +import numpy as np +from numpy.testing import IS_PYPY +from . import util + + +class TestF77Callback(util.F2PyTest): + sources = [util.getpath("tests", "src", "callback", "foo.f")] + + @pytest.mark.parametrize("name", "t,t2".split(",")) + @pytest.mark.slow + def test_all(self, name): + self.check_function(name) + + @pytest.mark.xfail(IS_PYPY, + reason="PyPy cannot modify tp_doc after PyType_Ready") + def test_docstring(self): + expected = textwrap.dedent("""\ + a = t(fun,[fun_extra_args]) + + Wrapper for ``t``. + + Parameters + ---------- + fun : call-back function + + Other Parameters + ---------------- + fun_extra_args : input tuple, optional + Default: () + + Returns + ------- + a : int + + Notes + ----- + Call-back functions:: + + def fun(): return a + Return objects: + a : int + """) + assert self.module.t.__doc__ == expected + + def check_function(self, name): + t = getattr(self.module, name) + r = t(lambda: 4) + assert r == 4 + r = t(lambda a: 5, fun_extra_args=(6, )) + assert r == 5 + r = t(lambda a: a, fun_extra_args=(6, )) + assert r == 6 + r = t(lambda a: 5 + a, fun_extra_args=(7, )) + assert r == 12 + r = t(lambda a: math.degrees(a), fun_extra_args=(math.pi, )) + assert r == 180 + r = t(math.degrees, fun_extra_args=(math.pi, )) + assert r == 180 + + r = t(self.module.func, fun_extra_args=(6, )) + assert r == 17 + r = t(self.module.func0) + assert r == 11 + r = t(self.module.func0._cpointer) + assert r == 11 + + class A: + def __call__(self): + return 7 + + def mth(self): + return 9 + + a = A() + r = t(a) + assert r == 7 + r = t(a.mth) + assert r == 9 + + @pytest.mark.skipif(sys.platform == 'win32', + reason='Fails with MinGW64 Gfortran (Issue #9673)') + def test_string_callback(self): + def callback(code): + if code == "r": + return 0 + else: + return 1 + + f = self.module.string_callback + r = f(callback) + assert r == 0 + + @pytest.mark.skipif(sys.platform == 'win32', + reason='Fails with MinGW64 Gfortran (Issue #9673)') + def test_string_callback_array(self): + # See gh-10027 + cu1 = np.zeros((1, ), "S8") + cu2 = np.zeros((1, 8), "c") + cu3 = np.array([""], "S8") + + def callback(cu, lencu): + if cu.shape != (lencu,): + return 1 + if cu.dtype != "S8": + return 2 + if not np.all(cu == b""): + return 3 + return 0 + + f = self.module.string_callback_array + for cu in [cu1, cu2, cu3]: + res = f(callback, cu, cu.size) + assert res == 0 + + def test_threadsafety(self): + # Segfaults if the callback handling is not threadsafe + + errors = [] + + def cb(): + # Sleep here to make it more likely for another thread + # to call their callback at the same time. + time.sleep(1e-3) + + # Check reentrancy + r = self.module.t(lambda: 123) + assert r == 123 + + return 42 + + def runner(name): + try: + for j in range(50): + r = self.module.t(cb) + assert r == 42 + self.check_function(name) + except Exception: + errors.append(traceback.format_exc()) + + threads = [ + threading.Thread(target=runner, args=(arg, )) + for arg in ("t", "t2") for n in range(20) + ] + + for t in threads: + t.start() + + for t in threads: + t.join() + + errors = "\n\n".join(errors) + if errors: + raise AssertionError(errors) + + def test_hidden_callback(self): + try: + self.module.hidden_callback(2) + except Exception as msg: + assert str(msg).startswith("Callback global_f not defined") + + try: + self.module.hidden_callback2(2) + except Exception as msg: + assert str(msg).startswith("cb: Callback global_f not defined") + + self.module.global_f = lambda x: x + 1 + r = self.module.hidden_callback(2) + assert r == 3 + + self.module.global_f = lambda x: x + 2 + r = self.module.hidden_callback(2) + assert r == 4 + + del self.module.global_f + try: + self.module.hidden_callback(2) + except Exception as msg: + assert str(msg).startswith("Callback global_f not defined") + + self.module.global_f = lambda x=0: x + 3 + r = self.module.hidden_callback(2) + assert r == 5 + + # reproducer of gh18341 + r = self.module.hidden_callback2(2) + assert r == 3 + + +class TestF77CallbackPythonTLS(TestF77Callback): + """ + Callback tests using Python thread-local storage instead of + compiler-provided + """ + + options = ["-DF2PY_USE_PYTHON_TLS"] + + +class TestF90Callback(util.F2PyTest): + sources = [util.getpath("tests", "src", "callback", "gh17797.f90")] + + @pytest.mark.slow + def test_gh17797(self): + def incr(x): + return x + 123 + + y = np.array([1, 2, 3], dtype=np.int64) + r = self.module.gh17797(incr, y) + assert r == 123 + 1 + 2 + 3 + + +class TestGH18335(util.F2PyTest): + """The reproduction of the reported issue requires specific input that + extensions may break the issue conditions, so the reproducer is + implemented as a separate test class. Do not extend this test with + other tests! + """ + sources = [util.getpath("tests", "src", "callback", "gh18335.f90")] + + @pytest.mark.slow + def test_gh18335(self): + def foo(x): + x[0] += 1 + + r = self.module.gh18335(foo) + assert r == 123 + 1 + + +class TestGH25211(util.F2PyTest): + sources = [util.getpath("tests", "src", "callback", "gh25211.f"), + util.getpath("tests", "src", "callback", "gh25211.pyf")] + module_name = "callback2" + + def test_gh25211(self): + def bar(x): + return x*x + + res = self.module.foo(bar) + assert res == 110 + + +@pytest.mark.slow +@pytest.mark.xfail(condition=(platform.system().lower() == 'darwin'), + run=False, + reason="Callback aborts cause CI failures on macOS") +class TestCBFortranCallstatement(util.F2PyTest): + sources = [util.getpath("tests", "src", "callback", "gh26681.f90")] + options = ['--lower'] + + def test_callstatement_fortran(self): + with pytest.raises(ValueError, match='helpme') as exc: + self.module.mypy_abort = self.module.utils.my_abort + self.module.utils.do_something('helpme') diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_character.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_character.py new file mode 100644 index 0000000000000000000000000000000000000000..da00fa9e27cd4dff6d7af9642067a0a162cfbbb0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_character.py @@ -0,0 +1,639 @@ +import pytest +import textwrap +from numpy.testing import assert_array_equal, assert_equal, assert_raises +import numpy as np +from numpy.f2py.tests import util + + +@pytest.mark.slow +class TestCharacterString(util.F2PyTest): + # options = ['--debug-capi', '--build-dir', '/tmp/test-build-f2py'] + suffix = '.f90' + fprefix = 'test_character_string' + length_list = ['1', '3', 'star'] + + code = '' + for length in length_list: + fsuffix = length + clength = dict(star='(*)').get(length, length) + + code += textwrap.dedent(f""" + + subroutine {fprefix}_input_{fsuffix}(c, o, n) + character*{clength}, intent(in) :: c + integer n + !f2py integer, depend(c), intent(hide) :: n = slen(c) + integer*1, dimension(n) :: o + !f2py intent(out) o + o = transfer(c, o) + end subroutine {fprefix}_input_{fsuffix} + + subroutine {fprefix}_output_{fsuffix}(c, o, n) + character*{clength}, intent(out) :: c + integer n + integer*1, dimension(n), intent(in) :: o + !f2py integer, depend(o), intent(hide) :: n = len(o) + c = transfer(o, c) + end subroutine {fprefix}_output_{fsuffix} + + subroutine {fprefix}_array_input_{fsuffix}(c, o, m, n) + integer m, i, n + character*{clength}, intent(in), dimension(m) :: c + !f2py integer, depend(c), intent(hide) :: m = len(c) + !f2py integer, depend(c), intent(hide) :: n = f2py_itemsize(c) + integer*1, dimension(m, n), intent(out) :: o + do i=1,m + o(i, :) = transfer(c(i), o(i, :)) + end do + end subroutine {fprefix}_array_input_{fsuffix} + + subroutine {fprefix}_array_output_{fsuffix}(c, o, m, n) + character*{clength}, intent(out), dimension(m) :: c + integer n + integer*1, dimension(m, n), intent(in) :: o + !f2py character(f2py_len=n) :: c + !f2py integer, depend(o), intent(hide) :: m = len(o) + !f2py integer, depend(o), intent(hide) :: n = shape(o, 1) + do i=1,m + c(i) = transfer(o(i, :), c(i)) + end do + end subroutine {fprefix}_array_output_{fsuffix} + + subroutine {fprefix}_2d_array_input_{fsuffix}(c, o, m1, m2, n) + integer m1, m2, i, j, n + character*{clength}, intent(in), dimension(m1, m2) :: c + !f2py integer, depend(c), intent(hide) :: m1 = len(c) + !f2py integer, depend(c), intent(hide) :: m2 = shape(c, 1) + !f2py integer, depend(c), intent(hide) :: n = f2py_itemsize(c) + integer*1, dimension(m1, m2, n), intent(out) :: o + do i=1,m1 + do j=1,m2 + o(i, j, :) = transfer(c(i, j), o(i, j, :)) + end do + end do + end subroutine {fprefix}_2d_array_input_{fsuffix} + """) + + @pytest.mark.parametrize("length", length_list) + def test_input(self, length): + fsuffix = {'(*)': 'star'}.get(length, length) + f = getattr(self.module, self.fprefix + '_input_' + fsuffix) + + a = {'1': 'a', '3': 'abc', 'star': 'abcde' * 3}[length] + + assert_array_equal(f(a), np.array(list(map(ord, a)), dtype='u1')) + + @pytest.mark.parametrize("length", length_list[:-1]) + def test_output(self, length): + fsuffix = length + f = getattr(self.module, self.fprefix + '_output_' + fsuffix) + + a = {'1': 'a', '3': 'abc'}[length] + + assert_array_equal(f(np.array(list(map(ord, a)), dtype='u1')), + a.encode()) + + @pytest.mark.parametrize("length", length_list) + def test_array_input(self, length): + fsuffix = length + f = getattr(self.module, self.fprefix + '_array_input_' + fsuffix) + + a = np.array([{'1': 'a', '3': 'abc', 'star': 'abcde' * 3}[length], + {'1': 'A', '3': 'ABC', 'star': 'ABCDE' * 3}[length], + ], dtype='S') + + expected = np.array([list(s) for s in a], dtype='u1') + assert_array_equal(f(a), expected) + + @pytest.mark.parametrize("length", length_list) + def test_array_output(self, length): + fsuffix = length + f = getattr(self.module, self.fprefix + '_array_output_' + fsuffix) + + expected = np.array( + [{'1': 'a', '3': 'abc', 'star': 'abcde' * 3}[length], + {'1': 'A', '3': 'ABC', 'star': 'ABCDE' * 3}[length]], dtype='S') + + a = np.array([list(s) for s in expected], dtype='u1') + assert_array_equal(f(a), expected) + + @pytest.mark.parametrize("length", length_list) + def test_2d_array_input(self, length): + fsuffix = length + f = getattr(self.module, self.fprefix + '_2d_array_input_' + fsuffix) + + a = np.array([[{'1': 'a', '3': 'abc', 'star': 'abcde' * 3}[length], + {'1': 'A', '3': 'ABC', 'star': 'ABCDE' * 3}[length]], + [{'1': 'f', '3': 'fgh', 'star': 'fghij' * 3}[length], + {'1': 'F', '3': 'FGH', 'star': 'FGHIJ' * 3}[length]]], + dtype='S') + expected = np.array([[list(item) for item in row] for row in a], + dtype='u1', order='F') + assert_array_equal(f(a), expected) + + +class TestCharacter(util.F2PyTest): + # options = ['--debug-capi', '--build-dir', '/tmp/test-build-f2py'] + suffix = '.f90' + fprefix = 'test_character' + + code = textwrap.dedent(f""" + subroutine {fprefix}_input(c, o) + character, intent(in) :: c + integer*1 o + !f2py intent(out) o + o = transfer(c, o) + end subroutine {fprefix}_input + + subroutine {fprefix}_output(c, o) + character :: c + integer*1, intent(in) :: o + !f2py intent(out) c + c = transfer(o, c) + end subroutine {fprefix}_output + + subroutine {fprefix}_input_output(c, o) + character, intent(in) :: c + character o + !f2py intent(out) o + o = c + end subroutine {fprefix}_input_output + + subroutine {fprefix}_inout(c, n) + character :: c, n + !f2py intent(in) n + !f2py intent(inout) c + c = n + end subroutine {fprefix}_inout + + function {fprefix}_return(o) result (c) + character :: c + character, intent(in) :: o + c = transfer(o, c) + end function {fprefix}_return + + subroutine {fprefix}_array_input(c, o) + character, intent(in) :: c(3) + integer*1 o(3) + !f2py intent(out) o + integer i + do i=1,3 + o(i) = transfer(c(i), o(i)) + end do + end subroutine {fprefix}_array_input + + subroutine {fprefix}_2d_array_input(c, o) + character, intent(in) :: c(2, 3) + integer*1 o(2, 3) + !f2py intent(out) o + integer i, j + do i=1,2 + do j=1,3 + o(i, j) = transfer(c(i, j), o(i, j)) + end do + end do + end subroutine {fprefix}_2d_array_input + + subroutine {fprefix}_array_output(c, o) + character :: c(3) + integer*1, intent(in) :: o(3) + !f2py intent(out) c + do i=1,3 + c(i) = transfer(o(i), c(i)) + end do + end subroutine {fprefix}_array_output + + subroutine {fprefix}_array_inout(c, n) + character :: c(3), n(3) + !f2py intent(in) n(3) + !f2py intent(inout) c(3) + do i=1,3 + c(i) = n(i) + end do + end subroutine {fprefix}_array_inout + + subroutine {fprefix}_2d_array_inout(c, n) + character :: c(2, 3), n(2, 3) + !f2py intent(in) n(2, 3) + !f2py intent(inout) c(2. 3) + integer i, j + do i=1,2 + do j=1,3 + c(i, j) = n(i, j) + end do + end do + end subroutine {fprefix}_2d_array_inout + + function {fprefix}_array_return(o) result (c) + character, dimension(3) :: c + character, intent(in) :: o(3) + do i=1,3 + c(i) = o(i) + end do + end function {fprefix}_array_return + + function {fprefix}_optional(o) result (c) + character, intent(in) :: o + !f2py character o = "a" + character :: c + c = o + end function {fprefix}_optional + """) + + @pytest.mark.parametrize("dtype", ['c', 'S1']) + def test_input(self, dtype): + f = getattr(self.module, self.fprefix + '_input') + + assert_equal(f(np.array('a', dtype=dtype)), ord('a')) + assert_equal(f(np.array(b'a', dtype=dtype)), ord('a')) + assert_equal(f(np.array(['a'], dtype=dtype)), ord('a')) + assert_equal(f(np.array('abc', dtype=dtype)), ord('a')) + assert_equal(f(np.array([['a']], dtype=dtype)), ord('a')) + + def test_input_varia(self): + f = getattr(self.module, self.fprefix + '_input') + + assert_equal(f('a'), ord('a')) + assert_equal(f(b'a'), ord(b'a')) + assert_equal(f(''), 0) + assert_equal(f(b''), 0) + assert_equal(f(b'\0'), 0) + assert_equal(f('ab'), ord('a')) + assert_equal(f(b'ab'), ord('a')) + assert_equal(f(['a']), ord('a')) + + assert_equal(f(np.array(b'a')), ord('a')) + assert_equal(f(np.array([b'a'])), ord('a')) + a = np.array('a') + assert_equal(f(a), ord('a')) + a = np.array(['a']) + assert_equal(f(a), ord('a')) + + try: + f([]) + except IndexError as msg: + if not str(msg).endswith(' got 0-list'): + raise + else: + raise SystemError(f'{f.__name__} should have failed on empty list') + + try: + f(97) + except TypeError as msg: + if not str(msg).endswith(' got int instance'): + raise + else: + raise SystemError(f'{f.__name__} should have failed on int value') + + @pytest.mark.parametrize("dtype", ['c', 'S1', 'U1']) + def test_array_input(self, dtype): + f = getattr(self.module, self.fprefix + '_array_input') + + assert_array_equal(f(np.array(['a', 'b', 'c'], dtype=dtype)), + np.array(list(map(ord, 'abc')), dtype='i1')) + assert_array_equal(f(np.array([b'a', b'b', b'c'], dtype=dtype)), + np.array(list(map(ord, 'abc')), dtype='i1')) + + def test_array_input_varia(self): + f = getattr(self.module, self.fprefix + '_array_input') + assert_array_equal(f(['a', 'b', 'c']), + np.array(list(map(ord, 'abc')), dtype='i1')) + assert_array_equal(f([b'a', b'b', b'c']), + np.array(list(map(ord, 'abc')), dtype='i1')) + + try: + f(['a', 'b', 'c', 'd']) + except ValueError as msg: + if not str(msg).endswith( + 'th dimension must be fixed to 3 but got 4'): + raise + else: + raise SystemError( + f'{f.__name__} should have failed on wrong input') + + @pytest.mark.parametrize("dtype", ['c', 'S1', 'U1']) + def test_2d_array_input(self, dtype): + f = getattr(self.module, self.fprefix + '_2d_array_input') + + a = np.array([['a', 'b', 'c'], + ['d', 'e', 'f']], dtype=dtype, order='F') + expected = a.view(np.uint32 if dtype == 'U1' else np.uint8) + assert_array_equal(f(a), expected) + + def test_output(self): + f = getattr(self.module, self.fprefix + '_output') + + assert_equal(f(ord(b'a')), b'a') + assert_equal(f(0), b'\0') + + def test_array_output(self): + f = getattr(self.module, self.fprefix + '_array_output') + + assert_array_equal(f(list(map(ord, 'abc'))), + np.array(list('abc'), dtype='S1')) + + def test_input_output(self): + f = getattr(self.module, self.fprefix + '_input_output') + + assert_equal(f(b'a'), b'a') + assert_equal(f('a'), b'a') + assert_equal(f(''), b'\0') + + @pytest.mark.parametrize("dtype", ['c', 'S1']) + def test_inout(self, dtype): + f = getattr(self.module, self.fprefix + '_inout') + + a = np.array(list('abc'), dtype=dtype) + f(a, 'A') + assert_array_equal(a, np.array(list('Abc'), dtype=a.dtype)) + f(a[1:], 'B') + assert_array_equal(a, np.array(list('ABc'), dtype=a.dtype)) + + a = np.array(['abc'], dtype=dtype) + f(a, 'A') + assert_array_equal(a, np.array(['Abc'], dtype=a.dtype)) + + def test_inout_varia(self): + f = getattr(self.module, self.fprefix + '_inout') + a = np.array('abc', dtype='S3') + f(a, 'A') + assert_array_equal(a, np.array('Abc', dtype=a.dtype)) + + a = np.array(['abc'], dtype='S3') + f(a, 'A') + assert_array_equal(a, np.array(['Abc'], dtype=a.dtype)) + + try: + f('abc', 'A') + except ValueError as msg: + if not str(msg).endswith(' got 3-str'): + raise + else: + raise SystemError(f'{f.__name__} should have failed on str value') + + @pytest.mark.parametrize("dtype", ['c', 'S1']) + def test_array_inout(self, dtype): + f = getattr(self.module, self.fprefix + '_array_inout') + n = np.array(['A', 'B', 'C'], dtype=dtype, order='F') + + a = np.array(['a', 'b', 'c'], dtype=dtype, order='F') + f(a, n) + assert_array_equal(a, n) + + a = np.array(['a', 'b', 'c', 'd'], dtype=dtype) + f(a[1:], n) + assert_array_equal(a, np.array(['a', 'A', 'B', 'C'], dtype=dtype)) + + a = np.array([['a', 'b', 'c']], dtype=dtype, order='F') + f(a, n) + assert_array_equal(a, np.array([['A', 'B', 'C']], dtype=dtype)) + + a = np.array(['a', 'b', 'c', 'd'], dtype=dtype, order='F') + try: + f(a, n) + except ValueError as msg: + if not str(msg).endswith( + 'th dimension must be fixed to 3 but got 4'): + raise + else: + raise SystemError( + f'{f.__name__} should have failed on wrong input') + + @pytest.mark.parametrize("dtype", ['c', 'S1']) + def test_2d_array_inout(self, dtype): + f = getattr(self.module, self.fprefix + '_2d_array_inout') + n = np.array([['A', 'B', 'C'], + ['D', 'E', 'F']], + dtype=dtype, order='F') + a = np.array([['a', 'b', 'c'], + ['d', 'e', 'f']], + dtype=dtype, order='F') + f(a, n) + assert_array_equal(a, n) + + def test_return(self): + f = getattr(self.module, self.fprefix + '_return') + + assert_equal(f('a'), b'a') + + @pytest.mark.skip('fortran function returning array segfaults') + def test_array_return(self): + f = getattr(self.module, self.fprefix + '_array_return') + + a = np.array(list('abc'), dtype='S1') + assert_array_equal(f(a), a) + + def test_optional(self): + f = getattr(self.module, self.fprefix + '_optional') + + assert_equal(f(), b"a") + assert_equal(f(b'B'), b"B") + + +class TestMiscCharacter(util.F2PyTest): + # options = ['--debug-capi', '--build-dir', '/tmp/test-build-f2py'] + suffix = '.f90' + fprefix = 'test_misc_character' + + code = textwrap.dedent(f""" + subroutine {fprefix}_gh18684(x, y, m) + character(len=5), dimension(m), intent(in) :: x + character*5, dimension(m), intent(out) :: y + integer i, m + !f2py integer, intent(hide), depend(x) :: m = f2py_len(x) + do i=1,m + y(i) = x(i) + end do + end subroutine {fprefix}_gh18684 + + subroutine {fprefix}_gh6308(x, i) + integer i + !f2py check(i>=0 && i<12) i + character*5 name, x + common name(12) + name(i + 1) = x + end subroutine {fprefix}_gh6308 + + subroutine {fprefix}_gh4519(x) + character(len=*), intent(in) :: x(:) + !f2py intent(out) x + integer :: i + ! Uncomment for debug printing: + !do i=1, size(x) + ! print*, "x(",i,")=", x(i) + !end do + end subroutine {fprefix}_gh4519 + + pure function {fprefix}_gh3425(x) result (y) + character(len=*), intent(in) :: x + character(len=len(x)) :: y + integer :: i + do i = 1, len(x) + j = iachar(x(i:i)) + if (j>=iachar("a") .and. j<=iachar("z") ) then + y(i:i) = achar(j-32) + else + y(i:i) = x(i:i) + endif + end do + end function {fprefix}_gh3425 + + subroutine {fprefix}_character_bc_new(x, y, z) + character, intent(in) :: x + character, intent(out) :: y + !f2py character, depend(x) :: y = x + !f2py character, dimension((x=='a'?1:2)), depend(x), intent(out) :: z + character, dimension(*) :: z + !f2py character, optional, check(x == 'a' || x == 'b') :: x = 'a' + !f2py callstatement (*f2py_func)(&x, &y, z) + !f2py callprotoargument character*, character*, character* + if (y.eq.x) then + y = x + else + y = 'e' + endif + z(1) = 'c' + end subroutine {fprefix}_character_bc_new + + subroutine {fprefix}_character_bc_old(x, y, z) + character, intent(in) :: x + character, intent(out) :: y + !f2py character, depend(x) :: y = x[0] + !f2py character, dimension((*x=='a'?1:2)), depend(x), intent(out) :: z + character, dimension(*) :: z + !f2py character, optional, check(*x == 'a' || x[0] == 'b') :: x = 'a' + !f2py callstatement (*f2py_func)(x, y, z) + !f2py callprotoargument char*, char*, char* + if (y.eq.x) then + y = x + else + y = 'e' + endif + z(1) = 'c' + end subroutine {fprefix}_character_bc_old + """) + + @pytest.mark.slow + def test_gh18684(self): + # Test character(len=5) and character*5 usages + f = getattr(self.module, self.fprefix + '_gh18684') + x = np.array(["abcde", "fghij"], dtype='S5') + y = f(x) + + assert_array_equal(x, y) + + def test_gh6308(self): + # Test character string array in a common block + f = getattr(self.module, self.fprefix + '_gh6308') + + assert_equal(self.module._BLNK_.name.dtype, np.dtype('S5')) + assert_equal(len(self.module._BLNK_.name), 12) + f("abcde", 0) + assert_equal(self.module._BLNK_.name[0], b"abcde") + f("12345", 5) + assert_equal(self.module._BLNK_.name[5], b"12345") + + def test_gh4519(self): + # Test array of assumed length strings + f = getattr(self.module, self.fprefix + '_gh4519') + + for x, expected in [ + ('a', dict(shape=(), dtype=np.dtype('S1'))), + ('text', dict(shape=(), dtype=np.dtype('S4'))), + (np.array(['1', '2', '3'], dtype='S1'), + dict(shape=(3,), dtype=np.dtype('S1'))), + (['1', '2', '34'], + dict(shape=(3,), dtype=np.dtype('S2'))), + (['', ''], dict(shape=(2,), dtype=np.dtype('S1')))]: + r = f(x) + for k, v in expected.items(): + assert_equal(getattr(r, k), v) + + def test_gh3425(self): + # Test returning a copy of assumed length string + f = getattr(self.module, self.fprefix + '_gh3425') + # f is equivalent to bytes.upper + + assert_equal(f('abC'), b'ABC') + assert_equal(f(''), b'') + assert_equal(f('abC12d'), b'ABC12D') + + @pytest.mark.parametrize("state", ['new', 'old']) + def test_character_bc(self, state): + f = getattr(self.module, self.fprefix + '_character_bc_' + state) + + c, a = f() + assert_equal(c, b'a') + assert_equal(len(a), 1) + + c, a = f(b'b') + assert_equal(c, b'b') + assert_equal(len(a), 2) + + assert_raises(Exception, lambda: f(b'c')) + + +class TestStringScalarArr(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "scalar_string.f90")] + + def test_char(self): + for out in (self.module.string_test.string, + self.module.string_test.string77): + expected = () + assert out.shape == expected + expected = '|S8' + assert out.dtype == expected + + def test_char_arr(self): + for out in (self.module.string_test.strarr, + self.module.string_test.strarr77): + expected = (5,7) + assert out.shape == expected + expected = '|S12' + assert out.dtype == expected + +class TestStringAssumedLength(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "gh24008.f")] + + def test_gh24008(self): + self.module.greet("joe", "bob") + +@pytest.mark.slow +class TestStringOptionalInOut(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "gh24662.f90")] + + def test_gh24662(self): + self.module.string_inout_optional() + a = np.array('hi', dtype='S32') + self.module.string_inout_optional(a) + assert "output string" in a.tobytes().decode() + with pytest.raises(Exception): + aa = "Hi" + self.module.string_inout_optional(aa) + + +@pytest.mark.slow +class TestNewCharHandling(util.F2PyTest): + # from v1.24 onwards, gh-19388 + sources = [ + util.getpath("tests", "src", "string", "gh25286.pyf"), + util.getpath("tests", "src", "string", "gh25286.f90") + ] + module_name = "_char_handling_test" + + def test_gh25286(self): + info = self.module.charint('T') + assert info == 2 + +@pytest.mark.slow +class TestBCCharHandling(util.F2PyTest): + # SciPy style, "incorrect" bindings with a hook + sources = [ + util.getpath("tests", "src", "string", "gh25286_bc.pyf"), + util.getpath("tests", "src", "string", "gh25286.f90") + ] + module_name = "_char_handling_test" + + def test_gh25286(self): + info = self.module.charint('T') + assert info == 2 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_common.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_common.py new file mode 100644 index 0000000000000000000000000000000000000000..09bd6147f0f3e72fb0eaca4c9e60d8be2b5b2b57 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_common.py @@ -0,0 +1,20 @@ +import pytest +import numpy as np +from . import util + +@pytest.mark.slow +class TestCommonBlock(util.F2PyTest): + sources = [util.getpath("tests", "src", "common", "block.f")] + + def test_common_block(self): + self.module.initcb() + assert self.module.block.long_bn == np.array(1.0, dtype=np.float64) + assert self.module.block.string_bn == np.array("2", dtype="|S1") + assert self.module.block.ok == np.array(3, dtype=np.int32) + + +class TestCommonWithUse(util.F2PyTest): + sources = [util.getpath("tests", "src", "common", "gh19161.f90")] + + def test_common_gh19161(self): + assert self.module.data.x == 0 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_crackfortran.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_crackfortran.py new file mode 100644 index 0000000000000000000000000000000000000000..965a6b0f87e8005b22f2e205ffa0579dfaf62b08 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_crackfortran.py @@ -0,0 +1,418 @@ +import importlib +import time +import pytest +import numpy as np +from numpy.f2py.crackfortran import markinnerspaces, nameargspattern +from . import util +from numpy.f2py import crackfortran +import textwrap +import contextlib +import io + + +class TestNoSpace(util.F2PyTest): + # issue gh-15035: add handling for endsubroutine, endfunction with no space + # between "end" and the block name + sources = [util.getpath("tests", "src", "crackfortran", "gh15035.f")] + + def test_module(self): + k = np.array([1, 2, 3], dtype=np.float64) + w = np.array([1, 2, 3], dtype=np.float64) + self.module.subb(k) + assert np.allclose(k, w + 1) + self.module.subc([w, k]) + assert np.allclose(k, w + 1) + assert self.module.t0("23") == b"2" + + +class TestPublicPrivate: + def test_defaultPrivate(self): + fpath = util.getpath("tests", "src", "crackfortran", "privatemod.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + mod = mod[0] + assert "private" in mod["vars"]["a"]["attrspec"] + assert "public" not in mod["vars"]["a"]["attrspec"] + assert "private" in mod["vars"]["b"]["attrspec"] + assert "public" not in mod["vars"]["b"]["attrspec"] + assert "private" not in mod["vars"]["seta"]["attrspec"] + assert "public" in mod["vars"]["seta"]["attrspec"] + + def test_defaultPublic(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "publicmod.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + mod = mod[0] + assert "private" in mod["vars"]["a"]["attrspec"] + assert "public" not in mod["vars"]["a"]["attrspec"] + assert "private" not in mod["vars"]["seta"]["attrspec"] + assert "public" in mod["vars"]["seta"]["attrspec"] + + def test_access_type(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "accesstype.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + tt = mod[0]['vars'] + assert set(tt['a']['attrspec']) == {'private', 'bind(c)'} + assert set(tt['b_']['attrspec']) == {'public', 'bind(c)'} + assert set(tt['c']['attrspec']) == {'public'} + + def test_nowrap_private_proceedures(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "gh23879.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + pyf = crackfortran.crack2fortran(mod) + assert 'bar' not in pyf + +class TestModuleProcedure: + def test_moduleOperators(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "operators.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + mod = mod[0] + assert "body" in mod and len(mod["body"]) == 9 + assert mod["body"][1]["name"] == "operator(.item.)" + assert "implementedby" in mod["body"][1] + assert mod["body"][1]["implementedby"] == \ + ["item_int", "item_real"] + assert mod["body"][2]["name"] == "operator(==)" + assert "implementedby" in mod["body"][2] + assert mod["body"][2]["implementedby"] == ["items_are_equal"] + assert mod["body"][3]["name"] == "assignment(=)" + assert "implementedby" in mod["body"][3] + assert mod["body"][3]["implementedby"] == \ + ["get_int", "get_real"] + + def test_notPublicPrivate(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "pubprivmod.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + mod = mod[0] + assert mod['vars']['a']['attrspec'] == ['private', ] + assert mod['vars']['b']['attrspec'] == ['public', ] + assert mod['vars']['seta']['attrspec'] == ['public', ] + + +class TestExternal(util.F2PyTest): + # issue gh-17859: add external attribute support + sources = [util.getpath("tests", "src", "crackfortran", "gh17859.f")] + + def test_external_as_statement(self): + def incr(x): + return x + 123 + + r = self.module.external_as_statement(incr) + assert r == 123 + + def test_external_as_attribute(self): + def incr(x): + return x + 123 + + r = self.module.external_as_attribute(incr) + assert r == 123 + + +class TestCrackFortran(util.F2PyTest): + # gh-2848: commented lines between parameters in subroutine parameter lists + sources = [util.getpath("tests", "src", "crackfortran", "gh2848.f90"), + util.getpath("tests", "src", "crackfortran", "common_with_division.f") + ] + + def test_gh2848(self): + r = self.module.gh2848(1, 2) + assert r == (1, 2) + + def test_common_with_division(self): + assert len(self.module.mortmp.ctmp) == 11 + +class TestMarkinnerspaces: + # gh-14118: markinnerspaces does not handle multiple quotations + + def test_do_not_touch_normal_spaces(self): + test_list = ["a ", " a", "a b c", "'abcdefghij'"] + for i in test_list: + assert markinnerspaces(i) == i + + def test_one_relevant_space(self): + assert markinnerspaces("a 'b c' \\' \\'") == "a 'b@_@c' \\' \\'" + assert markinnerspaces(r'a "b c" \" \"') == r'a "b@_@c" \" \"' + + def test_ignore_inner_quotes(self): + assert markinnerspaces("a 'b c\" \" d' e") == "a 'b@_@c\"@_@\"@_@d' e" + assert markinnerspaces("a \"b c' ' d\" e") == "a \"b@_@c'@_@'@_@d\" e" + + def test_multiple_relevant_spaces(self): + assert markinnerspaces("a 'b c' 'd e'") == "a 'b@_@c' 'd@_@e'" + assert markinnerspaces(r'a "b c" "d e"') == r'a "b@_@c" "d@_@e"' + + +class TestDimSpec(util.F2PyTest): + """This test suite tests various expressions that are used as dimension + specifications. + + There exists two usage cases where analyzing dimensions + specifications are important. + + In the first case, the size of output arrays must be defined based + on the inputs to a Fortran function. Because Fortran supports + arbitrary bases for indexing, for instance, `arr(lower:upper)`, + f2py has to evaluate an expression `upper - lower + 1` where + `lower` and `upper` are arbitrary expressions of input parameters. + The evaluation is performed in C, so f2py has to translate Fortran + expressions to valid C expressions (an alternative approach is + that a developer specifies the corresponding C expressions in a + .pyf file). + + In the second case, when user provides an input array with a given + size but some hidden parameters used in dimensions specifications + need to be determined based on the input array size. This is a + harder problem because f2py has to solve the inverse problem: find + a parameter `p` such that `upper(p) - lower(p) + 1` equals to the + size of input array. In the case when this equation cannot be + solved (e.g. because the input array size is wrong), raise an + error before calling the Fortran function (that otherwise would + likely crash Python process when the size of input arrays is + wrong). f2py currently supports this case only when the equation + is linear with respect to unknown parameter. + + """ + + suffix = ".f90" + + code_template = textwrap.dedent(""" + function get_arr_size_{count}(a, n) result (length) + integer, intent(in) :: n + integer, dimension({dimspec}), intent(out) :: a + integer length + length = size(a) + end function + + subroutine get_inv_arr_size_{count}(a, n) + integer :: n + ! the value of n is computed in f2py wrapper + !f2py intent(out) n + integer, dimension({dimspec}), intent(in) :: a + if (a({first}).gt.0) then + ! print*, "a=", a + endif + end subroutine + """) + + linear_dimspecs = [ + "n", "2*n", "2:n", "n/2", "5 - n/2", "3*n:20", "n*(n+1):n*(n+5)", + "2*n, n" + ] + nonlinear_dimspecs = ["2*n:3*n*n+2*n"] + all_dimspecs = linear_dimspecs + nonlinear_dimspecs + + code = "" + for count, dimspec in enumerate(all_dimspecs): + lst = [(d.split(":")[0] if ":" in d else "1") for d in dimspec.split(',')] + code += code_template.format( + count=count, + dimspec=dimspec, + first=", ".join(lst), + ) + + @pytest.mark.parametrize("dimspec", all_dimspecs) + @pytest.mark.slow + def test_array_size(self, dimspec): + + count = self.all_dimspecs.index(dimspec) + get_arr_size = getattr(self.module, f"get_arr_size_{count}") + + for n in [1, 2, 3, 4, 5]: + sz, a = get_arr_size(n) + assert a.size == sz + + @pytest.mark.parametrize("dimspec", all_dimspecs) + def test_inv_array_size(self, dimspec): + + count = self.all_dimspecs.index(dimspec) + get_arr_size = getattr(self.module, f"get_arr_size_{count}") + get_inv_arr_size = getattr(self.module, f"get_inv_arr_size_{count}") + + for n in [1, 2, 3, 4, 5]: + sz, a = get_arr_size(n) + if dimspec in self.nonlinear_dimspecs: + # one must specify n as input, the call we'll ensure + # that a and n are compatible: + n1 = get_inv_arr_size(a, n) + else: + # in case of linear dependence, n can be determined + # from the shape of a: + n1 = get_inv_arr_size(a) + # n1 may be different from n (for instance, when `a` size + # is a function of some `n` fraction) but it must produce + # the same sized array + sz1, _ = get_arr_size(n1) + assert sz == sz1, (n, n1, sz, sz1) + + +class TestModuleDeclaration: + def test_dependencies(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "foo_deps.f90") + mod = crackfortran.crackfortran([str(fpath)]) + assert len(mod) == 1 + assert mod[0]["vars"]["abar"]["="] == "bar('abar')" + + +class TestEval(util.F2PyTest): + def test_eval_scalar(self): + eval_scalar = crackfortran._eval_scalar + + assert eval_scalar('123', {}) == '123' + assert eval_scalar('12 + 3', {}) == '15' + assert eval_scalar('a + b', dict(a=1, b=2)) == '3' + assert eval_scalar('"123"', {}) == "'123'" + + +class TestFortranReader(util.F2PyTest): + @pytest.mark.parametrize("encoding", + ['ascii', 'utf-8', 'utf-16', 'utf-32']) + def test_input_encoding(self, tmp_path, encoding): + # gh-635 + f_path = tmp_path / f"input_with_{encoding}_encoding.f90" + with f_path.open('w', encoding=encoding) as ff: + ff.write(""" + subroutine foo() + end subroutine foo + """) + mod = crackfortran.crackfortran([str(f_path)]) + assert mod[0]['name'] == 'foo' + + +@pytest.mark.slow +class TestUnicodeComment(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "unicode_comment.f90")] + + @pytest.mark.skipif( + (importlib.util.find_spec("charset_normalizer") is None), + reason="test requires charset_normalizer which is not installed", + ) + def test_encoding_comment(self): + self.module.foo(3) + + +class TestNameArgsPatternBacktracking: + @pytest.mark.parametrize( + ['adversary'], + [ + ('@)@bind@(@',), + ('@)@bind @(@',), + ('@)@bind foo bar baz@(@',) + ] + ) + def test_nameargspattern_backtracking(self, adversary): + '''address ReDOS vulnerability: + https://github.com/numpy/numpy/issues/23338''' + trials_per_batch = 12 + batches_per_regex = 4 + start_reps, end_reps = 15, 25 + for ii in range(start_reps, end_reps): + repeated_adversary = adversary * ii + # test times in small batches. + # this gives us more chances to catch a bad regex + # while still catching it before too long if it is bad + for _ in range(batches_per_regex): + times = [] + for _ in range(trials_per_batch): + t0 = time.perf_counter() + mtch = nameargspattern.search(repeated_adversary) + times.append(time.perf_counter() - t0) + # our pattern should be much faster than 0.2s per search + # it's unlikely that a bad regex will pass even on fast CPUs + assert np.median(times) < 0.2 + assert not mtch + # if the adversary is capped with @)@, it becomes acceptable + # according to the old version of the regex. + # that should still be true. + good_version_of_adversary = repeated_adversary + '@)@' + assert nameargspattern.search(good_version_of_adversary) + +class TestFunctionReturn(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "gh23598.f90")] + + @pytest.mark.slow + def test_function_rettype(self): + # gh-23598 + assert self.module.intproduct(3, 4) == 12 + + +class TestFortranGroupCounters(util.F2PyTest): + def test_end_if_comment(self): + # gh-23533 + fpath = util.getpath("tests", "src", "crackfortran", "gh23533.f") + try: + crackfortran.crackfortran([str(fpath)]) + except Exception as exc: + assert False, f"'crackfortran.crackfortran' raised an exception {exc}" + + +class TestF77CommonBlockReader: + def test_gh22648(self, tmp_path): + fpath = util.getpath("tests", "src", "crackfortran", "gh22648.pyf") + with contextlib.redirect_stdout(io.StringIO()) as stdout_f2py: + mod = crackfortran.crackfortran([str(fpath)]) + assert "Mismatch" not in stdout_f2py.getvalue() + +class TestParamEval: + # issue gh-11612, array parameter parsing + def test_param_eval_nested(self): + v = '(/3.14, 4./)' + g_params = dict(kind=crackfortran._kind_func, + selected_int_kind=crackfortran._selected_int_kind_func, + selected_real_kind=crackfortran._selected_real_kind_func) + params = {'dp': 8, 'intparamarray': {1: 3, 2: 5}, + 'nested': {1: 1, 2: 2, 3: 3}} + dimspec = '(2)' + ret = crackfortran.param_eval(v, g_params, params, dimspec=dimspec) + assert ret == {1: 3.14, 2: 4.0} + + def test_param_eval_nonstandard_range(self): + v = '(/ 6, 3, 1 /)' + g_params = dict(kind=crackfortran._kind_func, + selected_int_kind=crackfortran._selected_int_kind_func, + selected_real_kind=crackfortran._selected_real_kind_func) + params = {} + dimspec = '(-1:1)' + ret = crackfortran.param_eval(v, g_params, params, dimspec=dimspec) + assert ret == {-1: 6, 0: 3, 1: 1} + + def test_param_eval_empty_range(self): + v = '6' + g_params = dict(kind=crackfortran._kind_func, + selected_int_kind=crackfortran._selected_int_kind_func, + selected_real_kind=crackfortran._selected_real_kind_func) + params = {} + dimspec = '' + pytest.raises(ValueError, crackfortran.param_eval, v, g_params, params, + dimspec=dimspec) + + def test_param_eval_non_array_param(self): + v = '3.14_dp' + g_params = dict(kind=crackfortran._kind_func, + selected_int_kind=crackfortran._selected_int_kind_func, + selected_real_kind=crackfortran._selected_real_kind_func) + params = {} + ret = crackfortran.param_eval(v, g_params, params, dimspec=None) + assert ret == '3.14_dp' + + def test_param_eval_too_many_dims(self): + v = 'reshape((/ (i, i=1, 250) /), (/5, 10, 5/))' + g_params = dict(kind=crackfortran._kind_func, + selected_int_kind=crackfortran._selected_int_kind_func, + selected_real_kind=crackfortran._selected_real_kind_func) + params = {} + dimspec = '(0:4, 3:12, 5)' + pytest.raises(ValueError, crackfortran.param_eval, v, g_params, params, + dimspec=dimspec) + +@pytest.mark.slow +class TestLowerF2PYDirective(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "gh27697.f90")] + options = ['--lower'] + + def test_no_lower_fail(self): + with pytest.raises(ValueError, match='aborting directly') as exc: + self.module.utils.my_abort('aborting directly') diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_data.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_data.py new file mode 100644 index 0000000000000000000000000000000000000000..e2a425084a55198adc3075cbb11347334b08ed76 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_data.py @@ -0,0 +1,70 @@ +import pytest +import numpy as np + +from . import util +from numpy.f2py.crackfortran import crackfortran + + +class TestData(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "data_stmts.f90")] + + # For gh-23276 + @pytest.mark.slow + def test_data_stmts(self): + assert self.module.cmplxdat.i == 2 + assert self.module.cmplxdat.j == 3 + assert self.module.cmplxdat.x == 1.5 + assert self.module.cmplxdat.y == 2.0 + assert self.module.cmplxdat.pi == 3.1415926535897932384626433832795028841971693993751058209749445923078164062 + assert self.module.cmplxdat.medium_ref_index == np.array(1.+0.j) + assert np.all(self.module.cmplxdat.z == np.array([3.5, 7.0])) + assert np.all(self.module.cmplxdat.my_array == np.array([ 1.+2.j, -3.+4.j])) + assert np.all(self.module.cmplxdat.my_real_array == np.array([ 1., 2., 3.])) + assert np.all(self.module.cmplxdat.ref_index_one == np.array([13.0 + 21.0j])) + assert np.all(self.module.cmplxdat.ref_index_two == np.array([-30.0 + 43.0j])) + + def test_crackedlines(self): + mod = crackfortran(self.sources) + assert mod[0]['vars']['x']['='] == '1.5' + assert mod[0]['vars']['y']['='] == '2.0' + assert mod[0]['vars']['pi']['='] == '3.1415926535897932384626433832795028841971693993751058209749445923078164062d0' + assert mod[0]['vars']['my_real_array']['='] == '(/1.0d0, 2.0d0, 3.0d0/)' + assert mod[0]['vars']['ref_index_one']['='] == '(13.0d0, 21.0d0)' + assert mod[0]['vars']['ref_index_two']['='] == '(-30.0d0, 43.0d0)' + assert mod[0]['vars']['my_array']['='] == '(/(1.0d0, 2.0d0), (-3.0d0, 4.0d0)/)' + assert mod[0]['vars']['z']['='] == '(/3.5, 7.0/)' + +class TestDataF77(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "data_common.f")] + + # For gh-23276 + def test_data_stmts(self): + assert self.module.mycom.mydata == 0 + + def test_crackedlines(self): + mod = crackfortran(str(self.sources[0])) + print(mod[0]['vars']) + assert mod[0]['vars']['mydata']['='] == '0' + + +class TestDataMultiplierF77(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "data_multiplier.f")] + + # For gh-23276 + def test_data_stmts(self): + assert self.module.mycom.ivar1 == 3 + assert self.module.mycom.ivar2 == 3 + assert self.module.mycom.ivar3 == 2 + assert self.module.mycom.ivar4 == 2 + assert self.module.mycom.evar5 == 0 + + +class TestDataWithCommentsF77(util.F2PyTest): + sources = [util.getpath("tests", "src", "crackfortran", "data_with_comments.f")] + + # For gh-23276 + def test_data_stmts(self): + assert len(self.module.mycom.mytab) == 3 + assert self.module.mycom.mytab[0] == 0 + assert self.module.mycom.mytab[1] == 4 + assert self.module.mycom.mytab[2] == 0 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_docs.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..efba7ea40ee661b4e42a05766504d378bcf3d7e2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_docs.py @@ -0,0 +1,59 @@ +import pytest +import numpy as np +from numpy.testing import assert_array_equal, assert_equal +from . import util +from pathlib import Path + +def get_docdir(): + parents = Path(__file__).resolve().parents + try: + # Assumes that spin is used to run tests + nproot = parents[8] + except IndexError: + docdir = None + else: + docdir = nproot / "doc" / "source" / "f2py" / "code" + if docdir and docdir.is_dir(): + return docdir + # Assumes that an editable install is used to run tests + return parents[3] / "doc" / "source" / "f2py" / "code" + +pytestmark = pytest.mark.skipif( + not get_docdir().is_dir(), + reason=f"Could not find f2py documentation sources" + f"({get_docdir()} does not exist)", +) + +def _path(*args): + return get_docdir().joinpath(*args) + +@pytest.mark.slow +class TestDocAdvanced(util.F2PyTest): + # options = ['--debug-capi', '--build-dir', '/tmp/build-f2py'] + sources = [_path('asterisk1.f90'), _path('asterisk2.f90'), + _path('ftype.f')] + + def test_asterisk1(self): + foo = self.module.foo1 + assert_equal(foo(), b'123456789A12') + + def test_asterisk2(self): + foo = self.module.foo2 + assert_equal(foo(2), b'12') + assert_equal(foo(12), b'123456789A12') + assert_equal(foo(20), b'123456789A123456789B') + + def test_ftype(self): + ftype = self.module + ftype.foo() + assert_equal(ftype.data.a, 0) + ftype.data.a = 3 + ftype.data.x = [1, 2, 3] + assert_equal(ftype.data.a, 3) + assert_array_equal(ftype.data.x, + np.array([1, 2, 3], dtype=np.float32)) + ftype.data.x[1] = 45 + assert_array_equal(ftype.data.x, + np.array([1, 45, 3], dtype=np.float32)) + + # TODO: implement test methods for other example Fortran codes diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2cmap.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2cmap.py new file mode 100644 index 0000000000000000000000000000000000000000..6596ada33a5454398427f3b605862dd1bae9cab4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2cmap.py @@ -0,0 +1,15 @@ +from . import util +import numpy as np + +class TestF2Cmap(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "f2cmap", "isoFortranEnvMap.f90"), + util.getpath("tests", "src", "f2cmap", ".f2py_f2cmap") + ] + + # gh-15095 + def test_gh15095(self): + inp = np.ones(3) + out = self.module.func1(inp) + exp_out = 3 + assert out == exp_out diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2py2e.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2py2e.py new file mode 100644 index 0000000000000000000000000000000000000000..3f321418f403ff628b859e3037fdd3e0f6bbd320 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_f2py2e.py @@ -0,0 +1,964 @@ +import re +import shlex +import subprocess +import sys +import textwrap +from pathlib import Path +from collections import namedtuple + +import platform + +import pytest + +from . import util +from numpy.f2py.f2py2e import main as f2pycli +from numpy.testing._private.utils import NOGIL_BUILD + +####################### +# F2PY Test utilities # +###################### + +# Tests for CLI commands which call meson will fail if no compilers are present, these are to be skipped + +def compiler_check_f2pycli(): + if not util.has_fortran_compiler(): + pytest.skip("CLI command needs a Fortran compiler") + else: + f2pycli() + +######################### +# CLI utils and classes # +######################### + +PPaths = namedtuple("PPaths", "finp, f90inp, pyf, wrap77, wrap90, cmodf") + + +def get_io_paths(fname_inp, mname="untitled"): + """Takes in a temporary file for testing and returns the expected output and input paths + + Here expected output is essentially one of any of the possible generated + files. + + ..note:: + + Since this does not actually run f2py, none of these are guaranteed to + exist, and module names are typically incorrect + + Parameters + ---------- + fname_inp : str + The input filename + mname : str, optional + The name of the module, untitled by default + + Returns + ------- + genp : NamedTuple PPaths + The possible paths which are generated, not all of which exist + """ + bpath = Path(fname_inp) + return PPaths( + finp=bpath.with_suffix(".f"), + f90inp=bpath.with_suffix(".f90"), + pyf=bpath.with_suffix(".pyf"), + wrap77=bpath.with_name(f"{mname}-f2pywrappers.f"), + wrap90=bpath.with_name(f"{mname}-f2pywrappers2.f90"), + cmodf=bpath.with_name(f"{mname}module.c"), + ) + + +################ +# CLI Fixtures # +################ + + +@pytest.fixture(scope="session") +def hello_world_f90(tmpdir_factory): + """Generates a single f90 file for testing""" + fdat = util.getpath("tests", "src", "cli", "hiworld.f90").read_text() + fn = tmpdir_factory.getbasetemp() / "hello.f90" + fn.write_text(fdat, encoding="ascii") + return fn + + +@pytest.fixture(scope="session") +def gh23598_warn(tmpdir_factory): + """F90 file for testing warnings in gh23598""" + fdat = util.getpath("tests", "src", "crackfortran", "gh23598Warn.f90").read_text() + fn = tmpdir_factory.getbasetemp() / "gh23598Warn.f90" + fn.write_text(fdat, encoding="ascii") + return fn + + +@pytest.fixture(scope="session") +def gh22819_cli(tmpdir_factory): + """F90 file for testing disallowed CLI arguments in ghff819""" + fdat = util.getpath("tests", "src", "cli", "gh_22819.pyf").read_text() + fn = tmpdir_factory.getbasetemp() / "gh_22819.pyf" + fn.write_text(fdat, encoding="ascii") + return fn + + +@pytest.fixture(scope="session") +def hello_world_f77(tmpdir_factory): + """Generates a single f77 file for testing""" + fdat = util.getpath("tests", "src", "cli", "hi77.f").read_text() + fn = tmpdir_factory.getbasetemp() / "hello.f" + fn.write_text(fdat, encoding="ascii") + return fn + + +@pytest.fixture(scope="session") +def retreal_f77(tmpdir_factory): + """Generates a single f77 file for testing""" + fdat = util.getpath("tests", "src", "return_real", "foo77.f").read_text() + fn = tmpdir_factory.getbasetemp() / "foo.f" + fn.write_text(fdat, encoding="ascii") + return fn + +@pytest.fixture(scope="session") +def f2cmap_f90(tmpdir_factory): + """Generates a single f90 file for testing""" + fdat = util.getpath("tests", "src", "f2cmap", "isoFortranEnvMap.f90").read_text() + f2cmap = util.getpath("tests", "src", "f2cmap", ".f2py_f2cmap").read_text() + fn = tmpdir_factory.getbasetemp() / "f2cmap.f90" + fmap = tmpdir_factory.getbasetemp() / "mapfile" + fn.write_text(fdat, encoding="ascii") + fmap.write_text(f2cmap, encoding="ascii") + return fn + +######### +# Tests # +######### + +def test_gh22819_cli(capfd, gh22819_cli, monkeypatch): + """Check that module names are handled correctly + gh-22819 + Essentially, the -m name cannot be used to import the module, so the module + named in the .pyf needs to be used instead + + CLI :: -m and a .pyf file + """ + ipath = Path(gh22819_cli) + monkeypatch.setattr(sys, "argv", f"f2py -m blah {ipath}".split()) + with util.switchdir(ipath.parent): + f2pycli() + gen_paths = [item.name for item in ipath.parent.rglob("*") if item.is_file()] + assert "blahmodule.c" not in gen_paths # shouldn't be generated + assert "blah-f2pywrappers.f" not in gen_paths + assert "test_22819-f2pywrappers.f" in gen_paths + assert "test_22819module.c" in gen_paths + assert "Ignoring blah" + + +def test_gh22819_many_pyf(capfd, gh22819_cli, monkeypatch): + """Only one .pyf file allowed + gh-22819 + CLI :: .pyf files + """ + ipath = Path(gh22819_cli) + monkeypatch.setattr(sys, "argv", f"f2py -m blah {ipath} hello.pyf".split()) + with util.switchdir(ipath.parent): + with pytest.raises(ValueError, match="Only one .pyf file per call"): + f2pycli() + + +def test_gh23598_warn(capfd, gh23598_warn, monkeypatch): + foutl = get_io_paths(gh23598_warn, mname="test") + ipath = foutl.f90inp + monkeypatch.setattr( + sys, "argv", + f'f2py {ipath} -m test'.split()) + + with util.switchdir(ipath.parent): + f2pycli() # Generate files + wrapper = foutl.wrap90.read_text() + assert "intproductf2pywrap, intpr" not in wrapper + + +def test_gen_pyf(capfd, hello_world_f90, monkeypatch): + """Ensures that a signature file is generated via the CLI + CLI :: -h + """ + ipath = Path(hello_world_f90) + opath = Path(hello_world_f90).stem + ".pyf" + monkeypatch.setattr(sys, "argv", f'f2py -h {opath} {ipath}'.split()) + + with util.switchdir(ipath.parent): + f2pycli() # Generate wrappers + out, _ = capfd.readouterr() + assert "Saving signatures to file" in out + assert Path(f'{opath}').exists() + + +def test_gen_pyf_stdout(capfd, hello_world_f90, monkeypatch): + """Ensures that a signature file can be dumped to stdout + CLI :: -h + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -h stdout {ipath}'.split()) + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Saving signatures to file" in out + assert "function hi() ! in " in out + + +def test_gen_pyf_no_overwrite(capfd, hello_world_f90, monkeypatch): + """Ensures that the CLI refuses to overwrite signature files + CLI :: -h without --overwrite-signature + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -h faker.pyf {ipath}'.split()) + + with util.switchdir(ipath.parent): + Path("faker.pyf").write_text("Fake news", encoding="ascii") + with pytest.raises(SystemExit): + f2pycli() # Refuse to overwrite + _, err = capfd.readouterr() + assert "Use --overwrite-signature to overwrite" in err + + +@pytest.mark.skipif(sys.version_info <= (3, 12), reason="Python 3.12 required") +def test_untitled_cli(capfd, hello_world_f90, monkeypatch): + """Check that modules are named correctly + + CLI :: defaults + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f"f2py --backend meson -c {ipath}".split()) + with util.switchdir(ipath.parent): + compiler_check_f2pycli() + out, _ = capfd.readouterr() + assert "untitledmodule.c" in out + + +@pytest.mark.skipif((platform.system() != 'Linux') or (sys.version_info <= (3, 12)), reason='Compiler and 3.12 required') +def test_no_py312_distutils_fcompiler(capfd, hello_world_f90, monkeypatch): + """Check that no distutils imports are performed on 3.12 + CLI :: --fcompiler --help-link --backend distutils + """ + MNAME = "hi" + foutl = get_io_paths(hello_world_f90, mname=MNAME) + ipath = foutl.f90inp + monkeypatch.setattr( + sys, "argv", f"f2py {ipath} -c --fcompiler=gfortran -m {MNAME}".split() + ) + with util.switchdir(ipath.parent): + compiler_check_f2pycli() + out, _ = capfd.readouterr() + assert "--fcompiler cannot be used with meson" in out + monkeypatch.setattr( + sys, "argv", "f2py --help-link".split() + ) + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Use --dep for meson builds" in out + MNAME = "hi2" # Needs to be different for a new -c + monkeypatch.setattr( + sys, "argv", f"f2py {ipath} -c -m {MNAME} --backend distutils".split() + ) + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Cannot use distutils backend with Python>=3.12" in out + + +@pytest.mark.xfail +def test_f2py_skip(capfd, retreal_f77, monkeypatch): + """Tests that functions can be skipped + CLI :: skip: + """ + foutl = get_io_paths(retreal_f77, mname="test") + ipath = foutl.finp + toskip = "t0 t4 t8 sd s8 s4" + remaining = "td s0" + monkeypatch.setattr( + sys, "argv", + f'f2py {ipath} -m test skip: {toskip}'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, err = capfd.readouterr() + for skey in toskip.split(): + assert ( + f'buildmodule: Could not found the body of interfaced routine "{skey}". Skipping.' + in err) + for rkey in remaining.split(): + assert f'Constructing wrapper function "{rkey}"' in out + + +def test_f2py_only(capfd, retreal_f77, monkeypatch): + """Test that functions can be kept by only: + CLI :: only: + """ + foutl = get_io_paths(retreal_f77, mname="test") + ipath = foutl.finp + toskip = "t0 t4 t8 sd s8 s4" + tokeep = "td s0" + monkeypatch.setattr( + sys, "argv", + f'f2py {ipath} -m test only: {tokeep}'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, err = capfd.readouterr() + for skey in toskip.split(): + assert ( + f'buildmodule: Could not find the body of interfaced routine "{skey}". Skipping.' + in err) + for rkey in tokeep.split(): + assert f'Constructing wrapper function "{rkey}"' in out + + +def test_file_processing_switch(capfd, hello_world_f90, retreal_f77, + monkeypatch): + """Tests that it is possible to return to file processing mode + CLI :: : + BUG: numpy-gh #20520 + """ + foutl = get_io_paths(retreal_f77, mname="test") + ipath = foutl.finp + toskip = "t0 t4 t8 sd s8 s4" + ipath2 = Path(hello_world_f90) + tokeep = "td s0 hi" # hi is in ipath2 + mname = "blah" + monkeypatch.setattr( + sys, + "argv", + f'f2py {ipath} -m {mname} only: {tokeep} : {ipath2}'.split( + ), + ) + + with util.switchdir(ipath.parent): + f2pycli() + out, err = capfd.readouterr() + for skey in toskip.split(): + assert ( + f'buildmodule: Could not find the body of interfaced routine "{skey}". Skipping.' + in err) + for rkey in tokeep.split(): + assert f'Constructing wrapper function "{rkey}"' in out + + +def test_mod_gen_f77(capfd, hello_world_f90, monkeypatch): + """Checks the generation of files based on a module name + CLI :: -m + """ + MNAME = "hi" + foutl = get_io_paths(hello_world_f90, mname=MNAME) + ipath = foutl.f90inp + monkeypatch.setattr(sys, "argv", f'f2py {ipath} -m {MNAME}'.split()) + with util.switchdir(ipath.parent): + f2pycli() + + # Always generate C module + assert Path.exists(foutl.cmodf) + # File contains a function, check for F77 wrappers + assert Path.exists(foutl.wrap77) + + +def test_mod_gen_gh25263(capfd, hello_world_f77, monkeypatch): + """Check that pyf files are correctly generated with module structure + CLI :: -m -h pyf_file + BUG: numpy-gh #20520 + """ + MNAME = "hi" + foutl = get_io_paths(hello_world_f77, mname=MNAME) + ipath = foutl.finp + monkeypatch.setattr(sys, "argv", f'f2py {ipath} -m {MNAME} -h hi.pyf'.split()) + with util.switchdir(ipath.parent): + f2pycli() + with Path('hi.pyf').open() as hipyf: + pyfdat = hipyf.read() + assert "python module hi" in pyfdat + + +def test_lower_cmod(capfd, hello_world_f77, monkeypatch): + """Lowers cases by flag or when -h is present + + CLI :: --[no-]lower + """ + foutl = get_io_paths(hello_world_f77, mname="test") + ipath = foutl.finp + capshi = re.compile(r"HI\(\)") + capslo = re.compile(r"hi\(\)") + # Case I: --lower is passed + monkeypatch.setattr(sys, "argv", f'f2py {ipath} -m test --lower'.split()) + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert capslo.search(out) is not None + assert capshi.search(out) is None + # Case II: --no-lower is passed + monkeypatch.setattr(sys, "argv", + f'f2py {ipath} -m test --no-lower'.split()) + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert capslo.search(out) is None + assert capshi.search(out) is not None + + +def test_lower_sig(capfd, hello_world_f77, monkeypatch): + """Lowers cases in signature files by flag or when -h is present + + CLI :: --[no-]lower -h + """ + foutl = get_io_paths(hello_world_f77, mname="test") + ipath = foutl.finp + # Signature files + capshi = re.compile(r"Block: HI") + capslo = re.compile(r"Block: hi") + # Case I: --lower is implied by -h + # TODO: Clean up to prevent passing --overwrite-signature + monkeypatch.setattr( + sys, + "argv", + f'f2py {ipath} -h {foutl.pyf} -m test --overwrite-signature'.split(), + ) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert capslo.search(out) is not None + assert capshi.search(out) is None + + # Case II: --no-lower overrides -h + monkeypatch.setattr( + sys, + "argv", + f'f2py {ipath} -h {foutl.pyf} -m test --overwrite-signature --no-lower' + .split(), + ) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert capslo.search(out) is None + assert capshi.search(out) is not None + + +def test_build_dir(capfd, hello_world_f90, monkeypatch): + """Ensures that the build directory can be specified + + CLI :: --build-dir + """ + ipath = Path(hello_world_f90) + mname = "blah" + odir = "tttmp" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --build-dir {odir}'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert f"Wrote C/API module \"{mname}\"" in out + + +def test_overwrite(capfd, hello_world_f90, monkeypatch): + """Ensures that the build directory can be specified + + CLI :: --overwrite-signature + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr( + sys, "argv", + f'f2py -h faker.pyf {ipath} --overwrite-signature'.split()) + + with util.switchdir(ipath.parent): + Path("faker.pyf").write_text("Fake news", encoding="ascii") + f2pycli() + out, _ = capfd.readouterr() + assert "Saving signatures to file" in out + + +def test_latexdoc(capfd, hello_world_f90, monkeypatch): + """Ensures that TeX documentation is written out + + CLI :: --latex-doc + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --latex-doc'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Documentation is saved to file" in out + with Path(f"{mname}module.tex").open() as otex: + assert "\\documentclass" in otex.read() + + +def test_nolatexdoc(capfd, hello_world_f90, monkeypatch): + """Ensures that TeX documentation is written out + + CLI :: --no-latex-doc + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --no-latex-doc'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Documentation is saved to file" not in out + + +def test_shortlatex(capfd, hello_world_f90, monkeypatch): + """Ensures that truncated documentation is written out + + TODO: Test to ensure this has no effect without --latex-doc + CLI :: --latex-doc --short-latex + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr( + sys, + "argv", + f'f2py -m {mname} {ipath} --latex-doc --short-latex'.split(), + ) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Documentation is saved to file" in out + with Path(f"./{mname}module.tex").open() as otex: + assert "\\documentclass" not in otex.read() + + +def test_restdoc(capfd, hello_world_f90, monkeypatch): + """Ensures that RsT documentation is written out + + CLI :: --rest-doc + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --rest-doc'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "ReST Documentation is saved to file" in out + with Path(f"./{mname}module.rest").open() as orst: + assert r".. -*- rest -*-" in orst.read() + + +def test_norestexdoc(capfd, hello_world_f90, monkeypatch): + """Ensures that TeX documentation is written out + + CLI :: --no-rest-doc + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --no-rest-doc'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "ReST Documentation is saved to file" not in out + + +def test_debugcapi(capfd, hello_world_f90, monkeypatch): + """Ensures that debugging wrappers are written + + CLI :: --debug-capi + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --debug-capi'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + with Path(f"./{mname}module.c").open() as ocmod: + assert r"#define DEBUGCFUNCS" in ocmod.read() + + +@pytest.mark.skip(reason="Consistently fails on CI; noisy so skip not xfail.") +def test_debugcapi_bld(hello_world_f90, monkeypatch): + """Ensures that debugging wrappers work + + CLI :: --debug-capi -c + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} -c --debug-capi'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + cmd_run = shlex.split(f"{sys.executable} -c \"import blah; blah.hi()\"") + rout = subprocess.run(cmd_run, capture_output=True, encoding='UTF-8') + eout = ' Hello World\n' + eerr = textwrap.dedent("""\ +debug-capi:Python C/API function blah.hi() +debug-capi:float hi=:output,hidden,scalar +debug-capi:hi=0 +debug-capi:Fortran subroutine `f2pywraphi(&hi)' +debug-capi:hi=0 +debug-capi:Building return value. +debug-capi:Python C/API function blah.hi: successful. +debug-capi:Freeing memory. + """) + assert rout.stdout == eout + assert rout.stderr == eerr + + +def test_wrapfunc_def(capfd, hello_world_f90, monkeypatch): + """Ensures that fortran subroutine wrappers for F77 are included by default + + CLI :: --[no]-wrap-functions + """ + # Implied + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", f'f2py -m {mname} {ipath}'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert r"Fortran 77 wrappers are saved to" in out + + # Explicit + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --wrap-functions'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert r"Fortran 77 wrappers are saved to" in out + + +def test_nowrapfunc(capfd, hello_world_f90, monkeypatch): + """Ensures that fortran subroutine wrappers for F77 can be disabled + + CLI :: --no-wrap-functions + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr(sys, "argv", + f'f2py -m {mname} {ipath} --no-wrap-functions'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert r"Fortran 77 wrappers are saved to" not in out + + +def test_inclheader(capfd, hello_world_f90, monkeypatch): + """Add to the include directories + + CLI :: -include + TODO: Document this in the help string + """ + ipath = Path(hello_world_f90) + mname = "blah" + monkeypatch.setattr( + sys, + "argv", + f'f2py -m {mname} {ipath} -include -include '. + split(), + ) + + with util.switchdir(ipath.parent): + f2pycli() + with Path(f"./{mname}module.c").open() as ocmod: + ocmr = ocmod.read() + assert "#include " in ocmr + assert "#include " in ocmr + + +def test_inclpath(): + """Add to the include directories + + CLI :: --include-paths + """ + # TODO: populate + pass + + +def test_hlink(): + """Add to the include directories + + CLI :: --help-link + """ + # TODO: populate + pass + + +def test_f2cmap(capfd, f2cmap_f90, monkeypatch): + """Check that Fortran-to-Python KIND specs can be passed + + CLI :: --f2cmap + """ + ipath = Path(f2cmap_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} --f2cmap mapfile'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "Reading f2cmap from 'mapfile' ..." in out + assert "Mapping \"real(kind=real32)\" to \"float\"" in out + assert "Mapping \"real(kind=real64)\" to \"double\"" in out + assert "Mapping \"integer(kind=int64)\" to \"long_long\"" in out + assert "Successfully applied user defined f2cmap changes" in out + + +def test_quiet(capfd, hello_world_f90, monkeypatch): + """Reduce verbosity + + CLI :: --quiet + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} --quiet'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert len(out) == 0 + + +def test_verbose(capfd, hello_world_f90, monkeypatch): + """Increase verbosity + + CLI :: --verbose + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} --verbose'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + out, _ = capfd.readouterr() + assert "analyzeline" in out + + +def test_version(capfd, monkeypatch): + """Ensure version + + CLI :: -v + """ + monkeypatch.setattr(sys, "argv", 'f2py -v'.split()) + # TODO: f2py2e should not call sys.exit() after printing the version + with pytest.raises(SystemExit): + f2pycli() + out, _ = capfd.readouterr() + import numpy as np + assert np.__version__ == out.strip() + + +@pytest.mark.skip(reason="Consistently fails on CI; noisy so skip not xfail.") +def test_npdistop(hello_world_f90, monkeypatch): + """ + CLI :: -c + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} -c'.split()) + + with util.switchdir(ipath.parent): + f2pycli() + cmd_run = shlex.split(f"{sys.executable} -c \"import blah; blah.hi()\"") + rout = subprocess.run(cmd_run, capture_output=True, encoding='UTF-8') + eout = ' Hello World\n' + assert rout.stdout == eout + + +@pytest.mark.skipif((platform.system() != 'Linux') or sys.version_info <= (3, 12), + reason='Compiler and Python 3.12 or newer required') +def test_no_freethreading_compatible(hello_world_f90, monkeypatch): + """ + CLI :: --no-freethreading-compatible + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} -c --no-freethreading-compatible'.split()) + + with util.switchdir(ipath.parent): + compiler_check_f2pycli() + cmd = f"{sys.executable} -c \"import blah; blah.hi();" + if NOGIL_BUILD: + cmd += "import sys; assert sys._is_gil_enabled() is True\"" + else: + cmd += "\"" + cmd_run = shlex.split(cmd) + rout = subprocess.run(cmd_run, capture_output=True, encoding='UTF-8') + eout = ' Hello World\n' + assert rout.stdout == eout + if NOGIL_BUILD: + assert "The global interpreter lock (GIL) has been enabled to load module 'blah'" in rout.stderr + assert rout.returncode == 0 + + +@pytest.mark.skipif((platform.system() != 'Linux') or sys.version_info <= (3, 12), + reason='Compiler and Python 3.12 or newer required') +def test_freethreading_compatible(hello_world_f90, monkeypatch): + """ + CLI :: --freethreading_compatible + """ + ipath = Path(hello_world_f90) + monkeypatch.setattr(sys, "argv", f'f2py -m blah {ipath} -c --freethreading-compatible'.split()) + + with util.switchdir(ipath.parent): + compiler_check_f2pycli() + cmd = f"{sys.executable} -c \"import blah; blah.hi();" + if NOGIL_BUILD: + cmd += "import sys; assert sys._is_gil_enabled() is False\"" + else: + cmd += "\"" + cmd_run = shlex.split(cmd) + rout = subprocess.run(cmd_run, capture_output=True, encoding='UTF-8') + eout = ' Hello World\n' + assert rout.stdout == eout + assert rout.stderr == "" + assert rout.returncode == 0 + + +# Numpy distutils flags +# TODO: These should be tested separately + +def test_npd_fcompiler(): + """ + CLI :: -c --fcompiler + """ + # TODO: populate + pass + + +def test_npd_compiler(): + """ + CLI :: -c --compiler + """ + # TODO: populate + pass + + +def test_npd_help_fcompiler(): + """ + CLI :: -c --help-fcompiler + """ + # TODO: populate + pass + + +def test_npd_f77exec(): + """ + CLI :: -c --f77exec + """ + # TODO: populate + pass + + +def test_npd_f90exec(): + """ + CLI :: -c --f90exec + """ + # TODO: populate + pass + + +def test_npd_f77flags(): + """ + CLI :: -c --f77flags + """ + # TODO: populate + pass + + +def test_npd_f90flags(): + """ + CLI :: -c --f90flags + """ + # TODO: populate + pass + + +def test_npd_opt(): + """ + CLI :: -c --opt + """ + # TODO: populate + pass + + +def test_npd_arch(): + """ + CLI :: -c --arch + """ + # TODO: populate + pass + + +def test_npd_noopt(): + """ + CLI :: -c --noopt + """ + # TODO: populate + pass + + +def test_npd_noarch(): + """ + CLI :: -c --noarch + """ + # TODO: populate + pass + + +def test_npd_debug(): + """ + CLI :: -c --debug + """ + # TODO: populate + pass + + +def test_npd_link_auto(): + """ + CLI :: -c --link- + """ + # TODO: populate + pass + + +def test_npd_lib(): + """ + CLI :: -c -L/path/to/lib/ -l + """ + # TODO: populate + pass + + +def test_npd_define(): + """ + CLI :: -D + """ + # TODO: populate + pass + + +def test_npd_undefine(): + """ + CLI :: -U + """ + # TODO: populate + pass + + +def test_npd_incl(): + """ + CLI :: -I/path/to/include/ + """ + # TODO: populate + pass + + +def test_npd_linker(): + """ + CLI :: .o .so .a + """ + # TODO: populate + pass diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_isoc.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_isoc.py new file mode 100644 index 0000000000000000000000000000000000000000..97f71e6c854cb1d9f9a9acc508b15df0b8123822 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_isoc.py @@ -0,0 +1,53 @@ +from . import util +import numpy as np +import pytest +from numpy.testing import assert_allclose + +class TestISOC(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "isocintrin", "isoCtests.f90"), + ] + + # gh-24553 + @pytest.mark.slow + def test_c_double(self): + out = self.module.coddity.c_add(1, 2) + exp_out = 3 + assert out == exp_out + + # gh-9693 + def test_bindc_function(self): + out = self.module.coddity.wat(1, 20) + exp_out = 8 + assert out == exp_out + + # gh-25207 + def test_bindc_kinds(self): + out = self.module.coddity.c_add_int64(1, 20) + exp_out = 21 + assert out == exp_out + + # gh-25207 + def test_bindc_add_arr(self): + a = np.array([1,2,3]) + b = np.array([1,2,3]) + out = self.module.coddity.add_arr(a, b) + exp_out = a*2 + assert_allclose(out, exp_out) + + +def test_process_f2cmap_dict(): + from numpy.f2py.auxfuncs import process_f2cmap_dict + + f2cmap_all = {"integer": {"8": "rubbish_type"}} + new_map = {"INTEGER": {"4": "int"}} + c2py_map = {"int": "int", "rubbish_type": "long"} + + exp_map, exp_maptyp = ({"integer": {"8": "rubbish_type", "4": "int"}}, ["int"]) + + # Call the function + res_map, res_maptyp = process_f2cmap_dict(f2cmap_all, new_map, c2py_map) + + # Assert the result is as expected + assert res_map == exp_map + assert res_maptyp == exp_maptyp diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_kind.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_kind.py new file mode 100644 index 0000000000000000000000000000000000000000..a8403ca3660686df4a0d612218cdb80b92cec0c8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_kind.py @@ -0,0 +1,49 @@ +import sys +import pytest +import platform + +from numpy.f2py.crackfortran import ( + _selected_int_kind_func as selected_int_kind, + _selected_real_kind_func as selected_real_kind, +) +from . import util + + +class TestKind(util.F2PyTest): + sources = [util.getpath("tests", "src", "kind", "foo.f90")] + + @pytest.mark.skipif(sys.maxsize < 2 ** 31 + 1, + reason="Fails for 32 bit machines") + def test_int(self): + """Test `int` kind_func for integers up to 10**40.""" + selectedintkind = self.module.selectedintkind + + for i in range(40): + assert selectedintkind(i) == selected_int_kind( + i + ), f"selectedintkind({i}): expected {selected_int_kind(i)!r} but got {selectedintkind(i)!r}" + + def test_real(self): + """ + Test (processor-dependent) `real` kind_func for real numbers + of up to 31 digits precision (extended/quadruple). + """ + selectedrealkind = self.module.selectedrealkind + + for i in range(32): + assert selectedrealkind(i) == selected_real_kind( + i + ), f"selectedrealkind({i}): expected {selected_real_kind(i)!r} but got {selectedrealkind(i)!r}" + + @pytest.mark.xfail(platform.machine().lower().startswith("ppc"), + reason="Some PowerPC may not support full IEEE 754 precision") + def test_quad_precision(self): + """ + Test kind_func for quadruple precision [`real(16)`] of 32+ digits . + """ + selectedrealkind = self.module.selectedrealkind + + for i in range(32, 40): + assert selectedrealkind(i) == selected_real_kind( + i + ), f"selectedrealkind({i}): expected {selected_real_kind(i)!r} but got {selectedrealkind(i)!r}" diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_mixed.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_mixed.py new file mode 100644 index 0000000000000000000000000000000000000000..688c1630fda60853d4213107e33e90f17be70d80 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_mixed.py @@ -0,0 +1,33 @@ +import textwrap +import pytest + +from numpy.testing import IS_PYPY +from . import util + + +class TestMixed(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "mixed", "foo.f"), + util.getpath("tests", "src", "mixed", "foo_fixed.f90"), + util.getpath("tests", "src", "mixed", "foo_free.f90"), + ] + + @pytest.mark.slow + def test_all(self): + assert self.module.bar11() == 11 + assert self.module.foo_fixed.bar12() == 12 + assert self.module.foo_free.bar13() == 13 + + @pytest.mark.xfail(IS_PYPY, + reason="PyPy cannot modify tp_doc after PyType_Ready") + def test_docstring(self): + expected = textwrap.dedent("""\ + a = bar11() + + Wrapper for ``bar11``. + + Returns + ------- + a : int + """) + assert self.module.bar11.__doc__ == expected diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_modules.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..436e0c70001795e292cbb6a5cb022dc0378301ca --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_modules.py @@ -0,0 +1,81 @@ +import pytest +import textwrap + +from . import util +from numpy.testing import IS_PYPY + + +@pytest.mark.slow +class TestModuleFilterPublicEntities(util.F2PyTest): + sources = [ + util.getpath( + "tests", "src", "modules", "gh26920", + "two_mods_with_one_public_routine.f90" + ) + ] + # we filter the only public function mod2 + only = ["mod1_func1", ] + + def test_gh26920(self): + # if it compiles and can be loaded, things are fine + pass + + +@pytest.mark.slow +class TestModuleWithoutPublicEntities(util.F2PyTest): + sources = [ + util.getpath( + "tests", "src", "modules", "gh26920", + "two_mods_with_no_public_entities.f90" + ) + ] + only = ["mod1_func1", ] + + def test_gh26920(self): + # if it compiles and can be loaded, things are fine + pass + + +@pytest.mark.slow +class TestModuleDocString(util.F2PyTest): + sources = [util.getpath("tests", "src", "modules", "module_data_docstring.f90")] + + @pytest.mark.xfail(IS_PYPY, reason="PyPy cannot modify tp_doc after PyType_Ready") + def test_module_docstring(self): + assert self.module.mod.__doc__ == textwrap.dedent( + """\ + i : 'i'-scalar + x : 'i'-array(4) + a : 'f'-array(2,3) + b : 'f'-array(-1,-1), not allocated\x00 + foo()\n + Wrapper for ``foo``.\n\n""" + ) + + +@pytest.mark.slow +class TestModuleAndSubroutine(util.F2PyTest): + module_name = "example" + sources = [ + util.getpath("tests", "src", "modules", "gh25337", "data.f90"), + util.getpath("tests", "src", "modules", "gh25337", "use_data.f90"), + ] + + def test_gh25337(self): + self.module.data.set_shift(3) + assert "data" in dir(self.module) + + +@pytest.mark.slow +class TestUsedModule(util.F2PyTest): + module_name = "fmath" + sources = [ + util.getpath("tests", "src", "modules", "use_modules.f90"), + ] + + def test_gh25867(self): + compiled_mods = [x for x in dir(self.module) if "__" not in x] + assert "useops" in compiled_mods + assert self.module.useops.sum_and_double(3, 7) == 20 + assert "mathops" in compiled_mods + assert self.module.mathops.add(3, 7) == 10 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_parameter.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_parameter.py new file mode 100644 index 0000000000000000000000000000000000000000..154131f49f7bc2b0bcee8f9e9104c98e98d7af3c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_parameter.py @@ -0,0 +1,130 @@ +import pytest + +import numpy as np + +from . import util + + +class TestParameters(util.F2PyTest): + # Check that intent(in out) translates as intent(inout) + sources = [ + util.getpath("tests", "src", "parameter", "constant_real.f90"), + util.getpath("tests", "src", "parameter", "constant_integer.f90"), + util.getpath("tests", "src", "parameter", "constant_both.f90"), + util.getpath("tests", "src", "parameter", "constant_compound.f90"), + util.getpath("tests", "src", "parameter", "constant_non_compound.f90"), + util.getpath("tests", "src", "parameter", "constant_array.f90"), + ] + + @pytest.mark.slow + def test_constant_real_single(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float32)[::2] + pytest.raises(ValueError, self.module.foo_single, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float32) + self.module.foo_single(x) + assert np.allclose(x, [0 + 1 + 2 * 3, 1, 2]) + + @pytest.mark.slow + def test_constant_real_double(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float64)[::2] + pytest.raises(ValueError, self.module.foo_double, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float64) + self.module.foo_double(x) + assert np.allclose(x, [0 + 1 + 2 * 3, 1, 2]) + + @pytest.mark.slow + def test_constant_compound_int(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.int32)[::2] + pytest.raises(ValueError, self.module.foo_compound_int, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.int32) + self.module.foo_compound_int(x) + assert np.allclose(x, [0 + 1 + 2 * 6, 1, 2]) + + @pytest.mark.slow + def test_constant_non_compound_int(self): + # check values + x = np.arange(4, dtype=np.int32) + self.module.foo_non_compound_int(x) + assert np.allclose(x, [0 + 1 + 2 + 3 * 4, 1, 2, 3]) + + @pytest.mark.slow + def test_constant_integer_int(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.int32)[::2] + pytest.raises(ValueError, self.module.foo_int, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.int32) + self.module.foo_int(x) + assert np.allclose(x, [0 + 1 + 2 * 3, 1, 2]) + + @pytest.mark.slow + def test_constant_integer_long(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.int64)[::2] + pytest.raises(ValueError, self.module.foo_long, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.int64) + self.module.foo_long(x) + assert np.allclose(x, [0 + 1 + 2 * 3, 1, 2]) + + @pytest.mark.slow + def test_constant_both(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float64)[::2] + pytest.raises(ValueError, self.module.foo, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float64) + self.module.foo(x) + assert np.allclose(x, [0 + 1 * 3 * 3 + 2 * 3 * 3, 1 * 3, 2 * 3]) + + @pytest.mark.slow + def test_constant_no(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float64)[::2] + pytest.raises(ValueError, self.module.foo_no, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float64) + self.module.foo_no(x) + assert np.allclose(x, [0 + 1 * 3 * 3 + 2 * 3 * 3, 1 * 3, 2 * 3]) + + @pytest.mark.slow + def test_constant_sum(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float64)[::2] + pytest.raises(ValueError, self.module.foo_sum, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float64) + self.module.foo_sum(x) + assert np.allclose(x, [0 + 1 * 3 * 3 + 2 * 3 * 3, 1 * 3, 2 * 3]) + + def test_constant_array(self): + x = np.arange(3, dtype=np.float64) + y = np.arange(5, dtype=np.float64) + z = self.module.foo_array(x, y) + assert np.allclose(x, [0.0, 1./10, 2./10]) + assert np.allclose(y, [0.0, 1.*10, 2.*10, 3.*10, 4.*10]) + assert np.allclose(z, 19.0) + + def test_constant_array_any_index(self): + x = np.arange(6, dtype=np.float64) + y = self.module.foo_array_any_index(x) + assert np.allclose(y, x.reshape((2, 3), order='F')) + + def test_constant_array_delims(self): + x = self.module.foo_array_delims() + assert x == 9 + diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_pyf_src.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_pyf_src.py new file mode 100644 index 0000000000000000000000000000000000000000..f77ded2f31d4443c1bda42bb1c21f79fa100ce23 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_pyf_src.py @@ -0,0 +1,44 @@ +# This test is ported from numpy.distutils +from numpy.f2py._src_pyf import process_str +from numpy.testing import assert_equal + + +pyf_src = """ +python module foo + <_rd=real,double precision> + interface + subroutine foosub(tol) + <_rd>, intent(in,out) :: tol + end subroutine foosub + end interface +end python module foo +""" + +expected_pyf = """ +python module foo + interface + subroutine sfoosub(tol) + real, intent(in,out) :: tol + end subroutine sfoosub + subroutine dfoosub(tol) + double precision, intent(in,out) :: tol + end subroutine dfoosub + end interface +end python module foo +""" + + +def normalize_whitespace(s): + """ + Remove leading and trailing whitespace, and convert internal + stretches of whitespace to a single space. + """ + return ' '.join(s.split()) + + +def test_from_template(): + """Regression test for gh-10712.""" + pyf = process_str(pyf_src) + normalized_pyf = normalize_whitespace(pyf) + normalized_expected_pyf = normalize_whitespace(expected_pyf) + assert_equal(normalized_pyf, normalized_expected_pyf) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_quoted_character.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_quoted_character.py new file mode 100644 index 0000000000000000000000000000000000000000..85e83a781e7b55b8e14e780de12d694858b4a236 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_quoted_character.py @@ -0,0 +1,17 @@ +"""See https://github.com/numpy/numpy/pull/10676. + +""" +import sys +import pytest + +from . import util + + +class TestQuotedCharacter(util.F2PyTest): + sources = [util.getpath("tests", "src", "quoted_character", "foo.f")] + + @pytest.mark.skipif(sys.platform == "win32", + reason="Fails with MinGW64 Gfortran (Issue #9673)") + @pytest.mark.slow + def test_quoted_character(self): + assert self.module.foo() == (b"'", b'"', b";", b"!", b"(", b")") diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_regression.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_regression.py new file mode 100644 index 0000000000000000000000000000000000000000..c62f82ac3fc0f3292af475f2f725674a17e96fd1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_regression.py @@ -0,0 +1,174 @@ +import os +import pytest +import platform + +import numpy as np +import numpy.testing as npt + +from . import util + + +class TestIntentInOut(util.F2PyTest): + # Check that intent(in out) translates as intent(inout) + sources = [util.getpath("tests", "src", "regression", "inout.f90")] + + @pytest.mark.slow + def test_inout(self): + # non-contiguous should raise error + x = np.arange(6, dtype=np.float32)[::2] + pytest.raises(ValueError, self.module.foo, x) + + # check values with contiguous array + x = np.arange(3, dtype=np.float32) + self.module.foo(x) + assert np.allclose(x, [3, 1, 2]) + + +class TestDataOnlyMultiModule(util.F2PyTest): + # Check that modules without subroutines work + sources = [util.getpath("tests", "src", "regression", "datonly.f90")] + + @pytest.mark.slow + def test_mdat(self): + assert self.module.datonly.max_value == 100 + assert self.module.dat.max_ == 1009 + int_in = 5 + assert self.module.simple_subroutine(5) == 1014 + + +class TestNegativeBounds(util.F2PyTest): + # Check that negative bounds work correctly + sources = [util.getpath("tests", "src", "negative_bounds", "issue_20853.f90")] + + @pytest.mark.slow + def test_negbound(self): + xvec = np.arange(12) + xlow = -6 + xhigh = 4 + # Calculate the upper bound, + # Keeping the 1 index in mind + def ubound(xl, xh): + return xh - xl + 1 + rval = self.module.foo(is_=xlow, ie_=xhigh, + arr=xvec[:ubound(xlow, xhigh)]) + expval = np.arange(11, dtype = np.float32) + assert np.allclose(rval, expval) + + +class TestNumpyVersionAttribute(util.F2PyTest): + # Check that th attribute __f2py_numpy_version__ is present + # in the compiled module and that has the value np.__version__. + sources = [util.getpath("tests", "src", "regression", "inout.f90")] + + @pytest.mark.slow + def test_numpy_version_attribute(self): + + # Check that self.module has an attribute named "__f2py_numpy_version__" + assert hasattr(self.module, "__f2py_numpy_version__") + + # Check that the attribute __f2py_numpy_version__ is a string + assert isinstance(self.module.__f2py_numpy_version__, str) + + # Check that __f2py_numpy_version__ has the value numpy.__version__ + assert np.__version__ == self.module.__f2py_numpy_version__ + + +def test_include_path(): + incdir = np.f2py.get_include() + fnames_in_dir = os.listdir(incdir) + for fname in ("fortranobject.c", "fortranobject.h"): + assert fname in fnames_in_dir + + +class TestIncludeFiles(util.F2PyTest): + sources = [util.getpath("tests", "src", "regression", "incfile.f90")] + options = [f"-I{util.getpath('tests', 'src', 'regression')}", + f"--include-paths {util.getpath('tests', 'src', 'regression')}"] + + @pytest.mark.slow + def test_gh25344(self): + exp = 7.0 + res = self.module.add(3.0, 4.0) + assert exp == res + +class TestF77Comments(util.F2PyTest): + # Check that comments are stripped from F77 continuation lines + sources = [util.getpath("tests", "src", "regression", "f77comments.f")] + + @pytest.mark.slow + def test_gh26148(self): + x1 = np.array(3, dtype=np.int32) + x2 = np.array(5, dtype=np.int32) + res=self.module.testsub(x1, x2) + assert(res[0] == 8) + assert(res[1] == 15) + + @pytest.mark.slow + def test_gh26466(self): + # Check that comments after PARAMETER directions are stripped + expected = np.arange(1, 11, dtype=np.float32)*2 + res=self.module.testsub2() + npt.assert_allclose(expected, res) + +class TestF90Contiuation(util.F2PyTest): + # Check that comments are stripped from F90 continuation lines + sources = [util.getpath("tests", "src", "regression", "f90continuation.f90")] + + @pytest.mark.slow + def test_gh26148b(self): + x1 = np.array(3, dtype=np.int32) + x2 = np.array(5, dtype=np.int32) + res=self.module.testsub(x1, x2) + assert(res[0] == 8) + assert(res[1] == 15) + +class TestLowerF2PYDirectives(util.F2PyTest): + # Check variables are cased correctly + sources = [util.getpath("tests", "src", "regression", "lower_f2py_fortran.f90")] + + @pytest.mark.slow + def test_gh28014(self): + self.module.inquire_next(3) + assert True + +@pytest.mark.slow +def test_gh26623(): + # Including libraries with . should not generate an incorrect meson.build + try: + aa = util.build_module( + [util.getpath("tests", "src", "regression", "f90continuation.f90")], + ["-lfoo.bar"], + module_name="Blah", + ) + except RuntimeError as rerr: + assert "lparen got assign" not in str(rerr) + + +@pytest.mark.slow +@pytest.mark.skipif(platform.system() not in ['Linux', 'Darwin'], reason='Unsupported on this platform for now') +def test_gh25784(): + # Compile dubious file using passed flags + try: + aa = util.build_module( + [util.getpath("tests", "src", "regression", "f77fixedform.f95")], + options=[ + # Meson will collect and dedup these to pass to fortran_args: + "--f77flags='-ffixed-form -O2'", + "--f90flags=\"-ffixed-form -Og\"", + ], + module_name="Blah", + ) + except ImportError as rerr: + assert "unknown_subroutine_" in str(rerr) + + +@pytest.mark.slow +class TestAssignmentOnlyModules(util.F2PyTest): + # Ensure that variables are exposed without functions or subroutines in a module + sources = [util.getpath("tests", "src", "regression", "assignOnlyModule.f90")] + + @pytest.mark.slow + def test_gh27167(self): + assert (self.module.f_globals.n_max == 16) + assert (self.module.f_globals.i_max == 18) + assert (self.module.f_globals.j_max == 72) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_character.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_character.py new file mode 100644 index 0000000000000000000000000000000000000000..078d445a6df6eb37c36e1dc73c4b2a13e250825c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_character.py @@ -0,0 +1,46 @@ +import pytest + +from numpy import array +from . import util +import platform + +IS_S390X = platform.machine() == "s390x" + + +@pytest.mark.slow +class TestReturnCharacter(util.F2PyTest): + def check_function(self, t, tname): + if tname in ["t0", "t1", "s0", "s1"]: + assert t("23") == b"2" + r = t("ab") + assert r == b"a" + r = t(array("ab")) + assert r == b"a" + r = t(array(77, "u1")) + assert r == b"M" + elif tname in ["ts", "ss"]: + assert t(23) == b"23" + assert t("123456789abcdef") == b"123456789a" + elif tname in ["t5", "s5"]: + assert t(23) == b"23" + assert t("ab") == b"ab" + assert t("123456789abcdef") == b"12345" + else: + raise NotImplementedError + + +class TestFReturnCharacter(TestReturnCharacter): + sources = [ + util.getpath("tests", "src", "return_character", "foo77.f"), + util.getpath("tests", "src", "return_character", "foo90.f90"), + ] + + @pytest.mark.xfail(IS_S390X, reason="callback returns ' '") + @pytest.mark.parametrize("name", "t0,t1,t5,s0,s1,s5,ss".split(",")) + def test_all_f77(self, name): + self.check_function(getattr(self.module, name), name) + + @pytest.mark.xfail(IS_S390X, reason="callback returns ' '") + @pytest.mark.parametrize("name", "t0,t1,t5,ts,s0,s1,s5,ss".split(",")) + def test_all_f90(self, name): + self.check_function(getattr(self.module.f90_return_char, name), name) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_complex.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_complex.py new file mode 100644 index 0000000000000000000000000000000000000000..17811f5d98f94ef1d865e49ccd22046051fd68e5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_complex.py @@ -0,0 +1,66 @@ +import pytest + +from numpy import array +from . import util + + +@pytest.mark.slow +class TestReturnComplex(util.F2PyTest): + def check_function(self, t, tname): + if tname in ["t0", "t8", "s0", "s8"]: + err = 1e-5 + else: + err = 0.0 + assert abs(t(234j) - 234.0j) <= err + assert abs(t(234.6) - 234.6) <= err + assert abs(t(234) - 234.0) <= err + assert abs(t(234.6 + 3j) - (234.6 + 3j)) <= err + # assert abs(t('234')-234.)<=err + # assert abs(t('234.6')-234.6)<=err + assert abs(t(-234) + 234.0) <= err + assert abs(t([234]) - 234.0) <= err + assert abs(t((234, )) - 234.0) <= err + assert abs(t(array(234)) - 234.0) <= err + assert abs(t(array(23 + 4j, "F")) - (23 + 4j)) <= err + assert abs(t(array([234])) - 234.0) <= err + assert abs(t(array([[234]])) - 234.0) <= err + assert abs(t(array([234]).astype("b")) + 22.0) <= err + assert abs(t(array([234], "h")) - 234.0) <= err + assert abs(t(array([234], "i")) - 234.0) <= err + assert abs(t(array([234], "l")) - 234.0) <= err + assert abs(t(array([234], "q")) - 234.0) <= err + assert abs(t(array([234], "f")) - 234.0) <= err + assert abs(t(array([234], "d")) - 234.0) <= err + assert abs(t(array([234 + 3j], "F")) - (234 + 3j)) <= err + assert abs(t(array([234], "D")) - 234.0) <= err + + # pytest.raises(TypeError, t, array([234], 'S1')) + pytest.raises(TypeError, t, "abc") + + pytest.raises(IndexError, t, []) + pytest.raises(IndexError, t, ()) + + pytest.raises(TypeError, t, t) + pytest.raises(TypeError, t, {}) + + try: + r = t(10**400) + assert repr(r) in ["(inf+0j)", "(Infinity+0j)"] + except OverflowError: + pass + + +class TestFReturnComplex(TestReturnComplex): + sources = [ + util.getpath("tests", "src", "return_complex", "foo77.f"), + util.getpath("tests", "src", "return_complex", "foo90.f90"), + ] + + @pytest.mark.parametrize("name", "t0,t8,t16,td,s0,s8,s16,sd".split(",")) + def test_all_f77(self, name): + self.check_function(getattr(self.module, name), name) + + @pytest.mark.parametrize("name", "t0,t8,t16,td,s0,s8,s16,sd".split(",")) + def test_all_f90(self, name): + self.check_function(getattr(self.module.f90_return_complex, name), + name) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_integer.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_integer.py new file mode 100644 index 0000000000000000000000000000000000000000..428afec4a0efa9f70acd158f610fdee104dbe82a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_integer.py @@ -0,0 +1,54 @@ +import pytest + +from numpy import array +from . import util + + +@pytest.mark.slow +class TestReturnInteger(util.F2PyTest): + def check_function(self, t, tname): + assert t(123) == 123 + assert t(123.6) == 123 + assert t("123") == 123 + assert t(-123) == -123 + assert t([123]) == 123 + assert t((123, )) == 123 + assert t(array(123)) == 123 + assert t(array(123, "b")) == 123 + assert t(array(123, "h")) == 123 + assert t(array(123, "i")) == 123 + assert t(array(123, "l")) == 123 + assert t(array(123, "B")) == 123 + assert t(array(123, "f")) == 123 + assert t(array(123, "d")) == 123 + + # pytest.raises(ValueError, t, array([123],'S3')) + pytest.raises(ValueError, t, "abc") + + pytest.raises(IndexError, t, []) + pytest.raises(IndexError, t, ()) + + pytest.raises(Exception, t, t) + pytest.raises(Exception, t, {}) + + if tname in ["t8", "s8"]: + pytest.raises(OverflowError, t, 100000000000000000000000) + pytest.raises(OverflowError, t, 10000000011111111111111.23) + + +class TestFReturnInteger(TestReturnInteger): + sources = [ + util.getpath("tests", "src", "return_integer", "foo77.f"), + util.getpath("tests", "src", "return_integer", "foo90.f90"), + ] + + @pytest.mark.parametrize("name", + "t0,t1,t2,t4,t8,s0,s1,s2,s4,s8".split(",")) + def test_all_f77(self, name): + self.check_function(getattr(self.module, name), name) + + @pytest.mark.parametrize("name", + "t0,t1,t2,t4,t8,s0,s1,s2,s4,s8".split(",")) + def test_all_f90(self, name): + self.check_function(getattr(self.module.f90_return_integer, name), + name) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_logical.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_logical.py new file mode 100644 index 0000000000000000000000000000000000000000..92fb902af4ddd269d67c427bc5090aabc35513dd --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_logical.py @@ -0,0 +1,64 @@ +import pytest + +from numpy import array +from . import util + + +class TestReturnLogical(util.F2PyTest): + def check_function(self, t): + assert t(True) == 1 + assert t(False) == 0 + assert t(0) == 0 + assert t(None) == 0 + assert t(0.0) == 0 + assert t(0j) == 0 + assert t(1j) == 1 + assert t(234) == 1 + assert t(234.6) == 1 + assert t(234.6 + 3j) == 1 + assert t("234") == 1 + assert t("aaa") == 1 + assert t("") == 0 + assert t([]) == 0 + assert t(()) == 0 + assert t({}) == 0 + assert t(t) == 1 + assert t(-234) == 1 + assert t(10**100) == 1 + assert t([234]) == 1 + assert t((234, )) == 1 + assert t(array(234)) == 1 + assert t(array([234])) == 1 + assert t(array([[234]])) == 1 + assert t(array([127], "b")) == 1 + assert t(array([234], "h")) == 1 + assert t(array([234], "i")) == 1 + assert t(array([234], "l")) == 1 + assert t(array([234], "f")) == 1 + assert t(array([234], "d")) == 1 + assert t(array([234 + 3j], "F")) == 1 + assert t(array([234], "D")) == 1 + assert t(array(0)) == 0 + assert t(array([0])) == 0 + assert t(array([[0]])) == 0 + assert t(array([0j])) == 0 + assert t(array([1])) == 1 + pytest.raises(ValueError, t, array([0, 0])) + + +class TestFReturnLogical(TestReturnLogical): + sources = [ + util.getpath("tests", "src", "return_logical", "foo77.f"), + util.getpath("tests", "src", "return_logical", "foo90.f90"), + ] + + @pytest.mark.slow + @pytest.mark.parametrize("name", "t0,t1,t2,t4,s0,s1,s2,s4".split(",")) + def test_all_f77(self, name): + self.check_function(getattr(self.module, name)) + + @pytest.mark.slow + @pytest.mark.parametrize("name", + "t0,t1,t2,t4,t8,s0,s1,s2,s4,s8".split(",")) + def test_all_f90(self, name): + self.check_function(getattr(self.module.f90_return_logical, name)) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_real.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_real.py new file mode 100644 index 0000000000000000000000000000000000000000..25b638890a961dd69441cc5b27ea1b384c2811de --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_return_real.py @@ -0,0 +1,107 @@ +import platform +import pytest + +from numpy import array +from numpy.testing import IS_64BIT +from . import util + + +@pytest.mark.slow +class TestReturnReal(util.F2PyTest): + def check_function(self, t, tname): + if tname in ["t0", "t4", "s0", "s4"]: + err = 1e-5 + else: + err = 0.0 + assert abs(t(234) - 234.0) <= err + assert abs(t(234.6) - 234.6) <= err + assert abs(t("234") - 234) <= err + assert abs(t("234.6") - 234.6) <= err + assert abs(t(-234) + 234) <= err + assert abs(t([234]) - 234) <= err + assert abs(t((234, )) - 234.0) <= err + assert abs(t(array(234)) - 234.0) <= err + assert abs(t(array(234).astype("b")) + 22) <= err + assert abs(t(array(234, "h")) - 234.0) <= err + assert abs(t(array(234, "i")) - 234.0) <= err + assert abs(t(array(234, "l")) - 234.0) <= err + assert abs(t(array(234, "B")) - 234.0) <= err + assert abs(t(array(234, "f")) - 234.0) <= err + assert abs(t(array(234, "d")) - 234.0) <= err + if tname in ["t0", "t4", "s0", "s4"]: + assert t(1e200) == t(1e300) # inf + + # pytest.raises(ValueError, t, array([234], 'S1')) + pytest.raises(ValueError, t, "abc") + + pytest.raises(IndexError, t, []) + pytest.raises(IndexError, t, ()) + + pytest.raises(Exception, t, t) + pytest.raises(Exception, t, {}) + + try: + r = t(10**400) + assert repr(r) in ["inf", "Infinity"] + except OverflowError: + pass + + +@pytest.mark.skipif( + platform.system() == "Darwin", + reason="Prone to error when run with numpy/f2py/tests on mac os, " + "but not when run in isolation", +) +@pytest.mark.skipif( + not IS_64BIT, reason="32-bit builds are buggy" +) +class TestCReturnReal(TestReturnReal): + suffix = ".pyf" + module_name = "c_ext_return_real" + code = """ +python module c_ext_return_real +usercode \'\'\' +float t4(float value) { return value; } +void s4(float *t4, float value) { *t4 = value; } +double t8(double value) { return value; } +void s8(double *t8, double value) { *t8 = value; } +\'\'\' +interface + function t4(value) + real*4 intent(c) :: t4,value + end + function t8(value) + real*8 intent(c) :: t8,value + end + subroutine s4(t4,value) + intent(c) s4 + real*4 intent(out) :: t4 + real*4 intent(c) :: value + end + subroutine s8(t8,value) + intent(c) s8 + real*8 intent(out) :: t8 + real*8 intent(c) :: value + end +end interface +end python module c_ext_return_real + """ + + @pytest.mark.parametrize("name", "t4,t8,s4,s8".split(",")) + def test_all(self, name): + self.check_function(getattr(self.module, name), name) + + +class TestFReturnReal(TestReturnReal): + sources = [ + util.getpath("tests", "src", "return_real", "foo77.f"), + util.getpath("tests", "src", "return_real", "foo90.f90"), + ] + + @pytest.mark.parametrize("name", "t0,t4,t8,td,s0,s4,s8,sd".split(",")) + def test_all_f77(self, name): + self.check_function(getattr(self.module, name), name) + + @pytest.mark.parametrize("name", "t0,t4,t8,td,s0,s4,s8,sd".split(",")) + def test_all_f90(self, name): + self.check_function(getattr(self.module.f90_return_real, name), name) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_routines.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_routines.py new file mode 100644 index 0000000000000000000000000000000000000000..d6ab475d899e105c5194537cf537c1fe3b71e1e9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_routines.py @@ -0,0 +1,28 @@ +import pytest +from . import util + + +@pytest.mark.slow +class TestRenamedFunc(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "routines", "funcfortranname.f"), + util.getpath("tests", "src", "routines", "funcfortranname.pyf"), + ] + module_name = "funcfortranname" + + def test_gh25799(self): + assert dir(self.module) + assert self.module.funcfortranname_default(200, 12) == 212 + + +@pytest.mark.slow +class TestRenamedSubroutine(util.F2PyTest): + sources = [ + util.getpath("tests", "src", "routines", "subrout.f"), + util.getpath("tests", "src", "routines", "subrout.pyf"), + ] + module_name = "subrout" + + def test_renamed_subroutine(self): + assert dir(self.module) + assert self.module.subrout_default(200, 12) == 212 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_semicolon_split.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_semicolon_split.py new file mode 100644 index 0000000000000000000000000000000000000000..8a9eb87435016ac69be071253ea25335b66f61a9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_semicolon_split.py @@ -0,0 +1,74 @@ +import platform +import pytest + +from numpy.testing import IS_64BIT + +from . import util + + +@pytest.mark.skipif( + platform.system() == "Darwin", + reason="Prone to error when run with numpy/f2py/tests on mac os, " + "but not when run in isolation", +) +@pytest.mark.skipif( + not IS_64BIT, reason="32-bit builds are buggy" +) +class TestMultiline(util.F2PyTest): + suffix = ".pyf" + module_name = "multiline" + code = f""" +python module {module_name} + usercode ''' +void foo(int* x) {{ + char dummy = ';'; + *x = 42; +}} +''' + interface + subroutine foo(x) + intent(c) foo + integer intent(out) :: x + end subroutine foo + end interface +end python module {module_name} + """ + + def test_multiline(self): + assert self.module.foo() == 42 + + +@pytest.mark.skipif( + platform.system() == "Darwin", + reason="Prone to error when run with numpy/f2py/tests on mac os, " + "but not when run in isolation", +) +@pytest.mark.skipif( + not IS_64BIT, reason="32-bit builds are buggy" +) +@pytest.mark.slow +class TestCallstatement(util.F2PyTest): + suffix = ".pyf" + module_name = "callstatement" + code = f""" +python module {module_name} + usercode ''' +void foo(int* x) {{ +}} +''' + interface + subroutine foo(x) + intent(c) foo + integer intent(out) :: x + callprotoargument int* + callstatement {{ & + ; & + x = 42; & + }} + end subroutine foo + end interface +end python module {module_name} + """ + + def test_callstatement(self): + assert self.module.foo() == 42 diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_size.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_size.py new file mode 100644 index 0000000000000000000000000000000000000000..b354711b457f5ed9920ba1d0118aba64ac90fc74 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_size.py @@ -0,0 +1,44 @@ +import pytest +import numpy as np + +from . import util + + +class TestSizeSumExample(util.F2PyTest): + sources = [util.getpath("tests", "src", "size", "foo.f90")] + + @pytest.mark.slow + def test_all(self): + r = self.module.foo([[]]) + assert r == [0] + + r = self.module.foo([[1, 2]]) + assert r == [3] + + r = self.module.foo([[1, 2], [3, 4]]) + assert np.allclose(r, [3, 7]) + + r = self.module.foo([[1, 2], [3, 4], [5, 6]]) + assert np.allclose(r, [3, 7, 11]) + + @pytest.mark.slow + def test_transpose(self): + r = self.module.trans([[]]) + assert np.allclose(r.T, np.array([[]])) + + r = self.module.trans([[1, 2]]) + assert np.allclose(r, [[1.], [2.]]) + + r = self.module.trans([[1, 2, 3], [4, 5, 6]]) + assert np.allclose(r, [[1, 4], [2, 5], [3, 6]]) + + @pytest.mark.slow + def test_flatten(self): + r = self.module.flatten([[]]) + assert np.allclose(r, []) + + r = self.module.flatten([[1, 2]]) + assert np.allclose(r, [1, 2]) + + r = self.module.flatten([[1, 2, 3], [4, 5, 6]]) + assert np.allclose(r, [1, 2, 3, 4, 5, 6]) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_string.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_string.py new file mode 100644 index 0000000000000000000000000000000000000000..1888f649f543c9965461ef79bb40b1ca18118ea6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_string.py @@ -0,0 +1,98 @@ +import pytest +import numpy as np +from . import util + + +class TestString(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "char.f90")] + + @pytest.mark.slow + def test_char(self): + strings = np.array(["ab", "cd", "ef"], dtype="c").T + inp, out = self.module.char_test.change_strings( + strings, strings.shape[1]) + assert inp == pytest.approx(strings) + expected = strings.copy() + expected[1, :] = "AAA" + assert out == pytest.approx(expected) + + +class TestDocStringArguments(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "string.f")] + + def test_example(self): + a = np.array(b"123\0\0") + b = np.array(b"123\0\0") + c = np.array(b"123") + d = np.array(b"123") + + self.module.foo(a, b, c, d) + + assert a.tobytes() == b"123\0\0" + assert b.tobytes() == b"B23\0\0" + assert c.tobytes() == b"123" + assert d.tobytes() == b"D23" + + +class TestFixedString(util.F2PyTest): + sources = [util.getpath("tests", "src", "string", "fixed_string.f90")] + + @staticmethod + def _sint(s, start=0, end=None): + """Return the content of a string buffer as integer value. + + For example: + _sint('1234') -> 4321 + _sint('123A') -> 17321 + """ + if isinstance(s, np.ndarray): + s = s.tobytes() + elif isinstance(s, str): + s = s.encode() + assert isinstance(s, bytes) + if end is None: + end = len(s) + i = 0 + for j in range(start, min(end, len(s))): + i += s[j] * 10**j + return i + + def _get_input(self, intent="in"): + if intent in ["in"]: + yield "" + yield "1" + yield "1234" + yield "12345" + yield b"" + yield b"\0" + yield b"1" + yield b"\01" + yield b"1\0" + yield b"1234" + yield b"12345" + yield np.ndarray((), np.bytes_, buffer=b"") # array(b'', dtype='|S0') + yield np.array(b"") # array(b'', dtype='|S1') + yield np.array(b"\0") + yield np.array(b"1") + yield np.array(b"1\0") + yield np.array(b"\01") + yield np.array(b"1234") + yield np.array(b"123\0") + yield np.array(b"12345") + + def test_intent_in(self): + for s in self._get_input(): + r = self.module.test_in_bytes4(s) + # also checks that s is not changed inplace + expected = self._sint(s, end=4) + assert r == expected, s + + def test_intent_inout(self): + for s in self._get_input(intent="inout"): + rest = self._sint(s, start=4) + r = self.module.test_inout_bytes4(s) + expected = self._sint(s, end=4) + assert r == expected + + # check that the rest of input string is preserved + assert rest == self._sint(s, start=4) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_symbolic.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_symbolic.py new file mode 100644 index 0000000000000000000000000000000000000000..8452783111ebe7130d17301d228eb5708e9eced7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_symbolic.py @@ -0,0 +1,494 @@ +import pytest + +from numpy.f2py.symbolic import ( + Expr, + Op, + ArithOp, + Language, + as_symbol, + as_number, + as_string, + as_array, + as_complex, + as_terms, + as_factors, + eliminate_quotes, + insert_quotes, + fromstring, + as_expr, + as_apply, + as_numer_denom, + as_ternary, + as_ref, + as_deref, + normalize, + as_eq, + as_ne, + as_lt, + as_gt, + as_le, + as_ge, +) +from . import util + + +class TestSymbolic(util.F2PyTest): + def test_eliminate_quotes(self): + def worker(s): + r, d = eliminate_quotes(s) + s1 = insert_quotes(r, d) + assert s1 == s + + for kind in ["", "mykind_"]: + worker(kind + '"1234" // "ABCD"') + worker(kind + '"1234" // ' + kind + '"ABCD"') + worker(kind + "\"1234\" // 'ABCD'") + worker(kind + '"1234" // ' + kind + "'ABCD'") + worker(kind + '"1\\"2\'AB\'34"') + worker("a = " + kind + "'1\\'2\"AB\"34'") + + def test_sanity(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + + assert x.op == Op.SYMBOL + assert repr(x) == "Expr(Op.SYMBOL, 'x')" + assert x == x + assert x != y + assert hash(x) is not None + + n = as_number(123) + m = as_number(456) + assert n.op == Op.INTEGER + assert repr(n) == "Expr(Op.INTEGER, (123, 4))" + assert n == n + assert n != m + assert hash(n) is not None + + fn = as_number(12.3) + fm = as_number(45.6) + assert fn.op == Op.REAL + assert repr(fn) == "Expr(Op.REAL, (12.3, 4))" + assert fn == fn + assert fn != fm + assert hash(fn) is not None + + c = as_complex(1, 2) + c2 = as_complex(3, 4) + assert c.op == Op.COMPLEX + assert repr(c) == ("Expr(Op.COMPLEX, (Expr(Op.INTEGER, (1, 4))," + " Expr(Op.INTEGER, (2, 4))))") + assert c == c + assert c != c2 + assert hash(c) is not None + + s = as_string("'123'") + s2 = as_string('"ABC"') + assert s.op == Op.STRING + assert repr(s) == "Expr(Op.STRING, (\"'123'\", 1))", repr(s) + assert s == s + assert s != s2 + + a = as_array((n, m)) + b = as_array((n, )) + assert a.op == Op.ARRAY + assert repr(a) == ("Expr(Op.ARRAY, (Expr(Op.INTEGER, (123, 4))," + " Expr(Op.INTEGER, (456, 4))))") + assert a == a + assert a != b + + t = as_terms(x) + u = as_terms(y) + assert t.op == Op.TERMS + assert repr(t) == "Expr(Op.TERMS, {Expr(Op.SYMBOL, 'x'): 1})" + assert t == t + assert t != u + assert hash(t) is not None + + v = as_factors(x) + w = as_factors(y) + assert v.op == Op.FACTORS + assert repr(v) == "Expr(Op.FACTORS, {Expr(Op.SYMBOL, 'x'): 1})" + assert v == v + assert w != v + assert hash(v) is not None + + t = as_ternary(x, y, z) + u = as_ternary(x, z, y) + assert t.op == Op.TERNARY + assert t == t + assert t != u + assert hash(t) is not None + + e = as_eq(x, y) + f = as_lt(x, y) + assert e.op == Op.RELATIONAL + assert e == e + assert e != f + assert hash(e) is not None + + def test_tostring_fortran(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + n = as_number(123) + m = as_number(456) + a = as_array((n, m)) + c = as_complex(n, m) + + assert str(x) == "x" + assert str(n) == "123" + assert str(a) == "[123, 456]" + assert str(c) == "(123, 456)" + + assert str(Expr(Op.TERMS, {x: 1})) == "x" + assert str(Expr(Op.TERMS, {x: 2})) == "2 * x" + assert str(Expr(Op.TERMS, {x: -1})) == "-x" + assert str(Expr(Op.TERMS, {x: -2})) == "-2 * x" + assert str(Expr(Op.TERMS, {x: 1, y: 1})) == "x + y" + assert str(Expr(Op.TERMS, {x: -1, y: -1})) == "-x - y" + assert str(Expr(Op.TERMS, {x: 2, y: 3})) == "2 * x + 3 * y" + assert str(Expr(Op.TERMS, {x: -2, y: 3})) == "-2 * x + 3 * y" + assert str(Expr(Op.TERMS, {x: 2, y: -3})) == "2 * x - 3 * y" + + assert str(Expr(Op.FACTORS, {x: 1})) == "x" + assert str(Expr(Op.FACTORS, {x: 2})) == "x ** 2" + assert str(Expr(Op.FACTORS, {x: -1})) == "x ** -1" + assert str(Expr(Op.FACTORS, {x: -2})) == "x ** -2" + assert str(Expr(Op.FACTORS, {x: 1, y: 1})) == "x * y" + assert str(Expr(Op.FACTORS, {x: 2, y: 3})) == "x ** 2 * y ** 3" + + v = Expr(Op.FACTORS, {x: 2, Expr(Op.TERMS, {x: 1, y: 1}): 3}) + assert str(v) == "x ** 2 * (x + y) ** 3", str(v) + v = Expr(Op.FACTORS, {x: 2, Expr(Op.FACTORS, {x: 1, y: 1}): 3}) + assert str(v) == "x ** 2 * (x * y) ** 3", str(v) + + assert str(Expr(Op.APPLY, ("f", (), {}))) == "f()" + assert str(Expr(Op.APPLY, ("f", (x, ), {}))) == "f(x)" + assert str(Expr(Op.APPLY, ("f", (x, y), {}))) == "f(x, y)" + assert str(Expr(Op.INDEXING, ("f", x))) == "f[x]" + + assert str(as_ternary(x, y, z)) == "merge(y, z, x)" + assert str(as_eq(x, y)) == "x .eq. y" + assert str(as_ne(x, y)) == "x .ne. y" + assert str(as_lt(x, y)) == "x .lt. y" + assert str(as_le(x, y)) == "x .le. y" + assert str(as_gt(x, y)) == "x .gt. y" + assert str(as_ge(x, y)) == "x .ge. y" + + def test_tostring_c(self): + language = Language.C + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + n = as_number(123) + + assert Expr(Op.FACTORS, {x: 2}).tostring(language=language) == "x * x" + assert (Expr(Op.FACTORS, { + x + y: 2 + }).tostring(language=language) == "(x + y) * (x + y)") + assert Expr(Op.FACTORS, { + x: 12 + }).tostring(language=language) == "pow(x, 12)" + + assert as_apply(ArithOp.DIV, x, + y).tostring(language=language) == "x / y" + assert (as_apply(ArithOp.DIV, x, + x + y).tostring(language=language) == "x / (x + y)") + assert (as_apply(ArithOp.DIV, x - y, x + + y).tostring(language=language) == "(x - y) / (x + y)") + assert (x + (x - y) / (x + y) + + n).tostring(language=language) == "123 + x + (x - y) / (x + y)" + + assert as_ternary(x, y, z).tostring(language=language) == "(x?y:z)" + assert as_eq(x, y).tostring(language=language) == "x == y" + assert as_ne(x, y).tostring(language=language) == "x != y" + assert as_lt(x, y).tostring(language=language) == "x < y" + assert as_le(x, y).tostring(language=language) == "x <= y" + assert as_gt(x, y).tostring(language=language) == "x > y" + assert as_ge(x, y).tostring(language=language) == "x >= y" + + def test_operations(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + + assert x + x == Expr(Op.TERMS, {x: 2}) + assert x - x == Expr(Op.INTEGER, (0, 4)) + assert x + y == Expr(Op.TERMS, {x: 1, y: 1}) + assert x - y == Expr(Op.TERMS, {x: 1, y: -1}) + assert x * x == Expr(Op.FACTORS, {x: 2}) + assert x * y == Expr(Op.FACTORS, {x: 1, y: 1}) + + assert +x == x + assert -x == Expr(Op.TERMS, {x: -1}), repr(-x) + assert 2 * x == Expr(Op.TERMS, {x: 2}) + assert 2 + x == Expr(Op.TERMS, {x: 1, as_number(1): 2}) + assert 2 * x + 3 * y == Expr(Op.TERMS, {x: 2, y: 3}) + assert (x + y) * 2 == Expr(Op.TERMS, {x: 2, y: 2}) + + assert x**2 == Expr(Op.FACTORS, {x: 2}) + assert (x + y)**2 == Expr( + Op.TERMS, + { + Expr(Op.FACTORS, {x: 2}): 1, + Expr(Op.FACTORS, {y: 2}): 1, + Expr(Op.FACTORS, { + x: 1, + y: 1 + }): 2, + }, + ) + assert (x + y) * x == x**2 + x * y + assert (x + y)**2 == x**2 + 2 * x * y + y**2 + assert (x + y)**2 + (x - y)**2 == 2 * x**2 + 2 * y**2 + assert (x + y) * z == x * z + y * z + assert z * (x + y) == x * z + y * z + + assert (x / 2) == as_apply(ArithOp.DIV, x, as_number(2)) + assert (2 * x / 2) == x + assert (3 * x / 2) == as_apply(ArithOp.DIV, 3 * x, as_number(2)) + assert (4 * x / 2) == 2 * x + assert (5 * x / 2) == as_apply(ArithOp.DIV, 5 * x, as_number(2)) + assert (6 * x / 2) == 3 * x + assert ((3 * 5) * x / 6) == as_apply(ArithOp.DIV, 5 * x, as_number(2)) + assert (30 * x**2 * y**4 / (24 * x**3 * y**3)) == as_apply( + ArithOp.DIV, 5 * y, 4 * x) + assert ((15 * x / 6) / 5) == as_apply(ArithOp.DIV, x, + as_number(2)), (15 * x / 6) / 5 + assert (x / (5 / x)) == as_apply(ArithOp.DIV, x**2, as_number(5)) + + assert (x / 2.0) == Expr(Op.TERMS, {x: 0.5}) + + s = as_string('"ABC"') + t = as_string('"123"') + + assert s // t == Expr(Op.STRING, ('"ABC123"', 1)) + assert s // x == Expr(Op.CONCAT, (s, x)) + assert x // s == Expr(Op.CONCAT, (x, s)) + + c = as_complex(1.0, 2.0) + assert -c == as_complex(-1.0, -2.0) + assert c + c == as_expr((1 + 2j) * 2) + assert c * c == as_expr((1 + 2j)**2) + + def test_substitute(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + a = as_array((x, y)) + + assert x.substitute({x: y}) == y + assert (x + y).substitute({x: z}) == y + z + assert (x * y).substitute({x: z}) == y * z + assert (x**4).substitute({x: z}) == z**4 + assert (x / y).substitute({x: z}) == z / y + assert x.substitute({x: y + z}) == y + z + assert a.substitute({x: y + z}) == as_array((y + z, y)) + + assert as_ternary(x, y, + z).substitute({x: y + z}) == as_ternary(y + z, y, z) + assert as_eq(x, y).substitute({x: y + z}) == as_eq(y + z, y) + + def test_fromstring(self): + + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + f = as_symbol("f") + s = as_string('"ABC"') + t = as_string('"123"') + a = as_array((x, y)) + + assert fromstring("x") == x + assert fromstring("+ x") == x + assert fromstring("- x") == -x + assert fromstring("x + y") == x + y + assert fromstring("x + 1") == x + 1 + assert fromstring("x * y") == x * y + assert fromstring("x * 2") == x * 2 + assert fromstring("x / y") == x / y + assert fromstring("x ** 2", language=Language.Python) == x**2 + assert fromstring("x ** 2 ** 3", language=Language.Python) == x**2**3 + assert fromstring("(x + y) * z") == (x + y) * z + + assert fromstring("f(x)") == f(x) + assert fromstring("f(x,y)") == f(x, y) + assert fromstring("f[x]") == f[x] + assert fromstring("f[x][y]") == f[x][y] + + assert fromstring('"ABC"') == s + assert (normalize( + fromstring('"ABC" // "123" ', + language=Language.Fortran)) == s // t) + assert fromstring('f("ABC")') == f(s) + assert fromstring('MYSTRKIND_"ABC"') == as_string('"ABC"', "MYSTRKIND") + + assert fromstring("(/x, y/)") == a, fromstring("(/x, y/)") + assert fromstring("f((/x, y/))") == f(a) + assert fromstring("(/(x+y)*z/)") == as_array(((x + y) * z, )) + + assert fromstring("123") == as_number(123) + assert fromstring("123_2") == as_number(123, 2) + assert fromstring("123_myintkind") == as_number(123, "myintkind") + + assert fromstring("123.0") == as_number(123.0, 4) + assert fromstring("123.0_4") == as_number(123.0, 4) + assert fromstring("123.0_8") == as_number(123.0, 8) + assert fromstring("123.0e0") == as_number(123.0, 4) + assert fromstring("123.0d0") == as_number(123.0, 8) + assert fromstring("123d0") == as_number(123.0, 8) + assert fromstring("123e-0") == as_number(123.0, 4) + assert fromstring("123d+0") == as_number(123.0, 8) + assert fromstring("123.0_myrealkind") == as_number(123.0, "myrealkind") + assert fromstring("3E4") == as_number(30000.0, 4) + + assert fromstring("(1, 2)") == as_complex(1, 2) + assert fromstring("(1e2, PI)") == as_complex(as_number(100.0), + as_symbol("PI")) + + assert fromstring("[1, 2]") == as_array((as_number(1), as_number(2))) + + assert fromstring("POINT(x, y=1)") == as_apply(as_symbol("POINT"), + x, + y=as_number(1)) + assert fromstring( + 'PERSON(name="John", age=50, shape=(/34, 23/))') == as_apply( + as_symbol("PERSON"), + name=as_string('"John"'), + age=as_number(50), + shape=as_array((as_number(34), as_number(23))), + ) + + assert fromstring("x?y:z") == as_ternary(x, y, z) + + assert fromstring("*x") == as_deref(x) + assert fromstring("**x") == as_deref(as_deref(x)) + assert fromstring("&x") == as_ref(x) + assert fromstring("(*x) * (*y)") == as_deref(x) * as_deref(y) + assert fromstring("(*x) * *y") == as_deref(x) * as_deref(y) + assert fromstring("*x * *y") == as_deref(x) * as_deref(y) + assert fromstring("*x**y") == as_deref(x) * as_deref(y) + + assert fromstring("x == y") == as_eq(x, y) + assert fromstring("x != y") == as_ne(x, y) + assert fromstring("x < y") == as_lt(x, y) + assert fromstring("x > y") == as_gt(x, y) + assert fromstring("x <= y") == as_le(x, y) + assert fromstring("x >= y") == as_ge(x, y) + + assert fromstring("x .eq. y", language=Language.Fortran) == as_eq(x, y) + assert fromstring("x .ne. y", language=Language.Fortran) == as_ne(x, y) + assert fromstring("x .lt. y", language=Language.Fortran) == as_lt(x, y) + assert fromstring("x .gt. y", language=Language.Fortran) == as_gt(x, y) + assert fromstring("x .le. y", language=Language.Fortran) == as_le(x, y) + assert fromstring("x .ge. y", language=Language.Fortran) == as_ge(x, y) + + def test_traverse(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + f = as_symbol("f") + + # Use traverse to substitute a symbol + def replace_visit(s, r=z): + if s == x: + return r + + assert x.traverse(replace_visit) == z + assert y.traverse(replace_visit) == y + assert z.traverse(replace_visit) == z + assert (f(y)).traverse(replace_visit) == f(y) + assert (f(x)).traverse(replace_visit) == f(z) + assert (f[y]).traverse(replace_visit) == f[y] + assert (f[z]).traverse(replace_visit) == f[z] + assert (x + y + z).traverse(replace_visit) == (2 * z + y) + assert (x + + f(y, x - z)).traverse(replace_visit) == (z + + f(y, as_number(0))) + assert as_eq(x, y).traverse(replace_visit) == as_eq(z, y) + + # Use traverse to collect symbols, method 1 + function_symbols = set() + symbols = set() + + def collect_symbols(s): + if s.op is Op.APPLY: + oper = s.data[0] + function_symbols.add(oper) + if oper in symbols: + symbols.remove(oper) + elif s.op is Op.SYMBOL and s not in function_symbols: + symbols.add(s) + + (x + f(y, x - z)).traverse(collect_symbols) + assert function_symbols == {f} + assert symbols == {x, y, z} + + # Use traverse to collect symbols, method 2 + def collect_symbols2(expr, symbols): + if expr.op is Op.SYMBOL: + symbols.add(expr) + + symbols = set() + (x + f(y, x - z)).traverse(collect_symbols2, symbols) + assert symbols == {x, y, z, f} + + # Use traverse to partially collect symbols + def collect_symbols3(expr, symbols): + if expr.op is Op.APPLY: + # skip traversing function calls + return expr + if expr.op is Op.SYMBOL: + symbols.add(expr) + + symbols = set() + (x + f(y, x - z)).traverse(collect_symbols3, symbols) + assert symbols == {x} + + def test_linear_solve(self): + x = as_symbol("x") + y = as_symbol("y") + z = as_symbol("z") + + assert x.linear_solve(x) == (as_number(1), as_number(0)) + assert (x + 1).linear_solve(x) == (as_number(1), as_number(1)) + assert (2 * x).linear_solve(x) == (as_number(2), as_number(0)) + assert (2 * x + 3).linear_solve(x) == (as_number(2), as_number(3)) + assert as_number(3).linear_solve(x) == (as_number(0), as_number(3)) + assert y.linear_solve(x) == (as_number(0), y) + assert (y * z).linear_solve(x) == (as_number(0), y * z) + + assert (x + y).linear_solve(x) == (as_number(1), y) + assert (z * x + y).linear_solve(x) == (z, y) + assert ((z + y) * x + y).linear_solve(x) == (z + y, y) + assert (z * y * x + y).linear_solve(x) == (z * y, y) + + pytest.raises(RuntimeError, lambda: (x * x).linear_solve(x)) + + def test_as_numer_denom(self): + x = as_symbol("x") + y = as_symbol("y") + n = as_number(123) + + assert as_numer_denom(x) == (x, as_number(1)) + assert as_numer_denom(x / n) == (x, n) + assert as_numer_denom(n / x) == (n, x) + assert as_numer_denom(x / y) == (x, y) + assert as_numer_denom(x * y) == (x * y, as_number(1)) + assert as_numer_denom(n + x / y) == (x + n * y, y) + assert as_numer_denom(n + x / (y - x / n)) == (y * n**2, y * n - x) + + def test_polynomial_atoms(self): + x = as_symbol("x") + y = as_symbol("y") + n = as_number(123) + + assert x.polynomial_atoms() == {x} + assert n.polynomial_atoms() == set() + assert (y[x]).polynomial_atoms() == {y[x]} + assert (y(x)).polynomial_atoms() == {y(x)} + assert (y(x) + x).polynomial_atoms() == {y(x), x} + assert (y(x) * x[y]).polynomial_atoms() == {y(x), x[y]} + assert (y(x)**x).polynomial_atoms() == {y(x)} diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_value_attrspec.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_value_attrspec.py new file mode 100644 index 0000000000000000000000000000000000000000..1f3fa676ba8cf37e443b3a4e06f31d8f8306bfe7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/test_value_attrspec.py @@ -0,0 +1,14 @@ +import pytest + +from . import util + +class TestValueAttr(util.F2PyTest): + sources = [util.getpath("tests", "src", "value_attrspec", "gh21665.f90")] + + # gh-21665 + @pytest.mark.slow + def test_gh21665(self): + inp = 2 + out = self.module.fortfuncs.square(inp) + exp_out = 4 + assert out == exp_out diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/tests/util.py b/venv/lib/python3.10/site-packages/numpy/f2py/tests/util.py new file mode 100644 index 0000000000000000000000000000000000000000..e2fcc1ba39d49fd77045b5d037a3ac91c3352e1b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/tests/util.py @@ -0,0 +1,441 @@ +""" +Utility functions for + +- building and importing modules on test time, using a temporary location +- detecting if compilers are present +- determining paths to tests + +""" +import glob +import os +import sys +import subprocess +import tempfile +import shutil +import atexit +import pytest +import contextlib +import numpy +import concurrent.futures + +from pathlib import Path +from numpy._utils import asunicode +from numpy.testing import temppath, IS_WASM +from importlib import import_module +from numpy.f2py._backends._meson import MesonBackend + +# +# Check if compilers are available at all... +# + +def check_language(lang, code_snippet=None): + if sys.platform == "win32": + pytest.skip("No Fortran tests on Windows (Issue #25134)", allow_module_level=True) + tmpdir = tempfile.mkdtemp() + try: + meson_file = os.path.join(tmpdir, "meson.build") + with open(meson_file, "w") as f: + f.write("project('check_compilers')\n") + f.write(f"add_languages('{lang}')\n") + if code_snippet: + f.write(f"{lang}_compiler = meson.get_compiler('{lang}')\n") + f.write(f"{lang}_code = '''{code_snippet}'''\n") + f.write( + f"_have_{lang}_feature =" + f"{lang}_compiler.compiles({lang}_code," + f" name: '{lang} feature check')\n" + ) + try: + runmeson = subprocess.run( + ["meson", "setup", "btmp"], + check=False, + cwd=tmpdir, + capture_output=True, + ) + except subprocess.CalledProcessError: + pytest.skip("meson not present, skipping compiler dependent test", allow_module_level=True) + return runmeson.returncode == 0 + finally: + shutil.rmtree(tmpdir) + + +fortran77_code = ''' +C Example Fortran 77 code + PROGRAM HELLO + PRINT *, 'Hello, Fortran 77!' + END +''' + +fortran90_code = ''' +! Example Fortran 90 code +program hello90 + type :: greeting + character(len=20) :: text + end type greeting + + type(greeting) :: greet + greet%text = 'hello, fortran 90!' + print *, greet%text +end program hello90 +''' + +# Dummy class for caching relevant checks +class CompilerChecker: + def __init__(self): + self.compilers_checked = False + self.has_c = False + self.has_f77 = False + self.has_f90 = False + + def check_compilers(self): + if (not self.compilers_checked) and (not sys.platform == "cygwin"): + with concurrent.futures.ThreadPoolExecutor() as executor: + futures = [ + executor.submit(check_language, "c"), + executor.submit(check_language, "fortran", fortran77_code), + executor.submit(check_language, "fortran", fortran90_code) + ] + + self.has_c = futures[0].result() + self.has_f77 = futures[1].result() + self.has_f90 = futures[2].result() + + self.compilers_checked = True + +if not IS_WASM: + checker = CompilerChecker() + checker.check_compilers() + +def has_c_compiler(): + return checker.has_c + +def has_f77_compiler(): + return checker.has_f77 + +def has_f90_compiler(): + return checker.has_f90 + +def has_fortran_compiler(): + return (checker.has_f90 and checker.has_f77) + + +# +# Maintaining a temporary module directory +# + +_module_dir = None +_module_num = 5403 + +if sys.platform == "cygwin": + NUMPY_INSTALL_ROOT = Path(__file__).parent.parent.parent + _module_list = list(NUMPY_INSTALL_ROOT.glob("**/*.dll")) + + +def _cleanup(): + global _module_dir + if _module_dir is not None: + try: + sys.path.remove(_module_dir) + except ValueError: + pass + try: + shutil.rmtree(_module_dir) + except OSError: + pass + _module_dir = None + + +def get_module_dir(): + global _module_dir + if _module_dir is None: + _module_dir = tempfile.mkdtemp() + atexit.register(_cleanup) + if _module_dir not in sys.path: + sys.path.insert(0, _module_dir) + return _module_dir + + +def get_temp_module_name(): + # Assume single-threaded, and the module dir usable only by this thread + global _module_num + get_module_dir() + name = "_test_ext_module_%d" % _module_num + _module_num += 1 + if name in sys.modules: + # this should not be possible, but check anyway + raise RuntimeError("Temporary module name already in use.") + return name + + +def _memoize(func): + memo = {} + + def wrapper(*a, **kw): + key = repr((a, kw)) + if key not in memo: + try: + memo[key] = func(*a, **kw) + except Exception as e: + memo[key] = e + raise + ret = memo[key] + if isinstance(ret, Exception): + raise ret + return ret + + wrapper.__name__ = func.__name__ + return wrapper + + +# +# Building modules +# + + +@_memoize +def build_module(source_files, options=[], skip=[], only=[], module_name=None): + """ + Compile and import a f2py module, built from the given files. + + """ + + code = f"import sys; sys.path = {sys.path!r}; import numpy.f2py; numpy.f2py.main()" + + d = get_module_dir() + # gh-27045 : Skip if no compilers are found + if not has_fortran_compiler(): + pytest.skip("No Fortran compiler available") + + # Copy files + dst_sources = [] + f2py_sources = [] + for fn in source_files: + if not os.path.isfile(fn): + raise RuntimeError("%s is not a file" % fn) + dst = os.path.join(d, os.path.basename(fn)) + shutil.copyfile(fn, dst) + dst_sources.append(dst) + + base, ext = os.path.splitext(dst) + if ext in (".f90", ".f95", ".f", ".c", ".pyf"): + f2py_sources.append(dst) + + assert f2py_sources + + # Prepare options + if module_name is None: + module_name = get_temp_module_name() + gil_options = [] + if '--freethreading-compatible' not in options and '--no-freethreading-compatible' not in options: + # default to disabling the GIL if unset in options + gil_options = ['--freethreading-compatible'] + f2py_opts = ["-c", "-m", module_name] + options + gil_options + f2py_sources + f2py_opts += ["--backend", "meson"] + if skip: + f2py_opts += ["skip:"] + skip + if only: + f2py_opts += ["only:"] + only + + # Build + cwd = os.getcwd() + try: + os.chdir(d) + cmd = [sys.executable, "-c", code] + f2py_opts + p = subprocess.Popen(cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT) + out, err = p.communicate() + if p.returncode != 0: + raise RuntimeError("Running f2py failed: %s\n%s" % + (cmd[4:], asunicode(out))) + finally: + os.chdir(cwd) + + # Partial cleanup + for fn in dst_sources: + os.unlink(fn) + + # Rebase (Cygwin-only) + if sys.platform == "cygwin": + # If someone starts deleting modules after import, this will + # need to change to record how big each module is, rather than + # relying on rebase being able to find that from the files. + _module_list.extend( + glob.glob(os.path.join(d, "{:s}*".format(module_name))) + ) + subprocess.check_call( + ["/usr/bin/rebase", "--database", "--oblivious", "--verbose"] + + _module_list + ) + + # Import + return import_module(module_name) + + +@_memoize +def build_code(source_code, + options=[], + skip=[], + only=[], + suffix=None, + module_name=None): + """ + Compile and import Fortran code using f2py. + + """ + if suffix is None: + suffix = ".f" + with temppath(suffix=suffix) as path: + with open(path, "w") as f: + f.write(source_code) + return build_module([path], + options=options, + skip=skip, + only=only, + module_name=module_name) + + +# +# Building with meson +# + + +class SimplifiedMesonBackend(MesonBackend): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def compile(self): + self.write_meson_build(self.build_dir) + self.run_meson(self.build_dir) + + +def build_meson(source_files, module_name=None, **kwargs): + """ + Build a module via Meson and import it. + """ + + # gh-27045 : Skip if no compilers are found + if not has_fortran_compiler(): + pytest.skip("No Fortran compiler available") + + build_dir = get_module_dir() + if module_name is None: + module_name = get_temp_module_name() + + # Initialize the MesonBackend + backend = SimplifiedMesonBackend( + modulename=module_name, + sources=source_files, + extra_objects=kwargs.get("extra_objects", []), + build_dir=build_dir, + include_dirs=kwargs.get("include_dirs", []), + library_dirs=kwargs.get("library_dirs", []), + libraries=kwargs.get("libraries", []), + define_macros=kwargs.get("define_macros", []), + undef_macros=kwargs.get("undef_macros", []), + f2py_flags=kwargs.get("f2py_flags", []), + sysinfo_flags=kwargs.get("sysinfo_flags", []), + fc_flags=kwargs.get("fc_flags", []), + flib_flags=kwargs.get("flib_flags", []), + setup_flags=kwargs.get("setup_flags", []), + remove_build_dir=kwargs.get("remove_build_dir", False), + extra_dat=kwargs.get("extra_dat", {}), + ) + + backend.compile() + + # Import the compiled module + sys.path.insert(0, f"{build_dir}/{backend.meson_build_dir}") + return import_module(module_name) + + +# +# Unittest convenience +# + + +class F2PyTest: + code = None + sources = None + options = [] + skip = [] + only = [] + suffix = ".f" + module = None + _has_c_compiler = None + _has_f77_compiler = None + _has_f90_compiler = None + + @property + def module_name(self): + cls = type(self) + return f'_{cls.__module__.rsplit(".",1)[-1]}_{cls.__name__}_ext_module' + + @classmethod + def setup_class(cls): + if sys.platform == "win32": + pytest.skip("Fails with MinGW64 Gfortran (Issue #9673)") + F2PyTest._has_c_compiler = has_c_compiler() + F2PyTest._has_f77_compiler = has_f77_compiler() + F2PyTest._has_f90_compiler = has_f90_compiler() + F2PyTest._has_fortran_compiler = has_fortran_compiler() + + def setup_method(self): + if self.module is not None: + return + + codes = self.sources if self.sources else [] + if self.code: + codes.append(self.suffix) + + needs_f77 = any(str(fn).endswith(".f") for fn in codes) + needs_f90 = any(str(fn).endswith(".f90") for fn in codes) + needs_pyf = any(str(fn).endswith(".pyf") for fn in codes) + + if needs_f77 and not self._has_f77_compiler: + pytest.skip("No Fortran 77 compiler available") + if needs_f90 and not self._has_f90_compiler: + pytest.skip("No Fortran 90 compiler available") + if needs_pyf and not self._has_fortran_compiler: + pytest.skip("No Fortran compiler available") + + # Build the module + if self.code is not None: + self.module = build_code( + self.code, + options=self.options, + skip=self.skip, + only=self.only, + suffix=self.suffix, + module_name=self.module_name, + ) + + if self.sources is not None: + self.module = build_module( + self.sources, + options=self.options, + skip=self.skip, + only=self.only, + module_name=self.module_name, + ) + + +# +# Helper functions +# + + +def getpath(*a): + # Package root + d = Path(numpy.f2py.__file__).parent.resolve() + return d.joinpath(*a) + + +@contextlib.contextmanager +def switchdir(path): + curpath = Path.cwd() + os.chdir(path) + try: + yield + finally: + os.chdir(curpath) diff --git a/venv/lib/python3.10/site-packages/numpy/f2py/use_rules.py b/venv/lib/python3.10/site-packages/numpy/f2py/use_rules.py new file mode 100644 index 0000000000000000000000000000000000000000..19c111aae56d4919d6e72e17ebaece31d5bafc82 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/f2py/use_rules.py @@ -0,0 +1,106 @@ +""" +Build 'use others module data' mechanism for f2py2e. + +Copyright 1999 -- 2011 Pearu Peterson all rights reserved. +Copyright 2011 -- present NumPy Developers. +Permission to use, modify, and distribute this software is given under the +terms of the NumPy License. + +NO WARRANTY IS EXPRESSED OR IMPLIED. USE AT YOUR OWN RISK. +""" +__version__ = "$Revision: 1.3 $"[10:-1] + +f2py_version = 'See `f2py -v`' + + +from .auxfuncs import ( + applyrules, dictappend, gentitle, hasnote, outmess +) + + +usemodule_rules = { + 'body': """ +#begintitle# +static char doc_#apiname#[] = \"\\\nVariable wrapper signature:\\n\\ +\t #name# = get_#name#()\\n\\ +Arguments:\\n\\ +#docstr#\"; +extern F_MODFUNC(#usemodulename#,#USEMODULENAME#,#realname#,#REALNAME#); +static PyObject *#apiname#(PyObject *capi_self, PyObject *capi_args) { +/*#decl#*/ +\tif (!PyArg_ParseTuple(capi_args, \"\")) goto capi_fail; +printf(\"c: %d\\n\",F_MODFUNC(#usemodulename#,#USEMODULENAME#,#realname#,#REALNAME#)); +\treturn Py_BuildValue(\"\"); +capi_fail: +\treturn NULL; +} +""", + 'method': '\t{\"get_#name#\",#apiname#,METH_VARARGS|METH_KEYWORDS,doc_#apiname#},', + 'need': ['F_MODFUNC'] +} + +################ + + +def buildusevars(m, r): + ret = {} + outmess( + '\t\tBuilding use variable hooks for module "%s" (feature only for F90/F95)...\n' % (m['name'])) + varsmap = {} + revmap = {} + if 'map' in r: + for k in r['map'].keys(): + if r['map'][k] in revmap: + outmess('\t\t\tVariable "%s<=%s" is already mapped by "%s". Skipping.\n' % ( + r['map'][k], k, revmap[r['map'][k]])) + else: + revmap[r['map'][k]] = k + if r.get('only'): + for v in r['map'].keys(): + if r['map'][v] in m['vars']: + + if revmap[r['map'][v]] == v: + varsmap[v] = r['map'][v] + else: + outmess('\t\t\tIgnoring map "%s=>%s". See above.\n' % + (v, r['map'][v])) + else: + outmess( + '\t\t\tNo definition for variable "%s=>%s". Skipping.\n' % (v, r['map'][v])) + else: + for v in m['vars'].keys(): + if v in revmap: + varsmap[v] = revmap[v] + else: + varsmap[v] = v + for v in varsmap.keys(): + ret = dictappend(ret, buildusevar(v, varsmap[v], m['vars'], m['name'])) + return ret + + +def buildusevar(name, realname, vars, usemodulename): + outmess('\t\t\tConstructing wrapper function for variable "%s=>%s"...\n' % ( + name, realname)) + ret = {} + vrd = {'name': name, + 'realname': realname, + 'REALNAME': realname.upper(), + 'usemodulename': usemodulename, + 'USEMODULENAME': usemodulename.upper(), + 'texname': name.replace('_', '\\_'), + 'begintitle': gentitle('%s=>%s' % (name, realname)), + 'endtitle': gentitle('end of %s=>%s' % (name, realname)), + 'apiname': '#modulename#_use_%s_from_%s' % (realname, usemodulename) + } + nummap = {0: 'Ro', 1: 'Ri', 2: 'Rii', 3: 'Riii', 4: 'Riv', + 5: 'Rv', 6: 'Rvi', 7: 'Rvii', 8: 'Rviii', 9: 'Rix'} + vrd['texnamename'] = name + for i in nummap.keys(): + vrd['texnamename'] = vrd['texnamename'].replace(repr(i), nummap[i]) + if hasnote(vars[realname]): + vrd['note'] = vars[realname]['note'] + rd = dictappend({}, vrd) + + print(name, realname, vars[realname]) + ret = applyrules(usemodule_rules, rd) + return ret diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__init__.py b/venv/lib/python3.10/site-packages/numpy/lib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..928121ce8f28bf49ab30283d5901aa8bb0414d29 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/__init__.py @@ -0,0 +1,94 @@ +""" +``numpy.lib`` is mostly a space for implementing functions that don't +belong in core or in another NumPy submodule with a clear purpose +(e.g. ``random``, ``fft``, ``linalg``, ``ma``). + +``numpy.lib``'s private submodules contain basic functions that are used by +other public modules and are useful to have in the main name-space. + +""" + +# Public submodules +# Note: recfunctions and (maybe) format are public too, but not imported +from . import array_utils +from . import introspect +from . import mixins +from . import npyio +from . import scimath +from . import stride_tricks + +# Private submodules +# load module names. See https://github.com/networkx/networkx/issues/5838 +from . import _type_check_impl +from . import _index_tricks_impl +from . import _nanfunctions_impl +from . import _function_base_impl +from . import _stride_tricks_impl +from . import _shape_base_impl +from . import _twodim_base_impl +from . import _ufunclike_impl +from . import _histograms_impl +from . import _utils_impl +from . import _arraysetops_impl +from . import _polynomial_impl +from . import _npyio_impl +from . import _arrayterator_impl +from . import _arraypad_impl +from . import _version + +# numpy.lib namespace members +from ._arrayterator_impl import Arrayterator +from ._version import NumpyVersion +from numpy._core._multiarray_umath import add_docstring, tracemalloc_domain +from numpy._core.function_base import add_newdoc + +__all__ = [ + "Arrayterator", "add_docstring", "add_newdoc", "array_utils", + "introspect", "mixins", "NumpyVersion", "npyio", "scimath", + "stride_tricks", "tracemalloc_domain" +] + +add_newdoc.__module__ = "numpy.lib" + +from numpy._pytesttester import PytestTester +test = PytestTester(__name__) +del PytestTester + +def __getattr__(attr): + # Warn for deprecated/removed aliases + import math + import warnings + + if attr == "math": + warnings.warn( + "`np.lib.math` is a deprecated alias for the standard library " + "`math` module (Deprecated Numpy 1.25). Replace usages of " + "`numpy.lib.math` with `math`", DeprecationWarning, stacklevel=2) + return math + elif attr == "emath": + raise AttributeError( + "numpy.lib.emath was an alias for emath module that was removed " + "in NumPy 2.0. Replace usages of numpy.lib.emath with " + "numpy.emath.", + name=None + ) + elif attr in ( + "histograms", "type_check", "nanfunctions", "function_base", + "arraypad", "arraysetops", "ufunclike", "utils", "twodim_base", + "shape_base", "polynomial", "index_tricks", + ): + raise AttributeError( + f"numpy.lib.{attr} is now private. If you are using a public " + "function, it should be available in the main numpy namespace, " + "otherwise check the NumPy 2.0 migration guide.", + name=None + ) + elif attr == "arrayterator": + raise AttributeError( + "numpy.lib.arrayterator submodule is now private. To access " + "Arrayterator class use numpy.lib.Arrayterator.", + name=None + ) + else: + raise AttributeError("module {!r} has no attribute " + "{!r}".format(__name__, attr)) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__init__.pyi b/venv/lib/python3.10/site-packages/numpy/lib/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..19d6ea7a4d3f52a70916e9ca25bb2816cdd5e974 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/__init__.pyi @@ -0,0 +1,20 @@ +from numpy._core.multiarray import add_docstring, tracemalloc_domain +from numpy._core.function_base import add_newdoc + +from . import array_utils, format, introspect, mixins, npyio, scimath, stride_tricks # noqa: F401 +from ._version import NumpyVersion +from ._arrayterator_impl import Arrayterator + +__all__ = [ + "Arrayterator", + "add_docstring", + "add_newdoc", + "array_utils", + "introspect", + "mixins", + "NumpyVersion", + "npyio", + "scimath", + "stride_tricks", + "tracemalloc_domain", +] diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb5fcffe82efdd3729ba2cc8c197dbd80e7f48c4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_array_utils_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_array_utils_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70cd9553bf744856f42c385d2fdb63d5a8b6b70c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_array_utils_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraypad_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraypad_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c84ef91e7df5484561b40663668c41e91cc4033a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraypad_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraysetops_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraysetops_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2fe12e35d6f43dd4d0eb23196bc97c98dea6ac62 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arraysetops_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arrayterator_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arrayterator_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e125f338c0ba38a10147852e67becca4245e7e93 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_arrayterator_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_datasource.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_datasource.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..877a8773fef3106b6738a85fa24e09d9140e12e5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_datasource.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_histograms_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_histograms_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2560892936a202786080b9089b9c1d4a25becb55 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_histograms_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_index_tricks_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_index_tricks_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9990e7c976b94c045955d0770bfd1df3183e407 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_index_tricks_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_iotools.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_iotools.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71b947fb4fe6e53c7a8c5e1abca5ef3318a2e32f Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_iotools.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_nanfunctions_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_nanfunctions_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..342ce47cc577fa5b2a8ab2596988edc7218eb6c0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_nanfunctions_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_npyio_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_npyio_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52698a106818f65c69018099da0b5762f01afeaa Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_npyio_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_polynomial_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_polynomial_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4743d534d1fdbd1ac55bf90fb70e84a0e533be94 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_polynomial_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_scimath_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_scimath_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f78cf381a5583855d6142cd3675404f83838ee52 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_scimath_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_shape_base_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_shape_base_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b3a5303fcd270abb052d716e3cbc8dab2ffc4a78 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_shape_base_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_stride_tricks_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_stride_tricks_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..764c6ba0cfcd55ccd3f2e3f88b202f84326408c5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_stride_tricks_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_twodim_base_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_twodim_base_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b97b1ac1234dc5f0340cda32655e1e7d390d84d0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_twodim_base_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_type_check_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_type_check_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b559594edc86a948f024f052fb8dffddc147b32 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_type_check_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_ufunclike_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_ufunclike_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a60fe591dae2e021d5005c5675b0ce45b3ee3314 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_ufunclike_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_user_array_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_user_array_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fbea877a15ebd200eb57c4dd0eb17296050c523 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_user_array_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_utils_impl.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_utils_impl.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..537c009923d77fd7c8297a789c7fa03f950b335a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_utils_impl.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_version.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bbf2838d12b220e9613265cd5428c7961780172 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/_version.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/array_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/array_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee4d18e1be49a6587550ba012ca78e0ce5cb49b5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/array_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/format.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/format.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce25a704bedd0d57695c3d1f65b4089def9f881f Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/format.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/introspect.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/introspect.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1def102e9877925fd3b220cb24bbfad2393cda3b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/introspect.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/mixins.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/mixins.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..104d5e1466a9885a33c560c56ae9968c47da3a39 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/mixins.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/npyio.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/npyio.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d237e72f8ebd4919d676fd8cc259096ff8a1576 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/npyio.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/recfunctions.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/recfunctions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a94f306a63ae821499b6db8cea0f76913e26657 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/recfunctions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/scimath.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/scimath.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9afb9c045491ade8607c07db1e7cbf93c425a1c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/scimath.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/stride_tricks.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/stride_tricks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fe9472174ac3dbc99bd273cb197fb166f33585b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/stride_tricks.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/user_array.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/user_array.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb5b6ad6bf3920df7251b3d8a115b30e304111d4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/__pycache__/user_array.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f778160358a05cc642ef43ebd523cff19e2206 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.py @@ -0,0 +1,62 @@ +""" +Miscellaneous utils. +""" +from numpy._core import asarray +from numpy._core.numeric import normalize_axis_tuple, normalize_axis_index +from numpy._utils import set_module + +__all__ = ["byte_bounds", "normalize_axis_tuple", "normalize_axis_index"] + + +@set_module("numpy.lib.array_utils") +def byte_bounds(a): + """ + Returns pointers to the end-points of an array. + + Parameters + ---------- + a : ndarray + Input array. It must conform to the Python-side of the array + interface. + + Returns + ------- + (low, high) : tuple of 2 integers + The first integer is the first byte of the array, the second + integer is just past the last byte of the array. If `a` is not + contiguous it will not use every byte between the (`low`, `high`) + values. + + Examples + -------- + >>> import numpy as np + >>> I = np.eye(2, dtype='f'); I.dtype + dtype('float32') + >>> low, high = np.lib.array_utils.byte_bounds(I) + >>> high - low == I.size*I.itemsize + True + >>> I = np.eye(2); I.dtype + dtype('float64') + >>> low, high = np.lib.array_utils.byte_bounds(I) + >>> high - low == I.size*I.itemsize + True + + """ + ai = a.__array_interface__ + a_data = ai['data'][0] + astrides = ai['strides'] + ashape = ai['shape'] + bytes_a = asarray(a).dtype.itemsize + + a_low = a_high = a_data + if astrides is None: + # contiguous case + a_high += a.size * bytes_a + else: + for shape, stride in zip(ashape, astrides): + if stride < 0: + a_low += (shape-1)*stride + else: + a_high += (shape-1)*stride + a_high += bytes_a + return a_low, a_high diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..11a2aafb8837dd1ac31f048743202933482da4a3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_array_utils_impl.pyi @@ -0,0 +1,25 @@ +from typing import Any, Iterable + +from numpy import generic +from numpy.typing import NDArray + +__all__ = ["byte_bounds", "normalize_axis_tuple", "normalize_axis_index"] + +# NOTE: In practice `byte_bounds` can (potentially) take any object +# implementing the `__array_interface__` protocol. The caveat is +# that certain keys, marked as optional in the spec, must be present for +# `byte_bounds`. This concerns `"strides"` and `"data"`. +def byte_bounds(a: generic | NDArray[Any]) -> tuple[int, int]: ... + +def normalize_axis_tuple( + axis: int | Iterable[int], + ndim: int = ..., + argname: None | str = ..., + allow_duplicate: None | bool = ..., +) -> tuple[int, int]: ... + +def normalize_axis_index( + axis: int = ..., + ndim: int = ..., + msg_prefix: None | str = ..., +) -> int: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..2e190871722ba8d9fba81bafe66e07a1e5462cdb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.py @@ -0,0 +1,891 @@ +""" +The arraypad module contains a group of functions to pad values onto the edges +of an n-dimensional array. + +""" +import numpy as np +from numpy._core.overrides import array_function_dispatch +from numpy.lib._index_tricks_impl import ndindex + + +__all__ = ['pad'] + + +############################################################################### +# Private utility functions. + + +def _round_if_needed(arr, dtype): + """ + Rounds arr inplace if destination dtype is integer. + + Parameters + ---------- + arr : ndarray + Input array. + dtype : dtype + The dtype of the destination array. + """ + if np.issubdtype(dtype, np.integer): + arr.round(out=arr) + + +def _slice_at_axis(sl, axis): + """ + Construct tuple of slices to slice an array in the given dimension. + + Parameters + ---------- + sl : slice + The slice for the given dimension. + axis : int + The axis to which `sl` is applied. All other dimensions are left + "unsliced". + + Returns + ------- + sl : tuple of slices + A tuple with slices matching `shape` in length. + + Examples + -------- + >>> np._slice_at_axis(slice(None, 3, -1), 1) + (slice(None, None, None), slice(None, 3, -1), (...,)) + """ + return (slice(None),) * axis + (sl,) + (...,) + + +def _view_roi(array, original_area_slice, axis): + """ + Get a view of the current region of interest during iterative padding. + + When padding multiple dimensions iteratively corner values are + unnecessarily overwritten multiple times. This function reduces the + working area for the first dimensions so that corners are excluded. + + Parameters + ---------- + array : ndarray + The array with the region of interest. + original_area_slice : tuple of slices + Denotes the area with original values of the unpadded array. + axis : int + The currently padded dimension assuming that `axis` is padded before + `axis` + 1. + + Returns + ------- + roi : ndarray + The region of interest of the original `array`. + """ + axis += 1 + sl = (slice(None),) * axis + original_area_slice[axis:] + return array[sl] + + +def _pad_simple(array, pad_width, fill_value=None): + """ + Pad array on all sides with either a single value or undefined values. + + Parameters + ---------- + array : ndarray + Array to grow. + pad_width : sequence of tuple[int, int] + Pad width on both sides for each dimension in `arr`. + fill_value : scalar, optional + If provided the padded area is filled with this value, otherwise + the pad area left undefined. + + Returns + ------- + padded : ndarray + The padded array with the same dtype as`array`. Its order will default + to C-style if `array` is not F-contiguous. + original_area_slice : tuple + A tuple of slices pointing to the area of the original array. + """ + # Allocate grown array + new_shape = tuple( + left + size + right + for size, (left, right) in zip(array.shape, pad_width) + ) + order = 'F' if array.flags.fnc else 'C' # Fortran and not also C-order + padded = np.empty(new_shape, dtype=array.dtype, order=order) + + if fill_value is not None: + padded.fill(fill_value) + + # Copy old array into correct space + original_area_slice = tuple( + slice(left, left + size) + for size, (left, right) in zip(array.shape, pad_width) + ) + padded[original_area_slice] = array + + return padded, original_area_slice + + +def _set_pad_area(padded, axis, width_pair, value_pair): + """ + Set empty-padded area in given dimension. + + Parameters + ---------- + padded : ndarray + Array with the pad area which is modified inplace. + axis : int + Dimension with the pad area to set. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + value_pair : tuple of scalars or ndarrays + Values inserted into the pad area on each side. It must match or be + broadcastable to the shape of `arr`. + """ + left_slice = _slice_at_axis(slice(None, width_pair[0]), axis) + padded[left_slice] = value_pair[0] + + right_slice = _slice_at_axis( + slice(padded.shape[axis] - width_pair[1], None), axis) + padded[right_slice] = value_pair[1] + + +def _get_edges(padded, axis, width_pair): + """ + Retrieve edge values from empty-padded array in given dimension. + + Parameters + ---------- + padded : ndarray + Empty-padded array. + axis : int + Dimension in which the edges are considered. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + + Returns + ------- + left_edge, right_edge : ndarray + Edge values of the valid area in `padded` in the given dimension. Its + shape will always match `padded` except for the dimension given by + `axis` which will have a length of 1. + """ + left_index = width_pair[0] + left_slice = _slice_at_axis(slice(left_index, left_index + 1), axis) + left_edge = padded[left_slice] + + right_index = padded.shape[axis] - width_pair[1] + right_slice = _slice_at_axis(slice(right_index - 1, right_index), axis) + right_edge = padded[right_slice] + + return left_edge, right_edge + + +def _get_linear_ramps(padded, axis, width_pair, end_value_pair): + """ + Construct linear ramps for empty-padded array in given dimension. + + Parameters + ---------- + padded : ndarray + Empty-padded array. + axis : int + Dimension in which the ramps are constructed. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + end_value_pair : (scalar, scalar) + End values for the linear ramps which form the edge of the fully padded + array. These values are included in the linear ramps. + + Returns + ------- + left_ramp, right_ramp : ndarray + Linear ramps to set on both sides of `padded`. + """ + edge_pair = _get_edges(padded, axis, width_pair) + + left_ramp, right_ramp = ( + np.linspace( + start=end_value, + stop=edge.squeeze(axis), # Dimension is replaced by linspace + num=width, + endpoint=False, + dtype=padded.dtype, + axis=axis + ) + for end_value, edge, width in zip( + end_value_pair, edge_pair, width_pair + ) + ) + + # Reverse linear space in appropriate dimension + right_ramp = right_ramp[_slice_at_axis(slice(None, None, -1), axis)] + + return left_ramp, right_ramp + + +def _get_stats(padded, axis, width_pair, length_pair, stat_func): + """ + Calculate statistic for the empty-padded array in given dimension. + + Parameters + ---------- + padded : ndarray + Empty-padded array. + axis : int + Dimension in which the statistic is calculated. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + length_pair : 2-element sequence of None or int + Gives the number of values in valid area from each side that is + taken into account when calculating the statistic. If None the entire + valid area in `padded` is considered. + stat_func : function + Function to compute statistic. The expected signature is + ``stat_func(x: ndarray, axis: int, keepdims: bool) -> ndarray``. + + Returns + ------- + left_stat, right_stat : ndarray + Calculated statistic for both sides of `padded`. + """ + # Calculate indices of the edges of the area with original values + left_index = width_pair[0] + right_index = padded.shape[axis] - width_pair[1] + # as well as its length + max_length = right_index - left_index + + # Limit stat_lengths to max_length + left_length, right_length = length_pair + if left_length is None or max_length < left_length: + left_length = max_length + if right_length is None or max_length < right_length: + right_length = max_length + + if (left_length == 0 or right_length == 0) \ + and stat_func in {np.amax, np.amin}: + # amax and amin can't operate on an empty array, + # raise a more descriptive warning here instead of the default one + raise ValueError("stat_length of 0 yields no value for padding") + + # Calculate statistic for the left side + left_slice = _slice_at_axis( + slice(left_index, left_index + left_length), axis) + left_chunk = padded[left_slice] + left_stat = stat_func(left_chunk, axis=axis, keepdims=True) + _round_if_needed(left_stat, padded.dtype) + + if left_length == right_length == max_length: + # return early as right_stat must be identical to left_stat + return left_stat, left_stat + + # Calculate statistic for the right side + right_slice = _slice_at_axis( + slice(right_index - right_length, right_index), axis) + right_chunk = padded[right_slice] + right_stat = stat_func(right_chunk, axis=axis, keepdims=True) + _round_if_needed(right_stat, padded.dtype) + + return left_stat, right_stat + + +def _set_reflect_both(padded, axis, width_pair, method, + original_period, include_edge=False): + """ + Pad `axis` of `arr` with reflection. + + Parameters + ---------- + padded : ndarray + Input array of arbitrary shape. + axis : int + Axis along which to pad `arr`. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + method : str + Controls method of reflection; options are 'even' or 'odd'. + original_period : int + Original length of data on `axis` of `arr`. + include_edge : bool + If true, edge value is included in reflection, otherwise the edge + value forms the symmetric axis to the reflection. + + Returns + ------- + pad_amt : tuple of ints, length 2 + New index positions of padding to do along the `axis`. If these are + both 0, padding is done in this dimension. + """ + left_pad, right_pad = width_pair + old_length = padded.shape[axis] - right_pad - left_pad + + if include_edge: + # Avoid wrapping with only a subset of the original area + # by ensuring period can only be a multiple of the original + # area's length. + old_length = old_length // original_period * original_period + # Edge is included, we need to offset the pad amount by 1 + edge_offset = 1 + else: + # Avoid wrapping with only a subset of the original area + # by ensuring period can only be a multiple of the original + # area's length. + old_length = ((old_length - 1) // (original_period - 1) + * (original_period - 1) + 1) + edge_offset = 0 # Edge is not included, no need to offset pad amount + old_length -= 1 # but must be omitted from the chunk + + if left_pad > 0: + # Pad with reflected values on left side: + # First limit chunk size which can't be larger than pad area + chunk_length = min(old_length, left_pad) + # Slice right to left, stop on or next to edge, start relative to stop + stop = left_pad - edge_offset + start = stop + chunk_length + left_slice = _slice_at_axis(slice(start, stop, -1), axis) + left_chunk = padded[left_slice] + + if method == "odd": + # Negate chunk and align with edge + edge_slice = _slice_at_axis(slice(left_pad, left_pad + 1), axis) + left_chunk = 2 * padded[edge_slice] - left_chunk + + # Insert chunk into padded area + start = left_pad - chunk_length + stop = left_pad + pad_area = _slice_at_axis(slice(start, stop), axis) + padded[pad_area] = left_chunk + # Adjust pointer to left edge for next iteration + left_pad -= chunk_length + + if right_pad > 0: + # Pad with reflected values on right side: + # First limit chunk size which can't be larger than pad area + chunk_length = min(old_length, right_pad) + # Slice right to left, start on or next to edge, stop relative to start + start = -right_pad + edge_offset - 2 + stop = start - chunk_length + right_slice = _slice_at_axis(slice(start, stop, -1), axis) + right_chunk = padded[right_slice] + + if method == "odd": + # Negate chunk and align with edge + edge_slice = _slice_at_axis( + slice(-right_pad - 1, -right_pad), axis) + right_chunk = 2 * padded[edge_slice] - right_chunk + + # Insert chunk into padded area + start = padded.shape[axis] - right_pad + stop = start + chunk_length + pad_area = _slice_at_axis(slice(start, stop), axis) + padded[pad_area] = right_chunk + # Adjust pointer to right edge for next iteration + right_pad -= chunk_length + + return left_pad, right_pad + + +def _set_wrap_both(padded, axis, width_pair, original_period): + """ + Pad `axis` of `arr` with wrapped values. + + Parameters + ---------- + padded : ndarray + Input array of arbitrary shape. + axis : int + Axis along which to pad `arr`. + width_pair : (int, int) + Pair of widths that mark the pad area on both sides in the given + dimension. + original_period : int + Original length of data on `axis` of `arr`. + + Returns + ------- + pad_amt : tuple of ints, length 2 + New index positions of padding to do along the `axis`. If these are + both 0, padding is done in this dimension. + """ + left_pad, right_pad = width_pair + period = padded.shape[axis] - right_pad - left_pad + # Avoid wrapping with only a subset of the original area by ensuring period + # can only be a multiple of the original area's length. + period = period // original_period * original_period + + # If the current dimension of `arr` doesn't contain enough valid values + # (not part of the undefined pad area) we need to pad multiple times. + # Each time the pad area shrinks on both sides which is communicated with + # these variables. + new_left_pad = 0 + new_right_pad = 0 + + if left_pad > 0: + # Pad with wrapped values on left side + # First slice chunk from left side of the non-pad area. + # Use min(period, left_pad) to ensure that chunk is not larger than + # pad area. + slice_end = left_pad + period + slice_start = slice_end - min(period, left_pad) + right_slice = _slice_at_axis(slice(slice_start, slice_end), axis) + right_chunk = padded[right_slice] + + if left_pad > period: + # Chunk is smaller than pad area + pad_area = _slice_at_axis(slice(left_pad - period, left_pad), axis) + new_left_pad = left_pad - period + else: + # Chunk matches pad area + pad_area = _slice_at_axis(slice(None, left_pad), axis) + padded[pad_area] = right_chunk + + if right_pad > 0: + # Pad with wrapped values on right side + # First slice chunk from right side of the non-pad area. + # Use min(period, right_pad) to ensure that chunk is not larger than + # pad area. + slice_start = -right_pad - period + slice_end = slice_start + min(period, right_pad) + left_slice = _slice_at_axis(slice(slice_start, slice_end), axis) + left_chunk = padded[left_slice] + + if right_pad > period: + # Chunk is smaller than pad area + pad_area = _slice_at_axis( + slice(-right_pad, -right_pad + period), axis) + new_right_pad = right_pad - period + else: + # Chunk matches pad area + pad_area = _slice_at_axis(slice(-right_pad, None), axis) + padded[pad_area] = left_chunk + + return new_left_pad, new_right_pad + + +def _as_pairs(x, ndim, as_index=False): + """ + Broadcast `x` to an array with the shape (`ndim`, 2). + + A helper function for `pad` that prepares and validates arguments like + `pad_width` for iteration in pairs. + + Parameters + ---------- + x : {None, scalar, array-like} + The object to broadcast to the shape (`ndim`, 2). + ndim : int + Number of pairs the broadcasted `x` will have. + as_index : bool, optional + If `x` is not None, try to round each element of `x` to an integer + (dtype `np.intp`) and ensure every element is positive. + + Returns + ------- + pairs : nested iterables, shape (`ndim`, 2) + The broadcasted version of `x`. + + Raises + ------ + ValueError + If `as_index` is True and `x` contains negative elements. + Or if `x` is not broadcastable to the shape (`ndim`, 2). + """ + if x is None: + # Pass through None as a special case, otherwise np.round(x) fails + # with an AttributeError + return ((None, None),) * ndim + + x = np.array(x) + if as_index: + x = np.round(x).astype(np.intp, copy=False) + + if x.ndim < 3: + # Optimization: Possibly use faster paths for cases where `x` has + # only 1 or 2 elements. `np.broadcast_to` could handle these as well + # but is currently slower + + if x.size == 1: + # x was supplied as a single value + x = x.ravel() # Ensure x[0] works for x.ndim == 0, 1, 2 + if as_index and x < 0: + raise ValueError("index can't contain negative values") + return ((x[0], x[0]),) * ndim + + if x.size == 2 and x.shape != (2, 1): + # x was supplied with a single value for each side + # but except case when each dimension has a single value + # which should be broadcasted to a pair, + # e.g. [[1], [2]] -> [[1, 1], [2, 2]] not [[1, 2], [1, 2]] + x = x.ravel() # Ensure x[0], x[1] works + if as_index and (x[0] < 0 or x[1] < 0): + raise ValueError("index can't contain negative values") + return ((x[0], x[1]),) * ndim + + if as_index and x.min() < 0: + raise ValueError("index can't contain negative values") + + # Converting the array with `tolist` seems to improve performance + # when iterating and indexing the result (see usage in `pad`) + return np.broadcast_to(x, (ndim, 2)).tolist() + + +def _pad_dispatcher(array, pad_width, mode=None, **kwargs): + return (array,) + + +############################################################################### +# Public functions + + +@array_function_dispatch(_pad_dispatcher, module='numpy') +def pad(array, pad_width, mode='constant', **kwargs): + """ + Pad an array. + + Parameters + ---------- + array : array_like of rank N + The array to pad. + pad_width : {sequence, array_like, int} + Number of values padded to the edges of each axis. + ``((before_1, after_1), ... (before_N, after_N))`` unique pad widths + for each axis. + ``(before, after)`` or ``((before, after),)`` yields same before + and after pad for each axis. + ``(pad,)`` or ``int`` is a shortcut for before = after = pad width + for all axes. + mode : str or function, optional + One of the following string values or a user supplied function. + + 'constant' (default) + Pads with a constant value. + 'edge' + Pads with the edge values of array. + 'linear_ramp' + Pads with the linear ramp between end_value and the + array edge value. + 'maximum' + Pads with the maximum value of all or part of the + vector along each axis. + 'mean' + Pads with the mean value of all or part of the + vector along each axis. + 'median' + Pads with the median value of all or part of the + vector along each axis. + 'minimum' + Pads with the minimum value of all or part of the + vector along each axis. + 'reflect' + Pads with the reflection of the vector mirrored on + the first and last values of the vector along each + axis. + 'symmetric' + Pads with the reflection of the vector mirrored + along the edge of the array. + 'wrap' + Pads with the wrap of the vector along the axis. + The first values are used to pad the end and the + end values are used to pad the beginning. + 'empty' + Pads with undefined values. + + + Padding function, see Notes. + stat_length : sequence or int, optional + Used in 'maximum', 'mean', 'median', and 'minimum'. Number of + values at edge of each axis used to calculate the statistic value. + + ``((before_1, after_1), ... (before_N, after_N))`` unique statistic + lengths for each axis. + + ``(before, after)`` or ``((before, after),)`` yields same before + and after statistic lengths for each axis. + + ``(stat_length,)`` or ``int`` is a shortcut for + ``before = after = statistic`` length for all axes. + + Default is ``None``, to use the entire axis. + constant_values : sequence or scalar, optional + Used in 'constant'. The values to set the padded values for each + axis. + + ``((before_1, after_1), ... (before_N, after_N))`` unique pad constants + for each axis. + + ``(before, after)`` or ``((before, after),)`` yields same before + and after constants for each axis. + + ``(constant,)`` or ``constant`` is a shortcut for + ``before = after = constant`` for all axes. + + Default is 0. + end_values : sequence or scalar, optional + Used in 'linear_ramp'. The values used for the ending value of the + linear_ramp and that will form the edge of the padded array. + + ``((before_1, after_1), ... (before_N, after_N))`` unique end values + for each axis. + + ``(before, after)`` or ``((before, after),)`` yields same before + and after end values for each axis. + + ``(constant,)`` or ``constant`` is a shortcut for + ``before = after = constant`` for all axes. + + Default is 0. + reflect_type : {'even', 'odd'}, optional + Used in 'reflect', and 'symmetric'. The 'even' style is the + default with an unaltered reflection around the edge value. For + the 'odd' style, the extended part of the array is created by + subtracting the reflected values from two times the edge value. + + Returns + ------- + pad : ndarray + Padded array of rank equal to `array` with shape increased + according to `pad_width`. + + Notes + ----- + For an array with rank greater than 1, some of the padding of later + axes is calculated from padding of previous axes. This is easiest to + think about with a rank 2 array where the corners of the padded array + are calculated by using padded values from the first axis. + + The padding function, if used, should modify a rank 1 array in-place. It + has the following signature:: + + padding_func(vector, iaxis_pad_width, iaxis, kwargs) + + where + + vector : ndarray + A rank 1 array already padded with zeros. Padded values are + vector[:iaxis_pad_width[0]] and vector[-iaxis_pad_width[1]:]. + iaxis_pad_width : tuple + A 2-tuple of ints, iaxis_pad_width[0] represents the number of + values padded at the beginning of vector where + iaxis_pad_width[1] represents the number of values padded at + the end of vector. + iaxis : int + The axis currently being calculated. + kwargs : dict + Any keyword arguments the function requires. + + Examples + -------- + >>> import numpy as np + >>> a = [1, 2, 3, 4, 5] + >>> np.pad(a, (2, 3), 'constant', constant_values=(4, 6)) + array([4, 4, 1, ..., 6, 6, 6]) + + >>> np.pad(a, (2, 3), 'edge') + array([1, 1, 1, ..., 5, 5, 5]) + + >>> np.pad(a, (2, 3), 'linear_ramp', end_values=(5, -4)) + array([ 5, 3, 1, 2, 3, 4, 5, 2, -1, -4]) + + >>> np.pad(a, (2,), 'maximum') + array([5, 5, 1, 2, 3, 4, 5, 5, 5]) + + >>> np.pad(a, (2,), 'mean') + array([3, 3, 1, 2, 3, 4, 5, 3, 3]) + + >>> np.pad(a, (2,), 'median') + array([3, 3, 1, 2, 3, 4, 5, 3, 3]) + + >>> a = [[1, 2], [3, 4]] + >>> np.pad(a, ((3, 2), (2, 3)), 'minimum') + array([[1, 1, 1, 2, 1, 1, 1], + [1, 1, 1, 2, 1, 1, 1], + [1, 1, 1, 2, 1, 1, 1], + [1, 1, 1, 2, 1, 1, 1], + [3, 3, 3, 4, 3, 3, 3], + [1, 1, 1, 2, 1, 1, 1], + [1, 1, 1, 2, 1, 1, 1]]) + + >>> a = [1, 2, 3, 4, 5] + >>> np.pad(a, (2, 3), 'reflect') + array([3, 2, 1, 2, 3, 4, 5, 4, 3, 2]) + + >>> np.pad(a, (2, 3), 'reflect', reflect_type='odd') + array([-1, 0, 1, 2, 3, 4, 5, 6, 7, 8]) + + >>> np.pad(a, (2, 3), 'symmetric') + array([2, 1, 1, 2, 3, 4, 5, 5, 4, 3]) + + >>> np.pad(a, (2, 3), 'symmetric', reflect_type='odd') + array([0, 1, 1, 2, 3, 4, 5, 5, 6, 7]) + + >>> np.pad(a, (2, 3), 'wrap') + array([4, 5, 1, 2, 3, 4, 5, 1, 2, 3]) + + >>> def pad_with(vector, pad_width, iaxis, kwargs): + ... pad_value = kwargs.get('padder', 10) + ... vector[:pad_width[0]] = pad_value + ... vector[-pad_width[1]:] = pad_value + >>> a = np.arange(6) + >>> a = a.reshape((2, 3)) + >>> np.pad(a, 2, pad_with) + array([[10, 10, 10, 10, 10, 10, 10], + [10, 10, 10, 10, 10, 10, 10], + [10, 10, 0, 1, 2, 10, 10], + [10, 10, 3, 4, 5, 10, 10], + [10, 10, 10, 10, 10, 10, 10], + [10, 10, 10, 10, 10, 10, 10]]) + >>> np.pad(a, 2, pad_with, padder=100) + array([[100, 100, 100, 100, 100, 100, 100], + [100, 100, 100, 100, 100, 100, 100], + [100, 100, 0, 1, 2, 100, 100], + [100, 100, 3, 4, 5, 100, 100], + [100, 100, 100, 100, 100, 100, 100], + [100, 100, 100, 100, 100, 100, 100]]) + """ + array = np.asarray(array) + pad_width = np.asarray(pad_width) + + if not pad_width.dtype.kind == 'i': + raise TypeError('`pad_width` must be of integral type.') + + # Broadcast to shape (array.ndim, 2) + pad_width = _as_pairs(pad_width, array.ndim, as_index=True) + + if callable(mode): + # Old behavior: Use user-supplied function with np.apply_along_axis + function = mode + # Create a new zero padded array + padded, _ = _pad_simple(array, pad_width, fill_value=0) + # And apply along each axis + + for axis in range(padded.ndim): + # Iterate using ndindex as in apply_along_axis, but assuming that + # function operates inplace on the padded array. + + # view with the iteration axis at the end + view = np.moveaxis(padded, axis, -1) + + # compute indices for the iteration axes, and append a trailing + # ellipsis to prevent 0d arrays decaying to scalars (gh-8642) + inds = ndindex(view.shape[:-1]) + inds = (ind + (Ellipsis,) for ind in inds) + for ind in inds: + function(view[ind], pad_width[axis], axis, kwargs) + + return padded + + # Make sure that no unsupported keywords were passed for the current mode + allowed_kwargs = { + 'empty': [], 'edge': [], 'wrap': [], + 'constant': ['constant_values'], + 'linear_ramp': ['end_values'], + 'maximum': ['stat_length'], + 'mean': ['stat_length'], + 'median': ['stat_length'], + 'minimum': ['stat_length'], + 'reflect': ['reflect_type'], + 'symmetric': ['reflect_type'], + } + try: + unsupported_kwargs = set(kwargs) - set(allowed_kwargs[mode]) + except KeyError: + raise ValueError("mode '{}' is not supported".format(mode)) from None + if unsupported_kwargs: + raise ValueError("unsupported keyword arguments for mode '{}': {}" + .format(mode, unsupported_kwargs)) + + stat_functions = {"maximum": np.amax, "minimum": np.amin, + "mean": np.mean, "median": np.median} + + # Create array with final shape and original values + # (padded area is undefined) + padded, original_area_slice = _pad_simple(array, pad_width) + # And prepare iteration over all dimensions + # (zipping may be more readable than using enumerate) + axes = range(padded.ndim) + + if mode == "constant": + values = kwargs.get("constant_values", 0) + values = _as_pairs(values, padded.ndim) + for axis, width_pair, value_pair in zip(axes, pad_width, values): + roi = _view_roi(padded, original_area_slice, axis) + _set_pad_area(roi, axis, width_pair, value_pair) + + elif mode == "empty": + pass # Do nothing as _pad_simple already returned the correct result + + elif array.size == 0: + # Only modes "constant" and "empty" can extend empty axes, all other + # modes depend on `array` not being empty + # -> ensure every empty axis is only "padded with 0" + for axis, width_pair in zip(axes, pad_width): + if array.shape[axis] == 0 and any(width_pair): + raise ValueError( + "can't extend empty axis {} using modes other than " + "'constant' or 'empty'".format(axis) + ) + # passed, don't need to do anything more as _pad_simple already + # returned the correct result + + elif mode == "edge": + for axis, width_pair in zip(axes, pad_width): + roi = _view_roi(padded, original_area_slice, axis) + edge_pair = _get_edges(roi, axis, width_pair) + _set_pad_area(roi, axis, width_pair, edge_pair) + + elif mode == "linear_ramp": + end_values = kwargs.get("end_values", 0) + end_values = _as_pairs(end_values, padded.ndim) + for axis, width_pair, value_pair in zip(axes, pad_width, end_values): + roi = _view_roi(padded, original_area_slice, axis) + ramp_pair = _get_linear_ramps(roi, axis, width_pair, value_pair) + _set_pad_area(roi, axis, width_pair, ramp_pair) + + elif mode in stat_functions: + func = stat_functions[mode] + length = kwargs.get("stat_length", None) + length = _as_pairs(length, padded.ndim, as_index=True) + for axis, width_pair, length_pair in zip(axes, pad_width, length): + roi = _view_roi(padded, original_area_slice, axis) + stat_pair = _get_stats(roi, axis, width_pair, length_pair, func) + _set_pad_area(roi, axis, width_pair, stat_pair) + + elif mode in {"reflect", "symmetric"}: + method = kwargs.get("reflect_type", "even") + include_edge = mode == "symmetric" + for axis, (left_index, right_index) in zip(axes, pad_width): + if array.shape[axis] == 1 and (left_index > 0 or right_index > 0): + # Extending singleton dimension for 'reflect' is legacy + # behavior; it really should raise an error. + edge_pair = _get_edges(padded, axis, (left_index, right_index)) + _set_pad_area( + padded, axis, (left_index, right_index), edge_pair) + continue + + roi = _view_roi(padded, original_area_slice, axis) + while left_index > 0 or right_index > 0: + # Iteratively pad until dimension is filled with reflected + # values. This is necessary if the pad area is larger than + # the length of the original values in the current dimension. + left_index, right_index = _set_reflect_both( + roi, axis, (left_index, right_index), + method, array.shape[axis], include_edge + ) + + elif mode == "wrap": + for axis, (left_index, right_index) in zip(axes, pad_width): + roi = _view_roi(padded, original_area_slice, axis) + original_period = padded.shape[axis] - right_index - left_index + while left_index > 0 or right_index > 0: + # Iteratively pad until dimension is filled with wrapped + # values. This is necessary if the pad area is larger than + # the length of the original values in the current dimension. + left_index, right_index = _set_wrap_both( + roi, axis, (left_index, right_index), original_period) + + return padded diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3a2c433c338a2eadd2afd83182bb9738ab37fffb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arraypad_impl.pyi @@ -0,0 +1,89 @@ +from typing import ( + Literal as L, + Any, + TypeAlias, + overload, + TypeVar, + Protocol, + type_check_only, +) + +from numpy import generic + +from numpy._typing import ( + ArrayLike, + NDArray, + _ArrayLikeInt, + _ArrayLike, +) + +__all__ = ["pad"] + +_SCT = TypeVar("_SCT", bound=generic) + +@type_check_only +class _ModeFunc(Protocol): + def __call__( + self, + vector: NDArray[Any], + iaxis_pad_width: tuple[int, int], + iaxis: int, + kwargs: dict[str, Any], + /, + ) -> None: ... + +_ModeKind: TypeAlias = L[ + "constant", + "edge", + "linear_ramp", + "maximum", + "mean", + "median", + "minimum", + "reflect", + "symmetric", + "wrap", + "empty", +] + + +# TODO: In practice each keyword argument is exclusive to one or more +# specific modes. Consider adding more overloads to express this in the future. + +# Expand `**kwargs` into explicit keyword-only arguments +@overload +def pad( + array: _ArrayLike[_SCT], + pad_width: _ArrayLikeInt, + mode: _ModeKind = ..., + *, + stat_length: None | _ArrayLikeInt = ..., + constant_values: ArrayLike = ..., + end_values: ArrayLike = ..., + reflect_type: L["odd", "even"] = ..., +) -> NDArray[_SCT]: ... +@overload +def pad( + array: ArrayLike, + pad_width: _ArrayLikeInt, + mode: _ModeKind = ..., + *, + stat_length: None | _ArrayLikeInt = ..., + constant_values: ArrayLike = ..., + end_values: ArrayLike = ..., + reflect_type: L["odd", "even"] = ..., +) -> NDArray[Any]: ... +@overload +def pad( + array: _ArrayLike[_SCT], + pad_width: _ArrayLikeInt, + mode: _ModeFunc, + **kwargs: Any, +) -> NDArray[_SCT]: ... +@overload +def pad( + array: ArrayLike, + pad_width: _ArrayLikeInt, + mode: _ModeFunc, + **kwargs: Any, +) -> NDArray[Any]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..60b3425682fb2e1de1fb4e7a7c99132b80608986 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.py @@ -0,0 +1,1215 @@ +""" +Set operations for arrays based on sorting. + +Notes +----- + +For floating point arrays, inaccurate results may appear due to usual round-off +and floating point comparison issues. + +Speed could be gained in some operations by an implementation of +`numpy.sort`, that can provide directly the permutation vectors, thus avoiding +calls to `numpy.argsort`. + +Original author: Robert Cimrman + +""" +import functools +import warnings +from typing import NamedTuple + +import numpy as np +from numpy._core import overrides +from numpy._core._multiarray_umath import _array_converter + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +__all__ = [ + "ediff1d", "in1d", "intersect1d", "isin", "setdiff1d", "setxor1d", + "union1d", "unique", "unique_all", "unique_counts", "unique_inverse", + "unique_values" +] + + +def _ediff1d_dispatcher(ary, to_end=None, to_begin=None): + return (ary, to_end, to_begin) + + +@array_function_dispatch(_ediff1d_dispatcher) +def ediff1d(ary, to_end=None, to_begin=None): + """ + The differences between consecutive elements of an array. + + Parameters + ---------- + ary : array_like + If necessary, will be flattened before the differences are taken. + to_end : array_like, optional + Number(s) to append at the end of the returned differences. + to_begin : array_like, optional + Number(s) to prepend at the beginning of the returned differences. + + Returns + ------- + ediff1d : ndarray + The differences. Loosely, this is ``ary.flat[1:] - ary.flat[:-1]``. + + See Also + -------- + diff, gradient + + Notes + ----- + When applied to masked arrays, this function drops the mask information + if the `to_begin` and/or `to_end` parameters are used. + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2, 4, 7, 0]) + >>> np.ediff1d(x) + array([ 1, 2, 3, -7]) + + >>> np.ediff1d(x, to_begin=-99, to_end=np.array([88, 99])) + array([-99, 1, 2, ..., -7, 88, 99]) + + The returned array is always 1D. + + >>> y = [[1, 2, 4], [1, 6, 24]] + >>> np.ediff1d(y) + array([ 1, 2, -3, 5, 18]) + + """ + conv = _array_converter(ary) + # Convert to (any) array and ravel: + ary = conv[0].ravel() + + # enforce that the dtype of `ary` is used for the output + dtype_req = ary.dtype + + # fast track default case + if to_begin is None and to_end is None: + return ary[1:] - ary[:-1] + + if to_begin is None: + l_begin = 0 + else: + to_begin = np.asanyarray(to_begin) + if not np.can_cast(to_begin, dtype_req, casting="same_kind"): + raise TypeError("dtype of `to_begin` must be compatible " + "with input `ary` under the `same_kind` rule.") + + to_begin = to_begin.ravel() + l_begin = len(to_begin) + + if to_end is None: + l_end = 0 + else: + to_end = np.asanyarray(to_end) + if not np.can_cast(to_end, dtype_req, casting="same_kind"): + raise TypeError("dtype of `to_end` must be compatible " + "with input `ary` under the `same_kind` rule.") + + to_end = to_end.ravel() + l_end = len(to_end) + + # do the calculation in place and copy to_begin and to_end + l_diff = max(len(ary) - 1, 0) + result = np.empty_like(ary, shape=l_diff + l_begin + l_end) + + if l_begin > 0: + result[:l_begin] = to_begin + if l_end > 0: + result[l_begin + l_diff:] = to_end + np.subtract(ary[1:], ary[:-1], result[l_begin:l_begin + l_diff]) + + return conv.wrap(result) + + +def _unpack_tuple(x): + """ Unpacks one-element tuples for use as return values """ + if len(x) == 1: + return x[0] + else: + return x + + +def _unique_dispatcher(ar, return_index=None, return_inverse=None, + return_counts=None, axis=None, *, equal_nan=None): + return (ar,) + + +@array_function_dispatch(_unique_dispatcher) +def unique(ar, return_index=False, return_inverse=False, + return_counts=False, axis=None, *, equal_nan=True): + """ + Find the unique elements of an array. + + Returns the sorted unique elements of an array. There are three optional + outputs in addition to the unique elements: + + * the indices of the input array that give the unique values + * the indices of the unique array that reconstruct the input array + * the number of times each unique value comes up in the input array + + Parameters + ---------- + ar : array_like + Input array. Unless `axis` is specified, this will be flattened if it + is not already 1-D. + return_index : bool, optional + If True, also return the indices of `ar` (along the specified axis, + if provided, or in the flattened array) that result in the unique array. + return_inverse : bool, optional + If True, also return the indices of the unique array (for the specified + axis, if provided) that can be used to reconstruct `ar`. + return_counts : bool, optional + If True, also return the number of times each unique item appears + in `ar`. + axis : int or None, optional + The axis to operate on. If None, `ar` will be flattened. If an integer, + the subarrays indexed by the given axis will be flattened and treated + as the elements of a 1-D array with the dimension of the given axis, + see the notes for more details. Object arrays or structured arrays + that contain objects are not supported if the `axis` kwarg is used. The + default is None. + + equal_nan : bool, optional + If True, collapses multiple NaN values in the return array into one. + + .. versionadded:: 1.24 + + Returns + ------- + unique : ndarray + The sorted unique values. + unique_indices : ndarray, optional + The indices of the first occurrences of the unique values in the + original array. Only provided if `return_index` is True. + unique_inverse : ndarray, optional + The indices to reconstruct the original array from the + unique array. Only provided if `return_inverse` is True. + unique_counts : ndarray, optional + The number of times each of the unique values comes up in the + original array. Only provided if `return_counts` is True. + + See Also + -------- + repeat : Repeat elements of an array. + sort : Return a sorted copy of an array. + + Notes + ----- + When an axis is specified the subarrays indexed by the axis are sorted. + This is done by making the specified axis the first dimension of the array + (move the axis to the first dimension to keep the order of the other axes) + and then flattening the subarrays in C order. The flattened subarrays are + then viewed as a structured type with each element given a label, with the + effect that we end up with a 1-D array of structured types that can be + treated in the same way as any other 1-D array. The result is that the + flattened subarrays are sorted in lexicographic order starting with the + first element. + + .. versionchanged:: 1.21 + Like np.sort, NaN will sort to the end of the values. + For complex arrays all NaN values are considered equivalent + (no matter whether the NaN is in the real or imaginary part). + As the representant for the returned array the smallest one in the + lexicographical order is chosen - see np.sort for how the lexicographical + order is defined for complex arrays. + + .. versionchanged:: 2.0 + For multi-dimensional inputs, ``unique_inverse`` is reshaped + such that the input can be reconstructed using + ``np.take(unique, unique_inverse, axis=axis)``. The result is + now not 1-dimensional when ``axis=None``. + + Note that in NumPy 2.0.0 a higher dimensional array was returned also + when ``axis`` was not ``None``. This was reverted, but + ``inverse.reshape(-1)`` can be used to ensure compatibility with both + versions. + + Examples + -------- + >>> import numpy as np + >>> np.unique([1, 1, 2, 2, 3, 3]) + array([1, 2, 3]) + >>> a = np.array([[1, 1], [2, 3]]) + >>> np.unique(a) + array([1, 2, 3]) + + Return the unique rows of a 2D array + + >>> a = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]]) + >>> np.unique(a, axis=0) + array([[1, 0, 0], [2, 3, 4]]) + + Return the indices of the original array that give the unique values: + + >>> a = np.array(['a', 'b', 'b', 'c', 'a']) + >>> u, indices = np.unique(a, return_index=True) + >>> u + array(['a', 'b', 'c'], dtype='>> indices + array([0, 1, 3]) + >>> a[indices] + array(['a', 'b', 'c'], dtype='>> a = np.array([1, 2, 6, 4, 2, 3, 2]) + >>> u, indices = np.unique(a, return_inverse=True) + >>> u + array([1, 2, 3, 4, 6]) + >>> indices + array([0, 1, 4, 3, 1, 2, 1]) + >>> u[indices] + array([1, 2, 6, 4, 2, 3, 2]) + + Reconstruct the input values from the unique values and counts: + + >>> a = np.array([1, 2, 6, 4, 2, 3, 2]) + >>> values, counts = np.unique(a, return_counts=True) + >>> values + array([1, 2, 3, 4, 6]) + >>> counts + array([1, 3, 1, 1, 1]) + >>> np.repeat(values, counts) + array([1, 2, 2, 2, 3, 4, 6]) # original order not preserved + + """ + ar = np.asanyarray(ar) + if axis is None: + ret = _unique1d(ar, return_index, return_inverse, return_counts, + equal_nan=equal_nan, inverse_shape=ar.shape, axis=None) + return _unpack_tuple(ret) + + # axis was specified and not None + try: + ar = np.moveaxis(ar, axis, 0) + except np.exceptions.AxisError: + # this removes the "axis1" or "axis2" prefix from the error message + raise np.exceptions.AxisError(axis, ar.ndim) from None + inverse_shape = [1] * ar.ndim + inverse_shape[axis] = ar.shape[0] + + # Must reshape to a contiguous 2D array for this to work... + orig_shape, orig_dtype = ar.shape, ar.dtype + ar = ar.reshape(orig_shape[0], np.prod(orig_shape[1:], dtype=np.intp)) + ar = np.ascontiguousarray(ar) + dtype = [('f{i}'.format(i=i), ar.dtype) for i in range(ar.shape[1])] + + # At this point, `ar` has shape `(n, m)`, and `dtype` is a structured + # data type with `m` fields where each field has the data type of `ar`. + # In the following, we create the array `consolidated`, which has + # shape `(n,)` with data type `dtype`. + try: + if ar.shape[1] > 0: + consolidated = ar.view(dtype) + else: + # If ar.shape[1] == 0, then dtype will be `np.dtype([])`, which is + # a data type with itemsize 0, and the call `ar.view(dtype)` will + # fail. Instead, we'll use `np.empty` to explicitly create the + # array with shape `(len(ar),)`. Because `dtype` in this case has + # itemsize 0, the total size of the result is still 0 bytes. + consolidated = np.empty(len(ar), dtype=dtype) + except TypeError as e: + # There's no good way to do this for object arrays, etc... + msg = 'The axis argument to unique is not supported for dtype {dt}' + raise TypeError(msg.format(dt=ar.dtype)) from e + + def reshape_uniq(uniq): + n = len(uniq) + uniq = uniq.view(orig_dtype) + uniq = uniq.reshape(n, *orig_shape[1:]) + uniq = np.moveaxis(uniq, 0, axis) + return uniq + + output = _unique1d(consolidated, return_index, + return_inverse, return_counts, + equal_nan=equal_nan, inverse_shape=inverse_shape, + axis=axis) + output = (reshape_uniq(output[0]),) + output[1:] + return _unpack_tuple(output) + + +def _unique1d(ar, return_index=False, return_inverse=False, + return_counts=False, *, equal_nan=True, inverse_shape=None, + axis=None): + """ + Find the unique elements of an array, ignoring shape. + """ + ar = np.asanyarray(ar).flatten() + + optional_indices = return_index or return_inverse + + if optional_indices: + perm = ar.argsort(kind='mergesort' if return_index else 'quicksort') + aux = ar[perm] + else: + ar.sort() + aux = ar + mask = np.empty(aux.shape, dtype=np.bool) + mask[:1] = True + if (equal_nan and aux.shape[0] > 0 and aux.dtype.kind in "cfmM" and + np.isnan(aux[-1])): + if aux.dtype.kind == "c": # for complex all NaNs are considered equivalent + aux_firstnan = np.searchsorted(np.isnan(aux), True, side='left') + else: + aux_firstnan = np.searchsorted(aux, aux[-1], side='left') + if aux_firstnan > 0: + mask[1:aux_firstnan] = ( + aux[1:aux_firstnan] != aux[:aux_firstnan - 1]) + mask[aux_firstnan] = True + mask[aux_firstnan + 1:] = False + else: + mask[1:] = aux[1:] != aux[:-1] + + ret = (aux[mask],) + if return_index: + ret += (perm[mask],) + if return_inverse: + imask = np.cumsum(mask) - 1 + inv_idx = np.empty(mask.shape, dtype=np.intp) + inv_idx[perm] = imask + ret += (inv_idx.reshape(inverse_shape) if axis is None else inv_idx,) + if return_counts: + idx = np.concatenate(np.nonzero(mask) + ([mask.size],)) + ret += (np.diff(idx),) + return ret + + +# Array API set functions + +class UniqueAllResult(NamedTuple): + values: np.ndarray + indices: np.ndarray + inverse_indices: np.ndarray + counts: np.ndarray + + +class UniqueCountsResult(NamedTuple): + values: np.ndarray + counts: np.ndarray + + +class UniqueInverseResult(NamedTuple): + values: np.ndarray + inverse_indices: np.ndarray + + +def _unique_all_dispatcher(x, /): + return (x,) + + +@array_function_dispatch(_unique_all_dispatcher) +def unique_all(x): + """ + Find the unique elements of an array, and counts, inverse, and indices. + + This function is an Array API compatible alternative to:: + + np.unique(x, return_index=True, return_inverse=True, + return_counts=True, equal_nan=False) + + but returns a namedtuple for easier access to each output. + + Parameters + ---------- + x : array_like + Input array. It will be flattened if it is not already 1-D. + + Returns + ------- + out : namedtuple + The result containing: + + * values - The unique elements of an input array. + * indices - The first occurring indices for each unique element. + * inverse_indices - The indices from the set of unique elements + that reconstruct `x`. + * counts - The corresponding counts for each unique element. + + See Also + -------- + unique : Find the unique elements of an array. + + Examples + -------- + >>> import numpy as np + >>> x = [1, 1, 2] + >>> uniq = np.unique_all(x) + >>> uniq.values + array([1, 2]) + >>> uniq.indices + array([0, 2]) + >>> uniq.inverse_indices + array([0, 0, 1]) + >>> uniq.counts + array([2, 1]) + """ + result = unique( + x, + return_index=True, + return_inverse=True, + return_counts=True, + equal_nan=False + ) + return UniqueAllResult(*result) + + +def _unique_counts_dispatcher(x, /): + return (x,) + + +@array_function_dispatch(_unique_counts_dispatcher) +def unique_counts(x): + """ + Find the unique elements and counts of an input array `x`. + + This function is an Array API compatible alternative to:: + + np.unique(x, return_counts=True, equal_nan=False) + + but returns a namedtuple for easier access to each output. + + Parameters + ---------- + x : array_like + Input array. It will be flattened if it is not already 1-D. + + Returns + ------- + out : namedtuple + The result containing: + + * values - The unique elements of an input array. + * counts - The corresponding counts for each unique element. + + See Also + -------- + unique : Find the unique elements of an array. + + Examples + -------- + >>> import numpy as np + >>> x = [1, 1, 2] + >>> uniq = np.unique_counts(x) + >>> uniq.values + array([1, 2]) + >>> uniq.counts + array([2, 1]) + """ + result = unique( + x, + return_index=False, + return_inverse=False, + return_counts=True, + equal_nan=False + ) + return UniqueCountsResult(*result) + + +def _unique_inverse_dispatcher(x, /): + return (x,) + + +@array_function_dispatch(_unique_inverse_dispatcher) +def unique_inverse(x): + """ + Find the unique elements of `x` and indices to reconstruct `x`. + + This function is an Array API compatible alternative to:: + + np.unique(x, return_inverse=True, equal_nan=False) + + but returns a namedtuple for easier access to each output. + + Parameters + ---------- + x : array_like + Input array. It will be flattened if it is not already 1-D. + + Returns + ------- + out : namedtuple + The result containing: + + * values - The unique elements of an input array. + * inverse_indices - The indices from the set of unique elements + that reconstruct `x`. + + See Also + -------- + unique : Find the unique elements of an array. + + Examples + -------- + >>> import numpy as np + >>> x = [1, 1, 2] + >>> uniq = np.unique_inverse(x) + >>> uniq.values + array([1, 2]) + >>> uniq.inverse_indices + array([0, 0, 1]) + """ + result = unique( + x, + return_index=False, + return_inverse=True, + return_counts=False, + equal_nan=False + ) + return UniqueInverseResult(*result) + + +def _unique_values_dispatcher(x, /): + return (x,) + + +@array_function_dispatch(_unique_values_dispatcher) +def unique_values(x): + """ + Returns the unique elements of an input array `x`. + + This function is an Array API compatible alternative to:: + + np.unique(x, equal_nan=False) + + Parameters + ---------- + x : array_like + Input array. It will be flattened if it is not already 1-D. + + Returns + ------- + out : ndarray + The unique elements of an input array. + + See Also + -------- + unique : Find the unique elements of an array. + + Examples + -------- + >>> import numpy as np + >>> np.unique_values([1, 1, 2]) + array([1, 2]) + + """ + return unique( + x, + return_index=False, + return_inverse=False, + return_counts=False, + equal_nan=False + ) + + +def _intersect1d_dispatcher( + ar1, ar2, assume_unique=None, return_indices=None): + return (ar1, ar2) + + +@array_function_dispatch(_intersect1d_dispatcher) +def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): + """ + Find the intersection of two arrays. + + Return the sorted, unique values that are in both of the input arrays. + + Parameters + ---------- + ar1, ar2 : array_like + Input arrays. Will be flattened if not already 1D. + assume_unique : bool + If True, the input arrays are both assumed to be unique, which + can speed up the calculation. If True but ``ar1`` or ``ar2`` are not + unique, incorrect results and out-of-bounds indices could result. + Default is False. + return_indices : bool + If True, the indices which correspond to the intersection of the two + arrays are returned. The first instance of a value is used if there are + multiple. Default is False. + + Returns + ------- + intersect1d : ndarray + Sorted 1D array of common and unique elements. + comm1 : ndarray + The indices of the first occurrences of the common values in `ar1`. + Only provided if `return_indices` is True. + comm2 : ndarray + The indices of the first occurrences of the common values in `ar2`. + Only provided if `return_indices` is True. + + Examples + -------- + >>> import numpy as np + >>> np.intersect1d([1, 3, 4, 3], [3, 1, 2, 1]) + array([1, 3]) + + To intersect more than two arrays, use functools.reduce: + + >>> from functools import reduce + >>> reduce(np.intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) + array([3]) + + To return the indices of the values common to the input arrays + along with the intersected values: + + >>> x = np.array([1, 1, 2, 3, 4]) + >>> y = np.array([2, 1, 4, 6]) + >>> xy, x_ind, y_ind = np.intersect1d(x, y, return_indices=True) + >>> x_ind, y_ind + (array([0, 2, 4]), array([1, 0, 2])) + >>> xy, x[x_ind], y[y_ind] + (array([1, 2, 4]), array([1, 2, 4]), array([1, 2, 4])) + + """ + ar1 = np.asanyarray(ar1) + ar2 = np.asanyarray(ar2) + + if not assume_unique: + if return_indices: + ar1, ind1 = unique(ar1, return_index=True) + ar2, ind2 = unique(ar2, return_index=True) + else: + ar1 = unique(ar1) + ar2 = unique(ar2) + else: + ar1 = ar1.ravel() + ar2 = ar2.ravel() + + aux = np.concatenate((ar1, ar2)) + if return_indices: + aux_sort_indices = np.argsort(aux, kind='mergesort') + aux = aux[aux_sort_indices] + else: + aux.sort() + + mask = aux[1:] == aux[:-1] + int1d = aux[:-1][mask] + + if return_indices: + ar1_indices = aux_sort_indices[:-1][mask] + ar2_indices = aux_sort_indices[1:][mask] - ar1.size + if not assume_unique: + ar1_indices = ind1[ar1_indices] + ar2_indices = ind2[ar2_indices] + + return int1d, ar1_indices, ar2_indices + else: + return int1d + + +def _setxor1d_dispatcher(ar1, ar2, assume_unique=None): + return (ar1, ar2) + + +@array_function_dispatch(_setxor1d_dispatcher) +def setxor1d(ar1, ar2, assume_unique=False): + """ + Find the set exclusive-or of two arrays. + + Return the sorted, unique values that are in only one (not both) of the + input arrays. + + Parameters + ---------- + ar1, ar2 : array_like + Input arrays. + assume_unique : bool + If True, the input arrays are both assumed to be unique, which + can speed up the calculation. Default is False. + + Returns + ------- + setxor1d : ndarray + Sorted 1D array of unique values that are in only one of the input + arrays. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1, 2, 3, 2, 4]) + >>> b = np.array([2, 3, 5, 7, 5]) + >>> np.setxor1d(a,b) + array([1, 4, 5, 7]) + + """ + if not assume_unique: + ar1 = unique(ar1) + ar2 = unique(ar2) + + aux = np.concatenate((ar1, ar2), axis=None) + if aux.size == 0: + return aux + + aux.sort() + flag = np.concatenate(([True], aux[1:] != aux[:-1], [True])) + return aux[flag[1:] & flag[:-1]] + + +def _in1d_dispatcher(ar1, ar2, assume_unique=None, invert=None, *, + kind=None): + return (ar1, ar2) + + +@array_function_dispatch(_in1d_dispatcher) +def in1d(ar1, ar2, assume_unique=False, invert=False, *, kind=None): + """ + Test whether each element of a 1-D array is also present in a second array. + + .. deprecated:: 2.0 + Use :func:`isin` instead of `in1d` for new code. + + Returns a boolean array the same length as `ar1` that is True + where an element of `ar1` is in `ar2` and False otherwise. + + Parameters + ---------- + ar1 : (M,) array_like + Input array. + ar2 : array_like + The values against which to test each value of `ar1`. + assume_unique : bool, optional + If True, the input arrays are both assumed to be unique, which + can speed up the calculation. Default is False. + invert : bool, optional + If True, the values in the returned array are inverted (that is, + False where an element of `ar1` is in `ar2` and True otherwise). + Default is False. ``np.in1d(a, b, invert=True)`` is equivalent + to (but is faster than) ``np.invert(in1d(a, b))``. + kind : {None, 'sort', 'table'}, optional + The algorithm to use. This will not affect the final result, + but will affect the speed and memory use. The default, None, + will select automatically based on memory considerations. + + * If 'sort', will use a mergesort-based approach. This will have + a memory usage of roughly 6 times the sum of the sizes of + `ar1` and `ar2`, not accounting for size of dtypes. + * If 'table', will use a lookup table approach similar + to a counting sort. This is only available for boolean and + integer arrays. This will have a memory usage of the + size of `ar1` plus the max-min value of `ar2`. `assume_unique` + has no effect when the 'table' option is used. + * If None, will automatically choose 'table' if + the required memory allocation is less than or equal to + 6 times the sum of the sizes of `ar1` and `ar2`, + otherwise will use 'sort'. This is done to not use + a large amount of memory by default, even though + 'table' may be faster in most cases. If 'table' is chosen, + `assume_unique` will have no effect. + + Returns + ------- + in1d : (M,) ndarray, bool + The values `ar1[in1d]` are in `ar2`. + + See Also + -------- + isin : Version of this function that preserves the + shape of ar1. + + Notes + ----- + `in1d` can be considered as an element-wise function version of the + python keyword `in`, for 1-D sequences. ``in1d(a, b)`` is roughly + equivalent to ``np.array([item in b for item in a])``. + However, this idea fails if `ar2` is a set, or similar (non-sequence) + container: As ``ar2`` is converted to an array, in those cases + ``asarray(ar2)`` is an object array rather than the expected array of + contained values. + + Using ``kind='table'`` tends to be faster than `kind='sort'` if the + following relationship is true: + ``log10(len(ar2)) > (log10(max(ar2)-min(ar2)) - 2.27) / 0.927``, + but may use greater memory. The default value for `kind` will + be automatically selected based only on memory usage, so one may + manually set ``kind='table'`` if memory constraints can be relaxed. + + Examples + -------- + >>> import numpy as np + >>> test = np.array([0, 1, 2, 5, 0]) + >>> states = [0, 2] + >>> mask = np.in1d(test, states) + >>> mask + array([ True, False, True, False, True]) + >>> test[mask] + array([0, 2, 0]) + >>> mask = np.in1d(test, states, invert=True) + >>> mask + array([False, True, False, True, False]) + >>> test[mask] + array([1, 5]) + """ + + # Deprecated in NumPy 2.0, 2023-08-18 + warnings.warn( + "`in1d` is deprecated. Use `np.isin` instead.", + DeprecationWarning, + stacklevel=2 + ) + + return _in1d(ar1, ar2, assume_unique, invert, kind=kind) + + +def _in1d(ar1, ar2, assume_unique=False, invert=False, *, kind=None): + # Ravel both arrays, behavior for the first array could be different + ar1 = np.asarray(ar1).ravel() + ar2 = np.asarray(ar2).ravel() + + # Ensure that iteration through object arrays yields size-1 arrays + if ar2.dtype == object: + ar2 = ar2.reshape(-1, 1) + + if kind not in {None, 'sort', 'table'}: + raise ValueError( + f"Invalid kind: '{kind}'. Please use None, 'sort' or 'table'.") + + # Can use the table method if all arrays are integers or boolean: + is_int_arrays = all(ar.dtype.kind in ("u", "i", "b") for ar in (ar1, ar2)) + use_table_method = is_int_arrays and kind in {None, 'table'} + + if use_table_method: + if ar2.size == 0: + if invert: + return np.ones_like(ar1, dtype=bool) + else: + return np.zeros_like(ar1, dtype=bool) + + # Convert booleans to uint8 so we can use the fast integer algorithm + if ar1.dtype == bool: + ar1 = ar1.astype(np.uint8) + if ar2.dtype == bool: + ar2 = ar2.astype(np.uint8) + + ar2_min = int(np.min(ar2)) + ar2_max = int(np.max(ar2)) + + ar2_range = ar2_max - ar2_min + + # Constraints on whether we can actually use the table method: + # 1. Assert memory usage is not too large + below_memory_constraint = ar2_range <= 6 * (ar1.size + ar2.size) + # 2. Check overflows for (ar2 - ar2_min); dtype=ar2.dtype + range_safe_from_overflow = ar2_range <= np.iinfo(ar2.dtype).max + + # Optimal performance is for approximately + # log10(size) > (log10(range) - 2.27) / 0.927. + # However, here we set the requirement that by default + # the intermediate array can only be 6x + # the combined memory allocation of the original + # arrays. See discussion on + # https://github.com/numpy/numpy/pull/12065. + + if ( + range_safe_from_overflow and + (below_memory_constraint or kind == 'table') + ): + + if invert: + outgoing_array = np.ones_like(ar1, dtype=bool) + else: + outgoing_array = np.zeros_like(ar1, dtype=bool) + + # Make elements 1 where the integer exists in ar2 + if invert: + isin_helper_ar = np.ones(ar2_range + 1, dtype=bool) + isin_helper_ar[ar2 - ar2_min] = 0 + else: + isin_helper_ar = np.zeros(ar2_range + 1, dtype=bool) + isin_helper_ar[ar2 - ar2_min] = 1 + + # Mask out elements we know won't work + basic_mask = (ar1 <= ar2_max) & (ar1 >= ar2_min) + in_range_ar1 = ar1[basic_mask] + if in_range_ar1.size == 0: + # Nothing more to do, since all values are out of range. + return outgoing_array + + # Unfortunately, ar2_min can be out of range for `intp` even + # if the calculation result must fit in range (and be positive). + # In that case, use ar2.dtype which must work for all unmasked + # values. + try: + ar2_min = np.array(ar2_min, dtype=np.intp) + dtype = np.intp + except OverflowError: + dtype = ar2.dtype + + out = np.empty_like(in_range_ar1, dtype=np.intp) + outgoing_array[basic_mask] = isin_helper_ar[ + np.subtract(in_range_ar1, ar2_min, dtype=dtype, + out=out, casting="unsafe")] + + return outgoing_array + elif kind == 'table': # not range_safe_from_overflow + raise RuntimeError( + "You have specified kind='table', " + "but the range of values in `ar2` or `ar1` exceed the " + "maximum integer of the datatype. " + "Please set `kind` to None or 'sort'." + ) + elif kind == 'table': + raise ValueError( + "The 'table' method is only " + "supported for boolean or integer arrays. " + "Please select 'sort' or None for kind." + ) + + + # Check if one of the arrays may contain arbitrary objects + contains_object = ar1.dtype.hasobject or ar2.dtype.hasobject + + # This code is run when + # a) the first condition is true, making the code significantly faster + # b) the second condition is true (i.e. `ar1` or `ar2` may contain + # arbitrary objects), since then sorting is not guaranteed to work + if len(ar2) < 10 * len(ar1) ** 0.145 or contains_object: + if invert: + mask = np.ones(len(ar1), dtype=bool) + for a in ar2: + mask &= (ar1 != a) + else: + mask = np.zeros(len(ar1), dtype=bool) + for a in ar2: + mask |= (ar1 == a) + return mask + + # Otherwise use sorting + if not assume_unique: + ar1, rev_idx = np.unique(ar1, return_inverse=True) + ar2 = np.unique(ar2) + + ar = np.concatenate((ar1, ar2)) + # We need this to be a stable sort, so always use 'mergesort' + # here. The values from the first array should always come before + # the values from the second array. + order = ar.argsort(kind='mergesort') + sar = ar[order] + if invert: + bool_ar = (sar[1:] != sar[:-1]) + else: + bool_ar = (sar[1:] == sar[:-1]) + flag = np.concatenate((bool_ar, [invert])) + ret = np.empty(ar.shape, dtype=bool) + ret[order] = flag + + if assume_unique: + return ret[:len(ar1)] + else: + return ret[rev_idx] + + +def _isin_dispatcher(element, test_elements, assume_unique=None, invert=None, + *, kind=None): + return (element, test_elements) + + +@array_function_dispatch(_isin_dispatcher) +def isin(element, test_elements, assume_unique=False, invert=False, *, + kind=None): + """ + Calculates ``element in test_elements``, broadcasting over `element` only. + Returns a boolean array of the same shape as `element` that is True + where an element of `element` is in `test_elements` and False otherwise. + + Parameters + ---------- + element : array_like + Input array. + test_elements : array_like + The values against which to test each value of `element`. + This argument is flattened if it is an array or array_like. + See notes for behavior with non-array-like parameters. + assume_unique : bool, optional + If True, the input arrays are both assumed to be unique, which + can speed up the calculation. Default is False. + invert : bool, optional + If True, the values in the returned array are inverted, as if + calculating `element not in test_elements`. Default is False. + ``np.isin(a, b, invert=True)`` is equivalent to (but faster + than) ``np.invert(np.isin(a, b))``. + kind : {None, 'sort', 'table'}, optional + The algorithm to use. This will not affect the final result, + but will affect the speed and memory use. The default, None, + will select automatically based on memory considerations. + + * If 'sort', will use a mergesort-based approach. This will have + a memory usage of roughly 6 times the sum of the sizes of + `element` and `test_elements`, not accounting for size of dtypes. + * If 'table', will use a lookup table approach similar + to a counting sort. This is only available for boolean and + integer arrays. This will have a memory usage of the + size of `element` plus the max-min value of `test_elements`. + `assume_unique` has no effect when the 'table' option is used. + * If None, will automatically choose 'table' if + the required memory allocation is less than or equal to + 6 times the sum of the sizes of `element` and `test_elements`, + otherwise will use 'sort'. This is done to not use + a large amount of memory by default, even though + 'table' may be faster in most cases. If 'table' is chosen, + `assume_unique` will have no effect. + + + Returns + ------- + isin : ndarray, bool + Has the same shape as `element`. The values `element[isin]` + are in `test_elements`. + + Notes + ----- + `isin` is an element-wise function version of the python keyword `in`. + ``isin(a, b)`` is roughly equivalent to + ``np.array([item in b for item in a])`` if `a` and `b` are 1-D sequences. + + `element` and `test_elements` are converted to arrays if they are not + already. If `test_elements` is a set (or other non-sequence collection) + it will be converted to an object array with one element, rather than an + array of the values contained in `test_elements`. This is a consequence + of the `array` constructor's way of handling non-sequence collections. + Converting the set to a list usually gives the desired behavior. + + Using ``kind='table'`` tends to be faster than `kind='sort'` if the + following relationship is true: + ``log10(len(test_elements)) > + (log10(max(test_elements)-min(test_elements)) - 2.27) / 0.927``, + but may use greater memory. The default value for `kind` will + be automatically selected based only on memory usage, so one may + manually set ``kind='table'`` if memory constraints can be relaxed. + + Examples + -------- + >>> import numpy as np + >>> element = 2*np.arange(4).reshape((2, 2)) + >>> element + array([[0, 2], + [4, 6]]) + >>> test_elements = [1, 2, 4, 8] + >>> mask = np.isin(element, test_elements) + >>> mask + array([[False, True], + [ True, False]]) + >>> element[mask] + array([2, 4]) + + The indices of the matched values can be obtained with `nonzero`: + + >>> np.nonzero(mask) + (array([0, 1]), array([1, 0])) + + The test can also be inverted: + + >>> mask = np.isin(element, test_elements, invert=True) + >>> mask + array([[ True, False], + [False, True]]) + >>> element[mask] + array([0, 6]) + + Because of how `array` handles sets, the following does not + work as expected: + + >>> test_set = {1, 2, 4, 8} + >>> np.isin(element, test_set) + array([[False, False], + [False, False]]) + + Casting the set to a list gives the expected result: + + >>> np.isin(element, list(test_set)) + array([[False, True], + [ True, False]]) + """ + element = np.asarray(element) + return _in1d(element, test_elements, assume_unique=assume_unique, + invert=invert, kind=kind).reshape(element.shape) + + +def _union1d_dispatcher(ar1, ar2): + return (ar1, ar2) + + +@array_function_dispatch(_union1d_dispatcher) +def union1d(ar1, ar2): + """ + Find the union of two arrays. + + Return the unique, sorted array of values that are in either of the two + input arrays. + + Parameters + ---------- + ar1, ar2 : array_like + Input arrays. They are flattened if they are not already 1D. + + Returns + ------- + union1d : ndarray + Unique, sorted union of the input arrays. + + Examples + -------- + >>> import numpy as np + >>> np.union1d([-1, 0, 1], [-2, 0, 2]) + array([-2, -1, 0, 1, 2]) + + To find the union of more than two arrays, use functools.reduce: + + >>> from functools import reduce + >>> reduce(np.union1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) + array([1, 2, 3, 4, 6]) + """ + return unique(np.concatenate((ar1, ar2), axis=None)) + + +def _setdiff1d_dispatcher(ar1, ar2, assume_unique=None): + return (ar1, ar2) + + +@array_function_dispatch(_setdiff1d_dispatcher) +def setdiff1d(ar1, ar2, assume_unique=False): + """ + Find the set difference of two arrays. + + Return the unique values in `ar1` that are not in `ar2`. + + Parameters + ---------- + ar1 : array_like + Input array. + ar2 : array_like + Input comparison array. + assume_unique : bool + If True, the input arrays are both assumed to be unique, which + can speed up the calculation. Default is False. + + Returns + ------- + setdiff1d : ndarray + 1D array of values in `ar1` that are not in `ar2`. The result + is sorted when `assume_unique=False`, but otherwise only sorted + if the input is sorted. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1, 2, 3, 2, 4, 1]) + >>> b = np.array([3, 4, 5, 6]) + >>> np.setdiff1d(a, b) + array([1, 2]) + + """ + if assume_unique: + ar1 = np.asarray(ar1).ravel() + else: + ar1 = unique(ar1) + ar2 = unique(ar2) + return ar1[_in1d(ar1, ar2, assume_unique=True, invert=True)] diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..20f2d576bf00b8627a001c81933356746952d258 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arraysetops_impl.pyi @@ -0,0 +1,379 @@ +from typing import ( + Any, + Generic, + Literal as L, + NamedTuple, + overload, + SupportsIndex, + TypeVar, +) +from typing_extensions import deprecated + +import numpy as np +from numpy import generic, number, int8, intp, timedelta64, object_ + +from numpy._typing import ( + ArrayLike, + NDArray, + _ArrayLike, + _ArrayLikeBool_co, + _ArrayLikeDT64_co, + _ArrayLikeTD64_co, + _ArrayLikeObject_co, + _ArrayLikeNumber_co, +) + +__all__ = [ + "ediff1d", + "in1d", + "intersect1d", + "isin", + "setdiff1d", + "setxor1d", + "union1d", + "unique", + "unique_all", + "unique_counts", + "unique_inverse", + "unique_values", +] + +_SCT = TypeVar("_SCT", bound=generic) +_NumberType = TypeVar("_NumberType", bound=number[Any]) + +# Explicitly set all allowed values to prevent accidental castings to +# abstract dtypes (their common super-type). +# +# Only relevant if two or more arguments are parametrized, (e.g. `setdiff1d`) +# which could result in, for example, `int64` and `float64`producing a +# `number[_64Bit]` array +_EitherSCT = TypeVar( + "_EitherSCT", + np.bool, + np.int8, np.int16, np.int32, np.int64, np.intp, + np.uint8, np.uint16, np.uint32, np.uint64, np.uintp, + np.float16, np.float32, np.float64, np.longdouble, + np.complex64, np.complex128, np.clongdouble, + np.timedelta64, np.datetime64, + np.bytes_, np.str_, np.void, np.object_, + np.integer, np.floating, np.complexfloating, np.character, +) # fmt: skip + +class UniqueAllResult(NamedTuple, Generic[_SCT]): + values: NDArray[_SCT] + indices: NDArray[intp] + inverse_indices: NDArray[intp] + counts: NDArray[intp] + +class UniqueCountsResult(NamedTuple, Generic[_SCT]): + values: NDArray[_SCT] + counts: NDArray[intp] + +class UniqueInverseResult(NamedTuple, Generic[_SCT]): + values: NDArray[_SCT] + inverse_indices: NDArray[intp] + +@overload +def ediff1d( + ary: _ArrayLikeBool_co, + to_end: None | ArrayLike = ..., + to_begin: None | ArrayLike = ..., +) -> NDArray[int8]: ... +@overload +def ediff1d( + ary: _ArrayLike[_NumberType], + to_end: None | ArrayLike = ..., + to_begin: None | ArrayLike = ..., +) -> NDArray[_NumberType]: ... +@overload +def ediff1d( + ary: _ArrayLikeNumber_co, + to_end: None | ArrayLike = ..., + to_begin: None | ArrayLike = ..., +) -> NDArray[Any]: ... +@overload +def ediff1d( + ary: _ArrayLikeDT64_co | _ArrayLikeTD64_co, + to_end: None | ArrayLike = ..., + to_begin: None | ArrayLike = ..., +) -> NDArray[timedelta64]: ... +@overload +def ediff1d( + ary: _ArrayLikeObject_co, + to_end: None | ArrayLike = ..., + to_begin: None | ArrayLike = ..., +) -> NDArray[object_]: ... + +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[False] = ..., + return_inverse: L[False] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False] = ..., + return_inverse: L[False] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> NDArray[Any]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[True] = ..., + return_inverse: L[False] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[True] = ..., + return_inverse: L[False] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[False] = ..., + return_inverse: L[True] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False] = ..., + return_inverse: L[True] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[False] = ..., + return_inverse: L[False] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False] = ..., + return_inverse: L[False] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[True] = ..., + return_inverse: L[True] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[True] = ..., + return_inverse: L[True] = ..., + return_counts: L[False] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[True] = ..., + return_inverse: L[False] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[True] = ..., + return_inverse: L[False] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[False] = ..., + return_inverse: L[True] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False] = ..., + return_inverse: L[True] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: _ArrayLike[_SCT], + return_index: L[True] = ..., + return_inverse: L[True] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[_SCT], NDArray[intp], NDArray[intp], NDArray[intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[True] = ..., + return_inverse: L[True] = ..., + return_counts: L[True] = ..., + axis: None | SupportsIndex = ..., + *, + equal_nan: bool = ..., +) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp], NDArray[intp]]: ... + +@overload +def unique_all( + x: _ArrayLike[_SCT], / +) -> UniqueAllResult[_SCT]: ... +@overload +def unique_all( + x: ArrayLike, / +) -> UniqueAllResult[Any]: ... + +@overload +def unique_counts( + x: _ArrayLike[_SCT], / +) -> UniqueCountsResult[_SCT]: ... +@overload +def unique_counts( + x: ArrayLike, / +) -> UniqueCountsResult[Any]: ... + +@overload +def unique_inverse(x: _ArrayLike[_SCT], /) -> UniqueInverseResult[_SCT]: ... +@overload +def unique_inverse(x: ArrayLike, /) -> UniqueInverseResult[Any]: ... + +@overload +def unique_values(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +@overload +def unique_values(x: ArrayLike, /) -> NDArray[Any]: ... + +@overload +def intersect1d( + ar1: _ArrayLike[_EitherSCT], + ar2: _ArrayLike[_EitherSCT], + assume_unique: bool = ..., + return_indices: L[False] = ..., +) -> NDArray[_EitherSCT]: ... +@overload +def intersect1d( + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = ..., + return_indices: L[False] = ..., +) -> NDArray[Any]: ... +@overload +def intersect1d( + ar1: _ArrayLike[_EitherSCT], + ar2: _ArrayLike[_EitherSCT], + assume_unique: bool = ..., + return_indices: L[True] = ..., +) -> tuple[NDArray[_EitherSCT], NDArray[intp], NDArray[intp]]: ... +@overload +def intersect1d( + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = ..., + return_indices: L[True] = ..., +) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp]]: ... + +@overload +def setxor1d( + ar1: _ArrayLike[_EitherSCT], + ar2: _ArrayLike[_EitherSCT], + assume_unique: bool = ..., +) -> NDArray[_EitherSCT]: ... +@overload +def setxor1d( + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = ..., +) -> NDArray[Any]: ... + +def isin( + element: ArrayLike, + test_elements: ArrayLike, + assume_unique: bool = ..., + invert: bool = ..., + *, + kind: None | str = ..., +) -> NDArray[np.bool]: ... + +@deprecated("Use 'isin' instead") +def in1d( + element: ArrayLike, + test_elements: ArrayLike, + assume_unique: bool = ..., + invert: bool = ..., + *, + kind: None | str = ..., +) -> NDArray[np.bool]: ... + +@overload +def union1d( + ar1: _ArrayLike[_EitherSCT], + ar2: _ArrayLike[_EitherSCT], +) -> NDArray[_EitherSCT]: ... +@overload +def union1d( + ar1: ArrayLike, + ar2: ArrayLike, +) -> NDArray[Any]: ... + +@overload +def setdiff1d( + ar1: _ArrayLike[_EitherSCT], + ar2: _ArrayLike[_EitherSCT], + assume_unique: bool = ..., +) -> NDArray[_EitherSCT]: ... +@overload +def setdiff1d( + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = ..., +) -> NDArray[Any]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..efc529de5cff5c051e5df5b5847a4317d60d75cb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.py @@ -0,0 +1,224 @@ +""" +A buffered iterator for big arrays. + +This module solves the problem of iterating over a big file-based array +without having to read it into memory. The `Arrayterator` class wraps +an array object, and when iterated it will return sub-arrays with at most +a user-specified number of elements. + +""" +from operator import mul +from functools import reduce + +__all__ = ['Arrayterator'] + + +class Arrayterator: + """ + Buffered iterator for big arrays. + + `Arrayterator` creates a buffered iterator for reading big arrays in small + contiguous blocks. The class is useful for objects stored in the + file system. It allows iteration over the object *without* reading + everything in memory; instead, small blocks are read and iterated over. + + `Arrayterator` can be used with any object that supports multidimensional + slices. This includes NumPy arrays, but also variables from + Scientific.IO.NetCDF or pynetcdf for example. + + Parameters + ---------- + var : array_like + The object to iterate over. + buf_size : int, optional + The buffer size. If `buf_size` is supplied, the maximum amount of + data that will be read into memory is `buf_size` elements. + Default is None, which will read as many element as possible + into memory. + + Attributes + ---------- + var + buf_size + start + stop + step + shape + flat + + See Also + -------- + numpy.ndenumerate : Multidimensional array iterator. + numpy.flatiter : Flat array iterator. + numpy.memmap : Create a memory-map to an array stored + in a binary file on disk. + + Notes + ----- + The algorithm works by first finding a "running dimension", along which + the blocks will be extracted. Given an array of dimensions + ``(d1, d2, ..., dn)``, e.g. if `buf_size` is smaller than ``d1``, the + first dimension will be used. If, on the other hand, + ``d1 < buf_size < d1*d2`` the second dimension will be used, and so on. + Blocks are extracted along this dimension, and when the last block is + returned the process continues from the next dimension, until all + elements have been read. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(3 * 4 * 5 * 6).reshape(3, 4, 5, 6) + >>> a_itor = np.lib.Arrayterator(a, 2) + >>> a_itor.shape + (3, 4, 5, 6) + + Now we can iterate over ``a_itor``, and it will return arrays of size + two. Since `buf_size` was smaller than any dimension, the first + dimension will be iterated over first: + + >>> for subarr in a_itor: + ... if not subarr.all(): + ... print(subarr, subarr.shape) # doctest: +SKIP + >>> # [[[[0 1]]]] (1, 1, 1, 2) + + """ + + __module__ = "numpy.lib" + + def __init__(self, var, buf_size=None): + self.var = var + self.buf_size = buf_size + + self.start = [0 for dim in var.shape] + self.stop = list(var.shape) + self.step = [1 for dim in var.shape] + + def __getattr__(self, attr): + return getattr(self.var, attr) + + def __getitem__(self, index): + """ + Return a new arrayterator. + + """ + # Fix index, handling ellipsis and incomplete slices. + if not isinstance(index, tuple): + index = (index,) + fixed = [] + length, dims = len(index), self.ndim + for slice_ in index: + if slice_ is Ellipsis: + fixed.extend([slice(None)] * (dims-length+1)) + length = len(fixed) + elif isinstance(slice_, int): + fixed.append(slice(slice_, slice_+1, 1)) + else: + fixed.append(slice_) + index = tuple(fixed) + if len(index) < dims: + index += (slice(None),) * (dims-len(index)) + + # Return a new arrayterator object. + out = self.__class__(self.var, self.buf_size) + for i, (start, stop, step, slice_) in enumerate( + zip(self.start, self.stop, self.step, index)): + out.start[i] = start + (slice_.start or 0) + out.step[i] = step * (slice_.step or 1) + out.stop[i] = start + (slice_.stop or stop-start) + out.stop[i] = min(stop, out.stop[i]) + return out + + def __array__(self, dtype=None, copy=None): + """ + Return corresponding data. + + """ + slice_ = tuple(slice(*t) for t in zip( + self.start, self.stop, self.step)) + return self.var[slice_] + + @property + def flat(self): + """ + A 1-D flat iterator for Arrayterator objects. + + This iterator returns elements of the array to be iterated over in + `~lib.Arrayterator` one by one. + It is similar to `flatiter`. + + See Also + -------- + lib.Arrayterator + flatiter + + Examples + -------- + >>> a = np.arange(3 * 4 * 5 * 6).reshape(3, 4, 5, 6) + >>> a_itor = np.lib.Arrayterator(a, 2) + + >>> for subarr in a_itor.flat: + ... if not subarr: + ... print(subarr, type(subarr)) + ... + 0 + + """ + for block in self: + yield from block.flat + + @property + def shape(self): + """ + The shape of the array to be iterated over. + + For an example, see `Arrayterator`. + + """ + return tuple(((stop-start-1)//step+1) for start, stop, step in + zip(self.start, self.stop, self.step)) + + def __iter__(self): + # Skip arrays with degenerate dimensions + if [dim for dim in self.shape if dim <= 0]: + return + + start = self.start[:] + stop = self.stop[:] + step = self.step[:] + ndims = self.var.ndim + + while True: + count = self.buf_size or reduce(mul, self.shape) + + # iterate over each dimension, looking for the + # running dimension (ie, the dimension along which + # the blocks will be built from) + rundim = 0 + for i in range(ndims-1, -1, -1): + # if count is zero we ran out of elements to read + # along higher dimensions, so we read only a single position + if count == 0: + stop[i] = start[i]+1 + elif count <= self.shape[i]: + # limit along this dimension + stop[i] = start[i] + count*step[i] + rundim = i + else: + # read everything along this dimension + stop[i] = self.stop[i] + stop[i] = min(self.stop[i], stop[i]) + count = count//self.shape[i] + + # yield a block + slice_ = tuple(slice(*t) for t in zip(start, stop, step)) + yield self.var[slice_] + + # Update start position, taking care of overflow to + # other dimensions + start[rundim] = stop[rundim] # start where we stopped + for i in range(ndims-1, 0, -1): + if start[i] >= self.stop[i]: + start[i] = self.start[i] + start[i-1] += self.step[i-1] + if start[0] >= self.stop[0]: + return diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c24fe56ac8a911c38b6d423a274e10fb0838fb98 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_arrayterator_impl.pyi @@ -0,0 +1,46 @@ +# pyright: reportIncompatibleMethodOverride=false + +from collections.abc import Generator +from types import EllipsisType +from typing import Any, Final, TypeAlias, overload + +from typing_extensions import TypeVar + +import numpy as np + +__all__ = ["Arrayterator"] + +_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], covariant=True) +_DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any]) +_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype[Any], covariant=True) +_ScalarT = TypeVar("_ScalarT", bound=np.generic) + +_AnyIndex: TypeAlias = EllipsisType | int | slice | tuple[EllipsisType | int | slice, ...] + + +# NOTE: In reality `Arrayterator` does not actually inherit from `ndarray`, +# but its ``__getattr__` method does wrap around the former and thus has +# access to all its methods + +class Arrayterator(np.ndarray[_ShapeT_co, _DTypeT_co]): + var: np.ndarray[_ShapeT_co, _DTypeT_co] # type: ignore[assignment] + buf_size: Final[int | None] + start: Final[list[int]] + stop: Final[list[int]] + step: Final[list[int]] + + @property # type: ignore[misc] + def shape(self) -> _ShapeT_co: ... + @property + def flat(self: Arrayterator[Any, np.dtype[_ScalarT]]) -> Generator[_ScalarT]: ... # type: ignore[override] + + # + def __init__(self, /, var: np.ndarray[_ShapeT_co, _DTypeT_co], buf_size: int | None = None) -> None: ... + def __getitem__(self, index: _AnyIndex, /) -> Arrayterator[tuple[int, ...], _DTypeT_co]: ... # type: ignore[override] + def __iter__(self) -> Generator[np.ndarray[tuple[int, ...], _DTypeT_co]]: ... + + # + @overload # type: ignore[override] + def __array__(self, /, dtype: None = None, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ... + @overload + def __array__(self, /, dtype: _DTypeT, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_datasource.py b/venv/lib/python3.10/site-packages/numpy/lib/_datasource.py new file mode 100644 index 0000000000000000000000000000000000000000..e3d85b8549415c93d3abe33009b22f4f1e6d55fc --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_datasource.py @@ -0,0 +1,700 @@ +"""A file interface for handling local and remote data files. + +The goal of datasource is to abstract some of the file system operations +when dealing with data files so the researcher doesn't have to know all the +low-level details. Through datasource, a researcher can obtain and use a +file with one function call, regardless of location of the file. + +DataSource is meant to augment standard python libraries, not replace them. +It should work seamlessly with standard file IO operations and the os +module. + +DataSource files can originate locally or remotely: + +- local files : '/home/guido/src/local/data.txt' +- URLs (http, ftp, ...) : 'http://www.scipy.org/not/real/data.txt' + +DataSource files can also be compressed or uncompressed. Currently only +gzip, bz2 and xz are supported. + +Example:: + + >>> # Create a DataSource, use os.curdir (default) for local storage. + >>> from numpy import DataSource + >>> ds = DataSource() + >>> + >>> # Open a remote file. + >>> # DataSource downloads the file, stores it locally in: + >>> # './www.google.com/index.html' + >>> # opens the file and returns a file object. + >>> fp = ds.open('http://www.google.com/') # doctest: +SKIP + >>> + >>> # Use the file as you normally would + >>> fp.read() # doctest: +SKIP + >>> fp.close() # doctest: +SKIP + +""" +import os + +from .._utils import set_module + + +_open = open + + +def _check_mode(mode, encoding, newline): + """Check mode and that encoding and newline are compatible. + + Parameters + ---------- + mode : str + File open mode. + encoding : str + File encoding. + newline : str + Newline for text files. + + """ + if "t" in mode: + if "b" in mode: + raise ValueError("Invalid mode: %r" % (mode,)) + else: + if encoding is not None: + raise ValueError("Argument 'encoding' not supported in binary mode") + if newline is not None: + raise ValueError("Argument 'newline' not supported in binary mode") + + +# Using a class instead of a module-level dictionary +# to reduce the initial 'import numpy' overhead by +# deferring the import of lzma, bz2 and gzip until needed + +# TODO: .zip support, .tar support? +class _FileOpeners: + """ + Container for different methods to open (un-)compressed files. + + `_FileOpeners` contains a dictionary that holds one method for each + supported file format. Attribute lookup is implemented in such a way + that an instance of `_FileOpeners` itself can be indexed with the keys + of that dictionary. Currently uncompressed files as well as files + compressed with ``gzip``, ``bz2`` or ``xz`` compression are supported. + + Notes + ----- + `_file_openers`, an instance of `_FileOpeners`, is made available for + use in the `_datasource` module. + + Examples + -------- + >>> import gzip + >>> np.lib._datasource._file_openers.keys() + [None, '.bz2', '.gz', '.xz', '.lzma'] + >>> np.lib._datasource._file_openers['.gz'] is gzip.open + True + + """ + + def __init__(self): + self._loaded = False + self._file_openers = {None: open} + + def _load(self): + if self._loaded: + return + + try: + import bz2 + self._file_openers[".bz2"] = bz2.open + except ImportError: + pass + + try: + import gzip + self._file_openers[".gz"] = gzip.open + except ImportError: + pass + + try: + import lzma + self._file_openers[".xz"] = lzma.open + self._file_openers[".lzma"] = lzma.open + except (ImportError, AttributeError): + # There are incompatible backports of lzma that do not have the + # lzma.open attribute, so catch that as well as ImportError. + pass + + self._loaded = True + + def keys(self): + """ + Return the keys of currently supported file openers. + + Parameters + ---------- + None + + Returns + ------- + keys : list + The keys are None for uncompressed files and the file extension + strings (i.e. ``'.gz'``, ``'.xz'``) for supported compression + methods. + + """ + self._load() + return list(self._file_openers.keys()) + + def __getitem__(self, key): + self._load() + return self._file_openers[key] + +_file_openers = _FileOpeners() + +def open(path, mode='r', destpath=os.curdir, encoding=None, newline=None): + """ + Open `path` with `mode` and return the file object. + + If ``path`` is an URL, it will be downloaded, stored in the + `DataSource` `destpath` directory and opened from there. + + Parameters + ---------- + path : str or pathlib.Path + Local file path or URL to open. + mode : str, optional + Mode to open `path`. Mode 'r' for reading, 'w' for writing, 'a' to + append. Available modes depend on the type of object specified by + path. Default is 'r'. + destpath : str, optional + Path to the directory where the source file gets downloaded to for + use. If `destpath` is None, a temporary directory will be created. + The default path is the current directory. + encoding : {None, str}, optional + Open text file with given encoding. The default encoding will be + what `open` uses. + newline : {None, str}, optional + Newline to use when reading text file. + + Returns + ------- + out : file object + The opened file. + + Notes + ----- + This is a convenience function that instantiates a `DataSource` and + returns the file object from ``DataSource.open(path)``. + + """ + + ds = DataSource(destpath) + return ds.open(path, mode, encoding=encoding, newline=newline) + + +@set_module('numpy.lib.npyio') +class DataSource: + """ + DataSource(destpath='.') + + A generic data source file (file, http, ftp, ...). + + DataSources can be local files or remote files/URLs. The files may + also be compressed or uncompressed. DataSource hides some of the + low-level details of downloading the file, allowing you to simply pass + in a valid file path (or URL) and obtain a file object. + + Parameters + ---------- + destpath : str or None, optional + Path to the directory where the source file gets downloaded to for + use. If `destpath` is None, a temporary directory will be created. + The default path is the current directory. + + Notes + ----- + URLs require a scheme string (``http://``) to be used, without it they + will fail:: + + >>> repos = np.lib.npyio.DataSource() + >>> repos.exists('www.google.com/index.html') + False + >>> repos.exists('http://www.google.com/index.html') + True + + Temporary directories are deleted when the DataSource is deleted. + + Examples + -------- + :: + + >>> ds = np.lib.npyio.DataSource('/home/guido') + >>> urlname = 'http://www.google.com/' + >>> gfile = ds.open('http://www.google.com/') + >>> ds.abspath(urlname) + '/home/guido/www.google.com/index.html' + + >>> ds = np.lib.npyio.DataSource(None) # use with temporary file + >>> ds.open('/home/guido/foobar.txt') + + >>> ds.abspath('/home/guido/foobar.txt') + '/tmp/.../home/guido/foobar.txt' + + """ + + def __init__(self, destpath=os.curdir): + """Create a DataSource with a local path at destpath.""" + if destpath: + self._destpath = os.path.abspath(destpath) + self._istmpdest = False + else: + import tempfile # deferring import to improve startup time + self._destpath = tempfile.mkdtemp() + self._istmpdest = True + + def __del__(self): + # Remove temp directories + if hasattr(self, '_istmpdest') and self._istmpdest: + import shutil + + shutil.rmtree(self._destpath) + + def _iszip(self, filename): + """Test if the filename is a zip file by looking at the file extension. + + """ + fname, ext = os.path.splitext(filename) + return ext in _file_openers.keys() + + def _iswritemode(self, mode): + """Test if the given mode will open a file for writing.""" + + # Currently only used to test the bz2 files. + _writemodes = ("w", "+") + return any(c in _writemodes for c in mode) + + def _splitzipext(self, filename): + """Split zip extension from filename and return filename. + + Returns + ------- + base, zip_ext : {tuple} + + """ + + if self._iszip(filename): + return os.path.splitext(filename) + else: + return filename, None + + def _possible_names(self, filename): + """Return a tuple containing compressed filename variations.""" + names = [filename] + if not self._iszip(filename): + for zipext in _file_openers.keys(): + if zipext: + names.append(filename+zipext) + return names + + def _isurl(self, path): + """Test if path is a net location. Tests the scheme and netloc.""" + + # We do this here to reduce the 'import numpy' initial import time. + from urllib.parse import urlparse + + # BUG : URLs require a scheme string ('http://') to be used. + # www.google.com will fail. + # Should we prepend the scheme for those that don't have it and + # test that also? Similar to the way we append .gz and test for + # for compressed versions of files. + + scheme, netloc, upath, uparams, uquery, ufrag = urlparse(path) + return bool(scheme and netloc) + + def _cache(self, path): + """Cache the file specified by path. + + Creates a copy of the file in the datasource cache. + + """ + # We import these here because importing them is slow and + # a significant fraction of numpy's total import time. + import shutil + from urllib.request import urlopen + + upath = self.abspath(path) + + # ensure directory exists + if not os.path.exists(os.path.dirname(upath)): + os.makedirs(os.path.dirname(upath)) + + # TODO: Doesn't handle compressed files! + if self._isurl(path): + with urlopen(path) as openedurl: + with _open(upath, 'wb') as f: + shutil.copyfileobj(openedurl, f) + else: + shutil.copyfile(path, upath) + return upath + + def _findfile(self, path): + """Searches for ``path`` and returns full path if found. + + If path is an URL, _findfile will cache a local copy and return the + path to the cached file. If path is a local file, _findfile will + return a path to that local file. + + The search will include possible compressed versions of the file + and return the first occurrence found. + + """ + + # Build list of possible local file paths + if not self._isurl(path): + # Valid local paths + filelist = self._possible_names(path) + # Paths in self._destpath + filelist += self._possible_names(self.abspath(path)) + else: + # Cached URLs in self._destpath + filelist = self._possible_names(self.abspath(path)) + # Remote URLs + filelist = filelist + self._possible_names(path) + + for name in filelist: + if self.exists(name): + if self._isurl(name): + name = self._cache(name) + return name + return None + + def abspath(self, path): + """ + Return absolute path of file in the DataSource directory. + + If `path` is an URL, then `abspath` will return either the location + the file exists locally or the location it would exist when opened + using the `open` method. + + Parameters + ---------- + path : str or pathlib.Path + Can be a local file or a remote URL. + + Returns + ------- + out : str + Complete path, including the `DataSource` destination directory. + + Notes + ----- + The functionality is based on `os.path.abspath`. + + """ + # We do this here to reduce the 'import numpy' initial import time. + from urllib.parse import urlparse + + # TODO: This should be more robust. Handles case where path includes + # the destpath, but not other sub-paths. Failing case: + # path = /home/guido/datafile.txt + # destpath = /home/alex/ + # upath = self.abspath(path) + # upath == '/home/alex/home/guido/datafile.txt' + + # handle case where path includes self._destpath + splitpath = path.split(self._destpath, 2) + if len(splitpath) > 1: + path = splitpath[1] + scheme, netloc, upath, uparams, uquery, ufrag = urlparse(path) + netloc = self._sanitize_relative_path(netloc) + upath = self._sanitize_relative_path(upath) + return os.path.join(self._destpath, netloc, upath) + + def _sanitize_relative_path(self, path): + """Return a sanitised relative path for which + os.path.abspath(os.path.join(base, path)).startswith(base) + """ + last = None + path = os.path.normpath(path) + while path != last: + last = path + # Note: os.path.join treats '/' as os.sep on Windows + path = path.lstrip(os.sep).lstrip('/') + path = path.lstrip(os.pardir).removeprefix('..') + drive, path = os.path.splitdrive(path) # for Windows + return path + + def exists(self, path): + """ + Test if path exists. + + Test if `path` exists as (and in this order): + + - a local file. + - a remote URL that has been downloaded and stored locally in the + `DataSource` directory. + - a remote URL that has not been downloaded, but is valid and + accessible. + + Parameters + ---------- + path : str or pathlib.Path + Can be a local file or a remote URL. + + Returns + ------- + out : bool + True if `path` exists. + + Notes + ----- + When `path` is an URL, `exists` will return True if it's either + stored locally in the `DataSource` directory, or is a valid remote + URL. `DataSource` does not discriminate between the two, the file + is accessible if it exists in either location. + + """ + + # First test for local path + if os.path.exists(path): + return True + + # We import this here because importing urllib is slow and + # a significant fraction of numpy's total import time. + from urllib.request import urlopen + from urllib.error import URLError + + # Test cached url + upath = self.abspath(path) + if os.path.exists(upath): + return True + + # Test remote url + if self._isurl(path): + try: + netfile = urlopen(path) + netfile.close() + del(netfile) + return True + except URLError: + return False + return False + + def open(self, path, mode='r', encoding=None, newline=None): + """ + Open and return file-like object. + + If `path` is an URL, it will be downloaded, stored in the + `DataSource` directory and opened from there. + + Parameters + ---------- + path : str or pathlib.Path + Local file path or URL to open. + mode : {'r', 'w', 'a'}, optional + Mode to open `path`. Mode 'r' for reading, 'w' for writing, + 'a' to append. Available modes depend on the type of object + specified by `path`. Default is 'r'. + encoding : {None, str}, optional + Open text file with given encoding. The default encoding will be + what `open` uses. + newline : {None, str}, optional + Newline to use when reading text file. + + Returns + ------- + out : file object + File object. + + """ + + # TODO: There is no support for opening a file for writing which + # doesn't exist yet (creating a file). Should there be? + + # TODO: Add a ``subdir`` parameter for specifying the subdirectory + # used to store URLs in self._destpath. + + if self._isurl(path) and self._iswritemode(mode): + raise ValueError("URLs are not writeable") + + # NOTE: _findfile will fail on a new file opened for writing. + found = self._findfile(path) + if found: + _fname, ext = self._splitzipext(found) + if ext == 'bz2': + mode.replace("+", "") + return _file_openers[ext](found, mode=mode, + encoding=encoding, newline=newline) + else: + raise FileNotFoundError(f"{path} not found.") + + +class Repository (DataSource): + """ + Repository(baseurl, destpath='.') + + A data repository where multiple DataSource's share a base + URL/directory. + + `Repository` extends `DataSource` by prepending a base URL (or + directory) to all the files it handles. Use `Repository` when you will + be working with multiple files from one base URL. Initialize + `Repository` with the base URL, then refer to each file by its filename + only. + + Parameters + ---------- + baseurl : str + Path to the local directory or remote location that contains the + data files. + destpath : str or None, optional + Path to the directory where the source file gets downloaded to for + use. If `destpath` is None, a temporary directory will be created. + The default path is the current directory. + + Examples + -------- + To analyze all files in the repository, do something like this + (note: this is not self-contained code):: + + >>> repos = np.lib._datasource.Repository('/home/user/data/dir/') + >>> for filename in filelist: + ... fp = repos.open(filename) + ... fp.analyze() + ... fp.close() + + Similarly you could use a URL for a repository:: + + >>> repos = np.lib._datasource.Repository('http://www.xyz.edu/data') + + """ + + def __init__(self, baseurl, destpath=os.curdir): + """Create a Repository with a shared url or directory of baseurl.""" + DataSource.__init__(self, destpath=destpath) + self._baseurl = baseurl + + def __del__(self): + DataSource.__del__(self) + + def _fullpath(self, path): + """Return complete path for path. Prepends baseurl if necessary.""" + splitpath = path.split(self._baseurl, 2) + if len(splitpath) == 1: + result = os.path.join(self._baseurl, path) + else: + result = path # path contains baseurl already + return result + + def _findfile(self, path): + """Extend DataSource method to prepend baseurl to ``path``.""" + return DataSource._findfile(self, self._fullpath(path)) + + def abspath(self, path): + """ + Return absolute path of file in the Repository directory. + + If `path` is an URL, then `abspath` will return either the location + the file exists locally or the location it would exist when opened + using the `open` method. + + Parameters + ---------- + path : str or pathlib.Path + Can be a local file or a remote URL. This may, but does not + have to, include the `baseurl` with which the `Repository` was + initialized. + + Returns + ------- + out : str + Complete path, including the `DataSource` destination directory. + + """ + return DataSource.abspath(self, self._fullpath(path)) + + def exists(self, path): + """ + Test if path exists prepending Repository base URL to path. + + Test if `path` exists as (and in this order): + + - a local file. + - a remote URL that has been downloaded and stored locally in the + `DataSource` directory. + - a remote URL that has not been downloaded, but is valid and + accessible. + + Parameters + ---------- + path : str or pathlib.Path + Can be a local file or a remote URL. This may, but does not + have to, include the `baseurl` with which the `Repository` was + initialized. + + Returns + ------- + out : bool + True if `path` exists. + + Notes + ----- + When `path` is an URL, `exists` will return True if it's either + stored locally in the `DataSource` directory, or is a valid remote + URL. `DataSource` does not discriminate between the two, the file + is accessible if it exists in either location. + + """ + return DataSource.exists(self, self._fullpath(path)) + + def open(self, path, mode='r', encoding=None, newline=None): + """ + Open and return file-like object prepending Repository base URL. + + If `path` is an URL, it will be downloaded, stored in the + DataSource directory and opened from there. + + Parameters + ---------- + path : str or pathlib.Path + Local file path or URL to open. This may, but does not have to, + include the `baseurl` with which the `Repository` was + initialized. + mode : {'r', 'w', 'a'}, optional + Mode to open `path`. Mode 'r' for reading, 'w' for writing, + 'a' to append. Available modes depend on the type of object + specified by `path`. Default is 'r'. + encoding : {None, str}, optional + Open text file with given encoding. The default encoding will be + what `open` uses. + newline : {None, str}, optional + Newline to use when reading text file. + + Returns + ------- + out : file object + File object. + + """ + return DataSource.open(self, self._fullpath(path), mode, + encoding=encoding, newline=newline) + + def listdir(self): + """ + List files in the source Repository. + + Returns + ------- + files : list of str or pathlib.Path + List of file names (not containing a directory part). + + Notes + ----- + Does not currently work for remote repositories. + + """ + if self._isurl(self._baseurl): + raise NotImplementedError( + "Directory listing of URLs, not supported yet.") + else: + return os.listdir(self._baseurl) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_datasource.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_datasource.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9f91fdf893a07a3bd3398ae43e2604a60d04d903 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_datasource.pyi @@ -0,0 +1,31 @@ +from pathlib import Path +from typing import IO, Any, TypeAlias + +from _typeshed import OpenBinaryMode, OpenTextMode + +_Mode: TypeAlias = OpenBinaryMode | OpenTextMode + +### + +# exported in numpy.lib.nppyio +class DataSource: + def __init__(self, /, destpath: Path | str | None = ...) -> 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: _Mode = "r", encoding: str | None = None, newline: str | None = None) -> IO[Any]: ... + +class Repository(DataSource): + def __init__(self, /, baseurl: str, destpath: str | None = ...) -> None: ... + def listdir(self, /) -> list[str]: ... + +def open( + path: str, + mode: _Mode = "r", + destpath: str | None = ..., + encoding: str | None = None, + newline: str | None = None, +) -> IO[Any]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..3fa9c5f99d95352ffc29ade5d33531ccc1f90ead --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.py @@ -0,0 +1,5827 @@ +import builtins +import collections.abc +import functools +import re +import sys +import warnings + +import numpy as np +import numpy._core.numeric as _nx +from numpy._core import transpose, overrides +from numpy._core.numeric import ( + ones, zeros_like, arange, concatenate, array, asarray, asanyarray, empty, + ndarray, take, dot, where, intp, integer, isscalar, absolute + ) +from numpy._core.umath import ( + pi, add, arctan2, frompyfunc, cos, less_equal, sqrt, sin, + mod, exp, not_equal, subtract, minimum + ) +from numpy._core.fromnumeric import ( + ravel, nonzero, partition, mean, any, sum + ) +from numpy._core.numerictypes import typecodes +from numpy.lib._twodim_base_impl import diag +from numpy._core.multiarray import ( + _place, bincount, normalize_axis_index, _monotonicity, + interp as compiled_interp, interp_complex as compiled_interp_complex + ) +from numpy._core._multiarray_umath import _array_converter +from numpy._utils import set_module + +# needed in this module for compatibility +from numpy.lib._histograms_impl import histogram, histogramdd # noqa: F401 + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +__all__ = [ + 'select', 'piecewise', 'trim_zeros', 'copy', 'iterable', 'percentile', + 'diff', 'gradient', 'angle', 'unwrap', 'sort_complex', 'flip', + 'rot90', 'extract', 'place', 'vectorize', 'asarray_chkfinite', 'average', + 'bincount', 'digitize', 'cov', 'corrcoef', + 'median', 'sinc', 'hamming', 'hanning', 'bartlett', + 'blackman', 'kaiser', 'trapezoid', 'trapz', 'i0', + 'meshgrid', 'delete', 'insert', 'append', 'interp', + 'quantile' + ] + +# _QuantileMethods is a dictionary listing all the supported methods to +# compute quantile/percentile. +# +# Below virtual_index refers to the index of the element where the percentile +# would be found in the sorted sample. +# When the sample contains exactly the percentile wanted, the virtual_index is +# an integer to the index of this element. +# When the percentile wanted is in between two elements, the virtual_index +# is made of a integer part (a.k.a 'i' or 'left') and a fractional part +# (a.k.a 'g' or 'gamma') +# +# Each method in _QuantileMethods has two properties +# get_virtual_index : Callable +# The function used to compute the virtual_index. +# fix_gamma : Callable +# A function used for discrete methods to force the index to a specific value. +_QuantileMethods = dict( + # --- HYNDMAN and FAN METHODS + # Discrete methods + inverted_cdf=dict( + get_virtual_index=lambda n, quantiles: _inverted_cdf(n, quantiles), + fix_gamma=None, # should never be called + ), + averaged_inverted_cdf=dict( + get_virtual_index=lambda n, quantiles: (n * quantiles) - 1, + fix_gamma=lambda gamma, _: _get_gamma_mask( + shape=gamma.shape, + default_value=1., + conditioned_value=0.5, + where=gamma == 0), + ), + closest_observation=dict( + get_virtual_index=lambda n, quantiles: _closest_observation(n, + quantiles), + fix_gamma=None, # should never be called + ), + # Continuous methods + interpolated_inverted_cdf=dict( + get_virtual_index=lambda n, quantiles: + _compute_virtual_index(n, quantiles, 0, 1), + fix_gamma=lambda gamma, _: gamma, + ), + hazen=dict( + get_virtual_index=lambda n, quantiles: + _compute_virtual_index(n, quantiles, 0.5, 0.5), + fix_gamma=lambda gamma, _: gamma, + ), + weibull=dict( + get_virtual_index=lambda n, quantiles: + _compute_virtual_index(n, quantiles, 0, 0), + fix_gamma=lambda gamma, _: gamma, + ), + # Default method. + # To avoid some rounding issues, `(n-1) * quantiles` is preferred to + # `_compute_virtual_index(n, quantiles, 1, 1)`. + # They are mathematically equivalent. + linear=dict( + get_virtual_index=lambda n, quantiles: (n - 1) * quantiles, + fix_gamma=lambda gamma, _: gamma, + ), + median_unbiased=dict( + get_virtual_index=lambda n, quantiles: + _compute_virtual_index(n, quantiles, 1 / 3.0, 1 / 3.0), + fix_gamma=lambda gamma, _: gamma, + ), + normal_unbiased=dict( + get_virtual_index=lambda n, quantiles: + _compute_virtual_index(n, quantiles, 3 / 8.0, 3 / 8.0), + fix_gamma=lambda gamma, _: gamma, + ), + # --- OTHER METHODS + lower=dict( + get_virtual_index=lambda n, quantiles: np.floor( + (n - 1) * quantiles).astype(np.intp), + fix_gamma=None, # should never be called, index dtype is int + ), + higher=dict( + get_virtual_index=lambda n, quantiles: np.ceil( + (n - 1) * quantiles).astype(np.intp), + fix_gamma=None, # should never be called, index dtype is int + ), + midpoint=dict( + get_virtual_index=lambda n, quantiles: 0.5 * ( + np.floor((n - 1) * quantiles) + + np.ceil((n - 1) * quantiles)), + fix_gamma=lambda gamma, index: _get_gamma_mask( + shape=gamma.shape, + default_value=0.5, + conditioned_value=0., + where=index % 1 == 0), + ), + nearest=dict( + get_virtual_index=lambda n, quantiles: np.around( + (n - 1) * quantiles).astype(np.intp), + fix_gamma=None, + # should never be called, index dtype is int + )) + + +def _rot90_dispatcher(m, k=None, axes=None): + return (m,) + + +@array_function_dispatch(_rot90_dispatcher) +def rot90(m, k=1, axes=(0, 1)): + """ + Rotate an array by 90 degrees in the plane specified by axes. + + Rotation direction is from the first towards the second axis. + This means for a 2D array with the default `k` and `axes`, the + rotation will be counterclockwise. + + Parameters + ---------- + m : array_like + Array of two or more dimensions. + k : integer + Number of times the array is rotated by 90 degrees. + axes : (2,) array_like + The array is rotated in the plane defined by the axes. + Axes must be different. + + Returns + ------- + y : ndarray + A rotated view of `m`. + + See Also + -------- + flip : Reverse the order of elements in an array along the given axis. + fliplr : Flip an array horizontally. + flipud : Flip an array vertically. + + Notes + ----- + ``rot90(m, k=1, axes=(1,0))`` is the reverse of + ``rot90(m, k=1, axes=(0,1))`` + + ``rot90(m, k=1, axes=(1,0))`` is equivalent to + ``rot90(m, k=-1, axes=(0,1))`` + + Examples + -------- + >>> import numpy as np + >>> m = np.array([[1,2],[3,4]], int) + >>> m + array([[1, 2], + [3, 4]]) + >>> np.rot90(m) + array([[2, 4], + [1, 3]]) + >>> np.rot90(m, 2) + array([[4, 3], + [2, 1]]) + >>> m = np.arange(8).reshape((2,2,2)) + >>> np.rot90(m, 1, (1,2)) + array([[[1, 3], + [0, 2]], + [[5, 7], + [4, 6]]]) + + """ + axes = tuple(axes) + if len(axes) != 2: + raise ValueError("len(axes) must be 2.") + + m = asanyarray(m) + + if axes[0] == axes[1] or absolute(axes[0] - axes[1]) == m.ndim: + raise ValueError("Axes must be different.") + + if (axes[0] >= m.ndim or axes[0] < -m.ndim + or axes[1] >= m.ndim or axes[1] < -m.ndim): + raise ValueError("Axes={} out of range for array of ndim={}." + .format(axes, m.ndim)) + + k %= 4 + + if k == 0: + return m[:] + if k == 2: + return flip(flip(m, axes[0]), axes[1]) + + axes_list = arange(0, m.ndim) + (axes_list[axes[0]], axes_list[axes[1]]) = (axes_list[axes[1]], + axes_list[axes[0]]) + + if k == 1: + return transpose(flip(m, axes[1]), axes_list) + else: + # k == 3 + return flip(transpose(m, axes_list), axes[1]) + + +def _flip_dispatcher(m, axis=None): + return (m,) + + +@array_function_dispatch(_flip_dispatcher) +def flip(m, axis=None): + """ + Reverse the order of elements in an array along the given axis. + + The shape of the array is preserved, but the elements are reordered. + + Parameters + ---------- + m : array_like + Input array. + axis : None or int or tuple of ints, optional + Axis or axes along which to flip over. The default, + axis=None, will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + + Returns + ------- + out : array_like + A view of `m` with the entries of axis reversed. Since a view is + returned, this operation is done in constant time. + + See Also + -------- + flipud : Flip an array vertically (axis=0). + fliplr : Flip an array horizontally (axis=1). + + Notes + ----- + flip(m, 0) is equivalent to flipud(m). + + flip(m, 1) is equivalent to fliplr(m). + + flip(m, n) corresponds to ``m[...,::-1,...]`` with ``::-1`` at position n. + + flip(m) corresponds to ``m[::-1,::-1,...,::-1]`` with ``::-1`` at all + positions. + + flip(m, (0, 1)) corresponds to ``m[::-1,::-1,...]`` with ``::-1`` at + position 0 and position 1. + + Examples + -------- + >>> import numpy as np + >>> A = np.arange(8).reshape((2,2,2)) + >>> A + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + >>> np.flip(A, 0) + array([[[4, 5], + [6, 7]], + [[0, 1], + [2, 3]]]) + >>> np.flip(A, 1) + array([[[2, 3], + [0, 1]], + [[6, 7], + [4, 5]]]) + >>> np.flip(A) + array([[[7, 6], + [5, 4]], + [[3, 2], + [1, 0]]]) + >>> np.flip(A, (0, 2)) + array([[[5, 4], + [7, 6]], + [[1, 0], + [3, 2]]]) + >>> rng = np.random.default_rng() + >>> A = rng.normal(size=(3,4,5)) + >>> np.all(np.flip(A,2) == A[:,:,::-1,...]) + True + """ + if not hasattr(m, 'ndim'): + m = asarray(m) + if axis is None: + indexer = (np.s_[::-1],) * m.ndim + else: + axis = _nx.normalize_axis_tuple(axis, m.ndim) + indexer = [np.s_[:]] * m.ndim + for ax in axis: + indexer[ax] = np.s_[::-1] + indexer = tuple(indexer) + return m[indexer] + + +@set_module('numpy') +def iterable(y): + """ + Check whether or not an object can be iterated over. + + Parameters + ---------- + y : object + Input object. + + Returns + ------- + b : bool + Return ``True`` if the object has an iterator method or is a + sequence and ``False`` otherwise. + + + Examples + -------- + >>> import numpy as np + >>> np.iterable([1, 2, 3]) + True + >>> np.iterable(2) + False + + Notes + ----- + In most cases, the results of ``np.iterable(obj)`` are consistent with + ``isinstance(obj, collections.abc.Iterable)``. One notable exception is + the treatment of 0-dimensional arrays:: + + >>> from collections.abc import Iterable + >>> a = np.array(1.0) # 0-dimensional numpy array + >>> isinstance(a, Iterable) + True + >>> np.iterable(a) + False + + """ + try: + iter(y) + except TypeError: + return False + return True + + +def _weights_are_valid(weights, a, axis): + """Validate weights array. + + We assume, weights is not None. + """ + wgt = np.asanyarray(weights) + + # Sanity checks + if a.shape != wgt.shape: + if axis is None: + raise TypeError( + "Axis must be specified when shapes of a and weights " + "differ.") + if wgt.shape != tuple(a.shape[ax] for ax in axis): + raise ValueError( + "Shape of weights must be consistent with " + "shape of a along specified axis.") + + # setup wgt to broadcast along axis + wgt = wgt.transpose(np.argsort(axis)) + wgt = wgt.reshape(tuple((s if ax in axis else 1) + for ax, s in enumerate(a.shape))) + return wgt + + +def _average_dispatcher(a, axis=None, weights=None, returned=None, *, + keepdims=None): + return (a, weights) + + +@array_function_dispatch(_average_dispatcher) +def average(a, axis=None, weights=None, returned=False, *, + keepdims=np._NoValue): + """ + Compute the weighted average along the specified axis. + + Parameters + ---------- + a : array_like + Array containing data to be averaged. If `a` is not an array, a + conversion is attempted. + axis : None or int or tuple of ints, optional + Axis or axes along which to average `a`. The default, + `axis=None`, will average over all of the elements of the input array. + If axis is negative it counts from the last to the first axis. + If axis is a tuple of ints, averaging is performed on all of the axes + specified in the tuple instead of a single axis or all the axes as + before. + weights : array_like, optional + An array of weights associated with the values in `a`. Each value in + `a` contributes to the average according to its associated weight. + The array of weights must be the same shape as `a` if no axis is + specified, otherwise the weights must have dimensions and shape + consistent with `a` along the specified axis. + If `weights=None`, then all data in `a` are assumed to have a + weight equal to one. + The calculation is:: + + avg = sum(a * weights) / sum(weights) + + where the sum is over all included elements. + The only constraint on the values of `weights` is that `sum(weights)` + must not be 0. + returned : bool, optional + Default is `False`. If `True`, the tuple (`average`, `sum_of_weights`) + is returned, otherwise only the average is returned. + If `weights=None`, `sum_of_weights` is equivalent to the number of + elements over which the average is taken. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + *Note:* `keepdims` will not work with instances of `numpy.matrix` + or other classes whose methods do not support `keepdims`. + + .. versionadded:: 1.23.0 + + Returns + ------- + retval, [sum_of_weights] : array_type or double + Return the average along the specified axis. When `returned` is `True`, + return a tuple with the average as the first element and the sum + of the weights as the second element. `sum_of_weights` is of the + same type as `retval`. The result dtype follows a general pattern. + If `weights` is None, the result dtype will be that of `a` , or ``float64`` + if `a` is integral. Otherwise, if `weights` is not None and `a` is non- + integral, the result type will be the type of lowest precision capable of + representing values of both `a` and `weights`. If `a` happens to be + integral, the previous rules still applies but the result dtype will + at least be ``float64``. + + Raises + ------ + ZeroDivisionError + When all weights along axis are zero. See `numpy.ma.average` for a + version robust to this type of error. + TypeError + When `weights` does not have the same shape as `a`, and `axis=None`. + ValueError + When `weights` does not have dimensions and shape consistent with `a` + along specified `axis`. + + See Also + -------- + mean + + ma.average : average for masked arrays -- useful if your data contains + "missing" values + numpy.result_type : Returns the type that results from applying the + numpy type promotion rules to the arguments. + + Examples + -------- + >>> import numpy as np + >>> data = np.arange(1, 5) + >>> data + array([1, 2, 3, 4]) + >>> np.average(data) + 2.5 + >>> np.average(np.arange(1, 11), weights=np.arange(10, 0, -1)) + 4.0 + + >>> data = np.arange(6).reshape((3, 2)) + >>> data + array([[0, 1], + [2, 3], + [4, 5]]) + >>> np.average(data, axis=1, weights=[1./4, 3./4]) + array([0.75, 2.75, 4.75]) + >>> np.average(data, weights=[1./4, 3./4]) + Traceback (most recent call last): + ... + TypeError: Axis must be specified when shapes of a and weights differ. + + With ``keepdims=True``, the following result has shape (3, 1). + + >>> np.average(data, axis=1, keepdims=True) + array([[0.5], + [2.5], + [4.5]]) + + >>> data = np.arange(8).reshape((2, 2, 2)) + >>> data + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + >>> np.average(data, axis=(0, 1), weights=[[1./4, 3./4], [1., 1./2]]) + array([3.4, 4.4]) + >>> np.average(data, axis=0, weights=[[1./4, 3./4], [1., 1./2]]) + Traceback (most recent call last): + ... + ValueError: Shape of weights must be consistent + with shape of a along specified axis. + """ + a = np.asanyarray(a) + + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, a.ndim, argname="axis") + + if keepdims is np._NoValue: + # Don't pass on the keepdims argument if one wasn't given. + keepdims_kw = {} + else: + keepdims_kw = {'keepdims': keepdims} + + if weights is None: + avg = a.mean(axis, **keepdims_kw) + avg_as_array = np.asanyarray(avg) + scl = avg_as_array.dtype.type(a.size/avg_as_array.size) + else: + wgt = _weights_are_valid(weights=weights, a=a, axis=axis) + + if issubclass(a.dtype.type, (np.integer, np.bool)): + result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8') + else: + result_dtype = np.result_type(a.dtype, wgt.dtype) + + scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) + if np.any(scl == 0.0): + raise ZeroDivisionError( + "Weights sum to zero, can't be normalized") + + avg = avg_as_array = np.multiply(a, wgt, + dtype=result_dtype).sum(axis, **keepdims_kw) / scl + + if returned: + if scl.shape != avg_as_array.shape: + scl = np.broadcast_to(scl, avg_as_array.shape).copy() + return avg, scl + else: + return avg + + +@set_module('numpy') +def asarray_chkfinite(a, dtype=None, order=None): + """Convert the input to an array, checking for NaNs or Infs. + + Parameters + ---------- + a : array_like + Input data, in any form that can be converted to an array. This + includes lists, lists of tuples, tuples, tuples of tuples, tuples + of lists and ndarrays. Success requires no NaNs or Infs. + dtype : data-type, optional + By default, the data-type is inferred from the input data. + order : {'C', 'F', 'A', 'K'}, optional + Memory layout. 'A' and 'K' depend on the order of input array a. + 'C' row-major (C-style), + 'F' column-major (Fortran-style) memory representation. + 'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise + 'K' (keep) preserve input order + Defaults to 'C'. + + Returns + ------- + out : ndarray + Array interpretation of `a`. No copy is performed if the input + is already an ndarray. If `a` is a subclass of ndarray, a base + class ndarray is returned. + + Raises + ------ + ValueError + Raises ValueError if `a` contains NaN (Not a Number) or Inf (Infinity). + + See Also + -------- + asarray : Create and array. + asanyarray : Similar function which passes through subclasses. + ascontiguousarray : Convert input to a contiguous array. + asfortranarray : Convert input to an ndarray with column-major + memory order. + fromiter : Create an array from an iterator. + fromfunction : Construct an array by executing a function on grid + positions. + + Examples + -------- + >>> import numpy as np + + Convert a list into an array. If all elements are finite, then + ``asarray_chkfinite`` is identical to ``asarray``. + + >>> a = [1, 2] + >>> np.asarray_chkfinite(a, dtype=float) + array([1., 2.]) + + Raises ValueError if array_like contains Nans or Infs. + + >>> a = [1, 2, np.inf] + >>> try: + ... np.asarray_chkfinite(a) + ... except ValueError: + ... print('ValueError') + ... + ValueError + + """ + a = asarray(a, dtype=dtype, order=order) + if a.dtype.char in typecodes['AllFloat'] and not np.isfinite(a).all(): + raise ValueError( + "array must not contain infs or NaNs") + return a + + +def _piecewise_dispatcher(x, condlist, funclist, *args, **kw): + yield x + # support the undocumented behavior of allowing scalars + if np.iterable(condlist): + yield from condlist + + +@array_function_dispatch(_piecewise_dispatcher) +def piecewise(x, condlist, funclist, *args, **kw): + """ + Evaluate a piecewise-defined function. + + Given a set of conditions and corresponding functions, evaluate each + function on the input data wherever its condition is true. + + Parameters + ---------- + x : ndarray or scalar + The input domain. + condlist : list of bool arrays or bool scalars + Each boolean array corresponds to a function in `funclist`. Wherever + `condlist[i]` is True, `funclist[i](x)` is used as the output value. + + Each boolean array in `condlist` selects a piece of `x`, + and should therefore be of the same shape as `x`. + + The length of `condlist` must correspond to that of `funclist`. + If one extra function is given, i.e. if + ``len(funclist) == len(condlist) + 1``, then that extra function + is the default value, used wherever all conditions are false. + funclist : list of callables, f(x,*args,**kw), or scalars + Each function is evaluated over `x` wherever its corresponding + condition is True. It should take a 1d array as input and give an 1d + array or a scalar value as output. If, instead of a callable, + a scalar is provided then a constant function (``lambda x: scalar``) is + assumed. + args : tuple, optional + Any further arguments given to `piecewise` are passed to the functions + upon execution, i.e., if called ``piecewise(..., ..., 1, 'a')``, then + each function is called as ``f(x, 1, 'a')``. + kw : dict, optional + Keyword arguments used in calling `piecewise` are passed to the + functions upon execution, i.e., if called + ``piecewise(..., ..., alpha=1)``, then each function is called as + ``f(x, alpha=1)``. + + Returns + ------- + out : ndarray + The output is the same shape and type as x and is found by + calling the functions in `funclist` on the appropriate portions of `x`, + as defined by the boolean arrays in `condlist`. Portions not covered + by any condition have a default value of 0. + + + See Also + -------- + choose, select, where + + Notes + ----- + This is similar to choose or select, except that functions are + evaluated on elements of `x` that satisfy the corresponding condition from + `condlist`. + + The result is:: + + |-- + |funclist[0](x[condlist[0]]) + out = |funclist[1](x[condlist[1]]) + |... + |funclist[n2](x[condlist[n2]]) + |-- + + Examples + -------- + >>> import numpy as np + + Define the signum function, which is -1 for ``x < 0`` and +1 for ``x >= 0``. + + >>> x = np.linspace(-2.5, 2.5, 6) + >>> np.piecewise(x, [x < 0, x >= 0], [-1, 1]) + array([-1., -1., -1., 1., 1., 1.]) + + Define the absolute value, which is ``-x`` for ``x <0`` and ``x`` for + ``x >= 0``. + + >>> np.piecewise(x, [x < 0, x >= 0], [lambda x: -x, lambda x: x]) + array([2.5, 1.5, 0.5, 0.5, 1.5, 2.5]) + + Apply the same function to a scalar value. + + >>> y = -2 + >>> np.piecewise(y, [y < 0, y >= 0], [lambda x: -x, lambda x: x]) + array(2) + + """ + x = asanyarray(x) + n2 = len(funclist) + + # undocumented: single condition is promoted to a list of one condition + if isscalar(condlist) or ( + not isinstance(condlist[0], (list, ndarray)) and x.ndim != 0): + condlist = [condlist] + + condlist = asarray(condlist, dtype=bool) + n = len(condlist) + + if n == n2 - 1: # compute the "otherwise" condition. + condelse = ~np.any(condlist, axis=0, keepdims=True) + condlist = np.concatenate([condlist, condelse], axis=0) + n += 1 + elif n != n2: + raise ValueError( + "with {} condition(s), either {} or {} functions are expected" + .format(n, n, n+1) + ) + + y = zeros_like(x) + for cond, func in zip(condlist, funclist): + if not isinstance(func, collections.abc.Callable): + y[cond] = func + else: + vals = x[cond] + if vals.size > 0: + y[cond] = func(vals, *args, **kw) + + return y + + +def _select_dispatcher(condlist, choicelist, default=None): + yield from condlist + yield from choicelist + + +@array_function_dispatch(_select_dispatcher) +def select(condlist, choicelist, default=0): + """ + Return an array drawn from elements in choicelist, depending on conditions. + + Parameters + ---------- + condlist : list of bool ndarrays + The list of conditions which determine from which array in `choicelist` + the output elements are taken. When multiple conditions are satisfied, + the first one encountered in `condlist` is used. + choicelist : list of ndarrays + The list of arrays from which the output elements are taken. It has + to be of the same length as `condlist`. + default : scalar, optional + The element inserted in `output` when all conditions evaluate to False. + + Returns + ------- + output : ndarray + The output at position m is the m-th element of the array in + `choicelist` where the m-th element of the corresponding array in + `condlist` is True. + + See Also + -------- + where : Return elements from one of two arrays depending on condition. + take, choose, compress, diag, diagonal + + Examples + -------- + >>> import numpy as np + + Beginning with an array of integers from 0 to 5 (inclusive), + elements less than ``3`` are negated, elements greater than ``3`` + are squared, and elements not meeting either of these conditions + (exactly ``3``) are replaced with a `default` value of ``42``. + + >>> x = np.arange(6) + >>> condlist = [x<3, x>3] + >>> choicelist = [x, x**2] + >>> np.select(condlist, choicelist, 42) + array([ 0, 1, 2, 42, 16, 25]) + + When multiple conditions are satisfied, the first one encountered in + `condlist` is used. + + >>> condlist = [x<=4, x>3] + >>> choicelist = [x, x**2] + >>> np.select(condlist, choicelist, 55) + array([ 0, 1, 2, 3, 4, 25]) + + """ + # Check the size of condlist and choicelist are the same, or abort. + if len(condlist) != len(choicelist): + raise ValueError( + 'list of cases must be same length as list of conditions') + + # Now that the dtype is known, handle the deprecated select([], []) case + if len(condlist) == 0: + raise ValueError("select with an empty condition list is not possible") + + # TODO: This preserves the Python int, float, complex manually to get the + # right `result_type` with NEP 50. Most likely we will grow a better + # way to spell this (and this can be replaced). + choicelist = [ + choice if type(choice) in (int, float, complex) else np.asarray(choice) + for choice in choicelist] + choicelist.append(default if type(default) in (int, float, complex) + else np.asarray(default)) + + try: + dtype = np.result_type(*choicelist) + except TypeError as e: + msg = f'Choicelist and default value do not have a common dtype: {e}' + raise TypeError(msg) from None + + # Convert conditions to arrays and broadcast conditions and choices + # as the shape is needed for the result. Doing it separately optimizes + # for example when all choices are scalars. + condlist = np.broadcast_arrays(*condlist) + choicelist = np.broadcast_arrays(*choicelist) + + # If cond array is not an ndarray in boolean format or scalar bool, abort. + for i, cond in enumerate(condlist): + if cond.dtype.type is not np.bool: + raise TypeError( + 'invalid entry {} in condlist: should be boolean ndarray'.format(i)) + + if choicelist[0].ndim == 0: + # This may be common, so avoid the call. + result_shape = condlist[0].shape + else: + result_shape = np.broadcast_arrays(condlist[0], choicelist[0])[0].shape + + result = np.full(result_shape, choicelist[-1], dtype) + + # Use np.copyto to burn each choicelist array onto result, using the + # corresponding condlist as a boolean mask. This is done in reverse + # order since the first choice should take precedence. + choicelist = choicelist[-2::-1] + condlist = condlist[::-1] + for choice, cond in zip(choicelist, condlist): + np.copyto(result, choice, where=cond) + + return result + + +def _copy_dispatcher(a, order=None, subok=None): + return (a,) + + +@array_function_dispatch(_copy_dispatcher) +def copy(a, order='K', subok=False): + """ + Return an array copy of the given object. + + Parameters + ---------- + a : array_like + Input data. + order : {'C', 'F', 'A', 'K'}, optional + Controls the memory layout of the copy. 'C' means C-order, + 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous, + 'C' otherwise. 'K' means match the layout of `a` as closely + as possible. (Note that this function and :meth:`ndarray.copy` are very + similar, but have different default values for their order= + arguments.) + subok : bool, optional + If True, then sub-classes will be passed-through, otherwise the + returned array will be forced to be a base-class array (defaults to False). + + Returns + ------- + arr : ndarray + Array interpretation of `a`. + + See Also + -------- + ndarray.copy : Preferred method for creating an array copy + + Notes + ----- + This is equivalent to: + + >>> np.array(a, copy=True) #doctest: +SKIP + + The copy made of the data is shallow, i.e., for arrays with object dtype, + the new array will point to the same objects. + See Examples from `ndarray.copy`. + + Examples + -------- + >>> import numpy as np + + Create an array x, with a reference y and a copy z: + + >>> x = np.array([1, 2, 3]) + >>> y = x + >>> z = np.copy(x) + + Note that, when we modify x, y changes, but not z: + + >>> x[0] = 10 + >>> x[0] == y[0] + True + >>> x[0] == z[0] + False + + Note that, np.copy clears previously set WRITEABLE=False flag. + + >>> a = np.array([1, 2, 3]) + >>> a.flags["WRITEABLE"] = False + >>> b = np.copy(a) + >>> b.flags["WRITEABLE"] + True + >>> b[0] = 3 + >>> b + array([3, 2, 3]) + """ + return array(a, order=order, subok=subok, copy=True) + +# Basic operations + + +def _gradient_dispatcher(f, *varargs, axis=None, edge_order=None): + yield f + yield from varargs + + +@array_function_dispatch(_gradient_dispatcher) +def gradient(f, *varargs, axis=None, edge_order=1): + """ + Return the gradient of an N-dimensional array. + + The gradient is computed using second order accurate central differences + in the interior points and either first or second order accurate one-sides + (forward or backwards) differences at the boundaries. + The returned gradient hence has the same shape as the input array. + + Parameters + ---------- + f : array_like + An N-dimensional array containing samples of a scalar function. + varargs : list of scalar or array, optional + Spacing between f values. Default unitary spacing for all dimensions. + Spacing can be specified using: + + 1. single scalar to specify a sample distance for all dimensions. + 2. N scalars to specify a constant sample distance for each dimension. + i.e. `dx`, `dy`, `dz`, ... + 3. N arrays to specify the coordinates of the values along each + dimension of F. The length of the array must match the size of + the corresponding dimension + 4. Any combination of N scalars/arrays with the meaning of 2. and 3. + + If `axis` is given, the number of varargs must equal the number of axes. + Default: 1. (see Examples below). + + edge_order : {1, 2}, optional + Gradient is calculated using N-th order accurate differences + at the boundaries. Default: 1. + axis : None or int or tuple of ints, optional + Gradient is calculated only along the given axis or axes + The default (axis = None) is to calculate the gradient for all the axes + of the input array. axis may be negative, in which case it counts from + the last to the first axis. + + Returns + ------- + gradient : ndarray or tuple of ndarray + A tuple of ndarrays (or a single ndarray if there is only one + dimension) corresponding to the derivatives of f with respect + to each dimension. Each derivative has the same shape as f. + + Examples + -------- + >>> import numpy as np + >>> f = np.array([1, 2, 4, 7, 11, 16]) + >>> np.gradient(f) + array([1. , 1.5, 2.5, 3.5, 4.5, 5. ]) + >>> np.gradient(f, 2) + array([0.5 , 0.75, 1.25, 1.75, 2.25, 2.5 ]) + + Spacing can be also specified with an array that represents the coordinates + of the values F along the dimensions. + For instance a uniform spacing: + + >>> x = np.arange(f.size) + >>> np.gradient(f, x) + array([1. , 1.5, 2.5, 3.5, 4.5, 5. ]) + + Or a non uniform one: + + >>> x = np.array([0., 1., 1.5, 3.5, 4., 6.]) + >>> np.gradient(f, x) + array([1. , 3. , 3.5, 6.7, 6.9, 2.5]) + + For two dimensional arrays, the return will be two arrays ordered by + axis. In this example the first array stands for the gradient in + rows and the second one in columns direction: + + >>> np.gradient(np.array([[1, 2, 6], [3, 4, 5]])) + (array([[ 2., 2., -1.], + [ 2., 2., -1.]]), + array([[1. , 2.5, 4. ], + [1. , 1. , 1. ]])) + + In this example the spacing is also specified: + uniform for axis=0 and non uniform for axis=1 + + >>> dx = 2. + >>> y = [1., 1.5, 3.5] + >>> np.gradient(np.array([[1, 2, 6], [3, 4, 5]]), dx, y) + (array([[ 1. , 1. , -0.5], + [ 1. , 1. , -0.5]]), + array([[2. , 2. , 2. ], + [2. , 1.7, 0.5]])) + + It is possible to specify how boundaries are treated using `edge_order` + + >>> x = np.array([0, 1, 2, 3, 4]) + >>> f = x**2 + >>> np.gradient(f, edge_order=1) + array([1., 2., 4., 6., 7.]) + >>> np.gradient(f, edge_order=2) + array([0., 2., 4., 6., 8.]) + + The `axis` keyword can be used to specify a subset of axes of which the + gradient is calculated + + >>> np.gradient(np.array([[1, 2, 6], [3, 4, 5]]), axis=0) + array([[ 2., 2., -1.], + [ 2., 2., -1.]]) + + The `varargs` argument defines the spacing between sample points in the + input array. It can take two forms: + + 1. An array, specifying coordinates, which may be unevenly spaced: + + >>> x = np.array([0., 2., 3., 6., 8.]) + >>> y = x ** 2 + >>> np.gradient(y, x, edge_order=2) + array([ 0., 4., 6., 12., 16.]) + + 2. A scalar, representing the fixed sample distance: + + >>> dx = 2 + >>> x = np.array([0., 2., 4., 6., 8.]) + >>> y = x ** 2 + >>> np.gradient(y, dx, edge_order=2) + array([ 0., 4., 8., 12., 16.]) + + It's possible to provide different data for spacing along each dimension. + The number of arguments must match the number of dimensions in the input + data. + + >>> dx = 2 + >>> dy = 3 + >>> x = np.arange(0, 6, dx) + >>> y = np.arange(0, 9, dy) + >>> xs, ys = np.meshgrid(x, y) + >>> zs = xs + 2 * ys + >>> np.gradient(zs, dy, dx) # Passing two scalars + (array([[2., 2., 2.], + [2., 2., 2.], + [2., 2., 2.]]), + array([[1., 1., 1.], + [1., 1., 1.], + [1., 1., 1.]])) + + Mixing scalars and arrays is also allowed: + + >>> np.gradient(zs, y, dx) # Passing one array and one scalar + (array([[2., 2., 2.], + [2., 2., 2.], + [2., 2., 2.]]), + array([[1., 1., 1.], + [1., 1., 1.], + [1., 1., 1.]])) + + Notes + ----- + Assuming that :math:`f\\in C^{3}` (i.e., :math:`f` has at least 3 continuous + derivatives) and let :math:`h_{*}` be a non-homogeneous stepsize, we + minimize the "consistency error" :math:`\\eta_{i}` between the true gradient + and its estimate from a linear combination of the neighboring grid-points: + + .. math:: + + \\eta_{i} = f_{i}^{\\left(1\\right)} - + \\left[ \\alpha f\\left(x_{i}\\right) + + \\beta f\\left(x_{i} + h_{d}\\right) + + \\gamma f\\left(x_{i}-h_{s}\\right) + \\right] + + By substituting :math:`f(x_{i} + h_{d})` and :math:`f(x_{i} - h_{s})` + with their Taylor series expansion, this translates into solving + the following the linear system: + + .. math:: + + \\left\\{ + \\begin{array}{r} + \\alpha+\\beta+\\gamma=0 \\\\ + \\beta h_{d}-\\gamma h_{s}=1 \\\\ + \\beta h_{d}^{2}+\\gamma h_{s}^{2}=0 + \\end{array} + \\right. + + The resulting approximation of :math:`f_{i}^{(1)}` is the following: + + .. math:: + + \\hat f_{i}^{(1)} = + \\frac{ + h_{s}^{2}f\\left(x_{i} + h_{d}\\right) + + \\left(h_{d}^{2} - h_{s}^{2}\\right)f\\left(x_{i}\\right) + - h_{d}^{2}f\\left(x_{i}-h_{s}\\right)} + { h_{s}h_{d}\\left(h_{d} + h_{s}\\right)} + + \\mathcal{O}\\left(\\frac{h_{d}h_{s}^{2} + + h_{s}h_{d}^{2}}{h_{d} + + h_{s}}\\right) + + It is worth noting that if :math:`h_{s}=h_{d}` + (i.e., data are evenly spaced) + we find the standard second order approximation: + + .. math:: + + \\hat f_{i}^{(1)}= + \\frac{f\\left(x_{i+1}\\right) - f\\left(x_{i-1}\\right)}{2h} + + \\mathcal{O}\\left(h^{2}\\right) + + With a similar procedure the forward/backward approximations used for + boundaries can be derived. + + References + ---------- + .. [1] Quarteroni A., Sacco R., Saleri F. (2007) Numerical Mathematics + (Texts in Applied Mathematics). New York: Springer. + .. [2] Durran D. R. (1999) Numerical Methods for Wave Equations + in Geophysical Fluid Dynamics. New York: Springer. + .. [3] Fornberg B. (1988) Generation of Finite Difference Formulas on + Arbitrarily Spaced Grids, + Mathematics of Computation 51, no. 184 : 699-706. + `PDF `_. + """ + f = np.asanyarray(f) + N = f.ndim # number of dimensions + + if axis is None: + axes = tuple(range(N)) + else: + axes = _nx.normalize_axis_tuple(axis, N) + + len_axes = len(axes) + n = len(varargs) + if n == 0: + # no spacing argument - use 1 in all axes + dx = [1.0] * len_axes + elif n == 1 and np.ndim(varargs[0]) == 0: + # single scalar for all axes + dx = varargs * len_axes + elif n == len_axes: + # scalar or 1d array for each axis + dx = list(varargs) + for i, distances in enumerate(dx): + distances = np.asanyarray(distances) + if distances.ndim == 0: + continue + elif distances.ndim != 1: + raise ValueError("distances must be either scalars or 1d") + if len(distances) != f.shape[axes[i]]: + raise ValueError("when 1d, distances must match " + "the length of the corresponding dimension") + if np.issubdtype(distances.dtype, np.integer): + # Convert numpy integer types to float64 to avoid modular + # arithmetic in np.diff(distances). + distances = distances.astype(np.float64) + diffx = np.diff(distances) + # if distances are constant reduce to the scalar case + # since it brings a consistent speedup + if (diffx == diffx[0]).all(): + diffx = diffx[0] + dx[i] = diffx + else: + raise TypeError("invalid number of arguments") + + if edge_order > 2: + raise ValueError("'edge_order' greater than 2 not supported") + + # use central differences on interior and one-sided differences on the + # endpoints. This preserves second order-accuracy over the full domain. + + outvals = [] + + # create slice objects --- initially all are [:, :, ..., :] + slice1 = [slice(None)]*N + slice2 = [slice(None)]*N + slice3 = [slice(None)]*N + slice4 = [slice(None)]*N + + otype = f.dtype + if otype.type is np.datetime64: + # the timedelta dtype with the same unit information + otype = np.dtype(otype.name.replace('datetime', 'timedelta')) + # view as timedelta to allow addition + f = f.view(otype) + elif otype.type is np.timedelta64: + pass + elif np.issubdtype(otype, np.inexact): + pass + else: + # All other types convert to floating point. + # First check if f is a numpy integer type; if so, convert f to float64 + # to avoid modular arithmetic when computing the changes in f. + if np.issubdtype(otype, np.integer): + f = f.astype(np.float64) + otype = np.float64 + + for axis, ax_dx in zip(axes, dx): + if f.shape[axis] < edge_order + 1: + raise ValueError( + "Shape of array too small to calculate a numerical gradient, " + "at least (edge_order + 1) elements are required.") + # result allocation + out = np.empty_like(f, dtype=otype) + + # spacing for the current axis + uniform_spacing = np.ndim(ax_dx) == 0 + + # Numerical differentiation: 2nd order interior + slice1[axis] = slice(1, -1) + slice2[axis] = slice(None, -2) + slice3[axis] = slice(1, -1) + slice4[axis] = slice(2, None) + + if uniform_spacing: + out[tuple(slice1)] = (f[tuple(slice4)] - f[tuple(slice2)]) / (2. * ax_dx) + else: + dx1 = ax_dx[0:-1] + dx2 = ax_dx[1:] + a = -(dx2)/(dx1 * (dx1 + dx2)) + b = (dx2 - dx1) / (dx1 * dx2) + c = dx1 / (dx2 * (dx1 + dx2)) + # fix the shape for broadcasting + shape = np.ones(N, dtype=int) + shape[axis] = -1 + a.shape = b.shape = c.shape = shape + # 1D equivalent -- out[1:-1] = a * f[:-2] + b * f[1:-1] + c * f[2:] + out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)] + + # Numerical differentiation: 1st order edges + if edge_order == 1: + slice1[axis] = 0 + slice2[axis] = 1 + slice3[axis] = 0 + dx_0 = ax_dx if uniform_spacing else ax_dx[0] + # 1D equivalent -- out[0] = (f[1] - f[0]) / (x[1] - x[0]) + out[tuple(slice1)] = (f[tuple(slice2)] - f[tuple(slice3)]) / dx_0 + + slice1[axis] = -1 + slice2[axis] = -1 + slice3[axis] = -2 + dx_n = ax_dx if uniform_spacing else ax_dx[-1] + # 1D equivalent -- out[-1] = (f[-1] - f[-2]) / (x[-1] - x[-2]) + out[tuple(slice1)] = (f[tuple(slice2)] - f[tuple(slice3)]) / dx_n + + # Numerical differentiation: 2nd order edges + else: + slice1[axis] = 0 + slice2[axis] = 0 + slice3[axis] = 1 + slice4[axis] = 2 + if uniform_spacing: + a = -1.5 / ax_dx + b = 2. / ax_dx + c = -0.5 / ax_dx + else: + dx1 = ax_dx[0] + dx2 = ax_dx[1] + a = -(2. * dx1 + dx2)/(dx1 * (dx1 + dx2)) + b = (dx1 + dx2) / (dx1 * dx2) + c = - dx1 / (dx2 * (dx1 + dx2)) + # 1D equivalent -- out[0] = a * f[0] + b * f[1] + c * f[2] + out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)] + + slice1[axis] = -1 + slice2[axis] = -3 + slice3[axis] = -2 + slice4[axis] = -1 + if uniform_spacing: + a = 0.5 / ax_dx + b = -2. / ax_dx + c = 1.5 / ax_dx + else: + dx1 = ax_dx[-2] + dx2 = ax_dx[-1] + a = (dx2) / (dx1 * (dx1 + dx2)) + b = - (dx2 + dx1) / (dx1 * dx2) + c = (2. * dx2 + dx1) / (dx2 * (dx1 + dx2)) + # 1D equivalent -- out[-1] = a * f[-3] + b * f[-2] + c * f[-1] + out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)] + + outvals.append(out) + + # reset the slice object in this dimension to ":" + slice1[axis] = slice(None) + slice2[axis] = slice(None) + slice3[axis] = slice(None) + slice4[axis] = slice(None) + + if len_axes == 1: + return outvals[0] + return tuple(outvals) + + +def _diff_dispatcher(a, n=None, axis=None, prepend=None, append=None): + return (a, prepend, append) + + +@array_function_dispatch(_diff_dispatcher) +def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): + """ + Calculate the n-th discrete difference along the given axis. + + The first difference is given by ``out[i] = a[i+1] - a[i]`` along + the given axis, higher differences are calculated by using `diff` + recursively. + + Parameters + ---------- + a : array_like + Input array + n : int, optional + The number of times values are differenced. If zero, the input + is returned as-is. + axis : int, optional + The axis along which the difference is taken, default is the + last axis. + prepend, append : array_like, optional + Values to prepend or append to `a` along axis prior to + performing the difference. Scalar values are expanded to + arrays with length 1 in the direction of axis and the shape + of the input array in along all other axes. Otherwise the + dimension and shape must match `a` except along axis. + + Returns + ------- + diff : ndarray + The n-th differences. The shape of the output is the same as `a` + except along `axis` where the dimension is smaller by `n`. The + type of the output is the same as the type of the difference + between any two elements of `a`. This is the same as the type of + `a` in most cases. A notable exception is `datetime64`, which + results in a `timedelta64` output array. + + See Also + -------- + gradient, ediff1d, cumsum + + Notes + ----- + Type is preserved for boolean arrays, so the result will contain + `False` when consecutive elements are the same and `True` when they + differ. + + For unsigned integer arrays, the results will also be unsigned. This + should not be surprising, as the result is consistent with + calculating the difference directly: + + >>> u8_arr = np.array([1, 0], dtype=np.uint8) + >>> np.diff(u8_arr) + array([255], dtype=uint8) + >>> u8_arr[1,...] - u8_arr[0,...] + np.uint8(255) + + If this is not desirable, then the array should be cast to a larger + integer type first: + + >>> i16_arr = u8_arr.astype(np.int16) + >>> np.diff(i16_arr) + array([-1], dtype=int16) + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2, 4, 7, 0]) + >>> np.diff(x) + array([ 1, 2, 3, -7]) + >>> np.diff(x, n=2) + array([ 1, 1, -10]) + + >>> x = np.array([[1, 3, 6, 10], [0, 5, 6, 8]]) + >>> np.diff(x) + array([[2, 3, 4], + [5, 1, 2]]) + >>> np.diff(x, axis=0) + array([[-1, 2, 0, -2]]) + + >>> x = np.arange('1066-10-13', '1066-10-16', dtype=np.datetime64) + >>> np.diff(x) + array([1, 1], dtype='timedelta64[D]') + + """ + if n == 0: + return a + if n < 0: + raise ValueError( + "order must be non-negative but got " + repr(n)) + + a = asanyarray(a) + nd = a.ndim + if nd == 0: + raise ValueError("diff requires input that is at least one dimensional") + axis = normalize_axis_index(axis, nd) + + combined = [] + if prepend is not np._NoValue: + prepend = np.asanyarray(prepend) + if prepend.ndim == 0: + shape = list(a.shape) + shape[axis] = 1 + prepend = np.broadcast_to(prepend, tuple(shape)) + combined.append(prepend) + + combined.append(a) + + if append is not np._NoValue: + append = np.asanyarray(append) + if append.ndim == 0: + shape = list(a.shape) + shape[axis] = 1 + append = np.broadcast_to(append, tuple(shape)) + combined.append(append) + + if len(combined) > 1: + a = np.concatenate(combined, axis) + + slice1 = [slice(None)] * nd + slice2 = [slice(None)] * nd + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) + slice1 = tuple(slice1) + slice2 = tuple(slice2) + + op = not_equal if a.dtype == np.bool else subtract + for _ in range(n): + a = op(a[slice1], a[slice2]) + + return a + + +def _interp_dispatcher(x, xp, fp, left=None, right=None, period=None): + return (x, xp, fp) + + +@array_function_dispatch(_interp_dispatcher) +def interp(x, xp, fp, left=None, right=None, period=None): + """ + One-dimensional linear interpolation for monotonically increasing sample points. + + Returns the one-dimensional piecewise linear interpolant to a function + with given discrete data points (`xp`, `fp`), evaluated at `x`. + + Parameters + ---------- + x : array_like + The x-coordinates at which to evaluate the interpolated values. + + xp : 1-D sequence of floats + The x-coordinates of the data points, must be increasing if argument + `period` is not specified. Otherwise, `xp` is internally sorted after + normalizing the periodic boundaries with ``xp = xp % period``. + + fp : 1-D sequence of float or complex + The y-coordinates of the data points, same length as `xp`. + + left : optional float or complex corresponding to fp + Value to return for `x < xp[0]`, default is `fp[0]`. + + right : optional float or complex corresponding to fp + Value to return for `x > xp[-1]`, default is `fp[-1]`. + + period : None or float, optional + A period for the x-coordinates. This parameter allows the proper + interpolation of angular x-coordinates. Parameters `left` and `right` + are ignored if `period` is specified. + + Returns + ------- + y : float or complex (corresponding to fp) or ndarray + The interpolated values, same shape as `x`. + + Raises + ------ + ValueError + If `xp` and `fp` have different length + If `xp` or `fp` are not 1-D sequences + If `period == 0` + + See Also + -------- + scipy.interpolate + + Warnings + -------- + The x-coordinate sequence is expected to be increasing, but this is not + explicitly enforced. However, if the sequence `xp` is non-increasing, + interpolation results are meaningless. + + Note that, since NaN is unsortable, `xp` also cannot contain NaNs. + + A simple check for `xp` being strictly increasing is:: + + np.all(np.diff(xp) > 0) + + Examples + -------- + >>> import numpy as np + >>> xp = [1, 2, 3] + >>> fp = [3, 2, 0] + >>> np.interp(2.5, xp, fp) + 1.0 + >>> np.interp([0, 1, 1.5, 2.72, 3.14], xp, fp) + array([3. , 3. , 2.5 , 0.56, 0. ]) + >>> UNDEF = -99.0 + >>> np.interp(3.14, xp, fp, right=UNDEF) + -99.0 + + Plot an interpolant to the sine function: + + >>> x = np.linspace(0, 2*np.pi, 10) + >>> y = np.sin(x) + >>> xvals = np.linspace(0, 2*np.pi, 50) + >>> yinterp = np.interp(xvals, x, y) + >>> import matplotlib.pyplot as plt + >>> plt.plot(x, y, 'o') + [] + >>> plt.plot(xvals, yinterp, '-x') + [] + >>> plt.show() + + Interpolation with periodic x-coordinates: + + >>> x = [-180, -170, -185, 185, -10, -5, 0, 365] + >>> xp = [190, -190, 350, -350] + >>> fp = [5, 10, 3, 4] + >>> np.interp(x, xp, fp, period=360) + array([7.5 , 5. , 8.75, 6.25, 3. , 3.25, 3.5 , 3.75]) + + Complex interpolation: + + >>> x = [1.5, 4.0] + >>> xp = [2,3,5] + >>> fp = [1.0j, 0, 2+3j] + >>> np.interp(x, xp, fp) + array([0.+1.j , 1.+1.5j]) + + """ + + fp = np.asarray(fp) + + if np.iscomplexobj(fp): + interp_func = compiled_interp_complex + input_dtype = np.complex128 + else: + interp_func = compiled_interp + input_dtype = np.float64 + + if period is not None: + if period == 0: + raise ValueError("period must be a non-zero value") + period = abs(period) + left = None + right = None + + x = np.asarray(x, dtype=np.float64) + xp = np.asarray(xp, dtype=np.float64) + fp = np.asarray(fp, dtype=input_dtype) + + if xp.ndim != 1 or fp.ndim != 1: + raise ValueError("Data points must be 1-D sequences") + if xp.shape[0] != fp.shape[0]: + raise ValueError("fp and xp are not of the same length") + # normalizing periodic boundaries + x = x % period + xp = xp % period + asort_xp = np.argsort(xp) + xp = xp[asort_xp] + fp = fp[asort_xp] + xp = np.concatenate((xp[-1:]-period, xp, xp[0:1]+period)) + fp = np.concatenate((fp[-1:], fp, fp[0:1])) + + return interp_func(x, xp, fp, left, right) + + +def _angle_dispatcher(z, deg=None): + return (z,) + + +@array_function_dispatch(_angle_dispatcher) +def angle(z, deg=False): + """ + Return the angle of the complex argument. + + Parameters + ---------- + z : array_like + A complex number or sequence of complex numbers. + deg : bool, optional + Return angle in degrees if True, radians if False (default). + + Returns + ------- + angle : ndarray or scalar + The counterclockwise angle from the positive real axis on the complex + plane in the range ``(-pi, pi]``, with dtype as numpy.float64. + + See Also + -------- + arctan2 + absolute + + Notes + ----- + This function passes the imaginary and real parts of the argument to + `arctan2` to compute the result; consequently, it follows the convention + of `arctan2` when the magnitude of the argument is zero. See example. + + Examples + -------- + >>> import numpy as np + >>> np.angle([1.0, 1.0j, 1+1j]) # in radians + array([ 0. , 1.57079633, 0.78539816]) # may vary + >>> np.angle(1+1j, deg=True) # in degrees + 45.0 + >>> np.angle([0., -0., complex(0., -0.), complex(-0., -0.)]) # convention + array([ 0. , 3.14159265, -0. , -3.14159265]) + + """ + z = asanyarray(z) + if issubclass(z.dtype.type, _nx.complexfloating): + zimag = z.imag + zreal = z.real + else: + zimag = 0 + zreal = z + + a = arctan2(zimag, zreal) + if deg: + a *= 180/pi + return a + + +def _unwrap_dispatcher(p, discont=None, axis=None, *, period=None): + return (p,) + + +@array_function_dispatch(_unwrap_dispatcher) +def unwrap(p, discont=None, axis=-1, *, period=2*pi): + r""" + Unwrap by taking the complement of large deltas with respect to the period. + + This unwraps a signal `p` by changing elements which have an absolute + difference from their predecessor of more than ``max(discont, period/2)`` + to their `period`-complementary values. + + For the default case where `period` is :math:`2\pi` and `discont` is + :math:`\pi`, this unwraps a radian phase `p` such that adjacent differences + are never greater than :math:`\pi` by adding :math:`2k\pi` for some + integer :math:`k`. + + Parameters + ---------- + p : array_like + Input array. + discont : float, optional + Maximum discontinuity between values, default is ``period/2``. + Values below ``period/2`` are treated as if they were ``period/2``. + To have an effect different from the default, `discont` should be + larger than ``period/2``. + axis : int, optional + Axis along which unwrap will operate, default is the last axis. + period : float, optional + Size of the range over which the input wraps. By default, it is + ``2 pi``. + + .. versionadded:: 1.21.0 + + Returns + ------- + out : ndarray + Output array. + + See Also + -------- + rad2deg, deg2rad + + Notes + ----- + If the discontinuity in `p` is smaller than ``period/2``, + but larger than `discont`, no unwrapping is done because taking + the complement would only make the discontinuity larger. + + Examples + -------- + >>> import numpy as np + >>> phase = np.linspace(0, np.pi, num=5) + >>> phase[3:] += np.pi + >>> phase + array([ 0. , 0.78539816, 1.57079633, 5.49778714, 6.28318531]) # may vary + >>> np.unwrap(phase) + array([ 0. , 0.78539816, 1.57079633, -0.78539816, 0. ]) # may vary + >>> np.unwrap([0, 1, 2, -1, 0], period=4) + array([0, 1, 2, 3, 4]) + >>> np.unwrap([ 1, 2, 3, 4, 5, 6, 1, 2, 3], period=6) + array([1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> np.unwrap([2, 3, 4, 5, 2, 3, 4, 5], period=4) + array([2, 3, 4, 5, 6, 7, 8, 9]) + >>> phase_deg = np.mod(np.linspace(0 ,720, 19), 360) - 180 + >>> np.unwrap(phase_deg, period=360) + array([-180., -140., -100., -60., -20., 20., 60., 100., 140., + 180., 220., 260., 300., 340., 380., 420., 460., 500., + 540.]) + """ + p = asarray(p) + nd = p.ndim + dd = diff(p, axis=axis) + if discont is None: + discont = period/2 + slice1 = [slice(None, None)]*nd # full slices + slice1[axis] = slice(1, None) + slice1 = tuple(slice1) + dtype = np.result_type(dd, period) + if _nx.issubdtype(dtype, _nx.integer): + interval_high, rem = divmod(period, 2) + boundary_ambiguous = rem == 0 + else: + interval_high = period / 2 + boundary_ambiguous = True + interval_low = -interval_high + ddmod = mod(dd - interval_low, period) + interval_low + if boundary_ambiguous: + # for `mask = (abs(dd) == period/2)`, the above line made + # `ddmod[mask] == -period/2`. correct these such that + # `ddmod[mask] == sign(dd[mask])*period/2`. + _nx.copyto(ddmod, interval_high, + where=(ddmod == interval_low) & (dd > 0)) + ph_correct = ddmod - dd + _nx.copyto(ph_correct, 0, where=abs(dd) < discont) + up = array(p, copy=True, dtype=dtype) + up[slice1] = p[slice1] + ph_correct.cumsum(axis) + return up + + +def _sort_complex(a): + return (a,) + + +@array_function_dispatch(_sort_complex) +def sort_complex(a): + """ + Sort a complex array using the real part first, then the imaginary part. + + Parameters + ---------- + a : array_like + Input array + + Returns + ------- + out : complex ndarray + Always returns a sorted complex array. + + Examples + -------- + >>> import numpy as np + >>> np.sort_complex([5, 3, 6, 2, 1]) + array([1.+0.j, 2.+0.j, 3.+0.j, 5.+0.j, 6.+0.j]) + + >>> np.sort_complex([1 + 2j, 2 - 1j, 3 - 2j, 3 - 3j, 3 + 5j]) + array([1.+2.j, 2.-1.j, 3.-3.j, 3.-2.j, 3.+5.j]) + + """ + b = array(a, copy=True) + b.sort() + if not issubclass(b.dtype.type, _nx.complexfloating): + if b.dtype.char in 'bhBH': + return b.astype('F') + elif b.dtype.char == 'g': + return b.astype('G') + else: + return b.astype('D') + else: + return b + + +def _arg_trim_zeros(filt): + """Return indices of the first and last non-zero element. + + Parameters + ---------- + filt : array_like + Input array. + + Returns + ------- + start, stop : ndarray + Two arrays containing the indices of the first and last non-zero + element in each dimension. + + See also + -------- + trim_zeros + + Examples + -------- + >>> import numpy as np + >>> _arg_trim_zeros(np.array([0, 0, 1, 1, 0])) + (array([2]), array([3])) + """ + nonzero = ( + np.argwhere(filt) + if filt.dtype != np.object_ + # Historically, `trim_zeros` treats `None` in an object array + # as non-zero while argwhere doesn't, account for that + else np.argwhere(filt != 0) + ) + if nonzero.size == 0: + start = stop = np.array([], dtype=np.intp) + else: + start = nonzero.min(axis=0) + stop = nonzero.max(axis=0) + return start, stop + + +def _trim_zeros(filt, trim=None, axis=None): + return (filt,) + + +@array_function_dispatch(_trim_zeros) +def trim_zeros(filt, trim='fb', axis=None): + """Remove values along a dimension which are zero along all other. + + Parameters + ---------- + filt : array_like + Input array. + trim : {"fb", "f", "b"}, optional + A string with 'f' representing trim from front and 'b' to trim from + back. By default, zeros are trimmed on both sides. + Front and back refer to the edges of a dimension, with "front" refering + to the side with the lowest index 0, and "back" refering to the highest + index (or index -1). + axis : int or sequence, optional + If None, `filt` is cropped such, that the smallest bounding box is + returned that still contains all values which are not zero. + If an axis is specified, `filt` will be sliced in that dimension only + on the sides specified by `trim`. The remaining area will be the + smallest that still contains all values wich are not zero. + + Returns + ------- + trimmed : ndarray or sequence + The result of trimming the input. The number of dimensions and the + input data type are preserved. + + Notes + ----- + For all-zero arrays, the first axis is trimmed first. + + Examples + -------- + >>> import numpy as np + >>> a = np.array((0, 0, 0, 1, 2, 3, 0, 2, 1, 0)) + >>> np.trim_zeros(a) + array([1, 2, 3, 0, 2, 1]) + + >>> np.trim_zeros(a, trim='b') + array([0, 0, 0, ..., 0, 2, 1]) + + Multiple dimensions are supported. + + >>> b = np.array([[0, 0, 2, 3, 0, 0], + ... [0, 1, 0, 3, 0, 0], + ... [0, 0, 0, 0, 0, 0]]) + >>> np.trim_zeros(b) + array([[0, 2, 3], + [1, 0, 3]]) + + >>> np.trim_zeros(b, axis=-1) + array([[0, 2, 3], + [1, 0, 3], + [0, 0, 0]]) + + The input data type is preserved, list/tuple in means list/tuple out. + + >>> np.trim_zeros([0, 1, 2, 0]) + [1, 2] + + """ + filt_ = np.asarray(filt) + + trim = trim.lower() + if trim not in {"fb", "bf", "f", "b"}: + raise ValueError(f"unexpected character(s) in `trim`: {trim!r}") + + start, stop = _arg_trim_zeros(filt_) + stop += 1 # Adjust for slicing + + if start.size == 0: + # filt is all-zero -> assign same values to start and stop so that + # resulting slice will be empty + start = stop = np.zeros(filt_.ndim, dtype=np.intp) + else: + if 'f' not in trim: + start = (None,) * filt_.ndim + if 'b' not in trim: + stop = (None,) * filt_.ndim + + if len(start) == 1: + # filt is 1D -> don't use multi-dimensional slicing to preserve + # non-array input types + sl = slice(start[0], stop[0]) + elif axis is None: + # trim all axes + sl = tuple(slice(*x) for x in zip(start, stop)) + else: + # only trim single axis + axis = normalize_axis_index(axis, filt_.ndim) + sl = (slice(None),) * axis + (slice(start[axis], stop[axis]),) + (...,) + + trimmed = filt[sl] + return trimmed + + + +def _extract_dispatcher(condition, arr): + return (condition, arr) + + +@array_function_dispatch(_extract_dispatcher) +def extract(condition, arr): + """ + Return the elements of an array that satisfy some condition. + + This is equivalent to ``np.compress(ravel(condition), ravel(arr))``. If + `condition` is boolean ``np.extract`` is equivalent to ``arr[condition]``. + + Note that `place` does the exact opposite of `extract`. + + Parameters + ---------- + condition : array_like + An array whose nonzero or True entries indicate the elements of `arr` + to extract. + arr : array_like + Input array of the same size as `condition`. + + Returns + ------- + extract : ndarray + Rank 1 array of values from `arr` where `condition` is True. + + See Also + -------- + take, put, copyto, compress, place + + Examples + -------- + >>> import numpy as np + >>> arr = np.arange(12).reshape((3, 4)) + >>> arr + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]]) + >>> condition = np.mod(arr, 3)==0 + >>> condition + array([[ True, False, False, True], + [False, False, True, False], + [False, True, False, False]]) + >>> np.extract(condition, arr) + array([0, 3, 6, 9]) + + + If `condition` is boolean: + + >>> arr[condition] + array([0, 3, 6, 9]) + + """ + return _nx.take(ravel(arr), nonzero(ravel(condition))[0]) + + +def _place_dispatcher(arr, mask, vals): + return (arr, mask, vals) + + +@array_function_dispatch(_place_dispatcher) +def place(arr, mask, vals): + """ + Change elements of an array based on conditional and input values. + + Similar to ``np.copyto(arr, vals, where=mask)``, the difference is that + `place` uses the first N elements of `vals`, where N is the number of + True values in `mask`, while `copyto` uses the elements where `mask` + is True. + + Note that `extract` does the exact opposite of `place`. + + Parameters + ---------- + arr : ndarray + Array to put data into. + mask : array_like + Boolean mask array. Must have the same size as `a`. + vals : 1-D sequence + Values to put into `a`. Only the first N elements are used, where + N is the number of True values in `mask`. If `vals` is smaller + than N, it will be repeated, and if elements of `a` are to be masked, + this sequence must be non-empty. + + See Also + -------- + copyto, put, take, extract + + Examples + -------- + >>> import numpy as np + >>> arr = np.arange(6).reshape(2, 3) + >>> np.place(arr, arr>2, [44, 55]) + >>> arr + array([[ 0, 1, 2], + [44, 55, 44]]) + + """ + return _place(arr, mask, vals) + + +def disp(mesg, device=None, linefeed=True): + """ + Display a message on a device. + + .. deprecated:: 2.0 + Use your own printing function instead. + + Parameters + ---------- + mesg : str + Message to display. + device : object + Device to write message. If None, defaults to ``sys.stdout`` which is + very similar to ``print``. `device` needs to have ``write()`` and + ``flush()`` methods. + linefeed : bool, optional + Option whether to print a line feed or not. Defaults to True. + + Raises + ------ + AttributeError + If `device` does not have a ``write()`` or ``flush()`` method. + + Examples + -------- + >>> import numpy as np + + Besides ``sys.stdout``, a file-like object can also be used as it has + both required methods: + + >>> from io import StringIO + >>> buf = StringIO() + >>> np.disp('"Display" in a file', device=buf) + >>> buf.getvalue() + '"Display" in a file\\n' + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`disp` is deprecated, " + "use your own printing function instead. " + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + if device is None: + device = sys.stdout + if linefeed: + device.write('%s\n' % mesg) + else: + device.write('%s' % mesg) + device.flush() + return + + +# See https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html +_DIMENSION_NAME = r'\w+' +_CORE_DIMENSION_LIST = '(?:{0:}(?:,{0:})*)?'.format(_DIMENSION_NAME) +_ARGUMENT = r'\({}\)'.format(_CORE_DIMENSION_LIST) +_ARGUMENT_LIST = '{0:}(?:,{0:})*'.format(_ARGUMENT) +_SIGNATURE = '^{0:}->{0:}$'.format(_ARGUMENT_LIST) + + +def _parse_gufunc_signature(signature): + """ + Parse string signatures for a generalized universal function. + + Arguments + --------- + signature : string + Generalized universal function signature, e.g., ``(m,n),(n,p)->(m,p)`` + for ``np.matmul``. + + Returns + ------- + Tuple of input and output core dimensions parsed from the signature, each + of the form List[Tuple[str, ...]]. + """ + signature = re.sub(r'\s+', '', signature) + + if not re.match(_SIGNATURE, signature): + raise ValueError( + 'not a valid gufunc signature: {}'.format(signature)) + return tuple([tuple(re.findall(_DIMENSION_NAME, arg)) + for arg in re.findall(_ARGUMENT, arg_list)] + for arg_list in signature.split('->')) + + +def _update_dim_sizes(dim_sizes, arg, core_dims): + """ + Incrementally check and update core dimension sizes for a single argument. + + Arguments + --------- + dim_sizes : Dict[str, int] + Sizes of existing core dimensions. Will be updated in-place. + arg : ndarray + Argument to examine. + core_dims : Tuple[str, ...] + Core dimensions for this argument. + """ + if not core_dims: + return + + num_core_dims = len(core_dims) + if arg.ndim < num_core_dims: + raise ValueError( + '%d-dimensional argument does not have enough ' + 'dimensions for all core dimensions %r' + % (arg.ndim, core_dims)) + + core_shape = arg.shape[-num_core_dims:] + for dim, size in zip(core_dims, core_shape): + if dim in dim_sizes: + if size != dim_sizes[dim]: + raise ValueError( + 'inconsistent size for core dimension %r: %r vs %r' + % (dim, size, dim_sizes[dim])) + else: + dim_sizes[dim] = size + + +def _parse_input_dimensions(args, input_core_dims): + """ + Parse broadcast and core dimensions for vectorize with a signature. + + Arguments + --------- + args : Tuple[ndarray, ...] + Tuple of input arguments to examine. + input_core_dims : List[Tuple[str, ...]] + List of core dimensions corresponding to each input. + + Returns + ------- + broadcast_shape : Tuple[int, ...] + Common shape to broadcast all non-core dimensions to. + dim_sizes : Dict[str, int] + Common sizes for named core dimensions. + """ + broadcast_args = [] + dim_sizes = {} + for arg, core_dims in zip(args, input_core_dims): + _update_dim_sizes(dim_sizes, arg, core_dims) + ndim = arg.ndim - len(core_dims) + dummy_array = np.lib.stride_tricks.as_strided(0, arg.shape[:ndim]) + broadcast_args.append(dummy_array) + broadcast_shape = np.lib._stride_tricks_impl._broadcast_shape( + *broadcast_args + ) + return broadcast_shape, dim_sizes + + +def _calculate_shapes(broadcast_shape, dim_sizes, list_of_core_dims): + """Helper for calculating broadcast shapes with core dimensions.""" + return [broadcast_shape + tuple(dim_sizes[dim] for dim in core_dims) + for core_dims in list_of_core_dims] + + +def _create_arrays(broadcast_shape, dim_sizes, list_of_core_dims, dtypes, + results=None): + """Helper for creating output arrays in vectorize.""" + shapes = _calculate_shapes(broadcast_shape, dim_sizes, list_of_core_dims) + if dtypes is None: + dtypes = [None] * len(shapes) + if results is None: + arrays = tuple(np.empty(shape=shape, dtype=dtype) + for shape, dtype in zip(shapes, dtypes)) + else: + arrays = tuple(np.empty_like(result, shape=shape, dtype=dtype) + for result, shape, dtype + in zip(results, shapes, dtypes)) + return arrays + + +def _get_vectorize_dtype(dtype): + if dtype.char in "SU": + return dtype.char + return dtype + + +@set_module('numpy') +class vectorize: + """ + vectorize(pyfunc=np._NoValue, otypes=None, doc=None, excluded=None, + cache=False, signature=None) + + Returns an object that acts like pyfunc, but takes arrays as input. + + Define a vectorized function which takes a nested sequence of objects or + numpy arrays as inputs and returns a single numpy array or a tuple of numpy + arrays. The vectorized function evaluates `pyfunc` over successive tuples + of the input arrays like the python map function, except it uses the + broadcasting rules of numpy. + + The data type of the output of `vectorized` is determined by calling + the function with the first element of the input. This can be avoided + by specifying the `otypes` argument. + + Parameters + ---------- + pyfunc : callable, optional + A python function or method. + Can be omitted to produce a decorator with keyword arguments. + otypes : str or list of dtypes, optional + The output data type. It must be specified as either a string of + typecode characters or a list of data type specifiers. There should + be one data type specifier for each output. + doc : str, optional + The docstring for the function. If None, the docstring will be the + ``pyfunc.__doc__``. + excluded : set, optional + Set of strings or integers representing the positional or keyword + arguments for which the function will not be vectorized. These will be + passed directly to `pyfunc` unmodified. + + cache : bool, optional + If `True`, then cache the first function call that determines the number + of outputs if `otypes` is not provided. + + signature : string, optional + Generalized universal function signature, e.g., ``(m,n),(n)->(m)`` for + vectorized matrix-vector multiplication. If provided, ``pyfunc`` will + be called with (and expected to return) arrays with shapes given by the + size of corresponding core dimensions. By default, ``pyfunc`` is + assumed to take scalars as input and output. + + Returns + ------- + out : callable + A vectorized function if ``pyfunc`` was provided, + a decorator otherwise. + + See Also + -------- + frompyfunc : Takes an arbitrary Python function and returns a ufunc + + Notes + ----- + The `vectorize` function is provided primarily for convenience, not for + performance. The implementation is essentially a for loop. + + If `otypes` is not specified, then a call to the function with the + first argument will be used to determine the number of outputs. The + results of this call will be cached if `cache` is `True` to prevent + calling the function twice. However, to implement the cache, the + original function must be wrapped which will slow down subsequent + calls, so only do this if your function is expensive. + + The new keyword argument interface and `excluded` argument support + further degrades performance. + + References + ---------- + .. [1] :doc:`/reference/c-api/generalized-ufuncs` + + Examples + -------- + >>> import numpy as np + >>> def myfunc(a, b): + ... "Return a-b if a>b, otherwise return a+b" + ... if a > b: + ... return a - b + ... else: + ... return a + b + + >>> vfunc = np.vectorize(myfunc) + >>> vfunc([1, 2, 3, 4], 2) + array([3, 4, 1, 2]) + + The docstring is taken from the input function to `vectorize` unless it + is specified: + + >>> vfunc.__doc__ + 'Return a-b if a>b, otherwise return a+b' + >>> vfunc = np.vectorize(myfunc, doc='Vectorized `myfunc`') + >>> vfunc.__doc__ + 'Vectorized `myfunc`' + + The output type is determined by evaluating the first element of the input, + unless it is specified: + + >>> out = vfunc([1, 2, 3, 4], 2) + >>> type(out[0]) + + >>> vfunc = np.vectorize(myfunc, otypes=[float]) + >>> out = vfunc([1, 2, 3, 4], 2) + >>> type(out[0]) + + + The `excluded` argument can be used to prevent vectorizing over certain + arguments. This can be useful for array-like arguments of a fixed length + such as the coefficients for a polynomial as in `polyval`: + + >>> def mypolyval(p, x): + ... _p = list(p) + ... res = _p.pop(0) + ... while _p: + ... res = res*x + _p.pop(0) + ... return res + + Here, we exclude the zeroth argument from vectorization whether it is + passed by position or keyword. + + >>> vpolyval = np.vectorize(mypolyval, excluded={0, 'p'}) + >>> vpolyval([1, 2, 3], x=[0, 1]) + array([3, 6]) + >>> vpolyval(p=[1, 2, 3], x=[0, 1]) + array([3, 6]) + + The `signature` argument allows for vectorizing functions that act on + non-scalar arrays of fixed length. For example, you can use it for a + vectorized calculation of Pearson correlation coefficient and its p-value: + + >>> import scipy.stats + >>> pearsonr = np.vectorize(scipy.stats.pearsonr, + ... signature='(n),(n)->(),()') + >>> pearsonr([[0, 1, 2, 3]], [[1, 2, 3, 4], [4, 3, 2, 1]]) + (array([ 1., -1.]), array([ 0., 0.])) + + Or for a vectorized convolution: + + >>> convolve = np.vectorize(np.convolve, signature='(n),(m)->(k)') + >>> convolve(np.eye(4), [1, 2, 1]) + array([[1., 2., 1., 0., 0., 0.], + [0., 1., 2., 1., 0., 0.], + [0., 0., 1., 2., 1., 0.], + [0., 0., 0., 1., 2., 1.]]) + + Decorator syntax is supported. The decorator can be called as + a function to provide keyword arguments: + + >>> @np.vectorize + ... def identity(x): + ... return x + ... + >>> identity([0, 1, 2]) + array([0, 1, 2]) + >>> @np.vectorize(otypes=[float]) + ... def as_float(x): + ... return x + ... + >>> as_float([0, 1, 2]) + array([0., 1., 2.]) + """ + def __init__(self, pyfunc=np._NoValue, otypes=None, doc=None, + excluded=None, cache=False, signature=None): + + if (pyfunc != np._NoValue) and (not callable(pyfunc)): + #Splitting the error message to keep + #the length below 79 characters. + part1 = "When used as a decorator, " + part2 = "only accepts keyword arguments." + raise TypeError(part1 + part2) + + self.pyfunc = pyfunc + self.cache = cache + self.signature = signature + if pyfunc != np._NoValue and hasattr(pyfunc, '__name__'): + self.__name__ = pyfunc.__name__ + + self._ufunc = {} # Caching to improve default performance + self._doc = None + self.__doc__ = doc + if doc is None and hasattr(pyfunc, '__doc__'): + self.__doc__ = pyfunc.__doc__ + else: + self._doc = doc + + if isinstance(otypes, str): + for char in otypes: + if char not in typecodes['All']: + raise ValueError("Invalid otype specified: %s" % (char,)) + elif iterable(otypes): + otypes = [_get_vectorize_dtype(_nx.dtype(x)) for x in otypes] + elif otypes is not None: + raise ValueError("Invalid otype specification") + self.otypes = otypes + + # Excluded variable support + if excluded is None: + excluded = set() + self.excluded = set(excluded) + + if signature is not None: + self._in_and_out_core_dims = _parse_gufunc_signature(signature) + else: + self._in_and_out_core_dims = None + + def _init_stage_2(self, pyfunc, *args, **kwargs): + self.__name__ = pyfunc.__name__ + self.pyfunc = pyfunc + if self._doc is None: + self.__doc__ = pyfunc.__doc__ + else: + self.__doc__ = self._doc + + def _call_as_normal(self, *args, **kwargs): + """ + Return arrays with the results of `pyfunc` broadcast (vectorized) over + `args` and `kwargs` not in `excluded`. + """ + excluded = self.excluded + if not kwargs and not excluded: + func = self.pyfunc + vargs = args + else: + # The wrapper accepts only positional arguments: we use `names` and + # `inds` to mutate `the_args` and `kwargs` to pass to the original + # function. + nargs = len(args) + + names = [_n for _n in kwargs if _n not in excluded] + inds = [_i for _i in range(nargs) if _i not in excluded] + the_args = list(args) + + def func(*vargs): + for _n, _i in enumerate(inds): + the_args[_i] = vargs[_n] + kwargs.update(zip(names, vargs[len(inds):])) + return self.pyfunc(*the_args, **kwargs) + + vargs = [args[_i] for _i in inds] + vargs.extend([kwargs[_n] for _n in names]) + + return self._vectorize_call(func=func, args=vargs) + + def __call__(self, *args, **kwargs): + if self.pyfunc is np._NoValue: + self._init_stage_2(*args, **kwargs) + return self + + return self._call_as_normal(*args, **kwargs) + + def _get_ufunc_and_otypes(self, func, args): + """Return (ufunc, otypes).""" + # frompyfunc will fail if args is empty + if not args: + raise ValueError('args can not be empty') + + if self.otypes is not None: + otypes = self.otypes + + # self._ufunc is a dictionary whose keys are the number of + # arguments (i.e. len(args)) and whose values are ufuncs created + # by frompyfunc. len(args) can be different for different calls if + # self.pyfunc has parameters with default values. We only use the + # cache when func is self.pyfunc, which occurs when the call uses + # only positional arguments and no arguments are excluded. + + nin = len(args) + nout = len(self.otypes) + if func is not self.pyfunc or nin not in self._ufunc: + ufunc = frompyfunc(func, nin, nout) + else: + ufunc = None # We'll get it from self._ufunc + if func is self.pyfunc: + ufunc = self._ufunc.setdefault(nin, ufunc) + else: + # Get number of outputs and output types by calling the function on + # the first entries of args. We also cache the result to prevent + # the subsequent call when the ufunc is evaluated. + # Assumes that ufunc first evaluates the 0th elements in the input + # arrays (the input values are not checked to ensure this) + args = [asarray(arg) for arg in args] + if builtins.any(arg.size == 0 for arg in args): + raise ValueError('cannot call `vectorize` on size 0 inputs ' + 'unless `otypes` is set') + + inputs = [arg.flat[0] for arg in args] + outputs = func(*inputs) + + # Performance note: profiling indicates that -- for simple + # functions at least -- this wrapping can almost double the + # execution time. + # Hence we make it optional. + if self.cache: + _cache = [outputs] + + def _func(*vargs): + if _cache: + return _cache.pop() + else: + return func(*vargs) + else: + _func = func + + if isinstance(outputs, tuple): + nout = len(outputs) + else: + nout = 1 + outputs = (outputs,) + + otypes = ''.join([asarray(outputs[_k]).dtype.char + for _k in range(nout)]) + + # Performance note: profiling indicates that creating the ufunc is + # not a significant cost compared with wrapping so it seems not + # worth trying to cache this. + ufunc = frompyfunc(_func, len(args), nout) + + return ufunc, otypes + + def _vectorize_call(self, func, args): + """Vectorized call to `func` over positional `args`.""" + if self.signature is not None: + res = self._vectorize_call_with_signature(func, args) + elif not args: + res = func() + else: + ufunc, otypes = self._get_ufunc_and_otypes(func=func, args=args) + + # Convert args to object arrays first + inputs = [asanyarray(a, dtype=object) for a in args] + + outputs = ufunc(*inputs) + + if ufunc.nout == 1: + res = asanyarray(outputs, dtype=otypes[0]) + else: + res = tuple([asanyarray(x, dtype=t) + for x, t in zip(outputs, otypes)]) + return res + + def _vectorize_call_with_signature(self, func, args): + """Vectorized call over positional arguments with a signature.""" + input_core_dims, output_core_dims = self._in_and_out_core_dims + + if len(args) != len(input_core_dims): + raise TypeError('wrong number of positional arguments: ' + 'expected %r, got %r' + % (len(input_core_dims), len(args))) + args = tuple(asanyarray(arg) for arg in args) + + broadcast_shape, dim_sizes = _parse_input_dimensions( + args, input_core_dims) + input_shapes = _calculate_shapes(broadcast_shape, dim_sizes, + input_core_dims) + args = [np.broadcast_to(arg, shape, subok=True) + for arg, shape in zip(args, input_shapes)] + + outputs = None + otypes = self.otypes + nout = len(output_core_dims) + + for index in np.ndindex(*broadcast_shape): + results = func(*(arg[index] for arg in args)) + + n_results = len(results) if isinstance(results, tuple) else 1 + + if nout != n_results: + raise ValueError( + 'wrong number of outputs from pyfunc: expected %r, got %r' + % (nout, n_results)) + + if nout == 1: + results = (results,) + + if outputs is None: + for result, core_dims in zip(results, output_core_dims): + _update_dim_sizes(dim_sizes, result, core_dims) + + outputs = _create_arrays(broadcast_shape, dim_sizes, + output_core_dims, otypes, results) + + for output, result in zip(outputs, results): + output[index] = result + + if outputs is None: + # did not call the function even once + if otypes is None: + raise ValueError('cannot call `vectorize` on size 0 inputs ' + 'unless `otypes` is set') + if builtins.any(dim not in dim_sizes + for dims in output_core_dims + for dim in dims): + raise ValueError('cannot call `vectorize` with a signature ' + 'including new output dimensions on size 0 ' + 'inputs') + outputs = _create_arrays(broadcast_shape, dim_sizes, + output_core_dims, otypes) + + return outputs[0] if nout == 1 else outputs + + +def _cov_dispatcher(m, y=None, rowvar=None, bias=None, ddof=None, + fweights=None, aweights=None, *, dtype=None): + return (m, y, fweights, aweights) + + +@array_function_dispatch(_cov_dispatcher) +def cov(m, y=None, rowvar=True, bias=False, ddof=None, fweights=None, + aweights=None, *, dtype=None): + """ + Estimate a covariance matrix, given data and weights. + + Covariance indicates the level to which two variables vary together. + If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, + then the covariance matrix element :math:`C_{ij}` is the covariance of + :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance + of :math:`x_i`. + + See the notes for an outline of the algorithm. + + Parameters + ---------- + m : array_like + A 1-D or 2-D array containing multiple variables and observations. + Each row of `m` represents a variable, and each column a single + observation of all those variables. Also see `rowvar` below. + y : array_like, optional + An additional set of variables and observations. `y` has the same form + as that of `m`. + rowvar : bool, optional + If `rowvar` is True (default), then each row represents a + variable, with observations in the columns. Otherwise, the relationship + is transposed: each column represents a variable, while the rows + contain observations. + bias : bool, optional + Default normalization (False) is by ``(N - 1)``, where ``N`` is the + number of observations given (unbiased estimate). If `bias` is True, + then normalization is by ``N``. These values can be overridden by using + the keyword ``ddof`` in numpy versions >= 1.5. + ddof : int, optional + If not ``None`` the default value implied by `bias` is overridden. + Note that ``ddof=1`` will return the unbiased estimate, even if both + `fweights` and `aweights` are specified, and ``ddof=0`` will return + the simple average. See the notes for the details. The default value + is ``None``. + fweights : array_like, int, optional + 1-D array of integer frequency weights; the number of times each + observation vector should be repeated. + aweights : array_like, optional + 1-D array of observation vector weights. These relative weights are + typically large for observations considered "important" and smaller for + observations considered less "important". If ``ddof=0`` the array of + weights can be used to assign probabilities to observation vectors. + dtype : data-type, optional + Data-type of the result. By default, the return data-type will have + at least `numpy.float64` precision. + + .. versionadded:: 1.20 + + Returns + ------- + out : ndarray + The covariance matrix of the variables. + + See Also + -------- + corrcoef : Normalized covariance matrix + + Notes + ----- + Assume that the observations are in the columns of the observation + array `m` and let ``f = fweights`` and ``a = aweights`` for brevity. The + steps to compute the weighted covariance are as follows:: + + >>> m = np.arange(10, dtype=np.float64) + >>> f = np.arange(10) * 2 + >>> a = np.arange(10) ** 2. + >>> ddof = 1 + >>> w = f * a + >>> v1 = np.sum(w) + >>> v2 = np.sum(w * a) + >>> m -= np.sum(m * w, axis=None, keepdims=True) / v1 + >>> cov = np.dot(m * w, m.T) * v1 / (v1**2 - ddof * v2) + + Note that when ``a == 1``, the normalization factor + ``v1 / (v1**2 - ddof * v2)`` goes over to ``1 / (np.sum(f) - ddof)`` + as it should. + + Examples + -------- + >>> import numpy as np + + Consider two variables, :math:`x_0` and :math:`x_1`, which + correlate perfectly, but in opposite directions: + + >>> x = np.array([[0, 2], [1, 1], [2, 0]]).T + >>> x + array([[0, 1, 2], + [2, 1, 0]]) + + Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance + matrix shows this clearly: + + >>> np.cov(x) + array([[ 1., -1.], + [-1., 1.]]) + + Note that element :math:`C_{0,1}`, which shows the correlation between + :math:`x_0` and :math:`x_1`, is negative. + + Further, note how `x` and `y` are combined: + + >>> x = [-2.1, -1, 4.3] + >>> y = [3, 1.1, 0.12] + >>> X = np.stack((x, y), axis=0) + >>> np.cov(X) + array([[11.71 , -4.286 ], # may vary + [-4.286 , 2.144133]]) + >>> np.cov(x, y) + array([[11.71 , -4.286 ], # may vary + [-4.286 , 2.144133]]) + >>> np.cov(x) + array(11.71) + + """ + # Check inputs + if ddof is not None and ddof != int(ddof): + raise ValueError( + "ddof must be integer") + + # Handles complex arrays too + m = np.asarray(m) + if m.ndim > 2: + raise ValueError("m has more than 2 dimensions") + + if y is not None: + y = np.asarray(y) + if y.ndim > 2: + raise ValueError("y has more than 2 dimensions") + + if dtype is None: + if y is None: + dtype = np.result_type(m, np.float64) + else: + dtype = np.result_type(m, y, np.float64) + + X = array(m, ndmin=2, dtype=dtype) + if not rowvar and m.ndim != 1: + X = X.T + if X.shape[0] == 0: + return np.array([]).reshape(0, 0) + if y is not None: + y = array(y, copy=None, ndmin=2, dtype=dtype) + if not rowvar and y.shape[0] != 1: + y = y.T + X = np.concatenate((X, y), axis=0) + + if ddof is None: + if bias == 0: + ddof = 1 + else: + ddof = 0 + + # Get the product of frequencies and weights + w = None + if fweights is not None: + fweights = np.asarray(fweights, dtype=float) + if not np.all(fweights == np.around(fweights)): + raise TypeError( + "fweights must be integer") + if fweights.ndim > 1: + raise RuntimeError( + "cannot handle multidimensional fweights") + if fweights.shape[0] != X.shape[1]: + raise RuntimeError( + "incompatible numbers of samples and fweights") + if any(fweights < 0): + raise ValueError( + "fweights cannot be negative") + w = fweights + if aweights is not None: + aweights = np.asarray(aweights, dtype=float) + if aweights.ndim > 1: + raise RuntimeError( + "cannot handle multidimensional aweights") + if aweights.shape[0] != X.shape[1]: + raise RuntimeError( + "incompatible numbers of samples and aweights") + if any(aweights < 0): + raise ValueError( + "aweights cannot be negative") + if w is None: + w = aweights + else: + w *= aweights + + avg, w_sum = average(X, axis=1, weights=w, returned=True) + w_sum = w_sum[0] + + # Determine the normalization + if w is None: + fact = X.shape[1] - ddof + elif ddof == 0: + fact = w_sum + elif aweights is None: + fact = w_sum - ddof + else: + fact = w_sum - ddof*sum(w*aweights)/w_sum + + if fact <= 0: + warnings.warn("Degrees of freedom <= 0 for slice", + RuntimeWarning, stacklevel=2) + fact = 0.0 + + X -= avg[:, None] + if w is None: + X_T = X.T + else: + X_T = (X*w).T + c = dot(X, X_T.conj()) + c *= np.true_divide(1, fact) + return c.squeeze() + + +def _corrcoef_dispatcher(x, y=None, rowvar=None, bias=None, ddof=None, *, + dtype=None): + return (x, y) + + +@array_function_dispatch(_corrcoef_dispatcher) +def corrcoef(x, y=None, rowvar=True, bias=np._NoValue, ddof=np._NoValue, *, + dtype=None): + """ + Return Pearson product-moment correlation coefficients. + + Please refer to the documentation for `cov` for more detail. The + relationship between the correlation coefficient matrix, `R`, and the + covariance matrix, `C`, is + + .. math:: R_{ij} = \\frac{ C_{ij} } { \\sqrt{ C_{ii} C_{jj} } } + + The values of `R` are between -1 and 1, inclusive. + + Parameters + ---------- + x : array_like + A 1-D or 2-D array containing multiple variables and observations. + Each row of `x` represents a variable, and each column a single + observation of all those variables. Also see `rowvar` below. + y : array_like, optional + An additional set of variables and observations. `y` has the same + shape as `x`. + rowvar : bool, optional + If `rowvar` is True (default), then each row represents a + variable, with observations in the columns. Otherwise, the relationship + is transposed: each column represents a variable, while the rows + contain observations. + bias : _NoValue, optional + Has no effect, do not use. + + .. deprecated:: 1.10.0 + ddof : _NoValue, optional + Has no effect, do not use. + + .. deprecated:: 1.10.0 + dtype : data-type, optional + Data-type of the result. By default, the return data-type will have + at least `numpy.float64` precision. + + .. versionadded:: 1.20 + + Returns + ------- + R : ndarray + The correlation coefficient matrix of the variables. + + See Also + -------- + cov : Covariance matrix + + Notes + ----- + Due to floating point rounding the resulting array may not be Hermitian, + the diagonal elements may not be 1, and the elements may not satisfy the + inequality abs(a) <= 1. The real and imaginary parts are clipped to the + interval [-1, 1] in an attempt to improve on that situation but is not + much help in the complex case. + + This function accepts but discards arguments `bias` and `ddof`. This is + for backwards compatibility with previous versions of this function. These + arguments had no effect on the return values of the function and can be + safely ignored in this and previous versions of numpy. + + Examples + -------- + >>> import numpy as np + + In this example we generate two random arrays, ``xarr`` and ``yarr``, and + compute the row-wise and column-wise Pearson correlation coefficients, + ``R``. Since ``rowvar`` is true by default, we first find the row-wise + Pearson correlation coefficients between the variables of ``xarr``. + + >>> import numpy as np + >>> rng = np.random.default_rng(seed=42) + >>> xarr = rng.random((3, 3)) + >>> xarr + array([[0.77395605, 0.43887844, 0.85859792], + [0.69736803, 0.09417735, 0.97562235], + [0.7611397 , 0.78606431, 0.12811363]]) + >>> R1 = np.corrcoef(xarr) + >>> R1 + array([[ 1. , 0.99256089, -0.68080986], + [ 0.99256089, 1. , -0.76492172], + [-0.68080986, -0.76492172, 1. ]]) + + If we add another set of variables and observations ``yarr``, we can + compute the row-wise Pearson correlation coefficients between the + variables in ``xarr`` and ``yarr``. + + >>> yarr = rng.random((3, 3)) + >>> yarr + array([[0.45038594, 0.37079802, 0.92676499], + [0.64386512, 0.82276161, 0.4434142 ], + [0.22723872, 0.55458479, 0.06381726]]) + >>> R2 = np.corrcoef(xarr, yarr) + >>> R2 + array([[ 1. , 0.99256089, -0.68080986, 0.75008178, -0.934284 , + -0.99004057], + [ 0.99256089, 1. , -0.76492172, 0.82502011, -0.97074098, + -0.99981569], + [-0.68080986, -0.76492172, 1. , -0.99507202, 0.89721355, + 0.77714685], + [ 0.75008178, 0.82502011, -0.99507202, 1. , -0.93657855, + -0.83571711], + [-0.934284 , -0.97074098, 0.89721355, -0.93657855, 1. , + 0.97517215], + [-0.99004057, -0.99981569, 0.77714685, -0.83571711, 0.97517215, + 1. ]]) + + Finally if we use the option ``rowvar=False``, the columns are now + being treated as the variables and we will find the column-wise Pearson + correlation coefficients between variables in ``xarr`` and ``yarr``. + + >>> R3 = np.corrcoef(xarr, yarr, rowvar=False) + >>> R3 + array([[ 1. , 0.77598074, -0.47458546, -0.75078643, -0.9665554 , + 0.22423734], + [ 0.77598074, 1. , -0.92346708, -0.99923895, -0.58826587, + -0.44069024], + [-0.47458546, -0.92346708, 1. , 0.93773029, 0.23297648, + 0.75137473], + [-0.75078643, -0.99923895, 0.93773029, 1. , 0.55627469, + 0.47536961], + [-0.9665554 , -0.58826587, 0.23297648, 0.55627469, 1. , + -0.46666491], + [ 0.22423734, -0.44069024, 0.75137473, 0.47536961, -0.46666491, + 1. ]]) + + """ + if bias is not np._NoValue or ddof is not np._NoValue: + # 2015-03-15, 1.10 + warnings.warn('bias and ddof have no effect and are deprecated', + DeprecationWarning, stacklevel=2) + c = cov(x, y, rowvar, dtype=dtype) + try: + d = diag(c) + except ValueError: + # scalar covariance + # nan if incorrect value (nan, inf, 0), 1 otherwise + return c / c + stddev = sqrt(d.real) + c /= stddev[:, None] + c /= stddev[None, :] + + # Clip real and imaginary parts to [-1, 1]. This does not guarantee + # abs(a[i,j]) <= 1 for complex arrays, but is the best we can do without + # excessive work. + np.clip(c.real, -1, 1, out=c.real) + if np.iscomplexobj(c): + np.clip(c.imag, -1, 1, out=c.imag) + + return c + + +@set_module('numpy') +def blackman(M): + """ + Return the Blackman window. + + The Blackman window is a taper formed by using the first three + terms of a summation of cosines. It was designed to have close to the + minimal leakage possible. It is close to optimal, only slightly worse + than a Kaiser window. + + Parameters + ---------- + M : int + Number of points in the output window. If zero or less, an empty + array is returned. + + Returns + ------- + out : ndarray + The window, with the maximum value normalized to one (the value one + appears only if the number of samples is odd). + + See Also + -------- + bartlett, hamming, hanning, kaiser + + Notes + ----- + The Blackman window is defined as + + .. math:: w(n) = 0.42 - 0.5 \\cos(2\\pi n/M) + 0.08 \\cos(4\\pi n/M) + + Most references to the Blackman window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. It is known as a + "near optimal" tapering function, almost as good (by some measures) + as the kaiser window. + + References + ---------- + Blackman, R.B. and Tukey, J.W., (1958) The measurement of power spectra, + Dover Publications, New York. + + Oppenheim, A.V., and R.W. Schafer. Discrete-Time Signal Processing. + Upper Saddle River, NJ: Prentice-Hall, 1999, pp. 468-471. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> np.blackman(12) + array([-1.38777878e-17, 3.26064346e-02, 1.59903635e-01, # may vary + 4.14397981e-01, 7.36045180e-01, 9.67046769e-01, + 9.67046769e-01, 7.36045180e-01, 4.14397981e-01, + 1.59903635e-01, 3.26064346e-02, -1.38777878e-17]) + + Plot the window and the frequency response. + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + from numpy.fft import fft, fftshift + window = np.blackman(51) + plt.plot(window) + plt.title("Blackman window") + plt.ylabel("Amplitude") + plt.xlabel("Sample") + plt.show() # doctest: +SKIP + + plt.figure() + A = fft(window, 2048) / 25.5 + mag = np.abs(fftshift(A)) + freq = np.linspace(-0.5, 0.5, len(A)) + with np.errstate(divide='ignore', invalid='ignore'): + response = 20 * np.log10(mag) + response = np.clip(response, -100, 100) + plt.plot(freq, response) + plt.title("Frequency response of Blackman window") + plt.ylabel("Magnitude [dB]") + plt.xlabel("Normalized frequency [cycles per sample]") + plt.axis('tight') + plt.show() + + """ + # Ensures at least float64 via 0.0. M should be an integer, but conversion + # to double is safe for a range. + values = np.array([0.0, M]) + M = values[1] + + if M < 1: + return array([], dtype=values.dtype) + if M == 1: + return ones(1, dtype=values.dtype) + n = arange(1-M, M, 2) + return 0.42 + 0.5*cos(pi*n/(M-1)) + 0.08*cos(2.0*pi*n/(M-1)) + + +@set_module('numpy') +def bartlett(M): + """ + Return the Bartlett window. + + The Bartlett window is very similar to a triangular window, except + that the end points are at zero. It is often used in signal + processing for tapering a signal, without generating too much + ripple in the frequency domain. + + Parameters + ---------- + M : int + Number of points in the output window. If zero or less, an + empty array is returned. + + Returns + ------- + out : array + The triangular window, with the maximum value normalized to one + (the value one appears only if the number of samples is odd), with + the first and last samples equal to zero. + + See Also + -------- + blackman, hamming, hanning, kaiser + + Notes + ----- + The Bartlett window is defined as + + .. math:: w(n) = \\frac{2}{M-1} \\left( + \\frac{M-1}{2} - \\left|n - \\frac{M-1}{2}\\right| + \\right) + + Most references to the Bartlett window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. Note that convolution with this window produces linear + interpolation. It is also known as an apodization (which means "removing + the foot", i.e. smoothing discontinuities at the beginning and end of the + sampled signal) or tapering function. The Fourier transform of the + Bartlett window is the product of two sinc functions. Note the excellent + discussion in Kanasewich [2]_. + + References + ---------- + .. [1] M.S. Bartlett, "Periodogram Analysis and Continuous Spectra", + Biometrika 37, 1-16, 1950. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", + The University of Alberta Press, 1975, pp. 109-110. + .. [3] A.V. Oppenheim and R.W. Schafer, "Discrete-Time Signal + Processing", Prentice-Hall, 1999, pp. 468-471. + .. [4] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [5] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 429. + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> np.bartlett(12) + array([ 0. , 0.18181818, 0.36363636, 0.54545455, 0.72727273, # may vary + 0.90909091, 0.90909091, 0.72727273, 0.54545455, 0.36363636, + 0.18181818, 0. ]) + + Plot the window and its frequency response (requires SciPy and matplotlib). + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + from numpy.fft import fft, fftshift + window = np.bartlett(51) + plt.plot(window) + plt.title("Bartlett window") + plt.ylabel("Amplitude") + plt.xlabel("Sample") + plt.show() + plt.figure() + A = fft(window, 2048) / 25.5 + mag = np.abs(fftshift(A)) + freq = np.linspace(-0.5, 0.5, len(A)) + with np.errstate(divide='ignore', invalid='ignore'): + response = 20 * np.log10(mag) + response = np.clip(response, -100, 100) + plt.plot(freq, response) + plt.title("Frequency response of Bartlett window") + plt.ylabel("Magnitude [dB]") + plt.xlabel("Normalized frequency [cycles per sample]") + plt.axis('tight') + plt.show() + + """ + # Ensures at least float64 via 0.0. M should be an integer, but conversion + # to double is safe for a range. + values = np.array([0.0, M]) + M = values[1] + + if M < 1: + return array([], dtype=values.dtype) + if M == 1: + return ones(1, dtype=values.dtype) + n = arange(1-M, M, 2) + return where(less_equal(n, 0), 1 + n/(M-1), 1 - n/(M-1)) + + +@set_module('numpy') +def hanning(M): + """ + Return the Hanning window. + + The Hanning window is a taper formed by using a weighted cosine. + + Parameters + ---------- + M : int + Number of points in the output window. If zero or less, an + empty array is returned. + + Returns + ------- + out : ndarray, shape(M,) + The window, with the maximum value normalized to one (the value + one appears only if `M` is odd). + + See Also + -------- + bartlett, blackman, hamming, kaiser + + Notes + ----- + The Hanning window is defined as + + .. math:: w(n) = 0.5 - 0.5\\cos\\left(\\frac{2\\pi{n}}{M-1}\\right) + \\qquad 0 \\leq n \\leq M-1 + + The Hanning was named for Julius von Hann, an Austrian meteorologist. + It is also known as the Cosine Bell. Some authors prefer that it be + called a Hann window, to help avoid confusion with the very similar + Hamming window. + + Most references to the Hanning window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] Blackman, R.B. and Tukey, J.W., (1958) The measurement of power + spectra, Dover Publications, New York. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", + The University of Alberta Press, 1975, pp. 106-108. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [4] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 425. + + Examples + -------- + >>> import numpy as np + >>> np.hanning(12) + array([0. , 0.07937323, 0.29229249, 0.57115742, 0.82743037, + 0.97974649, 0.97974649, 0.82743037, 0.57115742, 0.29229249, + 0.07937323, 0. ]) + + Plot the window and its frequency response. + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + from numpy.fft import fft, fftshift + window = np.hanning(51) + plt.plot(window) + plt.title("Hann window") + plt.ylabel("Amplitude") + plt.xlabel("Sample") + plt.show() + + plt.figure() + A = fft(window, 2048) / 25.5 + mag = np.abs(fftshift(A)) + freq = np.linspace(-0.5, 0.5, len(A)) + with np.errstate(divide='ignore', invalid='ignore'): + response = 20 * np.log10(mag) + response = np.clip(response, -100, 100) + plt.plot(freq, response) + plt.title("Frequency response of the Hann window") + plt.ylabel("Magnitude [dB]") + plt.xlabel("Normalized frequency [cycles per sample]") + plt.axis('tight') + plt.show() + + """ + # Ensures at least float64 via 0.0. M should be an integer, but conversion + # to double is safe for a range. + values = np.array([0.0, M]) + M = values[1] + + if M < 1: + return array([], dtype=values.dtype) + if M == 1: + return ones(1, dtype=values.dtype) + n = arange(1-M, M, 2) + return 0.5 + 0.5*cos(pi*n/(M-1)) + + +@set_module('numpy') +def hamming(M): + """ + Return the Hamming window. + + The Hamming window is a taper formed by using a weighted cosine. + + Parameters + ---------- + M : int + Number of points in the output window. If zero or less, an + empty array is returned. + + Returns + ------- + out : ndarray + The window, with the maximum value normalized to one (the value + one appears only if the number of samples is odd). + + See Also + -------- + bartlett, blackman, hanning, kaiser + + Notes + ----- + The Hamming window is defined as + + .. math:: w(n) = 0.54 - 0.46\\cos\\left(\\frac{2\\pi{n}}{M-1}\\right) + \\qquad 0 \\leq n \\leq M-1 + + The Hamming was named for R. W. Hamming, an associate of J. W. Tukey + and is described in Blackman and Tukey. It was recommended for + smoothing the truncated autocovariance function in the time domain. + Most references to the Hamming window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] Blackman, R.B. and Tukey, J.W., (1958) The measurement of power + spectra, Dover Publications, New York. + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", The + University of Alberta Press, 1975, pp. 109-110. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + .. [4] W.H. Press, B.P. Flannery, S.A. Teukolsky, and W.T. Vetterling, + "Numerical Recipes", Cambridge University Press, 1986, page 425. + + Examples + -------- + >>> import numpy as np + >>> np.hamming(12) + array([ 0.08 , 0.15302337, 0.34890909, 0.60546483, 0.84123594, # may vary + 0.98136677, 0.98136677, 0.84123594, 0.60546483, 0.34890909, + 0.15302337, 0.08 ]) + + Plot the window and the frequency response. + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + from numpy.fft import fft, fftshift + window = np.hamming(51) + plt.plot(window) + plt.title("Hamming window") + plt.ylabel("Amplitude") + plt.xlabel("Sample") + plt.show() + + plt.figure() + A = fft(window, 2048) / 25.5 + mag = np.abs(fftshift(A)) + freq = np.linspace(-0.5, 0.5, len(A)) + response = 20 * np.log10(mag) + response = np.clip(response, -100, 100) + plt.plot(freq, response) + plt.title("Frequency response of Hamming window") + plt.ylabel("Magnitude [dB]") + plt.xlabel("Normalized frequency [cycles per sample]") + plt.axis('tight') + plt.show() + + """ + # Ensures at least float64 via 0.0. M should be an integer, but conversion + # to double is safe for a range. + values = np.array([0.0, M]) + M = values[1] + + if M < 1: + return array([], dtype=values.dtype) + if M == 1: + return ones(1, dtype=values.dtype) + n = arange(1-M, M, 2) + return 0.54 + 0.46*cos(pi*n/(M-1)) + + +## Code from cephes for i0 + +_i0A = [ + -4.41534164647933937950E-18, + 3.33079451882223809783E-17, + -2.43127984654795469359E-16, + 1.71539128555513303061E-15, + -1.16853328779934516808E-14, + 7.67618549860493561688E-14, + -4.85644678311192946090E-13, + 2.95505266312963983461E-12, + -1.72682629144155570723E-11, + 9.67580903537323691224E-11, + -5.18979560163526290666E-10, + 2.65982372468238665035E-9, + -1.30002500998624804212E-8, + 6.04699502254191894932E-8, + -2.67079385394061173391E-7, + 1.11738753912010371815E-6, + -4.41673835845875056359E-6, + 1.64484480707288970893E-5, + -5.75419501008210370398E-5, + 1.88502885095841655729E-4, + -5.76375574538582365885E-4, + 1.63947561694133579842E-3, + -4.32430999505057594430E-3, + 1.05464603945949983183E-2, + -2.37374148058994688156E-2, + 4.93052842396707084878E-2, + -9.49010970480476444210E-2, + 1.71620901522208775349E-1, + -3.04682672343198398683E-1, + 6.76795274409476084995E-1 + ] + +_i0B = [ + -7.23318048787475395456E-18, + -4.83050448594418207126E-18, + 4.46562142029675999901E-17, + 3.46122286769746109310E-17, + -2.82762398051658348494E-16, + -3.42548561967721913462E-16, + 1.77256013305652638360E-15, + 3.81168066935262242075E-15, + -9.55484669882830764870E-15, + -4.15056934728722208663E-14, + 1.54008621752140982691E-14, + 3.85277838274214270114E-13, + 7.18012445138366623367E-13, + -1.79417853150680611778E-12, + -1.32158118404477131188E-11, + -3.14991652796324136454E-11, + 1.18891471078464383424E-11, + 4.94060238822496958910E-10, + 3.39623202570838634515E-9, + 2.26666899049817806459E-8, + 2.04891858946906374183E-7, + 2.89137052083475648297E-6, + 6.88975834691682398426E-5, + 3.36911647825569408990E-3, + 8.04490411014108831608E-1 + ] + + +def _chbevl(x, vals): + b0 = vals[0] + b1 = 0.0 + + for i in range(1, len(vals)): + b2 = b1 + b1 = b0 + b0 = x*b1 - b2 + vals[i] + + return 0.5*(b0 - b2) + + +def _i0_1(x): + return exp(x) * _chbevl(x/2.0-2, _i0A) + + +def _i0_2(x): + return exp(x) * _chbevl(32.0/x - 2.0, _i0B) / sqrt(x) + + +def _i0_dispatcher(x): + return (x,) + + +@array_function_dispatch(_i0_dispatcher) +def i0(x): + """ + Modified Bessel function of the first kind, order 0. + + Usually denoted :math:`I_0`. + + Parameters + ---------- + x : array_like of float + Argument of the Bessel function. + + Returns + ------- + out : ndarray, shape = x.shape, dtype = float + The modified Bessel function evaluated at each of the elements of `x`. + + See Also + -------- + scipy.special.i0, scipy.special.iv, scipy.special.ive + + Notes + ----- + The scipy implementation is recommended over this function: it is a + proper ufunc written in C, and more than an order of magnitude faster. + + We use the algorithm published by Clenshaw [1]_ and referenced by + Abramowitz and Stegun [2]_, for which the function domain is + partitioned into the two intervals [0,8] and (8,inf), and Chebyshev + polynomial expansions are employed in each interval. Relative error on + the domain [0,30] using IEEE arithmetic is documented [3]_ as having a + peak of 5.8e-16 with an rms of 1.4e-16 (n = 30000). + + References + ---------- + .. [1] C. W. Clenshaw, "Chebyshev series for mathematical functions", in + *National Physical Laboratory Mathematical Tables*, vol. 5, London: + Her Majesty's Stationery Office, 1962. + .. [2] M. Abramowitz and I. A. Stegun, *Handbook of Mathematical + Functions*, 10th printing, New York: Dover, 1964, pp. 379. + https://personal.math.ubc.ca/~cbm/aands/page_379.htm + .. [3] https://metacpan.org/pod/distribution/Math-Cephes/lib/Math/Cephes.pod#i0:-Modified-Bessel-function-of-order-zero + + Examples + -------- + >>> import numpy as np + >>> np.i0(0.) + array(1.0) + >>> np.i0([0, 1, 2, 3]) + array([1. , 1.26606588, 2.2795853 , 4.88079259]) + + """ + x = np.asanyarray(x) + if x.dtype.kind == 'c': + raise TypeError("i0 not supported for complex values") + if x.dtype.kind != 'f': + x = x.astype(float) + x = np.abs(x) + return piecewise(x, [x <= 8.0], [_i0_1, _i0_2]) + +## End of cephes code for i0 + + +@set_module('numpy') +def kaiser(M, beta): + """ + Return the Kaiser window. + + The Kaiser window is a taper formed by using a Bessel function. + + Parameters + ---------- + M : int + Number of points in the output window. If zero or less, an + empty array is returned. + beta : float + Shape parameter for window. + + Returns + ------- + out : array + The window, with the maximum value normalized to one (the value + one appears only if the number of samples is odd). + + See Also + -------- + bartlett, blackman, hamming, hanning + + Notes + ----- + The Kaiser window is defined as + + .. math:: w(n) = I_0\\left( \\beta \\sqrt{1-\\frac{4n^2}{(M-1)^2}} + \\right)/I_0(\\beta) + + with + + .. math:: \\quad -\\frac{M-1}{2} \\leq n \\leq \\frac{M-1}{2}, + + where :math:`I_0` is the modified zeroth-order Bessel function. + + The Kaiser was named for Jim Kaiser, who discovered a simple + approximation to the DPSS window based on Bessel functions. The Kaiser + window is a very good approximation to the Digital Prolate Spheroidal + Sequence, or Slepian window, which is the transform which maximizes the + energy in the main lobe of the window relative to total energy. + + The Kaiser can approximate many other windows by varying the beta + parameter. + + ==== ======================= + beta Window shape + ==== ======================= + 0 Rectangular + 5 Similar to a Hamming + 6 Similar to a Hanning + 8.6 Similar to a Blackman + ==== ======================= + + A beta value of 14 is probably a good starting point. Note that as beta + gets large, the window narrows, and so the number of samples needs to be + large enough to sample the increasingly narrow spike, otherwise NaNs will + get returned. + + Most references to the Kaiser window come from the signal processing + literature, where it is used as one of many windowing functions for + smoothing values. It is also known as an apodization (which means + "removing the foot", i.e. smoothing discontinuities at the beginning + and end of the sampled signal) or tapering function. + + References + ---------- + .. [1] J. F. Kaiser, "Digital Filters" - Ch 7 in "Systems analysis by + digital computer", Editors: F.F. Kuo and J.F. Kaiser, p 218-285. + John Wiley and Sons, New York, (1966). + .. [2] E.R. Kanasewich, "Time Sequence Analysis in Geophysics", The + University of Alberta Press, 1975, pp. 177-178. + .. [3] Wikipedia, "Window function", + https://en.wikipedia.org/wiki/Window_function + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> np.kaiser(12, 14) + array([7.72686684e-06, 3.46009194e-03, 4.65200189e-02, # may vary + 2.29737120e-01, 5.99885316e-01, 9.45674898e-01, + 9.45674898e-01, 5.99885316e-01, 2.29737120e-01, + 4.65200189e-02, 3.46009194e-03, 7.72686684e-06]) + + + Plot the window and the frequency response. + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + from numpy.fft import fft, fftshift + window = np.kaiser(51, 14) + plt.plot(window) + plt.title("Kaiser window") + plt.ylabel("Amplitude") + plt.xlabel("Sample") + plt.show() + + plt.figure() + A = fft(window, 2048) / 25.5 + mag = np.abs(fftshift(A)) + freq = np.linspace(-0.5, 0.5, len(A)) + response = 20 * np.log10(mag) + response = np.clip(response, -100, 100) + plt.plot(freq, response) + plt.title("Frequency response of Kaiser window") + plt.ylabel("Magnitude [dB]") + plt.xlabel("Normalized frequency [cycles per sample]") + plt.axis('tight') + plt.show() + + """ + # Ensures at least float64 via 0.0. M should be an integer, but conversion + # to double is safe for a range. (Simplified result_type with 0.0 + # strongly typed. result-type is not/less order sensitive, but that mainly + # matters for integers anyway.) + values = np.array([0.0, M, beta]) + M = values[1] + beta = values[2] + + if M == 1: + return np.ones(1, dtype=values.dtype) + n = arange(0, M) + alpha = (M-1)/2.0 + return i0(beta * sqrt(1-((n-alpha)/alpha)**2.0))/i0(beta) + + +def _sinc_dispatcher(x): + return (x,) + + +@array_function_dispatch(_sinc_dispatcher) +def sinc(x): + r""" + Return the normalized sinc function. + + The sinc function is equal to :math:`\sin(\pi x)/(\pi x)` for any argument + :math:`x\ne 0`. ``sinc(0)`` takes the limit value 1, making ``sinc`` not + only everywhere continuous but also infinitely differentiable. + + .. note:: + + Note the normalization factor of ``pi`` used in the definition. + This is the most commonly used definition in signal processing. + Use ``sinc(x / np.pi)`` to obtain the unnormalized sinc function + :math:`\sin(x)/x` that is more common in mathematics. + + Parameters + ---------- + x : ndarray + Array (possibly multi-dimensional) of values for which to calculate + ``sinc(x)``. + + Returns + ------- + out : ndarray + ``sinc(x)``, which has the same shape as the input. + + Notes + ----- + The name sinc is short for "sine cardinal" or "sinus cardinalis". + + The sinc function is used in various signal processing applications, + including in anti-aliasing, in the construction of a Lanczos resampling + filter, and in interpolation. + + For bandlimited interpolation of discrete-time signals, the ideal + interpolation kernel is proportional to the sinc function. + + References + ---------- + .. [1] Weisstein, Eric W. "Sinc Function." From MathWorld--A Wolfram Web + Resource. https://mathworld.wolfram.com/SincFunction.html + .. [2] Wikipedia, "Sinc function", + https://en.wikipedia.org/wiki/Sinc_function + + Examples + -------- + >>> import numpy as np + >>> import matplotlib.pyplot as plt + >>> x = np.linspace(-4, 4, 41) + >>> np.sinc(x) + array([-3.89804309e-17, -4.92362781e-02, -8.40918587e-02, # may vary + -8.90384387e-02, -5.84680802e-02, 3.89804309e-17, + 6.68206631e-02, 1.16434881e-01, 1.26137788e-01, + 8.50444803e-02, -3.89804309e-17, -1.03943254e-01, + -1.89206682e-01, -2.16236208e-01, -1.55914881e-01, + 3.89804309e-17, 2.33872321e-01, 5.04551152e-01, + 7.56826729e-01, 9.35489284e-01, 1.00000000e+00, + 9.35489284e-01, 7.56826729e-01, 5.04551152e-01, + 2.33872321e-01, 3.89804309e-17, -1.55914881e-01, + -2.16236208e-01, -1.89206682e-01, -1.03943254e-01, + -3.89804309e-17, 8.50444803e-02, 1.26137788e-01, + 1.16434881e-01, 6.68206631e-02, 3.89804309e-17, + -5.84680802e-02, -8.90384387e-02, -8.40918587e-02, + -4.92362781e-02, -3.89804309e-17]) + + >>> plt.plot(x, np.sinc(x)) + [] + >>> plt.title("Sinc Function") + Text(0.5, 1.0, 'Sinc Function') + >>> plt.ylabel("Amplitude") + Text(0, 0.5, 'Amplitude') + >>> plt.xlabel("X") + Text(0.5, 0, 'X') + >>> plt.show() + + """ + x = np.asanyarray(x) + y = pi * where(x == 0, 1.0e-20, x) + return sin(y)/y + + +def _ureduce(a, func, keepdims=False, **kwargs): + """ + Internal Function. + Call `func` with `a` as first argument swapping the axes to use extended + axis on functions that don't support it natively. + + Returns result and a.shape with axis dims set to 1. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + func : callable + Reduction function capable of receiving a single axis argument. + It is called with `a` as first argument followed by `kwargs`. + kwargs : keyword arguments + additional keyword arguments to pass to `func`. + + Returns + ------- + result : tuple + Result of func(a, **kwargs) and a.shape with axis dims set to 1 + which can be used to reshape the result to the same shape a ufunc with + keepdims=True would produce. + + """ + a = np.asanyarray(a) + axis = kwargs.get('axis', None) + out = kwargs.get('out', None) + + if keepdims is np._NoValue: + keepdims = False + + nd = a.ndim + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, nd) + + if keepdims: + if out is not None: + index_out = tuple( + 0 if i in axis else slice(None) for i in range(nd)) + kwargs['out'] = out[(Ellipsis, ) + index_out] + + if len(axis) == 1: + kwargs['axis'] = axis[0] + else: + keep = set(range(nd)) - set(axis) + nkeep = len(keep) + # swap axis that should not be reduced to front + for i, s in enumerate(sorted(keep)): + a = a.swapaxes(i, s) + # merge reduced axis + a = a.reshape(a.shape[:nkeep] + (-1,)) + kwargs['axis'] = -1 + else: + if keepdims: + if out is not None: + index_out = (0, ) * nd + kwargs['out'] = out[(Ellipsis, ) + index_out] + + r = func(a, **kwargs) + + if out is not None: + return out + + if keepdims: + if axis is None: + index_r = (np.newaxis, ) * nd + else: + index_r = tuple( + np.newaxis if i in axis else slice(None) + for i in range(nd)) + r = r[(Ellipsis, ) + index_r] + + return r + + +def _median_dispatcher( + a, axis=None, out=None, overwrite_input=None, keepdims=None): + return (a, out) + + +@array_function_dispatch(_median_dispatcher) +def median(a, axis=None, out=None, overwrite_input=False, keepdims=False): + """ + Compute the median along the specified axis. + + Returns the median of the array elements. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + axis : {int, sequence of int, None}, optional + Axis or axes along which the medians are computed. The default, + axis=None, will compute the median along a flattened version of + the array. If a sequence of axes, the array is first flattened + along the given axes, then the median is computed along the + resulting flattened axis. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output, + but the type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow use of memory of input array `a` for + calculations. The input array will be modified by the call to + `median`. This will save memory when you do not need to preserve + the contents of the input array. Treat the input as undefined, + but it will probably be fully or partially sorted. Default is + False. If `overwrite_input` is ``True`` and `a` is not already an + `ndarray`, an error will be raised. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. + + Returns + ------- + median : ndarray + A new array holding the result. If the input contains integers + or floats smaller than ``float64``, then the output data-type is + ``np.float64``. Otherwise, the data-type of the output is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + mean, percentile + + Notes + ----- + Given a vector ``V`` of length ``N``, the median of ``V`` is the + middle value of a sorted copy of ``V``, ``V_sorted`` - i + e., ``V_sorted[(N-1)/2]``, when ``N`` is odd, and the average of the + two middle values of ``V_sorted`` when ``N`` is even. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10, 7, 4], [3, 2, 1]]) + >>> a + array([[10, 7, 4], + [ 3, 2, 1]]) + >>> np.median(a) + np.float64(3.5) + >>> np.median(a, axis=0) + array([6.5, 4.5, 2.5]) + >>> np.median(a, axis=1) + array([7., 2.]) + >>> np.median(a, axis=(0, 1)) + np.float64(3.5) + >>> m = np.median(a, axis=0) + >>> out = np.zeros_like(m) + >>> np.median(a, axis=0, out=m) + array([6.5, 4.5, 2.5]) + >>> m + array([6.5, 4.5, 2.5]) + >>> b = a.copy() + >>> np.median(b, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a==b) + >>> b = a.copy() + >>> np.median(b, axis=None, overwrite_input=True) + np.float64(3.5) + >>> assert not np.all(a==b) + + """ + return _ureduce(a, func=_median, keepdims=keepdims, axis=axis, out=out, + overwrite_input=overwrite_input) + + +def _median(a, axis=None, out=None, overwrite_input=False): + # can't be reasonably be implemented in terms of percentile as we have to + # call mean to not break astropy + a = np.asanyarray(a) + + # Set the partition indexes + if axis is None: + sz = a.size + else: + sz = a.shape[axis] + if sz % 2 == 0: + szh = sz // 2 + kth = [szh - 1, szh] + else: + kth = [(sz - 1) // 2] + + # We have to check for NaNs (as of writing 'M' doesn't actually work). + supports_nans = np.issubdtype(a.dtype, np.inexact) or a.dtype.kind in 'Mm' + if supports_nans: + kth.append(-1) + + if overwrite_input: + if axis is None: + part = a.ravel() + part.partition(kth) + else: + a.partition(kth, axis=axis) + part = a + else: + part = partition(a, kth, axis=axis) + + if part.shape == (): + # make 0-D arrays work + return part.item() + if axis is None: + axis = 0 + + indexer = [slice(None)] * part.ndim + index = part.shape[axis] // 2 + if part.shape[axis] % 2 == 1: + # index with slice to allow mean (below) to work + indexer[axis] = slice(index, index+1) + else: + indexer[axis] = slice(index-1, index+1) + indexer = tuple(indexer) + + # Use mean in both odd and even case to coerce data type, + # using out array if needed. + rout = mean(part[indexer], axis=axis, out=out) + if supports_nans and sz > 0: + # If nans are possible, warn and replace by nans like mean would. + rout = np.lib._utils_impl._median_nancheck(part, rout, axis) + + return rout + + +def _percentile_dispatcher(a, q, axis=None, out=None, overwrite_input=None, + method=None, keepdims=None, *, weights=None, + interpolation=None): + return (a, q, out, weights) + + +@array_function_dispatch(_percentile_dispatcher) +def percentile(a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None): + """ + Compute the q-th percentile of the data along the specified axis. + + Returns the q-th percentile(s) of the array elements. + + Parameters + ---------- + a : array_like of real numbers + Input array or object that can be converted to an array. + q : array_like of float + Percentage or sequence of percentages for the percentiles to compute. + Values must be between 0 and 100 inclusive. + axis : {int, tuple of int, None}, optional + Axis or axes along which the percentiles are computed. The + default is to compute the percentile(s) along a flattened + version of the array. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output, + but the type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow the input array `a` to be modified by intermediate + calculations, to save memory. In this case, the contents of the input + `a` after this function completes is undefined. + method : str, optional + This parameter specifies the method to use for estimating the + percentile. There are many different methods, some unique to NumPy. + See the notes for explanation. The options sorted by their R type + as summarized in the H&F paper [1]_ are: + + 1. 'inverted_cdf' + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + The first three methods are discontinuous. NumPy further defines the + following discontinuous variations of the default 'linear' (7.) option: + + * 'lower' + * 'higher', + * 'midpoint' + * 'nearest' + + .. versionchanged:: 1.22.0 + This argument was previously called "interpolation" and only + offered the "linear" default and last four options. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in + the result as dimensions with size one. With this option, the + result will broadcast correctly against the original array `a`. + + weights : array_like, optional + An array of weights associated with the values in `a`. Each value in + `a` contributes to the percentile according to its associated weight. + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given axis) or of the same shape as `a`. + If `weights=None`, then all data in `a` are assumed to have a + weight equal to one. + Only `method="inverted_cdf"` supports weights. + See the notes for more details. + + .. versionadded:: 2.0.0 + + interpolation : str, optional + Deprecated name for the method keyword argument. + + .. deprecated:: 1.22.0 + + Returns + ------- + percentile : scalar or ndarray + If `q` is a single percentile and `axis=None`, then the result + is a scalar. If multiple percentiles are given, first axis of + the result corresponds to the percentiles. The other axes are + the axes that remain after the reduction of `a`. If the input + contains integers or floats smaller than ``float64``, the output + data-type is ``float64``. Otherwise, the output data-type is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + mean + median : equivalent to ``percentile(..., 50)`` + nanpercentile + quantile : equivalent to percentile, except q in the range [0, 1]. + + Notes + ----- + The behavior of `numpy.percentile` with percentage `q` is + that of `numpy.quantile` with argument ``q/100``. + For more information, please see `numpy.quantile`. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10, 7, 4], [3, 2, 1]]) + >>> a + array([[10, 7, 4], + [ 3, 2, 1]]) + >>> np.percentile(a, 50) + 3.5 + >>> np.percentile(a, 50, axis=0) + array([6.5, 4.5, 2.5]) + >>> np.percentile(a, 50, axis=1) + array([7., 2.]) + >>> np.percentile(a, 50, axis=1, keepdims=True) + array([[7.], + [2.]]) + + >>> m = np.percentile(a, 50, axis=0) + >>> out = np.zeros_like(m) + >>> np.percentile(a, 50, axis=0, out=out) + array([6.5, 4.5, 2.5]) + >>> m + array([6.5, 4.5, 2.5]) + + >>> b = a.copy() + >>> np.percentile(b, 50, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a == b) + + The different methods can be visualized graphically: + + .. plot:: + + import matplotlib.pyplot as plt + + a = np.arange(4) + p = np.linspace(0, 100, 6001) + ax = plt.gca() + lines = [ + ('linear', '-', 'C0'), + ('inverted_cdf', ':', 'C1'), + # Almost the same as `inverted_cdf`: + ('averaged_inverted_cdf', '-.', 'C1'), + ('closest_observation', ':', 'C2'), + ('interpolated_inverted_cdf', '--', 'C1'), + ('hazen', '--', 'C3'), + ('weibull', '-.', 'C4'), + ('median_unbiased', '--', 'C5'), + ('normal_unbiased', '-.', 'C6'), + ] + for method, style, color in lines: + ax.plot( + p, np.percentile(a, p, method=method), + label=method, linestyle=style, color=color) + ax.set( + title='Percentiles for different methods and data: ' + str(a), + xlabel='Percentile', + ylabel='Estimated percentile value', + yticks=a) + ax.legend(bbox_to_anchor=(1.03, 1)) + plt.tight_layout() + plt.show() + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + + """ + if interpolation is not None: + method = _check_interpolation_as_method( + method, interpolation, "percentile") + + a = np.asanyarray(a) + if a.dtype.kind == "c": + raise TypeError("a must be an array of real numbers") + + # Use dtype of array if possible (e.g., if q is a python int or float) + # by making the divisor have the dtype of the data array. + q = np.true_divide(q, a.dtype.type(100) if a.dtype.kind == "f" else 100) + q = asanyarray(q) # undo any decay that the ufunc performed (see gh-13105) + if not _quantile_is_valid(q): + raise ValueError("Percentiles must be in the range [0, 100]") + + if weights is not None: + if method != "inverted_cdf": + msg = ("Only method 'inverted_cdf' supports weights. " + f"Got: {method}.") + raise ValueError(msg) + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, a.ndim, argname="axis") + weights = _weights_are_valid(weights=weights, a=a, axis=axis) + if np.any(weights < 0): + raise ValueError("Weights must be non-negative.") + + return _quantile_unchecked( + a, q, axis, out, overwrite_input, method, keepdims, weights) + + +def _quantile_dispatcher(a, q, axis=None, out=None, overwrite_input=None, + method=None, keepdims=None, *, weights=None, + interpolation=None): + return (a, q, out, weights) + + +@array_function_dispatch(_quantile_dispatcher) +def quantile(a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None): + """ + Compute the q-th quantile of the data along the specified axis. + + Parameters + ---------- + a : array_like of real numbers + Input array or object that can be converted to an array. + q : array_like of float + Probability or sequence of probabilities of the quantiles to compute. + Values must be between 0 and 1 inclusive. + axis : {int, tuple of int, None}, optional + Axis or axes along which the quantiles are computed. The default is + to compute the quantile(s) along a flattened version of the array. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape and buffer length as the expected output, but the + type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow the input array `a` to be modified by + intermediate calculations, to save memory. In this case, the + contents of the input `a` after this function completes is + undefined. + method : str, optional + This parameter specifies the method to use for estimating the + quantile. There are many different methods, some unique to NumPy. + The recommended options, numbered as they appear in [1]_, are: + + 1. 'inverted_cdf' + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + The first three methods are discontinuous. For backward compatibility + with previous versions of NumPy, the following discontinuous variations + of the default 'linear' (7.) option are available: + + * 'lower' + * 'higher', + * 'midpoint' + * 'nearest' + + See Notes for details. + + .. versionchanged:: 1.22.0 + This argument was previously called "interpolation" and only + offered the "linear" default and last four options. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in + the result as dimensions with size one. With this option, the + result will broadcast correctly against the original array `a`. + + weights : array_like, optional + An array of weights associated with the values in `a`. Each value in + `a` contributes to the quantile according to its associated weight. + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given axis) or of the same shape as `a`. + If `weights=None`, then all data in `a` are assumed to have a + weight equal to one. + Only `method="inverted_cdf"` supports weights. + See the notes for more details. + + .. versionadded:: 2.0.0 + + interpolation : str, optional + Deprecated name for the method keyword argument. + + .. deprecated:: 1.22.0 + + Returns + ------- + quantile : scalar or ndarray + If `q` is a single probability and `axis=None`, then the result + is a scalar. If multiple probability levels are given, first axis + of the result corresponds to the quantiles. The other axes are + the axes that remain after the reduction of `a`. If the input + contains integers or floats smaller than ``float64``, the output + data-type is ``float64``. Otherwise, the output data-type is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + mean + percentile : equivalent to quantile, but with q in the range [0, 100]. + median : equivalent to ``quantile(..., 0.5)`` + nanquantile + + Notes + ----- + Given a sample `a` from an underlying distribution, `quantile` provides a + nonparametric estimate of the inverse cumulative distribution function. + + By default, this is done by interpolating between adjacent elements in + ``y``, a sorted copy of `a`:: + + (1-g)*y[j] + g*y[j+1] + + where the index ``j`` and coefficient ``g`` are the integral and + fractional components of ``q * (n-1)``, and ``n`` is the number of + elements in the sample. + + This is a special case of Equation 1 of H&F [1]_. More generally, + + - ``j = (q*n + m - 1) // 1``, and + - ``g = (q*n + m - 1) % 1``, + + where ``m`` may be defined according to several different conventions. + The preferred convention may be selected using the ``method`` parameter: + + =============================== =============== =============== + ``method`` number in H&F ``m`` + =============================== =============== =============== + ``interpolated_inverted_cdf`` 4 ``0`` + ``hazen`` 5 ``1/2`` + ``weibull`` 6 ``q`` + ``linear`` (default) 7 ``1 - q`` + ``median_unbiased`` 8 ``q/3 + 1/3`` + ``normal_unbiased`` 9 ``q/4 + 3/8`` + =============================== =============== =============== + + Note that indices ``j`` and ``j + 1`` are clipped to the range ``0`` to + ``n - 1`` when the results of the formula would be outside the allowed + range of non-negative indices. The ``- 1`` in the formulas for ``j`` and + ``g`` accounts for Python's 0-based indexing. + + The table above includes only the estimators from H&F that are continuous + functions of probability `q` (estimators 4-9). NumPy also provides the + three discontinuous estimators from H&F (estimators 1-3), where ``j`` is + defined as above, ``m`` is defined as follows, and ``g`` is a function + of the real-valued ``index = q*n + m - 1`` and ``j``. + + 1. ``inverted_cdf``: ``m = 0`` and ``g = int(index - j > 0)`` + 2. ``averaged_inverted_cdf``: ``m = 0`` and + ``g = (1 + int(index - j > 0)) / 2`` + 3. ``closest_observation``: ``m = -1/2`` and + ``g = 1 - int((index == j) & (j%2 == 1))`` + + For backward compatibility with previous versions of NumPy, `quantile` + provides four additional discontinuous estimators. Like + ``method='linear'``, all have ``m = 1 - q`` so that ``j = q*(n-1) // 1``, + but ``g`` is defined as follows. + + - ``lower``: ``g = 0`` + - ``midpoint``: ``g = 0.5`` + - ``higher``: ``g = 1`` + - ``nearest``: ``g = (q*(n-1) % 1) > 0.5`` + + **Weighted quantiles:** + More formally, the quantile at probability level :math:`q` of a cumulative + distribution function :math:`F(y)=P(Y \\leq y)` with probability measure + :math:`P` is defined as any number :math:`x` that fulfills the + *coverage conditions* + + .. math:: P(Y < x) \\leq q \\quad\\text{and}\\quad P(Y \\leq x) \\geq q + + with random variable :math:`Y\\sim P`. + Sample quantiles, the result of `quantile`, provide nonparametric + estimation of the underlying population counterparts, represented by the + unknown :math:`F`, given a data vector `a` of length ``n``. + + Some of the estimators above arise when one considers :math:`F` as the + empirical distribution function of the data, i.e. + :math:`F(y) = \\frac{1}{n} \\sum_i 1_{a_i \\leq y}`. + Then, different methods correspond to different choices of :math:`x` that + fulfill the above coverage conditions. Methods that follow this approach + are ``inverted_cdf`` and ``averaged_inverted_cdf``. + + For weighted quantiles, the coverage conditions still hold. The + empirical cumulative distribution is simply replaced by its weighted + version, i.e. + :math:`P(Y \\leq t) = \\frac{1}{\\sum_i w_i} \\sum_i w_i 1_{x_i \\leq t}`. + Only ``method="inverted_cdf"`` supports weights. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10, 7, 4], [3, 2, 1]]) + >>> a + array([[10, 7, 4], + [ 3, 2, 1]]) + >>> np.quantile(a, 0.5) + 3.5 + >>> np.quantile(a, 0.5, axis=0) + array([6.5, 4.5, 2.5]) + >>> np.quantile(a, 0.5, axis=1) + array([7., 2.]) + >>> np.quantile(a, 0.5, axis=1, keepdims=True) + array([[7.], + [2.]]) + >>> m = np.quantile(a, 0.5, axis=0) + >>> out = np.zeros_like(m) + >>> np.quantile(a, 0.5, axis=0, out=out) + array([6.5, 4.5, 2.5]) + >>> m + array([6.5, 4.5, 2.5]) + >>> b = a.copy() + >>> np.quantile(b, 0.5, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a == b) + + See also `numpy.percentile` for a visualization of most methods. + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + + """ + if interpolation is not None: + method = _check_interpolation_as_method( + method, interpolation, "quantile") + + a = np.asanyarray(a) + if a.dtype.kind == "c": + raise TypeError("a must be an array of real numbers") + + # Use dtype of array if possible (e.g., if q is a python int or float). + if isinstance(q, (int, float)) and a.dtype.kind == "f": + q = np.asanyarray(q, dtype=a.dtype) + else: + q = np.asanyarray(q) + + if not _quantile_is_valid(q): + raise ValueError("Quantiles must be in the range [0, 1]") + + if weights is not None: + if method != "inverted_cdf": + msg = ("Only method 'inverted_cdf' supports weights. " + f"Got: {method}.") + raise ValueError(msg) + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, a.ndim, argname="axis") + weights = _weights_are_valid(weights=weights, a=a, axis=axis) + if np.any(weights < 0): + raise ValueError("Weights must be non-negative.") + + return _quantile_unchecked( + a, q, axis, out, overwrite_input, method, keepdims, weights) + + +def _quantile_unchecked(a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + weights=None): + """Assumes that q is in [0, 1], and is an ndarray""" + return _ureduce(a, + func=_quantile_ureduce_func, + q=q, + weights=weights, + keepdims=keepdims, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method) + + +def _quantile_is_valid(q): + # avoid expensive reductions, relevant for arrays with < O(1000) elements + if q.ndim == 1 and q.size < 10: + for i in range(q.size): + if not (0.0 <= q[i] <= 1.0): + return False + else: + if not (q.min() >= 0 and q.max() <= 1): + return False + return True + + +def _check_interpolation_as_method(method, interpolation, fname): + # Deprecated NumPy 1.22, 2021-11-08 + warnings.warn( + f"the `interpolation=` argument to {fname} was renamed to " + "`method=`, which has additional options.\n" + "Users of the modes 'nearest', 'lower', 'higher', or " + "'midpoint' are encouraged to review the method they used. " + "(Deprecated NumPy 1.22)", + DeprecationWarning, stacklevel=4) + if method != "linear": + # sanity check, we assume this basically never happens + raise TypeError( + "You shall not pass both `method` and `interpolation`!\n" + "(`interpolation` is Deprecated in favor of `method`)") + return interpolation + + +def _compute_virtual_index(n, quantiles, alpha: float, beta: float): + """ + Compute the floating point indexes of an array for the linear + interpolation of quantiles. + n : array_like + The sample sizes. + quantiles : array_like + The quantiles values. + alpha : float + A constant used to correct the index computed. + beta : float + A constant used to correct the index computed. + + alpha and beta values depend on the chosen method + (see quantile documentation) + + Reference: + Hyndman&Fan paper "Sample Quantiles in Statistical Packages", + DOI: 10.1080/00031305.1996.10473566 + """ + return n * quantiles + ( + alpha + quantiles * (1 - alpha - beta) + ) - 1 + + +def _get_gamma(virtual_indexes, previous_indexes, method): + """ + Compute gamma (a.k.a 'm' or 'weight') for the linear interpolation + of quantiles. + + virtual_indexes : array_like + The indexes where the percentile is supposed to be found in the sorted + sample. + previous_indexes : array_like + The floor values of virtual_indexes. + interpolation : dict + The interpolation method chosen, which may have a specific rule + modifying gamma. + + gamma is usually the fractional part of virtual_indexes but can be modified + by the interpolation method. + """ + gamma = np.asanyarray(virtual_indexes - previous_indexes) + gamma = method["fix_gamma"](gamma, virtual_indexes) + # Ensure both that we have an array, and that we keep the dtype + # (which may have been matched to the input array). + return np.asanyarray(gamma, dtype=virtual_indexes.dtype) + + +def _lerp(a, b, t, out=None): + """ + Compute the linear interpolation weighted by gamma on each point of + two same shape array. + + a : array_like + Left bound. + b : array_like + Right bound. + t : array_like + The interpolation weight. + out : array_like + Output array. + """ + diff_b_a = subtract(b, a) + # asanyarray is a stop-gap until gh-13105 + lerp_interpolation = asanyarray(add(a, diff_b_a * t, out=out)) + subtract(b, diff_b_a * (1 - t), out=lerp_interpolation, where=t >= 0.5, + casting='unsafe', dtype=type(lerp_interpolation.dtype)) + if lerp_interpolation.ndim == 0 and out is None: + lerp_interpolation = lerp_interpolation[()] # unpack 0d arrays + return lerp_interpolation + + +def _get_gamma_mask(shape, default_value, conditioned_value, where): + out = np.full(shape, default_value) + np.copyto(out, conditioned_value, where=where, casting="unsafe") + return out + + +def _discrete_interpolation_to_boundaries(index, gamma_condition_fun): + previous = np.floor(index) + next = previous + 1 + gamma = index - previous + res = _get_gamma_mask(shape=index.shape, + default_value=next, + conditioned_value=previous, + where=gamma_condition_fun(gamma, index) + ).astype(np.intp) + # Some methods can lead to out-of-bound integers, clip them: + res[res < 0] = 0 + return res + + +def _closest_observation(n, quantiles): + # "choose the nearest even order statistic at g=0" (H&F (1996) pp. 362). + # Order is 1-based so for zero-based indexing round to nearest odd index. + gamma_fun = lambda gamma, index: (gamma == 0) & (np.floor(index) % 2 == 1) + return _discrete_interpolation_to_boundaries((n * quantiles) - 1 - 0.5, + gamma_fun) + + +def _inverted_cdf(n, quantiles): + gamma_fun = lambda gamma, _: (gamma == 0) + return _discrete_interpolation_to_boundaries((n * quantiles) - 1, + gamma_fun) + + +def _quantile_ureduce_func( + a: np.array, + q: np.array, + weights: np.array, + axis: int | None = None, + out=None, + overwrite_input: bool = False, + method="linear", +) -> np.array: + if q.ndim > 2: + # The code below works fine for nd, but it might not have useful + # semantics. For now, keep the supported dimensions the same as it was + # before. + raise ValueError("q must be a scalar or 1d") + if overwrite_input: + if axis is None: + axis = 0 + arr = a.ravel() + wgt = None if weights is None else weights.ravel() + else: + arr = a + wgt = weights + else: + if axis is None: + axis = 0 + arr = a.flatten() + wgt = None if weights is None else weights.flatten() + else: + arr = a.copy() + wgt = weights + result = _quantile(arr, + quantiles=q, + axis=axis, + method=method, + out=out, + weights=wgt) + return result + + +def _get_indexes(arr, virtual_indexes, valid_values_count): + """ + Get the valid indexes of arr neighbouring virtual_indexes. + Note + This is a companion function to linear interpolation of + Quantiles + + Returns + ------- + (previous_indexes, next_indexes): Tuple + A Tuple of virtual_indexes neighbouring indexes + """ + previous_indexes = np.asanyarray(np.floor(virtual_indexes)) + next_indexes = np.asanyarray(previous_indexes + 1) + indexes_above_bounds = virtual_indexes >= valid_values_count - 1 + # When indexes is above max index, take the max value of the array + if indexes_above_bounds.any(): + previous_indexes[indexes_above_bounds] = -1 + next_indexes[indexes_above_bounds] = -1 + # When indexes is below min index, take the min value of the array + indexes_below_bounds = virtual_indexes < 0 + if indexes_below_bounds.any(): + previous_indexes[indexes_below_bounds] = 0 + next_indexes[indexes_below_bounds] = 0 + if np.issubdtype(arr.dtype, np.inexact): + # After the sort, slices having NaNs will have for last element a NaN + virtual_indexes_nans = np.isnan(virtual_indexes) + if virtual_indexes_nans.any(): + previous_indexes[virtual_indexes_nans] = -1 + next_indexes[virtual_indexes_nans] = -1 + previous_indexes = previous_indexes.astype(np.intp) + next_indexes = next_indexes.astype(np.intp) + return previous_indexes, next_indexes + + +def _quantile( + arr: np.array, + quantiles: np.array, + axis: int = -1, + method="linear", + out=None, + weights=None, +): + """ + Private function that doesn't support extended axis or keepdims. + These methods are extended to this function using _ureduce + See nanpercentile for parameter usage + It computes the quantiles of the array for the given axis. + A linear interpolation is performed based on the `interpolation`. + + By default, the method is "linear" where alpha == beta == 1 which + performs the 7th method of Hyndman&Fan. + With "median_unbiased" we get alpha == beta == 1/3 + thus the 8th method of Hyndman&Fan. + """ + # --- Setup + arr = np.asanyarray(arr) + values_count = arr.shape[axis] + # The dimensions of `q` are prepended to the output shape, so we need the + # axis being sampled from `arr` to be last. + if axis != 0: # But moveaxis is slow, so only call it if necessary. + arr = np.moveaxis(arr, axis, destination=0) + supports_nans = ( + np.issubdtype(arr.dtype, np.inexact) or arr.dtype.kind in 'Mm' + ) + + if weights is None: + # --- Computation of indexes + # Index where to find the value in the sorted array. + # Virtual because it is a floating point value, not an valid index. + # The nearest neighbours are used for interpolation + try: + method_props = _QuantileMethods[method] + except KeyError: + raise ValueError( + f"{method!r} is not a valid method. Use one of: " + f"{_QuantileMethods.keys()}") from None + virtual_indexes = method_props["get_virtual_index"](values_count, + quantiles) + virtual_indexes = np.asanyarray(virtual_indexes) + + if method_props["fix_gamma"] is None: + supports_integers = True + else: + int_virtual_indices = np.issubdtype(virtual_indexes.dtype, + np.integer) + supports_integers = method == 'linear' and int_virtual_indices + + if supports_integers: + # No interpolation needed, take the points along axis + if supports_nans: + # may contain nan, which would sort to the end + arr.partition( + concatenate((virtual_indexes.ravel(), [-1])), axis=0, + ) + slices_having_nans = np.isnan(arr[-1, ...]) + else: + # cannot contain nan + arr.partition(virtual_indexes.ravel(), axis=0) + slices_having_nans = np.array(False, dtype=bool) + result = take(arr, virtual_indexes, axis=0, out=out) + else: + previous_indexes, next_indexes = _get_indexes(arr, + virtual_indexes, + values_count) + # --- Sorting + arr.partition( + np.unique(np.concatenate(([0, -1], + previous_indexes.ravel(), + next_indexes.ravel(), + ))), + axis=0) + if supports_nans: + slices_having_nans = np.isnan(arr[-1, ...]) + else: + slices_having_nans = None + # --- Get values from indexes + previous = arr[previous_indexes] + next = arr[next_indexes] + # --- Linear interpolation + gamma = _get_gamma(virtual_indexes, previous_indexes, method_props) + result_shape = virtual_indexes.shape + (1,) * (arr.ndim - 1) + gamma = gamma.reshape(result_shape) + result = _lerp(previous, + next, + gamma, + out=out) + else: + # Weighted case + # This implements method="inverted_cdf", the only supported weighted + # method, which needs to sort anyway. + weights = np.asanyarray(weights) + if axis != 0: + weights = np.moveaxis(weights, axis, destination=0) + index_array = np.argsort(arr, axis=0, kind="stable") + + # arr = arr[index_array, ...] # but this adds trailing dimensions of + # 1. + arr = np.take_along_axis(arr, index_array, axis=0) + if weights.shape == arr.shape: + weights = np.take_along_axis(weights, index_array, axis=0) + else: + # weights is 1d + weights = weights.reshape(-1)[index_array, ...] + + if supports_nans: + # may contain nan, which would sort to the end + slices_having_nans = np.isnan(arr[-1, ...]) + else: + # cannot contain nan + slices_having_nans = np.array(False, dtype=bool) + + # We use the weights to calculate the empirical cumulative + # distribution function cdf + cdf = weights.cumsum(axis=0, dtype=np.float64) + cdf /= cdf[-1, ...] # normalization to 1 + # Search index i such that + # sum(weights[j], j=0..i-1) < quantile <= sum(weights[j], j=0..i) + # is then equivalent to + # cdf[i-1] < quantile <= cdf[i] + # Unfortunately, searchsorted only accepts 1-d arrays as first + # argument, so we will need to iterate over dimensions. + + # Without the following cast, searchsorted can return surprising + # results, e.g. + # np.searchsorted(np.array([0.2, 0.4, 0.6, 0.8, 1.]), + # np.array(0.4, dtype=np.float32), side="left") + # returns 2 instead of 1 because 0.4 is not binary representable. + if quantiles.dtype.kind == "f": + cdf = cdf.astype(quantiles.dtype) + # Weights must be non-negative, so we might have zero weights at the + # beginning leading to some leading zeros in cdf. The call to + # np.searchsorted for quantiles=0 will then pick the first element, + # but should pick the first one larger than zero. We + # therefore simply set 0 values in cdf to -1. + if np.any(cdf[0, ...] == 0): + cdf[cdf == 0] = -1 + + def find_cdf_1d(arr, cdf): + indices = np.searchsorted(cdf, quantiles, side="left") + # We might have reached the maximum with i = len(arr), e.g. for + # quantiles = 1, and need to cut it to len(arr) - 1. + indices = minimum(indices, values_count - 1) + result = take(arr, indices, axis=0) + return result + + r_shape = arr.shape[1:] + if quantiles.ndim > 0: + r_shape = quantiles.shape + r_shape + if out is None: + result = np.empty_like(arr, shape=r_shape) + else: + if out.shape != r_shape: + msg = (f"Wrong shape of argument 'out', shape={r_shape} is " + f"required; got shape={out.shape}.") + raise ValueError(msg) + result = out + + # See apply_along_axis, which we do for axis=0. Note that Ni = (,) + # always, so we remove it here. + Nk = arr.shape[1:] + for kk in np.ndindex(Nk): + result[(...,) + kk] = find_cdf_1d( + arr[np.s_[:, ] + kk], cdf[np.s_[:, ] + kk] + ) + + # Make result the same as in unweighted inverted_cdf. + if result.shape == () and result.dtype == np.dtype("O"): + result = result.item() + + if np.any(slices_having_nans): + if result.ndim == 0 and out is None: + # can't write to a scalar, but indexing will be correct + result = arr[-1] + else: + np.copyto(result, arr[-1, ...], where=slices_having_nans) + return result + + +def _trapezoid_dispatcher(y, x=None, dx=None, axis=None): + return (y, x) + + +@array_function_dispatch(_trapezoid_dispatcher) +def trapezoid(y, x=None, dx=1.0, axis=-1): + r""" + Integrate along the given axis using the composite trapezoidal rule. + + If `x` is provided, the integration happens in sequence along its + elements - they are not sorted. + + Integrate `y` (`x`) along each 1d slice on the given axis, compute + :math:`\int y(x) dx`. + When `x` is specified, this integrates along the parametric curve, + computing :math:`\int_t y(t) dt = + \int_t y(t) \left.\frac{dx}{dt}\right|_{x=x(t)} dt`. + + .. versionadded:: 2.0.0 + + Parameters + ---------- + y : array_like + Input array to integrate. + x : array_like, optional + The sample points corresponding to the `y` values. If `x` is None, + the sample points are assumed to be evenly spaced `dx` apart. The + default is None. + dx : scalar, optional + The spacing between sample points when `x` is None. The default is 1. + axis : int, optional + The axis along which to integrate. + + Returns + ------- + trapezoid : float or ndarray + Definite integral of `y` = n-dimensional array as approximated along + a single axis by the trapezoidal rule. If `y` is a 1-dimensional array, + then the result is a float. If `n` is greater than 1, then the result + is an `n`-1 dimensional array. + + See Also + -------- + sum, cumsum + + Notes + ----- + Image [2]_ illustrates trapezoidal rule -- y-axis locations of points + will be taken from `y` array, by default x-axis distances between + points will be 1.0, alternatively they can be provided with `x` array + or with `dx` scalar. Return value will be equal to combined area under + the red lines. + + + References + ---------- + .. [1] Wikipedia page: https://en.wikipedia.org/wiki/Trapezoidal_rule + + .. [2] Illustration image: + https://en.wikipedia.org/wiki/File:Composite_trapezoidal_rule_illustration.png + + Examples + -------- + >>> import numpy as np + + Use the trapezoidal rule on evenly spaced points: + + >>> np.trapezoid([1, 2, 3]) + 4.0 + + The spacing between sample points can be selected by either the + ``x`` or ``dx`` arguments: + + >>> np.trapezoid([1, 2, 3], x=[4, 6, 8]) + 8.0 + >>> np.trapezoid([1, 2, 3], dx=2) + 8.0 + + Using a decreasing ``x`` corresponds to integrating in reverse: + + >>> np.trapezoid([1, 2, 3], x=[8, 6, 4]) + -8.0 + + More generally ``x`` is used to integrate along a parametric curve. We can + estimate the integral :math:`\int_0^1 x^2 = 1/3` using: + + >>> x = np.linspace(0, 1, num=50) + >>> y = x**2 + >>> np.trapezoid(y, x) + 0.33340274885464394 + + Or estimate the area of a circle, noting we repeat the sample which closes + the curve: + + >>> theta = np.linspace(0, 2 * np.pi, num=1000, endpoint=True) + >>> np.trapezoid(np.cos(theta), x=np.sin(theta)) + 3.141571941375841 + + ``np.trapezoid`` can be applied along a specified axis to do multiple + computations in one call: + + >>> a = np.arange(6).reshape(2, 3) + >>> a + array([[0, 1, 2], + [3, 4, 5]]) + >>> np.trapezoid(a, axis=0) + array([1.5, 2.5, 3.5]) + >>> np.trapezoid(a, axis=1) + array([2., 8.]) + """ + + y = asanyarray(y) + if x is None: + d = dx + else: + x = asanyarray(x) + if x.ndim == 1: + d = diff(x) + # reshape to correct shape + shape = [1]*y.ndim + shape[axis] = d.shape[0] + d = d.reshape(shape) + else: + d = diff(x, axis=axis) + nd = y.ndim + slice1 = [slice(None)]*nd + slice2 = [slice(None)]*nd + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) + try: + ret = (d * (y[tuple(slice1)] + y[tuple(slice2)]) / 2.0).sum(axis) + except ValueError: + # Operations didn't work, cast to ndarray + d = np.asarray(d) + y = np.asarray(y) + ret = add.reduce(d * (y[tuple(slice1)]+y[tuple(slice2)])/2.0, axis) + return ret + + +@set_module('numpy') +def trapz(y, x=None, dx=1.0, axis=-1): + """ + `trapz` is deprecated in NumPy 2.0. + + Please use `trapezoid` instead, or one of the numerical integration + functions in `scipy.integrate`. + """ + # Deprecated in NumPy 2.0, 2023-08-18 + warnings.warn( + "`trapz` is deprecated. Use `trapezoid` instead, or one of the " + "numerical integration functions in `scipy.integrate`.", + DeprecationWarning, + stacklevel=2 + ) + return trapezoid(y, x=x, dx=dx, axis=axis) + + +def _meshgrid_dispatcher(*xi, copy=None, sparse=None, indexing=None): + return xi + + +# Based on scitools meshgrid +@array_function_dispatch(_meshgrid_dispatcher) +def meshgrid(*xi, copy=True, sparse=False, indexing='xy'): + """ + Return a tuple of coordinate matrices from coordinate vectors. + + Make N-D coordinate arrays for vectorized evaluations of + N-D scalar/vector fields over N-D grids, given + one-dimensional coordinate arrays x1, x2,..., xn. + + Parameters + ---------- + x1, x2,..., xn : array_like + 1-D arrays representing the coordinates of a grid. + indexing : {'xy', 'ij'}, optional + Cartesian ('xy', default) or matrix ('ij') indexing of output. + See Notes for more details. + sparse : bool, optional + If True the shape of the returned coordinate array for dimension *i* + is reduced from ``(N1, ..., Ni, ... Nn)`` to + ``(1, ..., 1, Ni, 1, ..., 1)``. These sparse coordinate grids are + intended to be use with :ref:`basics.broadcasting`. When all + coordinates are used in an expression, broadcasting still leads to a + fully-dimensonal result array. + + Default is False. + + copy : bool, optional + If False, a view into the original arrays are returned in order to + conserve memory. Default is True. Please note that + ``sparse=False, copy=False`` will likely return non-contiguous + arrays. Furthermore, more than one element of a broadcast array + may refer to a single memory location. If you need to write to the + arrays, make copies first. + + Returns + ------- + X1, X2,..., XN : tuple of ndarrays + For vectors `x1`, `x2`,..., `xn` with lengths ``Ni=len(xi)``, + returns ``(N1, N2, N3,..., Nn)`` shaped arrays if indexing='ij' + or ``(N2, N1, N3,..., Nn)`` shaped arrays if indexing='xy' + with the elements of `xi` repeated to fill the matrix along + the first dimension for `x1`, the second for `x2` and so on. + + Notes + ----- + This function supports both indexing conventions through the indexing + keyword argument. Giving the string 'ij' returns a meshgrid with + matrix indexing, while 'xy' returns a meshgrid with Cartesian indexing. + In the 2-D case with inputs of length M and N, the outputs are of shape + (N, M) for 'xy' indexing and (M, N) for 'ij' indexing. In the 3-D case + with inputs of length M, N and P, outputs are of shape (N, M, P) for + 'xy' indexing and (M, N, P) for 'ij' indexing. The difference is + illustrated by the following code snippet:: + + xv, yv = np.meshgrid(x, y, indexing='ij') + for i in range(nx): + for j in range(ny): + # treat xv[i,j], yv[i,j] + + xv, yv = np.meshgrid(x, y, indexing='xy') + for i in range(nx): + for j in range(ny): + # treat xv[j,i], yv[j,i] + + In the 1-D and 0-D case, the indexing and sparse keywords have no effect. + + See Also + -------- + mgrid : Construct a multi-dimensional "meshgrid" using indexing notation. + ogrid : Construct an open multi-dimensional "meshgrid" using indexing + notation. + :ref:`how-to-index` + + Examples + -------- + >>> import numpy as np + >>> nx, ny = (3, 2) + >>> x = np.linspace(0, 1, nx) + >>> y = np.linspace(0, 1, ny) + >>> xv, yv = np.meshgrid(x, y) + >>> xv + array([[0. , 0.5, 1. ], + [0. , 0.5, 1. ]]) + >>> yv + array([[0., 0., 0.], + [1., 1., 1.]]) + + The result of `meshgrid` is a coordinate grid: + + >>> import matplotlib.pyplot as plt + >>> plt.plot(xv, yv, marker='o', color='k', linestyle='none') + >>> plt.show() + + You can create sparse output arrays to save memory and computation time. + + >>> xv, yv = np.meshgrid(x, y, sparse=True) + >>> xv + array([[0. , 0.5, 1. ]]) + >>> yv + array([[0.], + [1.]]) + + `meshgrid` is very useful to evaluate functions on a grid. If the + function depends on all coordinates, both dense and sparse outputs can be + used. + + >>> x = np.linspace(-5, 5, 101) + >>> y = np.linspace(-5, 5, 101) + >>> # full coordinate arrays + >>> xx, yy = np.meshgrid(x, y) + >>> zz = np.sqrt(xx**2 + yy**2) + >>> xx.shape, yy.shape, zz.shape + ((101, 101), (101, 101), (101, 101)) + >>> # sparse coordinate arrays + >>> xs, ys = np.meshgrid(x, y, sparse=True) + >>> zs = np.sqrt(xs**2 + ys**2) + >>> xs.shape, ys.shape, zs.shape + ((1, 101), (101, 1), (101, 101)) + >>> np.array_equal(zz, zs) + True + + >>> h = plt.contourf(x, y, zs) + >>> plt.axis('scaled') + >>> plt.colorbar() + >>> plt.show() + """ + ndim = len(xi) + + if indexing not in ['xy', 'ij']: + raise ValueError( + "Valid values for `indexing` are 'xy' and 'ij'.") + + s0 = (1,) * ndim + output = [np.asanyarray(x).reshape(s0[:i] + (-1,) + s0[i + 1:]) + for i, x in enumerate(xi)] + + if indexing == 'xy' and ndim > 1: + # switch first and second axis + output[0].shape = (1, -1) + s0[2:] + output[1].shape = (-1, 1) + s0[2:] + + if not sparse: + # Return the full N-D matrix (not only the 1-D vector) + output = np.broadcast_arrays(*output, subok=True) + + if copy: + output = tuple(x.copy() for x in output) + + return output + + +def _delete_dispatcher(arr, obj, axis=None): + return (arr, obj) + + +@array_function_dispatch(_delete_dispatcher) +def delete(arr, obj, axis=None): + """ + Return a new array with sub-arrays along an axis deleted. For a one + dimensional array, this returns those entries not returned by + `arr[obj]`. + + Parameters + ---------- + arr : array_like + Input array. + obj : slice, int, array-like of ints or bools + Indicate indices of sub-arrays to remove along the specified axis. + + .. versionchanged:: 1.19.0 + Boolean indices are now treated as a mask of elements to remove, + rather than being cast to the integers 0 and 1. + + axis : int, optional + The axis along which to delete the subarray defined by `obj`. + If `axis` is None, `obj` is applied to the flattened array. + + Returns + ------- + out : ndarray + A copy of `arr` with the elements specified by `obj` removed. Note + that `delete` does not occur in-place. If `axis` is None, `out` is + a flattened array. + + See Also + -------- + insert : Insert elements into an array. + append : Append elements at the end of an array. + + Notes + ----- + Often it is preferable to use a boolean mask. For example: + + >>> arr = np.arange(12) + 1 + >>> mask = np.ones(len(arr), dtype=bool) + >>> mask[[0,2,4]] = False + >>> result = arr[mask,...] + + Is equivalent to ``np.delete(arr, [0,2,4], axis=0)``, but allows further + use of `mask`. + + Examples + -------- + >>> import numpy as np + >>> arr = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]]) + >>> arr + array([[ 1, 2, 3, 4], + [ 5, 6, 7, 8], + [ 9, 10, 11, 12]]) + >>> np.delete(arr, 1, 0) + array([[ 1, 2, 3, 4], + [ 9, 10, 11, 12]]) + + >>> np.delete(arr, np.s_[::2], 1) + array([[ 2, 4], + [ 6, 8], + [10, 12]]) + >>> np.delete(arr, [1,3,5], None) + array([ 1, 3, 5, 7, 8, 9, 10, 11, 12]) + + """ + conv = _array_converter(arr) + arr, = conv.as_arrays(subok=False) + + ndim = arr.ndim + arrorder = 'F' if arr.flags.fnc else 'C' + if axis is None: + if ndim != 1: + arr = arr.ravel() + # needed for np.matrix, which is still not 1d after being ravelled + ndim = arr.ndim + axis = ndim - 1 + else: + axis = normalize_axis_index(axis, ndim) + + slobj = [slice(None)]*ndim + N = arr.shape[axis] + newshape = list(arr.shape) + + if isinstance(obj, slice): + start, stop, step = obj.indices(N) + xr = range(start, stop, step) + numtodel = len(xr) + + if numtodel <= 0: + return conv.wrap(arr.copy(order=arrorder), to_scalar=False) + + # Invert if step is negative: + if step < 0: + step = -step + start = xr[-1] + stop = xr[0] + 1 + + newshape[axis] -= numtodel + new = empty(newshape, arr.dtype, arrorder) + # copy initial chunk + if start == 0: + pass + else: + slobj[axis] = slice(None, start) + new[tuple(slobj)] = arr[tuple(slobj)] + # copy end chunk + if stop == N: + pass + else: + slobj[axis] = slice(stop-numtodel, None) + slobj2 = [slice(None)]*ndim + slobj2[axis] = slice(stop, None) + new[tuple(slobj)] = arr[tuple(slobj2)] + # copy middle pieces + if step == 1: + pass + else: # use array indexing. + keep = ones(stop-start, dtype=bool) + keep[:stop-start:step] = False + slobj[axis] = slice(start, stop-numtodel) + slobj2 = [slice(None)]*ndim + slobj2[axis] = slice(start, stop) + arr = arr[tuple(slobj2)] + slobj2[axis] = keep + new[tuple(slobj)] = arr[tuple(slobj2)] + + return conv.wrap(new, to_scalar=False) + + if isinstance(obj, (int, integer)) and not isinstance(obj, bool): + single_value = True + else: + single_value = False + _obj = obj + obj = np.asarray(obj) + # `size == 0` to allow empty lists similar to indexing, but (as there) + # is really too generic: + if obj.size == 0 and not isinstance(_obj, np.ndarray): + obj = obj.astype(intp) + elif obj.size == 1 and obj.dtype.kind in "ui": + # For a size 1 integer array we can use the single-value path + # (most dtypes, except boolean, should just fail later). + obj = obj.item() + single_value = True + + if single_value: + # optimization for a single value + if (obj < -N or obj >= N): + raise IndexError( + "index %i is out of bounds for axis %i with " + "size %i" % (obj, axis, N)) + if (obj < 0): + obj += N + newshape[axis] -= 1 + new = empty(newshape, arr.dtype, arrorder) + slobj[axis] = slice(None, obj) + new[tuple(slobj)] = arr[tuple(slobj)] + slobj[axis] = slice(obj, None) + slobj2 = [slice(None)]*ndim + slobj2[axis] = slice(obj+1, None) + new[tuple(slobj)] = arr[tuple(slobj2)] + else: + if obj.dtype == bool: + if obj.shape != (N,): + raise ValueError('boolean array argument obj to delete ' + 'must be one dimensional and match the axis ' + 'length of {}'.format(N)) + + # optimization, the other branch is slower + keep = ~obj + else: + keep = ones(N, dtype=bool) + keep[obj,] = False + + slobj[axis] = keep + new = arr[tuple(slobj)] + + return conv.wrap(new, to_scalar=False) + + +def _insert_dispatcher(arr, obj, values, axis=None): + return (arr, obj, values) + + +@array_function_dispatch(_insert_dispatcher) +def insert(arr, obj, values, axis=None): + """ + Insert values along the given axis before the given indices. + + Parameters + ---------- + arr : array_like + Input array. + obj : slice, int, array-like of ints or bools + Object that defines the index or indices before which `values` is + inserted. + + .. versionchanged:: 2.1.2 + Boolean indices are now treated as a mask of elements to insert, + rather than being cast to the integers 0 and 1. + + Support for multiple insertions when `obj` is a single scalar or a + sequence with one element (similar to calling insert multiple + times). + values : array_like + Values to insert into `arr`. If the type of `values` is different + from that of `arr`, `values` is converted to the type of `arr`. + `values` should be shaped so that ``arr[...,obj,...] = values`` + is legal. + axis : int, optional + Axis along which to insert `values`. If `axis` is None then `arr` + is flattened first. + + Returns + ------- + out : ndarray + A copy of `arr` with `values` inserted. Note that `insert` + does not occur in-place: a new array is returned. If + `axis` is None, `out` is a flattened array. + + See Also + -------- + append : Append elements at the end of an array. + concatenate : Join a sequence of arrays along an existing axis. + delete : Delete elements from an array. + + Notes + ----- + Note that for higher dimensional inserts ``obj=0`` behaves very different + from ``obj=[0]`` just like ``arr[:,0,:] = values`` is different from + ``arr[:,[0],:] = values``. This is because of the difference between basic + and advanced :ref:`indexing `. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(6).reshape(3, 2) + >>> a + array([[0, 1], + [2, 3], + [4, 5]]) + >>> np.insert(a, 1, 6) + array([0, 6, 1, 2, 3, 4, 5]) + >>> np.insert(a, 1, 6, axis=1) + array([[0, 6, 1], + [2, 6, 3], + [4, 6, 5]]) + + Difference between sequence and scalars, + showing how ``obj=[1]`` behaves different from ``obj=1``: + + >>> np.insert(a, [1], [[7],[8],[9]], axis=1) + array([[0, 7, 1], + [2, 8, 3], + [4, 9, 5]]) + >>> np.insert(a, 1, [[7],[8],[9]], axis=1) + array([[0, 7, 8, 9, 1], + [2, 7, 8, 9, 3], + [4, 7, 8, 9, 5]]) + >>> np.array_equal(np.insert(a, 1, [7, 8, 9], axis=1), + ... np.insert(a, [1], [[7],[8],[9]], axis=1)) + True + + >>> b = a.flatten() + >>> b + array([0, 1, 2, 3, 4, 5]) + >>> np.insert(b, [2, 2], [6, 7]) + array([0, 1, 6, 7, 2, 3, 4, 5]) + + >>> np.insert(b, slice(2, 4), [7, 8]) + array([0, 1, 7, 2, 8, 3, 4, 5]) + + >>> np.insert(b, [2, 2], [7.13, False]) # type casting + array([0, 1, 7, 0, 2, 3, 4, 5]) + + >>> x = np.arange(8).reshape(2, 4) + >>> idx = (1, 3) + >>> np.insert(x, idx, 999, axis=1) + array([[ 0, 999, 1, 2, 999, 3], + [ 4, 999, 5, 6, 999, 7]]) + + """ + conv = _array_converter(arr) + arr, = conv.as_arrays(subok=False) + + ndim = arr.ndim + arrorder = 'F' if arr.flags.fnc else 'C' + if axis is None: + if ndim != 1: + arr = arr.ravel() + # needed for np.matrix, which is still not 1d after being ravelled + ndim = arr.ndim + axis = ndim - 1 + else: + axis = normalize_axis_index(axis, ndim) + slobj = [slice(None)]*ndim + N = arr.shape[axis] + newshape = list(arr.shape) + + if isinstance(obj, slice): + # turn it into a range object + indices = arange(*obj.indices(N), dtype=intp) + else: + # need to copy obj, because indices will be changed in-place + indices = np.array(obj) + if indices.dtype == bool: + if obj.ndim != 1: + raise ValueError('boolean array argument obj to insert ' + 'must be one dimensional') + indices = np.flatnonzero(obj) + elif indices.ndim > 1: + raise ValueError( + "index array argument obj to insert must be one dimensional " + "or scalar") + if indices.size == 1: + index = indices.item() + if index < -N or index > N: + raise IndexError(f"index {obj} is out of bounds for axis {axis} " + f"with size {N}") + if (index < 0): + index += N + + # There are some object array corner cases here, but we cannot avoid + # that: + values = array(values, copy=None, ndmin=arr.ndim, dtype=arr.dtype) + if indices.ndim == 0: + # broadcasting is very different here, since a[:,0,:] = ... behaves + # very different from a[:,[0],:] = ...! This changes values so that + # it works likes the second case. (here a[:,0:1,:]) + values = np.moveaxis(values, 0, axis) + numnew = values.shape[axis] + newshape[axis] += numnew + new = empty(newshape, arr.dtype, arrorder) + slobj[axis] = slice(None, index) + new[tuple(slobj)] = arr[tuple(slobj)] + slobj[axis] = slice(index, index+numnew) + new[tuple(slobj)] = values + slobj[axis] = slice(index+numnew, None) + slobj2 = [slice(None)] * ndim + slobj2[axis] = slice(index, None) + new[tuple(slobj)] = arr[tuple(slobj2)] + + return conv.wrap(new, to_scalar=False) + + elif indices.size == 0 and not isinstance(obj, np.ndarray): + # Can safely cast the empty list to intp + indices = indices.astype(intp) + + indices[indices < 0] += N + + numnew = len(indices) + order = indices.argsort(kind='mergesort') # stable sort + indices[order] += np.arange(numnew) + + newshape[axis] += numnew + old_mask = ones(newshape[axis], dtype=bool) + old_mask[indices] = False + + new = empty(newshape, arr.dtype, arrorder) + slobj2 = [slice(None)]*ndim + slobj[axis] = indices + slobj2[axis] = old_mask + new[tuple(slobj)] = values + new[tuple(slobj2)] = arr + + return conv.wrap(new, to_scalar=False) + + +def _append_dispatcher(arr, values, axis=None): + return (arr, values) + + +@array_function_dispatch(_append_dispatcher) +def append(arr, values, axis=None): + """ + Append values to the end of an array. + + Parameters + ---------- + arr : array_like + Values are appended to a copy of this array. + values : array_like + These values are appended to a copy of `arr`. It must be of the + correct shape (the same shape as `arr`, excluding `axis`). If + `axis` is not specified, `values` can be any shape and will be + flattened before use. + axis : int, optional + The axis along which `values` are appended. If `axis` is not + given, both `arr` and `values` are flattened before use. + + Returns + ------- + append : ndarray + A copy of `arr` with `values` appended to `axis`. Note that + `append` does not occur in-place: a new array is allocated and + filled. If `axis` is None, `out` is a flattened array. + + See Also + -------- + insert : Insert elements into an array. + delete : Delete elements from an array. + + Examples + -------- + >>> import numpy as np + >>> np.append([1, 2, 3], [[4, 5, 6], [7, 8, 9]]) + array([1, 2, 3, ..., 7, 8, 9]) + + When `axis` is specified, `values` must have the correct shape. + + >>> np.append([[1, 2, 3], [4, 5, 6]], [[7, 8, 9]], axis=0) + array([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + + >>> np.append([[1, 2, 3], [4, 5, 6]], [7, 8, 9], axis=0) + Traceback (most recent call last): + ... + ValueError: all the input arrays must have same number of dimensions, but + the array at index 0 has 2 dimension(s) and the array at index 1 has 1 + dimension(s) + + >>> a = np.array([1, 2], dtype=int) + >>> c = np.append(a, []) + >>> c + array([1., 2.]) + >>> c.dtype + float64 + + Default dtype for empty ndarrays is `float64` thus making the output of dtype + `float64` when appended with dtype `int64` + + """ + arr = asanyarray(arr) + if axis is None: + if arr.ndim != 1: + arr = arr.ravel() + values = ravel(values) + axis = arr.ndim-1 + return concatenate((arr, values), axis=axis) + + +def _digitize_dispatcher(x, bins, right=None): + return (x, bins) + + +@array_function_dispatch(_digitize_dispatcher) +def digitize(x, bins, right=False): + """ + Return the indices of the bins to which each value in input array belongs. + + ========= ============= ============================ + `right` order of bins returned index `i` satisfies + ========= ============= ============================ + ``False`` increasing ``bins[i-1] <= x < bins[i]`` + ``True`` increasing ``bins[i-1] < x <= bins[i]`` + ``False`` decreasing ``bins[i-1] > x >= bins[i]`` + ``True`` decreasing ``bins[i-1] >= x > bins[i]`` + ========= ============= ============================ + + If values in `x` are beyond the bounds of `bins`, 0 or ``len(bins)`` is + returned as appropriate. + + Parameters + ---------- + x : array_like + Input array to be binned. Prior to NumPy 1.10.0, this array had to + be 1-dimensional, but can now have any shape. + bins : array_like + Array of bins. It has to be 1-dimensional and monotonic. + right : bool, optional + Indicating whether the intervals include the right or the left bin + edge. Default behavior is (right==False) indicating that the interval + does not include the right edge. The left bin end is open in this + case, i.e., bins[i-1] <= x < bins[i] is the default behavior for + monotonically increasing bins. + + Returns + ------- + indices : ndarray of ints + Output array of indices, of same shape as `x`. + + Raises + ------ + ValueError + If `bins` is not monotonic. + TypeError + If the type of the input is complex. + + See Also + -------- + bincount, histogram, unique, searchsorted + + Notes + ----- + If values in `x` are such that they fall outside the bin range, + attempting to index `bins` with the indices that `digitize` returns + will result in an IndexError. + + .. versionadded:: 1.10.0 + + `numpy.digitize` is implemented in terms of `numpy.searchsorted`. + This means that a binary search is used to bin the values, which scales + much better for larger number of bins than the previous linear search. + It also removes the requirement for the input array to be 1-dimensional. + + For monotonically *increasing* `bins`, the following are equivalent:: + + np.digitize(x, bins, right=True) + np.searchsorted(bins, x, side='left') + + Note that as the order of the arguments are reversed, the side must be too. + The `searchsorted` call is marginally faster, as it does not do any + monotonicity checks. Perhaps more importantly, it supports all dtypes. + + Examples + -------- + >>> import numpy as np + >>> x = np.array([0.2, 6.4, 3.0, 1.6]) + >>> bins = np.array([0.0, 1.0, 2.5, 4.0, 10.0]) + >>> inds = np.digitize(x, bins) + >>> inds + array([1, 4, 3, 2]) + >>> for n in range(x.size): + ... print(bins[inds[n]-1], "<=", x[n], "<", bins[inds[n]]) + ... + 0.0 <= 0.2 < 1.0 + 4.0 <= 6.4 < 10.0 + 2.5 <= 3.0 < 4.0 + 1.0 <= 1.6 < 2.5 + + >>> x = np.array([1.2, 10.0, 12.4, 15.5, 20.]) + >>> bins = np.array([0, 5, 10, 15, 20]) + >>> np.digitize(x,bins,right=True) + array([1, 2, 3, 4, 4]) + >>> np.digitize(x,bins,right=False) + array([1, 3, 3, 4, 5]) + """ + x = _nx.asarray(x) + bins = _nx.asarray(bins) + + # here for compatibility, searchsorted below is happy to take this + if np.issubdtype(x.dtype, _nx.complexfloating): + raise TypeError("x may not be complex") + + mono = _monotonicity(bins) + if mono == 0: + raise ValueError("bins must be monotonically increasing or decreasing") + + # this is backwards because the arguments below are swapped + side = 'left' if right else 'right' + if mono == -1: + # reverse the bins, and invert the results + return len(bins) - _nx.searchsorted(bins[::-1], x, side=side) + else: + return _nx.searchsorted(bins, x, side=side) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e98dcbb7e74155442daead8e16b6f888c10617fb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_function_base_impl.pyi @@ -0,0 +1,894 @@ +# ruff: noqa: ANN401 +from collections.abc import Callable, Iterable, Sequence +from typing import ( + Any, + Concatenate, + ParamSpec, + Protocol, + SupportsIndex, + SupportsInt, + TypeAlias, + TypeVar, + overload, + type_check_only, +) +from typing import Literal as L + +from _typeshed import Incomplete +from typing_extensions import TypeIs, deprecated + +import numpy as np +from numpy import ( + _OrderKACF, + bool_, + complex128, + complexfloating, + datetime64, + float64, + floating, + generic, + integer, + intp, + object_, + timedelta64, + vectorize, +) +from numpy._core.multiarray import bincount +from numpy._globals import _NoValueType +from numpy._typing import ( + ArrayLike, + DTypeLike, + NDArray, + _ArrayLike, + _ArrayLikeBool_co, + _ArrayLikeComplex_co, + _ArrayLikeDT64_co, + _ArrayLikeFloat_co, + _ArrayLikeInt_co, + _ArrayLikeNumber_co, + _ArrayLikeObject_co, + _ArrayLikeTD64_co, + _ComplexLike_co, + _DTypeLike, + _FloatLike_co, + _NestedSequence, + _NumberLike_co, + _ScalarLike_co, + _ShapeLike, +) + +__all__ = [ + "select", + "piecewise", + "trim_zeros", + "copy", + "iterable", + "percentile", + "diff", + "gradient", + "angle", + "unwrap", + "sort_complex", + "flip", + "rot90", + "extract", + "place", + "vectorize", + "asarray_chkfinite", + "average", + "bincount", + "digitize", + "cov", + "corrcoef", + "median", + "sinc", + "hamming", + "hanning", + "bartlett", + "blackman", + "kaiser", + "trapezoid", + "trapz", + "i0", + "meshgrid", + "delete", + "insert", + "append", + "interp", + "quantile", +] + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +# The `{}ss` suffix refers to the Python 3.12 syntax: `**P` +_Pss = ParamSpec("_Pss") +_SCT = TypeVar("_SCT", bound=generic) +_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) + +_2Tuple: TypeAlias = tuple[_T, _T] + +@type_check_only +class _TrimZerosSequence(Protocol[_T_co]): + def __len__(self, /) -> int: ... + @overload + def __getitem__(self, key: int, /) -> object: ... + @overload + def __getitem__(self, key: slice, /) -> _T_co: ... + +### + +@overload +def rot90( + m: _ArrayLike[_SCT], + k: int = ..., + axes: tuple[int, int] = ..., +) -> NDArray[_SCT]: ... +@overload +def rot90( + m: ArrayLike, + k: int = ..., + axes: tuple[int, int] = ..., +) -> NDArray[Any]: ... + +@overload +def flip(m: _SCT, axis: None = ...) -> _SCT: ... +@overload +def flip(m: _ScalarLike_co, axis: None = ...) -> Any: ... +@overload +def flip(m: _ArrayLike[_SCT], axis: None | _ShapeLike = ...) -> NDArray[_SCT]: ... +@overload +def flip(m: ArrayLike, axis: None | _ShapeLike = ...) -> NDArray[Any]: ... + +def iterable(y: object) -> TypeIs[Iterable[Any]]: ... + +@overload +def average( + a: _ArrayLikeFloat_co, + axis: None = None, + weights: _ArrayLikeFloat_co | None = None, + returned: L[False] = False, + *, + keepdims: L[False] | _NoValueType = ..., +) -> floating: ... +@overload +def average( + a: _ArrayLikeFloat_co, + axis: None = None, + weights: _ArrayLikeFloat_co | None = None, + *, + returned: L[True], + keepdims: L[False] | _NoValueType = ..., +) -> _2Tuple[floating]: ... +@overload +def average( + a: _ArrayLikeComplex_co, + axis: None = None, + weights: _ArrayLikeComplex_co | None = None, + returned: L[False] = False, + *, + keepdims: L[False] | _NoValueType = ..., +) -> complexfloating: ... +@overload +def average( + a: _ArrayLikeComplex_co, + axis: None = None, + weights: _ArrayLikeComplex_co | None = None, + *, + returned: L[True], + keepdims: L[False] | _NoValueType = ..., +) -> _2Tuple[complexfloating]: ... +@overload +def average( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = None, + weights: object | None = None, + *, + returned: L[True], + keepdims: bool | bool_ | _NoValueType = ..., +) -> _2Tuple[Incomplete]: ... +@overload +def average( + a: _ArrayLikeComplex_co | _ArrayLikeObject_co, + axis: _ShapeLike | None = None, + weights: object | None = None, + returned: bool | bool_ = False, + *, + keepdims: bool | bool_ | _NoValueType = ..., +) -> Incomplete: ... + +@overload +def asarray_chkfinite( + a: _ArrayLike[_SCT], + dtype: None = ..., + order: _OrderKACF = ..., +) -> NDArray[_SCT]: ... +@overload +def asarray_chkfinite( + a: object, + dtype: None = ..., + order: _OrderKACF = ..., +) -> NDArray[Any]: ... +@overload +def asarray_chkfinite( + a: Any, + dtype: _DTypeLike[_SCT], + order: _OrderKACF = ..., +) -> NDArray[_SCT]: ... +@overload +def asarray_chkfinite( + a: Any, + dtype: DTypeLike, + order: _OrderKACF = ..., +) -> NDArray[Any]: ... + +@overload +def piecewise( + x: _ArrayLike[_SCT], + condlist: _ArrayLike[bool_] | Sequence[_ArrayLikeBool_co], + funclist: Sequence[ + Callable[Concatenate[NDArray[_SCT], _Pss], NDArray[_SCT | Any]] + | _SCT | object + ], + /, + *args: _Pss.args, + **kw: _Pss.kwargs, +) -> NDArray[_SCT]: ... +@overload +def piecewise( + x: ArrayLike, + condlist: _ArrayLike[bool_] | Sequence[_ArrayLikeBool_co], + funclist: Sequence[ + Callable[Concatenate[NDArray[Any], _Pss], NDArray[Any]] + | object + ], + /, + *args: _Pss.args, + **kw: _Pss.kwargs, +) -> NDArray[Any]: ... + +def select( + condlist: Sequence[ArrayLike], + choicelist: Sequence[ArrayLike], + default: ArrayLike = ..., +) -> NDArray[Any]: ... + +@overload +def copy( + a: _ArrayType, + order: _OrderKACF, + subok: L[True], +) -> _ArrayType: ... +@overload +def copy( + a: _ArrayType, + order: _OrderKACF = ..., + *, + subok: L[True], +) -> _ArrayType: ... +@overload +def copy( + a: _ArrayLike[_SCT], + order: _OrderKACF = ..., + subok: L[False] = ..., +) -> NDArray[_SCT]: ... +@overload +def copy( + a: ArrayLike, + order: _OrderKACF = ..., + subok: L[False] = ..., +) -> NDArray[Any]: ... + +def gradient( + f: ArrayLike, + *varargs: ArrayLike, + axis: None | _ShapeLike = ..., + edge_order: L[1, 2] = ..., +) -> Any: ... + +@overload +def diff( + a: _T, + n: L[0], + axis: SupportsIndex = ..., + prepend: ArrayLike = ..., + append: ArrayLike = ..., +) -> _T: ... +@overload +def diff( + a: ArrayLike, + n: int = ..., + axis: SupportsIndex = ..., + prepend: ArrayLike = ..., + append: ArrayLike = ..., +) -> NDArray[Any]: ... + +@overload # float scalar +def interp( + x: _FloatLike_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeFloat_co, + left: _FloatLike_co | None = None, + right: _FloatLike_co | None = None, + period: _FloatLike_co | None = None, +) -> float64: ... +@overload # float array +def interp( + x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co], + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeFloat_co, + left: _FloatLike_co | None = None, + right: _FloatLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[float64]: ... +@overload # float scalar or array +def interp( + x: _ArrayLikeFloat_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeFloat_co, + left: _FloatLike_co | None = None, + right: _FloatLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[float64] | float64: ... +@overload # complex scalar +def interp( + x: _FloatLike_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLike[complexfloating], + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> complex128: ... +@overload # complex or float scalar +def interp( + x: _FloatLike_co, + xp: _ArrayLikeFloat_co, + fp: Sequence[complex | complexfloating], + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> complex128 | float64: ... +@overload # complex array +def interp( + x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co], + xp: _ArrayLikeFloat_co, + fp: _ArrayLike[complexfloating], + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[complex128]: ... +@overload # complex or float array +def interp( + x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co], + xp: _ArrayLikeFloat_co, + fp: Sequence[complex | complexfloating], + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[complex128 | float64]: ... +@overload # complex scalar or array +def interp( + x: _ArrayLikeFloat_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLike[complexfloating], + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[complex128] | complex128: ... +@overload # complex or float scalar or array +def interp( + x: _ArrayLikeFloat_co, + xp: _ArrayLikeFloat_co, + fp: _ArrayLikeNumber_co, + left: _NumberLike_co | None = None, + right: _NumberLike_co | None = None, + period: _FloatLike_co | None = None, +) -> NDArray[complex128 | float64] | complex128 | float64: ... + +@overload +def angle(z: _ComplexLike_co, deg: bool = ...) -> floating[Any]: ... +@overload +def angle(z: object_, deg: bool = ...) -> Any: ... +@overload +def angle(z: _ArrayLikeComplex_co, deg: bool = ...) -> NDArray[floating[Any]]: ... +@overload +def angle(z: _ArrayLikeObject_co, deg: bool = ...) -> NDArray[object_]: ... + +@overload +def unwrap( + p: _ArrayLikeFloat_co, + discont: None | float = ..., + axis: int = ..., + *, + period: float = ..., +) -> NDArray[floating[Any]]: ... +@overload +def unwrap( + p: _ArrayLikeObject_co, + discont: None | float = ..., + axis: int = ..., + *, + period: float = ..., +) -> NDArray[object_]: ... + +def sort_complex(a: ArrayLike) -> NDArray[complexfloating[Any, Any]]: ... + +def trim_zeros( + filt: _TrimZerosSequence[_T], + trim: L["f", "b", "fb", "bf"] = ..., +) -> _T: ... + +@overload +def extract(condition: ArrayLike, arr: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def extract(condition: ArrayLike, arr: ArrayLike) -> NDArray[Any]: ... + +def place(arr: NDArray[Any], mask: ArrayLike, vals: Any) -> None: ... + +@overload +def cov( + m: _ArrayLikeFloat_co, + y: None | _ArrayLikeFloat_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: None = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: _DTypeLike[_SCT], +) -> NDArray[_SCT]: ... +@overload +def cov( + m: _ArrayLikeComplex_co, + y: None | _ArrayLikeComplex_co = ..., + rowvar: bool = ..., + bias: bool = ..., + ddof: None | SupportsIndex | SupportsInt = ..., + fweights: None | ArrayLike = ..., + aweights: None | ArrayLike = ..., + *, + dtype: DTypeLike, +) -> NDArray[Any]: ... + +# NOTE `bias` and `ddof` are deprecated and ignored +@overload +def corrcoef( + m: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co | None = None, + rowvar: bool = True, + bias: _NoValueType = ..., + ddof: _NoValueType = ..., + *, + dtype: None = None, +) -> NDArray[floating]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co | None = None, + rowvar: bool = True, + bias: _NoValueType = ..., + ddof: _NoValueType = ..., + *, + dtype: None = None, +) -> NDArray[complexfloating]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co | None = None, + rowvar: bool = True, + bias: _NoValueType = ..., + ddof: _NoValueType = ..., + *, + dtype: _DTypeLike[_SCT], +) -> NDArray[_SCT]: ... +@overload +def corrcoef( + m: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co | None = None, + rowvar: bool = True, + bias: _NoValueType = ..., + ddof: _NoValueType = ..., + *, + dtype: DTypeLike | None = None, +) -> NDArray[Any]: ... + +def blackman(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def bartlett(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def hanning(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def hamming(M: _FloatLike_co) -> NDArray[floating[Any]]: ... + +def i0(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... + +def kaiser( + M: _FloatLike_co, + beta: _FloatLike_co, +) -> NDArray[floating[Any]]: ... + +@overload +def sinc(x: _FloatLike_co) -> floating[Any]: ... +@overload +def sinc(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def sinc(x: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... +@overload +def sinc(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def median( + a: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> floating[Any]: ... +@overload +def median( + a: _ArrayLikeComplex_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> complexfloating[Any, Any]: ... +@overload +def median( + a: _ArrayLikeTD64_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> timedelta64: ... +@overload +def median( + a: _ArrayLikeObject_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: L[False] = ..., +) -> Any: ... +@overload +def median( + a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + axis: None | _ShapeLike = ..., + out: None = ..., + overwrite_input: bool = ..., + keepdims: bool = ..., +) -> Any: ... +@overload +def median( + a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + axis: None | _ShapeLike, + out: _ArrayType, + overwrite_input: bool = ..., + keepdims: bool = ..., +) -> _ArrayType: ... +@overload +def median( + a: _ArrayLikeFloat_co | _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + axis: None | _ShapeLike = ..., + *, + out: _ArrayType, + overwrite_input: bool = ..., + keepdims: bool = ..., +) -> _ArrayType: ... + +_MethodKind = L[ + "inverted_cdf", + "averaged_inverted_cdf", + "closest_observation", + "interpolated_inverted_cdf", + "hazen", + "weibull", + "linear", + "median_unbiased", + "normal_unbiased", + "lower", + "higher", + "midpoint", + "nearest", +] + +@overload +def percentile( + a: _ArrayLikeFloat_co, + q: _FloatLike_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> floating[Any]: ... +@overload +def percentile( + a: _ArrayLikeComplex_co, + q: _FloatLike_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> complexfloating[Any, Any]: ... +@overload +def percentile( + a: _ArrayLikeTD64_co, + q: _FloatLike_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> timedelta64: ... +@overload +def percentile( + a: _ArrayLikeDT64_co, + q: _FloatLike_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> datetime64: ... +@overload +def percentile( + a: _ArrayLikeObject_co, + q: _FloatLike_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> Any: ... +@overload +def percentile( + a: _ArrayLikeFloat_co, + q: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> NDArray[floating[Any]]: ... +@overload +def percentile( + a: _ArrayLikeComplex_co, + q: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def percentile( + a: _ArrayLikeTD64_co, + q: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> NDArray[timedelta64]: ... +@overload +def percentile( + a: _ArrayLikeDT64_co, + q: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> NDArray[datetime64]: ... +@overload +def percentile( + a: _ArrayLikeObject_co, + q: _ArrayLikeFloat_co, + axis: None = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: L[False] = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> NDArray[object_]: ... +@overload +def percentile( + a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + q: _ArrayLikeFloat_co, + axis: None | _ShapeLike = ..., + out: None = ..., + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: bool = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> Any: ... +@overload +def percentile( + a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + q: _ArrayLikeFloat_co, + axis: None | _ShapeLike, + out: _ArrayType, + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: bool = ..., + *, + weights: None | _ArrayLikeFloat_co = ..., +) -> _ArrayType: ... +@overload +def percentile( + a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + q: _ArrayLikeFloat_co, + axis: None | _ShapeLike = ..., + *, + out: _ArrayType, + overwrite_input: bool = ..., + method: _MethodKind = ..., + keepdims: bool = ..., + weights: None | _ArrayLikeFloat_co = ..., +) -> _ArrayType: ... + +# NOTE: Not an alias, but they do have identical signatures +# (that we can reuse) +quantile = percentile + + +_SCT_fm = TypeVar( + "_SCT_fm", + bound=floating[Any] | complexfloating[Any, Any] | timedelta64, +) + +class _SupportsRMulFloat(Protocol[_T_co]): + def __rmul__(self, other: float, /) -> _T_co: ... + +@overload +def trapezoid( # type: ignore[overload-overlap] + y: Sequence[_FloatLike_co], + x: Sequence[_FloatLike_co] | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> float64: ... +@overload +def trapezoid( + y: Sequence[_ComplexLike_co], + x: Sequence[_ComplexLike_co] | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> complex128: ... +@overload +def trapezoid( + y: _ArrayLike[bool_ | integer[Any]], + x: _ArrayLike[bool_ | integer[Any]] | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> float64 | NDArray[float64]: ... +@overload +def trapezoid( # type: ignore[overload-overlap] + y: _ArrayLikeObject_co, + x: _ArrayLikeFloat_co | _ArrayLikeObject_co | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> float | NDArray[object_]: ... +@overload +def trapezoid( + y: _ArrayLike[_SCT_fm], + x: _ArrayLike[_SCT_fm] | _ArrayLikeInt_co | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> _SCT_fm | NDArray[_SCT_fm]: ... +@overload +def trapezoid( + y: Sequence[_SupportsRMulFloat[_T]], + x: Sequence[_SupportsRMulFloat[_T] | _T] | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> _T: ... +@overload +def trapezoid( + y: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, + x: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co | None = ..., + dx: float = ..., + axis: SupportsIndex = ..., +) -> ( + floating[Any] | complexfloating[Any, Any] | timedelta64 + | NDArray[floating[Any] | complexfloating[Any, Any] | timedelta64 | object_] +): ... + +@deprecated("Use 'trapezoid' instead") +def trapz(y: ArrayLike, x: ArrayLike | None = None, dx: float = 1.0, axis: int = -1) -> generic | NDArray[generic]: ... + +def meshgrid( + *xi: ArrayLike, + copy: bool = ..., + sparse: bool = ..., + indexing: L["xy", "ij"] = ..., +) -> tuple[NDArray[Any], ...]: ... + +@overload +def delete( + arr: _ArrayLike[_SCT], + obj: slice | _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def delete( + arr: ArrayLike, + obj: slice | _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., +) -> NDArray[Any]: ... + +@overload +def insert( + arr: _ArrayLike[_SCT], + obj: slice | _ArrayLikeInt_co, + values: ArrayLike, + axis: None | SupportsIndex = ..., +) -> NDArray[_SCT]: ... +@overload +def insert( + arr: ArrayLike, + obj: slice | _ArrayLikeInt_co, + values: ArrayLike, + axis: None | SupportsIndex = ..., +) -> NDArray[Any]: ... + +def append( + arr: ArrayLike, + values: ArrayLike, + axis: None | SupportsIndex = ..., +) -> NDArray[Any]: ... + +@overload +def digitize( + x: _FloatLike_co, + bins: _ArrayLikeFloat_co, + right: bool = ..., +) -> intp: ... +@overload +def digitize( + x: _ArrayLikeFloat_co, + bins: _ArrayLikeFloat_co, + right: bool = ..., +) -> NDArray[intp]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..b361bb4f91ac0723f65399cd6b73048649179c2b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.py @@ -0,0 +1,1090 @@ +""" +Histogram-related functions +""" +import contextlib +import functools +import operator +import warnings + +import numpy as np +from numpy._core import overrides + +__all__ = ['histogram', 'histogramdd', 'histogram_bin_edges'] + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + +# range is a keyword argument to many functions, so save the builtin so they can +# use it. +_range = range + + +def _ptp(x): + """Peak-to-peak value of x. + + This implementation avoids the problem of signed integer arrays having a + peak-to-peak value that cannot be represented with the array's data type. + This function returns an unsigned value for signed integer arrays. + """ + return _unsigned_subtract(x.max(), x.min()) + + +def _hist_bin_sqrt(x, range): + """ + Square root histogram bin estimator. + + Bin width is inversely proportional to the data size. Used by many + programs for its simplicity. + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + return _ptp(x) / np.sqrt(x.size) + + +def _hist_bin_sturges(x, range): + """ + Sturges histogram bin estimator. + + A very simplistic estimator based on the assumption of normality of + the data. This estimator has poor performance for non-normal data, + which becomes especially obvious for large data sets. The estimate + depends only on size of the data. + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + return _ptp(x) / (np.log2(x.size) + 1.0) + + +def _hist_bin_rice(x, range): + """ + Rice histogram bin estimator. + + Another simple estimator with no normality assumption. It has better + performance for large data than Sturges, but tends to overestimate + the number of bins. The number of bins is proportional to the cube + root of data size (asymptotically optimal). The estimate depends + only on size of the data. + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + return _ptp(x) / (2.0 * x.size ** (1.0 / 3)) + + +def _hist_bin_scott(x, range): + """ + Scott histogram bin estimator. + + The binwidth is proportional to the standard deviation of the data + and inversely proportional to the cube root of data size + (asymptotically optimal). + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + return (24.0 * np.pi**0.5 / x.size)**(1.0 / 3.0) * np.std(x) + + +def _hist_bin_stone(x, range): + """ + Histogram bin estimator based on minimizing the estimated integrated squared error (ISE). + + The number of bins is chosen by minimizing the estimated ISE against the unknown true distribution. + The ISE is estimated using cross-validation and can be regarded as a generalization of Scott's rule. + https://en.wikipedia.org/wiki/Histogram#Scott.27s_normal_reference_rule + + This paper by Stone appears to be the origination of this rule. + https://digitalassets.lib.berkeley.edu/sdtr/ucb/text/34.pdf + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + range : (float, float) + The lower and upper range of the bins. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + + n = x.size + ptp_x = _ptp(x) + if n <= 1 or ptp_x == 0: + return 0 + + def jhat(nbins): + hh = ptp_x / nbins + p_k = np.histogram(x, bins=nbins, range=range)[0] / n + return (2 - (n + 1) * p_k.dot(p_k)) / hh + + nbins_upper_bound = max(100, int(np.sqrt(n))) + nbins = min(_range(1, nbins_upper_bound + 1), key=jhat) + if nbins == nbins_upper_bound: + warnings.warn("The number of bins estimated may be suboptimal.", + RuntimeWarning, stacklevel=3) + return ptp_x / nbins + + +def _hist_bin_doane(x, range): + """ + Doane's histogram bin estimator. + + Improved version of Sturges' formula which works better for + non-normal data. See + stats.stackexchange.com/questions/55134/doanes-formula-for-histogram-binning + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + if x.size > 2: + sg1 = np.sqrt(6.0 * (x.size - 2) / ((x.size + 1.0) * (x.size + 3))) + sigma = np.std(x) + if sigma > 0.0: + # These three operations add up to + # g1 = np.mean(((x - np.mean(x)) / sigma)**3) + # but use only one temp array instead of three + temp = x - np.mean(x) + np.true_divide(temp, sigma, temp) + np.power(temp, 3, temp) + g1 = np.mean(temp) + return _ptp(x) / (1.0 + np.log2(x.size) + + np.log2(1.0 + np.absolute(g1) / sg1)) + return 0.0 + + +def _hist_bin_fd(x, range): + """ + The Freedman-Diaconis histogram bin estimator. + + The Freedman-Diaconis rule uses interquartile range (IQR) to + estimate binwidth. It is considered a variation of the Scott rule + with more robustness as the IQR is less affected by outliers than + the standard deviation. However, the IQR depends on fewer points + than the standard deviation, so it is less accurate, especially for + long tailed distributions. + + If the IQR is 0, this function returns 0 for the bin width. + Binwidth is inversely proportional to the cube root of data size + (asymptotically optimal). + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + """ + del range # unused + iqr = np.subtract(*np.percentile(x, [75, 25])) + return 2.0 * iqr * x.size ** (-1.0 / 3.0) + + +def _hist_bin_auto(x, range): + """ + Histogram bin estimator that uses the minimum width of the + Freedman-Diaconis and Sturges estimators if the FD bin width is non-zero. + If the bin width from the FD estimator is 0, the Sturges estimator is used. + + The FD estimator is usually the most robust method, but its width + estimate tends to be too large for small `x` and bad for data with limited + variance. The Sturges estimator is quite good for small (<1000) datasets + and is the default in the R language. This method gives good off-the-shelf + behaviour. + + If there is limited variance the IQR can be 0, which results in the + FD bin width being 0 too. This is not a valid bin width, so + ``np.histogram_bin_edges`` chooses 1 bin instead, which may not be optimal. + If the IQR is 0, it's unlikely any variance-based estimators will be of + use, so we revert to the Sturges estimator, which only uses the size of the + dataset in its calculation. + + Parameters + ---------- + x : array_like + Input data that is to be histogrammed, trimmed to range. May not + be empty. + + Returns + ------- + h : An estimate of the optimal bin width for the given data. + + See Also + -------- + _hist_bin_fd, _hist_bin_sturges + """ + fd_bw = _hist_bin_fd(x, range) + sturges_bw = _hist_bin_sturges(x, range) + del range # unused + if fd_bw: + return min(fd_bw, sturges_bw) + else: + # limited variance, so we return a len dependent bw estimator + return sturges_bw + +# Private dict initialized at module load time +_hist_bin_selectors = {'stone': _hist_bin_stone, + 'auto': _hist_bin_auto, + 'doane': _hist_bin_doane, + 'fd': _hist_bin_fd, + 'rice': _hist_bin_rice, + 'scott': _hist_bin_scott, + 'sqrt': _hist_bin_sqrt, + 'sturges': _hist_bin_sturges} + + +def _ravel_and_check_weights(a, weights): + """ Check a and weights have matching shapes, and ravel both """ + a = np.asarray(a) + + # Ensure that the array is a "subtractable" dtype + if a.dtype == np.bool: + warnings.warn("Converting input from {} to {} for compatibility." + .format(a.dtype, np.uint8), + RuntimeWarning, stacklevel=3) + a = a.astype(np.uint8) + + if weights is not None: + weights = np.asarray(weights) + if weights.shape != a.shape: + raise ValueError( + 'weights should have the same shape as a.') + weights = weights.ravel() + a = a.ravel() + return a, weights + + +def _get_outer_edges(a, range): + """ + Determine the outer bin edges to use, from either the data or the range + argument + """ + if range is not None: + first_edge, last_edge = range + if first_edge > last_edge: + raise ValueError( + 'max must be larger than min in range parameter.') + if not (np.isfinite(first_edge) and np.isfinite(last_edge)): + raise ValueError( + "supplied range of [{}, {}] is not finite".format(first_edge, last_edge)) + elif a.size == 0: + # handle empty arrays. Can't determine range, so use 0-1. + first_edge, last_edge = 0, 1 + else: + first_edge, last_edge = a.min(), a.max() + if not (np.isfinite(first_edge) and np.isfinite(last_edge)): + raise ValueError( + "autodetected range of [{}, {}] is not finite".format(first_edge, last_edge)) + + # expand empty range to avoid divide by zero + if first_edge == last_edge: + first_edge = first_edge - 0.5 + last_edge = last_edge + 0.5 + + return first_edge, last_edge + + +def _unsigned_subtract(a, b): + """ + Subtract two values where a >= b, and produce an unsigned result + + This is needed when finding the difference between the upper and lower + bound of an int16 histogram + """ + # coerce to a single type + signed_to_unsigned = { + np.byte: np.ubyte, + np.short: np.ushort, + np.intc: np.uintc, + np.int_: np.uint, + np.longlong: np.ulonglong + } + dt = np.result_type(a, b) + try: + unsigned_dt = signed_to_unsigned[dt.type] + except KeyError: + return np.subtract(a, b, dtype=dt) + else: + # we know the inputs are integers, and we are deliberately casting + # signed to unsigned. The input may be negative python integers so + # ensure we pass in arrays with the initial dtype (related to NEP 50). + return np.subtract(np.asarray(a, dtype=dt), np.asarray(b, dtype=dt), + casting='unsafe', dtype=unsigned_dt) + + +def _get_bin_edges(a, bins, range, weights): + """ + Computes the bins used internally by `histogram`. + + Parameters + ========== + a : ndarray + Ravelled data array + bins, range + Forwarded arguments from `histogram`. + weights : ndarray, optional + Ravelled weights array, or None + + Returns + ======= + bin_edges : ndarray + Array of bin edges + uniform_bins : (Number, Number, int): + The upper bound, lowerbound, and number of bins, used in the optimized + implementation of `histogram` that works on uniform bins. + """ + # parse the overloaded bins argument + n_equal_bins = None + bin_edges = None + + if isinstance(bins, str): + bin_name = bins + # if `bins` is a string for an automatic method, + # this will replace it with the number of bins calculated + if bin_name not in _hist_bin_selectors: + raise ValueError( + "{!r} is not a valid estimator for `bins`".format(bin_name)) + if weights is not None: + raise TypeError("Automated estimation of the number of " + "bins is not supported for weighted data") + + first_edge, last_edge = _get_outer_edges(a, range) + + # truncate the range if needed + if range is not None: + keep = (a >= first_edge) + keep &= (a <= last_edge) + if not np.logical_and.reduce(keep): + a = a[keep] + + if a.size == 0: + n_equal_bins = 1 + else: + # Do not call selectors on empty arrays + width = _hist_bin_selectors[bin_name](a, (first_edge, last_edge)) + if width: + if np.issubdtype(a.dtype, np.integer) and width < 1: + width = 1 + n_equal_bins = int(np.ceil(_unsigned_subtract(last_edge, first_edge) / width)) + else: + # Width can be zero for some estimators, e.g. FD when + # the IQR of the data is zero. + n_equal_bins = 1 + + elif np.ndim(bins) == 0: + try: + n_equal_bins = operator.index(bins) + except TypeError as e: + raise TypeError( + '`bins` must be an integer, a string, or an array') from e + if n_equal_bins < 1: + raise ValueError('`bins` must be positive, when an integer') + + first_edge, last_edge = _get_outer_edges(a, range) + + elif np.ndim(bins) == 1: + bin_edges = np.asarray(bins) + if np.any(bin_edges[:-1] > bin_edges[1:]): + raise ValueError( + '`bins` must increase monotonically, when an array') + + else: + raise ValueError('`bins` must be 1d, when an array') + + if n_equal_bins is not None: + # gh-10322 means that type resolution rules are dependent on array + # shapes. To avoid this causing problems, we pick a type now and stick + # with it throughout. + bin_type = np.result_type(first_edge, last_edge, a) + if np.issubdtype(bin_type, np.integer): + bin_type = np.result_type(bin_type, float) + + # bin edges must be computed + bin_edges = np.linspace( + first_edge, last_edge, n_equal_bins + 1, + endpoint=True, dtype=bin_type) + if np.any(bin_edges[:-1] >= bin_edges[1:]): + raise ValueError( + f'Too many bins for data range. Cannot create {n_equal_bins} ' + f'finite-sized bins.') + return bin_edges, (first_edge, last_edge, n_equal_bins) + else: + return bin_edges, None + + +def _search_sorted_inclusive(a, v): + """ + Like `searchsorted`, but where the last item in `v` is placed on the right. + + In the context of a histogram, this makes the last bin edge inclusive + """ + return np.concatenate(( + a.searchsorted(v[:-1], 'left'), + a.searchsorted(v[-1:], 'right') + )) + + +def _histogram_bin_edges_dispatcher(a, bins=None, range=None, weights=None): + return (a, bins, weights) + + +@array_function_dispatch(_histogram_bin_edges_dispatcher) +def histogram_bin_edges(a, bins=10, range=None, weights=None): + r""" + Function to calculate only the edges of the bins used by the `histogram` + function. + + Parameters + ---------- + a : array_like + Input data. The histogram is computed over the flattened array. + bins : int or sequence of scalars or str, optional + If `bins` is an int, it defines the number of equal-width + bins in the given range (10, by default). If `bins` is a + sequence, it defines the bin edges, including the rightmost + edge, allowing for non-uniform bin widths. + + If `bins` is a string from the list below, `histogram_bin_edges` will + use the method chosen to calculate the optimal bin width and + consequently the number of bins (see the Notes section for more detail + on the estimators) from the data that falls within the requested range. + While the bin width will be optimal for the actual data + in the range, the number of bins will be computed to fill the + entire range, including the empty portions. For visualisation, + using the 'auto' option is suggested. Weighted data is not + supported for automated bin size selection. + + 'auto' + Minimum bin width between the 'sturges' and 'fd' estimators. + Provides good all-around performance. + + 'fd' (Freedman Diaconis Estimator) + Robust (resilient to outliers) estimator that takes into + account data variability and data size. + + 'doane' + An improved version of Sturges' estimator that works better + with non-normal datasets. + + 'scott' + Less robust estimator that takes into account data variability + and data size. + + 'stone' + Estimator based on leave-one-out cross-validation estimate of + the integrated squared error. Can be regarded as a generalization + of Scott's rule. + + 'rice' + Estimator does not take variability into account, only data + size. Commonly overestimates number of bins required. + + 'sturges' + R's default method, only accounts for data size. Only + optimal for gaussian data and underestimates number of bins + for large non-gaussian datasets. + + 'sqrt' + Square root (of data size) estimator, used by Excel and + other programs for its speed and simplicity. + + range : (float, float), optional + The lower and upper range of the bins. If not provided, range + is simply ``(a.min(), a.max())``. Values outside the range are + ignored. The first element of the range must be less than or + equal to the second. `range` affects the automatic bin + computation as well. While bin width is computed to be optimal + based on the actual data within `range`, the bin count will fill + the entire range including portions containing no data. + + weights : array_like, optional + An array of weights, of the same shape as `a`. Each value in + `a` only contributes its associated weight towards the bin count + (instead of 1). This is currently not used by any of the bin estimators, + but may be in the future. + + Returns + ------- + bin_edges : array of dtype float + The edges to pass into `histogram` + + See Also + -------- + histogram + + Notes + ----- + The methods to estimate the optimal number of bins are well founded + in literature, and are inspired by the choices R provides for + histogram visualisation. Note that having the number of bins + proportional to :math:`n^{1/3}` is asymptotically optimal, which is + why it appears in most estimators. These are simply plug-in methods + that give good starting points for number of bins. In the equations + below, :math:`h` is the binwidth and :math:`n_h` is the number of + bins. All estimators that compute bin counts are recast to bin width + using the `ptp` of the data. The final bin count is obtained from + ``np.round(np.ceil(range / h))``. The final bin width is often less + than what is returned by the estimators below. + + 'auto' (minimum bin width of the 'sturges' and 'fd' estimators) + A compromise to get a good value. For small datasets the Sturges + value will usually be chosen, while larger datasets will usually + default to FD. Avoids the overly conservative behaviour of FD + and Sturges for small and large datasets respectively. + Switchover point is usually :math:`a.size \approx 1000`. + + 'fd' (Freedman Diaconis Estimator) + .. math:: h = 2 \frac{IQR}{n^{1/3}} + + The binwidth is proportional to the interquartile range (IQR) + and inversely proportional to cube root of a.size. Can be too + conservative for small datasets, but is quite good for large + datasets. The IQR is very robust to outliers. + + 'scott' + .. math:: h = \sigma \sqrt[3]{\frac{24 \sqrt{\pi}}{n}} + + The binwidth is proportional to the standard deviation of the + data and inversely proportional to cube root of ``x.size``. Can + be too conservative for small datasets, but is quite good for + large datasets. The standard deviation is not very robust to + outliers. Values are very similar to the Freedman-Diaconis + estimator in the absence of outliers. + + 'rice' + .. math:: n_h = 2n^{1/3} + + The number of bins is only proportional to cube root of + ``a.size``. It tends to overestimate the number of bins and it + does not take into account data variability. + + 'sturges' + .. math:: n_h = \log _{2}(n) + 1 + + The number of bins is the base 2 log of ``a.size``. This + estimator assumes normality of data and is too conservative for + larger, non-normal datasets. This is the default method in R's + ``hist`` method. + + 'doane' + .. math:: n_h = 1 + \log_{2}(n) + + \log_{2}\left(1 + \frac{|g_1|}{\sigma_{g_1}}\right) + + g_1 = mean\left[\left(\frac{x - \mu}{\sigma}\right)^3\right] + + \sigma_{g_1} = \sqrt{\frac{6(n - 2)}{(n + 1)(n + 3)}} + + An improved version of Sturges' formula that produces better + estimates for non-normal datasets. This estimator attempts to + account for the skew of the data. + + 'sqrt' + .. math:: n_h = \sqrt n + + The simplest and fastest estimator. Only takes into account the + data size. + + Additionally, if the data is of integer dtype, then the binwidth will never + be less than 1. + + Examples + -------- + >>> import numpy as np + >>> arr = np.array([0, 0, 0, 1, 2, 3, 3, 4, 5]) + >>> np.histogram_bin_edges(arr, bins='auto', range=(0, 1)) + array([0. , 0.25, 0.5 , 0.75, 1. ]) + >>> np.histogram_bin_edges(arr, bins=2) + array([0. , 2.5, 5. ]) + + For consistency with histogram, an array of pre-computed bins is + passed through unmodified: + + >>> np.histogram_bin_edges(arr, [1, 2]) + array([1, 2]) + + This function allows one set of bins to be computed, and reused across + multiple histograms: + + >>> shared_bins = np.histogram_bin_edges(arr, bins='auto') + >>> shared_bins + array([0., 1., 2., 3., 4., 5.]) + + >>> group_id = np.array([0, 1, 1, 0, 1, 1, 0, 1, 1]) + >>> hist_0, _ = np.histogram(arr[group_id == 0], bins=shared_bins) + >>> hist_1, _ = np.histogram(arr[group_id == 1], bins=shared_bins) + + >>> hist_0; hist_1 + array([1, 1, 0, 1, 0]) + array([2, 0, 1, 1, 2]) + + Which gives more easily comparable results than using separate bins for + each histogram: + + >>> hist_0, bins_0 = np.histogram(arr[group_id == 0], bins='auto') + >>> hist_1, bins_1 = np.histogram(arr[group_id == 1], bins='auto') + >>> hist_0; hist_1 + array([1, 1, 1]) + array([2, 1, 1, 2]) + >>> bins_0; bins_1 + array([0., 1., 2., 3.]) + array([0. , 1.25, 2.5 , 3.75, 5. ]) + + """ + a, weights = _ravel_and_check_weights(a, weights) + bin_edges, _ = _get_bin_edges(a, bins, range, weights) + return bin_edges + + +def _histogram_dispatcher( + a, bins=None, range=None, density=None, weights=None): + return (a, bins, weights) + + +@array_function_dispatch(_histogram_dispatcher) +def histogram(a, bins=10, range=None, density=None, weights=None): + r""" + Compute the histogram of a dataset. + + Parameters + ---------- + a : array_like + Input data. The histogram is computed over the flattened array. + bins : int or sequence of scalars or str, optional + If `bins` is an int, it defines the number of equal-width + bins in the given range (10, by default). If `bins` is a + sequence, it defines a monotonically increasing array of bin edges, + including the rightmost edge, allowing for non-uniform bin widths. + + If `bins` is a string, it defines the method used to calculate the + optimal bin width, as defined by `histogram_bin_edges`. + + range : (float, float), optional + The lower and upper range of the bins. If not provided, range + is simply ``(a.min(), a.max())``. Values outside the range are + ignored. The first element of the range must be less than or + equal to the second. `range` affects the automatic bin + computation as well. While bin width is computed to be optimal + based on the actual data within `range`, the bin count will fill + the entire range including portions containing no data. + weights : array_like, optional + An array of weights, of the same shape as `a`. Each value in + `a` only contributes its associated weight towards the bin count + (instead of 1). If `density` is True, the weights are + normalized, so that the integral of the density over the range + remains 1. + Please note that the ``dtype`` of `weights` will also become the + ``dtype`` of the returned accumulator (`hist`), so it must be + large enough to hold accumulated values as well. + density : bool, optional + If ``False``, the result will contain the number of samples in + each bin. If ``True``, the result is the value of the + probability *density* function at the bin, normalized such that + the *integral* over the range is 1. Note that the sum of the + histogram values will not be equal to 1 unless bins of unity + width are chosen; it is not a probability *mass* function. + + Returns + ------- + hist : array + The values of the histogram. See `density` and `weights` for a + description of the possible semantics. If `weights` are given, + ``hist.dtype`` will be taken from `weights`. + bin_edges : array of dtype float + Return the bin edges ``(length(hist)+1)``. + + + See Also + -------- + histogramdd, bincount, searchsorted, digitize, histogram_bin_edges + + Notes + ----- + All but the last (righthand-most) bin is half-open. In other words, + if `bins` is:: + + [1, 2, 3, 4] + + then the first bin is ``[1, 2)`` (including 1, but excluding 2) and + the second ``[2, 3)``. The last bin, however, is ``[3, 4]``, which + *includes* 4. + + + Examples + -------- + >>> import numpy as np + >>> np.histogram([1, 2, 1], bins=[0, 1, 2, 3]) + (array([0, 2, 1]), array([0, 1, 2, 3])) + >>> np.histogram(np.arange(4), bins=np.arange(5), density=True) + (array([0.25, 0.25, 0.25, 0.25]), array([0, 1, 2, 3, 4])) + >>> np.histogram([[1, 2, 1], [1, 0, 1]], bins=[0,1,2,3]) + (array([1, 4, 1]), array([0, 1, 2, 3])) + + >>> a = np.arange(5) + >>> hist, bin_edges = np.histogram(a, density=True) + >>> hist + array([0.5, 0. , 0.5, 0. , 0. , 0.5, 0. , 0.5, 0. , 0.5]) + >>> hist.sum() + 2.4999999999999996 + >>> np.sum(hist * np.diff(bin_edges)) + 1.0 + + Automated Bin Selection Methods example, using 2 peak random data + with 2000 points. + + .. plot:: + :include-source: + + import matplotlib.pyplot as plt + import numpy as np + + rng = np.random.RandomState(10) # deterministic random data + a = np.hstack((rng.normal(size=1000), + rng.normal(loc=5, scale=2, size=1000))) + plt.hist(a, bins='auto') # arguments are passed to np.histogram + plt.title("Histogram with 'auto' bins") + plt.show() + + """ + a, weights = _ravel_and_check_weights(a, weights) + + bin_edges, uniform_bins = _get_bin_edges(a, bins, range, weights) + + # Histogram is an integer or a float array depending on the weights. + if weights is None: + ntype = np.dtype(np.intp) + else: + ntype = weights.dtype + + # We set a block size, as this allows us to iterate over chunks when + # computing histograms, to minimize memory usage. + BLOCK = 65536 + + # The fast path uses bincount, but that only works for certain types + # of weight + simple_weights = ( + weights is None or + np.can_cast(weights.dtype, np.double) or + np.can_cast(weights.dtype, complex) + ) + + if uniform_bins is not None and simple_weights: + # Fast algorithm for equal bins + # We now convert values of a to bin indices, under the assumption of + # equal bin widths (which is valid here). + first_edge, last_edge, n_equal_bins = uniform_bins + + # Initialize empty histogram + n = np.zeros(n_equal_bins, ntype) + + # Pre-compute histogram scaling factor + norm_numerator = n_equal_bins + norm_denom = _unsigned_subtract(last_edge, first_edge) + + # We iterate over blocks here for two reasons: the first is that for + # large arrays, it is actually faster (for example for a 10^8 array it + # is 2x as fast) and it results in a memory footprint 3x lower in the + # limit of large arrays. + for i in _range(0, len(a), BLOCK): + tmp_a = a[i:i+BLOCK] + if weights is None: + tmp_w = None + else: + tmp_w = weights[i:i + BLOCK] + + # Only include values in the right range + keep = (tmp_a >= first_edge) + keep &= (tmp_a <= last_edge) + if not np.logical_and.reduce(keep): + tmp_a = tmp_a[keep] + if tmp_w is not None: + tmp_w = tmp_w[keep] + + # This cast ensures no type promotions occur below, which gh-10322 + # make unpredictable. Getting it wrong leads to precision errors + # like gh-8123. + tmp_a = tmp_a.astype(bin_edges.dtype, copy=False) + + # Compute the bin indices, and for values that lie exactly on + # last_edge we need to subtract one + f_indices = ((_unsigned_subtract(tmp_a, first_edge) / norm_denom) + * norm_numerator) + indices = f_indices.astype(np.intp) + indices[indices == n_equal_bins] -= 1 + + # The index computation is not guaranteed to give exactly + # consistent results within ~1 ULP of the bin edges. + decrement = tmp_a < bin_edges[indices] + indices[decrement] -= 1 + # The last bin includes the right edge. The other bins do not. + increment = ((tmp_a >= bin_edges[indices + 1]) + & (indices != n_equal_bins - 1)) + indices[increment] += 1 + + # We now compute the histogram using bincount + if ntype.kind == 'c': + n.real += np.bincount(indices, weights=tmp_w.real, + minlength=n_equal_bins) + n.imag += np.bincount(indices, weights=tmp_w.imag, + minlength=n_equal_bins) + else: + n += np.bincount(indices, weights=tmp_w, + minlength=n_equal_bins).astype(ntype) + else: + # Compute via cumulative histogram + cum_n = np.zeros(bin_edges.shape, ntype) + if weights is None: + for i in _range(0, len(a), BLOCK): + sa = np.sort(a[i:i+BLOCK]) + cum_n += _search_sorted_inclusive(sa, bin_edges) + else: + zero = np.zeros(1, dtype=ntype) + for i in _range(0, len(a), BLOCK): + tmp_a = a[i:i+BLOCK] + tmp_w = weights[i:i+BLOCK] + sorting_index = np.argsort(tmp_a) + sa = tmp_a[sorting_index] + sw = tmp_w[sorting_index] + cw = np.concatenate((zero, sw.cumsum())) + bin_index = _search_sorted_inclusive(sa, bin_edges) + cum_n += cw[bin_index] + + n = np.diff(cum_n) + + if density: + db = np.array(np.diff(bin_edges), float) + return n/db/n.sum(), bin_edges + + return n, bin_edges + + +def _histogramdd_dispatcher(sample, bins=None, range=None, density=None, + weights=None): + if hasattr(sample, 'shape'): # same condition as used in histogramdd + yield sample + else: + yield from sample + with contextlib.suppress(TypeError): + yield from bins + yield weights + + +@array_function_dispatch(_histogramdd_dispatcher) +def histogramdd(sample, bins=10, range=None, density=None, weights=None): + """ + Compute the multidimensional histogram of some data. + + Parameters + ---------- + sample : (N, D) array, or (N, D) array_like + The data to be histogrammed. + + Note the unusual interpretation of sample when an array_like: + + * When an array, each row is a coordinate in a D-dimensional space - + such as ``histogramdd(np.array([p1, p2, p3]))``. + * When an array_like, each element is the list of values for single + coordinate - such as ``histogramdd((X, Y, Z))``. + + The first form should be preferred. + + bins : sequence or int, optional + The bin specification: + + * A sequence of arrays describing the monotonically increasing bin + edges along each dimension. + * The number of bins for each dimension (nx, ny, ... =bins) + * The number of bins for all dimensions (nx=ny=...=bins). + + range : sequence, optional + A sequence of length D, each an optional (lower, upper) tuple giving + the outer bin edges to be used if the edges are not given explicitly in + `bins`. + An entry of None in the sequence results in the minimum and maximum + values being used for the corresponding dimension. + The default, None, is equivalent to passing a tuple of D None values. + density : bool, optional + If False, the default, returns the number of samples in each bin. + If True, returns the probability *density* function at the bin, + ``bin_count / sample_count / bin_volume``. + weights : (N,) array_like, optional + An array of values `w_i` weighing each sample `(x_i, y_i, z_i, ...)`. + Weights are normalized to 1 if density is True. If density is False, + the values of the returned histogram are equal to the sum of the + weights belonging to the samples falling into each bin. + + Returns + ------- + H : ndarray + The multidimensional histogram of sample x. See density and weights + for the different possible semantics. + edges : tuple of ndarrays + A tuple of D arrays describing the bin edges for each dimension. + + See Also + -------- + histogram: 1-D histogram + histogram2d: 2-D histogram + + Examples + -------- + >>> import numpy as np + >>> rng = np.random.default_rng() + >>> r = rng.normal(size=(100,3)) + >>> H, edges = np.histogramdd(r, bins = (5, 8, 4)) + >>> H.shape, edges[0].size, edges[1].size, edges[2].size + ((5, 8, 4), 6, 9, 5) + + """ + + try: + # Sample is an ND-array. + N, D = sample.shape + except (AttributeError, ValueError): + # Sample is a sequence of 1D arrays. + sample = np.atleast_2d(sample).T + N, D = sample.shape + + nbin = np.empty(D, np.intp) + edges = D*[None] + dedges = D*[None] + if weights is not None: + weights = np.asarray(weights) + + try: + M = len(bins) + if M != D: + raise ValueError( + 'The dimension of bins must be equal to the dimension of the ' + 'sample x.') + except TypeError: + # bins is an integer + bins = D*[bins] + + # normalize the range argument + if range is None: + range = (None,) * D + elif len(range) != D: + raise ValueError('range argument must have one entry per dimension') + + # Create edge arrays + for i in _range(D): + if np.ndim(bins[i]) == 0: + if bins[i] < 1: + raise ValueError( + '`bins[{}]` must be positive, when an integer'.format(i)) + smin, smax = _get_outer_edges(sample[:,i], range[i]) + try: + n = operator.index(bins[i]) + + except TypeError as e: + raise TypeError( + "`bins[{}]` must be an integer, when a scalar".format(i) + ) from e + + edges[i] = np.linspace(smin, smax, n + 1) + elif np.ndim(bins[i]) == 1: + edges[i] = np.asarray(bins[i]) + if np.any(edges[i][:-1] > edges[i][1:]): + raise ValueError( + '`bins[{}]` must be monotonically increasing, when an array' + .format(i)) + else: + raise ValueError( + '`bins[{}]` must be a scalar or 1d array'.format(i)) + + nbin[i] = len(edges[i]) + 1 # includes an outlier on each end + dedges[i] = np.diff(edges[i]) + + # Compute the bin number each sample falls into. + Ncount = tuple( + # avoid np.digitize to work around gh-11022 + np.searchsorted(edges[i], sample[:, i], side='right') + for i in _range(D) + ) + + # Using digitize, values that fall on an edge are put in the right bin. + # For the rightmost bin, we want values equal to the right edge to be + # counted in the last bin, and not as an outlier. + for i in _range(D): + # Find which points are on the rightmost edge. + on_edge = (sample[:, i] == edges[i][-1]) + # Shift these points one bin to the left. + Ncount[i][on_edge] -= 1 + + # Compute the sample indices in the flattened histogram matrix. + # This raises an error if the array is too large. + xy = np.ravel_multi_index(Ncount, nbin) + + # Compute the number of repetitions in xy and assign it to the + # flattened histmat. + hist = np.bincount(xy, weights, minlength=nbin.prod()) + + # Shape into a proper matrix + hist = hist.reshape(nbin) + + # This preserves the (bad) behavior observed in gh-7845, for now. + hist = hist.astype(float, casting='safe') + + # Remove outliers (indices 0 and -1 for each dimension). + core = D*(slice(1, -1),) + hist = hist[core] + + if density: + # calculate the probability density function + s = hist.sum() + for i in _range(D): + shape = np.ones(D, int) + shape[i] = nbin[i] - 2 + hist = hist / dedges[i].reshape(shape) + hist /= s + + if (hist.shape != nbin - 2).any(): + raise RuntimeError( + "Internal Shape Error") + return hist, edges diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e18ab99035b48f9b5d8d85ffe6bb226b5c284ae9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_histograms_impl.pyi @@ -0,0 +1,48 @@ +from collections.abc import Sequence +from typing import ( + Literal as L, + Any, + SupportsIndex, + TypeAlias, +) + +from numpy._typing import ( + NDArray, + ArrayLike, +) + +__all__ = ["histogram", "histogramdd", "histogram_bin_edges"] + +_BinKind: TypeAlias = L[ + "stone", + "auto", + "doane", + "fd", + "rice", + "scott", + "sqrt", + "sturges", +] + +def histogram_bin_edges( + a: ArrayLike, + bins: _BinKind | SupportsIndex | ArrayLike = ..., + range: None | tuple[float, float] = ..., + weights: None | ArrayLike = ..., +) -> NDArray[Any]: ... + +def histogram( + a: ArrayLike, + bins: _BinKind | SupportsIndex | ArrayLike = ..., + range: None | tuple[float, float] = ..., + density: bool = ..., + weights: None | ArrayLike = ..., +) -> tuple[NDArray[Any], NDArray[Any]]: ... + +def histogramdd( + sample: ArrayLike, + bins: SupportsIndex | ArrayLike = ..., + range: Sequence[tuple[float, float]] = ..., + density: None | bool = ..., + weights: None | ArrayLike = ..., +) -> tuple[NDArray[Any], tuple[NDArray[Any], ...]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..da8fbedc8072b4d73ebc9b2e25229fff761ff073 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.py @@ -0,0 +1,1069 @@ +import functools +import sys +import math +import warnings + +import numpy as np +from .._utils import set_module +import numpy._core.numeric as _nx +from numpy._core.numeric import ScalarType, array +from numpy._core.numerictypes import issubdtype + +import numpy.matrixlib as matrixlib +from numpy._core.multiarray import ravel_multi_index, unravel_index +from numpy._core import overrides, linspace +from numpy.lib.stride_tricks import as_strided +from numpy.lib._function_base_impl import diff + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +__all__ = [ + 'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_', + 's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal', + 'diag_indices', 'diag_indices_from' +] + + +def _ix__dispatcher(*args): + return args + + +@array_function_dispatch(_ix__dispatcher) +def ix_(*args): + """ + Construct an open mesh from multiple sequences. + + This function takes N 1-D sequences and returns N outputs with N + dimensions each, such that the shape is 1 in all but one dimension + and the dimension with the non-unit shape value cycles through all + N dimensions. + + Using `ix_` one can quickly construct index arrays that will index + the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array + ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``. + + Parameters + ---------- + args : 1-D sequences + Each sequence should be of integer or boolean type. + Boolean sequences will be interpreted as boolean masks for the + corresponding dimension (equivalent to passing in + ``np.nonzero(boolean_sequence)``). + + Returns + ------- + out : tuple of ndarrays + N arrays with N dimensions each, with N the number of input + sequences. Together these arrays form an open mesh. + + See Also + -------- + ogrid, mgrid, meshgrid + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(10).reshape(2, 5) + >>> a + array([[0, 1, 2, 3, 4], + [5, 6, 7, 8, 9]]) + >>> ixgrid = np.ix_([0, 1], [2, 4]) + >>> ixgrid + (array([[0], + [1]]), array([[2, 4]])) + >>> ixgrid[0].shape, ixgrid[1].shape + ((2, 1), (1, 2)) + >>> a[ixgrid] + array([[2, 4], + [7, 9]]) + + >>> ixgrid = np.ix_([True, True], [2, 4]) + >>> a[ixgrid] + array([[2, 4], + [7, 9]]) + >>> ixgrid = np.ix_([True, True], [False, False, True, False, True]) + >>> a[ixgrid] + array([[2, 4], + [7, 9]]) + + """ + out = [] + nd = len(args) + for k, new in enumerate(args): + if not isinstance(new, _nx.ndarray): + new = np.asarray(new) + if new.size == 0: + # Explicitly type empty arrays to avoid float default + new = new.astype(_nx.intp) + if new.ndim != 1: + raise ValueError("Cross index must be 1 dimensional") + if issubdtype(new.dtype, _nx.bool): + new, = new.nonzero() + new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1)) + out.append(new) + return tuple(out) + + +class nd_grid: + """ + Construct a multi-dimensional "meshgrid". + + ``grid = nd_grid()`` creates an instance which will return a mesh-grid + when indexed. The dimension and number of the output arrays are equal + to the number of indexing dimensions. If the step length is not a + complex number, then the stop is not inclusive. + + However, if the step length is a **complex number** (e.g. 5j), then the + integer part of its magnitude is interpreted as specifying the + number of points to create between the start and stop values, where + the stop value **is inclusive**. + + If instantiated with an argument of ``sparse=True``, the mesh-grid is + open (or not fleshed out) so that only one-dimension of each returned + argument is greater than 1. + + Parameters + ---------- + sparse : bool, optional + Whether the grid is sparse or not. Default is False. + + Notes + ----- + Two instances of `nd_grid` are made available in the NumPy namespace, + `mgrid` and `ogrid`, approximately defined as:: + + mgrid = nd_grid(sparse=False) + ogrid = nd_grid(sparse=True) + + Users should use these pre-defined instances instead of using `nd_grid` + directly. + """ + __slots__ = ('sparse',) + + def __init__(self, sparse=False): + self.sparse = sparse + + def __getitem__(self, key): + try: + size = [] + # Mimic the behavior of `np.arange` and use a data type + # which is at least as large as `np.int_` + num_list = [0] + for k in range(len(key)): + step = key[k].step + start = key[k].start + stop = key[k].stop + if start is None: + start = 0 + if step is None: + step = 1 + if isinstance(step, (_nx.complexfloating, complex)): + step = abs(step) + size.append(int(step)) + else: + size.append( + int(math.ceil((stop - start) / (step*1.0)))) + num_list += [start, stop, step] + typ = _nx.result_type(*num_list) + if self.sparse: + nn = [_nx.arange(_x, dtype=_t) + for _x, _t in zip(size, (typ,)*len(size))] + else: + nn = _nx.indices(size, typ) + for k, kk in enumerate(key): + step = kk.step + start = kk.start + if start is None: + start = 0 + if step is None: + step = 1 + if isinstance(step, (_nx.complexfloating, complex)): + step = int(abs(step)) + if step != 1: + step = (kk.stop - start) / float(step - 1) + nn[k] = (nn[k]*step+start) + if self.sparse: + slobj = [_nx.newaxis]*len(size) + for k in range(len(size)): + slobj[k] = slice(None, None) + nn[k] = nn[k][tuple(slobj)] + slobj[k] = _nx.newaxis + return tuple(nn) # ogrid -> tuple of arrays + return nn # mgrid -> ndarray + except (IndexError, TypeError): + step = key.step + stop = key.stop + start = key.start + if start is None: + start = 0 + if isinstance(step, (_nx.complexfloating, complex)): + # Prevent the (potential) creation of integer arrays + step_float = abs(step) + step = length = int(step_float) + if step != 1: + step = (key.stop-start)/float(step-1) + typ = _nx.result_type(start, stop, step_float) + return _nx.arange(0, length, 1, dtype=typ)*step + start + else: + return _nx.arange(start, stop, step) + + +class MGridClass(nd_grid): + """ + An instance which returns a dense multi-dimensional "meshgrid". + + An instance which returns a dense (or fleshed out) mesh-grid + when indexed, so that each returned argument has the same shape. + The dimensions and number of the output arrays are equal to the + number of indexing dimensions. If the step length is not a complex + number, then the stop is not inclusive. + + However, if the step length is a **complex number** (e.g. 5j), then + the integer part of its magnitude is interpreted as specifying the + number of points to create between the start and stop values, where + the stop value **is inclusive**. + + Returns + ------- + mesh-grid : ndarray + A single array, containing a set of `ndarray`\\ s all of the same + dimensions. stacked along the first axis. + + See Also + -------- + ogrid : like `mgrid` but returns open (not fleshed out) mesh grids + meshgrid: return coordinate matrices from coordinate vectors + r_ : array concatenator + :ref:`how-to-partition` + + Examples + -------- + >>> import numpy as np + >>> np.mgrid[0:5, 0:5] + array([[[0, 0, 0, 0, 0], + [1, 1, 1, 1, 1], + [2, 2, 2, 2, 2], + [3, 3, 3, 3, 3], + [4, 4, 4, 4, 4]], + [[0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4]]]) + >>> np.mgrid[-1:1:5j] + array([-1. , -0.5, 0. , 0.5, 1. ]) + + >>> np.mgrid[0:4].shape + (4,) + >>> np.mgrid[0:4, 0:5].shape + (2, 4, 5) + >>> np.mgrid[0:4, 0:5, 0:6].shape + (3, 4, 5, 6) + + """ + __slots__ = () + + def __init__(self): + super().__init__(sparse=False) + + +mgrid = MGridClass() + + +class OGridClass(nd_grid): + """ + An instance which returns an open multi-dimensional "meshgrid". + + An instance which returns an open (i.e. not fleshed out) mesh-grid + when indexed, so that only one dimension of each returned array is + greater than 1. The dimension and number of the output arrays are + equal to the number of indexing dimensions. If the step length is + not a complex number, then the stop is not inclusive. + + However, if the step length is a **complex number** (e.g. 5j), then + the integer part of its magnitude is interpreted as specifying the + number of points to create between the start and stop values, where + the stop value **is inclusive**. + + Returns + ------- + mesh-grid : ndarray or tuple of ndarrays + If the input is a single slice, returns an array. + If the input is multiple slices, returns a tuple of arrays, with + only one dimension not equal to 1. + + See Also + -------- + mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids + meshgrid: return coordinate matrices from coordinate vectors + r_ : array concatenator + :ref:`how-to-partition` + + Examples + -------- + >>> from numpy import ogrid + >>> ogrid[-1:1:5j] + array([-1. , -0.5, 0. , 0.5, 1. ]) + >>> ogrid[0:5, 0:5] + (array([[0], + [1], + [2], + [3], + [4]]), + array([[0, 1, 2, 3, 4]])) + + """ + __slots__ = () + + def __init__(self): + super().__init__(sparse=True) + + +ogrid = OGridClass() + + +class AxisConcatenator: + """ + Translates slice objects to concatenation along an axis. + + For detailed documentation on usage, see `r_`. + """ + __slots__ = ('axis', 'matrix', 'trans1d', 'ndmin') + + # allow ma.mr_ to override this + concatenate = staticmethod(_nx.concatenate) + makemat = staticmethod(matrixlib.matrix) + + def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1): + self.axis = axis + self.matrix = matrix + self.trans1d = trans1d + self.ndmin = ndmin + + def __getitem__(self, key): + # handle matrix builder syntax + if isinstance(key, str): + frame = sys._getframe().f_back + mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals) + return mymat + + if not isinstance(key, tuple): + key = (key,) + + # copy attributes, since they can be overridden in the first argument + trans1d = self.trans1d + ndmin = self.ndmin + matrix = self.matrix + axis = self.axis + + objs = [] + # dtypes or scalars for weak scalar handling in result_type + result_type_objs = [] + + for k, item in enumerate(key): + scalar = False + if isinstance(item, slice): + step = item.step + start = item.start + stop = item.stop + if start is None: + start = 0 + if step is None: + step = 1 + if isinstance(step, (_nx.complexfloating, complex)): + size = int(abs(step)) + newobj = linspace(start, stop, num=size) + else: + newobj = _nx.arange(start, stop, step) + if ndmin > 1: + newobj = array(newobj, copy=None, ndmin=ndmin) + if trans1d != -1: + newobj = newobj.swapaxes(-1, trans1d) + elif isinstance(item, str): + if k != 0: + raise ValueError("special directives must be the " + "first entry.") + if item in ('r', 'c'): + matrix = True + col = (item == 'c') + continue + if ',' in item: + vec = item.split(',') + try: + axis, ndmin = [int(x) for x in vec[:2]] + if len(vec) == 3: + trans1d = int(vec[2]) + continue + except Exception as e: + raise ValueError( + "unknown special directive {!r}".format(item) + ) from e + try: + axis = int(item) + continue + except (ValueError, TypeError) as e: + raise ValueError("unknown special directive") from e + elif type(item) in ScalarType: + scalar = True + newobj = item + else: + item_ndim = np.ndim(item) + newobj = array(item, copy=None, subok=True, ndmin=ndmin) + if trans1d != -1 and item_ndim < ndmin: + k2 = ndmin - item_ndim + k1 = trans1d + if k1 < 0: + k1 += k2 + 1 + defaxes = list(range(ndmin)) + axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2] + newobj = newobj.transpose(axes) + + objs.append(newobj) + if scalar: + result_type_objs.append(item) + else: + result_type_objs.append(newobj.dtype) + + # Ensure that scalars won't up-cast unless warranted, for 0, drops + # through to error in concatenate. + if len(result_type_objs) != 0: + final_dtype = _nx.result_type(*result_type_objs) + # concatenate could do cast, but that can be overridden: + objs = [array(obj, copy=None, subok=True, + ndmin=ndmin, dtype=final_dtype) for obj in objs] + + res = self.concatenate(tuple(objs), axis=axis) + + if matrix: + oldndim = res.ndim + res = self.makemat(res) + if oldndim == 1 and col: + res = res.T + return res + + def __len__(self): + return 0 + +# separate classes are used here instead of just making r_ = concatenator(0), +# etc. because otherwise we couldn't get the doc string to come out right +# in help(r_) + + +class RClass(AxisConcatenator): + """ + Translates slice objects to concatenation along the first axis. + + This is a simple way to build up arrays quickly. There are two use cases. + + 1. If the index expression contains comma separated arrays, then stack + them along their first axis. + 2. If the index expression contains slice notation or scalars then create + a 1-D array with a range indicated by the slice notation. + + If slice notation is used, the syntax ``start:stop:step`` is equivalent + to ``np.arange(start, stop, step)`` inside of the brackets. However, if + ``step`` is an imaginary number (i.e. 100j) then its integer portion is + interpreted as a number-of-points desired and the start and stop are + inclusive. In other words ``start:stop:stepj`` is interpreted as + ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets. + After expansion of slice notation, all comma separated sequences are + concatenated together. + + Optional character strings placed as the first element of the index + expression can be used to change the output. The strings 'r' or 'c' result + in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row) + matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1 + (column) matrix is produced. If the result is 2-D then both provide the + same matrix result. + + A string integer specifies which axis to stack multiple comma separated + arrays along. A string of two comma-separated integers allows indication + of the minimum number of dimensions to force each entry into as the + second integer (the axis to concatenate along is still the first integer). + + A string with three comma-separated integers allows specification of the + axis to concatenate along, the minimum number of dimensions to force the + entries to, and which axis should contain the start of the arrays which + are less than the specified number of dimensions. In other words the third + integer allows you to specify where the 1's should be placed in the shape + of the arrays that have their shapes upgraded. By default, they are placed + in the front of the shape tuple. The third argument allows you to specify + where the start of the array should be instead. Thus, a third argument of + '0' would place the 1's at the end of the array shape. Negative integers + specify where in the new shape tuple the last dimension of upgraded arrays + should be placed, so the default is '-1'. + + Parameters + ---------- + Not a function, so takes no parameters + + + Returns + ------- + A concatenated ndarray or matrix. + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + c_ : Translates slice objects to concatenation along the second axis. + + Examples + -------- + >>> import numpy as np + >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])] + array([1, 2, 3, ..., 4, 5, 6]) + >>> np.r_[-1:1:6j, [0]*3, 5, 6] + array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ]) + + String integers specify the axis to concatenate along or the minimum + number of dimensions to force entries into. + + >>> a = np.array([[0, 1, 2], [3, 4, 5]]) + >>> np.r_['-1', a, a] # concatenate along last axis + array([[0, 1, 2, 0, 1, 2], + [3, 4, 5, 3, 4, 5]]) + >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2 + array([[1, 2, 3], + [4, 5, 6]]) + + >>> np.r_['0,2,0', [1,2,3], [4,5,6]] + array([[1], + [2], + [3], + [4], + [5], + [6]]) + >>> np.r_['1,2,0', [1,2,3], [4,5,6]] + array([[1, 4], + [2, 5], + [3, 6]]) + + Using 'r' or 'c' as a first string argument creates a matrix. + + >>> np.r_['r',[1,2,3], [4,5,6]] + matrix([[1, 2, 3, 4, 5, 6]]) + + """ + __slots__ = () + + def __init__(self): + AxisConcatenator.__init__(self, 0) + + +r_ = RClass() + + +class CClass(AxisConcatenator): + """ + Translates slice objects to concatenation along the second axis. + + This is short-hand for ``np.r_['-1,2,0', index expression]``, which is + useful because of its common occurrence. In particular, arrays will be + stacked along their last axis after being upgraded to at least 2-D with + 1's post-pended to the shape (column vectors made out of 1-D arrays). + + See Also + -------- + column_stack : Stack 1-D arrays as columns into a 2-D array. + r_ : For more detailed documentation. + + Examples + -------- + >>> import numpy as np + >>> np.c_[np.array([1,2,3]), np.array([4,5,6])] + array([[1, 4], + [2, 5], + [3, 6]]) + >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])] + array([[1, 2, 3, ..., 4, 5, 6]]) + + """ + __slots__ = () + + def __init__(self): + AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0) + + +c_ = CClass() + + +@set_module('numpy') +class ndenumerate: + """ + Multidimensional index iterator. + + Return an iterator yielding pairs of array coordinates and values. + + Parameters + ---------- + arr : ndarray + Input array. + + See Also + -------- + ndindex, flatiter + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, 4]]) + >>> for index, x in np.ndenumerate(a): + ... print(index, x) + (0, 0) 1 + (0, 1) 2 + (1, 0) 3 + (1, 1) 4 + + """ + + def __init__(self, arr): + self.iter = np.asarray(arr).flat + + def __next__(self): + """ + Standard iterator method, returns the index tuple and array value. + + Returns + ------- + coords : tuple of ints + The indices of the current iteration. + val : scalar + The array element of the current iteration. + + """ + return self.iter.coords, next(self.iter) + + def __iter__(self): + return self + + +@set_module('numpy') +class ndindex: + """ + An N-dimensional iterator object to index arrays. + + Given the shape of an array, an `ndindex` instance iterates over + the N-dimensional index of the array. At each iteration a tuple + of indices is returned, the last dimension is iterated over first. + + Parameters + ---------- + shape : ints, or a single tuple of ints + The size of each dimension of the array can be passed as + individual parameters or as the elements of a tuple. + + See Also + -------- + ndenumerate, flatiter + + Examples + -------- + >>> import numpy as np + + Dimensions as individual arguments + + >>> for index in np.ndindex(3, 2, 1): + ... print(index) + (0, 0, 0) + (0, 1, 0) + (1, 0, 0) + (1, 1, 0) + (2, 0, 0) + (2, 1, 0) + + Same dimensions - but in a tuple ``(3, 2, 1)`` + + >>> for index in np.ndindex((3, 2, 1)): + ... print(index) + (0, 0, 0) + (0, 1, 0) + (1, 0, 0) + (1, 1, 0) + (2, 0, 0) + (2, 1, 0) + + """ + + def __init__(self, *shape): + if len(shape) == 1 and isinstance(shape[0], tuple): + shape = shape[0] + x = as_strided(_nx.zeros(1), shape=shape, + strides=_nx.zeros_like(shape)) + self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'], + order='C') + + def __iter__(self): + return self + + def ndincr(self): + """ + Increment the multi-dimensional index by one. + + This method is for backward compatibility only: do not use. + + .. deprecated:: 1.20.0 + This method has been advised against since numpy 1.8.0, but only + started emitting DeprecationWarning as of this version. + """ + # NumPy 1.20.0, 2020-09-08 + warnings.warn( + "`ndindex.ndincr()` is deprecated, use `next(ndindex)` instead", + DeprecationWarning, stacklevel=2) + next(self) + + def __next__(self): + """ + Standard iterator method, updates the index and returns the index + tuple. + + Returns + ------- + val : tuple of ints + Returns a tuple containing the indices of the current + iteration. + + """ + next(self._it) + return self._it.multi_index + + +# You can do all this with slice() plus a few special objects, +# but there's a lot to remember. This version is simpler because +# it uses the standard array indexing syntax. +# +# Written by Konrad Hinsen +# last revision: 1999-7-23 +# +# Cosmetic changes by T. Oliphant 2001 +# +# + +class IndexExpression: + """ + A nicer way to build up index tuples for arrays. + + .. note:: + Use one of the two predefined instances ``index_exp`` or `s_` + rather than directly using `IndexExpression`. + + For any index combination, including slicing and axis insertion, + ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any + array `a`. However, ``np.index_exp[indices]`` can be used anywhere + in Python code and returns a tuple of slice objects that can be + used in the construction of complex index expressions. + + Parameters + ---------- + maketuple : bool + If True, always returns a tuple. + + See Also + -------- + s_ : Predefined instance without tuple conversion: + `s_ = IndexExpression(maketuple=False)`. + The ``index_exp`` is another predefined instance that + always returns a tuple: + `index_exp = IndexExpression(maketuple=True)`. + + Notes + ----- + You can do all this with :class:`slice` plus a few special objects, + but there's a lot to remember and this version is simpler because + it uses the standard array indexing syntax. + + Examples + -------- + >>> import numpy as np + >>> np.s_[2::2] + slice(2, None, 2) + >>> np.index_exp[2::2] + (slice(2, None, 2),) + + >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]] + array([2, 4]) + + """ + __slots__ = ('maketuple',) + + def __init__(self, maketuple): + self.maketuple = maketuple + + def __getitem__(self, item): + if self.maketuple and not isinstance(item, tuple): + return (item,) + else: + return item + + +index_exp = IndexExpression(maketuple=True) +s_ = IndexExpression(maketuple=False) + +# End contribution from Konrad. + + +# The following functions complement those in twodim_base, but are +# applicable to N-dimensions. + + +def _fill_diagonal_dispatcher(a, val, wrap=None): + return (a,) + + +@array_function_dispatch(_fill_diagonal_dispatcher) +def fill_diagonal(a, val, wrap=False): + """Fill the main diagonal of the given array of any dimensionality. + + For an array `a` with ``a.ndim >= 2``, the diagonal is the list of + values ``a[i, ..., i]`` with indices ``i`` all identical. This function + modifies the input array in-place without returning a value. + + Parameters + ---------- + a : array, at least 2-D. + Array whose diagonal is to be filled in-place. + val : scalar or array_like + Value(s) to write on the diagonal. If `val` is scalar, the value is + written along the diagonal. If array-like, the flattened `val` is + written along the diagonal, repeating if necessary to fill all + diagonal entries. + + wrap : bool + For tall matrices in NumPy version up to 1.6.2, the + diagonal "wrapped" after N columns. You can have this behavior + with this option. This affects only tall matrices. + + See also + -------- + diag_indices, diag_indices_from + + Notes + ----- + This functionality can be obtained via `diag_indices`, but internally + this version uses a much faster implementation that never constructs the + indices and uses simple slicing. + + Examples + -------- + >>> import numpy as np + >>> a = np.zeros((3, 3), int) + >>> np.fill_diagonal(a, 5) + >>> a + array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5]]) + + The same function can operate on a 4-D array: + + >>> a = np.zeros((3, 3, 3, 3), int) + >>> np.fill_diagonal(a, 4) + + We only show a few blocks for clarity: + + >>> a[0, 0] + array([[4, 0, 0], + [0, 0, 0], + [0, 0, 0]]) + >>> a[1, 1] + array([[0, 0, 0], + [0, 4, 0], + [0, 0, 0]]) + >>> a[2, 2] + array([[0, 0, 0], + [0, 0, 0], + [0, 0, 4]]) + + The wrap option affects only tall matrices: + + >>> # tall matrices no wrap + >>> a = np.zeros((5, 3), int) + >>> np.fill_diagonal(a, 4) + >>> a + array([[4, 0, 0], + [0, 4, 0], + [0, 0, 4], + [0, 0, 0], + [0, 0, 0]]) + + >>> # tall matrices wrap + >>> a = np.zeros((5, 3), int) + >>> np.fill_diagonal(a, 4, wrap=True) + >>> a + array([[4, 0, 0], + [0, 4, 0], + [0, 0, 4], + [0, 0, 0], + [4, 0, 0]]) + + >>> # wide matrices + >>> a = np.zeros((3, 5), int) + >>> np.fill_diagonal(a, 4, wrap=True) + >>> a + array([[4, 0, 0, 0, 0], + [0, 4, 0, 0, 0], + [0, 0, 4, 0, 0]]) + + The anti-diagonal can be filled by reversing the order of elements + using either `numpy.flipud` or `numpy.fliplr`. + + >>> a = np.zeros((3, 3), int); + >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip + >>> a + array([[0, 0, 1], + [0, 2, 0], + [3, 0, 0]]) + >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip + >>> a + array([[0, 0, 3], + [0, 2, 0], + [1, 0, 0]]) + + Note that the order in which the diagonal is filled varies depending + on the flip function. + """ + if a.ndim < 2: + raise ValueError("array must be at least 2-d") + end = None + if a.ndim == 2: + # Explicit, fast formula for the common case. For 2-d arrays, we + # accept rectangular ones. + step = a.shape[1] + 1 + # This is needed to don't have tall matrix have the diagonal wrap. + if not wrap: + end = a.shape[1] * a.shape[1] + else: + # For more than d=2, the strided formula is only valid for arrays with + # all dimensions equal, so we check first. + if not np.all(diff(a.shape) == 0): + raise ValueError("All dimensions of input must be of equal length") + step = 1 + (np.cumprod(a.shape[:-1])).sum() + + # Write the value out into the diagonal. + a.flat[:end:step] = val + + +@set_module('numpy') +def diag_indices(n, ndim=2): + """ + Return the indices to access the main diagonal of an array. + + This returns a tuple of indices that can be used to access the main + diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape + (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for + ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]`` + for ``i = [0..n-1]``. + + Parameters + ---------- + n : int + The size, along each dimension, of the arrays for which the returned + indices can be used. + + ndim : int, optional + The number of dimensions. + + See Also + -------- + diag_indices_from + + Examples + -------- + >>> import numpy as np + + Create a set of indices to access the diagonal of a (4, 4) array: + + >>> di = np.diag_indices(4) + >>> di + (array([0, 1, 2, 3]), array([0, 1, 2, 3])) + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + >>> a[di] = 100 + >>> a + array([[100, 1, 2, 3], + [ 4, 100, 6, 7], + [ 8, 9, 100, 11], + [ 12, 13, 14, 100]]) + + Now, we create indices to manipulate a 3-D array: + + >>> d3 = np.diag_indices(2, 3) + >>> d3 + (array([0, 1]), array([0, 1]), array([0, 1])) + + And use it to set the diagonal of an array of zeros to 1: + + >>> a = np.zeros((2, 2, 2), dtype=int) + >>> a[d3] = 1 + >>> a + array([[[1, 0], + [0, 0]], + [[0, 0], + [0, 1]]]) + + """ + idx = np.arange(n) + return (idx,) * ndim + + +def _diag_indices_from(arr): + return (arr,) + + +@array_function_dispatch(_diag_indices_from) +def diag_indices_from(arr): + """ + Return the indices to access the main diagonal of an n-dimensional array. + + See `diag_indices` for full details. + + Parameters + ---------- + arr : array, at least 2-D + + See Also + -------- + diag_indices + + Examples + -------- + >>> import numpy as np + + Create a 4 by 4 array. + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Get the indices of the diagonal elements. + + >>> di = np.diag_indices_from(a) + >>> di + (array([0, 1, 2, 3]), array([0, 1, 2, 3])) + + >>> a[di] + array([ 0, 5, 10, 15]) + + This is simply syntactic sugar for diag_indices. + + >>> np.diag_indices(a.shape[0]) + (array([0, 1, 2, 3]), array([0, 1, 2, 3])) + + """ + + if not arr.ndim >= 2: + raise ValueError("input array must be at least 2-d") + # For more than d=2, the strided formula is only valid for arrays with + # all dimensions equal, so we check first. + if not np.all(diff(arr.shape) == 0): + raise ValueError("All dimensions of input must be of equal length") + + return diag_indices(arr.shape[0], arr.ndim) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..4a1426fd4d6cb2eeaf7327b8485efc4ca14670af --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_index_tricks_impl.pyi @@ -0,0 +1,196 @@ +from collections.abc import Sequence +from typing import Any, ClassVar, Final, Generic, SupportsIndex, final, overload +from typing import Literal as L + +from _typeshed import Incomplete +from typing_extensions import Self, TypeVar, deprecated + +import numpy as np +from numpy._core.multiarray import ravel_multi_index, unravel_index +from numpy._typing import ( + ArrayLike, + NDArray, + _FiniteNestedSequence, + _NestedSequence, + _Shape, + _SupportsArray, + _SupportsDType, +) + +__all__ = [ # noqa: RUF022 + "ravel_multi_index", + "unravel_index", + "mgrid", + "ogrid", + "r_", + "c_", + "s_", + "index_exp", + "ix_", + "ndenumerate", + "ndindex", + "fill_diagonal", + "diag_indices", + "diag_indices_from", +] + +### + +_T = TypeVar("_T") +_TupleT = TypeVar("_TupleT", bound=tuple[Any, ...]) +_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any]) +_DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any]) +_ScalarT = TypeVar("_ScalarT", bound=np.generic) +_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, covariant=True) +_BoolT_co = TypeVar("_BoolT_co", bound=bool, default=bool, covariant=True) + +_AxisT_co = TypeVar("_AxisT_co", bound=int, default=L[0], covariant=True) +_MatrixT_co = TypeVar("_MatrixT_co", bound=bool, default=L[False], covariant=True) +_NDMinT_co = TypeVar("_NDMinT_co", bound=int, default=L[1], covariant=True) +_Trans1DT_co = TypeVar("_Trans1DT_co", bound=int, default=L[-1], covariant=True) + +### + +class ndenumerate(Generic[_ScalarT_co]): + @overload + def __new__(cls, arr: _FiniteNestedSequence[_SupportsArray[np.dtype[_ScalarT]]]) -> ndenumerate[_ScalarT]: ... + @overload + def __new__(cls, arr: str | _NestedSequence[str]) -> ndenumerate[np.str_]: ... + @overload + def __new__(cls, arr: bytes | _NestedSequence[bytes]) -> ndenumerate[np.bytes_]: ... + @overload + def __new__(cls, arr: bool | _NestedSequence[bool]) -> ndenumerate[np.bool]: ... + @overload + def __new__(cls, arr: int | _NestedSequence[int]) -> ndenumerate[np.intp]: ... + @overload + def __new__(cls, arr: float | _NestedSequence[float]) -> ndenumerate[np.float64]: ... + @overload + def __new__(cls, arr: complex | _NestedSequence[complex]) -> ndenumerate[np.complex128]: ... + @overload + def __new__(cls, arr: object) -> ndenumerate[Any]: ... + + # The first overload is a (semi-)workaround for a mypy bug (tested with v1.10 and v1.11) + @overload + def __next__( + self: ndenumerate[np.bool | np.number | np.flexible | np.datetime64 | np.timedelta64], + /, + ) -> tuple[tuple[int, ...], _ScalarT_co]: ... + @overload + def __next__(self: ndenumerate[np.object_], /) -> tuple[tuple[int, ...], Any]: ... + @overload + def __next__(self, /) -> tuple[tuple[int, ...], _ScalarT_co]: ... + + # + def __iter__(self) -> Self: ... + +class ndindex: + @overload + def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ... + @overload + def __init__(self, /, *shape: SupportsIndex) -> None: ... + + # + def __iter__(self) -> Self: ... + def __next__(self) -> tuple[int, ...]: ... + + # + @deprecated("Deprecated since 1.20.0.") + def ndincr(self, /) -> None: ... + +class nd_grid(Generic[_BoolT_co]): + sparse: _BoolT_co + def __init__(self, sparse: _BoolT_co = ...) -> None: ... + @overload + def __getitem__(self: nd_grid[L[False]], key: slice | Sequence[slice]) -> NDArray[Any]: ... + @overload + def __getitem__(self: nd_grid[L[True]], key: slice | Sequence[slice]) -> tuple[NDArray[Any], ...]: ... + +@final +class MGridClass(nd_grid[L[False]]): + def __init__(self) -> None: ... + +@final +class OGridClass(nd_grid[L[True]]): + def __init__(self) -> None: ... + +class AxisConcatenator(Generic[_AxisT_co, _MatrixT_co, _NDMinT_co, _Trans1DT_co]): + __slots__ = "axis", "matrix", "ndmin", "trans1d" + + makemat: ClassVar[type[np.matrix[tuple[int, int], np.dtype[Any]]]] + + axis: _AxisT_co + matrix: _MatrixT_co + ndmin: _NDMinT_co + trans1d: _Trans1DT_co + + # + def __init__( + self, + /, + axis: _AxisT_co = ..., + matrix: _MatrixT_co = ..., + ndmin: _NDMinT_co = ..., + trans1d: _Trans1DT_co = ..., + ) -> None: ... + + # TODO(jorenham): annotate this + def __getitem__(self, key: Incomplete, /) -> Incomplete: ... + def __len__(self, /) -> L[0]: ... + + # + @staticmethod + @overload + def concatenate(*a: ArrayLike, axis: SupportsIndex | None = 0, out: _ArrayT) -> _ArrayT: ... + @staticmethod + @overload + def concatenate(*a: ArrayLike, axis: SupportsIndex | None = 0, out: None = None) -> NDArray[Any]: ... + +@final +class RClass(AxisConcatenator[L[0], L[False], L[1], L[-1]]): + def __init__(self, /) -> None: ... + +@final +class CClass(AxisConcatenator[L[-1], L[False], L[2], L[0]]): + def __init__(self, /) -> None: ... + +class IndexExpression(Generic[_BoolT_co]): + maketuple: _BoolT_co + def __init__(self, maketuple: _BoolT_co) -> None: ... + @overload + def __getitem__(self, item: _TupleT) -> _TupleT: ... + @overload + def __getitem__(self: IndexExpression[L[True]], item: _T) -> tuple[_T]: ... + @overload + def __getitem__(self: IndexExpression[L[False]], item: _T) -> _T: ... + +@overload +def ix_(*args: _FiniteNestedSequence[_SupportsDType[_DTypeT]]) -> tuple[np.ndarray[_Shape, _DTypeT], ...]: ... +@overload +def ix_(*args: str | _NestedSequence[str]) -> tuple[NDArray[np.str_], ...]: ... +@overload +def ix_(*args: bytes | _NestedSequence[bytes]) -> tuple[NDArray[np.bytes_], ...]: ... +@overload +def ix_(*args: bool | _NestedSequence[bool]) -> tuple[NDArray[np.bool], ...]: ... +@overload +def ix_(*args: int | _NestedSequence[int]) -> tuple[NDArray[np.intp], ...]: ... +@overload +def ix_(*args: float | _NestedSequence[float]) -> tuple[NDArray[np.float64], ...]: ... +@overload +def ix_(*args: complex | _NestedSequence[complex]) -> tuple[NDArray[np.complex128], ...]: ... + +# +def fill_diagonal(a: NDArray[Any], val: object, wrap: bool = ...) -> None: ... + +# +def diag_indices(n: int, ndim: int = ...) -> tuple[NDArray[np.intp], ...]: ... +def diag_indices_from(arr: ArrayLike) -> tuple[NDArray[np.intp], ...]: ... + +# +mgrid: Final[MGridClass] = ... +ogrid: Final[OGridClass] = ... + +r_: Final[RClass] = ... +c_: Final[CClass] = ... + +index_exp: Final[IndexExpression[L[True]]] = ... +s_: Final[IndexExpression[L[False]]] = ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_iotools.py b/venv/lib/python3.10/site-packages/numpy/lib/_iotools.py new file mode 100644 index 0000000000000000000000000000000000000000..908ca7762fdd16d9b725e29c7851366ea4d4bec7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_iotools.py @@ -0,0 +1,899 @@ +"""A collection of functions designed to help I/O with ascii files. + +""" +__docformat__ = "restructuredtext en" + +import numpy as np +import numpy._core.numeric as nx +from numpy._utils import asbytes, asunicode + + +def _decode_line(line, encoding=None): + """Decode bytes from binary input streams. + + Defaults to decoding from 'latin1'. That differs from the behavior of + np.compat.asunicode that decodes from 'ascii'. + + Parameters + ---------- + line : str or bytes + Line to be decoded. + encoding : str + Encoding used to decode `line`. + + Returns + ------- + decoded_line : str + + """ + if type(line) is bytes: + if encoding is None: + encoding = "latin1" + line = line.decode(encoding) + + return line + + +def _is_string_like(obj): + """ + Check whether obj behaves like a string. + """ + try: + obj + '' + except (TypeError, ValueError): + return False + return True + + +def _is_bytes_like(obj): + """ + Check whether obj behaves like a bytes object. + """ + try: + obj + b'' + except (TypeError, ValueError): + return False + return True + + +def has_nested_fields(ndtype): + """ + Returns whether one or several fields of a dtype are nested. + + Parameters + ---------- + ndtype : dtype + Data-type of a structured array. + + Raises + ------ + AttributeError + If `ndtype` does not have a `names` attribute. + + Examples + -------- + >>> import numpy as np + >>> dt = np.dtype([('name', 'S4'), ('x', float), ('y', float)]) + >>> np.lib._iotools.has_nested_fields(dt) + False + + """ + return any(ndtype[name].names is not None for name in ndtype.names or ()) + + +def flatten_dtype(ndtype, flatten_base=False): + """ + Unpack a structured data-type by collapsing nested fields and/or fields + with a shape. + + Note that the field names are lost. + + Parameters + ---------- + ndtype : dtype + The datatype to collapse + flatten_base : bool, optional + If True, transform a field with a shape into several fields. Default is + False. + + Examples + -------- + >>> import numpy as np + >>> dt = np.dtype([('name', 'S4'), ('x', float), ('y', float), + ... ('block', int, (2, 3))]) + >>> np.lib._iotools.flatten_dtype(dt) + [dtype('S4'), dtype('float64'), dtype('float64'), dtype('int64')] + >>> np.lib._iotools.flatten_dtype(dt, flatten_base=True) + [dtype('S4'), + dtype('float64'), + dtype('float64'), + dtype('int64'), + dtype('int64'), + dtype('int64'), + dtype('int64'), + dtype('int64'), + dtype('int64')] + + """ + names = ndtype.names + if names is None: + if flatten_base: + return [ndtype.base] * int(np.prod(ndtype.shape)) + return [ndtype.base] + else: + types = [] + for field in names: + info = ndtype.fields[field] + flat_dt = flatten_dtype(info[0], flatten_base) + types.extend(flat_dt) + return types + + +class LineSplitter: + """ + Object to split a string at a given delimiter or at given places. + + Parameters + ---------- + delimiter : str, int, or sequence of ints, optional + If a string, character used to delimit consecutive fields. + If an integer or a sequence of integers, width(s) of each field. + comments : str, optional + Character used to mark the beginning of a comment. Default is '#'. + autostrip : bool, optional + Whether to strip each individual field. Default is True. + + """ + + def autostrip(self, method): + """ + Wrapper to strip each member of the output of `method`. + + Parameters + ---------- + method : function + Function that takes a single argument and returns a sequence of + strings. + + Returns + ------- + wrapped : function + The result of wrapping `method`. `wrapped` takes a single input + argument and returns a list of strings that are stripped of + white-space. + + """ + return lambda input: [_.strip() for _ in method(input)] + + def __init__(self, delimiter=None, comments='#', autostrip=True, + encoding=None): + delimiter = _decode_line(delimiter) + comments = _decode_line(comments) + + self.comments = comments + + # Delimiter is a character + if (delimiter is None) or isinstance(delimiter, str): + delimiter = delimiter or None + _handyman = self._delimited_splitter + # Delimiter is a list of field widths + elif hasattr(delimiter, '__iter__'): + _handyman = self._variablewidth_splitter + idx = np.cumsum([0] + list(delimiter)) + delimiter = [slice(i, j) for (i, j) in zip(idx[:-1], idx[1:])] + # Delimiter is a single integer + elif int(delimiter): + (_handyman, delimiter) = ( + self._fixedwidth_splitter, int(delimiter)) + else: + (_handyman, delimiter) = (self._delimited_splitter, None) + self.delimiter = delimiter + if autostrip: + self._handyman = self.autostrip(_handyman) + else: + self._handyman = _handyman + self.encoding = encoding + + def _delimited_splitter(self, line): + """Chop off comments, strip, and split at delimiter. """ + if self.comments is not None: + line = line.split(self.comments)[0] + line = line.strip(" \r\n") + if not line: + return [] + return line.split(self.delimiter) + + def _fixedwidth_splitter(self, line): + if self.comments is not None: + line = line.split(self.comments)[0] + line = line.strip("\r\n") + if not line: + return [] + fixed = self.delimiter + slices = [slice(i, i + fixed) for i in range(0, len(line), fixed)] + return [line[s] for s in slices] + + def _variablewidth_splitter(self, line): + if self.comments is not None: + line = line.split(self.comments)[0] + if not line: + return [] + slices = self.delimiter + return [line[s] for s in slices] + + def __call__(self, line): + return self._handyman(_decode_line(line, self.encoding)) + + +class NameValidator: + """ + Object to validate a list of strings to use as field names. + + The strings are stripped of any non alphanumeric character, and spaces + are replaced by '_'. During instantiation, the user can define a list + of names to exclude, as well as a list of invalid characters. Names in + the exclusion list are appended a '_' character. + + Once an instance has been created, it can be called with a list of + names, and a list of valid names will be created. The `__call__` + method accepts an optional keyword "default" that sets the default name + in case of ambiguity. By default this is 'f', so that names will + default to `f0`, `f1`, etc. + + Parameters + ---------- + excludelist : sequence, optional + A list of names to exclude. This list is appended to the default + list ['return', 'file', 'print']. Excluded names are appended an + underscore: for example, `file` becomes `file_` if supplied. + deletechars : str, optional + A string combining invalid characters that must be deleted from the + names. + case_sensitive : {True, False, 'upper', 'lower'}, optional + * If True, field names are case-sensitive. + * If False or 'upper', field names are converted to upper case. + * If 'lower', field names are converted to lower case. + + The default value is True. + replace_space : '_', optional + Character(s) used in replacement of white spaces. + + Notes + ----- + Calling an instance of `NameValidator` is the same as calling its + method `validate`. + + Examples + -------- + >>> import numpy as np + >>> validator = np.lib._iotools.NameValidator() + >>> validator(['file', 'field2', 'with space', 'CaSe']) + ('file_', 'field2', 'with_space', 'CaSe') + + >>> validator = np.lib._iotools.NameValidator(excludelist=['excl'], + ... deletechars='q', + ... case_sensitive=False) + >>> validator(['excl', 'field2', 'no_q', 'with space', 'CaSe']) + ('EXCL', 'FIELD2', 'NO_Q', 'WITH_SPACE', 'CASE') + + """ + + defaultexcludelist = ['return', 'file', 'print'] + defaultdeletechars = set(r"""~!@#$%^&*()-=+~\|]}[{';: /?.>,<""") + + def __init__(self, excludelist=None, deletechars=None, + case_sensitive=None, replace_space='_'): + # Process the exclusion list .. + if excludelist is None: + excludelist = [] + excludelist.extend(self.defaultexcludelist) + self.excludelist = excludelist + # Process the list of characters to delete + if deletechars is None: + delete = self.defaultdeletechars + else: + delete = set(deletechars) + delete.add('"') + self.deletechars = delete + # Process the case option ..... + if (case_sensitive is None) or (case_sensitive is True): + self.case_converter = lambda x: x + elif (case_sensitive is False) or case_sensitive.startswith('u'): + self.case_converter = lambda x: x.upper() + elif case_sensitive.startswith('l'): + self.case_converter = lambda x: x.lower() + else: + msg = 'unrecognized case_sensitive value %s.' % case_sensitive + raise ValueError(msg) + + self.replace_space = replace_space + + def validate(self, names, defaultfmt="f%i", nbfields=None): + """ + Validate a list of strings as field names for a structured array. + + Parameters + ---------- + names : sequence of str + Strings to be validated. + defaultfmt : str, optional + Default format string, used if validating a given string + reduces its length to zero. + nbfields : integer, optional + Final number of validated names, used to expand or shrink the + initial list of names. + + Returns + ------- + validatednames : list of str + The list of validated field names. + + Notes + ----- + A `NameValidator` instance can be called directly, which is the + same as calling `validate`. For examples, see `NameValidator`. + + """ + # Initial checks .............. + if (names is None): + if (nbfields is None): + return None + names = [] + if isinstance(names, str): + names = [names, ] + if nbfields is not None: + nbnames = len(names) + if (nbnames < nbfields): + names = list(names) + [''] * (nbfields - nbnames) + elif (nbnames > nbfields): + names = names[:nbfields] + # Set some shortcuts ........... + deletechars = self.deletechars + excludelist = self.excludelist + case_converter = self.case_converter + replace_space = self.replace_space + # Initializes some variables ... + validatednames = [] + seen = dict() + nbempty = 0 + + for item in names: + item = case_converter(item).strip() + if replace_space: + item = item.replace(' ', replace_space) + item = ''.join([c for c in item if c not in deletechars]) + if item == '': + item = defaultfmt % nbempty + while item in names: + nbempty += 1 + item = defaultfmt % nbempty + nbempty += 1 + elif item in excludelist: + item += '_' + cnt = seen.get(item, 0) + if cnt > 0: + validatednames.append(item + '_%d' % cnt) + else: + validatednames.append(item) + seen[item] = cnt + 1 + return tuple(validatednames) + + def __call__(self, names, defaultfmt="f%i", nbfields=None): + return self.validate(names, defaultfmt=defaultfmt, nbfields=nbfields) + + +def str2bool(value): + """ + Tries to transform a string supposed to represent a boolean to a boolean. + + Parameters + ---------- + value : str + The string that is transformed to a boolean. + + Returns + ------- + boolval : bool + The boolean representation of `value`. + + Raises + ------ + ValueError + If the string is not 'True' or 'False' (case independent) + + Examples + -------- + >>> import numpy as np + >>> np.lib._iotools.str2bool('TRUE') + True + >>> np.lib._iotools.str2bool('false') + False + + """ + value = value.upper() + if value == 'TRUE': + return True + elif value == 'FALSE': + return False + else: + raise ValueError("Invalid boolean") + + +class ConverterError(Exception): + """ + Exception raised when an error occurs in a converter for string values. + + """ + pass + + +class ConverterLockError(ConverterError): + """ + Exception raised when an attempt is made to upgrade a locked converter. + + """ + pass + + +class ConversionWarning(UserWarning): + """ + Warning issued when a string converter has a problem. + + Notes + ----- + In `genfromtxt` a `ConversionWarning` is issued if raising exceptions + is explicitly suppressed with the "invalid_raise" keyword. + + """ + pass + + +class StringConverter: + """ + Factory class for function transforming a string into another object + (int, float). + + After initialization, an instance can be called to transform a string + into another object. If the string is recognized as representing a + missing value, a default value is returned. + + Attributes + ---------- + func : function + Function used for the conversion. + default : any + Default value to return when the input corresponds to a missing + value. + type : type + Type of the output. + _status : int + Integer representing the order of the conversion. + _mapper : sequence of tuples + Sequence of tuples (dtype, function, default value) to evaluate in + order. + _locked : bool + Holds `locked` parameter. + + Parameters + ---------- + dtype_or_func : {None, dtype, function}, optional + If a `dtype`, specifies the input data type, used to define a basic + function and a default value for missing data. For example, when + `dtype` is float, the `func` attribute is set to `float` and the + default value to `np.nan`. If a function, this function is used to + convert a string to another object. In this case, it is recommended + to give an associated default value as input. + default : any, optional + Value to return by default, that is, when the string to be + converted is flagged as missing. If not given, `StringConverter` + tries to supply a reasonable default value. + missing_values : {None, sequence of str}, optional + ``None`` or sequence of strings indicating a missing value. If ``None`` + then missing values are indicated by empty entries. The default is + ``None``. + locked : bool, optional + Whether the StringConverter should be locked to prevent automatic + upgrade or not. Default is False. + + """ + _mapper = [(nx.bool, str2bool, False), + (nx.int_, int, -1),] + + # On 32-bit systems, we need to make sure that we explicitly include + # nx.int64 since ns.int_ is nx.int32. + if nx.dtype(nx.int_).itemsize < nx.dtype(nx.int64).itemsize: + _mapper.append((nx.int64, int, -1)) + + _mapper.extend([(nx.float64, float, nx.nan), + (nx.complex128, complex, nx.nan + 0j), + (nx.longdouble, nx.longdouble, nx.nan), + # If a non-default dtype is passed, fall back to generic + # ones (should only be used for the converter) + (nx.integer, int, -1), + (nx.floating, float, nx.nan), + (nx.complexfloating, complex, nx.nan + 0j), + # Last, try with the string types (must be last, because + # `_mapper[-1]` is used as default in some cases) + (nx.str_, asunicode, '???'), + (nx.bytes_, asbytes, '???'), + ]) + + @classmethod + def _getdtype(cls, val): + """Returns the dtype of the input variable.""" + return np.array(val).dtype + + @classmethod + def _getsubdtype(cls, val): + """Returns the type of the dtype of the input variable.""" + return np.array(val).dtype.type + + @classmethod + def _dtypeortype(cls, dtype): + """Returns dtype for datetime64 and type of dtype otherwise.""" + + # This is a bit annoying. We want to return the "general" type in most + # cases (ie. "string" rather than "S10"), but we want to return the + # specific type for datetime64 (ie. "datetime64[us]" rather than + # "datetime64"). + if dtype.type == np.datetime64: + return dtype + return dtype.type + + @classmethod + def upgrade_mapper(cls, func, default=None): + """ + Upgrade the mapper of a StringConverter by adding a new function and + its corresponding default. + + The input function (or sequence of functions) and its associated + default value (if any) is inserted in penultimate position of the + mapper. The corresponding type is estimated from the dtype of the + default value. + + Parameters + ---------- + func : var + Function, or sequence of functions + + Examples + -------- + >>> import dateutil.parser + >>> import datetime + >>> dateparser = dateutil.parser.parse + >>> defaultdate = datetime.date(2000, 1, 1) + >>> StringConverter.upgrade_mapper(dateparser, default=defaultdate) + """ + # Func is a single functions + if callable(func): + cls._mapper.insert(-1, (cls._getsubdtype(default), func, default)) + return + elif hasattr(func, '__iter__'): + if isinstance(func[0], (tuple, list)): + for _ in func: + cls._mapper.insert(-1, _) + return + if default is None: + default = [None] * len(func) + else: + default = list(default) + default.append([None] * (len(func) - len(default))) + for fct, dft in zip(func, default): + cls._mapper.insert(-1, (cls._getsubdtype(dft), fct, dft)) + + @classmethod + def _find_map_entry(cls, dtype): + # if a converter for the specific dtype is available use that + for i, (deftype, func, default_def) in enumerate(cls._mapper): + if dtype.type == deftype: + return i, (deftype, func, default_def) + + # otherwise find an inexact match + for i, (deftype, func, default_def) in enumerate(cls._mapper): + if np.issubdtype(dtype.type, deftype): + return i, (deftype, func, default_def) + + raise LookupError + + def __init__(self, dtype_or_func=None, default=None, missing_values=None, + locked=False): + # Defines a lock for upgrade + self._locked = bool(locked) + # No input dtype: minimal initialization + if dtype_or_func is None: + self.func = str2bool + self._status = 0 + self.default = default or False + dtype = np.dtype('bool') + else: + # Is the input a np.dtype ? + try: + self.func = None + dtype = np.dtype(dtype_or_func) + except TypeError: + # dtype_or_func must be a function, then + if not callable(dtype_or_func): + errmsg = ("The input argument `dtype` is neither a" + " function nor a dtype (got '%s' instead)") + raise TypeError(errmsg % type(dtype_or_func)) + # Set the function + self.func = dtype_or_func + # If we don't have a default, try to guess it or set it to + # None + if default is None: + try: + default = self.func('0') + except ValueError: + default = None + dtype = self._getdtype(default) + + # find the best match in our mapper + try: + self._status, (_, func, default_def) = self._find_map_entry(dtype) + except LookupError: + # no match + self.default = default + _, func, _ = self._mapper[-1] + self._status = 0 + else: + # use the found default only if we did not already have one + if default is None: + self.default = default_def + else: + self.default = default + + # If the input was a dtype, set the function to the last we saw + if self.func is None: + self.func = func + + # If the status is 1 (int), change the function to + # something more robust. + if self.func == self._mapper[1][1]: + if issubclass(dtype.type, np.uint64): + self.func = np.uint64 + elif issubclass(dtype.type, np.int64): + self.func = np.int64 + else: + self.func = lambda x: int(float(x)) + # Store the list of strings corresponding to missing values. + if missing_values is None: + self.missing_values = {''} + else: + if isinstance(missing_values, str): + missing_values = missing_values.split(",") + self.missing_values = set(list(missing_values) + ['']) + + self._callingfunction = self._strict_call + self.type = self._dtypeortype(dtype) + self._checked = False + self._initial_default = default + + def _loose_call(self, value): + try: + return self.func(value) + except ValueError: + return self.default + + def _strict_call(self, value): + try: + + # We check if we can convert the value using the current function + new_value = self.func(value) + + # In addition to having to check whether func can convert the + # value, we also have to make sure that we don't get overflow + # errors for integers. + if self.func is int: + try: + np.array(value, dtype=self.type) + except OverflowError: + raise ValueError + + # We're still here so we can now return the new value + return new_value + + except ValueError: + if value.strip() in self.missing_values: + if not self._status: + self._checked = False + return self.default + raise ValueError("Cannot convert string '%s'" % value) + + def __call__(self, value): + return self._callingfunction(value) + + def _do_upgrade(self): + # Raise an exception if we locked the converter... + if self._locked: + errmsg = "Converter is locked and cannot be upgraded" + raise ConverterLockError(errmsg) + _statusmax = len(self._mapper) + # Complains if we try to upgrade by the maximum + _status = self._status + if _status == _statusmax: + errmsg = "Could not find a valid conversion function" + raise ConverterError(errmsg) + elif _status < _statusmax - 1: + _status += 1 + self.type, self.func, default = self._mapper[_status] + self._status = _status + if self._initial_default is not None: + self.default = self._initial_default + else: + self.default = default + + def upgrade(self, value): + """ + Find the best converter for a given string, and return the result. + + The supplied string `value` is converted by testing different + converters in order. First the `func` method of the + `StringConverter` instance is tried, if this fails other available + converters are tried. The order in which these other converters + are tried is determined by the `_status` attribute of the instance. + + Parameters + ---------- + value : str + The string to convert. + + Returns + ------- + out : any + The result of converting `value` with the appropriate converter. + + """ + self._checked = True + try: + return self._strict_call(value) + except ValueError: + self._do_upgrade() + return self.upgrade(value) + + def iterupgrade(self, value): + self._checked = True + if not hasattr(value, '__iter__'): + value = (value,) + _strict_call = self._strict_call + try: + for _m in value: + _strict_call(_m) + except ValueError: + self._do_upgrade() + self.iterupgrade(value) + + def update(self, func, default=None, testing_value=None, + missing_values='', locked=False): + """ + Set StringConverter attributes directly. + + Parameters + ---------- + func : function + Conversion function. + default : any, optional + Value to return by default, that is, when the string to be + converted is flagged as missing. If not given, + `StringConverter` tries to supply a reasonable default value. + testing_value : str, optional + A string representing a standard input value of the converter. + This string is used to help defining a reasonable default + value. + missing_values : {sequence of str, None}, optional + Sequence of strings indicating a missing value. If ``None``, then + the existing `missing_values` are cleared. The default is ``''``. + locked : bool, optional + Whether the StringConverter should be locked to prevent + automatic upgrade or not. Default is False. + + Notes + ----- + `update` takes the same parameters as the constructor of + `StringConverter`, except that `func` does not accept a `dtype` + whereas `dtype_or_func` in the constructor does. + + """ + self.func = func + self._locked = locked + + # Don't reset the default to None if we can avoid it + if default is not None: + self.default = default + self.type = self._dtypeortype(self._getdtype(default)) + else: + try: + tester = func(testing_value or '1') + except (TypeError, ValueError): + tester = None + self.type = self._dtypeortype(self._getdtype(tester)) + + # Add the missing values to the existing set or clear it. + if missing_values is None: + # Clear all missing values even though the ctor initializes it to + # set(['']) when the argument is None. + self.missing_values = set() + else: + if not np.iterable(missing_values): + missing_values = [missing_values] + if not all(isinstance(v, str) for v in missing_values): + raise TypeError("missing_values must be strings or unicode") + self.missing_values.update(missing_values) + + +def easy_dtype(ndtype, names=None, defaultfmt="f%i", **validationargs): + """ + Convenience function to create a `np.dtype` object. + + The function processes the input `dtype` and matches it with the given + names. + + Parameters + ---------- + ndtype : var + Definition of the dtype. Can be any string or dictionary recognized + by the `np.dtype` function, or a sequence of types. + names : str or sequence, optional + Sequence of strings to use as field names for a structured dtype. + For convenience, `names` can be a string of a comma-separated list + of names. + defaultfmt : str, optional + Format string used to define missing names, such as ``"f%i"`` + (default) or ``"fields_%02i"``. + validationargs : optional + A series of optional arguments used to initialize a + `NameValidator`. + + Examples + -------- + >>> import numpy as np + >>> np.lib._iotools.easy_dtype(float) + dtype('float64') + >>> np.lib._iotools.easy_dtype("i4, f8") + dtype([('f0', '>> np.lib._iotools.easy_dtype("i4, f8", defaultfmt="field_%03i") + dtype([('field_000', '>> np.lib._iotools.easy_dtype((int, float, float), names="a,b,c") + dtype([('a', '>> np.lib._iotools.easy_dtype(float, names="a,b,c") + dtype([('a', ' None: ... + def __call__(self, /, line: str | bytes) -> list[str]: ... + def autostrip(self, /, method: Callable[[_T], Iterable[str]]) -> Callable[[_T], list[str]]: ... + +class NameValidator: + defaultexcludelist: ClassVar[Sequence[str]] + defaultdeletechars: ClassVar[Sequence[str]] + excludelist: list[str] + deletechars: set[str] + case_converter: Callable[[str], str] + replace_space: str + + def __init__( + self, + /, + excludelist: Iterable[str] | None = None, + deletechars: Iterable[str] | None = None, + case_sensitive: Literal["upper", "lower"] | bool | None = None, + replace_space: str = "_", + ) -> None: ... + def __call__(self, /, names: Iterable[str], defaultfmt: str = "f%i", nbfields: int | None = None) -> tuple[str, ...]: ... + def validate(self, /, names: Iterable[str], defaultfmt: str = "f%i", nbfields: int | None = None) -> tuple[str, ...]: ... + +class StringConverter: + func: Callable[[str], Any] | None + default: Any + missing_values: set[str] + type: np.dtype[np.datetime64] | np.generic + + def __init__( + self, + /, + dtype_or_func: npt.DTypeLike | None = None, + default: None = None, + missing_values: Iterable[str] | None = None, + locked: bool = False, + ) -> None: ... + def update( + self, + /, + func: Callable[[str], Any], + default: object | None = None, + testing_value: str | None = None, + missing_values: str = "", + locked: bool = False, + ) -> None: ... + # + def __call__(self, /, value: str) -> Any: ... + def upgrade(self, /, value: str) -> Any: ... + def iterupgrade(self, /, value: Iterable[str] | str) -> None: ... + + # + @classmethod + def upgrade_mapper(cls, func: Callable[[str], Any], default: object | None = None) -> None: ... + +@overload +def str2bool(value: Literal["false", "False", "FALSE"]) -> Literal[False]: ... +@overload +def str2bool(value: Literal["true", "True", "TRUE"]) -> Literal[True]: ... + +# +def has_nested_fields(ndtype: np.dtype[np.void]) -> bool: ... +def flatten_dtype(ndtype: np.dtype[np.void], flatten_base: bool = False) -> type[np.dtype[Any]]: ... +def easy_dtype( + ndtype: npt.DTypeLike, + names: Iterable[str] | None = None, + defaultfmt: str = "f%i", + **validationargs: Unpack[_ValidationKwargs], +) -> np.dtype[np.void]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..9d0173dbe34018659ab4d9d06bc71ab2189b2fae --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.py @@ -0,0 +1,2028 @@ +""" +Functions that ignore NaN. + +Functions +--------- + +- `nanmin` -- minimum non-NaN value +- `nanmax` -- maximum non-NaN value +- `nanargmin` -- index of minimum non-NaN value +- `nanargmax` -- index of maximum non-NaN value +- `nansum` -- sum of non-NaN values +- `nanprod` -- product of non-NaN values +- `nancumsum` -- cumulative sum of non-NaN values +- `nancumprod` -- cumulative product of non-NaN values +- `nanmean` -- mean of non-NaN values +- `nanvar` -- variance of non-NaN values +- `nanstd` -- standard deviation of non-NaN values +- `nanmedian` -- median of non-NaN values +- `nanquantile` -- qth quantile of non-NaN values +- `nanpercentile` -- qth percentile of non-NaN values + +""" +import functools +import warnings +import numpy as np +import numpy._core.numeric as _nx +from numpy.lib import _function_base_impl as fnb +from numpy.lib._function_base_impl import _weights_are_valid +from numpy._core import overrides + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +__all__ = [ + 'nansum', 'nanmax', 'nanmin', 'nanargmax', 'nanargmin', 'nanmean', + 'nanmedian', 'nanpercentile', 'nanvar', 'nanstd', 'nanprod', + 'nancumsum', 'nancumprod', 'nanquantile' + ] + + +def _nan_mask(a, out=None): + """ + Parameters + ---------- + a : array-like + Input array with at least 1 dimension. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output and will prevent the allocation of a new array. + + Returns + ------- + y : bool ndarray or True + A bool array where ``np.nan`` positions are marked with ``False`` + and other positions are marked with ``True``. If the type of ``a`` + is such that it can't possibly contain ``np.nan``, returns ``True``. + """ + # we assume that a is an array for this private function + + if a.dtype.kind not in 'fc': + return True + + y = np.isnan(a, out=out) + y = np.invert(y, out=y) + return y + +def _replace_nan(a, val): + """ + If `a` is of inexact type, make a copy of `a`, replace NaNs with + the `val` value, and return the copy together with a boolean mask + marking the locations where NaNs were present. If `a` is not of + inexact type, do nothing and return `a` together with a mask of None. + + Note that scalars will end up as array scalars, which is important + for using the result as the value of the out argument in some + operations. + + Parameters + ---------- + a : array-like + Input array. + val : float + NaN values are set to val before doing the operation. + + Returns + ------- + y : ndarray + If `a` is of inexact type, return a copy of `a` with the NaNs + replaced by the fill value, otherwise return `a`. + mask: {bool, None} + If `a` is of inexact type, return a boolean mask marking locations of + NaNs, otherwise return None. + + """ + a = np.asanyarray(a) + + if a.dtype == np.object_: + # object arrays do not support `isnan` (gh-9009), so make a guess + mask = np.not_equal(a, a, dtype=bool) + elif issubclass(a.dtype.type, np.inexact): + mask = np.isnan(a) + else: + mask = None + + if mask is not None: + a = np.array(a, subok=True, copy=True) + np.copyto(a, val, where=mask) + + return a, mask + + +def _copyto(a, val, mask): + """ + Replace values in `a` with NaN where `mask` is True. This differs from + copyto in that it will deal with the case where `a` is a numpy scalar. + + Parameters + ---------- + a : ndarray or numpy scalar + Array or numpy scalar some of whose values are to be replaced + by val. + val : numpy scalar + Value used a replacement. + mask : ndarray, scalar + Boolean array. Where True the corresponding element of `a` is + replaced by `val`. Broadcasts. + + Returns + ------- + res : ndarray, scalar + Array with elements replaced or scalar `val`. + + """ + if isinstance(a, np.ndarray): + np.copyto(a, val, where=mask, casting='unsafe') + else: + a = a.dtype.type(val) + return a + + +def _remove_nan_1d(arr1d, second_arr1d=None, overwrite_input=False): + """ + Equivalent to arr1d[~arr1d.isnan()], but in a different order + + Presumably faster as it incurs fewer copies + + Parameters + ---------- + arr1d : ndarray + Array to remove nans from + second_arr1d : ndarray or None + A second array which will have the same positions removed as arr1d. + overwrite_input : bool + True if `arr1d` can be modified in place + + Returns + ------- + res : ndarray + Array with nan elements removed + second_res : ndarray or None + Second array with nan element positions of first array removed. + overwrite_input : bool + True if `res` can be modified in place, given the constraint on the + input + """ + if arr1d.dtype == object: + # object arrays do not support `isnan` (gh-9009), so make a guess + c = np.not_equal(arr1d, arr1d, dtype=bool) + else: + c = np.isnan(arr1d) + + s = np.nonzero(c)[0] + if s.size == arr1d.size: + warnings.warn("All-NaN slice encountered", RuntimeWarning, + stacklevel=6) + if second_arr1d is None: + return arr1d[:0], None, True + else: + return arr1d[:0], second_arr1d[:0], True + elif s.size == 0: + return arr1d, second_arr1d, overwrite_input + else: + if not overwrite_input: + arr1d = arr1d.copy() + # select non-nans at end of array + enonan = arr1d[-s.size:][~c[-s.size:]] + # fill nans in beginning of array with non-nans of end + arr1d[s[:enonan.size]] = enonan + + if second_arr1d is None: + return arr1d[:-s.size], None, True + else: + if not overwrite_input: + second_arr1d = second_arr1d.copy() + enonan = second_arr1d[-s.size:][~c[-s.size:]] + second_arr1d[s[:enonan.size]] = enonan + + return arr1d[:-s.size], second_arr1d[:-s.size], True + + +def _divide_by_count(a, b, out=None): + """ + Compute a/b ignoring invalid results. If `a` is an array the division + is done in place. If `a` is a scalar, then its type is preserved in the + output. If out is None, then a is used instead so that the division + is in place. Note that this is only called with `a` an inexact type. + + Parameters + ---------- + a : {ndarray, numpy scalar} + Numerator. Expected to be of inexact type but not checked. + b : {ndarray, numpy scalar} + Denominator. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output, but the type will be cast if necessary. + + Returns + ------- + ret : {ndarray, numpy scalar} + The return value is a/b. If `a` was an ndarray the division is done + in place. If `a` is a numpy scalar, the division preserves its type. + + """ + with np.errstate(invalid='ignore', divide='ignore'): + if isinstance(a, np.ndarray): + if out is None: + return np.divide(a, b, out=a, casting='unsafe') + else: + return np.divide(a, b, out=out, casting='unsafe') + else: + if out is None: + # Precaution against reduced object arrays + try: + return a.dtype.type(a / b) + except AttributeError: + return a / b + else: + # This is questionable, but currently a numpy scalar can + # be output to a zero dimensional array. + return np.divide(a, b, out=out, casting='unsafe') + + +def _nanmin_dispatcher(a, axis=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_nanmin_dispatcher) +def nanmin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return minimum of an array or minimum along an axis, ignoring any NaNs. + When all-NaN slices are encountered a ``RuntimeWarning`` is raised and + Nan is returned for that slice. + + Parameters + ---------- + a : array_like + Array containing numbers whose minimum is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the minimum is computed. The default is to compute + the minimum of the flattened array. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output, but the type will be cast if necessary. See + :ref:`ufuncs-output-type` for more details. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + + If the value is anything but the default, then + `keepdims` will be passed through to the `min` method + of sub-classes of `ndarray`. If the sub-classes methods + does not implement `keepdims` any exceptions will be raised. + initial : scalar, optional + The maximum value of an output element. Must be present to allow + computation on empty slice. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + where : array_like of bool, optional + Elements to compare for the minimum. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.22.0 + + Returns + ------- + nanmin : ndarray + An array with the same shape as `a`, with the specified axis + removed. If `a` is a 0-d array, or if axis is None, an ndarray + scalar is returned. The same dtype as `a` is returned. + + See Also + -------- + nanmax : + The maximum value of an array along a given axis, ignoring any NaNs. + amin : + The minimum value of an array along a given axis, propagating any NaNs. + fmin : + Element-wise minimum of two arrays, ignoring any NaNs. + minimum : + Element-wise minimum of two arrays, propagating any NaNs. + isnan : + Shows which elements are Not a Number (NaN). + isfinite: + Shows which elements are neither NaN nor infinity. + + amax, fmax, maximum + + Notes + ----- + NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic + (IEEE 754). This means that Not a Number is not equivalent to infinity. + Positive infinity is treated as a very large number and negative + infinity is treated as a very small (i.e. negative) number. + + If the input has a integer type the function is equivalent to np.min. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, np.nan]]) + >>> np.nanmin(a) + 1.0 + >>> np.nanmin(a, axis=0) + array([1., 2.]) + >>> np.nanmin(a, axis=1) + array([1., 3.]) + + When positive infinity and negative infinity are present: + + >>> np.nanmin([1, 2, np.nan, np.inf]) + 1.0 + >>> np.nanmin([1, 2, np.nan, -np.inf]) + -inf + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + if initial is not np._NoValue: + kwargs['initial'] = initial + if where is not np._NoValue: + kwargs['where'] = where + + if type(a) is np.ndarray and a.dtype != np.object_: + # Fast, but not safe for subclasses of ndarray, or object arrays, + # which do not implement isnan (gh-9009), or fmin correctly (gh-8975) + res = np.fmin.reduce(a, axis=axis, out=out, **kwargs) + if np.isnan(res).any(): + warnings.warn("All-NaN slice encountered", RuntimeWarning, + stacklevel=2) + else: + # Slow, but safe for subclasses of ndarray + a, mask = _replace_nan(a, +np.inf) + res = np.amin(a, axis=axis, out=out, **kwargs) + if mask is None: + return res + + # Check for all-NaN axis + kwargs.pop("initial", None) + mask = np.all(mask, axis=axis, **kwargs) + if np.any(mask): + res = _copyto(res, np.nan, mask) + warnings.warn("All-NaN axis encountered", RuntimeWarning, + stacklevel=2) + return res + + +def _nanmax_dispatcher(a, axis=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_nanmax_dispatcher) +def nanmax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): + """ + Return the maximum of an array or maximum along an axis, ignoring any + NaNs. When all-NaN slices are encountered a ``RuntimeWarning`` is + raised and NaN is returned for that slice. + + Parameters + ---------- + a : array_like + Array containing numbers whose maximum is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the maximum is computed. The default is to compute + the maximum of the flattened array. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output, but the type will be cast if necessary. See + :ref:`ufuncs-output-type` for more details. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + If the value is anything but the default, then + `keepdims` will be passed through to the `max` method + of sub-classes of `ndarray`. If the sub-classes methods + does not implement `keepdims` any exceptions will be raised. + initial : scalar, optional + The minimum value of an output element. Must be present to allow + computation on empty slice. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + where : array_like of bool, optional + Elements to compare for the maximum. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.22.0 + + Returns + ------- + nanmax : ndarray + An array with the same shape as `a`, with the specified axis removed. + If `a` is a 0-d array, or if axis is None, an ndarray scalar is + returned. The same dtype as `a` is returned. + + See Also + -------- + nanmin : + The minimum value of an array along a given axis, ignoring any NaNs. + amax : + The maximum value of an array along a given axis, propagating any NaNs. + fmax : + Element-wise maximum of two arrays, ignoring any NaNs. + maximum : + Element-wise maximum of two arrays, propagating any NaNs. + isnan : + Shows which elements are Not a Number (NaN). + isfinite: + Shows which elements are neither NaN nor infinity. + + amin, fmin, minimum + + Notes + ----- + NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic + (IEEE 754). This means that Not a Number is not equivalent to infinity. + Positive infinity is treated as a very large number and negative + infinity is treated as a very small (i.e. negative) number. + + If the input has a integer type the function is equivalent to np.max. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, 2], [3, np.nan]]) + >>> np.nanmax(a) + 3.0 + >>> np.nanmax(a, axis=0) + array([3., 2.]) + >>> np.nanmax(a, axis=1) + array([2., 3.]) + + When positive infinity and negative infinity are present: + + >>> np.nanmax([1, 2, np.nan, -np.inf]) + 2.0 + >>> np.nanmax([1, 2, np.nan, np.inf]) + inf + + """ + kwargs = {} + if keepdims is not np._NoValue: + kwargs['keepdims'] = keepdims + if initial is not np._NoValue: + kwargs['initial'] = initial + if where is not np._NoValue: + kwargs['where'] = where + + if type(a) is np.ndarray and a.dtype != np.object_: + # Fast, but not safe for subclasses of ndarray, or object arrays, + # which do not implement isnan (gh-9009), or fmax correctly (gh-8975) + res = np.fmax.reduce(a, axis=axis, out=out, **kwargs) + if np.isnan(res).any(): + warnings.warn("All-NaN slice encountered", RuntimeWarning, + stacklevel=2) + else: + # Slow, but safe for subclasses of ndarray + a, mask = _replace_nan(a, -np.inf) + res = np.amax(a, axis=axis, out=out, **kwargs) + if mask is None: + return res + + # Check for all-NaN axis + kwargs.pop("initial", None) + mask = np.all(mask, axis=axis, **kwargs) + if np.any(mask): + res = _copyto(res, np.nan, mask) + warnings.warn("All-NaN axis encountered", RuntimeWarning, + stacklevel=2) + return res + + +def _nanargmin_dispatcher(a, axis=None, out=None, *, keepdims=None): + return (a,) + + +@array_function_dispatch(_nanargmin_dispatcher) +def nanargmin(a, axis=None, out=None, *, keepdims=np._NoValue): + """ + Return the indices of the minimum values in the specified axis ignoring + NaNs. For all-NaN slices ``ValueError`` is raised. Warning: the results + cannot be trusted if a slice contains only NaNs and Infs. + + Parameters + ---------- + a : array_like + Input data. + axis : int, optional + Axis along which to operate. By default flattened input is used. + out : array, optional + If provided, the result will be inserted into this array. It should + be of the appropriate shape and dtype. + + .. versionadded:: 1.22.0 + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. + + .. versionadded:: 1.22.0 + + Returns + ------- + index_array : ndarray + An array of indices or a single index value. + + See Also + -------- + argmin, nanargmax + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[np.nan, 4], [2, 3]]) + >>> np.argmin(a) + 0 + >>> np.nanargmin(a) + 2 + >>> np.nanargmin(a, axis=0) + array([1, 1]) + >>> np.nanargmin(a, axis=1) + array([1, 0]) + + """ + a, mask = _replace_nan(a, np.inf) + if mask is not None and mask.size: + mask = np.all(mask, axis=axis) + if np.any(mask): + raise ValueError("All-NaN slice encountered") + res = np.argmin(a, axis=axis, out=out, keepdims=keepdims) + return res + + +def _nanargmax_dispatcher(a, axis=None, out=None, *, keepdims=None): + return (a,) + + +@array_function_dispatch(_nanargmax_dispatcher) +def nanargmax(a, axis=None, out=None, *, keepdims=np._NoValue): + """ + Return the indices of the maximum values in the specified axis ignoring + NaNs. For all-NaN slices ``ValueError`` is raised. Warning: the + results cannot be trusted if a slice contains only NaNs and -Infs. + + + Parameters + ---------- + a : array_like + Input data. + axis : int, optional + Axis along which to operate. By default flattened input is used. + out : array, optional + If provided, the result will be inserted into this array. It should + be of the appropriate shape and dtype. + + .. versionadded:: 1.22.0 + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. + + .. versionadded:: 1.22.0 + + Returns + ------- + index_array : ndarray + An array of indices or a single index value. + + See Also + -------- + argmax, nanargmin + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[np.nan, 4], [2, 3]]) + >>> np.argmax(a) + 0 + >>> np.nanargmax(a) + 1 + >>> np.nanargmax(a, axis=0) + array([1, 0]) + >>> np.nanargmax(a, axis=1) + array([1, 1]) + + """ + a, mask = _replace_nan(a, -np.inf) + if mask is not None and mask.size: + mask = np.all(mask, axis=axis) + if np.any(mask): + raise ValueError("All-NaN slice encountered") + res = np.argmax(a, axis=axis, out=out, keepdims=keepdims) + return res + + +def _nansum_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_nansum_dispatcher) +def nansum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): + """ + Return the sum of array elements over a given axis treating Not a + Numbers (NaNs) as zero. + + In NumPy versions <= 1.9.0 Nan is returned for slices that are all-NaN or + empty. In later versions zero is returned. + + Parameters + ---------- + a : array_like + Array containing numbers whose sum is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the sum is computed. The default is to compute the + sum of the flattened array. + dtype : data-type, optional + The type of the returned array and of the accumulator in which the + elements are summed. By default, the dtype of `a` is used. An + exception is when `a` has an integer type with less precision than + the platform (u)intp. In that case, the default will be either + (u)int32 or (u)int64 depending on whether the platform is 32 or 64 + bits. For inexact inputs, dtype must be inexact. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``. If provided, it must have the same shape as the + expected output, but the type will be cast if necessary. See + :ref:`ufuncs-output-type` for more details. The casting of NaN to integer + can yield unexpected results. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + + If the value is anything but the default, then + `keepdims` will be passed through to the `mean` or `sum` methods + of sub-classes of `ndarray`. If the sub-classes methods + does not implement `keepdims` any exceptions will be raised. + initial : scalar, optional + Starting value for the sum. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + where : array_like of bool, optional + Elements to include in the sum. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + + Returns + ------- + nansum : ndarray. + A new array holding the result is returned unless `out` is + specified, in which it is returned. The result has the same + size as `a`, and the same shape as `a` if `axis` is not None + or `a` is a 1-d array. + + See Also + -------- + numpy.sum : Sum across array propagating NaNs. + isnan : Show which elements are NaN. + isfinite : Show which elements are not NaN or +/-inf. + + Notes + ----- + If both positive and negative infinity are present, the sum will be Not + A Number (NaN). + + Examples + -------- + >>> import numpy as np + >>> np.nansum(1) + 1 + >>> np.nansum([1]) + 1 + >>> np.nansum([1, np.nan]) + 1.0 + >>> a = np.array([[1, 1], [1, np.nan]]) + >>> np.nansum(a) + 3.0 + >>> np.nansum(a, axis=0) + array([2., 1.]) + >>> np.nansum([1, np.nan, np.inf]) + inf + >>> np.nansum([1, np.nan, -np.inf]) + -inf + >>> from numpy.testing import suppress_warnings + >>> with np.errstate(invalid="ignore"): + ... np.nansum([1, np.nan, np.inf, -np.inf]) # both +/- infinity present + np.float64(nan) + + """ + a, mask = _replace_nan(a, 0) + return np.sum(a, axis=axis, dtype=dtype, out=out, keepdims=keepdims, + initial=initial, where=where) + + +def _nanprod_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + initial=None, where=None): + return (a, out) + + +@array_function_dispatch(_nanprod_dispatcher) +def nanprod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): + """ + Return the product of array elements over a given axis treating Not a + Numbers (NaNs) as ones. + + One is returned for slices that are all-NaN or empty. + + Parameters + ---------- + a : array_like + Array containing numbers whose product is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the product is computed. The default is to compute + the product of the flattened array. + dtype : data-type, optional + The type of the returned array and of the accumulator in which the + elements are summed. By default, the dtype of `a` is used. An + exception is when `a` has an integer type with less precision than + the platform (u)intp. In that case, the default will be either + (u)int32 or (u)int64 depending on whether the platform is 32 or 64 + bits. For inexact inputs, dtype must be inexact. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``. If provided, it must have the same shape as the + expected output, but the type will be cast if necessary. See + :ref:`ufuncs-output-type` for more details. The casting of NaN to integer + can yield unexpected results. + keepdims : bool, optional + If True, the axes which are reduced are left in the result as + dimensions with size one. With this option, the result will + broadcast correctly against the original `arr`. + initial : scalar, optional + The starting value for this product. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.22.0 + where : array_like of bool, optional + Elements to include in the product. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.22.0 + + Returns + ------- + nanprod : ndarray + A new array holding the result is returned unless `out` is + specified, in which case it is returned. + + See Also + -------- + numpy.prod : Product across array propagating NaNs. + isnan : Show which elements are NaN. + + Examples + -------- + >>> import numpy as np + >>> np.nanprod(1) + 1 + >>> np.nanprod([1]) + 1 + >>> np.nanprod([1, np.nan]) + 1.0 + >>> a = np.array([[1, 2], [3, np.nan]]) + >>> np.nanprod(a) + 6.0 + >>> np.nanprod(a, axis=0) + array([3., 2.]) + + """ + a, mask = _replace_nan(a, 1) + return np.prod(a, axis=axis, dtype=dtype, out=out, keepdims=keepdims, + initial=initial, where=where) + + +def _nancumsum_dispatcher(a, axis=None, dtype=None, out=None): + return (a, out) + + +@array_function_dispatch(_nancumsum_dispatcher) +def nancumsum(a, axis=None, dtype=None, out=None): + """ + Return the cumulative sum of array elements over a given axis treating Not a + Numbers (NaNs) as zero. The cumulative sum does not change when NaNs are + encountered and leading NaNs are replaced by zeros. + + Zeros are returned for slices that are all-NaN or empty. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + Axis along which the cumulative sum is computed. The default + (None) is to compute the cumsum over the flattened array. + dtype : dtype, optional + Type of the returned array and of the accumulator in which the + elements are summed. If `dtype` is not specified, it defaults + to the dtype of `a`, unless `a` has an integer dtype with a + precision less than that of the default platform integer. In + that case, the default platform integer is used. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type will be cast if necessary. See :ref:`ufuncs-output-type` for + more details. + + Returns + ------- + nancumsum : ndarray. + A new array holding the result is returned unless `out` is + specified, in which it is returned. The result has the same + size as `a`, and the same shape as `a` if `axis` is not None + or `a` is a 1-d array. + + See Also + -------- + numpy.cumsum : Cumulative sum across array propagating NaNs. + isnan : Show which elements are NaN. + + Examples + -------- + >>> import numpy as np + >>> np.nancumsum(1) + array([1]) + >>> np.nancumsum([1]) + array([1]) + >>> np.nancumsum([1, np.nan]) + array([1., 1.]) + >>> a = np.array([[1, 2], [3, np.nan]]) + >>> np.nancumsum(a) + array([1., 3., 6., 6.]) + >>> np.nancumsum(a, axis=0) + array([[1., 2.], + [4., 2.]]) + >>> np.nancumsum(a, axis=1) + array([[1., 3.], + [3., 3.]]) + + """ + a, mask = _replace_nan(a, 0) + return np.cumsum(a, axis=axis, dtype=dtype, out=out) + + +def _nancumprod_dispatcher(a, axis=None, dtype=None, out=None): + return (a, out) + + +@array_function_dispatch(_nancumprod_dispatcher) +def nancumprod(a, axis=None, dtype=None, out=None): + """ + Return the cumulative product of array elements over a given axis treating Not a + Numbers (NaNs) as one. The cumulative product does not change when NaNs are + encountered and leading NaNs are replaced by ones. + + Ones are returned for slices that are all-NaN or empty. + + Parameters + ---------- + a : array_like + Input array. + axis : int, optional + Axis along which the cumulative product is computed. By default + the input is flattened. + dtype : dtype, optional + Type of the returned array, as well as of the accumulator in which + the elements are multiplied. If *dtype* is not specified, it + defaults to the dtype of `a`, unless `a` has an integer dtype with + a precision less than that of the default platform integer. In + that case, the default platform integer is used instead. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output + but the type of the resulting values will be cast if necessary. + + Returns + ------- + nancumprod : ndarray + A new array holding the result is returned unless `out` is + specified, in which case it is returned. + + See Also + -------- + numpy.cumprod : Cumulative product across array propagating NaNs. + isnan : Show which elements are NaN. + + Examples + -------- + >>> import numpy as np + >>> np.nancumprod(1) + array([1]) + >>> np.nancumprod([1]) + array([1]) + >>> np.nancumprod([1, np.nan]) + array([1., 1.]) + >>> a = np.array([[1, 2], [3, np.nan]]) + >>> np.nancumprod(a) + array([1., 2., 6., 6.]) + >>> np.nancumprod(a, axis=0) + array([[1., 2.], + [3., 2.]]) + >>> np.nancumprod(a, axis=1) + array([[1., 2.], + [3., 3.]]) + + """ + a, mask = _replace_nan(a, 1) + return np.cumprod(a, axis=axis, dtype=dtype, out=out) + + +def _nanmean_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + *, where=None): + return (a, out) + + +@array_function_dispatch(_nanmean_dispatcher) +def nanmean(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + *, where=np._NoValue): + """ + Compute the arithmetic mean along the specified axis, ignoring NaNs. + + Returns the average of the array elements. The average is taken over + the flattened array by default, otherwise over the specified axis. + `float64` intermediate and return values are used for integer inputs. + + For all-NaN slices, NaN is returned and a `RuntimeWarning` is raised. + + Parameters + ---------- + a : array_like + Array containing numbers whose mean is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the means are computed. The default is to compute + the mean of the flattened array. + dtype : data-type, optional + Type to use in computing the mean. For integer inputs, the default + is `float64`; for inexact inputs, it is the same as the input + dtype. + out : ndarray, optional + Alternate output array in which to place the result. The default + is ``None``; if provided, it must have the same shape as the + expected output, but the type will be cast if necessary. + See :ref:`ufuncs-output-type` for more details. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + + If the value is anything but the default, then + `keepdims` will be passed through to the `mean` or `sum` methods + of sub-classes of `ndarray`. If the sub-classes methods + does not implement `keepdims` any exceptions will be raised. + where : array_like of bool, optional + Elements to include in the mean. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + + Returns + ------- + m : ndarray, see dtype parameter above + If `out=None`, returns a new array containing the mean values, + otherwise a reference to the output array is returned. Nan is + returned for slices that contain only NaNs. + + See Also + -------- + average : Weighted average + mean : Arithmetic mean taken while not ignoring NaNs + var, nanvar + + Notes + ----- + The arithmetic mean is the sum of the non-NaN elements along the axis + divided by the number of non-NaN elements. + + Note that for floating-point input, the mean is computed using the same + precision the input has. Depending on the input data, this can cause + the results to be inaccurate, especially for `float32`. Specifying a + higher-precision accumulator using the `dtype` keyword can alleviate + this issue. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, np.nan], [3, 4]]) + >>> np.nanmean(a) + 2.6666666666666665 + >>> np.nanmean(a, axis=0) + array([2., 4.]) + >>> np.nanmean(a, axis=1) + array([1., 3.5]) # may vary + + """ + arr, mask = _replace_nan(a, 0) + if mask is None: + return np.mean(arr, axis=axis, dtype=dtype, out=out, keepdims=keepdims, + where=where) + + if dtype is not None: + dtype = np.dtype(dtype) + if dtype is not None and not issubclass(dtype.type, np.inexact): + raise TypeError("If a is inexact, then dtype must be inexact") + if out is not None and not issubclass(out.dtype.type, np.inexact): + raise TypeError("If a is inexact, then out must be inexact") + + cnt = np.sum(~mask, axis=axis, dtype=np.intp, keepdims=keepdims, + where=where) + tot = np.sum(arr, axis=axis, dtype=dtype, out=out, keepdims=keepdims, + where=where) + avg = _divide_by_count(tot, cnt, out=out) + + isbad = (cnt == 0) + if isbad.any(): + warnings.warn("Mean of empty slice", RuntimeWarning, stacklevel=2) + # NaN is the only possible bad value, so no further + # action is needed to handle bad results. + return avg + + +def _nanmedian1d(arr1d, overwrite_input=False): + """ + Private function for rank 1 arrays. Compute the median ignoring NaNs. + See nanmedian for parameter usage + """ + arr1d_parsed, _, overwrite_input = _remove_nan_1d( + arr1d, overwrite_input=overwrite_input, + ) + + if arr1d_parsed.size == 0: + # Ensure that a nan-esque scalar of the appropriate type (and unit) + # is returned for `timedelta64` and `complexfloating` + return arr1d[-1] + + return np.median(arr1d_parsed, overwrite_input=overwrite_input) + + +def _nanmedian(a, axis=None, out=None, overwrite_input=False): + """ + Private function that doesn't support extended axis or keepdims. + These methods are extended to this function using _ureduce + See nanmedian for parameter usage + + """ + if axis is None or a.ndim == 1: + part = a.ravel() + if out is None: + return _nanmedian1d(part, overwrite_input) + else: + out[...] = _nanmedian1d(part, overwrite_input) + return out + else: + # for small medians use sort + indexing which is still faster than + # apply_along_axis + # benchmarked with shuffled (50, 50, x) containing a few NaN + if a.shape[axis] < 600: + return _nanmedian_small(a, axis, out, overwrite_input) + result = np.apply_along_axis(_nanmedian1d, axis, a, overwrite_input) + if out is not None: + out[...] = result + return result + + +def _nanmedian_small(a, axis=None, out=None, overwrite_input=False): + """ + sort + indexing median, faster for small medians along multiple + dimensions due to the high overhead of apply_along_axis + + see nanmedian for parameter usage + """ + a = np.ma.masked_array(a, np.isnan(a)) + m = np.ma.median(a, axis=axis, overwrite_input=overwrite_input) + for i in range(np.count_nonzero(m.mask.ravel())): + warnings.warn("All-NaN slice encountered", RuntimeWarning, + stacklevel=5) + + fill_value = np.timedelta64("NaT") if m.dtype.kind == "m" else np.nan + if out is not None: + out[...] = m.filled(fill_value) + return out + return m.filled(fill_value) + + +def _nanmedian_dispatcher( + a, axis=None, out=None, overwrite_input=None, keepdims=None): + return (a, out) + + +@array_function_dispatch(_nanmedian_dispatcher) +def nanmedian(a, axis=None, out=None, overwrite_input=False, keepdims=np._NoValue): + """ + Compute the median along the specified axis, while ignoring NaNs. + + Returns the median of the array elements. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array. + axis : {int, sequence of int, None}, optional + Axis or axes along which the medians are computed. The default + is to compute the median along a flattened version of the array. + A sequence of axes is supported since version 1.9.0. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output, + but the type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow use of memory of input array `a` for + calculations. The input array will be modified by the call to + `median`. This will save memory when you do not need to preserve + the contents of the input array. Treat the input as undefined, + but it will probably be fully or partially sorted. Default is + False. If `overwrite_input` is ``True`` and `a` is not already an + `ndarray`, an error will be raised. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + + If this is anything but the default value it will be passed + through (in the special case of an empty array) to the + `mean` function of the underlying array. If the array is + a sub-class and `mean` does not have the kwarg `keepdims` this + will raise a RuntimeError. + + Returns + ------- + median : ndarray + A new array holding the result. If the input contains integers + or floats smaller than ``float64``, then the output data-type is + ``np.float64``. Otherwise, the data-type of the output is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + mean, median, percentile + + Notes + ----- + Given a vector ``V`` of length ``N``, the median of ``V`` is the + middle value of a sorted copy of ``V``, ``V_sorted`` - i.e., + ``V_sorted[(N-1)/2]``, when ``N`` is odd and the average of the two + middle values of ``V_sorted`` when ``N`` is even. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10.0, 7, 4], [3, 2, 1]]) + >>> a[0, 1] = np.nan + >>> a + array([[10., nan, 4.], + [ 3., 2., 1.]]) + >>> np.median(a) + np.float64(nan) + >>> np.nanmedian(a) + 3.0 + >>> np.nanmedian(a, axis=0) + array([6.5, 2. , 2.5]) + >>> np.median(a, axis=1) + array([nan, 2.]) + >>> b = a.copy() + >>> np.nanmedian(b, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a==b) + >>> b = a.copy() + >>> np.nanmedian(b, axis=None, overwrite_input=True) + 3.0 + >>> assert not np.all(a==b) + + """ + a = np.asanyarray(a) + # apply_along_axis in _nanmedian doesn't handle empty arrays well, + # so deal them upfront + if a.size == 0: + return np.nanmean(a, axis, out=out, keepdims=keepdims) + + return fnb._ureduce(a, func=_nanmedian, keepdims=keepdims, + axis=axis, out=out, + overwrite_input=overwrite_input) + + +def _nanpercentile_dispatcher( + a, q, axis=None, out=None, overwrite_input=None, + method=None, keepdims=None, *, weights=None, interpolation=None): + return (a, q, out, weights) + + +@array_function_dispatch(_nanpercentile_dispatcher) +def nanpercentile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=np._NoValue, + *, + weights=None, + interpolation=None, +): + """ + Compute the qth percentile of the data along the specified axis, + while ignoring nan values. + + Returns the qth percentile(s) of the array elements. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array, containing + nan values to be ignored. + q : array_like of float + Percentile or sequence of percentiles to compute, which must be + between 0 and 100 inclusive. + axis : {int, tuple of int, None}, optional + Axis or axes along which the percentiles are computed. The default + is to compute the percentile(s) along a flattened version of the + array. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape and buffer length as the expected output, but the + type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow the input array `a` to be modified by + intermediate calculations, to save memory. In this case, the + contents of the input `a` after this function completes is + undefined. + method : str, optional + This parameter specifies the method to use for estimating the + percentile. There are many different methods, some unique to NumPy. + See the notes for explanation. The options sorted by their R type + as summarized in the H&F paper [1]_ are: + + 1. 'inverted_cdf' + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + The first three methods are discontinuous. NumPy further defines the + following discontinuous variations of the default 'linear' (7.) option: + + * 'lower' + * 'higher', + * 'midpoint' + * 'nearest' + + .. versionchanged:: 1.22.0 + This argument was previously called "interpolation" and only + offered the "linear" default and last four options. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in + the result as dimensions with size one. With this option, the + result will broadcast correctly against the original array `a`. + + If this is anything but the default value it will be passed + through (in the special case of an empty array) to the + `mean` function of the underlying array. If the array is + a sub-class and `mean` does not have the kwarg `keepdims` this + will raise a RuntimeError. + + weights : array_like, optional + An array of weights associated with the values in `a`. Each value in + `a` contributes to the percentile according to its associated weight. + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given axis) or of the same shape as `a`. + If `weights=None`, then all data in `a` are assumed to have a + weight equal to one. + Only `method="inverted_cdf"` supports weights. + + .. versionadded:: 2.0.0 + + interpolation : str, optional + Deprecated name for the method keyword argument. + + .. deprecated:: 1.22.0 + + Returns + ------- + percentile : scalar or ndarray + If `q` is a single percentile and `axis=None`, then the result + is a scalar. If multiple percentiles are given, first axis of + the result corresponds to the percentiles. The other axes are + the axes that remain after the reduction of `a`. If the input + contains integers or floats smaller than ``float64``, the output + data-type is ``float64``. Otherwise, the output data-type is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + nanmean + nanmedian : equivalent to ``nanpercentile(..., 50)`` + percentile, median, mean + nanquantile : equivalent to nanpercentile, except q in range [0, 1]. + + Notes + ----- + The behavior of `numpy.nanpercentile` with percentage `q` is that of + `numpy.quantile` with argument ``q/100`` (ignoring nan values). + For more information, please see `numpy.quantile`. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10., 7., 4.], [3., 2., 1.]]) + >>> a[0][1] = np.nan + >>> a + array([[10., nan, 4.], + [ 3., 2., 1.]]) + >>> np.percentile(a, 50) + np.float64(nan) + >>> np.nanpercentile(a, 50) + 3.0 + >>> np.nanpercentile(a, 50, axis=0) + array([6.5, 2. , 2.5]) + >>> np.nanpercentile(a, 50, axis=1, keepdims=True) + array([[7.], + [2.]]) + >>> m = np.nanpercentile(a, 50, axis=0) + >>> out = np.zeros_like(m) + >>> np.nanpercentile(a, 50, axis=0, out=out) + array([6.5, 2. , 2.5]) + >>> m + array([6.5, 2. , 2.5]) + + >>> b = a.copy() + >>> np.nanpercentile(b, 50, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a==b) + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + + """ + if interpolation is not None: + method = fnb._check_interpolation_as_method( + method, interpolation, "nanpercentile") + + a = np.asanyarray(a) + if a.dtype.kind == "c": + raise TypeError("a must be an array of real numbers") + + q = np.true_divide(q, a.dtype.type(100) if a.dtype.kind == "f" else 100) + # undo any decay that the ufunc performed (see gh-13105) + q = np.asanyarray(q) + if not fnb._quantile_is_valid(q): + raise ValueError("Percentiles must be in the range [0, 100]") + + if weights is not None: + if method != "inverted_cdf": + msg = ("Only method 'inverted_cdf' supports weights. " + f"Got: {method}.") + raise ValueError(msg) + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, a.ndim, argname="axis") + weights = _weights_are_valid(weights=weights, a=a, axis=axis) + if np.any(weights < 0): + raise ValueError("Weights must be non-negative.") + + return _nanquantile_unchecked( + a, q, axis, out, overwrite_input, method, keepdims, weights) + + +def _nanquantile_dispatcher(a, q, axis=None, out=None, overwrite_input=None, + method=None, keepdims=None, *, weights=None, + interpolation=None): + return (a, q, out, weights) + + +@array_function_dispatch(_nanquantile_dispatcher) +def nanquantile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=np._NoValue, + *, + weights=None, + interpolation=None, +): + """ + Compute the qth quantile of the data along the specified axis, + while ignoring nan values. + Returns the qth quantile(s) of the array elements. + + Parameters + ---------- + a : array_like + Input array or object that can be converted to an array, containing + nan values to be ignored + q : array_like of float + Probability or sequence of probabilities for the quantiles to compute. + Values must be between 0 and 1 inclusive. + axis : {int, tuple of int, None}, optional + Axis or axes along which the quantiles are computed. The + default is to compute the quantile(s) along a flattened + version of the array. + out : ndarray, optional + Alternative output array in which to place the result. It must + have the same shape and buffer length as the expected output, + but the type (of the output) will be cast if necessary. + overwrite_input : bool, optional + If True, then allow the input array `a` to be modified by intermediate + calculations, to save memory. In this case, the contents of the input + `a` after this function completes is undefined. + method : str, optional + This parameter specifies the method to use for estimating the + quantile. There are many different methods, some unique to NumPy. + See the notes for explanation. The options sorted by their R type + as summarized in the H&F paper [1]_ are: + + 1. 'inverted_cdf' + 2. 'averaged_inverted_cdf' + 3. 'closest_observation' + 4. 'interpolated_inverted_cdf' + 5. 'hazen' + 6. 'weibull' + 7. 'linear' (default) + 8. 'median_unbiased' + 9. 'normal_unbiased' + + The first three methods are discontinuous. NumPy further defines the + following discontinuous variations of the default 'linear' (7.) option: + + * 'lower' + * 'higher', + * 'midpoint' + * 'nearest' + + .. versionchanged:: 1.22.0 + This argument was previously called "interpolation" and only + offered the "linear" default and last four options. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left in + the result as dimensions with size one. With this option, the + result will broadcast correctly against the original array `a`. + + If this is anything but the default value it will be passed + through (in the special case of an empty array) to the + `mean` function of the underlying array. If the array is + a sub-class and `mean` does not have the kwarg `keepdims` this + will raise a RuntimeError. + + weights : array_like, optional + An array of weights associated with the values in `a`. Each value in + `a` contributes to the quantile according to its associated weight. + The weights array can either be 1-D (in which case its length must be + the size of `a` along the given axis) or of the same shape as `a`. + If `weights=None`, then all data in `a` are assumed to have a + weight equal to one. + Only `method="inverted_cdf"` supports weights. + + .. versionadded:: 2.0.0 + + interpolation : str, optional + Deprecated name for the method keyword argument. + + .. deprecated:: 1.22.0 + + Returns + ------- + quantile : scalar or ndarray + If `q` is a single probability and `axis=None`, then the result + is a scalar. If multiple probability levels are given, first axis of + the result corresponds to the quantiles. The other axes are + the axes that remain after the reduction of `a`. If the input + contains integers or floats smaller than ``float64``, the output + data-type is ``float64``. Otherwise, the output data-type is the + same as that of the input. If `out` is specified, that array is + returned instead. + + See Also + -------- + quantile + nanmean, nanmedian + nanmedian : equivalent to ``nanquantile(..., 0.5)`` + nanpercentile : same as nanquantile, but with q in the range [0, 100]. + + Notes + ----- + The behavior of `numpy.nanquantile` is the same as that of + `numpy.quantile` (ignoring nan values). + For more information, please see `numpy.quantile`. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[10., 7., 4.], [3., 2., 1.]]) + >>> a[0][1] = np.nan + >>> a + array([[10., nan, 4.], + [ 3., 2., 1.]]) + >>> np.quantile(a, 0.5) + np.float64(nan) + >>> np.nanquantile(a, 0.5) + 3.0 + >>> np.nanquantile(a, 0.5, axis=0) + array([6.5, 2. , 2.5]) + >>> np.nanquantile(a, 0.5, axis=1, keepdims=True) + array([[7.], + [2.]]) + >>> m = np.nanquantile(a, 0.5, axis=0) + >>> out = np.zeros_like(m) + >>> np.nanquantile(a, 0.5, axis=0, out=out) + array([6.5, 2. , 2.5]) + >>> m + array([6.5, 2. , 2.5]) + >>> b = a.copy() + >>> np.nanquantile(b, 0.5, axis=1, overwrite_input=True) + array([7., 2.]) + >>> assert not np.all(a==b) + + References + ---------- + .. [1] R. J. Hyndman and Y. Fan, + "Sample quantiles in statistical packages," + The American Statistician, 50(4), pp. 361-365, 1996 + + """ + + if interpolation is not None: + method = fnb._check_interpolation_as_method( + method, interpolation, "nanquantile") + + a = np.asanyarray(a) + if a.dtype.kind == "c": + raise TypeError("a must be an array of real numbers") + + # Use dtype of array if possible (e.g., if q is a python int or float). + if isinstance(q, (int, float)) and a.dtype.kind == "f": + q = np.asanyarray(q, dtype=a.dtype) + else: + q = np.asanyarray(q) + + if not fnb._quantile_is_valid(q): + raise ValueError("Quantiles must be in the range [0, 1]") + + if weights is not None: + if method != "inverted_cdf": + msg = ("Only method 'inverted_cdf' supports weights. " + f"Got: {method}.") + raise ValueError(msg) + if axis is not None: + axis = _nx.normalize_axis_tuple(axis, a.ndim, argname="axis") + weights = _weights_are_valid(weights=weights, a=a, axis=axis) + if np.any(weights < 0): + raise ValueError("Weights must be non-negative.") + + return _nanquantile_unchecked( + a, q, axis, out, overwrite_input, method, keepdims, weights) + + +def _nanquantile_unchecked( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=np._NoValue, + weights=None, +): + """Assumes that q is in [0, 1], and is an ndarray""" + # apply_along_axis in _nanpercentile doesn't handle empty arrays well, + # so deal them upfront + if a.size == 0: + return np.nanmean(a, axis, out=out, keepdims=keepdims) + return fnb._ureduce(a, + func=_nanquantile_ureduce_func, + q=q, + weights=weights, + keepdims=keepdims, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method) + + +def _nanquantile_ureduce_func( + a: np.array, + q: np.array, + weights: np.array, + axis: int | None = None, + out=None, + overwrite_input: bool = False, + method="linear", +): + """ + Private function that doesn't support extended axis or keepdims. + These methods are extended to this function using _ureduce + See nanpercentile for parameter usage + """ + if axis is None or a.ndim == 1: + part = a.ravel() + wgt = None if weights is None else weights.ravel() + result = _nanquantile_1d(part, q, overwrite_input, method, weights=wgt) + else: + # Note that this code could try to fill in `out` right away + if weights is None: + result = np.apply_along_axis(_nanquantile_1d, axis, a, q, + overwrite_input, method, weights) + # apply_along_axis fills in collapsed axis with results. + # Move those axes to the beginning to match percentile's + # convention. + if q.ndim != 0: + from_ax = [axis + i for i in range(q.ndim)] + result = np.moveaxis(result, from_ax, list(range(q.ndim))) + else: + # We need to apply along axis over 2 arrays, a and weights. + # move operation axes to end for simplicity: + a = np.moveaxis(a, axis, -1) + if weights is not None: + weights = np.moveaxis(weights, axis, -1) + if out is not None: + result = out + else: + # weights are limited to `inverted_cdf` so the result dtype + # is known to be identical to that of `a` here: + result = np.empty_like(a, shape=q.shape + a.shape[:-1]) + + for ii in np.ndindex(a.shape[:-1]): + result[(...,) + ii] = _nanquantile_1d( + a[ii], q, weights=weights[ii], + overwrite_input=overwrite_input, method=method, + ) + # This path dealt with `out` already... + return result + + if out is not None: + out[...] = result + return result + + +def _nanquantile_1d( + arr1d, q, overwrite_input=False, method="linear", weights=None, +): + """ + Private function for rank 1 arrays. Compute quantile ignoring NaNs. + See nanpercentile for parameter usage + """ + # TODO: What to do when arr1d = [1, np.nan] and weights = [0, 1]? + arr1d, weights, overwrite_input = _remove_nan_1d(arr1d, + second_arr1d=weights, overwrite_input=overwrite_input) + if arr1d.size == 0: + # convert to scalar + return np.full(q.shape, np.nan, dtype=arr1d.dtype)[()] + + return fnb._quantile_unchecked( + arr1d, + q, + overwrite_input=overwrite_input, + method=method, + weights=weights, + ) + + +def _nanvar_dispatcher(a, axis=None, dtype=None, out=None, ddof=None, + keepdims=None, *, where=None, mean=None, + correction=None): + return (a, out) + + +@array_function_dispatch(_nanvar_dispatcher) +def nanvar(a, axis=None, dtype=None, out=None, ddof=0, keepdims=np._NoValue, + *, where=np._NoValue, mean=np._NoValue, correction=np._NoValue): + """ + Compute the variance along the specified axis, while ignoring NaNs. + + Returns the variance of the array elements, a measure of the spread of + a distribution. The variance is computed for the flattened array by + default, otherwise over the specified axis. + + For all-NaN slices or slices with zero degrees of freedom, NaN is + returned and a `RuntimeWarning` is raised. + + Parameters + ---------- + a : array_like + Array containing numbers whose variance is desired. If `a` is not an + array, a conversion is attempted. + axis : {int, tuple of int, None}, optional + Axis or axes along which the variance is computed. The default is to compute + the variance of the flattened array. + dtype : data-type, optional + Type to use in computing the variance. For arrays of integer type + the default is `float64`; for arrays of float types it is the same as + the array type. + out : ndarray, optional + Alternate output array in which to place the result. It must have + the same shape as the expected output, but the type is cast if + necessary. + ddof : {int, float}, optional + "Delta Degrees of Freedom": the divisor used in the calculation is + ``N - ddof``, where ``N`` represents the number of non-NaN + elements. By default `ddof` is zero. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + where : array_like of bool, optional + Elements to include in the variance. See `~numpy.ufunc.reduce` for + details. + + .. versionadded:: 1.22.0 + + mean : array_like, optional + Provide the mean to prevent its recalculation. The mean should have + a shape as if it was calculated with ``keepdims=True``. + The axis for the calculation of the mean should be the same as used in + the call to this var function. + + .. versionadded:: 2.0.0 + + correction : {int, float}, optional + Array API compatible name for the ``ddof`` parameter. Only one of them + can be provided at the same time. + + .. versionadded:: 2.0.0 + + Returns + ------- + variance : ndarray, see dtype parameter above + If `out` is None, return a new array containing the variance, + otherwise return a reference to the output array. If ddof is >= the + number of non-NaN elements in a slice or the slice contains only + NaNs, then the result for that slice is NaN. + + See Also + -------- + std : Standard deviation + mean : Average + var : Variance while not ignoring NaNs + nanstd, nanmean + :ref:`ufuncs-output-type` + + Notes + ----- + The variance is the average of the squared deviations from the mean, + i.e., ``var = mean(abs(x - x.mean())**2)``. + + The mean is normally calculated as ``x.sum() / N``, where ``N = len(x)``. + If, however, `ddof` is specified, the divisor ``N - ddof`` is used + instead. In standard statistical practice, ``ddof=1`` provides an + unbiased estimator of the variance of a hypothetical infinite + population. ``ddof=0`` provides a maximum likelihood estimate of the + variance for normally distributed variables. + + Note that for complex numbers, the absolute value is taken before + squaring, so that the result is always real and nonnegative. + + For floating-point input, the variance is computed using the same + precision the input has. Depending on the input data, this can cause + the results to be inaccurate, especially for `float32` (see example + below). Specifying a higher-accuracy accumulator using the ``dtype`` + keyword can alleviate this issue. + + For this function to work on sub-classes of ndarray, they must define + `sum` with the kwarg `keepdims` + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, np.nan], [3, 4]]) + >>> np.nanvar(a) + 1.5555555555555554 + >>> np.nanvar(a, axis=0) + array([1., 0.]) + >>> np.nanvar(a, axis=1) + array([0., 0.25]) # may vary + + """ + arr, mask = _replace_nan(a, 0) + if mask is None: + return np.var(arr, axis=axis, dtype=dtype, out=out, ddof=ddof, + keepdims=keepdims, where=where, mean=mean, + correction=correction) + + if dtype is not None: + dtype = np.dtype(dtype) + if dtype is not None and not issubclass(dtype.type, np.inexact): + raise TypeError("If a is inexact, then dtype must be inexact") + if out is not None and not issubclass(out.dtype.type, np.inexact): + raise TypeError("If a is inexact, then out must be inexact") + + if correction != np._NoValue: + if ddof != 0: + raise ValueError( + "ddof and correction can't be provided simultaneously." + ) + else: + ddof = correction + + # Compute mean + if type(arr) is np.matrix: + _keepdims = np._NoValue + else: + _keepdims = True + + cnt = np.sum(~mask, axis=axis, dtype=np.intp, keepdims=_keepdims, + where=where) + + if mean is not np._NoValue: + avg = mean + else: + # we need to special case matrix for reverse compatibility + # in order for this to work, these sums need to be called with + # keepdims=True, however matrix now raises an error in this case, but + # the reason that it drops the keepdims kwarg is to force keepdims=True + # so this used to work by serendipity. + avg = np.sum(arr, axis=axis, dtype=dtype, + keepdims=_keepdims, where=where) + avg = _divide_by_count(avg, cnt) + + # Compute squared deviation from mean. + np.subtract(arr, avg, out=arr, casting='unsafe', where=where) + arr = _copyto(arr, 0, mask) + if issubclass(arr.dtype.type, np.complexfloating): + sqr = np.multiply(arr, arr.conj(), out=arr, where=where).real + else: + sqr = np.multiply(arr, arr, out=arr, where=where) + + # Compute variance. + var = np.sum(sqr, axis=axis, dtype=dtype, out=out, keepdims=keepdims, + where=where) + + # Precaution against reduced object arrays + try: + var_ndim = var.ndim + except AttributeError: + var_ndim = np.ndim(var) + if var_ndim < cnt.ndim: + # Subclasses of ndarray may ignore keepdims, so check here. + cnt = cnt.squeeze(axis) + dof = cnt - ddof + var = _divide_by_count(var, dof) + + isbad = (dof <= 0) + if np.any(isbad): + warnings.warn("Degrees of freedom <= 0 for slice.", RuntimeWarning, + stacklevel=2) + # NaN, inf, or negative numbers are all possible bad + # values, so explicitly replace them with NaN. + var = _copyto(var, np.nan, isbad) + return var + + +def _nanstd_dispatcher(a, axis=None, dtype=None, out=None, ddof=None, + keepdims=None, *, where=None, mean=None, + correction=None): + return (a, out) + + +@array_function_dispatch(_nanstd_dispatcher) +def nanstd(a, axis=None, dtype=None, out=None, ddof=0, keepdims=np._NoValue, + *, where=np._NoValue, mean=np._NoValue, correction=np._NoValue): + """ + Compute the standard deviation along the specified axis, while + ignoring NaNs. + + Returns the standard deviation, a measure of the spread of a + distribution, of the non-NaN array elements. The standard deviation is + computed for the flattened array by default, otherwise over the + specified axis. + + For all-NaN slices or slices with zero degrees of freedom, NaN is + returned and a `RuntimeWarning` is raised. + + Parameters + ---------- + a : array_like + Calculate the standard deviation of the non-NaN values. + axis : {int, tuple of int, None}, optional + Axis or axes along which the standard deviation is computed. The default is + to compute the standard deviation of the flattened array. + dtype : dtype, optional + Type to use in computing the standard deviation. For arrays of + integer type the default is float64, for arrays of float types it + is the same as the array type. + out : ndarray, optional + Alternative output array in which to place the result. It must have + the same shape as the expected output but the type (of the + calculated values) will be cast if necessary. + ddof : {int, float}, optional + Means Delta Degrees of Freedom. The divisor used in calculations + is ``N - ddof``, where ``N`` represents the number of non-NaN + elements. By default `ddof` is zero. + + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + + If this value is anything but the default it is passed through + as-is to the relevant functions of the sub-classes. If these + functions do not have a `keepdims` kwarg, a RuntimeError will + be raised. + where : array_like of bool, optional + Elements to include in the standard deviation. + See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.22.0 + + mean : array_like, optional + Provide the mean to prevent its recalculation. The mean should have + a shape as if it was calculated with ``keepdims=True``. + The axis for the calculation of the mean should be the same as used in + the call to this std function. + + .. versionadded:: 2.0.0 + + correction : {int, float}, optional + Array API compatible name for the ``ddof`` parameter. Only one of them + can be provided at the same time. + + .. versionadded:: 2.0.0 + + Returns + ------- + standard_deviation : ndarray, see dtype parameter above. + If `out` is None, return a new array containing the standard + deviation, otherwise return a reference to the output array. If + ddof is >= the number of non-NaN elements in a slice or the slice + contains only NaNs, then the result for that slice is NaN. + + See Also + -------- + var, mean, std + nanvar, nanmean + :ref:`ufuncs-output-type` + + Notes + ----- + The standard deviation is the square root of the average of the squared + deviations from the mean: ``std = sqrt(mean(abs(x - x.mean())**2))``. + + The average squared deviation is normally calculated as + ``x.sum() / N``, where ``N = len(x)``. If, however, `ddof` is + specified, the divisor ``N - ddof`` is used instead. In standard + statistical practice, ``ddof=1`` provides an unbiased estimator of the + variance of the infinite population. ``ddof=0`` provides a maximum + likelihood estimate of the variance for normally distributed variables. + The standard deviation computed in this function is the square root of + the estimated variance, so even with ``ddof=1``, it will not be an + unbiased estimate of the standard deviation per se. + + Note that, for complex numbers, `std` takes the absolute value before + squaring, so that the result is always real and nonnegative. + + For floating-point input, the *std* is computed using the same + precision the input has. Depending on the input data, this can cause + the results to be inaccurate, especially for float32 (see example + below). Specifying a higher-accuracy accumulator using the `dtype` + keyword can alleviate this issue. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([[1, np.nan], [3, 4]]) + >>> np.nanstd(a) + 1.247219128924647 + >>> np.nanstd(a, axis=0) + array([1., 0.]) + >>> np.nanstd(a, axis=1) + array([0., 0.5]) # may vary + + """ + var = nanvar(a, axis=axis, dtype=dtype, out=out, ddof=ddof, + keepdims=keepdims, where=where, mean=mean, + correction=correction) + if isinstance(var, np.ndarray): + std = np.sqrt(var, out=var) + elif hasattr(var, 'dtype'): + std = var.dtype.type(np.sqrt(var)) + else: + std = np.sqrt(var) + return std diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..081b53d8ea44b53a5ee1470ab2725e6b8fc7274a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_nanfunctions_impl.pyi @@ -0,0 +1,53 @@ +from numpy._core.fromnumeric import ( + amin, + amax, + argmin, + argmax, + sum, + prod, + cumsum, + cumprod, + mean, + var, + std +) + +from numpy.lib._function_base_impl import ( + median, + percentile, + quantile, +) + +__all__ = [ + "nansum", + "nanmax", + "nanmin", + "nanargmax", + "nanargmin", + "nanmean", + "nanmedian", + "nanpercentile", + "nanvar", + "nanstd", + "nanprod", + "nancumsum", + "nancumprod", + "nanquantile", +] + +# NOTE: In reality these functions are not aliases but distinct functions +# with identical signatures. +nanmin = amin +nanmax = amax +nanargmin = argmin +nanargmax = argmax +nansum = sum +nanprod = prod +nancumsum = cumsum +nancumprod = cumprod +nanmean = mean +nanvar = var +nanstd = std +nanmedian = median +nanpercentile = percentile +nanquantile = quantile diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..4dc3a4b9b7e2e117b033b103c8ebaf165a1b160c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.py @@ -0,0 +1,2595 @@ +""" +IO related functions. +""" +import os +import re +import functools +import itertools +import warnings +import weakref +import contextlib +import operator +from operator import itemgetter +from collections.abc import Mapping +import pickle + +import numpy as np +from . import format +from ._datasource import DataSource +from numpy._core import overrides +from numpy._core.multiarray import packbits, unpackbits +from numpy._core._multiarray_umath import _load_from_filelike +from numpy._core.overrides import finalize_array_function_like, set_module +from ._iotools import ( + LineSplitter, NameValidator, StringConverter, ConverterError, + ConverterLockError, ConversionWarning, _is_string_like, + has_nested_fields, flatten_dtype, easy_dtype, _decode_line + ) +from numpy._utils import asunicode, asbytes + + +__all__ = [ + 'savetxt', 'loadtxt', 'genfromtxt', 'load', 'save', 'savez', + 'savez_compressed', 'packbits', 'unpackbits', 'fromregex' + ] + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +class BagObj: + """ + BagObj(obj) + + Convert attribute look-ups to getitems on the object passed in. + + Parameters + ---------- + obj : class instance + Object on which attribute look-up is performed. + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib._npyio_impl import BagObj as BO + >>> class BagDemo: + ... def __getitem__(self, key): # An instance of BagObj(BagDemo) + ... # will call this method when any + ... # attribute look-up is required + ... result = "Doesn't matter what you want, " + ... return result + "you're gonna get this" + ... + >>> demo_obj = BagDemo() + >>> bagobj = BO(demo_obj) + >>> bagobj.hello_there + "Doesn't matter what you want, you're gonna get this" + >>> bagobj.I_can_be_anything + "Doesn't matter what you want, you're gonna get this" + + """ + + def __init__(self, obj): + # Use weakref to make NpzFile objects collectable by refcount + self._obj = weakref.proxy(obj) + + def __getattribute__(self, key): + try: + return object.__getattribute__(self, '_obj')[key] + except KeyError: + raise AttributeError(key) from None + + def __dir__(self): + """ + Enables dir(bagobj) to list the files in an NpzFile. + + This also enables tab-completion in an interpreter or IPython. + """ + return list(object.__getattribute__(self, '_obj').keys()) + + +def zipfile_factory(file, *args, **kwargs): + """ + Create a ZipFile. + + Allows for Zip64, and the `file` argument can accept file, str, or + pathlib.Path objects. `args` and `kwargs` are passed to the zipfile.ZipFile + constructor. + """ + if not hasattr(file, 'read'): + file = os.fspath(file) + import zipfile + kwargs['allowZip64'] = True + return zipfile.ZipFile(file, *args, **kwargs) + + +@set_module('numpy.lib.npyio') +class NpzFile(Mapping): + """ + NpzFile(fid) + + A dictionary-like object with lazy-loading of files in the zipped + archive provided on construction. + + `NpzFile` is used to load files in the NumPy ``.npz`` data archive + format. It assumes that files in the archive have a ``.npy`` extension, + other files are ignored. + + The arrays and file strings are lazily loaded on either + getitem access using ``obj['key']`` or attribute lookup using + ``obj.f.key``. A list of all files (without ``.npy`` extensions) can + be obtained with ``obj.files`` and the ZipFile object itself using + ``obj.zip``. + + Attributes + ---------- + files : list of str + List of all files in the archive with a ``.npy`` extension. + zip : ZipFile instance + The ZipFile object initialized with the zipped archive. + f : BagObj instance + An object on which attribute can be performed as an alternative + to getitem access on the `NpzFile` instance itself. + allow_pickle : bool, optional + Allow loading pickled data. Default: False + pickle_kwargs : dict, optional + Additional keyword arguments to pass on to pickle.load. + These are only useful when loading object arrays saved on + Python 2 when using Python 3. + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + This option is ignored when `allow_pickle` is passed. In that case + the file is by definition trusted and the limit is unnecessary. + + Parameters + ---------- + fid : file, str, or pathlib.Path + The zipped archive to open. This is either a file-like object + or a string containing the path to the archive. + own_fid : bool, optional + Whether NpzFile should close the file handle. + Requires that `fid` is a file-like object. + + Examples + -------- + >>> import numpy as np + >>> from tempfile import TemporaryFile + >>> outfile = TemporaryFile() + >>> x = np.arange(10) + >>> y = np.sin(x) + >>> np.savez(outfile, x=x, y=y) + >>> _ = outfile.seek(0) + + >>> npz = np.load(outfile) + >>> isinstance(npz, np.lib.npyio.NpzFile) + True + >>> npz + NpzFile 'object' with keys: x, y + >>> sorted(npz.files) + ['x', 'y'] + >>> npz['x'] # getitem access + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + >>> npz.f.x # attribute lookup + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + """ + # Make __exit__ safe if zipfile_factory raises an exception + zip = None + fid = None + _MAX_REPR_ARRAY_COUNT = 5 + + def __init__(self, fid, own_fid=False, allow_pickle=False, + pickle_kwargs=None, *, + max_header_size=format._MAX_HEADER_SIZE): + # Import is postponed to here since zipfile depends on gzip, an + # optional component of the so-called standard library. + _zip = zipfile_factory(fid) + self._files = _zip.namelist() + self.files = [] + self.allow_pickle = allow_pickle + self.max_header_size = max_header_size + self.pickle_kwargs = pickle_kwargs + for x in self._files: + if x.endswith('.npy'): + self.files.append(x[:-4]) + else: + self.files.append(x) + self.zip = _zip + self.f = BagObj(self) + if own_fid: + self.fid = fid + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + def close(self): + """ + Close the file. + + """ + if self.zip is not None: + self.zip.close() + self.zip = None + if self.fid is not None: + self.fid.close() + self.fid = None + self.f = None # break reference cycle + + def __del__(self): + self.close() + + # Implement the Mapping ABC + def __iter__(self): + return iter(self.files) + + def __len__(self): + return len(self.files) + + def __getitem__(self, key): + # FIXME: This seems like it will copy strings around + # more than is strictly necessary. The zipfile + # will read the string and then + # the format.read_array will copy the string + # to another place in memory. + # It would be better if the zipfile could read + # (or at least uncompress) the data + # directly into the array memory. + member = False + if key in self._files: + member = True + elif key in self.files: + member = True + key += '.npy' + if member: + bytes = self.zip.open(key) + magic = bytes.read(len(format.MAGIC_PREFIX)) + bytes.close() + if magic == format.MAGIC_PREFIX: + bytes = self.zip.open(key) + return format.read_array(bytes, + allow_pickle=self.allow_pickle, + pickle_kwargs=self.pickle_kwargs, + max_header_size=self.max_header_size) + else: + return self.zip.read(key) + else: + raise KeyError(f"{key} is not a file in the archive") + + def __contains__(self, key): + return (key in self._files or key in self.files) + + def __repr__(self): + # Get filename or default to `object` + if isinstance(self.fid, str): + filename = self.fid + else: + filename = getattr(self.fid, "name", "object") + + # Get the name of arrays + array_names = ', '.join(self.files[:self._MAX_REPR_ARRAY_COUNT]) + if len(self.files) > self._MAX_REPR_ARRAY_COUNT: + array_names += "..." + return f"NpzFile {filename!r} with keys: {array_names}" + + # Work around problems with the docstrings in the Mapping methods + # They contain a `->`, which confuses the type annotation interpretations + # of sphinx-docs. See gh-25964 + + def get(self, key, default=None, /): + """ + D.get(k,[,d]) returns D[k] if k in D, else d. d defaults to None. + """ + return Mapping.get(self, key, default) + + def items(self): + """ + D.items() returns a set-like object providing a view on the items + """ + return Mapping.items(self) + + def keys(self): + """ + D.keys() returns a set-like object providing a view on the keys + """ + return Mapping.keys(self) + + def values(self): + """ + D.values() returns a set-like object providing a view on the values + """ + return Mapping.values(self) + + +@set_module('numpy') +def load(file, mmap_mode=None, allow_pickle=False, fix_imports=True, + encoding='ASCII', *, max_header_size=format._MAX_HEADER_SIZE): + """ + Load arrays or pickled objects from ``.npy``, ``.npz`` or pickled files. + + .. warning:: Loading files that contain object arrays uses the ``pickle`` + module, which is not secure against erroneous or maliciously + constructed data. Consider passing ``allow_pickle=False`` to + load data that is known not to contain object arrays for the + safer handling of untrusted sources. + + Parameters + ---------- + file : file-like object, string, or pathlib.Path + The file to read. File-like objects must support the + ``seek()`` and ``read()`` methods and must always + be opened in binary mode. Pickled files require that the + file-like object support the ``readline()`` method as well. + mmap_mode : {None, 'r+', 'r', 'w+', 'c'}, optional + If not None, then memory-map the file, using the given mode (see + `numpy.memmap` for a detailed description of the modes). A + memory-mapped array is kept on disk. However, it can be accessed + and sliced like any ndarray. Memory mapping is especially useful + for accessing small fragments of large files without reading the + entire file into memory. + allow_pickle : bool, optional + Allow loading pickled object arrays stored in npy files. Reasons for + disallowing pickles include security, as loading pickled data can + execute arbitrary code. If pickles are disallowed, loading object + arrays will fail. Default: False + fix_imports : bool, optional + Only useful when loading Python 2 generated pickled files on Python 3, + which includes npy/npz files containing object arrays. If `fix_imports` + is True, pickle will try to map the old Python 2 names to the new names + used in Python 3. + encoding : str, optional + What encoding to use when reading Python 2 strings. Only useful when + loading Python 2 generated pickled files in Python 3, which includes + npy/npz files containing object arrays. Values other than 'latin1', + 'ASCII', and 'bytes' are not allowed, as they can corrupt numerical + data. Default: 'ASCII' + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + This option is ignored when `allow_pickle` is passed. In that case + the file is by definition trusted and the limit is unnecessary. + + Returns + ------- + result : array, tuple, dict, etc. + Data stored in the file. For ``.npz`` files, the returned instance + of NpzFile class must be closed to avoid leaking file descriptors. + + Raises + ------ + OSError + If the input file does not exist or cannot be read. + UnpicklingError + If ``allow_pickle=True``, but the file cannot be loaded as a pickle. + ValueError + The file contains an object array, but ``allow_pickle=False`` given. + EOFError + When calling ``np.load`` multiple times on the same file handle, + if all data has already been read + + See Also + -------- + save, savez, savez_compressed, loadtxt + memmap : Create a memory-map to an array stored in a file on disk. + lib.format.open_memmap : Create or load a memory-mapped ``.npy`` file. + + Notes + ----- + - If the file contains pickle data, then whatever object is stored + in the pickle is returned. + - If the file is a ``.npy`` file, then a single array is returned. + - If the file is a ``.npz`` file, then a dictionary-like object is + returned, containing ``{filename: array}`` key-value pairs, one for + each file in the archive. + - If the file is a ``.npz`` file, the returned value supports the + context manager protocol in a similar fashion to the open function:: + + with load('foo.npz') as data: + a = data['a'] + + The underlying file descriptor is closed when exiting the 'with' + block. + + Examples + -------- + >>> import numpy as np + + Store data to disk, and load it again: + + >>> np.save('/tmp/123', np.array([[1, 2, 3], [4, 5, 6]])) + >>> np.load('/tmp/123.npy') + array([[1, 2, 3], + [4, 5, 6]]) + + Store compressed data to disk, and load it again: + + >>> a=np.array([[1, 2, 3], [4, 5, 6]]) + >>> b=np.array([1, 2]) + >>> np.savez('/tmp/123.npz', a=a, b=b) + >>> data = np.load('/tmp/123.npz') + >>> data['a'] + array([[1, 2, 3], + [4, 5, 6]]) + >>> data['b'] + array([1, 2]) + >>> data.close() + + Mem-map the stored array, and then access the second row + directly from disk: + + >>> X = np.load('/tmp/123.npy', mmap_mode='r') + >>> X[1, :] + memmap([4, 5, 6]) + + """ + if encoding not in ('ASCII', 'latin1', 'bytes'): + # The 'encoding' value for pickle also affects what encoding + # the serialized binary data of NumPy arrays is loaded + # in. Pickle does not pass on the encoding information to + # NumPy. The unpickling code in numpy._core.multiarray is + # written to assume that unicode data appearing where binary + # should be is in 'latin1'. 'bytes' is also safe, as is 'ASCII'. + # + # Other encoding values can corrupt binary data, and we + # purposefully disallow them. For the same reason, the errors= + # argument is not exposed, as values other than 'strict' + # result can similarly silently corrupt numerical data. + raise ValueError("encoding must be 'ASCII', 'latin1', or 'bytes'") + + pickle_kwargs = dict(encoding=encoding, fix_imports=fix_imports) + + with contextlib.ExitStack() as stack: + if hasattr(file, 'read'): + fid = file + own_fid = False + else: + fid = stack.enter_context(open(os.fspath(file), "rb")) + own_fid = True + + # Code to distinguish from NumPy binary files and pickles. + _ZIP_PREFIX = b'PK\x03\x04' + _ZIP_SUFFIX = b'PK\x05\x06' # empty zip files start with this + N = len(format.MAGIC_PREFIX) + magic = fid.read(N) + if not magic: + raise EOFError("No data left in file") + # If the file size is less than N, we need to make sure not + # to seek past the beginning of the file + fid.seek(-min(N, len(magic)), 1) # back-up + if magic.startswith((_ZIP_PREFIX, _ZIP_SUFFIX)): + # zip-file (assume .npz) + # Potentially transfer file ownership to NpzFile + stack.pop_all() + ret = NpzFile(fid, own_fid=own_fid, allow_pickle=allow_pickle, + pickle_kwargs=pickle_kwargs, + max_header_size=max_header_size) + return ret + elif magic == format.MAGIC_PREFIX: + # .npy file + if mmap_mode: + if allow_pickle: + max_header_size = 2**64 + return format.open_memmap(file, mode=mmap_mode, + max_header_size=max_header_size) + else: + return format.read_array(fid, allow_pickle=allow_pickle, + pickle_kwargs=pickle_kwargs, + max_header_size=max_header_size) + else: + # Try a pickle + if not allow_pickle: + raise ValueError( + "This file contains pickled (object) data. If you trust " + "the file you can load it unsafely using the " + "`allow_pickle=` keyword argument or `pickle.load()`.") + try: + return pickle.load(fid, **pickle_kwargs) + except Exception as e: + raise pickle.UnpicklingError( + f"Failed to interpret file {file!r} as a pickle") from e + + +def _save_dispatcher(file, arr, allow_pickle=None, fix_imports=None): + return (arr,) + + +@array_function_dispatch(_save_dispatcher) +def save(file, arr, allow_pickle=True, fix_imports=np._NoValue): + """ + Save an array to a binary file in NumPy ``.npy`` format. + + Parameters + ---------- + file : file, str, or pathlib.Path + File or filename to which the data is saved. If file is a file-object, + then the filename is unchanged. If file is a string or Path, + a ``.npy`` extension will be appended to the filename if it does not + already have one. + arr : array_like + Array data to be saved. + allow_pickle : bool, optional + Allow saving object arrays using Python pickles. Reasons for + disallowing pickles include security (loading pickled data can execute + arbitrary code) and portability (pickled objects may not be loadable + on different Python installations, for example if the stored objects + require libraries that are not available, and not all pickled data is + compatible between different versions of Python). + Default: True + fix_imports : bool, optional + The `fix_imports` flag is deprecated and has no effect. + + .. deprecated:: 2.1 + This flag is ignored since NumPy 1.17 and was only needed to + support loading some files in Python 2 written in Python 3. + + See Also + -------- + savez : Save several arrays into a ``.npz`` archive + savetxt, load + + Notes + ----- + For a description of the ``.npy`` format, see :py:mod:`numpy.lib.format`. + + Any data saved to the file is appended to the end of the file. + + Examples + -------- + >>> import numpy as np + + >>> from tempfile import TemporaryFile + >>> outfile = TemporaryFile() + + >>> x = np.arange(10) + >>> np.save(outfile, x) + + >>> _ = outfile.seek(0) # Only needed to simulate closing & reopening file + >>> np.load(outfile) + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + + >>> with open('test.npy', 'wb') as f: + ... np.save(f, np.array([1, 2])) + ... np.save(f, np.array([1, 3])) + >>> with open('test.npy', 'rb') as f: + ... a = np.load(f) + ... b = np.load(f) + >>> print(a, b) + # [1 2] [1 3] + """ + if fix_imports is not np._NoValue: + # Deprecated 2024-05-16, NumPy 2.1 + warnings.warn( + "The 'fix_imports' flag is deprecated and has no effect. " + "(Deprecated in NumPy 2.1)", + DeprecationWarning, stacklevel=2) + if hasattr(file, 'write'): + file_ctx = contextlib.nullcontext(file) + else: + file = os.fspath(file) + if not file.endswith('.npy'): + file = file + '.npy' + file_ctx = open(file, "wb") + + with file_ctx as fid: + arr = np.asanyarray(arr) + format.write_array(fid, arr, allow_pickle=allow_pickle, + pickle_kwargs=dict(fix_imports=fix_imports)) + + +def _savez_dispatcher(file, *args, allow_pickle=True, **kwds): + yield from args + yield from kwds.values() + + +@array_function_dispatch(_savez_dispatcher) +def savez(file, *args, allow_pickle=True, **kwds): + """Save several arrays into a single file in uncompressed ``.npz`` format. + + Provide arrays as keyword arguments to store them under the + corresponding name in the output file: ``savez(fn, x=x, y=y)``. + + If arrays are specified as positional arguments, i.e., ``savez(fn, + x, y)``, their names will be `arr_0`, `arr_1`, etc. + + Parameters + ---------- + file : file, str, or pathlib.Path + Either the filename (string) or an open file (file-like object) + where the data will be saved. If file is a string or a Path, the + ``.npz`` extension will be appended to the filename if it is not + already there. + args : Arguments, optional + Arrays to save to the file. Please use keyword arguments (see + `kwds` below) to assign names to arrays. Arrays specified as + args will be named "arr_0", "arr_1", and so on. + allow_pickle : bool, optional + Allow saving object arrays using Python pickles. Reasons for + disallowing pickles include security (loading pickled data can execute + arbitrary code) and portability (pickled objects may not be loadable + on different Python installations, for example if the stored objects + require libraries that are not available, and not all pickled data is + compatible between different versions of Python). + Default: True + kwds : Keyword arguments, optional + Arrays to save to the file. Each array will be saved to the + output file with its corresponding keyword name. + + Returns + ------- + None + + See Also + -------- + save : Save a single array to a binary file in NumPy format. + savetxt : Save an array to a file as plain text. + savez_compressed : Save several arrays into a compressed ``.npz`` archive + + Notes + ----- + The ``.npz`` file format is a zipped archive of files named after the + variables they contain. The archive is not compressed and each file + in the archive contains one variable in ``.npy`` format. For a + description of the ``.npy`` format, see :py:mod:`numpy.lib.format`. + + When opening the saved ``.npz`` file with `load` a `~lib.npyio.NpzFile` + object is returned. This is a dictionary-like object which can be queried + for its list of arrays (with the ``.files`` attribute), and for the arrays + themselves. + + Keys passed in `kwds` are used as filenames inside the ZIP archive. + Therefore, keys should be valid filenames; e.g., avoid keys that begin with + ``/`` or contain ``.``. + + When naming variables with keyword arguments, it is not possible to name a + variable ``file``, as this would cause the ``file`` argument to be defined + twice in the call to ``savez``. + + Examples + -------- + >>> import numpy as np + >>> from tempfile import TemporaryFile + >>> outfile = TemporaryFile() + >>> x = np.arange(10) + >>> y = np.sin(x) + + Using `savez` with \\*args, the arrays are saved with default names. + + >>> np.savez(outfile, x, y) + >>> _ = outfile.seek(0) # Only needed to simulate closing & reopening file + >>> npzfile = np.load(outfile) + >>> npzfile.files + ['arr_0', 'arr_1'] + >>> npzfile['arr_0'] + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + Using `savez` with \\**kwds, the arrays are saved with the keyword names. + + >>> outfile = TemporaryFile() + >>> np.savez(outfile, x=x, y=y) + >>> _ = outfile.seek(0) + >>> npzfile = np.load(outfile) + >>> sorted(npzfile.files) + ['x', 'y'] + >>> npzfile['x'] + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + """ + _savez(file, args, kwds, False, allow_pickle=allow_pickle) + + +def _savez_compressed_dispatcher(file, *args, allow_pickle=True, **kwds): + yield from args + yield from kwds.values() + + +@array_function_dispatch(_savez_compressed_dispatcher) +def savez_compressed(file, *args, allow_pickle=True, **kwds): + """ + Save several arrays into a single file in compressed ``.npz`` format. + + Provide arrays as keyword arguments to store them under the + corresponding name in the output file: ``savez_compressed(fn, x=x, y=y)``. + + If arrays are specified as positional arguments, i.e., + ``savez_compressed(fn, x, y)``, their names will be `arr_0`, `arr_1`, etc. + + Parameters + ---------- + file : file, str, or pathlib.Path + Either the filename (string) or an open file (file-like object) + where the data will be saved. If file is a string or a Path, the + ``.npz`` extension will be appended to the filename if it is not + already there. + args : Arguments, optional + Arrays to save to the file. Please use keyword arguments (see + `kwds` below) to assign names to arrays. Arrays specified as + args will be named "arr_0", "arr_1", and so on. + allow_pickle : bool, optional + Allow saving object arrays using Python pickles. Reasons for + disallowing pickles include security (loading pickled data can execute + arbitrary code) and portability (pickled objects may not be loadable + on different Python installations, for example if the stored objects + require libraries that are not available, and not all pickled data is + compatible between different versions of Python). + Default: True + kwds : Keyword arguments, optional + Arrays to save to the file. Each array will be saved to the + output file with its corresponding keyword name. + + Returns + ------- + None + + See Also + -------- + numpy.save : Save a single array to a binary file in NumPy format. + numpy.savetxt : Save an array to a file as plain text. + numpy.savez : Save several arrays into an uncompressed ``.npz`` file format + numpy.load : Load the files created by savez_compressed. + + Notes + ----- + The ``.npz`` file format is a zipped archive of files named after the + variables they contain. The archive is compressed with + ``zipfile.ZIP_DEFLATED`` and each file in the archive contains one variable + in ``.npy`` format. For a description of the ``.npy`` format, see + :py:mod:`numpy.lib.format`. + + + When opening the saved ``.npz`` file with `load` a `~lib.npyio.NpzFile` + object is returned. This is a dictionary-like object which can be queried + for its list of arrays (with the ``.files`` attribute), and for the arrays + themselves. + + Examples + -------- + >>> import numpy as np + >>> test_array = np.random.rand(3, 2) + >>> test_vector = np.random.rand(4) + >>> np.savez_compressed('/tmp/123', a=test_array, b=test_vector) + >>> loaded = np.load('/tmp/123.npz') + >>> print(np.array_equal(test_array, loaded['a'])) + True + >>> print(np.array_equal(test_vector, loaded['b'])) + True + + """ + _savez(file, args, kwds, True, allow_pickle=allow_pickle) + + +def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None): + # Import is postponed to here since zipfile depends on gzip, an optional + # component of the so-called standard library. + import zipfile + + if not hasattr(file, 'write'): + file = os.fspath(file) + if not file.endswith('.npz'): + file = file + '.npz' + + namedict = kwds + for i, val in enumerate(args): + key = 'arr_%d' % i + if key in namedict.keys(): + raise ValueError( + "Cannot use un-named variables and keyword %s" % key) + namedict[key] = val + + if compress: + compression = zipfile.ZIP_DEFLATED + else: + compression = zipfile.ZIP_STORED + + zipf = zipfile_factory(file, mode="w", compression=compression) + try: + for key, val in namedict.items(): + fname = key + '.npy' + val = np.asanyarray(val) + # always force zip64, gh-10776 + with zipf.open(fname, 'w', force_zip64=True) as fid: + format.write_array(fid, val, + allow_pickle=allow_pickle, + pickle_kwargs=pickle_kwargs) + finally: + zipf.close() + + +def _ensure_ndmin_ndarray_check_param(ndmin): + """Just checks if the param ndmin is supported on + _ensure_ndmin_ndarray. It is intended to be used as + verification before running anything expensive. + e.g. loadtxt, genfromtxt + """ + # Check correctness of the values of `ndmin` + if ndmin not in [0, 1, 2]: + raise ValueError(f"Illegal value of ndmin keyword: {ndmin}") + +def _ensure_ndmin_ndarray(a, *, ndmin: int): + """This is a helper function of loadtxt and genfromtxt to ensure + proper minimum dimension as requested + + ndim : int. Supported values 1, 2, 3 + ^^ whenever this changes, keep in sync with + _ensure_ndmin_ndarray_check_param + """ + # Verify that the array has at least dimensions `ndmin`. + # Tweak the size and shape of the arrays - remove extraneous dimensions + if a.ndim > ndmin: + a = np.squeeze(a) + # and ensure we have the minimum number of dimensions asked for + # - has to be in this order for the odd case ndmin=1, a.squeeze().ndim=0 + if a.ndim < ndmin: + if ndmin == 1: + a = np.atleast_1d(a) + elif ndmin == 2: + a = np.atleast_2d(a).T + + return a + + +# amount of lines loadtxt reads in one chunk, can be overridden for testing +_loadtxt_chunksize = 50000 + + +def _check_nonneg_int(value, name="argument"): + try: + operator.index(value) + except TypeError: + raise TypeError(f"{name} must be an integer") from None + if value < 0: + raise ValueError(f"{name} must be nonnegative") + + +def _preprocess_comments(iterable, comments, encoding): + """ + Generator that consumes a line iterated iterable and strips out the + multiple (or multi-character) comments from lines. + This is a pre-processing step to achieve feature parity with loadtxt + (we assume that this feature is a nieche feature). + """ + for line in iterable: + if isinstance(line, bytes): + # Need to handle conversion here, or the splitting would fail + line = line.decode(encoding) + + for c in comments: + line = line.split(c, 1)[0] + + yield line + + +# The number of rows we read in one go if confronted with a parametric dtype +_loadtxt_chunksize = 50000 + + +def _read(fname, *, delimiter=',', comment='#', quote='"', + imaginary_unit='j', usecols=None, skiplines=0, + max_rows=None, converters=None, ndmin=None, unpack=False, + dtype=np.float64, encoding=None): + r""" + Read a NumPy array from a text file. + This is a helper function for loadtxt. + + Parameters + ---------- + fname : file, str, or pathlib.Path + The filename or the file to be read. + delimiter : str, optional + Field delimiter of the fields in line of the file. + Default is a comma, ','. If None any sequence of whitespace is + considered a delimiter. + comment : str or sequence of str or None, optional + Character that begins a comment. All text from the comment + character to the end of the line is ignored. + Multiple comments or multiple-character comment strings are supported, + but may be slower and `quote` must be empty if used. + Use None to disable all use of comments. + quote : str or None, optional + Character that is used to quote string fields. Default is '"' + (a double quote). Use None to disable quote support. + imaginary_unit : str, optional + Character that represent the imaginary unit `sqrt(-1)`. + Default is 'j'. + usecols : array_like, optional + A one-dimensional array of integer column numbers. These are the + columns from the file to be included in the array. If this value + is not given, all the columns are used. + skiplines : int, optional + Number of lines to skip before interpreting the data in the file. + max_rows : int, optional + Maximum number of rows of data to read. Default is to read the + entire file. + converters : dict or callable, optional + A function to parse all columns strings into the desired value, or + a dictionary mapping column number to a parser function. + E.g. if column 0 is a date string: ``converters = {0: datestr2num}``. + Converters can also be used to provide a default value for missing + data, e.g. ``converters = lambda s: float(s.strip() or 0)`` will + convert empty fields to 0. + Default: None + ndmin : int, optional + Minimum dimension of the array returned. + Allowed values are 0, 1 or 2. Default is 0. + unpack : bool, optional + If True, the returned array is transposed, so that arguments may be + unpacked using ``x, y, z = read(...)``. When used with a structured + data-type, arrays are returned for each field. Default is False. + dtype : numpy data type + A NumPy dtype instance, can be a structured dtype to map to the + columns of the file. + encoding : str, optional + Encoding used to decode the inputfile. The special value 'bytes' + (the default) enables backwards-compatible behavior for `converters`, + ensuring that inputs to the converter functions are encoded + bytes objects. The special value 'bytes' has no additional effect if + ``converters=None``. If encoding is ``'bytes'`` or ``None``, the + default system encoding is used. + + Returns + ------- + ndarray + NumPy array. + """ + # Handle special 'bytes' keyword for encoding + byte_converters = False + if encoding == 'bytes': + encoding = None + byte_converters = True + + if dtype is None: + raise TypeError("a dtype must be provided.") + dtype = np.dtype(dtype) + + read_dtype_via_object_chunks = None + if dtype.kind in 'SUM' and ( + dtype == "S0" or dtype == "U0" or dtype == "M8" or dtype == 'm8'): + # This is a legacy "flexible" dtype. We do not truly support + # parametric dtypes currently (no dtype discovery step in the core), + # but have to support these for backward compatibility. + read_dtype_via_object_chunks = dtype + dtype = np.dtype(object) + + if usecols is not None: + # Allow usecols to be a single int or a sequence of ints, the C-code + # handles the rest + try: + usecols = list(usecols) + except TypeError: + usecols = [usecols] + + _ensure_ndmin_ndarray_check_param(ndmin) + + if comment is None: + comments = None + else: + # assume comments are a sequence of strings + if "" in comment: + raise ValueError( + "comments cannot be an empty string. Use comments=None to " + "disable comments." + ) + comments = tuple(comment) + comment = None + if len(comments) == 0: + comments = None # No comments at all + elif len(comments) == 1: + # If there is only one comment, and that comment has one character, + # the normal parsing can deal with it just fine. + if isinstance(comments[0], str) and len(comments[0]) == 1: + comment = comments[0] + comments = None + else: + # Input validation if there are multiple comment characters + if delimiter in comments: + raise TypeError( + f"Comment characters '{comments}' cannot include the " + f"delimiter '{delimiter}'" + ) + + # comment is now either a 1 or 0 character string or a tuple: + if comments is not None: + # Note: An earlier version support two character comments (and could + # have been extended to multiple characters, we assume this is + # rare enough to not optimize for. + if quote is not None: + raise ValueError( + "when multiple comments or a multi-character comment is " + "given, quotes are not supported. In this case quotechar " + "must be set to None.") + + if len(imaginary_unit) != 1: + raise ValueError('len(imaginary_unit) must be 1.') + + _check_nonneg_int(skiplines) + if max_rows is not None: + _check_nonneg_int(max_rows) + else: + # Passing -1 to the C code means "read the entire file". + max_rows = -1 + + fh_closing_ctx = contextlib.nullcontext() + filelike = False + try: + if isinstance(fname, os.PathLike): + fname = os.fspath(fname) + if isinstance(fname, str): + fh = np.lib._datasource.open(fname, 'rt', encoding=encoding) + if encoding is None: + encoding = getattr(fh, 'encoding', 'latin1') + + fh_closing_ctx = contextlib.closing(fh) + data = fh + filelike = True + else: + if encoding is None: + encoding = getattr(fname, 'encoding', 'latin1') + data = iter(fname) + except TypeError as e: + raise ValueError( + f"fname must be a string, filehandle, list of strings,\n" + f"or generator. Got {type(fname)} instead.") from e + + with fh_closing_ctx: + if comments is not None: + if filelike: + data = iter(data) + filelike = False + data = _preprocess_comments(data, comments, encoding) + + if read_dtype_via_object_chunks is None: + arr = _load_from_filelike( + data, delimiter=delimiter, comment=comment, quote=quote, + imaginary_unit=imaginary_unit, + usecols=usecols, skiplines=skiplines, max_rows=max_rows, + converters=converters, dtype=dtype, + encoding=encoding, filelike=filelike, + byte_converters=byte_converters) + + else: + # This branch reads the file into chunks of object arrays and then + # casts them to the desired actual dtype. This ensures correct + # string-length and datetime-unit discovery (like `arr.astype()`). + # Due to chunking, certain error reports are less clear, currently. + if filelike: + data = iter(data) # cannot chunk when reading from file + filelike = False + + c_byte_converters = False + if read_dtype_via_object_chunks == "S": + c_byte_converters = True # Use latin1 rather than ascii + + chunks = [] + while max_rows != 0: + if max_rows < 0: + chunk_size = _loadtxt_chunksize + else: + chunk_size = min(_loadtxt_chunksize, max_rows) + + next_arr = _load_from_filelike( + data, delimiter=delimiter, comment=comment, quote=quote, + imaginary_unit=imaginary_unit, + usecols=usecols, skiplines=skiplines, max_rows=chunk_size, + converters=converters, dtype=dtype, + encoding=encoding, filelike=filelike, + byte_converters=byte_converters, + c_byte_converters=c_byte_converters) + # Cast here already. We hope that this is better even for + # large files because the storage is more compact. It could + # be adapted (in principle the concatenate could cast). + chunks.append(next_arr.astype(read_dtype_via_object_chunks)) + + skiplines = 0 # Only have to skip for first chunk + if max_rows >= 0: + max_rows -= chunk_size + if len(next_arr) < chunk_size: + # There was less data than requested, so we are done. + break + + # Need at least one chunk, but if empty, the last one may have + # the wrong shape. + if len(chunks) > 1 and len(chunks[-1]) == 0: + del chunks[-1] + if len(chunks) == 1: + arr = chunks[0] + else: + arr = np.concatenate(chunks, axis=0) + + # NOTE: ndmin works as advertised for structured dtypes, but normally + # these would return a 1D result plus the structured dimension, + # so ndmin=2 adds a third dimension even when no squeezing occurs. + # A `squeeze=False` could be a better solution (pandas uses squeeze). + arr = _ensure_ndmin_ndarray(arr, ndmin=ndmin) + + if arr.shape: + if arr.shape[0] == 0: + warnings.warn( + f'loadtxt: input contained no data: "{fname}"', + category=UserWarning, + stacklevel=3 + ) + + if unpack: + # Unpack structured dtypes if requested: + dt = arr.dtype + if dt.names is not None: + # For structured arrays, return an array for each field. + return [arr[field] for field in dt.names] + else: + return arr.T + else: + return arr + + +@finalize_array_function_like +@set_module('numpy') +def loadtxt(fname, dtype=float, comments='#', delimiter=None, + converters=None, skiprows=0, usecols=None, unpack=False, + ndmin=0, encoding=None, max_rows=None, *, quotechar=None, + like=None): + r""" + Load data from a text file. + + Parameters + ---------- + fname : file, str, pathlib.Path, list of str, generator + File, filename, list, or generator to read. If the filename + extension is ``.gz`` or ``.bz2``, the file is first decompressed. Note + that generators must return bytes or strings. The strings + in a list or produced by a generator are treated as lines. + dtype : data-type, optional + Data-type of the resulting array; default: float. If this is a + structured data-type, the resulting array will be 1-dimensional, and + each row will be interpreted as an element of the array. In this + case, the number of columns used must match the number of fields in + the data-type. + comments : str or sequence of str or None, optional + The characters or list of characters used to indicate the start of a + comment. None implies no comments. For backwards compatibility, byte + strings will be decoded as 'latin1'. The default is '#'. + delimiter : str, optional + The character used to separate the values. For backwards compatibility, + byte strings will be decoded as 'latin1'. The default is whitespace. + + .. versionchanged:: 1.23.0 + Only single character delimiters are supported. Newline characters + cannot be used as the delimiter. + + converters : dict or callable, optional + Converter functions to customize value parsing. If `converters` is + callable, the function is applied to all columns, else it must be a + dict that maps column number to a parser function. + See examples for further details. + Default: None. + + .. versionchanged:: 1.23.0 + The ability to pass a single callable to be applied to all columns + was added. + + skiprows : int, optional + Skip the first `skiprows` lines, including comments; default: 0. + usecols : int or sequence, optional + Which columns to read, with 0 being the first. For example, + ``usecols = (1,4,5)`` will extract the 2nd, 5th and 6th columns. + The default, None, results in all columns being read. + unpack : bool, optional + If True, the returned array is transposed, so that arguments may be + unpacked using ``x, y, z = loadtxt(...)``. When used with a + structured data-type, arrays are returned for each field. + Default is False. + ndmin : int, optional + The returned array will have at least `ndmin` dimensions. + Otherwise mono-dimensional axes will be squeezed. + Legal values: 0 (default), 1 or 2. + encoding : str, optional + Encoding used to decode the inputfile. Does not apply to input streams. + The special value 'bytes' enables backward compatibility workarounds + that ensures you receive byte arrays as results if possible and passes + 'latin1' encoded strings to converters. Override this value to receive + unicode arrays and pass strings as input to converters. If set to None + the system default is used. The default value is 'bytes'. + + .. versionchanged:: 2.0 + Before NumPy 2, the default was ``'bytes'`` for Python 2 + compatibility. The default is now ``None``. + + max_rows : int, optional + Read `max_rows` rows of content after `skiprows` lines. The default is + to read all the rows. Note that empty rows containing no data such as + empty lines and comment lines are not counted towards `max_rows`, + while such lines are counted in `skiprows`. + + .. versionchanged:: 1.23.0 + Lines containing no data, including comment lines (e.g., lines + starting with '#' or as specified via `comments`) are not counted + towards `max_rows`. + quotechar : unicode character or None, optional + The character used to denote the start and end of a quoted item. + Occurrences of the delimiter or comment characters are ignored within + a quoted item. The default value is ``quotechar=None``, which means + quoting support is disabled. + + If two consecutive instances of `quotechar` are found within a quoted + field, the first is treated as an escape character. See examples. + + .. versionadded:: 1.23.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + out : ndarray + Data read from the text file. + + See Also + -------- + load, fromstring, fromregex + genfromtxt : Load data with missing values handled as specified. + scipy.io.loadmat : reads MATLAB data files + + Notes + ----- + This function aims to be a fast reader for simply formatted files. The + `genfromtxt` function provides more sophisticated handling of, e.g., + lines with missing values. + + Each row in the input text file must have the same number of values to be + able to read all values. If all rows do not have same number of values, a + subset of up to n columns (where n is the least number of values present + in all rows) can be read by specifying the columns via `usecols`. + + The strings produced by the Python float.hex method can be used as + input for floats. + + Examples + -------- + >>> import numpy as np + >>> from io import StringIO # StringIO behaves like a file object + >>> c = StringIO("0 1\n2 3") + >>> np.loadtxt(c) + array([[0., 1.], + [2., 3.]]) + + >>> d = StringIO("M 21 72\nF 35 58") + >>> np.loadtxt(d, dtype={'names': ('gender', 'age', 'weight'), + ... 'formats': ('S1', 'i4', 'f4')}) + array([(b'M', 21, 72.), (b'F', 35, 58.)], + dtype=[('gender', 'S1'), ('age', '>> c = StringIO("1,0,2\n3,0,4") + >>> x, y = np.loadtxt(c, delimiter=',', usecols=(0, 2), unpack=True) + >>> x + array([1., 3.]) + >>> y + array([2., 4.]) + + The `converters` argument is used to specify functions to preprocess the + text prior to parsing. `converters` can be a dictionary that maps + preprocessing functions to each column: + + >>> s = StringIO("1.618, 2.296\n3.141, 4.669\n") + >>> conv = { + ... 0: lambda x: np.floor(float(x)), # conversion fn for column 0 + ... 1: lambda x: np.ceil(float(x)), # conversion fn for column 1 + ... } + >>> np.loadtxt(s, delimiter=",", converters=conv) + array([[1., 3.], + [3., 5.]]) + + `converters` can be a callable instead of a dictionary, in which case it + is applied to all columns: + + >>> s = StringIO("0xDE 0xAD\n0xC0 0xDE") + >>> import functools + >>> conv = functools.partial(int, base=16) + >>> np.loadtxt(s, converters=conv) + array([[222., 173.], + [192., 222.]]) + + This example shows how `converters` can be used to convert a field + with a trailing minus sign into a negative number. + + >>> s = StringIO("10.01 31.25-\n19.22 64.31\n17.57- 63.94") + >>> def conv(fld): + ... return -float(fld[:-1]) if fld.endswith("-") else float(fld) + ... + >>> np.loadtxt(s, converters=conv) + array([[ 10.01, -31.25], + [ 19.22, 64.31], + [-17.57, 63.94]]) + + Using a callable as the converter can be particularly useful for handling + values with different formatting, e.g. floats with underscores: + + >>> s = StringIO("1 2.7 100_000") + >>> np.loadtxt(s, converters=float) + array([1.e+00, 2.7e+00, 1.e+05]) + + This idea can be extended to automatically handle values specified in + many different formats, such as hex values: + + >>> def conv(val): + ... try: + ... return float(val) + ... except ValueError: + ... return float.fromhex(val) + >>> s = StringIO("1, 2.5, 3_000, 0b4, 0x1.4000000000000p+2") + >>> np.loadtxt(s, delimiter=",", converters=conv) + array([1.0e+00, 2.5e+00, 3.0e+03, 1.8e+02, 5.0e+00]) + + Or a format where the ``-`` sign comes after the number: + + >>> s = StringIO("10.01 31.25-\n19.22 64.31\n17.57- 63.94") + >>> conv = lambda x: -float(x[:-1]) if x.endswith("-") else float(x) + >>> np.loadtxt(s, converters=conv) + array([[ 10.01, -31.25], + [ 19.22, 64.31], + [-17.57, 63.94]]) + + Support for quoted fields is enabled with the `quotechar` parameter. + Comment and delimiter characters are ignored when they appear within a + quoted item delineated by `quotechar`: + + >>> s = StringIO('"alpha, #42", 10.0\n"beta, #64", 2.0\n') + >>> dtype = np.dtype([("label", "U12"), ("value", float)]) + >>> np.loadtxt(s, dtype=dtype, delimiter=",", quotechar='"') + array([('alpha, #42', 10.), ('beta, #64', 2.)], + dtype=[('label', '>> s = StringIO('"alpha, #42" 10.0\n"beta, #64" 2.0\n') + >>> dtype = np.dtype([("label", "U12"), ("value", float)]) + >>> np.loadtxt(s, dtype=dtype, delimiter=None, quotechar='"') + array([('alpha, #42', 10.), ('beta, #64', 2.)], + dtype=[('label', '>> s = StringIO('"Hello, my name is ""Monty""!"') + >>> np.loadtxt(s, dtype="U", delimiter=",", quotechar='"') + array('Hello, my name is "Monty"!', dtype='>> d = StringIO("1 2\n2 4\n3 9 12\n4 16 20") + >>> np.loadtxt(d, usecols=(0, 1)) + array([[ 1., 2.], + [ 2., 4.], + [ 3., 9.], + [ 4., 16.]]) + + """ + + if like is not None: + return _loadtxt_with_like( + like, fname, dtype=dtype, comments=comments, delimiter=delimiter, + converters=converters, skiprows=skiprows, usecols=usecols, + unpack=unpack, ndmin=ndmin, encoding=encoding, + max_rows=max_rows + ) + + if isinstance(delimiter, bytes): + delimiter.decode("latin1") + + if dtype is None: + dtype = np.float64 + + comment = comments + # Control character type conversions for Py3 convenience + if comment is not None: + if isinstance(comment, (str, bytes)): + comment = [comment] + comment = [ + x.decode('latin1') if isinstance(x, bytes) else x for x in comment] + if isinstance(delimiter, bytes): + delimiter = delimiter.decode('latin1') + + arr = _read(fname, dtype=dtype, comment=comment, delimiter=delimiter, + converters=converters, skiplines=skiprows, usecols=usecols, + unpack=unpack, ndmin=ndmin, encoding=encoding, + max_rows=max_rows, quote=quotechar) + + return arr + + +_loadtxt_with_like = array_function_dispatch()(loadtxt) + + +def _savetxt_dispatcher(fname, X, fmt=None, delimiter=None, newline=None, + header=None, footer=None, comments=None, + encoding=None): + return (X,) + + +@array_function_dispatch(_savetxt_dispatcher) +def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', + footer='', comments='# ', encoding=None): + """ + Save an array to a text file. + + Parameters + ---------- + fname : filename, file handle or pathlib.Path + If the filename ends in ``.gz``, the file is automatically saved in + compressed gzip format. `loadtxt` understands gzipped files + transparently. + X : 1D or 2D array_like + Data to be saved to a text file. + fmt : str or sequence of strs, optional + A single format (%10.5f), a sequence of formats, or a + multi-format string, e.g. 'Iteration %d -- %10.5f', in which + case `delimiter` is ignored. For complex `X`, the legal options + for `fmt` are: + + * a single specifier, ``fmt='%.4e'``, resulting in numbers formatted + like ``' (%s+%sj)' % (fmt, fmt)`` + * a full string specifying every real and imaginary part, e.g. + ``' %.4e %+.4ej %.4e %+.4ej %.4e %+.4ej'`` for 3 columns + * a list of specifiers, one per column - in this case, the real + and imaginary part must have separate specifiers, + e.g. ``['%.3e + %.3ej', '(%.15e%+.15ej)']`` for 2 columns + delimiter : str, optional + String or character separating columns. + newline : str, optional + String or character separating lines. + header : str, optional + String that will be written at the beginning of the file. + footer : str, optional + String that will be written at the end of the file. + comments : str, optional + String that will be prepended to the ``header`` and ``footer`` strings, + to mark them as comments. Default: '# ', as expected by e.g. + ``numpy.loadtxt``. + encoding : {None, str}, optional + Encoding used to encode the outputfile. Does not apply to output + streams. If the encoding is something other than 'bytes' or 'latin1' + you will not be able to load the file in NumPy versions < 1.14. Default + is 'latin1'. + + See Also + -------- + save : Save an array to a binary file in NumPy ``.npy`` format + savez : Save several arrays into an uncompressed ``.npz`` archive + savez_compressed : Save several arrays into a compressed ``.npz`` archive + + Notes + ----- + Further explanation of the `fmt` parameter + (``%[flag]width[.precision]specifier``): + + flags: + ``-`` : left justify + + ``+`` : Forces to precede result with + or -. + + ``0`` : Left pad the number with zeros instead of space (see width). + + width: + Minimum number of characters to be printed. The value is not truncated + if it has more characters. + + precision: + - For integer specifiers (eg. ``d,i,o,x``), the minimum number of + digits. + - For ``e, E`` and ``f`` specifiers, the number of digits to print + after the decimal point. + - For ``g`` and ``G``, the maximum number of significant digits. + - For ``s``, the maximum number of characters. + + specifiers: + ``c`` : character + + ``d`` or ``i`` : signed decimal integer + + ``e`` or ``E`` : scientific notation with ``e`` or ``E``. + + ``f`` : decimal floating point + + ``g,G`` : use the shorter of ``e,E`` or ``f`` + + ``o`` : signed octal + + ``s`` : string of characters + + ``u`` : unsigned decimal integer + + ``x,X`` : unsigned hexadecimal integer + + This explanation of ``fmt`` is not complete, for an exhaustive + specification see [1]_. + + References + ---------- + .. [1] `Format Specification Mini-Language + `_, + Python Documentation. + + Examples + -------- + >>> import numpy as np + >>> x = y = z = np.arange(0.0,5.0,1.0) + >>> np.savetxt('test.out', x, delimiter=',') # X is an array + >>> np.savetxt('test.out', (x,y,z)) # x,y,z equal sized 1D arrays + >>> np.savetxt('test.out', x, fmt='%1.4e') # use exponential notation + + """ + + class WriteWrap: + """Convert to bytes on bytestream inputs. + + """ + def __init__(self, fh, encoding): + self.fh = fh + self.encoding = encoding + self.do_write = self.first_write + + def close(self): + self.fh.close() + + def write(self, v): + self.do_write(v) + + def write_bytes(self, v): + if isinstance(v, bytes): + self.fh.write(v) + else: + self.fh.write(v.encode(self.encoding)) + + def write_normal(self, v): + self.fh.write(asunicode(v)) + + def first_write(self, v): + try: + self.write_normal(v) + self.write = self.write_normal + except TypeError: + # input is probably a bytestream + self.write_bytes(v) + self.write = self.write_bytes + + own_fh = False + if isinstance(fname, os.PathLike): + fname = os.fspath(fname) + if _is_string_like(fname): + # datasource doesn't support creating a new file ... + open(fname, 'wt').close() + fh = np.lib._datasource.open(fname, 'wt', encoding=encoding) + own_fh = True + elif hasattr(fname, 'write'): + # wrap to handle byte output streams + fh = WriteWrap(fname, encoding or 'latin1') + else: + raise ValueError('fname must be a string or file handle') + + try: + X = np.asarray(X) + + # Handle 1-dimensional arrays + if X.ndim == 0 or X.ndim > 2: + raise ValueError( + "Expected 1D or 2D array, got %dD array instead" % X.ndim) + elif X.ndim == 1: + # Common case -- 1d array of numbers + if X.dtype.names is None: + X = np.atleast_2d(X).T + ncol = 1 + + # Complex dtype -- each field indicates a separate column + else: + ncol = len(X.dtype.names) + else: + ncol = X.shape[1] + + iscomplex_X = np.iscomplexobj(X) + # `fmt` can be a string with multiple insertion points or a + # list of formats. E.g. '%10.5f\t%10d' or ('%10.5f', '$10d') + if type(fmt) in (list, tuple): + if len(fmt) != ncol: + raise AttributeError('fmt has wrong shape. %s' % str(fmt)) + format = delimiter.join(fmt) + elif isinstance(fmt, str): + n_fmt_chars = fmt.count('%') + error = ValueError('fmt has wrong number of %% formats: %s' % fmt) + if n_fmt_chars == 1: + if iscomplex_X: + fmt = [' (%s+%sj)' % (fmt, fmt), ] * ncol + else: + fmt = [fmt, ] * ncol + format = delimiter.join(fmt) + elif iscomplex_X and n_fmt_chars != (2 * ncol): + raise error + elif ((not iscomplex_X) and n_fmt_chars != ncol): + raise error + else: + format = fmt + else: + raise ValueError('invalid fmt: %r' % (fmt,)) + + if len(header) > 0: + header = header.replace('\n', '\n' + comments) + fh.write(comments + header + newline) + if iscomplex_X: + for row in X: + row2 = [] + for number in row: + row2.append(number.real) + row2.append(number.imag) + s = format % tuple(row2) + newline + fh.write(s.replace('+-', '-')) + else: + for row in X: + try: + v = format % tuple(row) + newline + except TypeError as e: + raise TypeError("Mismatch between array dtype ('%s') and " + "format specifier ('%s')" + % (str(X.dtype), format)) from e + fh.write(v) + + if len(footer) > 0: + footer = footer.replace('\n', '\n' + comments) + fh.write(comments + footer + newline) + finally: + if own_fh: + fh.close() + + +@set_module('numpy') +def fromregex(file, regexp, dtype, encoding=None): + r""" + Construct an array from a text file, using regular expression parsing. + + The returned array is always a structured array, and is constructed from + all matches of the regular expression in the file. Groups in the regular + expression are converted to fields of the structured array. + + Parameters + ---------- + file : file, str, or pathlib.Path + Filename or file object to read. + + .. versionchanged:: 1.22.0 + Now accepts `os.PathLike` implementations. + + regexp : str or regexp + Regular expression used to parse the file. + Groups in the regular expression correspond to fields in the dtype. + dtype : dtype or list of dtypes + Dtype for the structured array; must be a structured datatype. + encoding : str, optional + Encoding used to decode the inputfile. Does not apply to input streams. + + Returns + ------- + output : ndarray + The output array, containing the part of the content of `file` that + was matched by `regexp`. `output` is always a structured array. + + Raises + ------ + TypeError + When `dtype` is not a valid dtype for a structured array. + + See Also + -------- + fromstring, loadtxt + + Notes + ----- + Dtypes for structured arrays can be specified in several forms, but all + forms specify at least the data type and field name. For details see + `basics.rec`. + + Examples + -------- + >>> import numpy as np + >>> from io import StringIO + >>> text = StringIO("1312 foo\n1534 bar\n444 qux") + + >>> regexp = r"(\d+)\s+(...)" # match [digits, whitespace, anything] + >>> output = np.fromregex(text, regexp, + ... [('num', np.int64), ('key', 'S3')]) + >>> output + array([(1312, b'foo'), (1534, b'bar'), ( 444, b'qux')], + dtype=[('num', '>> output['num'] + array([1312, 1534, 444]) + + """ + own_fh = False + if not hasattr(file, "read"): + file = os.fspath(file) + file = np.lib._datasource.open(file, 'rt', encoding=encoding) + own_fh = True + + try: + if not isinstance(dtype, np.dtype): + dtype = np.dtype(dtype) + if dtype.names is None: + raise TypeError('dtype must be a structured datatype.') + + content = file.read() + if isinstance(content, bytes) and isinstance(regexp, str): + regexp = asbytes(regexp) + + if not hasattr(regexp, 'match'): + regexp = re.compile(regexp) + seq = regexp.findall(content) + if seq and not isinstance(seq[0], tuple): + # Only one group is in the regexp. + # Create the new array as a single data-type and then + # re-interpret as a single-field structured array. + newdtype = np.dtype(dtype[dtype.names[0]]) + output = np.array(seq, dtype=newdtype) + output.dtype = dtype + else: + output = np.array(seq, dtype=dtype) + + return output + finally: + if own_fh: + file.close() + + +#####-------------------------------------------------------------------------- +#---- --- ASCII functions --- +#####-------------------------------------------------------------------------- + + +@finalize_array_function_like +@set_module('numpy') +def genfromtxt(fname, dtype=float, comments='#', delimiter=None, + skip_header=0, skip_footer=0, converters=None, + missing_values=None, filling_values=None, usecols=None, + names=None, excludelist=None, + deletechars=''.join(sorted(NameValidator.defaultdeletechars)), + replace_space='_', autostrip=False, case_sensitive=True, + defaultfmt="f%i", unpack=None, usemask=False, loose=True, + invalid_raise=True, max_rows=None, encoding=None, + *, ndmin=0, like=None): + """ + Load data from a text file, with missing values handled as specified. + + Each line past the first `skip_header` lines is split at the `delimiter` + character, and characters following the `comments` character are discarded. + + Parameters + ---------- + fname : file, str, pathlib.Path, list of str, generator + File, filename, list, or generator to read. If the filename + extension is ``.gz`` or ``.bz2``, the file is first decompressed. Note + that generators must return bytes or strings. The strings + in a list or produced by a generator are treated as lines. + dtype : dtype, optional + Data type of the resulting array. + If None, the dtypes will be determined by the contents of each + column, individually. + comments : str, optional + The character used to indicate the start of a comment. + All the characters occurring on a line after a comment are discarded. + delimiter : str, int, or sequence, optional + The string used to separate values. By default, any consecutive + whitespaces act as delimiter. An integer or sequence of integers + can also be provided as width(s) of each field. + skiprows : int, optional + `skiprows` was removed in numpy 1.10. Please use `skip_header` instead. + skip_header : int, optional + The number of lines to skip at the beginning of the file. + skip_footer : int, optional + The number of lines to skip at the end of the file. + converters : variable, optional + The set of functions that convert the data of a column to a value. + The converters can also be used to provide a default value + for missing data: ``converters = {3: lambda s: float(s or 0)}``. + missing : variable, optional + `missing` was removed in numpy 1.10. Please use `missing_values` + instead. + missing_values : variable, optional + The set of strings corresponding to missing data. + filling_values : variable, optional + The set of values to be used as default when the data are missing. + usecols : sequence, optional + Which columns to read, with 0 being the first. For example, + ``usecols = (1, 4, 5)`` will extract the 2nd, 5th and 6th columns. + names : {None, True, str, sequence}, optional + If `names` is True, the field names are read from the first line after + the first `skip_header` lines. This line can optionally be preceded + by a comment delimiter. Any content before the comment delimiter is + discarded. If `names` is a sequence or a single-string of + comma-separated names, the names will be used to define the field + names in a structured dtype. If `names` is None, the names of the + dtype fields will be used, if any. + excludelist : sequence, optional + A list of names to exclude. This list is appended to the default list + ['return','file','print']. Excluded names are appended with an + underscore: for example, `file` would become `file_`. + deletechars : str, optional + A string combining invalid characters that must be deleted from the + names. + defaultfmt : str, optional + A format used to define default field names, such as "f%i" or "f_%02i". + autostrip : bool, optional + Whether to automatically strip white spaces from the variables. + replace_space : char, optional + Character(s) used in replacement of white spaces in the variable + names. By default, use a '_'. + case_sensitive : {True, False, 'upper', 'lower'}, optional + If True, field names are case sensitive. + If False or 'upper', field names are converted to upper case. + If 'lower', field names are converted to lower case. + unpack : bool, optional + If True, the returned array is transposed, so that arguments may be + unpacked using ``x, y, z = genfromtxt(...)``. When used with a + structured data-type, arrays are returned for each field. + Default is False. + usemask : bool, optional + If True, return a masked array. + If False, return a regular array. + loose : bool, optional + If True, do not raise errors for invalid values. + invalid_raise : bool, optional + If True, an exception is raised if an inconsistency is detected in the + number of columns. + If False, a warning is emitted and the offending lines are skipped. + max_rows : int, optional + The maximum number of rows to read. Must not be used with skip_footer + at the same time. If given, the value must be at least 1. Default is + to read the entire file. + encoding : str, optional + Encoding used to decode the inputfile. Does not apply when `fname` + is a file object. The special value 'bytes' enables backward + compatibility workarounds that ensure that you receive byte arrays + when possible and passes latin1 encoded strings to converters. + Override this value to receive unicode arrays and pass strings + as input to converters. If set to None the system default is used. + The default value is 'bytes'. + + .. versionchanged:: 2.0 + Before NumPy 2, the default was ``'bytes'`` for Python 2 + compatibility. The default is now ``None``. + + ndmin : int, optional + Same parameter as `loadtxt` + + .. versionadded:: 1.23.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + out : ndarray + Data read from the text file. If `usemask` is True, this is a + masked array. + + See Also + -------- + numpy.loadtxt : equivalent function when no data is missing. + + Notes + ----- + * When spaces are used as delimiters, or when no delimiter has been given + as input, there should not be any missing data between two fields. + * When variables are named (either by a flexible dtype or with a `names` + sequence), there must not be any header in the file (else a ValueError + exception is raised). + * Individual values are not stripped of spaces by default. + When using a custom converter, make sure the function does remove spaces. + * Custom converters may receive unexpected values due to dtype + discovery. + + References + ---------- + .. [1] NumPy User Guide, section `I/O with NumPy + `_. + + Examples + -------- + >>> from io import StringIO + >>> import numpy as np + + Comma delimited file with mixed dtype + + >>> s = StringIO("1,1.3,abcde") + >>> data = np.genfromtxt(s, dtype=[('myint','i8'),('myfloat','f8'), + ... ('mystring','S5')], delimiter=",") + >>> data + array((1, 1.3, b'abcde'), + dtype=[('myint', '>> _ = s.seek(0) # needed for StringIO example only + >>> data = np.genfromtxt(s, dtype=None, + ... names = ['myint','myfloat','mystring'], delimiter=",") + >>> data + array((1, 1.3, 'abcde'), + dtype=[('myint', '>> _ = s.seek(0) + >>> data = np.genfromtxt(s, dtype="i8,f8,S5", + ... names=['myint','myfloat','mystring'], delimiter=",") + >>> data + array((1, 1.3, b'abcde'), + dtype=[('myint', '>> s = StringIO("11.3abcde") + >>> data = np.genfromtxt(s, dtype=None, names=['intvar','fltvar','strvar'], + ... delimiter=[1,3,5]) + >>> data + array((1, 1.3, 'abcde'), + dtype=[('intvar', '>> f = StringIO(''' + ... text,# of chars + ... hello world,11 + ... numpy,5''') + >>> np.genfromtxt(f, dtype='S12,S12', delimiter=',') + array([(b'text', b''), (b'hello world', b'11'), (b'numpy', b'5')], + dtype=[('f0', 'S12'), ('f1', 'S12')]) + + """ + + if like is not None: + return _genfromtxt_with_like( + like, fname, dtype=dtype, comments=comments, delimiter=delimiter, + skip_header=skip_header, skip_footer=skip_footer, + converters=converters, missing_values=missing_values, + filling_values=filling_values, usecols=usecols, names=names, + excludelist=excludelist, deletechars=deletechars, + replace_space=replace_space, autostrip=autostrip, + case_sensitive=case_sensitive, defaultfmt=defaultfmt, + unpack=unpack, usemask=usemask, loose=loose, + invalid_raise=invalid_raise, max_rows=max_rows, encoding=encoding, + ndmin=ndmin, + ) + + _ensure_ndmin_ndarray_check_param(ndmin) + + if max_rows is not None: + if skip_footer: + raise ValueError( + "The keywords 'skip_footer' and 'max_rows' can not be " + "specified at the same time.") + if max_rows < 1: + raise ValueError("'max_rows' must be at least 1.") + + if usemask: + from numpy.ma import MaskedArray, make_mask_descr + # Check the input dictionary of converters + user_converters = converters or {} + if not isinstance(user_converters, dict): + raise TypeError( + "The input argument 'converter' should be a valid dictionary " + "(got '%s' instead)" % type(user_converters)) + + if encoding == 'bytes': + encoding = None + byte_converters = True + else: + byte_converters = False + + # Initialize the filehandle, the LineSplitter and the NameValidator + if isinstance(fname, os.PathLike): + fname = os.fspath(fname) + if isinstance(fname, str): + fid = np.lib._datasource.open(fname, 'rt', encoding=encoding) + fid_ctx = contextlib.closing(fid) + else: + fid = fname + fid_ctx = contextlib.nullcontext(fid) + try: + fhd = iter(fid) + except TypeError as e: + raise TypeError( + "fname must be a string, a filehandle, a sequence of strings,\n" + f"or an iterator of strings. Got {type(fname)} instead." + ) from e + with fid_ctx: + split_line = LineSplitter(delimiter=delimiter, comments=comments, + autostrip=autostrip, encoding=encoding) + validate_names = NameValidator(excludelist=excludelist, + deletechars=deletechars, + case_sensitive=case_sensitive, + replace_space=replace_space) + + # Skip the first `skip_header` rows + try: + for i in range(skip_header): + next(fhd) + + # Keep on until we find the first valid values + first_values = None + + while not first_values: + first_line = _decode_line(next(fhd), encoding) + if (names is True) and (comments is not None): + if comments in first_line: + first_line = ( + ''.join(first_line.split(comments)[1:])) + first_values = split_line(first_line) + except StopIteration: + # return an empty array if the datafile is empty + first_line = '' + first_values = [] + warnings.warn( + 'genfromtxt: Empty input file: "%s"' % fname, stacklevel=2 + ) + + # Should we take the first values as names ? + if names is True: + fval = first_values[0].strip() + if comments is not None: + if fval in comments: + del first_values[0] + + # Check the columns to use: make sure `usecols` is a list + if usecols is not None: + try: + usecols = [_.strip() for _ in usecols.split(",")] + except AttributeError: + try: + usecols = list(usecols) + except TypeError: + usecols = [usecols, ] + nbcols = len(usecols or first_values) + + # Check the names and overwrite the dtype.names if needed + if names is True: + names = validate_names([str(_.strip()) for _ in first_values]) + first_line = '' + elif _is_string_like(names): + names = validate_names([_.strip() for _ in names.split(',')]) + elif names: + names = validate_names(names) + # Get the dtype + if dtype is not None: + dtype = easy_dtype(dtype, defaultfmt=defaultfmt, names=names, + excludelist=excludelist, + deletechars=deletechars, + case_sensitive=case_sensitive, + replace_space=replace_space) + # Make sure the names is a list (for 2.5) + if names is not None: + names = list(names) + + if usecols: + for (i, current) in enumerate(usecols): + # if usecols is a list of names, convert to a list of indices + if _is_string_like(current): + usecols[i] = names.index(current) + elif current < 0: + usecols[i] = current + len(first_values) + # If the dtype is not None, make sure we update it + if (dtype is not None) and (len(dtype) > nbcols): + descr = dtype.descr + dtype = np.dtype([descr[_] for _ in usecols]) + names = list(dtype.names) + # If `names` is not None, update the names + elif (names is not None) and (len(names) > nbcols): + names = [names[_] for _ in usecols] + elif (names is not None) and (dtype is not None): + names = list(dtype.names) + + # Process the missing values ............................... + # Rename missing_values for convenience + user_missing_values = missing_values or () + if isinstance(user_missing_values, bytes): + user_missing_values = user_missing_values.decode('latin1') + + # Define the list of missing_values (one column: one list) + missing_values = [[''] for _ in range(nbcols)] + + # We have a dictionary: process it field by field + if isinstance(user_missing_values, dict): + # Loop on the items + for (key, val) in user_missing_values.items(): + # Is the key a string ? + if _is_string_like(key): + try: + # Transform it into an integer + key = names.index(key) + except ValueError: + # We couldn't find it: the name must have been dropped + continue + # Redefine the key as needed if it's a column number + if usecols: + try: + key = usecols.index(key) + except ValueError: + pass + # Transform the value as a list of string + if isinstance(val, (list, tuple)): + val = [str(_) for _ in val] + else: + val = [str(val), ] + # Add the value(s) to the current list of missing + if key is None: + # None acts as default + for miss in missing_values: + miss.extend(val) + else: + missing_values[key].extend(val) + # We have a sequence : each item matches a column + elif isinstance(user_missing_values, (list, tuple)): + for (value, entry) in zip(user_missing_values, missing_values): + value = str(value) + if value not in entry: + entry.append(value) + # We have a string : apply it to all entries + elif isinstance(user_missing_values, str): + user_value = user_missing_values.split(",") + for entry in missing_values: + entry.extend(user_value) + # We have something else: apply it to all entries + else: + for entry in missing_values: + entry.extend([str(user_missing_values)]) + + # Process the filling_values ............................... + # Rename the input for convenience + user_filling_values = filling_values + if user_filling_values is None: + user_filling_values = [] + # Define the default + filling_values = [None] * nbcols + # We have a dictionary : update each entry individually + if isinstance(user_filling_values, dict): + for (key, val) in user_filling_values.items(): + if _is_string_like(key): + try: + # Transform it into an integer + key = names.index(key) + except ValueError: + # We couldn't find it: the name must have been dropped + continue + # Redefine the key if it's a column number + # and usecols is defined + if usecols: + try: + key = usecols.index(key) + except ValueError: + pass + # Add the value to the list + filling_values[key] = val + # We have a sequence : update on a one-to-one basis + elif isinstance(user_filling_values, (list, tuple)): + n = len(user_filling_values) + if (n <= nbcols): + filling_values[:n] = user_filling_values + else: + filling_values = user_filling_values[:nbcols] + # We have something else : use it for all entries + else: + filling_values = [user_filling_values] * nbcols + + # Initialize the converters ................................ + if dtype is None: + # Note: we can't use a [...]*nbcols, as we would have 3 times + # the same converter, instead of 3 different converters. + converters = [ + StringConverter(None, missing_values=miss, default=fill) + for (miss, fill) in zip(missing_values, filling_values) + ] + else: + dtype_flat = flatten_dtype(dtype, flatten_base=True) + # Initialize the converters + if len(dtype_flat) > 1: + # Flexible type : get a converter from each dtype + zipit = zip(dtype_flat, missing_values, filling_values) + converters = [StringConverter(dt, + locked=True, + missing_values=miss, + default=fill) + for (dt, miss, fill) in zipit] + else: + # Set to a default converter (but w/ different missing values) + zipit = zip(missing_values, filling_values) + converters = [StringConverter(dtype, + locked=True, + missing_values=miss, + default=fill) + for (miss, fill) in zipit] + # Update the converters to use the user-defined ones + uc_update = [] + for (j, conv) in user_converters.items(): + # If the converter is specified by column names, + # use the index instead + if _is_string_like(j): + try: + j = names.index(j) + i = j + except ValueError: + continue + elif usecols: + try: + i = usecols.index(j) + except ValueError: + # Unused converter specified + continue + else: + i = j + # Find the value to test - first_line is not filtered by usecols: + if len(first_line): + testing_value = first_values[j] + else: + testing_value = None + if conv is bytes: + user_conv = asbytes + elif byte_converters: + # Converters may use decode to workaround numpy's old + # behavior, so encode the string again before passing + # to the user converter. + def tobytes_first(x, conv): + if type(x) is bytes: + return conv(x) + return conv(x.encode("latin1")) + user_conv = functools.partial(tobytes_first, conv=conv) + else: + user_conv = conv + converters[i].update(user_conv, locked=True, + testing_value=testing_value, + default=filling_values[i], + missing_values=missing_values[i],) + uc_update.append((i, user_conv)) + # Make sure we have the corrected keys in user_converters... + user_converters.update(uc_update) + + # Fixme: possible error as following variable never used. + # miss_chars = [_.missing_values for _ in converters] + + # Initialize the output lists ... + # ... rows + rows = [] + append_to_rows = rows.append + # ... masks + if usemask: + masks = [] + append_to_masks = masks.append + # ... invalid + invalid = [] + append_to_invalid = invalid.append + + # Parse each line + for (i, line) in enumerate(itertools.chain([first_line, ], fhd)): + values = split_line(line) + nbvalues = len(values) + # Skip an empty line + if nbvalues == 0: + continue + if usecols: + # Select only the columns we need + try: + values = [values[_] for _ in usecols] + except IndexError: + append_to_invalid((i + skip_header + 1, nbvalues)) + continue + elif nbvalues != nbcols: + append_to_invalid((i + skip_header + 1, nbvalues)) + continue + # Store the values + append_to_rows(tuple(values)) + if usemask: + append_to_masks(tuple([v.strip() in m + for (v, m) in zip(values, + missing_values)])) + if len(rows) == max_rows: + break + + # Upgrade the converters (if needed) + if dtype is None: + for (i, converter) in enumerate(converters): + current_column = [itemgetter(i)(_m) for _m in rows] + try: + converter.iterupgrade(current_column) + except ConverterLockError: + errmsg = "Converter #%i is locked and cannot be upgraded: " % i + current_column = map(itemgetter(i), rows) + for (j, value) in enumerate(current_column): + try: + converter.upgrade(value) + except (ConverterError, ValueError): + errmsg += "(occurred line #%i for value '%s')" + errmsg %= (j + 1 + skip_header, value) + raise ConverterError(errmsg) + + # Check that we don't have invalid values + nbinvalid = len(invalid) + if nbinvalid > 0: + nbrows = len(rows) + nbinvalid - skip_footer + # Construct the error message + template = " Line #%%i (got %%i columns instead of %i)" % nbcols + if skip_footer > 0: + nbinvalid_skipped = len([_ for _ in invalid + if _[0] > nbrows + skip_header]) + invalid = invalid[:nbinvalid - nbinvalid_skipped] + skip_footer -= nbinvalid_skipped +# +# nbrows -= skip_footer +# errmsg = [template % (i, nb) +# for (i, nb) in invalid if i < nbrows] +# else: + errmsg = [template % (i, nb) + for (i, nb) in invalid] + if len(errmsg): + errmsg.insert(0, "Some errors were detected !") + errmsg = "\n".join(errmsg) + # Raise an exception ? + if invalid_raise: + raise ValueError(errmsg) + # Issue a warning ? + else: + warnings.warn(errmsg, ConversionWarning, stacklevel=2) + + # Strip the last skip_footer data + if skip_footer > 0: + rows = rows[:-skip_footer] + if usemask: + masks = masks[:-skip_footer] + + # Convert each value according to the converter: + # We want to modify the list in place to avoid creating a new one... + if loose: + rows = list( + zip(*[[conv._loose_call(_r) for _r in map(itemgetter(i), rows)] + for (i, conv) in enumerate(converters)])) + else: + rows = list( + zip(*[[conv._strict_call(_r) for _r in map(itemgetter(i), rows)] + for (i, conv) in enumerate(converters)])) + + # Reset the dtype + data = rows + if dtype is None: + # Get the dtypes from the types of the converters + column_types = [conv.type for conv in converters] + # Find the columns with strings... + strcolidx = [i for (i, v) in enumerate(column_types) + if v == np.str_] + + if byte_converters and strcolidx: + # convert strings back to bytes for backward compatibility + warnings.warn( + "Reading unicode strings without specifying the encoding " + "argument is deprecated. Set the encoding, use None for the " + "system default.", + np.exceptions.VisibleDeprecationWarning, stacklevel=2) + + def encode_unicode_cols(row_tup): + row = list(row_tup) + for i in strcolidx: + row[i] = row[i].encode('latin1') + return tuple(row) + + try: + data = [encode_unicode_cols(r) for r in data] + except UnicodeEncodeError: + pass + else: + for i in strcolidx: + column_types[i] = np.bytes_ + + # Update string types to be the right length + sized_column_types = column_types[:] + for i, col_type in enumerate(column_types): + if np.issubdtype(col_type, np.character): + n_chars = max(len(row[i]) for row in data) + sized_column_types[i] = (col_type, n_chars) + + if names is None: + # If the dtype is uniform (before sizing strings) + base = { + c_type + for c, c_type in zip(converters, column_types) + if c._checked} + if len(base) == 1: + uniform_type, = base + (ddtype, mdtype) = (uniform_type, bool) + else: + ddtype = [(defaultfmt % i, dt) + for (i, dt) in enumerate(sized_column_types)] + if usemask: + mdtype = [(defaultfmt % i, bool) + for (i, dt) in enumerate(sized_column_types)] + else: + ddtype = list(zip(names, sized_column_types)) + mdtype = list(zip(names, [bool] * len(sized_column_types))) + output = np.array(data, dtype=ddtype) + if usemask: + outputmask = np.array(masks, dtype=mdtype) + else: + # Overwrite the initial dtype names if needed + if names and dtype.names is not None: + dtype.names = names + # Case 1. We have a structured type + if len(dtype_flat) > 1: + # Nested dtype, eg [('a', int), ('b', [('b0', int), ('b1', 'f4')])] + # First, create the array using a flattened dtype: + # [('a', int), ('b1', int), ('b2', float)] + # Then, view the array using the specified dtype. + if 'O' in (_.char for _ in dtype_flat): + if has_nested_fields(dtype): + raise NotImplementedError( + "Nested fields involving objects are not supported...") + else: + output = np.array(data, dtype=dtype) + else: + rows = np.array(data, dtype=[('', _) for _ in dtype_flat]) + output = rows.view(dtype) + # Now, process the rowmasks the same way + if usemask: + rowmasks = np.array( + masks, dtype=np.dtype([('', bool) for t in dtype_flat])) + # Construct the new dtype + mdtype = make_mask_descr(dtype) + outputmask = rowmasks.view(mdtype) + # Case #2. We have a basic dtype + else: + # We used some user-defined converters + if user_converters: + ishomogeneous = True + descr = [] + for i, ttype in enumerate([conv.type for conv in converters]): + # Keep the dtype of the current converter + if i in user_converters: + ishomogeneous &= (ttype == dtype.type) + if np.issubdtype(ttype, np.character): + ttype = (ttype, max(len(row[i]) for row in data)) + descr.append(('', ttype)) + else: + descr.append(('', dtype)) + # So we changed the dtype ? + if not ishomogeneous: + # We have more than one field + if len(descr) > 1: + dtype = np.dtype(descr) + # We have only one field: drop the name if not needed. + else: + dtype = np.dtype(ttype) + # + output = np.array(data, dtype) + if usemask: + if dtype.names is not None: + mdtype = [(_, bool) for _ in dtype.names] + else: + mdtype = bool + outputmask = np.array(masks, dtype=mdtype) + # Try to take care of the missing data we missed + names = output.dtype.names + if usemask and names: + for (name, conv) in zip(names, converters): + missing_values = [conv(_) for _ in conv.missing_values + if _ != ''] + for mval in missing_values: + outputmask[name] |= (output[name] == mval) + # Construct the final array + if usemask: + output = output.view(MaskedArray) + output._mask = outputmask + + output = _ensure_ndmin_ndarray(output, ndmin=ndmin) + + if unpack: + if names is None: + return output.T + elif len(names) == 1: + # squeeze single-name dtypes too + return output[names[0]] + else: + # For structured arrays with multiple fields, + # return an array for each field. + return [output[field] for field in names] + return output + + +_genfromtxt_with_like = array_function_dispatch()(genfromtxt) + + +def recfromtxt(fname, **kwargs): + """ + Load ASCII data from a file and return it in a record array. + + If ``usemask=False`` a standard `recarray` is returned, + if ``usemask=True`` a MaskedRecords array is returned. + + .. deprecated:: 2.0 + Use `numpy.genfromtxt` instead. + + Parameters + ---------- + fname, kwargs : For a description of input parameters, see `genfromtxt`. + + See Also + -------- + numpy.genfromtxt : generic function + + Notes + ----- + By default, `dtype` is None, which means that the data-type of the output + array will be determined from the data. + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`recfromtxt` is deprecated, " + "use `numpy.genfromtxt` instead." + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + kwargs.setdefault("dtype", None) + usemask = kwargs.get('usemask', False) + output = genfromtxt(fname, **kwargs) + if usemask: + from numpy.ma.mrecords import MaskedRecords + output = output.view(MaskedRecords) + else: + output = output.view(np.recarray) + return output + + +def recfromcsv(fname, **kwargs): + """ + Load ASCII data stored in a comma-separated file. + + The returned array is a record array (if ``usemask=False``, see + `recarray`) or a masked record array (if ``usemask=True``, + see `ma.mrecords.MaskedRecords`). + + .. deprecated:: 2.0 + Use `numpy.genfromtxt` with comma as `delimiter` instead. + + Parameters + ---------- + fname, kwargs : For a description of input parameters, see `genfromtxt`. + + See Also + -------- + numpy.genfromtxt : generic function to load ASCII data. + + Notes + ----- + By default, `dtype` is None, which means that the data-type of the output + array will be determined from the data. + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`recfromcsv` is deprecated, " + "use `numpy.genfromtxt` with comma as `delimiter` instead. " + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + # Set default kwargs for genfromtxt as relevant to csv import. + kwargs.setdefault("case_sensitive", "lower") + kwargs.setdefault("names", True) + kwargs.setdefault("delimiter", ",") + kwargs.setdefault("dtype", None) + output = genfromtxt(fname, **kwargs) + + usemask = kwargs.get("usemask", False) + if usemask: + from numpy.ma.mrecords import MaskedRecords + output = output.view(MaskedRecords) + else: + output = output.view(np.recarray) + return output diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..16d009524875c092a7bd209c6b4432301f947080 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_npyio_impl.pyi @@ -0,0 +1,285 @@ +import types +import zipfile +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence +from re import Pattern +from typing import IO, Any, ClassVar, Generic, Protocol, TypeAlias, overload, type_check_only +from typing import Literal as L + +from _typeshed import StrOrBytesPath, StrPath, SupportsKeysAndGetItem, SupportsRead, SupportsWrite +from typing_extensions import Self, TypeVar, deprecated, override + +import numpy as np +from numpy._core.multiarray import packbits, unpackbits +from numpy._typing import ArrayLike, DTypeLike, NDArray, _DTypeLike, _SupportsArrayFunc +from numpy.ma.mrecords import MaskedRecords + +from ._datasource import DataSource as DataSource + +__all__ = [ + "fromregex", + "genfromtxt", + "load", + "loadtxt", + "packbits", + "save", + "savetxt", + "savez", + "savez_compressed", + "unpackbits", +] + +_T_co = TypeVar("_T_co", covariant=True) +_SCT = TypeVar("_SCT", bound=np.generic) +_SCT_co = TypeVar("_SCT_co", bound=np.generic, default=Any, covariant=True) + +_FName: TypeAlias = StrPath | Iterable[str] | Iterable[bytes] +_FNameRead: TypeAlias = StrPath | SupportsRead[str] | SupportsRead[bytes] +_FNameWriteBytes: TypeAlias = StrPath | SupportsWrite[bytes] +_FNameWrite: TypeAlias = _FNameWriteBytes | SupportsWrite[str] + +@type_check_only +class _SupportsReadSeek(SupportsRead[_T_co], Protocol[_T_co]): + def seek(self, offset: int, whence: int, /) -> object: ... + +class BagObj(Generic[_T_co]): + def __init__(self, /, obj: SupportsKeysAndGetItem[str, _T_co]) -> None: ... + def __getattribute__(self, key: str, /) -> _T_co: ... + def __dir__(self) -> list[str]: ... + +class NpzFile(Mapping[str, NDArray[_SCT_co]]): + _MAX_REPR_ARRAY_COUNT: ClassVar[int] = 5 + + zip: zipfile.ZipFile + fid: IO[str] | None + files: list[str] + allow_pickle: bool + pickle_kwargs: Mapping[str, Any] | None + f: BagObj[NpzFile[_SCT_co]] + + # + def __init__( + self, + /, + fid: IO[Any], + own_fid: bool = False, + allow_pickle: bool = False, + pickle_kwargs: Mapping[str, object] | None = None, + *, + max_header_size: int = 10_000, + ) -> None: ... + def __del__(self) -> None: ... + def __enter__(self) -> Self: ... + def __exit__(self, cls: type[BaseException] | None, e: BaseException | None, tb: types.TracebackType | None, /) -> None: ... + @override + def __len__(self) -> int: ... + @override + def __iter__(self) -> Iterator[str]: ... + @override + def __getitem__(self, key: str, /) -> NDArray[_SCT_co]: ... + def close(self) -> None: ... + +# NOTE: Returns a `NpzFile` if file is a zip file; +# returns an `ndarray`/`memmap` otherwise +def load( + file: StrOrBytesPath | _SupportsReadSeek[bytes], + mmap_mode: L["r+", "r", "w+", "c"] | None = None, + allow_pickle: bool = False, + fix_imports: bool = True, + encoding: L["ASCII", "latin1", "bytes"] = "ASCII", + *, + max_header_size: int = 10_000, +) -> Any: ... + +@overload +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool = True) -> None: ... +@overload +@deprecated("The 'fix_imports' flag is deprecated in NumPy 2.1.") +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool, fix_imports: bool) -> None: ... +@overload +@deprecated("The 'fix_imports' flag is deprecated in NumPy 2.1.") +def save(file: _FNameWriteBytes, arr: ArrayLike, allow_pickle: bool = True, *, fix_imports: bool) -> None: ... + +# +def savez(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ... + +# +def savez_compressed(file: _FNameWriteBytes, *args: ArrayLike, allow_pickle: bool = True, **kwds: ArrayLike) -> None: ... + +# File-like objects only have to implement `__iter__` and, +# optionally, `encoding` +@overload +def loadtxt( + fname: _FName, + dtype: None = None, + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, + *, + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, +) -> NDArray[np.float64]: ... +@overload +def loadtxt( + fname: _FName, + dtype: _DTypeLike[_SCT], + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, + *, + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, +) -> NDArray[_SCT]: ... +@overload +def loadtxt( + fname: _FName, + dtype: DTypeLike, + comments: str | Sequence[str] | None = "#", + delimiter: str | None = None, + converters: Mapping[int | str, Callable[[str], Any]] | Callable[[str], Any] | None = None, + skiprows: int = 0, + usecols: int | Sequence[int] | None = None, + unpack: bool = False, + ndmin: L[0, 1, 2] = 0, + encoding: str | None = None, + max_rows: int | None = None, + *, + quotechar: str | None = None, + like: _SupportsArrayFunc | None = None, +) -> NDArray[Any]: ... + +def savetxt( + fname: _FNameWrite, + X: ArrayLike, + fmt: str | Sequence[str] = "%.18e", + delimiter: str = " ", + newline: str = "\n", + header: str = "", + footer: str = "", + comments: str = "# ", + encoding: str | None = None, +) -> None: ... + +@overload +def fromregex( + file: _FNameRead, + regexp: str | bytes | Pattern[Any], + dtype: _DTypeLike[_SCT], + encoding: str | None = None, +) -> NDArray[_SCT]: ... +@overload +def fromregex( + file: _FNameRead, + regexp: str | bytes | Pattern[Any], + dtype: DTypeLike, + encoding: str | None = None, +) -> NDArray[Any]: ... + +@overload +def genfromtxt( + fname: _FName, + dtype: None = None, + comments: str = ..., + delimiter: str | int | Iterable[int] | None = ..., + skip_header: int = ..., + skip_footer: int = ..., + converters: Mapping[int | str, Callable[[str], Any]] | None = ..., + missing_values: Any = ..., + filling_values: Any = ..., + usecols: Sequence[int] | None = ..., + names: L[True] | str | Collection[str] | None = ..., + excludelist: Sequence[str] | None = ..., + deletechars: str = ..., + replace_space: str = ..., + autostrip: bool = ..., + case_sensitive: bool | L["upper", "lower"] = ..., + defaultfmt: str = ..., + unpack: bool | None = ..., + usemask: bool = ..., + loose: bool = ..., + invalid_raise: bool = ..., + max_rows: int | None = ..., + encoding: str = ..., + *, + ndmin: L[0, 1, 2] = ..., + like: _SupportsArrayFunc | None = ..., +) -> NDArray[Any]: ... +@overload +def genfromtxt( + fname: _FName, + dtype: _DTypeLike[_SCT], + comments: str = ..., + delimiter: str | int | Iterable[int] | None = ..., + skip_header: int = ..., + skip_footer: int = ..., + converters: Mapping[int | str, Callable[[str], Any]] | None = ..., + missing_values: Any = ..., + filling_values: Any = ..., + usecols: Sequence[int] | None = ..., + names: L[True] | str | Collection[str] | None = ..., + excludelist: Sequence[str] | None = ..., + deletechars: str = ..., + replace_space: str = ..., + autostrip: bool = ..., + case_sensitive: bool | L["upper", "lower"] = ..., + defaultfmt: str = ..., + unpack: bool | None = ..., + usemask: bool = ..., + loose: bool = ..., + invalid_raise: bool = ..., + max_rows: int | None = ..., + encoding: str = ..., + *, + ndmin: L[0, 1, 2] = ..., + like: _SupportsArrayFunc | None = ..., +) -> NDArray[_SCT]: ... +@overload +def genfromtxt( + fname: _FName, + dtype: DTypeLike, + comments: str = ..., + delimiter: str | int | Iterable[int] | None = ..., + skip_header: int = ..., + skip_footer: int = ..., + converters: Mapping[int | str, Callable[[str], Any]] | None = ..., + missing_values: Any = ..., + filling_values: Any = ..., + usecols: Sequence[int] | None = ..., + names: L[True] | str | Collection[str] | None = ..., + excludelist: Sequence[str] | None = ..., + deletechars: str = ..., + replace_space: str = ..., + autostrip: bool = ..., + case_sensitive: bool | L["upper", "lower"] = ..., + defaultfmt: str = ..., + unpack: bool | None = ..., + usemask: bool = ..., + loose: bool = ..., + invalid_raise: bool = ..., + max_rows: int | None = ..., + encoding: str = ..., + *, + ndmin: L[0, 1, 2] = ..., + like: _SupportsArrayFunc | None = ..., +) -> NDArray[Any]: ... + +@overload +def recfromtxt(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ... +@overload +def recfromtxt(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ... + +@overload +def recfromcsv(fname: _FName, *, usemask: L[False] = False, **kwargs: object) -> np.recarray[Any, np.dtype[np.record]]: ... +@overload +def recfromcsv(fname: _FName, *, usemask: L[True], **kwargs: object) -> MaskedRecords[Any, np.dtype[np.void]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..9bcf0a3d92a6c8301b25e480f11994ff6feb1efa --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.py @@ -0,0 +1,1458 @@ +""" +Functions to operate on polynomials. + +""" +__all__ = ['poly', 'roots', 'polyint', 'polyder', 'polyadd', + 'polysub', 'polymul', 'polydiv', 'polyval', 'poly1d', + 'polyfit'] + +import functools +import re +import warnings + +from .._utils import set_module +import numpy._core.numeric as NX + +from numpy._core import (isscalar, abs, finfo, atleast_1d, hstack, dot, array, + ones) +from numpy._core import overrides +from numpy.exceptions import RankWarning +from numpy.lib._twodim_base_impl import diag, vander +from numpy.lib._function_base_impl import trim_zeros +from numpy.lib._type_check_impl import iscomplex, real, imag, mintypecode +from numpy.linalg import eigvals, lstsq, inv + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +def _poly_dispatcher(seq_of_zeros): + return seq_of_zeros + + +@array_function_dispatch(_poly_dispatcher) +def poly(seq_of_zeros): + """ + Find the coefficients of a polynomial with the given sequence of roots. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Returns the coefficients of the polynomial whose leading coefficient + is one for the given sequence of zeros (multiple roots must be included + in the sequence as many times as their multiplicity; see Examples). + A square matrix (or array, which will be treated as a matrix) can also + be given, in which case the coefficients of the characteristic polynomial + of the matrix are returned. + + Parameters + ---------- + seq_of_zeros : array_like, shape (N,) or (N, N) + A sequence of polynomial roots, or a square array or matrix object. + + Returns + ------- + c : ndarray + 1D array of polynomial coefficients from highest to lowest degree: + + ``c[0] * x**(N) + c[1] * x**(N-1) + ... + c[N-1] * x + c[N]`` + where c[0] always equals 1. + + Raises + ------ + ValueError + If input is the wrong shape (the input must be a 1-D or square + 2-D array). + + See Also + -------- + polyval : Compute polynomial values. + roots : Return the roots of a polynomial. + polyfit : Least squares polynomial fit. + poly1d : A one-dimensional polynomial class. + + Notes + ----- + Specifying the roots of a polynomial still leaves one degree of + freedom, typically represented by an undetermined leading + coefficient. [1]_ In the case of this function, that coefficient - + the first one in the returned array - is always taken as one. (If + for some reason you have one other point, the only automatic way + presently to leverage that information is to use ``polyfit``.) + + The characteristic polynomial, :math:`p_a(t)`, of an `n`-by-`n` + matrix **A** is given by + + :math:`p_a(t) = \\mathrm{det}(t\\, \\mathbf{I} - \\mathbf{A})`, + + where **I** is the `n`-by-`n` identity matrix. [2]_ + + References + ---------- + .. [1] M. Sullivan and M. Sullivan, III, "Algebra and Trigonometry, + Enhanced With Graphing Utilities," Prentice-Hall, pg. 318, 1996. + + .. [2] G. Strang, "Linear Algebra and Its Applications, 2nd Edition," + Academic Press, pg. 182, 1980. + + Examples + -------- + Given a sequence of a polynomial's zeros: + + >>> import numpy as np + + >>> np.poly((0, 0, 0)) # Multiple root example + array([1., 0., 0., 0.]) + + The line above represents z**3 + 0*z**2 + 0*z + 0. + + >>> np.poly((-1./2, 0, 1./2)) + array([ 1. , 0. , -0.25, 0. ]) + + The line above represents z**3 - z/4 + + >>> np.poly((np.random.random(1)[0], 0, np.random.random(1)[0])) + array([ 1. , -0.77086955, 0.08618131, 0. ]) # random + + Given a square array object: + + >>> P = np.array([[0, 1./3], [-1./2, 0]]) + >>> np.poly(P) + array([1. , 0. , 0.16666667]) + + Note how in all cases the leading coefficient is always 1. + + """ + seq_of_zeros = atleast_1d(seq_of_zeros) + sh = seq_of_zeros.shape + + if len(sh) == 2 and sh[0] == sh[1] and sh[0] != 0: + seq_of_zeros = eigvals(seq_of_zeros) + elif len(sh) == 1: + dt = seq_of_zeros.dtype + # Let object arrays slip through, e.g. for arbitrary precision + if dt != object: + seq_of_zeros = seq_of_zeros.astype(mintypecode(dt.char)) + else: + raise ValueError("input must be 1d or non-empty square 2d array.") + + if len(seq_of_zeros) == 0: + return 1.0 + dt = seq_of_zeros.dtype + a = ones((1,), dtype=dt) + for zero in seq_of_zeros: + a = NX.convolve(a, array([1, -zero], dtype=dt), mode='full') + + if issubclass(a.dtype.type, NX.complexfloating): + # if complex roots are all complex conjugates, the roots are real. + roots = NX.asarray(seq_of_zeros, complex) + if NX.all(NX.sort(roots) == NX.sort(roots.conjugate())): + a = a.real.copy() + + return a + + +def _roots_dispatcher(p): + return p + + +@array_function_dispatch(_roots_dispatcher) +def roots(p): + """ + Return the roots of a polynomial with coefficients given in p. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + The values in the rank-1 array `p` are coefficients of a polynomial. + If the length of `p` is n+1 then the polynomial is described by:: + + p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n] + + Parameters + ---------- + p : array_like + Rank-1 array of polynomial coefficients. + + Returns + ------- + out : ndarray + An array containing the roots of the polynomial. + + Raises + ------ + ValueError + When `p` cannot be converted to a rank-1 array. + + See also + -------- + poly : Find the coefficients of a polynomial with a given sequence + of roots. + polyval : Compute polynomial values. + polyfit : Least squares polynomial fit. + poly1d : A one-dimensional polynomial class. + + Notes + ----- + The algorithm relies on computing the eigenvalues of the + companion matrix [1]_. + + References + ---------- + .. [1] R. A. Horn & C. R. Johnson, *Matrix Analysis*. Cambridge, UK: + Cambridge University Press, 1999, pp. 146-7. + + Examples + -------- + >>> import numpy as np + >>> coeff = [3.2, 2, 1] + >>> np.roots(coeff) + array([-0.3125+0.46351241j, -0.3125-0.46351241j]) + + """ + # If input is scalar, this makes it an array + p = atleast_1d(p) + if p.ndim != 1: + raise ValueError("Input must be a rank-1 array.") + + # find non-zero array entries + non_zero = NX.nonzero(NX.ravel(p))[0] + + # Return an empty array if polynomial is all zeros + if len(non_zero) == 0: + return NX.array([]) + + # find the number of trailing zeros -- this is the number of roots at 0. + trailing_zeros = len(p) - non_zero[-1] - 1 + + # strip leading and trailing zeros + p = p[int(non_zero[0]):int(non_zero[-1])+1] + + # casting: if incoming array isn't floating point, make it floating point. + if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)): + p = p.astype(float) + + N = len(p) + if N > 1: + # build companion matrix and find its eigenvalues (the roots) + A = diag(NX.ones((N-2,), p.dtype), -1) + A[0,:] = -p[1:] / p[0] + roots = eigvals(A) + else: + roots = NX.array([]) + + # tack any zeros onto the back of the array + roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype))) + return roots + + +def _polyint_dispatcher(p, m=None, k=None): + return (p,) + + +@array_function_dispatch(_polyint_dispatcher) +def polyint(p, m=1, k=None): + """ + Return an antiderivative (indefinite integral) of a polynomial. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + The returned order `m` antiderivative `P` of polynomial `p` satisfies + :math:`\\frac{d^m}{dx^m}P(x) = p(x)` and is defined up to `m - 1` + integration constants `k`. The constants determine the low-order + polynomial part + + .. math:: \\frac{k_{m-1}}{0!} x^0 + \\ldots + \\frac{k_0}{(m-1)!}x^{m-1} + + of `P` so that :math:`P^{(j)}(0) = k_{m-j-1}`. + + Parameters + ---------- + p : array_like or poly1d + Polynomial to integrate. + A sequence is interpreted as polynomial coefficients, see `poly1d`. + m : int, optional + Order of the antiderivative. (Default: 1) + k : list of `m` scalars or scalar, optional + Integration constants. They are given in the order of integration: + those corresponding to highest-order terms come first. + + If ``None`` (default), all constants are assumed to be zero. + If `m = 1`, a single scalar can be given instead of a list. + + See Also + -------- + polyder : derivative of a polynomial + poly1d.integ : equivalent method + + Examples + -------- + The defining property of the antiderivative: + + >>> import numpy as np + + >>> p = np.poly1d([1,1,1]) + >>> P = np.polyint(p) + >>> P + poly1d([ 0.33333333, 0.5 , 1. , 0. ]) # may vary + >>> np.polyder(P) == p + True + + The integration constants default to zero, but can be specified: + + >>> P = np.polyint(p, 3) + >>> P(0) + 0.0 + >>> np.polyder(P)(0) + 0.0 + >>> np.polyder(P, 2)(0) + 0.0 + >>> P = np.polyint(p, 3, k=[6,5,3]) + >>> P + poly1d([ 0.01666667, 0.04166667, 0.16666667, 3. , 5. , 3. ]) # may vary + + Note that 3 = 6 / 2!, and that the constants are given in the order of + integrations. Constant of the highest-order polynomial term comes first: + + >>> np.polyder(P, 2)(0) + 6.0 + >>> np.polyder(P, 1)(0) + 5.0 + >>> P(0) + 3.0 + + """ + m = int(m) + if m < 0: + raise ValueError("Order of integral must be positive (see polyder)") + if k is None: + k = NX.zeros(m, float) + k = atleast_1d(k) + if len(k) == 1 and m > 1: + k = k[0]*NX.ones(m, float) + if len(k) < m: + raise ValueError( + "k must be a scalar or a rank-1 array of length 1 or >m.") + + truepoly = isinstance(p, poly1d) + p = NX.asarray(p) + if m == 0: + if truepoly: + return poly1d(p) + return p + else: + # Note: this must work also with object and integer arrays + y = NX.concatenate((p.__truediv__(NX.arange(len(p), 0, -1)), [k[0]])) + val = polyint(y, m - 1, k=k[1:]) + if truepoly: + return poly1d(val) + return val + + +def _polyder_dispatcher(p, m=None): + return (p,) + + +@array_function_dispatch(_polyder_dispatcher) +def polyder(p, m=1): + """ + Return the derivative of the specified order of a polynomial. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Parameters + ---------- + p : poly1d or sequence + Polynomial to differentiate. + A sequence is interpreted as polynomial coefficients, see `poly1d`. + m : int, optional + Order of differentiation (default: 1) + + Returns + ------- + der : poly1d + A new polynomial representing the derivative. + + See Also + -------- + polyint : Anti-derivative of a polynomial. + poly1d : Class for one-dimensional polynomials. + + Examples + -------- + The derivative of the polynomial :math:`x^3 + x^2 + x^1 + 1` is: + + >>> import numpy as np + + >>> p = np.poly1d([1,1,1,1]) + >>> p2 = np.polyder(p) + >>> p2 + poly1d([3, 2, 1]) + + which evaluates to: + + >>> p2(2.) + 17.0 + + We can verify this, approximating the derivative with + ``(f(x + h) - f(x))/h``: + + >>> (p(2. + 0.001) - p(2.)) / 0.001 + 17.007000999997857 + + The fourth-order derivative of a 3rd-order polynomial is zero: + + >>> np.polyder(p, 2) + poly1d([6, 2]) + >>> np.polyder(p, 3) + poly1d([6]) + >>> np.polyder(p, 4) + poly1d([0]) + + """ + m = int(m) + if m < 0: + raise ValueError("Order of derivative must be positive (see polyint)") + + truepoly = isinstance(p, poly1d) + p = NX.asarray(p) + n = len(p) - 1 + y = p[:-1] * NX.arange(n, 0, -1) + if m == 0: + val = p + else: + val = polyder(y, m - 1) + if truepoly: + val = poly1d(val) + return val + + +def _polyfit_dispatcher(x, y, deg, rcond=None, full=None, w=None, cov=None): + return (x, y, w) + + +@array_function_dispatch(_polyfit_dispatcher) +def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False): + """ + Least squares polynomial fit. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Fit a polynomial ``p(x) = p[0] * x**deg + ... + p[deg]`` of degree `deg` + to points `(x, y)`. Returns a vector of coefficients `p` that minimises + the squared error in the order `deg`, `deg-1`, ... `0`. + + The `Polynomial.fit ` class + method is recommended for new code as it is more stable numerically. See + the documentation of the method for more information. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int + Degree of the fitting polynomial + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (M,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + cov : bool or str, optional + If given and not `False`, return not just the estimate but also its + covariance matrix. By default, the covariance are scaled by + chi2/dof, where dof = M - (deg + 1), i.e., the weights are presumed + to be unreliable except in a relative sense and everything is scaled + such that the reduced chi2 is unity. This scaling is omitted if + ``cov='unscaled'``, as is relevant for the case that the weights are + w = 1/sigma, with sigma known to be a reliable estimate of the + uncertainty. + + Returns + ------- + p : ndarray, shape (deg + 1,) or (deg + 1, K) + Polynomial coefficients, highest power first. If `y` was 2-D, the + coefficients for `k`-th data set are in ``p[:,k]``. + + residuals, rank, singular_values, rcond + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the effective rank of the scaled Vandermonde + coefficient matrix + - singular_values -- singular values of the scaled Vandermonde + coefficient matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + V : ndarray, shape (deg + 1, deg + 1) or (deg + 1, deg + 1, K) + Present only if ``full == False`` and ``cov == True``. The covariance + matrix of the polynomial coefficient estimates. The diagonal of + this matrix are the variance estimates for each coefficient. If y + is a 2-D array, then the covariance matrix for the `k`-th data set + are in ``V[:,:,k]`` + + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. + + The warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + polyval : Compute polynomial values. + linalg.lstsq : Computes a least-squares fit. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution minimizes the squared error + + .. math:: + E = \\sum_{j=0}^k |p(x_j) - y_j|^2 + + in the equations:: + + x[0]**n * p[0] + ... + x[0] * p[n-1] + p[n] = y[0] + x[1]**n * p[0] + ... + x[1] * p[n-1] + p[n] = y[1] + ... + x[k]**n * p[0] + ... + x[k] * p[n-1] + p[n] = y[k] + + The coefficient matrix of the coefficients `p` is a Vandermonde matrix. + + `polyfit` issues a `~exceptions.RankWarning` when the least-squares fit is + badly conditioned. This implies that the best fit is not well-defined due + to numerical error. The results may be improved by lowering the polynomial + degree or by replacing `x` by `x` - `x`.mean(). The `rcond` parameter + can also be set to a value smaller than its default, but the resulting + fit may be spurious: including contributions from the small singular + values can add numerical noise to the result. + + Note that fitting polynomial coefficients is inherently badly conditioned + when the degree of the polynomial is large or the interval of sample points + is badly centered. The quality of the fit should always be checked in these + cases. When polynomial fits are not satisfactory, splines may be a good + alternative. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + .. [2] Wikipedia, "Polynomial interpolation", + https://en.wikipedia.org/wiki/Polynomial_interpolation + + Examples + -------- + >>> import numpy as np + >>> import warnings + >>> x = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + >>> y = np.array([0.0, 0.8, 0.9, 0.1, -0.8, -1.0]) + >>> z = np.polyfit(x, y, 3) + >>> z + array([ 0.08703704, -0.81349206, 1.69312169, -0.03968254]) # may vary + + It is convenient to use `poly1d` objects for dealing with polynomials: + + >>> p = np.poly1d(z) + >>> p(0.5) + 0.6143849206349179 # may vary + >>> p(3.5) + -0.34732142857143039 # may vary + >>> p(10) + 22.579365079365115 # may vary + + High-order polynomials may oscillate wildly: + + >>> with warnings.catch_warnings(): + ... warnings.simplefilter('ignore', np.exceptions.RankWarning) + ... p30 = np.poly1d(np.polyfit(x, y, 30)) + ... + >>> p30(4) + -0.80000000000000204 # may vary + >>> p30(5) + -0.99999999999999445 # may vary + >>> p30(4.5) + -0.10547061179440398 # may vary + + Illustration: + + >>> import matplotlib.pyplot as plt + >>> xp = np.linspace(-2, 6, 100) + >>> _ = plt.plot(x, y, '.', xp, p(xp), '-', xp, p30(xp), '--') + >>> plt.ylim(-2,2) + (-2, 2) + >>> plt.show() + + """ + order = int(deg) + 1 + x = NX.asarray(x) + 0.0 + y = NX.asarray(y) + 0.0 + + # check arguments. + if deg < 0: + raise ValueError("expected deg >= 0") + if x.ndim != 1: + raise TypeError("expected 1D vector for x") + if x.size == 0: + raise TypeError("expected non-empty vector for x") + if y.ndim < 1 or y.ndim > 2: + raise TypeError("expected 1D or 2D array for y") + if x.shape[0] != y.shape[0]: + raise TypeError("expected x and y to have same length") + + # set rcond + if rcond is None: + rcond = len(x)*finfo(x.dtype).eps + + # set up least squares equation for powers of x + lhs = vander(x, order) + rhs = y + + # apply weighting + if w is not None: + w = NX.asarray(w) + 0.0 + if w.ndim != 1: + raise TypeError("expected a 1-d array for weights") + if w.shape[0] != y.shape[0]: + raise TypeError("expected w and y to have the same length") + lhs *= w[:, NX.newaxis] + if rhs.ndim == 2: + rhs *= w[:, NX.newaxis] + else: + rhs *= w + + # scale lhs to improve condition number and solve + scale = NX.sqrt((lhs*lhs).sum(axis=0)) + lhs /= scale + c, resids, rank, s = lstsq(lhs, rhs, rcond) + c = (c.T/scale).T # broadcast scale coefficients + + # warn on rank reduction, which indicates an ill conditioned matrix + if rank != order and not full: + msg = "Polyfit may be poorly conditioned" + warnings.warn(msg, RankWarning, stacklevel=2) + + if full: + return c, resids, rank, s, rcond + elif cov: + Vbase = inv(dot(lhs.T, lhs)) + Vbase /= NX.outer(scale, scale) + if cov == "unscaled": + fac = 1 + else: + if len(x) <= order: + raise ValueError("the number of data points must exceed order " + "to scale the covariance matrix") + # note, this used to be: fac = resids / (len(x) - order - 2.0) + # it was decided that the "- 2" (originally justified by "Bayesian + # uncertainty analysis") is not what the user expects + # (see gh-11196 and gh-11197) + fac = resids / (len(x) - order) + if y.ndim == 1: + return c, Vbase * fac + else: + return c, Vbase[:,:, NX.newaxis] * fac + else: + return c + + +def _polyval_dispatcher(p, x): + return (p, x) + + +@array_function_dispatch(_polyval_dispatcher) +def polyval(p, x): + """ + Evaluate a polynomial at specific values. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + If `p` is of length N, this function returns the value:: + + p[0]*x**(N-1) + p[1]*x**(N-2) + ... + p[N-2]*x + p[N-1] + + If `x` is a sequence, then ``p(x)`` is returned for each element of ``x``. + If `x` is another polynomial then the composite polynomial ``p(x(t))`` + is returned. + + Parameters + ---------- + p : array_like or poly1d object + 1D array of polynomial coefficients (including coefficients equal + to zero) from highest degree to the constant term, or an + instance of poly1d. + x : array_like or poly1d object + A number, an array of numbers, or an instance of poly1d, at + which to evaluate `p`. + + Returns + ------- + values : ndarray or poly1d + If `x` is a poly1d instance, the result is the composition of the two + polynomials, i.e., `x` is "substituted" in `p` and the simplified + result is returned. In addition, the type of `x` - array_like or + poly1d - governs the type of the output: `x` array_like => `values` + array_like, `x` a poly1d object => `values` is also. + + See Also + -------- + poly1d: A polynomial class. + + Notes + ----- + Horner's scheme [1]_ is used to evaluate the polynomial. Even so, + for polynomials of high degree the values may be inaccurate due to + rounding errors. Use carefully. + + If `x` is a subtype of `ndarray` the return value will be of the same type. + + References + ---------- + .. [1] I. N. Bronshtein, K. A. Semendyayev, and K. A. Hirsch (Eng. + trans. Ed.), *Handbook of Mathematics*, New York, Van Nostrand + Reinhold Co., 1985, pg. 720. + + Examples + -------- + >>> import numpy as np + >>> np.polyval([3,0,1], 5) # 3 * 5**2 + 0 * 5**1 + 1 + 76 + >>> np.polyval([3,0,1], np.poly1d(5)) + poly1d([76]) + >>> np.polyval(np.poly1d([3,0,1]), 5) + 76 + >>> np.polyval(np.poly1d([3,0,1]), np.poly1d(5)) + poly1d([76]) + + """ + p = NX.asarray(p) + if isinstance(x, poly1d): + y = 0 + else: + x = NX.asanyarray(x) + y = NX.zeros_like(x) + for pv in p: + y = y * x + pv + return y + + +def _binary_op_dispatcher(a1, a2): + return (a1, a2) + + +@array_function_dispatch(_binary_op_dispatcher) +def polyadd(a1, a2): + """ + Find the sum of two polynomials. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Returns the polynomial resulting from the sum of two input polynomials. + Each input must be either a poly1d object or a 1D sequence of polynomial + coefficients, from highest to lowest degree. + + Parameters + ---------- + a1, a2 : array_like or poly1d object + Input polynomials. + + Returns + ------- + out : ndarray or poly1d object + The sum of the inputs. If either input is a poly1d object, then the + output is also a poly1d object. Otherwise, it is a 1D array of + polynomial coefficients from highest to lowest degree. + + See Also + -------- + poly1d : A one-dimensional polynomial class. + poly, polyadd, polyder, polydiv, polyfit, polyint, polysub, polyval + + Examples + -------- + >>> import numpy as np + >>> np.polyadd([1, 2], [9, 5, 4]) + array([9, 6, 6]) + + Using poly1d objects: + + >>> p1 = np.poly1d([1, 2]) + >>> p2 = np.poly1d([9, 5, 4]) + >>> print(p1) + 1 x + 2 + >>> print(p2) + 2 + 9 x + 5 x + 4 + >>> print(np.polyadd(p1, p2)) + 2 + 9 x + 6 x + 6 + + """ + truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d)) + a1 = atleast_1d(a1) + a2 = atleast_1d(a2) + diff = len(a2) - len(a1) + if diff == 0: + val = a1 + a2 + elif diff > 0: + zr = NX.zeros(diff, a1.dtype) + val = NX.concatenate((zr, a1)) + a2 + else: + zr = NX.zeros(abs(diff), a2.dtype) + val = a1 + NX.concatenate((zr, a2)) + if truepoly: + val = poly1d(val) + return val + + +@array_function_dispatch(_binary_op_dispatcher) +def polysub(a1, a2): + """ + Difference (subtraction) of two polynomials. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Given two polynomials `a1` and `a2`, returns ``a1 - a2``. + `a1` and `a2` can be either array_like sequences of the polynomials' + coefficients (including coefficients equal to zero), or `poly1d` objects. + + Parameters + ---------- + a1, a2 : array_like or poly1d + Minuend and subtrahend polynomials, respectively. + + Returns + ------- + out : ndarray or poly1d + Array or `poly1d` object of the difference polynomial's coefficients. + + See Also + -------- + polyval, polydiv, polymul, polyadd + + Examples + -------- + .. math:: (2 x^2 + 10 x - 2) - (3 x^2 + 10 x -4) = (-x^2 + 2) + + >>> import numpy as np + + >>> np.polysub([2, 10, -2], [3, 10, -4]) + array([-1, 0, 2]) + + """ + truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d)) + a1 = atleast_1d(a1) + a2 = atleast_1d(a2) + diff = len(a2) - len(a1) + if diff == 0: + val = a1 - a2 + elif diff > 0: + zr = NX.zeros(diff, a1.dtype) + val = NX.concatenate((zr, a1)) - a2 + else: + zr = NX.zeros(abs(diff), a2.dtype) + val = a1 - NX.concatenate((zr, a2)) + if truepoly: + val = poly1d(val) + return val + + +@array_function_dispatch(_binary_op_dispatcher) +def polymul(a1, a2): + """ + Find the product of two polynomials. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + Finds the polynomial resulting from the multiplication of the two input + polynomials. Each input must be either a poly1d object or a 1D sequence + of polynomial coefficients, from highest to lowest degree. + + Parameters + ---------- + a1, a2 : array_like or poly1d object + Input polynomials. + + Returns + ------- + out : ndarray or poly1d object + The polynomial resulting from the multiplication of the inputs. If + either inputs is a poly1d object, then the output is also a poly1d + object. Otherwise, it is a 1D array of polynomial coefficients from + highest to lowest degree. + + See Also + -------- + poly1d : A one-dimensional polynomial class. + poly, polyadd, polyder, polydiv, polyfit, polyint, polysub, polyval + convolve : Array convolution. Same output as polymul, but has parameter + for overlap mode. + + Examples + -------- + >>> import numpy as np + >>> np.polymul([1, 2, 3], [9, 5, 1]) + array([ 9, 23, 38, 17, 3]) + + Using poly1d objects: + + >>> p1 = np.poly1d([1, 2, 3]) + >>> p2 = np.poly1d([9, 5, 1]) + >>> print(p1) + 2 + 1 x + 2 x + 3 + >>> print(p2) + 2 + 9 x + 5 x + 1 + >>> print(np.polymul(p1, p2)) + 4 3 2 + 9 x + 23 x + 38 x + 17 x + 3 + + """ + truepoly = (isinstance(a1, poly1d) or isinstance(a2, poly1d)) + a1, a2 = poly1d(a1), poly1d(a2) + val = NX.convolve(a1, a2) + if truepoly: + val = poly1d(val) + return val + + +def _polydiv_dispatcher(u, v): + return (u, v) + + +@array_function_dispatch(_polydiv_dispatcher) +def polydiv(u, v): + """ + Returns the quotient and remainder of polynomial division. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + The input arrays are the coefficients (including any coefficients + equal to zero) of the "numerator" (dividend) and "denominator" + (divisor) polynomials, respectively. + + Parameters + ---------- + u : array_like or poly1d + Dividend polynomial's coefficients. + + v : array_like or poly1d + Divisor polynomial's coefficients. + + Returns + ------- + q : ndarray + Coefficients, including those equal to zero, of the quotient. + r : ndarray + Coefficients, including those equal to zero, of the remainder. + + See Also + -------- + poly, polyadd, polyder, polydiv, polyfit, polyint, polymul, polysub + polyval + + Notes + ----- + Both `u` and `v` must be 0-d or 1-d (ndim = 0 or 1), but `u.ndim` need + not equal `v.ndim`. In other words, all four possible combinations - + ``u.ndim = v.ndim = 0``, ``u.ndim = v.ndim = 1``, + ``u.ndim = 1, v.ndim = 0``, and ``u.ndim = 0, v.ndim = 1`` - work. + + Examples + -------- + .. math:: \\frac{3x^2 + 5x + 2}{2x + 1} = 1.5x + 1.75, remainder 0.25 + + >>> import numpy as np + >>> x = np.array([3.0, 5.0, 2.0]) + >>> y = np.array([2.0, 1.0]) + >>> np.polydiv(x, y) + (array([1.5 , 1.75]), array([0.25])) + + """ + truepoly = (isinstance(u, poly1d) or isinstance(v, poly1d)) + u = atleast_1d(u) + 0.0 + v = atleast_1d(v) + 0.0 + # w has the common type + w = u[0] + v[0] + m = len(u) - 1 + n = len(v) - 1 + scale = 1. / v[0] + q = NX.zeros((max(m - n + 1, 1),), w.dtype) + r = u.astype(w.dtype) + for k in range(0, m-n+1): + d = scale * r[k] + q[k] = d + r[k:k+n+1] -= d*v + while NX.allclose(r[0], 0, rtol=1e-14) and (r.shape[-1] > 1): + r = r[1:] + if truepoly: + return poly1d(q), poly1d(r) + return q, r + +_poly_mat = re.compile(r"\*\*([0-9]*)") +def _raise_power(astr, wrap=70): + n = 0 + line1 = '' + line2 = '' + output = ' ' + while True: + mat = _poly_mat.search(astr, n) + if mat is None: + break + span = mat.span() + power = mat.groups()[0] + partstr = astr[n:span[0]] + n = span[1] + toadd2 = partstr + ' '*(len(power)-1) + toadd1 = ' '*(len(partstr)-1) + power + if ((len(line2) + len(toadd2) > wrap) or + (len(line1) + len(toadd1) > wrap)): + output += line1 + "\n" + line2 + "\n " + line1 = toadd1 + line2 = toadd2 + else: + line2 += partstr + ' '*(len(power)-1) + line1 += ' '*(len(partstr)-1) + power + output += line1 + "\n" + line2 + return output + astr[n:] + + +@set_module('numpy') +class poly1d: + """ + A one-dimensional polynomial class. + + .. note:: + This forms part of the old polynomial API. Since version 1.4, the + new polynomial API defined in `numpy.polynomial` is preferred. + A summary of the differences can be found in the + :doc:`transition guide `. + + A convenience class, used to encapsulate "natural" operations on + polynomials so that said operations may take on their customary + form in code (see Examples). + + Parameters + ---------- + c_or_r : array_like + The polynomial's coefficients, in decreasing powers, or if + the value of the second parameter is True, the polynomial's + roots (values where the polynomial evaluates to 0). For example, + ``poly1d([1, 2, 3])`` returns an object that represents + :math:`x^2 + 2x + 3`, whereas ``poly1d([1, 2, 3], True)`` returns + one that represents :math:`(x-1)(x-2)(x-3) = x^3 - 6x^2 + 11x -6`. + r : bool, optional + If True, `c_or_r` specifies the polynomial's roots; the default + is False. + variable : str, optional + Changes the variable used when printing `p` from `x` to `variable` + (see Examples). + + Examples + -------- + Construct the polynomial :math:`x^2 + 2x + 3`: + + >>> import numpy as np + + >>> p = np.poly1d([1, 2, 3]) + >>> print(np.poly1d(p)) + 2 + 1 x + 2 x + 3 + + Evaluate the polynomial at :math:`x = 0.5`: + + >>> p(0.5) + 4.25 + + Find the roots: + + >>> p.r + array([-1.+1.41421356j, -1.-1.41421356j]) + >>> p(p.r) + array([ -4.44089210e-16+0.j, -4.44089210e-16+0.j]) # may vary + + These numbers in the previous line represent (0, 0) to machine precision + + Show the coefficients: + + >>> p.c + array([1, 2, 3]) + + Display the order (the leading zero-coefficients are removed): + + >>> p.order + 2 + + Show the coefficient of the k-th power in the polynomial + (which is equivalent to ``p.c[-(i+1)]``): + + >>> p[1] + 2 + + Polynomials can be added, subtracted, multiplied, and divided + (returns quotient and remainder): + + >>> p * p + poly1d([ 1, 4, 10, 12, 9]) + + >>> (p**3 + 4) / p + (poly1d([ 1., 4., 10., 12., 9.]), poly1d([4.])) + + ``asarray(p)`` gives the coefficient array, so polynomials can be + used in all functions that accept arrays: + + >>> p**2 # square of polynomial + poly1d([ 1, 4, 10, 12, 9]) + + >>> np.square(p) # square of individual coefficients + array([1, 4, 9]) + + The variable used in the string representation of `p` can be modified, + using the `variable` parameter: + + >>> p = np.poly1d([1,2,3], variable='z') + >>> print(p) + 2 + 1 z + 2 z + 3 + + Construct a polynomial from its roots: + + >>> np.poly1d([1, 2], True) + poly1d([ 1., -3., 2.]) + + This is the same polynomial as obtained by: + + >>> np.poly1d([1, -1]) * np.poly1d([1, -2]) + poly1d([ 1, -3, 2]) + + """ + __hash__ = None + + @property + def coeffs(self): + """ The polynomial coefficients """ + return self._coeffs + + @coeffs.setter + def coeffs(self, value): + # allowing this makes p.coeffs *= 2 legal + if value is not self._coeffs: + raise AttributeError("Cannot set attribute") + + @property + def variable(self): + """ The name of the polynomial variable """ + return self._variable + + # calculated attributes + @property + def order(self): + """ The order or degree of the polynomial """ + return len(self._coeffs) - 1 + + @property + def roots(self): + """ The roots of the polynomial, where self(x) == 0 """ + return roots(self._coeffs) + + # our internal _coeffs property need to be backed by __dict__['coeffs'] for + # scipy to work correctly. + @property + def _coeffs(self): + return self.__dict__['coeffs'] + @_coeffs.setter + def _coeffs(self, coeffs): + self.__dict__['coeffs'] = coeffs + + # alias attributes + r = roots + c = coef = coefficients = coeffs + o = order + + def __init__(self, c_or_r, r=False, variable=None): + if isinstance(c_or_r, poly1d): + self._variable = c_or_r._variable + self._coeffs = c_or_r._coeffs + + if set(c_or_r.__dict__) - set(self.__dict__): + msg = ("In the future extra properties will not be copied " + "across when constructing one poly1d from another") + warnings.warn(msg, FutureWarning, stacklevel=2) + self.__dict__.update(c_or_r.__dict__) + + if variable is not None: + self._variable = variable + return + if r: + c_or_r = poly(c_or_r) + c_or_r = atleast_1d(c_or_r) + if c_or_r.ndim > 1: + raise ValueError("Polynomial must be 1d only.") + c_or_r = trim_zeros(c_or_r, trim='f') + if len(c_or_r) == 0: + c_or_r = NX.array([0], dtype=c_or_r.dtype) + self._coeffs = c_or_r + if variable is None: + variable = 'x' + self._variable = variable + + def __array__(self, t=None, copy=None): + if t: + return NX.asarray(self.coeffs, t, copy=copy) + else: + return NX.asarray(self.coeffs, copy=copy) + + def __repr__(self): + vals = repr(self.coeffs) + vals = vals[6:-1] + return "poly1d(%s)" % vals + + def __len__(self): + return self.order + + def __str__(self): + thestr = "0" + var = self.variable + + # Remove leading zeros + coeffs = self.coeffs[NX.logical_or.accumulate(self.coeffs != 0)] + N = len(coeffs)-1 + + def fmt_float(q): + s = '%.4g' % q + if s.endswith('.0000'): + s = s[:-5] + return s + + for k, coeff in enumerate(coeffs): + if not iscomplex(coeff): + coefstr = fmt_float(real(coeff)) + elif real(coeff) == 0: + coefstr = '%sj' % fmt_float(imag(coeff)) + else: + coefstr = '(%s + %sj)' % (fmt_float(real(coeff)), + fmt_float(imag(coeff))) + + power = (N-k) + if power == 0: + if coefstr != '0': + newstr = '%s' % (coefstr,) + else: + if k == 0: + newstr = '0' + else: + newstr = '' + elif power == 1: + if coefstr == '0': + newstr = '' + elif coefstr == 'b': + newstr = var + else: + newstr = '%s %s' % (coefstr, var) + else: + if coefstr == '0': + newstr = '' + elif coefstr == 'b': + newstr = '%s**%d' % (var, power,) + else: + newstr = '%s %s**%d' % (coefstr, var, power) + + if k > 0: + if newstr != '': + if newstr.startswith('-'): + thestr = "%s - %s" % (thestr, newstr[1:]) + else: + thestr = "%s + %s" % (thestr, newstr) + else: + thestr = newstr + return _raise_power(thestr) + + def __call__(self, val): + return polyval(self.coeffs, val) + + def __neg__(self): + return poly1d(-self.coeffs) + + def __pos__(self): + return self + + def __mul__(self, other): + if isscalar(other): + return poly1d(self.coeffs * other) + else: + other = poly1d(other) + return poly1d(polymul(self.coeffs, other.coeffs)) + + def __rmul__(self, other): + if isscalar(other): + return poly1d(other * self.coeffs) + else: + other = poly1d(other) + return poly1d(polymul(self.coeffs, other.coeffs)) + + def __add__(self, other): + other = poly1d(other) + return poly1d(polyadd(self.coeffs, other.coeffs)) + + def __radd__(self, other): + other = poly1d(other) + return poly1d(polyadd(self.coeffs, other.coeffs)) + + def __pow__(self, val): + if not isscalar(val) or int(val) != val or val < 0: + raise ValueError("Power to non-negative integers only.") + res = [1] + for _ in range(val): + res = polymul(self.coeffs, res) + return poly1d(res) + + def __sub__(self, other): + other = poly1d(other) + return poly1d(polysub(self.coeffs, other.coeffs)) + + def __rsub__(self, other): + other = poly1d(other) + return poly1d(polysub(other.coeffs, self.coeffs)) + + def __div__(self, other): + if isscalar(other): + return poly1d(self.coeffs/other) + else: + other = poly1d(other) + return polydiv(self, other) + + __truediv__ = __div__ + + def __rdiv__(self, other): + if isscalar(other): + return poly1d(other/self.coeffs) + else: + other = poly1d(other) + return polydiv(other, self) + + __rtruediv__ = __rdiv__ + + def __eq__(self, other): + if not isinstance(other, poly1d): + return NotImplemented + if self.coeffs.shape != other.coeffs.shape: + return False + return (self.coeffs == other.coeffs).all() + + def __ne__(self, other): + if not isinstance(other, poly1d): + return NotImplemented + return not self.__eq__(other) + + + def __getitem__(self, val): + ind = self.order - val + if val > self.order: + return self.coeffs.dtype.type(0) + if val < 0: + return self.coeffs.dtype.type(0) + return self.coeffs[ind] + + def __setitem__(self, key, val): + ind = self.order - key + if key < 0: + raise ValueError("Does not support negative powers.") + if key > self.order: + zr = NX.zeros(key-self.order, self.coeffs.dtype) + self._coeffs = NX.concatenate((zr, self.coeffs)) + ind = 0 + self._coeffs[ind] = val + return + + def __iter__(self): + return iter(self.coeffs) + + def integ(self, m=1, k=0): + """ + Return an antiderivative (indefinite integral) of this polynomial. + + Refer to `polyint` for full documentation. + + See Also + -------- + polyint : equivalent function + + """ + return poly1d(polyint(self.coeffs, m=m, k=k)) + + def deriv(self, m=1): + """ + Return a derivative of this polynomial. + + Refer to `polyder` for full documentation. + + See Also + -------- + polyder : equivalent function + + """ + return poly1d(polyder(self.coeffs, m=m)) + +# Stuff to do on module import + +warnings.simplefilter('always', RankWarning) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..112ec33d2520426fdad8a5c034184c268fe68ea1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_polynomial_impl.pyi @@ -0,0 +1,315 @@ +from typing import ( + Literal as L, + TypeAlias, + overload, + Any, + SupportsInt, + SupportsIndex, + TypeVar, + NoReturn, +) + +import numpy as np +from numpy import ( + poly1d, + unsignedinteger, + signedinteger, + floating, + complexfloating, + int32, + int64, + float64, + complex128, + object_, +) + +from numpy._typing import ( + NDArray, + ArrayLike, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, +) + +_T = TypeVar("_T") + +_2Tup: TypeAlias = tuple[_T, _T] +_5Tup: TypeAlias = tuple[ + _T, + NDArray[float64], + NDArray[int32], + NDArray[float64], + NDArray[float64], +] + +__all__ = [ + "poly", + "roots", + "polyint", + "polyder", + "polyadd", + "polysub", + "polymul", + "polydiv", + "polyval", + "poly1d", + "polyfit", +] + +def poly(seq_of_zeros: ArrayLike) -> NDArray[floating[Any]]: ... + +# Returns either a float or complex array depending on the input values. +# See `np.linalg.eigvals`. +def roots(p: ArrayLike) -> NDArray[complexfloating[Any, Any]] | NDArray[floating[Any]]: ... + +@overload +def polyint( + p: poly1d, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeComplex_co | _ArrayLikeObject_co = ..., +) -> poly1d: ... +@overload +def polyint( + p: _ArrayLikeFloat_co, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeFloat_co = ..., +) -> NDArray[floating[Any]]: ... +@overload +def polyint( + p: _ArrayLikeComplex_co, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeComplex_co = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def polyint( + p: _ArrayLikeObject_co, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeObject_co = ..., +) -> NDArray[object_]: ... + +@overload +def polyder( + p: poly1d, + m: SupportsInt | SupportsIndex = ..., +) -> poly1d: ... +@overload +def polyder( + p: _ArrayLikeFloat_co, + m: SupportsInt | SupportsIndex = ..., +) -> NDArray[floating[Any]]: ... +@overload +def polyder( + p: _ArrayLikeComplex_co, + m: SupportsInt | SupportsIndex = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def polyder( + p: _ArrayLikeObject_co, + m: SupportsInt | SupportsIndex = ..., +) -> NDArray[object_]: ... + +@overload +def polyfit( + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[False] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: L[False] = ..., +) -> NDArray[float64]: ... +@overload +def polyfit( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[False] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: L[False] = ..., +) -> NDArray[complex128]: ... +@overload +def polyfit( + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[False] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: L[True, "unscaled"] = ..., +) -> _2Tup[NDArray[float64]]: ... +@overload +def polyfit( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[False] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: L[True, "unscaled"] = ..., +) -> _2Tup[NDArray[complex128]]: ... +@overload +def polyfit( + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[True] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: bool | L["unscaled"] = ..., +) -> _5Tup[NDArray[float64]]: ... +@overload +def polyfit( + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: SupportsIndex | SupportsInt, + rcond: None | float = ..., + full: L[True] = ..., + w: None | _ArrayLikeFloat_co = ..., + cov: bool | L["unscaled"] = ..., +) -> _5Tup[NDArray[complex128]]: ... + +@overload +def polyval( + p: _ArrayLikeBool_co, + x: _ArrayLikeBool_co, +) -> NDArray[int64]: ... +@overload +def polyval( + p: _ArrayLikeUInt_co, + x: _ArrayLikeUInt_co, +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def polyval( + p: _ArrayLikeInt_co, + x: _ArrayLikeInt_co, +) -> NDArray[signedinteger[Any]]: ... +@overload +def polyval( + p: _ArrayLikeFloat_co, + x: _ArrayLikeFloat_co, +) -> NDArray[floating[Any]]: ... +@overload +def polyval( + p: _ArrayLikeComplex_co, + x: _ArrayLikeComplex_co, +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def polyval( + p: _ArrayLikeObject_co, + x: _ArrayLikeObject_co, +) -> NDArray[object_]: ... + +@overload +def polyadd( + a1: poly1d, + a2: _ArrayLikeComplex_co | _ArrayLikeObject_co, +) -> poly1d: ... +@overload +def polyadd( + a1: _ArrayLikeComplex_co | _ArrayLikeObject_co, + a2: poly1d, +) -> poly1d: ... +@overload +def polyadd( + a1: _ArrayLikeBool_co, + a2: _ArrayLikeBool_co, +) -> NDArray[np.bool]: ... +@overload +def polyadd( + a1: _ArrayLikeUInt_co, + a2: _ArrayLikeUInt_co, +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def polyadd( + a1: _ArrayLikeInt_co, + a2: _ArrayLikeInt_co, +) -> NDArray[signedinteger[Any]]: ... +@overload +def polyadd( + a1: _ArrayLikeFloat_co, + a2: _ArrayLikeFloat_co, +) -> NDArray[floating[Any]]: ... +@overload +def polyadd( + a1: _ArrayLikeComplex_co, + a2: _ArrayLikeComplex_co, +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def polyadd( + a1: _ArrayLikeObject_co, + a2: _ArrayLikeObject_co, +) -> NDArray[object_]: ... + +@overload +def polysub( + a1: poly1d, + a2: _ArrayLikeComplex_co | _ArrayLikeObject_co, +) -> poly1d: ... +@overload +def polysub( + a1: _ArrayLikeComplex_co | _ArrayLikeObject_co, + a2: poly1d, +) -> poly1d: ... +@overload +def polysub( + a1: _ArrayLikeBool_co, + a2: _ArrayLikeBool_co, +) -> NoReturn: ... +@overload +def polysub( + a1: _ArrayLikeUInt_co, + a2: _ArrayLikeUInt_co, +) -> NDArray[unsignedinteger[Any]]: ... +@overload +def polysub( + a1: _ArrayLikeInt_co, + a2: _ArrayLikeInt_co, +) -> NDArray[signedinteger[Any]]: ... +@overload +def polysub( + a1: _ArrayLikeFloat_co, + a2: _ArrayLikeFloat_co, +) -> NDArray[floating[Any]]: ... +@overload +def polysub( + a1: _ArrayLikeComplex_co, + a2: _ArrayLikeComplex_co, +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def polysub( + a1: _ArrayLikeObject_co, + a2: _ArrayLikeObject_co, +) -> NDArray[object_]: ... + +# NOTE: Not an alias, but they do have the same signature (that we can reuse) +polymul = polyadd + +@overload +def polydiv( + u: poly1d, + v: _ArrayLikeComplex_co | _ArrayLikeObject_co, +) -> _2Tup[poly1d]: ... +@overload +def polydiv( + u: _ArrayLikeComplex_co | _ArrayLikeObject_co, + v: poly1d, +) -> _2Tup[poly1d]: ... +@overload +def polydiv( + u: _ArrayLikeFloat_co, + v: _ArrayLikeFloat_co, +) -> _2Tup[NDArray[floating[Any]]]: ... +@overload +def polydiv( + u: _ArrayLikeComplex_co, + v: _ArrayLikeComplex_co, +) -> _2Tup[NDArray[complexfloating[Any, Any]]]: ... +@overload +def polydiv( + u: _ArrayLikeObject_co, + v: _ArrayLikeObject_co, +) -> _2Tup[NDArray[Any]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..68e9cd2d5337ea8f7965b27e06f863a4e0cba7a6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.py @@ -0,0 +1,643 @@ +""" +Wrapper functions to more user-friendly calling of certain math functions +whose output data-type is different than the input data-type in certain +domains of the input. + +For example, for functions like `log` with branch cuts, the versions in this +module provide the mathematically valid answers in the complex plane:: + + >>> import math + >>> np.emath.log(-math.exp(1)) == (1+1j*math.pi) + True + +Similarly, `sqrt`, other base logarithms, `power` and trig functions are +correctly handled. See their respective docstrings for specific examples. + +""" +import numpy._core.numeric as nx +import numpy._core.numerictypes as nt +from numpy._core.numeric import asarray, any +from numpy._core.overrides import array_function_dispatch, set_module +from numpy.lib._type_check_impl import isreal + + +__all__ = [ + 'sqrt', 'log', 'log2', 'logn', 'log10', 'power', 'arccos', 'arcsin', + 'arctanh' + ] + + +_ln2 = nx.log(2.0) + + +def _tocomplex(arr): + """Convert its input `arr` to a complex array. + + The input is returned as a complex array of the smallest type that will fit + the original data: types like single, byte, short, etc. become csingle, + while others become cdouble. + + A copy of the input is always made. + + Parameters + ---------- + arr : array + + Returns + ------- + array + An array with the same input data as the input but in complex form. + + Examples + -------- + >>> import numpy as np + + First, consider an input of type short: + + >>> a = np.array([1,2,3],np.short) + + >>> ac = np.lib.scimath._tocomplex(a); ac + array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64) + + >>> ac.dtype + dtype('complex64') + + If the input is of type double, the output is correspondingly of the + complex double type as well: + + >>> b = np.array([1,2,3],np.double) + + >>> bc = np.lib.scimath._tocomplex(b); bc + array([1.+0.j, 2.+0.j, 3.+0.j]) + + >>> bc.dtype + dtype('complex128') + + Note that even if the input was complex to begin with, a copy is still + made, since the astype() method always copies: + + >>> c = np.array([1,2,3],np.csingle) + + >>> cc = np.lib.scimath._tocomplex(c); cc + array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64) + + >>> c *= 2; c + array([2.+0.j, 4.+0.j, 6.+0.j], dtype=complex64) + + >>> cc + array([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64) + """ + if issubclass(arr.dtype.type, (nt.single, nt.byte, nt.short, nt.ubyte, + nt.ushort, nt.csingle)): + return arr.astype(nt.csingle) + else: + return arr.astype(nt.cdouble) + + +def _fix_real_lt_zero(x): + """Convert `x` to complex if it has real, negative components. + + Otherwise, output is just the array version of the input (via asarray). + + Parameters + ---------- + x : array_like + + Returns + ------- + array + + Examples + -------- + >>> import numpy as np + >>> np.lib.scimath._fix_real_lt_zero([1,2]) + array([1, 2]) + + >>> np.lib.scimath._fix_real_lt_zero([-1,2]) + array([-1.+0.j, 2.+0.j]) + + """ + x = asarray(x) + if any(isreal(x) & (x < 0)): + x = _tocomplex(x) + return x + + +def _fix_int_lt_zero(x): + """Convert `x` to double if it has real, negative components. + + Otherwise, output is just the array version of the input (via asarray). + + Parameters + ---------- + x : array_like + + Returns + ------- + array + + Examples + -------- + >>> import numpy as np + >>> np.lib.scimath._fix_int_lt_zero([1,2]) + array([1, 2]) + + >>> np.lib.scimath._fix_int_lt_zero([-1,2]) + array([-1., 2.]) + """ + x = asarray(x) + if any(isreal(x) & (x < 0)): + x = x * 1.0 + return x + + +def _fix_real_abs_gt_1(x): + """Convert `x` to complex if it has real components x_i with abs(x_i)>1. + + Otherwise, output is just the array version of the input (via asarray). + + Parameters + ---------- + x : array_like + + Returns + ------- + array + + Examples + -------- + >>> import numpy as np + >>> np.lib.scimath._fix_real_abs_gt_1([0,1]) + array([0, 1]) + + >>> np.lib.scimath._fix_real_abs_gt_1([0,2]) + array([0.+0.j, 2.+0.j]) + """ + x = asarray(x) + if any(isreal(x) & (abs(x) > 1)): + x = _tocomplex(x) + return x + + +def _unary_dispatcher(x): + return (x,) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def sqrt(x): + """ + Compute the square root of x. + + For negative input elements, a complex value is returned + (unlike `numpy.sqrt` which returns NaN). + + Parameters + ---------- + x : array_like + The input value(s). + + Returns + ------- + out : ndarray or scalar + The square root of `x`. If `x` was a scalar, so is `out`, + otherwise an array is returned. + + See Also + -------- + numpy.sqrt + + Examples + -------- + For real, non-negative inputs this works just like `numpy.sqrt`: + + >>> import numpy as np + + >>> np.emath.sqrt(1) + 1.0 + >>> np.emath.sqrt([1, 4]) + array([1., 2.]) + + But it automatically handles negative inputs: + + >>> np.emath.sqrt(-1) + 1j + >>> np.emath.sqrt([-1,4]) + array([0.+1.j, 2.+0.j]) + + Different results are expected because: + floating point 0.0 and -0.0 are distinct. + + For more control, explicitly use complex() as follows: + + >>> np.emath.sqrt(complex(-4.0, 0.0)) + 2j + >>> np.emath.sqrt(complex(-4.0, -0.0)) + -2j + """ + x = _fix_real_lt_zero(x) + return nx.sqrt(x) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def log(x): + """ + Compute the natural logarithm of `x`. + + Return the "principal value" (for a description of this, see `numpy.log`) + of :math:`log_e(x)`. For real `x > 0`, this is a real number (``log(0)`` + returns ``-inf`` and ``log(np.inf)`` returns ``inf``). Otherwise, the + complex principle value is returned. + + Parameters + ---------- + x : array_like + The value(s) whose log is (are) required. + + Returns + ------- + out : ndarray or scalar + The log of the `x` value(s). If `x` was a scalar, so is `out`, + otherwise an array is returned. + + See Also + -------- + numpy.log + + Notes + ----- + For a log() that returns ``NAN`` when real `x < 0`, use `numpy.log` + (note, however, that otherwise `numpy.log` and this `log` are identical, + i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`, and, + notably, the complex principle value if ``x.imag != 0``). + + Examples + -------- + >>> import numpy as np + >>> np.emath.log(np.exp(1)) + 1.0 + + Negative arguments are handled "correctly" (recall that + ``exp(log(x)) == x`` does *not* hold for real ``x < 0``): + + >>> np.emath.log(-np.exp(1)) == (1 + np.pi * 1j) + True + + """ + x = _fix_real_lt_zero(x) + return nx.log(x) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def log10(x): + """ + Compute the logarithm base 10 of `x`. + + Return the "principal value" (for a description of this, see + `numpy.log10`) of :math:`log_{10}(x)`. For real `x > 0`, this + is a real number (``log10(0)`` returns ``-inf`` and ``log10(np.inf)`` + returns ``inf``). Otherwise, the complex principle value is returned. + + Parameters + ---------- + x : array_like or scalar + The value(s) whose log base 10 is (are) required. + + Returns + ------- + out : ndarray or scalar + The log base 10 of the `x` value(s). If `x` was a scalar, so is `out`, + otherwise an array object is returned. + + See Also + -------- + numpy.log10 + + Notes + ----- + For a log10() that returns ``NAN`` when real `x < 0`, use `numpy.log10` + (note, however, that otherwise `numpy.log10` and this `log10` are + identical, i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`, + and, notably, the complex principle value if ``x.imag != 0``). + + Examples + -------- + >>> import numpy as np + + (We set the printing precision so the example can be auto-tested) + + >>> np.set_printoptions(precision=4) + + >>> np.emath.log10(10**1) + 1.0 + + >>> np.emath.log10([-10**1, -10**2, 10**2]) + array([1.+1.3644j, 2.+1.3644j, 2.+0.j ]) + + """ + x = _fix_real_lt_zero(x) + return nx.log10(x) + + +def _logn_dispatcher(n, x): + return (n, x,) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_logn_dispatcher) +def logn(n, x): + """ + Take log base n of x. + + If `x` contains negative inputs, the answer is computed and returned in the + complex domain. + + Parameters + ---------- + n : array_like + The integer base(s) in which the log is taken. + x : array_like + The value(s) whose log base `n` is (are) required. + + Returns + ------- + out : ndarray or scalar + The log base `n` of the `x` value(s). If `x` was a scalar, so is + `out`, otherwise an array is returned. + + Examples + -------- + >>> import numpy as np + >>> np.set_printoptions(precision=4) + + >>> np.emath.logn(2, [4, 8]) + array([2., 3.]) + >>> np.emath.logn(2, [-4, -8, 8]) + array([2.+4.5324j, 3.+4.5324j, 3.+0.j ]) + + """ + x = _fix_real_lt_zero(x) + n = _fix_real_lt_zero(n) + return nx.log(x)/nx.log(n) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def log2(x): + """ + Compute the logarithm base 2 of `x`. + + Return the "principal value" (for a description of this, see + `numpy.log2`) of :math:`log_2(x)`. For real `x > 0`, this is + a real number (``log2(0)`` returns ``-inf`` and ``log2(np.inf)`` returns + ``inf``). Otherwise, the complex principle value is returned. + + Parameters + ---------- + x : array_like + The value(s) whose log base 2 is (are) required. + + Returns + ------- + out : ndarray or scalar + The log base 2 of the `x` value(s). If `x` was a scalar, so is `out`, + otherwise an array is returned. + + See Also + -------- + numpy.log2 + + Notes + ----- + For a log2() that returns ``NAN`` when real `x < 0`, use `numpy.log2` + (note, however, that otherwise `numpy.log2` and this `log2` are + identical, i.e., both return ``-inf`` for `x = 0`, ``inf`` for `x = inf`, + and, notably, the complex principle value if ``x.imag != 0``). + + Examples + -------- + + We set the printing precision so the example can be auto-tested: + + >>> np.set_printoptions(precision=4) + + >>> np.emath.log2(8) + 3.0 + >>> np.emath.log2([-4, -8, 8]) + array([2.+4.5324j, 3.+4.5324j, 3.+0.j ]) + + """ + x = _fix_real_lt_zero(x) + return nx.log2(x) + + +def _power_dispatcher(x, p): + return (x, p) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_power_dispatcher) +def power(x, p): + """ + Return x to the power p, (x**p). + + If `x` contains negative values, the output is converted to the + complex domain. + + Parameters + ---------- + x : array_like + The input value(s). + p : array_like of ints + The power(s) to which `x` is raised. If `x` contains multiple values, + `p` has to either be a scalar, or contain the same number of values + as `x`. In the latter case, the result is + ``x[0]**p[0], x[1]**p[1], ...``. + + Returns + ------- + out : ndarray or scalar + The result of ``x**p``. If `x` and `p` are scalars, so is `out`, + otherwise an array is returned. + + See Also + -------- + numpy.power + + Examples + -------- + >>> import numpy as np + >>> np.set_printoptions(precision=4) + + >>> np.emath.power(2, 2) + 4 + + >>> np.emath.power([2, 4], 2) + array([ 4, 16]) + + >>> np.emath.power([2, 4], -2) + array([0.25 , 0.0625]) + + >>> np.emath.power([-2, 4], 2) + array([ 4.-0.j, 16.+0.j]) + + >>> np.emath.power([2, 4], [2, 4]) + array([ 4, 256]) + + """ + x = _fix_real_lt_zero(x) + p = _fix_int_lt_zero(p) + return nx.power(x, p) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def arccos(x): + """ + Compute the inverse cosine of x. + + Return the "principal value" (for a description of this, see + `numpy.arccos`) of the inverse cosine of `x`. For real `x` such that + `abs(x) <= 1`, this is a real number in the closed interval + :math:`[0, \\pi]`. Otherwise, the complex principle value is returned. + + Parameters + ---------- + x : array_like or scalar + The value(s) whose arccos is (are) required. + + Returns + ------- + out : ndarray or scalar + The inverse cosine(s) of the `x` value(s). If `x` was a scalar, so + is `out`, otherwise an array object is returned. + + See Also + -------- + numpy.arccos + + Notes + ----- + For an arccos() that returns ``NAN`` when real `x` is not in the + interval ``[-1,1]``, use `numpy.arccos`. + + Examples + -------- + >>> import numpy as np + >>> np.set_printoptions(precision=4) + + >>> np.emath.arccos(1) # a scalar is returned + 0.0 + + >>> np.emath.arccos([1,2]) + array([0.-0.j , 0.-1.317j]) + + """ + x = _fix_real_abs_gt_1(x) + return nx.arccos(x) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def arcsin(x): + """ + Compute the inverse sine of x. + + Return the "principal value" (for a description of this, see + `numpy.arcsin`) of the inverse sine of `x`. For real `x` such that + `abs(x) <= 1`, this is a real number in the closed interval + :math:`[-\\pi/2, \\pi/2]`. Otherwise, the complex principle value is + returned. + + Parameters + ---------- + x : array_like or scalar + The value(s) whose arcsin is (are) required. + + Returns + ------- + out : ndarray or scalar + The inverse sine(s) of the `x` value(s). If `x` was a scalar, so + is `out`, otherwise an array object is returned. + + See Also + -------- + numpy.arcsin + + Notes + ----- + For an arcsin() that returns ``NAN`` when real `x` is not in the + interval ``[-1,1]``, use `numpy.arcsin`. + + Examples + -------- + >>> import numpy as np + >>> np.set_printoptions(precision=4) + + >>> np.emath.arcsin(0) + 0.0 + + >>> np.emath.arcsin([0,1]) + array([0. , 1.5708]) + + """ + x = _fix_real_abs_gt_1(x) + return nx.arcsin(x) + + +@set_module('numpy.lib.scimath') +@array_function_dispatch(_unary_dispatcher) +def arctanh(x): + """ + Compute the inverse hyperbolic tangent of `x`. + + Return the "principal value" (for a description of this, see + `numpy.arctanh`) of ``arctanh(x)``. For real `x` such that + ``abs(x) < 1``, this is a real number. If `abs(x) > 1`, or if `x` is + complex, the result is complex. Finally, `x = 1` returns``inf`` and + ``x=-1`` returns ``-inf``. + + Parameters + ---------- + x : array_like + The value(s) whose arctanh is (are) required. + + Returns + ------- + out : ndarray or scalar + The inverse hyperbolic tangent(s) of the `x` value(s). If `x` was + a scalar so is `out`, otherwise an array is returned. + + + See Also + -------- + numpy.arctanh + + Notes + ----- + For an arctanh() that returns ``NAN`` when real `x` is not in the + interval ``(-1,1)``, use `numpy.arctanh` (this latter, however, does + return +/-inf for ``x = +/-1``). + + Examples + -------- + >>> import numpy as np + >>> np.set_printoptions(precision=4) + + >>> np.emath.arctanh(0.5) + 0.5493061443340549 + + >>> from numpy.testing import suppress_warnings + >>> with suppress_warnings() as sup: + ... sup.filter(RuntimeWarning) + ... np.emath.arctanh(np.eye(2)) + array([[inf, 0.], + [ 0., inf]]) + >>> np.emath.arctanh([1j]) + array([0.+0.7854j]) + + """ + x = _fix_real_abs_gt_1(x) + return nx.arctanh(x) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..43b7110b2923453110edce9942a0659508c03eb6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_scimath_impl.pyi @@ -0,0 +1,94 @@ +from typing import overload, Any + +from numpy import complexfloating + +from numpy._typing import ( + NDArray, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ComplexLike_co, + _FloatLike_co, +) + +__all__ = ["sqrt", "log", "log2", "logn", "log10", "power", "arccos", "arcsin", "arctanh"] + +@overload +def sqrt(x: _FloatLike_co) -> Any: ... +@overload +def sqrt(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def sqrt(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def sqrt(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def log(x: _FloatLike_co) -> Any: ... +@overload +def log(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def log(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def log(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def log10(x: _FloatLike_co) -> Any: ... +@overload +def log10(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def log10(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def log10(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def log2(x: _FloatLike_co) -> Any: ... +@overload +def log2(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def log2(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def log2(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def logn(n: _FloatLike_co, x: _FloatLike_co) -> Any: ... +@overload +def logn(n: _ComplexLike_co, x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def logn(n: _ArrayLikeFloat_co, x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def logn(n: _ArrayLikeComplex_co, x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def power(x: _FloatLike_co, p: _FloatLike_co) -> Any: ... +@overload +def power(x: _ComplexLike_co, p: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def power(x: _ArrayLikeFloat_co, p: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def power(x: _ArrayLikeComplex_co, p: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def arccos(x: _FloatLike_co) -> Any: ... +@overload +def arccos(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def arccos(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def arccos(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def arcsin(x: _FloatLike_co) -> Any: ... +@overload +def arcsin(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def arcsin(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def arcsin(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + +@overload +def arctanh(x: _FloatLike_co) -> Any: ... +@overload +def arctanh(x: _ComplexLike_co) -> complexfloating[Any, Any]: ... +@overload +def arctanh(x: _ArrayLikeFloat_co) -> NDArray[Any]: ... +@overload +def arctanh(x: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..7d861bb6f2e00d1c39fb47701b511b2cb2319de1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.py @@ -0,0 +1,1294 @@ +import functools +import warnings + +import numpy._core.numeric as _nx +from numpy._core.numeric import asarray, zeros, zeros_like, array, asanyarray +from numpy._core.fromnumeric import reshape, transpose +from numpy._core.multiarray import normalize_axis_index +from numpy._core._multiarray_umath import _array_converter +from numpy._core import overrides +from numpy._core import vstack, atleast_3d +from numpy._core.numeric import normalize_axis_tuple +from numpy._core.overrides import set_module +from numpy._core.shape_base import _arrays_for_stack_dispatcher +from numpy.lib._index_tricks_impl import ndindex +from numpy.matrixlib.defmatrix import matrix # this raises all the right alarm bells + + +__all__ = [ + 'column_stack', 'row_stack', 'dstack', 'array_split', 'split', + 'hsplit', 'vsplit', 'dsplit', 'apply_over_axes', 'expand_dims', + 'apply_along_axis', 'kron', 'tile', 'take_along_axis', + 'put_along_axis' + ] + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +def _make_along_axis_idx(arr_shape, indices, axis): + # compute dimensions to iterate over + if not _nx.issubdtype(indices.dtype, _nx.integer): + raise IndexError('`indices` must be an integer array') + if len(arr_shape) != indices.ndim: + raise ValueError( + "`indices` and `arr` must have the same number of dimensions") + shape_ones = (1,) * indices.ndim + dest_dims = list(range(axis)) + [None] + list(range(axis+1, indices.ndim)) + + # build a fancy index, consisting of orthogonal aranges, with the + # requested index inserted at the right location + fancy_index = [] + for dim, n in zip(dest_dims, arr_shape): + if dim is None: + fancy_index.append(indices) + else: + ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim+1:] + fancy_index.append(_nx.arange(n).reshape(ind_shape)) + + return tuple(fancy_index) + + +def _take_along_axis_dispatcher(arr, indices, axis): + return (arr, indices) + + +@array_function_dispatch(_take_along_axis_dispatcher) +def take_along_axis(arr, indices, axis): + """ + Take values from the input array by matching 1d index and data slices. + + This iterates over matching 1d slices oriented along the specified axis in + the index and data arrays, and uses the former to look up values in the + latter. These slices can be different lengths. + + Functions returning an index along an axis, like `argsort` and + `argpartition`, produce suitable indices for this function. + + Parameters + ---------- + arr : ndarray (Ni..., M, Nk...) + Source array + indices : ndarray (Ni..., J, Nk...) + Indices to take along each 1d slice of `arr`. This must match the + dimension of arr, but dimensions Ni and Nj only need to broadcast + against `arr`. + axis : int + The axis to take 1d slices along. If axis is None, the input array is + treated as if it had first been flattened to 1d, for consistency with + `sort` and `argsort`. + + Returns + ------- + out: ndarray (Ni..., J, Nk...) + The indexed result. + + Notes + ----- + This is equivalent to (but faster than) the following use of `ndindex` and + `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices:: + + Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:] + J = indices.shape[axis] # Need not equal M + out = np.empty(Ni + (J,) + Nk) + + for ii in ndindex(Ni): + for kk in ndindex(Nk): + a_1d = a [ii + s_[:,] + kk] + indices_1d = indices[ii + s_[:,] + kk] + out_1d = out [ii + s_[:,] + kk] + for j in range(J): + out_1d[j] = a_1d[indices_1d[j]] + + Equivalently, eliminating the inner loop, the last two lines would be:: + + out_1d[:] = a_1d[indices_1d] + + See Also + -------- + take : Take along an axis, using the same indices for every 1d slice + put_along_axis : + Put values into the destination array by matching 1d index and data slices + + Examples + -------- + >>> import numpy as np + + For this sample array + + >>> a = np.array([[10, 30, 20], [60, 40, 50]]) + + We can sort either by using sort directly, or argsort and this function + + >>> np.sort(a, axis=1) + array([[10, 20, 30], + [40, 50, 60]]) + >>> ai = np.argsort(a, axis=1) + >>> ai + array([[0, 2, 1], + [1, 2, 0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[10, 20, 30], + [40, 50, 60]]) + + The same works for max and min, if you maintain the trivial dimension + with ``keepdims``: + + >>> np.max(a, axis=1, keepdims=True) + array([[30], + [60]]) + >>> ai = np.argmax(a, axis=1, keepdims=True) + >>> ai + array([[1], + [0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[30], + [60]]) + + If we want to get the max and min at the same time, we can stack the + indices first + + >>> ai_min = np.argmin(a, axis=1, keepdims=True) + >>> ai_max = np.argmax(a, axis=1, keepdims=True) + >>> ai = np.concatenate([ai_min, ai_max], axis=1) + >>> ai + array([[0, 1], + [1, 0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[10, 30], + [40, 60]]) + """ + # normalize inputs + if axis is None: + if indices.ndim != 1: + raise ValueError( + 'when axis=None, `indices` must have a single dimension.') + arr = arr.flat + arr_shape = (len(arr),) # flatiter has no .shape + axis = 0 + else: + axis = normalize_axis_index(axis, arr.ndim) + arr_shape = arr.shape + + # use the fancy index + return arr[_make_along_axis_idx(arr_shape, indices, axis)] + + +def _put_along_axis_dispatcher(arr, indices, values, axis): + return (arr, indices, values) + + +@array_function_dispatch(_put_along_axis_dispatcher) +def put_along_axis(arr, indices, values, axis): + """ + Put values into the destination array by matching 1d index and data slices. + + This iterates over matching 1d slices oriented along the specified axis in + the index and data arrays, and uses the former to place values into the + latter. These slices can be different lengths. + + Functions returning an index along an axis, like `argsort` and + `argpartition`, produce suitable indices for this function. + + Parameters + ---------- + arr : ndarray (Ni..., M, Nk...) + Destination array. + indices : ndarray (Ni..., J, Nk...) + Indices to change along each 1d slice of `arr`. This must match the + dimension of arr, but dimensions in Ni and Nj may be 1 to broadcast + against `arr`. + values : array_like (Ni..., J, Nk...) + values to insert at those indices. Its shape and dimension are + broadcast to match that of `indices`. + axis : int + The axis to take 1d slices along. If axis is None, the destination + array is treated as if a flattened 1d view had been created of it. + + Notes + ----- + This is equivalent to (but faster than) the following use of `ndindex` and + `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices:: + + Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:] + J = indices.shape[axis] # Need not equal M + + for ii in ndindex(Ni): + for kk in ndindex(Nk): + a_1d = a [ii + s_[:,] + kk] + indices_1d = indices[ii + s_[:,] + kk] + values_1d = values [ii + s_[:,] + kk] + for j in range(J): + a_1d[indices_1d[j]] = values_1d[j] + + Equivalently, eliminating the inner loop, the last two lines would be:: + + a_1d[indices_1d] = values_1d + + See Also + -------- + take_along_axis : + Take values from the input array by matching 1d index and data slices + + Examples + -------- + >>> import numpy as np + + For this sample array + + >>> a = np.array([[10, 30, 20], [60, 40, 50]]) + + We can replace the maximum values with: + + >>> ai = np.argmax(a, axis=1, keepdims=True) + >>> ai + array([[1], + [0]]) + >>> np.put_along_axis(a, ai, 99, axis=1) + >>> a + array([[10, 99, 20], + [99, 40, 50]]) + + """ + # normalize inputs + if axis is None: + if indices.ndim != 1: + raise ValueError( + 'when axis=None, `indices` must have a single dimension.') + arr = arr.flat + axis = 0 + arr_shape = (len(arr),) # flatiter has no .shape + else: + axis = normalize_axis_index(axis, arr.ndim) + arr_shape = arr.shape + + # use the fancy index + arr[_make_along_axis_idx(arr_shape, indices, axis)] = values + + +def _apply_along_axis_dispatcher(func1d, axis, arr, *args, **kwargs): + return (arr,) + + +@array_function_dispatch(_apply_along_axis_dispatcher) +def apply_along_axis(func1d, axis, arr, *args, **kwargs): + """ + Apply a function to 1-D slices along the given axis. + + Execute `func1d(a, *args, **kwargs)` where `func1d` operates on 1-D arrays + and `a` is a 1-D slice of `arr` along `axis`. + + This is equivalent to (but faster than) the following use of `ndindex` and + `s_`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of indices:: + + Ni, Nk = a.shape[:axis], a.shape[axis+1:] + for ii in ndindex(Ni): + for kk in ndindex(Nk): + f = func1d(arr[ii + s_[:,] + kk]) + Nj = f.shape + for jj in ndindex(Nj): + out[ii + jj + kk] = f[jj] + + Equivalently, eliminating the inner loop, this can be expressed as:: + + Ni, Nk = a.shape[:axis], a.shape[axis+1:] + for ii in ndindex(Ni): + for kk in ndindex(Nk): + out[ii + s_[...,] + kk] = func1d(arr[ii + s_[:,] + kk]) + + Parameters + ---------- + func1d : function (M,) -> (Nj...) + This function should accept 1-D arrays. It is applied to 1-D + slices of `arr` along the specified axis. + axis : integer + Axis along which `arr` is sliced. + arr : ndarray (Ni..., M, Nk...) + Input array. + args : any + Additional arguments to `func1d`. + kwargs : any + Additional named arguments to `func1d`. + + Returns + ------- + out : ndarray (Ni..., Nj..., Nk...) + The output array. The shape of `out` is identical to the shape of + `arr`, except along the `axis` dimension. This axis is removed, and + replaced with new dimensions equal to the shape of the return value + of `func1d`. So if `func1d` returns a scalar `out` will have one + fewer dimensions than `arr`. + + See Also + -------- + apply_over_axes : Apply a function repeatedly over multiple axes. + + Examples + -------- + >>> import numpy as np + >>> def my_func(a): + ... \"\"\"Average first and last element of a 1-D array\"\"\" + ... return (a[0] + a[-1]) * 0.5 + >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]]) + >>> np.apply_along_axis(my_func, 0, b) + array([4., 5., 6.]) + >>> np.apply_along_axis(my_func, 1, b) + array([2., 5., 8.]) + + For a function that returns a 1D array, the number of dimensions in + `outarr` is the same as `arr`. + + >>> b = np.array([[8,1,7], [4,3,9], [5,2,6]]) + >>> np.apply_along_axis(sorted, 1, b) + array([[1, 7, 8], + [3, 4, 9], + [2, 5, 6]]) + + For a function that returns a higher dimensional array, those dimensions + are inserted in place of the `axis` dimension. + + >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]]) + >>> np.apply_along_axis(np.diag, -1, b) + array([[[1, 0, 0], + [0, 2, 0], + [0, 0, 3]], + [[4, 0, 0], + [0, 5, 0], + [0, 0, 6]], + [[7, 0, 0], + [0, 8, 0], + [0, 0, 9]]]) + """ + # handle negative axes + conv = _array_converter(arr) + arr = conv[0] + + nd = arr.ndim + axis = normalize_axis_index(axis, nd) + + # arr, with the iteration axis at the end + in_dims = list(range(nd)) + inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis+1:] + [axis]) + + # compute indices for the iteration axes, and append a trailing ellipsis to + # prevent 0d arrays decaying to scalars, which fixes gh-8642 + inds = ndindex(inarr_view.shape[:-1]) + inds = (ind + (Ellipsis,) for ind in inds) + + # invoke the function on the first item + try: + ind0 = next(inds) + except StopIteration: + raise ValueError( + 'Cannot apply_along_axis when any iteration dimensions are 0' + ) from None + res = asanyarray(func1d(inarr_view[ind0], *args, **kwargs)) + + # build a buffer for storing evaluations of func1d. + # remove the requested axis, and add the new ones on the end. + # laid out so that each write is contiguous. + # for a tuple index inds, buff[inds] = func1d(inarr_view[inds]) + if not isinstance(res, matrix): + buff = zeros_like(res, shape=inarr_view.shape[:-1] + res.shape) + else: + # Matrices are nasty with reshaping, so do not preserve them here. + buff = zeros(inarr_view.shape[:-1] + res.shape, dtype=res.dtype) + + # permutation of axes such that out = buff.transpose(buff_permute) + buff_dims = list(range(buff.ndim)) + buff_permute = ( + buff_dims[0 : axis] + + buff_dims[buff.ndim-res.ndim : buff.ndim] + + buff_dims[axis : buff.ndim-res.ndim] + ) + + # save the first result, then compute and save all remaining results + buff[ind0] = res + for ind in inds: + buff[ind] = asanyarray(func1d(inarr_view[ind], *args, **kwargs)) + + res = transpose(buff, buff_permute) + return conv.wrap(res) + + +def _apply_over_axes_dispatcher(func, a, axes): + return (a,) + + +@array_function_dispatch(_apply_over_axes_dispatcher) +def apply_over_axes(func, a, axes): + """ + Apply a function repeatedly over multiple axes. + + `func` is called as `res = func(a, axis)`, where `axis` is the first + element of `axes`. The result `res` of the function call must have + either the same dimensions as `a` or one less dimension. If `res` + has one less dimension than `a`, a dimension is inserted before + `axis`. The call to `func` is then repeated for each axis in `axes`, + with `res` as the first argument. + + Parameters + ---------- + func : function + This function must take two arguments, `func(a, axis)`. + a : array_like + Input array. + axes : array_like + Axes over which `func` is applied; the elements must be integers. + + Returns + ------- + apply_over_axis : ndarray + The output array. The number of dimensions is the same as `a`, + but the shape can be different. This depends on whether `func` + changes the shape of its output with respect to its input. + + See Also + -------- + apply_along_axis : + Apply a function to 1-D slices of an array along the given axis. + + Notes + ----- + This function is equivalent to tuple axis arguments to reorderable ufuncs + with keepdims=True. Tuple axis arguments to ufuncs have been available since + version 1.7.0. + + Examples + -------- + >>> import numpy as np + >>> a = np.arange(24).reshape(2,3,4) + >>> a + array([[[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11]], + [[12, 13, 14, 15], + [16, 17, 18, 19], + [20, 21, 22, 23]]]) + + Sum over axes 0 and 2. The result has same number of dimensions + as the original array: + + >>> np.apply_over_axes(np.sum, a, [0,2]) + array([[[ 60], + [ 92], + [124]]]) + + Tuple axis arguments to ufuncs are equivalent: + + >>> np.sum(a, axis=(0,2), keepdims=True) + array([[[ 60], + [ 92], + [124]]]) + + """ + val = asarray(a) + N = a.ndim + if array(axes).ndim == 0: + axes = (axes,) + for axis in axes: + if axis < 0: + axis = N + axis + args = (val, axis) + res = func(*args) + if res.ndim == val.ndim: + val = res + else: + res = expand_dims(res, axis) + if res.ndim == val.ndim: + val = res + else: + raise ValueError("function is not returning " + "an array of the correct shape") + return val + + +def _expand_dims_dispatcher(a, axis): + return (a,) + + +@array_function_dispatch(_expand_dims_dispatcher) +def expand_dims(a, axis): + """ + Expand the shape of an array. + + Insert a new axis that will appear at the `axis` position in the expanded + array shape. + + Parameters + ---------- + a : array_like + Input array. + axis : int or tuple of ints + Position in the expanded axes where the new axis (or axes) is placed. + + .. deprecated:: 1.13.0 + Passing an axis where ``axis > a.ndim`` will be treated as + ``axis == a.ndim``, and passing ``axis < -a.ndim - 1`` will + be treated as ``axis == 0``. This behavior is deprecated. + + Returns + ------- + result : ndarray + View of `a` with the number of dimensions increased. + + See Also + -------- + squeeze : The inverse operation, removing singleton dimensions + reshape : Insert, remove, and combine dimensions, and resize existing ones + atleast_1d, atleast_2d, atleast_3d + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2]) + >>> x.shape + (2,) + + The following is equivalent to ``x[np.newaxis, :]`` or ``x[np.newaxis]``: + + >>> y = np.expand_dims(x, axis=0) + >>> y + array([[1, 2]]) + >>> y.shape + (1, 2) + + The following is equivalent to ``x[:, np.newaxis]``: + + >>> y = np.expand_dims(x, axis=1) + >>> y + array([[1], + [2]]) + >>> y.shape + (2, 1) + + ``axis`` may also be a tuple: + + >>> y = np.expand_dims(x, axis=(0, 1)) + >>> y + array([[[1, 2]]]) + + >>> y = np.expand_dims(x, axis=(2, 0)) + >>> y + array([[[1], + [2]]]) + + Note that some examples may use ``None`` instead of ``np.newaxis``. These + are the same objects: + + >>> np.newaxis is None + True + + """ + if isinstance(a, matrix): + a = asarray(a) + else: + a = asanyarray(a) + + if type(axis) not in (tuple, list): + axis = (axis,) + + out_ndim = len(axis) + a.ndim + axis = normalize_axis_tuple(axis, out_ndim) + + shape_it = iter(a.shape) + shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)] + + return a.reshape(shape) + + +# NOTE: Remove once deprecation period passes +@set_module("numpy") +def row_stack(tup, *, dtype=None, casting="same_kind"): + # Deprecated in NumPy 2.0, 2023-08-18 + warnings.warn( + "`row_stack` alias is deprecated. " + "Use `np.vstack` directly.", + DeprecationWarning, + stacklevel=2 + ) + return vstack(tup, dtype=dtype, casting=casting) + + +row_stack.__doc__ = vstack.__doc__ + + +def _column_stack_dispatcher(tup): + return _arrays_for_stack_dispatcher(tup) + + +@array_function_dispatch(_column_stack_dispatcher) +def column_stack(tup): + """ + Stack 1-D arrays as columns into a 2-D array. + + Take a sequence of 1-D arrays and stack them as columns + to make a single 2-D array. 2-D arrays are stacked as-is, + just like with `hstack`. 1-D arrays are turned into 2-D columns + first. + + Parameters + ---------- + tup : sequence of 1-D or 2-D arrays. + Arrays to stack. All of them must have the same first dimension. + + Returns + ------- + stacked : 2-D array + The array formed by stacking the given arrays. + + See Also + -------- + stack, hstack, vstack, concatenate + + Examples + -------- + >>> import numpy as np + >>> a = np.array((1,2,3)) + >>> b = np.array((2,3,4)) + >>> np.column_stack((a,b)) + array([[1, 2], + [2, 3], + [3, 4]]) + + """ + arrays = [] + for v in tup: + arr = asanyarray(v) + if arr.ndim < 2: + arr = array(arr, copy=None, subok=True, ndmin=2).T + arrays.append(arr) + return _nx.concatenate(arrays, 1) + + +def _dstack_dispatcher(tup): + return _arrays_for_stack_dispatcher(tup) + + +@array_function_dispatch(_dstack_dispatcher) +def dstack(tup): + """ + Stack arrays in sequence depth wise (along third axis). + + This is equivalent to concatenation along the third axis after 2-D arrays + of shape `(M,N)` have been reshaped to `(M,N,1)` and 1-D arrays of shape + `(N,)` have been reshaped to `(1,N,1)`. Rebuilds arrays divided by + `dsplit`. + + This function makes most sense for arrays with up to 3 dimensions. For + instance, for pixel-data with a height (first axis), width (second axis), + and r/g/b channels (third axis). The functions `concatenate`, `stack` and + `block` provide more general stacking and concatenation operations. + + Parameters + ---------- + tup : sequence of arrays + The arrays must have the same shape along all but the third axis. + 1-D or 2-D arrays must have the same shape. + + Returns + ------- + stacked : ndarray + The array formed by stacking the given arrays, will be at least 3-D. + + See Also + -------- + concatenate : Join a sequence of arrays along an existing axis. + stack : Join a sequence of arrays along a new axis. + block : Assemble an nd-array from nested lists of blocks. + vstack : Stack arrays in sequence vertically (row wise). + hstack : Stack arrays in sequence horizontally (column wise). + column_stack : Stack 1-D arrays as columns into a 2-D array. + dsplit : Split array along third axis. + + Examples + -------- + >>> import numpy as np + >>> a = np.array((1,2,3)) + >>> b = np.array((2,3,4)) + >>> np.dstack((a,b)) + array([[[1, 2], + [2, 3], + [3, 4]]]) + + >>> a = np.array([[1],[2],[3]]) + >>> b = np.array([[2],[3],[4]]) + >>> np.dstack((a,b)) + array([[[1, 2]], + [[2, 3]], + [[3, 4]]]) + + """ + arrs = atleast_3d(*tup) + if not isinstance(arrs, tuple): + arrs = (arrs,) + return _nx.concatenate(arrs, 2) + + +def _replace_zero_by_x_arrays(sub_arys): + for i in range(len(sub_arys)): + if _nx.ndim(sub_arys[i]) == 0: + sub_arys[i] = _nx.empty(0, dtype=sub_arys[i].dtype) + elif _nx.sometrue(_nx.equal(_nx.shape(sub_arys[i]), 0)): + sub_arys[i] = _nx.empty(0, dtype=sub_arys[i].dtype) + return sub_arys + + +def _array_split_dispatcher(ary, indices_or_sections, axis=None): + return (ary, indices_or_sections) + + +@array_function_dispatch(_array_split_dispatcher) +def array_split(ary, indices_or_sections, axis=0): + """ + Split an array into multiple sub-arrays. + + Please refer to the ``split`` documentation. The only difference + between these functions is that ``array_split`` allows + `indices_or_sections` to be an integer that does *not* equally + divide the axis. For an array of length l that should be split + into n sections, it returns l % n sub-arrays of size l//n + 1 + and the rest of size l//n. + + See Also + -------- + split : Split array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(8.0) + >>> np.array_split(x, 3) + [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])] + + >>> x = np.arange(9) + >>> np.array_split(x, 4) + [array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8])] + + """ + try: + Ntotal = ary.shape[axis] + except AttributeError: + Ntotal = len(ary) + try: + # handle array case. + Nsections = len(indices_or_sections) + 1 + div_points = [0] + list(indices_or_sections) + [Ntotal] + except TypeError: + # indices_or_sections is a scalar, not an array. + Nsections = int(indices_or_sections) + if Nsections <= 0: + raise ValueError('number sections must be larger than 0.') from None + Neach_section, extras = divmod(Ntotal, Nsections) + section_sizes = ([0] + + extras * [Neach_section+1] + + (Nsections-extras) * [Neach_section]) + div_points = _nx.array(section_sizes, dtype=_nx.intp).cumsum() + + sub_arys = [] + sary = _nx.swapaxes(ary, axis, 0) + for i in range(Nsections): + st = div_points[i] + end = div_points[i + 1] + sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0)) + + return sub_arys + + +def _split_dispatcher(ary, indices_or_sections, axis=None): + return (ary, indices_or_sections) + + +@array_function_dispatch(_split_dispatcher) +def split(ary, indices_or_sections, axis=0): + """ + Split an array into multiple sub-arrays as views into `ary`. + + Parameters + ---------- + ary : ndarray + Array to be divided into sub-arrays. + indices_or_sections : int or 1-D array + If `indices_or_sections` is an integer, N, the array will be divided + into N equal arrays along `axis`. If such a split is not possible, + an error is raised. + + If `indices_or_sections` is a 1-D array of sorted integers, the entries + indicate where along `axis` the array is split. For example, + ``[2, 3]`` would, for ``axis=0``, result in + + - ary[:2] + - ary[2:3] + - ary[3:] + + If an index exceeds the dimension of the array along `axis`, + an empty sub-array is returned correspondingly. + axis : int, optional + The axis along which to split, default is 0. + + Returns + ------- + sub-arrays : list of ndarrays + A list of sub-arrays as views into `ary`. + + Raises + ------ + ValueError + If `indices_or_sections` is given as an integer, but + a split does not result in equal division. + + See Also + -------- + array_split : Split an array into multiple sub-arrays of equal or + near-equal size. Does not raise an exception if + an equal division cannot be made. + hsplit : Split array into multiple sub-arrays horizontally (column-wise). + vsplit : Split array into multiple sub-arrays vertically (row wise). + dsplit : Split array into multiple sub-arrays along the 3rd axis (depth). + concatenate : Join a sequence of arrays along an existing axis. + stack : Join a sequence of arrays along a new axis. + hstack : Stack arrays in sequence horizontally (column wise). + vstack : Stack arrays in sequence vertically (row wise). + dstack : Stack arrays in sequence depth wise (along third dimension). + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(9.0) + >>> np.split(x, 3) + [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])] + + >>> x = np.arange(8.0) + >>> np.split(x, [3, 5, 6, 10]) + [array([0., 1., 2.]), + array([3., 4.]), + array([5.]), + array([6., 7.]), + array([], dtype=float64)] + + """ + try: + len(indices_or_sections) + except TypeError: + sections = indices_or_sections + N = ary.shape[axis] + if N % sections: + raise ValueError( + 'array split does not result in an equal division') from None + return array_split(ary, indices_or_sections, axis) + + +def _hvdsplit_dispatcher(ary, indices_or_sections): + return (ary, indices_or_sections) + + +@array_function_dispatch(_hvdsplit_dispatcher) +def hsplit(ary, indices_or_sections): + """ + Split an array into multiple sub-arrays horizontally (column-wise). + + Please refer to the `split` documentation. `hsplit` is equivalent + to `split` with ``axis=1``, the array is always split along the second + axis except for 1-D arrays, where it is split at ``axis=0``. + + See Also + -------- + split : Split an array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(16.0).reshape(4, 4) + >>> x + array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> np.hsplit(x, 2) + [array([[ 0., 1.], + [ 4., 5.], + [ 8., 9.], + [12., 13.]]), + array([[ 2., 3.], + [ 6., 7.], + [10., 11.], + [14., 15.]])] + >>> np.hsplit(x, np.array([3, 6])) + [array([[ 0., 1., 2.], + [ 4., 5., 6.], + [ 8., 9., 10.], + [12., 13., 14.]]), + array([[ 3.], + [ 7.], + [11.], + [15.]]), + array([], shape=(4, 0), dtype=float64)] + + With a higher dimensional array the split is still along the second axis. + + >>> x = np.arange(8.0).reshape(2, 2, 2) + >>> x + array([[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]]) + >>> np.hsplit(x, 2) + [array([[[0., 1.]], + [[4., 5.]]]), + array([[[2., 3.]], + [[6., 7.]]])] + + With a 1-D array, the split is along axis 0. + + >>> x = np.array([0, 1, 2, 3, 4, 5]) + >>> np.hsplit(x, 2) + [array([0, 1, 2]), array([3, 4, 5])] + + """ + if _nx.ndim(ary) == 0: + raise ValueError('hsplit only works on arrays of 1 or more dimensions') + if ary.ndim > 1: + return split(ary, indices_or_sections, 1) + else: + return split(ary, indices_or_sections, 0) + + +@array_function_dispatch(_hvdsplit_dispatcher) +def vsplit(ary, indices_or_sections): + """ + Split an array into multiple sub-arrays vertically (row-wise). + + Please refer to the ``split`` documentation. ``vsplit`` is equivalent + to ``split`` with `axis=0` (default), the array is always split along the + first axis regardless of the array dimension. + + See Also + -------- + split : Split an array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(16.0).reshape(4, 4) + >>> x + array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> np.vsplit(x, 2) + [array([[0., 1., 2., 3.], + [4., 5., 6., 7.]]), + array([[ 8., 9., 10., 11.], + [12., 13., 14., 15.]])] + >>> np.vsplit(x, np.array([3, 6])) + [array([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]]), + array([[12., 13., 14., 15.]]), + array([], shape=(0, 4), dtype=float64)] + + With a higher dimensional array the split is still along the first axis. + + >>> x = np.arange(8.0).reshape(2, 2, 2) + >>> x + array([[[0., 1.], + [2., 3.]], + [[4., 5.], + [6., 7.]]]) + >>> np.vsplit(x, 2) + [array([[[0., 1.], + [2., 3.]]]), + array([[[4., 5.], + [6., 7.]]])] + + """ + if _nx.ndim(ary) < 2: + raise ValueError('vsplit only works on arrays of 2 or more dimensions') + return split(ary, indices_or_sections, 0) + + +@array_function_dispatch(_hvdsplit_dispatcher) +def dsplit(ary, indices_or_sections): + """ + Split array into multiple sub-arrays along the 3rd axis (depth). + + Please refer to the `split` documentation. `dsplit` is equivalent + to `split` with ``axis=2``, the array is always split along the third + axis provided the array dimension is greater than or equal to 3. + + See Also + -------- + split : Split an array into multiple sub-arrays of equal size. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(16.0).reshape(2, 2, 4) + >>> x + array([[[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.]], + [[ 8., 9., 10., 11.], + [12., 13., 14., 15.]]]) + >>> np.dsplit(x, 2) + [array([[[ 0., 1.], + [ 4., 5.]], + [[ 8., 9.], + [12., 13.]]]), array([[[ 2., 3.], + [ 6., 7.]], + [[10., 11.], + [14., 15.]]])] + >>> np.dsplit(x, np.array([3, 6])) + [array([[[ 0., 1., 2.], + [ 4., 5., 6.]], + [[ 8., 9., 10.], + [12., 13., 14.]]]), + array([[[ 3.], + [ 7.]], + [[11.], + [15.]]]), + array([], shape=(2, 2, 0), dtype=float64)] + """ + if _nx.ndim(ary) < 3: + raise ValueError('dsplit only works on arrays of 3 or more dimensions') + return split(ary, indices_or_sections, 2) + + +def get_array_wrap(*args): + """Find the wrapper for the array with the highest priority. + + In case of ties, leftmost wins. If no wrapper is found, return None. + + .. deprecated:: 2.0 + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`get_array_wrap` is deprecated. " + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + wrappers = sorted((getattr(x, '__array_priority__', 0), -i, + x.__array_wrap__) for i, x in enumerate(args) + if hasattr(x, '__array_wrap__')) + if wrappers: + return wrappers[-1][-1] + return None + + +def _kron_dispatcher(a, b): + return (a, b) + + +@array_function_dispatch(_kron_dispatcher) +def kron(a, b): + """ + Kronecker product of two arrays. + + Computes the Kronecker product, a composite array made of blocks of the + second array scaled by the first. + + Parameters + ---------- + a, b : array_like + + Returns + ------- + out : ndarray + + See Also + -------- + outer : The outer product + + Notes + ----- + The function assumes that the number of dimensions of `a` and `b` + are the same, if necessary prepending the smallest with ones. + If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``, + the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``. + The elements are products of elements from `a` and `b`, organized + explicitly by:: + + kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN] + + where:: + + kt = it * st + jt, t = 0,...,N + + In the common 2-D case (N=1), the block structure can be visualized:: + + [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ], + [ ... ... ], + [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]] + + + Examples + -------- + >>> import numpy as np + >>> np.kron([1,10,100], [5,6,7]) + array([ 5, 6, 7, ..., 500, 600, 700]) + >>> np.kron([5,6,7], [1,10,100]) + array([ 5, 50, 500, ..., 7, 70, 700]) + + >>> np.kron(np.eye(2), np.ones((2,2))) + array([[1., 1., 0., 0.], + [1., 1., 0., 0.], + [0., 0., 1., 1.], + [0., 0., 1., 1.]]) + + >>> a = np.arange(100).reshape((2,5,2,5)) + >>> b = np.arange(24).reshape((2,3,4)) + >>> c = np.kron(a,b) + >>> c.shape + (2, 10, 6, 20) + >>> I = (1,3,0,2) + >>> J = (0,2,1) + >>> J1 = (0,) + J # extend to ndim=4 + >>> S1 = (1,) + b.shape + >>> K = tuple(np.array(I) * np.array(S1) + np.array(J1)) + >>> c[K] == a[I]*b[J] + True + + """ + # Working: + # 1. Equalise the shapes by prepending smaller array with 1s + # 2. Expand shapes of both the arrays by adding new axes at + # odd positions for 1st array and even positions for 2nd + # 3. Compute the product of the modified array + # 4. The inner most array elements now contain the rows of + # the Kronecker product + # 5. Reshape the result to kron's shape, which is same as + # product of shapes of the two arrays. + b = asanyarray(b) + a = array(a, copy=None, subok=True, ndmin=b.ndim) + is_any_mat = isinstance(a, matrix) or isinstance(b, matrix) + ndb, nda = b.ndim, a.ndim + nd = max(ndb, nda) + + if (nda == 0 or ndb == 0): + return _nx.multiply(a, b) + + as_ = a.shape + bs = b.shape + if not a.flags.contiguous: + a = reshape(a, as_) + if not b.flags.contiguous: + b = reshape(b, bs) + + # Equalise the shapes by prepending smaller one with 1s + as_ = (1,)*max(0, ndb-nda) + as_ + bs = (1,)*max(0, nda-ndb) + bs + + # Insert empty dimensions + a_arr = expand_dims(a, axis=tuple(range(ndb-nda))) + b_arr = expand_dims(b, axis=tuple(range(nda-ndb))) + + # Compute the product + a_arr = expand_dims(a_arr, axis=tuple(range(1, nd*2, 2))) + b_arr = expand_dims(b_arr, axis=tuple(range(0, nd*2, 2))) + # In case of `mat`, convert result to `array` + result = _nx.multiply(a_arr, b_arr, subok=(not is_any_mat)) + + # Reshape back + result = result.reshape(_nx.multiply(as_, bs)) + + return result if not is_any_mat else matrix(result, copy=False) + + +def _tile_dispatcher(A, reps): + return (A, reps) + + +@array_function_dispatch(_tile_dispatcher) +def tile(A, reps): + """ + Construct an array by repeating A the number of times given by reps. + + If `reps` has length ``d``, the result will have dimension of + ``max(d, A.ndim)``. + + If ``A.ndim < d``, `A` is promoted to be d-dimensional by prepending new + axes. So a shape (3,) array is promoted to (1, 3) for 2-D replication, + or shape (1, 1, 3) for 3-D replication. If this is not the desired + behavior, promote `A` to d-dimensions manually before calling this + function. + + If ``A.ndim > d``, `reps` is promoted to `A`.ndim by prepending 1's to it. + Thus for an `A` of shape (2, 3, 4, 5), a `reps` of (2, 2) is treated as + (1, 1, 2, 2). + + Note : Although tile may be used for broadcasting, it is strongly + recommended to use numpy's broadcasting operations and functions. + + Parameters + ---------- + A : array_like + The input array. + reps : array_like + The number of repetitions of `A` along each axis. + + Returns + ------- + c : ndarray + The tiled output array. + + See Also + -------- + repeat : Repeat elements of an array. + broadcast_to : Broadcast an array to a new shape + + Examples + -------- + >>> import numpy as np + >>> a = np.array([0, 1, 2]) + >>> np.tile(a, 2) + array([0, 1, 2, 0, 1, 2]) + >>> np.tile(a, (2, 2)) + array([[0, 1, 2, 0, 1, 2], + [0, 1, 2, 0, 1, 2]]) + >>> np.tile(a, (2, 1, 2)) + array([[[0, 1, 2, 0, 1, 2]], + [[0, 1, 2, 0, 1, 2]]]) + + >>> b = np.array([[1, 2], [3, 4]]) + >>> np.tile(b, 2) + array([[1, 2, 1, 2], + [3, 4, 3, 4]]) + >>> np.tile(b, (2, 1)) + array([[1, 2], + [3, 4], + [1, 2], + [3, 4]]) + + >>> c = np.array([1,2,3,4]) + >>> np.tile(c,(4,1)) + array([[1, 2, 3, 4], + [1, 2, 3, 4], + [1, 2, 3, 4], + [1, 2, 3, 4]]) + """ + try: + tup = tuple(reps) + except TypeError: + tup = (reps,) + d = len(tup) + if all(x == 1 for x in tup) and isinstance(A, _nx.ndarray): + # Fixes the problem that the function does not make a copy if A is a + # numpy array and the repetitions are 1 in all dimensions + return _nx.array(A, copy=True, subok=True, ndmin=d) + else: + # Note that no copy of zero-sized arrays is made. However since they + # have no data there is no risk of an inadvertent overwrite. + c = _nx.array(A, copy=None, subok=True, ndmin=d) + if (d < c.ndim): + tup = (1,)*(c.ndim-d) + tup + shape_out = tuple(s*t for s, t in zip(c.shape, tup)) + n = c.size + if n > 0: + for dim_in, nrep in zip(c.shape, tup): + if nrep != 1: + c = c.reshape(-1, n).repeat(nrep, 0) + n //= dim_in + return c.reshape(shape_out) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..77e5d2de9cb9259efba6d9128705f1a4597d3ef3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.pyi @@ -0,0 +1,225 @@ +from collections.abc import Callable, Sequence +from typing import ( + TypeVar, + Any, + overload, + SupportsIndex, + Protocol, + ParamSpec, + Concatenate, + type_check_only, +) + +from typing_extensions import deprecated + +import numpy as np +from numpy import _CastingKind, generic, integer, ufunc, unsignedinteger, signedinteger, floating, complexfloating, object_ +from numpy._typing import ( + ArrayLike, + DTypeLike, + NDArray, + _ShapeLike, + _ArrayLike, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, +) + +__all__ = [ + "column_stack", + "row_stack", + "dstack", + "array_split", + "split", + "hsplit", + "vsplit", + "dsplit", + "apply_over_axes", + "expand_dims", + "apply_along_axis", + "kron", + "tile", + "take_along_axis", + "put_along_axis", +] + +_P = ParamSpec("_P") +_SCT = TypeVar("_SCT", bound=generic) + +# Signature of `__array_wrap__` +@type_check_only +class _ArrayWrap(Protocol): + def __call__( + self, + array: NDArray[Any], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + return_scalar: bool = ..., + /, + ) -> Any: ... + +@type_check_only +class _SupportsArrayWrap(Protocol): + @property + def __array_wrap__(self) -> _ArrayWrap: ... + +### + +def take_along_axis( + arr: _SCT | NDArray[_SCT], + indices: NDArray[integer[Any]], + axis: None | int, +) -> NDArray[_SCT]: ... + +def put_along_axis( + arr: NDArray[_SCT], + indices: NDArray[integer[Any]], + values: ArrayLike, + axis: None | int, +) -> None: ... + +@overload +def apply_along_axis( + func1d: Callable[Concatenate[NDArray[Any], _P], _ArrayLike[_SCT]], + axis: SupportsIndex, + arr: ArrayLike, + *args: _P.args, + **kwargs: _P.kwargs, +) -> NDArray[_SCT]: ... +@overload +def apply_along_axis( + func1d: Callable[Concatenate[NDArray[Any], _P], Any], + axis: SupportsIndex, + arr: ArrayLike, + *args: _P.args, + **kwargs: _P.kwargs, +) -> NDArray[Any]: ... + +def apply_over_axes( + func: Callable[[NDArray[Any], int], NDArray[_SCT]], + a: ArrayLike, + axes: int | Sequence[int], +) -> NDArray[_SCT]: ... + +@overload +def expand_dims( + a: _ArrayLike[_SCT], + axis: _ShapeLike, +) -> NDArray[_SCT]: ... +@overload +def expand_dims( + a: ArrayLike, + axis: _ShapeLike, +) -> NDArray[Any]: ... + +# Deprecated in NumPy 2.0, 2023-08-18 +@deprecated("`row_stack` alias is deprecated. Use `np.vstack` directly.") +def row_stack( + tup: Sequence[ArrayLike], + *, + dtype: DTypeLike | None = None, + casting: _CastingKind = "same_kind", +) -> NDArray[Any]: ... + +# +@overload +def column_stack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... +@overload +def column_stack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... + +@overload +def dstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... +@overload +def dstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... + +@overload +def array_split( + ary: _ArrayLike[_SCT], + indices_or_sections: _ShapeLike, + axis: SupportsIndex = ..., +) -> list[NDArray[_SCT]]: ... +@overload +def array_split( + ary: ArrayLike, + indices_or_sections: _ShapeLike, + axis: SupportsIndex = ..., +) -> list[NDArray[Any]]: ... + +@overload +def split( + ary: _ArrayLike[_SCT], + indices_or_sections: _ShapeLike, + axis: SupportsIndex = ..., +) -> list[NDArray[_SCT]]: ... +@overload +def split( + ary: ArrayLike, + indices_or_sections: _ShapeLike, + axis: SupportsIndex = ..., +) -> list[NDArray[Any]]: ... + +@overload +def hsplit( + ary: _ArrayLike[_SCT], + indices_or_sections: _ShapeLike, +) -> list[NDArray[_SCT]]: ... +@overload +def hsplit( + ary: ArrayLike, + indices_or_sections: _ShapeLike, +) -> list[NDArray[Any]]: ... + +@overload +def vsplit( + ary: _ArrayLike[_SCT], + indices_or_sections: _ShapeLike, +) -> list[NDArray[_SCT]]: ... +@overload +def vsplit( + ary: ArrayLike, + indices_or_sections: _ShapeLike, +) -> list[NDArray[Any]]: ... + +@overload +def dsplit( + ary: _ArrayLike[_SCT], + indices_or_sections: _ShapeLike, +) -> list[NDArray[_SCT]]: ... +@overload +def dsplit( + ary: ArrayLike, + indices_or_sections: _ShapeLike, +) -> list[NDArray[Any]]: ... + +@overload +def get_array_wrap(*args: _SupportsArrayWrap) -> _ArrayWrap: ... +@overload +def get_array_wrap(*args: object) -> None | _ArrayWrap: ... + +@overload +def kron(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co) -> NDArray[np.bool]: ... # type: ignore[misc] +@overload +def kron(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] +@overload +def kron(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] +@overload +def kron(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] +@overload +def kron(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def kron(a: _ArrayLikeObject_co, b: Any) -> NDArray[object_]: ... +@overload +def kron(a: Any, b: _ArrayLikeObject_co) -> NDArray[object_]: ... + +@overload +def tile( + A: _ArrayLike[_SCT], + reps: int | Sequence[int], +) -> NDArray[_SCT]: ... +@overload +def tile( + A: ArrayLike, + reps: int | Sequence[int], +) -> NDArray[Any]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..d4780783a63812a8e5d6ffe1eb3b0f1ff0b1266c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.py @@ -0,0 +1,549 @@ +""" +Utilities that manipulate strides to achieve desirable effects. + +An explanation of strides can be found in the :ref:`arrays.ndarray`. + +""" +import numpy as np +from numpy._core.numeric import normalize_axis_tuple +from numpy._core.overrides import array_function_dispatch, set_module + +__all__ = ['broadcast_to', 'broadcast_arrays', 'broadcast_shapes'] + + +class DummyArray: + """Dummy object that just exists to hang __array_interface__ dictionaries + and possibly keep alive a reference to a base array. + """ + + def __init__(self, interface, base=None): + self.__array_interface__ = interface + self.base = base + + +def _maybe_view_as_subclass(original_array, new_array): + if type(original_array) is not type(new_array): + # if input was an ndarray subclass and subclasses were OK, + # then view the result as that subclass. + new_array = new_array.view(type=type(original_array)) + # Since we have done something akin to a view from original_array, we + # should let the subclass finalize (if it has it implemented, i.e., is + # not None). + if new_array.__array_finalize__: + new_array.__array_finalize__(original_array) + return new_array + + +@set_module("numpy.lib.stride_tricks") +def as_strided(x, shape=None, strides=None, subok=False, writeable=True): + """ + Create a view into the array with the given shape and strides. + + .. warning:: This function has to be used with extreme care, see notes. + + Parameters + ---------- + x : ndarray + Array to create a new. + shape : sequence of int, optional + The shape of the new array. Defaults to ``x.shape``. + strides : sequence of int, optional + The strides of the new array. Defaults to ``x.strides``. + subok : bool, optional + If True, subclasses are preserved. + writeable : bool, optional + If set to False, the returned array will always be readonly. + Otherwise it will be writable if the original array was. It + is advisable to set this to False if possible (see Notes). + + Returns + ------- + view : ndarray + + See also + -------- + broadcast_to : broadcast an array to a given shape. + reshape : reshape an array. + lib.stride_tricks.sliding_window_view : + userfriendly and safe function for a creation of sliding window views. + + Notes + ----- + ``as_strided`` creates a view into the array given the exact strides + and shape. This means it manipulates the internal data structure of + ndarray and, if done incorrectly, the array elements can point to + invalid memory and can corrupt results or crash your program. + It is advisable to always use the original ``x.strides`` when + calculating new strides to avoid reliance on a contiguous memory + layout. + + Furthermore, arrays created with this function often contain self + overlapping memory, so that two elements are identical. + Vectorized write operations on such arrays will typically be + unpredictable. They may even give different results for small, large, + or transposed arrays. + + Since writing to these arrays has to be tested and done with great + care, you may want to use ``writeable=False`` to avoid accidental write + operations. + + For these reasons it is advisable to avoid ``as_strided`` when + possible. + """ + # first convert input to array, possibly keeping subclass + x = np.array(x, copy=None, subok=subok) + interface = dict(x.__array_interface__) + if shape is not None: + interface['shape'] = tuple(shape) + if strides is not None: + interface['strides'] = tuple(strides) + + array = np.asarray(DummyArray(interface, base=x)) + # The route via `__interface__` does not preserve structured + # dtypes. Since dtype should remain unchanged, we set it explicitly. + array.dtype = x.dtype + + view = _maybe_view_as_subclass(x, array) + + if view.flags.writeable and not writeable: + view.flags.writeable = False + + return view + + +def _sliding_window_view_dispatcher(x, window_shape, axis=None, *, + subok=None, writeable=None): + return (x,) + + +@array_function_dispatch( + _sliding_window_view_dispatcher, module="numpy.lib.stride_tricks" +) +def sliding_window_view(x, window_shape, axis=None, *, + subok=False, writeable=False): + """ + Create a sliding window view into the array with the given window shape. + + Also known as rolling or moving window, the window slides across all + dimensions of the array and extracts subsets of the array at all window + positions. + + .. versionadded:: 1.20.0 + + Parameters + ---------- + x : array_like + Array to create the sliding window view from. + window_shape : int or tuple of int + Size of window over each axis that takes part in the sliding window. + If `axis` is not present, must have same length as the number of input + array dimensions. Single integers `i` are treated as if they were the + tuple `(i,)`. + axis : int or tuple of int, optional + Axis or axes along which the sliding window is applied. + By default, the sliding window is applied to all axes and + `window_shape[i]` will refer to axis `i` of `x`. + If `axis` is given as a `tuple of int`, `window_shape[i]` will refer to + the axis `axis[i]` of `x`. + Single integers `i` are treated as if they were the tuple `(i,)`. + subok : bool, optional + If True, sub-classes will be passed-through, otherwise the returned + array will be forced to be a base-class array (default). + writeable : bool, optional + When true, allow writing to the returned view. The default is false, + as this should be used with caution: the returned view contains the + same memory location multiple times, so writing to one location will + cause others to change. + + Returns + ------- + view : ndarray + Sliding window view of the array. The sliding window dimensions are + inserted at the end, and the original dimensions are trimmed as + required by the size of the sliding window. + That is, ``view.shape = x_shape_trimmed + window_shape``, where + ``x_shape_trimmed`` is ``x.shape`` with every entry reduced by one less + than the corresponding window size. + + See Also + -------- + lib.stride_tricks.as_strided: A lower-level and less safe routine for + creating arbitrary views from custom shape and strides. + broadcast_to: broadcast an array to a given shape. + + Notes + ----- + For many applications using a sliding window view can be convenient, but + potentially very slow. Often specialized solutions exist, for example: + + - `scipy.signal.fftconvolve` + + - filtering functions in `scipy.ndimage` + + - moving window functions provided by + `bottleneck `_. + + As a rough estimate, a sliding window approach with an input size of `N` + and a window size of `W` will scale as `O(N*W)` where frequently a special + algorithm can achieve `O(N)`. That means that the sliding window variant + for a window size of 100 can be a 100 times slower than a more specialized + version. + + Nevertheless, for small window sizes, when no custom algorithm exists, or + as a prototyping and developing tool, this function can be a good solution. + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib.stride_tricks import sliding_window_view + >>> x = np.arange(6) + >>> x.shape + (6,) + >>> v = sliding_window_view(x, 3) + >>> v.shape + (4, 3) + >>> v + array([[0, 1, 2], + [1, 2, 3], + [2, 3, 4], + [3, 4, 5]]) + + This also works in more dimensions, e.g. + + >>> i, j = np.ogrid[:3, :4] + >>> x = 10*i + j + >>> x.shape + (3, 4) + >>> x + array([[ 0, 1, 2, 3], + [10, 11, 12, 13], + [20, 21, 22, 23]]) + >>> shape = (2,2) + >>> v = sliding_window_view(x, shape) + >>> v.shape + (2, 3, 2, 2) + >>> v + array([[[[ 0, 1], + [10, 11]], + [[ 1, 2], + [11, 12]], + [[ 2, 3], + [12, 13]]], + [[[10, 11], + [20, 21]], + [[11, 12], + [21, 22]], + [[12, 13], + [22, 23]]]]) + + The axis can be specified explicitly: + + >>> v = sliding_window_view(x, 3, 0) + >>> v.shape + (1, 4, 3) + >>> v + array([[[ 0, 10, 20], + [ 1, 11, 21], + [ 2, 12, 22], + [ 3, 13, 23]]]) + + The same axis can be used several times. In that case, every use reduces + the corresponding original dimension: + + >>> v = sliding_window_view(x, (2, 3), (1, 1)) + >>> v.shape + (3, 1, 2, 3) + >>> v + array([[[[ 0, 1, 2], + [ 1, 2, 3]]], + [[[10, 11, 12], + [11, 12, 13]]], + [[[20, 21, 22], + [21, 22, 23]]]]) + + Combining with stepped slicing (`::step`), this can be used to take sliding + views which skip elements: + + >>> x = np.arange(7) + >>> sliding_window_view(x, 5)[:, ::2] + array([[0, 2, 4], + [1, 3, 5], + [2, 4, 6]]) + + or views which move by multiple elements + + >>> x = np.arange(7) + >>> sliding_window_view(x, 3)[::2, :] + array([[0, 1, 2], + [2, 3, 4], + [4, 5, 6]]) + + A common application of `sliding_window_view` is the calculation of running + statistics. The simplest example is the + `moving average `_: + + >>> x = np.arange(6) + >>> x.shape + (6,) + >>> v = sliding_window_view(x, 3) + >>> v.shape + (4, 3) + >>> v + array([[0, 1, 2], + [1, 2, 3], + [2, 3, 4], + [3, 4, 5]]) + >>> moving_average = v.mean(axis=-1) + >>> moving_average + array([1., 2., 3., 4.]) + + Note that a sliding window approach is often **not** optimal (see Notes). + """ + window_shape = (tuple(window_shape) + if np.iterable(window_shape) + else (window_shape,)) + # first convert input to array, possibly keeping subclass + x = np.array(x, copy=None, subok=subok) + + window_shape_array = np.array(window_shape) + if np.any(window_shape_array < 0): + raise ValueError('`window_shape` cannot contain negative values') + + if axis is None: + axis = tuple(range(x.ndim)) + if len(window_shape) != len(axis): + raise ValueError(f'Since axis is `None`, must provide ' + f'window_shape for all dimensions of `x`; ' + f'got {len(window_shape)} window_shape elements ' + f'and `x.ndim` is {x.ndim}.') + else: + axis = normalize_axis_tuple(axis, x.ndim, allow_duplicate=True) + if len(window_shape) != len(axis): + raise ValueError(f'Must provide matching length window_shape and ' + f'axis; got {len(window_shape)} window_shape ' + f'elements and {len(axis)} axes elements.') + + out_strides = x.strides + tuple(x.strides[ax] for ax in axis) + + # note: same axis can be windowed repeatedly + x_shape_trimmed = list(x.shape) + for ax, dim in zip(axis, window_shape): + if x_shape_trimmed[ax] < dim: + raise ValueError( + 'window shape cannot be larger than input array shape') + x_shape_trimmed[ax] -= dim - 1 + out_shape = tuple(x_shape_trimmed) + window_shape + return as_strided(x, strides=out_strides, shape=out_shape, + subok=subok, writeable=writeable) + + +def _broadcast_to(array, shape, subok, readonly): + shape = tuple(shape) if np.iterable(shape) else (shape,) + array = np.array(array, copy=None, subok=subok) + if not shape and array.shape: + raise ValueError('cannot broadcast a non-scalar to a scalar array') + if any(size < 0 for size in shape): + raise ValueError('all elements of broadcast shape must be non-' + 'negative') + extras = [] + it = np.nditer( + (array,), flags=['multi_index', 'refs_ok', 'zerosize_ok'] + extras, + op_flags=['readonly'], itershape=shape, order='C') + with it: + # never really has writebackifcopy semantics + broadcast = it.itviews[0] + result = _maybe_view_as_subclass(array, broadcast) + # In a future version this will go away + if not readonly and array.flags._writeable_no_warn: + result.flags.writeable = True + result.flags._warn_on_write = True + return result + + +def _broadcast_to_dispatcher(array, shape, subok=None): + return (array,) + + +@array_function_dispatch(_broadcast_to_dispatcher, module='numpy') +def broadcast_to(array, shape, subok=False): + """Broadcast an array to a new shape. + + Parameters + ---------- + array : array_like + The array to broadcast. + shape : tuple or int + The shape of the desired array. A single integer ``i`` is interpreted + as ``(i,)``. + subok : bool, optional + If True, then sub-classes will be passed-through, otherwise + the returned array will be forced to be a base-class array (default). + + Returns + ------- + broadcast : array + A readonly view on the original array with the given shape. It is + typically not contiguous. Furthermore, more than one element of a + broadcasted array may refer to a single memory location. + + Raises + ------ + ValueError + If the array is not compatible with the new shape according to NumPy's + broadcasting rules. + + See Also + -------- + broadcast + broadcast_arrays + broadcast_shapes + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2, 3]) + >>> np.broadcast_to(x, (3, 3)) + array([[1, 2, 3], + [1, 2, 3], + [1, 2, 3]]) + """ + return _broadcast_to(array, shape, subok=subok, readonly=True) + + +def _broadcast_shape(*args): + """Returns the shape of the arrays that would result from broadcasting the + supplied arrays against each other. + """ + # use the old-iterator because np.nditer does not handle size 0 arrays + # consistently + b = np.broadcast(*args[:32]) + # unfortunately, it cannot handle 32 or more arguments directly + for pos in range(32, len(args), 31): + # ironically, np.broadcast does not properly handle np.broadcast + # objects (it treats them as scalars) + # use broadcasting to avoid allocating the full array + b = broadcast_to(0, b.shape) + b = np.broadcast(b, *args[pos:(pos + 31)]) + return b.shape + + +_size0_dtype = np.dtype([]) + + +@set_module('numpy') +def broadcast_shapes(*args): + """ + Broadcast the input shapes into a single shape. + + :ref:`Learn more about broadcasting here `. + + .. versionadded:: 1.20.0 + + Parameters + ---------- + *args : tuples of ints, or ints + The shapes to be broadcast against each other. + + Returns + ------- + tuple + Broadcasted shape. + + Raises + ------ + ValueError + If the shapes are not compatible and cannot be broadcast according + to NumPy's broadcasting rules. + + See Also + -------- + broadcast + broadcast_arrays + broadcast_to + + Examples + -------- + >>> import numpy as np + >>> np.broadcast_shapes((1, 2), (3, 1), (3, 2)) + (3, 2) + + >>> np.broadcast_shapes((6, 7), (5, 6, 1), (7,), (5, 1, 7)) + (5, 6, 7) + """ + arrays = [np.empty(x, dtype=_size0_dtype) for x in args] + return _broadcast_shape(*arrays) + + +def _broadcast_arrays_dispatcher(*args, subok=None): + return args + + +@array_function_dispatch(_broadcast_arrays_dispatcher, module='numpy') +def broadcast_arrays(*args, subok=False): + """ + Broadcast any number of arrays against each other. + + Parameters + ---------- + *args : array_likes + The arrays to broadcast. + + subok : bool, optional + If True, then sub-classes will be passed-through, otherwise + the returned arrays will be forced to be a base-class array (default). + + Returns + ------- + broadcasted : tuple of arrays + These arrays are views on the original arrays. They are typically + not contiguous. Furthermore, more than one element of a + broadcasted array may refer to a single memory location. If you need + to write to the arrays, make copies first. While you can set the + ``writable`` flag True, writing to a single output value may end up + changing more than one location in the output array. + + .. deprecated:: 1.17 + The output is currently marked so that if written to, a deprecation + warning will be emitted. A future version will set the + ``writable`` flag False so writing to it will raise an error. + + See Also + -------- + broadcast + broadcast_to + broadcast_shapes + + Examples + -------- + >>> import numpy as np + >>> x = np.array([[1,2,3]]) + >>> y = np.array([[4],[5]]) + >>> np.broadcast_arrays(x, y) + (array([[1, 2, 3], + [1, 2, 3]]), + array([[4, 4, 4], + [5, 5, 5]])) + + Here is a useful idiom for getting contiguous copies instead of + non-contiguous views. + + >>> [np.array(a) for a in np.broadcast_arrays(x, y)] + [array([[1, 2, 3], + [1, 2, 3]]), + array([[4, 4, 4], + [5, 5, 5]])] + + """ + # nditer is not used here to avoid the limit of 32 arrays. + # Otherwise, something like the following one-liner would suffice: + # return np.nditer(args, flags=['multi_index', 'zerosize_ok'], + # order='C').itviews + + args = [np.array(_m, copy=None, subok=subok) for _m in args] + + shape = _broadcast_shape(*args) + + result = [array if array.shape == shape + else _broadcast_to(array, shape, subok=subok, readonly=False) + for array in args] + return tuple(result) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e2284115eeb40276bf4003bb420f3ce83f004fb8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_stride_tricks_impl.pyi @@ -0,0 +1,80 @@ +from collections.abc import Iterable +from typing import Any, TypeVar, overload, SupportsIndex + +from numpy import generic +from numpy._typing import ( + NDArray, + ArrayLike, + _ShapeLike, + _Shape, + _ArrayLike +) + +__all__ = ["broadcast_to", "broadcast_arrays", "broadcast_shapes"] + +_SCT = TypeVar("_SCT", bound=generic) + +class DummyArray: + __array_interface__: dict[str, Any] + base: None | NDArray[Any] + def __init__( + self, + interface: dict[str, Any], + base: None | NDArray[Any] = ..., + ) -> None: ... + +@overload +def as_strided( + x: _ArrayLike[_SCT], + shape: None | Iterable[int] = ..., + strides: None | Iterable[int] = ..., + subok: bool = ..., + writeable: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def as_strided( + x: ArrayLike, + shape: None | Iterable[int] = ..., + strides: None | Iterable[int] = ..., + subok: bool = ..., + writeable: bool = ..., +) -> NDArray[Any]: ... + +@overload +def sliding_window_view( + x: _ArrayLike[_SCT], + window_shape: int | Iterable[int], + axis: None | SupportsIndex = ..., + *, + subok: bool = ..., + writeable: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def sliding_window_view( + x: ArrayLike, + window_shape: int | Iterable[int], + axis: None | SupportsIndex = ..., + *, + subok: bool = ..., + writeable: bool = ..., +) -> NDArray[Any]: ... + +@overload +def broadcast_to( + array: _ArrayLike[_SCT], + shape: int | Iterable[int], + subok: bool = ..., +) -> NDArray[_SCT]: ... +@overload +def broadcast_to( + array: ArrayLike, + shape: int | Iterable[int], + subok: bool = ..., +) -> NDArray[Any]: ... + +def broadcast_shapes(*args: _ShapeLike) -> _Shape: ... + +def broadcast_arrays( + *args: ArrayLike, + subok: bool = ..., +) -> tuple[NDArray[Any], ...]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..e8815bede8919c5658243146bd06f4d325126051 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.py @@ -0,0 +1,1188 @@ +""" Basic functions for manipulating 2d arrays + +""" +import functools +import operator + +from numpy._core._multiarray_umath import _array_converter +from numpy._core.numeric import ( + asanyarray, arange, zeros, greater_equal, multiply, ones, + asarray, where, int8, int16, int32, int64, intp, empty, promote_types, + diagonal, nonzero, indices + ) +from numpy._core.overrides import finalize_array_function_like, set_module +from numpy._core import overrides +from numpy._core import iinfo +from numpy.lib._stride_tricks_impl import broadcast_to + + +__all__ = [ + 'diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'tri', 'triu', + 'tril', 'vander', 'histogram2d', 'mask_indices', 'tril_indices', + 'tril_indices_from', 'triu_indices', 'triu_indices_from', ] + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +i1 = iinfo(int8) +i2 = iinfo(int16) +i4 = iinfo(int32) + + +def _min_int(low, high): + """ get small int that fits the range """ + if high <= i1.max and low >= i1.min: + return int8 + if high <= i2.max and low >= i2.min: + return int16 + if high <= i4.max and low >= i4.min: + return int32 + return int64 + + +def _flip_dispatcher(m): + return (m,) + + +@array_function_dispatch(_flip_dispatcher) +def fliplr(m): + """ + Reverse the order of elements along axis 1 (left/right). + + For a 2-D array, this flips the entries in each row in the left/right + direction. Columns are preserved, but appear in a different order than + before. + + Parameters + ---------- + m : array_like + Input array, must be at least 2-D. + + Returns + ------- + f : ndarray + A view of `m` with the columns reversed. Since a view + is returned, this operation is :math:`\\mathcal O(1)`. + + See Also + -------- + flipud : Flip array in the up/down direction. + flip : Flip array in one or more dimensions. + rot90 : Rotate array counterclockwise. + + Notes + ----- + Equivalent to ``m[:,::-1]`` or ``np.flip(m, axis=1)``. + Requires the array to be at least 2-D. + + Examples + -------- + >>> import numpy as np + >>> A = np.diag([1.,2.,3.]) + >>> A + array([[1., 0., 0.], + [0., 2., 0.], + [0., 0., 3.]]) + >>> np.fliplr(A) + array([[0., 0., 1.], + [0., 2., 0.], + [3., 0., 0.]]) + + >>> rng = np.random.default_rng() + >>> A = rng.normal(size=(2,3,5)) + >>> np.all(np.fliplr(A) == A[:,::-1,...]) + True + + """ + m = asanyarray(m) + if m.ndim < 2: + raise ValueError("Input must be >= 2-d.") + return m[:, ::-1] + + +@array_function_dispatch(_flip_dispatcher) +def flipud(m): + """ + Reverse the order of elements along axis 0 (up/down). + + For a 2-D array, this flips the entries in each column in the up/down + direction. Rows are preserved, but appear in a different order than before. + + Parameters + ---------- + m : array_like + Input array. + + Returns + ------- + out : array_like + A view of `m` with the rows reversed. Since a view is + returned, this operation is :math:`\\mathcal O(1)`. + + See Also + -------- + fliplr : Flip array in the left/right direction. + flip : Flip array in one or more dimensions. + rot90 : Rotate array counterclockwise. + + Notes + ----- + Equivalent to ``m[::-1, ...]`` or ``np.flip(m, axis=0)``. + Requires the array to be at least 1-D. + + Examples + -------- + >>> import numpy as np + >>> A = np.diag([1.0, 2, 3]) + >>> A + array([[1., 0., 0.], + [0., 2., 0.], + [0., 0., 3.]]) + >>> np.flipud(A) + array([[0., 0., 3.], + [0., 2., 0.], + [1., 0., 0.]]) + + >>> rng = np.random.default_rng() + >>> A = rng.normal(size=(2,3,5)) + >>> np.all(np.flipud(A) == A[::-1,...]) + True + + >>> np.flipud([1,2]) + array([2, 1]) + + """ + m = asanyarray(m) + if m.ndim < 1: + raise ValueError("Input must be >= 1-d.") + return m[::-1, ...] + + +@finalize_array_function_like +@set_module('numpy') +def eye(N, M=None, k=0, dtype=float, order='C', *, device=None, like=None): + """ + Return a 2-D array 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. If None, defaults to `N`. + k : int, optional + Index of the diagonal: 0 (the default) refers to the main diagonal, + a positive value refers to an upper diagonal, and a negative value + to a lower diagonal. + dtype : data-type, optional + Data-type of the returned array. + order : {'C', 'F'}, optional + Whether the output should be stored in row-major (C-style) or + column-major (Fortran-style) order in memory. + device : str, optional + The device on which to place the created array. Default: None. + For Array-API interoperability only, so must be ``"cpu"`` if passed. + + .. versionadded:: 2.0.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + I : ndarray of shape (N,M) + An array where all elements are equal to zero, except for the `k`-th + diagonal, whose values are equal to one. + + See Also + -------- + identity : (almost) equivalent function + diag : diagonal 2-D array from a 1-D array specified by the user. + + Examples + -------- + >>> import numpy as np + >>> np.eye(2, dtype=int) + array([[1, 0], + [0, 1]]) + >>> np.eye(3, k=1) + array([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) + + """ + if like is not None: + return _eye_with_like( + like, N, M=M, k=k, dtype=dtype, order=order, device=device + ) + if M is None: + M = N + m = zeros((N, M), dtype=dtype, order=order, device=device) + if k >= M: + return m + # Ensure M and k are integers, so we don't get any surprise casting + # results in the expressions `M-k` and `M+1` used below. This avoids + # a problem with inputs with type (for example) np.uint64. + M = operator.index(M) + k = operator.index(k) + if k >= 0: + i = k + else: + i = (-k) * M + m[:M-k].flat[i::M+1] = 1 + return m + + +_eye_with_like = array_function_dispatch()(eye) + + +def _diag_dispatcher(v, k=None): + return (v,) + + +@array_function_dispatch(_diag_dispatcher) +def diag(v, k=0): + """ + Extract a diagonal or construct a diagonal array. + + See the more detailed documentation for ``numpy.diagonal`` if you use this + function to extract a diagonal and wish to write to the resulting array; + whether it returns a copy or a view depends on what version of numpy you + are using. + + Parameters + ---------- + v : array_like + If `v` is a 2-D array, return a copy of its `k`-th diagonal. + If `v` is a 1-D array, return a 2-D array with `v` on the `k`-th + diagonal. + k : int, optional + Diagonal in question. The default is 0. Use `k>0` for diagonals + above the main diagonal, and `k<0` for diagonals below the main + diagonal. + + Returns + ------- + out : ndarray + The extracted diagonal or constructed diagonal array. + + See Also + -------- + diagonal : Return specified diagonals. + diagflat : Create a 2-D array with the flattened input as a diagonal. + trace : Sum along diagonals. + triu : Upper triangle of an array. + tril : Lower triangle of an array. + + Examples + -------- + >>> import numpy as np + >>> x = np.arange(9).reshape((3,3)) + >>> x + array([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + + >>> np.diag(x) + array([0, 4, 8]) + >>> np.diag(x, k=1) + array([1, 5]) + >>> np.diag(x, k=-1) + array([3, 7]) + + >>> np.diag(np.diag(x)) + array([[0, 0, 0], + [0, 4, 0], + [0, 0, 8]]) + + """ + v = asanyarray(v) + s = v.shape + if len(s) == 1: + n = s[0]+abs(k) + res = zeros((n, n), v.dtype) + if k >= 0: + i = k + else: + i = (-k) * n + res[:n-k].flat[i::n+1] = v + return res + elif len(s) == 2: + return diagonal(v, k) + else: + raise ValueError("Input must be 1- or 2-d.") + + +@array_function_dispatch(_diag_dispatcher) +def diagflat(v, k=0): + """ + Create a two-dimensional array with the flattened input as a diagonal. + + Parameters + ---------- + v : array_like + Input data, which is flattened and set as the `k`-th + diagonal of the output. + k : int, optional + Diagonal to set; 0, the default, corresponds to the "main" diagonal, + a positive (negative) `k` giving the number of the diagonal above + (below) the main. + + Returns + ------- + out : ndarray + The 2-D output array. + + See Also + -------- + diag : MATLAB work-alike for 1-D and 2-D arrays. + diagonal : Return specified diagonals. + trace : Sum along diagonals. + + Examples + -------- + >>> import numpy as np + >>> np.diagflat([[1,2], [3,4]]) + array([[1, 0, 0, 0], + [0, 2, 0, 0], + [0, 0, 3, 0], + [0, 0, 0, 4]]) + + >>> np.diagflat([1,2], 1) + array([[0, 1, 0], + [0, 0, 2], + [0, 0, 0]]) + + """ + conv = _array_converter(v) + v, = conv.as_arrays(subok=False) + v = v.ravel() + s = len(v) + n = s + abs(k) + res = zeros((n, n), v.dtype) + if (k >= 0): + i = arange(0, n-k, dtype=intp) + fi = i+k+i*n + else: + i = arange(0, n+k, dtype=intp) + fi = i+(i-k)*n + res.flat[fi] = v + + return conv.wrap(res) + + +@finalize_array_function_like +@set_module('numpy') +def tri(N, M=None, k=0, dtype=float, *, like=None): + """ + An array with ones at and below the given diagonal and zeros elsewhere. + + Parameters + ---------- + N : int + Number of rows in the array. + M : int, optional + Number of columns in the array. + By default, `M` is taken equal to `N`. + k : int, optional + The sub-diagonal at and below which the array is filled. + `k` = 0 is the main diagonal, while `k` < 0 is below it, + and `k` > 0 is above. The default is 0. + dtype : dtype, optional + Data type of the returned array. The default is float. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 + + Returns + ------- + tri : ndarray of shape (N, M) + Array with its lower triangle filled with ones and zero elsewhere; + in other words ``T[i,j] == 1`` for ``j <= i + k``, 0 otherwise. + + Examples + -------- + >>> import numpy as np + >>> np.tri(3, 5, 2, dtype=int) + array([[1, 1, 1, 0, 0], + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 1]]) + + >>> np.tri(3, 5, -1) + array([[0., 0., 0., 0., 0.], + [1., 0., 0., 0., 0.], + [1., 1., 0., 0., 0.]]) + + """ + if like is not None: + return _tri_with_like(like, N, M=M, k=k, dtype=dtype) + + if M is None: + M = N + + m = greater_equal.outer(arange(N, dtype=_min_int(0, N)), + arange(-k, M-k, dtype=_min_int(-k, M - k))) + + # Avoid making a copy if the requested type is already bool + m = m.astype(dtype, copy=False) + + return m + + +_tri_with_like = array_function_dispatch()(tri) + + +def _trilu_dispatcher(m, k=None): + return (m,) + + +@array_function_dispatch(_trilu_dispatcher) +def tril(m, k=0): + """ + Lower triangle of an array. + + Return a copy of an array with elements above the `k`-th diagonal zeroed. + For arrays with ``ndim`` exceeding 2, `tril` will apply to the final two + axes. + + Parameters + ---------- + m : array_like, shape (..., M, N) + Input array. + k : int, optional + Diagonal above which to zero elements. `k = 0` (the default) is the + main diagonal, `k < 0` is below it and `k > 0` is above. + + Returns + ------- + tril : ndarray, shape (..., M, N) + Lower triangle of `m`, of same shape and data-type as `m`. + + See Also + -------- + triu : same thing, only for the upper triangle + + Examples + -------- + >>> import numpy as np + >>> np.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) + array([[ 0, 0, 0], + [ 4, 0, 0], + [ 7, 8, 0], + [10, 11, 12]]) + + >>> np.tril(np.arange(3*4*5).reshape(3, 4, 5)) + array([[[ 0, 0, 0, 0, 0], + [ 5, 6, 0, 0, 0], + [10, 11, 12, 0, 0], + [15, 16, 17, 18, 0]], + [[20, 0, 0, 0, 0], + [25, 26, 0, 0, 0], + [30, 31, 32, 0, 0], + [35, 36, 37, 38, 0]], + [[40, 0, 0, 0, 0], + [45, 46, 0, 0, 0], + [50, 51, 52, 0, 0], + [55, 56, 57, 58, 0]]]) + + """ + m = asanyarray(m) + mask = tri(*m.shape[-2:], k=k, dtype=bool) + + return where(mask, m, zeros(1, m.dtype)) + + +@array_function_dispatch(_trilu_dispatcher) +def triu(m, k=0): + """ + Upper triangle of an array. + + Return a copy of an array with the elements below the `k`-th diagonal + zeroed. For arrays with ``ndim`` exceeding 2, `triu` will apply to the + final two axes. + + Please refer to the documentation for `tril` for further details. + + See Also + -------- + tril : lower triangle of an array + + Examples + -------- + >>> import numpy as np + >>> np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1) + array([[ 1, 2, 3], + [ 4, 5, 6], + [ 0, 8, 9], + [ 0, 0, 12]]) + + >>> np.triu(np.arange(3*4*5).reshape(3, 4, 5)) + array([[[ 0, 1, 2, 3, 4], + [ 0, 6, 7, 8, 9], + [ 0, 0, 12, 13, 14], + [ 0, 0, 0, 18, 19]], + [[20, 21, 22, 23, 24], + [ 0, 26, 27, 28, 29], + [ 0, 0, 32, 33, 34], + [ 0, 0, 0, 38, 39]], + [[40, 41, 42, 43, 44], + [ 0, 46, 47, 48, 49], + [ 0, 0, 52, 53, 54], + [ 0, 0, 0, 58, 59]]]) + + """ + m = asanyarray(m) + mask = tri(*m.shape[-2:], k=k-1, dtype=bool) + + return where(mask, zeros(1, m.dtype), m) + + +def _vander_dispatcher(x, N=None, increasing=None): + return (x,) + + +# Originally borrowed from John Hunter and matplotlib +@array_function_dispatch(_vander_dispatcher) +def vander(x, N=None, increasing=False): + """ + Generate a Vandermonde matrix. + + The columns of the output matrix are powers of the input vector. The + order of the powers is determined by the `increasing` boolean argument. + Specifically, when `increasing` is False, the `i`-th output column is + the input vector raised element-wise to the power of ``N - i - 1``. Such + a matrix with a geometric progression in each row is named for Alexandre- + Theophile Vandermonde. + + Parameters + ---------- + x : array_like + 1-D input array. + N : int, optional + Number of columns in the output. If `N` is not specified, a square + array is returned (``N = len(x)``). + increasing : bool, optional + Order of the powers of the columns. If True, the powers increase + from left to right, if False (the default) they are reversed. + + Returns + ------- + out : ndarray + Vandermonde matrix. If `increasing` is False, the first column is + ``x^(N-1)``, the second ``x^(N-2)`` and so forth. If `increasing` is + True, the columns are ``x^0, x^1, ..., x^(N-1)``. + + See Also + -------- + polynomial.polynomial.polyvander + + Examples + -------- + >>> import numpy as np + >>> x = np.array([1, 2, 3, 5]) + >>> N = 3 + >>> np.vander(x, N) + array([[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]]) + + >>> np.column_stack([x**(N-1-i) for i in range(N)]) + array([[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]]) + + >>> x = np.array([1, 2, 3, 5]) + >>> np.vander(x) + array([[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]]) + >>> np.vander(x, increasing=True) + array([[ 1, 1, 1, 1], + [ 1, 2, 4, 8], + [ 1, 3, 9, 27], + [ 1, 5, 25, 125]]) + + The determinant of a square Vandermonde matrix is the product + of the differences between the values of the input vector: + + >>> np.linalg.det(np.vander(x)) + 48.000000000000043 # may vary + >>> (5-3)*(5-2)*(5-1)*(3-2)*(3-1)*(2-1) + 48 + + """ + x = asarray(x) + if x.ndim != 1: + raise ValueError("x must be a one-dimensional array or sequence.") + if N is None: + N = len(x) + + v = empty((len(x), N), dtype=promote_types(x.dtype, int)) + tmp = v[:, ::-1] if not increasing else v + + if N > 0: + tmp[:, 0] = 1 + if N > 1: + tmp[:, 1:] = x[:, None] + multiply.accumulate(tmp[:, 1:], out=tmp[:, 1:], axis=1) + + return v + + +def _histogram2d_dispatcher(x, y, bins=None, range=None, density=None, + weights=None): + yield x + yield y + + # This terrible logic is adapted from the checks in histogram2d + try: + N = len(bins) + except TypeError: + N = 1 + if N == 2: + yield from bins # bins=[x, y] + else: + yield bins + + yield weights + + +@array_function_dispatch(_histogram2d_dispatcher) +def histogram2d(x, y, bins=10, range=None, density=None, weights=None): + """ + Compute the bi-dimensional histogram of two data samples. + + Parameters + ---------- + x : array_like, shape (N,) + An array containing the x coordinates of the points to be + histogrammed. + y : array_like, shape (N,) + An array containing the y coordinates of the points to be + histogrammed. + bins : int or array_like or [int, int] or [array, array], optional + The bin specification: + + * If int, the number of bins for the two dimensions (nx=ny=bins). + * If array_like, the bin edges for the two dimensions + (x_edges=y_edges=bins). + * If [int, int], the number of bins in each dimension + (nx, ny = bins). + * If [array, array], the bin edges in each dimension + (x_edges, y_edges = bins). + * A combination [int, array] or [array, int], where int + is the number of bins and array is the bin edges. + + range : array_like, shape(2,2), optional + The leftmost and rightmost edges of the bins along each dimension + (if not specified explicitly in the `bins` parameters): + ``[[xmin, xmax], [ymin, ymax]]``. All values outside of this range + will be considered outliers and not tallied in the histogram. + density : bool, optional + If False, the default, returns the number of samples in each bin. + If True, returns the probability *density* function at the bin, + ``bin_count / sample_count / bin_area``. + weights : array_like, shape(N,), optional + An array of values ``w_i`` weighing each sample ``(x_i, y_i)``. + Weights are normalized to 1 if `density` is True. If `density` is + False, the values of the returned histogram are equal to the sum of + the weights belonging to the samples falling into each bin. + + Returns + ------- + H : ndarray, shape(nx, ny) + The bi-dimensional histogram of samples `x` and `y`. Values in `x` + are histogrammed along the first dimension and values in `y` are + histogrammed along the second dimension. + xedges : ndarray, shape(nx+1,) + The bin edges along the first dimension. + yedges : ndarray, shape(ny+1,) + The bin edges along the second dimension. + + See Also + -------- + histogram : 1D histogram + histogramdd : Multidimensional histogram + + Notes + ----- + When `density` is True, then the returned histogram is the sample + density, defined such that the sum over bins of the product + ``bin_value * bin_area`` is 1. + + Please note that the histogram does not follow the Cartesian convention + where `x` values are on the abscissa and `y` values on the ordinate + axis. Rather, `x` is histogrammed along the first dimension of the + array (vertical), and `y` along the second dimension of the array + (horizontal). This ensures compatibility with `histogramdd`. + + Examples + -------- + >>> import numpy as np + >>> from matplotlib.image import NonUniformImage + >>> import matplotlib.pyplot as plt + + Construct a 2-D histogram with variable bin width. First define the bin + edges: + + >>> xedges = [0, 1, 3, 5] + >>> yedges = [0, 2, 3, 4, 6] + + Next we create a histogram H with random bin content: + + >>> x = np.random.normal(2, 1, 100) + >>> y = np.random.normal(1, 1, 100) + >>> H, xedges, yedges = np.histogram2d(x, y, bins=(xedges, yedges)) + >>> # Histogram does not follow Cartesian convention (see Notes), + >>> # therefore transpose H for visualization purposes. + >>> H = H.T + + :func:`imshow ` can only display square bins: + + >>> fig = plt.figure(figsize=(7, 3)) + >>> ax = fig.add_subplot(131, title='imshow: square bins') + >>> plt.imshow(H, interpolation='nearest', origin='lower', + ... extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]]) + + + :func:`pcolormesh ` can display actual edges: + + >>> ax = fig.add_subplot(132, title='pcolormesh: actual edges', + ... aspect='equal') + >>> X, Y = np.meshgrid(xedges, yedges) + >>> ax.pcolormesh(X, Y, H) + + + :class:`NonUniformImage ` can be used to + display actual bin edges with interpolation: + + >>> ax = fig.add_subplot(133, title='NonUniformImage: interpolated', + ... aspect='equal', xlim=xedges[[0, -1]], ylim=yedges[[0, -1]]) + >>> im = NonUniformImage(ax, interpolation='bilinear') + >>> xcenters = (xedges[:-1] + xedges[1:]) / 2 + >>> ycenters = (yedges[:-1] + yedges[1:]) / 2 + >>> im.set_data(xcenters, ycenters, H) + >>> ax.add_image(im) + >>> plt.show() + + It is also possible to construct a 2-D histogram without specifying bin + edges: + + >>> # Generate non-symmetric test data + >>> n = 10000 + >>> x = np.linspace(1, 100, n) + >>> y = 2*np.log(x) + np.random.rand(n) - 0.5 + >>> # Compute 2d histogram. Note the order of x/y and xedges/yedges + >>> H, yedges, xedges = np.histogram2d(y, x, bins=20) + + Now we can plot the histogram using + :func:`pcolormesh `, and a + :func:`hexbin ` for comparison. + + >>> # Plot histogram using pcolormesh + >>> fig, (ax1, ax2) = plt.subplots(ncols=2, sharey=True) + >>> ax1.pcolormesh(xedges, yedges, H, cmap='rainbow') + >>> ax1.plot(x, 2*np.log(x), 'k-') + >>> ax1.set_xlim(x.min(), x.max()) + >>> ax1.set_ylim(y.min(), y.max()) + >>> ax1.set_xlabel('x') + >>> ax1.set_ylabel('y') + >>> ax1.set_title('histogram2d') + >>> ax1.grid() + + >>> # Create hexbin plot for comparison + >>> ax2.hexbin(x, y, gridsize=20, cmap='rainbow') + >>> ax2.plot(x, 2*np.log(x), 'k-') + >>> ax2.set_title('hexbin') + >>> ax2.set_xlim(x.min(), x.max()) + >>> ax2.set_xlabel('x') + >>> ax2.grid() + + >>> plt.show() + """ + from numpy import histogramdd + + if len(x) != len(y): + raise ValueError('x and y must have the same length.') + + try: + N = len(bins) + except TypeError: + N = 1 + + if N != 1 and N != 2: + xedges = yedges = asarray(bins) + bins = [xedges, yedges] + hist, edges = histogramdd([x, y], bins, range, density, weights) + return hist, edges[0], edges[1] + + +@set_module('numpy') +def mask_indices(n, mask_func, k=0): + """ + Return the indices to access (n, n) arrays, given a masking function. + + Assume `mask_func` is a function that, for a square array a of size + ``(n, n)`` with a possible offset argument `k`, when called as + ``mask_func(a, k)`` returns a new array with zeros in certain locations + (functions like `triu` or `tril` do precisely this). Then this function + returns the indices where the non-zero values would be located. + + Parameters + ---------- + n : int + The returned indices will be valid to access arrays of shape (n, n). + mask_func : callable + A function whose call signature is similar to that of `triu`, `tril`. + That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. + `k` is an optional argument to the function. + k : scalar + An optional argument which is passed through to `mask_func`. Functions + like `triu`, `tril` take a second argument that is interpreted as an + offset. + + Returns + ------- + indices : tuple of arrays. + The `n` arrays of indices corresponding to the locations where + ``mask_func(np.ones((n, n)), k)`` is True. + + See Also + -------- + triu, tril, triu_indices, tril_indices + + Examples + -------- + >>> import numpy as np + + These are the indices that would allow you to access the upper triangular + part of any 3x3 array: + + >>> iu = np.mask_indices(3, np.triu) + + For example, if `a` is a 3x3 array: + + >>> a = np.arange(9).reshape(3, 3) + >>> a + array([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> a[iu] + array([0, 1, 2, 4, 5, 8]) + + An offset can be passed also to the masking function. This gets us the + indices starting on the first diagonal right of the main one: + + >>> iu1 = np.mask_indices(3, np.triu, 1) + + with which we now extract only three elements: + + >>> a[iu1] + array([1, 2, 5]) + + """ + m = ones((n, n), int) + a = mask_func(m, k) + return nonzero(a != 0) + + +@set_module('numpy') +def tril_indices(n, k=0, m=None): + """ + Return the indices for the lower-triangle of an (n, m) array. + + Parameters + ---------- + n : int + The row dimension of the arrays for which the returned + indices will be valid. + k : int, optional + Diagonal offset (see `tril` for details). + m : int, optional + The column dimension of the arrays for which the returned + arrays will be valid. + By default `m` is taken equal to `n`. + + + Returns + ------- + inds : tuple of arrays + The row and column indices, respectively. The row indices are sorted + in non-decreasing order, and the correspdonding column indices are + strictly increasing for each row. + + See also + -------- + triu_indices : similar function, for upper-triangular. + mask_indices : generic function accepting an arbitrary mask function. + tril, triu + + Examples + -------- + >>> import numpy as np + + Compute two different sets of indices to access 4x4 arrays, one for the + lower triangular part starting at the main diagonal, and one starting two + diagonals further right: + + >>> il1 = np.tril_indices(4) + >>> il1 + (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + Note that row indices (first array) are non-decreasing, and the corresponding + column indices (second array) are strictly increasing for each row. + Here is how they can be used with a sample array: + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Both for indexing: + + >>> a[il1] + array([ 0, 4, 5, ..., 13, 14, 15]) + + And for assigning values: + + >>> a[il1] = -1 + >>> a + array([[-1, 1, 2, 3], + [-1, -1, 6, 7], + [-1, -1, -1, 11], + [-1, -1, -1, -1]]) + + These cover almost the whole array (two diagonals right of the main one): + + >>> il2 = np.tril_indices(4, 2) + >>> a[il2] = -10 + >>> a + array([[-10, -10, -10, 3], + [-10, -10, -10, -10], + [-10, -10, -10, -10], + [-10, -10, -10, -10]]) + + """ + tri_ = tri(n, m, k=k, dtype=bool) + + return tuple(broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True)) + + +def _trilu_indices_form_dispatcher(arr, k=None): + return (arr,) + + +@array_function_dispatch(_trilu_indices_form_dispatcher) +def tril_indices_from(arr, k=0): + """ + Return the indices for the lower-triangle of arr. + + See `tril_indices` for full details. + + Parameters + ---------- + arr : array_like + The indices will be valid for square arrays whose dimensions are + the same as arr. + k : int, optional + Diagonal offset (see `tril` for details). + + Examples + -------- + >>> import numpy as np + + Create a 4 by 4 array + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Pass the array to get the indices of the lower triangular elements. + + >>> trili = np.tril_indices_from(a) + >>> trili + (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + >>> a[trili] + array([ 0, 4, 5, 8, 9, 10, 12, 13, 14, 15]) + + This is syntactic sugar for tril_indices(). + + >>> np.tril_indices(a.shape[0]) + (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + Use the `k` parameter to return the indices for the lower triangular array + up to the k-th diagonal. + + >>> trili1 = np.tril_indices_from(a, k=1) + >>> a[trili1] + array([ 0, 1, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]) + + See Also + -------- + tril_indices, tril, triu_indices_from + """ + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) + + +@set_module('numpy') +def triu_indices(n, k=0, m=None): + """ + Return the indices for the upper-triangle of an (n, m) array. + + Parameters + ---------- + n : int + The size of the arrays for which the returned indices will + be valid. + k : int, optional + Diagonal offset (see `triu` for details). + m : int, optional + The column dimension of the arrays for which the returned + arrays will be valid. + By default `m` is taken equal to `n`. + + + Returns + ------- + inds : tuple, shape(2) of ndarrays, shape(`n`) + The row and column indices, respectively. The row indices are sorted + in non-decreasing order, and the correspdonding column indices are + strictly increasing for each row. + + See also + -------- + tril_indices : similar function, for lower-triangular. + mask_indices : generic function accepting an arbitrary mask function. + triu, tril + + Examples + -------- + >>> import numpy as np + + Compute two different sets of indices to access 4x4 arrays, one for the + upper triangular part starting at the main diagonal, and one starting two + diagonals further right: + + >>> iu1 = np.triu_indices(4) + >>> iu1 + (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])) + + Note that row indices (first array) are non-decreasing, and the corresponding + column indices (second array) are strictly increasing for each row. + + Here is how they can be used with a sample array: + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Both for indexing: + + >>> a[iu1] + array([ 0, 1, 2, ..., 10, 11, 15]) + + And for assigning values: + + >>> a[iu1] = -1 + >>> a + array([[-1, -1, -1, -1], + [ 4, -1, -1, -1], + [ 8, 9, -1, -1], + [12, 13, 14, -1]]) + + These cover only a small part of the whole array (two diagonals right + of the main one): + + >>> iu2 = np.triu_indices(4, 2) + >>> a[iu2] = -10 + >>> a + array([[ -1, -1, -10, -10], + [ 4, -1, -1, -10], + [ 8, 9, -1, -1], + [ 12, 13, 14, -1]]) + + """ + tri_ = ~tri(n, m, k=k - 1, dtype=bool) + + return tuple(broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True)) + + +@array_function_dispatch(_trilu_indices_form_dispatcher) +def triu_indices_from(arr, k=0): + """ + Return the indices for the upper-triangle of arr. + + See `triu_indices` for full details. + + Parameters + ---------- + arr : ndarray, shape(N, N) + The indices will be valid for square arrays. + k : int, optional + Diagonal offset (see `triu` for details). + + Returns + ------- + triu_indices_from : tuple, shape(2) of ndarray, shape(N) + Indices for the upper-triangle of `arr`. + + Examples + -------- + >>> import numpy as np + + Create a 4 by 4 array + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Pass the array to get the indices of the upper triangular elements. + + >>> triui = np.triu_indices_from(a) + >>> triui + (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])) + + >>> a[triui] + array([ 0, 1, 2, 3, 5, 6, 7, 10, 11, 15]) + + This is syntactic sugar for triu_indices(). + + >>> np.triu_indices(a.shape[0]) + (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])) + + Use the `k` parameter to return the indices for the upper triangular array + from the k-th diagonal. + + >>> triuim1 = np.triu_indices_from(a, k=1) + >>> a[triuim1] + array([ 1, 2, 3, 6, 7, 11]) + + + See Also + -------- + triu_indices, triu, tril_indices_from + """ + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..5d3ea54511b8c54efaa47ae4b7ecb71cd53b394a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_twodim_base_impl.pyi @@ -0,0 +1,437 @@ +from collections.abc import Callable, Sequence +from typing import ( + Any, + TypeAlias, + overload, + TypeVar, + Literal as L, +) + +import numpy as np +from numpy import ( + generic, + timedelta64, + datetime64, + int_, + intp, + float64, + complex128, + signedinteger, + floating, + complexfloating, + object_, + _OrderCF, +) + +from numpy._typing import ( + DTypeLike, + _DTypeLike, + ArrayLike, + _ArrayLike, + NDArray, + _SupportsArray, + _SupportsArrayFunc, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeObject_co, +) + +__all__ = [ + "diag", + "diagflat", + "eye", + "fliplr", + "flipud", + "tri", + "triu", + "tril", + "vander", + "histogram2d", + "mask_indices", + "tril_indices", + "tril_indices_from", + "triu_indices", + "triu_indices_from", +] + +### + +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=generic) +_SCT_complex = TypeVar("_SCT_complex", bound=np.complexfloating) +_SCT_inexact = TypeVar("_SCT_inexact", bound=np.inexact) +_SCT_number_co = TypeVar("_SCT_number_co", bound=_Number_co) + +# The returned arrays dtype must be compatible with `np.equal` +_MaskFunc: TypeAlias = Callable[[NDArray[int_], _T], NDArray[_Number_co | timedelta64 | datetime64 | object_]] + +_Int_co: TypeAlias = np.integer | np.bool +_Float_co: TypeAlias = np.floating | _Int_co +_Number_co: TypeAlias = np.number | np.bool + +_ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_SCT]] | Sequence[_SCT] +_ArrayLike1DInt_co: TypeAlias = _SupportsArray[np.dtype[_Int_co]] | Sequence[int | _Int_co] +_ArrayLike1DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[float | _Float_co] +_ArrayLike2DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[_ArrayLike1DFloat_co] +_ArrayLike1DNumber_co: TypeAlias = _SupportsArray[np.dtype[_Number_co]] | Sequence[complex | _Number_co] + +### + +@overload +def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def fliplr(m: ArrayLike) -> NDArray[Any]: ... + +@overload +def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... +@overload +def flipud(m: ArrayLike) -> NDArray[Any]: ... + +@overload +def eye( + N: int, + M: None | int = ..., + k: int = ..., + dtype: None = ..., + order: _OrderCF = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[float64]: ... +@overload +def eye( + N: int, + M: None | int, + k: int, + dtype: _DTypeLike[_SCT], + order: _OrderCF = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def eye( + N: int, + M: None | int = ..., + k: int = ..., + *, + dtype: _DTypeLike[_SCT], + order: _OrderCF = ..., + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[_SCT]: ... +@overload +def eye( + N: int, + M: None | int = ..., + k: int = ..., + dtype: DTypeLike = ..., + order: _OrderCF = ..., + *, + device: None | L["cpu"] = ..., + like: None | _SupportsArrayFunc = ..., +) -> NDArray[Any]: ... + +@overload +def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... +@overload +def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... + +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: None = ..., + *, + like: None | _SupportsArrayFunc = ... +) -> NDArray[float64]: ... +@overload +def tri( + N: int, + M: None | int, + k: int, + dtype: _DTypeLike[_SCT], + *, + like: None | _SupportsArrayFunc = ... +) -> NDArray[_SCT]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + *, + dtype: _DTypeLike[_SCT], + like: None | _SupportsArrayFunc = ... +) -> NDArray[_SCT]: ... +@overload +def tri( + N: int, + M: None | int = ..., + k: int = ..., + dtype: DTypeLike = ..., + *, + like: None | _SupportsArrayFunc = ... +) -> NDArray[Any]: ... + +@overload +def tril(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ... +@overload +def tril(m: ArrayLike, k: int = 0) -> NDArray[Any]: ... + +@overload +def triu(m: _ArrayLike[_SCT], k: int = 0) -> NDArray[_SCT]: ... +@overload +def triu(m: ArrayLike, k: int = 0) -> NDArray[Any]: ... + +@overload +def vander( # type: ignore[misc] + x: _ArrayLikeInt_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[signedinteger[Any]]: ... +@overload +def vander( # type: ignore[misc] + x: _ArrayLikeFloat_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[floating[Any]]: ... +@overload +def vander( + x: _ArrayLikeComplex_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[complexfloating[Any, Any]]: ... +@overload +def vander( + x: _ArrayLikeObject_co, + N: None | int = ..., + increasing: bool = ..., +) -> NDArray[object_]: ... + +@overload +def histogram2d( + x: _ArrayLike1D[_SCT_complex], + y: _ArrayLike1D[_SCT_complex | _Float_co], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_complex], + NDArray[_SCT_complex], +]: ... +@overload +def histogram2d( + x: _ArrayLike1D[_SCT_complex | _Float_co], + y: _ArrayLike1D[_SCT_complex], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_complex], + NDArray[_SCT_complex], +]: ... +@overload +def histogram2d( + x: _ArrayLike1D[_SCT_inexact], + y: _ArrayLike1D[_SCT_inexact | _Int_co], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_inexact], + NDArray[_SCT_inexact], +]: ... +@overload +def histogram2d( + x: _ArrayLike1D[_SCT_inexact | _Int_co], + y: _ArrayLike1D[_SCT_inexact], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_inexact], + NDArray[_SCT_inexact], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DInt_co | Sequence[float | int], + y: _ArrayLike1DInt_co | Sequence[float | int], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[float64], + NDArray[float64], +]: ... +@overload +def histogram2d( + x: Sequence[complex | float | int], + y: Sequence[complex | float | int], + bins: int | Sequence[int] = ..., + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[complex128 | float64], + NDArray[complex128 | float64], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DNumber_co, + y: _ArrayLike1DNumber_co, + bins: _ArrayLike1D[_SCT_number_co] | Sequence[_ArrayLike1D[_SCT_number_co]], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_number_co], + NDArray[_SCT_number_co], +]: ... +@overload +def histogram2d( + x: _ArrayLike1D[_SCT_inexact], + y: _ArrayLike1D[_SCT_inexact], + bins: Sequence[_ArrayLike1D[_SCT_number_co] | int], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_number_co | _SCT_inexact], + NDArray[_SCT_number_co | _SCT_inexact], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DInt_co | Sequence[float | int], + y: _ArrayLike1DInt_co | Sequence[float | int], + bins: Sequence[_ArrayLike1D[_SCT_number_co] | int], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_number_co | float64], + NDArray[_SCT_number_co | float64], +]: ... +@overload +def histogram2d( + x: Sequence[complex | float | int], + y: Sequence[complex | float | int], + bins: Sequence[_ArrayLike1D[_SCT_number_co] | int], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[_SCT_number_co | complex128 | float64], + NDArray[_SCT_number_co | complex128 | float64] , +]: ... +@overload +def histogram2d( + x: _ArrayLike1DNumber_co, + y: _ArrayLike1DNumber_co, + bins: Sequence[Sequence[bool]], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[np.bool], + NDArray[np.bool], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DNumber_co, + y: _ArrayLike1DNumber_co, + bins: Sequence[Sequence[int | bool]], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[np.int_ | np.bool], + NDArray[np.int_ | np.bool], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DNumber_co, + y: _ArrayLike1DNumber_co, + bins: Sequence[Sequence[float | int | bool]], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[np.float64 | np.int_ | np.bool], + NDArray[np.float64 | np.int_ | np.bool], +]: ... +@overload +def histogram2d( + x: _ArrayLike1DNumber_co, + y: _ArrayLike1DNumber_co, + bins: Sequence[Sequence[complex | float | int | bool]], + range: None | _ArrayLike2DFloat_co = ..., + density: None | bool = ..., + weights: None | _ArrayLike1DFloat_co = ..., +) -> tuple[ + NDArray[float64], + NDArray[np.complex128 | np.float64 | np.int_ | np.bool], + NDArray[np.complex128 | np.float64 | np.int_ | np.bool], +]: ... + +# NOTE: we're assuming/demanding here the `mask_func` returns +# an ndarray of shape `(n, n)`; otherwise there is the possibility +# of the output tuple having more or less than 2 elements +@overload +def mask_indices( + n: int, + mask_func: _MaskFunc[int], + k: int = ..., +) -> tuple[NDArray[intp], NDArray[intp]]: ... +@overload +def mask_indices( + n: int, + mask_func: _MaskFunc[_T], + k: _T, +) -> tuple[NDArray[intp], NDArray[intp]]: ... + +def tril_indices( + n: int, + k: int = ..., + m: None | int = ..., +) -> tuple[NDArray[int_], NDArray[int_]]: ... + +def tril_indices_from( + arr: NDArray[Any], + k: int = ..., +) -> tuple[NDArray[int_], NDArray[int_]]: ... + +def triu_indices( + n: int, + k: int = ..., + m: None | int = ..., +) -> tuple[NDArray[int_], NDArray[int_]]: ... + +def triu_indices_from( + arr: NDArray[Any], + k: int = ..., +) -> tuple[NDArray[int_], NDArray[int_]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..e5c9ffbbb8d41049c17f65d56f5afe23a5d5bbcd --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.py @@ -0,0 +1,699 @@ +"""Automatically adapted for numpy Sep 19, 2005 by convertcode.py + +""" +import functools + +__all__ = ['iscomplexobj', 'isrealobj', 'imag', 'iscomplex', + 'isreal', 'nan_to_num', 'real', 'real_if_close', + 'typename', 'mintypecode', + 'common_type'] + +from .._utils import set_module +import numpy._core.numeric as _nx +from numpy._core.numeric import asarray, asanyarray, isnan, zeros +from numpy._core import overrides, getlimits +from ._ufunclike_impl import isneginf, isposinf + + +array_function_dispatch = functools.partial( + overrides.array_function_dispatch, module='numpy') + + +_typecodes_by_elsize = 'GDFgdfQqLlIiHhBb?' + + +@set_module('numpy') +def mintypecode(typechars, typeset='GDFgdf', default='d'): + """ + Return the character for the minimum-size type to which given types can + be safely cast. + + The returned type character must represent the smallest size dtype such + that an array of the returned type can handle the data from an array of + all types in `typechars` (or if `typechars` is an array, then its + dtype.char). + + Parameters + ---------- + typechars : list of str or array_like + If a list of strings, each string should represent a dtype. + If array_like, the character representation of the array dtype is used. + typeset : str or list of str, optional + The set of characters that the returned character is chosen from. + The default set is 'GDFgdf'. + default : str, optional + The default character, this is returned if none of the characters in + `typechars` matches a character in `typeset`. + + Returns + ------- + typechar : str + The character representing the minimum-size type that was found. + + See Also + -------- + dtype + + Examples + -------- + >>> import numpy as np + >>> np.mintypecode(['d', 'f', 'S']) + 'd' + >>> x = np.array([1.1, 2-3.j]) + >>> np.mintypecode(x) + 'D' + + >>> np.mintypecode('abceh', default='G') + 'G' + + """ + typecodes = ((isinstance(t, str) and t) or asarray(t).dtype.char + for t in typechars) + intersection = set(t for t in typecodes if t in typeset) + if not intersection: + return default + if 'F' in intersection and 'd' in intersection: + return 'D' + return min(intersection, key=_typecodes_by_elsize.index) + + +def _real_dispatcher(val): + return (val,) + + +@array_function_dispatch(_real_dispatcher) +def real(val): + """ + Return the real part of the complex argument. + + Parameters + ---------- + val : array_like + Input array. + + Returns + ------- + out : ndarray or scalar + The real component of the complex argument. If `val` is real, the type + of `val` is used for the output. If `val` has complex elements, the + returned type is float. + + See Also + -------- + real_if_close, imag, angle + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1+2j, 3+4j, 5+6j]) + >>> a.real + array([1., 3., 5.]) + >>> a.real = 9 + >>> a + array([9.+2.j, 9.+4.j, 9.+6.j]) + >>> a.real = np.array([9, 8, 7]) + >>> a + array([9.+2.j, 8.+4.j, 7.+6.j]) + >>> np.real(1 + 1j) + 1.0 + + """ + try: + return val.real + except AttributeError: + return asanyarray(val).real + + +def _imag_dispatcher(val): + return (val,) + + +@array_function_dispatch(_imag_dispatcher) +def imag(val): + """ + Return the imaginary part of the complex argument. + + Parameters + ---------- + val : array_like + Input array. + + Returns + ------- + out : ndarray or scalar + The imaginary component of the complex argument. If `val` is real, + the type of `val` is used for the output. If `val` has complex + elements, the returned type is float. + + See Also + -------- + real, angle, real_if_close + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1+2j, 3+4j, 5+6j]) + >>> a.imag + array([2., 4., 6.]) + >>> a.imag = np.array([8, 10, 12]) + >>> a + array([1. +8.j, 3.+10.j, 5.+12.j]) + >>> np.imag(1 + 1j) + 1.0 + + """ + try: + return val.imag + except AttributeError: + return asanyarray(val).imag + + +def _is_type_dispatcher(x): + return (x,) + + +@array_function_dispatch(_is_type_dispatcher) +def iscomplex(x): + """ + Returns a bool array, where True if input element is complex. + + What is tested is whether the input has a non-zero imaginary part, not if + the input type is complex. + + Parameters + ---------- + x : array_like + Input array. + + Returns + ------- + out : ndarray of bools + Output array. + + See Also + -------- + isreal + iscomplexobj : Return True if x is a complex type or an array of complex + numbers. + + Examples + -------- + >>> import numpy as np + >>> np.iscomplex([1+1j, 1+0j, 4.5, 3, 2, 2j]) + array([ True, False, False, False, False, True]) + + """ + ax = asanyarray(x) + if issubclass(ax.dtype.type, _nx.complexfloating): + return ax.imag != 0 + res = zeros(ax.shape, bool) + return res[()] # convert to scalar if needed + + +@array_function_dispatch(_is_type_dispatcher) +def isreal(x): + """ + Returns a bool array, where True if input element is real. + + If element has complex type with zero imaginary part, the return value + for that element is True. + + Parameters + ---------- + x : array_like + Input array. + + Returns + ------- + out : ndarray, bool + Boolean array of same shape as `x`. + + Notes + ----- + `isreal` may behave unexpectedly for string or object arrays (see examples) + + See Also + -------- + iscomplex + isrealobj : Return True if x is not a complex type. + + Examples + -------- + >>> import numpy as np + >>> a = np.array([1+1j, 1+0j, 4.5, 3, 2, 2j], dtype=complex) + >>> np.isreal(a) + array([False, True, True, True, True, False]) + + The function does not work on string arrays. + + >>> a = np.array([2j, "a"], dtype="U") + >>> np.isreal(a) # Warns about non-elementwise comparison + False + + Returns True for all elements in input array of ``dtype=object`` even if + any of the elements is complex. + + >>> a = np.array([1, "2", 3+4j], dtype=object) + >>> np.isreal(a) + array([ True, True, True]) + + isreal should not be used with object arrays + + >>> a = np.array([1+2j, 2+1j], dtype=object) + >>> np.isreal(a) + array([ True, True]) + + """ + return imag(x) == 0 + + +@array_function_dispatch(_is_type_dispatcher) +def iscomplexobj(x): + """ + Check for a complex type or an array of complex numbers. + + The type of the input is checked, not the value. Even if the input + has an imaginary part equal to zero, `iscomplexobj` evaluates to True. + + Parameters + ---------- + x : any + The input can be of any type and shape. + + Returns + ------- + iscomplexobj : bool + The return value, True if `x` is of a complex type or has at least + one complex element. + + See Also + -------- + isrealobj, iscomplex + + Examples + -------- + >>> import numpy as np + >>> np.iscomplexobj(1) + False + >>> np.iscomplexobj(1+0j) + True + >>> np.iscomplexobj([3, 1+0j, True]) + True + + """ + try: + dtype = x.dtype + type_ = dtype.type + except AttributeError: + type_ = asarray(x).dtype.type + return issubclass(type_, _nx.complexfloating) + + +@array_function_dispatch(_is_type_dispatcher) +def isrealobj(x): + """ + Return True if x is a not complex type or an array of complex numbers. + + The type of the input is checked, not the value. So even if the input + has an imaginary part equal to zero, `isrealobj` evaluates to False + if the data type is complex. + + Parameters + ---------- + x : any + The input can be of any type and shape. + + Returns + ------- + y : bool + The return value, False if `x` is of a complex type. + + See Also + -------- + iscomplexobj, isreal + + Notes + ----- + The function is only meant for arrays with numerical values but it + accepts all other objects. Since it assumes array input, the return + value of other objects may be True. + + >>> np.isrealobj('A string') + True + >>> np.isrealobj(False) + True + >>> np.isrealobj(None) + True + + Examples + -------- + >>> import numpy as np + >>> np.isrealobj(1) + True + >>> np.isrealobj(1+0j) + False + >>> np.isrealobj([3, 1+0j, True]) + False + + """ + return not iscomplexobj(x) + +#----------------------------------------------------------------------------- + +def _getmaxmin(t): + from numpy._core import getlimits + f = getlimits.finfo(t) + return f.max, f.min + + +def _nan_to_num_dispatcher(x, copy=None, nan=None, posinf=None, neginf=None): + return (x,) + + +@array_function_dispatch(_nan_to_num_dispatcher) +def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None): + """ + Replace NaN with zero and infinity with large finite numbers (default + behaviour) or with the numbers defined by the user using the `nan`, + `posinf` and/or `neginf` keywords. + + If `x` is inexact, NaN is replaced by zero or by the user defined value in + `nan` keyword, infinity is replaced by the largest finite floating point + values representable by ``x.dtype`` or by the user defined value in + `posinf` keyword and -infinity is replaced by the most negative finite + floating point values representable by ``x.dtype`` or by the user defined + value in `neginf` keyword. + + For complex dtypes, the above is applied to each of the real and + imaginary components of `x` separately. + + If `x` is not inexact, then no replacements are made. + + Parameters + ---------- + x : scalar or array_like + Input data. + copy : bool, optional + Whether to create a copy of `x` (True) or to replace values + in-place (False). The in-place operation only occurs if + casting to an array does not require a copy. + Default is True. + nan : int, float, optional + Value to be used to fill NaN values. If no value is passed + then NaN values will be replaced with 0.0. + posinf : int, float, optional + Value to be used to fill positive infinity values. If no value is + passed then positive infinity values will be replaced with a very + large number. + neginf : int, float, optional + Value to be used to fill negative infinity values. If no value is + passed then negative infinity values will be replaced with a very + small (or negative) number. + + Returns + ------- + out : ndarray + `x`, with the non-finite values replaced. If `copy` is False, this may + be `x` itself. + + See Also + -------- + isinf : Shows which elements are positive or negative infinity. + isneginf : Shows which elements are negative infinity. + isposinf : Shows which elements are positive infinity. + isnan : Shows which elements are Not a Number (NaN). + isfinite : Shows which elements are finite (not NaN, not infinity) + + Notes + ----- + NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic + (IEEE 754). This means that Not a Number is not equivalent to infinity. + + Examples + -------- + >>> import numpy as np + >>> np.nan_to_num(np.inf) + 1.7976931348623157e+308 + >>> np.nan_to_num(-np.inf) + -1.7976931348623157e+308 + >>> np.nan_to_num(np.nan) + 0.0 + >>> x = np.array([np.inf, -np.inf, np.nan, -128, 128]) + >>> np.nan_to_num(x) + array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary + -1.28000000e+002, 1.28000000e+002]) + >>> np.nan_to_num(x, nan=-9999, posinf=33333333, neginf=33333333) + array([ 3.3333333e+07, 3.3333333e+07, -9.9990000e+03, + -1.2800000e+02, 1.2800000e+02]) + >>> y = np.array([complex(np.inf, np.nan), np.nan, complex(np.nan, np.inf)]) + array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary + -1.28000000e+002, 1.28000000e+002]) + >>> np.nan_to_num(y) + array([ 1.79769313e+308 +0.00000000e+000j, # may vary + 0.00000000e+000 +0.00000000e+000j, + 0.00000000e+000 +1.79769313e+308j]) + >>> np.nan_to_num(y, nan=111111, posinf=222222) + array([222222.+111111.j, 111111. +0.j, 111111.+222222.j]) + """ + x = _nx.array(x, subok=True, copy=copy) + xtype = x.dtype.type + + isscalar = (x.ndim == 0) + + if not issubclass(xtype, _nx.inexact): + return x[()] if isscalar else x + + iscomplex = issubclass(xtype, _nx.complexfloating) + + dest = (x.real, x.imag) if iscomplex else (x,) + maxf, minf = _getmaxmin(x.real.dtype) + if posinf is not None: + maxf = posinf + if neginf is not None: + minf = neginf + for d in dest: + idx_nan = isnan(d) + idx_posinf = isposinf(d) + idx_neginf = isneginf(d) + _nx.copyto(d, nan, where=idx_nan) + _nx.copyto(d, maxf, where=idx_posinf) + _nx.copyto(d, minf, where=idx_neginf) + return x[()] if isscalar else x + +#----------------------------------------------------------------------------- + +def _real_if_close_dispatcher(a, tol=None): + return (a,) + + +@array_function_dispatch(_real_if_close_dispatcher) +def real_if_close(a, tol=100): + """ + If input is complex with all imaginary parts close to zero, return + real parts. + + "Close to zero" is defined as `tol` * (machine epsilon of the type for + `a`). + + Parameters + ---------- + a : array_like + Input array. + tol : float + Tolerance in machine epsilons for the complex part of the elements + in the array. If the tolerance is <=1, then the absolute tolerance + is used. + + Returns + ------- + out : ndarray + If `a` is real, the type of `a` is used for the output. If `a` + has complex elements, the returned type is float. + + See Also + -------- + real, imag, angle + + Notes + ----- + Machine epsilon varies from machine to machine and between data types + but Python floats on most platforms have a machine epsilon equal to + 2.2204460492503131e-16. You can use 'np.finfo(float).eps' to print + out the machine epsilon for floats. + + Examples + -------- + >>> import numpy as np + >>> np.finfo(float).eps + 2.2204460492503131e-16 # may vary + + >>> np.real_if_close([2.1 + 4e-14j, 5.2 + 3e-15j], tol=1000) + array([2.1, 5.2]) + >>> np.real_if_close([2.1 + 4e-13j, 5.2 + 3e-15j], tol=1000) + array([2.1+4.e-13j, 5.2 + 3e-15j]) + + """ + a = asanyarray(a) + type_ = a.dtype.type + if not issubclass(type_, _nx.complexfloating): + return a + if tol > 1: + f = getlimits.finfo(type_) + tol = f.eps * tol + if _nx.all(_nx.absolute(a.imag) < tol): + a = a.real + return a + + +#----------------------------------------------------------------------------- + +_namefromtype = {'S1': 'character', + '?': 'bool', + 'b': 'signed char', + 'B': 'unsigned char', + 'h': 'short', + 'H': 'unsigned short', + 'i': 'integer', + 'I': 'unsigned integer', + 'l': 'long integer', + 'L': 'unsigned long integer', + 'q': 'long long integer', + 'Q': 'unsigned long long integer', + 'f': 'single precision', + 'd': 'double precision', + 'g': 'long precision', + 'F': 'complex single precision', + 'D': 'complex double precision', + 'G': 'complex long double precision', + 'S': 'string', + 'U': 'unicode', + 'V': 'void', + 'O': 'object' + } + +@set_module('numpy') +def typename(char): + """ + Return a description for the given data type code. + + Parameters + ---------- + char : str + Data type code. + + Returns + ------- + out : str + Description of the input data type code. + + See Also + -------- + dtype + + Examples + -------- + >>> import numpy as np + >>> typechars = ['S1', '?', 'B', 'D', 'G', 'F', 'I', 'H', 'L', 'O', 'Q', + ... 'S', 'U', 'V', 'b', 'd', 'g', 'f', 'i', 'h', 'l', 'q'] + >>> for typechar in typechars: + ... print(typechar, ' : ', np.typename(typechar)) + ... + S1 : character + ? : bool + B : unsigned char + D : complex double precision + G : complex long double precision + F : complex single precision + I : unsigned integer + H : unsigned short + L : unsigned long integer + O : object + Q : unsigned long long integer + S : string + U : unicode + V : void + b : signed char + d : double precision + g : long precision + f : single precision + i : integer + h : short + l : long integer + q : long long integer + + """ + return _namefromtype[char] + +#----------------------------------------------------------------------------- + + +#determine the "minimum common type" for a group of arrays. +array_type = [[_nx.float16, _nx.float32, _nx.float64, _nx.longdouble], + [None, _nx.complex64, _nx.complex128, _nx.clongdouble]] +array_precision = {_nx.float16: 0, + _nx.float32: 1, + _nx.float64: 2, + _nx.longdouble: 3, + _nx.complex64: 1, + _nx.complex128: 2, + _nx.clongdouble: 3} + + +def _common_type_dispatcher(*arrays): + return arrays + + +@array_function_dispatch(_common_type_dispatcher) +def common_type(*arrays): + """ + Return a scalar type which is common to the input arrays. + + The return type will always be an inexact (i.e. floating point) scalar + type, even if all the arrays are integer arrays. If one of the inputs is + an integer array, the minimum precision type that is returned is a + 64-bit floating point dtype. + + All input arrays except int64 and uint64 can be safely cast to the + returned dtype without loss of information. + + Parameters + ---------- + array1, array2, ... : ndarrays + Input arrays. + + Returns + ------- + out : data type code + Data type code. + + See Also + -------- + dtype, mintypecode + + Examples + -------- + >>> np.common_type(np.arange(2, dtype=np.float32)) + + >>> np.common_type(np.arange(2, dtype=np.float32), np.arange(2)) + + >>> np.common_type(np.arange(4), np.array([45, 6.j]), np.array([45.0])) + + + """ + is_complex = False + precision = 0 + for a in arrays: + t = a.dtype.type + if iscomplexobj(a): + is_complex = True + if issubclass(t, _nx.integer): + p = 2 # array_precision[_nx.double] + else: + p = array_precision.get(t) + if p is None: + raise TypeError("can't get common type for non-numeric array") + precision = max(precision, p) + if is_complex: + return array_type[1][precision] + else: + return array_type[0][precision] diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e195238103fa16d7051b7aa6b2ac5054a9c9a76b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_type_check_impl.pyi @@ -0,0 +1,201 @@ +from collections.abc import Container, Iterable +from typing import Literal as L, Any, overload, TypeVar + +import numpy as np +from numpy import ( + _HasRealAndImag, + dtype, + generic, + floating, + complexfloating, + integer, +) + +from numpy._typing import ( + ArrayLike, + NBitBase, + NDArray, + _64Bit, + _SupportsDType, + _ScalarLike_co, + _ArrayLike, +) + +__all__ = [ + "iscomplexobj", + "isrealobj", + "imag", + "iscomplex", + "isreal", + "nan_to_num", + "real", + "real_if_close", + "typename", + "mintypecode", + "common_type", +] + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +_SCT = TypeVar("_SCT", bound=generic) +_NBit1 = TypeVar("_NBit1", bound=NBitBase) +_NBit2 = TypeVar("_NBit2", bound=NBitBase) + + +def mintypecode( + typechars: Iterable[str | ArrayLike], + typeset: Container[str] = ..., + default: str = ..., +) -> str: ... + +@overload +def real(val: _HasRealAndImag[_T, Any]) -> _T: ... +@overload +def real(val: ArrayLike) -> NDArray[Any]: ... + +@overload +def imag(val: _HasRealAndImag[Any, _T]) -> _T: ... +@overload +def imag(val: ArrayLike) -> NDArray[Any]: ... + +@overload +def iscomplex(x: _ScalarLike_co) -> np.bool: ... # type: ignore[misc] +@overload +def iscomplex(x: ArrayLike) -> NDArray[np.bool]: ... + +@overload +def isreal(x: _ScalarLike_co) -> np.bool: ... # type: ignore[misc] +@overload +def isreal(x: ArrayLike) -> NDArray[np.bool]: ... + +def iscomplexobj(x: _SupportsDType[dtype[Any]] | ArrayLike) -> bool: ... + +def isrealobj(x: _SupportsDType[dtype[Any]] | ArrayLike) -> bool: ... + +@overload +def nan_to_num( # type: ignore[misc] + x: _SCT, + copy: bool = ..., + nan: float = ..., + posinf: None | float = ..., + neginf: None | float = ..., +) -> _SCT: ... +@overload +def nan_to_num( + x: _ScalarLike_co, + copy: bool = ..., + nan: float = ..., + posinf: None | float = ..., + neginf: None | float = ..., +) -> Any: ... +@overload +def nan_to_num( + x: _ArrayLike[_SCT], + copy: bool = ..., + nan: float = ..., + posinf: None | float = ..., + neginf: None | float = ..., +) -> NDArray[_SCT]: ... +@overload +def nan_to_num( + x: ArrayLike, + copy: bool = ..., + nan: float = ..., + posinf: None | float = ..., + neginf: None | float = ..., +) -> NDArray[Any]: ... + +# If one passes a complex array to `real_if_close`, then one is reasonably +# expected to verify the output dtype (so we can return an unsafe union here) + +@overload +def real_if_close( # type: ignore[misc] + a: _ArrayLike[complexfloating[_NBit1, _NBit1]], + tol: float = ..., +) -> NDArray[floating[_NBit1]] | NDArray[complexfloating[_NBit1, _NBit1]]: ... +@overload +def real_if_close( + a: _ArrayLike[_SCT], + tol: float = ..., +) -> NDArray[_SCT]: ... +@overload +def real_if_close( + a: ArrayLike, + tol: float = ..., +) -> NDArray[Any]: ... + +@overload +def typename(char: L['S1']) -> L['character']: ... +@overload +def typename(char: L['?']) -> L['bool']: ... +@overload +def typename(char: L['b']) -> L['signed char']: ... +@overload +def typename(char: L['B']) -> L['unsigned char']: ... +@overload +def typename(char: L['h']) -> L['short']: ... +@overload +def typename(char: L['H']) -> L['unsigned short']: ... +@overload +def typename(char: L['i']) -> L['integer']: ... +@overload +def typename(char: L['I']) -> L['unsigned integer']: ... +@overload +def typename(char: L['l']) -> L['long integer']: ... +@overload +def typename(char: L['L']) -> L['unsigned long integer']: ... +@overload +def typename(char: L['q']) -> L['long long integer']: ... +@overload +def typename(char: L['Q']) -> L['unsigned long long integer']: ... +@overload +def typename(char: L['f']) -> L['single precision']: ... +@overload +def typename(char: L['d']) -> L['double precision']: ... +@overload +def typename(char: L['g']) -> L['long precision']: ... +@overload +def typename(char: L['F']) -> L['complex single precision']: ... +@overload +def typename(char: L['D']) -> L['complex double precision']: ... +@overload +def typename(char: L['G']) -> L['complex long double precision']: ... +@overload +def typename(char: L['S']) -> L['string']: ... +@overload +def typename(char: L['U']) -> L['unicode']: ... +@overload +def typename(char: L['V']) -> L['void']: ... +@overload +def typename(char: L['O']) -> L['object']: ... + +@overload +def common_type( # type: ignore[misc] + *arrays: _SupportsDType[dtype[ + integer[Any] + ]] +) -> type[floating[_64Bit]]: ... +@overload +def common_type( # type: ignore[misc] + *arrays: _SupportsDType[dtype[ + floating[_NBit1] + ]] +) -> type[floating[_NBit1]]: ... +@overload +def common_type( # type: ignore[misc] + *arrays: _SupportsDType[dtype[ + integer[Any] | floating[_NBit1] + ]] +) -> type[floating[_NBit1 | _64Bit]]: ... +@overload +def common_type( # type: ignore[misc] + *arrays: _SupportsDType[dtype[ + floating[_NBit1] | complexfloating[_NBit2, _NBit2] + ]] +) -> type[complexfloating[_NBit1 | _NBit2, _NBit1 | _NBit2]]: ... +@overload +def common_type( + *arrays: _SupportsDType[dtype[ + integer[Any] | floating[_NBit1] | complexfloating[_NBit2, _NBit2] + ]] +) -> type[complexfloating[_64Bit | _NBit1 | _NBit2, _64Bit | _NBit1 | _NBit2]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..695aab1b8922385da6ecb8eb4c0bfcba9742dd9c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.py @@ -0,0 +1,207 @@ +""" +Module of functions that are like ufuncs in acting on arrays and optionally +storing results in an output array. + +""" +__all__ = ['fix', 'isneginf', 'isposinf'] + +import numpy._core.numeric as nx +from numpy._core.overrides import array_function_dispatch + + +def _dispatcher(x, out=None): + return (x, out) + + +@array_function_dispatch(_dispatcher, verify=False, module='numpy') +def fix(x, out=None): + """ + Round to nearest integer towards zero. + + Round an array of floats element-wise to nearest integer towards zero. + The rounded values have the same data-type as the input. + + Parameters + ---------- + x : array_like + An array to be rounded + out : ndarray, optional + A location into which the result is stored. If provided, it must have + a shape that the input broadcasts to. If not provided or None, a + freshly-allocated array is returned. + + Returns + ------- + out : ndarray of floats + An array with the same dimensions and data-type as the input. + If second argument is not supplied then a new array is returned + with the rounded values. + + If a second argument is supplied the result is stored there. + The return value ``out`` is then a reference to that array. + + See Also + -------- + rint, trunc, floor, ceil + around : Round to given number of decimals + + Examples + -------- + >>> import numpy as np + >>> np.fix(3.14) + 3.0 + >>> np.fix(3) + 3 + >>> np.fix([2.1, 2.9, -2.1, -2.9]) + array([ 2., 2., -2., -2.]) + + """ + # promote back to an array if flattened + res = nx.asanyarray(nx.ceil(x, out=out)) + res = nx.floor(x, out=res, where=nx.greater_equal(x, 0)) + + # when no out argument is passed and no subclasses are involved, flatten + # scalars + if out is None and type(res) is nx.ndarray: + res = res[()] + return res + + +@array_function_dispatch(_dispatcher, verify=False, module='numpy') +def isposinf(x, out=None): + """ + Test element-wise for positive infinity, return result as bool array. + + Parameters + ---------- + x : array_like + The input array. + out : array_like, optional + A location into which the result is stored. If provided, it must have a + shape that the input broadcasts to. If not provided or None, a + freshly-allocated boolean array is returned. + + Returns + ------- + out : ndarray + A boolean array with the same dimensions as the input. + If second argument is not supplied then a boolean array is returned + with values True where the corresponding element of the input is + positive infinity and values False where the element of the input is + not positive infinity. + + If a second argument is supplied the result is stored there. If the + type of that array is a numeric type the result is represented as zeros + and ones, if the type is boolean then as False and True. + The return value `out` is then a reference to that array. + + See Also + -------- + isinf, isneginf, isfinite, isnan + + Notes + ----- + NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic + (IEEE 754). + + Errors result if the second argument is also supplied when x is a scalar + input, if first and second arguments have different shapes, or if the + first argument has complex values + + Examples + -------- + >>> import numpy as np + >>> np.isposinf(np.inf) + True + >>> np.isposinf(-np.inf) + False + >>> np.isposinf([-np.inf, 0., np.inf]) + array([False, False, True]) + + >>> x = np.array([-np.inf, 0., np.inf]) + >>> y = np.array([2, 2, 2]) + >>> np.isposinf(x, y) + array([0, 0, 1]) + >>> y + array([0, 0, 1]) + + """ + is_inf = nx.isinf(x) + try: + signbit = ~nx.signbit(x) + except TypeError as e: + dtype = nx.asanyarray(x).dtype + raise TypeError(f'This operation is not supported for {dtype} values ' + 'because it would be ambiguous.') from e + else: + return nx.logical_and(is_inf, signbit, out) + + +@array_function_dispatch(_dispatcher, verify=False, module='numpy') +def isneginf(x, out=None): + """ + Test element-wise for negative infinity, return result as bool array. + + Parameters + ---------- + x : array_like + The input array. + out : array_like, optional + A location into which the result is stored. If provided, it must have a + shape that the input broadcasts to. If not provided or None, a + freshly-allocated boolean array is returned. + + Returns + ------- + out : ndarray + A boolean array with the same dimensions as the input. + If second argument is not supplied then a numpy boolean array is + returned with values True where the corresponding element of the + input is negative infinity and values False where the element of + the input is not negative infinity. + + If a second argument is supplied the result is stored there. If the + type of that array is a numeric type the result is represented as + zeros and ones, if the type is boolean then as False and True. The + return value `out` is then a reference to that array. + + See Also + -------- + isinf, isposinf, isnan, isfinite + + Notes + ----- + NumPy uses the IEEE Standard for Binary Floating-Point for Arithmetic + (IEEE 754). + + Errors result if the second argument is also supplied when x is a scalar + input, if first and second arguments have different shapes, or if the + first argument has complex values. + + Examples + -------- + >>> import numpy as np + >>> np.isneginf(-np.inf) + True + >>> np.isneginf(np.inf) + False + >>> np.isneginf([-np.inf, 0., np.inf]) + array([ True, False, False]) + + >>> x = np.array([-np.inf, 0., np.inf]) + >>> y = np.array([2, 2, 2]) + >>> np.isneginf(x, y) + array([1, 0, 0]) + >>> y + array([1, 0, 0]) + + """ + is_inf = nx.isinf(x) + try: + signbit = nx.signbit(x) + except TypeError as e: + dtype = nx.asanyarray(x).dtype + raise TypeError(f'This operation is not supported for {dtype} values ' + 'because it would be ambiguous.') from e + else: + return nx.logical_and(is_inf, signbit, out) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..8d87ae8bf4c62067e1b19e8b83ac6b951620168d --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_ufunclike_impl.pyi @@ -0,0 +1,67 @@ +from typing import Any, overload, TypeVar + +import numpy as np +from numpy import floating, object_ +from numpy._typing import ( + NDArray, + _FloatLike_co, + _ArrayLikeFloat_co, + _ArrayLikeObject_co, +) + +__all__ = ["fix", "isneginf", "isposinf"] + +_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) + +@overload +def fix( # type: ignore[misc] + x: _FloatLike_co, + out: None = ..., +) -> floating[Any]: ... +@overload +def fix( + x: _ArrayLikeFloat_co, + out: None = ..., +) -> NDArray[floating[Any]]: ... +@overload +def fix( + x: _ArrayLikeObject_co, + out: None = ..., +) -> NDArray[object_]: ... +@overload +def fix( + x: _ArrayLikeFloat_co | _ArrayLikeObject_co, + out: _ArrayType, +) -> _ArrayType: ... + +@overload +def isposinf( # type: ignore[misc] + x: _FloatLike_co, + out: None = ..., +) -> np.bool: ... +@overload +def isposinf( + x: _ArrayLikeFloat_co, + out: None = ..., +) -> NDArray[np.bool]: ... +@overload +def isposinf( + x: _ArrayLikeFloat_co, + out: _ArrayType, +) -> _ArrayType: ... + +@overload +def isneginf( # type: ignore[misc] + x: _FloatLike_co, + out: None = ..., +) -> np.bool: ... +@overload +def isneginf( + x: _ArrayLikeFloat_co, + out: None = ..., +) -> NDArray[np.bool]: ... +@overload +def isneginf( + x: _ArrayLikeFloat_co, + out: _ArrayType, +) -> _ArrayType: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..cae6e0556687a471f28558fb13dd013b966328e4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.py @@ -0,0 +1,291 @@ +""" +Container class for backward compatibility with NumArray. + +The user_array.container class exists for backward compatibility with NumArray +and is not meant to be used in new code. If you need to create an array +container class, we recommend either creating a class that wraps an ndarray +or subclasses ndarray. + +""" +from numpy._core import ( + array, asarray, absolute, add, subtract, multiply, divide, + remainder, power, left_shift, right_shift, bitwise_and, bitwise_or, + bitwise_xor, invert, less, less_equal, not_equal, equal, greater, + greater_equal, shape, reshape, arange, sin, sqrt, transpose +) +from numpy._core.overrides import set_module + + +@set_module("numpy.lib.user_array") +class container: + """ + container(data, dtype=None, copy=True) + + Standard container-class for easy multiple-inheritance. + + Methods + ------- + copy + tostring + byteswap + astype + + """ + def __init__(self, data, dtype=None, copy=True): + self.array = array(data, dtype, copy=copy) + + def __repr__(self): + if self.ndim > 0: + return self.__class__.__name__ + repr(self.array)[len("array"):] + else: + return self.__class__.__name__ + "(" + repr(self.array) + ")" + + def __array__(self, t=None): + if t: + return self.array.astype(t) + return self.array + + # Array as sequence + def __len__(self): + return len(self.array) + + def __getitem__(self, index): + return self._rc(self.array[index]) + + def __setitem__(self, index, value): + self.array[index] = asarray(value, self.dtype) + + def __abs__(self): + return self._rc(absolute(self.array)) + + def __neg__(self): + return self._rc(-self.array) + + def __add__(self, other): + return self._rc(self.array + asarray(other)) + + __radd__ = __add__ + + def __iadd__(self, other): + add(self.array, other, self.array) + return self + + def __sub__(self, other): + return self._rc(self.array - asarray(other)) + + def __rsub__(self, other): + return self._rc(asarray(other) - self.array) + + def __isub__(self, other): + subtract(self.array, other, self.array) + return self + + def __mul__(self, other): + return self._rc(multiply(self.array, asarray(other))) + + __rmul__ = __mul__ + + def __imul__(self, other): + multiply(self.array, other, self.array) + return self + + def __div__(self, other): + return self._rc(divide(self.array, asarray(other))) + + def __rdiv__(self, other): + return self._rc(divide(asarray(other), self.array)) + + def __idiv__(self, other): + divide(self.array, other, self.array) + return self + + def __mod__(self, other): + return self._rc(remainder(self.array, other)) + + def __rmod__(self, other): + return self._rc(remainder(other, self.array)) + + def __imod__(self, other): + remainder(self.array, other, self.array) + return self + + def __divmod__(self, other): + return (self._rc(divide(self.array, other)), + self._rc(remainder(self.array, other))) + + def __rdivmod__(self, other): + return (self._rc(divide(other, self.array)), + self._rc(remainder(other, self.array))) + + def __pow__(self, other): + return self._rc(power(self.array, asarray(other))) + + def __rpow__(self, other): + return self._rc(power(asarray(other), self.array)) + + def __ipow__(self, other): + power(self.array, other, self.array) + return self + + def __lshift__(self, other): + return self._rc(left_shift(self.array, other)) + + def __rshift__(self, other): + return self._rc(right_shift(self.array, other)) + + def __rlshift__(self, other): + return self._rc(left_shift(other, self.array)) + + def __rrshift__(self, other): + return self._rc(right_shift(other, self.array)) + + def __ilshift__(self, other): + left_shift(self.array, other, self.array) + return self + + def __irshift__(self, other): + right_shift(self.array, other, self.array) + return self + + def __and__(self, other): + return self._rc(bitwise_and(self.array, other)) + + def __rand__(self, other): + return self._rc(bitwise_and(other, self.array)) + + def __iand__(self, other): + bitwise_and(self.array, other, self.array) + return self + + def __xor__(self, other): + return self._rc(bitwise_xor(self.array, other)) + + def __rxor__(self, other): + return self._rc(bitwise_xor(other, self.array)) + + def __ixor__(self, other): + bitwise_xor(self.array, other, self.array) + return self + + def __or__(self, other): + return self._rc(bitwise_or(self.array, other)) + + def __ror__(self, other): + return self._rc(bitwise_or(other, self.array)) + + def __ior__(self, other): + bitwise_or(self.array, other, self.array) + return self + + def __pos__(self): + return self._rc(self.array) + + def __invert__(self): + return self._rc(invert(self.array)) + + def _scalarfunc(self, func): + if self.ndim == 0: + return func(self[0]) + else: + raise TypeError( + "only rank-0 arrays can be converted to Python scalars.") + + def __complex__(self): + return self._scalarfunc(complex) + + def __float__(self): + return self._scalarfunc(float) + + def __int__(self): + return self._scalarfunc(int) + + def __hex__(self): + return self._scalarfunc(hex) + + def __oct__(self): + return self._scalarfunc(oct) + + def __lt__(self, other): + return self._rc(less(self.array, other)) + + def __le__(self, other): + return self._rc(less_equal(self.array, other)) + + def __eq__(self, other): + return self._rc(equal(self.array, other)) + + def __ne__(self, other): + return self._rc(not_equal(self.array, other)) + + def __gt__(self, other): + return self._rc(greater(self.array, other)) + + def __ge__(self, other): + return self._rc(greater_equal(self.array, other)) + + def copy(self): + "" + return self._rc(self.array.copy()) + + def tostring(self): + "" + return self.array.tostring() + + def tobytes(self): + "" + return self.array.tobytes() + + def byteswap(self): + "" + return self._rc(self.array.byteswap()) + + def astype(self, typecode): + "" + return self._rc(self.array.astype(typecode)) + + def _rc(self, a): + if len(shape(a)) == 0: + return a + else: + return self.__class__(a) + + def __array_wrap__(self, *args): + return self.__class__(args[0]) + + def __setattr__(self, attr, value): + if attr == 'array': + object.__setattr__(self, attr, value) + return + try: + self.array.__setattr__(attr, value) + except AttributeError: + object.__setattr__(self, attr, value) + + # Only called after other approaches fail. + def __getattr__(self, attr): + if (attr == 'array'): + return object.__getattribute__(self, attr) + return self.array.__getattribute__(attr) + + +############################################################# +# Test of class container +############################################################# +if __name__ == '__main__': + temp = reshape(arange(10000), (100, 100)) + + ua = container(temp) + # new object created begin test + print(dir(ua)) + print(shape(ua), ua.shape) # I have changed Numeric.py + + ua_small = ua[:3, :5] + print(ua_small) + # this did not change ua[0,0], which is not normal behavior + ua_small[0, 0] = 10 + print(ua_small[0, 0], ua[0, 0]) + print(sin(ua_small) / 3. * 6. + sqrt(ua_small ** 2)) + print(less(ua_small, 103), type(less(ua_small, 103))) + print(type(ua_small * reshape(arange(15), shape(ua_small)))) + print(reshape(ua_small, (5, 3))) + print(transpose(ua_small)) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..d5dfb0573c7150a9287e3ff5497546b0ad6a59f1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_user_array_impl.pyi @@ -0,0 +1,220 @@ +from types import EllipsisType +from typing import Any, Generic, SupportsIndex, TypeAlias, TypeVar, overload + +from _typeshed import Incomplete +from typing_extensions import Self, deprecated, override + +import numpy as np +import numpy.typing as npt +from numpy._typing import _ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co, _DTypeLike + +### + +_ScalarT = TypeVar("_ScalarT", bound=np.generic) +_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) +_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=Any, covariant=True) +_DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any]) +_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype[Any], default=np.dtype[Any], covariant=True) + +_BoolArrayT = TypeVar("_BoolArrayT", bound=container[Any, np.dtype[np.bool]]) +_IntegralArrayT = TypeVar("_IntegralArrayT", bound=container[Any, np.dtype[np.bool | np.integer | np.object_]]) +_RealContainerT = TypeVar( + "_RealContainerT", + bound=container[Any, np.dtype[np.bool | np.integer | np.floating | np.timedelta64 | np.object_]], +) +_NumericContainerT = TypeVar("_NumericContainerT", bound=container[Any, np.dtype[np.number | np.timedelta64 | np.object_]]) + +_ArrayInt_co: TypeAlias = npt.NDArray[np.integer | np.bool] + +_ToIndexSlice: TypeAlias = slice | EllipsisType | _ArrayInt_co | None +_ToIndexSlices: TypeAlias = _ToIndexSlice | tuple[_ToIndexSlice, ...] +_ToIndex: TypeAlias = SupportsIndex | _ToIndexSlice +_ToIndices: TypeAlias = _ToIndex | tuple[_ToIndex, ...] + +### + +class container(Generic[_ShapeT_co, _DTypeT_co]): + array: np.ndarray[_ShapeT_co, _DTypeT_co] + + @overload + def __init__( + self, + /, + data: container[_ShapeT_co, _DTypeT_co] | np.ndarray[_ShapeT_co, _DTypeT_co], + dtype: None = None, + copy: bool = True, + ) -> None: ... + @overload + def __init__( + self: container[Any, np.dtype[_ScalarT]], + /, + data: _ArrayLike[_ScalarT], + dtype: None = None, + copy: bool = True, + ) -> None: ... + @overload + def __init__( + self: container[Any, np.dtype[_ScalarT]], + /, + data: npt.ArrayLike, + dtype: _DTypeLike[_ScalarT], + copy: bool = True, + ) -> None: ... + @overload + def __init__(self, /, data: npt.ArrayLike, dtype: npt.DTypeLike | None = None, copy: bool = True) -> None: ... + + # + def __complex__(self, /) -> complex: ... + def __float__(self, /) -> float: ... + def __int__(self, /) -> int: ... + def __hex__(self, /) -> str: ... + def __oct__(self, /) -> str: ... + + # + @override + def __eq__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + @override + def __ne__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + + # + def __lt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... + def __le__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... + def __gt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... + def __ge__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... + + # + def __len__(self, /) -> int: ... + + # keep in sync with np.ndarray + @overload + def __getitem__(self, key: _ArrayInt_co | tuple[_ArrayInt_co, ...], /) -> container[_ShapeT_co, _DTypeT_co]: ... + @overload + def __getitem__(self, key: _ToIndexSlices, /) -> container[Any, _DTypeT_co]: ... + @overload + def __getitem__(self, key: _ToIndices, /) -> Any: ... + @overload + def __getitem__(self: container[Any, np.dtype[np.void]], key: list[str], /) -> container[_ShapeT_co, np.dtype[np.void]]: ... + @overload + def __getitem__(self: container[Any, np.dtype[np.void]], key: str, /) -> container[_ShapeT_co, np.dtype[Any]]: ... + + # keep in sync with np.ndarray + @overload + def __setitem__(self, index: _ToIndices, value: object, /) -> None: ... + @overload + def __setitem__(self: container[Any, np.dtype[np.void]], key: str | list[str], value: object, /) -> None: ... + + # keep in sync with np.ndarray + @overload + def __abs__(self: container[_ShapeT, np.dtype[np.complex64]], /) -> container[_ShapeT, np.dtype[np.float32]]: ... # type: ignore[overload-overlap] + @overload + def __abs__(self: container[_ShapeT, np.dtype[np.complex128]], /) -> container[_ShapeT, np.dtype[np.float64]]: ... + @overload + def __abs__(self: container[_ShapeT, np.dtype[np.complex192]], /) -> container[_ShapeT, np.dtype[np.float96]]: ... + @overload + def __abs__(self: container[_ShapeT, np.dtype[np.complex256]], /) -> container[_ShapeT, np.dtype[np.float128]]: ... + @overload + def __abs__(self: _RealContainerT, /) -> _RealContainerT: ... + + # + def __neg__(self: _NumericContainerT, /) -> _NumericContainerT: ... # noqa: PYI019 + def __pos__(self: _NumericContainerT, /) -> _NumericContainerT: ... # noqa: PYI019 + def __invert__(self: _IntegralArrayT, /) -> _IntegralArrayT: ... # noqa: PYI019 + + # TODO(jorenham): complete these binary ops + + # + def __add__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __radd__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __iadd__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __sub__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __rsub__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __isub__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __mul__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __rmul__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __imul__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __div__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __rdiv__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __idiv__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __mod__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __rmod__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __imod__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __divmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ... + def __rdivmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ... + + # + def __pow__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __rpow__(self, other: npt.ArrayLike, /) -> Incomplete: ... + def __ipow__(self, other: npt.ArrayLike, /) -> Self: ... + + # + def __lshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ... + def __rlshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ... + def __ilshift__(self, other: _ArrayLikeInt_co, /) -> Self: ... + + # + def __rshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ... + def __rrshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ... + def __irshift__(self, other: _ArrayLikeInt_co, /) -> Self: ... + + # + @overload + def __and__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ... + @overload + def __and__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ... + __rand__ = __and__ + @overload + def __iand__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ... + @overload + def __iand__(self, other: _ArrayLikeInt_co, /) -> Self: ... + + # + @overload + def __xor__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ... + @overload + def __xor__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ... + __rxor__ = __xor__ + @overload + def __ixor__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ... + @overload + def __ixor__(self, other: _ArrayLikeInt_co, /) -> Self: ... + + # + @overload + def __or__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ... + @overload + def __or__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ... + __ror__ = __or__ + @overload + def __ior__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ... + @overload + def __ior__(self, other: _ArrayLikeInt_co, /) -> Self: ... + + # + @overload + def __array__(self, /, t: None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ... + @overload + def __array__(self, /, t: _DTypeT) -> np.ndarray[_ShapeT_co, _DTypeT]: ... + + # + @overload + def __array_wrap__(self, arg0: npt.ArrayLike, /) -> container[_ShapeT_co, _DTypeT_co]: ... + @overload + def __array_wrap__(self, a: np.ndarray[_ShapeT, _DTypeT], c: Any = ..., s: Any = ..., /) -> container[_ShapeT, _DTypeT]: ... + + # + def copy(self, /) -> Self: ... + @deprecated("tostring() is deprecated. Use tobytes() instead.") + def tostring(self, /) -> bytes: ... + def tobytes(self, /) -> bytes: ... + def byteswap(self, /) -> Self: ... + def astype(self, /, typecode: _DTypeLike[_ScalarT]) -> container[_ShapeT_co, np.dtype[_ScalarT]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.py b/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..c2f0f31d7bfcc7e96bf7936887788817319ddc5e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.py @@ -0,0 +1,775 @@ +import os +import sys +import textwrap +import types +import warnings +import functools +import platform + +from numpy._core import ndarray +from numpy._utils import set_module +import numpy as np + +__all__ = [ + 'get_include', 'info', 'show_runtime' +] + + +@set_module('numpy') +def show_runtime(): + """ + Print information about various resources in the system + including available intrinsic support and BLAS/LAPACK library + in use + + .. versionadded:: 1.24.0 + + See Also + -------- + show_config : Show libraries in the system on which NumPy was built. + + Notes + ----- + 1. Information is derived with the help of `threadpoolctl `_ + library if available. + 2. SIMD related information is derived from ``__cpu_features__``, + ``__cpu_baseline__`` and ``__cpu_dispatch__`` + + """ + from numpy._core._multiarray_umath import ( + __cpu_features__, __cpu_baseline__, __cpu_dispatch__ + ) + from pprint import pprint + config_found = [{ + "numpy_version": np.__version__, + "python": sys.version, + "uname": platform.uname(), + }] + features_found, features_not_found = [], [] + for feature in __cpu_dispatch__: + if __cpu_features__[feature]: + features_found.append(feature) + else: + features_not_found.append(feature) + config_found.append({ + "simd_extensions": { + "baseline": __cpu_baseline__, + "found": features_found, + "not_found": features_not_found + } + }) + try: + from threadpoolctl import threadpool_info + config_found.extend(threadpool_info()) + except ImportError: + print("WARNING: `threadpoolctl` not found in system!" + " Install it by `pip install threadpoolctl`." + " Once installed, try `np.show_runtime` again" + " for more detailed build information") + pprint(config_found) + + +@set_module('numpy') +def get_include(): + """ + Return the directory that contains the NumPy \\*.h header files. + + Extension modules that need to compile against NumPy may need to use this + function to locate the appropriate include directory. + + Notes + ----- + When using ``setuptools``, for example in ``setup.py``:: + + import numpy as np + ... + Extension('extension_name', ... + include_dirs=[np.get_include()]) + ... + + Note that a CLI tool ``numpy-config`` was introduced in NumPy 2.0, using + that is likely preferred for build systems other than ``setuptools``:: + + $ numpy-config --cflags + -I/path/to/site-packages/numpy/_core/include + + # Or rely on pkg-config: + $ export PKG_CONFIG_PATH=$(numpy-config --pkgconfigdir) + $ pkg-config --cflags + -I/path/to/site-packages/numpy/_core/include + + Examples + -------- + >>> np.get_include() + '.../site-packages/numpy/core/include' # may vary + + """ + import numpy + if numpy.show_config is None: + # running from numpy source directory + d = os.path.join(os.path.dirname(numpy.__file__), '_core', 'include') + else: + # using installed numpy core headers + import numpy._core as _core + d = os.path.join(os.path.dirname(_core.__file__), 'include') + return d + + +class _Deprecate: + """ + Decorator class to deprecate old functions. + + Refer to `deprecate` for details. + + See Also + -------- + deprecate + + """ + + def __init__(self, old_name=None, new_name=None, message=None): + self.old_name = old_name + self.new_name = new_name + self.message = message + + def __call__(self, func, *args, **kwargs): + """ + Decorator call. Refer to ``decorate``. + + """ + old_name = self.old_name + new_name = self.new_name + message = self.message + + if old_name is None: + old_name = func.__name__ + if new_name is None: + depdoc = "`%s` is deprecated!" % old_name + else: + depdoc = "`%s` is deprecated, use `%s` instead!" % \ + (old_name, new_name) + + if message is not None: + depdoc += "\n" + message + + @functools.wraps(func) + def newfunc(*args, **kwds): + warnings.warn(depdoc, DeprecationWarning, stacklevel=2) + return func(*args, **kwds) + + newfunc.__name__ = old_name + doc = func.__doc__ + if doc is None: + doc = depdoc + else: + lines = doc.expandtabs().split('\n') + indent = _get_indent(lines[1:]) + if lines[0].lstrip(): + # Indent the original first line to let inspect.cleandoc() + # dedent the docstring despite the deprecation notice. + doc = indent * ' ' + doc + else: + # Remove the same leading blank lines as cleandoc() would. + skip = len(lines[0]) + 1 + for line in lines[1:]: + if len(line) > indent: + break + skip += len(line) + 1 + doc = doc[skip:] + depdoc = textwrap.indent(depdoc, ' ' * indent) + doc = f'{depdoc}\n\n{doc}' + newfunc.__doc__ = doc + + return newfunc + + +def _get_indent(lines): + """ + Determines the leading whitespace that could be removed from all the lines. + """ + indent = sys.maxsize + for line in lines: + content = len(line.lstrip()) + if content: + indent = min(indent, len(line) - content) + if indent == sys.maxsize: + indent = 0 + return indent + + +def deprecate(*args, **kwargs): + """ + Issues a DeprecationWarning, adds warning to `old_name`'s + docstring, rebinds ``old_name.__name__`` and returns the new + function object. + + This function may also be used as a decorator. + + .. deprecated:: 2.0 + Use `~warnings.warn` with :exc:`DeprecationWarning` instead. + + Parameters + ---------- + func : function + The function to be deprecated. + old_name : str, optional + The name of the function to be deprecated. Default is None, in + which case the name of `func` is used. + new_name : str, optional + The new name for the function. Default is None, in which case the + deprecation message is that `old_name` is deprecated. If given, the + deprecation message is that `old_name` is deprecated and `new_name` + should be used instead. + message : str, optional + Additional explanation of the deprecation. Displayed in the + docstring after the warning. + + Returns + ------- + old_func : function + The deprecated function. + + Examples + -------- + Note that ``olduint`` returns a value after printing Deprecation + Warning: + + >>> olduint = np.lib.utils.deprecate(np.uint) + DeprecationWarning: `uint64` is deprecated! # may vary + >>> olduint(6) + 6 + + """ + # Deprecate may be run as a function or as a decorator + # If run as a function, we initialise the decorator class + # and execute its __call__ method. + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`deprecate` is deprecated, " + "use `warn` with `DeprecationWarning` instead. " + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + if args: + fn = args[0] + args = args[1:] + + return _Deprecate(*args, **kwargs)(fn) + else: + return _Deprecate(*args, **kwargs) + + +def deprecate_with_doc(msg): + """ + Deprecates a function and includes the deprecation in its docstring. + + .. deprecated:: 2.0 + Use `~warnings.warn` with :exc:`DeprecationWarning` instead. + + This function is used as a decorator. It returns an object that can be + used to issue a DeprecationWarning, by passing the to-be decorated + function as argument, this adds warning to the to-be decorated function's + docstring and returns the new function object. + + See Also + -------- + deprecate : Decorate a function such that it issues a + :exc:`DeprecationWarning` + + Parameters + ---------- + msg : str + Additional explanation of the deprecation. Displayed in the + docstring after the warning. + + Returns + ------- + obj : object + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`deprecate` is deprecated, " + "use `warn` with `DeprecationWarning` instead. " + "(deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + return _Deprecate(message=msg) + + +#----------------------------------------------------------------------------- + + +# NOTE: pydoc defines a help function which works similarly to this +# except it uses a pager to take over the screen. + +# combine name and arguments and split to multiple lines of width +# characters. End lines on a comma and begin argument list indented with +# the rest of the arguments. +def _split_line(name, arguments, width): + firstwidth = len(name) + k = firstwidth + newstr = name + sepstr = ", " + arglist = arguments.split(sepstr) + for argument in arglist: + if k == firstwidth: + addstr = "" + else: + addstr = sepstr + k = k + len(argument) + len(addstr) + if k > width: + k = firstwidth + 1 + len(argument) + newstr = newstr + ",\n" + " "*(firstwidth+2) + argument + else: + newstr = newstr + addstr + argument + return newstr + +_namedict = None +_dictlist = None + +# Traverse all module directories underneath globals +# to see if something is defined +def _makenamedict(module='numpy'): + module = __import__(module, globals(), locals(), []) + thedict = {module.__name__:module.__dict__} + dictlist = [module.__name__] + totraverse = [module.__dict__] + while True: + if len(totraverse) == 0: + break + thisdict = totraverse.pop(0) + for x in thisdict.keys(): + if isinstance(thisdict[x], types.ModuleType): + modname = thisdict[x].__name__ + if modname not in dictlist: + moddict = thisdict[x].__dict__ + dictlist.append(modname) + totraverse.append(moddict) + thedict[modname] = moddict + return thedict, dictlist + + +def _info(obj, output=None): + """Provide information about ndarray obj. + + Parameters + ---------- + obj : ndarray + Must be ndarray, not checked. + output + Where printed output goes. + + Notes + ----- + Copied over from the numarray module prior to its removal. + Adapted somewhat as only numpy is an option now. + + Called by info. + + """ + extra = "" + tic = "" + bp = lambda x: x + cls = getattr(obj, '__class__', type(obj)) + nm = getattr(cls, '__name__', cls) + strides = obj.strides + endian = obj.dtype.byteorder + + if output is None: + output = sys.stdout + + print("class: ", nm, file=output) + print("shape: ", obj.shape, file=output) + print("strides: ", strides, file=output) + print("itemsize: ", obj.itemsize, file=output) + print("aligned: ", bp(obj.flags.aligned), file=output) + print("contiguous: ", bp(obj.flags.contiguous), file=output) + print("fortran: ", obj.flags.fortran, file=output) + print( + "data pointer: %s%s" % (hex(obj.ctypes._as_parameter_.value), extra), + file=output + ) + print("byteorder: ", end=' ', file=output) + if endian in ['|', '=']: + print("%s%s%s" % (tic, sys.byteorder, tic), file=output) + byteswap = False + elif endian == '>': + print("%sbig%s" % (tic, tic), file=output) + byteswap = sys.byteorder != "big" + else: + print("%slittle%s" % (tic, tic), file=output) + byteswap = sys.byteorder != "little" + print("byteswap: ", bp(byteswap), file=output) + print("type: %s" % obj.dtype, file=output) + + +@set_module('numpy') +def info(object=None, maxwidth=76, output=None, toplevel='numpy'): + """ + Get help information for an array, function, class, or module. + + Parameters + ---------- + object : object or str, optional + Input object or name to get information about. If `object` is + an `ndarray` instance, information about the array is printed. + If `object` is a numpy object, its docstring is given. If it is + a string, available modules are searched for matching objects. + If None, information about `info` itself is returned. + maxwidth : int, optional + Printing width. + output : file like object, optional + File like object that the output is written to, default is + ``None``, in which case ``sys.stdout`` will be used. + The object has to be opened in 'w' or 'a' mode. + toplevel : str, optional + Start search at this level. + + Notes + ----- + When used interactively with an object, ``np.info(obj)`` is equivalent + to ``help(obj)`` on the Python prompt or ``obj?`` on the IPython + prompt. + + Examples + -------- + >>> np.info(np.polyval) # doctest: +SKIP + polyval(p, x) + Evaluate the polynomial p at x. + ... + + When using a string for `object` it is possible to get multiple results. + + >>> np.info('fft') # doctest: +SKIP + *** Found in numpy *** + Core FFT routines + ... + *** Found in numpy.fft *** + fft(a, n=None, axis=-1) + ... + *** Repeat reference found in numpy.fft.fftpack *** + *** Total of 3 references found. *** + + When the argument is an array, information about the array is printed. + + >>> a = np.array([[1 + 2j, 3, -4], [-5j, 6, 0]], dtype=np.complex64) + >>> np.info(a) + class: ndarray + shape: (2, 3) + strides: (24, 8) + itemsize: 8 + aligned: True + contiguous: True + fortran: False + data pointer: 0x562b6e0d2860 # may vary + byteorder: little + byteswap: False + type: complex64 + + """ + global _namedict, _dictlist + # Local import to speed up numpy's import time. + import pydoc + import inspect + + if (hasattr(object, '_ppimport_importer') or + hasattr(object, '_ppimport_module')): + object = object._ppimport_module + elif hasattr(object, '_ppimport_attr'): + object = object._ppimport_attr + + if output is None: + output = sys.stdout + + if object is None: + info(info) + elif isinstance(object, ndarray): + _info(object, output=output) + elif isinstance(object, str): + if _namedict is None: + _namedict, _dictlist = _makenamedict(toplevel) + numfound = 0 + objlist = [] + for namestr in _dictlist: + try: + obj = _namedict[namestr][object] + if id(obj) in objlist: + print("\n " + "*** Repeat reference found in %s *** " % namestr, + file=output + ) + else: + objlist.append(id(obj)) + print(" *** Found in %s ***" % namestr, file=output) + info(obj) + print("-"*maxwidth, file=output) + numfound += 1 + except KeyError: + pass + if numfound == 0: + print("Help for %s not found." % object, file=output) + else: + print("\n " + "*** Total of %d references found. ***" % numfound, + file=output + ) + + elif inspect.isfunction(object) or inspect.ismethod(object): + name = object.__name__ + try: + arguments = str(inspect.signature(object)) + except Exception: + arguments = "()" + + if len(name+arguments) > maxwidth: + argstr = _split_line(name, arguments, maxwidth) + else: + argstr = name + arguments + + print(" " + argstr + "\n", file=output) + print(inspect.getdoc(object), file=output) + + elif inspect.isclass(object): + name = object.__name__ + try: + arguments = str(inspect.signature(object)) + except Exception: + arguments = "()" + + if len(name+arguments) > maxwidth: + argstr = _split_line(name, arguments, maxwidth) + else: + argstr = name + arguments + + print(" " + argstr + "\n", file=output) + doc1 = inspect.getdoc(object) + if doc1 is None: + if hasattr(object, '__init__'): + print(inspect.getdoc(object.__init__), file=output) + else: + print(inspect.getdoc(object), file=output) + + methods = pydoc.allmethods(object) + + public_methods = [meth for meth in methods if meth[0] != '_'] + if public_methods: + print("\n\nMethods:\n", file=output) + for meth in public_methods: + thisobj = getattr(object, meth, None) + if thisobj is not None: + methstr, other = pydoc.splitdoc( + inspect.getdoc(thisobj) or "None" + ) + print(" %s -- %s" % (meth, methstr), file=output) + + elif hasattr(object, '__doc__'): + print(inspect.getdoc(object), file=output) + + +def safe_eval(source): + """ + Protected string evaluation. + + .. deprecated:: 2.0 + Use `ast.literal_eval` instead. + + Evaluate a string containing a Python literal expression without + allowing the execution of arbitrary non-literal code. + + .. warning:: + + This function is identical to :py:meth:`ast.literal_eval` and + has the same security implications. It may not always be safe + to evaluate large input strings. + + Parameters + ---------- + source : str + The string to evaluate. + + Returns + ------- + obj : object + The result of evaluating `source`. + + Raises + ------ + SyntaxError + If the code has invalid Python syntax, or if it contains + non-literal code. + + Examples + -------- + >>> np.safe_eval('1') + 1 + >>> np.safe_eval('[1, 2, 3]') + [1, 2, 3] + >>> np.safe_eval('{"foo": ("bar", 10.0)}') + {'foo': ('bar', 10.0)} + + >>> np.safe_eval('import os') + Traceback (most recent call last): + ... + SyntaxError: invalid syntax + + >>> np.safe_eval('open("/home/user/.ssh/id_dsa").read()') + Traceback (most recent call last): + ... + ValueError: malformed node or string: <_ast.Call object at 0x...> + + """ + + # Deprecated in NumPy 2.0, 2023-07-11 + warnings.warn( + "`safe_eval` is deprecated. Use `ast.literal_eval` instead. " + "Be aware of security implications, such as memory exhaustion " + "based attacks (deprecated in NumPy 2.0)", + DeprecationWarning, + stacklevel=2 + ) + + # Local import to speed up numpy's import time. + import ast + return ast.literal_eval(source) + + +def _median_nancheck(data, result, axis): + """ + Utility function to check median result from data for NaN values at the end + and return NaN in that case. Input result can also be a MaskedArray. + + Parameters + ---------- + data : array + Sorted input data to median function + result : Array or MaskedArray + Result of median function. + axis : int + Axis along which the median was computed. + + Returns + ------- + result : scalar or ndarray + Median or NaN in axes which contained NaN in the input. If the input + was an array, NaN will be inserted in-place. If a scalar, either the + input itself or a scalar NaN. + """ + if data.size == 0: + return result + potential_nans = data.take(-1, axis=axis) + n = np.isnan(potential_nans) + # masked NaN values are ok, although for masked the copyto may fail for + # unmasked ones (this was always broken) when the result is a scalar. + if np.ma.isMaskedArray(n): + n = n.filled(False) + + if not n.any(): + return result + + # Without given output, it is possible that the current result is a + # numpy scalar, which is not writeable. If so, just return nan. + if isinstance(result, np.generic): + return potential_nans + + # Otherwise copy NaNs (if there are any) + np.copyto(result, potential_nans, where=n) + return result + +def _opt_info(): + """ + Returns a string containing the CPU features supported + by the current build. + + The format of the string can be explained as follows: + - Dispatched features supported by the running machine end with `*`. + - Dispatched features not supported by the running machine + end with `?`. + - Remaining features represent the baseline. + + Returns: + str: A formatted string indicating the supported CPU features. + """ + from numpy._core._multiarray_umath import ( + __cpu_features__, __cpu_baseline__, __cpu_dispatch__ + ) + + if len(__cpu_baseline__) == 0 and len(__cpu_dispatch__) == 0: + return '' + + enabled_features = ' '.join(__cpu_baseline__) + for feature in __cpu_dispatch__: + if __cpu_features__[feature]: + enabled_features += f" {feature}*" + else: + enabled_features += f" {feature}?" + + return enabled_features + +def drop_metadata(dtype, /): + """ + Returns the dtype unchanged if it contained no metadata or a copy of the + dtype if it (or any of its structure dtypes) contained metadata. + + This utility is used by `np.save` and `np.savez` to drop metadata before + saving. + + .. note:: + + Due to its limitation this function may move to a more appropriate + home or change in the future and is considered semi-public API only. + + .. warning:: + + This function does not preserve more strange things like record dtypes + and user dtypes may simply return the wrong thing. If you need to be + sure about the latter, check the result with: + ``np.can_cast(new_dtype, dtype, casting="no")``. + + """ + if dtype.fields is not None: + found_metadata = dtype.metadata is not None + + names = [] + formats = [] + offsets = [] + titles = [] + for name, field in dtype.fields.items(): + field_dt = drop_metadata(field[0]) + if field_dt is not field[0]: + found_metadata = True + + names.append(name) + formats.append(field_dt) + offsets.append(field[1]) + titles.append(None if len(field) < 3 else field[2]) + + if not found_metadata: + return dtype + + structure = dict( + names=names, formats=formats, offsets=offsets, titles=titles, + itemsize=dtype.itemsize) + + # NOTE: Could pass (dtype.type, structure) to preserve record dtypes... + return np.dtype(structure, align=dtype.isalignedstruct) + elif dtype.subdtype is not None: + # subarray dtype + subdtype, shape = dtype.subdtype + new_subdtype = drop_metadata(subdtype) + if dtype.metadata is None and new_subdtype is subdtype: + return dtype + + return np.dtype((new_subdtype, shape)) + else: + # Normal unstructured dtype + if dtype.metadata is None: + return dtype + # Note that `dt.str` doesn't round-trip e.g. for user-dtypes. + return np.dtype(dtype.str) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.pyi new file mode 100644 index 0000000000000000000000000000000000000000..2a9eb76a5b381ecb476af0e6fdf6d839c4a10f66 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_utils_impl.pyi @@ -0,0 +1,7 @@ +from _typeshed import SupportsWrite + +__all__ = ["get_include", "info", "show_runtime"] + +def get_include() -> str: ... +def show_runtime() -> None: ... +def info(object: object = ..., maxwidth: int = ..., output: SupportsWrite[str] | None = ..., toplevel: str = ...) -> None: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_version.py b/venv/lib/python3.10/site-packages/numpy/lib/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..929f8a1c6685d587663086a78f6379dd6ba684ee --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_version.py @@ -0,0 +1,155 @@ +"""Utility to compare (NumPy) version strings. + +The NumpyVersion class allows properly comparing numpy version strings. +The LooseVersion and StrictVersion classes that distutils provides don't +work; they don't recognize anything like alpha/beta/rc/dev versions. + +""" +import re + + +__all__ = ['NumpyVersion'] + + +class NumpyVersion: + """Parse and compare numpy version strings. + + NumPy has the following versioning scheme (numbers given are examples; they + can be > 9 in principle): + + - Released version: '1.8.0', '1.8.1', etc. + - Alpha: '1.8.0a1', '1.8.0a2', etc. + - Beta: '1.8.0b1', '1.8.0b2', etc. + - Release candidates: '1.8.0rc1', '1.8.0rc2', etc. + - Development versions: '1.8.0.dev-f1234afa' (git commit hash appended) + - Development versions after a1: '1.8.0a1.dev-f1234afa', + '1.8.0b2.dev-f1234afa', + '1.8.1rc1.dev-f1234afa', etc. + - Development versions (no git hash available): '1.8.0.dev-Unknown' + + Comparing needs to be done against a valid version string or other + `NumpyVersion` instance. Note that all development versions of the same + (pre-)release compare equal. + + Parameters + ---------- + vstring : str + NumPy version string (``np.__version__``). + + Examples + -------- + >>> from numpy.lib import NumpyVersion + >>> if NumpyVersion(np.__version__) < '1.7.0': + ... print('skip') + >>> # skip + + >>> NumpyVersion('1.7') # raises ValueError, add ".0" + Traceback (most recent call last): + ... + ValueError: Not a valid numpy version string + + """ + + __module__ = "numpy.lib" + + def __init__(self, vstring): + self.vstring = vstring + ver_main = re.match(r'\d+\.\d+\.\d+', vstring) + if not ver_main: + raise ValueError("Not a valid numpy version string") + + self.version = ver_main.group() + self.major, self.minor, self.bugfix = [int(x) for x in + self.version.split('.')] + if len(vstring) == ver_main.end(): + self.pre_release = 'final' + else: + alpha = re.match(r'a\d', vstring[ver_main.end():]) + beta = re.match(r'b\d', vstring[ver_main.end():]) + rc = re.match(r'rc\d', vstring[ver_main.end():]) + pre_rel = [m for m in [alpha, beta, rc] if m is not None] + if pre_rel: + self.pre_release = pre_rel[0].group() + else: + self.pre_release = '' + + self.is_devversion = bool(re.search(r'.dev', vstring)) + + def _compare_version(self, other): + """Compare major.minor.bugfix""" + if self.major == other.major: + if self.minor == other.minor: + if self.bugfix == other.bugfix: + vercmp = 0 + elif self.bugfix > other.bugfix: + vercmp = 1 + else: + vercmp = -1 + elif self.minor > other.minor: + vercmp = 1 + else: + vercmp = -1 + elif self.major > other.major: + vercmp = 1 + else: + vercmp = -1 + + return vercmp + + def _compare_pre_release(self, other): + """Compare alpha/beta/rc/final.""" + if self.pre_release == other.pre_release: + vercmp = 0 + elif self.pre_release == 'final': + vercmp = 1 + elif other.pre_release == 'final': + vercmp = -1 + elif self.pre_release > other.pre_release: + vercmp = 1 + else: + vercmp = -1 + + return vercmp + + def _compare(self, other): + if not isinstance(other, (str, NumpyVersion)): + raise ValueError("Invalid object to compare with NumpyVersion.") + + if isinstance(other, str): + other = NumpyVersion(other) + + vercmp = self._compare_version(other) + if vercmp == 0: + # Same x.y.z version, check for alpha/beta/rc + vercmp = self._compare_pre_release(other) + if vercmp == 0: + # Same version and same pre-release, check if dev version + if self.is_devversion is other.is_devversion: + vercmp = 0 + elif self.is_devversion: + vercmp = -1 + else: + vercmp = 1 + + return vercmp + + def __lt__(self, other): + return self._compare(other) < 0 + + def __le__(self, other): + return self._compare(other) <= 0 + + def __eq__(self, other): + return self._compare(other) == 0 + + def __ne__(self, other): + return self._compare(other) != 0 + + def __gt__(self, other): + return self._compare(other) > 0 + + def __ge__(self, other): + return self._compare(other) >= 0 + + def __repr__(self): + return "NumpyVersion(%s)" % self.vstring diff --git a/venv/lib/python3.10/site-packages/numpy/lib/_version.pyi b/venv/lib/python3.10/site-packages/numpy/lib/_version.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c53ef795f9266f5233d8c86e696bd8e1f3699557 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/_version.pyi @@ -0,0 +1,17 @@ +__all__ = ["NumpyVersion"] + +class NumpyVersion: + vstring: str + version: str + major: int + minor: int + bugfix: int + pre_release: str + is_devversion: bool + def __init__(self, vstring: str) -> None: ... + def __lt__(self, other: str | NumpyVersion) -> bool: ... + def __le__(self, other: str | NumpyVersion) -> bool: ... + def __eq__(self, other: str | NumpyVersion) -> bool: ... # type: ignore[override] + def __ne__(self, other: str | NumpyVersion) -> bool: ... # type: ignore[override] + def __gt__(self, other: str | NumpyVersion) -> bool: ... + def __ge__(self, other: str | NumpyVersion) -> bool: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/array_utils.py b/venv/lib/python3.10/site-packages/numpy/lib/array_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e7976131d268b307d8ac3bb13034d9e425f6e6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/array_utils.py @@ -0,0 +1,7 @@ +from ._array_utils_impl import ( + __all__, + __doc__, + byte_bounds, + normalize_axis_index, + normalize_axis_tuple, +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/array_utils.pyi b/venv/lib/python3.10/site-packages/numpy/lib/array_utils.pyi new file mode 100644 index 0000000000000000000000000000000000000000..4b9ebe334a1f211a21aac0cfd1605166ca909d1b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/array_utils.pyi @@ -0,0 +1,6 @@ +from ._array_utils_impl import ( + __all__ as __all__, + byte_bounds as byte_bounds, + normalize_axis_index as normalize_axis_index, + normalize_axis_tuple as normalize_axis_tuple, +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/format.py b/venv/lib/python3.10/site-packages/numpy/lib/format.py new file mode 100644 index 0000000000000000000000000000000000000000..a22c096b246ce7109ca7b39f7735e4e38ab1e036 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/format.py @@ -0,0 +1,1008 @@ +""" +Binary serialization + +NPY format +========== + +A simple format for saving numpy arrays to disk with the full +information about them. + +The ``.npy`` format is the standard binary file format in NumPy for +persisting a *single* arbitrary NumPy array on disk. The format stores all +of the shape and dtype information necessary to reconstruct the array +correctly even on another machine with a different architecture. +The format is designed to be as simple as possible while achieving +its limited goals. + +The ``.npz`` format is the standard format for persisting *multiple* NumPy +arrays on disk. A ``.npz`` file is a zip file containing multiple ``.npy`` +files, one for each array. + +Capabilities +------------ + +- Can represent all NumPy arrays including nested record arrays and + object arrays. + +- Represents the data in its native binary form. + +- Supports Fortran-contiguous arrays directly. + +- Stores all of the necessary information to reconstruct the array + including shape and dtype on a machine of a different + architecture. Both little-endian and big-endian arrays are + supported, and a file with little-endian numbers will yield + a little-endian array on any machine reading the file. The + types are described in terms of their actual sizes. For example, + if a machine with a 64-bit C "long int" writes out an array with + "long ints", a reading machine with 32-bit C "long ints" will yield + an array with 64-bit integers. + +- Is straightforward to reverse engineer. Datasets often live longer than + the programs that created them. A competent developer should be + able to create a solution in their preferred programming language to + read most ``.npy`` files that they have been given without much + documentation. + +- Allows memory-mapping of the data. See `open_memmap`. + +- Can be read from a filelike stream object instead of an actual file. + +- Stores object arrays, i.e. arrays containing elements that are arbitrary + Python objects. Files with object arrays are not to be mmapable, but + can be read and written to disk. + +Limitations +----------- + +- Arbitrary subclasses of numpy.ndarray are not completely preserved. + Subclasses will be accepted for writing, but only the array data will + be written out. A regular numpy.ndarray object will be created + upon reading the file. + +.. warning:: + + Due to limitations in the interpretation of structured dtypes, dtypes + with fields with empty names will have the names replaced by 'f0', 'f1', + etc. Such arrays will not round-trip through the format entirely + accurately. The data is intact; only the field names will differ. We are + working on a fix for this. This fix will not require a change in the + file format. The arrays with such structures can still be saved and + restored, and the correct dtype may be restored by using the + ``loadedarray.view(correct_dtype)`` method. + +File extensions +--------------- + +We recommend using the ``.npy`` and ``.npz`` extensions for files saved +in this format. This is by no means a requirement; applications may wish +to use these file formats but use an extension specific to the +application. In the absence of an obvious alternative, however, +we suggest using ``.npy`` and ``.npz``. + +Version numbering +----------------- + +The version numbering of these formats is independent of NumPy version +numbering. If the format is upgraded, the code in `numpy.io` will still +be able to read and write Version 1.0 files. + +Format Version 1.0 +------------------ + +The first 6 bytes are a magic string: exactly ``\\x93NUMPY``. + +The next 1 byte is an unsigned byte: the major version number of the file +format, e.g. ``\\x01``. + +The next 1 byte is an unsigned byte: the minor version number of the file +format, e.g. ``\\x00``. Note: the version of the file format is not tied +to the version of the numpy package. + +The next 2 bytes form a little-endian unsigned short int: the length of +the header data HEADER_LEN. + +The next HEADER_LEN bytes form the header data describing the array's +format. It is an ASCII string which contains a Python literal expression +of a dictionary. It is terminated by a newline (``\\n``) and padded with +spaces (``\\x20``) to make the total of +``len(magic string) + 2 + len(length) + HEADER_LEN`` be evenly divisible +by 64 for alignment purposes. + +The dictionary contains three keys: + + "descr" : dtype.descr + An object that can be passed as an argument to the `numpy.dtype` + constructor to create the array's dtype. + "fortran_order" : bool + Whether the array data is Fortran-contiguous or not. Since + Fortran-contiguous arrays are a common form of non-C-contiguity, + we allow them to be written directly to disk for efficiency. + "shape" : tuple of int + The shape of the array. + +For repeatability and readability, the dictionary keys are sorted in +alphabetic order. This is for convenience only. A writer SHOULD implement +this if possible. A reader MUST NOT depend on this. + +Following the header comes the array data. If the dtype contains Python +objects (i.e. ``dtype.hasobject is True``), then the data is a Python +pickle of the array. Otherwise the data is the contiguous (either C- +or Fortran-, depending on ``fortran_order``) bytes of the array. +Consumers can figure out the number of bytes by multiplying the number +of elements given by the shape (noting that ``shape=()`` means there is +1 element) by ``dtype.itemsize``. + +Format Version 2.0 +------------------ + +The version 1.0 format only allowed the array header to have a total size of +65535 bytes. This can be exceeded by structured arrays with a large number of +columns. The version 2.0 format extends the header size to 4 GiB. +`numpy.save` will automatically save in 2.0 format if the data requires it, +else it will always use the more compatible 1.0 format. + +The description of the fourth element of the header therefore has become: +"The next 4 bytes form a little-endian unsigned int: the length of the header +data HEADER_LEN." + +Format Version 3.0 +------------------ + +This version replaces the ASCII string (which in practice was latin1) with +a utf8-encoded string, so supports structured types with any unicode field +names. + +Notes +----- +The ``.npy`` format, including motivation for creating it and a comparison of +alternatives, is described in the +:doc:`"npy-format" NEP `, however details have +evolved with time and this document is more current. + +""" +import io +import os +import pickle +import warnings + +import numpy +from numpy.lib._utils_impl import drop_metadata + + +__all__ = [] + +drop_metadata.__module__ = "numpy.lib.format" + +EXPECTED_KEYS = {'descr', 'fortran_order', 'shape'} +MAGIC_PREFIX = b'\x93NUMPY' +MAGIC_LEN = len(MAGIC_PREFIX) + 2 +ARRAY_ALIGN = 64 # plausible values are powers of 2 between 16 and 4096 +BUFFER_SIZE = 2**18 # size of buffer for reading npz files in bytes +# allow growth within the address space of a 64 bit machine along one axis +GROWTH_AXIS_MAX_DIGITS = 21 # = len(str(8*2**64-1)) hypothetical int1 dtype + +# difference between version 1.0 and 2.0 is a 4 byte (I) header length +# instead of 2 bytes (H) allowing storage of large structured arrays +_header_size_info = { + (1, 0): (' 255: + raise ValueError("major version must be 0 <= major < 256") + if minor < 0 or minor > 255: + raise ValueError("minor version must be 0 <= minor < 256") + return MAGIC_PREFIX + bytes([major, minor]) + +def read_magic(fp): + """ Read the magic string to get the version of the file format. + + Parameters + ---------- + fp : filelike object + + Returns + ------- + major : int + minor : int + """ + magic_str = _read_bytes(fp, MAGIC_LEN, "magic string") + if magic_str[:-2] != MAGIC_PREFIX: + msg = "the magic string is not correct; expected %r, got %r" + raise ValueError(msg % (MAGIC_PREFIX, magic_str[:-2])) + major, minor = magic_str[-2:] + return major, minor + + +def dtype_to_descr(dtype): + """ + Get a serializable descriptor from the dtype. + + The .descr attribute of a dtype object cannot be round-tripped through + the dtype() constructor. Simple types, like dtype('float32'), have + a descr which looks like a record array with one field with '' as + a name. The dtype() constructor interprets this as a request to give + a default name. Instead, we construct descriptor that can be passed to + dtype(). + + Parameters + ---------- + dtype : dtype + The dtype of the array that will be written to disk. + + Returns + ------- + descr : object + An object that can be passed to `numpy.dtype()` in order to + replicate the input dtype. + + """ + # NOTE: that drop_metadata may not return the right dtype e.g. for user + # dtypes. In that case our code below would fail the same, though. + new_dtype = drop_metadata(dtype) + if new_dtype is not dtype: + warnings.warn("metadata on a dtype is not saved to an npy/npz. " + "Use another format (such as pickle) to store it.", + UserWarning, stacklevel=2) + dtype = new_dtype + + if dtype.names is not None: + # This is a record array. The .descr is fine. XXX: parts of the + # record array with an empty name, like padding bytes, still get + # fiddled with. This needs to be fixed in the C implementation of + # dtype(). + return dtype.descr + elif not type(dtype)._legacy: + # this must be a user-defined dtype since numpy does not yet expose any + # non-legacy dtypes in the public API + # + # non-legacy dtypes don't yet have __array_interface__ + # support. Instead, as a hack, we use pickle to save the array, and lie + # that the dtype is object. When the array is loaded, the descriptor is + # unpickled with the array and the object dtype in the header is + # discarded. + # + # a future NEP should define a way to serialize user-defined + # descriptors and ideally work out the possible security implications + warnings.warn("Custom dtypes are saved as python objects using the " + "pickle protocol. Loading this file requires " + "allow_pickle=True to be set.", + UserWarning, stacklevel=2) + return "|O" + else: + return dtype.str + +def descr_to_dtype(descr): + """ + Returns a dtype based off the given description. + + This is essentially the reverse of `~lib.format.dtype_to_descr`. It will + remove the valueless padding fields created by, i.e. simple fields like + dtype('float32'), and then convert the description to its corresponding + dtype. + + Parameters + ---------- + descr : object + The object retrieved by dtype.descr. Can be passed to + `numpy.dtype` in order to replicate the input dtype. + + Returns + ------- + dtype : dtype + The dtype constructed by the description. + + """ + if isinstance(descr, str): + # No padding removal needed + return numpy.dtype(descr) + elif isinstance(descr, tuple): + # subtype, will always have a shape descr[1] + dt = descr_to_dtype(descr[0]) + return numpy.dtype((dt, descr[1])) + + titles = [] + names = [] + formats = [] + offsets = [] + offset = 0 + for field in descr: + if len(field) == 2: + name, descr_str = field + dt = descr_to_dtype(descr_str) + else: + name, descr_str, shape = field + dt = numpy.dtype((descr_to_dtype(descr_str), shape)) + + # Ignore padding bytes, which will be void bytes with '' as name + # Once support for blank names is removed, only "if name == ''" needed) + is_pad = (name == '' and dt.type is numpy.void and dt.names is None) + if not is_pad: + title, name = name if isinstance(name, tuple) else (None, name) + titles.append(title) + names.append(name) + formats.append(dt) + offsets.append(offset) + offset += dt.itemsize + + return numpy.dtype({'names': names, 'formats': formats, 'titles': titles, + 'offsets': offsets, 'itemsize': offset}) + +def header_data_from_array_1_0(array): + """ Get the dictionary of header metadata from a numpy.ndarray. + + Parameters + ---------- + array : numpy.ndarray + + Returns + ------- + d : dict + This has the appropriate entries for writing its string representation + to the header of the file. + """ + d = {'shape': array.shape} + if array.flags.c_contiguous: + d['fortran_order'] = False + elif array.flags.f_contiguous: + d['fortran_order'] = True + else: + # Totally non-contiguous data. We will have to make it C-contiguous + # before writing. Note that we need to test for C_CONTIGUOUS first + # because a 1-D array is both C_CONTIGUOUS and F_CONTIGUOUS. + d['fortran_order'] = False + + d['descr'] = dtype_to_descr(array.dtype) + return d + + +def _wrap_header(header, version): + """ + Takes a stringified header, and attaches the prefix and padding to it + """ + import struct + assert version is not None + fmt, encoding = _header_size_info[version] + header = header.encode(encoding) + hlen = len(header) + 1 + padlen = ARRAY_ALIGN - ((MAGIC_LEN + struct.calcsize(fmt) + hlen) % ARRAY_ALIGN) + try: + header_prefix = magic(*version) + struct.pack(fmt, hlen + padlen) + except struct.error: + msg = "Header length {} too big for version={}".format(hlen, version) + raise ValueError(msg) from None + + # Pad the header with spaces and a final newline such that the magic + # string, the header-length short and the header are aligned on a + # ARRAY_ALIGN byte boundary. This supports memory mapping of dtypes + # aligned up to ARRAY_ALIGN on systems like Linux where mmap() + # offset must be page-aligned (i.e. the beginning of the file). + return header_prefix + header + b' '*padlen + b'\n' + + +def _wrap_header_guess_version(header): + """ + Like `_wrap_header`, but chooses an appropriate version given the contents + """ + try: + return _wrap_header(header, (1, 0)) + except ValueError: + pass + + try: + ret = _wrap_header(header, (2, 0)) + except UnicodeEncodeError: + pass + else: + warnings.warn("Stored array in format 2.0. It can only be" + "read by NumPy >= 1.9", UserWarning, stacklevel=2) + return ret + + header = _wrap_header(header, (3, 0)) + warnings.warn("Stored array in format 3.0. It can only be " + "read by NumPy >= 1.17", UserWarning, stacklevel=2) + return header + + +def _write_array_header(fp, d, version=None): + """ Write the header for an array and returns the version used + + Parameters + ---------- + fp : filelike object + d : dict + This has the appropriate entries for writing its string representation + to the header of the file. + version : tuple or None + None means use oldest that works. Providing an explicit version will + raise a ValueError if the format does not allow saving this data. + Default: None + """ + header = ["{"] + for key, value in sorted(d.items()): + # Need to use repr here, since we eval these when reading + header.append("'%s': %s, " % (key, repr(value))) + header.append("}") + header = "".join(header) + + # Add some spare space so that the array header can be modified in-place + # when changing the array size, e.g. when growing it by appending data at + # the end. + shape = d['shape'] + header += " " * ((GROWTH_AXIS_MAX_DIGITS - len(repr( + shape[-1 if d['fortran_order'] else 0] + ))) if len(shape) > 0 else 0) + + if version is None: + header = _wrap_header_guess_version(header) + else: + header = _wrap_header(header, version) + fp.write(header) + +def write_array_header_1_0(fp, d): + """ Write the header for an array using the 1.0 format. + + Parameters + ---------- + fp : filelike object + d : dict + This has the appropriate entries for writing its string + representation to the header of the file. + """ + _write_array_header(fp, d, (1, 0)) + + +def write_array_header_2_0(fp, d): + """ Write the header for an array using the 2.0 format. + The 2.0 format allows storing very large structured arrays. + + Parameters + ---------- + fp : filelike object + d : dict + This has the appropriate entries for writing its string + representation to the header of the file. + """ + _write_array_header(fp, d, (2, 0)) + +def read_array_header_1_0(fp, max_header_size=_MAX_HEADER_SIZE): + """ + Read an array header from a filelike object using the 1.0 file format + version. + + This will leave the file object located just after the header. + + Parameters + ---------- + fp : filelike object + A file object or something with a `.read()` method like a file. + + Returns + ------- + shape : tuple of int + The shape of the array. + fortran_order : bool + The array data will be written out directly if it is either + C-contiguous or Fortran-contiguous. Otherwise, it will be made + contiguous before writing it out. + dtype : dtype + The dtype of the file's data. + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + + Raises + ------ + ValueError + If the data is invalid. + + """ + return _read_array_header( + fp, version=(1, 0), max_header_size=max_header_size) + +def read_array_header_2_0(fp, max_header_size=_MAX_HEADER_SIZE): + """ + Read an array header from a filelike object using the 2.0 file format + version. + + This will leave the file object located just after the header. + + Parameters + ---------- + fp : filelike object + A file object or something with a `.read()` method like a file. + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + + Returns + ------- + shape : tuple of int + The shape of the array. + fortran_order : bool + The array data will be written out directly if it is either + C-contiguous or Fortran-contiguous. Otherwise, it will be made + contiguous before writing it out. + dtype : dtype + The dtype of the file's data. + + Raises + ------ + ValueError + If the data is invalid. + + """ + return _read_array_header( + fp, version=(2, 0), max_header_size=max_header_size) + + +def _filter_header(s): + """Clean up 'L' in npz header ints. + + Cleans up the 'L' in strings representing integers. Needed to allow npz + headers produced in Python2 to be read in Python3. + + Parameters + ---------- + s : string + Npy file header. + + Returns + ------- + header : str + Cleaned up header. + + """ + import tokenize + from io import StringIO + + tokens = [] + last_token_was_number = False + for token in tokenize.generate_tokens(StringIO(s).readline): + token_type = token[0] + token_string = token[1] + if (last_token_was_number and + token_type == tokenize.NAME and + token_string == "L"): + continue + else: + tokens.append(token) + last_token_was_number = (token_type == tokenize.NUMBER) + return tokenize.untokenize(tokens) + + +def _read_array_header(fp, version, max_header_size=_MAX_HEADER_SIZE): + """ + see read_array_header_1_0 + """ + # Read an unsigned, little-endian short int which has the length of the + # header. + import ast + import struct + hinfo = _header_size_info.get(version) + if hinfo is None: + raise ValueError("Invalid version {!r}".format(version)) + hlength_type, encoding = hinfo + + hlength_str = _read_bytes(fp, struct.calcsize(hlength_type), "array header length") + header_length = struct.unpack(hlength_type, hlength_str)[0] + header = _read_bytes(fp, header_length, "array header") + header = header.decode(encoding) + if len(header) > max_header_size: + raise ValueError( + f"Header info length ({len(header)}) is large and may not be safe " + "to load securely.\n" + "To allow loading, adjust `max_header_size` or fully trust " + "the `.npy` file using `allow_pickle=True`.\n" + "For safety against large resource use or crashes, sandboxing " + "may be necessary.") + + # The header is a pretty-printed string representation of a literal + # Python dictionary with trailing newlines padded to a ARRAY_ALIGN byte + # boundary. The keys are strings. + # "shape" : tuple of int + # "fortran_order" : bool + # "descr" : dtype.descr + # Versions (2, 0) and (1, 0) could have been created by a Python 2 + # implementation before header filtering was implemented. + # + # For performance reasons, we try without _filter_header first though + try: + d = ast.literal_eval(header) + except SyntaxError as e: + if version <= (2, 0): + header = _filter_header(header) + try: + d = ast.literal_eval(header) + except SyntaxError as e2: + msg = "Cannot parse header: {!r}" + raise ValueError(msg.format(header)) from e2 + else: + warnings.warn( + "Reading `.npy` or `.npz` file required additional " + "header parsing as it was created on Python 2. Save the " + "file again to speed up loading and avoid this warning.", + UserWarning, stacklevel=4) + else: + msg = "Cannot parse header: {!r}" + raise ValueError(msg.format(header)) from e + if not isinstance(d, dict): + msg = "Header is not a dictionary: {!r}" + raise ValueError(msg.format(d)) + + if EXPECTED_KEYS != d.keys(): + keys = sorted(d.keys()) + msg = "Header does not contain the correct keys: {!r}" + raise ValueError(msg.format(keys)) + + # Sanity-check the values. + if (not isinstance(d['shape'], tuple) or + not all(isinstance(x, int) for x in d['shape'])): + msg = "shape is not valid: {!r}" + raise ValueError(msg.format(d['shape'])) + if not isinstance(d['fortran_order'], bool): + msg = "fortran_order is not a valid bool: {!r}" + raise ValueError(msg.format(d['fortran_order'])) + try: + dtype = descr_to_dtype(d['descr']) + except TypeError as e: + msg = "descr is not a valid dtype descriptor: {!r}" + raise ValueError(msg.format(d['descr'])) from e + + return d['shape'], d['fortran_order'], dtype + +def write_array(fp, array, version=None, allow_pickle=True, pickle_kwargs=None): + """ + Write an array to an NPY file, including a header. + + If the array is neither C-contiguous nor Fortran-contiguous AND the + file_like object is not a real file object, this function will have to + copy data in memory. + + Parameters + ---------- + fp : file_like object + An open, writable file object, or similar object with a + ``.write()`` method. + array : ndarray + The array to write to disk. + version : (int, int) or None, optional + The version number of the format. None means use the oldest + supported version that is able to store the data. Default: None + allow_pickle : bool, optional + Whether to allow writing pickled data. Default: True + pickle_kwargs : dict, optional + Additional keyword arguments to pass to pickle.dump, excluding + 'protocol'. These are only useful when pickling objects in object + arrays on Python 3 to Python 2 compatible format. + + Raises + ------ + ValueError + If the array cannot be persisted. This includes the case of + allow_pickle=False and array being an object array. + Various other errors + If the array contains Python objects as part of its dtype, the + process of pickling them may raise various errors if the objects + are not picklable. + + """ + _check_version(version) + _write_array_header(fp, header_data_from_array_1_0(array), version) + + if array.itemsize == 0: + buffersize = 0 + else: + # Set buffer size to 16 MiB to hide the Python loop overhead. + buffersize = max(16 * 1024 ** 2 // array.itemsize, 1) + + dtype_class = type(array.dtype) + + if array.dtype.hasobject or not dtype_class._legacy: + # We contain Python objects so we cannot write out the data + # directly. Instead, we will pickle it out + if not allow_pickle: + if array.dtype.hasobject: + raise ValueError("Object arrays cannot be saved when " + "allow_pickle=False") + if not dtype_class._legacy: + raise ValueError("User-defined dtypes cannot be saved " + "when allow_pickle=False") + if pickle_kwargs is None: + pickle_kwargs = {} + pickle.dump(array, fp, protocol=4, **pickle_kwargs) + elif array.flags.f_contiguous and not array.flags.c_contiguous: + if isfileobj(fp): + array.T.tofile(fp) + else: + for chunk in numpy.nditer( + array, flags=['external_loop', 'buffered', 'zerosize_ok'], + buffersize=buffersize, order='F'): + fp.write(chunk.tobytes('C')) + else: + if isfileobj(fp): + array.tofile(fp) + else: + for chunk in numpy.nditer( + array, flags=['external_loop', 'buffered', 'zerosize_ok'], + buffersize=buffersize, order='C'): + fp.write(chunk.tobytes('C')) + + +def read_array(fp, allow_pickle=False, pickle_kwargs=None, *, + max_header_size=_MAX_HEADER_SIZE): + """ + Read an array from an NPY file. + + Parameters + ---------- + fp : file_like object + If this is not a real file object, then this may take extra memory + and time. + allow_pickle : bool, optional + Whether to allow writing pickled data. Default: False + pickle_kwargs : dict + Additional keyword arguments to pass to pickle.load. These are only + useful when loading object arrays saved on Python 2 when using + Python 3. + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + This option is ignored when `allow_pickle` is passed. In that case + the file is by definition trusted and the limit is unnecessary. + + Returns + ------- + array : ndarray + The array from the data on disk. + + Raises + ------ + ValueError + If the data is invalid, or allow_pickle=False and the file contains + an object array. + + """ + if allow_pickle: + # Effectively ignore max_header_size, since `allow_pickle` indicates + # that the input is fully trusted. + max_header_size = 2**64 + + version = read_magic(fp) + _check_version(version) + shape, fortran_order, dtype = _read_array_header( + fp, version, max_header_size=max_header_size) + if len(shape) == 0: + count = 1 + else: + count = numpy.multiply.reduce(shape, dtype=numpy.int64) + + # Now read the actual data. + if dtype.hasobject: + # The array contained Python objects. We need to unpickle the data. + if not allow_pickle: + raise ValueError("Object arrays cannot be loaded when " + "allow_pickle=False") + if pickle_kwargs is None: + pickle_kwargs = {} + try: + array = pickle.load(fp, **pickle_kwargs) + except UnicodeError as err: + # Friendlier error message + raise UnicodeError("Unpickling a python object failed: %r\n" + "You may need to pass the encoding= option " + "to numpy.load" % (err,)) from err + else: + if isfileobj(fp): + # We can use the fast fromfile() function. + array = numpy.fromfile(fp, dtype=dtype, count=count) + else: + # This is not a real file. We have to read it the + # memory-intensive way. + # crc32 module fails on reads greater than 2 ** 32 bytes, + # breaking large reads from gzip streams. Chunk reads to + # BUFFER_SIZE bytes to avoid issue and reduce memory overhead + # of the read. In non-chunked case count < max_read_count, so + # only one read is performed. + + # Use np.ndarray instead of np.empty since the latter does + # not correctly instantiate zero-width string dtypes; see + # https://github.com/numpy/numpy/pull/6430 + array = numpy.ndarray(count, dtype=dtype) + + if dtype.itemsize > 0: + # If dtype.itemsize == 0 then there's nothing more to read + max_read_count = BUFFER_SIZE // min(BUFFER_SIZE, dtype.itemsize) + + for i in range(0, count, max_read_count): + read_count = min(max_read_count, count - i) + read_size = int(read_count * dtype.itemsize) + data = _read_bytes(fp, read_size, "array data") + array[i:i+read_count] = numpy.frombuffer(data, dtype=dtype, + count=read_count) + + if fortran_order: + array.shape = shape[::-1] + array = array.transpose() + else: + array.shape = shape + + return array + + +def open_memmap(filename, mode='r+', dtype=None, shape=None, + fortran_order=False, version=None, *, + max_header_size=_MAX_HEADER_SIZE): + """ + Open a .npy file as a memory-mapped array. + + This may be used to read an existing file or create a new one. + + Parameters + ---------- + filename : str or path-like + The name of the file on disk. This may *not* be a file-like + object. + mode : str, optional + The mode in which to open the file; the default is 'r+'. In + addition to the standard file modes, 'c' is also accepted to mean + "copy on write." See `memmap` for the available mode strings. + dtype : data-type, optional + The data type of the array if we are creating a new file in "write" + mode, if not, `dtype` is ignored. The default value is None, which + results in a data-type of `float64`. + shape : tuple of int + The shape of the array if we are creating a new file in "write" + mode, in which case this parameter is required. Otherwise, this + parameter is ignored and is thus optional. + fortran_order : bool, optional + Whether the array should be Fortran-contiguous (True) or + C-contiguous (False, the default) if we are creating a new file in + "write" mode. + version : tuple of int (major, minor) or None + If the mode is a "write" mode, then this is the version of the file + format used to create the file. None means use the oldest + supported version that is able to store the data. Default: None + max_header_size : int, optional + Maximum allowed size of the header. Large headers may not be safe + to load securely and thus require explicitly passing a larger value. + See :py:func:`ast.literal_eval()` for details. + + Returns + ------- + marray : memmap + The memory-mapped array. + + Raises + ------ + ValueError + If the data or the mode is invalid. + OSError + If the file is not found or cannot be opened correctly. + + See Also + -------- + numpy.memmap + + """ + if isfileobj(filename): + raise ValueError("Filename must be a string or a path-like object." + " Memmap cannot use existing file handles.") + + if 'w' in mode: + # We are creating the file, not reading it. + # Check if we ought to create the file. + _check_version(version) + # Ensure that the given dtype is an authentic dtype object rather + # than just something that can be interpreted as a dtype object. + dtype = numpy.dtype(dtype) + if dtype.hasobject: + msg = "Array can't be memory-mapped: Python objects in dtype." + raise ValueError(msg) + d = dict( + descr=dtype_to_descr(dtype), + fortran_order=fortran_order, + shape=shape, + ) + # If we got here, then it should be safe to create the file. + with open(os.fspath(filename), mode+'b') as fp: + _write_array_header(fp, d, version) + offset = fp.tell() + else: + # Read the header of the file first. + with open(os.fspath(filename), 'rb') as fp: + version = read_magic(fp) + _check_version(version) + + shape, fortran_order, dtype = _read_array_header( + fp, version, max_header_size=max_header_size) + if dtype.hasobject: + msg = "Array can't be memory-mapped: Python objects in dtype." + raise ValueError(msg) + offset = fp.tell() + + if fortran_order: + order = 'F' + else: + order = 'C' + + # We need to change a write-only mode to a read-write mode since we've + # already written data to the file. + if mode == 'w+': + mode = 'r+' + + marray = numpy.memmap(filename, dtype=dtype, shape=shape, order=order, + mode=mode, offset=offset) + + return marray + + +def _read_bytes(fp, size, error_template="ran out of data"): + """ + Read from file-like object until size bytes are read. + Raises ValueError if not EOF is encountered before size bytes are read. + Non-blocking objects only supported if they derive from io objects. + + Required as e.g. ZipExtFile in python 2.6 can return less data than + requested. + """ + data = bytes() + while True: + # io files (default in python3) return None or raise on + # would-block, python2 file will truncate, probably nothing can be + # done about that. note that regular files can't be non-blocking + try: + r = fp.read(size - len(data)) + data += r + if len(r) == 0 or len(data) == size: + break + except BlockingIOError: + pass + if len(data) != size: + msg = "EOF: reading %s, expected %d bytes got %d" + raise ValueError(msg % (error_template, size, len(data))) + else: + return data + + +def isfileobj(f): + if not isinstance(f, (io.FileIO, io.BufferedReader, io.BufferedWriter)): + return False + try: + # BufferedReader/Writer may raise OSError when + # fetching `fileno()` (e.g. when wrapping BytesIO). + f.fileno() + return True + except OSError: + return False diff --git a/venv/lib/python3.10/site-packages/numpy/lib/format.pyi b/venv/lib/python3.10/site-packages/numpy/lib/format.pyi new file mode 100644 index 0000000000000000000000000000000000000000..57c7e1e206e0a576eb9966898ac439ab3c7c1969 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/format.pyi @@ -0,0 +1,22 @@ +from typing import Literal, Final + +__all__: list[str] = [] + +EXPECTED_KEYS: Final[set[str]] +MAGIC_PREFIX: Final[bytes] +MAGIC_LEN: Literal[8] +ARRAY_ALIGN: Literal[64] +BUFFER_SIZE: Literal[262144] # 2**18 + +def magic(major, minor): ... +def read_magic(fp): ... +def dtype_to_descr(dtype): ... +def descr_to_dtype(descr): ... +def header_data_from_array_1_0(array): ... +def write_array_header_1_0(fp, d): ... +def write_array_header_2_0(fp, d): ... +def read_array_header_1_0(fp): ... +def read_array_header_2_0(fp): ... +def write_array(fp, array, version=..., allow_pickle=..., pickle_kwargs=...): ... +def read_array(fp, allow_pickle=..., pickle_kwargs=...): ... +def open_memmap(filename, mode=..., dtype=..., shape=..., fortran_order=..., version=...): ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/introspect.py b/venv/lib/python3.10/site-packages/numpy/lib/introspect.py new file mode 100644 index 0000000000000000000000000000000000000000..4826440dd410948273832e402ccea583bcfb361b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/introspect.py @@ -0,0 +1,95 @@ +""" +Introspection helper functions. +""" +import re + +__all__ = ['opt_func_info'] + + +def opt_func_info(func_name=None, signature=None): + """ + Returns a dictionary containing the currently supported CPU dispatched + features for all optimized functions. + + Parameters + ---------- + func_name : str (optional) + Regular expression to filter by function name. + + signature : str (optional) + Regular expression to filter by data type. + + Returns + ------- + dict + A dictionary where keys are optimized function names and values are + nested dictionaries indicating supported targets based on data types. + + Examples + -------- + Retrieve dispatch information for functions named 'add' or 'sub' and + data types 'float64' or 'float32': + + >>> import numpy as np + >>> dict = np.lib.introspect.opt_func_info( + ... func_name="add|abs", signature="float64|complex64" + ... ) + >>> import json + >>> print(json.dumps(dict, indent=2)) + { + "absolute": { + "dd": { + "current": "SSE41", + "available": "SSE41 baseline(SSE SSE2 SSE3)" + }, + "Ff": { + "current": "FMA3__AVX2", + "available": "AVX512F FMA3__AVX2 baseline(SSE SSE2 SSE3)" + }, + "Dd": { + "current": "FMA3__AVX2", + "available": "AVX512F FMA3__AVX2 baseline(SSE SSE2 SSE3)" + } + }, + "add": { + "ddd": { + "current": "FMA3__AVX2", + "available": "FMA3__AVX2 baseline(SSE SSE2 SSE3)" + }, + "FFF": { + "current": "FMA3__AVX2", + "available": "FMA3__AVX2 baseline(SSE SSE2 SSE3)" + } + } + } + + """ + from numpy._core._multiarray_umath import ( + __cpu_targets_info__ as targets, dtype + ) + + if func_name is not None: + func_pattern = re.compile(func_name) + matching_funcs = { + k: v for k, v in targets.items() + if func_pattern.search(k) + } + else: + matching_funcs = targets + + if signature is not None: + sig_pattern = re.compile(signature) + matching_sigs = {} + for k, v in matching_funcs.items(): + matching_chars = {} + for chars, targets in v.items(): + if any( + sig_pattern.search(c) or sig_pattern.search(dtype(c).name) + for c in chars + ): + matching_chars[chars] = targets + if matching_chars: + matching_sigs[k] = matching_chars + else: + matching_sigs = matching_funcs + return matching_sigs diff --git a/venv/lib/python3.10/site-packages/numpy/lib/introspect.pyi b/venv/lib/python3.10/site-packages/numpy/lib/introspect.pyi new file mode 100644 index 0000000000000000000000000000000000000000..7929981cd6362ad7a992b50a74b37c700f00cbe2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/introspect.pyi @@ -0,0 +1,3 @@ +__all__ = ["opt_func_info"] + +def opt_func_info(func_name: str | None = None, signature: str | None = None) -> dict[str, dict[str, dict[str, str]]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/mixins.py b/venv/lib/python3.10/site-packages/numpy/lib/mixins.py new file mode 100644 index 0000000000000000000000000000000000000000..5e78ac0990b34c9926f7f5a819d9652605d5a0ca --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/mixins.py @@ -0,0 +1,182 @@ +""" +Mixin classes for custom array types that don't inherit from ndarray. +""" +from numpy._core import umath as um + + +__all__ = ['NDArrayOperatorsMixin'] + + +def _disables_array_ufunc(obj): + """True when __array_ufunc__ is set to None.""" + try: + return obj.__array_ufunc__ is None + except AttributeError: + return False + + +def _binary_method(ufunc, name): + """Implement a forward binary method with a ufunc, e.g., __add__.""" + def func(self, other): + if _disables_array_ufunc(other): + return NotImplemented + return ufunc(self, other) + func.__name__ = '__{}__'.format(name) + return func + + +def _reflected_binary_method(ufunc, name): + """Implement a reflected binary method with a ufunc, e.g., __radd__.""" + def func(self, other): + if _disables_array_ufunc(other): + return NotImplemented + return ufunc(other, self) + func.__name__ = '__r{}__'.format(name) + return func + + +def _inplace_binary_method(ufunc, name): + """Implement an in-place binary method with a ufunc, e.g., __iadd__.""" + def func(self, other): + return ufunc(self, other, out=(self,)) + func.__name__ = '__i{}__'.format(name) + return func + + +def _numeric_methods(ufunc, name): + """Implement forward, reflected and inplace binary methods with a ufunc.""" + return (_binary_method(ufunc, name), + _reflected_binary_method(ufunc, name), + _inplace_binary_method(ufunc, name)) + + +def _unary_method(ufunc, name): + """Implement a unary special method with a ufunc.""" + def func(self): + return ufunc(self) + func.__name__ = '__{}__'.format(name) + return func + + +class NDArrayOperatorsMixin: + """Mixin defining all operator special methods using __array_ufunc__. + + This class implements the special methods for almost all of Python's + builtin operators defined in the `operator` module, including comparisons + (``==``, ``>``, etc.) and arithmetic (``+``, ``*``, ``-``, etc.), by + deferring to the ``__array_ufunc__`` method, which subclasses must + implement. + + It is useful for writing classes that do not inherit from `numpy.ndarray`, + but that should support arithmetic and numpy universal functions like + arrays as described in `A Mechanism for Overriding Ufuncs + `_. + + As an trivial example, consider this implementation of an ``ArrayLike`` + class that simply wraps a NumPy array and ensures that the result of any + arithmetic operation is also an ``ArrayLike`` object: + + >>> import numbers + >>> class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin): + ... def __init__(self, value): + ... self.value = np.asarray(value) + ... + ... # One might also consider adding the built-in list type to this + ... # list, to support operations like np.add(array_like, list) + ... _HANDLED_TYPES = (np.ndarray, numbers.Number) + ... + ... def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + ... out = kwargs.get('out', ()) + ... for x in inputs + out: + ... # Only support operations with instances of + ... # _HANDLED_TYPES. Use ArrayLike instead of type(self) + ... # for isinstance to allow subclasses that don't + ... # override __array_ufunc__ to handle ArrayLike objects. + ... if not isinstance( + ... x, self._HANDLED_TYPES + (ArrayLike,) + ... ): + ... return NotImplemented + ... + ... # Defer to the implementation of the ufunc + ... # on unwrapped values. + ... inputs = tuple(x.value if isinstance(x, ArrayLike) else x + ... for x in inputs) + ... if out: + ... kwargs['out'] = tuple( + ... x.value if isinstance(x, ArrayLike) else x + ... for x in out) + ... result = getattr(ufunc, method)(*inputs, **kwargs) + ... + ... if type(result) is tuple: + ... # multiple return values + ... return tuple(type(self)(x) for x in result) + ... elif method == 'at': + ... # no return value + ... return None + ... else: + ... # one return value + ... return type(self)(result) + ... + ... def __repr__(self): + ... return '%s(%r)' % (type(self).__name__, self.value) + + In interactions between ``ArrayLike`` objects and numbers or numpy arrays, + the result is always another ``ArrayLike``: + + >>> x = ArrayLike([1, 2, 3]) + >>> x - 1 + ArrayLike(array([0, 1, 2])) + >>> 1 - x + ArrayLike(array([ 0, -1, -2])) + >>> np.arange(3) - x + ArrayLike(array([-1, -1, -1])) + >>> x - np.arange(3) + ArrayLike(array([1, 1, 1])) + + Note that unlike ``numpy.ndarray``, ``ArrayLike`` does not allow operations + with arbitrary, unrecognized types. This ensures that interactions with + ArrayLike preserve a well-defined casting hierarchy. + + """ + __slots__ = () + # Like np.ndarray, this mixin class implements "Option 1" from the ufunc + # overrides NEP. + + # comparisons don't have reflected and in-place versions + __lt__ = _binary_method(um.less, 'lt') + __le__ = _binary_method(um.less_equal, 'le') + __eq__ = _binary_method(um.equal, 'eq') + __ne__ = _binary_method(um.not_equal, 'ne') + __gt__ = _binary_method(um.greater, 'gt') + __ge__ = _binary_method(um.greater_equal, 'ge') + + # numeric methods + __add__, __radd__, __iadd__ = _numeric_methods(um.add, 'add') + __sub__, __rsub__, __isub__ = _numeric_methods(um.subtract, 'sub') + __mul__, __rmul__, __imul__ = _numeric_methods(um.multiply, 'mul') + __matmul__, __rmatmul__, __imatmul__ = _numeric_methods( + um.matmul, 'matmul') + # Python 3 does not use __div__, __rdiv__, or __idiv__ + __truediv__, __rtruediv__, __itruediv__ = _numeric_methods( + um.true_divide, 'truediv') + __floordiv__, __rfloordiv__, __ifloordiv__ = _numeric_methods( + um.floor_divide, 'floordiv') + __mod__, __rmod__, __imod__ = _numeric_methods(um.remainder, 'mod') + __divmod__ = _binary_method(um.divmod, 'divmod') + __rdivmod__ = _reflected_binary_method(um.divmod, 'divmod') + # __idivmod__ does not exist + # TODO: handle the optional third argument for __pow__? + __pow__, __rpow__, __ipow__ = _numeric_methods(um.power, 'pow') + __lshift__, __rlshift__, __ilshift__ = _numeric_methods( + um.left_shift, 'lshift') + __rshift__, __rrshift__, __irshift__ = _numeric_methods( + um.right_shift, 'rshift') + __and__, __rand__, __iand__ = _numeric_methods(um.bitwise_and, 'and') + __xor__, __rxor__, __ixor__ = _numeric_methods(um.bitwise_xor, 'xor') + __or__, __ror__, __ior__ = _numeric_methods(um.bitwise_or, 'or') + + # unary methods + __neg__ = _unary_method(um.negative, 'neg') + __pos__ = _unary_method(um.positive, 'pos') + __abs__ = _unary_method(um.absolute, 'abs') + __invert__ = _unary_method(um.invert, 'invert') diff --git a/venv/lib/python3.10/site-packages/numpy/lib/mixins.pyi b/venv/lib/python3.10/site-packages/numpy/lib/mixins.pyi new file mode 100644 index 0000000000000000000000000000000000000000..d13d0fe81df48c4e725e6090c699498a9518e4a6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/mixins.pyi @@ -0,0 +1,74 @@ +from abc import ABCMeta, abstractmethod +from typing import Literal as L, Any + +from numpy import ufunc + +__all__ = ["NDArrayOperatorsMixin"] + +# NOTE: `NDArrayOperatorsMixin` is not formally an abstract baseclass, +# even though it's reliant on subclasses implementing `__array_ufunc__` + +# NOTE: The accepted input- and output-types of the various dunders are +# completely dependent on how `__array_ufunc__` is implemented. +# As such, only little type safety can be provided here. + +class NDArrayOperatorsMixin(metaclass=ABCMeta): + @abstractmethod + def __array_ufunc__( + self, + ufunc: ufunc, + method: L["__call__", "reduce", "reduceat", "accumulate", "outer", "at"], + *inputs: Any, + **kwargs: Any, + ) -> Any: ... + def __lt__(self, other: Any) -> Any: ... + def __le__(self, other: Any) -> Any: ... + def __eq__(self, other: Any) -> Any: ... + def __ne__(self, other: Any) -> Any: ... + def __gt__(self, other: Any) -> Any: ... + def __ge__(self, other: Any) -> Any: ... + def __add__(self, other: Any) -> Any: ... + def __radd__(self, other: Any) -> Any: ... + def __iadd__(self, other: Any) -> Any: ... + def __sub__(self, other: Any) -> Any: ... + def __rsub__(self, other: Any) -> Any: ... + def __isub__(self, other: Any) -> Any: ... + def __mul__(self, other: Any) -> Any: ... + def __rmul__(self, other: Any) -> Any: ... + def __imul__(self, other: Any) -> Any: ... + def __matmul__(self, other: Any) -> Any: ... + def __rmatmul__(self, other: Any) -> Any: ... + def __imatmul__(self, other: Any) -> Any: ... + def __truediv__(self, other: Any) -> Any: ... + def __rtruediv__(self, other: Any) -> Any: ... + def __itruediv__(self, other: Any) -> Any: ... + def __floordiv__(self, other: Any) -> Any: ... + def __rfloordiv__(self, other: Any) -> Any: ... + def __ifloordiv__(self, other: Any) -> Any: ... + def __mod__(self, other: Any) -> Any: ... + def __rmod__(self, other: Any) -> Any: ... + def __imod__(self, other: Any) -> Any: ... + def __divmod__(self, other: Any) -> Any: ... + def __rdivmod__(self, other: Any) -> Any: ... + def __pow__(self, other: Any) -> Any: ... + def __rpow__(self, other: Any) -> Any: ... + def __ipow__(self, other: Any) -> Any: ... + def __lshift__(self, other: Any) -> Any: ... + def __rlshift__(self, other: Any) -> Any: ... + def __ilshift__(self, other: Any) -> Any: ... + def __rshift__(self, other: Any) -> Any: ... + def __rrshift__(self, other: Any) -> Any: ... + def __irshift__(self, other: Any) -> Any: ... + def __and__(self, other: Any) -> Any: ... + def __rand__(self, other: Any) -> Any: ... + def __iand__(self, other: Any) -> Any: ... + def __xor__(self, other: Any) -> Any: ... + def __rxor__(self, other: Any) -> Any: ... + def __ixor__(self, other: Any) -> Any: ... + def __or__(self, other: Any) -> Any: ... + def __ror__(self, other: Any) -> Any: ... + def __ior__(self, other: Any) -> Any: ... + def __neg__(self) -> Any: ... + def __pos__(self) -> Any: ... + def __abs__(self) -> Any: ... + def __invert__(self) -> Any: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/npyio.py b/venv/lib/python3.10/site-packages/numpy/lib/npyio.py new file mode 100644 index 0000000000000000000000000000000000000000..1003ef5be4b1940ddf6943a90d1bee786677c55e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/npyio.py @@ -0,0 +1,3 @@ +from ._npyio_impl import ( + __doc__, DataSource, NpzFile +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/npyio.pyi b/venv/lib/python3.10/site-packages/numpy/lib/npyio.pyi new file mode 100644 index 0000000000000000000000000000000000000000..fd3ae8f5a28778764a94fd791e72816b4db51320 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/npyio.pyi @@ -0,0 +1,5 @@ +from numpy.lib._npyio_impl import ( + DataSource as DataSource, + NpzFile as NpzFile, + __doc__ as __doc__, +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.py b/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.py new file mode 100644 index 0000000000000000000000000000000000000000..8f4bae4f4721a4c16d9eb5950a97d0d2d3bb2779 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.py @@ -0,0 +1,1685 @@ +""" +Collection of utilities to manipulate structured arrays. + +Most of these functions were initially implemented by John Hunter for +matplotlib. They have been rewritten and extended for convenience. + +""" +import itertools + +import numpy as np +import numpy.ma as ma +import numpy.ma.mrecords as mrec +from numpy._core.overrides import array_function_dispatch +from numpy.lib._iotools import _is_string_like + + +__all__ = [ + 'append_fields', 'apply_along_fields', 'assign_fields_by_name', + 'drop_fields', 'find_duplicates', 'flatten_descr', + 'get_fieldstructure', 'get_names', 'get_names_flat', + 'join_by', 'merge_arrays', 'rec_append_fields', + 'rec_drop_fields', 'rec_join', 'recursive_fill_fields', + 'rename_fields', 'repack_fields', 'require_fields', + 'stack_arrays', 'structured_to_unstructured', 'unstructured_to_structured', + ] + + +def _recursive_fill_fields_dispatcher(input, output): + return (input, output) + + +@array_function_dispatch(_recursive_fill_fields_dispatcher) +def recursive_fill_fields(input, output): + """ + Fills fields from output with fields from input, + with support for nested structures. + + Parameters + ---------- + input : ndarray + Input array. + output : ndarray + Output array. + + Notes + ----- + * `output` should be at least the same size as `input` + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> a = np.array([(1, 10.), (2, 20.)], dtype=[('A', np.int64), ('B', np.float64)]) + >>> b = np.zeros((3,), dtype=a.dtype) + >>> rfn.recursive_fill_fields(a, b) + array([(1, 10.), (2, 20.), (0, 0.)], dtype=[('A', '>> import numpy as np + >>> dt = np.dtype([(('a', 'A'), np.int64), ('b', np.double, 3)]) + >>> dt.descr + [(('a', 'A'), '>> _get_fieldspec(dt) + [(('a', 'A'), dtype('int64')), ('b', dtype(('>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> rfn.get_names(np.empty((1,), dtype=[('A', int)]).dtype) + ('A',) + >>> rfn.get_names(np.empty((1,), dtype=[('A',int), ('B', float)]).dtype) + ('A', 'B') + >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])]) + >>> rfn.get_names(adtype) + ('a', ('b', ('ba', 'bb'))) + """ + listnames = [] + names = adtype.names + for name in names: + current = adtype[name] + if current.names is not None: + listnames.append((name, tuple(get_names(current)))) + else: + listnames.append(name) + return tuple(listnames) + + +def get_names_flat(adtype): + """ + Returns the field names of the input datatype as a tuple. Input datatype + must have fields otherwise error is raised. + Nested structure are flattened beforehand. + + Parameters + ---------- + adtype : dtype + Input datatype + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> rfn.get_names_flat(np.empty((1,), dtype=[('A', int)]).dtype) is None + False + >>> rfn.get_names_flat(np.empty((1,), dtype=[('A',int), ('B', str)]).dtype) + ('A', 'B') + >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])]) + >>> rfn.get_names_flat(adtype) + ('a', 'b', 'ba', 'bb') + """ + listnames = [] + names = adtype.names + for name in names: + listnames.append(name) + current = adtype[name] + if current.names is not None: + listnames.extend(get_names_flat(current)) + return tuple(listnames) + + +def flatten_descr(ndtype): + """ + Flatten a structured data-type description. + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> ndtype = np.dtype([('a', '>> rfn.flatten_descr(ndtype) + (('a', dtype('int32')), ('ba', dtype('float64')), ('bb', dtype('int32'))) + + """ + names = ndtype.names + if names is None: + return (('', ndtype),) + else: + descr = [] + for field in names: + (typ, _) = ndtype.fields[field] + if typ.names is not None: + descr.extend(flatten_descr(typ)) + else: + descr.append((field, typ)) + return tuple(descr) + + +def _zip_dtype(seqarrays, flatten=False): + newdtype = [] + if flatten: + for a in seqarrays: + newdtype.extend(flatten_descr(a.dtype)) + else: + for a in seqarrays: + current = a.dtype + if current.names is not None and len(current.names) == 1: + # special case - dtypes of 1 field are flattened + newdtype.extend(_get_fieldspec(current)) + else: + newdtype.append(('', current)) + return np.dtype(newdtype) + + +def _zip_descr(seqarrays, flatten=False): + """ + Combine the dtype description of a series of arrays. + + Parameters + ---------- + seqarrays : sequence of arrays + Sequence of arrays + flatten : {boolean}, optional + Whether to collapse nested descriptions. + """ + return _zip_dtype(seqarrays, flatten=flatten).descr + + +def get_fieldstructure(adtype, lastname=None, parents=None,): + """ + Returns a dictionary with fields indexing lists of their parent fields. + + This function is used to simplify access to fields nested in other fields. + + Parameters + ---------- + adtype : np.dtype + Input datatype + lastname : optional + Last processed field name (used internally during recursion). + parents : dictionary + Dictionary of parent fields (used internally during recursion). + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> ndtype = np.dtype([('A', int), + ... ('B', [('BA', int), + ... ('BB', [('BBA', int), ('BBB', int)])])]) + >>> rfn.get_fieldstructure(ndtype) + ... # XXX: possible regression, order of BBA and BBB is swapped + {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'], 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']} + + """ + if parents is None: + parents = {} + names = adtype.names + for name in names: + current = adtype[name] + if current.names is not None: + if lastname: + parents[name] = [lastname, ] + else: + parents[name] = [] + parents.update(get_fieldstructure(current, name, parents)) + else: + lastparent = list((parents.get(lastname, []) or [])) + if lastparent: + lastparent.append(lastname) + elif lastname: + lastparent = [lastname, ] + parents[name] = lastparent or [] + return parents + + +def _izip_fields_flat(iterable): + """ + Returns an iterator of concatenated fields from a sequence of arrays, + collapsing any nested structure. + + """ + for element in iterable: + if isinstance(element, np.void): + yield from _izip_fields_flat(tuple(element)) + else: + yield element + + +def _izip_fields(iterable): + """ + Returns an iterator of concatenated fields from a sequence of arrays. + + """ + for element in iterable: + if (hasattr(element, '__iter__') and + not isinstance(element, str)): + yield from _izip_fields(element) + elif isinstance(element, np.void) and len(tuple(element)) == 1: + # this statement is the same from the previous expression + yield from _izip_fields(element) + else: + yield element + + +def _izip_records(seqarrays, fill_value=None, flatten=True): + """ + Returns an iterator of concatenated items from a sequence of arrays. + + Parameters + ---------- + seqarrays : sequence of arrays + Sequence of arrays. + fill_value : {None, integer} + Value used to pad shorter iterables. + flatten : {True, False}, + Whether to + """ + + # Should we flatten the items, or just use a nested approach + if flatten: + zipfunc = _izip_fields_flat + else: + zipfunc = _izip_fields + + for tup in itertools.zip_longest(*seqarrays, fillvalue=fill_value): + yield tuple(zipfunc(tup)) + + +def _fix_output(output, usemask=True, asrecarray=False): + """ + Private function: return a recarray, a ndarray, a MaskedArray + or a MaskedRecords depending on the input parameters + """ + if not isinstance(output, ma.MaskedArray): + usemask = False + if usemask: + if asrecarray: + output = output.view(mrec.MaskedRecords) + else: + output = ma.filled(output) + if asrecarray: + output = output.view(np.recarray) + return output + + +def _fix_defaults(output, defaults=None): + """ + Update the fill_value and masked data of `output` + from the default given in a dictionary defaults. + """ + names = output.dtype.names + (data, mask, fill_value) = (output.data, output.mask, output.fill_value) + for (k, v) in (defaults or {}).items(): + if k in names: + fill_value[k] = v + data[k][mask[k]] = v + return output + + +def _merge_arrays_dispatcher(seqarrays, fill_value=None, flatten=None, + usemask=None, asrecarray=None): + return seqarrays + + +@array_function_dispatch(_merge_arrays_dispatcher) +def merge_arrays(seqarrays, fill_value=-1, flatten=False, + usemask=False, asrecarray=False): + """ + Merge arrays field by field. + + Parameters + ---------- + seqarrays : sequence of ndarrays + Sequence of arrays + fill_value : {float}, optional + Filling value used to pad missing data on the shorter arrays. + flatten : {False, True}, optional + Whether to collapse nested fields. + usemask : {False, True}, optional + Whether to return a masked array or not. + asrecarray : {False, True}, optional + Whether to return a recarray (MaskedRecords) or not. + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> rfn.merge_arrays((np.array([1, 2]), np.array([10., 20., 30.]))) + array([( 1, 10.), ( 2, 20.), (-1, 30.)], + dtype=[('f0', '>> rfn.merge_arrays((np.array([1, 2], dtype=np.int64), + ... np.array([10., 20., 30.])), usemask=False) + array([(1, 10.0), (2, 20.0), (-1, 30.0)], + dtype=[('f0', '>> rfn.merge_arrays((np.array([1, 2]).view([('a', np.int64)]), + ... np.array([10., 20., 30.])), + ... usemask=False, asrecarray=True) + rec.array([( 1, 10.), ( 2, 20.), (-1, 30.)], + dtype=[('a', '>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> a = np.array([(1, (2, 3.0)), (4, (5, 6.0))], + ... dtype=[('a', np.int64), ('b', [('ba', np.double), ('bb', np.int64)])]) + >>> rfn.drop_fields(a, 'a') + array([((2., 3),), ((5., 6),)], + dtype=[('b', [('ba', '>> rfn.drop_fields(a, 'ba') + array([(1, (3,)), (4, (6,))], dtype=[('a', '>> rfn.drop_fields(a, ['ba', 'bb']) + array([(1,), (4,)], dtype=[('a', '>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))], + ... dtype=[('a', int),('b', [('ba', float), ('bb', (float, 2))])]) + >>> rfn.rename_fields(a, {'a':'A', 'bb':'BB'}) + array([(1, (2., [ 3., 30.])), (4, (5., [ 6., 60.]))], + dtype=[('A', ' 1: + data = merge_arrays(data, flatten=True, usemask=usemask, + fill_value=fill_value) + else: + data = data.pop() + # + output = ma.masked_all( + max(len(base), len(data)), + dtype=_get_fieldspec(base.dtype) + _get_fieldspec(data.dtype)) + output = recursive_fill_fields(base, output) + output = recursive_fill_fields(data, output) + # + return _fix_output(output, usemask=usemask, asrecarray=asrecarray) + + +def _rec_append_fields_dispatcher(base, names, data, dtypes=None): + yield base + yield from data + + +@array_function_dispatch(_rec_append_fields_dispatcher) +def rec_append_fields(base, names, data, dtypes=None): + """ + Add new fields to an existing array. + + The names of the fields are given with the `names` arguments, + the corresponding values with the `data` arguments. + If a single field is appended, `names`, `data` and `dtypes` do not have + to be lists but just values. + + Parameters + ---------- + base : array + Input array to extend. + names : string, sequence + String or sequence of strings corresponding to the names + of the new fields. + data : array or sequence of arrays + Array or sequence of arrays storing the fields to add to the base. + dtypes : sequence of datatypes, optional + Datatype or sequence of datatypes. + If None, the datatypes are estimated from the `data`. + + See Also + -------- + append_fields + + Returns + ------- + appended_array : np.recarray + """ + return append_fields(base, names, data=data, dtypes=dtypes, + asrecarray=True, usemask=False) + + +def _repack_fields_dispatcher(a, align=None, recurse=None): + return (a,) + + +@array_function_dispatch(_repack_fields_dispatcher) +def repack_fields(a, align=False, recurse=False): + """ + Re-pack the fields of a structured array or dtype in memory. + + The memory layout of structured datatypes allows fields at arbitrary + byte offsets. This means the fields can be separated by padding bytes, + their offsets can be non-monotonically increasing, and they can overlap. + + This method removes any overlaps and reorders the fields in memory so they + have increasing byte offsets, and adds or removes padding bytes depending + on the `align` option, which behaves like the `align` option to + `numpy.dtype`. + + If `align=False`, this method produces a "packed" memory layout in which + each field starts at the byte the previous field ended, and any padding + bytes are removed. + + If `align=True`, this methods produces an "aligned" memory layout in which + each field's offset is a multiple of its alignment, and the total itemsize + is a multiple of the largest alignment, by adding padding bytes as needed. + + Parameters + ---------- + a : ndarray or dtype + array or dtype for which to repack the fields. + align : boolean + If true, use an "aligned" memory layout, otherwise use a "packed" layout. + recurse : boolean + If True, also repack nested structures. + + Returns + ------- + repacked : ndarray or dtype + Copy of `a` with fields repacked, or `a` itself if no repacking was + needed. + + Examples + -------- + >>> import numpy as np + + >>> from numpy.lib import recfunctions as rfn + >>> def print_offsets(d): + ... print("offsets:", [d.fields[name][1] for name in d.names]) + ... print("itemsize:", d.itemsize) + ... + >>> dt = np.dtype('u1, >> dt + dtype({'names': ['f0', 'f1', 'f2'], 'formats': ['u1', '>> print_offsets(dt) + offsets: [0, 8, 16] + itemsize: 24 + >>> packed_dt = rfn.repack_fields(dt) + >>> packed_dt + dtype([('f0', 'u1'), ('f1', '>> print_offsets(packed_dt) + offsets: [0, 1, 9] + itemsize: 17 + + """ + if not isinstance(a, np.dtype): + dt = repack_fields(a.dtype, align=align, recurse=recurse) + return a.astype(dt, copy=False) + + if a.names is None: + return a + + fieldinfo = [] + for name in a.names: + tup = a.fields[name] + if recurse: + fmt = repack_fields(tup[0], align=align, recurse=True) + else: + fmt = tup[0] + + if len(tup) == 3: + name = (tup[2], name) + + fieldinfo.append((name, fmt)) + + dt = np.dtype(fieldinfo, align=align) + return np.dtype((a.type, dt)) + +def _get_fields_and_offsets(dt, offset=0): + """ + Returns a flat list of (dtype, count, offset) tuples of all the + scalar fields in the dtype "dt", including nested fields, in left + to right order. + """ + + # counts up elements in subarrays, including nested subarrays, and returns + # base dtype and count + def count_elem(dt): + count = 1 + while dt.shape != (): + for size in dt.shape: + count *= size + dt = dt.base + return dt, count + + fields = [] + for name in dt.names: + field = dt.fields[name] + f_dt, f_offset = field[0], field[1] + f_dt, n = count_elem(f_dt) + + if f_dt.names is None: + fields.append((np.dtype((f_dt, (n,))), n, f_offset + offset)) + else: + subfields = _get_fields_and_offsets(f_dt, f_offset + offset) + size = f_dt.itemsize + + for i in range(n): + if i == 0: + # optimization: avoid list comprehension if no subarray + fields.extend(subfields) + else: + fields.extend([(d, c, o + i*size) for d, c, o in subfields]) + return fields + +def _common_stride(offsets, counts, itemsize): + """ + Returns the stride between the fields, or None if the stride is not + constant. The values in "counts" designate the lengths of + subarrays. Subarrays are treated as many contiguous fields, with + always positive stride. + """ + if len(offsets) <= 1: + return itemsize + + negative = offsets[1] < offsets[0] # negative stride + if negative: + # reverse, so offsets will be ascending + it = zip(reversed(offsets), reversed(counts)) + else: + it = zip(offsets, counts) + + prev_offset = None + stride = None + for offset, count in it: + if count != 1: # subarray: always c-contiguous + if negative: + return None # subarrays can never have a negative stride + if stride is None: + stride = itemsize + if stride != itemsize: + return None + end_offset = offset + (count - 1) * itemsize + else: + end_offset = offset + + if prev_offset is not None: + new_stride = offset - prev_offset + if stride is None: + stride = new_stride + if stride != new_stride: + return None + + prev_offset = end_offset + + if negative: + return -stride + return stride + + +def _structured_to_unstructured_dispatcher(arr, dtype=None, copy=None, + casting=None): + return (arr,) + +@array_function_dispatch(_structured_to_unstructured_dispatcher) +def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'): + """ + Converts an n-D structured array into an (n+1)-D unstructured array. + + The new array will have a new last dimension equal in size to the + number of field-elements of the input array. If not supplied, the output + datatype is determined from the numpy type promotion rules applied to all + the field datatypes. + + Nested fields, as well as each element of any subarray fields, all count + as a single field-elements. + + Parameters + ---------- + arr : ndarray + Structured array or dtype to convert. Cannot contain object datatype. + dtype : dtype, optional + The dtype of the output unstructured array. + copy : bool, optional + If true, always return a copy. If false, a view is returned if + possible, such as when the `dtype` and strides of the fields are + suitable and the array subtype is one of `numpy.ndarray`, + `numpy.recarray` or `numpy.memmap`. + + .. versionchanged:: 1.25.0 + A view can now be returned if the fields are separated by a + uniform stride. + + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + See casting argument of `numpy.ndarray.astype`. Controls what kind of + data casting may occur. + + Returns + ------- + unstructured : ndarray + Unstructured array with one more dimension. + + Examples + -------- + >>> import numpy as np + + >>> from numpy.lib import recfunctions as rfn + >>> a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)]) + >>> a + array([(0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.]), + (0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.])], + dtype=[('a', '>> rfn.structured_to_unstructured(a) + array([[0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) + + >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)], + ... dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')]) + >>> np.mean(rfn.structured_to_unstructured(b[['x', 'z']]), axis=-1) + array([ 3. , 5.5, 9. , 11. ]) + + """ + if arr.dtype.names is None: + raise ValueError('arr must be a structured array') + + fields = _get_fields_and_offsets(arr.dtype) + n_fields = len(fields) + if n_fields == 0 and dtype is None: + raise ValueError("arr has no fields. Unable to guess dtype") + elif n_fields == 0: + # too many bugs elsewhere for this to work now + raise NotImplementedError("arr with no fields is not supported") + + dts, counts, offsets = zip(*fields) + names = ['f{}'.format(n) for n in range(n_fields)] + + if dtype is None: + out_dtype = np.result_type(*[dt.base for dt in dts]) + else: + out_dtype = np.dtype(dtype) + + # Use a series of views and casts to convert to an unstructured array: + + # first view using flattened fields (doesn't work for object arrays) + # Note: dts may include a shape for subarrays + flattened_fields = np.dtype({'names': names, + 'formats': dts, + 'offsets': offsets, + 'itemsize': arr.dtype.itemsize}) + arr = arr.view(flattened_fields) + + # we only allow a few types to be unstructured by manipulating the + # strides, because we know it won't work with, for example, np.matrix nor + # np.ma.MaskedArray. + can_view = type(arr) in (np.ndarray, np.recarray, np.memmap) + if (not copy) and can_view and all(dt.base == out_dtype for dt in dts): + # all elements have the right dtype already; if they have a common + # stride, we can just return a view + common_stride = _common_stride(offsets, counts, out_dtype.itemsize) + if common_stride is not None: + wrap = arr.__array_wrap__ + + new_shape = arr.shape + (sum(counts), out_dtype.itemsize) + new_strides = arr.strides + (abs(common_stride), 1) + + arr = arr[..., np.newaxis].view(np.uint8) # view as bytes + arr = arr[..., min(offsets):] # remove the leading unused data + arr = np.lib.stride_tricks.as_strided(arr, + new_shape, + new_strides, + subok=True) + + # cast and drop the last dimension again + arr = arr.view(out_dtype)[..., 0] + + if common_stride < 0: + arr = arr[..., ::-1] # reverse, if the stride was negative + if type(arr) is not type(wrap.__self__): + # Some types (e.g. recarray) turn into an ndarray along the + # way, so we have to wrap it again in order to match the + # behavior with copy=True. + arr = wrap(arr) + return arr + + # next cast to a packed format with all fields converted to new dtype + packed_fields = np.dtype({'names': names, + 'formats': [(out_dtype, dt.shape) for dt in dts]}) + arr = arr.astype(packed_fields, copy=copy, casting=casting) + + # finally is it safe to view the packed fields as the unstructured type + return arr.view((out_dtype, (sum(counts),))) + + +def _unstructured_to_structured_dispatcher(arr, dtype=None, names=None, + align=None, copy=None, casting=None): + return (arr,) + +@array_function_dispatch(_unstructured_to_structured_dispatcher) +def unstructured_to_structured(arr, dtype=None, names=None, align=False, + copy=False, casting='unsafe'): + """ + Converts an n-D unstructured array into an (n-1)-D structured array. + + The last dimension of the input array is converted into a structure, with + number of field-elements equal to the size of the last dimension of the + input array. By default all output fields have the input array's dtype, but + an output structured dtype with an equal number of fields-elements can be + supplied instead. + + Nested fields, as well as each element of any subarray fields, all count + towards the number of field-elements. + + Parameters + ---------- + arr : ndarray + Unstructured array or dtype to convert. + dtype : dtype, optional + The structured dtype of the output array + names : list of strings, optional + If dtype is not supplied, this specifies the field names for the output + dtype, in order. The field dtypes will be the same as the input array. + align : boolean, optional + Whether to create an aligned memory layout. + copy : bool, optional + See copy argument to `numpy.ndarray.astype`. If true, always return a + copy. If false, and `dtype` requirements are satisfied, a view is + returned. + casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional + See casting argument of `numpy.ndarray.astype`. Controls what kind of + data casting may occur. + + Returns + ------- + structured : ndarray + Structured array with fewer dimensions. + + Examples + -------- + >>> import numpy as np + + >>> from numpy.lib import recfunctions as rfn + >>> dt = np.dtype([('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)]) + >>> a = np.arange(20).reshape((4,5)) + >>> a + array([[ 0, 1, 2, 3, 4], + [ 5, 6, 7, 8, 9], + [10, 11, 12, 13, 14], + [15, 16, 17, 18, 19]]) + >>> rfn.unstructured_to_structured(a, dt) + array([( 0, ( 1., 2), [ 3., 4.]), ( 5, ( 6., 7), [ 8., 9.]), + (10, (11., 12), [13., 14.]), (15, (16., 17), [18., 19.])], + dtype=[('a', '>> import numpy as np + + >>> from numpy.lib import recfunctions as rfn + >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)], + ... dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')]) + >>> rfn.apply_along_fields(np.mean, b) + array([ 2.66666667, 5.33333333, 8.66666667, 11. ]) + >>> rfn.apply_along_fields(np.mean, b[['x', 'z']]) + array([ 3. , 5.5, 9. , 11. ]) + + """ + if arr.dtype.names is None: + raise ValueError('arr must be a structured array') + + uarr = structured_to_unstructured(arr) + return func(uarr, axis=-1) + # works and avoids axis requirement, but very, very slow: + #return np.apply_along_axis(func, -1, uarr) + +def _assign_fields_by_name_dispatcher(dst, src, zero_unassigned=None): + return dst, src + +@array_function_dispatch(_assign_fields_by_name_dispatcher) +def assign_fields_by_name(dst, src, zero_unassigned=True): + """ + Assigns values from one structured array to another by field name. + + Normally in numpy >= 1.14, assignment of one structured array to another + copies fields "by position", meaning that the first field from the src is + copied to the first field of the dst, and so on, regardless of field name. + + This function instead copies "by field name", such that fields in the dst + are assigned from the identically named field in the src. This applies + recursively for nested structures. This is how structure assignment worked + in numpy >= 1.6 to <= 1.13. + + Parameters + ---------- + dst : ndarray + src : ndarray + The source and destination arrays during assignment. + zero_unassigned : bool, optional + If True, fields in the dst for which there was no matching + field in the src are filled with the value 0 (zero). This + was the behavior of numpy <= 1.13. If False, those fields + are not modified. + """ + + if dst.dtype.names is None: + dst[...] = src + return + + for name in dst.dtype.names: + if name not in src.dtype.names: + if zero_unassigned: + dst[name] = 0 + else: + assign_fields_by_name(dst[name], src[name], + zero_unassigned) + +def _require_fields_dispatcher(array, required_dtype): + return (array,) + +@array_function_dispatch(_require_fields_dispatcher) +def require_fields(array, required_dtype): + """ + Casts a structured array to a new dtype using assignment by field-name. + + This function assigns from the old to the new array by name, so the + value of a field in the output array is the value of the field with the + same name in the source array. This has the effect of creating a new + ndarray containing only the fields "required" by the required_dtype. + + If a field name in the required_dtype does not exist in the + input array, that field is created and set to 0 in the output array. + + Parameters + ---------- + a : ndarray + array to cast + required_dtype : dtype + datatype for output array + + Returns + ------- + out : ndarray + array with the new dtype, with field values copied from the fields in + the input array with the same name + + Examples + -------- + >>> import numpy as np + + >>> from numpy.lib import recfunctions as rfn + >>> a = np.ones(4, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')]) + >>> rfn.require_fields(a, [('b', 'f4'), ('c', 'u1')]) + array([(1., 1), (1., 1), (1., 1), (1., 1)], + dtype=[('b', '>> rfn.require_fields(a, [('b', 'f4'), ('newf', 'u1')]) + array([(1., 0), (1., 0), (1., 0), (1., 0)], + dtype=[('b', '>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> x = np.array([1, 2,]) + >>> rfn.stack_arrays(x) is x + True + >>> z = np.array([('A', 1), ('B', 2)], dtype=[('A', '|S3'), ('B', float)]) + >>> zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)], + ... dtype=[('A', '|S3'), ('B', np.double), ('C', np.double)]) + >>> test = rfn.stack_arrays((z,zz)) + >>> test + masked_array(data=[(b'A', 1.0, --), (b'B', 2.0, --), (b'a', 10.0, 100.0), + (b'b', 20.0, 200.0), (b'c', 30.0, 300.0)], + mask=[(False, False, True), (False, False, True), + (False, False, False), (False, False, False), + (False, False, False)], + fill_value=(b'N/A', 1e+20, 1e+20), + dtype=[('A', 'S3'), ('B', ' '%s'" % + (cdtype, fdtype)) + # Only one field: use concatenate + if len(newdescr) == 1: + output = ma.concatenate(seqarrays) + else: + # + output = ma.masked_all((np.sum(nrecords),), newdescr) + offset = np.cumsum(np.r_[0, nrecords]) + seen = [] + for (a, n, i, j) in zip(seqarrays, fldnames, offset[:-1], offset[1:]): + names = a.dtype.names + if names is None: + output['f%i' % len(seen)][i:j] = a + else: + for name in n: + output[name][i:j] = a[name] + if name not in seen: + seen.append(name) + # + return _fix_output(_fix_defaults(output, defaults), + usemask=usemask, asrecarray=asrecarray) + + +def _find_duplicates_dispatcher( + a, key=None, ignoremask=None, return_index=None): + return (a,) + + +@array_function_dispatch(_find_duplicates_dispatcher) +def find_duplicates(a, key=None, ignoremask=True, return_index=False): + """ + Find the duplicates in a structured array along a given key + + Parameters + ---------- + a : array-like + Input array + key : {string, None}, optional + Name of the fields along which to check the duplicates. + If None, the search is performed by records + ignoremask : {True, False}, optional + Whether masked data should be discarded or considered as duplicates. + return_index : {False, True}, optional + Whether to return the indices of the duplicated values. + + Examples + -------- + >>> import numpy as np + >>> from numpy.lib import recfunctions as rfn + >>> ndtype = [('a', int)] + >>> a = np.ma.array([1, 1, 1, 2, 2, 3, 3], + ... mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype) + >>> rfn.find_duplicates(a, ignoremask=True, return_index=True) + (masked_array(data=[(1,), (1,), (2,), (2,)], + mask=[(False,), (False,), (False,), (False,)], + fill_value=(999999,), + dtype=[('a', '= nb1)] - nb1 + (r1cmn, r2cmn) = (len(idx_1), len(idx_2)) + if jointype == 'inner': + (r1spc, r2spc) = (0, 0) + elif jointype == 'outer': + idx_out = idx_sort[~flag_in] + idx_1 = np.concatenate((idx_1, idx_out[(idx_out < nb1)])) + idx_2 = np.concatenate((idx_2, idx_out[(idx_out >= nb1)] - nb1)) + (r1spc, r2spc) = (len(idx_1) - r1cmn, len(idx_2) - r2cmn) + elif jointype == 'leftouter': + idx_out = idx_sort[~flag_in] + idx_1 = np.concatenate((idx_1, idx_out[(idx_out < nb1)])) + (r1spc, r2spc) = (len(idx_1) - r1cmn, 0) + # Select the entries from each input + (s1, s2) = (r1[idx_1], r2[idx_2]) + # + # Build the new description of the output array ....... + # Start with the key fields + ndtype = _get_fieldspec(r1k.dtype) + + # Add the fields from r1 + for fname, fdtype in _get_fieldspec(r1.dtype): + if fname not in key: + ndtype.append((fname, fdtype)) + + # Add the fields from r2 + for fname, fdtype in _get_fieldspec(r2.dtype): + # Have we seen the current name already ? + # we need to rebuild this list every time + names = list(name for name, dtype in ndtype) + try: + nameidx = names.index(fname) + except ValueError: + #... we haven't: just add the description to the current list + ndtype.append((fname, fdtype)) + else: + # collision + _, cdtype = ndtype[nameidx] + if fname in key: + # The current field is part of the key: take the largest dtype + ndtype[nameidx] = (fname, max(fdtype, cdtype)) + else: + # The current field is not part of the key: add the suffixes, + # and place the new field adjacent to the old one + ndtype[nameidx:nameidx + 1] = [ + (fname + r1postfix, cdtype), + (fname + r2postfix, fdtype) + ] + # Rebuild a dtype from the new fields + ndtype = np.dtype(ndtype) + # Find the largest nb of common fields : + # r1cmn and r2cmn should be equal, but... + cmn = max(r1cmn, r2cmn) + # Construct an empty array + output = ma.masked_all((cmn + r1spc + r2spc,), dtype=ndtype) + names = output.dtype.names + for f in r1names: + selected = s1[f] + if f not in names or (f in r2names and not r2postfix and f not in key): + f += r1postfix + current = output[f] + current[:r1cmn] = selected[:r1cmn] + if jointype in ('outer', 'leftouter'): + current[cmn:cmn + r1spc] = selected[r1cmn:] + for f in r2names: + selected = s2[f] + if f not in names or (f in r1names and not r1postfix and f not in key): + f += r2postfix + current = output[f] + current[:r2cmn] = selected[:r2cmn] + if (jointype == 'outer') and r2spc: + current[-r2spc:] = selected[r2cmn:] + # Sort and finalize the output + output.sort(order=key) + kwargs = dict(usemask=usemask, asrecarray=asrecarray) + return _fix_output(_fix_defaults(output, defaults), **kwargs) + + +def _rec_join_dispatcher( + key, r1, r2, jointype=None, r1postfix=None, r2postfix=None, + defaults=None): + return (r1, r2) + + +@array_function_dispatch(_rec_join_dispatcher) +def rec_join(key, r1, r2, jointype='inner', r1postfix='1', r2postfix='2', + defaults=None): + """ + Join arrays `r1` and `r2` on keys. + Alternative to join_by, that always returns a np.recarray. + + See Also + -------- + join_by : equivalent function + """ + kwargs = dict(jointype=jointype, r1postfix=r1postfix, r2postfix=r2postfix, + defaults=defaults, usemask=False, asrecarray=True) + return join_by(key, r1, r2, **kwargs) + + +del array_function_dispatch diff --git a/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.pyi b/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.pyi new file mode 100644 index 0000000000000000000000000000000000000000..442530e9cd39803ad0411527dbb24c6b60de54c4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/recfunctions.pyi @@ -0,0 +1,435 @@ +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import Any, Literal, TypeAlias, overload + +from _typeshed import Incomplete +from typing_extensions import TypeVar + +import numpy as np +import numpy.typing as npt +from numpy._typing import _DTypeLike, _DTypeLikeVoid +from numpy.ma.mrecords import MaskedRecords + +__all__ = [ + "append_fields", + "apply_along_fields", + "assign_fields_by_name", + "drop_fields", + "find_duplicates", + "flatten_descr", + "get_fieldstructure", + "get_names", + "get_names_flat", + "join_by", + "merge_arrays", + "rec_append_fields", + "rec_drop_fields", + "rec_join", + "recursive_fill_fields", + "rename_fields", + "repack_fields", + "require_fields", + "stack_arrays", + "structured_to_unstructured", + "unstructured_to_structured", +] + +_T = TypeVar("_T") +_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...]) +_ScalarT = TypeVar("_ScalarT", bound=np.generic) +_DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any]) +_ArrayT = TypeVar("_ArrayT", bound=npt.NDArray[Any]) +_VoidArrayT = TypeVar("_VoidArrayT", bound=npt.NDArray[np.void]) +_NonVoidDTypeT = TypeVar("_NonVoidDTypeT", bound=_NonVoidDType) + +_OneOrMany: TypeAlias = _T | Iterable[_T] +_BuiltinSequence: TypeAlias = tuple[_T, ...] | list[_T] + +_NestedNames: TypeAlias = tuple[str | _NestedNames, ...] +_NonVoid: TypeAlias = np.bool | np.number | np.character | np.datetime64 | np.timedelta64 | np.object_ +_NonVoidDType: TypeAlias = np.dtype[_NonVoid] | np.dtypes.StringDType + +_JoinType: TypeAlias = Literal["inner", "outer", "leftouter"] + +### + +def recursive_fill_fields(input: npt.NDArray[np.void], output: _VoidArrayT) -> _VoidArrayT: ... + +# +def get_names(adtype: np.dtype[np.void]) -> _NestedNames: ... +def get_names_flat(adtype: np.dtype[np.void]) -> tuple[str, ...]: ... + +# +@overload +def flatten_descr(ndtype: _NonVoidDTypeT) -> tuple[tuple[Literal[""], _NonVoidDTypeT]]: ... +@overload +def flatten_descr(ndtype: np.dtype[np.void]) -> tuple[tuple[str, np.dtype[Any]]]: ... + +# +def get_fieldstructure( + adtype: np.dtype[np.void], + lastname: str | None = None, + parents: dict[str, list[str]] | None = None, +) -> dict[str, list[str]]: ... + +# +@overload +def merge_arrays( + seqarrays: Sequence[np.ndarray[_ShapeT, np.dtype[Any]]] | np.ndarray[_ShapeT, np.dtype[Any]], + fill_value: float = -1, + flatten: bool = False, + usemask: bool = False, + asrecarray: bool = False, +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def merge_arrays( + seqarrays: Sequence[npt.ArrayLike] | np.void, + fill_value: float = -1, + flatten: bool = False, + usemask: bool = False, + asrecarray: bool = False, +) -> np.recarray[Any, np.dtype[np.void]]: ... + +# +@overload +def drop_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + drop_names: str | Iterable[str], + usemask: bool = True, + asrecarray: Literal[False] = False, +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def drop_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + drop_names: str | Iterable[str], + usemask: bool, + asrecarray: Literal[True], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def drop_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + drop_names: str | Iterable[str], + usemask: bool = True, + *, + asrecarray: Literal[True], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... + +# +@overload +def rename_fields( + base: MaskedRecords[_ShapeT, np.dtype[np.void]], + namemapper: Mapping[str, str], +) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ... +@overload +def rename_fields( + base: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], + namemapper: Mapping[str, str], +) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ... +@overload +def rename_fields( + base: np.recarray[_ShapeT, np.dtype[np.void]], + namemapper: Mapping[str, str], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def rename_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + namemapper: Mapping[str, str], +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... + +# +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None, + fill_value: int, + usemask: Literal[False], + asrecarray: Literal[False] = False, +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None = None, + fill_value: int = -1, + *, + usemask: Literal[False], + asrecarray: Literal[False] = False, +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None, + fill_value: int, + usemask: Literal[False], + asrecarray: Literal[True], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None = None, + fill_value: int = -1, + *, + usemask: Literal[False], + asrecarray: Literal[True], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None = None, + fill_value: int = -1, + usemask: Literal[True] = True, + asrecarray: Literal[False] = False, +) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None, + fill_value: int, + usemask: Literal[True], + asrecarray: Literal[True], +) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ... +@overload +def append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None = None, + fill_value: int = -1, + usemask: Literal[True] = True, + *, + asrecarray: Literal[True], +) -> MaskedRecords[_ShapeT, np.dtype[np.void]]: ... + +# +def rec_drop_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + drop_names: str | Iterable[str], +) -> np.recarray[_ShapeT, np.dtype[np.void]]: ... + +# +def rec_append_fields( + base: np.ndarray[_ShapeT, np.dtype[np.void]], + names: _OneOrMany[str], + data: _OneOrMany[npt.NDArray[Any]], + dtypes: _BuiltinSequence[np.dtype[Any]] | None = None, +) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ... + +# TODO(jorenham): Stop passing `void` directly once structured dtypes are implemented, +# e.g. using a `TypeVar` with constraints. +# https://github.com/numpy/numtype/issues/92 +@overload +def repack_fields(a: _DTypeT, align: bool = False, recurse: bool = False) -> _DTypeT: ... +@overload +def repack_fields(a: _ScalarT, align: bool = False, recurse: bool = False) -> _ScalarT: ... +@overload +def repack_fields(a: _ArrayT, align: bool = False, recurse: bool = False) -> _ArrayT: ... + +# TODO(jorenham): Attempt shape-typing (return type has ndim == arr.ndim + 1) +@overload +def structured_to_unstructured( + arr: npt.NDArray[np.void], + dtype: _DTypeLike[_ScalarT], + copy: bool = False, + casting: np._CastingKind = "unsafe", +) -> npt.NDArray[_ScalarT]: ... +@overload +def structured_to_unstructured( + arr: npt.NDArray[np.void], + dtype: npt.DTypeLike | None = None, + copy: bool = False, + casting: np._CastingKind = "unsafe", +) -> npt.NDArray[Any]: ... + +# +@overload +def unstructured_to_structured( + arr: npt.NDArray[Any], + dtype: npt.DTypeLike, + names: None = None, + align: bool = False, + copy: bool = False, + casting: str = "unsafe", +) -> npt.NDArray[np.void]: ... +@overload +def unstructured_to_structured( + arr: npt.NDArray[Any], + dtype: None, + names: _OneOrMany[str], + align: bool = False, + copy: bool = False, + casting: str = "unsafe", +) -> npt.NDArray[np.void]: ... + +# +def apply_along_fields( + func: Callable[[np.ndarray[_ShapeT, Any]], npt.NDArray[Any]], + arr: np.ndarray[_ShapeT, np.dtype[np.void]], +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... + +# +def assign_fields_by_name(dst: npt.NDArray[np.void], src: npt.NDArray[np.void], zero_unassigned: bool = True) -> None: ... + +# +def require_fields( + array: np.ndarray[_ShapeT, np.dtype[np.void]], + required_dtype: _DTypeLikeVoid, +) -> np.ndarray[_ShapeT, np.dtype[np.void]]: ... + +# TODO(jorenham): Attempt shape-typing +@overload +def stack_arrays( + arrays: _ArrayT, + defaults: Mapping[str, object] | None = None, + usemask: bool = True, + asrecarray: bool = False, + autoconvert: bool = False, +) -> _ArrayT: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None, + usemask: Literal[False], + asrecarray: Literal[False] = False, + autoconvert: bool = False, +) -> npt.NDArray[np.void]: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None = None, + *, + usemask: Literal[False], + asrecarray: Literal[False] = False, + autoconvert: bool = False, +) -> npt.NDArray[np.void]: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None = None, + *, + usemask: Literal[False], + asrecarray: Literal[True], + autoconvert: bool = False, +) -> np.recarray[tuple[int, ...], np.dtype[np.void]]: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None = None, + usemask: Literal[True] = True, + asrecarray: Literal[False] = False, + autoconvert: bool = False, +) -> np.ma.MaskedArray[tuple[int, ...], np.dtype[np.void]]: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None, + usemask: Literal[True], + asrecarray: Literal[True], + autoconvert: bool = False, +) -> MaskedRecords[tuple[int, ...], np.dtype[np.void]]: ... +@overload +def stack_arrays( + arrays: Sequence[npt.NDArray[Any]], + defaults: Mapping[str, Incomplete] | None = None, + usemask: Literal[True] = True, + *, + asrecarray: Literal[True], + autoconvert: bool = False, +) -> MaskedRecords[tuple[int, ...], np.dtype[np.void]]: ... + +# +@overload +def find_duplicates( + a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], + key: str | None = None, + ignoremask: bool = True, + return_index: Literal[False] = False, +) -> np.ma.MaskedArray[_ShapeT, np.dtype[np.void]]: ... +@overload +def find_duplicates( + a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], + key: str | None, + ignoremask: bool, + return_index: Literal[True], +) -> tuple[np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], np.ndarray[_ShapeT, np.dtype[np.int_]]]: ... +@overload +def find_duplicates( + a: np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], + key: str | None = None, + ignoremask: bool = True, + *, + return_index: Literal[True], +) -> tuple[np.ma.MaskedArray[_ShapeT, np.dtype[np.void]], np.ndarray[_ShapeT, np.dtype[np.int_]]]: ... + +# +@overload +def join_by( + key: str | Sequence[str], + r1: npt.NDArray[np.void], + r2: npt.NDArray[np.void], + jointype: _JoinType = "inner", + r1postfix: str = "1", + r2postfix: str = "2", + defaults: Mapping[str, object] | None = None, + *, + usemask: Literal[False], + asrecarray: Literal[False] = False, +) -> np.ndarray[tuple[int], np.dtype[np.void]]: ... +@overload +def join_by( + key: str | Sequence[str], + r1: npt.NDArray[np.void], + r2: npt.NDArray[np.void], + jointype: _JoinType = "inner", + r1postfix: str = "1", + r2postfix: str = "2", + defaults: Mapping[str, object] | None = None, + *, + usemask: Literal[False], + asrecarray: Literal[True], +) -> np.recarray[tuple[int], np.dtype[np.void]]: ... +@overload +def join_by( + key: str | Sequence[str], + r1: npt.NDArray[np.void], + r2: npt.NDArray[np.void], + jointype: _JoinType = "inner", + r1postfix: str = "1", + r2postfix: str = "2", + defaults: Mapping[str, object] | None = None, + usemask: Literal[True] = True, + asrecarray: Literal[False] = False, +) -> np.ma.MaskedArray[tuple[int], np.dtype[np.void]]: ... +@overload +def join_by( + key: str | Sequence[str], + r1: npt.NDArray[np.void], + r2: npt.NDArray[np.void], + jointype: _JoinType = "inner", + r1postfix: str = "1", + r2postfix: str = "2", + defaults: Mapping[str, object] | None = None, + usemask: Literal[True] = True, + *, + asrecarray: Literal[True], +) -> MaskedRecords[tuple[int], np.dtype[np.void]]: ... + +# +def rec_join( + key: str | Sequence[str], + r1: npt.NDArray[np.void], + r2: npt.NDArray[np.void], + jointype: _JoinType = "inner", + r1postfix: str = "1", + r2postfix: str = "2", + defaults: Mapping[str, object] | None = None, +) -> np.recarray[tuple[int], np.dtype[np.void]]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/lib/scimath.py b/venv/lib/python3.10/site-packages/numpy/lib/scimath.py new file mode 100644 index 0000000000000000000000000000000000000000..ffd05ef9f364fd64fa9d91d5e6f41a9d90744307 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/scimath.py @@ -0,0 +1,4 @@ +from ._scimath_impl import ( + __all__, __doc__, sqrt, log, log2, logn, log10, power, arccos, arcsin, + arctanh +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/scimath.pyi b/venv/lib/python3.10/site-packages/numpy/lib/scimath.pyi new file mode 100644 index 0000000000000000000000000000000000000000..cff5b9097faec0c831a99098f3275c591e2be10a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/scimath.pyi @@ -0,0 +1,12 @@ +from ._scimath_impl import ( + __all__ as __all__, + sqrt as sqrt, + log as log, + log2 as log2, + logn as logn, + log10 as log10, + power as power, + arccos as arccos, + arcsin as arcsin, + arctanh as arctanh, +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.py b/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.py new file mode 100644 index 0000000000000000000000000000000000000000..ba567be0c823408f316178f28ed80a970ee9f516 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.py @@ -0,0 +1,3 @@ +from ._stride_tricks_impl import ( + __doc__, as_strided, sliding_window_view +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.pyi b/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.pyi new file mode 100644 index 0000000000000000000000000000000000000000..eb46f28ae5f43e92c7c14df31edb154cf6d1f60b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/stride_tricks.pyi @@ -0,0 +1,4 @@ +from numpy.lib._stride_tricks_impl import ( + as_strided as as_strided, + sliding_window_view as sliding_window_view, +) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__init__.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23e41210ae4de9175079a2b823911a75d98573af Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__datasource.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__datasource.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d6f76bd486fe98270ae87024e56cdeaedb44639 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__datasource.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__iotools.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__iotools.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e6e8e3a4c9aa1b5491e34401a0b26d9c28cf42c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__iotools.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__version.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__version.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8052eb5a311eb3095b9a483bc5c7002f34a6cf28 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test__version.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_array_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_array_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6adbab1e492ba1bfb9df0c520443c9bb33137969 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_array_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraypad.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraypad.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c8fc4b20ea94ee9209dceac85bbf9f7ffb4d588 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraypad.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraysetops.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraysetops.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a2cd93a383826390719a1bbc22996622d71f7b6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arraysetops.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arrayterator.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arrayterator.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..077490f9bb69c42d204c6364ca27e7a173b0a791 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_arrayterator.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_format.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_format.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aeb6b36490d0e352a81f2b19b4d01d218843913b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_format.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_histograms.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_histograms.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40fe29255d8530c8ed8389261828f9f1242e4a93 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_histograms.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_index_tricks.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_index_tricks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cae9b60dea90056c30c88613c4a13cb3a1f0a2a4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_index_tricks.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_io.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_io.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..743d617a35d0ab5eebe9ca20fb46869dd152cc01 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_io.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_loadtxt.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_loadtxt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17c401fed7e3e0187e4fb82b37e8aa4f69a67b79 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_loadtxt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_mixins.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_mixins.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eaa25a27a36b37206b21815eed92b504bc46351f Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_mixins.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_nanfunctions.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_nanfunctions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5edcbf2a930f68a83f393883b4bc36a366a7688 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_nanfunctions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_packbits.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_packbits.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..847d0a312c87fec0ae4183d1ad4179cdaceacbd1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_packbits.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_polynomial.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_polynomial.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b401a9ac8fe0f0943ee7107fd88d101d3bb361f4 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_polynomial.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_recfunctions.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_recfunctions.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0056dead1ed665fd6a70c7d58da08f4f8392e103 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_recfunctions.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_regression.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_regression.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23234a29649fa5b0b72837733843c326328d3b8a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_regression.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_shape_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_shape_base.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c94de5ea15ae32abbb9ec0c8a80360c5a5c5bd0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_shape_base.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_stride_tricks.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_stride_tricks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a29803574ddcbf36be8df33ede90acc5c86f52e Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_stride_tricks.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_twodim_base.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_twodim_base.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d7d0b51c0c6674f745f6e7a59ace753697a4787 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_twodim_base.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_type_check.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_type_check.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a11443157c54172a431f2118105e7f5b3f889d5a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_type_check.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_ufunclike.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_ufunclike.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcde36f3e860d2a397bf20015b94dea048534e41 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_ufunclike.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64cdd5b3f528fc27ea2848e9669ed4fa2f1135e7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/lib/tests/__pycache__/test_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test__datasource.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__datasource.py new file mode 100644 index 0000000000000000000000000000000000000000..c8149abc30c40d6b9528fa9f3a98b0340d5914d3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__datasource.py @@ -0,0 +1,350 @@ +import os +import pytest +from tempfile import mkdtemp, mkstemp, NamedTemporaryFile +from shutil import rmtree + +import numpy.lib._datasource as datasource +from numpy.testing import assert_, assert_equal, assert_raises + +import urllib.request as urllib_request +from urllib.parse import urlparse +from urllib.error import URLError + + +def urlopen_stub(url, data=None): + '''Stub to replace urlopen for testing.''' + if url == valid_httpurl(): + tmpfile = NamedTemporaryFile(prefix='urltmp_') + return tmpfile + else: + raise URLError('Name or service not known') + +# setup and teardown +old_urlopen = None + + +def setup_module(): + global old_urlopen + + old_urlopen = urllib_request.urlopen + urllib_request.urlopen = urlopen_stub + + +def teardown_module(): + urllib_request.urlopen = old_urlopen + +# A valid website for more robust testing +http_path = 'http://www.google.com/' +http_file = 'index.html' + +http_fakepath = 'http://fake.abc.web/site/' +http_fakefile = 'fake.txt' + +malicious_files = ['/etc/shadow', '../../shadow', + '..\\system.dat', 'c:\\windows\\system.dat'] + +magic_line = b'three is the magic number' + + +# Utility functions used by many tests +def valid_textfile(filedir): + # Generate and return a valid temporary file. + fd, path = mkstemp(suffix='.txt', prefix='dstmp_', dir=filedir, text=True) + os.close(fd) + return path + + +def invalid_textfile(filedir): + # Generate and return an invalid filename. + fd, path = mkstemp(suffix='.txt', prefix='dstmp_', dir=filedir) + os.close(fd) + os.remove(path) + return path + + +def valid_httpurl(): + return http_path+http_file + + +def invalid_httpurl(): + return http_fakepath+http_fakefile + + +def valid_baseurl(): + return http_path + + +def invalid_baseurl(): + return http_fakepath + + +def valid_httpfile(): + return http_file + + +def invalid_httpfile(): + return http_fakefile + + +class TestDataSourceOpen: + def setup_method(self): + self.tmpdir = mkdtemp() + self.ds = datasource.DataSource(self.tmpdir) + + def teardown_method(self): + rmtree(self.tmpdir) + del self.ds + + def test_ValidHTTP(self): + fh = self.ds.open(valid_httpurl()) + assert_(fh) + fh.close() + + def test_InvalidHTTP(self): + url = invalid_httpurl() + assert_raises(OSError, self.ds.open, url) + try: + self.ds.open(url) + except OSError as e: + # Regression test for bug fixed in r4342. + assert_(e.errno is None) + + def test_InvalidHTTPCacheURLError(self): + assert_raises(URLError, self.ds._cache, invalid_httpurl()) + + def test_ValidFile(self): + local_file = valid_textfile(self.tmpdir) + fh = self.ds.open(local_file) + assert_(fh) + fh.close() + + def test_InvalidFile(self): + invalid_file = invalid_textfile(self.tmpdir) + assert_raises(OSError, self.ds.open, invalid_file) + + def test_ValidGzipFile(self): + try: + import gzip + except ImportError: + # We don't have the gzip capabilities to test. + pytest.skip() + # Test datasource's internal file_opener for Gzip files. + filepath = os.path.join(self.tmpdir, 'foobar.txt.gz') + fp = gzip.open(filepath, 'w') + fp.write(magic_line) + fp.close() + fp = self.ds.open(filepath) + result = fp.readline() + fp.close() + assert_equal(magic_line, result) + + def test_ValidBz2File(self): + try: + import bz2 + except ImportError: + # We don't have the bz2 capabilities to test. + pytest.skip() + # Test datasource's internal file_opener for BZip2 files. + filepath = os.path.join(self.tmpdir, 'foobar.txt.bz2') + fp = bz2.BZ2File(filepath, 'w') + fp.write(magic_line) + fp.close() + fp = self.ds.open(filepath) + result = fp.readline() + fp.close() + assert_equal(magic_line, result) + + +class TestDataSourceExists: + def setup_method(self): + self.tmpdir = mkdtemp() + self.ds = datasource.DataSource(self.tmpdir) + + def teardown_method(self): + rmtree(self.tmpdir) + del self.ds + + def test_ValidHTTP(self): + assert_(self.ds.exists(valid_httpurl())) + + def test_InvalidHTTP(self): + assert_equal(self.ds.exists(invalid_httpurl()), False) + + def test_ValidFile(self): + # Test valid file in destpath + tmpfile = valid_textfile(self.tmpdir) + assert_(self.ds.exists(tmpfile)) + # Test valid local file not in destpath + localdir = mkdtemp() + tmpfile = valid_textfile(localdir) + assert_(self.ds.exists(tmpfile)) + rmtree(localdir) + + def test_InvalidFile(self): + tmpfile = invalid_textfile(self.tmpdir) + assert_equal(self.ds.exists(tmpfile), False) + + +class TestDataSourceAbspath: + def setup_method(self): + self.tmpdir = os.path.abspath(mkdtemp()) + self.ds = datasource.DataSource(self.tmpdir) + + def teardown_method(self): + rmtree(self.tmpdir) + del self.ds + + def test_ValidHTTP(self): + scheme, netloc, upath, pms, qry, frg = urlparse(valid_httpurl()) + local_path = os.path.join(self.tmpdir, netloc, + upath.strip(os.sep).strip('/')) + assert_equal(local_path, self.ds.abspath(valid_httpurl())) + + def test_ValidFile(self): + tmpfile = valid_textfile(self.tmpdir) + tmpfilename = os.path.split(tmpfile)[-1] + # Test with filename only + assert_equal(tmpfile, self.ds.abspath(tmpfilename)) + # Test filename with complete path + assert_equal(tmpfile, self.ds.abspath(tmpfile)) + + def test_InvalidHTTP(self): + scheme, netloc, upath, pms, qry, frg = urlparse(invalid_httpurl()) + invalidhttp = os.path.join(self.tmpdir, netloc, + upath.strip(os.sep).strip('/')) + assert_(invalidhttp != self.ds.abspath(valid_httpurl())) + + def test_InvalidFile(self): + invalidfile = valid_textfile(self.tmpdir) + tmpfile = valid_textfile(self.tmpdir) + tmpfilename = os.path.split(tmpfile)[-1] + # Test with filename only + assert_(invalidfile != self.ds.abspath(tmpfilename)) + # Test filename with complete path + assert_(invalidfile != self.ds.abspath(tmpfile)) + + def test_sandboxing(self): + tmpfile = valid_textfile(self.tmpdir) + tmpfilename = os.path.split(tmpfile)[-1] + + tmp_path = lambda x: os.path.abspath(self.ds.abspath(x)) + + assert_(tmp_path(valid_httpurl()).startswith(self.tmpdir)) + assert_(tmp_path(invalid_httpurl()).startswith(self.tmpdir)) + assert_(tmp_path(tmpfile).startswith(self.tmpdir)) + assert_(tmp_path(tmpfilename).startswith(self.tmpdir)) + for fn in malicious_files: + assert_(tmp_path(http_path+fn).startswith(self.tmpdir)) + assert_(tmp_path(fn).startswith(self.tmpdir)) + + def test_windows_os_sep(self): + orig_os_sep = os.sep + try: + os.sep = '\\' + self.test_ValidHTTP() + self.test_ValidFile() + self.test_InvalidHTTP() + self.test_InvalidFile() + self.test_sandboxing() + finally: + os.sep = orig_os_sep + + +class TestRepositoryAbspath: + def setup_method(self): + self.tmpdir = os.path.abspath(mkdtemp()) + self.repos = datasource.Repository(valid_baseurl(), self.tmpdir) + + def teardown_method(self): + rmtree(self.tmpdir) + del self.repos + + def test_ValidHTTP(self): + scheme, netloc, upath, pms, qry, frg = urlparse(valid_httpurl()) + local_path = os.path.join(self.repos._destpath, netloc, + upath.strip(os.sep).strip('/')) + filepath = self.repos.abspath(valid_httpfile()) + assert_equal(local_path, filepath) + + def test_sandboxing(self): + tmp_path = lambda x: os.path.abspath(self.repos.abspath(x)) + assert_(tmp_path(valid_httpfile()).startswith(self.tmpdir)) + for fn in malicious_files: + assert_(tmp_path(http_path+fn).startswith(self.tmpdir)) + assert_(tmp_path(fn).startswith(self.tmpdir)) + + def test_windows_os_sep(self): + orig_os_sep = os.sep + try: + os.sep = '\\' + self.test_ValidHTTP() + self.test_sandboxing() + finally: + os.sep = orig_os_sep + + +class TestRepositoryExists: + def setup_method(self): + self.tmpdir = mkdtemp() + self.repos = datasource.Repository(valid_baseurl(), self.tmpdir) + + def teardown_method(self): + rmtree(self.tmpdir) + del self.repos + + def test_ValidFile(self): + # Create local temp file + tmpfile = valid_textfile(self.tmpdir) + assert_(self.repos.exists(tmpfile)) + + def test_InvalidFile(self): + tmpfile = invalid_textfile(self.tmpdir) + assert_equal(self.repos.exists(tmpfile), False) + + def test_RemoveHTTPFile(self): + assert_(self.repos.exists(valid_httpurl())) + + def test_CachedHTTPFile(self): + localfile = valid_httpurl() + # Create a locally cached temp file with an URL based + # directory structure. This is similar to what Repository.open + # would do. + scheme, netloc, upath, pms, qry, frg = urlparse(localfile) + local_path = os.path.join(self.repos._destpath, netloc) + os.mkdir(local_path, 0o0700) + tmpfile = valid_textfile(local_path) + assert_(self.repos.exists(tmpfile)) + + +class TestOpenFunc: + def setup_method(self): + self.tmpdir = mkdtemp() + + def teardown_method(self): + rmtree(self.tmpdir) + + def test_DataSourceOpen(self): + local_file = valid_textfile(self.tmpdir) + # Test case where destpath is passed in + fp = datasource.open(local_file, destpath=self.tmpdir) + assert_(fp) + fp.close() + # Test case where default destpath is used + fp = datasource.open(local_file) + assert_(fp) + fp.close() + +def test_del_attr_handling(): + # DataSource __del__ can be called + # even if __init__ fails when the + # Exception object is caught by the + # caller as happens in refguide_check + # is_deprecated() function + + ds = datasource.DataSource() + # simulate failed __init__ by removing key attribute + # produced within __init__ and expected by __del__ + del ds._istmpdest + # should not raise an AttributeError if __del__ + # gracefully handles failed __init__: + ds.__del__() diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test__iotools.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__iotools.py new file mode 100644 index 0000000000000000000000000000000000000000..396d4147c6c55bf169cd66e3cb5472e59d5abdb9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__iotools.py @@ -0,0 +1,353 @@ +import time +from datetime import date + +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_allclose, assert_raises, + ) +from numpy.lib._iotools import ( + LineSplitter, NameValidator, StringConverter, + has_nested_fields, easy_dtype, flatten_dtype + ) + + +class TestLineSplitter: + "Tests the LineSplitter class." + + def test_no_delimiter(self): + "Test LineSplitter w/o delimiter" + strg = " 1 2 3 4 5 # test" + test = LineSplitter()(strg) + assert_equal(test, ['1', '2', '3', '4', '5']) + test = LineSplitter('')(strg) + assert_equal(test, ['1', '2', '3', '4', '5']) + + def test_space_delimiter(self): + "Test space delimiter" + strg = " 1 2 3 4 5 # test" + test = LineSplitter(' ')(strg) + assert_equal(test, ['1', '2', '3', '4', '', '5']) + test = LineSplitter(' ')(strg) + assert_equal(test, ['1 2 3 4', '5']) + + def test_tab_delimiter(self): + "Test tab delimiter" + strg = " 1\t 2\t 3\t 4\t 5 6" + test = LineSplitter('\t')(strg) + assert_equal(test, ['1', '2', '3', '4', '5 6']) + strg = " 1 2\t 3 4\t 5 6" + test = LineSplitter('\t')(strg) + assert_equal(test, ['1 2', '3 4', '5 6']) + + def test_other_delimiter(self): + "Test LineSplitter on delimiter" + strg = "1,2,3,4,,5" + test = LineSplitter(',')(strg) + assert_equal(test, ['1', '2', '3', '4', '', '5']) + # + strg = " 1,2,3,4,,5 # test" + test = LineSplitter(',')(strg) + assert_equal(test, ['1', '2', '3', '4', '', '5']) + + # gh-11028 bytes comment/delimiters should get encoded + strg = b" 1,2,3,4,,5 % test" + test = LineSplitter(delimiter=b',', comments=b'%')(strg) + assert_equal(test, ['1', '2', '3', '4', '', '5']) + + def test_constant_fixed_width(self): + "Test LineSplitter w/ fixed-width fields" + strg = " 1 2 3 4 5 # test" + test = LineSplitter(3)(strg) + assert_equal(test, ['1', '2', '3', '4', '', '5', '']) + # + strg = " 1 3 4 5 6# test" + test = LineSplitter(20)(strg) + assert_equal(test, ['1 3 4 5 6']) + # + strg = " 1 3 4 5 6# test" + test = LineSplitter(30)(strg) + assert_equal(test, ['1 3 4 5 6']) + + def test_variable_fixed_width(self): + strg = " 1 3 4 5 6# test" + test = LineSplitter((3, 6, 6, 3))(strg) + assert_equal(test, ['1', '3', '4 5', '6']) + # + strg = " 1 3 4 5 6# test" + test = LineSplitter((6, 6, 9))(strg) + assert_equal(test, ['1', '3 4', '5 6']) + +# ----------------------------------------------------------------------------- + + +class TestNameValidator: + + def test_case_sensitivity(self): + "Test case sensitivity" + names = ['A', 'a', 'b', 'c'] + test = NameValidator().validate(names) + assert_equal(test, ['A', 'a', 'b', 'c']) + test = NameValidator(case_sensitive=False).validate(names) + assert_equal(test, ['A', 'A_1', 'B', 'C']) + test = NameValidator(case_sensitive='upper').validate(names) + assert_equal(test, ['A', 'A_1', 'B', 'C']) + test = NameValidator(case_sensitive='lower').validate(names) + assert_equal(test, ['a', 'a_1', 'b', 'c']) + + # check exceptions + assert_raises(ValueError, NameValidator, case_sensitive='foobar') + + def test_excludelist(self): + "Test excludelist" + names = ['dates', 'data', 'Other Data', 'mask'] + validator = NameValidator(excludelist=['dates', 'data', 'mask']) + test = validator.validate(names) + assert_equal(test, ['dates_', 'data_', 'Other_Data', 'mask_']) + + def test_missing_names(self): + "Test validate missing names" + namelist = ('a', 'b', 'c') + validator = NameValidator() + assert_equal(validator(namelist), ['a', 'b', 'c']) + namelist = ('', 'b', 'c') + assert_equal(validator(namelist), ['f0', 'b', 'c']) + namelist = ('a', 'b', '') + assert_equal(validator(namelist), ['a', 'b', 'f0']) + namelist = ('', 'f0', '') + assert_equal(validator(namelist), ['f1', 'f0', 'f2']) + + def test_validate_nb_names(self): + "Test validate nb names" + namelist = ('a', 'b', 'c') + validator = NameValidator() + assert_equal(validator(namelist, nbfields=1), ('a',)) + assert_equal(validator(namelist, nbfields=5, defaultfmt="g%i"), + ['a', 'b', 'c', 'g0', 'g1']) + + def test_validate_wo_names(self): + "Test validate no names" + namelist = None + validator = NameValidator() + assert_(validator(namelist) is None) + assert_equal(validator(namelist, nbfields=3), ['f0', 'f1', 'f2']) + +# ----------------------------------------------------------------------------- + + +def _bytes_to_date(s): + return date(*time.strptime(s, "%Y-%m-%d")[:3]) + + +class TestStringConverter: + "Test StringConverter" + + def test_creation(self): + "Test creation of a StringConverter" + converter = StringConverter(int, -99999) + assert_equal(converter._status, 1) + assert_equal(converter.default, -99999) + + def test_upgrade(self): + "Tests the upgrade method." + + converter = StringConverter() + assert_equal(converter._status, 0) + + # test int + assert_equal(converter.upgrade('0'), 0) + assert_equal(converter._status, 1) + + # On systems where long defaults to 32-bit, the statuses will be + # offset by one, so we check for this here. + import numpy._core.numeric as nx + status_offset = int(nx.dtype(nx.int_).itemsize < nx.dtype(nx.int64).itemsize) + + # test int > 2**32 + assert_equal(converter.upgrade('17179869184'), 17179869184) + assert_equal(converter._status, 1 + status_offset) + + # test float + assert_allclose(converter.upgrade('0.'), 0.0) + assert_equal(converter._status, 2 + status_offset) + + # test complex + assert_equal(converter.upgrade('0j'), complex('0j')) + assert_equal(converter._status, 3 + status_offset) + + # test str + # note that the longdouble type has been skipped, so the + # _status increases by 2. Everything should succeed with + # unicode conversion (8). + for s in ['a', b'a']: + res = converter.upgrade(s) + assert_(type(res) is str) + assert_equal(res, 'a') + assert_equal(converter._status, 8 + status_offset) + + def test_missing(self): + "Tests the use of missing values." + converter = StringConverter(missing_values=('missing', + 'missed')) + converter.upgrade('0') + assert_equal(converter('0'), 0) + assert_equal(converter(''), converter.default) + assert_equal(converter('missing'), converter.default) + assert_equal(converter('missed'), converter.default) + try: + converter('miss') + except ValueError: + pass + + def test_upgrademapper(self): + "Tests updatemapper" + dateparser = _bytes_to_date + _original_mapper = StringConverter._mapper[:] + try: + StringConverter.upgrade_mapper(dateparser, date(2000, 1, 1)) + convert = StringConverter(dateparser, date(2000, 1, 1)) + test = convert('2001-01-01') + assert_equal(test, date(2001, 1, 1)) + test = convert('2009-01-01') + assert_equal(test, date(2009, 1, 1)) + test = convert('') + assert_equal(test, date(2000, 1, 1)) + finally: + StringConverter._mapper = _original_mapper + + def test_string_to_object(self): + "Make sure that string-to-object functions are properly recognized" + old_mapper = StringConverter._mapper[:] # copy of list + conv = StringConverter(_bytes_to_date) + assert_equal(conv._mapper, old_mapper) + assert_(hasattr(conv, 'default')) + + def test_keep_default(self): + "Make sure we don't lose an explicit default" + converter = StringConverter(None, missing_values='', + default=-999) + converter.upgrade('3.14159265') + assert_equal(converter.default, -999) + assert_equal(converter.type, np.dtype(float)) + # + converter = StringConverter( + None, missing_values='', default=0) + converter.upgrade('3.14159265') + assert_equal(converter.default, 0) + assert_equal(converter.type, np.dtype(float)) + + def test_keep_default_zero(self): + "Check that we don't lose a default of 0" + converter = StringConverter(int, default=0, + missing_values="N/A") + assert_equal(converter.default, 0) + + def test_keep_missing_values(self): + "Check that we're not losing missing values" + converter = StringConverter(int, default=0, + missing_values="N/A") + assert_equal( + converter.missing_values, {'', 'N/A'}) + + def test_int64_dtype(self): + "Check that int64 integer types can be specified" + converter = StringConverter(np.int64, default=0) + val = "-9223372036854775807" + assert_(converter(val) == -9223372036854775807) + val = "9223372036854775807" + assert_(converter(val) == 9223372036854775807) + + def test_uint64_dtype(self): + "Check that uint64 integer types can be specified" + converter = StringConverter(np.uint64, default=0) + val = "9223372043271415339" + assert_(converter(val) == 9223372043271415339) + + +class TestMiscFunctions: + + def test_has_nested_dtype(self): + "Test has_nested_dtype" + ndtype = np.dtype(float) + assert_equal(has_nested_fields(ndtype), False) + ndtype = np.dtype([('A', '|S3'), ('B', float)]) + assert_equal(has_nested_fields(ndtype), False) + ndtype = np.dtype([('A', int), ('B', [('BA', float), ('BB', '|S1')])]) + assert_equal(has_nested_fields(ndtype), True) + + def test_easy_dtype(self): + "Test ndtype on dtypes" + # Simple case + ndtype = float + assert_equal(easy_dtype(ndtype), np.dtype(float)) + # As string w/o names + ndtype = "i4, f8" + assert_equal(easy_dtype(ndtype), + np.dtype([('f0', "i4"), ('f1', "f8")])) + # As string w/o names but different default format + assert_equal(easy_dtype(ndtype, defaultfmt="field_%03i"), + np.dtype([('field_000', "i4"), ('field_001', "f8")])) + # As string w/ names + ndtype = "i4, f8" + assert_equal(easy_dtype(ndtype, names="a, b"), + np.dtype([('a', "i4"), ('b', "f8")])) + # As string w/ names (too many) + ndtype = "i4, f8" + assert_equal(easy_dtype(ndtype, names="a, b, c"), + np.dtype([('a', "i4"), ('b', "f8")])) + # As string w/ names (not enough) + ndtype = "i4, f8" + assert_equal(easy_dtype(ndtype, names=", b"), + np.dtype([('f0', "i4"), ('b', "f8")])) + # ... (with different default format) + assert_equal(easy_dtype(ndtype, names="a", defaultfmt="f%02i"), + np.dtype([('a', "i4"), ('f00', "f8")])) + # As list of tuples w/o names + ndtype = [('A', int), ('B', float)] + assert_equal(easy_dtype(ndtype), np.dtype([('A', int), ('B', float)])) + # As list of tuples w/ names + assert_equal(easy_dtype(ndtype, names="a,b"), + np.dtype([('a', int), ('b', float)])) + # As list of tuples w/ not enough names + assert_equal(easy_dtype(ndtype, names="a"), + np.dtype([('a', int), ('f0', float)])) + # As list of tuples w/ too many names + assert_equal(easy_dtype(ndtype, names="a,b,c"), + np.dtype([('a', int), ('b', float)])) + # As list of types w/o names + ndtype = (int, float, float) + assert_equal(easy_dtype(ndtype), + np.dtype([('f0', int), ('f1', float), ('f2', float)])) + # As list of types w names + ndtype = (int, float, float) + assert_equal(easy_dtype(ndtype, names="a, b, c"), + np.dtype([('a', int), ('b', float), ('c', float)])) + # As simple dtype w/ names + ndtype = np.dtype(float) + assert_equal(easy_dtype(ndtype, names="a, b, c"), + np.dtype([(_, float) for _ in ('a', 'b', 'c')])) + # As simple dtype w/o names (but multiple fields) + ndtype = np.dtype(float) + assert_equal( + easy_dtype(ndtype, names=['', '', ''], defaultfmt="f%02i"), + np.dtype([(_, float) for _ in ('f00', 'f01', 'f02')])) + + def test_flatten_dtype(self): + "Testing flatten_dtype" + # Standard dtype + dt = np.dtype([("a", "f8"), ("b", "f8")]) + dt_flat = flatten_dtype(dt) + assert_equal(dt_flat, [float, float]) + # Recursive dtype + dt = np.dtype([("a", [("aa", '|S1'), ("ab", '|S2')]), ("b", int)]) + dt_flat = flatten_dtype(dt) + assert_equal(dt_flat, [np.dtype('|S1'), np.dtype('|S2'), int]) + # dtype with shaped fields + dt = np.dtype([("a", (float, 2)), ("b", (int, 3))]) + dt_flat = flatten_dtype(dt) + assert_equal(dt_flat, [float, int]) + dt_flat = flatten_dtype(dt, True) + assert_equal(dt_flat, [float] * 2 + [int] * 3) + # dtype w/ titles + dt = np.dtype([(("a", "A"), "f8"), (("b", "B"), "f8")]) + dt_flat = flatten_dtype(dt) + assert_equal(dt_flat, [float, float]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test__version.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__version.py new file mode 100644 index 0000000000000000000000000000000000000000..e6d41ad939323792d31faa7ae517e6835ea851d1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test__version.py @@ -0,0 +1,64 @@ +"""Tests for the NumpyVersion class. + +""" +from numpy.testing import assert_, assert_raises +from numpy.lib import NumpyVersion + + +def test_main_versions(): + assert_(NumpyVersion('1.8.0') == '1.8.0') + for ver in ['1.9.0', '2.0.0', '1.8.1', '10.0.1']: + assert_(NumpyVersion('1.8.0') < ver) + + for ver in ['1.7.0', '1.7.1', '0.9.9']: + assert_(NumpyVersion('1.8.0') > ver) + + +def test_version_1_point_10(): + # regression test for gh-2998. + assert_(NumpyVersion('1.9.0') < '1.10.0') + assert_(NumpyVersion('1.11.0') < '1.11.1') + assert_(NumpyVersion('1.11.0') == '1.11.0') + assert_(NumpyVersion('1.99.11') < '1.99.12') + + +def test_alpha_beta_rc(): + assert_(NumpyVersion('1.8.0rc1') == '1.8.0rc1') + for ver in ['1.8.0', '1.8.0rc2']: + assert_(NumpyVersion('1.8.0rc1') < ver) + + for ver in ['1.8.0a2', '1.8.0b3', '1.7.2rc4']: + assert_(NumpyVersion('1.8.0rc1') > ver) + + assert_(NumpyVersion('1.8.0b1') > '1.8.0a2') + + +def test_dev_version(): + assert_(NumpyVersion('1.9.0.dev-Unknown') < '1.9.0') + for ver in ['1.9.0', '1.9.0a1', '1.9.0b2', '1.9.0b2.dev-ffffffff']: + assert_(NumpyVersion('1.9.0.dev-f16acvda') < ver) + + assert_(NumpyVersion('1.9.0.dev-f16acvda') == '1.9.0.dev-11111111') + + +def test_dev_a_b_rc_mixed(): + assert_(NumpyVersion('1.9.0a2.dev-f16acvda') == '1.9.0a2.dev-11111111') + assert_(NumpyVersion('1.9.0a2.dev-6acvda54') < '1.9.0a2') + + +def test_dev0_version(): + assert_(NumpyVersion('1.9.0.dev0+Unknown') < '1.9.0') + for ver in ['1.9.0', '1.9.0a1', '1.9.0b2', '1.9.0b2.dev0+ffffffff']: + assert_(NumpyVersion('1.9.0.dev0+f16acvda') < ver) + + assert_(NumpyVersion('1.9.0.dev0+f16acvda') == '1.9.0.dev0+11111111') + + +def test_dev0_a_b_rc_mixed(): + assert_(NumpyVersion('1.9.0a2.dev0+f16acvda') == '1.9.0a2.dev0+11111111') + assert_(NumpyVersion('1.9.0a2.dev0+6acvda54') < '1.9.0a2') + + +def test_raises(): + for ver in ['1.9', '1,9.0', '1.7.x']: + assert_raises(ValueError, NumpyVersion, ver) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_array_utils.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_array_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3d8b2bd4616e0a6d4956bb6c2b2396a33ee8836a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_array_utils.py @@ -0,0 +1,33 @@ +import numpy as np + +from numpy.lib import array_utils +from numpy.testing import assert_equal + + +class TestByteBounds: + def test_byte_bounds(self): + # pointer difference matches size * itemsize + # due to contiguity + a = np.arange(12).reshape(3, 4) + low, high = array_utils.byte_bounds(a) + assert_equal(high - low, a.size * a.itemsize) + + def test_unusual_order_positive_stride(self): + a = np.arange(12).reshape(3, 4) + b = a.T + low, high = array_utils.byte_bounds(b) + assert_equal(high - low, b.size * b.itemsize) + + def test_unusual_order_negative_stride(self): + a = np.arange(12).reshape(3, 4) + b = a.T[::-1] + low, high = array_utils.byte_bounds(b) + assert_equal(high - low, b.size * b.itemsize) + + def test_strided(self): + a = np.arange(12) + b = a[::2] + low, high = array_utils.byte_bounds(b) + # the largest pointer address is lost (even numbers only in the + # stride), and compensate addresses for striding by 2 + assert_equal(high - low, b.size * 2 * b.itemsize - b.itemsize) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraypad.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraypad.py new file mode 100644 index 0000000000000000000000000000000000000000..6c1247db8e0c8765a4ddeeee292ec326daeedda2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraypad.py @@ -0,0 +1,1416 @@ +"""Tests for the array padding functions. + +""" +import pytest + +import numpy as np +from numpy.testing import assert_array_equal, assert_allclose, assert_equal +from numpy.lib._arraypad_impl import _as_pairs + + +_numeric_dtypes = ( + np._core.sctypes["uint"] + + np._core.sctypes["int"] + + np._core.sctypes["float"] + + np._core.sctypes["complex"] +) +_all_modes = { + 'constant': {'constant_values': 0}, + 'edge': {}, + 'linear_ramp': {'end_values': 0}, + 'maximum': {'stat_length': None}, + 'mean': {'stat_length': None}, + 'median': {'stat_length': None}, + 'minimum': {'stat_length': None}, + 'reflect': {'reflect_type': 'even'}, + 'symmetric': {'reflect_type': 'even'}, + 'wrap': {}, + 'empty': {} +} + + +class TestAsPairs: + def test_single_value(self): + """Test casting for a single value.""" + expected = np.array([[3, 3]] * 10) + for x in (3, [3], [[3]]): + result = _as_pairs(x, 10) + assert_equal(result, expected) + # Test with dtype=object + obj = object() + assert_equal( + _as_pairs(obj, 10), + np.array([[obj, obj]] * 10) + ) + + def test_two_values(self): + """Test proper casting for two different values.""" + # Broadcasting in the first dimension with numbers + expected = np.array([[3, 4]] * 10) + for x in ([3, 4], [[3, 4]]): + result = _as_pairs(x, 10) + assert_equal(result, expected) + # and with dtype=object + obj = object() + assert_equal( + _as_pairs(["a", obj], 10), + np.array([["a", obj]] * 10) + ) + + # Broadcasting in the second / last dimension with numbers + assert_equal( + _as_pairs([[3], [4]], 2), + np.array([[3, 3], [4, 4]]) + ) + # and with dtype=object + assert_equal( + _as_pairs([["a"], [obj]], 2), + np.array([["a", "a"], [obj, obj]]) + ) + + def test_with_none(self): + expected = ((None, None), (None, None), (None, None)) + assert_equal( + _as_pairs(None, 3, as_index=False), + expected + ) + assert_equal( + _as_pairs(None, 3, as_index=True), + expected + ) + + def test_pass_through(self): + """Test if `x` already matching desired output are passed through.""" + expected = np.arange(12).reshape((6, 2)) + assert_equal( + _as_pairs(expected, 6), + expected + ) + + def test_as_index(self): + """Test results if `as_index=True`.""" + assert_equal( + _as_pairs([2.6, 3.3], 10, as_index=True), + np.array([[3, 3]] * 10, dtype=np.intp) + ) + assert_equal( + _as_pairs([2.6, 4.49], 10, as_index=True), + np.array([[3, 4]] * 10, dtype=np.intp) + ) + for x in (-3, [-3], [[-3]], [-3, 4], [3, -4], [[-3, 4]], [[4, -3]], + [[1, 2]] * 9 + [[1, -2]]): + with pytest.raises(ValueError, match="negative values"): + _as_pairs(x, 10, as_index=True) + + def test_exceptions(self): + """Ensure faulty usage is discovered.""" + with pytest.raises(ValueError, match="more dimensions than allowed"): + _as_pairs([[[3]]], 10) + with pytest.raises(ValueError, match="could not be broadcast"): + _as_pairs([[1, 2], [3, 4]], 3) + with pytest.raises(ValueError, match="could not be broadcast"): + _as_pairs(np.ones((2, 3)), 3) + + +class TestConditionalShortcuts: + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_zero_padding_shortcuts(self, mode): + test = np.arange(120).reshape(4, 5, 6) + pad_amt = [(0, 0) for _ in test.shape] + assert_array_equal(test, np.pad(test, pad_amt, mode=mode)) + + @pytest.mark.parametrize("mode", ['maximum', 'mean', 'median', 'minimum',]) + def test_shallow_statistic_range(self, mode): + test = np.arange(120).reshape(4, 5, 6) + pad_amt = [(1, 1) for _ in test.shape] + assert_array_equal(np.pad(test, pad_amt, mode='edge'), + np.pad(test, pad_amt, mode=mode, stat_length=1)) + + @pytest.mark.parametrize("mode", ['maximum', 'mean', 'median', 'minimum',]) + def test_clip_statistic_range(self, mode): + test = np.arange(30).reshape(5, 6) + pad_amt = [(3, 3) for _ in test.shape] + assert_array_equal(np.pad(test, pad_amt, mode=mode), + np.pad(test, pad_amt, mode=mode, stat_length=30)) + + +class TestStatistic: + def test_check_mean_stat_length(self): + a = np.arange(100).astype('f') + a = np.pad(a, ((25, 20), ), 'mean', stat_length=((2, 3), )) + b = np.array( + [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, + 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, + 0.5, 0.5, 0.5, 0.5, 0.5, + + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., + 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., + 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., + 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., + 40., 41., 42., 43., 44., 45., 46., 47., 48., 49., + 50., 51., 52., 53., 54., 55., 56., 57., 58., 59., + 60., 61., 62., 63., 64., 65., 66., 67., 68., 69., + 70., 71., 72., 73., 74., 75., 76., 77., 78., 79., + 80., 81., 82., 83., 84., 85., 86., 87., 88., 89., + 90., 91., 92., 93., 94., 95., 96., 97., 98., 99., + + 98., 98., 98., 98., 98., 98., 98., 98., 98., 98., + 98., 98., 98., 98., 98., 98., 98., 98., 98., 98. + ]) + assert_array_equal(a, b) + + def test_check_maximum_1(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'maximum') + b = np.array( + [99, 99, 99, 99, 99, 99, 99, 99, 99, 99, + 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, + 99, 99, 99, 99, 99, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, + 99, 99, 99, 99, 99, 99, 99, 99, 99, 99] + ) + assert_array_equal(a, b) + + def test_check_maximum_2(self): + a = np.arange(100) + 1 + a = np.pad(a, (25, 20), 'maximum') + b = np.array( + [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, + 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, + 100, 100, 100, 100, 100, + + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, + 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, + 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, + 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, + 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, + 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, + 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, + 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, + 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, + + 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, + 100, 100, 100, 100, 100, 100, 100, 100, 100, 100] + ) + assert_array_equal(a, b) + + def test_check_maximum_stat_length(self): + a = np.arange(100) + 1 + a = np.pad(a, (25, 20), 'maximum', stat_length=10) + b = np.array( + [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, + 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, + 10, 10, 10, 10, 10, + + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, + 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, + 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, + 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, + 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, + 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, + 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, + 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, + 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, + + 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, + 100, 100, 100, 100, 100, 100, 100, 100, 100, 100] + ) + assert_array_equal(a, b) + + def test_check_minimum_1(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'minimum') + b = np.array( + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + ) + assert_array_equal(a, b) + + def test_check_minimum_2(self): + a = np.arange(100) + 2 + a = np.pad(a, (25, 20), 'minimum') + b = np.array( + [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, + + 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, + 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, + 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, + 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, + 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, + 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, + 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, + 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, + + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2] + ) + assert_array_equal(a, b) + + def test_check_minimum_stat_length(self): + a = np.arange(100) + 1 + a = np.pad(a, (25, 20), 'minimum', stat_length=10) + b = np.array( + [ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, + + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, + 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, + 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, + 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, + 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, + 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, + 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, + 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, + 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, + 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, + + 91, 91, 91, 91, 91, 91, 91, 91, 91, 91, + 91, 91, 91, 91, 91, 91, 91, 91, 91, 91] + ) + assert_array_equal(a, b) + + def test_check_median(self): + a = np.arange(100).astype('f') + a = np.pad(a, (25, 20), 'median') + b = np.array( + [49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, + + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., + 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., + 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., + 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., + 40., 41., 42., 43., 44., 45., 46., 47., 48., 49., + 50., 51., 52., 53., 54., 55., 56., 57., 58., 59., + 60., 61., 62., 63., 64., 65., 66., 67., 68., 69., + 70., 71., 72., 73., 74., 75., 76., 77., 78., 79., + 80., 81., 82., 83., 84., 85., 86., 87., 88., 89., + 90., 91., 92., 93., 94., 95., 96., 97., 98., 99., + + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5] + ) + assert_array_equal(a, b) + + def test_check_median_01(self): + a = np.array([[3, 1, 4], [4, 5, 9], [9, 8, 2]]) + a = np.pad(a, 1, 'median') + b = np.array( + [[4, 4, 5, 4, 4], + + [3, 3, 1, 4, 3], + [5, 4, 5, 9, 5], + [8, 9, 8, 2, 8], + + [4, 4, 5, 4, 4]] + ) + assert_array_equal(a, b) + + def test_check_median_02(self): + a = np.array([[3, 1, 4], [4, 5, 9], [9, 8, 2]]) + a = np.pad(a.T, 1, 'median').T + b = np.array( + [[5, 4, 5, 4, 5], + + [3, 3, 1, 4, 3], + [5, 4, 5, 9, 5], + [8, 9, 8, 2, 8], + + [5, 4, 5, 4, 5]] + ) + assert_array_equal(a, b) + + def test_check_median_stat_length(self): + a = np.arange(100).astype('f') + a[1] = 2. + a[97] = 96. + a = np.pad(a, (25, 20), 'median', stat_length=(3, 5)) + b = np.array( + [ 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 2., + + 0., 2., 2., 3., 4., 5., 6., 7., 8., 9., + 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., + 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., + 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., + 40., 41., 42., 43., 44., 45., 46., 47., 48., 49., + 50., 51., 52., 53., 54., 55., 56., 57., 58., 59., + 60., 61., 62., 63., 64., 65., 66., 67., 68., 69., + 70., 71., 72., 73., 74., 75., 76., 77., 78., 79., + 80., 81., 82., 83., 84., 85., 86., 87., 88., 89., + 90., 91., 92., 93., 94., 95., 96., 96., 98., 99., + + 96., 96., 96., 96., 96., 96., 96., 96., 96., 96., + 96., 96., 96., 96., 96., 96., 96., 96., 96., 96.] + ) + assert_array_equal(a, b) + + def test_check_mean_shape_one(self): + a = [[4, 5, 6]] + a = np.pad(a, (5, 7), 'mean', stat_length=2) + b = np.array( + [[4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6], + [4, 4, 4, 4, 4, 4, 5, 6, 6, 6, 6, 6, 6, 6, 6]] + ) + assert_array_equal(a, b) + + def test_check_mean_2(self): + a = np.arange(100).astype('f') + a = np.pad(a, (25, 20), 'mean') + b = np.array( + [49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, + + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., + 10., 11., 12., 13., 14., 15., 16., 17., 18., 19., + 20., 21., 22., 23., 24., 25., 26., 27., 28., 29., + 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., + 40., 41., 42., 43., 44., 45., 46., 47., 48., 49., + 50., 51., 52., 53., 54., 55., 56., 57., 58., 59., + 60., 61., 62., 63., 64., 65., 66., 67., 68., 69., + 70., 71., 72., 73., 74., 75., 76., 77., 78., 79., + 80., 81., 82., 83., 84., 85., 86., 87., 88., 89., + 90., 91., 92., 93., 94., 95., 96., 97., 98., 99., + + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, + 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5, 49.5] + ) + assert_array_equal(a, b) + + @pytest.mark.parametrize("mode", [ + "mean", + "median", + "minimum", + "maximum" + ]) + def test_same_prepend_append(self, mode): + """ Test that appended and prepended values are equal """ + # This test is constructed to trigger floating point rounding errors in + # a way that caused gh-11216 for mode=='mean' + a = np.array([-1, 2, -1]) + np.array([0, 1e-12, 0], dtype=np.float64) + a = np.pad(a, (1, 1), mode) + assert_equal(a[0], a[-1]) + + @pytest.mark.parametrize("mode", ["mean", "median", "minimum", "maximum"]) + @pytest.mark.parametrize( + "stat_length", [-2, (-2,), (3, -1), ((5, 2), (-2, 3)), ((-4,), (2,))] + ) + def test_check_negative_stat_length(self, mode, stat_length): + arr = np.arange(30).reshape((6, 5)) + match = "index can't contain negative values" + with pytest.raises(ValueError, match=match): + np.pad(arr, 2, mode, stat_length=stat_length) + + def test_simple_stat_length(self): + a = np.arange(30) + a = np.reshape(a, (6, 5)) + a = np.pad(a, ((2, 3), (3, 2)), mode='mean', stat_length=(3,)) + b = np.array( + [[6, 6, 6, 5, 6, 7, 8, 9, 8, 8], + [6, 6, 6, 5, 6, 7, 8, 9, 8, 8], + + [1, 1, 1, 0, 1, 2, 3, 4, 3, 3], + [6, 6, 6, 5, 6, 7, 8, 9, 8, 8], + [11, 11, 11, 10, 11, 12, 13, 14, 13, 13], + [16, 16, 16, 15, 16, 17, 18, 19, 18, 18], + [21, 21, 21, 20, 21, 22, 23, 24, 23, 23], + [26, 26, 26, 25, 26, 27, 28, 29, 28, 28], + + [21, 21, 21, 20, 21, 22, 23, 24, 23, 23], + [21, 21, 21, 20, 21, 22, 23, 24, 23, 23], + [21, 21, 21, 20, 21, 22, 23, 24, 23, 23]] + ) + assert_array_equal(a, b) + + @pytest.mark.filterwarnings("ignore:Mean of empty slice:RuntimeWarning") + @pytest.mark.filterwarnings( + "ignore:invalid value encountered in( scalar)? divide:RuntimeWarning" + ) + @pytest.mark.parametrize("mode", ["mean", "median"]) + def test_zero_stat_length_valid(self, mode): + arr = np.pad([1., 2.], (1, 2), mode, stat_length=0) + expected = np.array([np.nan, 1., 2., np.nan, np.nan]) + assert_equal(arr, expected) + + @pytest.mark.parametrize("mode", ["minimum", "maximum"]) + def test_zero_stat_length_invalid(self, mode): + match = "stat_length of 0 yields no value for padding" + with pytest.raises(ValueError, match=match): + np.pad([1., 2.], 0, mode, stat_length=0) + with pytest.raises(ValueError, match=match): + np.pad([1., 2.], 0, mode, stat_length=(1, 0)) + with pytest.raises(ValueError, match=match): + np.pad([1., 2.], 1, mode, stat_length=0) + with pytest.raises(ValueError, match=match): + np.pad([1., 2.], 1, mode, stat_length=(1, 0)) + + +class TestConstant: + def test_check_constant(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'constant', constant_values=(10, 20)) + b = np.array( + [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, + 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, + 10, 10, 10, 10, 10, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, + 20, 20, 20, 20, 20, 20, 20, 20, 20, 20] + ) + assert_array_equal(a, b) + + def test_check_constant_zeros(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'constant') + b = np.array( + [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + ) + assert_array_equal(a, b) + + def test_check_constant_float(self): + # If input array is int, but constant_values are float, the dtype of + # the array to be padded is kept + arr = np.arange(30).reshape(5, 6) + test = np.pad(arr, (1, 2), mode='constant', + constant_values=1.1) + expected = np.array( + [[ 1, 1, 1, 1, 1, 1, 1, 1, 1], + + [ 1, 0, 1, 2, 3, 4, 5, 1, 1], + [ 1, 6, 7, 8, 9, 10, 11, 1, 1], + [ 1, 12, 13, 14, 15, 16, 17, 1, 1], + [ 1, 18, 19, 20, 21, 22, 23, 1, 1], + [ 1, 24, 25, 26, 27, 28, 29, 1, 1], + + [ 1, 1, 1, 1, 1, 1, 1, 1, 1], + [ 1, 1, 1, 1, 1, 1, 1, 1, 1]] + ) + assert_allclose(test, expected) + + def test_check_constant_float2(self): + # If input array is float, and constant_values are float, the dtype of + # the array to be padded is kept - here retaining the float constants + arr = np.arange(30).reshape(5, 6) + arr_float = arr.astype(np.float64) + test = np.pad(arr_float, ((1, 2), (1, 2)), mode='constant', + constant_values=1.1) + expected = np.array( + [[ 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + + [ 1.1, 0. , 1. , 2. , 3. , 4. , 5. , 1.1, 1.1], + [ 1.1, 6. , 7. , 8. , 9. , 10. , 11. , 1.1, 1.1], + [ 1.1, 12. , 13. , 14. , 15. , 16. , 17. , 1.1, 1.1], + [ 1.1, 18. , 19. , 20. , 21. , 22. , 23. , 1.1, 1.1], + [ 1.1, 24. , 25. , 26. , 27. , 28. , 29. , 1.1, 1.1], + + [ 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [ 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1]] + ) + assert_allclose(test, expected) + + def test_check_constant_float3(self): + a = np.arange(100, dtype=float) + a = np.pad(a, (25, 20), 'constant', constant_values=(-1.1, -1.2)) + b = np.array( + [-1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, + -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, -1.1, + -1.1, -1.1, -1.1, -1.1, -1.1, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, + -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2, -1.2] + ) + assert_allclose(a, b) + + def test_check_constant_odd_pad_amount(self): + arr = np.arange(30).reshape(5, 6) + test = np.pad(arr, ((1,), (2,)), mode='constant', + constant_values=3) + expected = np.array( + [[ 3, 3, 3, 3, 3, 3, 3, 3, 3, 3], + + [ 3, 3, 0, 1, 2, 3, 4, 5, 3, 3], + [ 3, 3, 6, 7, 8, 9, 10, 11, 3, 3], + [ 3, 3, 12, 13, 14, 15, 16, 17, 3, 3], + [ 3, 3, 18, 19, 20, 21, 22, 23, 3, 3], + [ 3, 3, 24, 25, 26, 27, 28, 29, 3, 3], + + [ 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]] + ) + assert_allclose(test, expected) + + def test_check_constant_pad_2d(self): + arr = np.arange(4).reshape(2, 2) + test = np.pad(arr, ((1, 2), (1, 3)), mode='constant', + constant_values=((1, 2), (3, 4))) + expected = np.array( + [[3, 1, 1, 4, 4, 4], + [3, 0, 1, 4, 4, 4], + [3, 2, 3, 4, 4, 4], + [3, 2, 2, 4, 4, 4], + [3, 2, 2, 4, 4, 4]] + ) + assert_allclose(test, expected) + + def test_check_large_integers(self): + uint64_max = 2 ** 64 - 1 + arr = np.full(5, uint64_max, dtype=np.uint64) + test = np.pad(arr, 1, mode="constant", constant_values=arr.min()) + expected = np.full(7, uint64_max, dtype=np.uint64) + assert_array_equal(test, expected) + + int64_max = 2 ** 63 - 1 + arr = np.full(5, int64_max, dtype=np.int64) + test = np.pad(arr, 1, mode="constant", constant_values=arr.min()) + expected = np.full(7, int64_max, dtype=np.int64) + assert_array_equal(test, expected) + + def test_check_object_array(self): + arr = np.empty(1, dtype=object) + obj_a = object() + arr[0] = obj_a + obj_b = object() + obj_c = object() + arr = np.pad(arr, pad_width=1, mode='constant', + constant_values=(obj_b, obj_c)) + + expected = np.empty((3,), dtype=object) + expected[0] = obj_b + expected[1] = obj_a + expected[2] = obj_c + + assert_array_equal(arr, expected) + + def test_pad_empty_dimension(self): + arr = np.zeros((3, 0, 2)) + result = np.pad(arr, [(0,), (2,), (1,)], mode="constant") + assert result.shape == (3, 4, 4) + + +class TestLinearRamp: + def test_check_simple(self): + a = np.arange(100).astype('f') + a = np.pad(a, (25, 20), 'linear_ramp', end_values=(4, 5)) + b = np.array( + [4.00, 3.84, 3.68, 3.52, 3.36, 3.20, 3.04, 2.88, 2.72, 2.56, + 2.40, 2.24, 2.08, 1.92, 1.76, 1.60, 1.44, 1.28, 1.12, 0.96, + 0.80, 0.64, 0.48, 0.32, 0.16, + + 0.00, 1.00, 2.00, 3.00, 4.00, 5.00, 6.00, 7.00, 8.00, 9.00, + 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, + 20.0, 21.0, 22.0, 23.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, + 30.0, 31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0, + 40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0, 48.0, 49.0, + 50.0, 51.0, 52.0, 53.0, 54.0, 55.0, 56.0, 57.0, 58.0, 59.0, + 60.0, 61.0, 62.0, 63.0, 64.0, 65.0, 66.0, 67.0, 68.0, 69.0, + 70.0, 71.0, 72.0, 73.0, 74.0, 75.0, 76.0, 77.0, 78.0, 79.0, + 80.0, 81.0, 82.0, 83.0, 84.0, 85.0, 86.0, 87.0, 88.0, 89.0, + 90.0, 91.0, 92.0, 93.0, 94.0, 95.0, 96.0, 97.0, 98.0, 99.0, + + 94.3, 89.6, 84.9, 80.2, 75.5, 70.8, 66.1, 61.4, 56.7, 52.0, + 47.3, 42.6, 37.9, 33.2, 28.5, 23.8, 19.1, 14.4, 9.7, 5.] + ) + assert_allclose(a, b, rtol=1e-5, atol=1e-5) + + def test_check_2d(self): + arr = np.arange(20).reshape(4, 5).astype(np.float64) + test = np.pad(arr, (2, 2), mode='linear_ramp', end_values=(0, 0)) + expected = np.array( + [[0., 0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0.5, 1., 1.5, 2., 1., 0.], + [0., 0., 0., 1., 2., 3., 4., 2., 0.], + [0., 2.5, 5., 6., 7., 8., 9., 4.5, 0.], + [0., 5., 10., 11., 12., 13., 14., 7., 0.], + [0., 7.5, 15., 16., 17., 18., 19., 9.5, 0.], + [0., 3.75, 7.5, 8., 8.5, 9., 9.5, 4.75, 0.], + [0., 0., 0., 0., 0., 0., 0., 0., 0.]]) + assert_allclose(test, expected) + + @pytest.mark.xfail(exceptions=(AssertionError,)) + def test_object_array(self): + from fractions import Fraction + arr = np.array([Fraction(1, 2), Fraction(-1, 2)]) + actual = np.pad(arr, (2, 3), mode='linear_ramp', end_values=0) + + # deliberately chosen to have a non-power-of-2 denominator such that + # rounding to floats causes a failure. + expected = np.array([ + Fraction( 0, 12), + Fraction( 3, 12), + Fraction( 6, 12), + Fraction(-6, 12), + Fraction(-4, 12), + Fraction(-2, 12), + Fraction(-0, 12), + ]) + assert_equal(actual, expected) + + def test_end_values(self): + """Ensure that end values are exact.""" + a = np.pad(np.ones(10).reshape(2, 5), (223, 123), mode="linear_ramp") + assert_equal(a[:, 0], 0.) + assert_equal(a[:, -1], 0.) + assert_equal(a[0, :], 0.) + assert_equal(a[-1, :], 0.) + + @pytest.mark.parametrize("dtype", _numeric_dtypes) + def test_negative_difference(self, dtype): + """ + Check correct behavior of unsigned dtypes if there is a negative + difference between the edge to pad and `end_values`. Check both cases + to be independent of implementation. Test behavior for all other dtypes + in case dtype casting interferes with complex dtypes. See gh-14191. + """ + x = np.array([3], dtype=dtype) + result = np.pad(x, 3, mode="linear_ramp", end_values=0) + expected = np.array([0, 1, 2, 3, 2, 1, 0], dtype=dtype) + assert_equal(result, expected) + + x = np.array([0], dtype=dtype) + result = np.pad(x, 3, mode="linear_ramp", end_values=3) + expected = np.array([3, 2, 1, 0, 1, 2, 3], dtype=dtype) + assert_equal(result, expected) + + +class TestReflect: + def test_check_simple(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'reflect') + b = np.array( + [25, 24, 23, 22, 21, 20, 19, 18, 17, 16, + 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, + 5, 4, 3, 2, 1, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 98, 97, 96, 95, 94, 93, 92, 91, 90, 89, + 88, 87, 86, 85, 84, 83, 82, 81, 80, 79] + ) + assert_array_equal(a, b) + + def test_check_odd_method(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'reflect', reflect_type='odd') + b = np.array( + [-25, -24, -23, -22, -21, -20, -19, -18, -17, -16, + -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, + -5, -4, -3, -2, -1, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, + 110, 111, 112, 113, 114, 115, 116, 117, 118, 119] + ) + assert_array_equal(a, b) + + def test_check_large_pad(self): + a = [[4, 5, 6], [6, 7, 8]] + a = np.pad(a, (5, 7), 'reflect') + b = np.array( + [[7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7, 8, 7, 6, 7], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5]] + ) + assert_array_equal(a, b) + + def test_check_shape(self): + a = [[4, 5, 6]] + a = np.pad(a, (5, 7), 'reflect') + b = np.array( + [[5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5], + [5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5, 6, 5, 4, 5]] + ) + assert_array_equal(a, b) + + def test_check_01(self): + a = np.pad([1, 2, 3], 2, 'reflect') + b = np.array([3, 2, 1, 2, 3, 2, 1]) + assert_array_equal(a, b) + + def test_check_02(self): + a = np.pad([1, 2, 3], 3, 'reflect') + b = np.array([2, 3, 2, 1, 2, 3, 2, 1, 2]) + assert_array_equal(a, b) + + def test_check_03(self): + a = np.pad([1, 2, 3], 4, 'reflect') + b = np.array([1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3]) + assert_array_equal(a, b) + + def test_check_04(self): + a = np.pad([1, 2, 3], [1, 10], 'reflect') + b = np.array([2, 1, 2, 3, 2, 1, 2, 3, 2, 1, 2, 3, 2, 1]) + assert_array_equal(a, b) + + def test_check_05(self): + a = np.pad([1, 2, 3, 4], [45, 10], 'reflect') + b = np.array( + [4, 3, 2, 1, 2, 3, 4, 3, 2, 1, + 2, 3, 4, 3, 2, 1, 2, 3, 4, 3, + 2, 1, 2, 3, 4, 3, 2, 1, 2, 3, + 4, 3, 2, 1, 2, 3, 4, 3, 2, 1, + 2, 3, 4, 3, 2, 1, 2, 3, 4, 3, + 2, 1, 2, 3, 4, 3, 2, 1, 2]) + assert_array_equal(a, b) + + def test_check_06(self): + a = np.pad([1, 2, 3, 4], [15, 2], 'symmetric') + b = np.array( + [2, 3, 4, 4, 3, 2, 1, 1, 2, 3, + 4, 4, 3, 2, 1, 1, 2, 3, 4, 4, + 3] + ) + assert_array_equal(a, b) + + def test_check_07(self): + a = np.pad([1, 2, 3, 4, 5, 6], [45, 3], 'symmetric') + b = np.array( + [4, 5, 6, 6, 5, 4, 3, 2, 1, 1, + 2, 3, 4, 5, 6, 6, 5, 4, 3, 2, + 1, 1, 2, 3, 4, 5, 6, 6, 5, 4, + 3, 2, 1, 1, 2, 3, 4, 5, 6, 6, + 5, 4, 3, 2, 1, 1, 2, 3, 4, 5, + 6, 6, 5, 4]) + assert_array_equal(a, b) + + +class TestEmptyArray: + """Check how padding behaves on arrays with an empty dimension.""" + + @pytest.mark.parametrize( + # Keep parametrization ordered, otherwise pytest-xdist might believe + # that different tests were collected during parallelization + "mode", sorted(_all_modes.keys() - {"constant", "empty"}) + ) + def test_pad_empty_dimension(self, mode): + match = ("can't extend empty axis 0 using modes other than 'constant' " + "or 'empty'") + with pytest.raises(ValueError, match=match): + np.pad([], 4, mode=mode) + with pytest.raises(ValueError, match=match): + np.pad(np.ndarray(0), 4, mode=mode) + with pytest.raises(ValueError, match=match): + np.pad(np.zeros((0, 3)), ((1,), (0,)), mode=mode) + + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_pad_non_empty_dimension(self, mode): + result = np.pad(np.ones((2, 0, 2)), ((3,), (0,), (1,)), mode=mode) + assert result.shape == (8, 0, 4) + + +class TestSymmetric: + def test_check_simple(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'symmetric') + b = np.array( + [24, 23, 22, 21, 20, 19, 18, 17, 16, 15, + 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, + 4, 3, 2, 1, 0, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 99, 98, 97, 96, 95, 94, 93, 92, 91, 90, + 89, 88, 87, 86, 85, 84, 83, 82, 81, 80] + ) + assert_array_equal(a, b) + + def test_check_odd_method(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'symmetric', reflect_type='odd') + b = np.array( + [-24, -23, -22, -21, -20, -19, -18, -17, -16, -15, + -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, + -4, -3, -2, -1, 0, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, + 109, 110, 111, 112, 113, 114, 115, 116, 117, 118] + ) + assert_array_equal(a, b) + + def test_check_large_pad(self): + a = [[4, 5, 6], [6, 7, 8]] + a = np.pad(a, (5, 7), 'symmetric') + b = np.array( + [[5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + [7, 8, 8, 7, 6, 6, 7, 8, 8, 7, 6, 6, 7, 8, 8], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6]] + ) + + assert_array_equal(a, b) + + def test_check_large_pad_odd(self): + a = [[4, 5, 6], [6, 7, 8]] + a = np.pad(a, (5, 7), 'symmetric', reflect_type='odd') + b = np.array( + [[-3, -2, -2, -1, 0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6], + [-3, -2, -2, -1, 0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6], + [-1, 0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8], + [-1, 0, 0, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8], + [ 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10], + + [ 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10], + [ 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12], + + [ 3, 4, 4, 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12], + [ 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14], + [ 5, 6, 6, 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14], + [ 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16], + [ 7, 8, 8, 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16], + [ 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16, 17, 18, 18], + [ 9, 10, 10, 11, 12, 12, 13, 14, 14, 15, 16, 16, 17, 18, 18]] + ) + assert_array_equal(a, b) + + def test_check_shape(self): + a = [[4, 5, 6]] + a = np.pad(a, (5, 7), 'symmetric') + b = np.array( + [[5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6], + [5, 6, 6, 5, 4, 4, 5, 6, 6, 5, 4, 4, 5, 6, 6]] + ) + assert_array_equal(a, b) + + def test_check_01(self): + a = np.pad([1, 2, 3], 2, 'symmetric') + b = np.array([2, 1, 1, 2, 3, 3, 2]) + assert_array_equal(a, b) + + def test_check_02(self): + a = np.pad([1, 2, 3], 3, 'symmetric') + b = np.array([3, 2, 1, 1, 2, 3, 3, 2, 1]) + assert_array_equal(a, b) + + def test_check_03(self): + a = np.pad([1, 2, 3], 6, 'symmetric') + b = np.array([1, 2, 3, 3, 2, 1, 1, 2, 3, 3, 2, 1, 1, 2, 3]) + assert_array_equal(a, b) + + +class TestWrap: + def test_check_simple(self): + a = np.arange(100) + a = np.pad(a, (25, 20), 'wrap') + b = np.array( + [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, + 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, + 95, 96, 97, 98, 99, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, + 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, + 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, + 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, + 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, + 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, + 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, + 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, + + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19] + ) + assert_array_equal(a, b) + + def test_check_large_pad(self): + a = np.arange(12) + a = np.reshape(a, (3, 4)) + a = np.pad(a, (10, 12), 'wrap') + b = np.array( + [[10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11], + [2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, + 3, 0, 1, 2, 3, 0, 1, 2, 3], + [6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, + 7, 4, 5, 6, 7, 4, 5, 6, 7], + [10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, 11, 8, 9, 10, + 11, 8, 9, 10, 11, 8, 9, 10, 11]] + ) + assert_array_equal(a, b) + + def test_check_01(self): + a = np.pad([1, 2, 3], 3, 'wrap') + b = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3]) + assert_array_equal(a, b) + + def test_check_02(self): + a = np.pad([1, 2, 3], 4, 'wrap') + b = np.array([3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1]) + assert_array_equal(a, b) + + def test_pad_with_zero(self): + a = np.ones((3, 5)) + b = np.pad(a, (0, 5), mode="wrap") + assert_array_equal(a, b[:-5, :-5]) + + def test_repeated_wrapping(self): + """ + Check wrapping on each side individually if the wrapped area is longer + than the original array. + """ + a = np.arange(5) + b = np.pad(a, (12, 0), mode="wrap") + assert_array_equal(np.r_[a, a, a, a][3:], b) + + a = np.arange(5) + b = np.pad(a, (0, 12), mode="wrap") + assert_array_equal(np.r_[a, a, a, a][:-3], b) + + def test_repeated_wrapping_multiple_origin(self): + """ + Assert that 'wrap' pads only with multiples of the original area if + the pad width is larger than the original array. + """ + a = np.arange(4).reshape(2, 2) + a = np.pad(a, [(1, 3), (3, 1)], mode='wrap') + b = np.array( + [[3, 2, 3, 2, 3, 2], + [1, 0, 1, 0, 1, 0], + [3, 2, 3, 2, 3, 2], + [1, 0, 1, 0, 1, 0], + [3, 2, 3, 2, 3, 2], + [1, 0, 1, 0, 1, 0]] + ) + assert_array_equal(a, b) + + +class TestEdge: + def test_check_simple(self): + a = np.arange(12) + a = np.reshape(a, (4, 3)) + a = np.pad(a, ((2, 3), (3, 2)), 'edge') + b = np.array( + [[0, 0, 0, 0, 1, 2, 2, 2], + [0, 0, 0, 0, 1, 2, 2, 2], + + [0, 0, 0, 0, 1, 2, 2, 2], + [3, 3, 3, 3, 4, 5, 5, 5], + [6, 6, 6, 6, 7, 8, 8, 8], + [9, 9, 9, 9, 10, 11, 11, 11], + + [9, 9, 9, 9, 10, 11, 11, 11], + [9, 9, 9, 9, 10, 11, 11, 11], + [9, 9, 9, 9, 10, 11, 11, 11]] + ) + assert_array_equal(a, b) + + def test_check_width_shape_1_2(self): + # Check a pad_width of the form ((1, 2),). + # Regression test for issue gh-7808. + a = np.array([1, 2, 3]) + padded = np.pad(a, ((1, 2),), 'edge') + expected = np.array([1, 1, 2, 3, 3, 3]) + assert_array_equal(padded, expected) + + a = np.array([[1, 2, 3], [4, 5, 6]]) + padded = np.pad(a, ((1, 2),), 'edge') + expected = np.pad(a, ((1, 2), (1, 2)), 'edge') + assert_array_equal(padded, expected) + + a = np.arange(24).reshape(2, 3, 4) + padded = np.pad(a, ((1, 2),), 'edge') + expected = np.pad(a, ((1, 2), (1, 2), (1, 2)), 'edge') + assert_array_equal(padded, expected) + + +class TestEmpty: + def test_simple(self): + arr = np.arange(24).reshape(4, 6) + result = np.pad(arr, [(2, 3), (3, 1)], mode="empty") + assert result.shape == (9, 10) + assert_equal(arr, result[2:-3, 3:-1]) + + def test_pad_empty_dimension(self): + arr = np.zeros((3, 0, 2)) + result = np.pad(arr, [(0,), (2,), (1,)], mode="empty") + assert result.shape == (3, 4, 4) + + +def test_legacy_vector_functionality(): + def _padwithtens(vector, pad_width, iaxis, kwargs): + vector[:pad_width[0]] = 10 + vector[-pad_width[1]:] = 10 + + a = np.arange(6).reshape(2, 3) + a = np.pad(a, 2, _padwithtens) + b = np.array( + [[10, 10, 10, 10, 10, 10, 10], + [10, 10, 10, 10, 10, 10, 10], + + [10, 10, 0, 1, 2, 10, 10], + [10, 10, 3, 4, 5, 10, 10], + + [10, 10, 10, 10, 10, 10, 10], + [10, 10, 10, 10, 10, 10, 10]] + ) + assert_array_equal(a, b) + + +def test_unicode_mode(): + a = np.pad([1], 2, mode='constant') + b = np.array([0, 0, 1, 0, 0]) + assert_array_equal(a, b) + + +@pytest.mark.parametrize("mode", ["edge", "symmetric", "reflect", "wrap"]) +def test_object_input(mode): + # Regression test for issue gh-11395. + a = np.full((4, 3), fill_value=None) + pad_amt = ((2, 3), (3, 2)) + b = np.full((9, 8), fill_value=None) + assert_array_equal(np.pad(a, pad_amt, mode=mode), b) + + +class TestPadWidth: + @pytest.mark.parametrize("pad_width", [ + (4, 5, 6, 7), + ((1,), (2,), (3,)), + ((1, 2), (3, 4), (5, 6)), + ((3, 4, 5), (0, 1, 2)), + ]) + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_misshaped_pad_width(self, pad_width, mode): + arr = np.arange(30).reshape((6, 5)) + match = "operands could not be broadcast together" + with pytest.raises(ValueError, match=match): + np.pad(arr, pad_width, mode) + + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_misshaped_pad_width_2(self, mode): + arr = np.arange(30).reshape((6, 5)) + match = ("input operand has more dimensions than allowed by the axis " + "remapping") + with pytest.raises(ValueError, match=match): + np.pad(arr, (((3,), (4,), (5,)), ((0,), (1,), (2,))), mode) + + @pytest.mark.parametrize( + "pad_width", [-2, (-2,), (3, -1), ((5, 2), (-2, 3)), ((-4,), (2,))]) + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_negative_pad_width(self, pad_width, mode): + arr = np.arange(30).reshape((6, 5)) + match = "index can't contain negative values" + with pytest.raises(ValueError, match=match): + np.pad(arr, pad_width, mode) + + @pytest.mark.parametrize("pad_width, dtype", [ + ("3", None), + ("word", None), + (None, None), + (object(), None), + (3.4, None), + (((2, 3, 4), (3, 2)), object), + (complex(1, -1), None), + (((-2.1, 3), (3, 2)), None), + ]) + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_bad_type(self, pad_width, dtype, mode): + arr = np.arange(30).reshape((6, 5)) + match = "`pad_width` must be of integral type." + if dtype is not None: + # avoid DeprecationWarning when not specifying dtype + with pytest.raises(TypeError, match=match): + np.pad(arr, np.array(pad_width, dtype=dtype), mode) + else: + with pytest.raises(TypeError, match=match): + np.pad(arr, pad_width, mode) + with pytest.raises(TypeError, match=match): + np.pad(arr, np.array(pad_width), mode) + + def test_pad_width_as_ndarray(self): + a = np.arange(12) + a = np.reshape(a, (4, 3)) + a = np.pad(a, np.array(((2, 3), (3, 2))), 'edge') + b = np.array( + [[0, 0, 0, 0, 1, 2, 2, 2], + [0, 0, 0, 0, 1, 2, 2, 2], + + [0, 0, 0, 0, 1, 2, 2, 2], + [3, 3, 3, 3, 4, 5, 5, 5], + [6, 6, 6, 6, 7, 8, 8, 8], + [9, 9, 9, 9, 10, 11, 11, 11], + + [9, 9, 9, 9, 10, 11, 11, 11], + [9, 9, 9, 9, 10, 11, 11, 11], + [9, 9, 9, 9, 10, 11, 11, 11]] + ) + assert_array_equal(a, b) + + @pytest.mark.parametrize("pad_width", [0, (0, 0), ((0, 0), (0, 0))]) + @pytest.mark.parametrize("mode", _all_modes.keys()) + def test_zero_pad_width(self, pad_width, mode): + arr = np.arange(30).reshape(6, 5) + assert_array_equal(arr, np.pad(arr, pad_width, mode=mode)) + + +@pytest.mark.parametrize("mode", _all_modes.keys()) +def test_kwargs(mode): + """Test behavior of pad's kwargs for the given mode.""" + allowed = _all_modes[mode] + not_allowed = {} + for kwargs in _all_modes.values(): + if kwargs != allowed: + not_allowed.update(kwargs) + # Test if allowed keyword arguments pass + np.pad([1, 2, 3], 1, mode, **allowed) + # Test if prohibited keyword arguments of other modes raise an error + for key, value in not_allowed.items(): + match = "unsupported keyword arguments for mode '{}'".format(mode) + with pytest.raises(ValueError, match=match): + np.pad([1, 2, 3], 1, mode, **{key: value}) + + +def test_constant_zero_default(): + arr = np.array([1, 1]) + assert_array_equal(np.pad(arr, 2), [0, 0, 1, 1, 0, 0]) + + +@pytest.mark.parametrize("mode", [1, "const", object(), None, True, False]) +def test_unsupported_mode(mode): + match= "mode '{}' is not supported".format(mode) + with pytest.raises(ValueError, match=match): + np.pad([1, 2, 3], 4, mode=mode) + + +@pytest.mark.parametrize("mode", _all_modes.keys()) +def test_non_contiguous_array(mode): + arr = np.arange(24).reshape(4, 6)[::2, ::2] + result = np.pad(arr, (2, 3), mode) + assert result.shape == (7, 8) + assert_equal(result[2:-3, 2:-3], arr) + + +@pytest.mark.parametrize("mode", _all_modes.keys()) +def test_memory_layout_persistence(mode): + """Test if C and F order is preserved for all pad modes.""" + x = np.ones((5, 10), order='C') + assert np.pad(x, 5, mode).flags["C_CONTIGUOUS"] + x = np.ones((5, 10), order='F') + assert np.pad(x, 5, mode).flags["F_CONTIGUOUS"] + + +@pytest.mark.parametrize("dtype", _numeric_dtypes) +@pytest.mark.parametrize("mode", _all_modes.keys()) +def test_dtype_persistence(dtype, mode): + arr = np.zeros((3, 2, 1), dtype=dtype) + result = np.pad(arr, 1, mode=mode) + assert result.dtype == dtype diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraysetops.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraysetops.py new file mode 100644 index 0000000000000000000000000000000000000000..d9721266036dd9a379efea599cc66108951a4d0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arraysetops.py @@ -0,0 +1,1000 @@ +"""Test functions for 1D array set operations. + +""" +import numpy as np + +from numpy import ( + ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, isin + ) +from numpy.exceptions import AxisError +from numpy.testing import (assert_array_equal, assert_equal, + assert_raises, assert_raises_regex) +import pytest + + +class TestSetOps: + + def test_intersect1d(self): + # unique inputs + a = np.array([5, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5]) + + ec = np.array([1, 2, 5]) + c = intersect1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + # non-unique inputs + a = np.array([5, 5, 7, 1, 2]) + b = np.array([2, 1, 4, 3, 3, 1, 5]) + + ed = np.array([1, 2, 5]) + c = intersect1d(a, b) + assert_array_equal(c, ed) + assert_array_equal([], intersect1d([], [])) + + def test_intersect1d_array_like(self): + # See gh-11772 + class Test: + def __array__(self, dtype=None, copy=None): + return np.arange(3) + + a = Test() + res = intersect1d(a, a) + assert_array_equal(res, a) + res = intersect1d([1, 2, 3], [1, 2, 3]) + assert_array_equal(res, [1, 2, 3]) + + def test_intersect1d_indices(self): + # unique inputs + a = np.array([1, 2, 3, 4]) + b = np.array([2, 1, 4, 6]) + c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) + ee = np.array([1, 2, 4]) + assert_array_equal(c, ee) + assert_array_equal(a[i1], ee) + assert_array_equal(b[i2], ee) + + # non-unique inputs + a = np.array([1, 2, 2, 3, 4, 3, 2]) + b = np.array([1, 8, 4, 2, 2, 3, 2, 3]) + c, i1, i2 = intersect1d(a, b, return_indices=True) + ef = np.array([1, 2, 3, 4]) + assert_array_equal(c, ef) + assert_array_equal(a[i1], ef) + assert_array_equal(b[i2], ef) + + # non1d, unique inputs + a = np.array([[2, 4, 5, 6], [7, 8, 1, 15]]) + b = np.array([[3, 2, 7, 6], [10, 12, 8, 9]]) + c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) + ui1 = np.unravel_index(i1, a.shape) + ui2 = np.unravel_index(i2, b.shape) + ea = np.array([2, 6, 7, 8]) + assert_array_equal(ea, a[ui1]) + assert_array_equal(ea, b[ui2]) + + # non1d, not assumed to be uniqueinputs + a = np.array([[2, 4, 5, 6, 6], [4, 7, 8, 7, 2]]) + b = np.array([[3, 2, 7, 7], [10, 12, 8, 7]]) + c, i1, i2 = intersect1d(a, b, return_indices=True) + ui1 = np.unravel_index(i1, a.shape) + ui2 = np.unravel_index(i2, b.shape) + ea = np.array([2, 7, 8]) + assert_array_equal(ea, a[ui1]) + assert_array_equal(ea, b[ui2]) + + def test_setxor1d(self): + a = np.array([5, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5]) + + ec = np.array([3, 4, 7]) + c = setxor1d(a, b) + assert_array_equal(c, ec) + + a = np.array([1, 2, 3]) + b = np.array([6, 5, 4]) + + ec = np.array([1, 2, 3, 4, 5, 6]) + c = setxor1d(a, b) + assert_array_equal(c, ec) + + a = np.array([1, 8, 2, 3]) + b = np.array([6, 5, 4, 8]) + + ec = np.array([1, 2, 3, 4, 5, 6]) + c = setxor1d(a, b) + assert_array_equal(c, ec) + + assert_array_equal([], setxor1d([], [])) + + def test_setxor1d_unique(self): + a = np.array([1, 8, 2, 3]) + b = np.array([6, 5, 4, 8]) + + ec = np.array([1, 2, 3, 4, 5, 6]) + c = setxor1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + a = np.array([[1], [8], [2], [3]]) + b = np.array([[6, 5], [4, 8]]) + + ec = np.array([1, 2, 3, 4, 5, 6]) + c = setxor1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + def test_ediff1d(self): + zero_elem = np.array([]) + one_elem = np.array([1]) + two_elem = np.array([1, 2]) + + assert_array_equal([], ediff1d(zero_elem)) + assert_array_equal([0], ediff1d(zero_elem, to_begin=0)) + assert_array_equal([0], ediff1d(zero_elem, to_end=0)) + assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0)) + assert_array_equal([], ediff1d(one_elem)) + assert_array_equal([1], ediff1d(two_elem)) + assert_array_equal([7, 1, 9], ediff1d(two_elem, to_begin=7, to_end=9)) + assert_array_equal([5, 6, 1, 7, 8], + ediff1d(two_elem, to_begin=[5, 6], to_end=[7, 8])) + assert_array_equal([1, 9], ediff1d(two_elem, to_end=9)) + assert_array_equal([1, 7, 8], ediff1d(two_elem, to_end=[7, 8])) + assert_array_equal([7, 1], ediff1d(two_elem, to_begin=7)) + assert_array_equal([5, 6, 1], ediff1d(two_elem, to_begin=[5, 6])) + + @pytest.mark.parametrize("ary, prepend, append, expected", [ + # should fail because trying to cast + # np.nan standard floating point value + # into an integer array: + (np.array([1, 2, 3], dtype=np.int64), + None, + np.nan, + 'to_end'), + # should fail because attempting + # to downcast to int type: + (np.array([1, 2, 3], dtype=np.int64), + np.array([5, 7, 2], dtype=np.float32), + None, + 'to_begin'), + # should fail because attempting to cast + # two special floating point values + # to integers (on both sides of ary), + # `to_begin` is in the error message as the impl checks this first: + (np.array([1., 3., 9.], dtype=np.int8), + np.nan, + np.nan, + 'to_begin'), + ]) + def test_ediff1d_forbidden_type_casts(self, ary, prepend, append, expected): + # verify resolution of gh-11490 + + # specifically, raise an appropriate + # Exception when attempting to append or + # prepend with an incompatible type + msg = 'dtype of `{}` must be compatible'.format(expected) + with assert_raises_regex(TypeError, msg): + ediff1d(ary=ary, + to_end=append, + to_begin=prepend) + + @pytest.mark.parametrize( + "ary,prepend,append,expected", + [ + (np.array([1, 2, 3], dtype=np.int16), + 2**16, # will be cast to int16 under same kind rule. + 2**16 + 4, + np.array([0, 1, 1, 4], dtype=np.int16)), + (np.array([1, 2, 3], dtype=np.float32), + np.array([5], dtype=np.float64), + None, + np.array([5, 1, 1], dtype=np.float32)), + (np.array([1, 2, 3], dtype=np.int32), + 0, + 0, + np.array([0, 1, 1, 0], dtype=np.int32)), + (np.array([1, 2, 3], dtype=np.int64), + 3, + -9, + np.array([3, 1, 1, -9], dtype=np.int64)), + ] + ) + def test_ediff1d_scalar_handling(self, + ary, + prepend, + append, + expected): + # maintain backwards-compatibility + # of scalar prepend / append behavior + # in ediff1d following fix for gh-11490 + actual = np.ediff1d(ary=ary, + to_end=append, + to_begin=prepend) + assert_equal(actual, expected) + assert actual.dtype == expected.dtype + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin(self, kind): + def _isin_slow(a, b): + b = np.asarray(b).flatten().tolist() + return a in b + isin_slow = np.vectorize(_isin_slow, otypes=[bool], excluded={1}) + + def assert_isin_equal(a, b): + x = isin(a, b, kind=kind) + y = isin_slow(a, b) + assert_array_equal(x, y) + + # multidimensional arrays in both arguments + a = np.arange(24).reshape([2, 3, 4]) + b = np.array([[10, 20, 30], [0, 1, 3], [11, 22, 33]]) + assert_isin_equal(a, b) + + # array-likes as both arguments + c = [(9, 8), (7, 6)] + d = (9, 7) + assert_isin_equal(c, d) + + # zero-d array: + f = np.array(3) + assert_isin_equal(f, b) + assert_isin_equal(a, f) + assert_isin_equal(f, f) + + # scalar: + assert_isin_equal(5, b) + assert_isin_equal(a, 6) + assert_isin_equal(5, 6) + + # empty array-like: + if kind != "table": + # An empty list will become float64, + # which is invalid for kind="table" + x = [] + assert_isin_equal(x, b) + assert_isin_equal(a, x) + assert_isin_equal(x, x) + + # empty array with various types: + for dtype in [bool, np.int64, np.float64]: + if kind == "table" and dtype == np.float64: + continue + + if dtype in {np.int64, np.float64}: + ar = np.array([10, 20, 30], dtype=dtype) + elif dtype in {bool}: + ar = np.array([True, False, False]) + + empty_array = np.array([], dtype=dtype) + + assert_isin_equal(empty_array, ar) + assert_isin_equal(ar, empty_array) + assert_isin_equal(empty_array, empty_array) + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_additional(self, kind): + # we use two different sizes for the b array here to test the + # two different paths in isin(). + for mult in (1, 10): + # One check without np.array to make sure lists are handled correct + a = [5, 7, 1, 2] + b = [2, 4, 3, 1, 5] * mult + ec = np.array([True, False, True, True]) + c = isin(a, b, assume_unique=True, kind=kind) + assert_array_equal(c, ec) + + a[0] = 8 + ec = np.array([False, False, True, True]) + c = isin(a, b, assume_unique=True, kind=kind) + assert_array_equal(c, ec) + + a[0], a[3] = 4, 8 + ec = np.array([True, False, True, False]) + c = isin(a, b, assume_unique=True, kind=kind) + assert_array_equal(c, ec) + + a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) + b = [2, 3, 4] * mult + ec = [False, True, False, True, True, True, True, True, True, + False, True, False, False, False] + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + b = b + [5, 5, 4] * mult + ec = [True, True, True, True, True, True, True, True, True, True, + True, False, True, True] + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5] * mult) + ec = np.array([True, False, True, True]) + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 1, 2]) + b = np.array([2, 4, 3, 3, 1, 5] * mult) + ec = np.array([True, False, True, True, True]) + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + a = np.array([5, 5]) + b = np.array([2, 2] * mult) + ec = np.array([False, False]) + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + a = np.array([5]) + b = np.array([2]) + ec = np.array([False]) + c = isin(a, b, kind=kind) + assert_array_equal(c, ec) + + if kind in {None, "sort"}: + assert_array_equal(isin([], [], kind=kind), []) + + def test_isin_char_array(self): + a = np.array(['a', 'b', 'c', 'd', 'e', 'c', 'e', 'b']) + b = np.array(['a', 'c']) + + ec = np.array([True, False, True, False, False, True, False, False]) + c = isin(a, b) + + assert_array_equal(c, ec) + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_invert(self, kind): + "Test isin's invert parameter" + # We use two different sizes for the b array here to test the + # two different paths in isin(). + for mult in (1, 10): + a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) + b = [2, 3, 4] * mult + assert_array_equal(np.invert(isin(a, b, kind=kind)), + isin(a, b, invert=True, kind=kind)) + + # float: + if kind in {None, "sort"}: + for mult in (1, 10): + a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5], + dtype=np.float32) + b = [2, 3, 4] * mult + b = np.array(b, dtype=np.float32) + assert_array_equal(np.invert(isin(a, b, kind=kind)), + isin(a, b, invert=True, kind=kind)) + + def test_isin_hit_alternate_algorithm(self): + """Hit the standard isin code with integers""" + # Need extreme range to hit standard code + # This hits it without the use of kind='table' + a = np.array([5, 4, 5, 3, 4, 4, 1e9], dtype=np.int64) + b = np.array([2, 3, 4, 1e9], dtype=np.int64) + expected = np.array([0, 1, 0, 1, 1, 1, 1], dtype=bool) + assert_array_equal(expected, isin(a, b)) + assert_array_equal(np.invert(expected), isin(a, b, invert=True)) + + a = np.array([5, 7, 1, 2], dtype=np.int64) + b = np.array([2, 4, 3, 1, 5, 1e9], dtype=np.int64) + ec = np.array([True, False, True, True]) + c = isin(a, b, assume_unique=True) + assert_array_equal(c, ec) + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_boolean(self, kind): + """Test that isin works for boolean input""" + a = np.array([True, False]) + b = np.array([False, False, False]) + expected = np.array([False, True]) + assert_array_equal(expected, + isin(a, b, kind=kind)) + assert_array_equal(np.invert(expected), + isin(a, b, invert=True, kind=kind)) + + @pytest.mark.parametrize("kind", [None, "sort"]) + def test_isin_timedelta(self, kind): + """Test that isin works for timedelta input""" + rstate = np.random.RandomState(0) + a = rstate.randint(0, 100, size=10) + b = rstate.randint(0, 100, size=10) + truth = isin(a, b) + a_timedelta = a.astype("timedelta64[s]") + b_timedelta = b.astype("timedelta64[s]") + assert_array_equal(truth, isin(a_timedelta, b_timedelta, kind=kind)) + + def test_isin_table_timedelta_fails(self): + a = np.array([0, 1, 2], dtype="timedelta64[s]") + b = a + # Make sure it raises a value error: + with pytest.raises(ValueError): + isin(a, b, kind="table") + + @pytest.mark.parametrize( + "dtype1,dtype2", + [ + (np.int8, np.int16), + (np.int16, np.int8), + (np.uint8, np.uint16), + (np.uint16, np.uint8), + (np.uint8, np.int16), + (np.int16, np.uint8), + (np.uint64, np.int64), + ] + ) + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_mixed_dtype(self, dtype1, dtype2, kind): + """Test that isin works as expected for mixed dtype input.""" + is_dtype2_signed = np.issubdtype(dtype2, np.signedinteger) + ar1 = np.array([0, 0, 1, 1], dtype=dtype1) + + if is_dtype2_signed: + ar2 = np.array([-128, 0, 127], dtype=dtype2) + else: + ar2 = np.array([127, 0, 255], dtype=dtype2) + + expected = np.array([True, True, False, False]) + + expect_failure = kind == "table" and ( + dtype1 == np.int16 and dtype2 == np.int8) + + if expect_failure: + with pytest.raises(RuntimeError, match="exceed the maximum"): + isin(ar1, ar2, kind=kind) + else: + assert_array_equal(isin(ar1, ar2, kind=kind), expected) + + @pytest.mark.parametrize("data", [ + np.array([2**63, 2**63+1], dtype=np.uint64), + np.array([-2**62, -2**62-1], dtype=np.int64), + ]) + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_mixed_huge_vals(self, kind, data): + """Test values outside intp range (negative ones if 32bit system)""" + query = data[1] + res = np.isin(data, query, kind=kind) + assert_array_equal(res, [False, True]) + # Also check that nothing weird happens for values can't possibly + # in range. + data = data.astype(np.int32) # clearly different values + res = np.isin(data, query, kind=kind) + assert_array_equal(res, [False, False]) + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_isin_mixed_boolean(self, kind): + """Test that isin works as expected for bool/int input.""" + for dtype in np.typecodes["AllInteger"]: + a = np.array([True, False, False], dtype=bool) + b = np.array([0, 0, 0, 0], dtype=dtype) + expected = np.array([False, True, True], dtype=bool) + assert_array_equal(isin(a, b, kind=kind), expected) + + a, b = b, a + expected = np.array([True, True, True, True], dtype=bool) + assert_array_equal(isin(a, b, kind=kind), expected) + + def test_isin_first_array_is_object(self): + ar1 = [None] + ar2 = np.array([1]*10) + expected = np.array([False]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + + def test_isin_second_array_is_object(self): + ar1 = 1 + ar2 = np.array([None]*10) + expected = np.array([False]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + + def test_isin_both_arrays_are_object(self): + ar1 = [None] + ar2 = np.array([None]*10) + expected = np.array([True]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + + def test_isin_both_arrays_have_structured_dtype(self): + # Test arrays of a structured data type containing an integer field + # and a field of dtype `object` allowing for arbitrary Python objects + dt = np.dtype([('field1', int), ('field2', object)]) + ar1 = np.array([(1, None)], dtype=dt) + ar2 = np.array([(1, None)]*10, dtype=dt) + expected = np.array([True]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + + def test_isin_with_arrays_containing_tuples(self): + ar1 = np.array([(1,), 2], dtype=object) + ar2 = np.array([(1,), 2], dtype=object) + expected = np.array([True, True]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + result = np.isin(ar1, ar2, invert=True) + assert_array_equal(result, np.invert(expected)) + + # An integer is added at the end of the array to make sure + # that the array builder will create the array with tuples + # and after it's created the integer is removed. + # There's a bug in the array constructor that doesn't handle + # tuples properly and adding the integer fixes that. + ar1 = np.array([(1,), (2, 1), 1], dtype=object) + ar1 = ar1[:-1] + ar2 = np.array([(1,), (2, 1), 1], dtype=object) + ar2 = ar2[:-1] + expected = np.array([True, True]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + result = np.isin(ar1, ar2, invert=True) + assert_array_equal(result, np.invert(expected)) + + ar1 = np.array([(1,), (2, 3), 1], dtype=object) + ar1 = ar1[:-1] + ar2 = np.array([(1,), 2], dtype=object) + expected = np.array([True, False]) + result = np.isin(ar1, ar2) + assert_array_equal(result, expected) + result = np.isin(ar1, ar2, invert=True) + assert_array_equal(result, np.invert(expected)) + + def test_isin_errors(self): + """Test that isin raises expected errors.""" + + # Error 1: `kind` is not one of 'sort' 'table' or None. + ar1 = np.array([1, 2, 3, 4, 5]) + ar2 = np.array([2, 4, 6, 8, 10]) + assert_raises(ValueError, isin, ar1, ar2, kind='quicksort') + + # Error 2: `kind="table"` does not work for non-integral arrays. + obj_ar1 = np.array([1, 'a', 3, 'b', 5], dtype=object) + obj_ar2 = np.array([1, 'a', 3, 'b', 5], dtype=object) + assert_raises(ValueError, isin, obj_ar1, obj_ar2, kind='table') + + for dtype in [np.int32, np.int64]: + ar1 = np.array([-1, 2, 3, 4, 5], dtype=dtype) + # The range of this array will overflow: + overflow_ar2 = np.array([-1, np.iinfo(dtype).max], dtype=dtype) + + # Error 3: `kind="table"` will trigger a runtime error + # if there is an integer overflow expected when computing the + # range of ar2 + assert_raises( + RuntimeError, + isin, ar1, overflow_ar2, kind='table' + ) + + # Non-error: `kind=None` will *not* trigger a runtime error + # if there is an integer overflow, it will switch to + # the `sort` algorithm. + result = np.isin(ar1, overflow_ar2, kind=None) + assert_array_equal(result, [True] + [False] * 4) + result = np.isin(ar1, overflow_ar2, kind='sort') + assert_array_equal(result, [True] + [False] * 4) + + def test_union1d(self): + a = np.array([5, 4, 7, 1, 2]) + b = np.array([2, 4, 3, 3, 2, 1, 5]) + + ec = np.array([1, 2, 3, 4, 5, 7]) + c = union1d(a, b) + assert_array_equal(c, ec) + + # Tests gh-10340, arguments to union1d should be + # flattened if they are not already 1D + x = np.array([[0, 1, 2], [3, 4, 5]]) + y = np.array([0, 1, 2, 3, 4]) + ez = np.array([0, 1, 2, 3, 4, 5]) + z = union1d(x, y) + assert_array_equal(z, ez) + + assert_array_equal([], union1d([], [])) + + def test_setdiff1d(self): + a = np.array([6, 5, 4, 7, 1, 2, 7, 4]) + b = np.array([2, 4, 3, 3, 2, 1, 5]) + + ec = np.array([6, 7]) + c = setdiff1d(a, b) + assert_array_equal(c, ec) + + a = np.arange(21) + b = np.arange(19) + ec = np.array([19, 20]) + c = setdiff1d(a, b) + assert_array_equal(c, ec) + + assert_array_equal([], setdiff1d([], [])) + a = np.array((), np.uint32) + assert_equal(setdiff1d(a, []).dtype, np.uint32) + + def test_setdiff1d_unique(self): + a = np.array([3, 2, 1]) + b = np.array([7, 5, 2]) + expected = np.array([3, 1]) + actual = setdiff1d(a, b, assume_unique=True) + assert_equal(actual, expected) + + def test_setdiff1d_char_array(self): + a = np.array(['a', 'b', 'c']) + b = np.array(['a', 'b', 's']) + assert_array_equal(setdiff1d(a, b), np.array(['c'])) + + def test_manyways(self): + a = np.array([5, 7, 1, 2, 8]) + b = np.array([9, 8, 2, 4, 3, 1, 5]) + + c1 = setxor1d(a, b) + aux1 = intersect1d(a, b) + aux2 = union1d(a, b) + c2 = setdiff1d(aux2, aux1) + assert_array_equal(c1, c2) + + +class TestUnique: + + def test_unique_1d(self): + + def check_all(a, b, i1, i2, c, dt): + base_msg = 'check {0} failed for type {1}' + + msg = base_msg.format('values', dt) + v = unique(a) + assert_array_equal(v, b, msg) + + msg = base_msg.format('return_index', dt) + v, j = unique(a, True, False, False) + assert_array_equal(v, b, msg) + assert_array_equal(j, i1, msg) + + msg = base_msg.format('return_inverse', dt) + v, j = unique(a, False, True, False) + assert_array_equal(v, b, msg) + assert_array_equal(j, i2, msg) + + msg = base_msg.format('return_counts', dt) + v, j = unique(a, False, False, True) + assert_array_equal(v, b, msg) + assert_array_equal(j, c, msg) + + msg = base_msg.format('return_index and return_inverse', dt) + v, j1, j2 = unique(a, True, True, False) + assert_array_equal(v, b, msg) + assert_array_equal(j1, i1, msg) + assert_array_equal(j2, i2, msg) + + msg = base_msg.format('return_index and return_counts', dt) + v, j1, j2 = unique(a, True, False, True) + assert_array_equal(v, b, msg) + assert_array_equal(j1, i1, msg) + assert_array_equal(j2, c, msg) + + msg = base_msg.format('return_inverse and return_counts', dt) + v, j1, j2 = unique(a, False, True, True) + assert_array_equal(v, b, msg) + assert_array_equal(j1, i2, msg) + assert_array_equal(j2, c, msg) + + msg = base_msg.format(('return_index, return_inverse ' + 'and return_counts'), dt) + v, j1, j2, j3 = unique(a, True, True, True) + assert_array_equal(v, b, msg) + assert_array_equal(j1, i1, msg) + assert_array_equal(j2, i2, msg) + assert_array_equal(j3, c, msg) + + a = [5, 7, 1, 2, 1, 5, 7]*10 + b = [1, 2, 5, 7] + i1 = [2, 3, 0, 1] + i2 = [2, 3, 0, 1, 0, 2, 3]*10 + c = np.multiply([2, 1, 2, 2], 10) + + # test for numeric arrays + types = [] + types.extend(np.typecodes['AllInteger']) + types.extend(np.typecodes['AllFloat']) + types.append('datetime64[D]') + types.append('timedelta64[D]') + for dt in types: + aa = np.array(a, dt) + bb = np.array(b, dt) + check_all(aa, bb, i1, i2, c, dt) + + # test for object arrays + dt = 'O' + aa = np.empty(len(a), dt) + aa[:] = a + bb = np.empty(len(b), dt) + bb[:] = b + check_all(aa, bb, i1, i2, c, dt) + + # test for structured arrays + dt = [('', 'i'), ('', 'i')] + aa = np.array(list(zip(a, a)), dt) + bb = np.array(list(zip(b, b)), dt) + check_all(aa, bb, i1, i2, c, dt) + + # test for ticket #2799 + aa = [1. + 0.j, 1 - 1.j, 1] + assert_array_equal(np.unique(aa), [1. - 1.j, 1. + 0.j]) + + # test for ticket #4785 + a = [(1, 2), (1, 2), (2, 3)] + unq = [1, 2, 3] + inv = [[0, 1], [0, 1], [1, 2]] + a1 = unique(a) + assert_array_equal(a1, unq) + a2, a2_inv = unique(a, return_inverse=True) + assert_array_equal(a2, unq) + assert_array_equal(a2_inv, inv) + + # test for chararrays with return_inverse (gh-5099) + a = np.char.chararray(5) + a[...] = '' + a2, a2_inv = np.unique(a, return_inverse=True) + assert_array_equal(a2_inv, np.zeros(5)) + + # test for ticket #9137 + a = [] + a1_idx = np.unique(a, return_index=True)[1] + a2_inv = np.unique(a, return_inverse=True)[1] + a3_idx, a3_inv = np.unique(a, return_index=True, + return_inverse=True)[1:] + assert_equal(a1_idx.dtype, np.intp) + assert_equal(a2_inv.dtype, np.intp) + assert_equal(a3_idx.dtype, np.intp) + assert_equal(a3_inv.dtype, np.intp) + + # test for ticket 2111 - float + a = [2.0, np.nan, 1.0, np.nan] + ua = [1.0, 2.0, np.nan] + ua_idx = [2, 0, 1] + ua_inv = [1, 2, 0, 2] + ua_cnt = [1, 1, 2] + assert_equal(np.unique(a), ua) + assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) + assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) + assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) + + # test for ticket 2111 - complex + a = [2.0-1j, np.nan, 1.0+1j, complex(0.0, np.nan), complex(1.0, np.nan)] + ua = [1.0+1j, 2.0-1j, complex(0.0, np.nan)] + ua_idx = [2, 0, 3] + ua_inv = [1, 2, 0, 2, 2] + ua_cnt = [1, 1, 3] + assert_equal(np.unique(a), ua) + assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) + assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) + assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) + + # test for ticket 2111 - datetime64 + nat = np.datetime64('nat') + a = [np.datetime64('2020-12-26'), nat, np.datetime64('2020-12-24'), nat] + ua = [np.datetime64('2020-12-24'), np.datetime64('2020-12-26'), nat] + ua_idx = [2, 0, 1] + ua_inv = [1, 2, 0, 2] + ua_cnt = [1, 1, 2] + assert_equal(np.unique(a), ua) + assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) + assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) + assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) + + # test for ticket 2111 - timedelta + nat = np.timedelta64('nat') + a = [np.timedelta64(1, 'D'), nat, np.timedelta64(1, 'h'), nat] + ua = [np.timedelta64(1, 'h'), np.timedelta64(1, 'D'), nat] + ua_idx = [2, 0, 1] + ua_inv = [1, 2, 0, 2] + ua_cnt = [1, 1, 2] + assert_equal(np.unique(a), ua) + assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) + assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) + assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) + + # test for gh-19300 + all_nans = [np.nan] * 4 + ua = [np.nan] + ua_idx = [0] + ua_inv = [0, 0, 0, 0] + ua_cnt = [4] + assert_equal(np.unique(all_nans), ua) + assert_equal(np.unique(all_nans, return_index=True), (ua, ua_idx)) + assert_equal(np.unique(all_nans, return_inverse=True), (ua, ua_inv)) + assert_equal(np.unique(all_nans, return_counts=True), (ua, ua_cnt)) + + def test_unique_axis_errors(self): + assert_raises(TypeError, self._run_axis_tests, object) + assert_raises(TypeError, self._run_axis_tests, + [('a', int), ('b', object)]) + + assert_raises(AxisError, unique, np.arange(10), axis=2) + assert_raises(AxisError, unique, np.arange(10), axis=-2) + + def test_unique_axis_list(self): + msg = "Unique failed on list of lists" + inp = [[0, 1, 0], [0, 1, 0]] + inp_arr = np.asarray(inp) + assert_array_equal(unique(inp, axis=0), unique(inp_arr, axis=0), msg) + assert_array_equal(unique(inp, axis=1), unique(inp_arr, axis=1), msg) + + def test_unique_axis(self): + types = [] + types.extend(np.typecodes['AllInteger']) + types.extend(np.typecodes['AllFloat']) + types.append('datetime64[D]') + types.append('timedelta64[D]') + types.append([('a', int), ('b', int)]) + types.append([('a', int), ('b', float)]) + + for dtype in types: + self._run_axis_tests(dtype) + + msg = 'Non-bitwise-equal booleans test failed' + data = np.arange(10, dtype=np.uint8).reshape(-1, 2).view(bool) + result = np.array([[False, True], [True, True]], dtype=bool) + assert_array_equal(unique(data, axis=0), result, msg) + + msg = 'Negative zero equality test failed' + data = np.array([[-0.0, 0.0], [0.0, -0.0], [-0.0, 0.0], [0.0, -0.0]]) + result = np.array([[-0.0, 0.0]]) + assert_array_equal(unique(data, axis=0), result, msg) + + @pytest.mark.parametrize("axis", [0, -1]) + def test_unique_1d_with_axis(self, axis): + x = np.array([4, 3, 2, 3, 2, 1, 2, 2]) + uniq = unique(x, axis=axis) + assert_array_equal(uniq, [1, 2, 3, 4]) + + @pytest.mark.parametrize("axis", [None, 0, -1]) + def test_unique_inverse_with_axis(self, axis): + x = np.array([[4, 4, 3], [2, 2, 1], [2, 2, 1], [4, 4, 3]]) + uniq, inv = unique(x, return_inverse=True, axis=axis) + assert_equal(inv.ndim, x.ndim if axis is None else 1) + assert_array_equal(x, np.take(uniq, inv, axis=axis)) + + def test_unique_axis_zeros(self): + # issue 15559 + single_zero = np.empty(shape=(2, 0), dtype=np.int8) + uniq, idx, inv, cnt = unique(single_zero, axis=0, return_index=True, + return_inverse=True, return_counts=True) + + # there's 1 element of shape (0,) along axis 0 + assert_equal(uniq.dtype, single_zero.dtype) + assert_array_equal(uniq, np.empty(shape=(1, 0))) + assert_array_equal(idx, np.array([0])) + assert_array_equal(inv, np.array([0, 0])) + assert_array_equal(cnt, np.array([2])) + + # there's 0 elements of shape (2,) along axis 1 + uniq, idx, inv, cnt = unique(single_zero, axis=1, return_index=True, + return_inverse=True, return_counts=True) + + assert_equal(uniq.dtype, single_zero.dtype) + assert_array_equal(uniq, np.empty(shape=(2, 0))) + assert_array_equal(idx, np.array([])) + assert_array_equal(inv, np.array([])) + assert_array_equal(cnt, np.array([])) + + # test a "complicated" shape + shape = (0, 2, 0, 3, 0, 4, 0) + multiple_zeros = np.empty(shape=shape) + for axis in range(len(shape)): + expected_shape = list(shape) + if shape[axis] == 0: + expected_shape[axis] = 0 + else: + expected_shape[axis] = 1 + + assert_array_equal(unique(multiple_zeros, axis=axis), + np.empty(shape=expected_shape)) + + def test_unique_masked(self): + # issue 8664 + x = np.array([64, 0, 1, 2, 3, 63, 63, 0, 0, 0, 1, 2, 0, 63, 0], + dtype='uint8') + y = np.ma.masked_equal(x, 0) + + v = np.unique(y) + v2, i, c = np.unique(y, return_index=True, return_counts=True) + + msg = 'Unique returned different results when asked for index' + assert_array_equal(v.data, v2.data, msg) + assert_array_equal(v.mask, v2.mask, msg) + + def test_unique_sort_order_with_axis(self): + # These tests fail if sorting along axis is done by treating subarrays + # as unsigned byte strings. See gh-10495. + fmt = "sort order incorrect for integer type '%s'" + for dt in 'bhilq': + a = np.array([[-1], [0]], dt) + b = np.unique(a, axis=0) + assert_array_equal(a, b, fmt % dt) + + def _run_axis_tests(self, dtype): + data = np.array([[0, 1, 0, 0], + [1, 0, 0, 0], + [0, 1, 0, 0], + [1, 0, 0, 0]]).astype(dtype) + + msg = 'Unique with 1d array and axis=0 failed' + result = np.array([0, 1]) + assert_array_equal(unique(data), result.astype(dtype), msg) + + msg = 'Unique with 2d array and axis=0 failed' + result = np.array([[0, 1, 0, 0], [1, 0, 0, 0]]) + assert_array_equal(unique(data, axis=0), result.astype(dtype), msg) + + msg = 'Unique with 2d array and axis=1 failed' + result = np.array([[0, 0, 1], [0, 1, 0], [0, 0, 1], [0, 1, 0]]) + assert_array_equal(unique(data, axis=1), result.astype(dtype), msg) + + msg = 'Unique with 3d array and axis=2 failed' + data3d = np.array([[[1, 1], + [1, 0]], + [[0, 1], + [0, 0]]]).astype(dtype) + result = np.take(data3d, [1, 0], axis=2) + assert_array_equal(unique(data3d, axis=2), result, msg) + + uniq, idx, inv, cnt = unique(data, axis=0, return_index=True, + return_inverse=True, return_counts=True) + msg = "Unique's return_index=True failed with axis=0" + assert_array_equal(data[idx], uniq, msg) + msg = "Unique's return_inverse=True failed with axis=0" + assert_array_equal(np.take(uniq, inv, axis=0), data) + msg = "Unique's return_counts=True failed with axis=0" + assert_array_equal(cnt, np.array([2, 2]), msg) + + uniq, idx, inv, cnt = unique(data, axis=1, return_index=True, + return_inverse=True, return_counts=True) + msg = "Unique's return_index=True failed with axis=1" + assert_array_equal(data[:, idx], uniq) + msg = "Unique's return_inverse=True failed with axis=1" + assert_array_equal(np.take(uniq, inv, axis=1), data) + msg = "Unique's return_counts=True failed with axis=1" + assert_array_equal(cnt, np.array([2, 1, 1]), msg) + + def test_unique_nanequals(self): + # issue 20326 + a = np.array([1, 1, np.nan, np.nan, np.nan]) + unq = np.unique(a) + not_unq = np.unique(a, equal_nan=False) + assert_array_equal(unq, np.array([1, np.nan])) + assert_array_equal(not_unq, np.array([1, np.nan, np.nan, np.nan])) + + def test_unique_array_api_functions(self): + arr = np.array([np.nan, 1, 4, 1, 3, 4, np.nan, 5, 1]) + + for res_unique_array_api, res_unique in [ + ( + np.unique_values(arr), + np.unique(arr, equal_nan=False) + ), + ( + np.unique_counts(arr), + np.unique(arr, return_counts=True, equal_nan=False) + ), + ( + np.unique_inverse(arr), + np.unique(arr, return_inverse=True, equal_nan=False) + ), + ( + np.unique_all(arr), + np.unique( + arr, + return_index=True, + return_inverse=True, + return_counts=True, + equal_nan=False + ) + ) + ]: + assert len(res_unique_array_api) == len(res_unique) + for actual, expected in zip(res_unique_array_api, res_unique): + assert_array_equal(actual, expected) + + def test_unique_inverse_shape(self): + # Regression test for https://github.com/numpy/numpy/issues/25552 + arr = np.array([[1, 2, 3], [2, 3, 1]]) + expected_values, expected_inverse = np.unique(arr, return_inverse=True) + expected_inverse = expected_inverse.reshape(arr.shape) + for func in np.unique_inverse, np.unique_all: + result = func(arr) + assert_array_equal(expected_values, result.values) + assert_array_equal(expected_inverse, result.inverse_indices) + assert_array_equal(arr, result.values[result.inverse_indices]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arrayterator.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arrayterator.py new file mode 100644 index 0000000000000000000000000000000000000000..c00ed13d7f3076d53ec080a46fe7e13ff7dfb5a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_arrayterator.py @@ -0,0 +1,46 @@ +from operator import mul +from functools import reduce + +import numpy as np +from numpy.random import randint +from numpy.lib import Arrayterator +from numpy.testing import assert_ + + +def test(): + np.random.seed(np.arange(10)) + + # Create a random array + ndims = randint(5)+1 + shape = tuple(randint(10)+1 for dim in range(ndims)) + els = reduce(mul, shape) + a = np.arange(els) + a.shape = shape + + buf_size = randint(2*els) + b = Arrayterator(a, buf_size) + + # Check that each block has at most ``buf_size`` elements + for block in b: + assert_(len(block.flat) <= (buf_size or els)) + + # Check that all elements are iterated correctly + assert_(list(b.flat) == list(a.flat)) + + # Slice arrayterator + start = [randint(dim) for dim in shape] + stop = [randint(dim)+1 for dim in shape] + step = [randint(dim)+1 for dim in shape] + slice_ = tuple(slice(*t) for t in zip(start, stop, step)) + c = b[slice_] + d = a[slice_] + + # Check that each block has at most ``buf_size`` elements + for block in c: + assert_(len(block.flat) <= (buf_size or els)) + + # Check that the arrayterator is sliced correctly + assert_(np.all(c.__array__() == d)) + + # Check that all elements are iterated correctly + assert_(list(c.flat) == list(d.flat)) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_format.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_format.py new file mode 100644 index 0000000000000000000000000000000000000000..0cac8819f5fd3c36f7292da5029b2765d45f56bb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_format.py @@ -0,0 +1,1026 @@ +# doctest +r''' Test the .npy file format. + +Set up: + + >>> import sys + >>> from io import BytesIO + >>> from numpy.lib import format + >>> + >>> scalars = [ + ... np.uint8, + ... np.int8, + ... np.uint16, + ... np.int16, + ... np.uint32, + ... np.int32, + ... np.uint64, + ... np.int64, + ... np.float32, + ... np.float64, + ... np.complex64, + ... np.complex128, + ... object, + ... ] + >>> + >>> basic_arrays = [] + >>> + >>> for scalar in scalars: + ... for endian in '<>': + ... dtype = np.dtype(scalar).newbyteorder(endian) + ... basic = np.arange(15).astype(dtype) + ... basic_arrays.extend([ + ... np.array([], dtype=dtype), + ... np.array(10, dtype=dtype), + ... basic, + ... basic.reshape((3,5)), + ... basic.reshape((3,5)).T, + ... basic.reshape((3,5))[::-1,::2], + ... ]) + ... + >>> + >>> Pdescr = [ + ... ('x', 'i4', (2,)), + ... ('y', 'f8', (2, 2)), + ... ('z', 'u1')] + >>> + >>> + >>> PbufferT = [ + ... ([3,2], [[6.,4.],[6.,4.]], 8), + ... ([4,3], [[7.,5.],[7.,5.]], 9), + ... ] + >>> + >>> + >>> Ndescr = [ + ... ('x', 'i4', (2,)), + ... ('Info', [ + ... ('value', 'c16'), + ... ('y2', 'f8'), + ... ('Info2', [ + ... ('name', 'S2'), + ... ('value', 'c16', (2,)), + ... ('y3', 'f8', (2,)), + ... ('z3', 'u4', (2,))]), + ... ('name', 'S2'), + ... ('z2', 'b1')]), + ... ('color', 'S2'), + ... ('info', [ + ... ('Name', 'U8'), + ... ('Value', 'c16')]), + ... ('y', 'f8', (2, 2)), + ... ('z', 'u1')] + >>> + >>> + >>> NbufferT = [ + ... ([3,2], (6j, 6., ('nn', [6j,4j], [6.,4.], [1,2]), 'NN', True), 'cc', ('NN', 6j), [[6.,4.],[6.,4.]], 8), + ... ([4,3], (7j, 7., ('oo', [7j,5j], [7.,5.], [2,1]), 'OO', False), 'dd', ('OO', 7j), [[7.,5.],[7.,5.]], 9), + ... ] + >>> + >>> + >>> record_arrays = [ + ... np.array(PbufferT, dtype=np.dtype(Pdescr).newbyteorder('<')), + ... np.array(NbufferT, dtype=np.dtype(Ndescr).newbyteorder('<')), + ... np.array(PbufferT, dtype=np.dtype(Pdescr).newbyteorder('>')), + ... np.array(NbufferT, dtype=np.dtype(Ndescr).newbyteorder('>')), + ... ] + +Test the magic string writing. + + >>> format.magic(1, 0) + '\x93NUMPY\x01\x00' + >>> format.magic(0, 0) + '\x93NUMPY\x00\x00' + >>> format.magic(255, 255) + '\x93NUMPY\xff\xff' + >>> format.magic(2, 5) + '\x93NUMPY\x02\x05' + +Test the magic string reading. + + >>> format.read_magic(BytesIO(format.magic(1, 0))) + (1, 0) + >>> format.read_magic(BytesIO(format.magic(0, 0))) + (0, 0) + >>> format.read_magic(BytesIO(format.magic(255, 255))) + (255, 255) + >>> format.read_magic(BytesIO(format.magic(2, 5))) + (2, 5) + +Test the header writing. + + >>> for arr in basic_arrays + record_arrays: + ... f = BytesIO() + ... format.write_array_header_1_0(f, arr) # XXX: arr is not a dict, items gets called on it + ... print(repr(f.getvalue())) + ... + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '|u1', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '|u1', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '|u1', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '|i1', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '|i1', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '|i1', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'u2', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>u2', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>u2', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>u2', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>u2', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>u2', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'i2', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>i2', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>i2', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>i2', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>i2', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>i2', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'u4', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>u4', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>u4', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>u4', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>u4', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>u4', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'i4', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>i4', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>i4', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>i4', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>i4', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>i4', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'u8', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>u8', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>u8', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>u8', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>u8', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>u8', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'i8', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>i8', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>i8', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>i8', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>i8', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>i8', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'f4', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>f4', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>f4', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>f4', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>f4', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>f4', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'f8', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>f8', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>f8', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>f8', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>f8', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>f8', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'c8', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>c8', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>c8', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>c8', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>c8', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>c8', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'c16', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': '>c16', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': '>c16', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': '>c16', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': '>c16', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': '>c16', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': 'O', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (3, 3)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (0,)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': ()} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (15,)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (3, 5)} \n" + "F\x00{'descr': 'O', 'fortran_order': True, 'shape': (5, 3)} \n" + "F\x00{'descr': 'O', 'fortran_order': False, 'shape': (3, 3)} \n" + "v\x00{'descr': [('x', 'i4', (2,)), ('y', '>f8', (2, 2)), ('z', '|u1')],\n 'fortran_order': False,\n 'shape': (2,)} \n" + "\x16\x02{'descr': [('x', '>i4', (2,)),\n ('Info',\n [('value', '>c16'),\n ('y2', '>f8'),\n ('Info2',\n [('name', '|S2'),\n ('value', '>c16', (2,)),\n ('y3', '>f8', (2,)),\n ('z3', '>u4', (2,))]),\n ('name', '|S2'),\n ('z2', '|b1')]),\n ('color', '|S2'),\n ('info', [('Name', '>U8'), ('Value', '>c16')]),\n ('y', '>f8', (2, 2)),\n ('z', '|u1')],\n 'fortran_order': False,\n 'shape': (2,)} \n" +''' +import sys +import os +import warnings +import pytest +from io import BytesIO + +import numpy as np +from numpy.testing import ( + assert_, assert_array_equal, assert_raises, assert_raises_regex, + assert_warns, IS_PYPY, IS_WASM, IS_64BIT + ) +from numpy.testing._private.utils import requires_memory +from numpy.lib import format + + +# Generate some basic arrays to test with. +scalars = [ + np.uint8, + np.int8, + np.uint16, + np.int16, + np.uint32, + np.int32, + np.uint64, + np.int64, + np.float32, + np.float64, + np.complex64, + np.complex128, + object, +] +basic_arrays = [] +for scalar in scalars: + for endian in '<>': + dtype = np.dtype(scalar).newbyteorder(endian) + basic = np.arange(1500).astype(dtype) + basic_arrays.extend([ + # Empty + np.array([], dtype=dtype), + # Rank-0 + np.array(10, dtype=dtype), + # 1-D + basic, + # 2-D C-contiguous + basic.reshape((30, 50)), + # 2-D F-contiguous + basic.reshape((30, 50)).T, + # 2-D non-contiguous + basic.reshape((30, 50))[::-1, ::2], + ]) + +# More complicated record arrays. +# This is the structure of the table used for plain objects: +# +# +-+-+-+ +# |x|y|z| +# +-+-+-+ + +# Structure of a plain array description: +Pdescr = [ + ('x', 'i4', (2,)), + ('y', 'f8', (2, 2)), + ('z', 'u1')] + +# A plain list of tuples with values for testing: +PbufferT = [ + # x y z + ([3, 2], [[6., 4.], [6., 4.]], 8), + ([4, 3], [[7., 5.], [7., 5.]], 9), + ] + + +# This is the structure of the table used for nested objects (DON'T PANIC!): +# +# +-+---------------------------------+-----+----------+-+-+ +# |x|Info |color|info |y|z| +# | +-----+--+----------------+----+--+ +----+-----+ | | +# | |value|y2|Info2 |name|z2| |Name|Value| | | +# | | | +----+-----+--+--+ | | | | | | | +# | | | |name|value|y3|z3| | | | | | | | +# +-+-----+--+----+-----+--+--+----+--+-----+----+-----+-+-+ +# + +# The corresponding nested array description: +Ndescr = [ + ('x', 'i4', (2,)), + ('Info', [ + ('value', 'c16'), + ('y2', 'f8'), + ('Info2', [ + ('name', 'S2'), + ('value', 'c16', (2,)), + ('y3', 'f8', (2,)), + ('z3', 'u4', (2,))]), + ('name', 'S2'), + ('z2', 'b1')]), + ('color', 'S2'), + ('info', [ + ('Name', 'U8'), + ('Value', 'c16')]), + ('y', 'f8', (2, 2)), + ('z', 'u1')] + +NbufferT = [ + # x Info color info y z + # value y2 Info2 name z2 Name Value + # name value y3 z3 + ([3, 2], (6j, 6., ('nn', [6j, 4j], [6., 4.], [1, 2]), 'NN', True), + 'cc', ('NN', 6j), [[6., 4.], [6., 4.]], 8), + ([4, 3], (7j, 7., ('oo', [7j, 5j], [7., 5.], [2, 1]), 'OO', False), + 'dd', ('OO', 7j), [[7., 5.], [7., 5.]], 9), + ] + +record_arrays = [ + np.array(PbufferT, dtype=np.dtype(Pdescr).newbyteorder('<')), + np.array(NbufferT, dtype=np.dtype(Ndescr).newbyteorder('<')), + np.array(PbufferT, dtype=np.dtype(Pdescr).newbyteorder('>')), + np.array(NbufferT, dtype=np.dtype(Ndescr).newbyteorder('>')), + np.zeros(1, dtype=[('c', ('= (3, 12), reason="see gh-23988") +@pytest.mark.xfail(IS_WASM, reason="Emscripten NODEFS has a buggy dup") +def test_python2_python3_interoperability(): + fname = 'win64python2.npy' + path = os.path.join(os.path.dirname(__file__), 'data', fname) + with pytest.warns(UserWarning, match="Reading.*this warning\\."): + data = np.load(path) + assert_array_equal(data, np.ones(2)) + + +def test_pickle_python2_python3(): + # Test that loading object arrays saved on Python 2 works both on + # Python 2 and Python 3 and vice versa + data_dir = os.path.join(os.path.dirname(__file__), 'data') + + expected = np.array([None, range, '\u512a\u826f', + b'\xe4\xb8\x8d\xe8\x89\xaf'], + dtype=object) + + for fname in ['py2-np0-objarr.npy', 'py2-objarr.npy', 'py2-objarr.npz', + 'py3-objarr.npy', 'py3-objarr.npz']: + path = os.path.join(data_dir, fname) + + for encoding in ['bytes', 'latin1']: + data_f = np.load(path, allow_pickle=True, encoding=encoding) + if fname.endswith('.npz'): + data = data_f['x'] + data_f.close() + else: + data = data_f + + if encoding == 'latin1' and fname.startswith('py2'): + assert_(isinstance(data[3], str)) + assert_array_equal(data[:-1], expected[:-1]) + # mojibake occurs + assert_array_equal(data[-1].encode(encoding), expected[-1]) + else: + assert_(isinstance(data[3], bytes)) + assert_array_equal(data, expected) + + if fname.startswith('py2'): + if fname.endswith('.npz'): + data = np.load(path, allow_pickle=True) + assert_raises(UnicodeError, data.__getitem__, 'x') + data.close() + data = np.load(path, allow_pickle=True, fix_imports=False, + encoding='latin1') + assert_raises(ImportError, data.__getitem__, 'x') + data.close() + else: + assert_raises(UnicodeError, np.load, path, + allow_pickle=True) + assert_raises(ImportError, np.load, path, + allow_pickle=True, fix_imports=False, + encoding='latin1') + + +def test_pickle_disallow(tmpdir): + data_dir = os.path.join(os.path.dirname(__file__), 'data') + + path = os.path.join(data_dir, 'py2-objarr.npy') + assert_raises(ValueError, np.load, path, + allow_pickle=False, encoding='latin1') + + path = os.path.join(data_dir, 'py2-objarr.npz') + with np.load(path, allow_pickle=False, encoding='latin1') as f: + assert_raises(ValueError, f.__getitem__, 'x') + + path = os.path.join(tmpdir, 'pickle-disabled.npy') + assert_raises(ValueError, np.save, path, np.array([None], dtype=object), + allow_pickle=False) + +@pytest.mark.parametrize('dt', [ + np.dtype(np.dtype([('a', np.int8), + ('b', np.int16), + ('c', np.int32), + ], align=True), + (3,)), + np.dtype([('x', np.dtype({'names':['a','b'], + 'formats':['i1','i1'], + 'offsets':[0,4], + 'itemsize':8, + }, + (3,)), + (4,), + )]), + np.dtype([('x', + (' 1, a) + assert_array_equal(b, [3, 2, 2, 3, 3]) + + def test_place(self): + # Make sure that non-np.ndarray objects + # raise an error instead of doing nothing + assert_raises(TypeError, place, [1, 2, 3], [True, False], [0, 1]) + + a = np.array([1, 4, 3, 2, 5, 8, 7]) + place(a, [0, 1, 0, 1, 0, 1, 0], [2, 4, 6]) + assert_array_equal(a, [1, 2, 3, 4, 5, 6, 7]) + + place(a, np.zeros(7), []) + assert_array_equal(a, np.arange(1, 8)) + + place(a, [1, 0, 1, 0, 1, 0, 1], [8, 9]) + assert_array_equal(a, [8, 2, 9, 4, 8, 6, 9]) + assert_raises_regex(ValueError, "Cannot insert from an empty array", + lambda: place(a, [0, 0, 0, 0, 0, 1, 0], [])) + + # See Issue #6974 + a = np.array(['12', '34']) + place(a, [0, 1], '9') + assert_array_equal(a, ['12', '9']) + + def test_both(self): + a = rand(10) + mask = a > 0.5 + ac = a.copy() + c = extract(mask, a) + place(a, mask, 0) + place(a, mask, c) + assert_array_equal(a, ac) + + +# _foo1 and _foo2 are used in some tests in TestVectorize. + +def _foo1(x, y=1.0): + return y*math.floor(x) + + +def _foo2(x, y=1.0, z=0.0): + return y*math.floor(x) + z + + +class TestVectorize: + + def test_simple(self): + def addsubtract(a, b): + if a > b: + return a - b + else: + return a + b + + f = vectorize(addsubtract) + r = f([0, 3, 6, 9], [1, 3, 5, 7]) + assert_array_equal(r, [1, 6, 1, 2]) + + def test_scalar(self): + def addsubtract(a, b): + if a > b: + return a - b + else: + return a + b + + f = vectorize(addsubtract) + r = f([0, 3, 6, 9], 5) + assert_array_equal(r, [5, 8, 1, 4]) + + def test_large(self): + x = np.linspace(-3, 2, 10000) + f = vectorize(lambda x: x) + y = f(x) + assert_array_equal(y, x) + + def test_ufunc(self): + f = vectorize(math.cos) + args = np.array([0, 0.5 * np.pi, np.pi, 1.5 * np.pi, 2 * np.pi]) + r1 = f(args) + r2 = np.cos(args) + assert_array_almost_equal(r1, r2) + + def test_keywords(self): + + def foo(a, b=1): + return a + b + + f = vectorize(foo) + args = np.array([1, 2, 3]) + r1 = f(args) + r2 = np.array([2, 3, 4]) + assert_array_equal(r1, r2) + r1 = f(args, 2) + r2 = np.array([3, 4, 5]) + assert_array_equal(r1, r2) + + def test_keywords_with_otypes_order1(self): + # gh-1620: The second call of f would crash with + # `ValueError: invalid number of arguments`. + f = vectorize(_foo1, otypes=[float]) + # We're testing the caching of ufuncs by vectorize, so the order + # of these function calls is an important part of the test. + r1 = f(np.arange(3.0), 1.0) + r2 = f(np.arange(3.0)) + assert_array_equal(r1, r2) + + def test_keywords_with_otypes_order2(self): + # gh-1620: The second call of f would crash with + # `ValueError: non-broadcastable output operand with shape () + # doesn't match the broadcast shape (3,)`. + f = vectorize(_foo1, otypes=[float]) + # We're testing the caching of ufuncs by vectorize, so the order + # of these function calls is an important part of the test. + r1 = f(np.arange(3.0)) + r2 = f(np.arange(3.0), 1.0) + assert_array_equal(r1, r2) + + def test_keywords_with_otypes_order3(self): + # gh-1620: The third call of f would crash with + # `ValueError: invalid number of arguments`. + f = vectorize(_foo1, otypes=[float]) + # We're testing the caching of ufuncs by vectorize, so the order + # of these function calls is an important part of the test. + r1 = f(np.arange(3.0)) + r2 = f(np.arange(3.0), y=1.0) + r3 = f(np.arange(3.0)) + assert_array_equal(r1, r2) + assert_array_equal(r1, r3) + + def test_keywords_with_otypes_several_kwd_args1(self): + # gh-1620 Make sure different uses of keyword arguments + # don't break the vectorized function. + f = vectorize(_foo2, otypes=[float]) + # We're testing the caching of ufuncs by vectorize, so the order + # of these function calls is an important part of the test. + r1 = f(10.4, z=100) + r2 = f(10.4, y=-1) + r3 = f(10.4) + assert_equal(r1, _foo2(10.4, z=100)) + assert_equal(r2, _foo2(10.4, y=-1)) + assert_equal(r3, _foo2(10.4)) + + def test_keywords_with_otypes_several_kwd_args2(self): + # gh-1620 Make sure different uses of keyword arguments + # don't break the vectorized function. + f = vectorize(_foo2, otypes=[float]) + # We're testing the caching of ufuncs by vectorize, so the order + # of these function calls is an important part of the test. + r1 = f(z=100, x=10.4, y=-1) + r2 = f(1, 2, 3) + assert_equal(r1, _foo2(z=100, x=10.4, y=-1)) + assert_equal(r2, _foo2(1, 2, 3)) + + def test_keywords_no_func_code(self): + # This needs to test a function that has keywords but + # no func_code attribute, since otherwise vectorize will + # inspect the func_code. + import random + try: + vectorize(random.randrange) # Should succeed + except Exception: + raise AssertionError + + def test_keywords2_ticket_2100(self): + # Test kwarg support: enhancement ticket 2100 + + def foo(a, b=1): + return a + b + + f = vectorize(foo) + args = np.array([1, 2, 3]) + r1 = f(a=args) + r2 = np.array([2, 3, 4]) + assert_array_equal(r1, r2) + r1 = f(b=1, a=args) + assert_array_equal(r1, r2) + r1 = f(args, b=2) + r2 = np.array([3, 4, 5]) + assert_array_equal(r1, r2) + + def test_keywords3_ticket_2100(self): + # Test excluded with mixed positional and kwargs: ticket 2100 + def mypolyval(x, p): + _p = list(p) + res = _p.pop(0) + while _p: + res = res * x + _p.pop(0) + return res + + vpolyval = np.vectorize(mypolyval, excluded=['p', 1]) + ans = [3, 6] + assert_array_equal(ans, vpolyval(x=[0, 1], p=[1, 2, 3])) + assert_array_equal(ans, vpolyval([0, 1], p=[1, 2, 3])) + assert_array_equal(ans, vpolyval([0, 1], [1, 2, 3])) + + def test_keywords4_ticket_2100(self): + # Test vectorizing function with no positional args. + @vectorize + def f(**kw): + res = 1.0 + for _k in kw: + res *= kw[_k] + return res + + assert_array_equal(f(a=[1, 2], b=[3, 4]), [3, 8]) + + def test_keywords5_ticket_2100(self): + # Test vectorizing function with no kwargs args. + @vectorize + def f(*v): + return np.prod(v) + + assert_array_equal(f([1, 2], [3, 4]), [3, 8]) + + def test_coverage1_ticket_2100(self): + def foo(): + return 1 + + f = vectorize(foo) + assert_array_equal(f(), 1) + + def test_assigning_docstring(self): + def foo(x): + """Original documentation""" + return x + + f = vectorize(foo) + assert_equal(f.__doc__, foo.__doc__) + + doc = "Provided documentation" + f = vectorize(foo, doc=doc) + assert_equal(f.__doc__, doc) + + def test_UnboundMethod_ticket_1156(self): + # Regression test for issue 1156 + class Foo: + b = 2 + + def bar(self, a): + return a ** self.b + + assert_array_equal(vectorize(Foo().bar)(np.arange(9)), + np.arange(9) ** 2) + assert_array_equal(vectorize(Foo.bar)(Foo(), np.arange(9)), + np.arange(9) ** 2) + + def test_execution_order_ticket_1487(self): + # Regression test for dependence on execution order: issue 1487 + f1 = vectorize(lambda x: x) + res1a = f1(np.arange(3)) + res1b = f1(np.arange(0.1, 3)) + f2 = vectorize(lambda x: x) + res2b = f2(np.arange(0.1, 3)) + res2a = f2(np.arange(3)) + assert_equal(res1a, res2a) + assert_equal(res1b, res2b) + + def test_string_ticket_1892(self): + # Test vectorization over strings: issue 1892. + f = np.vectorize(lambda x: x) + s = '0123456789' * 10 + assert_equal(s, f(s)) + + def test_cache(self): + # Ensure that vectorized func called exactly once per argument. + _calls = [0] + + @vectorize + def f(x): + _calls[0] += 1 + return x ** 2 + + f.cache = True + x = np.arange(5) + assert_array_equal(f(x), x * x) + assert_equal(_calls[0], len(x)) + + def test_otypes(self): + f = np.vectorize(lambda x: x) + f.otypes = 'i' + x = np.arange(5) + assert_array_equal(f(x), x) + + def test_parse_gufunc_signature(self): + assert_equal(nfb._parse_gufunc_signature('(x)->()'), ([('x',)], [()])) + assert_equal(nfb._parse_gufunc_signature('(x,y)->()'), + ([('x', 'y')], [()])) + assert_equal(nfb._parse_gufunc_signature('(x),(y)->()'), + ([('x',), ('y',)], [()])) + assert_equal(nfb._parse_gufunc_signature('(x)->(y)'), + ([('x',)], [('y',)])) + assert_equal(nfb._parse_gufunc_signature('(x)->(y),()'), + ([('x',)], [('y',), ()])) + assert_equal(nfb._parse_gufunc_signature('(),(a,b,c),(d)->(d,e)'), + ([(), ('a', 'b', 'c'), ('d',)], [('d', 'e')])) + + # Tests to check if whitespaces are ignored + assert_equal(nfb._parse_gufunc_signature('(x )->()'), ([('x',)], [()])) + assert_equal(nfb._parse_gufunc_signature('( x , y )->( )'), + ([('x', 'y')], [()])) + assert_equal(nfb._parse_gufunc_signature('(x),( y) ->()'), + ([('x',), ('y',)], [()])) + assert_equal(nfb._parse_gufunc_signature('( x)-> (y ) '), + ([('x',)], [('y',)])) + assert_equal(nfb._parse_gufunc_signature(' (x)->( y),( )'), + ([('x',)], [('y',), ()])) + assert_equal(nfb._parse_gufunc_signature( + '( ), ( a, b,c ) ,( d) -> (d , e)'), + ([(), ('a', 'b', 'c'), ('d',)], [('d', 'e')])) + + with assert_raises(ValueError): + nfb._parse_gufunc_signature('(x)(y)->()') + with assert_raises(ValueError): + nfb._parse_gufunc_signature('(x),(y)->') + with assert_raises(ValueError): + nfb._parse_gufunc_signature('((x))->(x)') + + def test_signature_simple(self): + def addsubtract(a, b): + if a > b: + return a - b + else: + return a + b + + f = vectorize(addsubtract, signature='(),()->()') + r = f([0, 3, 6, 9], [1, 3, 5, 7]) + assert_array_equal(r, [1, 6, 1, 2]) + + def test_signature_mean_last(self): + def mean(a): + return a.mean() + + f = vectorize(mean, signature='(n)->()') + r = f([[1, 3], [2, 4]]) + assert_array_equal(r, [2, 3]) + + def test_signature_center(self): + def center(a): + return a - a.mean() + + f = vectorize(center, signature='(n)->(n)') + r = f([[1, 3], [2, 4]]) + assert_array_equal(r, [[-1, 1], [-1, 1]]) + + def test_signature_two_outputs(self): + f = vectorize(lambda x: (x, x), signature='()->(),()') + r = f([1, 2, 3]) + assert_(isinstance(r, tuple) and len(r) == 2) + assert_array_equal(r[0], [1, 2, 3]) + assert_array_equal(r[1], [1, 2, 3]) + + def test_signature_outer(self): + f = vectorize(np.outer, signature='(a),(b)->(a,b)') + r = f([1, 2], [1, 2, 3]) + assert_array_equal(r, [[1, 2, 3], [2, 4, 6]]) + + r = f([[[1, 2]]], [1, 2, 3]) + assert_array_equal(r, [[[[1, 2, 3], [2, 4, 6]]]]) + + r = f([[1, 0], [2, 0]], [1, 2, 3]) + assert_array_equal(r, [[[1, 2, 3], [0, 0, 0]], + [[2, 4, 6], [0, 0, 0]]]) + + r = f([1, 2], [[1, 2, 3], [0, 0, 0]]) + assert_array_equal(r, [[[1, 2, 3], [2, 4, 6]], + [[0, 0, 0], [0, 0, 0]]]) + + def test_signature_computed_size(self): + f = vectorize(lambda x: x[:-1], signature='(n)->(m)') + r = f([1, 2, 3]) + assert_array_equal(r, [1, 2]) + + r = f([[1, 2, 3], [2, 3, 4]]) + assert_array_equal(r, [[1, 2], [2, 3]]) + + def test_signature_excluded(self): + + def foo(a, b=1): + return a + b + + f = vectorize(foo, signature='()->()', excluded={'b'}) + assert_array_equal(f([1, 2, 3]), [2, 3, 4]) + assert_array_equal(f([1, 2, 3], b=0), [1, 2, 3]) + + def test_signature_otypes(self): + f = vectorize(lambda x: x, signature='(n)->(n)', otypes=['float64']) + r = f([1, 2, 3]) + assert_equal(r.dtype, np.dtype('float64')) + assert_array_equal(r, [1, 2, 3]) + + def test_signature_invalid_inputs(self): + f = vectorize(operator.add, signature='(n),(n)->(n)') + with assert_raises_regex(TypeError, 'wrong number of positional'): + f([1, 2]) + with assert_raises_regex( + ValueError, 'does not have enough dimensions'): + f(1, 2) + with assert_raises_regex( + ValueError, 'inconsistent size for core dimension'): + f([1, 2], [1, 2, 3]) + + f = vectorize(operator.add, signature='()->()') + with assert_raises_regex(TypeError, 'wrong number of positional'): + f(1, 2) + + def test_signature_invalid_outputs(self): + + f = vectorize(lambda x: x[:-1], signature='(n)->(n)') + with assert_raises_regex( + ValueError, 'inconsistent size for core dimension'): + f([1, 2, 3]) + + f = vectorize(lambda x: x, signature='()->(),()') + with assert_raises_regex(ValueError, 'wrong number of outputs'): + f(1) + + f = vectorize(lambda x: (x, x), signature='()->()') + with assert_raises_regex(ValueError, 'wrong number of outputs'): + f([1, 2]) + + def test_size_zero_output(self): + # see issue 5868 + f = np.vectorize(lambda x: x) + x = np.zeros([0, 5], dtype=int) + with assert_raises_regex(ValueError, 'otypes'): + f(x) + + f.otypes = 'i' + assert_array_equal(f(x), x) + + f = np.vectorize(lambda x: x, signature='()->()') + with assert_raises_regex(ValueError, 'otypes'): + f(x) + + f = np.vectorize(lambda x: x, signature='()->()', otypes='i') + assert_array_equal(f(x), x) + + f = np.vectorize(lambda x: x, signature='(n)->(n)', otypes='i') + assert_array_equal(f(x), x) + + f = np.vectorize(lambda x: x, signature='(n)->(n)') + assert_array_equal(f(x.T), x.T) + + f = np.vectorize(lambda x: [x], signature='()->(n)', otypes='i') + with assert_raises_regex(ValueError, 'new output dimensions'): + f(x) + + def test_subclasses(self): + class subclass(np.ndarray): + pass + + m = np.array([[1., 0., 0.], + [0., 0., 1.], + [0., 1., 0.]]).view(subclass) + v = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]).view(subclass) + # generalized (gufunc) + matvec = np.vectorize(np.matmul, signature='(m,m),(m)->(m)') + r = matvec(m, v) + assert_equal(type(r), subclass) + assert_equal(r, [[1., 3., 2.], [4., 6., 5.], [7., 9., 8.]]) + + # element-wise (ufunc) + mult = np.vectorize(lambda x, y: x*y) + r = mult(m, v) + assert_equal(type(r), subclass) + assert_equal(r, m * v) + + def test_name(self): + #See gh-23021 + @np.vectorize + def f2(a, b): + return a + b + + assert f2.__name__ == 'f2' + + def test_decorator(self): + @vectorize + def addsubtract(a, b): + if a > b: + return a - b + else: + return a + b + + r = addsubtract([0, 3, 6, 9], [1, 3, 5, 7]) + assert_array_equal(r, [1, 6, 1, 2]) + + def test_docstring(self): + @vectorize + def f(x): + """Docstring""" + return x + + if sys.flags.optimize < 2: + assert f.__doc__ == "Docstring" + + def test_partial(self): + def foo(x, y): + return x + y + + bar = partial(foo, 3) + vbar = np.vectorize(bar) + assert vbar(1) == 4 + + def test_signature_otypes_decorator(self): + @vectorize(signature='(n)->(n)', otypes=['float64']) + def f(x): + return x + + r = f([1, 2, 3]) + assert_equal(r.dtype, np.dtype('float64')) + assert_array_equal(r, [1, 2, 3]) + assert f.__name__ == 'f' + + def test_bad_input(self): + with assert_raises(TypeError): + A = np.vectorize(pyfunc = 3) + + def test_no_keywords(self): + with assert_raises(TypeError): + @np.vectorize("string") + def foo(): + return "bar" + + def test_positional_regression_9477(self): + # This supplies the first keyword argument as a positional, + # to ensure that they are still properly forwarded after the + # enhancement for #9477 + f = vectorize((lambda x: x), ['float64']) + r = f([2]) + assert_equal(r.dtype, np.dtype('float64')) + + def test_datetime_conversion(self): + otype = "datetime64[ns]" + arr = np.array(['2024-01-01', '2024-01-02', '2024-01-03'], + dtype='datetime64[ns]') + assert_array_equal(np.vectorize(lambda x: x, signature="(i)->(j)", + otypes=[otype])(arr), arr) + + +class TestLeaks: + class A: + iters = 20 + + def bound(self, *args): + return 0 + + @staticmethod + def unbound(*args): + return 0 + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + @pytest.mark.skipif(NOGIL_BUILD, + reason=("Functions are immortalized if a thread is " + "launched, making this test flaky")) + @pytest.mark.parametrize('name, incr', [ + ('bound', A.iters), + ('unbound', 0), + ]) + def test_frompyfunc_leaks(self, name, incr): + # exposed in gh-11867 as np.vectorized, but the problem stems from + # frompyfunc. + # class.attribute = np.frompyfunc() creates a + # reference cycle if is a bound class method. It requires a + # gc collection cycle to break the cycle (on CPython 3) + import gc + A_func = getattr(self.A, name) + gc.disable() + try: + refcount = sys.getrefcount(A_func) + for i in range(self.A.iters): + a = self.A() + a.f = np.frompyfunc(getattr(a, name), 1, 1) + out = a.f(np.arange(10)) + a = None + # A.func is part of a reference cycle if incr is non-zero + assert_equal(sys.getrefcount(A_func), refcount + incr) + for i in range(5): + gc.collect() + assert_equal(sys.getrefcount(A_func), refcount) + finally: + gc.enable() + + +class TestDigitize: + + def test_forward(self): + x = np.arange(-6, 5) + bins = np.arange(-5, 5) + assert_array_equal(digitize(x, bins), np.arange(11)) + + def test_reverse(self): + x = np.arange(5, -6, -1) + bins = np.arange(5, -5, -1) + assert_array_equal(digitize(x, bins), np.arange(11)) + + def test_random(self): + x = rand(10) + bin = np.linspace(x.min(), x.max(), 10) + assert_(np.all(digitize(x, bin) != 0)) + + def test_right_basic(self): + x = [1, 5, 4, 10, 8, 11, 0] + bins = [1, 5, 10] + default_answer = [1, 2, 1, 3, 2, 3, 0] + assert_array_equal(digitize(x, bins), default_answer) + right_answer = [0, 1, 1, 2, 2, 3, 0] + assert_array_equal(digitize(x, bins, True), right_answer) + + def test_right_open(self): + x = np.arange(-6, 5) + bins = np.arange(-6, 4) + assert_array_equal(digitize(x, bins, True), np.arange(11)) + + def test_right_open_reverse(self): + x = np.arange(5, -6, -1) + bins = np.arange(4, -6, -1) + assert_array_equal(digitize(x, bins, True), np.arange(11)) + + def test_right_open_random(self): + x = rand(10) + bins = np.linspace(x.min(), x.max(), 10) + assert_(np.all(digitize(x, bins, True) != 10)) + + def test_monotonic(self): + x = [-1, 0, 1, 2] + bins = [0, 0, 1] + assert_array_equal(digitize(x, bins, False), [0, 2, 3, 3]) + assert_array_equal(digitize(x, bins, True), [0, 0, 2, 3]) + bins = [1, 1, 0] + assert_array_equal(digitize(x, bins, False), [3, 2, 0, 0]) + assert_array_equal(digitize(x, bins, True), [3, 3, 2, 0]) + bins = [1, 1, 1, 1] + assert_array_equal(digitize(x, bins, False), [0, 0, 4, 4]) + assert_array_equal(digitize(x, bins, True), [0, 0, 0, 4]) + bins = [0, 0, 1, 0] + assert_raises(ValueError, digitize, x, bins) + bins = [1, 1, 0, 1] + assert_raises(ValueError, digitize, x, bins) + + def test_casting_error(self): + x = [1, 2, 3 + 1.j] + bins = [1, 2, 3] + assert_raises(TypeError, digitize, x, bins) + x, bins = bins, x + assert_raises(TypeError, digitize, x, bins) + + def test_return_type(self): + # Functions returning indices should always return base ndarrays + class A(np.ndarray): + pass + a = np.arange(5).view(A) + b = np.arange(1, 3).view(A) + assert_(not isinstance(digitize(b, a, False), A)) + assert_(not isinstance(digitize(b, a, True), A)) + + def test_large_integers_increasing(self): + # gh-11022 + x = 2**54 # loses precision in a float + assert_equal(np.digitize(x, [x - 1, x + 1]), 1) + + @pytest.mark.xfail( + reason="gh-11022: np._core.multiarray._monoticity loses precision") + def test_large_integers_decreasing(self): + # gh-11022 + x = 2**54 # loses precision in a float + assert_equal(np.digitize(x, [x + 1, x - 1]), 1) + + +class TestUnwrap: + + def test_simple(self): + # check that unwrap removes jumps greater that 2*pi + assert_array_equal(unwrap([1, 1 + 2 * np.pi]), [1, 1]) + # check that unwrap maintains continuity + assert_(np.all(diff(unwrap(rand(10) * 100)) < np.pi)) + + def test_period(self): + # check that unwrap removes jumps greater that 255 + assert_array_equal(unwrap([1, 1 + 256], period=255), [1, 2]) + # check that unwrap maintains continuity + assert_(np.all(diff(unwrap(rand(10) * 1000, period=255)) < 255)) + # check simple case + simple_seq = np.array([0, 75, 150, 225, 300]) + wrap_seq = np.mod(simple_seq, 255) + assert_array_equal(unwrap(wrap_seq, period=255), simple_seq) + # check custom discont value + uneven_seq = np.array([0, 75, 150, 225, 300, 430]) + wrap_uneven = np.mod(uneven_seq, 250) + no_discont = unwrap(wrap_uneven, period=250) + assert_array_equal(no_discont, [0, 75, 150, 225, 300, 180]) + sm_discont = unwrap(wrap_uneven, period=250, discont=140) + assert_array_equal(sm_discont, [0, 75, 150, 225, 300, 430]) + assert sm_discont.dtype == wrap_uneven.dtype + + +@pytest.mark.parametrize( + "dtype", "O" + np.typecodes["AllInteger"] + np.typecodes["Float"] +) +@pytest.mark.parametrize("M", [0, 1, 10]) +class TestFilterwindows: + + def test_hanning(self, dtype: str, M: int) -> None: + scalar = np.array(M, dtype=dtype)[()] + + w = hanning(scalar) + if dtype == "O": + ref_dtype = np.float64 + else: + ref_dtype = np.result_type(scalar.dtype, np.float64) + assert w.dtype == ref_dtype + + # check symmetry + assert_equal(w, flipud(w)) + + # check known value + if scalar < 1: + assert_array_equal(w, np.array([])) + elif scalar == 1: + assert_array_equal(w, np.ones(1)) + else: + assert_almost_equal(np.sum(w, axis=0), 4.500, 4) + + def test_hamming(self, dtype: str, M: int) -> None: + scalar = np.array(M, dtype=dtype)[()] + + w = hamming(scalar) + if dtype == "O": + ref_dtype = np.float64 + else: + ref_dtype = np.result_type(scalar.dtype, np.float64) + assert w.dtype == ref_dtype + + # check symmetry + assert_equal(w, flipud(w)) + + # check known value + if scalar < 1: + assert_array_equal(w, np.array([])) + elif scalar == 1: + assert_array_equal(w, np.ones(1)) + else: + assert_almost_equal(np.sum(w, axis=0), 4.9400, 4) + + def test_bartlett(self, dtype: str, M: int) -> None: + scalar = np.array(M, dtype=dtype)[()] + + w = bartlett(scalar) + if dtype == "O": + ref_dtype = np.float64 + else: + ref_dtype = np.result_type(scalar.dtype, np.float64) + assert w.dtype == ref_dtype + + # check symmetry + assert_equal(w, flipud(w)) + + # check known value + if scalar < 1: + assert_array_equal(w, np.array([])) + elif scalar == 1: + assert_array_equal(w, np.ones(1)) + else: + assert_almost_equal(np.sum(w, axis=0), 4.4444, 4) + + def test_blackman(self, dtype: str, M: int) -> None: + scalar = np.array(M, dtype=dtype)[()] + + w = blackman(scalar) + if dtype == "O": + ref_dtype = np.float64 + else: + ref_dtype = np.result_type(scalar.dtype, np.float64) + assert w.dtype == ref_dtype + + # check symmetry + assert_equal(w, flipud(w)) + + # check known value + if scalar < 1: + assert_array_equal(w, np.array([])) + elif scalar == 1: + assert_array_equal(w, np.ones(1)) + else: + assert_almost_equal(np.sum(w, axis=0), 3.7800, 4) + + def test_kaiser(self, dtype: str, M: int) -> None: + scalar = np.array(M, dtype=dtype)[()] + + w = kaiser(scalar, 0) + if dtype == "O": + ref_dtype = np.float64 + else: + ref_dtype = np.result_type(scalar.dtype, np.float64) + assert w.dtype == ref_dtype + + # check symmetry + assert_equal(w, flipud(w)) + + # check known value + if scalar < 1: + assert_array_equal(w, np.array([])) + elif scalar == 1: + assert_array_equal(w, np.ones(1)) + else: + assert_almost_equal(np.sum(w, axis=0), 10, 15) + + +class TestTrapezoid: + + def test_simple(self): + x = np.arange(-10, 10, .1) + r = trapezoid(np.exp(-.5 * x ** 2) / np.sqrt(2 * np.pi), dx=0.1) + # check integral of normal equals 1 + assert_almost_equal(r, 1, 7) + + def test_ndim(self): + x = np.linspace(0, 1, 3) + y = np.linspace(0, 2, 8) + z = np.linspace(0, 3, 13) + + wx = np.ones_like(x) * (x[1] - x[0]) + wx[0] /= 2 + wx[-1] /= 2 + wy = np.ones_like(y) * (y[1] - y[0]) + wy[0] /= 2 + wy[-1] /= 2 + wz = np.ones_like(z) * (z[1] - z[0]) + wz[0] /= 2 + wz[-1] /= 2 + + q = x[:, None, None] + y[None,:, None] + z[None, None,:] + + qx = (q * wx[:, None, None]).sum(axis=0) + qy = (q * wy[None, :, None]).sum(axis=1) + qz = (q * wz[None, None, :]).sum(axis=2) + + # n-d `x` + r = trapezoid(q, x=x[:, None, None], axis=0) + assert_almost_equal(r, qx) + r = trapezoid(q, x=y[None, :, None], axis=1) + assert_almost_equal(r, qy) + r = trapezoid(q, x=z[None, None, :], axis=2) + assert_almost_equal(r, qz) + + # 1-d `x` + r = trapezoid(q, x=x, axis=0) + assert_almost_equal(r, qx) + r = trapezoid(q, x=y, axis=1) + assert_almost_equal(r, qy) + r = trapezoid(q, x=z, axis=2) + assert_almost_equal(r, qz) + + def test_masked(self): + # Testing that masked arrays behave as if the function is 0 where + # masked + x = np.arange(5) + y = x * x + mask = x == 2 + ym = np.ma.array(y, mask=mask) + r = 13.0 # sum(0.5 * (0 + 1) * 1.0 + 0.5 * (9 + 16)) + assert_almost_equal(trapezoid(ym, x), r) + + xm = np.ma.array(x, mask=mask) + assert_almost_equal(trapezoid(ym, xm), r) + + xm = np.ma.array(x, mask=mask) + assert_almost_equal(trapezoid(y, xm), r) + + +class TestSinc: + + def test_simple(self): + assert_(sinc(0) == 1) + w = sinc(np.linspace(-1, 1, 100)) + # check symmetry + assert_array_almost_equal(w, flipud(w), 7) + + def test_array_like(self): + x = [0, 0.5] + y1 = sinc(np.array(x)) + y2 = sinc(list(x)) + y3 = sinc(tuple(x)) + assert_array_equal(y1, y2) + assert_array_equal(y1, y3) + + +class TestUnique: + + def test_simple(self): + x = np.array([4, 3, 2, 1, 1, 2, 3, 4, 0]) + assert_(np.all(unique(x) == [0, 1, 2, 3, 4])) + assert_(unique(np.array([1, 1, 1, 1, 1])) == np.array([1])) + x = ['widget', 'ham', 'foo', 'bar', 'foo', 'ham'] + assert_(np.all(unique(x) == ['bar', 'foo', 'ham', 'widget'])) + x = np.array([5 + 6j, 1 + 1j, 1 + 10j, 10, 5 + 6j]) + assert_(np.all(unique(x) == [1 + 1j, 1 + 10j, 5 + 6j, 10])) + + +class TestCheckFinite: + + def test_simple(self): + a = [1, 2, 3] + b = [1, 2, np.inf] + c = [1, 2, np.nan] + np.asarray_chkfinite(a) + assert_raises(ValueError, np.asarray_chkfinite, b) + assert_raises(ValueError, np.asarray_chkfinite, c) + + def test_dtype_order(self): + # Regression test for missing dtype and order arguments + a = [1, 2, 3] + a = np.asarray_chkfinite(a, order='F', dtype=np.float64) + assert_(a.dtype == np.float64) + + +class TestCorrCoef: + A = np.array( + [[0.15391142, 0.18045767, 0.14197213], + [0.70461506, 0.96474128, 0.27906989], + [0.9297531, 0.32296769, 0.19267156]]) + B = np.array( + [[0.10377691, 0.5417086, 0.49807457], + [0.82872117, 0.77801674, 0.39226705], + [0.9314666, 0.66800209, 0.03538394]]) + res1 = np.array( + [[1., 0.9379533, -0.04931983], + [0.9379533, 1., 0.30007991], + [-0.04931983, 0.30007991, 1.]]) + res2 = np.array( + [[1., 0.9379533, -0.04931983, 0.30151751, 0.66318558, 0.51532523], + [0.9379533, 1., 0.30007991, -0.04781421, 0.88157256, 0.78052386], + [-0.04931983, 0.30007991, 1., -0.96717111, 0.71483595, 0.83053601], + [0.30151751, -0.04781421, -0.96717111, 1., -0.51366032, -0.66173113], + [0.66318558, 0.88157256, 0.71483595, -0.51366032, 1., 0.98317823], + [0.51532523, 0.78052386, 0.83053601, -0.66173113, 0.98317823, 1.]]) + + def test_non_array(self): + assert_almost_equal(np.corrcoef([0, 1, 0], [1, 0, 1]), + [[1., -1.], [-1., 1.]]) + + def test_simple(self): + tgt1 = corrcoef(self.A) + assert_almost_equal(tgt1, self.res1) + assert_(np.all(np.abs(tgt1) <= 1.0)) + + tgt2 = corrcoef(self.A, self.B) + assert_almost_equal(tgt2, self.res2) + assert_(np.all(np.abs(tgt2) <= 1.0)) + + def test_ddof(self): + # ddof raises DeprecationWarning + with suppress_warnings() as sup: + warnings.simplefilter("always") + assert_warns(DeprecationWarning, corrcoef, self.A, ddof=-1) + sup.filter(DeprecationWarning) + # ddof has no or negligible effect on the function + assert_almost_equal(corrcoef(self.A, ddof=-1), self.res1) + assert_almost_equal(corrcoef(self.A, self.B, ddof=-1), self.res2) + assert_almost_equal(corrcoef(self.A, ddof=3), self.res1) + assert_almost_equal(corrcoef(self.A, self.B, ddof=3), self.res2) + + def test_bias(self): + # bias raises DeprecationWarning + with suppress_warnings() as sup: + warnings.simplefilter("always") + assert_warns(DeprecationWarning, corrcoef, self.A, self.B, 1, 0) + assert_warns(DeprecationWarning, corrcoef, self.A, bias=0) + sup.filter(DeprecationWarning) + # bias has no or negligible effect on the function + assert_almost_equal(corrcoef(self.A, bias=1), self.res1) + + def test_complex(self): + x = np.array([[1, 2, 3], [1j, 2j, 3j]]) + res = corrcoef(x) + tgt = np.array([[1., -1.j], [1.j, 1.]]) + assert_allclose(res, tgt) + assert_(np.all(np.abs(res) <= 1.0)) + + def test_xy(self): + x = np.array([[1, 2, 3]]) + y = np.array([[1j, 2j, 3j]]) + assert_allclose(np.corrcoef(x, y), np.array([[1., -1.j], [1.j, 1.]])) + + def test_empty(self): + with warnings.catch_warnings(record=True): + warnings.simplefilter('always', RuntimeWarning) + assert_array_equal(corrcoef(np.array([])), np.nan) + assert_array_equal(corrcoef(np.array([]).reshape(0, 2)), + np.array([]).reshape(0, 0)) + assert_array_equal(corrcoef(np.array([]).reshape(2, 0)), + np.array([[np.nan, np.nan], [np.nan, np.nan]])) + + def test_extreme(self): + x = [[1e-100, 1e100], [1e100, 1e-100]] + with np.errstate(all='raise'): + c = corrcoef(x) + assert_array_almost_equal(c, np.array([[1., -1.], [-1., 1.]])) + assert_(np.all(np.abs(c) <= 1.0)) + + @pytest.mark.parametrize("test_type", [np.half, np.single, np.double, np.longdouble]) + def test_corrcoef_dtype(self, test_type): + cast_A = self.A.astype(test_type) + res = corrcoef(cast_A, dtype=test_type) + assert test_type == res.dtype + + +class TestCov: + x1 = np.array([[0, 2], [1, 1], [2, 0]]).T + res1 = np.array([[1., -1.], [-1., 1.]]) + x2 = np.array([0.0, 1.0, 2.0], ndmin=2) + frequencies = np.array([1, 4, 1]) + x2_repeats = np.array([[0.0], [1.0], [1.0], [1.0], [1.0], [2.0]]).T + res2 = np.array([[0.4, -0.4], [-0.4, 0.4]]) + unit_frequencies = np.ones(3, dtype=np.int_) + weights = np.array([1.0, 4.0, 1.0]) + res3 = np.array([[2. / 3., -2. / 3.], [-2. / 3., 2. / 3.]]) + unit_weights = np.ones(3) + x3 = np.array([0.3942, 0.5969, 0.7730, 0.9918, 0.7964]) + + def test_basic(self): + assert_allclose(cov(self.x1), self.res1) + + def test_complex(self): + x = np.array([[1, 2, 3], [1j, 2j, 3j]]) + res = np.array([[1., -1.j], [1.j, 1.]]) + assert_allclose(cov(x), res) + assert_allclose(cov(x, aweights=np.ones(3)), res) + + def test_xy(self): + x = np.array([[1, 2, 3]]) + y = np.array([[1j, 2j, 3j]]) + assert_allclose(cov(x, y), np.array([[1., -1.j], [1.j, 1.]])) + + def test_empty(self): + with warnings.catch_warnings(record=True): + warnings.simplefilter('always', RuntimeWarning) + assert_array_equal(cov(np.array([])), np.nan) + assert_array_equal(cov(np.array([]).reshape(0, 2)), + np.array([]).reshape(0, 0)) + assert_array_equal(cov(np.array([]).reshape(2, 0)), + np.array([[np.nan, np.nan], [np.nan, np.nan]])) + + def test_wrong_ddof(self): + with warnings.catch_warnings(record=True): + warnings.simplefilter('always', RuntimeWarning) + assert_array_equal(cov(self.x1, ddof=5), + np.array([[np.inf, -np.inf], + [-np.inf, np.inf]])) + + def test_1D_rowvar(self): + assert_allclose(cov(self.x3), cov(self.x3, rowvar=False)) + y = np.array([0.0780, 0.3107, 0.2111, 0.0334, 0.8501]) + assert_allclose(cov(self.x3, y), cov(self.x3, y, rowvar=False)) + + def test_1D_variance(self): + assert_allclose(cov(self.x3, ddof=1), np.var(self.x3, ddof=1)) + + def test_fweights(self): + assert_allclose(cov(self.x2, fweights=self.frequencies), + cov(self.x2_repeats)) + assert_allclose(cov(self.x1, fweights=self.frequencies), + self.res2) + assert_allclose(cov(self.x1, fweights=self.unit_frequencies), + self.res1) + nonint = self.frequencies + 0.5 + assert_raises(TypeError, cov, self.x1, fweights=nonint) + f = np.ones((2, 3), dtype=np.int_) + assert_raises(RuntimeError, cov, self.x1, fweights=f) + f = np.ones(2, dtype=np.int_) + assert_raises(RuntimeError, cov, self.x1, fweights=f) + f = -1 * np.ones(3, dtype=np.int_) + assert_raises(ValueError, cov, self.x1, fweights=f) + + def test_aweights(self): + assert_allclose(cov(self.x1, aweights=self.weights), self.res3) + assert_allclose(cov(self.x1, aweights=3.0 * self.weights), + cov(self.x1, aweights=self.weights)) + assert_allclose(cov(self.x1, aweights=self.unit_weights), self.res1) + w = np.ones((2, 3)) + assert_raises(RuntimeError, cov, self.x1, aweights=w) + w = np.ones(2) + assert_raises(RuntimeError, cov, self.x1, aweights=w) + w = -1.0 * np.ones(3) + assert_raises(ValueError, cov, self.x1, aweights=w) + + def test_unit_fweights_and_aweights(self): + assert_allclose(cov(self.x2, fweights=self.frequencies, + aweights=self.unit_weights), + cov(self.x2_repeats)) + assert_allclose(cov(self.x1, fweights=self.frequencies, + aweights=self.unit_weights), + self.res2) + assert_allclose(cov(self.x1, fweights=self.unit_frequencies, + aweights=self.unit_weights), + self.res1) + assert_allclose(cov(self.x1, fweights=self.unit_frequencies, + aweights=self.weights), + self.res3) + assert_allclose(cov(self.x1, fweights=self.unit_frequencies, + aweights=3.0 * self.weights), + cov(self.x1, aweights=self.weights)) + assert_allclose(cov(self.x1, fweights=self.unit_frequencies, + aweights=self.unit_weights), + self.res1) + + @pytest.mark.parametrize("test_type", [np.half, np.single, np.double, np.longdouble]) + def test_cov_dtype(self, test_type): + cast_x1 = self.x1.astype(test_type) + res = cov(cast_x1, dtype=test_type) + assert test_type == res.dtype + + def test_gh_27658(self): + x = np.ones((3, 1)) + expected = np.cov(x, ddof=0, rowvar=True) + actual = np.cov(x.T, ddof=0, rowvar=False) + assert_allclose(actual, expected, strict=True) + + +class Test_I0: + + def test_simple(self): + assert_almost_equal( + i0(0.5), + np.array(1.0634833707413234)) + + # need at least one test above 8, as the implementation is piecewise + A = np.array([0.49842636, 0.6969809, 0.22011976, 0.0155549, 10.0]) + expected = np.array([1.06307822, 1.12518299, 1.01214991, 1.00006049, 2815.71662847]) + assert_almost_equal(i0(A), expected) + assert_almost_equal(i0(-A), expected) + + B = np.array([[0.827002, 0.99959078], + [0.89694769, 0.39298162], + [0.37954418, 0.05206293], + [0.36465447, 0.72446427], + [0.48164949, 0.50324519]]) + assert_almost_equal( + i0(B), + np.array([[1.17843223, 1.26583466], + [1.21147086, 1.03898290], + [1.03633899, 1.00067775], + [1.03352052, 1.13557954], + [1.05884290, 1.06432317]])) + # Regression test for gh-11205 + i0_0 = np.i0([0.]) + assert_equal(i0_0.shape, (1,)) + assert_array_equal(np.i0([0.]), np.array([1.])) + + def test_non_array(self): + a = np.arange(4) + + class array_like: + __array_interface__ = a.__array_interface__ + + def __array_wrap__(self, arr, context, return_scalar): + return self + + # E.g. pandas series survive ufunc calls through array-wrap: + assert isinstance(np.abs(array_like()), array_like) + exp = np.i0(a) + res = np.i0(array_like()) + + assert_array_equal(exp, res) + + def test_complex(self): + a = np.array([0, 1 + 2j]) + with pytest.raises(TypeError, match="i0 not supported for complex values"): + res = i0(a) + + +class TestKaiser: + + def test_simple(self): + assert_(np.isfinite(kaiser(1, 1.0))) + assert_almost_equal(kaiser(0, 1.0), + np.array([])) + assert_almost_equal(kaiser(2, 1.0), + np.array([0.78984831, 0.78984831])) + assert_almost_equal(kaiser(5, 1.0), + np.array([0.78984831, 0.94503323, 1., + 0.94503323, 0.78984831])) + assert_almost_equal(kaiser(5, 1.56789), + np.array([0.58285404, 0.88409679, 1., + 0.88409679, 0.58285404])) + + def test_int_beta(self): + kaiser(3, 4) + + +class TestMeshgrid: + + def test_simple(self): + [X, Y] = meshgrid([1, 2, 3], [4, 5, 6, 7]) + assert_array_equal(X, np.array([[1, 2, 3], + [1, 2, 3], + [1, 2, 3], + [1, 2, 3]])) + assert_array_equal(Y, np.array([[4, 4, 4], + [5, 5, 5], + [6, 6, 6], + [7, 7, 7]])) + + def test_single_input(self): + [X] = meshgrid([1, 2, 3, 4]) + assert_array_equal(X, np.array([1, 2, 3, 4])) + + def test_no_input(self): + args = [] + assert_array_equal([], meshgrid(*args)) + assert_array_equal([], meshgrid(*args, copy=False)) + + def test_indexing(self): + x = [1, 2, 3] + y = [4, 5, 6, 7] + [X, Y] = meshgrid(x, y, indexing='ij') + assert_array_equal(X, np.array([[1, 1, 1, 1], + [2, 2, 2, 2], + [3, 3, 3, 3]])) + assert_array_equal(Y, np.array([[4, 5, 6, 7], + [4, 5, 6, 7], + [4, 5, 6, 7]])) + + # Test expected shapes: + z = [8, 9] + assert_(meshgrid(x, y)[0].shape == (4, 3)) + assert_(meshgrid(x, y, indexing='ij')[0].shape == (3, 4)) + assert_(meshgrid(x, y, z)[0].shape == (4, 3, 2)) + assert_(meshgrid(x, y, z, indexing='ij')[0].shape == (3, 4, 2)) + + assert_raises(ValueError, meshgrid, x, y, indexing='notvalid') + + def test_sparse(self): + [X, Y] = meshgrid([1, 2, 3], [4, 5, 6, 7], sparse=True) + assert_array_equal(X, np.array([[1, 2, 3]])) + assert_array_equal(Y, np.array([[4], [5], [6], [7]])) + + def test_invalid_arguments(self): + # Test that meshgrid complains about invalid arguments + # Regression test for issue #4755: + # https://github.com/numpy/numpy/issues/4755 + assert_raises(TypeError, meshgrid, + [1, 2, 3], [4, 5, 6, 7], indices='ij') + + def test_return_type(self): + # Test for appropriate dtype in returned arrays. + # Regression test for issue #5297 + # https://github.com/numpy/numpy/issues/5297 + x = np.arange(0, 10, dtype=np.float32) + y = np.arange(10, 20, dtype=np.float64) + + X, Y = np.meshgrid(x,y) + + assert_(X.dtype == x.dtype) + assert_(Y.dtype == y.dtype) + + # copy + X, Y = np.meshgrid(x,y, copy=True) + + assert_(X.dtype == x.dtype) + assert_(Y.dtype == y.dtype) + + # sparse + X, Y = np.meshgrid(x,y, sparse=True) + + assert_(X.dtype == x.dtype) + assert_(Y.dtype == y.dtype) + + def test_writeback(self): + # Issue 8561 + X = np.array([1.1, 2.2]) + Y = np.array([3.3, 4.4]) + x, y = np.meshgrid(X, Y, sparse=False, copy=True) + + x[0, :] = 0 + assert_equal(x[0, :], 0) + assert_equal(x[1, :], X) + + def test_nd_shape(self): + a, b, c, d, e = np.meshgrid(*([0] * i for i in range(1, 6))) + expected_shape = (2, 1, 3, 4, 5) + assert_equal(a.shape, expected_shape) + assert_equal(b.shape, expected_shape) + assert_equal(c.shape, expected_shape) + assert_equal(d.shape, expected_shape) + assert_equal(e.shape, expected_shape) + + def test_nd_values(self): + a, b, c = np.meshgrid([0], [1, 2], [3, 4, 5]) + assert_equal(a, [[[0, 0, 0]], [[0, 0, 0]]]) + assert_equal(b, [[[1, 1, 1]], [[2, 2, 2]]]) + assert_equal(c, [[[3, 4, 5]], [[3, 4, 5]]]) + + def test_nd_indexing(self): + a, b, c = np.meshgrid([0], [1, 2], [3, 4, 5], indexing='ij') + assert_equal(a, [[[0, 0, 0], [0, 0, 0]]]) + assert_equal(b, [[[1, 1, 1], [2, 2, 2]]]) + assert_equal(c, [[[3, 4, 5], [3, 4, 5]]]) + + +class TestPiecewise: + + def test_simple(self): + # Condition is single bool list + x = piecewise([0, 0], [True, False], [1]) + assert_array_equal(x, [1, 0]) + + # List of conditions: single bool list + x = piecewise([0, 0], [[True, False]], [1]) + assert_array_equal(x, [1, 0]) + + # Conditions is single bool array + x = piecewise([0, 0], np.array([True, False]), [1]) + assert_array_equal(x, [1, 0]) + + # Condition is single int array + x = piecewise([0, 0], np.array([1, 0]), [1]) + assert_array_equal(x, [1, 0]) + + # List of conditions: int array + x = piecewise([0, 0], [np.array([1, 0])], [1]) + assert_array_equal(x, [1, 0]) + + x = piecewise([0, 0], [[False, True]], [lambda x:-1]) + assert_array_equal(x, [0, -1]) + + assert_raises_regex(ValueError, '1 or 2 functions are expected', + piecewise, [0, 0], [[False, True]], []) + assert_raises_regex(ValueError, '1 or 2 functions are expected', + piecewise, [0, 0], [[False, True]], [1, 2, 3]) + + def test_two_conditions(self): + x = piecewise([1, 2], [[True, False], [False, True]], [3, 4]) + assert_array_equal(x, [3, 4]) + + def test_scalar_domains_three_conditions(self): + x = piecewise(3, [True, False, False], [4, 2, 0]) + assert_equal(x, 4) + + def test_default(self): + # No value specified for x[1], should be 0 + x = piecewise([1, 2], [True, False], [2]) + assert_array_equal(x, [2, 0]) + + # Should set x[1] to 3 + x = piecewise([1, 2], [True, False], [2, 3]) + assert_array_equal(x, [2, 3]) + + def test_0d(self): + x = np.array(3) + y = piecewise(x, x > 3, [4, 0]) + assert_(y.ndim == 0) + assert_(y == 0) + + x = 5 + y = piecewise(x, [True, False], [1, 0]) + assert_(y.ndim == 0) + assert_(y == 1) + + # With 3 ranges (It was failing, before) + y = piecewise(x, [False, False, True], [1, 2, 3]) + assert_array_equal(y, 3) + + def test_0d_comparison(self): + x = 3 + y = piecewise(x, [x <= 3, x > 3], [4, 0]) # Should succeed. + assert_equal(y, 4) + + # With 3 ranges (It was failing, before) + x = 4 + y = piecewise(x, [x <= 3, (x > 3) * (x <= 5), x > 5], [1, 2, 3]) + assert_array_equal(y, 2) + + assert_raises_regex(ValueError, '2 or 3 functions are expected', + piecewise, x, [x <= 3, x > 3], [1]) + assert_raises_regex(ValueError, '2 or 3 functions are expected', + piecewise, x, [x <= 3, x > 3], [1, 1, 1, 1]) + + def test_0d_0d_condition(self): + x = np.array(3) + c = np.array(x > 3) + y = piecewise(x, [c], [1, 2]) + assert_equal(y, 2) + + def test_multidimensional_extrafunc(self): + x = np.array([[-2.5, -1.5, -0.5], + [0.5, 1.5, 2.5]]) + y = piecewise(x, [x < 0, x >= 2], [-1, 1, 3]) + assert_array_equal(y, np.array([[-1., -1., -1.], + [3., 3., 1.]])) + + def test_subclasses(self): + class subclass(np.ndarray): + pass + x = np.arange(5.).view(subclass) + r = piecewise(x, [x<2., x>=4], [-1., 1., 0.]) + assert_equal(type(r), subclass) + assert_equal(r, [-1., -1., 0., 0., 1.]) + + +class TestBincount: + + def test_simple(self): + y = np.bincount(np.arange(4)) + assert_array_equal(y, np.ones(4)) + + def test_simple2(self): + y = np.bincount(np.array([1, 5, 2, 4, 1])) + assert_array_equal(y, np.array([0, 2, 1, 0, 1, 1])) + + def test_simple_weight(self): + x = np.arange(4) + w = np.array([0.2, 0.3, 0.5, 0.1]) + y = np.bincount(x, w) + assert_array_equal(y, w) + + def test_simple_weight2(self): + x = np.array([1, 2, 4, 5, 2]) + w = np.array([0.2, 0.3, 0.5, 0.1, 0.2]) + y = np.bincount(x, w) + assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1])) + + def test_with_minlength(self): + x = np.array([0, 1, 0, 1, 1]) + y = np.bincount(x, minlength=3) + assert_array_equal(y, np.array([2, 3, 0])) + x = [] + y = np.bincount(x, minlength=0) + assert_array_equal(y, np.array([])) + + def test_with_minlength_smaller_than_maxvalue(self): + x = np.array([0, 1, 1, 2, 2, 3, 3]) + y = np.bincount(x, minlength=2) + assert_array_equal(y, np.array([1, 2, 2, 2])) + y = np.bincount(x, minlength=0) + assert_array_equal(y, np.array([1, 2, 2, 2])) + + def test_with_minlength_and_weights(self): + x = np.array([1, 2, 4, 5, 2]) + w = np.array([0.2, 0.3, 0.5, 0.1, 0.2]) + y = np.bincount(x, w, 8) + assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1, 0, 0])) + + def test_empty(self): + x = np.array([], dtype=int) + y = np.bincount(x) + assert_array_equal(x, y) + + def test_empty_with_minlength(self): + x = np.array([], dtype=int) + y = np.bincount(x, minlength=5) + assert_array_equal(y, np.zeros(5, dtype=int)) + + @pytest.mark.parametrize('minlength', [0, 3]) + def test_empty_list(self, minlength): + assert_array_equal(np.bincount([], minlength=minlength), + np.zeros(minlength, dtype=int)) + + def test_with_incorrect_minlength(self): + x = np.array([], dtype=int) + assert_raises_regex(TypeError, + "'str' object cannot be interpreted", + lambda: np.bincount(x, minlength="foobar")) + assert_raises_regex(ValueError, + "must not be negative", + lambda: np.bincount(x, minlength=-1)) + + x = np.arange(5) + assert_raises_regex(TypeError, + "'str' object cannot be interpreted", + lambda: np.bincount(x, minlength="foobar")) + assert_raises_regex(ValueError, + "must not be negative", + lambda: np.bincount(x, minlength=-1)) + + @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") + def test_dtype_reference_leaks(self): + # gh-6805 + intp_refcount = sys.getrefcount(np.dtype(np.intp)) + double_refcount = sys.getrefcount(np.dtype(np.double)) + + for j in range(10): + np.bincount([1, 2, 3]) + assert_equal(sys.getrefcount(np.dtype(np.intp)), intp_refcount) + assert_equal(sys.getrefcount(np.dtype(np.double)), double_refcount) + + for j in range(10): + np.bincount([1, 2, 3], [4, 5, 6]) + assert_equal(sys.getrefcount(np.dtype(np.intp)), intp_refcount) + assert_equal(sys.getrefcount(np.dtype(np.double)), double_refcount) + + @pytest.mark.parametrize("vals", [[[2, 2]], 2]) + def test_error_not_1d(self, vals): + # Test that values has to be 1-D (both as array and nested list) + vals_arr = np.asarray(vals) + with assert_raises(ValueError): + np.bincount(vals_arr) + with assert_raises(ValueError): + np.bincount(vals) + + @pytest.mark.parametrize("dt", np.typecodes["AllInteger"]) + def test_gh_28354(self, dt): + a = np.array([0, 1, 1, 3, 2, 1, 7], dtype=dt) + actual = np.bincount(a) + expected = [1, 3, 1, 1, 0, 0, 0, 1] + assert_array_equal(actual, expected) + + def test_contiguous_handling(self): + # check for absence of hard crash + np.bincount(np.arange(10000)[::2]) + + def test_gh_28354_array_like(self): + class A: + def __array__(self): + return np.array([0, 1, 1, 3, 2, 1, 7], dtype=np.uint64) + + a = A() + actual = np.bincount(a) + expected = [1, 3, 1, 1, 0, 0, 0, 1] + assert_array_equal(actual, expected) + + +class TestInterp: + + def test_exceptions(self): + assert_raises(ValueError, interp, 0, [], []) + assert_raises(ValueError, interp, 0, [0], [1, 2]) + assert_raises(ValueError, interp, 0, [0, 1], [1, 2], period=0) + assert_raises(ValueError, interp, 0, [], [], period=360) + assert_raises(ValueError, interp, 0, [0], [1, 2], period=360) + + def test_basic(self): + x = np.linspace(0, 1, 5) + y = np.linspace(0, 1, 5) + x0 = np.linspace(0, 1, 50) + assert_almost_equal(np.interp(x0, x, y), x0) + + def test_right_left_behavior(self): + # Needs range of sizes to test different code paths. + # size ==1 is special cased, 1 < size < 5 is linear search, and + # size >= 5 goes through local search and possibly binary search. + for size in range(1, 10): + xp = np.arange(size, dtype=np.double) + yp = np.ones(size, dtype=np.double) + incpts = np.array([-1, 0, size - 1, size], dtype=np.double) + decpts = incpts[::-1] + + incres = interp(incpts, xp, yp) + decres = interp(decpts, xp, yp) + inctgt = np.array([1, 1, 1, 1], dtype=float) + dectgt = inctgt[::-1] + assert_equal(incres, inctgt) + assert_equal(decres, dectgt) + + incres = interp(incpts, xp, yp, left=0) + decres = interp(decpts, xp, yp, left=0) + inctgt = np.array([0, 1, 1, 1], dtype=float) + dectgt = inctgt[::-1] + assert_equal(incres, inctgt) + assert_equal(decres, dectgt) + + incres = interp(incpts, xp, yp, right=2) + decres = interp(decpts, xp, yp, right=2) + inctgt = np.array([1, 1, 1, 2], dtype=float) + dectgt = inctgt[::-1] + assert_equal(incres, inctgt) + assert_equal(decres, dectgt) + + incres = interp(incpts, xp, yp, left=0, right=2) + decres = interp(decpts, xp, yp, left=0, right=2) + inctgt = np.array([0, 1, 1, 2], dtype=float) + dectgt = inctgt[::-1] + assert_equal(incres, inctgt) + assert_equal(decres, dectgt) + + def test_scalar_interpolation_point(self): + x = np.linspace(0, 1, 5) + y = np.linspace(0, 1, 5) + x0 = 0 + assert_almost_equal(np.interp(x0, x, y), x0) + x0 = .3 + assert_almost_equal(np.interp(x0, x, y), x0) + x0 = np.float32(.3) + assert_almost_equal(np.interp(x0, x, y), x0) + x0 = np.float64(.3) + assert_almost_equal(np.interp(x0, x, y), x0) + x0 = np.nan + assert_almost_equal(np.interp(x0, x, y), x0) + + def test_non_finite_behavior_exact_x(self): + x = [1, 2, 2.5, 3, 4] + xp = [1, 2, 3, 4] + fp = [1, 2, np.inf, 4] + assert_almost_equal(np.interp(x, xp, fp), [1, 2, np.inf, np.inf, 4]) + fp = [1, 2, np.nan, 4] + assert_almost_equal(np.interp(x, xp, fp), [1, 2, np.nan, np.nan, 4]) + + @pytest.fixture(params=[ + lambda x: np.float64(x), + lambda x: _make_complex(x, 0), + lambda x: _make_complex(0, x), + lambda x: _make_complex(x, np.multiply(x, -2)) + ], ids=[ + 'real', + 'complex-real', + 'complex-imag', + 'complex-both' + ]) + def sc(self, request): + """ scale function used by the below tests """ + return request.param + + def test_non_finite_any_nan(self, sc): + """ test that nans are propagated """ + assert_equal(np.interp(0.5, [np.nan, 1], sc([ 0, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, np.nan], sc([ 0, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, 1], sc([np.nan, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, 1], sc([ 0, np.nan])), sc(np.nan)) + + def test_non_finite_inf(self, sc): + """ Test that interp between opposite infs gives nan """ + assert_equal(np.interp(0.5, [-np.inf, +np.inf], sc([ 0, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, 1], sc([-np.inf, +np.inf])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, 1], sc([+np.inf, -np.inf])), sc(np.nan)) + + # unless the y values are equal + assert_equal(np.interp(0.5, [-np.inf, +np.inf], sc([ 10, 10])), sc(10)) + + def test_non_finite_half_inf_xf(self, sc): + """ Test that interp where both axes have a bound at inf gives nan """ + assert_equal(np.interp(0.5, [-np.inf, 1], sc([-np.inf, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [-np.inf, 1], sc([+np.inf, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [-np.inf, 1], sc([ 0, -np.inf])), sc(np.nan)) + assert_equal(np.interp(0.5, [-np.inf, 1], sc([ 0, +np.inf])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, +np.inf], sc([-np.inf, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, +np.inf], sc([+np.inf, 10])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, +np.inf], sc([ 0, -np.inf])), sc(np.nan)) + assert_equal(np.interp(0.5, [ 0, +np.inf], sc([ 0, +np.inf])), sc(np.nan)) + + def test_non_finite_half_inf_x(self, sc): + """ Test interp where the x axis has a bound at inf """ + assert_equal(np.interp(0.5, [-np.inf, -np.inf], sc([0, 10])), sc(10)) + assert_equal(np.interp(0.5, [-np.inf, 1 ], sc([0, 10])), sc(10)) + assert_equal(np.interp(0.5, [ 0, +np.inf], sc([0, 10])), sc(0)) + assert_equal(np.interp(0.5, [+np.inf, +np.inf], sc([0, 10])), sc(0)) + + def test_non_finite_half_inf_f(self, sc): + """ Test interp where the f axis has a bound at inf """ + assert_equal(np.interp(0.5, [0, 1], sc([ 0, -np.inf])), sc(-np.inf)) + assert_equal(np.interp(0.5, [0, 1], sc([ 0, +np.inf])), sc(+np.inf)) + assert_equal(np.interp(0.5, [0, 1], sc([-np.inf, 10])), sc(-np.inf)) + assert_equal(np.interp(0.5, [0, 1], sc([+np.inf, 10])), sc(+np.inf)) + assert_equal(np.interp(0.5, [0, 1], sc([-np.inf, -np.inf])), sc(-np.inf)) + assert_equal(np.interp(0.5, [0, 1], sc([+np.inf, +np.inf])), sc(+np.inf)) + + def test_complex_interp(self): + # test complex interpolation + x = np.linspace(0, 1, 5) + y = np.linspace(0, 1, 5) + (1 + np.linspace(0, 1, 5))*1.0j + x0 = 0.3 + y0 = x0 + (1+x0)*1.0j + assert_almost_equal(np.interp(x0, x, y), y0) + # test complex left and right + x0 = -1 + left = 2 + 3.0j + assert_almost_equal(np.interp(x0, x, y, left=left), left) + x0 = 2.0 + right = 2 + 3.0j + assert_almost_equal(np.interp(x0, x, y, right=right), right) + # test complex non finite + x = [1, 2, 2.5, 3, 4] + xp = [1, 2, 3, 4] + fp = [1, 2+1j, np.inf, 4] + y = [1, 2+1j, np.inf+0.5j, np.inf, 4] + assert_almost_equal(np.interp(x, xp, fp), y) + # test complex periodic + x = [-180, -170, -185, 185, -10, -5, 0, 365] + xp = [190, -190, 350, -350] + fp = [5+1.0j, 10+2j, 3+3j, 4+4j] + y = [7.5+1.5j, 5.+1.0j, 8.75+1.75j, 6.25+1.25j, 3.+3j, 3.25+3.25j, + 3.5+3.5j, 3.75+3.75j] + assert_almost_equal(np.interp(x, xp, fp, period=360), y) + + def test_zero_dimensional_interpolation_point(self): + x = np.linspace(0, 1, 5) + y = np.linspace(0, 1, 5) + x0 = np.array(.3) + assert_almost_equal(np.interp(x0, x, y), x0) + + xp = np.array([0, 2, 4]) + fp = np.array([1, -1, 1]) + + actual = np.interp(np.array(1), xp, fp) + assert_equal(actual, 0) + assert_(isinstance(actual, np.float64)) + + actual = np.interp(np.array(4.5), xp, fp, period=4) + assert_equal(actual, 0.5) + assert_(isinstance(actual, np.float64)) + + def test_if_len_x_is_small(self): + xp = np.arange(0, 10, 0.0001) + fp = np.sin(xp) + assert_almost_equal(np.interp(np.pi, xp, fp), 0.0) + + def test_period(self): + x = [-180, -170, -185, 185, -10, -5, 0, 365] + xp = [190, -190, 350, -350] + fp = [5, 10, 3, 4] + y = [7.5, 5., 8.75, 6.25, 3., 3.25, 3.5, 3.75] + assert_almost_equal(np.interp(x, xp, fp, period=360), y) + x = np.array(x, order='F').reshape(2, -1) + y = np.array(y, order='C').reshape(2, -1) + assert_almost_equal(np.interp(x, xp, fp, period=360), y) + + +class TestPercentile: + + def test_basic(self): + x = np.arange(8) * 0.5 + assert_equal(np.percentile(x, 0), 0.) + assert_equal(np.percentile(x, 100), 3.5) + assert_equal(np.percentile(x, 50), 1.75) + x[1] = np.nan + assert_equal(np.percentile(x, 0), np.nan) + assert_equal(np.percentile(x, 0, method='nearest'), np.nan) + assert_equal(np.percentile(x, 0, method='inverted_cdf'), np.nan) + assert_equal( + np.percentile(x, 0, method='inverted_cdf', + weights=np.ones_like(x)), + np.nan, + ) + + def test_fraction(self): + x = [Fraction(i, 2) for i in range(8)] + + p = np.percentile(x, Fraction(0)) + assert_equal(p, Fraction(0)) + assert_equal(type(p), Fraction) + + p = np.percentile(x, Fraction(100)) + assert_equal(p, Fraction(7, 2)) + assert_equal(type(p), Fraction) + + p = np.percentile(x, Fraction(50)) + assert_equal(p, Fraction(7, 4)) + assert_equal(type(p), Fraction) + + p = np.percentile(x, [Fraction(50)]) + assert_equal(p, np.array([Fraction(7, 4)])) + assert_equal(type(p), np.ndarray) + + def test_api(self): + d = np.ones(5) + np.percentile(d, 5, None, None, False) + np.percentile(d, 5, None, None, False, 'linear') + o = np.ones((1,)) + np.percentile(d, 5, None, o, False, 'linear') + + def test_complex(self): + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='G') + assert_raises(TypeError, np.percentile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='D') + assert_raises(TypeError, np.percentile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='F') + assert_raises(TypeError, np.percentile, arr_c, 0.5) + + def test_2D(self): + x = np.array([[1, 1, 1], + [1, 1, 1], + [4, 4, 3], + [1, 1, 1], + [1, 1, 1]]) + assert_array_equal(np.percentile(x, 50, axis=0), [1, 1, 1]) + + @pytest.mark.parametrize("dtype", np.typecodes["Float"]) + def test_linear_nan_1D(self, dtype): + # METHOD 1 of H&F + arr = np.asarray([15.0, np.nan, 35.0, 40.0, 50.0], dtype=dtype) + res = np.percentile( + arr, + 40.0, + method="linear") + np.testing.assert_equal(res, np.nan) + np.testing.assert_equal(res.dtype, arr.dtype) + + H_F_TYPE_CODES = [(int_type, np.float64) + for int_type in np.typecodes["AllInteger"] + ] + [(np.float16, np.float16), + (np.float32, np.float32), + (np.float64, np.float64), + (np.longdouble, np.longdouble), + (np.dtype("O"), np.float64)] + + @pytest.mark.parametrize(["function", "quantile"], + [(np.quantile, 0.4), + (np.percentile, 40.0)]) + @pytest.mark.parametrize(["input_dtype", "expected_dtype"], H_F_TYPE_CODES) + @pytest.mark.parametrize(["method", "weighted", "expected"], + [("inverted_cdf", False, 20), + ("inverted_cdf", True, 20), + ("averaged_inverted_cdf", False, 27.5), + ("closest_observation", False, 20), + ("interpolated_inverted_cdf", False, 20), + ("hazen", False, 27.5), + ("weibull", False, 26), + ("linear", False, 29), + ("median_unbiased", False, 27), + ("normal_unbiased", False, 27.125), + ]) + def test_linear_interpolation(self, + function, + quantile, + method, + weighted, + expected, + input_dtype, + expected_dtype): + expected_dtype = np.dtype(expected_dtype) + + arr = np.asarray([15.0, 20.0, 35.0, 40.0, 50.0], dtype=input_dtype) + weights = np.ones_like(arr) if weighted else None + if input_dtype is np.longdouble: + if function is np.quantile: + # 0.4 is not exactly representable and it matters + # for "averaged_inverted_cdf", so we need to cheat. + quantile = input_dtype("0.4") + # We want to use nulp, but that does not work for longdouble + test_function = np.testing.assert_almost_equal + else: + test_function = np.testing.assert_array_almost_equal_nulp + + actual = function(arr, quantile, method=method, weights=weights) + + test_function(actual, expected_dtype.type(expected)) + + if method in ["inverted_cdf", "closest_observation"]: + if input_dtype == "O": + np.testing.assert_equal(np.asarray(actual).dtype, np.float64) + else: + np.testing.assert_equal(np.asarray(actual).dtype, + np.dtype(input_dtype)) + else: + np.testing.assert_equal(np.asarray(actual).dtype, + np.dtype(expected_dtype)) + + TYPE_CODES = np.typecodes["AllInteger"] + np.typecodes["Float"] + "O" + + @pytest.mark.parametrize("dtype", TYPE_CODES) + def test_lower_higher(self, dtype): + assert_equal(np.percentile(np.arange(10, dtype=dtype), 50, + method='lower'), 4) + assert_equal(np.percentile(np.arange(10, dtype=dtype), 50, + method='higher'), 5) + + @pytest.mark.parametrize("dtype", TYPE_CODES) + def test_midpoint(self, dtype): + assert_equal(np.percentile(np.arange(10, dtype=dtype), 51, + method='midpoint'), 4.5) + assert_equal(np.percentile(np.arange(9, dtype=dtype) + 1, 50, + method='midpoint'), 5) + assert_equal(np.percentile(np.arange(11, dtype=dtype), 51, + method='midpoint'), 5.5) + assert_equal(np.percentile(np.arange(11, dtype=dtype), 50, + method='midpoint'), 5) + + @pytest.mark.parametrize("dtype", TYPE_CODES) + def test_nearest(self, dtype): + assert_equal(np.percentile(np.arange(10, dtype=dtype), 51, + method='nearest'), 5) + assert_equal(np.percentile(np.arange(10, dtype=dtype), 49, + method='nearest'), 4) + + def test_linear_interpolation_extrapolation(self): + arr = np.random.rand(5) + + actual = np.percentile(arr, 100) + np.testing.assert_equal(actual, arr.max()) + + actual = np.percentile(arr, 0) + np.testing.assert_equal(actual, arr.min()) + + def test_sequence(self): + x = np.arange(8) * 0.5 + assert_equal(np.percentile(x, [0, 100, 50]), [0, 3.5, 1.75]) + + def test_axis(self): + x = np.arange(12).reshape(3, 4) + + assert_equal(np.percentile(x, (25, 50, 100)), [2.75, 5.5, 11.0]) + + r0 = [[2, 3, 4, 5], [4, 5, 6, 7], [8, 9, 10, 11]] + assert_equal(np.percentile(x, (25, 50, 100), axis=0), r0) + + r1 = [[0.75, 1.5, 3], [4.75, 5.5, 7], [8.75, 9.5, 11]] + assert_equal(np.percentile(x, (25, 50, 100), axis=1), np.array(r1).T) + + # ensure qth axis is always first as with np.array(old_percentile(..)) + x = np.arange(3 * 4 * 5 * 6).reshape(3, 4, 5, 6) + assert_equal(np.percentile(x, (25, 50)).shape, (2,)) + assert_equal(np.percentile(x, (25, 50, 75)).shape, (3,)) + assert_equal(np.percentile(x, (25, 50), axis=0).shape, (2, 4, 5, 6)) + assert_equal(np.percentile(x, (25, 50), axis=1).shape, (2, 3, 5, 6)) + assert_equal(np.percentile(x, (25, 50), axis=2).shape, (2, 3, 4, 6)) + assert_equal(np.percentile(x, (25, 50), axis=3).shape, (2, 3, 4, 5)) + assert_equal( + np.percentile(x, (25, 50, 75), axis=1).shape, (3, 3, 5, 6)) + assert_equal(np.percentile(x, (25, 50), + method="higher").shape, (2,)) + assert_equal(np.percentile(x, (25, 50, 75), + method="higher").shape, (3,)) + assert_equal(np.percentile(x, (25, 50), axis=0, + method="higher").shape, (2, 4, 5, 6)) + assert_equal(np.percentile(x, (25, 50), axis=1, + method="higher").shape, (2, 3, 5, 6)) + assert_equal(np.percentile(x, (25, 50), axis=2, + method="higher").shape, (2, 3, 4, 6)) + assert_equal(np.percentile(x, (25, 50), axis=3, + method="higher").shape, (2, 3, 4, 5)) + assert_equal(np.percentile(x, (25, 50, 75), axis=1, + method="higher").shape, (3, 3, 5, 6)) + + def test_scalar_q(self): + # test for no empty dimensions for compatibility with old percentile + x = np.arange(12).reshape(3, 4) + assert_equal(np.percentile(x, 50), 5.5) + assert_(np.isscalar(np.percentile(x, 50))) + r0 = np.array([4., 5., 6., 7.]) + assert_equal(np.percentile(x, 50, axis=0), r0) + assert_equal(np.percentile(x, 50, axis=0).shape, r0.shape) + r1 = np.array([1.5, 5.5, 9.5]) + assert_almost_equal(np.percentile(x, 50, axis=1), r1) + assert_equal(np.percentile(x, 50, axis=1).shape, r1.shape) + + out = np.empty(1) + assert_equal(np.percentile(x, 50, out=out), 5.5) + assert_equal(out, 5.5) + out = np.empty(4) + assert_equal(np.percentile(x, 50, axis=0, out=out), r0) + assert_equal(out, r0) + out = np.empty(3) + assert_equal(np.percentile(x, 50, axis=1, out=out), r1) + assert_equal(out, r1) + + # test for no empty dimensions for compatibility with old percentile + x = np.arange(12).reshape(3, 4) + assert_equal(np.percentile(x, 50, method='lower'), 5.) + assert_(np.isscalar(np.percentile(x, 50))) + r0 = np.array([4., 5., 6., 7.]) + c0 = np.percentile(x, 50, method='lower', axis=0) + assert_equal(c0, r0) + assert_equal(c0.shape, r0.shape) + r1 = np.array([1., 5., 9.]) + c1 = np.percentile(x, 50, method='lower', axis=1) + assert_almost_equal(c1, r1) + assert_equal(c1.shape, r1.shape) + + out = np.empty((), dtype=x.dtype) + c = np.percentile(x, 50, method='lower', out=out) + assert_equal(c, 5) + assert_equal(out, 5) + out = np.empty(4, dtype=x.dtype) + c = np.percentile(x, 50, method='lower', axis=0, out=out) + assert_equal(c, r0) + assert_equal(out, r0) + out = np.empty(3, dtype=x.dtype) + c = np.percentile(x, 50, method='lower', axis=1, out=out) + assert_equal(c, r1) + assert_equal(out, r1) + + def test_exception(self): + assert_raises(ValueError, np.percentile, [1, 2], 56, + method='foobar') + assert_raises(ValueError, np.percentile, [1], 101) + assert_raises(ValueError, np.percentile, [1], -1) + assert_raises(ValueError, np.percentile, [1], list(range(50)) + [101]) + assert_raises(ValueError, np.percentile, [1], list(range(50)) + [-0.1]) + + def test_percentile_list(self): + assert_equal(np.percentile([1, 2, 3], 0), 1) + + @pytest.mark.parametrize( + "percentile, with_weights", + [ + (np.percentile, False), + (partial(np.percentile, method="inverted_cdf"), True), + ] + ) + def test_percentile_out(self, percentile, with_weights): + out_dtype = int if with_weights else float + x = np.array([1, 2, 3]) + y = np.zeros((3,), dtype=out_dtype) + p = (1, 2, 3) + weights = np.ones_like(x) if with_weights else None + r = percentile(x, p, out=y, weights=weights) + assert r is y + assert_equal(percentile(x, p, weights=weights), y) + + x = np.array([[1, 2, 3], + [4, 5, 6]]) + y = np.zeros((3, 3), dtype=out_dtype) + weights = np.ones_like(x) if with_weights else None + r = percentile(x, p, axis=0, out=y, weights=weights) + assert r is y + assert_equal(percentile(x, p, weights=weights, axis=0), y) + + y = np.zeros((3, 2), dtype=out_dtype) + percentile(x, p, axis=1, out=y, weights=weights) + assert_equal(percentile(x, p, weights=weights, axis=1), y) + + x = np.arange(12).reshape(3, 4) + # q.dim > 1, float + if with_weights: + r0 = np.array([[0, 1, 2, 3], [4, 5, 6, 7]]) + else: + r0 = np.array([[2., 3., 4., 5.], [4., 5., 6., 7.]]) + out = np.empty((2, 4), dtype=out_dtype) + weights = np.ones_like(x) if with_weights else None + assert_equal( + percentile(x, (25, 50), axis=0, out=out, weights=weights), r0 + ) + assert_equal(out, r0) + r1 = np.array([[0.75, 4.75, 8.75], [1.5, 5.5, 9.5]]) + out = np.empty((2, 3)) + assert_equal(np.percentile(x, (25, 50), axis=1, out=out), r1) + assert_equal(out, r1) + + # q.dim > 1, int + r0 = np.array([[0, 1, 2, 3], [4, 5, 6, 7]]) + out = np.empty((2, 4), dtype=x.dtype) + c = np.percentile(x, (25, 50), method='lower', axis=0, out=out) + assert_equal(c, r0) + assert_equal(out, r0) + r1 = np.array([[0, 4, 8], [1, 5, 9]]) + out = np.empty((2, 3), dtype=x.dtype) + c = np.percentile(x, (25, 50), method='lower', axis=1, out=out) + assert_equal(c, r1) + assert_equal(out, r1) + + def test_percentile_empty_dim(self): + # empty dims are preserved + d = np.arange(11 * 2).reshape(11, 1, 2, 1) + assert_array_equal(np.percentile(d, 50, axis=0).shape, (1, 2, 1)) + assert_array_equal(np.percentile(d, 50, axis=1).shape, (11, 2, 1)) + assert_array_equal(np.percentile(d, 50, axis=2).shape, (11, 1, 1)) + assert_array_equal(np.percentile(d, 50, axis=3).shape, (11, 1, 2)) + assert_array_equal(np.percentile(d, 50, axis=-1).shape, (11, 1, 2)) + assert_array_equal(np.percentile(d, 50, axis=-2).shape, (11, 1, 1)) + assert_array_equal(np.percentile(d, 50, axis=-3).shape, (11, 2, 1)) + assert_array_equal(np.percentile(d, 50, axis=-4).shape, (1, 2, 1)) + + assert_array_equal(np.percentile(d, 50, axis=2, + method='midpoint').shape, + (11, 1, 1)) + assert_array_equal(np.percentile(d, 50, axis=-2, + method='midpoint').shape, + (11, 1, 1)) + + assert_array_equal(np.array(np.percentile(d, [10, 50], axis=0)).shape, + (2, 1, 2, 1)) + assert_array_equal(np.array(np.percentile(d, [10, 50], axis=1)).shape, + (2, 11, 2, 1)) + assert_array_equal(np.array(np.percentile(d, [10, 50], axis=2)).shape, + (2, 11, 1, 1)) + assert_array_equal(np.array(np.percentile(d, [10, 50], axis=3)).shape, + (2, 11, 1, 2)) + + def test_percentile_no_overwrite(self): + a = np.array([2, 3, 4, 1]) + np.percentile(a, [50], overwrite_input=False) + assert_equal(a, np.array([2, 3, 4, 1])) + + a = np.array([2, 3, 4, 1]) + np.percentile(a, [50]) + assert_equal(a, np.array([2, 3, 4, 1])) + + def test_no_p_overwrite(self): + p = np.linspace(0., 100., num=5) + np.percentile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, np.linspace(0., 100., num=5)) + p = np.linspace(0., 100., num=5).tolist() + np.percentile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, np.linspace(0., 100., num=5).tolist()) + + def test_percentile_overwrite(self): + a = np.array([2, 3, 4, 1]) + b = np.percentile(a, [50], overwrite_input=True) + assert_equal(b, np.array([2.5])) + + b = np.percentile([2, 3, 4, 1], [50], overwrite_input=True) + assert_equal(b, np.array([2.5])) + + def test_extended_axis(self): + o = np.random.normal(size=(71, 23)) + x = np.dstack([o] * 10) + assert_equal(np.percentile(x, 30, axis=(0, 1)), np.percentile(o, 30)) + x = np.moveaxis(x, -1, 0) + assert_equal(np.percentile(x, 30, axis=(-2, -1)), np.percentile(o, 30)) + x = x.swapaxes(0, 1).copy() + assert_equal(np.percentile(x, 30, axis=(0, -1)), np.percentile(o, 30)) + x = x.swapaxes(0, 1).copy() + + assert_equal(np.percentile(x, [25, 60], axis=(0, 1, 2)), + np.percentile(x, [25, 60], axis=None)) + assert_equal(np.percentile(x, [25, 60], axis=(0,)), + np.percentile(x, [25, 60], axis=0)) + + d = np.arange(3 * 5 * 7 * 11).reshape((3, 5, 7, 11)) + np.random.shuffle(d.ravel()) + assert_equal(np.percentile(d, 25, axis=(0, 1, 2))[0], + np.percentile(d[:,:,:, 0].flatten(), 25)) + assert_equal(np.percentile(d, [10, 90], axis=(0, 1, 3))[:, 1], + np.percentile(d[:,:, 1,:].flatten(), [10, 90])) + assert_equal(np.percentile(d, 25, axis=(3, 1, -4))[2], + np.percentile(d[:,:, 2,:].flatten(), 25)) + assert_equal(np.percentile(d, 25, axis=(3, 1, 2))[2], + np.percentile(d[2,:,:,:].flatten(), 25)) + assert_equal(np.percentile(d, 25, axis=(3, 2))[2, 1], + np.percentile(d[2, 1,:,:].flatten(), 25)) + assert_equal(np.percentile(d, 25, axis=(1, -2))[2, 1], + np.percentile(d[2,:,:, 1].flatten(), 25)) + assert_equal(np.percentile(d, 25, axis=(1, 3))[2, 2], + np.percentile(d[2,:, 2,:].flatten(), 25)) + + def test_extended_axis_invalid(self): + d = np.ones((3, 5, 7, 11)) + assert_raises(AxisError, np.percentile, d, axis=-5, q=25) + assert_raises(AxisError, np.percentile, d, axis=(0, -5), q=25) + assert_raises(AxisError, np.percentile, d, axis=4, q=25) + assert_raises(AxisError, np.percentile, d, axis=(0, 4), q=25) + # each of these refers to the same axis twice + assert_raises(ValueError, np.percentile, d, axis=(1, 1), q=25) + assert_raises(ValueError, np.percentile, d, axis=(-1, -1), q=25) + assert_raises(ValueError, np.percentile, d, axis=(3, -1), q=25) + + def test_keepdims(self): + d = np.ones((3, 5, 7, 11)) + assert_equal(np.percentile(d, 7, axis=None, keepdims=True).shape, + (1, 1, 1, 1)) + assert_equal(np.percentile(d, 7, axis=(0, 1), keepdims=True).shape, + (1, 1, 7, 11)) + assert_equal(np.percentile(d, 7, axis=(0, 3), keepdims=True).shape, + (1, 5, 7, 1)) + assert_equal(np.percentile(d, 7, axis=(1,), keepdims=True).shape, + (3, 1, 7, 11)) + assert_equal(np.percentile(d, 7, (0, 1, 2, 3), keepdims=True).shape, + (1, 1, 1, 1)) + assert_equal(np.percentile(d, 7, axis=(0, 1, 3), keepdims=True).shape, + (1, 1, 7, 1)) + + assert_equal(np.percentile(d, [1, 7], axis=(0, 1, 3), + keepdims=True).shape, (2, 1, 1, 7, 1)) + assert_equal(np.percentile(d, [1, 7], axis=(0, 3), + keepdims=True).shape, (2, 1, 5, 7, 1)) + + @pytest.mark.parametrize('q', [7, [1, 7]]) + @pytest.mark.parametrize( + argnames='axis', + argvalues=[ + None, + 1, + (1,), + (0, 1), + (-3, -1), + ] + ) + def test_keepdims_out(self, q, axis): + d = np.ones((3, 5, 7, 11)) + if axis is None: + shape_out = (1,) * d.ndim + else: + axis_norm = normalize_axis_tuple(axis, d.ndim) + shape_out = tuple( + 1 if i in axis_norm else d.shape[i] for i in range(d.ndim)) + shape_out = np.shape(q) + shape_out + + out = np.empty(shape_out) + result = np.percentile(d, q, axis=axis, keepdims=True, out=out) + assert result is out + assert_equal(result.shape, shape_out) + + def test_out(self): + o = np.zeros((4,)) + d = np.ones((3, 4)) + assert_equal(np.percentile(d, 0, 0, out=o), o) + assert_equal(np.percentile(d, 0, 0, method='nearest', out=o), o) + o = np.zeros((3,)) + assert_equal(np.percentile(d, 1, 1, out=o), o) + assert_equal(np.percentile(d, 1, 1, method='nearest', out=o), o) + + o = np.zeros(()) + assert_equal(np.percentile(d, 2, out=o), o) + assert_equal(np.percentile(d, 2, method='nearest', out=o), o) + + @pytest.mark.parametrize("method, weighted", [ + ("linear", False), + ("nearest", False), + ("inverted_cdf", False), + ("inverted_cdf", True), + ]) + def test_out_nan(self, method, weighted): + if weighted: + kwargs = {"weights": np.ones((3, 4)), "method": method} + else: + kwargs = {"method": method} + with warnings.catch_warnings(record=True): + warnings.filterwarnings('always', '', RuntimeWarning) + o = np.zeros((4,)) + d = np.ones((3, 4)) + d[2, 1] = np.nan + assert_equal(np.percentile(d, 0, 0, out=o, **kwargs), o) + + o = np.zeros((3,)) + assert_equal(np.percentile(d, 1, 1, out=o, **kwargs), o) + + o = np.zeros(()) + assert_equal(np.percentile(d, 1, out=o, **kwargs), o) + + def test_nan_behavior(self): + a = np.arange(24, dtype=float) + a[2] = np.nan + assert_equal(np.percentile(a, 0.3), np.nan) + assert_equal(np.percentile(a, 0.3, axis=0), np.nan) + assert_equal(np.percentile(a, [0.3, 0.6], axis=0), + np.array([np.nan] * 2)) + + a = np.arange(24, dtype=float).reshape(2, 3, 4) + a[1, 2, 3] = np.nan + a[1, 1, 2] = np.nan + + # no axis + assert_equal(np.percentile(a, 0.3), np.nan) + assert_equal(np.percentile(a, 0.3).ndim, 0) + + # axis0 zerod + b = np.percentile(np.arange(24, dtype=float).reshape(2, 3, 4), 0.3, 0) + b[2, 3] = np.nan + b[1, 2] = np.nan + assert_equal(np.percentile(a, 0.3, 0), b) + + # axis0 not zerod + b = np.percentile(np.arange(24, dtype=float).reshape(2, 3, 4), + [0.3, 0.6], 0) + b[:, 2, 3] = np.nan + b[:, 1, 2] = np.nan + assert_equal(np.percentile(a, [0.3, 0.6], 0), b) + + # axis1 zerod + b = np.percentile(np.arange(24, dtype=float).reshape(2, 3, 4), 0.3, 1) + b[1, 3] = np.nan + b[1, 2] = np.nan + assert_equal(np.percentile(a, 0.3, 1), b) + # axis1 not zerod + b = np.percentile( + np.arange(24, dtype=float).reshape(2, 3, 4), [0.3, 0.6], 1) + b[:, 1, 3] = np.nan + b[:, 1, 2] = np.nan + assert_equal(np.percentile(a, [0.3, 0.6], 1), b) + + # axis02 zerod + b = np.percentile( + np.arange(24, dtype=float).reshape(2, 3, 4), 0.3, (0, 2)) + b[1] = np.nan + b[2] = np.nan + assert_equal(np.percentile(a, 0.3, (0, 2)), b) + # axis02 not zerod + b = np.percentile(np.arange(24, dtype=float).reshape(2, 3, 4), + [0.3, 0.6], (0, 2)) + b[:, 1] = np.nan + b[:, 2] = np.nan + assert_equal(np.percentile(a, [0.3, 0.6], (0, 2)), b) + # axis02 not zerod with method='nearest' + b = np.percentile(np.arange(24, dtype=float).reshape(2, 3, 4), + [0.3, 0.6], (0, 2), method='nearest') + b[:, 1] = np.nan + b[:, 2] = np.nan + assert_equal(np.percentile( + a, [0.3, 0.6], (0, 2), method='nearest'), b) + + def test_nan_q(self): + # GH18830 + with pytest.raises(ValueError, match="Percentiles must be in"): + np.percentile([1, 2, 3, 4.0], np.nan) + with pytest.raises(ValueError, match="Percentiles must be in"): + np.percentile([1, 2, 3, 4.0], [np.nan]) + q = np.linspace(1.0, 99.0, 16) + q[0] = np.nan + with pytest.raises(ValueError, match="Percentiles must be in"): + np.percentile([1, 2, 3, 4.0], q) + + @pytest.mark.parametrize("dtype", ["m8[D]", "M8[s]"]) + @pytest.mark.parametrize("pos", [0, 23, 10]) + def test_nat_basic(self, dtype, pos): + # TODO: Note that times have dubious rounding as of fixing NaTs! + # NaT and NaN should behave the same, do basic tests for NaT: + a = np.arange(0, 24, dtype=dtype) + a[pos] = "NaT" + res = np.percentile(a, 30) + assert res.dtype == dtype + assert np.isnat(res) + res = np.percentile(a, [30, 60]) + assert res.dtype == dtype + assert np.isnat(res).all() + + a = np.arange(0, 24*3, dtype=dtype).reshape(-1, 3) + a[pos, 1] = "NaT" + res = np.percentile(a, 30, axis=0) + assert_array_equal(np.isnat(res), [False, True, False]) + + +quantile_methods = [ + 'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation', + 'interpolated_inverted_cdf', 'hazen', 'weibull', 'linear', + 'median_unbiased', 'normal_unbiased', 'nearest', 'lower', 'higher', + 'midpoint'] + + +methods_supporting_weights = ["inverted_cdf"] + + +class TestQuantile: + # most of this is already tested by TestPercentile + + def V(self, x, y, alpha): + # Identification function used in several tests. + return (x >= y) - alpha + + def test_max_ulp(self): + x = [0.0, 0.2, 0.4] + a = np.quantile(x, 0.45) + # The default linear method would result in 0 + 0.2 * (0.45/2) = 0.18. + # 0.18 is not exactly representable and the formula leads to a 1 ULP + # different result. Ensure it is this exact within 1 ULP, see gh-20331. + np.testing.assert_array_max_ulp(a, 0.18, maxulp=1) + + def test_basic(self): + x = np.arange(8) * 0.5 + assert_equal(np.quantile(x, 0), 0.) + assert_equal(np.quantile(x, 1), 3.5) + assert_equal(np.quantile(x, 0.5), 1.75) + + def test_correct_quantile_value(self): + a = np.array([True]) + tf_quant = np.quantile(True, False) + assert_equal(tf_quant, a[0]) + assert_equal(type(tf_quant), a.dtype) + a = np.array([False, True, True]) + quant_res = np.quantile(a, a) + assert_array_equal(quant_res, a) + assert_equal(quant_res.dtype, a.dtype) + + def test_fraction(self): + # fractional input, integral quantile + x = [Fraction(i, 2) for i in range(8)] + q = np.quantile(x, 0) + assert_equal(q, 0) + assert_equal(type(q), Fraction) + + q = np.quantile(x, 1) + assert_equal(q, Fraction(7, 2)) + assert_equal(type(q), Fraction) + + q = np.quantile(x, .5) + assert_equal(q, 1.75) + assert_equal(type(q), np.float64) + + q = np.quantile(x, Fraction(1, 2)) + assert_equal(q, Fraction(7, 4)) + assert_equal(type(q), Fraction) + + q = np.quantile(x, [Fraction(1, 2)]) + assert_equal(q, np.array([Fraction(7, 4)])) + assert_equal(type(q), np.ndarray) + + q = np.quantile(x, [[Fraction(1, 2)]]) + assert_equal(q, np.array([[Fraction(7, 4)]])) + assert_equal(type(q), np.ndarray) + + # repeat with integral input but fractional quantile + x = np.arange(8) + assert_equal(np.quantile(x, Fraction(1, 2)), Fraction(7, 2)) + + def test_complex(self): + #See gh-22652 + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='G') + assert_raises(TypeError, np.quantile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='D') + assert_raises(TypeError, np.quantile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='F') + assert_raises(TypeError, np.quantile, arr_c, 0.5) + + def test_no_p_overwrite(self): + # this is worth retesting, because quantile does not make a copy + p0 = np.array([0, 0.75, 0.25, 0.5, 1.0]) + p = p0.copy() + np.quantile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, p0) + + p0 = p0.tolist() + p = p.tolist() + np.quantile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, p0) + + @pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) + def test_quantile_preserve_int_type(self, dtype): + res = np.quantile(np.array([1, 2], dtype=dtype), [0.5], + method="nearest") + assert res.dtype == dtype + + @pytest.mark.parametrize("method", quantile_methods) + def test_q_zero_one(self, method): + # gh-24710 + arr = [10, 11, 12] + quantile = np.quantile(arr, q = [0, 1], method=method) + assert_equal(quantile, np.array([10, 12])) + + @pytest.mark.parametrize("method", quantile_methods) + def test_quantile_monotonic(self, method): + # GH 14685 + # test that the return value of quantile is monotonic if p0 is ordered + # Also tests that the boundary values are not mishandled. + p0 = np.linspace(0, 1, 101) + quantile = np.quantile(np.array([0, 1, 1, 2, 2, 3, 3, 4, 5, 5, 1, 1, 9, 9, 9, + 8, 8, 7]) * 0.1, p0, method=method) + assert_equal(np.sort(quantile), quantile) + + # Also test one where the number of data points is clearly divisible: + quantile = np.quantile([0., 1., 2., 3.], p0, method=method) + assert_equal(np.sort(quantile), quantile) + + @hypothesis.given( + arr=arrays(dtype=np.float64, + shape=st.integers(min_value=3, max_value=1000), + elements=st.floats(allow_infinity=False, allow_nan=False, + min_value=-1e300, max_value=1e300))) + def test_quantile_monotonic_hypo(self, arr): + p0 = np.arange(0, 1, 0.01) + quantile = np.quantile(arr, p0) + assert_equal(np.sort(quantile), quantile) + + def test_quantile_scalar_nan(self): + a = np.array([[10., 7., 4.], [3., 2., 1.]]) + a[0][1] = np.nan + actual = np.quantile(a, 0.5) + assert np.isscalar(actual) + assert_equal(np.quantile(a, 0.5), np.nan) + + @pytest.mark.parametrize("weights", [False, True]) + @pytest.mark.parametrize("method", quantile_methods) + @pytest.mark.parametrize("alpha", [0.2, 0.5, 0.9]) + def test_quantile_identification_equation(self, weights, method, alpha): + # Test that the identification equation holds for the empirical + # CDF: + # E[V(x, Y)] = 0 <=> x is quantile + # with Y the random variable for which we have observed values and + # V(x, y) the canonical identification function for the quantile (at + # level alpha), see + # https://doi.org/10.48550/arXiv.0912.0902 + if weights and method not in methods_supporting_weights: + pytest.skip("Weights not supported by method.") + rng = np.random.default_rng(4321) + # We choose n and alpha such that we cover 3 cases: + # - n * alpha is an integer + # - n * alpha is a float that gets rounded down + # - n * alpha is a float that gest rounded up + n = 102 # n * alpha = 20.4, 51. , 91.8 + y = rng.random(n) + w = rng.integers(low=0, high=10, size=n) if weights else None + x = np.quantile(y, alpha, method=method, weights=w) + + if method in ("higher",): + # These methods do not fulfill the identification equation. + assert np.abs(np.mean(self.V(x, y, alpha))) > 0.1 / n + elif int(n * alpha) == n * alpha and not weights: + # We can expect exact results, up to machine precision. + assert_allclose( + np.average(self.V(x, y, alpha), weights=w), 0, atol=1e-14, + ) + else: + # V = (x >= y) - alpha cannot sum to zero exactly but within + # "sample precision". + assert_allclose(np.average(self.V(x, y, alpha), weights=w), 0, + atol=1 / n / np.amin([alpha, 1 - alpha])) + + @pytest.mark.parametrize("weights", [False, True]) + @pytest.mark.parametrize("method", quantile_methods) + @pytest.mark.parametrize("alpha", [0.2, 0.5, 0.9]) + def test_quantile_add_and_multiply_constant(self, weights, method, alpha): + # Test that + # 1. quantile(c + x) = c + quantile(x) + # 2. quantile(c * x) = c * quantile(x) + # 3. quantile(-x) = -quantile(x, 1 - alpha) + # On empirical quantiles, this equation does not hold exactly. + # Koenker (2005) "Quantile Regression" Chapter 2.2.3 calls these + # properties equivariance. + if weights and method not in methods_supporting_weights: + pytest.skip("Weights not supported by method.") + rng = np.random.default_rng(4321) + # We choose n and alpha such that we have cases for + # - n * alpha is an integer + # - n * alpha is a float that gets rounded down + # - n * alpha is a float that gest rounded up + n = 102 # n * alpha = 20.4, 51. , 91.8 + y = rng.random(n) + w = rng.integers(low=0, high=10, size=n) if weights else None + q = np.quantile(y, alpha, method=method, weights=w) + c = 13.5 + + # 1 + assert_allclose(np.quantile(c + y, alpha, method=method, weights=w), + c + q) + # 2 + assert_allclose(np.quantile(c * y, alpha, method=method, weights=w), + c * q) + # 3 + if weights: + # From here on, we would need more methods to support weights. + return + q = -np.quantile(-y, 1 - alpha, method=method) + if method == "inverted_cdf": + if ( + n * alpha == int(n * alpha) + or np.round(n * alpha) == int(n * alpha) + 1 + ): + assert_allclose(q, np.quantile(y, alpha, method="higher")) + else: + assert_allclose(q, np.quantile(y, alpha, method="lower")) + elif method == "closest_observation": + if n * alpha == int(n * alpha): + assert_allclose(q, np.quantile(y, alpha, method="higher")) + elif np.round(n * alpha) == int(n * alpha) + 1: + assert_allclose( + q, np.quantile(y, alpha + 1/n, method="higher")) + else: + assert_allclose(q, np.quantile(y, alpha, method="lower")) + elif method == "interpolated_inverted_cdf": + assert_allclose(q, np.quantile(y, alpha + 1/n, method=method)) + elif method == "nearest": + if n * alpha == int(n * alpha): + assert_allclose(q, np.quantile(y, alpha + 1/n, method=method)) + else: + assert_allclose(q, np.quantile(y, alpha, method=method)) + elif method == "lower": + assert_allclose(q, np.quantile(y, alpha, method="higher")) + elif method == "higher": + assert_allclose(q, np.quantile(y, alpha, method="lower")) + else: + # "averaged_inverted_cdf", "hazen", "weibull", "linear", + # "median_unbiased", "normal_unbiased", "midpoint" + assert_allclose(q, np.quantile(y, alpha, method=method)) + + @pytest.mark.parametrize("method", methods_supporting_weights) + @pytest.mark.parametrize("alpha", [0.2, 0.5, 0.9]) + def test_quantile_constant_weights(self, method, alpha): + rng = np.random.default_rng(4321) + # We choose n and alpha such that we have cases for + # - n * alpha is an integer + # - n * alpha is a float that gets rounded down + # - n * alpha is a float that gest rounded up + n = 102 # n * alpha = 20.4, 51. , 91.8 + y = rng.random(n) + q = np.quantile(y, alpha, method=method) + + w = np.ones_like(y) + qw = np.quantile(y, alpha, method=method, weights=w) + assert_allclose(qw, q) + + w = 8.125 * np.ones_like(y) + qw = np.quantile(y, alpha, method=method, weights=w) + assert_allclose(qw, q) + + @pytest.mark.parametrize("method", methods_supporting_weights) + @pytest.mark.parametrize("alpha", [0, 0.2, 0.5, 0.9, 1]) + def test_quantile_with_integer_weights(self, method, alpha): + # Integer weights can be interpreted as repeated observations. + rng = np.random.default_rng(4321) + # We choose n and alpha such that we have cases for + # - n * alpha is an integer + # - n * alpha is a float that gets rounded down + # - n * alpha is a float that gest rounded up + n = 102 # n * alpha = 20.4, 51. , 91.8 + y = rng.random(n) + w = rng.integers(low=0, high=10, size=n, dtype=np.int32) + + qw = np.quantile(y, alpha, method=method, weights=w) + q = np.quantile(np.repeat(y, w), alpha, method=method) + assert_allclose(qw, q) + + @pytest.mark.parametrize("method", methods_supporting_weights) + def test_quantile_with_weights_and_axis(self, method): + rng = np.random.default_rng(4321) + + # 1d weight and single alpha + y = rng.random((2, 10, 3)) + w = np.abs(rng.random(10)) + alpha = 0.5 + q = np.quantile(y, alpha, weights=w, method=method, axis=1) + q_res = np.zeros(shape=(2, 3)) + for i in range(2): + for j in range(3): + q_res[i, j] = np.quantile( + y[i, :, j], alpha, method=method, weights=w + ) + assert_allclose(q, q_res) + + # 1d weight and 1d alpha + alpha = [0, 0.2, 0.4, 0.6, 0.8, 1] # shape (6,) + q = np.quantile(y, alpha, weights=w, method=method, axis=1) + q_res = np.zeros(shape=(6, 2, 3)) + for i in range(2): + for j in range(3): + q_res[:, i, j] = np.quantile( + y[i, :, j], alpha, method=method, weights=w + ) + assert_allclose(q, q_res) + + # 1d weight and 2d alpha + alpha = [[0, 0.2], [0.4, 0.6], [0.8, 1]] # shape (3, 2) + q = np.quantile(y, alpha, weights=w, method=method, axis=1) + q_res = q_res.reshape((3, 2, 2, 3)) + assert_allclose(q, q_res) + + # shape of weights equals shape of y + w = np.abs(rng.random((2, 10, 3))) + alpha = 0.5 + q = np.quantile(y, alpha, weights=w, method=method, axis=1) + q_res = np.zeros(shape=(2, 3)) + for i in range(2): + for j in range(3): + q_res[i, j] = np.quantile( + y[i, :, j], alpha, method=method, weights=w[i, :, j] + ) + assert_allclose(q, q_res) + + @pytest.mark.parametrize("method", methods_supporting_weights) + def test_quantile_weights_min_max(self, method): + # Test weighted quantile at 0 and 1 with leading and trailing zero + # weights. + w = [0, 0, 1, 2, 3, 0] + y = np.arange(6) + y_min = np.quantile(y, 0, weights=w, method="inverted_cdf") + y_max = np.quantile(y, 1, weights=w, method="inverted_cdf") + assert y_min == y[2] # == 2 + assert y_max == y[4] # == 4 + + def test_quantile_weights_raises_negative_weights(self): + y = [1, 2] + w = [-0.5, 1] + with pytest.raises(ValueError, match="Weights must be non-negative"): + np.quantile(y, 0.5, weights=w, method="inverted_cdf") + + @pytest.mark.parametrize( + "method", + sorted(set(quantile_methods) - set(methods_supporting_weights)), + ) + def test_quantile_weights_raises_unsupported_methods(self, method): + y = [1, 2] + w = [0.5, 1] + msg = "Only method 'inverted_cdf' supports weights" + with pytest.raises(ValueError, match=msg): + np.quantile(y, 0.5, weights=w, method=method) + + def test_weibull_fraction(self): + arr = [Fraction(0, 1), Fraction(1, 10)] + quantile = np.quantile(arr, [0, ], method='weibull') + assert_equal(quantile, np.array(Fraction(0, 1))) + quantile = np.quantile(arr, [Fraction(1, 2)], method='weibull') + assert_equal(quantile, np.array(Fraction(1, 20))) + + def test_closest_observation(self): + # Round ties to nearest even order statistic (see #26656) + m = 'closest_observation' + q = 0.5 + arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + assert_equal(2, np.quantile(arr[0:3], q, method=m)) + assert_equal(2, np.quantile(arr[0:4], q, method=m)) + assert_equal(2, np.quantile(arr[0:5], q, method=m)) + assert_equal(3, np.quantile(arr[0:6], q, method=m)) + assert_equal(4, np.quantile(arr[0:7], q, method=m)) + assert_equal(4, np.quantile(arr[0:8], q, method=m)) + assert_equal(4, np.quantile(arr[0:9], q, method=m)) + assert_equal(5, np.quantile(arr, q, method=m)) + + +class TestLerp: + @hypothesis.given(t0=st.floats(allow_nan=False, allow_infinity=False, + min_value=0, max_value=1), + t1=st.floats(allow_nan=False, allow_infinity=False, + min_value=0, max_value=1), + a = st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300), + b = st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300)) + def test_linear_interpolation_formula_monotonic(self, t0, t1, a, b): + l0 = nfb._lerp(a, b, t0) + l1 = nfb._lerp(a, b, t1) + if t0 == t1 or a == b: + assert l0 == l1 # uninteresting + elif (t0 < t1) == (a < b): + assert l0 <= l1 + else: + assert l0 >= l1 + + @hypothesis.given(t=st.floats(allow_nan=False, allow_infinity=False, + min_value=0, max_value=1), + a=st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300), + b=st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300)) + def test_linear_interpolation_formula_bounded(self, t, a, b): + if a <= b: + assert a <= nfb._lerp(a, b, t) <= b + else: + assert b <= nfb._lerp(a, b, t) <= a + + @hypothesis.given(t=st.floats(allow_nan=False, allow_infinity=False, + min_value=0, max_value=1), + a=st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300), + b=st.floats(allow_nan=False, allow_infinity=False, + min_value=-1e300, max_value=1e300)) + def test_linear_interpolation_formula_symmetric(self, t, a, b): + # double subtraction is needed to remove the extra precision of t < 0.5 + left = nfb._lerp(a, b, 1 - (1 - t)) + right = nfb._lerp(b, a, 1 - t) + assert_allclose(left, right) + + def test_linear_interpolation_formula_0d_inputs(self): + a = np.array(2) + b = np.array(5) + t = np.array(0.2) + assert nfb._lerp(a, b, t) == 2.6 + + +class TestMedian: + + def test_basic(self): + a0 = np.array(1) + a1 = np.arange(2) + a2 = np.arange(6).reshape(2, 3) + assert_equal(np.median(a0), 1) + assert_allclose(np.median(a1), 0.5) + assert_allclose(np.median(a2), 2.5) + assert_allclose(np.median(a2, axis=0), [1.5, 2.5, 3.5]) + assert_equal(np.median(a2, axis=1), [1, 4]) + assert_allclose(np.median(a2, axis=None), 2.5) + + a = np.array([0.0444502, 0.0463301, 0.141249, 0.0606775]) + assert_almost_equal((a[1] + a[3]) / 2., np.median(a)) + a = np.array([0.0463301, 0.0444502, 0.141249]) + assert_equal(a[0], np.median(a)) + a = np.array([0.0444502, 0.141249, 0.0463301]) + assert_equal(a[-1], np.median(a)) + # check array scalar result + assert_equal(np.median(a).ndim, 0) + a[1] = np.nan + assert_equal(np.median(a).ndim, 0) + + def test_axis_keyword(self): + a3 = np.array([[2, 3], + [0, 1], + [6, 7], + [4, 5]]) + for a in [a3, np.random.randint(0, 100, size=(2, 3, 4))]: + orig = a.copy() + np.median(a, axis=None) + for ax in range(a.ndim): + np.median(a, axis=ax) + assert_array_equal(a, orig) + + assert_allclose(np.median(a3, axis=0), [3, 4]) + assert_allclose(np.median(a3.T, axis=1), [3, 4]) + assert_allclose(np.median(a3), 3.5) + assert_allclose(np.median(a3, axis=None), 3.5) + assert_allclose(np.median(a3.T), 3.5) + + def test_overwrite_keyword(self): + a3 = np.array([[2, 3], + [0, 1], + [6, 7], + [4, 5]]) + a0 = np.array(1) + a1 = np.arange(2) + a2 = np.arange(6).reshape(2, 3) + assert_allclose(np.median(a0.copy(), overwrite_input=True), 1) + assert_allclose(np.median(a1.copy(), overwrite_input=True), 0.5) + assert_allclose(np.median(a2.copy(), overwrite_input=True), 2.5) + assert_allclose(np.median(a2.copy(), overwrite_input=True, axis=0), + [1.5, 2.5, 3.5]) + assert_allclose( + np.median(a2.copy(), overwrite_input=True, axis=1), [1, 4]) + assert_allclose( + np.median(a2.copy(), overwrite_input=True, axis=None), 2.5) + assert_allclose( + np.median(a3.copy(), overwrite_input=True, axis=0), [3, 4]) + assert_allclose(np.median(a3.T.copy(), overwrite_input=True, axis=1), + [3, 4]) + + a4 = np.arange(3 * 4 * 5, dtype=np.float32).reshape((3, 4, 5)) + np.random.shuffle(a4.ravel()) + assert_allclose(np.median(a4, axis=None), + np.median(a4.copy(), axis=None, overwrite_input=True)) + assert_allclose(np.median(a4, axis=0), + np.median(a4.copy(), axis=0, overwrite_input=True)) + assert_allclose(np.median(a4, axis=1), + np.median(a4.copy(), axis=1, overwrite_input=True)) + assert_allclose(np.median(a4, axis=2), + np.median(a4.copy(), axis=2, overwrite_input=True)) + + def test_array_like(self): + x = [1, 2, 3] + assert_almost_equal(np.median(x), 2) + x2 = [x] + assert_almost_equal(np.median(x2), 2) + assert_allclose(np.median(x2, axis=0), x) + + def test_subclass(self): + # gh-3846 + class MySubClass(np.ndarray): + + def __new__(cls, input_array, info=None): + obj = np.asarray(input_array).view(cls) + obj.info = info + return obj + + def mean(self, axis=None, dtype=None, out=None): + return -7 + + a = MySubClass([1, 2, 3]) + assert_equal(np.median(a), -7) + + @pytest.mark.parametrize('arr', + ([1., 2., 3.], [1., np.nan, 3.], np.nan, 0.)) + def test_subclass2(self, arr): + """Check that we return subclasses, even if a NaN scalar.""" + class MySubclass(np.ndarray): + pass + + m = np.median(np.array(arr).view(MySubclass)) + assert isinstance(m, MySubclass) + + def test_out(self): + o = np.zeros((4,)) + d = np.ones((3, 4)) + assert_equal(np.median(d, 0, out=o), o) + o = np.zeros((3,)) + assert_equal(np.median(d, 1, out=o), o) + o = np.zeros(()) + assert_equal(np.median(d, out=o), o) + + def test_out_nan(self): + with warnings.catch_warnings(record=True): + warnings.filterwarnings('always', '', RuntimeWarning) + o = np.zeros((4,)) + d = np.ones((3, 4)) + d[2, 1] = np.nan + assert_equal(np.median(d, 0, out=o), o) + o = np.zeros((3,)) + assert_equal(np.median(d, 1, out=o), o) + o = np.zeros(()) + assert_equal(np.median(d, out=o), o) + + def test_nan_behavior(self): + a = np.arange(24, dtype=float) + a[2] = np.nan + assert_equal(np.median(a), np.nan) + assert_equal(np.median(a, axis=0), np.nan) + + a = np.arange(24, dtype=float).reshape(2, 3, 4) + a[1, 2, 3] = np.nan + a[1, 1, 2] = np.nan + + # no axis + assert_equal(np.median(a), np.nan) + assert_equal(np.median(a).ndim, 0) + + # axis0 + b = np.median(np.arange(24, dtype=float).reshape(2, 3, 4), 0) + b[2, 3] = np.nan + b[1, 2] = np.nan + assert_equal(np.median(a, 0), b) + + # axis1 + b = np.median(np.arange(24, dtype=float).reshape(2, 3, 4), 1) + b[1, 3] = np.nan + b[1, 2] = np.nan + assert_equal(np.median(a, 1), b) + + # axis02 + b = np.median(np.arange(24, dtype=float).reshape(2, 3, 4), (0, 2)) + b[1] = np.nan + b[2] = np.nan + assert_equal(np.median(a, (0, 2)), b) + + @pytest.mark.skipif(IS_WASM, reason="fp errors don't work correctly") + def test_empty(self): + # mean(empty array) emits two warnings: empty slice and divide by 0 + a = np.array([], dtype=float) + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_equal(np.median(a), np.nan) + assert_(w[0].category is RuntimeWarning) + assert_equal(len(w), 2) + + # multiple dimensions + a = np.array([], dtype=float, ndmin=3) + # no axis + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_equal(np.median(a), np.nan) + assert_(w[0].category is RuntimeWarning) + + # axis 0 and 1 + b = np.array([], dtype=float, ndmin=2) + assert_equal(np.median(a, axis=0), b) + assert_equal(np.median(a, axis=1), b) + + # axis 2 + b = np.array(np.nan, dtype=float, ndmin=2) + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', RuntimeWarning) + assert_equal(np.median(a, axis=2), b) + assert_(w[0].category is RuntimeWarning) + + def test_object(self): + o = np.arange(7.) + assert_(type(np.median(o.astype(object))), float) + o[2] = np.nan + assert_(type(np.median(o.astype(object))), float) + + def test_extended_axis(self): + o = np.random.normal(size=(71, 23)) + x = np.dstack([o] * 10) + assert_equal(np.median(x, axis=(0, 1)), np.median(o)) + x = np.moveaxis(x, -1, 0) + assert_equal(np.median(x, axis=(-2, -1)), np.median(o)) + x = x.swapaxes(0, 1).copy() + assert_equal(np.median(x, axis=(0, -1)), np.median(o)) + + assert_equal(np.median(x, axis=(0, 1, 2)), np.median(x, axis=None)) + assert_equal(np.median(x, axis=(0, )), np.median(x, axis=0)) + assert_equal(np.median(x, axis=(-1, )), np.median(x, axis=-1)) + + d = np.arange(3 * 5 * 7 * 11).reshape((3, 5, 7, 11)) + np.random.shuffle(d.ravel()) + assert_equal(np.median(d, axis=(0, 1, 2))[0], + np.median(d[:,:,:, 0].flatten())) + assert_equal(np.median(d, axis=(0, 1, 3))[1], + np.median(d[:,:, 1,:].flatten())) + assert_equal(np.median(d, axis=(3, 1, -4))[2], + np.median(d[:,:, 2,:].flatten())) + assert_equal(np.median(d, axis=(3, 1, 2))[2], + np.median(d[2,:,:,:].flatten())) + assert_equal(np.median(d, axis=(3, 2))[2, 1], + np.median(d[2, 1,:,:].flatten())) + assert_equal(np.median(d, axis=(1, -2))[2, 1], + np.median(d[2,:,:, 1].flatten())) + assert_equal(np.median(d, axis=(1, 3))[2, 2], + np.median(d[2,:, 2,:].flatten())) + + def test_extended_axis_invalid(self): + d = np.ones((3, 5, 7, 11)) + assert_raises(AxisError, np.median, d, axis=-5) + assert_raises(AxisError, np.median, d, axis=(0, -5)) + assert_raises(AxisError, np.median, d, axis=4) + assert_raises(AxisError, np.median, d, axis=(0, 4)) + assert_raises(ValueError, np.median, d, axis=(1, 1)) + + def test_keepdims(self): + d = np.ones((3, 5, 7, 11)) + assert_equal(np.median(d, axis=None, keepdims=True).shape, + (1, 1, 1, 1)) + assert_equal(np.median(d, axis=(0, 1), keepdims=True).shape, + (1, 1, 7, 11)) + assert_equal(np.median(d, axis=(0, 3), keepdims=True).shape, + (1, 5, 7, 1)) + assert_equal(np.median(d, axis=(1,), keepdims=True).shape, + (3, 1, 7, 11)) + assert_equal(np.median(d, axis=(0, 1, 2, 3), keepdims=True).shape, + (1, 1, 1, 1)) + assert_equal(np.median(d, axis=(0, 1, 3), keepdims=True).shape, + (1, 1, 7, 1)) + + @pytest.mark.parametrize( + argnames='axis', + argvalues=[ + None, + 1, + (1, ), + (0, 1), + (-3, -1), + ] + ) + def test_keepdims_out(self, axis): + d = np.ones((3, 5, 7, 11)) + if axis is None: + shape_out = (1,) * d.ndim + else: + axis_norm = normalize_axis_tuple(axis, d.ndim) + shape_out = tuple( + 1 if i in axis_norm else d.shape[i] for i in range(d.ndim)) + out = np.empty(shape_out) + result = np.median(d, axis=axis, keepdims=True, out=out) + assert result is out + assert_equal(result.shape, shape_out) + + @pytest.mark.parametrize("dtype", ["m8[s]"]) + @pytest.mark.parametrize("pos", [0, 23, 10]) + def test_nat_behavior(self, dtype, pos): + # TODO: Median does not support Datetime, due to `mean`. + # NaT and NaN should behave the same, do basic tests for NaT. + a = np.arange(0, 24, dtype=dtype) + a[pos] = "NaT" + res = np.median(a) + assert res.dtype == dtype + assert np.isnat(res) + res = np.percentile(a, [30, 60]) + assert res.dtype == dtype + assert np.isnat(res).all() + + a = np.arange(0, 24*3, dtype=dtype).reshape(-1, 3) + a[pos, 1] = "NaT" + res = np.median(a, axis=0) + assert_array_equal(np.isnat(res), [False, True, False]) + + +class TestSortComplex: + + @pytest.mark.parametrize("type_in, type_out", [ + ('l', 'D'), + ('h', 'F'), + ('H', 'F'), + ('b', 'F'), + ('B', 'F'), + ('g', 'G'), + ]) + def test_sort_real(self, type_in, type_out): + # sort_complex() type casting for real input types + a = np.array([5, 3, 6, 2, 1], dtype=type_in) + actual = np.sort_complex(a) + expected = np.sort(a).astype(type_out) + assert_equal(actual, expected) + assert_equal(actual.dtype, expected.dtype) + + def test_sort_complex(self): + # sort_complex() handling of complex input + a = np.array([2 + 3j, 1 - 2j, 1 - 3j, 2 + 1j], dtype='D') + expected = np.array([1 - 3j, 1 - 2j, 2 + 1j, 2 + 3j], dtype='D') + actual = np.sort_complex(a) + assert_equal(actual, expected) + assert_equal(actual.dtype, expected.dtype) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_histograms.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_histograms.py new file mode 100644 index 0000000000000000000000000000000000000000..4b300624cac76dec3b247c7a7609da440746dde5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_histograms.py @@ -0,0 +1,841 @@ +import numpy as np + +from numpy import histogram, histogramdd, histogram_bin_edges +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_almost_equal, + assert_array_almost_equal, assert_raises, assert_allclose, + assert_array_max_ulp, assert_raises_regex, suppress_warnings, + ) +import pytest + + +class TestHistogram: + + def setup_method(self): + pass + + def teardown_method(self): + pass + + def test_simple(self): + n = 100 + v = np.random.rand(n) + (a, b) = histogram(v) + # check if the sum of the bins equals the number of samples + assert_equal(np.sum(a, axis=0), n) + # check that the bin counts are evenly spaced when the data is from + # a linear function + (a, b) = histogram(np.linspace(0, 10, 100)) + assert_array_equal(a, 10) + + def test_one_bin(self): + # Ticket 632 + hist, edges = histogram([1, 2, 3, 4], [1, 2]) + assert_array_equal(hist, [2, ]) + assert_array_equal(edges, [1, 2]) + assert_raises(ValueError, histogram, [1, 2], bins=0) + h, e = histogram([1, 2], bins=1) + assert_equal(h, np.array([2])) + assert_allclose(e, np.array([1., 2.])) + + def test_density(self): + # Check that the integral of the density equals 1. + n = 100 + v = np.random.rand(n) + a, b = histogram(v, density=True) + area = np.sum(a * np.diff(b)) + assert_almost_equal(area, 1) + + # Check with non-constant bin widths + v = np.arange(10) + bins = [0, 1, 3, 6, 10] + a, b = histogram(v, bins, density=True) + assert_array_equal(a, .1) + assert_equal(np.sum(a * np.diff(b)), 1) + + # Test that passing False works too + a, b = histogram(v, bins, density=False) + assert_array_equal(a, [1, 2, 3, 4]) + + # Variable bin widths are especially useful to deal with + # infinities. + v = np.arange(10) + bins = [0, 1, 3, 6, np.inf] + a, b = histogram(v, bins, density=True) + assert_array_equal(a, [.1, .1, .1, 0.]) + + # Taken from a bug report from N. Becker on the numpy-discussion + # mailing list Aug. 6, 2010. + counts, dmy = np.histogram( + [1, 2, 3, 4], [0.5, 1.5, np.inf], density=True) + assert_equal(counts, [.25, 0]) + + def test_outliers(self): + # Check that outliers are not tallied + a = np.arange(10) + .5 + + # Lower outliers + h, b = histogram(a, range=[0, 9]) + assert_equal(h.sum(), 9) + + # Upper outliers + h, b = histogram(a, range=[1, 10]) + assert_equal(h.sum(), 9) + + # Normalization + h, b = histogram(a, range=[1, 9], density=True) + assert_almost_equal((h * np.diff(b)).sum(), 1, decimal=15) + + # Weights + w = np.arange(10) + .5 + h, b = histogram(a, range=[1, 9], weights=w, density=True) + assert_equal((h * np.diff(b)).sum(), 1) + + h, b = histogram(a, bins=8, range=[1, 9], weights=w) + assert_equal(h, w[1:-1]) + + def test_arr_weights_mismatch(self): + a = np.arange(10) + .5 + w = np.arange(11) + .5 + with assert_raises_regex(ValueError, "same shape as"): + h, b = histogram(a, range=[1, 9], weights=w, density=True) + + + def test_type(self): + # Check the type of the returned histogram + a = np.arange(10) + .5 + h, b = histogram(a) + assert_(np.issubdtype(h.dtype, np.integer)) + + h, b = histogram(a, density=True) + assert_(np.issubdtype(h.dtype, np.floating)) + + h, b = histogram(a, weights=np.ones(10, int)) + assert_(np.issubdtype(h.dtype, np.integer)) + + h, b = histogram(a, weights=np.ones(10, float)) + assert_(np.issubdtype(h.dtype, np.floating)) + + def test_f32_rounding(self): + # gh-4799, check that the rounding of the edges works with float32 + x = np.array([276.318359, -69.593948, 21.329449], dtype=np.float32) + y = np.array([5005.689453, 4481.327637, 6010.369629], dtype=np.float32) + counts_hist, xedges, yedges = np.histogram2d(x, y, bins=100) + assert_equal(counts_hist.sum(), 3.) + + def test_bool_conversion(self): + # gh-12107 + # Reference integer histogram + a = np.array([1, 1, 0], dtype=np.uint8) + int_hist, int_edges = np.histogram(a) + + # Should raise an warning on booleans + # Ensure that the histograms are equivalent, need to suppress + # the warnings to get the actual outputs + with suppress_warnings() as sup: + rec = sup.record(RuntimeWarning, 'Converting input from .*') + hist, edges = np.histogram([True, True, False]) + # A warning should be issued + assert_equal(len(rec), 1) + assert_array_equal(hist, int_hist) + assert_array_equal(edges, int_edges) + + def test_weights(self): + v = np.random.rand(100) + w = np.ones(100) * 5 + a, b = histogram(v) + na, nb = histogram(v, density=True) + wa, wb = histogram(v, weights=w) + nwa, nwb = histogram(v, weights=w, density=True) + assert_array_almost_equal(a * 5, wa) + assert_array_almost_equal(na, nwa) + + # Check weights are properly applied. + v = np.linspace(0, 10, 10) + w = np.concatenate((np.zeros(5), np.ones(5))) + wa, wb = histogram(v, bins=np.arange(11), weights=w) + assert_array_almost_equal(wa, w) + + # Check with integer weights + wa, wb = histogram([1, 2, 2, 4], bins=4, weights=[4, 3, 2, 1]) + assert_array_equal(wa, [4, 5, 0, 1]) + wa, wb = histogram( + [1, 2, 2, 4], bins=4, weights=[4, 3, 2, 1], density=True) + assert_array_almost_equal(wa, np.array([4, 5, 0, 1]) / 10. / 3. * 4) + + # Check weights with non-uniform bin widths + a, b = histogram( + np.arange(9), [0, 1, 3, 6, 10], + weights=[2, 1, 1, 1, 1, 1, 1, 1, 1], density=True) + assert_almost_equal(a, [.2, .1, .1, .075]) + + def test_exotic_weights(self): + + # Test the use of weights that are not integer or floats, but e.g. + # complex numbers or object types. + + # Complex weights + values = np.array([1.3, 2.5, 2.3]) + weights = np.array([1, -1, 2]) + 1j * np.array([2, 1, 2]) + + # Check with custom bins + wa, wb = histogram(values, bins=[0, 2, 3], weights=weights) + assert_array_almost_equal(wa, np.array([1, 1]) + 1j * np.array([2, 3])) + + # Check with even bins + wa, wb = histogram(values, bins=2, range=[1, 3], weights=weights) + assert_array_almost_equal(wa, np.array([1, 1]) + 1j * np.array([2, 3])) + + # Decimal weights + from decimal import Decimal + values = np.array([1.3, 2.5, 2.3]) + weights = np.array([Decimal(1), Decimal(2), Decimal(3)]) + + # Check with custom bins + wa, wb = histogram(values, bins=[0, 2, 3], weights=weights) + assert_array_almost_equal(wa, [Decimal(1), Decimal(5)]) + + # Check with even bins + wa, wb = histogram(values, bins=2, range=[1, 3], weights=weights) + assert_array_almost_equal(wa, [Decimal(1), Decimal(5)]) + + def test_no_side_effects(self): + # This is a regression test that ensures that values passed to + # ``histogram`` are unchanged. + values = np.array([1.3, 2.5, 2.3]) + np.histogram(values, range=[-10, 10], bins=100) + assert_array_almost_equal(values, [1.3, 2.5, 2.3]) + + def test_empty(self): + a, b = histogram([], bins=([0, 1])) + assert_array_equal(a, np.array([0])) + assert_array_equal(b, np.array([0, 1])) + + def test_error_binnum_type (self): + # Tests if right Error is raised if bins argument is float + vals = np.linspace(0.0, 1.0, num=100) + histogram(vals, 5) + assert_raises(TypeError, histogram, vals, 2.4) + + def test_finite_range(self): + # Normal ranges should be fine + vals = np.linspace(0.0, 1.0, num=100) + histogram(vals, range=[0.25,0.75]) + assert_raises(ValueError, histogram, vals, range=[np.nan,0.75]) + assert_raises(ValueError, histogram, vals, range=[0.25,np.inf]) + + def test_invalid_range(self): + # start of range must be < end of range + vals = np.linspace(0.0, 1.0, num=100) + with assert_raises_regex(ValueError, "max must be larger than"): + np.histogram(vals, range=[0.1, 0.01]) + + def test_bin_edge_cases(self): + # Ensure that floating-point computations correctly place edge cases. + arr = np.array([337, 404, 739, 806, 1007, 1811, 2012]) + hist, edges = np.histogram(arr, bins=8296, range=(2, 2280)) + mask = hist > 0 + left_edges = edges[:-1][mask] + right_edges = edges[1:][mask] + for x, left, right in zip(arr, left_edges, right_edges): + assert_(x >= left) + assert_(x < right) + + def test_last_bin_inclusive_range(self): + arr = np.array([0., 0., 0., 1., 2., 3., 3., 4., 5.]) + hist, edges = np.histogram(arr, bins=30, range=(-0.5, 5)) + assert_equal(hist[-1], 1) + + def test_bin_array_dims(self): + # gracefully handle bins object > 1 dimension + vals = np.linspace(0.0, 1.0, num=100) + bins = np.array([[0, 0.5], [0.6, 1.0]]) + with assert_raises_regex(ValueError, "must be 1d"): + np.histogram(vals, bins=bins) + + def test_unsigned_monotonicity_check(self): + # Ensures ValueError is raised if bins not increasing monotonically + # when bins contain unsigned values (see #9222) + arr = np.array([2]) + bins = np.array([1, 3, 1], dtype='uint64') + with assert_raises(ValueError): + hist, edges = np.histogram(arr, bins=bins) + + def test_object_array_of_0d(self): + # gh-7864 + assert_raises(ValueError, + histogram, [np.array(0.4) for i in range(10)] + [-np.inf]) + assert_raises(ValueError, + histogram, [np.array(0.4) for i in range(10)] + [np.inf]) + + # these should not crash + np.histogram([np.array(0.5) for i in range(10)] + [.500000000000002]) + np.histogram([np.array(0.5) for i in range(10)] + [.5]) + + def test_some_nan_values(self): + # gh-7503 + one_nan = np.array([0, 1, np.nan]) + all_nan = np.array([np.nan, np.nan]) + + # the internal comparisons with NaN give warnings + sup = suppress_warnings() + sup.filter(RuntimeWarning) + with sup: + # can't infer range with nan + assert_raises(ValueError, histogram, one_nan, bins='auto') + assert_raises(ValueError, histogram, all_nan, bins='auto') + + # explicit range solves the problem + h, b = histogram(one_nan, bins='auto', range=(0, 1)) + assert_equal(h.sum(), 2) # nan is not counted + h, b = histogram(all_nan, bins='auto', range=(0, 1)) + assert_equal(h.sum(), 0) # nan is not counted + + # as does an explicit set of bins + h, b = histogram(one_nan, bins=[0, 1]) + assert_equal(h.sum(), 2) # nan is not counted + h, b = histogram(all_nan, bins=[0, 1]) + assert_equal(h.sum(), 0) # nan is not counted + + def test_datetime(self): + begin = np.datetime64('2000-01-01', 'D') + offsets = np.array([0, 0, 1, 1, 2, 3, 5, 10, 20]) + bins = np.array([0, 2, 7, 20]) + dates = begin + offsets + date_bins = begin + bins + + td = np.dtype('timedelta64[D]') + + # Results should be the same for integer offsets or datetime values. + # For now, only explicit bins are supported, since linspace does not + # work on datetimes or timedeltas + d_count, d_edge = histogram(dates, bins=date_bins) + t_count, t_edge = histogram(offsets.astype(td), bins=bins.astype(td)) + i_count, i_edge = histogram(offsets, bins=bins) + + assert_equal(d_count, i_count) + assert_equal(t_count, i_count) + + assert_equal((d_edge - begin).astype(int), i_edge) + assert_equal(t_edge.astype(int), i_edge) + + assert_equal(d_edge.dtype, dates.dtype) + assert_equal(t_edge.dtype, td) + + def do_signed_overflow_bounds(self, dtype): + exponent = 8 * np.dtype(dtype).itemsize - 1 + arr = np.array([-2**exponent + 4, 2**exponent - 4], dtype=dtype) + hist, e = histogram(arr, bins=2) + assert_equal(e, [-2**exponent + 4, 0, 2**exponent - 4]) + assert_equal(hist, [1, 1]) + + def test_signed_overflow_bounds(self): + self.do_signed_overflow_bounds(np.byte) + self.do_signed_overflow_bounds(np.short) + self.do_signed_overflow_bounds(np.intc) + self.do_signed_overflow_bounds(np.int_) + self.do_signed_overflow_bounds(np.longlong) + + def do_precision_lower_bound(self, float_small, float_large): + eps = np.finfo(float_large).eps + + arr = np.array([1.0], float_small) + range = np.array([1.0 + eps, 2.0], float_large) + + # test is looking for behavior when the bounds change between dtypes + if range.astype(float_small)[0] != 1: + return + + # previously crashed + count, x_loc = np.histogram(arr, bins=1, range=range) + assert_equal(count, [0]) + assert_equal(x_loc.dtype, float_large) + + def do_precision_upper_bound(self, float_small, float_large): + eps = np.finfo(float_large).eps + + arr = np.array([1.0], float_small) + range = np.array([0.0, 1.0 - eps], float_large) + + # test is looking for behavior when the bounds change between dtypes + if range.astype(float_small)[-1] != 1: + return + + # previously crashed + count, x_loc = np.histogram(arr, bins=1, range=range) + assert_equal(count, [0]) + + assert_equal(x_loc.dtype, float_large) + + def do_precision(self, float_small, float_large): + self.do_precision_lower_bound(float_small, float_large) + self.do_precision_upper_bound(float_small, float_large) + + def test_precision(self): + # not looping results in a useful stack trace upon failure + self.do_precision(np.half, np.single) + self.do_precision(np.half, np.double) + self.do_precision(np.half, np.longdouble) + self.do_precision(np.single, np.double) + self.do_precision(np.single, np.longdouble) + self.do_precision(np.double, np.longdouble) + + def test_histogram_bin_edges(self): + hist, e = histogram([1, 2, 3, 4], [1, 2]) + edges = histogram_bin_edges([1, 2, 3, 4], [1, 2]) + assert_array_equal(edges, e) + + arr = np.array([0., 0., 0., 1., 2., 3., 3., 4., 5.]) + hist, e = histogram(arr, bins=30, range=(-0.5, 5)) + edges = histogram_bin_edges(arr, bins=30, range=(-0.5, 5)) + assert_array_equal(edges, e) + + hist, e = histogram(arr, bins='auto', range=(0, 1)) + edges = histogram_bin_edges(arr, bins='auto', range=(0, 1)) + assert_array_equal(edges, e) + + def test_small_value_range(self): + arr = np.array([1, 1 + 2e-16] * 10) + with pytest.raises(ValueError, match="Too many bins for data range"): + histogram(arr, bins=10) + + # @requires_memory(free_bytes=1e10) + # @pytest.mark.slow + @pytest.mark.skip(reason="Bad memory reports lead to OOM in ci testing") + def test_big_arrays(self): + sample = np.zeros([100000000, 3]) + xbins = 400 + ybins = 400 + zbins = np.arange(16000) + hist = np.histogramdd(sample=sample, bins=(xbins, ybins, zbins)) + assert_equal(type(hist), type((1, 2))) + + def test_gh_23110(self): + hist, e = np.histogram(np.array([-0.9e-308], dtype='>f8'), + bins=2, + range=(-1e-308, -2e-313)) + expected_hist = np.array([1, 0]) + assert_array_equal(hist, expected_hist) + + +class TestHistogramOptimBinNums: + """ + Provide test coverage when using provided estimators for optimal number of + bins + """ + + def test_empty(self): + estimator_list = ['fd', 'scott', 'rice', 'sturges', + 'doane', 'sqrt', 'auto', 'stone'] + # check it can deal with empty data + for estimator in estimator_list: + a, b = histogram([], bins=estimator) + assert_array_equal(a, np.array([0])) + assert_array_equal(b, np.array([0, 1])) + + def test_simple(self): + """ + Straightforward testing with a mixture of linspace data (for + consistency). All test values have been precomputed and the values + shouldn't change + """ + # Some basic sanity checking, with some fixed data. + # Checking for the correct number of bins + basic_test = {50: {'fd': 4, 'scott': 4, 'rice': 8, 'sturges': 7, + 'doane': 8, 'sqrt': 8, 'auto': 7, 'stone': 2}, + 500: {'fd': 8, 'scott': 8, 'rice': 16, 'sturges': 10, + 'doane': 12, 'sqrt': 23, 'auto': 10, 'stone': 9}, + 5000: {'fd': 17, 'scott': 17, 'rice': 35, 'sturges': 14, + 'doane': 17, 'sqrt': 71, 'auto': 17, 'stone': 20}} + + for testlen, expectedResults in basic_test.items(): + # Create some sort of non uniform data to test with + # (2 peak uniform mixture) + x1 = np.linspace(-10, -1, testlen // 5 * 2) + x2 = np.linspace(1, 10, testlen // 5 * 3) + x = np.concatenate((x1, x2)) + for estimator, numbins in expectedResults.items(): + a, b = np.histogram(x, estimator) + assert_equal(len(a), numbins, err_msg="For the {0} estimator " + "with datasize of {1}".format(estimator, testlen)) + + def test_small(self): + """ + Smaller datasets have the potential to cause issues with the data + adaptive methods, especially the FD method. All bin numbers have been + precalculated. + """ + small_dat = {1: {'fd': 1, 'scott': 1, 'rice': 1, 'sturges': 1, + 'doane': 1, 'sqrt': 1, 'stone': 1}, + 2: {'fd': 2, 'scott': 1, 'rice': 3, 'sturges': 2, + 'doane': 1, 'sqrt': 2, 'stone': 1}, + 3: {'fd': 2, 'scott': 2, 'rice': 3, 'sturges': 3, + 'doane': 3, 'sqrt': 2, 'stone': 1}} + + for testlen, expectedResults in small_dat.items(): + testdat = np.arange(testlen).astype(float) + for estimator, expbins in expectedResults.items(): + a, b = np.histogram(testdat, estimator) + assert_equal(len(a), expbins, err_msg="For the {0} estimator " + "with datasize of {1}".format(estimator, testlen)) + + def test_incorrect_methods(self): + """ + Check a Value Error is thrown when an unknown string is passed in + """ + check_list = ['mad', 'freeman', 'histograms', 'IQR'] + for estimator in check_list: + assert_raises(ValueError, histogram, [1, 2, 3], estimator) + + def test_novariance(self): + """ + Check that methods handle no variance in data + Primarily for Scott and FD as the SD and IQR are both 0 in this case + """ + novar_dataset = np.ones(100) + novar_resultdict = {'fd': 1, 'scott': 1, 'rice': 1, 'sturges': 1, + 'doane': 1, 'sqrt': 1, 'auto': 1, 'stone': 1} + + for estimator, numbins in novar_resultdict.items(): + a, b = np.histogram(novar_dataset, estimator) + assert_equal(len(a), numbins, err_msg="{0} estimator, " + "No Variance test".format(estimator)) + + def test_limited_variance(self): + """ + Check when IQR is 0, but variance exists, we return the sturges value + and not the fd value. + """ + lim_var_data = np.ones(1000) + lim_var_data[:3] = 0 + lim_var_data[-4:] = 100 + + edges_auto = histogram_bin_edges(lim_var_data, 'auto') + assert_equal(edges_auto, np.linspace(0, 100, 12)) + + edges_fd = histogram_bin_edges(lim_var_data, 'fd') + assert_equal(edges_fd, np.array([0, 100])) + + edges_sturges = histogram_bin_edges(lim_var_data, 'sturges') + assert_equal(edges_sturges, np.linspace(0, 100, 12)) + + def test_outlier(self): + """ + Check the FD, Scott and Doane with outliers. + + The FD estimates a smaller binwidth since it's less affected by + outliers. Since the range is so (artificially) large, this means more + bins, most of which will be empty, but the data of interest usually is + unaffected. The Scott estimator is more affected and returns fewer bins, + despite most of the variance being in one area of the data. The Doane + estimator lies somewhere between the other two. + """ + xcenter = np.linspace(-10, 10, 50) + outlier_dataset = np.hstack((np.linspace(-110, -100, 5), xcenter)) + + outlier_resultdict = {'fd': 21, 'scott': 5, 'doane': 11, 'stone': 6} + + for estimator, numbins in outlier_resultdict.items(): + a, b = np.histogram(outlier_dataset, estimator) + assert_equal(len(a), numbins) + + def test_scott_vs_stone(self): + """Verify that Scott's rule and Stone's rule converges for normally distributed data""" + + def nbins_ratio(seed, size): + rng = np.random.RandomState(seed) + x = rng.normal(loc=0, scale=2, size=size) + a, b = len(np.histogram(x, 'stone')[0]), len(np.histogram(x, 'scott')[0]) + return a / (a + b) + + ll = [[nbins_ratio(seed, size) for size in np.geomspace(start=10, stop=100, num=4).round().astype(int)] + for seed in range(10)] + + # the average difference between the two methods decreases as the dataset size increases. + avg = abs(np.mean(ll, axis=0) - 0.5) + assert_almost_equal(avg, [0.15, 0.09, 0.08, 0.03], decimal=2) + + def test_simple_range(self): + """ + Straightforward testing with a mixture of linspace data (for + consistency). Adding in a 3rd mixture that will then be + completely ignored. All test values have been precomputed and + the shouldn't change. + """ + # some basic sanity checking, with some fixed data. + # Checking for the correct number of bins + basic_test = { + 50: {'fd': 8, 'scott': 8, 'rice': 15, + 'sturges': 14, 'auto': 14, 'stone': 8}, + 500: {'fd': 15, 'scott': 16, 'rice': 32, + 'sturges': 20, 'auto': 20, 'stone': 80}, + 5000: {'fd': 33, 'scott': 33, 'rice': 69, + 'sturges': 27, 'auto': 33, 'stone': 80} + } + + for testlen, expectedResults in basic_test.items(): + # create some sort of non uniform data to test with + # (3 peak uniform mixture) + x1 = np.linspace(-10, -1, testlen // 5 * 2) + x2 = np.linspace(1, 10, testlen // 5 * 3) + x3 = np.linspace(-100, -50, testlen) + x = np.hstack((x1, x2, x3)) + for estimator, numbins in expectedResults.items(): + a, b = np.histogram(x, estimator, range = (-20, 20)) + msg = "For the {0} estimator".format(estimator) + msg += " with datasize of {0}".format(testlen) + assert_equal(len(a), numbins, err_msg=msg) + + @pytest.mark.parametrize("bins", ['auto', 'fd', 'doane', 'scott', + 'stone', 'rice', 'sturges']) + def test_signed_integer_data(self, bins): + # Regression test for gh-14379. + a = np.array([-2, 0, 127], dtype=np.int8) + hist, edges = np.histogram(a, bins=bins) + hist32, edges32 = np.histogram(a.astype(np.int32), bins=bins) + assert_array_equal(hist, hist32) + assert_array_equal(edges, edges32) + + @pytest.mark.parametrize("bins", ['auto', 'fd', 'doane', 'scott', + 'stone', 'rice', 'sturges']) + def test_integer(self, bins): + """ + Test that bin width for integer data is at least 1. + """ + with suppress_warnings() as sup: + if bins == 'stone': + sup.filter(RuntimeWarning) + assert_equal( + np.histogram_bin_edges(np.tile(np.arange(9), 1000), bins), + np.arange(9)) + + def test_integer_non_auto(self): + """ + Test that the bin-width>=1 requirement *only* applies to auto binning. + """ + assert_equal( + np.histogram_bin_edges(np.tile(np.arange(9), 1000), 16), + np.arange(17) / 2) + assert_equal( + np.histogram_bin_edges(np.tile(np.arange(9), 1000), [.1, .2]), + [.1, .2]) + + def test_simple_weighted(self): + """ + Check that weighted data raises a TypeError + """ + estimator_list = ['fd', 'scott', 'rice', 'sturges', 'auto'] + for estimator in estimator_list: + assert_raises(TypeError, histogram, [1, 2, 3], + estimator, weights=[1, 2, 3]) + + +class TestHistogramdd: + + def test_simple(self): + x = np.array([[-.5, .5, 1.5], [-.5, 1.5, 2.5], [-.5, 2.5, .5], + [.5, .5, 1.5], [.5, 1.5, 2.5], [.5, 2.5, 2.5]]) + H, edges = histogramdd(x, (2, 3, 3), + range=[[-1, 1], [0, 3], [0, 3]]) + answer = np.array([[[0, 1, 0], [0, 0, 1], [1, 0, 0]], + [[0, 1, 0], [0, 0, 1], [0, 0, 1]]]) + assert_array_equal(H, answer) + + # Check normalization + ed = [[-2, 0, 2], [0, 1, 2, 3], [0, 1, 2, 3]] + H, edges = histogramdd(x, bins=ed, density=True) + assert_(np.all(H == answer / 12.)) + + # Check that H has the correct shape. + H, edges = histogramdd(x, (2, 3, 4), + range=[[-1, 1], [0, 3], [0, 4]], + density=True) + answer = np.array([[[0, 1, 0, 0], [0, 0, 1, 0], [1, 0, 0, 0]], + [[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 1, 0]]]) + assert_array_almost_equal(H, answer / 6., 4) + # Check that a sequence of arrays is accepted and H has the correct + # shape. + z = [np.squeeze(y) for y in np.split(x, 3, axis=1)] + H, edges = histogramdd( + z, bins=(4, 3, 2), range=[[-2, 2], [0, 3], [0, 2]]) + answer = np.array([[[0, 0], [0, 0], [0, 0]], + [[0, 1], [0, 0], [1, 0]], + [[0, 1], [0, 0], [0, 0]], + [[0, 0], [0, 0], [0, 0]]]) + assert_array_equal(H, answer) + + Z = np.zeros((5, 5, 5)) + Z[list(range(5)), list(range(5)), list(range(5))] = 1. + H, edges = histogramdd([np.arange(5), np.arange(5), np.arange(5)], 5) + assert_array_equal(H, Z) + + def test_shape_3d(self): + # All possible permutations for bins of different lengths in 3D. + bins = ((5, 4, 6), (6, 4, 5), (5, 6, 4), (4, 6, 5), (6, 5, 4), + (4, 5, 6)) + r = np.random.rand(10, 3) + for b in bins: + H, edges = histogramdd(r, b) + assert_(H.shape == b) + + def test_shape_4d(self): + # All possible permutations for bins of different lengths in 4D. + bins = ((7, 4, 5, 6), (4, 5, 7, 6), (5, 6, 4, 7), (7, 6, 5, 4), + (5, 7, 6, 4), (4, 6, 7, 5), (6, 5, 7, 4), (7, 5, 4, 6), + (7, 4, 6, 5), (6, 4, 7, 5), (6, 7, 5, 4), (4, 6, 5, 7), + (4, 7, 5, 6), (5, 4, 6, 7), (5, 7, 4, 6), (6, 7, 4, 5), + (6, 5, 4, 7), (4, 7, 6, 5), (4, 5, 6, 7), (7, 6, 4, 5), + (5, 4, 7, 6), (5, 6, 7, 4), (6, 4, 5, 7), (7, 5, 6, 4)) + + r = np.random.rand(10, 4) + for b in bins: + H, edges = histogramdd(r, b) + assert_(H.shape == b) + + def test_weights(self): + v = np.random.rand(100, 2) + hist, edges = histogramdd(v) + n_hist, edges = histogramdd(v, density=True) + w_hist, edges = histogramdd(v, weights=np.ones(100)) + assert_array_equal(w_hist, hist) + w_hist, edges = histogramdd(v, weights=np.ones(100) * 2, density=True) + assert_array_equal(w_hist, n_hist) + w_hist, edges = histogramdd(v, weights=np.ones(100, int) * 2) + assert_array_equal(w_hist, 2 * hist) + + def test_identical_samples(self): + x = np.zeros((10, 2), int) + hist, edges = histogramdd(x, bins=2) + assert_array_equal(edges[0], np.array([-0.5, 0., 0.5])) + + def test_empty(self): + a, b = histogramdd([[], []], bins=([0, 1], [0, 1])) + assert_array_max_ulp(a, np.array([[0.]])) + a, b = np.histogramdd([[], [], []], bins=2) + assert_array_max_ulp(a, np.zeros((2, 2, 2))) + + def test_bins_errors(self): + # There are two ways to specify bins. Check for the right errors + # when mixing those. + x = np.arange(8).reshape(2, 4) + assert_raises(ValueError, np.histogramdd, x, bins=[-1, 2, 4, 5]) + assert_raises(ValueError, np.histogramdd, x, bins=[1, 0.99, 1, 1]) + assert_raises( + ValueError, np.histogramdd, x, bins=[1, 1, 1, [1, 2, 3, -3]]) + assert_(np.histogramdd(x, bins=[1, 1, 1, [1, 2, 3, 4]])) + + def test_inf_edges(self): + # Test using +/-inf bin edges works. See #1788. + with np.errstate(invalid='ignore'): + x = np.arange(6).reshape(3, 2) + expected = np.array([[1, 0], [0, 1], [0, 1]]) + h, e = np.histogramdd(x, bins=[3, [-np.inf, 2, 10]]) + assert_allclose(h, expected) + h, e = np.histogramdd(x, bins=[3, np.array([-1, 2, np.inf])]) + assert_allclose(h, expected) + h, e = np.histogramdd(x, bins=[3, [-np.inf, 3, np.inf]]) + assert_allclose(h, expected) + + def test_rightmost_binedge(self): + # Test event very close to rightmost binedge. See Github issue #4266 + x = [0.9999999995] + bins = [[0., 0.5, 1.0]] + hist, _ = histogramdd(x, bins=bins) + assert_(hist[0] == 0.0) + assert_(hist[1] == 1.) + x = [1.0] + bins = [[0., 0.5, 1.0]] + hist, _ = histogramdd(x, bins=bins) + assert_(hist[0] == 0.0) + assert_(hist[1] == 1.) + x = [1.0000000001] + bins = [[0., 0.5, 1.0]] + hist, _ = histogramdd(x, bins=bins) + assert_(hist[0] == 0.0) + assert_(hist[1] == 0.0) + x = [1.0001] + bins = [[0., 0.5, 1.0]] + hist, _ = histogramdd(x, bins=bins) + assert_(hist[0] == 0.0) + assert_(hist[1] == 0.0) + + def test_finite_range(self): + vals = np.random.random((100, 3)) + histogramdd(vals, range=[[0.0, 1.0], [0.25, 0.75], [0.25, 0.5]]) + assert_raises(ValueError, histogramdd, vals, + range=[[0.0, 1.0], [0.25, 0.75], [0.25, np.inf]]) + assert_raises(ValueError, histogramdd, vals, + range=[[0.0, 1.0], [np.nan, 0.75], [0.25, 0.5]]) + + def test_equal_edges(self): + """ Test that adjacent entries in an edge array can be equal """ + x = np.array([0, 1, 2]) + y = np.array([0, 1, 2]) + x_edges = np.array([0, 2, 2]) + y_edges = 1 + hist, edges = histogramdd((x, y), bins=(x_edges, y_edges)) + + hist_expected = np.array([ + [2.], + [1.], # x == 2 falls in the final bin + ]) + assert_equal(hist, hist_expected) + + def test_edge_dtype(self): + """ Test that if an edge array is input, its type is preserved """ + x = np.array([0, 10, 20]) + y = x / 10 + x_edges = np.array([0, 5, 15, 20]) + y_edges = x_edges / 10 + hist, edges = histogramdd((x, y), bins=(x_edges, y_edges)) + + assert_equal(edges[0].dtype, x_edges.dtype) + assert_equal(edges[1].dtype, y_edges.dtype) + + def test_large_integers(self): + big = 2**60 # Too large to represent with a full precision float + + x = np.array([0], np.int64) + x_edges = np.array([-1, +1], np.int64) + y = big + x + y_edges = big + x_edges + + hist, edges = histogramdd((x, y), bins=(x_edges, y_edges)) + + assert_equal(hist[0, 0], 1) + + def test_density_non_uniform_2d(self): + # Defines the following grid: + # + # 0 2 8 + # 0+-+-----+ + # + | + + # + | + + # 6+-+-----+ + # 8+-+-----+ + x_edges = np.array([0, 2, 8]) + y_edges = np.array([0, 6, 8]) + relative_areas = np.array([ + [3, 9], + [1, 3]]) + + # ensure the number of points in each region is proportional to its area + x = np.array([1] + [1]*3 + [7]*3 + [7]*9) + y = np.array([7] + [1]*3 + [7]*3 + [1]*9) + + # sanity check that the above worked as intended + hist, edges = histogramdd((y, x), bins=(y_edges, x_edges)) + assert_equal(hist, relative_areas) + + # resulting histogram should be uniform, since counts and areas are proportional + hist, edges = histogramdd((y, x), bins=(y_edges, x_edges), density=True) + assert_equal(hist, 1 / (8*8)) + + def test_density_non_uniform_1d(self): + # compare to histogram to show the results are the same + v = np.arange(10) + bins = np.array([0, 1, 3, 6, 10]) + hist, edges = histogram(v, bins, density=True) + hist_dd, edges_dd = histogramdd((v,), (bins,), density=True) + assert_equal(hist, hist_dd) + assert_equal(edges, edges_dd[0]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_index_tricks.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_index_tricks.py new file mode 100644 index 0000000000000000000000000000000000000000..fe1cfce2eaf83272d0a01d3fb1911772693c69d7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_index_tricks.py @@ -0,0 +1,553 @@ +import pytest + +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_almost_equal, + assert_array_almost_equal, assert_raises, assert_raises_regex, + ) +from numpy.lib._index_tricks_impl import ( + mgrid, ogrid, ndenumerate, fill_diagonal, diag_indices, diag_indices_from, + index_exp, ndindex, c_, r_, s_, ix_ + ) + + +class TestRavelUnravelIndex: + def test_basic(self): + assert_equal(np.unravel_index(2, (2, 2)), (1, 0)) + + # test that new shape argument works properly + assert_equal(np.unravel_index(indices=2, + shape=(2, 2)), + (1, 0)) + + # test that an invalid second keyword argument + # is properly handled, including the old name `dims`. + with assert_raises(TypeError): + np.unravel_index(indices=2, hape=(2, 2)) + + with assert_raises(TypeError): + np.unravel_index(2, hape=(2, 2)) + + with assert_raises(TypeError): + np.unravel_index(254, ims=(17, 94)) + + with assert_raises(TypeError): + np.unravel_index(254, dims=(17, 94)) + + assert_equal(np.ravel_multi_index((1, 0), (2, 2)), 2) + assert_equal(np.unravel_index(254, (17, 94)), (2, 66)) + assert_equal(np.ravel_multi_index((2, 66), (17, 94)), 254) + assert_raises(ValueError, np.unravel_index, -1, (2, 2)) + assert_raises(TypeError, np.unravel_index, 0.5, (2, 2)) + assert_raises(ValueError, np.unravel_index, 4, (2, 2)) + assert_raises(ValueError, np.ravel_multi_index, (-3, 1), (2, 2)) + assert_raises(ValueError, np.ravel_multi_index, (2, 1), (2, 2)) + assert_raises(ValueError, np.ravel_multi_index, (0, -3), (2, 2)) + assert_raises(ValueError, np.ravel_multi_index, (0, 2), (2, 2)) + assert_raises(TypeError, np.ravel_multi_index, (0.1, 0.), (2, 2)) + + assert_equal(np.unravel_index((2*3 + 1)*6 + 4, (4, 3, 6)), [2, 1, 4]) + assert_equal( + np.ravel_multi_index([2, 1, 4], (4, 3, 6)), (2*3 + 1)*6 + 4) + + arr = np.array([[3, 6, 6], [4, 5, 1]]) + assert_equal(np.ravel_multi_index(arr, (7, 6)), [22, 41, 37]) + assert_equal( + np.ravel_multi_index(arr, (7, 6), order='F'), [31, 41, 13]) + assert_equal( + np.ravel_multi_index(arr, (4, 6), mode='clip'), [22, 23, 19]) + assert_equal(np.ravel_multi_index(arr, (4, 4), mode=('clip', 'wrap')), + [12, 13, 13]) + assert_equal(np.ravel_multi_index((3, 1, 4, 1), (6, 7, 8, 9)), 1621) + + assert_equal(np.unravel_index(np.array([22, 41, 37]), (7, 6)), + [[3, 6, 6], [4, 5, 1]]) + assert_equal( + np.unravel_index(np.array([31, 41, 13]), (7, 6), order='F'), + [[3, 6, 6], [4, 5, 1]]) + assert_equal(np.unravel_index(1621, (6, 7, 8, 9)), [3, 1, 4, 1]) + + def test_empty_indices(self): + msg1 = 'indices must be integral: the provided empty sequence was' + msg2 = 'only int indices permitted' + assert_raises_regex(TypeError, msg1, np.unravel_index, [], (10, 3, 5)) + assert_raises_regex(TypeError, msg1, np.unravel_index, (), (10, 3, 5)) + assert_raises_regex(TypeError, msg2, np.unravel_index, np.array([]), + (10, 3, 5)) + assert_equal(np.unravel_index(np.array([],dtype=int), (10, 3, 5)), + [[], [], []]) + assert_raises_regex(TypeError, msg1, np.ravel_multi_index, ([], []), + (10, 3)) + assert_raises_regex(TypeError, msg1, np.ravel_multi_index, ([], ['abc']), + (10, 3)) + assert_raises_regex(TypeError, msg2, np.ravel_multi_index, + (np.array([]), np.array([])), (5, 3)) + assert_equal(np.ravel_multi_index( + (np.array([], dtype=int), np.array([], dtype=int)), (5, 3)), []) + assert_equal(np.ravel_multi_index(np.array([[], []], dtype=int), + (5, 3)), []) + + def test_big_indices(self): + # ravel_multi_index for big indices (issue #7546) + if np.intp == np.int64: + arr = ([1, 29], [3, 5], [3, 117], [19, 2], + [2379, 1284], [2, 2], [0, 1]) + assert_equal( + np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)), + [5627771580, 117259570957]) + + # test unravel_index for big indices (issue #9538) + assert_raises(ValueError, np.unravel_index, 1, (2**32-1, 2**31+1)) + + # test overflow checking for too big array (issue #7546) + dummy_arr = ([0],[0]) + half_max = np.iinfo(np.intp).max // 2 + assert_equal( + np.ravel_multi_index(dummy_arr, (half_max, 2)), [0]) + assert_raises(ValueError, + np.ravel_multi_index, dummy_arr, (half_max+1, 2)) + assert_equal( + np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0]) + assert_raises(ValueError, + np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') + + def test_dtypes(self): + # Test with different data types + for dtype in [np.int16, np.uint16, np.int32, + np.uint32, np.int64, np.uint64]: + coords = np.array( + [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0]], dtype=dtype) + shape = (5, 8) + uncoords = 8*coords[0]+coords[1] + assert_equal(np.ravel_multi_index(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*coords[1] + assert_equal( + np.ravel_multi_index(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + coords = np.array( + [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0], [1, 3, 1, 0, 9, 5]], + dtype=dtype) + shape = (5, 8, 10) + uncoords = 10*(8*coords[0]+coords[1])+coords[2] + assert_equal(np.ravel_multi_index(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*(coords[1]+8*coords[2]) + assert_equal( + np.ravel_multi_index(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + def test_clipmodes(self): + # Test clipmodes + assert_equal( + np.ravel_multi_index([5, 1, -1, 2], (4, 3, 7, 12), mode='wrap'), + np.ravel_multi_index([1, 1, 6, 2], (4, 3, 7, 12))) + assert_equal(np.ravel_multi_index([5, 1, -1, 2], (4, 3, 7, 12), + mode=( + 'wrap', 'raise', 'clip', 'raise')), + np.ravel_multi_index([1, 1, 0, 2], (4, 3, 7, 12))) + assert_raises( + ValueError, np.ravel_multi_index, [5, 1, -1, 2], (4, 3, 7, 12)) + + def test_writeability(self): + # See gh-7269 + x, y = np.unravel_index([1, 2, 3], (4, 5)) + assert_(x.flags.writeable) + assert_(y.flags.writeable) + + def test_0d(self): + # gh-580 + x = np.unravel_index(0, ()) + assert_equal(x, ()) + + assert_raises_regex(ValueError, "0d array", np.unravel_index, [0], ()) + assert_raises_regex( + ValueError, "out of bounds", np.unravel_index, [1], ()) + + @pytest.mark.parametrize("mode", ["clip", "wrap", "raise"]) + def test_empty_array_ravel(self, mode): + res = np.ravel_multi_index( + np.zeros((3, 0), dtype=np.intp), (2, 1, 0), mode=mode) + assert(res.shape == (0,)) + + with assert_raises(ValueError): + np.ravel_multi_index( + np.zeros((3, 1), dtype=np.intp), (2, 1, 0), mode=mode) + + def test_empty_array_unravel(self): + res = np.unravel_index(np.zeros(0, dtype=np.intp), (2, 1, 0)) + # res is a tuple of three empty arrays + assert(len(res) == 3) + assert(all(a.shape == (0,) for a in res)) + + with assert_raises(ValueError): + np.unravel_index([1], (2, 1, 0)) + +class TestGrid: + def test_basic(self): + a = mgrid[-1:1:10j] + b = mgrid[-1:1:0.1] + assert_(a.shape == (10,)) + assert_(b.shape == (20,)) + assert_(a[0] == -1) + assert_almost_equal(a[-1], 1) + assert_(b[0] == -1) + assert_almost_equal(b[1]-b[0], 0.1, 11) + assert_almost_equal(b[-1], b[0]+19*0.1, 11) + assert_almost_equal(a[1]-a[0], 2.0/9.0, 11) + + def test_linspace_equivalence(self): + y, st = np.linspace(2, 10, retstep=True) + assert_almost_equal(st, 8/49.0) + assert_array_almost_equal(y, mgrid[2:10:50j], 13) + + def test_nd(self): + c = mgrid[-1:1:10j, -2:2:10j] + d = mgrid[-1:1:0.1, -2:2:0.2] + assert_(c.shape == (2, 10, 10)) + assert_(d.shape == (2, 20, 20)) + assert_array_equal(c[0][0, :], -np.ones(10, 'd')) + assert_array_equal(c[1][:, 0], -2*np.ones(10, 'd')) + assert_array_almost_equal(c[0][-1, :], np.ones(10, 'd'), 11) + assert_array_almost_equal(c[1][:, -1], 2*np.ones(10, 'd'), 11) + assert_array_almost_equal(d[0, 1, :] - d[0, 0, :], + 0.1*np.ones(20, 'd'), 11) + assert_array_almost_equal(d[1, :, 1] - d[1, :, 0], + 0.2*np.ones(20, 'd'), 11) + + def test_sparse(self): + grid_full = mgrid[-1:1:10j, -2:2:10j] + grid_sparse = ogrid[-1:1:10j, -2:2:10j] + + # sparse grids can be made dense by broadcasting + grid_broadcast = np.broadcast_arrays(*grid_sparse) + for f, b in zip(grid_full, grid_broadcast): + assert_equal(f, b) + + @pytest.mark.parametrize("start, stop, step, expected", [ + (None, 10, 10j, (200, 10)), + (-10, 20, None, (1800, 30)), + ]) + def test_mgrid_size_none_handling(self, start, stop, step, expected): + # regression test None value handling for + # start and step values used by mgrid; + # internally, this aims to cover previously + # unexplored code paths in nd_grid() + grid = mgrid[start:stop:step, start:stop:step] + # need a smaller grid to explore one of the + # untested code paths + grid_small = mgrid[start:stop:step] + assert_equal(grid.size, expected[0]) + assert_equal(grid_small.size, expected[1]) + + def test_accepts_npfloating(self): + # regression test for #16466 + grid64 = mgrid[0.1:0.33:0.1, ] + grid32 = mgrid[np.float32(0.1):np.float32(0.33):np.float32(0.1), ] + assert_array_almost_equal(grid64, grid32) + # At some point this was float64, but NEP 50 changed it: + assert grid32.dtype == np.float32 + assert grid64.dtype == np.float64 + + # different code path for single slice + grid64 = mgrid[0.1:0.33:0.1] + grid32 = mgrid[np.float32(0.1):np.float32(0.33):np.float32(0.1)] + assert_(grid32.dtype == np.float64) + assert_array_almost_equal(grid64, grid32) + + def test_accepts_longdouble(self): + # regression tests for #16945 + grid64 = mgrid[0.1:0.33:0.1, ] + grid128 = mgrid[ + np.longdouble(0.1):np.longdouble(0.33):np.longdouble(0.1), + ] + assert_(grid128.dtype == np.longdouble) + assert_array_almost_equal(grid64, grid128) + + grid128c_a = mgrid[0:np.longdouble(1):3.4j] + grid128c_b = mgrid[0:np.longdouble(1):3.4j, ] + assert_(grid128c_a.dtype == grid128c_b.dtype == np.longdouble) + assert_array_equal(grid128c_a, grid128c_b[0]) + + # different code path for single slice + grid64 = mgrid[0.1:0.33:0.1] + grid128 = mgrid[ + np.longdouble(0.1):np.longdouble(0.33):np.longdouble(0.1) + ] + assert_(grid128.dtype == np.longdouble) + assert_array_almost_equal(grid64, grid128) + + def test_accepts_npcomplexfloating(self): + # Related to #16466 + assert_array_almost_equal( + mgrid[0.1:0.3:3j, ], mgrid[0.1:0.3:np.complex64(3j), ] + ) + + # different code path for single slice + assert_array_almost_equal( + mgrid[0.1:0.3:3j], mgrid[0.1:0.3:np.complex64(3j)] + ) + + # Related to #16945 + grid64_a = mgrid[0.1:0.3:3.3j] + grid64_b = mgrid[0.1:0.3:3.3j, ][0] + assert_(grid64_a.dtype == grid64_b.dtype == np.float64) + assert_array_equal(grid64_a, grid64_b) + + grid128_a = mgrid[0.1:0.3:np.clongdouble(3.3j)] + grid128_b = mgrid[0.1:0.3:np.clongdouble(3.3j), ][0] + assert_(grid128_a.dtype == grid128_b.dtype == np.longdouble) + assert_array_equal(grid64_a, grid64_b) + + +class TestConcatenator: + def test_1d(self): + assert_array_equal(r_[1, 2, 3, 4, 5, 6], np.array([1, 2, 3, 4, 5, 6])) + b = np.ones(5) + c = r_[b, 0, 0, b] + assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1]) + + def test_mixed_type(self): + g = r_[10.1, 1:10] + assert_(g.dtype == 'f8') + + def test_more_mixed_type(self): + g = r_[-10.1, np.array([1]), np.array([2, 3, 4]), 10.0] + assert_(g.dtype == 'f8') + + def test_complex_step(self): + # Regression test for #12262 + g = r_[0:36:100j] + assert_(g.shape == (100,)) + + # Related to #16466 + g = r_[0:36:np.complex64(100j)] + assert_(g.shape == (100,)) + + def test_2d(self): + b = np.random.rand(5, 5) + c = np.random.rand(5, 5) + d = r_['1', b, c] # append columns + assert_(d.shape == (5, 10)) + assert_array_equal(d[:, :5], b) + assert_array_equal(d[:, 5:], c) + d = r_[b, c] + assert_(d.shape == (10, 5)) + assert_array_equal(d[:5, :], b) + assert_array_equal(d[5:, :], c) + + def test_0d(self): + assert_equal(r_[0, np.array(1), 2], [0, 1, 2]) + assert_equal(r_[[0, 1, 2], np.array(3)], [0, 1, 2, 3]) + assert_equal(r_[np.array(0), [1, 2, 3]], [0, 1, 2, 3]) + + +class TestNdenumerate: + def test_basic(self): + a = np.array([[1, 2], [3, 4]]) + assert_equal(list(ndenumerate(a)), + [((0, 0), 1), ((0, 1), 2), ((1, 0), 3), ((1, 1), 4)]) + + +class TestIndexExpression: + def test_regression_1(self): + # ticket #1196 + a = np.arange(2) + assert_equal(a[:-1], a[s_[:-1]]) + assert_equal(a[:-1], a[index_exp[:-1]]) + + def test_simple_1(self): + a = np.random.rand(4, 5, 6) + + assert_equal(a[:, :3, [1, 2]], a[index_exp[:, :3, [1, 2]]]) + assert_equal(a[:, :3, [1, 2]], a[s_[:, :3, [1, 2]]]) + + +class TestIx_: + def test_regression_1(self): + # Test empty untyped inputs create outputs of indexing type, gh-5804 + a, = np.ix_(range(0)) + assert_equal(a.dtype, np.intp) + + a, = np.ix_([]) + assert_equal(a.dtype, np.intp) + + # but if the type is specified, don't change it + a, = np.ix_(np.array([], dtype=np.float32)) + assert_equal(a.dtype, np.float32) + + def test_shape_and_dtype(self): + sizes = (4, 5, 3, 2) + # Test both lists and arrays + for func in (range, np.arange): + arrays = np.ix_(*[func(sz) for sz in sizes]) + for k, (a, sz) in enumerate(zip(arrays, sizes)): + assert_equal(a.shape[k], sz) + assert_(all(sh == 1 for j, sh in enumerate(a.shape) if j != k)) + assert_(np.issubdtype(a.dtype, np.integer)) + + def test_bool(self): + bool_a = [True, False, True, True] + int_a, = np.nonzero(bool_a) + assert_equal(np.ix_(bool_a)[0], int_a) + + def test_1d_only(self): + idx2d = [[1, 2, 3], [4, 5, 6]] + assert_raises(ValueError, np.ix_, idx2d) + + def test_repeated_input(self): + length_of_vector = 5 + x = np.arange(length_of_vector) + out = ix_(x, x) + assert_equal(out[0].shape, (length_of_vector, 1)) + assert_equal(out[1].shape, (1, length_of_vector)) + # check that input shape is not modified + assert_equal(x.shape, (length_of_vector,)) + + +def test_c_(): + a = c_[np.array([[1, 2, 3]]), 0, 0, np.array([[4, 5, 6]])] + assert_equal(a, [[1, 2, 3, 0, 0, 4, 5, 6]]) + + +class TestFillDiagonal: + def test_basic(self): + a = np.zeros((3, 3), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5]]) + ) + + def test_tall_matrix(self): + a = np.zeros((10, 3), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]) + ) + + def test_tall_matrix_wrap(self): + a = np.zeros((10, 3), int) + fill_diagonal(a, 5, True) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [5, 0, 0], + [0, 5, 0]]) + ) + + def test_wide_matrix(self): + a = np.zeros((3, 10), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 5, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 5, 0, 0, 0, 0, 0, 0, 0]]) + ) + + def test_operate_4d_array(self): + a = np.zeros((3, 3, 3, 3), int) + fill_diagonal(a, 4) + i = np.array([0, 1, 2]) + assert_equal(np.where(a != 0), (i, i, i, i)) + + def test_low_dim_handling(self): + # raise error with low dimensionality + a = np.zeros(3, int) + with assert_raises_regex(ValueError, "at least 2-d"): + fill_diagonal(a, 5) + + def test_hetero_shape_handling(self): + # raise error with high dimensionality and + # shape mismatch + a = np.zeros((3,3,7,3), int) + with assert_raises_regex(ValueError, "equal length"): + fill_diagonal(a, 2) + + +def test_diag_indices(): + di = diag_indices(4) + a = np.array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + a[di] = 100 + assert_array_equal( + a, np.array([[100, 2, 3, 4], + [5, 100, 7, 8], + [9, 10, 100, 12], + [13, 14, 15, 100]]) + ) + + # Now, we create indices to manipulate a 3-d array: + d3 = diag_indices(2, 3) + + # And use it to set the diagonal of a zeros array to 1: + a = np.zeros((2, 2, 2), int) + a[d3] = 1 + assert_array_equal( + a, np.array([[[1, 0], + [0, 0]], + [[0, 0], + [0, 1]]]) + ) + + +class TestDiagIndicesFrom: + + def test_diag_indices_from(self): + x = np.random.random((4, 4)) + r, c = diag_indices_from(x) + assert_array_equal(r, np.arange(4)) + assert_array_equal(c, np.arange(4)) + + def test_error_small_input(self): + x = np.ones(7) + with assert_raises_regex(ValueError, "at least 2-d"): + diag_indices_from(x) + + def test_error_shape_mismatch(self): + x = np.zeros((3, 3, 2, 3), int) + with assert_raises_regex(ValueError, "equal length"): + diag_indices_from(x) + + +def test_ndindex(): + x = list(ndindex(1, 2, 3)) + expected = [ix for ix, e in ndenumerate(np.zeros((1, 2, 3)))] + assert_array_equal(x, expected) + + x = list(ndindex((1, 2, 3))) + assert_array_equal(x, expected) + + # Test use of scalars and tuples + x = list(ndindex((3,))) + assert_array_equal(x, list(ndindex(3))) + + # Make sure size argument is optional + x = list(ndindex()) + assert_equal(x, [()]) + + x = list(ndindex(())) + assert_equal(x, [()]) + + # Make sure 0-sized ndindex works correctly + x = list(ndindex(*[0])) + assert_equal(x, []) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_io.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_io.py new file mode 100644 index 0000000000000000000000000000000000000000..742915e22ef09268d986aee3482a9da9a3602be1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_io.py @@ -0,0 +1,2821 @@ +import sys +import gc +import gzip +import os +import threading +import time +import warnings +import re +import pytest +from pathlib import Path +from tempfile import NamedTemporaryFile +from io import BytesIO, StringIO +from datetime import datetime +import locale +from multiprocessing import Value, get_context +from ctypes import c_bool + +import numpy as np +import numpy.ma as ma +from numpy.exceptions import VisibleDeprecationWarning +from numpy.lib._iotools import ConverterError, ConversionWarning +from numpy.lib import _npyio_impl +from numpy.lib._npyio_impl import recfromcsv, recfromtxt +from numpy.ma.testutils import assert_equal +from numpy.testing import ( + assert_warns, assert_, assert_raises_regex, assert_raises, + assert_allclose, assert_array_equal, temppath, tempdir, IS_PYPY, + HAS_REFCOUNT, suppress_warnings, assert_no_gc_cycles, assert_no_warnings, + break_cycles, IS_WASM + ) +from numpy.testing._private.utils import requires_memory +from numpy._utils import asbytes + + +class TextIO(BytesIO): + """Helper IO class. + + Writes encode strings to bytes if needed, reads return bytes. + This makes it easier to emulate files opened in binary mode + without needing to explicitly convert strings to bytes in + setting up the test data. + + """ + def __init__(self, s=""): + BytesIO.__init__(self, asbytes(s)) + + def write(self, s): + BytesIO.write(self, asbytes(s)) + + def writelines(self, lines): + BytesIO.writelines(self, [asbytes(s) for s in lines]) + + +IS_64BIT = sys.maxsize > 2**32 +try: + import bz2 + HAS_BZ2 = True +except ImportError: + HAS_BZ2 = False +try: + import lzma + HAS_LZMA = True +except ImportError: + HAS_LZMA = False + + +def strptime(s, fmt=None): + """ + This function is available in the datetime module only from Python >= + 2.5. + + """ + if type(s) == bytes: + s = s.decode("latin1") + return datetime(*time.strptime(s, fmt)[:3]) + + +class RoundtripTest: + def roundtrip(self, save_func, *args, **kwargs): + """ + save_func : callable + Function used to save arrays to file. + file_on_disk : bool + If true, store the file on disk, instead of in a + string buffer. + save_kwds : dict + Parameters passed to `save_func`. + load_kwds : dict + Parameters passed to `numpy.load`. + args : tuple of arrays + Arrays stored to file. + + """ + save_kwds = kwargs.get('save_kwds', {}) + load_kwds = kwargs.get('load_kwds', {"allow_pickle": True}) + file_on_disk = kwargs.get('file_on_disk', False) + + if file_on_disk: + target_file = NamedTemporaryFile(delete=False) + load_file = target_file.name + else: + target_file = BytesIO() + load_file = target_file + + try: + arr = args + + save_func(target_file, *arr, **save_kwds) + target_file.flush() + target_file.seek(0) + + if sys.platform == 'win32' and not isinstance(target_file, BytesIO): + target_file.close() + + arr_reloaded = np.load(load_file, **load_kwds) + + self.arr = arr + self.arr_reloaded = arr_reloaded + finally: + if not isinstance(target_file, BytesIO): + target_file.close() + # holds an open file descriptor so it can't be deleted on win + if 'arr_reloaded' in locals(): + if not isinstance(arr_reloaded, np.lib.npyio.NpzFile): + os.remove(target_file.name) + + def check_roundtrips(self, a): + self.roundtrip(a) + self.roundtrip(a, file_on_disk=True) + self.roundtrip(np.asfortranarray(a)) + self.roundtrip(np.asfortranarray(a), file_on_disk=True) + if a.shape[0] > 1: + # neither C nor Fortran contiguous for 2D arrays or more + self.roundtrip(np.asfortranarray(a)[1:]) + self.roundtrip(np.asfortranarray(a)[1:], file_on_disk=True) + + def test_array(self): + a = np.array([], float) + self.check_roundtrips(a) + + a = np.array([[1, 2], [3, 4]], float) + self.check_roundtrips(a) + + a = np.array([[1, 2], [3, 4]], int) + self.check_roundtrips(a) + + a = np.array([[1 + 5j, 2 + 6j], [3 + 7j, 4 + 8j]], dtype=np.csingle) + self.check_roundtrips(a) + + a = np.array([[1 + 5j, 2 + 6j], [3 + 7j, 4 + 8j]], dtype=np.cdouble) + self.check_roundtrips(a) + + def test_array_object(self): + a = np.array([], object) + self.check_roundtrips(a) + + a = np.array([[1, 2], [3, 4]], object) + self.check_roundtrips(a) + + def test_1D(self): + a = np.array([1, 2, 3, 4], int) + self.roundtrip(a) + + @pytest.mark.skipif(sys.platform == 'win32', reason="Fails on Win32") + def test_mmap(self): + a = np.array([[1, 2.5], [4, 7.3]]) + self.roundtrip(a, file_on_disk=True, load_kwds={'mmap_mode': 'r'}) + + a = np.asfortranarray([[1, 2.5], [4, 7.3]]) + self.roundtrip(a, file_on_disk=True, load_kwds={'mmap_mode': 'r'}) + + def test_record(self): + a = np.array([(1, 2), (3, 4)], dtype=[('x', 'i4'), ('y', 'i4')]) + self.check_roundtrips(a) + + @pytest.mark.slow + def test_format_2_0(self): + dt = [(("%d" % i) * 100, float) for i in range(500)] + a = np.ones(1000, dtype=dt) + with warnings.catch_warnings(record=True): + warnings.filterwarnings('always', '', UserWarning) + self.check_roundtrips(a) + + +class TestSaveLoad(RoundtripTest): + def roundtrip(self, *args, **kwargs): + RoundtripTest.roundtrip(self, np.save, *args, **kwargs) + assert_equal(self.arr[0], self.arr_reloaded) + assert_equal(self.arr[0].dtype, self.arr_reloaded.dtype) + assert_equal(self.arr[0].flags.fnc, self.arr_reloaded.flags.fnc) + + +class TestSavezLoad(RoundtripTest): + def roundtrip(self, *args, **kwargs): + RoundtripTest.roundtrip(self, np.savez, *args, **kwargs) + try: + for n, arr in enumerate(self.arr): + reloaded = self.arr_reloaded['arr_%d' % n] + assert_equal(arr, reloaded) + assert_equal(arr.dtype, reloaded.dtype) + assert_equal(arr.flags.fnc, reloaded.flags.fnc) + finally: + # delete tempfile, must be done here on windows + if self.arr_reloaded.fid: + self.arr_reloaded.fid.close() + os.remove(self.arr_reloaded.fid.name) + + @pytest.mark.skipif(IS_PYPY, reason="Hangs on PyPy") + @pytest.mark.skipif(not IS_64BIT, reason="Needs 64bit platform") + @pytest.mark.slow + def test_big_arrays(self): + L = (1 << 31) + 100000 + a = np.empty(L, dtype=np.uint8) + with temppath(prefix="numpy_test_big_arrays_", suffix=".npz") as tmp: + np.savez(tmp, a=a) + del a + npfile = np.load(tmp) + a = npfile['a'] # Should succeed + npfile.close() + del a # Avoid pyflakes unused variable warning. + + def test_multiple_arrays(self): + a = np.array([[1, 2], [3, 4]], float) + b = np.array([[1 + 2j, 2 + 7j], [3 - 6j, 4 + 12j]], complex) + self.roundtrip(a, b) + + def test_named_arrays(self): + a = np.array([[1, 2], [3, 4]], float) + b = np.array([[1 + 2j, 2 + 7j], [3 - 6j, 4 + 12j]], complex) + c = BytesIO() + np.savez(c, file_a=a, file_b=b) + c.seek(0) + l = np.load(c) + assert_equal(a, l['file_a']) + assert_equal(b, l['file_b']) + + + def test_tuple_getitem_raises(self): + # gh-23748 + a = np.array([1, 2, 3]) + f = BytesIO() + np.savez(f, a=a) + f.seek(0) + l = np.load(f) + with pytest.raises(KeyError, match="(1, 2)"): + l[1, 2] + + def test_BagObj(self): + a = np.array([[1, 2], [3, 4]], float) + b = np.array([[1 + 2j, 2 + 7j], [3 - 6j, 4 + 12j]], complex) + c = BytesIO() + np.savez(c, file_a=a, file_b=b) + c.seek(0) + l = np.load(c) + assert_equal(sorted(dir(l.f)), ['file_a','file_b']) + assert_equal(a, l.f.file_a) + assert_equal(b, l.f.file_b) + + @pytest.mark.skipif(IS_WASM, reason="Cannot start thread") + def test_savez_filename_clashes(self): + # Test that issue #852 is fixed + # and savez functions in multithreaded environment + + def writer(error_list): + with temppath(suffix='.npz') as tmp: + arr = np.random.randn(500, 500) + try: + np.savez(tmp, arr=arr) + except OSError as err: + error_list.append(err) + + errors = [] + threads = [threading.Thread(target=writer, args=(errors,)) + for j in range(3)] + for t in threads: + t.start() + for t in threads: + t.join() + + if errors: + raise AssertionError(errors) + + def test_not_closing_opened_fid(self): + # Test that issue #2178 is fixed: + # verify could seek on 'loaded' file + with temppath(suffix='.npz') as tmp: + with open(tmp, 'wb') as fp: + np.savez(fp, data='LOVELY LOAD') + with open(tmp, 'rb', 10000) as fp: + fp.seek(0) + assert_(not fp.closed) + np.load(fp)['data'] + # fp must not get closed by .load + assert_(not fp.closed) + fp.seek(0) + assert_(not fp.closed) + + @pytest.mark.slow_pypy + def test_closing_fid(self): + # Test that issue #1517 (too many opened files) remains closed + # It might be a "weak" test since failed to get triggered on + # e.g. Debian sid of 2012 Jul 05 but was reported to + # trigger the failure on Ubuntu 10.04: + # http://projects.scipy.org/numpy/ticket/1517#comment:2 + with temppath(suffix='.npz') as tmp: + np.savez(tmp, data='LOVELY LOAD') + # We need to check if the garbage collector can properly close + # numpy npz file returned by np.load when their reference count + # goes to zero. Python 3 running in debug mode raises a + # ResourceWarning when file closing is left to the garbage + # collector, so we catch the warnings. + with suppress_warnings() as sup: + sup.filter(ResourceWarning) # TODO: specify exact message + for i in range(1, 1025): + try: + np.load(tmp)["data"] + except Exception as e: + msg = "Failed to load data from a file: %s" % e + raise AssertionError(msg) + finally: + if IS_PYPY: + gc.collect() + + def test_closing_zipfile_after_load(self): + # Check that zipfile owns file and can close it. This needs to + # pass a file name to load for the test. On windows failure will + # cause a second error will be raised when the attempt to remove + # the open file is made. + prefix = 'numpy_test_closing_zipfile_after_load_' + with temppath(suffix='.npz', prefix=prefix) as tmp: + np.savez(tmp, lab='place holder') + data = np.load(tmp) + fp = data.zip.fp + data.close() + assert_(fp.closed) + + @pytest.mark.parametrize("count, expected_repr", [ + (1, "NpzFile {fname!r} with keys: arr_0"), + (5, "NpzFile {fname!r} with keys: arr_0, arr_1, arr_2, arr_3, arr_4"), + # _MAX_REPR_ARRAY_COUNT is 5, so files with more than 5 keys are + # expected to end in '...' + (6, "NpzFile {fname!r} with keys: arr_0, arr_1, arr_2, arr_3, arr_4..."), + ]) + def test_repr_lists_keys(self, count, expected_repr): + a = np.array([[1, 2], [3, 4]], float) + with temppath(suffix='.npz') as tmp: + np.savez(tmp, *[a]*count) + l = np.load(tmp) + assert repr(l) == expected_repr.format(fname=tmp) + l.close() + + +class TestSaveTxt: + def test_array(self): + a = np.array([[1, 2], [3, 4]], float) + fmt = "%.18e" + c = BytesIO() + np.savetxt(c, a, fmt=fmt) + c.seek(0) + assert_equal(c.readlines(), + [asbytes((fmt + ' ' + fmt + '\n') % (1, 2)), + asbytes((fmt + ' ' + fmt + '\n') % (3, 4))]) + + a = np.array([[1, 2], [3, 4]], int) + c = BytesIO() + np.savetxt(c, a, fmt='%d') + c.seek(0) + assert_equal(c.readlines(), [b'1 2\n', b'3 4\n']) + + def test_1D(self): + a = np.array([1, 2, 3, 4], int) + c = BytesIO() + np.savetxt(c, a, fmt='%d') + c.seek(0) + lines = c.readlines() + assert_equal(lines, [b'1\n', b'2\n', b'3\n', b'4\n']) + + def test_0D_3D(self): + c = BytesIO() + assert_raises(ValueError, np.savetxt, c, np.array(1)) + assert_raises(ValueError, np.savetxt, c, np.array([[[1], [2]]])) + + def test_structured(self): + a = np.array([(1, 2), (3, 4)], dtype=[('x', 'i4'), ('y', 'i4')]) + c = BytesIO() + np.savetxt(c, a, fmt='%d') + c.seek(0) + assert_equal(c.readlines(), [b'1 2\n', b'3 4\n']) + + def test_structured_padded(self): + # gh-13297 + a = np.array([(1, 2, 3),(4, 5, 6)], dtype=[ + ('foo', 'i4'), ('bar', 'i4'), ('baz', 'i4') + ]) + c = BytesIO() + np.savetxt(c, a[['foo', 'baz']], fmt='%d') + c.seek(0) + assert_equal(c.readlines(), [b'1 3\n', b'4 6\n']) + + def test_multifield_view(self): + a = np.ones(1, dtype=[('x', 'i4'), ('y', 'i4'), ('z', 'f4')]) + v = a[['x', 'z']] + with temppath(suffix='.npy') as path: + path = Path(path) + np.save(path, v) + data = np.load(path) + assert_array_equal(data, v) + + def test_delimiter(self): + a = np.array([[1., 2.], [3., 4.]]) + c = BytesIO() + np.savetxt(c, a, delimiter=',', fmt='%d') + c.seek(0) + assert_equal(c.readlines(), [b'1,2\n', b'3,4\n']) + + def test_format(self): + a = np.array([(1, 2), (3, 4)]) + c = BytesIO() + # Sequence of formats + np.savetxt(c, a, fmt=['%02d', '%3.1f']) + c.seek(0) + assert_equal(c.readlines(), [b'01 2.0\n', b'03 4.0\n']) + + # A single multiformat string + c = BytesIO() + np.savetxt(c, a, fmt='%02d : %3.1f') + c.seek(0) + lines = c.readlines() + assert_equal(lines, [b'01 : 2.0\n', b'03 : 4.0\n']) + + # Specify delimiter, should be overridden + c = BytesIO() + np.savetxt(c, a, fmt='%02d : %3.1f', delimiter=',') + c.seek(0) + lines = c.readlines() + assert_equal(lines, [b'01 : 2.0\n', b'03 : 4.0\n']) + + # Bad fmt, should raise a ValueError + c = BytesIO() + assert_raises(ValueError, np.savetxt, c, a, fmt=99) + + def test_header_footer(self): + # Test the functionality of the header and footer keyword argument. + + c = BytesIO() + a = np.array([(1, 2), (3, 4)], dtype=int) + test_header_footer = 'Test header / footer' + # Test the header keyword argument + np.savetxt(c, a, fmt='%1d', header=test_header_footer) + c.seek(0) + assert_equal(c.read(), + asbytes('# ' + test_header_footer + '\n1 2\n3 4\n')) + # Test the footer keyword argument + c = BytesIO() + np.savetxt(c, a, fmt='%1d', footer=test_header_footer) + c.seek(0) + assert_equal(c.read(), + asbytes('1 2\n3 4\n# ' + test_header_footer + '\n')) + # Test the commentstr keyword argument used on the header + c = BytesIO() + commentstr = '% ' + np.savetxt(c, a, fmt='%1d', + header=test_header_footer, comments=commentstr) + c.seek(0) + assert_equal(c.read(), + asbytes(commentstr + test_header_footer + '\n' + '1 2\n3 4\n')) + # Test the commentstr keyword argument used on the footer + c = BytesIO() + commentstr = '% ' + np.savetxt(c, a, fmt='%1d', + footer=test_header_footer, comments=commentstr) + c.seek(0) + assert_equal(c.read(), + asbytes('1 2\n3 4\n' + commentstr + test_header_footer + '\n')) + + @pytest.mark.parametrize("filename_type", [Path, str]) + def test_file_roundtrip(self, filename_type): + with temppath() as name: + a = np.array([(1, 2), (3, 4)]) + np.savetxt(filename_type(name), a) + b = np.loadtxt(filename_type(name)) + assert_array_equal(a, b) + + def test_complex_arrays(self): + ncols = 2 + nrows = 2 + a = np.zeros((ncols, nrows), dtype=np.complex128) + re = np.pi + im = np.e + a[:] = re + 1.0j * im + + # One format only + c = BytesIO() + np.savetxt(c, a, fmt=' %+.3e') + c.seek(0) + lines = c.readlines() + assert_equal( + lines, + [b' ( +3.142e+00+ +2.718e+00j) ( +3.142e+00+ +2.718e+00j)\n', + b' ( +3.142e+00+ +2.718e+00j) ( +3.142e+00+ +2.718e+00j)\n']) + + # One format for each real and imaginary part + c = BytesIO() + np.savetxt(c, a, fmt=' %+.3e' * 2 * ncols) + c.seek(0) + lines = c.readlines() + assert_equal( + lines, + [b' +3.142e+00 +2.718e+00 +3.142e+00 +2.718e+00\n', + b' +3.142e+00 +2.718e+00 +3.142e+00 +2.718e+00\n']) + + # One format for each complex number + c = BytesIO() + np.savetxt(c, a, fmt=['(%.3e%+.3ej)'] * ncols) + c.seek(0) + lines = c.readlines() + assert_equal( + lines, + [b'(3.142e+00+2.718e+00j) (3.142e+00+2.718e+00j)\n', + b'(3.142e+00+2.718e+00j) (3.142e+00+2.718e+00j)\n']) + + def test_complex_negative_exponent(self): + # Previous to 1.15, some formats generated x+-yj, gh 7895 + ncols = 2 + nrows = 2 + a = np.zeros((ncols, nrows), dtype=np.complex128) + re = np.pi + im = np.e + a[:] = re - 1.0j * im + c = BytesIO() + np.savetxt(c, a, fmt='%.3e') + c.seek(0) + lines = c.readlines() + assert_equal( + lines, + [b' (3.142e+00-2.718e+00j) (3.142e+00-2.718e+00j)\n', + b' (3.142e+00-2.718e+00j) (3.142e+00-2.718e+00j)\n']) + + + def test_custom_writer(self): + + class CustomWriter(list): + def write(self, text): + self.extend(text.split(b'\n')) + + w = CustomWriter() + a = np.array([(1, 2), (3, 4)]) + np.savetxt(w, a) + b = np.loadtxt(w) + assert_array_equal(a, b) + + def test_unicode(self): + utf8 = b'\xcf\x96'.decode('UTF-8') + a = np.array([utf8], dtype=np.str_) + with tempdir() as tmpdir: + # set encoding as on windows it may not be unicode even on py3 + np.savetxt(os.path.join(tmpdir, 'test.csv'), a, fmt=['%s'], + encoding='UTF-8') + + def test_unicode_roundtrip(self): + utf8 = b'\xcf\x96'.decode('UTF-8') + a = np.array([utf8], dtype=np.str_) + # our gz wrapper support encoding + suffixes = ['', '.gz'] + if HAS_BZ2: + suffixes.append('.bz2') + if HAS_LZMA: + suffixes.extend(['.xz', '.lzma']) + with tempdir() as tmpdir: + for suffix in suffixes: + np.savetxt(os.path.join(tmpdir, 'test.csv' + suffix), a, + fmt=['%s'], encoding='UTF-16-LE') + b = np.loadtxt(os.path.join(tmpdir, 'test.csv' + suffix), + encoding='UTF-16-LE', dtype=np.str_) + assert_array_equal(a, b) + + def test_unicode_bytestream(self): + utf8 = b'\xcf\x96'.decode('UTF-8') + a = np.array([utf8], dtype=np.str_) + s = BytesIO() + np.savetxt(s, a, fmt=['%s'], encoding='UTF-8') + s.seek(0) + assert_equal(s.read().decode('UTF-8'), utf8 + '\n') + + def test_unicode_stringstream(self): + utf8 = b'\xcf\x96'.decode('UTF-8') + a = np.array([utf8], dtype=np.str_) + s = StringIO() + np.savetxt(s, a, fmt=['%s'], encoding='UTF-8') + s.seek(0) + assert_equal(s.read(), utf8 + '\n') + + @pytest.mark.parametrize("iotype", [StringIO, BytesIO]) + def test_unicode_and_bytes_fmt(self, iotype): + # string type of fmt should not matter, see also gh-4053 + a = np.array([1.]) + s = iotype() + np.savetxt(s, a, fmt="%f") + s.seek(0) + if iotype is StringIO: + assert_equal(s.read(), "%f\n" % 1.) + else: + assert_equal(s.read(), b"%f\n" % 1.) + + @pytest.mark.skipif(sys.platform=='win32', reason="files>4GB may not work") + @pytest.mark.slow + @requires_memory(free_bytes=7e9) + def test_large_zip(self): + def check_large_zip(memoryerror_raised): + memoryerror_raised.value = False + try: + # The test takes at least 6GB of memory, writes a file larger + # than 4GB. This tests the ``allowZip64`` kwarg to ``zipfile`` + test_data = np.asarray([np.random.rand( + np.random.randint(50,100),4) + for i in range(800000)], dtype=object) + with tempdir() as tmpdir: + np.savez(os.path.join(tmpdir, 'test.npz'), + test_data=test_data) + except MemoryError: + memoryerror_raised.value = True + raise + # run in a subprocess to ensure memory is released on PyPy, see gh-15775 + # Use an object in shared memory to re-raise the MemoryError exception + # in our process if needed, see gh-16889 + memoryerror_raised = Value(c_bool) + + # Since Python 3.8, the default start method for multiprocessing has + # been changed from 'fork' to 'spawn' on macOS, causing inconsistency + # on memory sharing model, lead to failed test for check_large_zip + ctx = get_context('fork') + p = ctx.Process(target=check_large_zip, args=(memoryerror_raised,)) + p.start() + p.join() + if memoryerror_raised.value: + raise MemoryError("Child process raised a MemoryError exception") + # -9 indicates a SIGKILL, probably an OOM. + if p.exitcode == -9: + pytest.xfail("subprocess got a SIGKILL, apparently free memory was not sufficient") + assert p.exitcode == 0 + +class LoadTxtBase: + def check_compressed(self, fopen, suffixes): + # Test that we can load data from a compressed file + wanted = np.arange(6).reshape((2, 3)) + linesep = ('\n', '\r\n', '\r') + for sep in linesep: + data = '0 1 2' + sep + '3 4 5' + for suffix in suffixes: + with temppath(suffix=suffix) as name: + with fopen(name, mode='wt', encoding='UTF-32-LE') as f: + f.write(data) + res = self.loadfunc(name, encoding='UTF-32-LE') + assert_array_equal(res, wanted) + with fopen(name, "rt", encoding='UTF-32-LE') as f: + res = self.loadfunc(f) + assert_array_equal(res, wanted) + + def test_compressed_gzip(self): + self.check_compressed(gzip.open, ('.gz',)) + + @pytest.mark.skipif(not HAS_BZ2, reason="Needs bz2") + def test_compressed_bz2(self): + self.check_compressed(bz2.open, ('.bz2',)) + + @pytest.mark.skipif(not HAS_LZMA, reason="Needs lzma") + def test_compressed_lzma(self): + self.check_compressed(lzma.open, ('.xz', '.lzma')) + + def test_encoding(self): + with temppath() as path: + with open(path, "wb") as f: + f.write('0.\n1.\n2.'.encode("UTF-16")) + x = self.loadfunc(path, encoding="UTF-16") + assert_array_equal(x, [0., 1., 2.]) + + def test_stringload(self): + # umlaute + nonascii = b'\xc3\xb6\xc3\xbc\xc3\xb6'.decode("UTF-8") + with temppath() as path: + with open(path, "wb") as f: + f.write(nonascii.encode("UTF-16")) + x = self.loadfunc(path, encoding="UTF-16", dtype=np.str_) + assert_array_equal(x, nonascii) + + def test_binary_decode(self): + utf16 = b'\xff\xfeh\x04 \x00i\x04 \x00j\x04' + v = self.loadfunc(BytesIO(utf16), dtype=np.str_, encoding='UTF-16') + assert_array_equal(v, np.array(utf16.decode('UTF-16').split())) + + def test_converters_decode(self): + # test converters that decode strings + c = TextIO() + c.write(b'\xcf\x96') + c.seek(0) + x = self.loadfunc(c, dtype=np.str_, encoding="bytes", + converters={0: lambda x: x.decode('UTF-8')}) + a = np.array([b'\xcf\x96'.decode('UTF-8')]) + assert_array_equal(x, a) + + def test_converters_nodecode(self): + # test native string converters enabled by setting an encoding + utf8 = b'\xcf\x96'.decode('UTF-8') + with temppath() as path: + with open(path, 'wt', encoding='UTF-8') as f: + f.write(utf8) + x = self.loadfunc(path, dtype=np.str_, + converters={0: lambda x: x + 't'}, + encoding='UTF-8') + a = np.array([utf8 + 't']) + assert_array_equal(x, a) + + +class TestLoadTxt(LoadTxtBase): + loadfunc = staticmethod(np.loadtxt) + + def setup_method(self): + # lower chunksize for testing + self.orig_chunk = _npyio_impl._loadtxt_chunksize + _npyio_impl._loadtxt_chunksize = 1 + + def teardown_method(self): + _npyio_impl._loadtxt_chunksize = self.orig_chunk + + def test_record(self): + c = TextIO() + c.write('1 2\n3 4') + c.seek(0) + x = np.loadtxt(c, dtype=[('x', np.int32), ('y', np.int32)]) + a = np.array([(1, 2), (3, 4)], dtype=[('x', 'i4'), ('y', 'i4')]) + assert_array_equal(x, a) + + d = TextIO() + d.write('M 64 75.0\nF 25 60.0') + d.seek(0) + mydescriptor = {'names': ('gender', 'age', 'weight'), + 'formats': ('S1', 'i4', 'f4')} + b = np.array([('M', 64.0, 75.0), + ('F', 25.0, 60.0)], dtype=mydescriptor) + y = np.loadtxt(d, dtype=mydescriptor) + assert_array_equal(y, b) + + def test_array(self): + c = TextIO() + c.write('1 2\n3 4') + + c.seek(0) + x = np.loadtxt(c, dtype=int) + a = np.array([[1, 2], [3, 4]], int) + assert_array_equal(x, a) + + c.seek(0) + x = np.loadtxt(c, dtype=float) + a = np.array([[1, 2], [3, 4]], float) + assert_array_equal(x, a) + + def test_1D(self): + c = TextIO() + c.write('1\n2\n3\n4\n') + c.seek(0) + x = np.loadtxt(c, dtype=int) + a = np.array([1, 2, 3, 4], int) + assert_array_equal(x, a) + + c = TextIO() + c.write('1,2,3,4\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',') + a = np.array([1, 2, 3, 4], int) + assert_array_equal(x, a) + + def test_missing(self): + c = TextIO() + c.write('1,2,3,,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + converters={3: lambda s: int(s or - 999)}) + a = np.array([1, 2, 3, -999, 5], int) + assert_array_equal(x, a) + + def test_converters_with_usecols(self): + c = TextIO() + c.write('1,2,3,,5\n6,7,8,9,10\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + converters={3: lambda s: int(s or - 999)}, + usecols=(1, 3,)) + a = np.array([[2, -999], [7, 9]], int) + assert_array_equal(x, a) + + def test_comments_unicode(self): + c = TextIO() + c.write('# comment\n1,2,3,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + comments='#') + a = np.array([1, 2, 3, 5], int) + assert_array_equal(x, a) + + def test_comments_byte(self): + c = TextIO() + c.write('# comment\n1,2,3,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + comments=b'#') + a = np.array([1, 2, 3, 5], int) + assert_array_equal(x, a) + + def test_comments_multiple(self): + c = TextIO() + c.write('# comment\n1,2,3\n@ comment2\n4,5,6 // comment3') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + comments=['#', '@', '//']) + a = np.array([[1, 2, 3], [4, 5, 6]], int) + assert_array_equal(x, a) + + @pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") + def test_comments_multi_chars(self): + c = TextIO() + c.write('/* comment\n1,2,3,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + comments='/*') + a = np.array([1, 2, 3, 5], int) + assert_array_equal(x, a) + + # Check that '/*' is not transformed to ['/', '*'] + c = TextIO() + c.write('*/ comment\n1,2,3,5\n') + c.seek(0) + assert_raises(ValueError, np.loadtxt, c, dtype=int, delimiter=',', + comments='/*') + + def test_skiprows(self): + c = TextIO() + c.write('comment\n1,2,3,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + skiprows=1) + a = np.array([1, 2, 3, 5], int) + assert_array_equal(x, a) + + c = TextIO() + c.write('# comment\n1,2,3,5\n') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + skiprows=1) + a = np.array([1, 2, 3, 5], int) + assert_array_equal(x, a) + + def test_usecols(self): + a = np.array([[1, 2], [3, 4]], float) + c = BytesIO() + np.savetxt(c, a) + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=(1,)) + assert_array_equal(x, a[:, 1]) + + a = np.array([[1, 2, 3], [3, 4, 5]], float) + c = BytesIO() + np.savetxt(c, a) + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=(1, 2)) + assert_array_equal(x, a[:, 1:]) + + # Testing with arrays instead of tuples. + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=np.array([1, 2])) + assert_array_equal(x, a[:, 1:]) + + # Testing with an integer instead of a sequence + for int_type in [int, np.int8, np.int16, + np.int32, np.int64, np.uint8, np.uint16, + np.uint32, np.uint64]: + to_read = int_type(1) + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=to_read) + assert_array_equal(x, a[:, 1]) + + # Testing with some crazy custom integer type + class CrazyInt: + def __index__(self): + return 1 + + crazy_int = CrazyInt() + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=crazy_int) + assert_array_equal(x, a[:, 1]) + + c.seek(0) + x = np.loadtxt(c, dtype=float, usecols=(crazy_int,)) + assert_array_equal(x, a[:, 1]) + + # Checking with dtypes defined converters. + data = '''JOE 70.1 25.3 + BOB 60.5 27.9 + ''' + c = TextIO(data) + names = ['stid', 'temp'] + dtypes = ['S4', 'f8'] + arr = np.loadtxt(c, usecols=(0, 2), dtype=list(zip(names, dtypes))) + assert_equal(arr['stid'], [b"JOE", b"BOB"]) + assert_equal(arr['temp'], [25.3, 27.9]) + + # Testing non-ints in usecols + c.seek(0) + bogus_idx = 1.5 + assert_raises_regex( + TypeError, + '^usecols must be.*%s' % type(bogus_idx).__name__, + np.loadtxt, c, usecols=bogus_idx + ) + + assert_raises_regex( + TypeError, + '^usecols must be.*%s' % type(bogus_idx).__name__, + np.loadtxt, c, usecols=[0, bogus_idx, 0] + ) + + def test_bad_usecols(self): + with pytest.raises(OverflowError): + np.loadtxt(["1\n"], usecols=[2**64], delimiter=",") + with pytest.raises((ValueError, OverflowError)): + # Overflow error on 32bit platforms + np.loadtxt(["1\n"], usecols=[2**62], delimiter=",") + with pytest.raises(TypeError, + match="If a structured dtype .*. But 1 usecols were given and " + "the number of fields is 3."): + np.loadtxt(["1,1\n"], dtype="i,2i", usecols=[0], delimiter=",") + + def test_fancy_dtype(self): + c = TextIO() + c.write('1,2,3.0\n4,5,6.0\n') + c.seek(0) + dt = np.dtype([('x', int), ('y', [('t', int), ('s', float)])]) + x = np.loadtxt(c, dtype=dt, delimiter=',') + a = np.array([(1, (2, 3.0)), (4, (5, 6.0))], dt) + assert_array_equal(x, a) + + def test_shaped_dtype(self): + c = TextIO("aaaa 1.0 8.0 1 2 3 4 5 6") + dt = np.dtype([('name', 'S4'), ('x', float), ('y', float), + ('block', int, (2, 3))]) + x = np.loadtxt(c, dtype=dt) + a = np.array([('aaaa', 1.0, 8.0, [[1, 2, 3], [4, 5, 6]])], + dtype=dt) + assert_array_equal(x, a) + + def test_3d_shaped_dtype(self): + c = TextIO("aaaa 1.0 8.0 1 2 3 4 5 6 7 8 9 10 11 12") + dt = np.dtype([('name', 'S4'), ('x', float), ('y', float), + ('block', int, (2, 2, 3))]) + x = np.loadtxt(c, dtype=dt) + a = np.array([('aaaa', 1.0, 8.0, + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])], + dtype=dt) + assert_array_equal(x, a) + + def test_str_dtype(self): + # see gh-8033 + c = ["str1", "str2"] + + for dt in (str, np.bytes_): + a = np.array(["str1", "str2"], dtype=dt) + x = np.loadtxt(c, dtype=dt) + assert_array_equal(x, a) + + def test_empty_file(self): + with pytest.warns(UserWarning, match="input contained no data"): + c = TextIO() + x = np.loadtxt(c) + assert_equal(x.shape, (0,)) + x = np.loadtxt(c, dtype=np.int64) + assert_equal(x.shape, (0,)) + assert_(x.dtype == np.int64) + + def test_unused_converter(self): + c = TextIO() + c.writelines(['1 21\n', '3 42\n']) + c.seek(0) + data = np.loadtxt(c, usecols=(1,), + converters={0: lambda s: int(s, 16)}) + assert_array_equal(data, [21, 42]) + + c.seek(0) + data = np.loadtxt(c, usecols=(1,), + converters={1: lambda s: int(s, 16)}) + assert_array_equal(data, [33, 66]) + + def test_dtype_with_object(self): + # Test using an explicit dtype with an object + data = """ 1; 2001-01-01 + 2; 2002-01-31 """ + ndtype = [('idx', int), ('code', object)] + func = lambda s: strptime(s.strip(), "%Y-%m-%d") + converters = {1: func} + test = np.loadtxt(TextIO(data), delimiter=";", dtype=ndtype, + converters=converters) + control = np.array( + [(1, datetime(2001, 1, 1)), (2, datetime(2002, 1, 31))], + dtype=ndtype) + assert_equal(test, control) + + def test_uint64_type(self): + tgt = (9223372043271415339, 9223372043271415853) + c = TextIO() + c.write("%s %s" % tgt) + c.seek(0) + res = np.loadtxt(c, dtype=np.uint64) + assert_equal(res, tgt) + + def test_int64_type(self): + tgt = (-9223372036854775807, 9223372036854775807) + c = TextIO() + c.write("%s %s" % tgt) + c.seek(0) + res = np.loadtxt(c, dtype=np.int64) + assert_equal(res, tgt) + + def test_from_float_hex(self): + # IEEE doubles and floats only, otherwise the float32 + # conversion may fail. + tgt = np.logspace(-10, 10, 5).astype(np.float32) + tgt = np.hstack((tgt, -tgt)).astype(float) + inp = '\n'.join(map(float.hex, tgt)) + c = TextIO() + c.write(inp) + for dt in [float, np.float32]: + c.seek(0) + res = np.loadtxt( + c, dtype=dt, converters=float.fromhex, encoding="latin1") + assert_equal(res, tgt, err_msg="%s" % dt) + + @pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") + def test_default_float_converter_no_default_hex_conversion(self): + """ + Ensure that fromhex is only used for values with the correct prefix and + is not called by default. Regression test related to gh-19598. + """ + c = TextIO("a b c") + with pytest.raises(ValueError, + match=".*convert string 'a' to float64 at row 0, column 1"): + np.loadtxt(c) + + @pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") + def test_default_float_converter_exception(self): + """ + Ensure that the exception message raised during failed floating point + conversion is correct. Regression test related to gh-19598. + """ + c = TextIO("qrs tuv") # Invalid values for default float converter + with pytest.raises(ValueError, + match="could not convert string 'qrs' to float64"): + np.loadtxt(c) + + def test_from_complex(self): + tgt = (complex(1, 1), complex(1, -1)) + c = TextIO() + c.write("%s %s" % tgt) + c.seek(0) + res = np.loadtxt(c, dtype=complex) + assert_equal(res, tgt) + + def test_complex_misformatted(self): + # test for backward compatibility + # some complex formats used to generate x+-yj + a = np.zeros((2, 2), dtype=np.complex128) + re = np.pi + im = np.e + a[:] = re - 1.0j * im + c = BytesIO() + np.savetxt(c, a, fmt='%.16e') + c.seek(0) + txt = c.read() + c.seek(0) + # misformat the sign on the imaginary part, gh 7895 + txt_bad = txt.replace(b'e+00-', b'e00+-') + assert_(txt_bad != txt) + c.write(txt_bad) + c.seek(0) + res = np.loadtxt(c, dtype=complex) + assert_equal(res, a) + + def test_universal_newline(self): + with temppath() as name: + with open(name, 'w') as f: + f.write('1 21\r3 42\r') + data = np.loadtxt(name) + assert_array_equal(data, [[1, 21], [3, 42]]) + + def test_empty_field_after_tab(self): + c = TextIO() + c.write('1 \t2 \t3\tstart \n4\t5\t6\t \n7\t8\t9.5\t') + c.seek(0) + dt = {'names': ('x', 'y', 'z', 'comment'), + 'formats': (' num rows + c = TextIO() + c.write('comment\n1,2,3,5\n4,5,7,8\n2,1,4,5') + c.seek(0) + x = np.loadtxt(c, dtype=int, delimiter=',', + skiprows=1, max_rows=6) + a = np.array([[1, 2, 3, 5], [4, 5, 7, 8], [2, 1, 4, 5]], int) + assert_array_equal(x, a) + + @pytest.mark.parametrize(["skip", "data"], [ + (1, ["ignored\n", "1,2\n", "\n", "3,4\n"]), + # "Bad" lines that do not end in newlines: + (1, ["ignored", "1,2", "", "3,4"]), + (1, StringIO("ignored\n1,2\n\n3,4")), + # Same as above, but do not skip any lines: + (0, ["-1,0\n", "1,2\n", "\n", "3,4\n"]), + (0, ["-1,0", "1,2", "", "3,4"]), + (0, StringIO("-1,0\n1,2\n\n3,4"))]) + def test_max_rows_empty_lines(self, skip, data): + with pytest.warns(UserWarning, + match=f"Input line 3.*max_rows={3-skip}"): + res = np.loadtxt(data, dtype=int, skiprows=skip, delimiter=",", + max_rows=3-skip) + assert_array_equal(res, [[-1, 0], [1, 2], [3, 4]][skip:]) + + if isinstance(data, StringIO): + data.seek(0) + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + with pytest.raises(UserWarning): + np.loadtxt(data, dtype=int, skiprows=skip, delimiter=",", + max_rows=3-skip) + +class Testfromregex: + def test_record(self): + c = TextIO() + c.write('1.312 foo\n1.534 bar\n4.444 qux') + c.seek(0) + + dt = [('num', np.float64), ('val', 'S3')] + x = np.fromregex(c, r"([0-9.]+)\s+(...)", dt) + a = np.array([(1.312, 'foo'), (1.534, 'bar'), (4.444, 'qux')], + dtype=dt) + assert_array_equal(x, a) + + def test_record_2(self): + c = TextIO() + c.write('1312 foo\n1534 bar\n4444 qux') + c.seek(0) + + dt = [('num', np.int32), ('val', 'S3')] + x = np.fromregex(c, r"(\d+)\s+(...)", dt) + a = np.array([(1312, 'foo'), (1534, 'bar'), (4444, 'qux')], + dtype=dt) + assert_array_equal(x, a) + + def test_record_3(self): + c = TextIO() + c.write('1312 foo\n1534 bar\n4444 qux') + c.seek(0) + + dt = [('num', np.float64)] + x = np.fromregex(c, r"(\d+)\s+...", dt) + a = np.array([(1312,), (1534,), (4444,)], dtype=dt) + assert_array_equal(x, a) + + @pytest.mark.parametrize("path_type", [str, Path]) + def test_record_unicode(self, path_type): + utf8 = b'\xcf\x96' + with temppath() as str_path: + path = path_type(str_path) + with open(path, 'wb') as f: + f.write(b'1.312 foo' + utf8 + b' \n1.534 bar\n4.444 qux') + + dt = [('num', np.float64), ('val', 'U4')] + x = np.fromregex(path, r"(?u)([0-9.]+)\s+(\w+)", dt, encoding='UTF-8') + a = np.array([(1.312, 'foo' + utf8.decode('UTF-8')), (1.534, 'bar'), + (4.444, 'qux')], dtype=dt) + assert_array_equal(x, a) + + regexp = re.compile(r"([0-9.]+)\s+(\w+)", re.UNICODE) + x = np.fromregex(path, regexp, dt, encoding='UTF-8') + assert_array_equal(x, a) + + def test_compiled_bytes(self): + regexp = re.compile(b'(\\d)') + c = BytesIO(b'123') + dt = [('num', np.float64)] + a = np.array([1, 2, 3], dtype=dt) + x = np.fromregex(c, regexp, dt) + assert_array_equal(x, a) + + def test_bad_dtype_not_structured(self): + regexp = re.compile(b'(\\d)') + c = BytesIO(b'123') + with pytest.raises(TypeError, match='structured datatype'): + np.fromregex(c, regexp, dtype=np.float64) + + +#####-------------------------------------------------------------------------- + + +class TestFromTxt(LoadTxtBase): + loadfunc = staticmethod(np.genfromtxt) + + def test_record(self): + # Test w/ explicit dtype + data = TextIO('1 2\n3 4') + test = np.genfromtxt(data, dtype=[('x', np.int32), ('y', np.int32)]) + control = np.array([(1, 2), (3, 4)], dtype=[('x', 'i4'), ('y', 'i4')]) + assert_equal(test, control) + # + data = TextIO('M 64.0 75.0\nF 25.0 60.0') + descriptor = {'names': ('gender', 'age', 'weight'), + 'formats': ('S1', 'i4', 'f4')} + control = np.array([('M', 64.0, 75.0), ('F', 25.0, 60.0)], + dtype=descriptor) + test = np.genfromtxt(data, dtype=descriptor) + assert_equal(test, control) + + def test_array(self): + # Test outputting a standard ndarray + data = TextIO('1 2\n3 4') + control = np.array([[1, 2], [3, 4]], dtype=int) + test = np.genfromtxt(data, dtype=int) + assert_array_equal(test, control) + # + data.seek(0) + control = np.array([[1, 2], [3, 4]], dtype=float) + test = np.loadtxt(data, dtype=float) + assert_array_equal(test, control) + + def test_1D(self): + # Test squeezing to 1D + control = np.array([1, 2, 3, 4], int) + # + data = TextIO('1\n2\n3\n4\n') + test = np.genfromtxt(data, dtype=int) + assert_array_equal(test, control) + # + data = TextIO('1,2,3,4\n') + test = np.genfromtxt(data, dtype=int, delimiter=',') + assert_array_equal(test, control) + + def test_comments(self): + # Test the stripping of comments + control = np.array([1, 2, 3, 5], int) + # Comment on its own line + data = TextIO('# comment\n1,2,3,5\n') + test = np.genfromtxt(data, dtype=int, delimiter=',', comments='#') + assert_equal(test, control) + # Comment at the end of a line + data = TextIO('1,2,3,5# comment\n') + test = np.genfromtxt(data, dtype=int, delimiter=',', comments='#') + assert_equal(test, control) + + def test_skiprows(self): + # Test row skipping + control = np.array([1, 2, 3, 5], int) + kwargs = dict(dtype=int, delimiter=',') + # + data = TextIO('comment\n1,2,3,5\n') + test = np.genfromtxt(data, skip_header=1, **kwargs) + assert_equal(test, control) + # + data = TextIO('# comment\n1,2,3,5\n') + test = np.loadtxt(data, skiprows=1, **kwargs) + assert_equal(test, control) + + def test_skip_footer(self): + data = ["# %i" % i for i in range(1, 6)] + data.append("A, B, C") + data.extend(["%i,%3.1f,%03s" % (i, i, i) for i in range(51)]) + data[-1] = "99,99" + kwargs = dict(delimiter=",", names=True, skip_header=5, skip_footer=10) + test = np.genfromtxt(TextIO("\n".join(data)), **kwargs) + ctrl = np.array([("%f" % i, "%f" % i, "%f" % i) for i in range(41)], + dtype=[(_, float) for _ in "ABC"]) + assert_equal(test, ctrl) + + def test_skip_footer_with_invalid(self): + with suppress_warnings() as sup: + sup.filter(ConversionWarning) + basestr = '1 1\n2 2\n3 3\n4 4\n5 \n6 \n7 \n' + # Footer too small to get rid of all invalid values + assert_raises(ValueError, np.genfromtxt, + TextIO(basestr), skip_footer=1) + # except ValueError: + # pass + a = np.genfromtxt( + TextIO(basestr), skip_footer=1, invalid_raise=False) + assert_equal(a, np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.]])) + # + a = np.genfromtxt(TextIO(basestr), skip_footer=3) + assert_equal(a, np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.]])) + # + basestr = '1 1\n2 \n3 3\n4 4\n5 \n6 6\n7 7\n' + a = np.genfromtxt( + TextIO(basestr), skip_footer=1, invalid_raise=False) + assert_equal(a, np.array([[1., 1.], [3., 3.], [4., 4.], [6., 6.]])) + a = np.genfromtxt( + TextIO(basestr), skip_footer=3, invalid_raise=False) + assert_equal(a, np.array([[1., 1.], [3., 3.], [4., 4.]])) + + def test_header(self): + # Test retrieving a header + data = TextIO('gender age weight\nM 64.0 75.0\nF 25.0 60.0') + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, dtype=None, names=True, + encoding='bytes') + assert_(w[0].category is VisibleDeprecationWarning) + control = {'gender': np.array([b'M', b'F']), + 'age': np.array([64.0, 25.0]), + 'weight': np.array([75.0, 60.0])} + assert_equal(test['gender'], control['gender']) + assert_equal(test['age'], control['age']) + assert_equal(test['weight'], control['weight']) + + def test_auto_dtype(self): + # Test the automatic definition of the output dtype + data = TextIO('A 64 75.0 3+4j True\nBCD 25 60.0 5+6j False') + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, dtype=None, encoding='bytes') + assert_(w[0].category is VisibleDeprecationWarning) + control = [np.array([b'A', b'BCD']), + np.array([64, 25]), + np.array([75.0, 60.0]), + np.array([3 + 4j, 5 + 6j]), + np.array([True, False]), ] + assert_equal(test.dtype.names, ['f0', 'f1', 'f2', 'f3', 'f4']) + for (i, ctrl) in enumerate(control): + assert_equal(test['f%i' % i], ctrl) + + def test_auto_dtype_uniform(self): + # Tests whether the output dtype can be uniformized + data = TextIO('1 2 3 4\n5 6 7 8\n') + test = np.genfromtxt(data, dtype=None) + control = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) + assert_equal(test, control) + + def test_fancy_dtype(self): + # Check that a nested dtype isn't MIA + data = TextIO('1,2,3.0\n4,5,6.0\n') + fancydtype = np.dtype([('x', int), ('y', [('t', int), ('s', float)])]) + test = np.genfromtxt(data, dtype=fancydtype, delimiter=',') + control = np.array([(1, (2, 3.0)), (4, (5, 6.0))], dtype=fancydtype) + assert_equal(test, control) + + def test_names_overwrite(self): + # Test overwriting the names of the dtype + descriptor = {'names': ('g', 'a', 'w'), + 'formats': ('S1', 'i4', 'f4')} + data = TextIO(b'M 64.0 75.0\nF 25.0 60.0') + names = ('gender', 'age', 'weight') + test = np.genfromtxt(data, dtype=descriptor, names=names) + descriptor['names'] = names + control = np.array([('M', 64.0, 75.0), + ('F', 25.0, 60.0)], dtype=descriptor) + assert_equal(test, control) + + def test_bad_fname(self): + with pytest.raises(TypeError, match='fname must be a string,'): + np.genfromtxt(123) + + def test_commented_header(self): + # Check that names can be retrieved even if the line is commented out. + data = TextIO(""" +#gender age weight +M 21 72.100000 +F 35 58.330000 +M 33 21.99 + """) + # The # is part of the first name and should be deleted automatically. + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, names=True, dtype=None, + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + ctrl = np.array([('M', 21, 72.1), ('F', 35, 58.33), ('M', 33, 21.99)], + dtype=[('gender', '|S1'), ('age', int), ('weight', float)]) + assert_equal(test, ctrl) + # Ditto, but we should get rid of the first element + data = TextIO(b""" +# gender age weight +M 21 72.100000 +F 35 58.330000 +M 33 21.99 + """) + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, names=True, dtype=None, + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + assert_equal(test, ctrl) + + def test_names_and_comments_none(self): + # Tests case when names is true but comments is None (gh-10780) + data = TextIO('col1 col2\n 1 2\n 3 4') + test = np.genfromtxt(data, dtype=(int, int), comments=None, names=True) + control = np.array([(1, 2), (3, 4)], dtype=[('col1', int), ('col2', int)]) + assert_equal(test, control) + + def test_file_is_closed_on_error(self): + # gh-13200 + with tempdir() as tmpdir: + fpath = os.path.join(tmpdir, "test.csv") + with open(fpath, "wb") as f: + f.write('\N{GREEK PI SYMBOL}'.encode()) + + # ResourceWarnings are emitted from a destructor, so won't be + # detected by regular propagation to errors. + with assert_no_warnings(): + with pytest.raises(UnicodeDecodeError): + np.genfromtxt(fpath, encoding="ascii") + + def test_autonames_and_usecols(self): + # Tests names and usecols + data = TextIO('A B C D\n aaaa 121 45 9.1') + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, usecols=('A', 'C', 'D'), + names=True, dtype=None, encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + control = np.array(('aaaa', 45, 9.1), + dtype=[('A', '|S4'), ('C', int), ('D', float)]) + assert_equal(test, control) + + def test_converters_with_usecols(self): + # Test the combination user-defined converters and usecol + data = TextIO('1,2,3,,5\n6,7,8,9,10\n') + test = np.genfromtxt(data, dtype=int, delimiter=',', + converters={3: lambda s: int(s or - 999)}, + usecols=(1, 3,)) + control = np.array([[2, -999], [7, 9]], int) + assert_equal(test, control) + + def test_converters_with_usecols_and_names(self): + # Tests names and usecols + data = TextIO('A B C D\n aaaa 121 45 9.1') + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(data, usecols=('A', 'C', 'D'), names=True, + dtype=None, encoding="bytes", + converters={'C': lambda s: 2 * int(s)}) + assert_(w[0].category is VisibleDeprecationWarning) + control = np.array(('aaaa', 90, 9.1), + dtype=[('A', '|S4'), ('C', int), ('D', float)]) + assert_equal(test, control) + + def test_converters_cornercases(self): + # Test the conversion to datetime. + converter = { + 'date': lambda s: strptime(s, '%Y-%m-%d %H:%M:%SZ')} + data = TextIO('2009-02-03 12:00:00Z, 72214.0') + test = np.genfromtxt(data, delimiter=',', dtype=None, + names=['date', 'stid'], converters=converter) + control = np.array((datetime(2009, 2, 3), 72214.), + dtype=[('date', np.object_), ('stid', float)]) + assert_equal(test, control) + + def test_converters_cornercases2(self): + # Test the conversion to datetime64. + converter = { + 'date': lambda s: np.datetime64(strptime(s, '%Y-%m-%d %H:%M:%SZ'))} + data = TextIO('2009-02-03 12:00:00Z, 72214.0') + test = np.genfromtxt(data, delimiter=',', dtype=None, + names=['date', 'stid'], converters=converter) + control = np.array((datetime(2009, 2, 3), 72214.), + dtype=[('date', 'datetime64[us]'), ('stid', float)]) + assert_equal(test, control) + + def test_unused_converter(self): + # Test whether unused converters are forgotten + data = TextIO("1 21\n 3 42\n") + test = np.genfromtxt(data, usecols=(1,), + converters={0: lambda s: int(s, 16)}) + assert_equal(test, [21, 42]) + # + data.seek(0) + test = np.genfromtxt(data, usecols=(1,), + converters={1: lambda s: int(s, 16)}) + assert_equal(test, [33, 66]) + + def test_invalid_converter(self): + strip_rand = lambda x: float((b'r' in x.lower() and x.split()[-1]) or + (b'r' not in x.lower() and x.strip() or 0.0)) + strip_per = lambda x: float((b'%' in x.lower() and x.split()[0]) or + (b'%' not in x.lower() and x.strip() or 0.0)) + s = TextIO("D01N01,10/1/2003 ,1 %,R 75,400,600\r\n" + "L24U05,12/5/2003, 2 %,1,300, 150.5\r\n" + "D02N03,10/10/2004,R 1,,7,145.55") + kwargs = dict( + converters={2: strip_per, 3: strip_rand}, delimiter=",", + dtype=None, encoding="bytes") + assert_raises(ConverterError, np.genfromtxt, s, **kwargs) + + def test_tricky_converter_bug1666(self): + # Test some corner cases + s = TextIO('q1,2\nq3,4') + cnv = lambda s: float(s[1:]) + test = np.genfromtxt(s, delimiter=',', converters={0: cnv}) + control = np.array([[1., 2.], [3., 4.]]) + assert_equal(test, control) + + def test_dtype_with_converters(self): + dstr = "2009; 23; 46" + test = np.genfromtxt(TextIO(dstr,), + delimiter=";", dtype=float, converters={0: bytes}) + control = np.array([('2009', 23., 46)], + dtype=[('f0', '|S4'), ('f1', float), ('f2', float)]) + assert_equal(test, control) + test = np.genfromtxt(TextIO(dstr,), + delimiter=";", dtype=float, converters={0: float}) + control = np.array([2009., 23., 46],) + assert_equal(test, control) + + @pytest.mark.filterwarnings("ignore:.*recfromcsv.*:DeprecationWarning") + def test_dtype_with_converters_and_usecols(self): + dstr = "1,5,-1,1:1\n2,8,-1,1:n\n3,3,-2,m:n\n" + dmap = {'1:1':0, '1:n':1, 'm:1':2, 'm:n':3} + dtyp = [('e1','i4'),('e2','i4'),('e3','i2'),('n', 'i1')] + conv = {0: int, 1: int, 2: int, 3: lambda r: dmap[r.decode()]} + test = recfromcsv(TextIO(dstr,), dtype=dtyp, delimiter=',', + names=None, converters=conv, encoding="bytes") + control = np.rec.array([(1,5,-1,0), (2,8,-1,1), (3,3,-2,3)], dtype=dtyp) + assert_equal(test, control) + dtyp = [('e1', 'i4'), ('e2', 'i4'), ('n', 'i1')] + test = recfromcsv(TextIO(dstr,), dtype=dtyp, delimiter=',', + usecols=(0, 1, 3), names=None, converters=conv, + encoding="bytes") + control = np.rec.array([(1,5,0), (2,8,1), (3,3,3)], dtype=dtyp) + assert_equal(test, control) + + def test_dtype_with_object(self): + # Test using an explicit dtype with an object + data = """ 1; 2001-01-01 + 2; 2002-01-31 """ + ndtype = [('idx', int), ('code', object)] + func = lambda s: strptime(s.strip(), "%Y-%m-%d") + converters = {1: func} + test = np.genfromtxt(TextIO(data), delimiter=";", dtype=ndtype, + converters=converters) + control = np.array( + [(1, datetime(2001, 1, 1)), (2, datetime(2002, 1, 31))], + dtype=ndtype) + assert_equal(test, control) + + ndtype = [('nest', [('idx', int), ('code', object)])] + with assert_raises_regex(NotImplementedError, + 'Nested fields.* not supported.*'): + test = np.genfromtxt(TextIO(data), delimiter=";", + dtype=ndtype, converters=converters) + + # nested but empty fields also aren't supported + ndtype = [('idx', int), ('code', object), ('nest', [])] + with assert_raises_regex(NotImplementedError, + 'Nested fields.* not supported.*'): + test = np.genfromtxt(TextIO(data), delimiter=";", + dtype=ndtype, converters=converters) + + def test_dtype_with_object_no_converter(self): + # Object without a converter uses bytes: + parsed = np.genfromtxt(TextIO("1"), dtype=object) + assert parsed[()] == b"1" + parsed = np.genfromtxt(TextIO("string"), dtype=object) + assert parsed[()] == b"string" + + def test_userconverters_with_explicit_dtype(self): + # Test user_converters w/ explicit (standard) dtype + data = TextIO('skip,skip,2001-01-01,1.0,skip') + test = np.genfromtxt(data, delimiter=",", names=None, dtype=float, + usecols=(2, 3), converters={2: bytes}) + control = np.array([('2001-01-01', 1.)], + dtype=[('', '|S10'), ('', float)]) + assert_equal(test, control) + + def test_utf8_userconverters_with_explicit_dtype(self): + utf8 = b'\xcf\x96' + with temppath() as path: + with open(path, 'wb') as f: + f.write(b'skip,skip,2001-01-01' + utf8 + b',1.0,skip') + test = np.genfromtxt(path, delimiter=",", names=None, dtype=float, + usecols=(2, 3), converters={2: str}, + encoding='UTF-8') + control = np.array([('2001-01-01' + utf8.decode('UTF-8'), 1.)], + dtype=[('', '|U11'), ('', float)]) + assert_equal(test, control) + + def test_spacedelimiter(self): + # Test space delimiter + data = TextIO("1 2 3 4 5\n6 7 8 9 10") + test = np.genfromtxt(data) + control = np.array([[1., 2., 3., 4., 5.], + [6., 7., 8., 9., 10.]]) + assert_equal(test, control) + + def test_integer_delimiter(self): + # Test using an integer for delimiter + data = " 1 2 3\n 4 5 67\n890123 4" + test = np.genfromtxt(TextIO(data), delimiter=3) + control = np.array([[1, 2, 3], [4, 5, 67], [890, 123, 4]]) + assert_equal(test, control) + + def test_missing(self): + data = TextIO('1,2,3,,5\n') + test = np.genfromtxt(data, dtype=int, delimiter=',', + converters={3: lambda s: int(s or - 999)}) + control = np.array([1, 2, 3, -999, 5], int) + assert_equal(test, control) + + def test_missing_with_tabs(self): + # Test w/ a delimiter tab + txt = "1\t2\t3\n\t2\t\n1\t\t3" + test = np.genfromtxt(TextIO(txt), delimiter="\t", + usemask=True,) + ctrl_d = np.array([(1, 2, 3), (np.nan, 2, np.nan), (1, np.nan, 3)],) + ctrl_m = np.array([(0, 0, 0), (1, 0, 1), (0, 1, 0)], dtype=bool) + assert_equal(test.data, ctrl_d) + assert_equal(test.mask, ctrl_m) + + def test_usecols(self): + # Test the selection of columns + # Select 1 column + control = np.array([[1, 2], [3, 4]], float) + data = TextIO() + np.savetxt(data, control) + data.seek(0) + test = np.genfromtxt(data, dtype=float, usecols=(1,)) + assert_equal(test, control[:, 1]) + # + control = np.array([[1, 2, 3], [3, 4, 5]], float) + data = TextIO() + np.savetxt(data, control) + data.seek(0) + test = np.genfromtxt(data, dtype=float, usecols=(1, 2)) + assert_equal(test, control[:, 1:]) + # Testing with arrays instead of tuples. + data.seek(0) + test = np.genfromtxt(data, dtype=float, usecols=np.array([1, 2])) + assert_equal(test, control[:, 1:]) + + def test_usecols_as_css(self): + # Test giving usecols with a comma-separated string + data = "1 2 3\n4 5 6" + test = np.genfromtxt(TextIO(data), + names="a, b, c", usecols="a, c") + ctrl = np.array([(1, 3), (4, 6)], dtype=[(_, float) for _ in "ac"]) + assert_equal(test, ctrl) + + def test_usecols_with_structured_dtype(self): + # Test usecols with an explicit structured dtype + data = TextIO("JOE 70.1 25.3\nBOB 60.5 27.9") + names = ['stid', 'temp'] + dtypes = ['S4', 'f8'] + test = np.genfromtxt( + data, usecols=(0, 2), dtype=list(zip(names, dtypes))) + assert_equal(test['stid'], [b"JOE", b"BOB"]) + assert_equal(test['temp'], [25.3, 27.9]) + + def test_usecols_with_integer(self): + # Test usecols with an integer + test = np.genfromtxt(TextIO(b"1 2 3\n4 5 6"), usecols=0) + assert_equal(test, np.array([1., 4.])) + + def test_usecols_with_named_columns(self): + # Test usecols with named columns + ctrl = np.array([(1, 3), (4, 6)], dtype=[('a', float), ('c', float)]) + data = "1 2 3\n4 5 6" + kwargs = dict(names="a, b, c") + test = np.genfromtxt(TextIO(data), usecols=(0, -1), **kwargs) + assert_equal(test, ctrl) + test = np.genfromtxt(TextIO(data), + usecols=('a', 'c'), **kwargs) + assert_equal(test, ctrl) + + def test_empty_file(self): + # Test that an empty file raises the proper warning. + with suppress_warnings() as sup: + sup.filter(message="genfromtxt: Empty input file:") + data = TextIO() + test = np.genfromtxt(data) + assert_equal(test, np.array([])) + + # when skip_header > 0 + test = np.genfromtxt(data, skip_header=1) + assert_equal(test, np.array([])) + + def test_fancy_dtype_alt(self): + # Check that a nested dtype isn't MIA + data = TextIO('1,2,3.0\n4,5,6.0\n') + fancydtype = np.dtype([('x', int), ('y', [('t', int), ('s', float)])]) + test = np.genfromtxt(data, dtype=fancydtype, delimiter=',', usemask=True) + control = ma.array([(1, (2, 3.0)), (4, (5, 6.0))], dtype=fancydtype) + assert_equal(test, control) + + def test_shaped_dtype(self): + c = TextIO("aaaa 1.0 8.0 1 2 3 4 5 6") + dt = np.dtype([('name', 'S4'), ('x', float), ('y', float), + ('block', int, (2, 3))]) + x = np.genfromtxt(c, dtype=dt) + a = np.array([('aaaa', 1.0, 8.0, [[1, 2, 3], [4, 5, 6]])], + dtype=dt) + assert_array_equal(x, a) + + def test_withmissing(self): + data = TextIO('A,B\n0,1\n2,N/A') + kwargs = dict(delimiter=",", missing_values="N/A", names=True) + test = np.genfromtxt(data, dtype=None, usemask=True, **kwargs) + control = ma.array([(0, 1), (2, -1)], + mask=[(False, False), (False, True)], + dtype=[('A', int), ('B', int)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + # + data.seek(0) + test = np.genfromtxt(data, usemask=True, **kwargs) + control = ma.array([(0, 1), (2, -1)], + mask=[(False, False), (False, True)], + dtype=[('A', float), ('B', float)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + def test_user_missing_values(self): + data = "A, B, C\n0, 0., 0j\n1, N/A, 1j\n-9, 2.2, N/A\n3, -99, 3j" + basekwargs = dict(dtype=None, delimiter=",", names=True,) + mdtype = [('A', int), ('B', float), ('C', complex)] + # + test = np.genfromtxt(TextIO(data), missing_values="N/A", + **basekwargs) + control = ma.array([(0, 0.0, 0j), (1, -999, 1j), + (-9, 2.2, -999j), (3, -99, 3j)], + mask=[(0, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 0)], + dtype=mdtype) + assert_equal(test, control) + # + basekwargs['dtype'] = mdtype + test = np.genfromtxt(TextIO(data), + missing_values={0: -9, 1: -99, 2: -999j}, usemask=True, **basekwargs) + control = ma.array([(0, 0.0, 0j), (1, -999, 1j), + (-9, 2.2, -999j), (3, -99, 3j)], + mask=[(0, 0, 0), (0, 1, 0), (1, 0, 1), (0, 1, 0)], + dtype=mdtype) + assert_equal(test, control) + # + test = np.genfromtxt(TextIO(data), + missing_values={0: -9, 'B': -99, 'C': -999j}, + usemask=True, + **basekwargs) + control = ma.array([(0, 0.0, 0j), (1, -999, 1j), + (-9, 2.2, -999j), (3, -99, 3j)], + mask=[(0, 0, 0), (0, 1, 0), (1, 0, 1), (0, 1, 0)], + dtype=mdtype) + assert_equal(test, control) + + def test_user_filling_values(self): + # Test with missing and filling values + ctrl = np.array([(0, 3), (4, -999)], dtype=[('a', int), ('b', int)]) + data = "N/A, 2, 3\n4, ,???" + kwargs = dict(delimiter=",", + dtype=int, + names="a,b,c", + missing_values={0: "N/A", 'b': " ", 2: "???"}, + filling_values={0: 0, 'b': 0, 2: -999}) + test = np.genfromtxt(TextIO(data), **kwargs) + ctrl = np.array([(0, 2, 3), (4, 0, -999)], + dtype=[(_, int) for _ in "abc"]) + assert_equal(test, ctrl) + # + test = np.genfromtxt(TextIO(data), usecols=(0, -1), **kwargs) + ctrl = np.array([(0, 3), (4, -999)], dtype=[(_, int) for _ in "ac"]) + assert_equal(test, ctrl) + + data2 = "1,2,*,4\n5,*,7,8\n" + test = np.genfromtxt(TextIO(data2), delimiter=',', dtype=int, + missing_values="*", filling_values=0) + ctrl = np.array([[1, 2, 0, 4], [5, 0, 7, 8]]) + assert_equal(test, ctrl) + test = np.genfromtxt(TextIO(data2), delimiter=',', dtype=int, + missing_values="*", filling_values=-1) + ctrl = np.array([[1, 2, -1, 4], [5, -1, 7, 8]]) + assert_equal(test, ctrl) + + def test_withmissing_float(self): + data = TextIO('A,B\n0,1.5\n2,-999.00') + test = np.genfromtxt(data, dtype=None, delimiter=',', + missing_values='-999.0', names=True, usemask=True) + control = ma.array([(0, 1.5), (2, -1.)], + mask=[(False, False), (False, True)], + dtype=[('A', int), ('B', float)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + def test_with_masked_column_uniform(self): + # Test masked column + data = TextIO('1 2 3\n4 5 6\n') + test = np.genfromtxt(data, dtype=None, + missing_values='2,5', usemask=True) + control = ma.array([[1, 2, 3], [4, 5, 6]], mask=[[0, 1, 0], [0, 1, 0]]) + assert_equal(test, control) + + def test_with_masked_column_various(self): + # Test masked column + data = TextIO('True 2 3\nFalse 5 6\n') + test = np.genfromtxt(data, dtype=None, + missing_values='2,5', usemask=True) + control = ma.array([(1, 2, 3), (0, 5, 6)], + mask=[(0, 1, 0), (0, 1, 0)], + dtype=[('f0', bool), ('f1', bool), ('f2', int)]) + assert_equal(test, control) + + def test_invalid_raise(self): + # Test invalid raise + data = ["1, 1, 1, 1, 1"] * 50 + for i in range(5): + data[10 * i] = "2, 2, 2, 2 2" + data.insert(0, "a, b, c, d, e") + mdata = TextIO("\n".join(data)) + + kwargs = dict(delimiter=",", dtype=None, names=True) + def f(): + return np.genfromtxt(mdata, invalid_raise=False, **kwargs) + mtest = assert_warns(ConversionWarning, f) + assert_equal(len(mtest), 45) + assert_equal(mtest, np.ones(45, dtype=[(_, int) for _ in 'abcde'])) + # + mdata.seek(0) + assert_raises(ValueError, np.genfromtxt, mdata, + delimiter=",", names=True) + + def test_invalid_raise_with_usecols(self): + # Test invalid_raise with usecols + data = ["1, 1, 1, 1, 1"] * 50 + for i in range(5): + data[10 * i] = "2, 2, 2, 2 2" + data.insert(0, "a, b, c, d, e") + mdata = TextIO("\n".join(data)) + + kwargs = dict(delimiter=",", dtype=None, names=True, + invalid_raise=False) + def f(): + return np.genfromtxt(mdata, usecols=(0, 4), **kwargs) + mtest = assert_warns(ConversionWarning, f) + assert_equal(len(mtest), 45) + assert_equal(mtest, np.ones(45, dtype=[(_, int) for _ in 'ae'])) + # + mdata.seek(0) + mtest = np.genfromtxt(mdata, usecols=(0, 1), **kwargs) + assert_equal(len(mtest), 50) + control = np.ones(50, dtype=[(_, int) for _ in 'ab']) + control[[10 * _ for _ in range(5)]] = (2, 2) + assert_equal(mtest, control) + + def test_inconsistent_dtype(self): + # Test inconsistent dtype + data = ["1, 1, 1, 1, -1.1"] * 50 + mdata = TextIO("\n".join(data)) + + converters = {4: lambda x: "(%s)" % x.decode()} + kwargs = dict(delimiter=",", converters=converters, + dtype=[(_, int) for _ in 'abcde'], encoding="bytes") + assert_raises(ValueError, np.genfromtxt, mdata, **kwargs) + + def test_default_field_format(self): + # Test default format + data = "0, 1, 2.3\n4, 5, 6.7" + mtest = np.genfromtxt(TextIO(data), + delimiter=",", dtype=None, defaultfmt="f%02i") + ctrl = np.array([(0, 1, 2.3), (4, 5, 6.7)], + dtype=[("f00", int), ("f01", int), ("f02", float)]) + assert_equal(mtest, ctrl) + + def test_single_dtype_wo_names(self): + # Test single dtype w/o names + data = "0, 1, 2.3\n4, 5, 6.7" + mtest = np.genfromtxt(TextIO(data), + delimiter=",", dtype=float, defaultfmt="f%02i") + ctrl = np.array([[0., 1., 2.3], [4., 5., 6.7]], dtype=float) + assert_equal(mtest, ctrl) + + def test_single_dtype_w_explicit_names(self): + # Test single dtype w explicit names + data = "0, 1, 2.3\n4, 5, 6.7" + mtest = np.genfromtxt(TextIO(data), + delimiter=",", dtype=float, names="a, b, c") + ctrl = np.array([(0., 1., 2.3), (4., 5., 6.7)], + dtype=[(_, float) for _ in "abc"]) + assert_equal(mtest, ctrl) + + def test_single_dtype_w_implicit_names(self): + # Test single dtype w implicit names + data = "a, b, c\n0, 1, 2.3\n4, 5, 6.7" + mtest = np.genfromtxt(TextIO(data), + delimiter=",", dtype=float, names=True) + ctrl = np.array([(0., 1., 2.3), (4., 5., 6.7)], + dtype=[(_, float) for _ in "abc"]) + assert_equal(mtest, ctrl) + + def test_easy_structured_dtype(self): + # Test easy structured dtype + data = "0, 1, 2.3\n4, 5, 6.7" + mtest = np.genfromtxt(TextIO(data), delimiter=",", + dtype=(int, float, float), defaultfmt="f_%02i") + ctrl = np.array([(0, 1., 2.3), (4, 5., 6.7)], + dtype=[("f_00", int), ("f_01", float), ("f_02", float)]) + assert_equal(mtest, ctrl) + + def test_autostrip(self): + # Test autostrip + data = "01/01/2003 , 1.3, abcde" + kwargs = dict(delimiter=",", dtype=None, encoding="bytes") + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + mtest = np.genfromtxt(TextIO(data), **kwargs) + assert_(w[0].category is VisibleDeprecationWarning) + ctrl = np.array([('01/01/2003 ', 1.3, ' abcde')], + dtype=[('f0', '|S12'), ('f1', float), ('f2', '|S8')]) + assert_equal(mtest, ctrl) + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + mtest = np.genfromtxt(TextIO(data), autostrip=True, **kwargs) + assert_(w[0].category is VisibleDeprecationWarning) + ctrl = np.array([('01/01/2003', 1.3, 'abcde')], + dtype=[('f0', '|S10'), ('f1', float), ('f2', '|S5')]) + assert_equal(mtest, ctrl) + + def test_replace_space(self): + # Test the 'replace_space' option + txt = "A.A, B (B), C:C\n1, 2, 3.14" + # Test default: replace ' ' by '_' and delete non-alphanum chars + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=None) + ctrl_dtype = [("AA", int), ("B_B", int), ("CC", float)] + ctrl = np.array((1, 2, 3.14), dtype=ctrl_dtype) + assert_equal(test, ctrl) + # Test: no replace, no delete + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=None, + replace_space='', deletechars='') + ctrl_dtype = [("A.A", int), ("B (B)", int), ("C:C", float)] + ctrl = np.array((1, 2, 3.14), dtype=ctrl_dtype) + assert_equal(test, ctrl) + # Test: no delete (spaces are replaced by _) + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=None, + deletechars='') + ctrl_dtype = [("A.A", int), ("B_(B)", int), ("C:C", float)] + ctrl = np.array((1, 2, 3.14), dtype=ctrl_dtype) + assert_equal(test, ctrl) + + def test_replace_space_known_dtype(self): + # Test the 'replace_space' (and related) options when dtype != None + txt = "A.A, B (B), C:C\n1, 2, 3" + # Test default: replace ' ' by '_' and delete non-alphanum chars + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=int) + ctrl_dtype = [("AA", int), ("B_B", int), ("CC", int)] + ctrl = np.array((1, 2, 3), dtype=ctrl_dtype) + assert_equal(test, ctrl) + # Test: no replace, no delete + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=int, + replace_space='', deletechars='') + ctrl_dtype = [("A.A", int), ("B (B)", int), ("C:C", int)] + ctrl = np.array((1, 2, 3), dtype=ctrl_dtype) + assert_equal(test, ctrl) + # Test: no delete (spaces are replaced by _) + test = np.genfromtxt(TextIO(txt), + delimiter=",", names=True, dtype=int, + deletechars='') + ctrl_dtype = [("A.A", int), ("B_(B)", int), ("C:C", int)] + ctrl = np.array((1, 2, 3), dtype=ctrl_dtype) + assert_equal(test, ctrl) + + def test_incomplete_names(self): + # Test w/ incomplete names + data = "A,,C\n0,1,2\n3,4,5" + kwargs = dict(delimiter=",", names=True) + # w/ dtype=None + ctrl = np.array([(0, 1, 2), (3, 4, 5)], + dtype=[(_, int) for _ in ('A', 'f0', 'C')]) + test = np.genfromtxt(TextIO(data), dtype=None, **kwargs) + assert_equal(test, ctrl) + # w/ default dtype + ctrl = np.array([(0, 1, 2), (3, 4, 5)], + dtype=[(_, float) for _ in ('A', 'f0', 'C')]) + test = np.genfromtxt(TextIO(data), **kwargs) + + def test_names_auto_completion(self): + # Make sure that names are properly completed + data = "1 2 3\n 4 5 6" + test = np.genfromtxt(TextIO(data), + dtype=(int, float, int), names="a") + ctrl = np.array([(1, 2, 3), (4, 5, 6)], + dtype=[('a', int), ('f0', float), ('f1', int)]) + assert_equal(test, ctrl) + + def test_names_with_usecols_bug1636(self): + # Make sure we pick up the right names w/ usecols + data = "A,B,C,D,E\n0,1,2,3,4\n0,1,2,3,4\n0,1,2,3,4" + ctrl_names = ("A", "C", "E") + test = np.genfromtxt(TextIO(data), + dtype=(int, int, int), delimiter=",", + usecols=(0, 2, 4), names=True) + assert_equal(test.dtype.names, ctrl_names) + # + test = np.genfromtxt(TextIO(data), + dtype=(int, int, int), delimiter=",", + usecols=("A", "C", "E"), names=True) + assert_equal(test.dtype.names, ctrl_names) + # + test = np.genfromtxt(TextIO(data), + dtype=int, delimiter=",", + usecols=("A", "C", "E"), names=True) + assert_equal(test.dtype.names, ctrl_names) + + def test_fixed_width_names(self): + # Test fix-width w/ names + data = " A B C\n 0 1 2.3\n 45 67 9." + kwargs = dict(delimiter=(5, 5, 4), names=True, dtype=None) + ctrl = np.array([(0, 1, 2.3), (45, 67, 9.)], + dtype=[('A', int), ('B', int), ('C', float)]) + test = np.genfromtxt(TextIO(data), **kwargs) + assert_equal(test, ctrl) + # + kwargs = dict(delimiter=5, names=True, dtype=None) + ctrl = np.array([(0, 1, 2.3), (45, 67, 9.)], + dtype=[('A', int), ('B', int), ('C', float)]) + test = np.genfromtxt(TextIO(data), **kwargs) + assert_equal(test, ctrl) + + def test_filling_values(self): + # Test missing values + data = b"1, 2, 3\n1, , 5\n0, 6, \n" + kwargs = dict(delimiter=",", dtype=None, filling_values=-999) + ctrl = np.array([[1, 2, 3], [1, -999, 5], [0, 6, -999]], dtype=int) + test = np.genfromtxt(TextIO(data), **kwargs) + assert_equal(test, ctrl) + + def test_comments_is_none(self): + # Github issue 329 (None was previously being converted to 'None'). + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(TextIO("test1,testNonetherestofthedata"), + dtype=None, comments=None, delimiter=',', + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + assert_equal(test[1], b'testNonetherestofthedata') + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(TextIO("test1, testNonetherestofthedata"), + dtype=None, comments=None, delimiter=',', + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + assert_equal(test[1], b' testNonetherestofthedata') + + def test_latin1(self): + latin1 = b'\xf6\xfc\xf6' + norm = b"norm1,norm2,norm3\n" + enc = b"test1,testNonethe" + latin1 + b",test3\n" + s = norm + enc + norm + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(TextIO(s), + dtype=None, comments=None, delimiter=',', + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + assert_equal(test[1, 0], b"test1") + assert_equal(test[1, 1], b"testNonethe" + latin1) + assert_equal(test[1, 2], b"test3") + test = np.genfromtxt(TextIO(s), + dtype=None, comments=None, delimiter=',', + encoding='latin1') + assert_equal(test[1, 0], "test1") + assert_equal(test[1, 1], "testNonethe" + latin1.decode('latin1')) + assert_equal(test[1, 2], "test3") + + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(TextIO(b"0,testNonethe" + latin1), + dtype=None, comments=None, delimiter=',', + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + assert_equal(test['f0'], 0) + assert_equal(test['f1'], b"testNonethe" + latin1) + + def test_binary_decode_autodtype(self): + utf16 = b'\xff\xfeh\x04 \x00i\x04 \x00j\x04' + v = self.loadfunc(BytesIO(utf16), dtype=None, encoding='UTF-16') + assert_array_equal(v, np.array(utf16.decode('UTF-16').split())) + + def test_utf8_byte_encoding(self): + utf8 = b"\xcf\x96" + norm = b"norm1,norm2,norm3\n" + enc = b"test1,testNonethe" + utf8 + b",test3\n" + s = norm + enc + norm + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', VisibleDeprecationWarning) + test = np.genfromtxt(TextIO(s), + dtype=None, comments=None, delimiter=',', + encoding="bytes") + assert_(w[0].category is VisibleDeprecationWarning) + ctl = np.array([ + [b'norm1', b'norm2', b'norm3'], + [b'test1', b'testNonethe' + utf8, b'test3'], + [b'norm1', b'norm2', b'norm3']]) + assert_array_equal(test, ctl) + + def test_utf8_file(self): + utf8 = b"\xcf\x96" + with temppath() as path: + with open(path, "wb") as f: + f.write((b"test1,testNonethe" + utf8 + b",test3\n") * 2) + test = np.genfromtxt(path, dtype=None, comments=None, + delimiter=',', encoding="UTF-8") + ctl = np.array([ + ["test1", "testNonethe" + utf8.decode("UTF-8"), "test3"], + ["test1", "testNonethe" + utf8.decode("UTF-8"), "test3"]], + dtype=np.str_) + assert_array_equal(test, ctl) + + # test a mixed dtype + with open(path, "wb") as f: + f.write(b"0,testNonethe" + utf8) + test = np.genfromtxt(path, dtype=None, comments=None, + delimiter=',', encoding="UTF-8") + assert_equal(test['f0'], 0) + assert_equal(test['f1'], "testNonethe" + utf8.decode("UTF-8")) + + def test_utf8_file_nodtype_unicode(self): + # bytes encoding with non-latin1 -> unicode upcast + utf8 = '\u03d6' + latin1 = '\xf6\xfc\xf6' + + # skip test if cannot encode utf8 test string with preferred + # encoding. The preferred encoding is assumed to be the default + # encoding of open. Will need to change this for PyTest, maybe + # using pytest.mark.xfail(raises=***). + try: + encoding = locale.getpreferredencoding() + utf8.encode(encoding) + except (UnicodeError, ImportError): + pytest.skip('Skipping test_utf8_file_nodtype_unicode, ' + 'unable to encode utf8 in preferred encoding') + + with temppath() as path: + with open(path, "wt") as f: + f.write("norm1,norm2,norm3\n") + f.write("norm1," + latin1 + ",norm3\n") + f.write("test1,testNonethe" + utf8 + ",test3\n") + with warnings.catch_warnings(record=True) as w: + warnings.filterwarnings('always', '', + VisibleDeprecationWarning) + test = np.genfromtxt(path, dtype=None, comments=None, + delimiter=',', encoding="bytes") + # Check for warning when encoding not specified. + assert_(w[0].category is VisibleDeprecationWarning) + ctl = np.array([ + ["norm1", "norm2", "norm3"], + ["norm1", latin1, "norm3"], + ["test1", "testNonethe" + utf8, "test3"]], + dtype=np.str_) + assert_array_equal(test, ctl) + + @pytest.mark.filterwarnings("ignore:.*recfromtxt.*:DeprecationWarning") + def test_recfromtxt(self): + # + data = TextIO('A,B\n0,1\n2,3') + kwargs = dict(delimiter=",", missing_values="N/A", names=True) + test = recfromtxt(data, **kwargs) + control = np.array([(0, 1), (2, 3)], + dtype=[('A', int), ('B', int)]) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + # + data = TextIO('A,B\n0,1\n2,N/A') + test = recfromtxt(data, dtype=None, usemask=True, **kwargs) + control = ma.array([(0, 1), (2, -1)], + mask=[(False, False), (False, True)], + dtype=[('A', int), ('B', int)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + assert_equal(test.A, [0, 2]) + + @pytest.mark.filterwarnings("ignore:.*recfromcsv.*:DeprecationWarning") + def test_recfromcsv(self): + # + data = TextIO('A,B\n0,1\n2,3') + kwargs = dict(missing_values="N/A", names=True, case_sensitive=True, + encoding="bytes") + test = recfromcsv(data, dtype=None, **kwargs) + control = np.array([(0, 1), (2, 3)], + dtype=[('A', int), ('B', int)]) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + # + data = TextIO('A,B\n0,1\n2,N/A') + test = recfromcsv(data, dtype=None, usemask=True, **kwargs) + control = ma.array([(0, 1), (2, -1)], + mask=[(False, False), (False, True)], + dtype=[('A', int), ('B', int)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + assert_equal(test.A, [0, 2]) + # + data = TextIO('A,B\n0,1\n2,3') + test = recfromcsv(data, missing_values='N/A',) + control = np.array([(0, 1), (2, 3)], + dtype=[('a', int), ('b', int)]) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + # + data = TextIO('A,B\n0,1\n2,3') + dtype = [('a', int), ('b', float)] + test = recfromcsv(data, missing_values='N/A', dtype=dtype) + control = np.array([(0, 1), (2, 3)], + dtype=dtype) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + + #gh-10394 + data = TextIO('color\n"red"\n"blue"') + test = recfromcsv(data, converters={0: lambda x: x.strip('\"')}) + control = np.array([('red',), ('blue',)], dtype=[('color', (str, 4))]) + assert_equal(test.dtype, control.dtype) + assert_equal(test, control) + + def test_max_rows(self): + # Test the `max_rows` keyword argument. + data = '1 2\n3 4\n5 6\n7 8\n9 10\n' + txt = TextIO(data) + a1 = np.genfromtxt(txt, max_rows=3) + a2 = np.genfromtxt(txt) + assert_equal(a1, [[1, 2], [3, 4], [5, 6]]) + assert_equal(a2, [[7, 8], [9, 10]]) + + # max_rows must be at least 1. + assert_raises(ValueError, np.genfromtxt, TextIO(data), max_rows=0) + + # An input with several invalid rows. + data = '1 1\n2 2\n0 \n3 3\n4 4\n5 \n6 \n7 \n' + + test = np.genfromtxt(TextIO(data), max_rows=2) + control = np.array([[1., 1.], [2., 2.]]) + assert_equal(test, control) + + # Test keywords conflict + assert_raises(ValueError, np.genfromtxt, TextIO(data), skip_footer=1, + max_rows=4) + + # Test with invalid value + assert_raises(ValueError, np.genfromtxt, TextIO(data), max_rows=4) + + # Test with invalid not raise + with suppress_warnings() as sup: + sup.filter(ConversionWarning) + + test = np.genfromtxt(TextIO(data), max_rows=4, invalid_raise=False) + control = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.]]) + assert_equal(test, control) + + test = np.genfromtxt(TextIO(data), max_rows=5, invalid_raise=False) + control = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.]]) + assert_equal(test, control) + + # Structured array with field names. + data = 'a b\n#c d\n1 1\n2 2\n#0 \n3 3\n4 4\n5 5\n' + + # Test with header, names and comments + txt = TextIO(data) + test = np.genfromtxt(txt, skip_header=1, max_rows=3, names=True) + control = np.array([(1.0, 1.0), (2.0, 2.0), (3.0, 3.0)], + dtype=[('c', ' should convert to float + # 2**34 = 17179869184 => should convert to int64 + # 2**10 = 1024 => should convert to int (int32 on 32-bit systems, + # int64 on 64-bit systems) + + data = TextIO('73786976294838206464 17179869184 1024') + + test = np.genfromtxt(data, dtype=None) + + assert_equal(test.dtype.names, ['f0', 'f1', 'f2']) + + assert_(test.dtype['f0'] == float) + assert_(test.dtype['f1'] == np.int64) + assert_(test.dtype['f2'] == np.int_) + + assert_allclose(test['f0'], 73786976294838206464.) + assert_equal(test['f1'], 17179869184) + assert_equal(test['f2'], 1024) + + def test_unpack_float_data(self): + txt = TextIO("1,2,3\n4,5,6\n7,8,9\n0.0,1.0,2.0") + a, b, c = np.loadtxt(txt, delimiter=",", unpack=True) + assert_array_equal(a, np.array([1.0, 4.0, 7.0, 0.0])) + assert_array_equal(b, np.array([2.0, 5.0, 8.0, 1.0])) + assert_array_equal(c, np.array([3.0, 6.0, 9.0, 2.0])) + + def test_unpack_structured(self): + # Regression test for gh-4341 + # Unpacking should work on structured arrays + txt = TextIO("M 21 72\nF 35 58") + dt = {'names': ('a', 'b', 'c'), 'formats': ('S1', 'i4', 'f4')} + a, b, c = np.genfromtxt(txt, dtype=dt, unpack=True) + assert_equal(a.dtype, np.dtype('S1')) + assert_equal(b.dtype, np.dtype('i4')) + assert_equal(c.dtype, np.dtype('f4')) + assert_array_equal(a, np.array([b'M', b'F'])) + assert_array_equal(b, np.array([21, 35])) + assert_array_equal(c, np.array([72., 58.])) + + def test_unpack_auto_dtype(self): + # Regression test for gh-4341 + # Unpacking should work when dtype=None + txt = TextIO("M 21 72.\nF 35 58.") + expected = (np.array(["M", "F"]), np.array([21, 35]), np.array([72., 58.])) + test = np.genfromtxt(txt, dtype=None, unpack=True, encoding="utf-8") + for arr, result in zip(expected, test): + assert_array_equal(arr, result) + assert_equal(arr.dtype, result.dtype) + + def test_unpack_single_name(self): + # Regression test for gh-4341 + # Unpacking should work when structured dtype has only one field + txt = TextIO("21\n35") + dt = {'names': ('a',), 'formats': ('i4',)} + expected = np.array([21, 35], dtype=np.int32) + test = np.genfromtxt(txt, dtype=dt, unpack=True) + assert_array_equal(expected, test) + assert_equal(expected.dtype, test.dtype) + + def test_squeeze_scalar(self): + # Regression test for gh-4341 + # Unpacking a scalar should give zero-dim output, + # even if dtype is structured + txt = TextIO("1") + dt = {'names': ('a',), 'formats': ('i4',)} + expected = np.array((1,), dtype=np.int32) + test = np.genfromtxt(txt, dtype=dt, unpack=True) + assert_array_equal(expected, test) + assert_equal((), test.shape) + assert_equal(expected.dtype, test.dtype) + + @pytest.mark.parametrize("ndim", [0, 1, 2]) + def test_ndmin_keyword(self, ndim: int): + # lets have the same behaviour of ndmin as loadtxt + # as they should be the same for non-missing values + txt = "42" + + a = np.loadtxt(StringIO(txt), ndmin=ndim) + b = np.genfromtxt(StringIO(txt), ndmin=ndim) + + assert_array_equal(a, b) + + +class TestPathUsage: + # Test that pathlib.Path can be used + def test_loadtxt(self): + with temppath(suffix='.txt') as path: + path = Path(path) + a = np.array([[1.1, 2], [3, 4]]) + np.savetxt(path, a) + x = np.loadtxt(path) + assert_array_equal(x, a) + + def test_save_load(self): + # Test that pathlib.Path instances can be used with save. + with temppath(suffix='.npy') as path: + path = Path(path) + a = np.array([[1, 2], [3, 4]], int) + np.save(path, a) + data = np.load(path) + assert_array_equal(data, a) + + def test_save_load_memmap(self): + # Test that pathlib.Path instances can be loaded mem-mapped. + with temppath(suffix='.npy') as path: + path = Path(path) + a = np.array([[1, 2], [3, 4]], int) + np.save(path, a) + data = np.load(path, mmap_mode='r') + assert_array_equal(data, a) + # close the mem-mapped file + del data + if IS_PYPY: + break_cycles() + break_cycles() + + @pytest.mark.xfail(IS_WASM, reason="memmap doesn't work correctly") + @pytest.mark.parametrize("filename_type", [Path, str]) + def test_save_load_memmap_readwrite(self, filename_type): + with temppath(suffix='.npy') as path: + path = filename_type(path) + a = np.array([[1, 2], [3, 4]], int) + np.save(path, a) + b = np.load(path, mmap_mode='r+') + a[0][0] = 5 + b[0][0] = 5 + del b # closes the file + if IS_PYPY: + break_cycles() + break_cycles() + data = np.load(path) + assert_array_equal(data, a) + + @pytest.mark.parametrize("filename_type", [Path, str]) + def test_savez_load(self, filename_type): + with temppath(suffix='.npz') as path: + path = filename_type(path) + np.savez(path, lab='place holder') + with np.load(path) as data: + assert_array_equal(data['lab'], 'place holder') + + @pytest.mark.parametrize("filename_type", [Path, str]) + def test_savez_compressed_load(self, filename_type): + with temppath(suffix='.npz') as path: + path = filename_type(path) + np.savez_compressed(path, lab='place holder') + data = np.load(path) + assert_array_equal(data['lab'], 'place holder') + data.close() + + @pytest.mark.parametrize("filename_type", [Path, str]) + def test_genfromtxt(self, filename_type): + with temppath(suffix='.txt') as path: + path = filename_type(path) + a = np.array([(1, 2), (3, 4)]) + np.savetxt(path, a) + data = np.genfromtxt(path) + assert_array_equal(a, data) + + @pytest.mark.parametrize("filename_type", [Path, str]) + @pytest.mark.filterwarnings("ignore:.*recfromtxt.*:DeprecationWarning") + def test_recfromtxt(self, filename_type): + with temppath(suffix='.txt') as path: + path = filename_type(path) + with open(path, 'w') as f: + f.write('A,B\n0,1\n2,3') + + kwargs = dict(delimiter=",", missing_values="N/A", names=True) + test = recfromtxt(path, **kwargs) + control = np.array([(0, 1), (2, 3)], + dtype=[('A', int), ('B', int)]) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + + @pytest.mark.parametrize("filename_type", [Path, str]) + @pytest.mark.filterwarnings("ignore:.*recfromcsv.*:DeprecationWarning") + def test_recfromcsv(self, filename_type): + with temppath(suffix='.txt') as path: + path = filename_type(path) + with open(path, 'w') as f: + f.write('A,B\n0,1\n2,3') + + kwargs = dict( + missing_values="N/A", names=True, case_sensitive=True + ) + test = recfromcsv(path, dtype=None, **kwargs) + control = np.array([(0, 1), (2, 3)], + dtype=[('A', int), ('B', int)]) + assert_(isinstance(test, np.recarray)) + assert_equal(test, control) + + +def test_gzip_load(): + a = np.random.random((5, 5)) + + s = BytesIO() + f = gzip.GzipFile(fileobj=s, mode="w") + + np.save(f, a) + f.close() + s.seek(0) + + f = gzip.GzipFile(fileobj=s, mode="r") + assert_array_equal(np.load(f), a) + + +# These next two classes encode the minimal API needed to save()/load() arrays. +# The `test_ducktyping` ensures they work correctly +class JustWriter: + def __init__(self, base): + self.base = base + + def write(self, s): + return self.base.write(s) + + def flush(self): + return self.base.flush() + +class JustReader: + def __init__(self, base): + self.base = base + + def read(self, n): + return self.base.read(n) + + def seek(self, off, whence=0): + return self.base.seek(off, whence) + + +def test_ducktyping(): + a = np.random.random((5, 5)) + + s = BytesIO() + f = JustWriter(s) + + np.save(f, a) + f.flush() + s.seek(0) + + f = JustReader(s) + assert_array_equal(np.load(f), a) + + + +def test_gzip_loadtxt(): + # Thanks to another windows brokenness, we can't use + # NamedTemporaryFile: a file created from this function cannot be + # reopened by another open call. So we first put the gzipped string + # of the test reference array, write it to a securely opened file, + # which is then read from by the loadtxt function + s = BytesIO() + g = gzip.GzipFile(fileobj=s, mode='w') + g.write(b'1 2 3\n') + g.close() + + s.seek(0) + with temppath(suffix='.gz') as name: + with open(name, 'wb') as f: + f.write(s.read()) + res = np.loadtxt(name) + s.close() + + assert_array_equal(res, [1, 2, 3]) + + +def test_gzip_loadtxt_from_string(): + s = BytesIO() + f = gzip.GzipFile(fileobj=s, mode="w") + f.write(b'1 2 3\n') + f.close() + s.seek(0) + + f = gzip.GzipFile(fileobj=s, mode="r") + assert_array_equal(np.loadtxt(f), [1, 2, 3]) + + +def test_npzfile_dict(): + s = BytesIO() + x = np.zeros((3, 3)) + y = np.zeros((3, 3)) + + np.savez(s, x=x, y=y) + s.seek(0) + + z = np.load(s) + + assert_('x' in z) + assert_('y' in z) + assert_('x' in z.keys()) + assert_('y' in z.keys()) + + for f, a in z.items(): + assert_(f in ['x', 'y']) + assert_equal(a.shape, (3, 3)) + + for a in z.values(): + assert_equal(a.shape, (3, 3)) + + assert_(len(z.items()) == 2) + + for f in z: + assert_(f in ['x', 'y']) + + assert_('x' in z.keys()) + assert (z.get('x') == z['x']).all() + + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_load_refcount(): + # Check that objects returned by np.load are directly freed based on + # their refcount, rather than needing the gc to collect them. + + f = BytesIO() + np.savez(f, [1, 2, 3]) + f.seek(0) + + with assert_no_gc_cycles(): + np.load(f) + + f.seek(0) + dt = [("a", 'u1', 2), ("b", 'u1', 2)] + with assert_no_gc_cycles(): + x = np.loadtxt(TextIO("0 1 2 3"), dtype=dt) + assert_equal(x, np.array([((0, 1), (2, 3))], dtype=dt)) + + +def test_load_multiple_arrays_until_eof(): + f = BytesIO() + np.save(f, 1) + np.save(f, 2) + f.seek(0) + out1 = np.load(f) + assert out1 == 1 + out2 = np.load(f) + assert out2 == 2 + with pytest.raises(EOFError): + np.load(f) + + +def test_savez_nopickle(): + obj_array = np.array([1, 'hello'], dtype=object) + with temppath(suffix='.npz') as tmp: + np.savez(tmp, obj_array) + + with temppath(suffix='.npz') as tmp: + with pytest.raises(ValueError, match="Object arrays cannot be saved when.*"): + np.savez(tmp, obj_array, allow_pickle=False) + + with temppath(suffix='.npz') as tmp: + np.savez_compressed(tmp, obj_array) + + with temppath(suffix='.npz') as tmp: + with pytest.raises(ValueError, match="Object arrays cannot be saved when.*"): + np.savez_compressed(tmp, obj_array, allow_pickle=False) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_loadtxt.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_loadtxt.py new file mode 100644 index 0000000000000000000000000000000000000000..60717be3bd9af7c8159beee4900bec2ce411c7cb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_loadtxt.py @@ -0,0 +1,1100 @@ +""" +Tests specific to `np.loadtxt` added during the move of loadtxt to be backed +by C code. +These tests complement those found in `test_io.py`. +""" + +import sys +import os +import pytest +from tempfile import NamedTemporaryFile, mkstemp +from io import StringIO + +import numpy as np +from numpy.ma.testutils import assert_equal +from numpy.testing import assert_array_equal, HAS_REFCOUNT, IS_PYPY + + +def test_scientific_notation(): + """Test that both 'e' and 'E' are parsed correctly.""" + data = StringIO( + + "1.0e-1,2.0E1,3.0\n" + "4.0e-2,5.0E-1,6.0\n" + "7.0e-3,8.0E1,9.0\n" + "0.0e-4,1.0E-1,2.0" + + ) + expected = np.array( + [[0.1, 20., 3.0], [0.04, 0.5, 6], [0.007, 80., 9], [0, 0.1, 2]] + ) + assert_array_equal(np.loadtxt(data, delimiter=","), expected) + + +@pytest.mark.parametrize("comment", ["..", "//", "@-", "this is a comment:"]) +def test_comment_multiple_chars(comment): + content = "# IGNORE\n1.5, 2.5# ABC\n3.0,4.0# XXX\n5.5,6.0\n" + txt = StringIO(content.replace("#", comment)) + a = np.loadtxt(txt, delimiter=",", comments=comment) + assert_equal(a, [[1.5, 2.5], [3.0, 4.0], [5.5, 6.0]]) + + +@pytest.fixture +def mixed_types_structured(): + """ + Fixture providing heterogeneous input data with a structured dtype, along + with the associated structured array. + """ + data = StringIO( + + "1000;2.4;alpha;-34\n" + "2000;3.1;beta;29\n" + "3500;9.9;gamma;120\n" + "4090;8.1;delta;0\n" + "5001;4.4;epsilon;-99\n" + "6543;7.8;omega;-1\n" + + ) + dtype = np.dtype( + [('f0', np.uint16), ('f1', np.float64), ('f2', 'S7'), ('f3', np.int8)] + ) + expected = np.array( + [ + (1000, 2.4, "alpha", -34), + (2000, 3.1, "beta", 29), + (3500, 9.9, "gamma", 120), + (4090, 8.1, "delta", 0), + (5001, 4.4, "epsilon", -99), + (6543, 7.8, "omega", -1) + ], + dtype=dtype + ) + return data, dtype, expected + + +@pytest.mark.parametrize('skiprows', [0, 1, 2, 3]) +def test_structured_dtype_and_skiprows_no_empty_lines( + skiprows, mixed_types_structured): + data, dtype, expected = mixed_types_structured + a = np.loadtxt(data, dtype=dtype, delimiter=";", skiprows=skiprows) + assert_array_equal(a, expected[skiprows:]) + + +def test_unpack_structured(mixed_types_structured): + data, dtype, expected = mixed_types_structured + + a, b, c, d = np.loadtxt(data, dtype=dtype, delimiter=";", unpack=True) + assert_array_equal(a, expected["f0"]) + assert_array_equal(b, expected["f1"]) + assert_array_equal(c, expected["f2"]) + assert_array_equal(d, expected["f3"]) + + +def test_structured_dtype_with_shape(): + dtype = np.dtype([("a", "u1", 2), ("b", "u1", 2)]) + data = StringIO("0,1,2,3\n6,7,8,9\n") + expected = np.array([((0, 1), (2, 3)), ((6, 7), (8, 9))], dtype=dtype) + assert_array_equal(np.loadtxt(data, delimiter=",", dtype=dtype), expected) + + +def test_structured_dtype_with_multi_shape(): + dtype = np.dtype([("a", "u1", (2, 2))]) + data = StringIO("0 1 2 3\n") + expected = np.array([(((0, 1), (2, 3)),)], dtype=dtype) + assert_array_equal(np.loadtxt(data, dtype=dtype), expected) + + +def test_nested_structured_subarray(): + # Test from gh-16678 + point = np.dtype([('x', float), ('y', float)]) + dt = np.dtype([('code', int), ('points', point, (2,))]) + data = StringIO("100,1,2,3,4\n200,5,6,7,8\n") + expected = np.array( + [ + (100, [(1., 2.), (3., 4.)]), + (200, [(5., 6.), (7., 8.)]), + ], + dtype=dt + ) + assert_array_equal(np.loadtxt(data, dtype=dt, delimiter=","), expected) + + +def test_structured_dtype_offsets(): + # An aligned structured dtype will have additional padding + dt = np.dtype("i1, i4, i1, i4, i1, i4", align=True) + data = StringIO("1,2,3,4,5,6\n7,8,9,10,11,12\n") + expected = np.array([(1, 2, 3, 4, 5, 6), (7, 8, 9, 10, 11, 12)], dtype=dt) + assert_array_equal(np.loadtxt(data, delimiter=",", dtype=dt), expected) + + +@pytest.mark.parametrize("param", ("skiprows", "max_rows")) +def test_exception_negative_row_limits(param): + """skiprows and max_rows should raise for negative parameters.""" + with pytest.raises(ValueError, match="argument must be nonnegative"): + np.loadtxt("foo.bar", **{param: -3}) + + +@pytest.mark.parametrize("param", ("skiprows", "max_rows")) +def test_exception_noninteger_row_limits(param): + with pytest.raises(TypeError, match="argument must be an integer"): + np.loadtxt("foo.bar", **{param: 1.0}) + + +@pytest.mark.parametrize( + "data, shape", + [ + ("1 2 3 4 5\n", (1, 5)), # Single row + ("1\n2\n3\n4\n5\n", (5, 1)), # Single column + ] +) +def test_ndmin_single_row_or_col(data, shape): + arr = np.array([1, 2, 3, 4, 5]) + arr2d = arr.reshape(shape) + + assert_array_equal(np.loadtxt(StringIO(data), dtype=int), arr) + assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=0), arr) + assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=1), arr) + assert_array_equal(np.loadtxt(StringIO(data), dtype=int, ndmin=2), arr2d) + + +@pytest.mark.parametrize("badval", [-1, 3, None, "plate of shrimp"]) +def test_bad_ndmin(badval): + with pytest.raises(ValueError, match="Illegal value of ndmin keyword"): + np.loadtxt("foo.bar", ndmin=badval) + + +@pytest.mark.parametrize( + "ws", + ( + " ", # space + "\t", # tab + "\u2003", # em + "\u00A0", # non-break + "\u3000", # ideographic space + ) +) +def test_blank_lines_spaces_delimit(ws): + txt = StringIO( + f"1 2{ws}30\n\n{ws}\n" + f"4 5 60{ws}\n {ws} \n" + f"7 8 {ws} 90\n # comment\n" + f"3 2 1" + ) + # NOTE: It is unclear that the ` # comment` should succeed. Except + # for delimiter=None, which should use any whitespace (and maybe + # should just be implemented closer to Python + expected = np.array([[1, 2, 30], [4, 5, 60], [7, 8, 90], [3, 2, 1]]) + assert_equal( + np.loadtxt(txt, dtype=int, delimiter=None, comments="#"), expected + ) + + +def test_blank_lines_normal_delimiter(): + txt = StringIO('1,2,30\n\n4,5,60\n\n7,8,90\n# comment\n3,2,1') + expected = np.array([[1, 2, 30], [4, 5, 60], [7, 8, 90], [3, 2, 1]]) + assert_equal( + np.loadtxt(txt, dtype=int, delimiter=',', comments="#"), expected + ) + + +@pytest.mark.parametrize("dtype", (float, object)) +def test_maxrows_no_blank_lines(dtype): + txt = StringIO("1.5,2.5\n3.0,4.0\n5.5,6.0") + res = np.loadtxt(txt, dtype=dtype, delimiter=",", max_rows=2) + assert_equal(res.dtype, dtype) + assert_equal(res, np.array([["1.5", "2.5"], ["3.0", "4.0"]], dtype=dtype)) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", (np.dtype("f8"), np.dtype("i2"))) +def test_exception_message_bad_values(dtype): + txt = StringIO("1,2\n3,XXX\n5,6") + msg = f"could not convert string 'XXX' to {dtype} at row 1, column 2" + with pytest.raises(ValueError, match=msg): + np.loadtxt(txt, dtype=dtype, delimiter=",") + + +def test_converters_negative_indices(): + txt = StringIO('1.5,2.5\n3.0,XXX\n5.5,6.0') + conv = {-1: lambda s: np.nan if s == 'XXX' else float(s)} + expected = np.array([[1.5, 2.5], [3.0, np.nan], [5.5, 6.0]]) + res = np.loadtxt(txt, dtype=np.float64, delimiter=",", converters=conv) + assert_equal(res, expected) + + +def test_converters_negative_indices_with_usecols(): + txt = StringIO('1.5,2.5,3.5\n3.0,4.0,XXX\n5.5,6.0,7.5\n') + conv = {-1: lambda s: np.nan if s == 'XXX' else float(s)} + expected = np.array([[1.5, 3.5], [3.0, np.nan], [5.5, 7.5]]) + res = np.loadtxt( + txt, + dtype=np.float64, + delimiter=",", + converters=conv, + usecols=[0, -1], + ) + assert_equal(res, expected) + + # Second test with variable number of rows: + res = np.loadtxt(StringIO('''0,1,2\n0,1,2,3,4'''), delimiter=",", + usecols=[0, -1], converters={-1: (lambda x: -1)}) + assert_array_equal(res, [[0, -1], [0, -1]]) + + +def test_ragged_error(): + rows = ["1,2,3", "1,2,3", "4,3,2,1"] + with pytest.raises(ValueError, + match="the number of columns changed from 3 to 4 at row 3"): + np.loadtxt(rows, delimiter=",") + + +def test_ragged_usecols(): + # usecols, and negative ones, work even with varying number of columns. + txt = StringIO("0,0,XXX\n0,XXX,0,XXX\n0,XXX,XXX,0,XXX\n") + expected = np.array([[0, 0], [0, 0], [0, 0]]) + res = np.loadtxt(txt, dtype=float, delimiter=",", usecols=[0, -2]) + assert_equal(res, expected) + + txt = StringIO("0,0,XXX\n0\n0,XXX,XXX,0,XXX\n") + with pytest.raises(ValueError, + match="invalid column index -2 at row 2 with 1 columns"): + # There is no -2 column in the second row: + np.loadtxt(txt, dtype=float, delimiter=",", usecols=[0, -2]) + + +def test_empty_usecols(): + txt = StringIO("0,0,XXX\n0,XXX,0,XXX\n0,XXX,XXX,0,XXX\n") + res = np.loadtxt(txt, dtype=np.dtype([]), delimiter=",", usecols=[]) + assert res.shape == (3,) + assert res.dtype == np.dtype([]) + + +@pytest.mark.parametrize("c1", ["a", "の", "🫕"]) +@pytest.mark.parametrize("c2", ["a", "の", "🫕"]) +def test_large_unicode_characters(c1, c2): + # c1 and c2 span ascii, 16bit and 32bit range. + txt = StringIO(f"a,{c1},c,1.0\ne,{c2},2.0,g") + res = np.loadtxt(txt, dtype=np.dtype('U12'), delimiter=",") + expected = np.array( + [f"a,{c1},c,1.0".split(","), f"e,{c2},2.0,g".split(",")], + dtype=np.dtype('U12') + ) + assert_equal(res, expected) + + +def test_unicode_with_converter(): + txt = StringIO("cat,dog\nαβγ,δεζ\nabc,def\n") + conv = {0: lambda s: s.upper()} + res = np.loadtxt( + txt, + dtype=np.dtype("U12"), + converters=conv, + delimiter=",", + encoding=None + ) + expected = np.array([['CAT', 'dog'], ['ΑΒΓ', 'δεζ'], ['ABC', 'def']]) + assert_equal(res, expected) + + +def test_converter_with_structured_dtype(): + txt = StringIO('1.5,2.5,Abc\n3.0,4.0,dEf\n5.5,6.0,ghI\n') + dt = np.dtype([('m', np.int32), ('r', np.float32), ('code', 'U8')]) + conv = {0: lambda s: int(10*float(s)), -1: lambda s: s.upper()} + res = np.loadtxt(txt, dtype=dt, delimiter=",", converters=conv) + expected = np.array( + [(15, 2.5, 'ABC'), (30, 4.0, 'DEF'), (55, 6.0, 'GHI')], dtype=dt + ) + assert_equal(res, expected) + + +def test_converter_with_unicode_dtype(): + """ + With the 'bytes' encoding, tokens are encoded prior to being + passed to the converter. This means that the output of the converter may + be bytes instead of unicode as expected by `read_rows`. + + This test checks that outputs from the above scenario are properly decoded + prior to parsing by `read_rows`. + """ + txt = StringIO('abc,def\nrst,xyz') + conv = bytes.upper + res = np.loadtxt( + txt, dtype=np.dtype("U3"), converters=conv, delimiter=",", + encoding="bytes") + expected = np.array([['ABC', 'DEF'], ['RST', 'XYZ']]) + assert_equal(res, expected) + + +def test_read_huge_row(): + row = "1.5, 2.5," * 50000 + row = row[:-1] + "\n" + txt = StringIO(row * 2) + res = np.loadtxt(txt, delimiter=",", dtype=float) + assert_equal(res, np.tile([1.5, 2.5], (2, 50000))) + + +@pytest.mark.parametrize("dtype", "edfgFDG") +def test_huge_float(dtype): + # Covers a non-optimized path that is rarely taken: + field = "0" * 1000 + ".123456789" + dtype = np.dtype(dtype) + value = np.loadtxt([field], dtype=dtype)[()] + assert value == dtype.type("0.123456789") + + +@pytest.mark.parametrize( + ("given_dtype", "expected_dtype"), + [ + ("S", np.dtype("S5")), + ("U", np.dtype("U5")), + ], +) +def test_string_no_length_given(given_dtype, expected_dtype): + """ + The given dtype is just 'S' or 'U' with no length. In these cases, the + length of the resulting dtype is determined by the longest string found + in the file. + """ + txt = StringIO("AAA,5-1\nBBBBB,0-3\nC,4-9\n") + res = np.loadtxt(txt, dtype=given_dtype, delimiter=",") + expected = np.array( + [['AAA', '5-1'], ['BBBBB', '0-3'], ['C', '4-9']], dtype=expected_dtype + ) + assert_equal(res, expected) + assert_equal(res.dtype, expected_dtype) + + +def test_float_conversion(): + """ + Some tests that the conversion to float64 works as accurately as the + Python built-in `float` function. In a naive version of the float parser, + these strings resulted in values that were off by an ULP or two. + """ + strings = [ + '0.9999999999999999', + '9876543210.123456', + '5.43215432154321e+300', + '0.901', + '0.333', + ] + txt = StringIO('\n'.join(strings)) + res = np.loadtxt(txt) + expected = np.array([float(s) for s in strings]) + assert_equal(res, expected) + + +def test_bool(): + # Simple test for bool via integer + txt = StringIO("1, 0\n10, -1") + res = np.loadtxt(txt, dtype=bool, delimiter=",") + assert res.dtype == bool + assert_array_equal(res, [[True, False], [True, True]]) + # Make sure we use only 1 and 0 on the byte level: + assert_array_equal(res.view(np.uint8), [[1, 0], [1, 1]]) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) +@pytest.mark.filterwarnings("error:.*integer via a float.*:DeprecationWarning") +def test_integer_signs(dtype): + dtype = np.dtype(dtype) + assert np.loadtxt(["+2"], dtype=dtype) == 2 + if dtype.kind == "u": + with pytest.raises(ValueError): + np.loadtxt(["-1\n"], dtype=dtype) + else: + assert np.loadtxt(["-2\n"], dtype=dtype) == -2 + + for sign in ["++", "+-", "--", "-+"]: + with pytest.raises(ValueError): + np.loadtxt([f"{sign}2\n"], dtype=dtype) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", np.typecodes["AllInteger"]) +@pytest.mark.filterwarnings("error:.*integer via a float.*:DeprecationWarning") +def test_implicit_cast_float_to_int_fails(dtype): + txt = StringIO("1.0, 2.1, 3.7\n4, 5, 6") + with pytest.raises(ValueError): + np.loadtxt(txt, dtype=dtype, delimiter=",") + +@pytest.mark.parametrize("dtype", (np.complex64, np.complex128)) +@pytest.mark.parametrize("with_parens", (False, True)) +def test_complex_parsing(dtype, with_parens): + s = "(1.0-2.5j),3.75,(7+-5.0j)\n(4),(-19e2j),(0)" + if not with_parens: + s = s.replace("(", "").replace(")", "") + + res = np.loadtxt(StringIO(s), dtype=dtype, delimiter=",") + expected = np.array( + [[1.0-2.5j, 3.75, 7-5j], [4.0, -1900j, 0]], dtype=dtype + ) + assert_equal(res, expected) + + +def test_read_from_generator(): + def gen(): + for i in range(4): + yield f"{i},{2*i},{i**2}" + + res = np.loadtxt(gen(), dtype=int, delimiter=",") + expected = np.array([[0, 0, 0], [1, 2, 1], [2, 4, 4], [3, 6, 9]]) + assert_equal(res, expected) + + +def test_read_from_generator_multitype(): + def gen(): + for i in range(3): + yield f"{i} {i / 4}" + + res = np.loadtxt(gen(), dtype="i, d", delimiter=" ") + expected = np.array([(0, 0.0), (1, 0.25), (2, 0.5)], dtype="i, d") + assert_equal(res, expected) + + +def test_read_from_bad_generator(): + def gen(): + yield from ["1,2", b"3, 5", 12738] + + with pytest.raises( + TypeError, match=r"non-string returned while reading data"): + np.loadtxt(gen(), dtype="i, i", delimiter=",") + + +@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") +def test_object_cleanup_on_read_error(): + sentinel = object() + already_read = 0 + + def conv(x): + nonlocal already_read + if already_read > 4999: + raise ValueError("failed half-way through!") + already_read += 1 + return sentinel + + txt = StringIO("x\n" * 10000) + + with pytest.raises(ValueError, match="at row 5000, column 1"): + np.loadtxt(txt, dtype=object, converters={0: conv}) + + assert sys.getrefcount(sentinel) == 2 + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +def test_character_not_bytes_compatible(): + """Test exception when a character cannot be encoded as 'S'.""" + data = StringIO("–") # == \u2013 + with pytest.raises(ValueError): + np.loadtxt(data, dtype="S5") + + +@pytest.mark.parametrize("conv", (0, [float], "")) +def test_invalid_converter(conv): + msg = ( + "converters must be a dictionary mapping columns to converter " + "functions or a single callable." + ) + with pytest.raises(TypeError, match=msg): + np.loadtxt(StringIO("1 2\n3 4"), converters=conv) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +def test_converters_dict_raises_non_integer_key(): + with pytest.raises(TypeError, match="keys of the converters dict"): + np.loadtxt(StringIO("1 2\n3 4"), converters={"a": int}) + with pytest.raises(TypeError, match="keys of the converters dict"): + np.loadtxt(StringIO("1 2\n3 4"), converters={"a": int}, usecols=0) + + +@pytest.mark.parametrize("bad_col_ind", (3, -3)) +def test_converters_dict_raises_non_col_key(bad_col_ind): + data = StringIO("1 2\n3 4") + with pytest.raises(ValueError, match="converter specified for column"): + np.loadtxt(data, converters={bad_col_ind: int}) + + +def test_converters_dict_raises_val_not_callable(): + with pytest.raises(TypeError, + match="values of the converters dictionary must be callable"): + np.loadtxt(StringIO("1 2\n3 4"), converters={0: 1}) + + +@pytest.mark.parametrize("q", ('"', "'", "`")) +def test_quoted_field(q): + txt = StringIO( + f"{q}alpha, x{q}, 2.5\n{q}beta, y{q}, 4.5\n{q}gamma, z{q}, 5.0\n" + ) + dtype = np.dtype([('f0', 'U8'), ('f1', np.float64)]) + expected = np.array( + [("alpha, x", 2.5), ("beta, y", 4.5), ("gamma, z", 5.0)], dtype=dtype + ) + + res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar=q) + assert_array_equal(res, expected) + + +@pytest.mark.parametrize("q", ('"', "'", "`")) +def test_quoted_field_with_whitepace_delimiter(q): + txt = StringIO( + f"{q}alpha, x{q} 2.5\n{q}beta, y{q} 4.5\n{q}gamma, z{q} 5.0\n" + ) + dtype = np.dtype([('f0', 'U8'), ('f1', np.float64)]) + expected = np.array( + [("alpha, x", 2.5), ("beta, y", 4.5), ("gamma, z", 5.0)], dtype=dtype + ) + + res = np.loadtxt(txt, dtype=dtype, delimiter=None, quotechar=q) + assert_array_equal(res, expected) + + +def test_quote_support_default(): + """Support for quoted fields is disabled by default.""" + txt = StringIO('"lat,long", 45, 30\n') + dtype = np.dtype([('f0', 'U24'), ('f1', np.float64), ('f2', np.float64)]) + + with pytest.raises(ValueError, + match="the dtype passed requires 3 columns but 4 were"): + np.loadtxt(txt, dtype=dtype, delimiter=",") + + # Enable quoting support with non-None value for quotechar param + txt.seek(0) + expected = np.array([("lat,long", 45., 30.)], dtype=dtype) + + res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar='"') + assert_array_equal(res, expected) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +def test_quotechar_multichar_error(): + txt = StringIO("1,2\n3,4") + msg = r".*must be a single unicode character or None" + with pytest.raises(TypeError, match=msg): + np.loadtxt(txt, delimiter=",", quotechar="''") + + +def test_comment_multichar_error_with_quote(): + txt = StringIO("1,2\n3,4") + msg = ( + "when multiple comments or a multi-character comment is given, " + "quotes are not supported." + ) + with pytest.raises(ValueError, match=msg): + np.loadtxt(txt, delimiter=",", comments="123", quotechar='"') + with pytest.raises(ValueError, match=msg): + np.loadtxt(txt, delimiter=",", comments=["#", "%"], quotechar='"') + + # A single character string in a tuple is unpacked though: + res = np.loadtxt(txt, delimiter=",", comments=("#",), quotechar="'") + assert_equal(res, [[1, 2], [3, 4]]) + + +def test_structured_dtype_with_quotes(): + data = StringIO( + + "1000;2.4;'alpha';-34\n" + "2000;3.1;'beta';29\n" + "3500;9.9;'gamma';120\n" + "4090;8.1;'delta';0\n" + "5001;4.4;'epsilon';-99\n" + "6543;7.8;'omega';-1\n" + + ) + dtype = np.dtype( + [('f0', np.uint16), ('f1', np.float64), ('f2', 'S7'), ('f3', np.int8)] + ) + expected = np.array( + [ + (1000, 2.4, "alpha", -34), + (2000, 3.1, "beta", 29), + (3500, 9.9, "gamma", 120), + (4090, 8.1, "delta", 0), + (5001, 4.4, "epsilon", -99), + (6543, 7.8, "omega", -1) + ], + dtype=dtype + ) + res = np.loadtxt(data, dtype=dtype, delimiter=";", quotechar="'") + assert_array_equal(res, expected) + + +def test_quoted_field_is_not_empty(): + txt = StringIO('1\n\n"4"\n""') + expected = np.array(["1", "4", ""], dtype="U1") + res = np.loadtxt(txt, delimiter=",", dtype="U1", quotechar='"') + assert_equal(res, expected) + +def test_quoted_field_is_not_empty_nonstrict(): + # Same as test_quoted_field_is_not_empty but check that we are not strict + # about missing closing quote (this is the `csv.reader` default also) + txt = StringIO('1\n\n"4"\n"') + expected = np.array(["1", "4", ""], dtype="U1") + res = np.loadtxt(txt, delimiter=",", dtype="U1", quotechar='"') + assert_equal(res, expected) + +def test_consecutive_quotechar_escaped(): + txt = StringIO('"Hello, my name is ""Monty""!"') + expected = np.array('Hello, my name is "Monty"!', dtype="U40") + res = np.loadtxt(txt, dtype="U40", delimiter=",", quotechar='"') + assert_equal(res, expected) + + +@pytest.mark.parametrize("data", ("", "\n\n\n", "# 1 2 3\n# 4 5 6\n")) +@pytest.mark.parametrize("ndmin", (0, 1, 2)) +@pytest.mark.parametrize("usecols", [None, (1, 2, 3)]) +def test_warn_on_no_data(data, ndmin, usecols): + """Check that a UserWarning is emitted when no data is read from input.""" + if usecols is not None: + expected_shape = (0, 3) + elif ndmin == 2: + expected_shape = (0, 1) # guess a single column?! + else: + expected_shape = (0,) + + txt = StringIO(data) + with pytest.warns(UserWarning, match="input contained no data"): + res = np.loadtxt(txt, ndmin=ndmin, usecols=usecols) + assert res.shape == expected_shape + + with NamedTemporaryFile(mode="w") as fh: + fh.write(data) + fh.seek(0) + with pytest.warns(UserWarning, match="input contained no data"): + res = np.loadtxt(txt, ndmin=ndmin, usecols=usecols) + assert res.shape == expected_shape + +@pytest.mark.parametrize("skiprows", (2, 3)) +def test_warn_on_skipped_data(skiprows): + data = "1 2 3\n4 5 6" + txt = StringIO(data) + with pytest.warns(UserWarning, match="input contained no data"): + np.loadtxt(txt, skiprows=skiprows) + + +@pytest.mark.parametrize(["dtype", "value"], [ + ("i2", 0x0001), ("u2", 0x0001), + ("i4", 0x00010203), ("u4", 0x00010203), + ("i8", 0x0001020304050607), ("u8", 0x0001020304050607), + # The following values are constructed to lead to unique bytes: + ("float16", 3.07e-05), + ("float32", 9.2557e-41), ("complex64", 9.2557e-41+2.8622554e-29j), + ("float64", -1.758571353180402e-24), + # Here and below, the repr side-steps a small loss of precision in + # complex `str` in PyPy (which is probably fine, as repr works): + ("complex128", repr(5.406409232372729e-29-1.758571353180402e-24j)), + # Use integer values that fit into double. Everything else leads to + # problems due to longdoubles going via double and decimal strings + # causing rounding errors. + ("longdouble", 0x01020304050607), + ("clongdouble", repr(0x01020304050607 + (0x00121314151617 * 1j))), + ("U2", "\U00010203\U000a0b0c")]) +@pytest.mark.parametrize("swap", [True, False]) +def test_byteswapping_and_unaligned(dtype, value, swap): + # Try to create "interesting" values within the valid unicode range: + dtype = np.dtype(dtype) + data = [f"x,{value}\n"] # repr as PyPy `str` truncates some + if swap: + dtype = dtype.newbyteorder() + full_dt = np.dtype([("a", "S1"), ("b", dtype)], align=False) + # The above ensures that the interesting "b" field is unaligned: + assert full_dt.fields["b"][1] == 1 + res = np.loadtxt(data, dtype=full_dt, delimiter=",", + max_rows=1) # max-rows prevents over-allocation + assert res["b"] == dtype.type(value) + + +@pytest.mark.parametrize("dtype", + np.typecodes["AllInteger"] + "efdFD" + "?") +def test_unicode_whitespace_stripping(dtype): + # Test that all numeric types (and bool) strip whitespace correctly + # \u202F is a narrow no-break space, `\n` is just a whitespace if quoted. + # Currently, skip float128 as it did not always support this and has no + # "custom" parsing: + txt = StringIO(' 3 ,"\u202F2\n"') + res = np.loadtxt(txt, dtype=dtype, delimiter=",", quotechar='"') + assert_array_equal(res, np.array([3, 2]).astype(dtype)) + + +@pytest.mark.parametrize("dtype", "FD") +def test_unicode_whitespace_stripping_complex(dtype): + # Complex has a few extra cases since it has two components and + # parentheses + line = " 1 , 2+3j , ( 4+5j ), ( 6+-7j ) , 8j , ( 9j ) \n" + data = [line, line.replace(" ", "\u202F")] + res = np.loadtxt(data, dtype=dtype, delimiter=',') + assert_array_equal(res, np.array([[1, 2+3j, 4+5j, 6-7j, 8j, 9j]] * 2)) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", "FD") +@pytest.mark.parametrize("field", + ["1 +2j", "1+ 2j", "1+2 j", "1+-+3", "(1j", "(1", "(1+2j", "1+2j)"]) +def test_bad_complex(dtype, field): + with pytest.raises(ValueError): + np.loadtxt([field + "\n"], dtype=dtype, delimiter=",") + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", + np.typecodes["AllInteger"] + "efgdFDG" + "?") +def test_nul_character_error(dtype): + # Test that a \0 character is correctly recognized as an error even if + # what comes before is valid (not everything gets parsed internally). + if dtype.lower() == "g": + pytest.xfail("longdouble/clongdouble assignment may misbehave.") + with pytest.raises(ValueError): + np.loadtxt(["1\000"], dtype=dtype, delimiter=",", quotechar='"') + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +@pytest.mark.parametrize("dtype", + np.typecodes["AllInteger"] + "efgdFDG" + "?") +def test_no_thousands_support(dtype): + # Mainly to document behaviour, Python supports thousands like 1_1. + # (e and G may end up using different conversion and support it, this is + # a bug but happens...) + if dtype == "e": + pytest.skip("half assignment currently uses Python float converter") + if dtype in "eG": + pytest.xfail("clongdouble assignment is buggy (uses `complex`?).") + + assert int("1_1") == float("1_1") == complex("1_1") == 11 + with pytest.raises(ValueError): + np.loadtxt(["1_1\n"], dtype=dtype) + + +@pytest.mark.parametrize("data", [ + ["1,2\n", "2\n,3\n"], + ["1,2\n", "2\r,3\n"]]) +def test_bad_newline_in_iterator(data): + # In NumPy <=1.22 this was accepted, because newlines were completely + # ignored when the input was an iterable. This could be changed, but right + # now, we raise an error. + msg = "Found an unquoted embedded newline within a single line" + with pytest.raises(ValueError, match=msg): + np.loadtxt(data, delimiter=",") + + +@pytest.mark.parametrize("data", [ + ["1,2\n", "2,3\r\n"], # a universal newline + ["1,2\n", "'2\n',3\n"], # a quoted newline + ["1,2\n", "'2\r',3\n"], + ["1,2\n", "'2\r\n',3\n"], +]) +def test_good_newline_in_iterator(data): + # The quoted newlines will be untransformed here, but are just whitespace. + res = np.loadtxt(data, delimiter=",", quotechar="'") + assert_array_equal(res, [[1., 2.], [2., 3.]]) + + +@pytest.mark.parametrize("newline", ["\n", "\r", "\r\n"]) +def test_universal_newlines_quoted(newline): + # Check that universal newline support within the tokenizer is not applied + # to quoted fields. (note that lines must end in newline or quoted + # fields will not include a newline at all) + data = ['1,"2\n"\n', '3,"4\n', '1"\n'] + data = [row.replace("\n", newline) for row in data] + res = np.loadtxt(data, dtype=object, delimiter=",", quotechar='"') + assert_array_equal(res, [['1', f'2{newline}'], ['3', f'4{newline}1']]) + + +def test_null_character(): + # Basic tests to check that the NUL character is not special: + res = np.loadtxt(["1\0002\0003\n", "4\0005\0006"], delimiter="\000") + assert_array_equal(res, [[1, 2, 3], [4, 5, 6]]) + + # Also not as part of a field (avoid unicode/arrays as unicode strips \0) + res = np.loadtxt(["1\000,2\000,3\n", "4\000,5\000,6"], + delimiter=",", dtype=object) + assert res.tolist() == [["1\000", "2\000", "3"], ["4\000", "5\000", "6"]] + + +def test_iterator_fails_getting_next_line(): + class BadSequence: + def __len__(self): + return 100 + + def __getitem__(self, item): + if item == 50: + raise RuntimeError("Bad things happened!") + return f"{item}, {item+1}" + + with pytest.raises(RuntimeError, match="Bad things happened!"): + np.loadtxt(BadSequence(), dtype=int, delimiter=",") + + +class TestCReaderUnitTests: + # These are internal tests for path that should not be possible to hit + # unless things go very very wrong somewhere. + def test_not_an_filelike(self): + with pytest.raises(AttributeError, match=".*read"): + np._core._multiarray_umath._load_from_filelike( + object(), dtype=np.dtype("i"), filelike=True) + + def test_filelike_read_fails(self): + # Can only be reached if loadtxt opens the file, so it is hard to do + # via the public interface (although maybe not impossible considering + # the current "DataClass" backing). + class BadFileLike: + counter = 0 + + def read(self, size): + self.counter += 1 + if self.counter > 20: + raise RuntimeError("Bad bad bad!") + return "1,2,3\n" + + with pytest.raises(RuntimeError, match="Bad bad bad!"): + np._core._multiarray_umath._load_from_filelike( + BadFileLike(), dtype=np.dtype("i"), filelike=True) + + def test_filelike_bad_read(self): + # Can only be reached if loadtxt opens the file, so it is hard to do + # via the public interface (although maybe not impossible considering + # the current "DataClass" backing). + + class BadFileLike: + counter = 0 + + def read(self, size): + return 1234 # not a string! + + with pytest.raises(TypeError, + match="non-string returned while reading data"): + np._core._multiarray_umath._load_from_filelike( + BadFileLike(), dtype=np.dtype("i"), filelike=True) + + def test_not_an_iter(self): + with pytest.raises(TypeError, + match="error reading from object, expected an iterable"): + np._core._multiarray_umath._load_from_filelike( + object(), dtype=np.dtype("i"), filelike=False) + + def test_bad_type(self): + with pytest.raises(TypeError, match="internal error: dtype must"): + np._core._multiarray_umath._load_from_filelike( + object(), dtype="i", filelike=False) + + def test_bad_encoding(self): + with pytest.raises(TypeError, match="encoding must be a unicode"): + np._core._multiarray_umath._load_from_filelike( + object(), dtype=np.dtype("i"), filelike=False, encoding=123) + + @pytest.mark.parametrize("newline", ["\r", "\n", "\r\n"]) + def test_manual_universal_newlines(self, newline): + # This is currently not available to users, because we should always + # open files with universal newlines enabled `newlines=None`. + # (And reading from an iterator uses slightly different code paths.) + # We have no real support for `newline="\r"` or `newline="\n" as the + # user cannot specify those options. + data = StringIO('0\n1\n"2\n"\n3\n4 #\n'.replace("\n", newline), + newline="") + + res = np._core._multiarray_umath._load_from_filelike( + data, dtype=np.dtype("U10"), filelike=True, + quote='"', comment="#", skiplines=1) + assert_array_equal(res[:, 0], ["1", f"2{newline}", "3", "4 "]) + + +def test_delimiter_comment_collision_raises(): + with pytest.raises(TypeError, match=".*control characters.*incompatible"): + np.loadtxt(StringIO("1, 2, 3"), delimiter=",", comments=",") + + +def test_delimiter_quotechar_collision_raises(): + with pytest.raises(TypeError, match=".*control characters.*incompatible"): + np.loadtxt(StringIO("1, 2, 3"), delimiter=",", quotechar=",") + + +def test_comment_quotechar_collision_raises(): + with pytest.raises(TypeError, match=".*control characters.*incompatible"): + np.loadtxt(StringIO("1 2 3"), comments="#", quotechar="#") + + +def test_delimiter_and_multiple_comments_collision_raises(): + with pytest.raises( + TypeError, match="Comment characters.*cannot include the delimiter" + ): + np.loadtxt(StringIO("1, 2, 3"), delimiter=",", comments=["#", ","]) + + +@pytest.mark.parametrize( + "ws", + ( + " ", # space + "\t", # tab + "\u2003", # em + "\u00A0", # non-break + "\u3000", # ideographic space + ) +) +def test_collision_with_default_delimiter_raises(ws): + with pytest.raises(TypeError, match=".*control characters.*incompatible"): + np.loadtxt(StringIO(f"1{ws}2{ws}3\n4{ws}5{ws}6\n"), comments=ws) + with pytest.raises(TypeError, match=".*control characters.*incompatible"): + np.loadtxt(StringIO(f"1{ws}2{ws}3\n4{ws}5{ws}6\n"), quotechar=ws) + + +@pytest.mark.parametrize("nl", ("\n", "\r")) +def test_control_character_newline_raises(nl): + txt = StringIO(f"1{nl}2{nl}3{nl}{nl}4{nl}5{nl}6{nl}{nl}") + msg = "control character.*cannot be a newline" + with pytest.raises(TypeError, match=msg): + np.loadtxt(txt, delimiter=nl) + with pytest.raises(TypeError, match=msg): + np.loadtxt(txt, comments=nl) + with pytest.raises(TypeError, match=msg): + np.loadtxt(txt, quotechar=nl) + + +@pytest.mark.parametrize( + ("generic_data", "long_datum", "unitless_dtype", "expected_dtype"), + [ + ("2012-03", "2013-01-15", "M8", "M8[D]"), # Datetimes + ("spam-a-lot", "tis_but_a_scratch", "U", "U17"), # str + ], +) +@pytest.mark.parametrize("nrows", (10, 50000, 60000)) # lt, eq, gt chunksize +def test_parametric_unit_discovery( + generic_data, long_datum, unitless_dtype, expected_dtype, nrows +): + """Check that the correct unit (e.g. month, day, second) is discovered from + the data when a user specifies a unitless datetime.""" + # Unit should be "D" (days) due to last entry + data = [generic_data] * nrows + [long_datum] + expected = np.array(data, dtype=expected_dtype) + assert len(data) == nrows+1 + assert len(data) == len(expected) + + # file-like path + txt = StringIO("\n".join(data)) + a = np.loadtxt(txt, dtype=unitless_dtype) + assert len(a) == len(expected) + assert a.dtype == expected.dtype + assert_equal(a, expected) + + # file-obj path + fd, fname = mkstemp() + os.close(fd) + with open(fname, "w") as fh: + fh.write("\n".join(data)+"\n") + # loading the full file... + a = np.loadtxt(fname, dtype=unitless_dtype) + assert len(a) == len(expected) + assert a.dtype == expected.dtype + assert_equal(a, expected) + # loading half of the file... + a = np.loadtxt(fname, dtype=unitless_dtype, max_rows=int(nrows/2)) + os.remove(fname) + assert len(a) == int(nrows/2) + assert_equal(a, expected[:int(nrows/2)]) + + +def test_str_dtype_unit_discovery_with_converter(): + data = ["spam-a-lot"] * 60000 + ["XXXtis_but_a_scratch"] + expected = np.array( + ["spam-a-lot"] * 60000 + ["tis_but_a_scratch"], dtype="U17" + ) + conv = lambda s: s.removeprefix("XXX") + + # file-like path + txt = StringIO("\n".join(data)) + a = np.loadtxt(txt, dtype="U", converters=conv) + assert a.dtype == expected.dtype + assert_equal(a, expected) + + # file-obj path + fd, fname = mkstemp() + os.close(fd) + with open(fname, "w") as fh: + fh.write("\n".join(data)) + a = np.loadtxt(fname, dtype="U", converters=conv) + os.remove(fname) + assert a.dtype == expected.dtype + assert_equal(a, expected) + + +@pytest.mark.skipif(IS_PYPY and sys.implementation.version <= (7, 3, 8), + reason="PyPy bug in error formatting") +def test_control_character_empty(): + with pytest.raises(TypeError, match="Text reading control character must"): + np.loadtxt(StringIO("1 2 3"), delimiter="") + with pytest.raises(TypeError, match="Text reading control character must"): + np.loadtxt(StringIO("1 2 3"), quotechar="") + with pytest.raises(ValueError, match="comments cannot be an empty string"): + np.loadtxt(StringIO("1 2 3"), comments="") + with pytest.raises(ValueError, match="comments cannot be an empty string"): + np.loadtxt(StringIO("1 2 3"), comments=["#", ""]) + + +def test_control_characters_as_bytes(): + """Byte control characters (comments, delimiter) are supported.""" + a = np.loadtxt(StringIO("#header\n1,2,3"), comments=b"#", delimiter=b",") + assert_equal(a, [1, 2, 3]) + + +@pytest.mark.filterwarnings('ignore::UserWarning') +def test_field_growing_cases(): + # Test empty field appending/growing (each field still takes 1 character) + # to see if the final field appending does not create issues. + res = np.loadtxt([""], delimiter=",", dtype=bytes) + assert len(res) == 0 + + for i in range(1, 1024): + res = np.loadtxt(["," * i], delimiter=",", dtype=bytes, max_rows=10) + assert len(res) == i+1 + +@pytest.mark.parametrize("nmax", (10000, 50000, 55000, 60000)) +def test_maxrows_exceeding_chunksize(nmax): + # tries to read all of the file, + # or less, equal, greater than _loadtxt_chunksize + file_length = 60000 + + # file-like path + data = ["a 0.5 1"]*file_length + txt = StringIO("\n".join(data)) + res = np.loadtxt(txt, dtype=str, delimiter=" ", max_rows=nmax) + assert len(res) == nmax + + # file-obj path + fd, fname = mkstemp() + os.close(fd) + with open(fname, "w") as fh: + fh.write("\n".join(data)) + res = np.loadtxt(fname, dtype=str, delimiter=" ", max_rows=nmax) + os.remove(fname) + assert len(res) == nmax + +@pytest.mark.parametrize("nskip", (0, 10000, 12345, 50000, 67891, 100000)) +def test_skiprow_exceeding_maxrows_exceeding_chunksize(tmpdir, nskip): + # tries to read a file in chunks by skipping a variable amount of lines, + # less, equal, greater than max_rows + file_length = 110000 + data = "\n".join(f"{i} a 0.5 1" for i in range(1, file_length + 1)) + expected_length = min(60000, file_length - nskip) + expected = np.arange(nskip + 1, nskip + 1 + expected_length).astype(str) + + # file-like path + txt = StringIO(data) + res = np.loadtxt(txt, dtype='str', delimiter=" ", skiprows=nskip, max_rows=60000) + assert len(res) == expected_length + # are the right lines read in res? + assert_array_equal(expected, res[:, 0]) + + # file-obj path + tmp_file = tmpdir / "test_data.txt" + tmp_file.write(data) + fname = str(tmp_file) + res = np.loadtxt(fname, dtype='str', delimiter=" ", skiprows=nskip, max_rows=60000) + assert len(res) == expected_length + # are the right lines read in res? + assert_array_equal(expected, res[:, 0]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_mixins.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_mixins.py new file mode 100644 index 0000000000000000000000000000000000000000..632058763b7d9e826122af6834bb72d4bd970434 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_mixins.py @@ -0,0 +1,216 @@ +import numbers +import operator + +import numpy as np +from numpy.testing import assert_, assert_equal, assert_raises + + +# NOTE: This class should be kept as an exact copy of the example from the +# docstring for NDArrayOperatorsMixin. + +class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin): + def __init__(self, value): + self.value = np.asarray(value) + + # One might also consider adding the built-in list type to this + # list, to support operations like np.add(array_like, list) + _HANDLED_TYPES = (np.ndarray, numbers.Number) + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + out = kwargs.get('out', ()) + for x in inputs + out: + # Only support operations with instances of _HANDLED_TYPES. + # Use ArrayLike instead of type(self) for isinstance to + # allow subclasses that don't override __array_ufunc__ to + # handle ArrayLike objects. + if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)): + return NotImplemented + + # Defer to the implementation of the ufunc on unwrapped values. + inputs = tuple(x.value if isinstance(x, ArrayLike) else x + for x in inputs) + if out: + kwargs['out'] = tuple( + x.value if isinstance(x, ArrayLike) else x + for x in out) + result = getattr(ufunc, method)(*inputs, **kwargs) + + if type(result) is tuple: + # multiple return values + return tuple(type(self)(x) for x in result) + elif method == 'at': + # no return value + return None + else: + # one return value + return type(self)(result) + + def __repr__(self): + return '%s(%r)' % (type(self).__name__, self.value) + + +def wrap_array_like(result): + if type(result) is tuple: + return tuple(ArrayLike(r) for r in result) + else: + return ArrayLike(result) + + +def _assert_equal_type_and_value(result, expected, err_msg=None): + assert_equal(type(result), type(expected), err_msg=err_msg) + if isinstance(result, tuple): + assert_equal(len(result), len(expected), err_msg=err_msg) + for result_item, expected_item in zip(result, expected): + _assert_equal_type_and_value(result_item, expected_item, err_msg) + else: + assert_equal(result.value, expected.value, err_msg=err_msg) + assert_equal(getattr(result.value, 'dtype', None), + getattr(expected.value, 'dtype', None), err_msg=err_msg) + + +_ALL_BINARY_OPERATORS = [ + operator.lt, + operator.le, + operator.eq, + operator.ne, + operator.gt, + operator.ge, + operator.add, + operator.sub, + operator.mul, + operator.truediv, + operator.floordiv, + operator.mod, + divmod, + pow, + operator.lshift, + operator.rshift, + operator.and_, + operator.xor, + operator.or_, +] + + +class TestNDArrayOperatorsMixin: + + def test_array_like_add(self): + + def check(result): + _assert_equal_type_and_value(result, ArrayLike(0)) + + check(ArrayLike(0) + 0) + check(0 + ArrayLike(0)) + + check(ArrayLike(0) + np.array(0)) + check(np.array(0) + ArrayLike(0)) + + check(ArrayLike(np.array(0)) + 0) + check(0 + ArrayLike(np.array(0))) + + check(ArrayLike(np.array(0)) + np.array(0)) + check(np.array(0) + ArrayLike(np.array(0))) + + def test_inplace(self): + array_like = ArrayLike(np.array([0])) + array_like += 1 + _assert_equal_type_and_value(array_like, ArrayLike(np.array([1]))) + + array = np.array([0]) + array += ArrayLike(1) + _assert_equal_type_and_value(array, ArrayLike(np.array([1]))) + + def test_opt_out(self): + + class OptOut: + """Object that opts out of __array_ufunc__.""" + __array_ufunc__ = None + + def __add__(self, other): + return self + + def __radd__(self, other): + return self + + array_like = ArrayLike(1) + opt_out = OptOut() + + # supported operations + assert_(array_like + opt_out is opt_out) + assert_(opt_out + array_like is opt_out) + + # not supported + with assert_raises(TypeError): + # don't use the Python default, array_like = array_like + opt_out + array_like += opt_out + with assert_raises(TypeError): + array_like - opt_out + with assert_raises(TypeError): + opt_out - array_like + + def test_subclass(self): + + class SubArrayLike(ArrayLike): + """Should take precedence over ArrayLike.""" + + x = ArrayLike(0) + y = SubArrayLike(1) + _assert_equal_type_and_value(x + y, y) + _assert_equal_type_and_value(y + x, y) + + def test_object(self): + x = ArrayLike(0) + obj = object() + with assert_raises(TypeError): + x + obj + with assert_raises(TypeError): + obj + x + with assert_raises(TypeError): + x += obj + + def test_unary_methods(self): + array = np.array([-1, 0, 1, 2]) + array_like = ArrayLike(array) + for op in [operator.neg, + operator.pos, + abs, + operator.invert]: + _assert_equal_type_and_value(op(array_like), ArrayLike(op(array))) + + def test_forward_binary_methods(self): + array = np.array([-1, 0, 1, 2]) + array_like = ArrayLike(array) + for op in _ALL_BINARY_OPERATORS: + expected = wrap_array_like(op(array, 1)) + actual = op(array_like, 1) + err_msg = 'failed for operator {}'.format(op) + _assert_equal_type_and_value(expected, actual, err_msg=err_msg) + + def test_reflected_binary_methods(self): + for op in _ALL_BINARY_OPERATORS: + expected = wrap_array_like(op(2, 1)) + actual = op(2, ArrayLike(1)) + err_msg = 'failed for operator {}'.format(op) + _assert_equal_type_and_value(expected, actual, err_msg=err_msg) + + def test_matmul(self): + array = np.array([1, 2], dtype=np.float64) + array_like = ArrayLike(array) + expected = ArrayLike(np.float64(5)) + _assert_equal_type_and_value(expected, np.matmul(array_like, array)) + _assert_equal_type_and_value( + expected, operator.matmul(array_like, array)) + _assert_equal_type_and_value( + expected, operator.matmul(array, array_like)) + + def test_ufunc_at(self): + array = ArrayLike(np.array([1, 2, 3, 4])) + assert_(np.negative.at(array, np.array([0, 1])) is None) + _assert_equal_type_and_value(array, ArrayLike([-1, -2, 3, 4])) + + def test_ufunc_two_outputs(self): + mantissa, exponent = np.frexp(2 ** -3) + expected = (ArrayLike(mantissa), ArrayLike(exponent)) + _assert_equal_type_and_value( + np.frexp(ArrayLike(2 ** -3)), expected) + _assert_equal_type_and_value( + np.frexp(ArrayLike(np.array(2 ** -3))), expected) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_nanfunctions.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_nanfunctions.py new file mode 100644 index 0000000000000000000000000000000000000000..c8fa7df86b24602ae5ae9cf3ad465a17b4ce20ea --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_nanfunctions.py @@ -0,0 +1,1418 @@ +import warnings +import pytest +import inspect +from functools import partial + +import numpy as np +from numpy._core.numeric import normalize_axis_tuple +from numpy.exceptions import AxisError, ComplexWarning +from numpy.lib._nanfunctions_impl import _nan_mask, _replace_nan +from numpy.testing import ( + assert_, assert_equal, assert_almost_equal, assert_raises, + assert_raises_regex, assert_array_equal, suppress_warnings + ) + + +# Test data +_ndat = np.array([[0.6244, np.nan, 0.2692, 0.0116, np.nan, 0.1170], + [0.5351, -0.9403, np.nan, 0.2100, 0.4759, 0.2833], + [np.nan, np.nan, np.nan, 0.1042, np.nan, -0.5954], + [0.1610, np.nan, np.nan, 0.1859, 0.3146, np.nan]]) + + +# Rows of _ndat with nans removed +_rdat = [np.array([0.6244, 0.2692, 0.0116, 0.1170]), + np.array([0.5351, -0.9403, 0.2100, 0.4759, 0.2833]), + np.array([0.1042, -0.5954]), + np.array([0.1610, 0.1859, 0.3146])] + +# Rows of _ndat with nans converted to ones +_ndat_ones = np.array([[0.6244, 1.0, 0.2692, 0.0116, 1.0, 0.1170], + [0.5351, -0.9403, 1.0, 0.2100, 0.4759, 0.2833], + [1.0, 1.0, 1.0, 0.1042, 1.0, -0.5954], + [0.1610, 1.0, 1.0, 0.1859, 0.3146, 1.0]]) + +# Rows of _ndat with nans converted to zeros +_ndat_zeros = np.array([[0.6244, 0.0, 0.2692, 0.0116, 0.0, 0.1170], + [0.5351, -0.9403, 0.0, 0.2100, 0.4759, 0.2833], + [0.0, 0.0, 0.0, 0.1042, 0.0, -0.5954], + [0.1610, 0.0, 0.0, 0.1859, 0.3146, 0.0]]) + + +class TestSignatureMatch: + NANFUNCS = { + np.nanmin: np.amin, + np.nanmax: np.amax, + np.nanargmin: np.argmin, + np.nanargmax: np.argmax, + np.nansum: np.sum, + np.nanprod: np.prod, + np.nancumsum: np.cumsum, + np.nancumprod: np.cumprod, + np.nanmean: np.mean, + np.nanmedian: np.median, + np.nanpercentile: np.percentile, + np.nanquantile: np.quantile, + np.nanvar: np.var, + np.nanstd: np.std, + } + IDS = [k.__name__ for k in NANFUNCS] + + @staticmethod + def get_signature(func, default="..."): + """Construct a signature and replace all default parameter-values.""" + prm_list = [] + signature = inspect.signature(func) + for prm in signature.parameters.values(): + if prm.default is inspect.Parameter.empty: + prm_list.append(prm) + else: + prm_list.append(prm.replace(default=default)) + return inspect.Signature(prm_list) + + @pytest.mark.parametrize("nan_func,func", NANFUNCS.items(), ids=IDS) + def test_signature_match(self, nan_func, func): + # Ignore the default parameter-values as they can sometimes differ + # between the two functions (*e.g.* one has `False` while the other + # has `np._NoValue`) + signature = self.get_signature(func) + nan_signature = self.get_signature(nan_func) + np.testing.assert_equal(signature, nan_signature) + + def test_exhaustiveness(self): + """Validate that all nan functions are actually tested.""" + np.testing.assert_equal( + set(self.IDS), set(np.lib._nanfunctions_impl.__all__) + ) + + +class TestNanFunctions_MinMax: + + nanfuncs = [np.nanmin, np.nanmax] + stdfuncs = [np.min, np.max] + + def test_mutation(self): + # Check that passed array is not modified. + ndat = _ndat.copy() + for f in self.nanfuncs: + f(ndat) + assert_equal(ndat, _ndat) + + def test_keepdims(self): + mat = np.eye(3) + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for axis in [None, 0, 1]: + tgt = rf(mat, axis=axis, keepdims=True) + res = nf(mat, axis=axis, keepdims=True) + assert_(res.ndim == tgt.ndim) + + def test_out(self): + mat = np.eye(3) + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + resout = np.zeros(3) + tgt = rf(mat, axis=1) + res = nf(mat, axis=1, out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + + def test_dtype_from_input(self): + codes = 'efdgFDG' + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for c in codes: + mat = np.eye(3, dtype=c) + tgt = rf(mat, axis=1).dtype.type + res = nf(mat, axis=1).dtype.type + assert_(res is tgt) + # scalar case + tgt = rf(mat, axis=None).dtype.type + res = nf(mat, axis=None).dtype.type + assert_(res is tgt) + + def test_result_values(self): + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + tgt = [rf(d) for d in _rdat] + res = nf(_ndat, axis=1) + assert_almost_equal(res, tgt) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + match = "All-NaN slice encountered" + for func in self.nanfuncs: + with pytest.warns(RuntimeWarning, match=match): + out = func(array, axis=axis) + assert np.isnan(out).all() + assert out.dtype == array.dtype + + def test_masked(self): + mat = np.ma.fix_invalid(_ndat) + msk = mat._mask.copy() + for f in [np.nanmin]: + res = f(mat, axis=1) + tgt = f(_ndat, axis=1) + assert_equal(res, tgt) + assert_equal(mat._mask, msk) + assert_(not np.isinf(mat).any()) + + def test_scalar(self): + for f in self.nanfuncs: + assert_(f(0.) == 0.) + + def test_subclass(self): + class MyNDArray(np.ndarray): + pass + + # Check that it works and that type and + # shape are preserved + mine = np.eye(3).view(MyNDArray) + for f in self.nanfuncs: + res = f(mine, axis=0) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == (3,)) + res = f(mine, axis=1) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == (3,)) + res = f(mine) + assert_(res.shape == ()) + + # check that rows of nan are dealt with for subclasses (#4628) + mine[1] = np.nan + for f in self.nanfuncs: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + res = f(mine, axis=0) + assert_(isinstance(res, MyNDArray)) + assert_(not np.any(np.isnan(res))) + assert_(len(w) == 0) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + res = f(mine, axis=1) + assert_(isinstance(res, MyNDArray)) + assert_(np.isnan(res[1]) and not np.isnan(res[0]) + and not np.isnan(res[2])) + assert_(len(w) == 1, 'no warning raised') + assert_(issubclass(w[0].category, RuntimeWarning)) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + res = f(mine) + assert_(res.shape == ()) + assert_(res != np.nan) + assert_(len(w) == 0) + + def test_object_array(self): + arr = np.array([[1.0, 2.0], [np.nan, 4.0], [np.nan, np.nan]], dtype=object) + assert_equal(np.nanmin(arr), 1.0) + assert_equal(np.nanmin(arr, axis=0), [1.0, 2.0]) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + # assert_equal does not work on object arrays of nan + assert_equal(list(np.nanmin(arr, axis=1)), [1.0, 4.0, np.nan]) + assert_(len(w) == 1, 'no warning raised') + assert_(issubclass(w[0].category, RuntimeWarning)) + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_initial(self, dtype): + class MyNDArray(np.ndarray): + pass + + ar = np.arange(9).astype(dtype) + ar[:5] = np.nan + + for f in self.nanfuncs: + initial = 100 if f is np.nanmax else 0 + + ret1 = f(ar, initial=initial) + assert ret1.dtype == dtype + assert ret1 == initial + + ret2 = f(ar.view(MyNDArray), initial=initial) + assert ret2.dtype == dtype + assert ret2 == initial + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_where(self, dtype): + class MyNDArray(np.ndarray): + pass + + ar = np.arange(9).reshape(3, 3).astype(dtype) + ar[0, :] = np.nan + where = np.ones_like(ar, dtype=np.bool) + where[:, 0] = False + + for f in self.nanfuncs: + reference = 4 if f is np.nanmin else 8 + + ret1 = f(ar, where=where, initial=5) + assert ret1.dtype == dtype + assert ret1 == reference + + ret2 = f(ar.view(MyNDArray), where=where, initial=5) + assert ret2.dtype == dtype + assert ret2 == reference + + +class TestNanFunctions_ArgminArgmax: + + nanfuncs = [np.nanargmin, np.nanargmax] + + def test_mutation(self): + # Check that passed array is not modified. + ndat = _ndat.copy() + for f in self.nanfuncs: + f(ndat) + assert_equal(ndat, _ndat) + + def test_result_values(self): + for f, fcmp in zip(self.nanfuncs, [np.greater, np.less]): + for row in _ndat: + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "invalid value encountered in") + ind = f(row) + val = row[ind] + # comparing with NaN is tricky as the result + # is always false except for NaN != NaN + assert_(not np.isnan(val)) + assert_(not fcmp(val, row).any()) + assert_(not np.equal(val, row[:ind]).any()) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + for func in self.nanfuncs: + with pytest.raises(ValueError, match="All-NaN slice encountered"): + func(array, axis=axis) + + def test_empty(self): + mat = np.zeros((0, 3)) + for f in self.nanfuncs: + for axis in [0, None]: + assert_raises_regex( + ValueError, + "attempt to get argm.. of an empty sequence", + f, mat, axis=axis) + for axis in [1]: + res = f(mat, axis=axis) + assert_equal(res, np.zeros(0)) + + def test_scalar(self): + for f in self.nanfuncs: + assert_(f(0.) == 0.) + + def test_subclass(self): + class MyNDArray(np.ndarray): + pass + + # Check that it works and that type and + # shape are preserved + mine = np.eye(3).view(MyNDArray) + for f in self.nanfuncs: + res = f(mine, axis=0) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == (3,)) + res = f(mine, axis=1) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == (3,)) + res = f(mine) + assert_(res.shape == ()) + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_keepdims(self, dtype): + ar = np.arange(9).astype(dtype) + ar[:5] = np.nan + + for f in self.nanfuncs: + reference = 5 if f is np.nanargmin else 8 + ret = f(ar, keepdims=True) + assert ret.ndim == ar.ndim + assert ret == reference + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_out(self, dtype): + ar = np.arange(9).astype(dtype) + ar[:5] = np.nan + + for f in self.nanfuncs: + out = np.zeros((), dtype=np.intp) + reference = 5 if f is np.nanargmin else 8 + ret = f(ar, out=out) + assert ret is out + assert ret == reference + + + +_TEST_ARRAYS = { + "0d": np.array(5), + "1d": np.array([127, 39, 93, 87, 46]) +} +for _v in _TEST_ARRAYS.values(): + _v.setflags(write=False) + + +@pytest.mark.parametrize( + "dtype", + np.typecodes["AllInteger"] + np.typecodes["AllFloat"] + "O", +) +@pytest.mark.parametrize("mat", _TEST_ARRAYS.values(), ids=_TEST_ARRAYS.keys()) +class TestNanFunctions_NumberTypes: + nanfuncs = { + np.nanmin: np.min, + np.nanmax: np.max, + np.nanargmin: np.argmin, + np.nanargmax: np.argmax, + np.nansum: np.sum, + np.nanprod: np.prod, + np.nancumsum: np.cumsum, + np.nancumprod: np.cumprod, + np.nanmean: np.mean, + np.nanmedian: np.median, + np.nanvar: np.var, + np.nanstd: np.std, + } + nanfunc_ids = [i.__name__ for i in nanfuncs] + + @pytest.mark.parametrize("nanfunc,func", nanfuncs.items(), ids=nanfunc_ids) + @np.errstate(over="ignore") + def test_nanfunc(self, mat, dtype, nanfunc, func): + mat = mat.astype(dtype) + tgt = func(mat) + out = nanfunc(mat) + + assert_almost_equal(out, tgt) + if dtype == "O": + assert type(out) is type(tgt) + else: + assert out.dtype == tgt.dtype + + @pytest.mark.parametrize( + "nanfunc,func", + [(np.nanquantile, np.quantile), (np.nanpercentile, np.percentile)], + ids=["nanquantile", "nanpercentile"], + ) + def test_nanfunc_q(self, mat, dtype, nanfunc, func): + mat = mat.astype(dtype) + if mat.dtype.kind == "c": + assert_raises(TypeError, func, mat, q=1) + assert_raises(TypeError, nanfunc, mat, q=1) + + else: + tgt = func(mat, q=1) + out = nanfunc(mat, q=1) + + assert_almost_equal(out, tgt) + + if dtype == "O": + assert type(out) is type(tgt) + else: + assert out.dtype == tgt.dtype + + @pytest.mark.parametrize( + "nanfunc,func", + [(np.nanvar, np.var), (np.nanstd, np.std)], + ids=["nanvar", "nanstd"], + ) + def test_nanfunc_ddof(self, mat, dtype, nanfunc, func): + mat = mat.astype(dtype) + tgt = func(mat, ddof=0.5) + out = nanfunc(mat, ddof=0.5) + + assert_almost_equal(out, tgt) + if dtype == "O": + assert type(out) is type(tgt) + else: + assert out.dtype == tgt.dtype + + @pytest.mark.parametrize( + "nanfunc", [np.nanvar, np.nanstd] + ) + def test_nanfunc_correction(self, mat, dtype, nanfunc): + mat = mat.astype(dtype) + assert_almost_equal( + nanfunc(mat, correction=0.5), nanfunc(mat, ddof=0.5) + ) + + err_msg = "ddof and correction can't be provided simultaneously." + with assert_raises_regex(ValueError, err_msg): + nanfunc(mat, ddof=0.5, correction=0.5) + + with assert_raises_regex(ValueError, err_msg): + nanfunc(mat, ddof=1, correction=0) + + +class SharedNanFunctionsTestsMixin: + def test_mutation(self): + # Check that passed array is not modified. + ndat = _ndat.copy() + for f in self.nanfuncs: + f(ndat) + assert_equal(ndat, _ndat) + + def test_keepdims(self): + mat = np.eye(3) + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for axis in [None, 0, 1]: + tgt = rf(mat, axis=axis, keepdims=True) + res = nf(mat, axis=axis, keepdims=True) + assert_(res.ndim == tgt.ndim) + + def test_out(self): + mat = np.eye(3) + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + resout = np.zeros(3) + tgt = rf(mat, axis=1) + res = nf(mat, axis=1, out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + + def test_dtype_from_dtype(self): + mat = np.eye(3) + codes = 'efdgFDG' + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for c in codes: + with suppress_warnings() as sup: + if nf in {np.nanstd, np.nanvar} and c in 'FDG': + # Giving the warning is a small bug, see gh-8000 + sup.filter(ComplexWarning) + tgt = rf(mat, dtype=np.dtype(c), axis=1).dtype.type + res = nf(mat, dtype=np.dtype(c), axis=1).dtype.type + assert_(res is tgt) + # scalar case + tgt = rf(mat, dtype=np.dtype(c), axis=None).dtype.type + res = nf(mat, dtype=np.dtype(c), axis=None).dtype.type + assert_(res is tgt) + + def test_dtype_from_char(self): + mat = np.eye(3) + codes = 'efdgFDG' + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for c in codes: + with suppress_warnings() as sup: + if nf in {np.nanstd, np.nanvar} and c in 'FDG': + # Giving the warning is a small bug, see gh-8000 + sup.filter(ComplexWarning) + tgt = rf(mat, dtype=c, axis=1).dtype.type + res = nf(mat, dtype=c, axis=1).dtype.type + assert_(res is tgt) + # scalar case + tgt = rf(mat, dtype=c, axis=None).dtype.type + res = nf(mat, dtype=c, axis=None).dtype.type + assert_(res is tgt) + + def test_dtype_from_input(self): + codes = 'efdgFDG' + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + for c in codes: + mat = np.eye(3, dtype=c) + tgt = rf(mat, axis=1).dtype.type + res = nf(mat, axis=1).dtype.type + assert_(res is tgt, "res %s, tgt %s" % (res, tgt)) + # scalar case + tgt = rf(mat, axis=None).dtype.type + res = nf(mat, axis=None).dtype.type + assert_(res is tgt) + + def test_result_values(self): + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + tgt = [rf(d) for d in _rdat] + res = nf(_ndat, axis=1) + assert_almost_equal(res, tgt) + + def test_scalar(self): + for f in self.nanfuncs: + assert_(f(0.) == 0.) + + def test_subclass(self): + class MyNDArray(np.ndarray): + pass + + # Check that it works and that type and + # shape are preserved + array = np.eye(3) + mine = array.view(MyNDArray) + for f in self.nanfuncs: + expected_shape = f(array, axis=0).shape + res = f(mine, axis=0) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == expected_shape) + expected_shape = f(array, axis=1).shape + res = f(mine, axis=1) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == expected_shape) + expected_shape = f(array).shape + res = f(mine) + assert_(isinstance(res, MyNDArray)) + assert_(res.shape == expected_shape) + + +class TestNanFunctions_SumProd(SharedNanFunctionsTestsMixin): + + nanfuncs = [np.nansum, np.nanprod] + stdfuncs = [np.sum, np.prod] + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + for func, identity in zip(self.nanfuncs, [0, 1]): + out = func(array, axis=axis) + assert np.all(out == identity) + assert out.dtype == array.dtype + + def test_empty(self): + for f, tgt_value in zip([np.nansum, np.nanprod], [0, 1]): + mat = np.zeros((0, 3)) + tgt = [tgt_value]*3 + res = f(mat, axis=0) + assert_equal(res, tgt) + tgt = [] + res = f(mat, axis=1) + assert_equal(res, tgt) + tgt = tgt_value + res = f(mat, axis=None) + assert_equal(res, tgt) + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_initial(self, dtype): + ar = np.arange(9).astype(dtype) + ar[:5] = np.nan + + for f in self.nanfuncs: + reference = 28 if f is np.nansum else 3360 + ret = f(ar, initial=2) + assert ret.dtype == dtype + assert ret == reference + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_where(self, dtype): + ar = np.arange(9).reshape(3, 3).astype(dtype) + ar[0, :] = np.nan + where = np.ones_like(ar, dtype=np.bool) + where[:, 0] = False + + for f in self.nanfuncs: + reference = 26 if f is np.nansum else 2240 + ret = f(ar, where=where, initial=2) + assert ret.dtype == dtype + assert ret == reference + + +class TestNanFunctions_CumSumProd(SharedNanFunctionsTestsMixin): + + nanfuncs = [np.nancumsum, np.nancumprod] + stdfuncs = [np.cumsum, np.cumprod] + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan) + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + for func, identity in zip(self.nanfuncs, [0, 1]): + out = func(array) + assert np.all(out == identity) + assert out.dtype == array.dtype + + def test_empty(self): + for f, tgt_value in zip(self.nanfuncs, [0, 1]): + mat = np.zeros((0, 3)) + tgt = tgt_value*np.ones((0, 3)) + res = f(mat, axis=0) + assert_equal(res, tgt) + tgt = mat + res = f(mat, axis=1) + assert_equal(res, tgt) + tgt = np.zeros(0) + res = f(mat, axis=None) + assert_equal(res, tgt) + + def test_keepdims(self): + for f, g in zip(self.nanfuncs, self.stdfuncs): + mat = np.eye(3) + for axis in [None, 0, 1]: + tgt = f(mat, axis=axis, out=None) + res = g(mat, axis=axis, out=None) + assert_(res.ndim == tgt.ndim) + + for f in self.nanfuncs: + d = np.ones((3, 5, 7, 11)) + # Randomly set some elements to NaN: + rs = np.random.RandomState(0) + d[rs.rand(*d.shape) < 0.5] = np.nan + res = f(d, axis=None) + assert_equal(res.shape, (1155,)) + for axis in np.arange(4): + res = f(d, axis=axis) + assert_equal(res.shape, (3, 5, 7, 11)) + + def test_result_values(self): + for axis in (-2, -1, 0, 1, None): + tgt = np.cumprod(_ndat_ones, axis=axis) + res = np.nancumprod(_ndat, axis=axis) + assert_almost_equal(res, tgt) + tgt = np.cumsum(_ndat_zeros,axis=axis) + res = np.nancumsum(_ndat, axis=axis) + assert_almost_equal(res, tgt) + + def test_out(self): + mat = np.eye(3) + for nf, rf in zip(self.nanfuncs, self.stdfuncs): + resout = np.eye(3) + for axis in (-2, -1, 0, 1): + tgt = rf(mat, axis=axis) + res = nf(mat, axis=axis, out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + + +class TestNanFunctions_MeanVarStd(SharedNanFunctionsTestsMixin): + + nanfuncs = [np.nanmean, np.nanvar, np.nanstd] + stdfuncs = [np.mean, np.var, np.std] + + def test_dtype_error(self): + for f in self.nanfuncs: + for dtype in [np.bool, np.int_, np.object_]: + assert_raises(TypeError, f, _ndat, axis=1, dtype=dtype) + + def test_out_dtype_error(self): + for f in self.nanfuncs: + for dtype in [np.bool, np.int_, np.object_]: + out = np.empty(_ndat.shape[0], dtype=dtype) + assert_raises(TypeError, f, _ndat, axis=1, out=out) + + def test_ddof(self): + nanfuncs = [np.nanvar, np.nanstd] + stdfuncs = [np.var, np.std] + for nf, rf in zip(nanfuncs, stdfuncs): + for ddof in [0, 1]: + tgt = [rf(d, ddof=ddof) for d in _rdat] + res = nf(_ndat, axis=1, ddof=ddof) + assert_almost_equal(res, tgt) + + def test_ddof_too_big(self): + nanfuncs = [np.nanvar, np.nanstd] + stdfuncs = [np.var, np.std] + dsize = [len(d) for d in _rdat] + for nf, rf in zip(nanfuncs, stdfuncs): + for ddof in range(5): + with suppress_warnings() as sup: + sup.record(RuntimeWarning) + sup.filter(ComplexWarning) + tgt = [ddof >= d for d in dsize] + res = nf(_ndat, axis=1, ddof=ddof) + assert_equal(np.isnan(res), tgt) + if any(tgt): + assert_(len(sup.log) == 1) + else: + assert_(len(sup.log) == 0) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + match = "(Degrees of freedom <= 0 for slice.)|(Mean of empty slice)" + for func in self.nanfuncs: + with pytest.warns(RuntimeWarning, match=match): + out = func(array, axis=axis) + assert np.isnan(out).all() + + # `nanvar` and `nanstd` convert complex inputs to their + # corresponding floating dtype + if func is np.nanmean: + assert out.dtype == array.dtype + else: + assert out.dtype == np.abs(array).dtype + + def test_empty(self): + mat = np.zeros((0, 3)) + for f in self.nanfuncs: + for axis in [0, None]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_(np.isnan(f(mat, axis=axis)).all()) + assert_(len(w) == 1) + assert_(issubclass(w[0].category, RuntimeWarning)) + for axis in [1]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_equal(f(mat, axis=axis), np.zeros([])) + assert_(len(w) == 0) + + @pytest.mark.parametrize("dtype", np.typecodes["AllFloat"]) + def test_where(self, dtype): + ar = np.arange(9).reshape(3, 3).astype(dtype) + ar[0, :] = np.nan + where = np.ones_like(ar, dtype=np.bool) + where[:, 0] = False + + for f, f_std in zip(self.nanfuncs, self.stdfuncs): + reference = f_std(ar[where][2:]) + dtype_reference = dtype if f is np.nanmean else ar.real.dtype + + ret = f(ar, where=where) + assert ret.dtype == dtype_reference + np.testing.assert_allclose(ret, reference) + + def test_nanstd_with_mean_keyword(self): + # Setting the seed to make the test reproducible + rng = np.random.RandomState(1234) + A = rng.randn(10, 20, 5) + 0.5 + A[:, 5, :] = np.nan + + mean_out = np.zeros((10, 1, 5)) + std_out = np.zeros((10, 1, 5)) + + mean = np.nanmean(A, + out=mean_out, + axis=1, + keepdims=True) + + # The returned object should be the object specified during calling + assert mean_out is mean + + std = np.nanstd(A, + out=std_out, + axis=1, + keepdims=True, + mean=mean) + + # The returned object should be the object specified during calling + assert std_out is std + + # Shape of returned mean and std should be same + assert std.shape == mean.shape + assert std.shape == (10, 1, 5) + + # Output should be the same as from the individual algorithms + std_old = np.nanstd(A, axis=1, keepdims=True) + + assert std_old.shape == mean.shape + assert_almost_equal(std, std_old) + +_TIME_UNITS = ( + "Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns", "ps", "fs", "as" +) + +# All `inexact` + `timdelta64` type codes +_TYPE_CODES = list(np.typecodes["AllFloat"]) +_TYPE_CODES += [f"m8[{unit}]" for unit in _TIME_UNITS] + + +class TestNanFunctions_Median: + + def test_mutation(self): + # Check that passed array is not modified. + ndat = _ndat.copy() + np.nanmedian(ndat) + assert_equal(ndat, _ndat) + + def test_keepdims(self): + mat = np.eye(3) + for axis in [None, 0, 1]: + tgt = np.median(mat, axis=axis, out=None, overwrite_input=False) + res = np.nanmedian(mat, axis=axis, out=None, overwrite_input=False) + assert_(res.ndim == tgt.ndim) + + d = np.ones((3, 5, 7, 11)) + # Randomly set some elements to NaN: + w = np.random.random((4, 200)) * np.array(d.shape)[:, None] + w = w.astype(np.intp) + d[tuple(w)] = np.nan + with suppress_warnings() as sup: + sup.filter(RuntimeWarning) + res = np.nanmedian(d, axis=None, keepdims=True) + assert_equal(res.shape, (1, 1, 1, 1)) + res = np.nanmedian(d, axis=(0, 1), keepdims=True) + assert_equal(res.shape, (1, 1, 7, 11)) + res = np.nanmedian(d, axis=(0, 3), keepdims=True) + assert_equal(res.shape, (1, 5, 7, 1)) + res = np.nanmedian(d, axis=(1,), keepdims=True) + assert_equal(res.shape, (3, 1, 7, 11)) + res = np.nanmedian(d, axis=(0, 1, 2, 3), keepdims=True) + assert_equal(res.shape, (1, 1, 1, 1)) + res = np.nanmedian(d, axis=(0, 1, 3), keepdims=True) + assert_equal(res.shape, (1, 1, 7, 1)) + + @pytest.mark.parametrize( + argnames='axis', + argvalues=[ + None, + 1, + (1, ), + (0, 1), + (-3, -1), + ] + ) + @pytest.mark.filterwarnings("ignore:All-NaN slice:RuntimeWarning") + def test_keepdims_out(self, axis): + d = np.ones((3, 5, 7, 11)) + # Randomly set some elements to NaN: + w = np.random.random((4, 200)) * np.array(d.shape)[:, None] + w = w.astype(np.intp) + d[tuple(w)] = np.nan + if axis is None: + shape_out = (1,) * d.ndim + else: + axis_norm = normalize_axis_tuple(axis, d.ndim) + shape_out = tuple( + 1 if i in axis_norm else d.shape[i] for i in range(d.ndim)) + out = np.empty(shape_out) + result = np.nanmedian(d, axis=axis, keepdims=True, out=out) + assert result is out + assert_equal(result.shape, shape_out) + + def test_out(self): + mat = np.random.rand(3, 3) + nan_mat = np.insert(mat, [0, 2], np.nan, axis=1) + resout = np.zeros(3) + tgt = np.median(mat, axis=1) + res = np.nanmedian(nan_mat, axis=1, out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + # 0-d output: + resout = np.zeros(()) + tgt = np.median(mat, axis=None) + res = np.nanmedian(nan_mat, axis=None, out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + res = np.nanmedian(nan_mat, axis=(0, 1), out=resout) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + + def test_small_large(self): + # test the small and large code paths, current cutoff 400 elements + for s in [5, 20, 51, 200, 1000]: + d = np.random.randn(4, s) + # Randomly set some elements to NaN: + w = np.random.randint(0, d.size, size=d.size // 5) + d.ravel()[w] = np.nan + d[:,0] = 1. # ensure at least one good value + # use normal median without nans to compare + tgt = [] + for x in d: + nonan = np.compress(~np.isnan(x), x) + tgt.append(np.median(nonan, overwrite_input=True)) + + assert_array_equal(np.nanmedian(d, axis=-1), tgt) + + def test_result_values(self): + tgt = [np.median(d) for d in _rdat] + res = np.nanmedian(_ndat, axis=1) + assert_almost_equal(res, tgt) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", _TYPE_CODES) + def test_allnans(self, dtype, axis): + mat = np.full((3, 3), np.nan).astype(dtype) + with suppress_warnings() as sup: + sup.record(RuntimeWarning) + + output = np.nanmedian(mat, axis=axis) + assert output.dtype == mat.dtype + assert np.isnan(output).all() + + if axis is None: + assert_(len(sup.log) == 1) + else: + assert_(len(sup.log) == 3) + + # Check scalar + scalar = np.array(np.nan).astype(dtype)[()] + output_scalar = np.nanmedian(scalar) + assert output_scalar.dtype == scalar.dtype + assert np.isnan(output_scalar) + + if axis is None: + assert_(len(sup.log) == 2) + else: + assert_(len(sup.log) == 4) + + def test_empty(self): + mat = np.zeros((0, 3)) + for axis in [0, None]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_(np.isnan(np.nanmedian(mat, axis=axis)).all()) + assert_(len(w) == 1) + assert_(issubclass(w[0].category, RuntimeWarning)) + for axis in [1]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_equal(np.nanmedian(mat, axis=axis), np.zeros([])) + assert_(len(w) == 0) + + def test_scalar(self): + assert_(np.nanmedian(0.) == 0.) + + def test_extended_axis_invalid(self): + d = np.ones((3, 5, 7, 11)) + assert_raises(AxisError, np.nanmedian, d, axis=-5) + assert_raises(AxisError, np.nanmedian, d, axis=(0, -5)) + assert_raises(AxisError, np.nanmedian, d, axis=4) + assert_raises(AxisError, np.nanmedian, d, axis=(0, 4)) + assert_raises(ValueError, np.nanmedian, d, axis=(1, 1)) + + def test_float_special(self): + with suppress_warnings() as sup: + sup.filter(RuntimeWarning) + for inf in [np.inf, -np.inf]: + a = np.array([[inf, np.nan], [np.nan, np.nan]]) + assert_equal(np.nanmedian(a, axis=0), [inf, np.nan]) + assert_equal(np.nanmedian(a, axis=1), [inf, np.nan]) + assert_equal(np.nanmedian(a), inf) + + # minimum fill value check + a = np.array([[np.nan, np.nan, inf], + [np.nan, np.nan, inf]]) + assert_equal(np.nanmedian(a), inf) + assert_equal(np.nanmedian(a, axis=0), [np.nan, np.nan, inf]) + assert_equal(np.nanmedian(a, axis=1), inf) + + # no mask path + a = np.array([[inf, inf], [inf, inf]]) + assert_equal(np.nanmedian(a, axis=1), inf) + + a = np.array([[inf, 7, -inf, -9], + [-10, np.nan, np.nan, 5], + [4, np.nan, np.nan, inf]], + dtype=np.float32) + if inf > 0: + assert_equal(np.nanmedian(a, axis=0), [4., 7., -inf, 5.]) + assert_equal(np.nanmedian(a), 4.5) + else: + assert_equal(np.nanmedian(a, axis=0), [-10., 7., -inf, -9.]) + assert_equal(np.nanmedian(a), -2.5) + assert_equal(np.nanmedian(a, axis=-1), [-1., -2.5, inf]) + + for i in range(0, 10): + for j in range(1, 10): + a = np.array([([np.nan] * i) + ([inf] * j)] * 2) + assert_equal(np.nanmedian(a), inf) + assert_equal(np.nanmedian(a, axis=1), inf) + assert_equal(np.nanmedian(a, axis=0), + ([np.nan] * i) + [inf] * j) + + a = np.array([([np.nan] * i) + ([-inf] * j)] * 2) + assert_equal(np.nanmedian(a), -inf) + assert_equal(np.nanmedian(a, axis=1), -inf) + assert_equal(np.nanmedian(a, axis=0), + ([np.nan] * i) + [-inf] * j) + + +class TestNanFunctions_Percentile: + + def test_mutation(self): + # Check that passed array is not modified. + ndat = _ndat.copy() + np.nanpercentile(ndat, 30) + assert_equal(ndat, _ndat) + + def test_keepdims(self): + mat = np.eye(3) + for axis in [None, 0, 1]: + tgt = np.percentile(mat, 70, axis=axis, out=None, + overwrite_input=False) + res = np.nanpercentile(mat, 70, axis=axis, out=None, + overwrite_input=False) + assert_(res.ndim == tgt.ndim) + + d = np.ones((3, 5, 7, 11)) + # Randomly set some elements to NaN: + w = np.random.random((4, 200)) * np.array(d.shape)[:, None] + w = w.astype(np.intp) + d[tuple(w)] = np.nan + with suppress_warnings() as sup: + sup.filter(RuntimeWarning) + res = np.nanpercentile(d, 90, axis=None, keepdims=True) + assert_equal(res.shape, (1, 1, 1, 1)) + res = np.nanpercentile(d, 90, axis=(0, 1), keepdims=True) + assert_equal(res.shape, (1, 1, 7, 11)) + res = np.nanpercentile(d, 90, axis=(0, 3), keepdims=True) + assert_equal(res.shape, (1, 5, 7, 1)) + res = np.nanpercentile(d, 90, axis=(1,), keepdims=True) + assert_equal(res.shape, (3, 1, 7, 11)) + res = np.nanpercentile(d, 90, axis=(0, 1, 2, 3), keepdims=True) + assert_equal(res.shape, (1, 1, 1, 1)) + res = np.nanpercentile(d, 90, axis=(0, 1, 3), keepdims=True) + assert_equal(res.shape, (1, 1, 7, 1)) + + @pytest.mark.parametrize('q', [7, [1, 7]]) + @pytest.mark.parametrize( + argnames='axis', + argvalues=[ + None, + 1, + (1,), + (0, 1), + (-3, -1), + ] + ) + @pytest.mark.filterwarnings("ignore:All-NaN slice:RuntimeWarning") + def test_keepdims_out(self, q, axis): + d = np.ones((3, 5, 7, 11)) + # Randomly set some elements to NaN: + w = np.random.random((4, 200)) * np.array(d.shape)[:, None] + w = w.astype(np.intp) + d[tuple(w)] = np.nan + if axis is None: + shape_out = (1,) * d.ndim + else: + axis_norm = normalize_axis_tuple(axis, d.ndim) + shape_out = tuple( + 1 if i in axis_norm else d.shape[i] for i in range(d.ndim)) + shape_out = np.shape(q) + shape_out + + out = np.empty(shape_out) + result = np.nanpercentile(d, q, axis=axis, keepdims=True, out=out) + assert result is out + assert_equal(result.shape, shape_out) + + @pytest.mark.parametrize("weighted", [False, True]) + def test_out(self, weighted): + mat = np.random.rand(3, 3) + nan_mat = np.insert(mat, [0, 2], np.nan, axis=1) + resout = np.zeros(3) + if weighted: + w_args = {"weights": np.ones_like(mat), "method": "inverted_cdf"} + nan_w_args = { + "weights": np.ones_like(nan_mat), "method": "inverted_cdf" + } + else: + w_args = dict() + nan_w_args = dict() + tgt = np.percentile(mat, 42, axis=1, **w_args) + res = np.nanpercentile(nan_mat, 42, axis=1, out=resout, **nan_w_args) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + # 0-d output: + resout = np.zeros(()) + tgt = np.percentile(mat, 42, axis=None, **w_args) + res = np.nanpercentile( + nan_mat, 42, axis=None, out=resout, **nan_w_args + ) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + res = np.nanpercentile( + nan_mat, 42, axis=(0, 1), out=resout, **nan_w_args + ) + assert_almost_equal(res, resout) + assert_almost_equal(res, tgt) + + def test_complex(self): + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='G') + assert_raises(TypeError, np.nanpercentile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='D') + assert_raises(TypeError, np.nanpercentile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='F') + assert_raises(TypeError, np.nanpercentile, arr_c, 0.5) + + @pytest.mark.parametrize("weighted", [False, True]) + @pytest.mark.parametrize("use_out", [False, True]) + def test_result_values(self, weighted, use_out): + if weighted: + percentile = partial(np.percentile, method="inverted_cdf") + nanpercentile = partial(np.nanpercentile, method="inverted_cdf") + + def gen_weights(d): + return np.ones_like(d) + + else: + percentile = np.percentile + nanpercentile = np.nanpercentile + + def gen_weights(d): + return None + + tgt = [percentile(d, 28, weights=gen_weights(d)) for d in _rdat] + out = np.empty_like(tgt) if use_out else None + res = nanpercentile(_ndat, 28, axis=1, + weights=gen_weights(_ndat), out=out) + assert_almost_equal(res, tgt) + # Transpose the array to fit the output convention of numpy.percentile + tgt = np.transpose([percentile(d, (28, 98), weights=gen_weights(d)) + for d in _rdat]) + out = np.empty_like(tgt) if use_out else None + res = nanpercentile(_ndat, (28, 98), axis=1, + weights=gen_weights(_ndat), out=out) + assert_almost_equal(res, tgt) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["Float"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + with pytest.warns(RuntimeWarning, match="All-NaN slice encountered"): + out = np.nanpercentile(array, 60, axis=axis) + assert np.isnan(out).all() + assert out.dtype == array.dtype + + def test_empty(self): + mat = np.zeros((0, 3)) + for axis in [0, None]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_(np.isnan(np.nanpercentile(mat, 40, axis=axis)).all()) + assert_(len(w) == 1) + assert_(issubclass(w[0].category, RuntimeWarning)) + for axis in [1]: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + assert_equal(np.nanpercentile(mat, 40, axis=axis), np.zeros([])) + assert_(len(w) == 0) + + def test_scalar(self): + assert_equal(np.nanpercentile(0., 100), 0.) + a = np.arange(6) + r = np.nanpercentile(a, 50, axis=0) + assert_equal(r, 2.5) + assert_(np.isscalar(r)) + + def test_extended_axis_invalid(self): + d = np.ones((3, 5, 7, 11)) + assert_raises(AxisError, np.nanpercentile, d, q=5, axis=-5) + assert_raises(AxisError, np.nanpercentile, d, q=5, axis=(0, -5)) + assert_raises(AxisError, np.nanpercentile, d, q=5, axis=4) + assert_raises(AxisError, np.nanpercentile, d, q=5, axis=(0, 4)) + assert_raises(ValueError, np.nanpercentile, d, q=5, axis=(1, 1)) + + def test_multiple_percentiles(self): + perc = [50, 100] + mat = np.ones((4, 3)) + nan_mat = np.nan * mat + # For checking consistency in higher dimensional case + large_mat = np.ones((3, 4, 5)) + large_mat[:, 0:2:4, :] = 0 + large_mat[:, :, 3:] *= 2 + for axis in [None, 0, 1]: + for keepdim in [False, True]: + with suppress_warnings() as sup: + sup.filter(RuntimeWarning, "All-NaN slice encountered") + val = np.percentile(mat, perc, axis=axis, keepdims=keepdim) + nan_val = np.nanpercentile(nan_mat, perc, axis=axis, + keepdims=keepdim) + assert_equal(nan_val.shape, val.shape) + + val = np.percentile(large_mat, perc, axis=axis, + keepdims=keepdim) + nan_val = np.nanpercentile(large_mat, perc, axis=axis, + keepdims=keepdim) + assert_equal(nan_val, val) + + megamat = np.ones((3, 4, 5, 6)) + assert_equal( + np.nanpercentile(megamat, perc, axis=(1, 2)).shape, (2, 3, 6) + ) + + @pytest.mark.parametrize("nan_weight", [0, 1, 2, 3, 1e200]) + def test_nan_value_with_weight(self, nan_weight): + x = [1, np.nan, 2, 3] + result = np.float64(2.0) + q_unweighted = np.nanpercentile(x, 50, method="inverted_cdf") + assert_equal(q_unweighted, result) + + # The weight value at the nan position should not matter. + w = [1.0, nan_weight, 1.0, 1.0] + q_weighted = np.nanpercentile(x, 50, weights=w, method="inverted_cdf") + assert_equal(q_weighted, result) + + @pytest.mark.parametrize("axis", [0, 1, 2]) + def test_nan_value_with_weight_ndim(self, axis): + # Create a multi-dimensional array to test + np.random.seed(1) + x_no_nan = np.random.random(size=(100, 99, 2)) + # Set some places to NaN (not particularly smart) so there is always + # some non-Nan. + x = x_no_nan.copy() + x[np.arange(99), np.arange(99), 0] = np.nan + + p = np.array([[20., 50., 30], [70, 33, 80]]) + + # We just use ones as weights, but replace it with 0 or 1e200 at the + # NaN positions below. + weights = np.ones_like(x) + + # For comparison use weighted normal percentile with nan weights at + # 0 (and no NaNs); not sure this is strictly identical but should be + # sufficiently so (if a percentile lies exactly on a 0 value). + weights[np.isnan(x)] = 0 + p_expected = np.percentile( + x_no_nan, p, axis=axis, weights=weights, method="inverted_cdf") + + p_unweighted = np.nanpercentile( + x, p, axis=axis, method="inverted_cdf") + # The normal and unweighted versions should be identical: + assert_equal(p_unweighted, p_expected) + + weights[np.isnan(x)] = 1e200 # huge value, shouldn't matter + p_weighted = np.nanpercentile( + x, p, axis=axis, weights=weights, method="inverted_cdf") + assert_equal(p_weighted, p_expected) + # Also check with out passed: + out = np.empty_like(p_weighted) + res = np.nanpercentile( + x, p, axis=axis, weights=weights, out=out, method="inverted_cdf") + + assert res is out + assert_equal(out, p_expected) + + +class TestNanFunctions_Quantile: + # most of this is already tested by TestPercentile + + @pytest.mark.parametrize("weighted", [False, True]) + def test_regression(self, weighted): + ar = np.arange(24).reshape(2, 3, 4).astype(float) + ar[0][1] = np.nan + if weighted: + w_args = {"weights": np.ones_like(ar), "method": "inverted_cdf"} + else: + w_args = dict() + + assert_equal(np.nanquantile(ar, q=0.5, **w_args), + np.nanpercentile(ar, q=50, **w_args)) + assert_equal(np.nanquantile(ar, q=0.5, axis=0, **w_args), + np.nanpercentile(ar, q=50, axis=0, **w_args)) + assert_equal(np.nanquantile(ar, q=0.5, axis=1, **w_args), + np.nanpercentile(ar, q=50, axis=1, **w_args)) + assert_equal(np.nanquantile(ar, q=[0.5], axis=1, **w_args), + np.nanpercentile(ar, q=[50], axis=1, **w_args)) + assert_equal(np.nanquantile(ar, q=[0.25, 0.5, 0.75], axis=1, **w_args), + np.nanpercentile(ar, q=[25, 50, 75], axis=1, **w_args)) + + def test_basic(self): + x = np.arange(8) * 0.5 + assert_equal(np.nanquantile(x, 0), 0.) + assert_equal(np.nanquantile(x, 1), 3.5) + assert_equal(np.nanquantile(x, 0.5), 1.75) + + def test_complex(self): + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='G') + assert_raises(TypeError, np.nanquantile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='D') + assert_raises(TypeError, np.nanquantile, arr_c, 0.5) + arr_c = np.array([0.5+3.0j, 2.1+0.5j, 1.6+2.3j], dtype='F') + assert_raises(TypeError, np.nanquantile, arr_c, 0.5) + + def test_no_p_overwrite(self): + # this is worth retesting, because quantile does not make a copy + p0 = np.array([0, 0.75, 0.25, 0.5, 1.0]) + p = p0.copy() + np.nanquantile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, p0) + + p0 = p0.tolist() + p = p.tolist() + np.nanquantile(np.arange(100.), p, method="midpoint") + assert_array_equal(p, p0) + + @pytest.mark.parametrize("axis", [None, 0, 1]) + @pytest.mark.parametrize("dtype", np.typecodes["Float"]) + @pytest.mark.parametrize("array", [ + np.array(np.nan), + np.full((3, 3), np.nan), + ], ids=["0d", "2d"]) + def test_allnans(self, axis, dtype, array): + if axis is not None and array.ndim == 0: + pytest.skip("`axis != None` not supported for 0d arrays") + + array = array.astype(dtype) + with pytest.warns(RuntimeWarning, match="All-NaN slice encountered"): + out = np.nanquantile(array, 1, axis=axis) + assert np.isnan(out).all() + assert out.dtype == array.dtype + +@pytest.mark.parametrize("arr, expected", [ + # array of floats with some nans + (np.array([np.nan, 5.0, np.nan, np.inf]), + np.array([False, True, False, True])), + # int64 array that can't possibly have nans + (np.array([1, 5, 7, 9], dtype=np.int64), + True), + # bool array that can't possibly have nans + (np.array([False, True, False, True]), + True), + # 2-D complex array with nans + (np.array([[np.nan, 5.0], + [np.nan, np.inf]], dtype=np.complex64), + np.array([[False, True], + [False, True]])), + ]) +def test__nan_mask(arr, expected): + for out in [None, np.empty(arr.shape, dtype=np.bool)]: + actual = _nan_mask(arr, out=out) + assert_equal(actual, expected) + # the above won't distinguish between True proper + # and an array of True values; we want True proper + # for types that can't possibly contain NaN + if type(expected) is not np.ndarray: + assert actual is True + + +def test__replace_nan(): + """ Test that _replace_nan returns the original array if there are no + NaNs, not a copy. + """ + for dtype in [np.bool, np.int32, np.int64]: + arr = np.array([0, 1], dtype=dtype) + result, mask = _replace_nan(arr, 0) + assert mask is None + # do not make a copy if there are no nans + assert result is arr + + for dtype in [np.float32, np.float64]: + arr = np.array([0, 1], dtype=dtype) + result, mask = _replace_nan(arr, 2) + assert (mask == False).all() + # mask is not None, so we make a copy + assert result is not arr + assert_equal(result, arr) + + arr_nan = np.array([0, 1, np.nan], dtype=dtype) + result_nan, mask_nan = _replace_nan(arr_nan, 2) + assert_equal(mask_nan, np.array([False, False, True])) + assert result_nan is not arr_nan + assert_equal(result_nan, np.array([0, 1, 2])) + assert np.isnan(arr_nan[-1]) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_packbits.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_packbits.py new file mode 100644 index 0000000000000000000000000000000000000000..a446156327cd4f1fc0c088fbc61c3ca713f379e4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_packbits.py @@ -0,0 +1,376 @@ +import numpy as np +from numpy.testing import assert_array_equal, assert_equal, assert_raises +import pytest +from itertools import chain + +def test_packbits(): + # Copied from the docstring. + a = [[[1, 0, 1], [0, 1, 0]], + [[1, 1, 0], [0, 0, 1]]] + for dt in '?bBhHiIlLqQ': + arr = np.array(a, dtype=dt) + b = np.packbits(arr, axis=-1) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, np.array([[[160], [64]], [[192], [32]]])) + + assert_raises(TypeError, np.packbits, np.array(a, dtype=float)) + + +def test_packbits_empty(): + shapes = [ + (0,), (10, 20, 0), (10, 0, 20), (0, 10, 20), (20, 0, 0), (0, 20, 0), + (0, 0, 20), (0, 0, 0), + ] + for dt in '?bBhHiIlLqQ': + for shape in shapes: + a = np.empty(shape, dtype=dt) + b = np.packbits(a) + assert_equal(b.dtype, np.uint8) + assert_equal(b.shape, (0,)) + + +def test_packbits_empty_with_axis(): + # Original shapes and lists of packed shapes for different axes. + shapes = [ + ((0,), [(0,)]), + ((10, 20, 0), [(2, 20, 0), (10, 3, 0), (10, 20, 0)]), + ((10, 0, 20), [(2, 0, 20), (10, 0, 20), (10, 0, 3)]), + ((0, 10, 20), [(0, 10, 20), (0, 2, 20), (0, 10, 3)]), + ((20, 0, 0), [(3, 0, 0), (20, 0, 0), (20, 0, 0)]), + ((0, 20, 0), [(0, 20, 0), (0, 3, 0), (0, 20, 0)]), + ((0, 0, 20), [(0, 0, 20), (0, 0, 20), (0, 0, 3)]), + ((0, 0, 0), [(0, 0, 0), (0, 0, 0), (0, 0, 0)]), + ] + for dt in '?bBhHiIlLqQ': + for in_shape, out_shapes in shapes: + for ax, out_shape in enumerate(out_shapes): + a = np.empty(in_shape, dtype=dt) + b = np.packbits(a, axis=ax) + assert_equal(b.dtype, np.uint8) + assert_equal(b.shape, out_shape) + +@pytest.mark.parametrize('bitorder', ('little', 'big')) +def test_packbits_large(bitorder): + # test data large enough for 16 byte vectorization + a = np.array([1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, + 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, + 1, 1, 0, 1, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, + 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, + 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, + 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, + 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, + 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 1, + 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, + 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, + 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, + 0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, + 1, 1, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, + 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, + 1, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0]) + a = a.repeat(3) + for dtype in '?bBhHiIlLqQ': + arr = np.array(a, dtype=dtype) + b = np.packbits(arr, axis=None, bitorder=bitorder) + assert_equal(b.dtype, np.uint8) + r = [252, 127, 192, 3, 254, 7, 252, 0, 7, 31, 240, 0, 28, 1, 255, 252, + 113, 248, 3, 255, 192, 28, 15, 192, 28, 126, 0, 224, 127, 255, + 227, 142, 7, 31, 142, 63, 28, 126, 56, 227, 240, 0, 227, 128, 63, + 224, 14, 56, 252, 112, 56, 255, 241, 248, 3, 240, 56, 224, 112, + 63, 255, 255, 199, 224, 14, 0, 31, 143, 192, 3, 255, 199, 0, 1, + 255, 224, 1, 255, 252, 126, 63, 0, 1, 192, 252, 14, 63, 0, 15, + 199, 252, 113, 255, 3, 128, 56, 252, 14, 7, 0, 113, 255, 255, 142, 56, 227, + 129, 248, 227, 129, 199, 31, 128] + if bitorder == 'big': + assert_array_equal(b, r) + # equal for size being multiple of 8 + assert_array_equal(np.unpackbits(b, bitorder=bitorder)[:-4], a) + + # check last byte of different remainders (16 byte vectorization) + b = [np.packbits(arr[:-i], axis=None)[-1] for i in range(1, 16)] + assert_array_equal(b, [128, 128, 128, 31, 30, 28, 24, 16, 0, 0, 0, 199, + 198, 196, 192]) + + + arr = arr.reshape(36, 25) + b = np.packbits(arr, axis=0) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, [[190, 186, 178, 178, 150, 215, 87, 83, 83, 195, + 199, 206, 204, 204, 140, 140, 136, 136, 8, 40, 105, + 107, 75, 74, 88], + [72, 216, 248, 241, 227, 195, 202, 90, 90, 83, + 83, 119, 127, 109, 73, 64, 208, 244, 189, 45, + 41, 104, 122, 90, 18], + [113, 120, 248, 216, 152, 24, 60, 52, 182, 150, + 150, 150, 146, 210, 210, 246, 255, 255, 223, + 151, 21, 17, 17, 131, 163], + [214, 210, 210, 64, 68, 5, 5, 1, 72, 88, 92, + 92, 78, 110, 39, 181, 149, 220, 222, 218, 218, + 202, 234, 170, 168], + [0, 128, 128, 192, 80, 112, 48, 160, 160, 224, + 240, 208, 144, 128, 160, 224, 240, 208, 144, + 144, 176, 240, 224, 192, 128]]) + + b = np.packbits(arr, axis=1) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, [[252, 127, 192, 0], + [ 7, 252, 15, 128], + [240, 0, 28, 0], + [255, 128, 0, 128], + [192, 31, 255, 128], + [142, 63, 0, 0], + [255, 240, 7, 0], + [ 7, 224, 14, 0], + [126, 0, 224, 0], + [255, 255, 199, 0], + [ 56, 28, 126, 0], + [113, 248, 227, 128], + [227, 142, 63, 0], + [ 0, 28, 112, 0], + [ 15, 248, 3, 128], + [ 28, 126, 56, 0], + [ 56, 255, 241, 128], + [240, 7, 224, 0], + [227, 129, 192, 128], + [255, 255, 254, 0], + [126, 0, 224, 0], + [ 3, 241, 248, 0], + [ 0, 255, 241, 128], + [128, 0, 255, 128], + [224, 1, 255, 128], + [248, 252, 126, 0], + [ 0, 7, 3, 128], + [224, 113, 248, 0], + [ 0, 252, 127, 128], + [142, 63, 224, 0], + [224, 14, 63, 0], + [ 7, 3, 128, 0], + [113, 255, 255, 128], + [ 28, 113, 199, 0], + [ 7, 227, 142, 0], + [ 14, 56, 252, 0]]) + + arr = arr.T.copy() + b = np.packbits(arr, axis=0) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, [[252, 7, 240, 255, 192, 142, 255, 7, 126, 255, + 56, 113, 227, 0, 15, 28, 56, 240, 227, 255, + 126, 3, 0, 128, 224, 248, 0, 224, 0, 142, 224, + 7, 113, 28, 7, 14], + [127, 252, 0, 128, 31, 63, 240, 224, 0, 255, + 28, 248, 142, 28, 248, 126, 255, 7, 129, 255, + 0, 241, 255, 0, 1, 252, 7, 113, 252, 63, 14, + 3, 255, 113, 227, 56], + [192, 15, 28, 0, 255, 0, 7, 14, 224, 199, 126, + 227, 63, 112, 3, 56, 241, 224, 192, 254, 224, + 248, 241, 255, 255, 126, 3, 248, 127, 224, 63, + 128, 255, 199, 142, 252], + [0, 128, 0, 128, 128, 0, 0, 0, 0, 0, 0, 128, 0, + 0, 128, 0, 128, 0, 128, 0, 0, 0, 128, 128, + 128, 0, 128, 0, 128, 0, 0, 0, 128, 0, 0, 0]]) + + b = np.packbits(arr, axis=1) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, [[190, 72, 113, 214, 0], + [186, 216, 120, 210, 128], + [178, 248, 248, 210, 128], + [178, 241, 216, 64, 192], + [150, 227, 152, 68, 80], + [215, 195, 24, 5, 112], + [ 87, 202, 60, 5, 48], + [ 83, 90, 52, 1, 160], + [ 83, 90, 182, 72, 160], + [195, 83, 150, 88, 224], + [199, 83, 150, 92, 240], + [206, 119, 150, 92, 208], + [204, 127, 146, 78, 144], + [204, 109, 210, 110, 128], + [140, 73, 210, 39, 160], + [140, 64, 246, 181, 224], + [136, 208, 255, 149, 240], + [136, 244, 255, 220, 208], + [ 8, 189, 223, 222, 144], + [ 40, 45, 151, 218, 144], + [105, 41, 21, 218, 176], + [107, 104, 17, 202, 240], + [ 75, 122, 17, 234, 224], + [ 74, 90, 131, 170, 192], + [ 88, 18, 163, 168, 128]]) + + + # result is the same if input is multiplied with a nonzero value + for dtype in 'bBhHiIlLqQ': + arr = np.array(a, dtype=dtype) + rnd = np.random.randint(low=np.iinfo(dtype).min, + high=np.iinfo(dtype).max, size=arr.size, + dtype=dtype) + rnd[rnd == 0] = 1 + arr *= rnd.astype(dtype) + b = np.packbits(arr, axis=-1) + assert_array_equal(np.unpackbits(b)[:-4], a) + + assert_raises(TypeError, np.packbits, np.array(a, dtype=float)) + + +def test_packbits_very_large(): + # test some with a larger arrays gh-8637 + # code is covered earlier but larger array makes crash on bug more likely + for s in range(950, 1050): + for dt in '?bBhHiIlLqQ': + x = np.ones((200, s), dtype=bool) + np.packbits(x, axis=1) + + +def test_unpackbits(): + # Copied from the docstring. + a = np.array([[2], [7], [23]], dtype=np.uint8) + b = np.unpackbits(a, axis=1) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, np.array([[0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 0, 1, 1, 1]])) + +def test_pack_unpack_order(): + a = np.array([[2], [7], [23]], dtype=np.uint8) + b = np.unpackbits(a, axis=1) + assert_equal(b.dtype, np.uint8) + b_little = np.unpackbits(a, axis=1, bitorder='little') + b_big = np.unpackbits(a, axis=1, bitorder='big') + assert_array_equal(b, b_big) + assert_array_equal(a, np.packbits(b_little, axis=1, bitorder='little')) + assert_array_equal(b[:,::-1], b_little) + assert_array_equal(a, np.packbits(b_big, axis=1, bitorder='big')) + assert_raises(ValueError, np.unpackbits, a, bitorder='r') + assert_raises(TypeError, np.unpackbits, a, bitorder=10) + + + +def test_unpackbits_empty(): + a = np.empty((0,), dtype=np.uint8) + b = np.unpackbits(a) + assert_equal(b.dtype, np.uint8) + assert_array_equal(b, np.empty((0,))) + + +def test_unpackbits_empty_with_axis(): + # Lists of packed shapes for different axes and unpacked shapes. + shapes = [ + ([(0,)], (0,)), + ([(2, 24, 0), (16, 3, 0), (16, 24, 0)], (16, 24, 0)), + ([(2, 0, 24), (16, 0, 24), (16, 0, 3)], (16, 0, 24)), + ([(0, 16, 24), (0, 2, 24), (0, 16, 3)], (0, 16, 24)), + ([(3, 0, 0), (24, 0, 0), (24, 0, 0)], (24, 0, 0)), + ([(0, 24, 0), (0, 3, 0), (0, 24, 0)], (0, 24, 0)), + ([(0, 0, 24), (0, 0, 24), (0, 0, 3)], (0, 0, 24)), + ([(0, 0, 0), (0, 0, 0), (0, 0, 0)], (0, 0, 0)), + ] + for in_shapes, out_shape in shapes: + for ax, in_shape in enumerate(in_shapes): + a = np.empty(in_shape, dtype=np.uint8) + b = np.unpackbits(a, axis=ax) + assert_equal(b.dtype, np.uint8) + assert_equal(b.shape, out_shape) + + +def test_unpackbits_large(): + # test all possible numbers via comparison to already tested packbits + d = np.arange(277, dtype=np.uint8) + assert_array_equal(np.packbits(np.unpackbits(d)), d) + assert_array_equal(np.packbits(np.unpackbits(d[::2])), d[::2]) + d = np.tile(d, (3, 1)) + assert_array_equal(np.packbits(np.unpackbits(d, axis=1), axis=1), d) + d = d.T.copy() + assert_array_equal(np.packbits(np.unpackbits(d, axis=0), axis=0), d) + + +class TestCount: + x = np.array([ + [1, 0, 1, 0, 0, 1, 0], + [0, 1, 1, 1, 0, 0, 0], + [0, 0, 1, 0, 0, 1, 1], + [1, 1, 0, 0, 0, 1, 1], + [1, 0, 1, 0, 1, 0, 1], + [0, 0, 1, 1, 1, 0, 0], + [0, 1, 0, 1, 0, 1, 0], + ], dtype=np.uint8) + padded1 = np.zeros(57, dtype=np.uint8) + padded1[:49] = x.ravel() + padded1b = np.zeros(57, dtype=np.uint8) + padded1b[:49] = x[::-1].copy().ravel() + padded2 = np.zeros((9, 9), dtype=np.uint8) + padded2[:7, :7] = x + + @pytest.mark.parametrize('bitorder', ('little', 'big')) + @pytest.mark.parametrize('count', chain(range(58), range(-1, -57, -1))) + def test_roundtrip(self, bitorder, count): + if count < 0: + # one extra zero of padding + cutoff = count - 1 + else: + cutoff = count + # test complete invertibility of packbits and unpackbits with count + packed = np.packbits(self.x, bitorder=bitorder) + unpacked = np.unpackbits(packed, count=count, bitorder=bitorder) + assert_equal(unpacked.dtype, np.uint8) + assert_array_equal(unpacked, self.padded1[:cutoff]) + + @pytest.mark.parametrize('kwargs', [ + {}, {'count': None}, + ]) + def test_count(self, kwargs): + packed = np.packbits(self.x) + unpacked = np.unpackbits(packed, **kwargs) + assert_equal(unpacked.dtype, np.uint8) + assert_array_equal(unpacked, self.padded1[:-1]) + + @pytest.mark.parametrize('bitorder', ('little', 'big')) + # delta==-1 when count<0 because one extra zero of padding + @pytest.mark.parametrize('count', chain(range(8), range(-1, -9, -1))) + def test_roundtrip_axis(self, bitorder, count): + if count < 0: + # one extra zero of padding + cutoff = count - 1 + else: + cutoff = count + packed0 = np.packbits(self.x, axis=0, bitorder=bitorder) + unpacked0 = np.unpackbits(packed0, axis=0, count=count, + bitorder=bitorder) + assert_equal(unpacked0.dtype, np.uint8) + assert_array_equal(unpacked0, self.padded2[:cutoff, :self.x.shape[1]]) + + packed1 = np.packbits(self.x, axis=1, bitorder=bitorder) + unpacked1 = np.unpackbits(packed1, axis=1, count=count, + bitorder=bitorder) + assert_equal(unpacked1.dtype, np.uint8) + assert_array_equal(unpacked1, self.padded2[:self.x.shape[0], :cutoff]) + + @pytest.mark.parametrize('kwargs', [ + {}, {'count': None}, + {'bitorder' : 'little'}, + {'bitorder': 'little', 'count': None}, + {'bitorder' : 'big'}, + {'bitorder': 'big', 'count': None}, + ]) + def test_axis_count(self, kwargs): + packed0 = np.packbits(self.x, axis=0) + unpacked0 = np.unpackbits(packed0, axis=0, **kwargs) + assert_equal(unpacked0.dtype, np.uint8) + if kwargs.get('bitorder', 'big') == 'big': + assert_array_equal(unpacked0, self.padded2[:-1, :self.x.shape[1]]) + else: + assert_array_equal(unpacked0[::-1, :], self.padded2[:-1, :self.x.shape[1]]) + + packed1 = np.packbits(self.x, axis=1) + unpacked1 = np.unpackbits(packed1, axis=1, **kwargs) + assert_equal(unpacked1.dtype, np.uint8) + if kwargs.get('bitorder', 'big') == 'big': + assert_array_equal(unpacked1, self.padded2[:self.x.shape[0], :-1]) + else: + assert_array_equal(unpacked1[:, ::-1], self.padded2[:self.x.shape[0], :-1]) + + def test_bad_count(self): + packed0 = np.packbits(self.x, axis=0) + assert_raises(ValueError, np.unpackbits, packed0, axis=0, count=-9) + packed1 = np.packbits(self.x, axis=1) + assert_raises(ValueError, np.unpackbits, packed1, axis=1, count=-9) + packed = np.packbits(self.x) + assert_raises(ValueError, np.unpackbits, packed, count=-57) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_polynomial.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_polynomial.py new file mode 100644 index 0000000000000000000000000000000000000000..460de9985fa0c6d803f42661018672708a7e14dc --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_polynomial.py @@ -0,0 +1,303 @@ +import numpy as np +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_almost_equal, + assert_array_almost_equal, assert_raises, assert_allclose + ) + +import pytest + +# `poly1d` has some support for `np.bool` and `np.timedelta64`, +# but it is limited and they are therefore excluded here +TYPE_CODES = np.typecodes["AllInteger"] + np.typecodes["AllFloat"] + "O" + + +class TestPolynomial: + def test_poly1d_str_and_repr(self): + p = np.poly1d([1., 2, 3]) + assert_equal(repr(p), 'poly1d([1., 2., 3.])') + assert_equal(str(p), + ' 2\n' + '1 x + 2 x + 3') + + q = np.poly1d([3., 2, 1]) + assert_equal(repr(q), 'poly1d([3., 2., 1.])') + assert_equal(str(q), + ' 2\n' + '3 x + 2 x + 1') + + r = np.poly1d([1.89999 + 2j, -3j, -5.12345678, 2 + 1j]) + assert_equal(str(r), + ' 3 2\n' + '(1.9 + 2j) x - 3j x - 5.123 x + (2 + 1j)') + + assert_equal(str(np.poly1d([-3, -2, -1])), + ' 2\n' + '-3 x - 2 x - 1') + + def test_poly1d_resolution(self): + p = np.poly1d([1., 2, 3]) + q = np.poly1d([3., 2, 1]) + assert_equal(p(0), 3.0) + assert_equal(p(5), 38.0) + assert_equal(q(0), 1.0) + assert_equal(q(5), 86.0) + + def test_poly1d_math(self): + # here we use some simple coeffs to make calculations easier + p = np.poly1d([1., 2, 4]) + q = np.poly1d([4., 2, 1]) + assert_equal(p/q, (np.poly1d([0.25]), np.poly1d([1.5, 3.75]))) + assert_equal(p.integ(), np.poly1d([1/3, 1., 4., 0.])) + assert_equal(p.integ(1), np.poly1d([1/3, 1., 4., 0.])) + + p = np.poly1d([1., 2, 3]) + q = np.poly1d([3., 2, 1]) + assert_equal(p * q, np.poly1d([3., 8., 14., 8., 3.])) + assert_equal(p + q, np.poly1d([4., 4., 4.])) + assert_equal(p - q, np.poly1d([-2., 0., 2.])) + assert_equal(p ** 4, np.poly1d([1., 8., 36., 104., 214., 312., 324., 216., 81.])) + assert_equal(p(q), np.poly1d([9., 12., 16., 8., 6.])) + assert_equal(q(p), np.poly1d([3., 12., 32., 40., 34.])) + assert_equal(p.deriv(), np.poly1d([2., 2.])) + assert_equal(p.deriv(2), np.poly1d([2.])) + assert_equal(np.polydiv(np.poly1d([1, 0, -1]), np.poly1d([1, 1])), + (np.poly1d([1., -1.]), np.poly1d([0.]))) + + @pytest.mark.parametrize("type_code", TYPE_CODES) + def test_poly1d_misc(self, type_code: str) -> None: + dtype = np.dtype(type_code) + ar = np.array([1, 2, 3], dtype=dtype) + p = np.poly1d(ar) + + # `__eq__` + assert_equal(np.asarray(p), ar) + assert_equal(np.asarray(p).dtype, dtype) + assert_equal(len(p), 2) + + # `__getitem__` + comparison_dct = {-1: 0, 0: 3, 1: 2, 2: 1, 3: 0} + for index, ref in comparison_dct.items(): + scalar = p[index] + assert_equal(scalar, ref) + if dtype == np.object_: + assert isinstance(scalar, int) + else: + assert_equal(scalar.dtype, dtype) + + def test_poly1d_variable_arg(self): + q = np.poly1d([1., 2, 3], variable='y') + assert_equal(str(q), + ' 2\n' + '1 y + 2 y + 3') + q = np.poly1d([1., 2, 3], variable='lambda') + assert_equal(str(q), + ' 2\n' + '1 lambda + 2 lambda + 3') + + def test_poly(self): + assert_array_almost_equal(np.poly([3, -np.sqrt(2), np.sqrt(2)]), + [1, -3, -2, 6]) + + # From matlab docs + A = [[1, 2, 3], [4, 5, 6], [7, 8, 0]] + assert_array_almost_equal(np.poly(A), [1, -6, -72, -27]) + + # Should produce real output for perfect conjugates + assert_(np.isrealobj(np.poly([+1.082j, +2.613j, -2.613j, -1.082j]))) + assert_(np.isrealobj(np.poly([0+1j, -0+-1j, 1+2j, + 1-2j, 1.+3.5j, 1-3.5j]))) + assert_(np.isrealobj(np.poly([1j, -1j, 1+2j, 1-2j, 1+3j, 1-3.j]))) + assert_(np.isrealobj(np.poly([1j, -1j, 1+2j, 1-2j]))) + assert_(np.isrealobj(np.poly([1j, -1j, 2j, -2j]))) + assert_(np.isrealobj(np.poly([1j, -1j]))) + assert_(np.isrealobj(np.poly([1, -1]))) + + assert_(np.iscomplexobj(np.poly([1j, -1.0000001j]))) + + np.random.seed(42) + a = np.random.randn(100) + 1j*np.random.randn(100) + assert_(np.isrealobj(np.poly(np.concatenate((a, np.conjugate(a)))))) + + def test_roots(self): + assert_array_equal(np.roots([1, 0, 0]), [0, 0]) + + def test_str_leading_zeros(self): + p = np.poly1d([4, 3, 2, 1]) + p[3] = 0 + assert_equal(str(p), + " 2\n" + "3 x + 2 x + 1") + + p = np.poly1d([1, 2]) + p[0] = 0 + p[1] = 0 + assert_equal(str(p), " \n0") + + def test_polyfit(self): + c = np.array([3., 2., 1.]) + x = np.linspace(0, 2, 7) + y = np.polyval(c, x) + err = [1, -1, 1, -1, 1, -1, 1] + weights = np.arange(8, 1, -1)**2/7.0 + + # Check exception when too few points for variance estimate. Note that + # the estimate requires the number of data points to exceed + # degree + 1 + assert_raises(ValueError, np.polyfit, + [1], [1], deg=0, cov=True) + + # check 1D case + m, cov = np.polyfit(x, y+err, 2, cov=True) + est = [3.8571, 0.2857, 1.619] + assert_almost_equal(est, m, decimal=4) + val0 = [[ 1.4694, -2.9388, 0.8163], + [-2.9388, 6.3673, -2.1224], + [ 0.8163, -2.1224, 1.161 ]] + assert_almost_equal(val0, cov, decimal=4) + + m2, cov2 = np.polyfit(x, y+err, 2, w=weights, cov=True) + assert_almost_equal([4.8927, -1.0177, 1.7768], m2, decimal=4) + val = [[ 4.3964, -5.0052, 0.4878], + [-5.0052, 6.8067, -0.9089], + [ 0.4878, -0.9089, 0.3337]] + assert_almost_equal(val, cov2, decimal=4) + + m3, cov3 = np.polyfit(x, y+err, 2, w=weights, cov="unscaled") + assert_almost_equal([4.8927, -1.0177, 1.7768], m3, decimal=4) + val = [[ 0.1473, -0.1677, 0.0163], + [-0.1677, 0.228 , -0.0304], + [ 0.0163, -0.0304, 0.0112]] + assert_almost_equal(val, cov3, decimal=4) + + # check 2D (n,1) case + y = y[:, np.newaxis] + c = c[:, np.newaxis] + assert_almost_equal(c, np.polyfit(x, y, 2)) + # check 2D (n,2) case + yy = np.concatenate((y, y), axis=1) + cc = np.concatenate((c, c), axis=1) + assert_almost_equal(cc, np.polyfit(x, yy, 2)) + + m, cov = np.polyfit(x, yy + np.array(err)[:, np.newaxis], 2, cov=True) + assert_almost_equal(est, m[:, 0], decimal=4) + assert_almost_equal(est, m[:, 1], decimal=4) + assert_almost_equal(val0, cov[:, :, 0], decimal=4) + assert_almost_equal(val0, cov[:, :, 1], decimal=4) + + # check order 1 (deg=0) case, were the analytic results are simple + np.random.seed(123) + y = np.random.normal(size=(4, 10000)) + mean, cov = np.polyfit(np.zeros(y.shape[0]), y, deg=0, cov=True) + # Should get sigma_mean = sigma/sqrt(N) = 1./sqrt(4) = 0.5. + assert_allclose(mean.std(), 0.5, atol=0.01) + assert_allclose(np.sqrt(cov.mean()), 0.5, atol=0.01) + # Without scaling, since reduced chi2 is 1, the result should be the same. + mean, cov = np.polyfit(np.zeros(y.shape[0]), y, w=np.ones(y.shape[0]), + deg=0, cov="unscaled") + assert_allclose(mean.std(), 0.5, atol=0.01) + assert_almost_equal(np.sqrt(cov.mean()), 0.5) + # If we estimate our errors wrong, no change with scaling: + w = np.full(y.shape[0], 1./0.5) + mean, cov = np.polyfit(np.zeros(y.shape[0]), y, w=w, deg=0, cov=True) + assert_allclose(mean.std(), 0.5, atol=0.01) + assert_allclose(np.sqrt(cov.mean()), 0.5, atol=0.01) + # But if we do not scale, our estimate for the error in the mean will + # differ. + mean, cov = np.polyfit(np.zeros(y.shape[0]), y, w=w, deg=0, cov="unscaled") + assert_allclose(mean.std(), 0.5, atol=0.01) + assert_almost_equal(np.sqrt(cov.mean()), 0.25) + + def test_objects(self): + from decimal import Decimal + p = np.poly1d([Decimal('4.0'), Decimal('3.0'), Decimal('2.0')]) + p2 = p * Decimal('1.333333333333333') + assert_(p2[1] == Decimal("3.9999999999999990")) + p2 = p.deriv() + assert_(p2[1] == Decimal('8.0')) + p2 = p.integ() + assert_(p2[3] == Decimal("1.333333333333333333333333333")) + assert_(p2[2] == Decimal('1.5')) + assert_(np.issubdtype(p2.coeffs.dtype, np.object_)) + p = np.poly([Decimal(1), Decimal(2)]) + assert_equal(np.poly([Decimal(1), Decimal(2)]), + [1, Decimal(-3), Decimal(2)]) + + def test_complex(self): + p = np.poly1d([3j, 2j, 1j]) + p2 = p.integ() + assert_((p2.coeffs == [1j, 1j, 1j, 0]).all()) + p2 = p.deriv() + assert_((p2.coeffs == [6j, 2j]).all()) + + def test_integ_coeffs(self): + p = np.poly1d([3, 2, 1]) + p2 = p.integ(3, k=[9, 7, 6]) + assert_( + (p2.coeffs == [1/4./5., 1/3./4., 1/2./3., 9/1./2., 7, 6]).all()) + + def test_zero_dims(self): + try: + np.poly(np.zeros((0, 0))) + except ValueError: + pass + + def test_poly_int_overflow(self): + """ + Regression test for gh-5096. + """ + v = np.arange(1, 21) + assert_almost_equal(np.poly(v), np.poly(np.diag(v))) + + def test_zero_poly_dtype(self): + """ + Regression test for gh-16354. + """ + z = np.array([0, 0, 0]) + p = np.poly1d(z.astype(np.int64)) + assert_equal(p.coeffs.dtype, np.int64) + + p = np.poly1d(z.astype(np.float32)) + assert_equal(p.coeffs.dtype, np.float32) + + p = np.poly1d(z.astype(np.complex64)) + assert_equal(p.coeffs.dtype, np.complex64) + + def test_poly_eq(self): + p = np.poly1d([1, 2, 3]) + p2 = np.poly1d([1, 2, 4]) + assert_equal(p == None, False) # noqa: E711 + assert_equal(p != None, True) # noqa: E711 + assert_equal(p == p, True) + assert_equal(p == p2, False) + assert_equal(p != p2, True) + + def test_polydiv(self): + b = np.poly1d([2, 6, 6, 1]) + a = np.poly1d([-1j, (1+2j), -(2+1j), 1]) + q, r = np.polydiv(b, a) + assert_equal(q.coeffs.dtype, np.complex128) + assert_equal(r.coeffs.dtype, np.complex128) + assert_equal(q*a + r, b) + + c = [1, 2, 3] + d = np.poly1d([1, 2, 3]) + s, t = np.polydiv(c, d) + assert isinstance(s, np.poly1d) + assert isinstance(t, np.poly1d) + u, v = np.polydiv(d, c) + assert isinstance(u, np.poly1d) + assert isinstance(v, np.poly1d) + + def test_poly_coeffs_mutable(self): + """ Coefficients should be modifiable """ + p = np.poly1d([1, 2, 3]) + + p.coeffs += 1 + assert_equal(p.coeffs, [2, 3, 4]) + + p.coeffs[2] += 10 + assert_equal(p.coeffs, [2, 3, 14]) + + # this never used to be allowed - let's not add features to deprecated + # APIs + assert_raises(AttributeError, setattr, p, 'coeffs', np.array(1)) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_recfunctions.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_recfunctions.py new file mode 100644 index 0000000000000000000000000000000000000000..37ab6d390ac8d5b9692f419e61dc791743ef2469 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_recfunctions.py @@ -0,0 +1,1042 @@ + +import numpy as np +import numpy.ma as ma +from numpy.ma.mrecords import MaskedRecords +from numpy.ma.testutils import assert_equal +from numpy.testing import assert_, assert_raises +from numpy.lib.recfunctions import ( + drop_fields, rename_fields, get_fieldstructure, recursive_fill_fields, + find_duplicates, merge_arrays, append_fields, stack_arrays, join_by, + repack_fields, unstructured_to_structured, structured_to_unstructured, + apply_along_fields, require_fields, assign_fields_by_name) +get_fieldspec = np.lib.recfunctions._get_fieldspec +get_names = np.lib.recfunctions.get_names +get_names_flat = np.lib.recfunctions.get_names_flat +zip_descr = np.lib.recfunctions._zip_descr +zip_dtype = np.lib.recfunctions._zip_dtype + + +class TestRecFunctions: + # Misc tests + + def setup_method(self): + x = np.array([1, 2, ]) + y = np.array([10, 20, 30]) + z = np.array([('A', 1.), ('B', 2.)], + dtype=[('A', '|S3'), ('B', float)]) + w = np.array([(1, (2, 3.0)), (4, (5, 6.0))], + dtype=[('a', int), ('b', [('ba', float), ('bb', int)])]) + self.data = (w, x, y, z) + + def test_zip_descr(self): + # Test zip_descr + (w, x, y, z) = self.data + + # Std array + test = zip_descr((x, x), flatten=True) + assert_equal(test, + np.dtype([('', int), ('', int)])) + test = zip_descr((x, x), flatten=False) + assert_equal(test, + np.dtype([('', int), ('', int)])) + + # Std & flexible-dtype + test = zip_descr((x, z), flatten=True) + assert_equal(test, + np.dtype([('', int), ('A', '|S3'), ('B', float)])) + test = zip_descr((x, z), flatten=False) + assert_equal(test, + np.dtype([('', int), + ('', [('A', '|S3'), ('B', float)])])) + + # Standard & nested dtype + test = zip_descr((x, w), flatten=True) + assert_equal(test, + np.dtype([('', int), + ('a', int), + ('ba', float), ('bb', int)])) + test = zip_descr((x, w), flatten=False) + assert_equal(test, + np.dtype([('', int), + ('', [('a', int), + ('b', [('ba', float), ('bb', int)])])])) + + def test_drop_fields(self): + # Test drop_fields + a = np.array([(1, (2, 3.0)), (4, (5, 6.0))], + dtype=[('a', int), ('b', [('ba', float), ('bb', int)])]) + + # A basic field + test = drop_fields(a, 'a') + control = np.array([((2, 3.0),), ((5, 6.0),)], + dtype=[('b', [('ba', float), ('bb', int)])]) + assert_equal(test, control) + + # Another basic field (but nesting two fields) + test = drop_fields(a, 'b') + control = np.array([(1,), (4,)], dtype=[('a', int)]) + assert_equal(test, control) + + # A nested sub-field + test = drop_fields(a, ['ba', ]) + control = np.array([(1, (3.0,)), (4, (6.0,))], + dtype=[('a', int), ('b', [('bb', int)])]) + assert_equal(test, control) + + # All the nested sub-field from a field: zap that field + test = drop_fields(a, ['ba', 'bb']) + control = np.array([(1,), (4,)], dtype=[('a', int)]) + assert_equal(test, control) + + # dropping all fields results in an array with no fields + test = drop_fields(a, ['a', 'b']) + control = np.array([(), ()], dtype=[]) + assert_equal(test, control) + + def test_rename_fields(self): + # Test rename fields + a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))], + dtype=[('a', int), + ('b', [('ba', float), ('bb', (float, 2))])]) + test = rename_fields(a, {'a': 'A', 'bb': 'BB'}) + newdtype = [('A', int), ('b', [('ba', float), ('BB', (float, 2))])] + control = a.view(newdtype) + assert_equal(test.dtype, newdtype) + assert_equal(test, control) + + def test_get_names(self): + # Test get_names + ndtype = np.dtype([('A', '|S3'), ('B', float)]) + test = get_names(ndtype) + assert_equal(test, ('A', 'B')) + + ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])]) + test = get_names(ndtype) + assert_equal(test, ('a', ('b', ('ba', 'bb')))) + + ndtype = np.dtype([('a', int), ('b', [])]) + test = get_names(ndtype) + assert_equal(test, ('a', ('b', ()))) + + ndtype = np.dtype([]) + test = get_names(ndtype) + assert_equal(test, ()) + + def test_get_names_flat(self): + # Test get_names_flat + ndtype = np.dtype([('A', '|S3'), ('B', float)]) + test = get_names_flat(ndtype) + assert_equal(test, ('A', 'B')) + + ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])]) + test = get_names_flat(ndtype) + assert_equal(test, ('a', 'b', 'ba', 'bb')) + + ndtype = np.dtype([('a', int), ('b', [])]) + test = get_names_flat(ndtype) + assert_equal(test, ('a', 'b')) + + ndtype = np.dtype([]) + test = get_names_flat(ndtype) + assert_equal(test, ()) + + def test_get_fieldstructure(self): + # Test get_fieldstructure + + # No nested fields + ndtype = np.dtype([('A', '|S3'), ('B', float)]) + test = get_fieldstructure(ndtype) + assert_equal(test, {'A': [], 'B': []}) + + # One 1-nested field + ndtype = np.dtype([('A', int), ('B', [('BA', float), ('BB', '|S1')])]) + test = get_fieldstructure(ndtype) + assert_equal(test, {'A': [], 'B': [], 'BA': ['B', ], 'BB': ['B']}) + + # One 2-nested fields + ndtype = np.dtype([('A', int), + ('B', [('BA', int), + ('BB', [('BBA', int), ('BBB', int)])])]) + test = get_fieldstructure(ndtype) + control = {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'], + 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']} + assert_equal(test, control) + + # 0 fields + ndtype = np.dtype([]) + test = get_fieldstructure(ndtype) + assert_equal(test, {}) + + def test_find_duplicates(self): + # Test find_duplicates + a = ma.array([(2, (2., 'B')), (1, (2., 'B')), (2, (2., 'B')), + (1, (1., 'B')), (2, (2., 'B')), (2, (2., 'C'))], + mask=[(0, (0, 0)), (0, (0, 0)), (0, (0, 0)), + (0, (0, 0)), (1, (0, 0)), (0, (1, 0))], + dtype=[('A', int), ('B', [('BA', float), ('BB', '|S1')])]) + + test = find_duplicates(a, ignoremask=False, return_index=True) + control = [0, 2] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + test = find_duplicates(a, key='A', return_index=True) + control = [0, 1, 2, 3, 5] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + test = find_duplicates(a, key='B', return_index=True) + control = [0, 1, 2, 4] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + test = find_duplicates(a, key='BA', return_index=True) + control = [0, 1, 2, 4] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + test = find_duplicates(a, key='BB', return_index=True) + control = [0, 1, 2, 3, 4] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + def test_find_duplicates_ignoremask(self): + # Test the ignoremask option of find_duplicates + ndtype = [('a', int)] + a = ma.array([1, 1, 1, 2, 2, 3, 3], + mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype) + test = find_duplicates(a, ignoremask=True, return_index=True) + control = [0, 1, 3, 4] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + test = find_duplicates(a, ignoremask=False, return_index=True) + control = [0, 1, 2, 3, 4, 6] + assert_equal(sorted(test[-1]), control) + assert_equal(test[0], a[test[-1]]) + + def test_repack_fields(self): + dt = np.dtype('u1,f4,i8', align=True) + a = np.zeros(2, dtype=dt) + + assert_equal(repack_fields(dt), np.dtype('u1,f4,i8')) + assert_equal(repack_fields(a).itemsize, 13) + assert_equal(repack_fields(repack_fields(dt), align=True), dt) + + # make sure type is preserved + dt = np.dtype((np.record, dt)) + assert_(repack_fields(dt).type is np.record) + + def test_structured_to_unstructured(self, tmp_path): + a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)]) + out = structured_to_unstructured(a) + assert_equal(out, np.zeros((4,5), dtype='f8')) + + b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)], + dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')]) + out = np.mean(structured_to_unstructured(b[['x', 'z']]), axis=-1) + assert_equal(out, np.array([ 3. , 5.5, 9. , 11. ])) + out = np.mean(structured_to_unstructured(b[['x']]), axis=-1) + assert_equal(out, np.array([ 1. , 4. , 7. , 10. ])) + + c = np.arange(20).reshape((4,5)) + out = unstructured_to_structured(c, a.dtype) + want = np.array([( 0, ( 1., 2), [ 3., 4.]), + ( 5, ( 6., 7), [ 8., 9.]), + (10, (11., 12), [13., 14.]), + (15, (16., 17), [18., 19.])], + dtype=[('a', 'i4'), + ('b', [('f0', 'f4'), ('f1', 'u2')]), + ('c', 'f4', (2,))]) + assert_equal(out, want) + + d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)], + dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')]) + assert_equal(apply_along_fields(np.mean, d), + np.array([ 8.0/3, 16.0/3, 26.0/3, 11. ])) + assert_equal(apply_along_fields(np.mean, d[['x', 'z']]), + np.array([ 3. , 5.5, 9. , 11. ])) + + # check that for uniform field dtypes we get a view, not a copy: + d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)], + dtype=[('x', 'i4'), ('y', 'i4'), ('z', 'i4')]) + dd = structured_to_unstructured(d) + ddd = unstructured_to_structured(dd, d.dtype) + assert_(np.shares_memory(dd, d)) + assert_(np.shares_memory(ddd, d)) + + # check that reversing the order of attributes works + dd_attrib_rev = structured_to_unstructured(d[['z', 'x']]) + assert_equal(dd_attrib_rev, [[5, 1], [7, 4], [11, 7], [12, 10]]) + assert_(np.shares_memory(dd_attrib_rev, d)) + + # including uniform fields with subarrays unpacked + d = np.array([(1, [2, 3], [[ 4, 5], [ 6, 7]]), + (8, [9, 10], [[11, 12], [13, 14]])], + dtype=[('x0', 'i4'), ('x1', ('i4', 2)), + ('x2', ('i4', (2, 2)))]) + dd = structured_to_unstructured(d) + ddd = unstructured_to_structured(dd, d.dtype) + assert_(np.shares_memory(dd, d)) + assert_(np.shares_memory(ddd, d)) + + # check that reversing with sub-arrays works as expected + d_rev = d[::-1] + dd_rev = structured_to_unstructured(d_rev) + assert_equal(dd_rev, [[8, 9, 10, 11, 12, 13, 14], + [1, 2, 3, 4, 5, 6, 7]]) + + # check that sub-arrays keep the order of their values + d_attrib_rev = d[['x2', 'x1', 'x0']] + dd_attrib_rev = structured_to_unstructured(d_attrib_rev) + assert_equal(dd_attrib_rev, [[4, 5, 6, 7, 2, 3, 1], + [11, 12, 13, 14, 9, 10, 8]]) + + # with ignored field at the end + d = np.array([(1, [2, 3], [[4, 5], [6, 7]], 32), + (8, [9, 10], [[11, 12], [13, 14]], 64)], + dtype=[('x0', 'i4'), ('x1', ('i4', 2)), + ('x2', ('i4', (2, 2))), ('ignored', 'u1')]) + dd = structured_to_unstructured(d[['x0', 'x1', 'x2']]) + assert_(np.shares_memory(dd, d)) + assert_equal(dd, [[1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14]]) + + # test that nested fields with identical names don't break anything + point = np.dtype([('x', int), ('y', int)]) + triangle = np.dtype([('a', point), ('b', point), ('c', point)]) + arr = np.zeros(10, triangle) + res = structured_to_unstructured(arr, dtype=int) + assert_equal(res, np.zeros((10, 6), dtype=int)) + + + # test nested combinations of subarrays and structured arrays, gh-13333 + def subarray(dt, shape): + return np.dtype((dt, shape)) + + def structured(*dts): + return np.dtype([('x{}'.format(i), dt) for i, dt in enumerate(dts)]) + + def inspect(dt, dtype=None): + arr = np.zeros((), dt) + ret = structured_to_unstructured(arr, dtype=dtype) + backarr = unstructured_to_structured(ret, dt) + return ret.shape, ret.dtype, backarr.dtype + + dt = structured(subarray(structured(np.int32, np.int32), 3)) + assert_equal(inspect(dt), ((6,), np.int32, dt)) + + dt = structured(subarray(subarray(np.int32, 2), 2)) + assert_equal(inspect(dt), ((4,), np.int32, dt)) + + dt = structured(np.int32) + assert_equal(inspect(dt), ((1,), np.int32, dt)) + + dt = structured(np.int32, subarray(subarray(np.int32, 2), 2)) + assert_equal(inspect(dt), ((5,), np.int32, dt)) + + dt = structured() + assert_raises(ValueError, structured_to_unstructured, np.zeros(3, dt)) + + # these currently don't work, but we may make it work in the future + assert_raises(NotImplementedError, structured_to_unstructured, + np.zeros(3, dt), dtype=np.int32) + assert_raises(NotImplementedError, unstructured_to_structured, + np.zeros((3,0), dtype=np.int32)) + + # test supported ndarray subclasses + d_plain = np.array([(1, 2), (3, 4)], dtype=[('a', 'i4'), ('b', 'i4')]) + dd_expected = structured_to_unstructured(d_plain, copy=True) + + # recarray + d = d_plain.view(np.recarray) + + dd = structured_to_unstructured(d, copy=False) + ddd = structured_to_unstructured(d, copy=True) + assert_(np.shares_memory(d, dd)) + assert_(type(dd) is np.recarray) + assert_(type(ddd) is np.recarray) + assert_equal(dd, dd_expected) + assert_equal(ddd, dd_expected) + + # memmap + d = np.memmap(tmp_path / 'memmap', + mode='w+', + dtype=d_plain.dtype, + shape=d_plain.shape) + d[:] = d_plain + dd = structured_to_unstructured(d, copy=False) + ddd = structured_to_unstructured(d, copy=True) + assert_(np.shares_memory(d, dd)) + assert_(type(dd) is np.memmap) + assert_(type(ddd) is np.memmap) + assert_equal(dd, dd_expected) + assert_equal(ddd, dd_expected) + + def test_unstructured_to_structured(self): + # test if dtype is the args of np.dtype + a = np.zeros((20, 2)) + test_dtype_args = [('x', float), ('y', float)] + test_dtype = np.dtype(test_dtype_args) + field1 = unstructured_to_structured(a, dtype=test_dtype_args) # now + field2 = unstructured_to_structured(a, dtype=test_dtype) # before + assert_equal(field1, field2) + + def test_field_assignment_by_name(self): + a = np.ones(2, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')]) + newdt = [('b', 'f4'), ('c', 'u1')] + + assert_equal(require_fields(a, newdt), np.ones(2, newdt)) + + b = np.array([(1,2), (3,4)], dtype=newdt) + assign_fields_by_name(a, b, zero_unassigned=False) + assert_equal(a, np.array([(1,1,2),(1,3,4)], dtype=a.dtype)) + assign_fields_by_name(a, b) + assert_equal(a, np.array([(0,1,2),(0,3,4)], dtype=a.dtype)) + + # test nested fields + a = np.ones(2, dtype=[('a', [('b', 'f8'), ('c', 'u1')])]) + newdt = [('a', [('c', 'u1')])] + assert_equal(require_fields(a, newdt), np.ones(2, newdt)) + b = np.array([((2,),), ((3,),)], dtype=newdt) + assign_fields_by_name(a, b, zero_unassigned=False) + assert_equal(a, np.array([((1,2),), ((1,3),)], dtype=a.dtype)) + assign_fields_by_name(a, b) + assert_equal(a, np.array([((0,2),), ((0,3),)], dtype=a.dtype)) + + # test unstructured code path for 0d arrays + a, b = np.array(3), np.array(0) + assign_fields_by_name(b, a) + assert_equal(b[()], 3) + + +class TestRecursiveFillFields: + # Test recursive_fill_fields. + def test_simple_flexible(self): + # Test recursive_fill_fields on flexible-array + a = np.array([(1, 10.), (2, 20.)], dtype=[('A', int), ('B', float)]) + b = np.zeros((3,), dtype=a.dtype) + test = recursive_fill_fields(a, b) + control = np.array([(1, 10.), (2, 20.), (0, 0.)], + dtype=[('A', int), ('B', float)]) + assert_equal(test, control) + + def test_masked_flexible(self): + # Test recursive_fill_fields on masked flexible-array + a = ma.array([(1, 10.), (2, 20.)], mask=[(0, 1), (1, 0)], + dtype=[('A', int), ('B', float)]) + b = ma.zeros((3,), dtype=a.dtype) + test = recursive_fill_fields(a, b) + control = ma.array([(1, 10.), (2, 20.), (0, 0.)], + mask=[(0, 1), (1, 0), (0, 0)], + dtype=[('A', int), ('B', float)]) + assert_equal(test, control) + + +class TestMergeArrays: + # Test merge_arrays + + def setup_method(self): + x = np.array([1, 2, ]) + y = np.array([10, 20, 30]) + z = np.array( + [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)]) + w = np.array( + [(1, (2, 3.0, ())), (4, (5, 6.0, ()))], + dtype=[('a', int), ('b', [('ba', float), ('bb', int), ('bc', [])])]) + self.data = (w, x, y, z) + + def test_solo(self): + # Test merge_arrays on a single array. + (_, x, _, z) = self.data + + test = merge_arrays(x) + control = np.array([(1,), (2,)], dtype=[('f0', int)]) + assert_equal(test, control) + test = merge_arrays((x,)) + assert_equal(test, control) + + test = merge_arrays(z, flatten=False) + assert_equal(test, z) + test = merge_arrays(z, flatten=True) + assert_equal(test, z) + + def test_solo_w_flatten(self): + # Test merge_arrays on a single array w & w/o flattening + w = self.data[0] + test = merge_arrays(w, flatten=False) + assert_equal(test, w) + + test = merge_arrays(w, flatten=True) + control = np.array([(1, 2, 3.0), (4, 5, 6.0)], + dtype=[('a', int), ('ba', float), ('bb', int)]) + assert_equal(test, control) + + def test_standard(self): + # Test standard & standard + # Test merge arrays + (_, x, y, _) = self.data + test = merge_arrays((x, y), usemask=False) + control = np.array([(1, 10), (2, 20), (-1, 30)], + dtype=[('f0', int), ('f1', int)]) + assert_equal(test, control) + + test = merge_arrays((x, y), usemask=True) + control = ma.array([(1, 10), (2, 20), (-1, 30)], + mask=[(0, 0), (0, 0), (1, 0)], + dtype=[('f0', int), ('f1', int)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + def test_flatten(self): + # Test standard & flexible + (_, x, _, z) = self.data + test = merge_arrays((x, z), flatten=True) + control = np.array([(1, 'A', 1.), (2, 'B', 2.)], + dtype=[('f0', int), ('A', '|S3'), ('B', float)]) + assert_equal(test, control) + + test = merge_arrays((x, z), flatten=False) + control = np.array([(1, ('A', 1.)), (2, ('B', 2.))], + dtype=[('f0', int), + ('f1', [('A', '|S3'), ('B', float)])]) + assert_equal(test, control) + + def test_flatten_wflexible(self): + # Test flatten standard & nested + (w, x, _, _) = self.data + test = merge_arrays((x, w), flatten=True) + control = np.array([(1, 1, 2, 3.0), (2, 4, 5, 6.0)], + dtype=[('f0', int), + ('a', int), ('ba', float), ('bb', int)]) + assert_equal(test, control) + + test = merge_arrays((x, w), flatten=False) + controldtype = [('f0', int), + ('f1', [('a', int), + ('b', [('ba', float), ('bb', int), ('bc', [])])])] + control = np.array([(1., (1, (2, 3.0, ()))), (2, (4, (5, 6.0, ())))], + dtype=controldtype) + assert_equal(test, control) + + def test_wmasked_arrays(self): + # Test merge_arrays masked arrays + (_, x, _, _) = self.data + mx = ma.array([1, 2, 3], mask=[1, 0, 0]) + test = merge_arrays((x, mx), usemask=True) + control = ma.array([(1, 1), (2, 2), (-1, 3)], + mask=[(0, 1), (0, 0), (1, 0)], + dtype=[('f0', int), ('f1', int)]) + assert_equal(test, control) + test = merge_arrays((x, mx), usemask=True, asrecarray=True) + assert_equal(test, control) + assert_(isinstance(test, MaskedRecords)) + + def test_w_singlefield(self): + # Test single field + test = merge_arrays((np.array([1, 2]).view([('a', int)]), + np.array([10., 20., 30.])),) + control = ma.array([(1, 10.), (2, 20.), (-1, 30.)], + mask=[(0, 0), (0, 0), (1, 0)], + dtype=[('a', int), ('f1', float)]) + assert_equal(test, control) + + def test_w_shorter_flex(self): + # Test merge_arrays w/ a shorter flexndarray. + z = self.data[-1] + + # Fixme, this test looks incomplete and broken + #test = merge_arrays((z, np.array([10, 20, 30]).view([('C', int)]))) + #control = np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)], + # dtype=[('A', '|S3'), ('B', float), ('C', int)]) + #assert_equal(test, control) + + # Hack to avoid pyflakes warnings about unused variables + merge_arrays((z, np.array([10, 20, 30]).view([('C', int)]))) + np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)], + dtype=[('A', '|S3'), ('B', float), ('C', int)]) + + def test_singlerecord(self): + (_, x, y, z) = self.data + test = merge_arrays((x[0], y[0], z[0]), usemask=False) + control = np.array([(1, 10, ('A', 1))], + dtype=[('f0', int), + ('f1', int), + ('f2', [('A', '|S3'), ('B', float)])]) + assert_equal(test, control) + + +class TestAppendFields: + # Test append_fields + + def setup_method(self): + x = np.array([1, 2, ]) + y = np.array([10, 20, 30]) + z = np.array( + [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)]) + w = np.array([(1, (2, 3.0)), (4, (5, 6.0))], + dtype=[('a', int), ('b', [('ba', float), ('bb', int)])]) + self.data = (w, x, y, z) + + def test_append_single(self): + # Test simple case + (_, x, _, _) = self.data + test = append_fields(x, 'A', data=[10, 20, 30]) + control = ma.array([(1, 10), (2, 20), (-1, 30)], + mask=[(0, 0), (0, 0), (1, 0)], + dtype=[('f0', int), ('A', int)],) + assert_equal(test, control) + + def test_append_double(self): + # Test simple case + (_, x, _, _) = self.data + test = append_fields(x, ('A', 'B'), data=[[10, 20, 30], [100, 200]]) + control = ma.array([(1, 10, 100), (2, 20, 200), (-1, 30, -1)], + mask=[(0, 0, 0), (0, 0, 0), (1, 0, 1)], + dtype=[('f0', int), ('A', int), ('B', int)],) + assert_equal(test, control) + + def test_append_on_flex(self): + # Test append_fields on flexible type arrays + z = self.data[-1] + test = append_fields(z, 'C', data=[10, 20, 30]) + control = ma.array([('A', 1., 10), ('B', 2., 20), (-1, -1., 30)], + mask=[(0, 0, 0), (0, 0, 0), (1, 1, 0)], + dtype=[('A', '|S3'), ('B', float), ('C', int)],) + assert_equal(test, control) + + def test_append_on_nested(self): + # Test append_fields on nested fields + w = self.data[0] + test = append_fields(w, 'C', data=[10, 20, 30]) + control = ma.array([(1, (2, 3.0), 10), + (4, (5, 6.0), 20), + (-1, (-1, -1.), 30)], + mask=[( + 0, (0, 0), 0), (0, (0, 0), 0), (1, (1, 1), 0)], + dtype=[('a', int), + ('b', [('ba', float), ('bb', int)]), + ('C', int)],) + assert_equal(test, control) + + +class TestStackArrays: + # Test stack_arrays + def setup_method(self): + x = np.array([1, 2, ]) + y = np.array([10, 20, 30]) + z = np.array( + [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)]) + w = np.array([(1, (2, 3.0)), (4, (5, 6.0))], + dtype=[('a', int), ('b', [('ba', float), ('bb', int)])]) + self.data = (w, x, y, z) + + def test_solo(self): + # Test stack_arrays on single arrays + (_, x, _, _) = self.data + test = stack_arrays((x,)) + assert_equal(test, x) + assert_(test is x) + + test = stack_arrays(x) + assert_equal(test, x) + assert_(test is x) + + def test_unnamed_fields(self): + # Tests combinations of arrays w/o named fields + (_, x, y, _) = self.data + + test = stack_arrays((x, x), usemask=False) + control = np.array([1, 2, 1, 2]) + assert_equal(test, control) + + test = stack_arrays((x, y), usemask=False) + control = np.array([1, 2, 10, 20, 30]) + assert_equal(test, control) + + test = stack_arrays((y, x), usemask=False) + control = np.array([10, 20, 30, 1, 2]) + assert_equal(test, control) + + def test_unnamed_and_named_fields(self): + # Test combination of arrays w/ & w/o named fields + (_, x, _, z) = self.data + + test = stack_arrays((x, z)) + control = ma.array([(1, -1, -1), (2, -1, -1), + (-1, 'A', 1), (-1, 'B', 2)], + mask=[(0, 1, 1), (0, 1, 1), + (1, 0, 0), (1, 0, 0)], + dtype=[('f0', int), ('A', '|S3'), ('B', float)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + test = stack_arrays((z, x)) + control = ma.array([('A', 1, -1), ('B', 2, -1), + (-1, -1, 1), (-1, -1, 2), ], + mask=[(0, 0, 1), (0, 0, 1), + (1, 1, 0), (1, 1, 0)], + dtype=[('A', '|S3'), ('B', float), ('f2', int)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + test = stack_arrays((z, z, x)) + control = ma.array([('A', 1, -1), ('B', 2, -1), + ('A', 1, -1), ('B', 2, -1), + (-1, -1, 1), (-1, -1, 2), ], + mask=[(0, 0, 1), (0, 0, 1), + (0, 0, 1), (0, 0, 1), + (1, 1, 0), (1, 1, 0)], + dtype=[('A', '|S3'), ('B', float), ('f2', int)]) + assert_equal(test, control) + + def test_matching_named_fields(self): + # Test combination of arrays w/ matching field names + (_, x, _, z) = self.data + zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)], + dtype=[('A', '|S3'), ('B', float), ('C', float)]) + test = stack_arrays((z, zz)) + control = ma.array([('A', 1, -1), ('B', 2, -1), + ( + 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)], + dtype=[('A', '|S3'), ('B', float), ('C', float)], + mask=[(0, 0, 1), (0, 0, 1), + (0, 0, 0), (0, 0, 0), (0, 0, 0)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + test = stack_arrays((z, zz, x)) + ndtype = [('A', '|S3'), ('B', float), ('C', float), ('f3', int)] + control = ma.array([('A', 1, -1, -1), ('B', 2, -1, -1), + ('a', 10., 100., -1), ('b', 20., 200., -1), + ('c', 30., 300., -1), + (-1, -1, -1, 1), (-1, -1, -1, 2)], + dtype=ndtype, + mask=[(0, 0, 1, 1), (0, 0, 1, 1), + (0, 0, 0, 1), (0, 0, 0, 1), (0, 0, 0, 1), + (1, 1, 1, 0), (1, 1, 1, 0)]) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + def test_defaults(self): + # Test defaults: no exception raised if keys of defaults are not fields. + (_, _, _, z) = self.data + zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)], + dtype=[('A', '|S3'), ('B', float), ('C', float)]) + defaults = {'A': '???', 'B': -999., 'C': -9999., 'D': -99999.} + test = stack_arrays((z, zz), defaults=defaults) + control = ma.array([('A', 1, -9999.), ('B', 2, -9999.), + ( + 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)], + dtype=[('A', '|S3'), ('B', float), ('C', float)], + mask=[(0, 0, 1), (0, 0, 1), + (0, 0, 0), (0, 0, 0), (0, 0, 0)]) + assert_equal(test, control) + assert_equal(test.data, control.data) + assert_equal(test.mask, control.mask) + + def test_autoconversion(self): + # Tests autoconversion + adtype = [('A', int), ('B', bool), ('C', float)] + a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype) + bdtype = [('A', int), ('B', float), ('C', float)] + b = ma.array([(4, 5, 6)], dtype=bdtype) + control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)], + dtype=bdtype) + test = stack_arrays((a, b), autoconvert=True) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + with assert_raises(TypeError): + stack_arrays((a, b), autoconvert=False) + + def test_checktitles(self): + # Test using titles in the field names + adtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)] + a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype) + bdtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)] + b = ma.array([(4, 5, 6)], dtype=bdtype) + test = stack_arrays((a, b)) + control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)], + dtype=bdtype) + assert_equal(test, control) + assert_equal(test.mask, control.mask) + + def test_subdtype(self): + z = np.array([ + ('A', 1), ('B', 2) + ], dtype=[('A', '|S3'), ('B', float, (1,))]) + zz = np.array([ + ('a', [10.], 100.), ('b', [20.], 200.), ('c', [30.], 300.) + ], dtype=[('A', '|S3'), ('B', float, (1,)), ('C', float)]) + + res = stack_arrays((z, zz)) + expected = ma.array( + data=[ + (b'A', [1.0], 0), + (b'B', [2.0], 0), + (b'a', [10.0], 100.0), + (b'b', [20.0], 200.0), + (b'c', [30.0], 300.0)], + mask=[ + (False, [False], True), + (False, [False], True), + (False, [False], False), + (False, [False], False), + (False, [False], False) + ], + dtype=zz.dtype + ) + assert_equal(res.dtype, expected.dtype) + assert_equal(res, expected) + assert_equal(res.mask, expected.mask) + + +class TestJoinBy: + def setup_method(self): + self.a = np.array(list(zip(np.arange(10), np.arange(50, 60), + np.arange(100, 110))), + dtype=[('a', int), ('b', int), ('c', int)]) + self.b = np.array(list(zip(np.arange(5, 15), np.arange(65, 75), + np.arange(100, 110))), + dtype=[('a', int), ('b', int), ('d', int)]) + + def test_inner_join(self): + # Basic test of join_by + a, b = self.a, self.b + + test = join_by('a', a, b, jointype='inner') + control = np.array([(5, 55, 65, 105, 100), (6, 56, 66, 106, 101), + (7, 57, 67, 107, 102), (8, 58, 68, 108, 103), + (9, 59, 69, 109, 104)], + dtype=[('a', int), ('b1', int), ('b2', int), + ('c', int), ('d', int)]) + assert_equal(test, control) + + def test_join(self): + a, b = self.a, self.b + + # Fixme, this test is broken + #test = join_by(('a', 'b'), a, b) + #control = np.array([(5, 55, 105, 100), (6, 56, 106, 101), + # (7, 57, 107, 102), (8, 58, 108, 103), + # (9, 59, 109, 104)], + # dtype=[('a', int), ('b', int), + # ('c', int), ('d', int)]) + #assert_equal(test, control) + + # Hack to avoid pyflakes unused variable warnings + join_by(('a', 'b'), a, b) + np.array([(5, 55, 105, 100), (6, 56, 106, 101), + (7, 57, 107, 102), (8, 58, 108, 103), + (9, 59, 109, 104)], + dtype=[('a', int), ('b', int), + ('c', int), ('d', int)]) + + def test_join_subdtype(self): + # tests the bug in https://stackoverflow.com/q/44769632/102441 + foo = np.array([(1,)], + dtype=[('key', int)]) + bar = np.array([(1, np.array([1,2,3]))], + dtype=[('key', int), ('value', 'uint16', 3)]) + res = join_by('key', foo, bar) + assert_equal(res, bar.view(ma.MaskedArray)) + + def test_outer_join(self): + a, b = self.a, self.b + + test = join_by(('a', 'b'), a, b, 'outer') + control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1), + (2, 52, 102, -1), (3, 53, 103, -1), + (4, 54, 104, -1), (5, 55, 105, -1), + (5, 65, -1, 100), (6, 56, 106, -1), + (6, 66, -1, 101), (7, 57, 107, -1), + (7, 67, -1, 102), (8, 58, 108, -1), + (8, 68, -1, 103), (9, 59, 109, -1), + (9, 69, -1, 104), (10, 70, -1, 105), + (11, 71, -1, 106), (12, 72, -1, 107), + (13, 73, -1, 108), (14, 74, -1, 109)], + mask=[(0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 1, 0), (0, 0, 0, 1), + (0, 0, 1, 0), (0, 0, 0, 1), + (0, 0, 1, 0), (0, 0, 0, 1), + (0, 0, 1, 0), (0, 0, 0, 1), + (0, 0, 1, 0), (0, 0, 1, 0), + (0, 0, 1, 0), (0, 0, 1, 0), + (0, 0, 1, 0), (0, 0, 1, 0)], + dtype=[('a', int), ('b', int), + ('c', int), ('d', int)]) + assert_equal(test, control) + + def test_leftouter_join(self): + a, b = self.a, self.b + + test = join_by(('a', 'b'), a, b, 'leftouter') + control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1), + (2, 52, 102, -1), (3, 53, 103, -1), + (4, 54, 104, -1), (5, 55, 105, -1), + (6, 56, 106, -1), (7, 57, 107, -1), + (8, 58, 108, -1), (9, 59, 109, -1)], + mask=[(0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1), + (0, 0, 0, 1), (0, 0, 0, 1)], + dtype=[('a', int), ('b', int), ('c', int), ('d', int)]) + assert_equal(test, control) + + def test_different_field_order(self): + # gh-8940 + a = np.zeros(3, dtype=[('a', 'i4'), ('b', 'f4'), ('c', 'u1')]) + b = np.ones(3, dtype=[('c', 'u1'), ('b', 'f4'), ('a', 'i4')]) + # this should not give a FutureWarning: + j = join_by(['c', 'b'], a, b, jointype='inner', usemask=False) + assert_equal(j.dtype.names, ['b', 'c', 'a1', 'a2']) + + def test_duplicate_keys(self): + a = np.zeros(3, dtype=[('a', 'i4'), ('b', 'f4'), ('c', 'u1')]) + b = np.ones(3, dtype=[('c', 'u1'), ('b', 'f4'), ('a', 'i4')]) + assert_raises(ValueError, join_by, ['a', 'b', 'b'], a, b) + + def test_same_name_different_dtypes_key(self): + a_dtype = np.dtype([('key', 'S5'), ('value', ' 2**32 + + +def _add_keepdims(func): + """ hack in keepdims behavior into a function taking an axis """ + @functools.wraps(func) + def wrapped(a, axis, **kwargs): + res = func(a, axis=axis, **kwargs) + if axis is None: + axis = 0 # res is now a scalar, so we can insert this anywhere + return np.expand_dims(res, axis=axis) + return wrapped + + +class TestTakeAlongAxis: + def test_argequivalent(self): + """ Test it translates from arg to """ + from numpy.random import rand + a = rand(3, 4, 5) + + funcs = [ + (np.sort, np.argsort, dict()), + (_add_keepdims(np.min), _add_keepdims(np.argmin), dict()), + (_add_keepdims(np.max), _add_keepdims(np.argmax), dict()), + #(np.partition, np.argpartition, dict(kth=2)), + ] + + for func, argfunc, kwargs in funcs: + for axis in list(range(a.ndim)) + [None]: + a_func = func(a, axis=axis, **kwargs) + ai_func = argfunc(a, axis=axis, **kwargs) + assert_equal(a_func, take_along_axis(a, ai_func, axis=axis)) + + def test_invalid(self): + """ Test it errors when indices has too few dimensions """ + a = np.ones((10, 10)) + ai = np.ones((10, 2), dtype=np.intp) + + # sanity check + take_along_axis(a, ai, axis=1) + + # not enough indices + assert_raises(ValueError, take_along_axis, a, np.array(1), axis=1) + # bool arrays not allowed + assert_raises(IndexError, take_along_axis, a, ai.astype(bool), axis=1) + # float arrays not allowed + assert_raises(IndexError, take_along_axis, a, ai.astype(float), axis=1) + # invalid axis + assert_raises(AxisError, take_along_axis, a, ai, axis=10) + # invalid indices + assert_raises(ValueError, take_along_axis, a, ai, axis=None) + + def test_empty(self): + """ Test everything is ok with empty results, even with inserted dims """ + a = np.ones((3, 4, 5)) + ai = np.ones((3, 0, 5), dtype=np.intp) + + actual = take_along_axis(a, ai, axis=1) + assert_equal(actual.shape, ai.shape) + + def test_broadcast(self): + """ Test that non-indexing dimensions are broadcast in both directions """ + a = np.ones((3, 4, 1)) + ai = np.ones((1, 2, 5), dtype=np.intp) + actual = take_along_axis(a, ai, axis=1) + assert_equal(actual.shape, (3, 2, 5)) + + +class TestPutAlongAxis: + def test_replace_max(self): + a_base = np.array([[10, 30, 20], [60, 40, 50]]) + + for axis in list(range(a_base.ndim)) + [None]: + # we mutate this in the loop + a = a_base.copy() + + # replace the max with a small value + i_max = _add_keepdims(np.argmax)(a, axis=axis) + put_along_axis(a, i_max, -99, axis=axis) + + # find the new minimum, which should max + i_min = _add_keepdims(np.argmin)(a, axis=axis) + + assert_equal(i_min, i_max) + + def test_broadcast(self): + """ Test that non-indexing dimensions are broadcast in both directions """ + a = np.ones((3, 4, 1)) + ai = np.arange(10, dtype=np.intp).reshape((1, 2, 5)) % 4 + put_along_axis(a, ai, 20, axis=1) + assert_equal(take_along_axis(a, ai, axis=1), 20) + + def test_invalid(self): + """ Test invalid inputs """ + a_base = np.array([[10, 30, 20], [60, 40, 50]]) + indices = np.array([[0], [1]]) + values = np.array([[2], [1]]) + + # sanity check + a = a_base.copy() + put_along_axis(a, indices, values, axis=0) + assert np.all(a == [[2, 2, 2], [1, 1, 1]]) + + # invalid indices + a = a_base.copy() + with assert_raises(ValueError) as exc: + put_along_axis(a, indices, values, axis=None) + assert "single dimension" in str(exc.exception) + + + +class TestApplyAlongAxis: + def test_simple(self): + a = np.ones((20, 10), 'd') + assert_array_equal( + apply_along_axis(len, 0, a), len(a)*np.ones(a.shape[1])) + + def test_simple101(self): + a = np.ones((10, 101), 'd') + assert_array_equal( + apply_along_axis(len, 0, a), len(a)*np.ones(a.shape[1])) + + def test_3d(self): + a = np.arange(27).reshape((3, 3, 3)) + assert_array_equal(apply_along_axis(np.sum, 0, a), + [[27, 30, 33], [36, 39, 42], [45, 48, 51]]) + + def test_preserve_subclass(self): + def double(row): + return row * 2 + + class MyNDArray(np.ndarray): + pass + + m = np.array([[0, 1], [2, 3]]).view(MyNDArray) + expected = np.array([[0, 2], [4, 6]]).view(MyNDArray) + + result = apply_along_axis(double, 0, m) + assert_(isinstance(result, MyNDArray)) + assert_array_equal(result, expected) + + result = apply_along_axis(double, 1, m) + assert_(isinstance(result, MyNDArray)) + assert_array_equal(result, expected) + + def test_subclass(self): + class MinimalSubclass(np.ndarray): + data = 1 + + def minimal_function(array): + return array.data + + a = np.zeros((6, 3)).view(MinimalSubclass) + + assert_array_equal( + apply_along_axis(minimal_function, 0, a), np.array([1, 1, 1]) + ) + + def test_scalar_array(self, cls=np.ndarray): + a = np.ones((6, 3)).view(cls) + res = apply_along_axis(np.sum, 0, a) + assert_(isinstance(res, cls)) + assert_array_equal(res, np.array([6, 6, 6]).view(cls)) + + def test_0d_array(self, cls=np.ndarray): + def sum_to_0d(x): + """ Sum x, returning a 0d array of the same class """ + assert_equal(x.ndim, 1) + return np.squeeze(np.sum(x, keepdims=True)) + a = np.ones((6, 3)).view(cls) + res = apply_along_axis(sum_to_0d, 0, a) + assert_(isinstance(res, cls)) + assert_array_equal(res, np.array([6, 6, 6]).view(cls)) + + res = apply_along_axis(sum_to_0d, 1, a) + assert_(isinstance(res, cls)) + assert_array_equal(res, np.array([3, 3, 3, 3, 3, 3]).view(cls)) + + def test_axis_insertion(self, cls=np.ndarray): + def f1to2(x): + """produces an asymmetric non-square matrix from x""" + assert_equal(x.ndim, 1) + return (x[::-1] * x[1:,None]).view(cls) + + a2d = np.arange(6*3).reshape((6, 3)) + + # 2d insertion along first axis + actual = apply_along_axis(f1to2, 0, a2d) + expected = np.stack([ + f1to2(a2d[:,i]) for i in range(a2d.shape[1]) + ], axis=-1).view(cls) + assert_equal(type(actual), type(expected)) + assert_equal(actual, expected) + + # 2d insertion along last axis + actual = apply_along_axis(f1to2, 1, a2d) + expected = np.stack([ + f1to2(a2d[i,:]) for i in range(a2d.shape[0]) + ], axis=0).view(cls) + assert_equal(type(actual), type(expected)) + assert_equal(actual, expected) + + # 3d insertion along middle axis + a3d = np.arange(6*5*3).reshape((6, 5, 3)) + + actual = apply_along_axis(f1to2, 1, a3d) + expected = np.stack([ + np.stack([ + f1to2(a3d[i,:,j]) for i in range(a3d.shape[0]) + ], axis=0) + for j in range(a3d.shape[2]) + ], axis=-1).view(cls) + assert_equal(type(actual), type(expected)) + assert_equal(actual, expected) + + def test_subclass_preservation(self): + class MinimalSubclass(np.ndarray): + pass + self.test_scalar_array(MinimalSubclass) + self.test_0d_array(MinimalSubclass) + self.test_axis_insertion(MinimalSubclass) + + def test_axis_insertion_ma(self): + def f1to2(x): + """produces an asymmetric non-square matrix from x""" + assert_equal(x.ndim, 1) + res = x[::-1] * x[1:,None] + return np.ma.masked_where(res%5==0, res) + a = np.arange(6*3).reshape((6, 3)) + res = apply_along_axis(f1to2, 0, a) + assert_(isinstance(res, np.ma.masked_array)) + assert_equal(res.ndim, 3) + assert_array_equal(res[:,:,0].mask, f1to2(a[:,0]).mask) + assert_array_equal(res[:,:,1].mask, f1to2(a[:,1]).mask) + assert_array_equal(res[:,:,2].mask, f1to2(a[:,2]).mask) + + def test_tuple_func1d(self): + def sample_1d(x): + return x[1], x[0] + res = np.apply_along_axis(sample_1d, 1, np.array([[1, 2], [3, 4]])) + assert_array_equal(res, np.array([[2, 1], [4, 3]])) + + def test_empty(self): + # can't apply_along_axis when there's no chance to call the function + def never_call(x): + assert_(False) # should never be reached + + a = np.empty((0, 0)) + assert_raises(ValueError, np.apply_along_axis, never_call, 0, a) + assert_raises(ValueError, np.apply_along_axis, never_call, 1, a) + + # but it's sometimes ok with some non-zero dimensions + def empty_to_1(x): + assert_(len(x) == 0) + return 1 + + a = np.empty((10, 0)) + actual = np.apply_along_axis(empty_to_1, 1, a) + assert_equal(actual, np.ones(10)) + assert_raises(ValueError, np.apply_along_axis, empty_to_1, 0, a) + + def test_with_iterable_object(self): + # from issue 5248 + d = np.array([ + [{1, 11}, {2, 22}, {3, 33}], + [{4, 44}, {5, 55}, {6, 66}] + ]) + actual = np.apply_along_axis(lambda a: set.union(*a), 0, d) + expected = np.array([{1, 11, 4, 44}, {2, 22, 5, 55}, {3, 33, 6, 66}]) + + assert_equal(actual, expected) + + # issue 8642 - assert_equal doesn't detect this! + for i in np.ndindex(actual.shape): + assert_equal(type(actual[i]), type(expected[i])) + + +class TestApplyOverAxes: + def test_simple(self): + a = np.arange(24).reshape(2, 3, 4) + aoa_a = apply_over_axes(np.sum, a, [0, 2]) + assert_array_equal(aoa_a, np.array([[[60], [92], [124]]])) + + +class TestExpandDims: + def test_functionality(self): + s = (2, 3, 4, 5) + a = np.empty(s) + for axis in range(-5, 4): + b = expand_dims(a, axis) + assert_(b.shape[axis] == 1) + assert_(np.squeeze(b).shape == s) + + def test_axis_tuple(self): + a = np.empty((3, 3, 3)) + assert np.expand_dims(a, axis=(0, 1, 2)).shape == (1, 1, 1, 3, 3, 3) + assert np.expand_dims(a, axis=(0, -1, -2)).shape == (1, 3, 3, 3, 1, 1) + assert np.expand_dims(a, axis=(0, 3, 5)).shape == (1, 3, 3, 1, 3, 1) + assert np.expand_dims(a, axis=(0, -3, -5)).shape == (1, 1, 3, 1, 3, 3) + + def test_axis_out_of_range(self): + s = (2, 3, 4, 5) + a = np.empty(s) + assert_raises(AxisError, expand_dims, a, -6) + assert_raises(AxisError, expand_dims, a, 5) + + a = np.empty((3, 3, 3)) + assert_raises(AxisError, expand_dims, a, (0, -6)) + assert_raises(AxisError, expand_dims, a, (0, 5)) + + def test_repeated_axis(self): + a = np.empty((3, 3, 3)) + assert_raises(ValueError, expand_dims, a, axis=(1, 1)) + + def test_subclasses(self): + a = np.arange(10).reshape((2, 5)) + a = np.ma.array(a, mask=a%3 == 0) + + expanded = np.expand_dims(a, axis=1) + assert_(isinstance(expanded, np.ma.MaskedArray)) + assert_equal(expanded.shape, (2, 1, 5)) + assert_equal(expanded.mask.shape, (2, 1, 5)) + + +class TestArraySplit: + def test_integer_0_split(self): + a = np.arange(10) + assert_raises(ValueError, array_split, a, 0) + + def test_integer_split(self): + a = np.arange(10) + res = array_split(a, 1) + desired = [np.arange(10)] + compare_results(res, desired) + + res = array_split(a, 2) + desired = [np.arange(5), np.arange(5, 10)] + compare_results(res, desired) + + res = array_split(a, 3) + desired = [np.arange(4), np.arange(4, 7), np.arange(7, 10)] + compare_results(res, desired) + + res = array_split(a, 4) + desired = [np.arange(3), np.arange(3, 6), np.arange(6, 8), + np.arange(8, 10)] + compare_results(res, desired) + + res = array_split(a, 5) + desired = [np.arange(2), np.arange(2, 4), np.arange(4, 6), + np.arange(6, 8), np.arange(8, 10)] + compare_results(res, desired) + + res = array_split(a, 6) + desired = [np.arange(2), np.arange(2, 4), np.arange(4, 6), + np.arange(6, 8), np.arange(8, 9), np.arange(9, 10)] + compare_results(res, desired) + + res = array_split(a, 7) + desired = [np.arange(2), np.arange(2, 4), np.arange(4, 6), + np.arange(6, 7), np.arange(7, 8), np.arange(8, 9), + np.arange(9, 10)] + compare_results(res, desired) + + res = array_split(a, 8) + desired = [np.arange(2), np.arange(2, 4), np.arange(4, 5), + np.arange(5, 6), np.arange(6, 7), np.arange(7, 8), + np.arange(8, 9), np.arange(9, 10)] + compare_results(res, desired) + + res = array_split(a, 9) + desired = [np.arange(2), np.arange(2, 3), np.arange(3, 4), + np.arange(4, 5), np.arange(5, 6), np.arange(6, 7), + np.arange(7, 8), np.arange(8, 9), np.arange(9, 10)] + compare_results(res, desired) + + res = array_split(a, 10) + desired = [np.arange(1), np.arange(1, 2), np.arange(2, 3), + np.arange(3, 4), np.arange(4, 5), np.arange(5, 6), + np.arange(6, 7), np.arange(7, 8), np.arange(8, 9), + np.arange(9, 10)] + compare_results(res, desired) + + res = array_split(a, 11) + desired = [np.arange(1), np.arange(1, 2), np.arange(2, 3), + np.arange(3, 4), np.arange(4, 5), np.arange(5, 6), + np.arange(6, 7), np.arange(7, 8), np.arange(8, 9), + np.arange(9, 10), np.array([])] + compare_results(res, desired) + + def test_integer_split_2D_rows(self): + a = np.array([np.arange(10), np.arange(10)]) + res = array_split(a, 3, axis=0) + tgt = [np.array([np.arange(10)]), np.array([np.arange(10)]), + np.zeros((0, 10))] + compare_results(res, tgt) + assert_(a.dtype.type is res[-1].dtype.type) + + # Same thing for manual splits: + res = array_split(a, [0, 1], axis=0) + tgt = [np.zeros((0, 10)), np.array([np.arange(10)]), + np.array([np.arange(10)])] + compare_results(res, tgt) + assert_(a.dtype.type is res[-1].dtype.type) + + def test_integer_split_2D_cols(self): + a = np.array([np.arange(10), np.arange(10)]) + res = array_split(a, 3, axis=-1) + desired = [np.array([np.arange(4), np.arange(4)]), + np.array([np.arange(4, 7), np.arange(4, 7)]), + np.array([np.arange(7, 10), np.arange(7, 10)])] + compare_results(res, desired) + + def test_integer_split_2D_default(self): + """ This will fail if we change default axis + """ + a = np.array([np.arange(10), np.arange(10)]) + res = array_split(a, 3) + tgt = [np.array([np.arange(10)]), np.array([np.arange(10)]), + np.zeros((0, 10))] + compare_results(res, tgt) + assert_(a.dtype.type is res[-1].dtype.type) + # perhaps should check higher dimensions + + @pytest.mark.skipif(not IS_64BIT, reason="Needs 64bit platform") + def test_integer_split_2D_rows_greater_max_int32(self): + a = np.broadcast_to([0], (1 << 32, 2)) + res = array_split(a, 4) + chunk = np.broadcast_to([0], (1 << 30, 2)) + tgt = [chunk] * 4 + for i in range(len(tgt)): + assert_equal(res[i].shape, tgt[i].shape) + + def test_index_split_simple(self): + a = np.arange(10) + indices = [1, 5, 7] + res = array_split(a, indices, axis=-1) + desired = [np.arange(0, 1), np.arange(1, 5), np.arange(5, 7), + np.arange(7, 10)] + compare_results(res, desired) + + def test_index_split_low_bound(self): + a = np.arange(10) + indices = [0, 5, 7] + res = array_split(a, indices, axis=-1) + desired = [np.array([]), np.arange(0, 5), np.arange(5, 7), + np.arange(7, 10)] + compare_results(res, desired) + + def test_index_split_high_bound(self): + a = np.arange(10) + indices = [0, 5, 7, 10, 12] + res = array_split(a, indices, axis=-1) + desired = [np.array([]), np.arange(0, 5), np.arange(5, 7), + np.arange(7, 10), np.array([]), np.array([])] + compare_results(res, desired) + + +class TestSplit: + # The split function is essentially the same as array_split, + # except that it test if splitting will result in an + # equal split. Only test for this case. + + def test_equal_split(self): + a = np.arange(10) + res = split(a, 2) + desired = [np.arange(5), np.arange(5, 10)] + compare_results(res, desired) + + def test_unequal_split(self): + a = np.arange(10) + assert_raises(ValueError, split, a, 3) + + +class TestColumnStack: + def test_non_iterable(self): + assert_raises(TypeError, column_stack, 1) + + def test_1D_arrays(self): + # example from docstring + a = np.array((1, 2, 3)) + b = np.array((2, 3, 4)) + expected = np.array([[1, 2], + [2, 3], + [3, 4]]) + actual = np.column_stack((a, b)) + assert_equal(actual, expected) + + def test_2D_arrays(self): + # same as hstack 2D docstring example + a = np.array([[1], [2], [3]]) + b = np.array([[2], [3], [4]]) + expected = np.array([[1, 2], + [2, 3], + [3, 4]]) + actual = np.column_stack((a, b)) + assert_equal(actual, expected) + + def test_generator(self): + with pytest.raises(TypeError, match="arrays to stack must be"): + column_stack(np.arange(3) for _ in range(2)) + + +class TestDstack: + def test_non_iterable(self): + assert_raises(TypeError, dstack, 1) + + def test_0D_array(self): + a = np.array(1) + b = np.array(2) + res = dstack([a, b]) + desired = np.array([[[1, 2]]]) + assert_array_equal(res, desired) + + def test_1D_array(self): + a = np.array([1]) + b = np.array([2]) + res = dstack([a, b]) + desired = np.array([[[1, 2]]]) + assert_array_equal(res, desired) + + def test_2D_array(self): + a = np.array([[1], [2]]) + b = np.array([[1], [2]]) + res = dstack([a, b]) + desired = np.array([[[1, 1]], [[2, 2, ]]]) + assert_array_equal(res, desired) + + def test_2D_array2(self): + a = np.array([1, 2]) + b = np.array([1, 2]) + res = dstack([a, b]) + desired = np.array([[[1, 1], [2, 2]]]) + assert_array_equal(res, desired) + + def test_generator(self): + with pytest.raises(TypeError, match="arrays to stack must be"): + dstack(np.arange(3) for _ in range(2)) + + +# array_split has more comprehensive test of splitting. +# only do simple test on hsplit, vsplit, and dsplit +class TestHsplit: + """Only testing for integer splits. + + """ + def test_non_iterable(self): + assert_raises(ValueError, hsplit, 1, 1) + + def test_0D_array(self): + a = np.array(1) + try: + hsplit(a, 2) + assert_(0) + except ValueError: + pass + + def test_1D_array(self): + a = np.array([1, 2, 3, 4]) + res = hsplit(a, 2) + desired = [np.array([1, 2]), np.array([3, 4])] + compare_results(res, desired) + + def test_2D_array(self): + a = np.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) + res = hsplit(a, 2) + desired = [np.array([[1, 2], [1, 2]]), np.array([[3, 4], [3, 4]])] + compare_results(res, desired) + + +class TestVsplit: + """Only testing for integer splits. + + """ + def test_non_iterable(self): + assert_raises(ValueError, vsplit, 1, 1) + + def test_0D_array(self): + a = np.array(1) + assert_raises(ValueError, vsplit, a, 2) + + def test_1D_array(self): + a = np.array([1, 2, 3, 4]) + try: + vsplit(a, 2) + assert_(0) + except ValueError: + pass + + def test_2D_array(self): + a = np.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) + res = vsplit(a, 2) + desired = [np.array([[1, 2, 3, 4]]), np.array([[1, 2, 3, 4]])] + compare_results(res, desired) + + +class TestDsplit: + # Only testing for integer splits. + def test_non_iterable(self): + assert_raises(ValueError, dsplit, 1, 1) + + def test_0D_array(self): + a = np.array(1) + assert_raises(ValueError, dsplit, a, 2) + + def test_1D_array(self): + a = np.array([1, 2, 3, 4]) + assert_raises(ValueError, dsplit, a, 2) + + def test_2D_array(self): + a = np.array([[1, 2, 3, 4], + [1, 2, 3, 4]]) + try: + dsplit(a, 2) + assert_(0) + except ValueError: + pass + + def test_3D_array(self): + a = np.array([[[1, 2, 3, 4], + [1, 2, 3, 4]], + [[1, 2, 3, 4], + [1, 2, 3, 4]]]) + res = dsplit(a, 2) + desired = [np.array([[[1, 2], [1, 2]], [[1, 2], [1, 2]]]), + np.array([[[3, 4], [3, 4]], [[3, 4], [3, 4]]])] + compare_results(res, desired) + + +class TestSqueeze: + def test_basic(self): + from numpy.random import rand + + a = rand(20, 10, 10, 1, 1) + b = rand(20, 1, 10, 1, 20) + c = rand(1, 1, 20, 10) + assert_array_equal(np.squeeze(a), np.reshape(a, (20, 10, 10))) + assert_array_equal(np.squeeze(b), np.reshape(b, (20, 10, 20))) + assert_array_equal(np.squeeze(c), np.reshape(c, (20, 10))) + + # Squeezing to 0-dim should still give an ndarray + a = [[[1.5]]] + res = np.squeeze(a) + assert_equal(res, 1.5) + assert_equal(res.ndim, 0) + assert_equal(type(res), np.ndarray) + + +class TestKron: + def test_basic(self): + # Using 0-dimensional ndarray + a = np.array(1) + b = np.array([[1, 2], [3, 4]]) + k = np.array([[1, 2], [3, 4]]) + assert_array_equal(np.kron(a, b), k) + a = np.array([[1, 2], [3, 4]]) + b = np.array(1) + assert_array_equal(np.kron(a, b), k) + + # Using 1-dimensional ndarray + a = np.array([3]) + b = np.array([[1, 2], [3, 4]]) + k = np.array([[3, 6], [9, 12]]) + assert_array_equal(np.kron(a, b), k) + a = np.array([[1, 2], [3, 4]]) + b = np.array([3]) + assert_array_equal(np.kron(a, b), k) + + # Using 3-dimensional ndarray + a = np.array([[[1]], [[2]]]) + b = np.array([[1, 2], [3, 4]]) + k = np.array([[[1, 2], [3, 4]], [[2, 4], [6, 8]]]) + assert_array_equal(np.kron(a, b), k) + a = np.array([[1, 2], [3, 4]]) + b = np.array([[[1]], [[2]]]) + k = np.array([[[1, 2], [3, 4]], [[2, 4], [6, 8]]]) + assert_array_equal(np.kron(a, b), k) + + def test_return_type(self): + class myarray(np.ndarray): + __array_priority__ = 1.0 + + a = np.ones([2, 2]) + ma = myarray(a.shape, a.dtype, a.data) + assert_equal(type(kron(a, a)), np.ndarray) + assert_equal(type(kron(ma, ma)), myarray) + assert_equal(type(kron(a, ma)), myarray) + assert_equal(type(kron(ma, a)), myarray) + + @pytest.mark.parametrize( + "array_class", [np.asarray, np.asmatrix] + ) + def test_kron_smoke(self, array_class): + a = array_class(np.ones([3, 3])) + b = array_class(np.ones([3, 3])) + k = array_class(np.ones([9, 9])) + + assert_array_equal(np.kron(a, b), k) + + def test_kron_ma(self): + x = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]]) + k = np.ma.array(np.diag([1, 4, 4, 16]), + mask=~np.array(np.identity(4), dtype=bool)) + + assert_array_equal(k, np.kron(x, x)) + + @pytest.mark.parametrize( + "shape_a,shape_b", [ + ((1, 1), (1, 1)), + ((1, 2, 3), (4, 5, 6)), + ((2, 2), (2, 2, 2)), + ((1, 0), (1, 1)), + ((2, 0, 2), (2, 2)), + ((2, 0, 0, 2), (2, 0, 2)), + ]) + def test_kron_shape(self, shape_a, shape_b): + a = np.ones(shape_a) + b = np.ones(shape_b) + normalised_shape_a = (1,) * max(0, len(shape_b)-len(shape_a)) + shape_a + normalised_shape_b = (1,) * max(0, len(shape_a)-len(shape_b)) + shape_b + expected_shape = np.multiply(normalised_shape_a, normalised_shape_b) + + k = np.kron(a, b) + assert np.array_equal( + k.shape, expected_shape), "Unexpected shape from kron" + + +class TestTile: + def test_basic(self): + a = np.array([0, 1, 2]) + b = [[1, 2], [3, 4]] + assert_equal(tile(a, 2), [0, 1, 2, 0, 1, 2]) + assert_equal(tile(a, (2, 2)), [[0, 1, 2, 0, 1, 2], [0, 1, 2, 0, 1, 2]]) + assert_equal(tile(a, (1, 2)), [[0, 1, 2, 0, 1, 2]]) + assert_equal(tile(b, 2), [[1, 2, 1, 2], [3, 4, 3, 4]]) + assert_equal(tile(b, (2, 1)), [[1, 2], [3, 4], [1, 2], [3, 4]]) + assert_equal(tile(b, (2, 2)), [[1, 2, 1, 2], [3, 4, 3, 4], + [1, 2, 1, 2], [3, 4, 3, 4]]) + + def test_tile_one_repetition_on_array_gh4679(self): + a = np.arange(5) + b = tile(a, 1) + b += 2 + assert_equal(a, np.arange(5)) + + def test_empty(self): + a = np.array([[[]]]) + b = np.array([[], []]) + c = tile(b, 2).shape + d = tile(a, (3, 2, 5)).shape + assert_equal(c, (2, 0)) + assert_equal(d, (3, 2, 0)) + + def test_kroncompare(self): + from numpy.random import randint + + reps = [(2,), (1, 2), (2, 1), (2, 2), (2, 3, 2), (3, 2)] + shape = [(3,), (2, 3), (3, 4, 3), (3, 2, 3), (4, 3, 2, 4), (2, 2)] + for s in shape: + b = randint(0, 10, size=s) + for r in reps: + a = np.ones(r, b.dtype) + large = tile(b, r) + klarge = kron(a, b) + assert_equal(large, klarge) + + +class TestMayShareMemory: + def test_basic(self): + d = np.ones((50, 60)) + d2 = np.ones((30, 60, 6)) + assert_(np.may_share_memory(d, d)) + assert_(np.may_share_memory(d, d[::-1])) + assert_(np.may_share_memory(d, d[::2])) + assert_(np.may_share_memory(d, d[1:, ::-1])) + + assert_(not np.may_share_memory(d[::-1], d2)) + assert_(not np.may_share_memory(d[::2], d2)) + assert_(not np.may_share_memory(d[1:, ::-1], d2)) + assert_(np.may_share_memory(d2[1:, ::-1], d2)) + + +# Utility +def compare_results(res, desired): + """Compare lists of arrays.""" + if len(res) != len(desired): + raise ValueError("Iterables have different lengths") + # See also PEP 618 for Python 3.10 + for x, y in zip(res, desired): + assert_array_equal(x, y) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_stride_tricks.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_stride_tricks.py new file mode 100644 index 0000000000000000000000000000000000000000..3cbebbdd552eb92b3072b4e7286f29a09328358a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_stride_tricks.py @@ -0,0 +1,647 @@ +import numpy as np +from numpy._core._rational_tests import rational +from numpy.testing import ( + assert_equal, assert_array_equal, assert_raises, assert_, + assert_raises_regex, assert_warns, + ) +from numpy.lib._stride_tricks_impl import ( + as_strided, broadcast_arrays, _broadcast_shape, broadcast_to, + broadcast_shapes, sliding_window_view, + ) +import pytest + + +def assert_shapes_correct(input_shapes, expected_shape): + # Broadcast a list of arrays with the given input shapes and check the + # common output shape. + + inarrays = [np.zeros(s) for s in input_shapes] + outarrays = broadcast_arrays(*inarrays) + outshapes = [a.shape for a in outarrays] + expected = [expected_shape] * len(inarrays) + assert_equal(outshapes, expected) + + +def assert_incompatible_shapes_raise(input_shapes): + # Broadcast a list of arrays with the given (incompatible) input shapes + # and check that they raise a ValueError. + + inarrays = [np.zeros(s) for s in input_shapes] + assert_raises(ValueError, broadcast_arrays, *inarrays) + + +def assert_same_as_ufunc(shape0, shape1, transposed=False, flipped=False): + # Broadcast two shapes against each other and check that the data layout + # is the same as if a ufunc did the broadcasting. + + x0 = np.zeros(shape0, dtype=int) + # Note that multiply.reduce's identity element is 1.0, so when shape1==(), + # this gives the desired n==1. + n = int(np.multiply.reduce(shape1)) + x1 = np.arange(n).reshape(shape1) + if transposed: + x0 = x0.T + x1 = x1.T + if flipped: + x0 = x0[::-1] + x1 = x1[::-1] + # Use the add ufunc to do the broadcasting. Since we're adding 0s to x1, the + # result should be exactly the same as the broadcasted view of x1. + y = x0 + x1 + b0, b1 = broadcast_arrays(x0, x1) + assert_array_equal(y, b1) + + +def test_same(): + x = np.arange(10) + y = np.arange(10) + bx, by = broadcast_arrays(x, y) + assert_array_equal(x, bx) + assert_array_equal(y, by) + +def test_broadcast_kwargs(): + # ensure that a TypeError is appropriately raised when + # np.broadcast_arrays() is called with any keyword + # argument other than 'subok' + x = np.arange(10) + y = np.arange(10) + + with assert_raises_regex(TypeError, 'got an unexpected keyword'): + broadcast_arrays(x, y, dtype='float64') + + +def test_one_off(): + x = np.array([[1, 2, 3]]) + y = np.array([[1], [2], [3]]) + bx, by = broadcast_arrays(x, y) + bx0 = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]]) + by0 = bx0.T + assert_array_equal(bx0, bx) + assert_array_equal(by0, by) + + +def test_same_input_shapes(): + # Check that the final shape is just the input shape. + + data = [ + (), + (1,), + (3,), + (0, 1), + (0, 3), + (1, 0), + (3, 0), + (1, 3), + (3, 1), + (3, 3), + ] + for shape in data: + input_shapes = [shape] + # Single input. + assert_shapes_correct(input_shapes, shape) + # Double input. + input_shapes2 = [shape, shape] + assert_shapes_correct(input_shapes2, shape) + # Triple input. + input_shapes3 = [shape, shape, shape] + assert_shapes_correct(input_shapes3, shape) + + +def test_two_compatible_by_ones_input_shapes(): + # Check that two different input shapes of the same length, but some have + # ones, broadcast to the correct shape. + + data = [ + [[(1,), (3,)], (3,)], + [[(1, 3), (3, 3)], (3, 3)], + [[(3, 1), (3, 3)], (3, 3)], + [[(1, 3), (3, 1)], (3, 3)], + [[(1, 1), (3, 3)], (3, 3)], + [[(1, 1), (1, 3)], (1, 3)], + [[(1, 1), (3, 1)], (3, 1)], + [[(1, 0), (0, 0)], (0, 0)], + [[(0, 1), (0, 0)], (0, 0)], + [[(1, 0), (0, 1)], (0, 0)], + [[(1, 1), (0, 0)], (0, 0)], + [[(1, 1), (1, 0)], (1, 0)], + [[(1, 1), (0, 1)], (0, 1)], + ] + for input_shapes, expected_shape in data: + assert_shapes_correct(input_shapes, expected_shape) + # Reverse the input shapes since broadcasting should be symmetric. + assert_shapes_correct(input_shapes[::-1], expected_shape) + + +def test_two_compatible_by_prepending_ones_input_shapes(): + # Check that two different input shapes (of different lengths) broadcast + # to the correct shape. + + data = [ + [[(), (3,)], (3,)], + [[(3,), (3, 3)], (3, 3)], + [[(3,), (3, 1)], (3, 3)], + [[(1,), (3, 3)], (3, 3)], + [[(), (3, 3)], (3, 3)], + [[(1, 1), (3,)], (1, 3)], + [[(1,), (3, 1)], (3, 1)], + [[(1,), (1, 3)], (1, 3)], + [[(), (1, 3)], (1, 3)], + [[(), (3, 1)], (3, 1)], + [[(), (0,)], (0,)], + [[(0,), (0, 0)], (0, 0)], + [[(0,), (0, 1)], (0, 0)], + [[(1,), (0, 0)], (0, 0)], + [[(), (0, 0)], (0, 0)], + [[(1, 1), (0,)], (1, 0)], + [[(1,), (0, 1)], (0, 1)], + [[(1,), (1, 0)], (1, 0)], + [[(), (1, 0)], (1, 0)], + [[(), (0, 1)], (0, 1)], + ] + for input_shapes, expected_shape in data: + assert_shapes_correct(input_shapes, expected_shape) + # Reverse the input shapes since broadcasting should be symmetric. + assert_shapes_correct(input_shapes[::-1], expected_shape) + + +def test_incompatible_shapes_raise_valueerror(): + # Check that a ValueError is raised for incompatible shapes. + + data = [ + [(3,), (4,)], + [(2, 3), (2,)], + [(3,), (3,), (4,)], + [(1, 3, 4), (2, 3, 3)], + ] + for input_shapes in data: + assert_incompatible_shapes_raise(input_shapes) + # Reverse the input shapes since broadcasting should be symmetric. + assert_incompatible_shapes_raise(input_shapes[::-1]) + + +def test_same_as_ufunc(): + # Check that the data layout is the same as if a ufunc did the operation. + + data = [ + [[(1,), (3,)], (3,)], + [[(1, 3), (3, 3)], (3, 3)], + [[(3, 1), (3, 3)], (3, 3)], + [[(1, 3), (3, 1)], (3, 3)], + [[(1, 1), (3, 3)], (3, 3)], + [[(1, 1), (1, 3)], (1, 3)], + [[(1, 1), (3, 1)], (3, 1)], + [[(1, 0), (0, 0)], (0, 0)], + [[(0, 1), (0, 0)], (0, 0)], + [[(1, 0), (0, 1)], (0, 0)], + [[(1, 1), (0, 0)], (0, 0)], + [[(1, 1), (1, 0)], (1, 0)], + [[(1, 1), (0, 1)], (0, 1)], + [[(), (3,)], (3,)], + [[(3,), (3, 3)], (3, 3)], + [[(3,), (3, 1)], (3, 3)], + [[(1,), (3, 3)], (3, 3)], + [[(), (3, 3)], (3, 3)], + [[(1, 1), (3,)], (1, 3)], + [[(1,), (3, 1)], (3, 1)], + [[(1,), (1, 3)], (1, 3)], + [[(), (1, 3)], (1, 3)], + [[(), (3, 1)], (3, 1)], + [[(), (0,)], (0,)], + [[(0,), (0, 0)], (0, 0)], + [[(0,), (0, 1)], (0, 0)], + [[(1,), (0, 0)], (0, 0)], + [[(), (0, 0)], (0, 0)], + [[(1, 1), (0,)], (1, 0)], + [[(1,), (0, 1)], (0, 1)], + [[(1,), (1, 0)], (1, 0)], + [[(), (1, 0)], (1, 0)], + [[(), (0, 1)], (0, 1)], + ] + for input_shapes, expected_shape in data: + assert_same_as_ufunc(input_shapes[0], input_shapes[1], + "Shapes: %s %s" % (input_shapes[0], input_shapes[1])) + # Reverse the input shapes since broadcasting should be symmetric. + assert_same_as_ufunc(input_shapes[1], input_shapes[0]) + # Try them transposed, too. + assert_same_as_ufunc(input_shapes[0], input_shapes[1], True) + # ... and flipped for non-rank-0 inputs in order to test negative + # strides. + if () not in input_shapes: + assert_same_as_ufunc(input_shapes[0], input_shapes[1], False, True) + assert_same_as_ufunc(input_shapes[0], input_shapes[1], True, True) + + +def test_broadcast_to_succeeds(): + data = [ + [np.array(0), (0,), np.array(0)], + [np.array(0), (1,), np.zeros(1)], + [np.array(0), (3,), np.zeros(3)], + [np.ones(1), (1,), np.ones(1)], + [np.ones(1), (2,), np.ones(2)], + [np.ones(1), (1, 2, 3), np.ones((1, 2, 3))], + [np.arange(3), (3,), np.arange(3)], + [np.arange(3), (1, 3), np.arange(3).reshape(1, -1)], + [np.arange(3), (2, 3), np.array([[0, 1, 2], [0, 1, 2]])], + # test if shape is not a tuple + [np.ones(0), 0, np.ones(0)], + [np.ones(1), 1, np.ones(1)], + [np.ones(1), 2, np.ones(2)], + # these cases with size 0 are strange, but they reproduce the behavior + # of broadcasting with ufuncs (see test_same_as_ufunc above) + [np.ones(1), (0,), np.ones(0)], + [np.ones((1, 2)), (0, 2), np.ones((0, 2))], + [np.ones((2, 1)), (2, 0), np.ones((2, 0))], + ] + for input_array, shape, expected in data: + actual = broadcast_to(input_array, shape) + assert_array_equal(expected, actual) + + +def test_broadcast_to_raises(): + data = [ + [(0,), ()], + [(1,), ()], + [(3,), ()], + [(3,), (1,)], + [(3,), (2,)], + [(3,), (4,)], + [(1, 2), (2, 1)], + [(1, 1), (1,)], + [(1,), -1], + [(1,), (-1,)], + [(1, 2), (-1, 2)], + ] + for orig_shape, target_shape in data: + arr = np.zeros(orig_shape) + assert_raises(ValueError, lambda: broadcast_to(arr, target_shape)) + + +def test_broadcast_shape(): + # tests internal _broadcast_shape + # _broadcast_shape is already exercised indirectly by broadcast_arrays + # _broadcast_shape is also exercised by the public broadcast_shapes function + assert_equal(_broadcast_shape(), ()) + assert_equal(_broadcast_shape([1, 2]), (2,)) + assert_equal(_broadcast_shape(np.ones((1, 1))), (1, 1)) + assert_equal(_broadcast_shape(np.ones((1, 1)), np.ones((3, 4))), (3, 4)) + assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 32)), (1, 2)) + assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 100)), (1, 2)) + + # regression tests for gh-5862 + assert_equal(_broadcast_shape(*([np.ones(2)] * 32 + [1])), (2,)) + bad_args = [np.ones(2)] * 32 + [np.ones(3)] * 32 + assert_raises(ValueError, lambda: _broadcast_shape(*bad_args)) + + +def test_broadcast_shapes_succeeds(): + # tests public broadcast_shapes + data = [ + [[], ()], + [[()], ()], + [[(7,)], (7,)], + [[(1, 2), (2,)], (1, 2)], + [[(1, 1)], (1, 1)], + [[(1, 1), (3, 4)], (3, 4)], + [[(6, 7), (5, 6, 1), (7,), (5, 1, 7)], (5, 6, 7)], + [[(5, 6, 1)], (5, 6, 1)], + [[(1, 3), (3, 1)], (3, 3)], + [[(1, 0), (0, 0)], (0, 0)], + [[(0, 1), (0, 0)], (0, 0)], + [[(1, 0), (0, 1)], (0, 0)], + [[(1, 1), (0, 0)], (0, 0)], + [[(1, 1), (1, 0)], (1, 0)], + [[(1, 1), (0, 1)], (0, 1)], + [[(), (0,)], (0,)], + [[(0,), (0, 0)], (0, 0)], + [[(0,), (0, 1)], (0, 0)], + [[(1,), (0, 0)], (0, 0)], + [[(), (0, 0)], (0, 0)], + [[(1, 1), (0,)], (1, 0)], + [[(1,), (0, 1)], (0, 1)], + [[(1,), (1, 0)], (1, 0)], + [[(), (1, 0)], (1, 0)], + [[(), (0, 1)], (0, 1)], + [[(1,), (3,)], (3,)], + [[2, (3, 2)], (3, 2)], + ] + for input_shapes, target_shape in data: + assert_equal(broadcast_shapes(*input_shapes), target_shape) + + assert_equal(broadcast_shapes(*([(1, 2)] * 32)), (1, 2)) + assert_equal(broadcast_shapes(*([(1, 2)] * 100)), (1, 2)) + + # regression tests for gh-5862 + assert_equal(broadcast_shapes(*([(2,)] * 32)), (2,)) + + +def test_broadcast_shapes_raises(): + # tests public broadcast_shapes + data = [ + [(3,), (4,)], + [(2, 3), (2,)], + [(3,), (3,), (4,)], + [(1, 3, 4), (2, 3, 3)], + [(1, 2), (3, 1), (3, 2), (10, 5)], + [2, (2, 3)], + ] + for input_shapes in data: + assert_raises(ValueError, lambda: broadcast_shapes(*input_shapes)) + + bad_args = [(2,)] * 32 + [(3,)] * 32 + assert_raises(ValueError, lambda: broadcast_shapes(*bad_args)) + + +def test_as_strided(): + a = np.array([None]) + a_view = as_strided(a) + expected = np.array([None]) + assert_array_equal(a_view, np.array([None])) + + a = np.array([1, 2, 3, 4]) + a_view = as_strided(a, shape=(2,), strides=(2 * a.itemsize,)) + expected = np.array([1, 3]) + assert_array_equal(a_view, expected) + + a = np.array([1, 2, 3, 4]) + a_view = as_strided(a, shape=(3, 4), strides=(0, 1 * a.itemsize)) + expected = np.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) + assert_array_equal(a_view, expected) + + # Regression test for gh-5081 + dt = np.dtype([('num', 'i4'), ('obj', 'O')]) + a = np.empty((4,), dtype=dt) + a['num'] = np.arange(1, 5) + a_view = as_strided(a, shape=(3, 4), strides=(0, a.itemsize)) + expected_num = [[1, 2, 3, 4]] * 3 + expected_obj = [[None]*4]*3 + assert_equal(a_view.dtype, dt) + assert_array_equal(expected_num, a_view['num']) + assert_array_equal(expected_obj, a_view['obj']) + + # Make sure that void types without fields are kept unchanged + a = np.empty((4,), dtype='V4') + a_view = as_strided(a, shape=(3, 4), strides=(0, a.itemsize)) + assert_equal(a.dtype, a_view.dtype) + + # Make sure that the only type that could fail is properly handled + dt = np.dtype({'names': [''], 'formats': ['V4']}) + a = np.empty((4,), dtype=dt) + a_view = as_strided(a, shape=(3, 4), strides=(0, a.itemsize)) + assert_equal(a.dtype, a_view.dtype) + + # Custom dtypes should not be lost (gh-9161) + r = [rational(i) for i in range(4)] + a = np.array(r, dtype=rational) + a_view = as_strided(a, shape=(3, 4), strides=(0, a.itemsize)) + assert_equal(a.dtype, a_view.dtype) + assert_array_equal([r] * 3, a_view) + + +class TestSlidingWindowView: + def test_1d(self): + arr = np.arange(5) + arr_view = sliding_window_view(arr, 2) + expected = np.array([[0, 1], + [1, 2], + [2, 3], + [3, 4]]) + assert_array_equal(arr_view, expected) + + def test_2d(self): + i, j = np.ogrid[:3, :4] + arr = 10*i + j + shape = (2, 2) + arr_view = sliding_window_view(arr, shape) + expected = np.array([[[[0, 1], [10, 11]], + [[1, 2], [11, 12]], + [[2, 3], [12, 13]]], + [[[10, 11], [20, 21]], + [[11, 12], [21, 22]], + [[12, 13], [22, 23]]]]) + assert_array_equal(arr_view, expected) + + def test_2d_with_axis(self): + i, j = np.ogrid[:3, :4] + arr = 10*i + j + arr_view = sliding_window_view(arr, 3, 0) + expected = np.array([[[0, 10, 20], + [1, 11, 21], + [2, 12, 22], + [3, 13, 23]]]) + assert_array_equal(arr_view, expected) + + def test_2d_repeated_axis(self): + i, j = np.ogrid[:3, :4] + arr = 10*i + j + arr_view = sliding_window_view(arr, (2, 3), (1, 1)) + expected = np.array([[[[0, 1, 2], + [1, 2, 3]]], + [[[10, 11, 12], + [11, 12, 13]]], + [[[20, 21, 22], + [21, 22, 23]]]]) + assert_array_equal(arr_view, expected) + + def test_2d_without_axis(self): + i, j = np.ogrid[:4, :4] + arr = 10*i + j + shape = (2, 3) + arr_view = sliding_window_view(arr, shape) + expected = np.array([[[[0, 1, 2], [10, 11, 12]], + [[1, 2, 3], [11, 12, 13]]], + [[[10, 11, 12], [20, 21, 22]], + [[11, 12, 13], [21, 22, 23]]], + [[[20, 21, 22], [30, 31, 32]], + [[21, 22, 23], [31, 32, 33]]]]) + assert_array_equal(arr_view, expected) + + def test_errors(self): + i, j = np.ogrid[:4, :4] + arr = 10*i + j + with pytest.raises(ValueError, match='cannot contain negative values'): + sliding_window_view(arr, (-1, 3)) + with pytest.raises( + ValueError, + match='must provide window_shape for all dimensions of `x`'): + sliding_window_view(arr, (1,)) + with pytest.raises( + ValueError, + match='Must provide matching length window_shape and axis'): + sliding_window_view(arr, (1, 3, 4), axis=(0, 1)) + with pytest.raises( + ValueError, + match='window shape cannot be larger than input array'): + sliding_window_view(arr, (5, 5)) + + def test_writeable(self): + arr = np.arange(5) + view = sliding_window_view(arr, 2, writeable=False) + assert_(not view.flags.writeable) + with pytest.raises( + ValueError, + match='assignment destination is read-only'): + view[0, 0] = 3 + view = sliding_window_view(arr, 2, writeable=True) + assert_(view.flags.writeable) + view[0, 1] = 3 + assert_array_equal(arr, np.array([0, 3, 2, 3, 4])) + + def test_subok(self): + class MyArray(np.ndarray): + pass + + arr = np.arange(5).view(MyArray) + assert_(not isinstance(sliding_window_view(arr, 2, + subok=False), + MyArray)) + assert_(isinstance(sliding_window_view(arr, 2, subok=True), MyArray)) + # Default behavior + assert_(not isinstance(sliding_window_view(arr, 2), MyArray)) + + +def as_strided_writeable(): + arr = np.ones(10) + view = as_strided(arr, writeable=False) + assert_(not view.flags.writeable) + + # Check that writeable also is fine: + view = as_strided(arr, writeable=True) + assert_(view.flags.writeable) + view[...] = 3 + assert_array_equal(arr, np.full_like(arr, 3)) + + # Test that things do not break down for readonly: + arr.flags.writeable = False + view = as_strided(arr, writeable=False) + view = as_strided(arr, writeable=True) + assert_(not view.flags.writeable) + + +class VerySimpleSubClass(np.ndarray): + def __new__(cls, *args, **kwargs): + return np.array(*args, subok=True, **kwargs).view(cls) + + +class SimpleSubClass(VerySimpleSubClass): + def __new__(cls, *args, **kwargs): + self = np.array(*args, subok=True, **kwargs).view(cls) + self.info = 'simple' + return self + + def __array_finalize__(self, obj): + self.info = getattr(obj, 'info', '') + ' finalized' + + +def test_subclasses(): + # test that subclass is preserved only if subok=True + a = VerySimpleSubClass([1, 2, 3, 4]) + assert_(type(a) is VerySimpleSubClass) + a_view = as_strided(a, shape=(2,), strides=(2 * a.itemsize,)) + assert_(type(a_view) is np.ndarray) + a_view = as_strided(a, shape=(2,), strides=(2 * a.itemsize,), subok=True) + assert_(type(a_view) is VerySimpleSubClass) + # test that if a subclass has __array_finalize__, it is used + a = SimpleSubClass([1, 2, 3, 4]) + a_view = as_strided(a, shape=(2,), strides=(2 * a.itemsize,), subok=True) + assert_(type(a_view) is SimpleSubClass) + assert_(a_view.info == 'simple finalized') + + # similar tests for broadcast_arrays + b = np.arange(len(a)).reshape(-1, 1) + a_view, b_view = broadcast_arrays(a, b) + assert_(type(a_view) is np.ndarray) + assert_(type(b_view) is np.ndarray) + assert_(a_view.shape == b_view.shape) + a_view, b_view = broadcast_arrays(a, b, subok=True) + assert_(type(a_view) is SimpleSubClass) + assert_(a_view.info == 'simple finalized') + assert_(type(b_view) is np.ndarray) + assert_(a_view.shape == b_view.shape) + + # and for broadcast_to + shape = (2, 4) + a_view = broadcast_to(a, shape) + assert_(type(a_view) is np.ndarray) + assert_(a_view.shape == shape) + a_view = broadcast_to(a, shape, subok=True) + assert_(type(a_view) is SimpleSubClass) + assert_(a_view.info == 'simple finalized') + assert_(a_view.shape == shape) + + +def test_writeable(): + # broadcast_to should return a readonly array + original = np.array([1, 2, 3]) + result = broadcast_to(original, (2, 3)) + assert_equal(result.flags.writeable, False) + assert_raises(ValueError, result.__setitem__, slice(None), 0) + + # but the result of broadcast_arrays needs to be writeable, to + # preserve backwards compatibility + test_cases = [((False,), broadcast_arrays(original,)), + ((True, False), broadcast_arrays(0, original))] + for is_broadcast, results in test_cases: + for array_is_broadcast, result in zip(is_broadcast, results): + # This will change to False in a future version + if array_is_broadcast: + with assert_warns(FutureWarning): + assert_equal(result.flags.writeable, True) + with assert_warns(DeprecationWarning): + result[:] = 0 + # Warning not emitted, writing to the array resets it + assert_equal(result.flags.writeable, True) + else: + # No warning: + assert_equal(result.flags.writeable, True) + + for results in [broadcast_arrays(original), + broadcast_arrays(0, original)]: + for result in results: + # resets the warn_on_write DeprecationWarning + result.flags.writeable = True + # check: no warning emitted + assert_equal(result.flags.writeable, True) + result[:] = 0 + + # keep readonly input readonly + original.flags.writeable = False + _, result = broadcast_arrays(0, original) + assert_equal(result.flags.writeable, False) + + # regression test for GH6491 + shape = (2,) + strides = [0] + tricky_array = as_strided(np.array(0), shape, strides) + other = np.zeros((1,)) + first, second = broadcast_arrays(tricky_array, other) + assert_(first.shape == second.shape) + + +def test_writeable_memoryview(): + # The result of broadcast_arrays exports as a non-writeable memoryview + # because otherwise there is no good way to opt in to the new behaviour + # (i.e. you would need to set writeable to False explicitly). + # See gh-13929. + original = np.array([1, 2, 3]) + + test_cases = [((False, ), broadcast_arrays(original,)), + ((True, False), broadcast_arrays(0, original))] + for is_broadcast, results in test_cases: + for array_is_broadcast, result in zip(is_broadcast, results): + # This will change to False in a future version + if array_is_broadcast: + # memoryview(result, writable=True) will give warning but cannot + # be tested using the python API. + assert memoryview(result).readonly + else: + assert not memoryview(result).readonly + + +def test_reference_types(): + input_array = np.array('a', dtype=object) + expected = np.array(['a'] * 3, dtype=object) + actual = broadcast_to(input_array, (3,)) + assert_array_equal(expected, actual) + + actual, _ = broadcast_arrays(input_array, np.ones(3)) + assert_array_equal(expected, actual) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_twodim_base.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_twodim_base.py new file mode 100644 index 0000000000000000000000000000000000000000..eb008c6002c86c94b180533230f849c909d10f39 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_twodim_base.py @@ -0,0 +1,541 @@ +"""Test functions for matrix module + +""" +from numpy.testing import ( + assert_equal, assert_array_equal, assert_array_max_ulp, + assert_array_almost_equal, assert_raises, assert_ +) +from numpy import ( + arange, add, fliplr, flipud, zeros, ones, eye, array, diag, histogram2d, + tri, mask_indices, triu_indices, triu_indices_from, tril_indices, + tril_indices_from, vander, +) +import numpy as np + +import pytest + + +def get_mat(n): + data = arange(n) + data = add.outer(data, data) + return data + + +class TestEye: + def test_basic(self): + assert_equal(eye(4), + array([[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1]])) + + assert_equal(eye(4, dtype='f'), + array([[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1]], 'f')) + + assert_equal(eye(3) == 1, + eye(3, dtype=bool)) + + def test_uint64(self): + # Regression test for gh-9982 + assert_equal(eye(np.uint64(2), dtype=int), array([[1, 0], [0, 1]])) + assert_equal(eye(np.uint64(2), M=np.uint64(4), k=np.uint64(1)), + array([[0, 1, 0, 0], [0, 0, 1, 0]])) + + def test_diag(self): + assert_equal(eye(4, k=1), + array([[0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + [0, 0, 0, 0]])) + + assert_equal(eye(4, k=-1), + array([[0, 0, 0, 0], + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0]])) + + def test_2d(self): + assert_equal(eye(4, 3), + array([[1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, 0, 0]])) + + assert_equal(eye(3, 4), + array([[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0]])) + + def test_diag2d(self): + assert_equal(eye(3, 4, k=2), + array([[0, 0, 1, 0], + [0, 0, 0, 1], + [0, 0, 0, 0]])) + + assert_equal(eye(4, 3, k=-2), + array([[0, 0, 0], + [0, 0, 0], + [1, 0, 0], + [0, 1, 0]])) + + def test_eye_bounds(self): + assert_equal(eye(2, 2, 1), [[0, 1], [0, 0]]) + assert_equal(eye(2, 2, -1), [[0, 0], [1, 0]]) + assert_equal(eye(2, 2, 2), [[0, 0], [0, 0]]) + assert_equal(eye(2, 2, -2), [[0, 0], [0, 0]]) + assert_equal(eye(3, 2, 2), [[0, 0], [0, 0], [0, 0]]) + assert_equal(eye(3, 2, 1), [[0, 1], [0, 0], [0, 0]]) + assert_equal(eye(3, 2, -1), [[0, 0], [1, 0], [0, 1]]) + assert_equal(eye(3, 2, -2), [[0, 0], [0, 0], [1, 0]]) + assert_equal(eye(3, 2, -3), [[0, 0], [0, 0], [0, 0]]) + + def test_strings(self): + assert_equal(eye(2, 2, dtype='S3'), + [[b'1', b''], [b'', b'1']]) + + def test_bool(self): + assert_equal(eye(2, 2, dtype=bool), [[True, False], [False, True]]) + + def test_order(self): + mat_c = eye(4, 3, k=-1) + mat_f = eye(4, 3, k=-1, order='F') + assert_equal(mat_c, mat_f) + assert mat_c.flags.c_contiguous + assert not mat_c.flags.f_contiguous + assert not mat_f.flags.c_contiguous + assert mat_f.flags.f_contiguous + + +class TestDiag: + def test_vector(self): + vals = (100 * arange(5)).astype('l') + b = zeros((5, 5)) + for k in range(5): + b[k, k] = vals[k] + assert_equal(diag(vals), b) + b = zeros((7, 7)) + c = b.copy() + for k in range(5): + b[k, k + 2] = vals[k] + c[k + 2, k] = vals[k] + assert_equal(diag(vals, k=2), b) + assert_equal(diag(vals, k=-2), c) + + def test_matrix(self, vals=None): + if vals is None: + vals = (100 * get_mat(5) + 1).astype('l') + b = zeros((5,)) + for k in range(5): + b[k] = vals[k, k] + assert_equal(diag(vals), b) + b = b * 0 + for k in range(3): + b[k] = vals[k, k + 2] + assert_equal(diag(vals, 2), b[:3]) + for k in range(3): + b[k] = vals[k + 2, k] + assert_equal(diag(vals, -2), b[:3]) + + def test_fortran_order(self): + vals = array((100 * get_mat(5) + 1), order='F', dtype='l') + self.test_matrix(vals) + + def test_diag_bounds(self): + A = [[1, 2], [3, 4], [5, 6]] + assert_equal(diag(A, k=2), []) + assert_equal(diag(A, k=1), [2]) + assert_equal(diag(A, k=0), [1, 4]) + assert_equal(diag(A, k=-1), [3, 6]) + assert_equal(diag(A, k=-2), [5]) + assert_equal(diag(A, k=-3), []) + + def test_failure(self): + assert_raises(ValueError, diag, [[[1]]]) + + +class TestFliplr: + def test_basic(self): + assert_raises(ValueError, fliplr, ones(4)) + a = get_mat(4) + b = a[:, ::-1] + assert_equal(fliplr(a), b) + a = [[0, 1, 2], + [3, 4, 5]] + b = [[2, 1, 0], + [5, 4, 3]] + assert_equal(fliplr(a), b) + + +class TestFlipud: + def test_basic(self): + a = get_mat(4) + b = a[::-1, :] + assert_equal(flipud(a), b) + a = [[0, 1, 2], + [3, 4, 5]] + b = [[3, 4, 5], + [0, 1, 2]] + assert_equal(flipud(a), b) + + +class TestHistogram2d: + def test_simple(self): + x = array( + [0.41702200, 0.72032449, 1.1437481e-4, 0.302332573, 0.146755891]) + y = array( + [0.09233859, 0.18626021, 0.34556073, 0.39676747, 0.53881673]) + xedges = np.linspace(0, 1, 10) + yedges = np.linspace(0, 1, 10) + H = histogram2d(x, y, (xedges, yedges))[0] + answer = array( + [[0, 0, 0, 1, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0, 0, 0, 0], + [0, 1, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0]]) + assert_array_equal(H.T, answer) + H = histogram2d(x, y, xedges)[0] + assert_array_equal(H.T, answer) + H, xedges, yedges = histogram2d(list(range(10)), list(range(10))) + assert_array_equal(H, eye(10, 10)) + assert_array_equal(xedges, np.linspace(0, 9, 11)) + assert_array_equal(yedges, np.linspace(0, 9, 11)) + + def test_asym(self): + x = array([1, 1, 2, 3, 4, 4, 4, 5]) + y = array([1, 3, 2, 0, 1, 2, 3, 4]) + H, xed, yed = histogram2d( + x, y, (6, 5), range=[[0, 6], [0, 5]], density=True) + answer = array( + [[0., 0, 0, 0, 0], + [0, 1, 0, 1, 0], + [0, 0, 1, 0, 0], + [1, 0, 0, 0, 0], + [0, 1, 1, 1, 0], + [0, 0, 0, 0, 1]]) + assert_array_almost_equal(H, answer/8., 3) + assert_array_equal(xed, np.linspace(0, 6, 7)) + assert_array_equal(yed, np.linspace(0, 5, 6)) + + def test_density(self): + x = array([1, 2, 3, 1, 2, 3, 1, 2, 3]) + y = array([1, 1, 1, 2, 2, 2, 3, 3, 3]) + H, xed, yed = histogram2d( + x, y, [[1, 2, 3, 5], [1, 2, 3, 5]], density=True) + answer = array([[1, 1, .5], + [1, 1, .5], + [.5, .5, .25]])/9. + assert_array_almost_equal(H, answer, 3) + + def test_all_outliers(self): + r = np.random.rand(100) + 1. + 1e6 # histogramdd rounds by decimal=6 + H, xed, yed = histogram2d(r, r, (4, 5), range=([0, 1], [0, 1])) + assert_array_equal(H, 0) + + def test_empty(self): + a, edge1, edge2 = histogram2d([], [], bins=([0, 1], [0, 1])) + assert_array_max_ulp(a, array([[0.]])) + + a, edge1, edge2 = histogram2d([], [], bins=4) + assert_array_max_ulp(a, np.zeros((4, 4))) + + def test_binparameter_combination(self): + x = array( + [0, 0.09207008, 0.64575234, 0.12875982, 0.47390599, + 0.59944483, 1]) + y = array( + [0, 0.14344267, 0.48988575, 0.30558665, 0.44700682, + 0.15886423, 1]) + edges = (0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1) + H, xe, ye = histogram2d(x, y, (edges, 4)) + answer = array( + [[2., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 1., 0., 0.], + [1., 0., 0., 0.], + [0., 1., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 1.]]) + assert_array_equal(H, answer) + assert_array_equal(ye, array([0., 0.25, 0.5, 0.75, 1])) + H, xe, ye = histogram2d(x, y, (4, edges)) + answer = array( + [[1., 1., 0., 1., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.], + [0., 1., 0., 0., 1., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]]) + assert_array_equal(H, answer) + assert_array_equal(xe, array([0., 0.25, 0.5, 0.75, 1])) + + def test_dispatch(self): + class ShouldDispatch: + def __array_function__(self, function, types, args, kwargs): + return types, args, kwargs + + xy = [1, 2] + s_d = ShouldDispatch() + r = histogram2d(s_d, xy) + # Cannot use assert_equal since that dispatches... + assert_(r == ((ShouldDispatch,), (s_d, xy), {})) + r = histogram2d(xy, s_d) + assert_(r == ((ShouldDispatch,), (xy, s_d), {})) + r = histogram2d(xy, xy, bins=s_d) + assert_(r, ((ShouldDispatch,), (xy, xy), dict(bins=s_d))) + r = histogram2d(xy, xy, bins=[s_d, 5]) + assert_(r, ((ShouldDispatch,), (xy, xy), dict(bins=[s_d, 5]))) + assert_raises(Exception, histogram2d, xy, xy, bins=[s_d]) + r = histogram2d(xy, xy, weights=s_d) + assert_(r, ((ShouldDispatch,), (xy, xy), dict(weights=s_d))) + + @pytest.mark.parametrize(("x_len", "y_len"), [(10, 11), (20, 19)]) + def test_bad_length(self, x_len, y_len): + x, y = np.ones(x_len), np.ones(y_len) + with pytest.raises(ValueError, + match='x and y must have the same length.'): + histogram2d(x, y) + + +class TestTri: + def test_dtype(self): + out = array([[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]) + assert_array_equal(tri(3), out) + assert_array_equal(tri(3, dtype=bool), out.astype(bool)) + + +def test_tril_triu_ndim2(): + for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']: + a = np.ones((2, 2), dtype=dtype) + b = np.tril(a) + c = np.triu(a) + assert_array_equal(b, [[1, 0], [1, 1]]) + assert_array_equal(c, b.T) + # should return the same dtype as the original array + assert_equal(b.dtype, a.dtype) + assert_equal(c.dtype, a.dtype) + + +def test_tril_triu_ndim3(): + for dtype in np.typecodes['AllFloat'] + np.typecodes['AllInteger']: + a = np.array([ + [[1, 1], [1, 1]], + [[1, 1], [1, 0]], + [[1, 1], [0, 0]], + ], dtype=dtype) + a_tril_desired = np.array([ + [[1, 0], [1, 1]], + [[1, 0], [1, 0]], + [[1, 0], [0, 0]], + ], dtype=dtype) + a_triu_desired = np.array([ + [[1, 1], [0, 1]], + [[1, 1], [0, 0]], + [[1, 1], [0, 0]], + ], dtype=dtype) + a_triu_observed = np.triu(a) + a_tril_observed = np.tril(a) + assert_array_equal(a_triu_observed, a_triu_desired) + assert_array_equal(a_tril_observed, a_tril_desired) + assert_equal(a_triu_observed.dtype, a.dtype) + assert_equal(a_tril_observed.dtype, a.dtype) + + +def test_tril_triu_with_inf(): + # Issue 4859 + arr = np.array([[1, 1, np.inf], + [1, 1, 1], + [np.inf, 1, 1]]) + out_tril = np.array([[1, 0, 0], + [1, 1, 0], + [np.inf, 1, 1]]) + out_triu = out_tril.T + assert_array_equal(np.triu(arr), out_triu) + assert_array_equal(np.tril(arr), out_tril) + + +def test_tril_triu_dtype(): + # Issue 4916 + # tril and triu should return the same dtype as input + for c in np.typecodes['All']: + if c == 'V': + continue + arr = np.zeros((3, 3), dtype=c) + assert_equal(np.triu(arr).dtype, arr.dtype) + assert_equal(np.tril(arr).dtype, arr.dtype) + + # check special cases + arr = np.array([['2001-01-01T12:00', '2002-02-03T13:56'], + ['2004-01-01T12:00', '2003-01-03T13:45']], + dtype='datetime64') + assert_equal(np.triu(arr).dtype, arr.dtype) + assert_equal(np.tril(arr).dtype, arr.dtype) + + arr = np.zeros((3, 3), dtype='f4,f4') + assert_equal(np.triu(arr).dtype, arr.dtype) + assert_equal(np.tril(arr).dtype, arr.dtype) + + +def test_mask_indices(): + # simple test without offset + iu = mask_indices(3, np.triu) + a = np.arange(9).reshape(3, 3) + assert_array_equal(a[iu], array([0, 1, 2, 4, 5, 8])) + # Now with an offset + iu1 = mask_indices(3, np.triu, 1) + assert_array_equal(a[iu1], array([1, 2, 5])) + + +def test_tril_indices(): + # indices without and with offset + il1 = tril_indices(4) + il2 = tril_indices(4, k=2) + il3 = tril_indices(4, m=5) + il4 = tril_indices(4, k=2, m=5) + + a = np.array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + b = np.arange(1, 21).reshape(4, 5) + + # indexing: + assert_array_equal(a[il1], + array([1, 5, 6, 9, 10, 11, 13, 14, 15, 16])) + assert_array_equal(b[il3], + array([1, 6, 7, 11, 12, 13, 16, 17, 18, 19])) + + # And for assigning values: + a[il1] = -1 + assert_array_equal(a, + array([[-1, 2, 3, 4], + [-1, -1, 7, 8], + [-1, -1, -1, 12], + [-1, -1, -1, -1]])) + b[il3] = -1 + assert_array_equal(b, + array([[-1, 2, 3, 4, 5], + [-1, -1, 8, 9, 10], + [-1, -1, -1, 14, 15], + [-1, -1, -1, -1, 20]])) + # These cover almost the whole array (two diagonals right of the main one): + a[il2] = -10 + assert_array_equal(a, + array([[-10, -10, -10, 4], + [-10, -10, -10, -10], + [-10, -10, -10, -10], + [-10, -10, -10, -10]])) + b[il4] = -10 + assert_array_equal(b, + array([[-10, -10, -10, 4, 5], + [-10, -10, -10, -10, 10], + [-10, -10, -10, -10, -10], + [-10, -10, -10, -10, -10]])) + + +class TestTriuIndices: + def test_triu_indices(self): + iu1 = triu_indices(4) + iu2 = triu_indices(4, k=2) + iu3 = triu_indices(4, m=5) + iu4 = triu_indices(4, k=2, m=5) + + a = np.array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + b = np.arange(1, 21).reshape(4, 5) + + # Both for indexing: + assert_array_equal(a[iu1], + array([1, 2, 3, 4, 6, 7, 8, 11, 12, 16])) + assert_array_equal(b[iu3], + array([1, 2, 3, 4, 5, 7, 8, 9, + 10, 13, 14, 15, 19, 20])) + + # And for assigning values: + a[iu1] = -1 + assert_array_equal(a, + array([[-1, -1, -1, -1], + [5, -1, -1, -1], + [9, 10, -1, -1], + [13, 14, 15, -1]])) + b[iu3] = -1 + assert_array_equal(b, + array([[-1, -1, -1, -1, -1], + [6, -1, -1, -1, -1], + [11, 12, -1, -1, -1], + [16, 17, 18, -1, -1]])) + + # These cover almost the whole array (two diagonals right of the + # main one): + a[iu2] = -10 + assert_array_equal(a, + array([[-1, -1, -10, -10], + [5, -1, -1, -10], + [9, 10, -1, -1], + [13, 14, 15, -1]])) + b[iu4] = -10 + assert_array_equal(b, + array([[-1, -1, -10, -10, -10], + [6, -1, -1, -10, -10], + [11, 12, -1, -1, -10], + [16, 17, 18, -1, -1]])) + + +class TestTrilIndicesFrom: + def test_exceptions(self): + assert_raises(ValueError, tril_indices_from, np.ones((2,))) + assert_raises(ValueError, tril_indices_from, np.ones((2, 2, 2))) + # assert_raises(ValueError, tril_indices_from, np.ones((2, 3))) + + +class TestTriuIndicesFrom: + def test_exceptions(self): + assert_raises(ValueError, triu_indices_from, np.ones((2,))) + assert_raises(ValueError, triu_indices_from, np.ones((2, 2, 2))) + # assert_raises(ValueError, triu_indices_from, np.ones((2, 3))) + + +class TestVander: + def test_basic(self): + c = np.array([0, 1, -2, 3]) + v = vander(c) + powers = np.array([[0, 0, 0, 0, 1], + [1, 1, 1, 1, 1], + [16, -8, 4, -2, 1], + [81, 27, 9, 3, 1]]) + # Check default value of N: + assert_array_equal(v, powers[:, 1:]) + # Check a range of N values, including 0 and 5 (greater than default) + m = powers.shape[1] + for n in range(6): + v = vander(c, N=n) + assert_array_equal(v, powers[:, m-n:m]) + + def test_dtypes(self): + c = array([11, -12, 13], dtype=np.int8) + v = vander(c) + expected = np.array([[121, 11, 1], + [144, -12, 1], + [169, 13, 1]]) + assert_array_equal(v, expected) + + c = array([1.0+1j, 1.0-1j]) + v = vander(c, N=3) + expected = np.array([[2j, 1+1j, 1], + [-2j, 1-1j, 1]]) + # The data is floating point, but the values are small integers, + # so assert_array_equal *should* be safe here (rather than, say, + # assert_array_almost_equal). + assert_array_equal(v, expected) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_type_check.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_type_check.py new file mode 100644 index 0000000000000000000000000000000000000000..01c888bef6f1bc524226142b12fda4bc149f5ce0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_type_check.py @@ -0,0 +1,465 @@ +import numpy as np +from numpy import ( + common_type, mintypecode, isreal, iscomplex, isposinf, isneginf, + nan_to_num, isrealobj, iscomplexobj, real_if_close + ) +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_raises + ) + + +def assert_all(x): + assert_(np.all(x), x) + + +class TestCommonType: + def test_basic(self): + ai32 = np.array([[1, 2], [3, 4]], dtype=np.int32) + af16 = np.array([[1, 2], [3, 4]], dtype=np.float16) + af32 = np.array([[1, 2], [3, 4]], dtype=np.float32) + af64 = np.array([[1, 2], [3, 4]], dtype=np.float64) + acs = np.array([[1+5j, 2+6j], [3+7j, 4+8j]], dtype=np.complex64) + acd = np.array([[1+5j, 2+6j], [3+7j, 4+8j]], dtype=np.complex128) + assert_(common_type(ai32) == np.float64) + assert_(common_type(af16) == np.float16) + assert_(common_type(af32) == np.float32) + assert_(common_type(af64) == np.float64) + assert_(common_type(acs) == np.complex64) + assert_(common_type(acd) == np.complex128) + + +class TestMintypecode: + + def test_default_1(self): + for itype in '1bcsuwil': + assert_equal(mintypecode(itype), 'd') + assert_equal(mintypecode('f'), 'f') + assert_equal(mintypecode('d'), 'd') + assert_equal(mintypecode('F'), 'F') + assert_equal(mintypecode('D'), 'D') + + def test_default_2(self): + for itype in '1bcsuwil': + assert_equal(mintypecode(itype+'f'), 'f') + assert_equal(mintypecode(itype+'d'), 'd') + assert_equal(mintypecode(itype+'F'), 'F') + assert_equal(mintypecode(itype+'D'), 'D') + assert_equal(mintypecode('ff'), 'f') + assert_equal(mintypecode('fd'), 'd') + assert_equal(mintypecode('fF'), 'F') + assert_equal(mintypecode('fD'), 'D') + assert_equal(mintypecode('df'), 'd') + assert_equal(mintypecode('dd'), 'd') + #assert_equal(mintypecode('dF',savespace=1),'F') + assert_equal(mintypecode('dF'), 'D') + assert_equal(mintypecode('dD'), 'D') + assert_equal(mintypecode('Ff'), 'F') + #assert_equal(mintypecode('Fd',savespace=1),'F') + assert_equal(mintypecode('Fd'), 'D') + assert_equal(mintypecode('FF'), 'F') + assert_equal(mintypecode('FD'), 'D') + assert_equal(mintypecode('Df'), 'D') + assert_equal(mintypecode('Dd'), 'D') + assert_equal(mintypecode('DF'), 'D') + assert_equal(mintypecode('DD'), 'D') + + def test_default_3(self): + assert_equal(mintypecode('fdF'), 'D') + #assert_equal(mintypecode('fdF',savespace=1),'F') + assert_equal(mintypecode('fdD'), 'D') + assert_equal(mintypecode('fFD'), 'D') + assert_equal(mintypecode('dFD'), 'D') + + assert_equal(mintypecode('ifd'), 'd') + assert_equal(mintypecode('ifF'), 'F') + assert_equal(mintypecode('ifD'), 'D') + assert_equal(mintypecode('idF'), 'D') + #assert_equal(mintypecode('idF',savespace=1),'F') + assert_equal(mintypecode('idD'), 'D') + + +class TestIsscalar: + + def test_basic(self): + assert_(np.isscalar(3)) + assert_(not np.isscalar([3])) + assert_(not np.isscalar((3,))) + assert_(np.isscalar(3j)) + assert_(np.isscalar(4.0)) + + +class TestReal: + + def test_real(self): + y = np.random.rand(10,) + assert_array_equal(y, np.real(y)) + + y = np.array(1) + out = np.real(y) + assert_array_equal(y, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + out = np.real(y) + assert_equal(y, out) + assert_(not isinstance(out, np.ndarray)) + + def test_cmplx(self): + y = np.random.rand(10,)+1j*np.random.rand(10,) + assert_array_equal(y.real, np.real(y)) + + y = np.array(1 + 1j) + out = np.real(y) + assert_array_equal(y.real, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + 1j + out = np.real(y) + assert_equal(1.0, out) + assert_(not isinstance(out, np.ndarray)) + + +class TestImag: + + def test_real(self): + y = np.random.rand(10,) + assert_array_equal(0, np.imag(y)) + + y = np.array(1) + out = np.imag(y) + assert_array_equal(0, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + out = np.imag(y) + assert_equal(0, out) + assert_(not isinstance(out, np.ndarray)) + + def test_cmplx(self): + y = np.random.rand(10,)+1j*np.random.rand(10,) + assert_array_equal(y.imag, np.imag(y)) + + y = np.array(1 + 1j) + out = np.imag(y) + assert_array_equal(y.imag, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + 1j + out = np.imag(y) + assert_equal(1.0, out) + assert_(not isinstance(out, np.ndarray)) + + +class TestIscomplex: + + def test_fail(self): + z = np.array([-1, 0, 1]) + res = iscomplex(z) + assert_(not np.any(res, axis=0)) + + def test_pass(self): + z = np.array([-1j, 1, 0]) + res = iscomplex(z) + assert_array_equal(res, [1, 0, 0]) + + +class TestIsreal: + + def test_pass(self): + z = np.array([-1, 0, 1j]) + res = isreal(z) + assert_array_equal(res, [1, 1, 0]) + + def test_fail(self): + z = np.array([-1j, 1, 0]) + res = isreal(z) + assert_array_equal(res, [0, 1, 1]) + + +class TestIscomplexobj: + + def test_basic(self): + z = np.array([-1, 0, 1]) + assert_(not iscomplexobj(z)) + z = np.array([-1j, 0, -1]) + assert_(iscomplexobj(z)) + + def test_scalar(self): + assert_(not iscomplexobj(1.0)) + assert_(iscomplexobj(1+0j)) + + def test_list(self): + assert_(iscomplexobj([3, 1+0j, True])) + assert_(not iscomplexobj([3, 1, True])) + + def test_duck(self): + class DummyComplexArray: + @property + def dtype(self): + return np.dtype(complex) + dummy = DummyComplexArray() + assert_(iscomplexobj(dummy)) + + def test_pandas_duck(self): + # This tests a custom np.dtype duck-typed class, such as used by pandas + # (pandas.core.dtypes) + class PdComplex(np.complex128): + pass + class PdDtype: + name = 'category' + names = None + type = PdComplex + kind = 'c' + str = ' 1e10) and assert_all(np.isfinite(vals[2])) + assert_equal(type(vals), np.ndarray) + + # perform the same tests but with nan, posinf and neginf keywords + with np.errstate(divide='ignore', invalid='ignore'): + vals = nan_to_num(np.array((-1., 0, 1))/0., + nan=10, posinf=20, neginf=30) + assert_equal(vals, [30, 10, 20]) + assert_all(np.isfinite(vals[[0, 2]])) + assert_equal(type(vals), np.ndarray) + + # perform the same test but in-place + with np.errstate(divide='ignore', invalid='ignore'): + vals = np.array((-1., 0, 1))/0. + result = nan_to_num(vals, copy=False) + + assert_(result is vals) + assert_all(vals[0] < -1e10) and assert_all(np.isfinite(vals[0])) + assert_(vals[1] == 0) + assert_all(vals[2] > 1e10) and assert_all(np.isfinite(vals[2])) + assert_equal(type(vals), np.ndarray) + + # perform the same test but in-place + with np.errstate(divide='ignore', invalid='ignore'): + vals = np.array((-1., 0, 1))/0. + result = nan_to_num(vals, copy=False, nan=10, posinf=20, neginf=30) + + assert_(result is vals) + assert_equal(vals, [30, 10, 20]) + assert_all(np.isfinite(vals[[0, 2]])) + assert_equal(type(vals), np.ndarray) + + def test_array(self): + vals = nan_to_num([1]) + assert_array_equal(vals, np.array([1], int)) + assert_equal(type(vals), np.ndarray) + vals = nan_to_num([1], nan=10, posinf=20, neginf=30) + assert_array_equal(vals, np.array([1], int)) + assert_equal(type(vals), np.ndarray) + + def test_integer(self): + vals = nan_to_num(1) + assert_all(vals == 1) + assert_equal(type(vals), np.int_) + vals = nan_to_num(1, nan=10, posinf=20, neginf=30) + assert_all(vals == 1) + assert_equal(type(vals), np.int_) + + def test_float(self): + vals = nan_to_num(1.0) + assert_all(vals == 1.0) + assert_equal(type(vals), np.float64) + vals = nan_to_num(1.1, nan=10, posinf=20, neginf=30) + assert_all(vals == 1.1) + assert_equal(type(vals), np.float64) + + def test_complex_good(self): + vals = nan_to_num(1+1j) + assert_all(vals == 1+1j) + assert_equal(type(vals), np.complex128) + vals = nan_to_num(1+1j, nan=10, posinf=20, neginf=30) + assert_all(vals == 1+1j) + assert_equal(type(vals), np.complex128) + + def test_complex_bad(self): + with np.errstate(divide='ignore', invalid='ignore'): + v = 1 + 1j + v += np.array(0+1.j)/0. + vals = nan_to_num(v) + # !! This is actually (unexpectedly) zero + assert_all(np.isfinite(vals)) + assert_equal(type(vals), np.complex128) + + def test_complex_bad2(self): + with np.errstate(divide='ignore', invalid='ignore'): + v = 1 + 1j + v += np.array(-1+1.j)/0. + vals = nan_to_num(v) + assert_all(np.isfinite(vals)) + assert_equal(type(vals), np.complex128) + # Fixme + #assert_all(vals.imag > 1e10) and assert_all(np.isfinite(vals)) + # !! This is actually (unexpectedly) positive + # !! inf. Comment out for now, and see if it + # !! changes + #assert_all(vals.real < -1e10) and assert_all(np.isfinite(vals)) + + def test_do_not_rewrite_previous_keyword(self): + # This is done to test that when, for instance, nan=np.inf then these + # values are not rewritten by posinf keyword to the posinf value. + with np.errstate(divide='ignore', invalid='ignore'): + vals = nan_to_num(np.array((-1., 0, 1))/0., nan=np.inf, posinf=999) + assert_all(np.isfinite(vals[[0, 2]])) + assert_all(vals[0] < -1e10) + assert_equal(vals[[1, 2]], [np.inf, 999]) + assert_equal(type(vals), np.ndarray) + + +class TestRealIfClose: + + def test_basic(self): + a = np.random.rand(10) + b = real_if_close(a+1e-15j) + assert_all(isrealobj(b)) + assert_array_equal(a, b) + b = real_if_close(a+1e-7j) + assert_all(iscomplexobj(b)) + b = real_if_close(a+1e-7j, tol=1e-6) + assert_all(isrealobj(b)) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_ufunclike.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_ufunclike.py new file mode 100644 index 0000000000000000000000000000000000000000..4b5d11010e0f0e4dcd94c3b0067caee7a849ea58 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_ufunclike.py @@ -0,0 +1,100 @@ +import numpy as np + +from numpy import fix, isposinf, isneginf +from numpy.testing import ( + assert_, assert_equal, assert_array_equal, assert_raises +) + + +class TestUfunclike: + + def test_isposinf(self): + a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0]) + out = np.zeros(a.shape, bool) + tgt = np.array([True, False, False, False, False, False]) + + res = isposinf(a) + assert_equal(res, tgt) + res = isposinf(a, out) + assert_equal(res, tgt) + assert_equal(out, tgt) + + a = a.astype(np.complex128) + with assert_raises(TypeError): + isposinf(a) + + def test_isneginf(self): + a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0]) + out = np.zeros(a.shape, bool) + tgt = np.array([False, True, False, False, False, False]) + + res = isneginf(a) + assert_equal(res, tgt) + res = isneginf(a, out) + assert_equal(res, tgt) + assert_equal(out, tgt) + + a = a.astype(np.complex128) + with assert_raises(TypeError): + isneginf(a) + + def test_fix(self): + a = np.array([[1.0, 1.1, 1.5, 1.8], [-1.0, -1.1, -1.5, -1.8]]) + out = np.zeros(a.shape, float) + tgt = np.array([[1., 1., 1., 1.], [-1., -1., -1., -1.]]) + + res = fix(a) + assert_equal(res, tgt) + res = fix(a, out) + assert_equal(res, tgt) + assert_equal(out, tgt) + assert_equal(fix(3.14), 3) + + def test_fix_with_subclass(self): + class MyArray(np.ndarray): + def __new__(cls, data, metadata=None): + res = np.array(data, copy=True).view(cls) + res.metadata = metadata + return res + + def __array_wrap__(self, obj, context=None, return_scalar=False): + if not isinstance(obj, MyArray): + obj = obj.view(MyArray) + if obj.metadata is None: + obj.metadata = self.metadata + return obj + + def __array_finalize__(self, obj): + self.metadata = getattr(obj, 'metadata', None) + return self + + a = np.array([1.1, -1.1]) + m = MyArray(a, metadata='foo') + f = fix(m) + assert_array_equal(f, np.array([1, -1])) + assert_(isinstance(f, MyArray)) + assert_equal(f.metadata, 'foo') + + # check 0d arrays don't decay to scalars + m0d = m[0,...] + m0d.metadata = 'bar' + f0d = fix(m0d) + assert_(isinstance(f0d, MyArray)) + assert_equal(f0d.metadata, 'bar') + + def test_scalar(self): + x = np.inf + actual = np.isposinf(x) + expected = np.True_ + assert_equal(actual, expected) + assert_equal(type(actual), type(expected)) + + x = -3.4 + actual = np.fix(x) + expected = np.float64(-3.0) + assert_equal(actual, expected) + assert_equal(type(actual), type(expected)) + + out = np.array(0.0) + actual = np.fix(x, out=out) + assert_(actual is out) diff --git a/venv/lib/python3.10/site-packages/numpy/lib/tests/test_utils.py b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..644912d941e3c233741ba415f6469b1cebf994b2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/lib/tests/test_utils.py @@ -0,0 +1,80 @@ +import pytest + +import numpy as np +from numpy.testing import assert_raises_regex +import numpy.lib._utils_impl as _utils_impl + +from io import StringIO + + +def test_assert_raises_regex_context_manager(): + with assert_raises_regex(ValueError, 'no deprecation warning'): + raise ValueError('no deprecation warning') + + +def test_info_method_heading(): + # info(class) should only print "Methods:" heading if methods exist + + class NoPublicMethods: + pass + + class WithPublicMethods: + def first_method(): + pass + + def _has_method_heading(cls): + out = StringIO() + np.info(cls, output=out) + return 'Methods:' in out.getvalue() + + assert _has_method_heading(WithPublicMethods) + assert not _has_method_heading(NoPublicMethods) + + +def test_drop_metadata(): + def _compare_dtypes(dt1, dt2): + return np.can_cast(dt1, dt2, casting='no') + + # structured dtype + dt = np.dtype([('l1', [('l2', np.dtype('S8', metadata={'msg': 'toto'}))])], + metadata={'msg': 'titi'}) + dt_m = _utils_impl.drop_metadata(dt) + assert _compare_dtypes(dt, dt_m) is True + assert dt_m.metadata is None + assert dt_m['l1'].metadata is None + assert dt_m['l1']['l2'].metadata is None + + # alignment + dt = np.dtype([('x', '>> from numpy.polynomial import Chebyshev + >>> xdata = [1, 2, 3, 4] + >>> ydata = [1, 4, 9, 16] + >>> c = Chebyshev.fit(xdata, ydata, deg=1) + +is preferred over the `chebyshev.chebfit` function from the +``np.polynomial.chebyshev`` module:: + + >>> from numpy.polynomial.chebyshev import chebfit + >>> c = chebfit(xdata, ydata, deg=1) + +See :doc:`routines.polynomials.classes` for more details. + +Convenience Classes +=================== + +The following lists the various constants and methods common to all of +the classes representing the various kinds of polynomials. In the following, +the term ``Poly`` represents any one of the convenience classes (e.g. +`~polynomial.Polynomial`, `~chebyshev.Chebyshev`, `~hermite.Hermite`, etc.) +while the lowercase ``p`` represents an **instance** of a polynomial class. + +Constants +--------- + +- ``Poly.domain`` -- Default domain +- ``Poly.window`` -- Default window +- ``Poly.basis_name`` -- String used to represent the basis +- ``Poly.maxpower`` -- Maximum value ``n`` such that ``p**n`` is allowed +- ``Poly.nickname`` -- String used in printing + +Creation +-------- + +Methods for creating polynomial instances. + +- ``Poly.basis(degree)`` -- Basis polynomial of given degree +- ``Poly.identity()`` -- ``p`` where ``p(x) = x`` for all ``x`` +- ``Poly.fit(x, y, deg)`` -- ``p`` of degree ``deg`` with coefficients + determined by the least-squares fit to the data ``x``, ``y`` +- ``Poly.fromroots(roots)`` -- ``p`` with specified roots +- ``p.copy()`` -- Create a copy of ``p`` + +Conversion +---------- + +Methods for converting a polynomial instance of one kind to another. + +- ``p.cast(Poly)`` -- Convert ``p`` to instance of kind ``Poly`` +- ``p.convert(Poly)`` -- Convert ``p`` to instance of kind ``Poly`` or map + between ``domain`` and ``window`` + +Calculus +-------- +- ``p.deriv()`` -- Take the derivative of ``p`` +- ``p.integ()`` -- Integrate ``p`` + +Validation +---------- +- ``Poly.has_samecoef(p1, p2)`` -- Check if coefficients match +- ``Poly.has_samedomain(p1, p2)`` -- Check if domains match +- ``Poly.has_sametype(p1, p2)`` -- Check if types match +- ``Poly.has_samewindow(p1, p2)`` -- Check if windows match + +Misc +---- +- ``p.linspace()`` -- Return ``x, p(x)`` at equally-spaced points in ``domain`` +- ``p.mapparms()`` -- Return the parameters for the linear mapping between + ``domain`` and ``window``. +- ``p.roots()`` -- Return the roots of ``p``. +- ``p.trim()`` -- Remove trailing coefficients. +- ``p.cutdeg(degree)`` -- Truncate ``p`` to given degree +- ``p.truncate(size)`` -- Truncate ``p`` to given size + +""" +from .polynomial import Polynomial +from .chebyshev import Chebyshev +from .legendre import Legendre +from .hermite import Hermite +from .hermite_e import HermiteE +from .laguerre import Laguerre + +__all__ = [ + "set_default_printstyle", + "polynomial", "Polynomial", + "chebyshev", "Chebyshev", + "legendre", "Legendre", + "hermite", "Hermite", + "hermite_e", "HermiteE", + "laguerre", "Laguerre", +] + + +def set_default_printstyle(style): + """ + Set the default format for the string representation of polynomials. + + Values for ``style`` must be valid inputs to ``__format__``, i.e. 'ascii' + or 'unicode'. + + Parameters + ---------- + style : str + Format string for default printing style. Must be either 'ascii' or + 'unicode'. + + Notes + ----- + The default format depends on the platform: 'unicode' is used on + Unix-based systems and 'ascii' on Windows. This determination is based on + default font support for the unicode superscript and subscript ranges. + + Examples + -------- + >>> p = np.polynomial.Polynomial([1, 2, 3]) + >>> c = np.polynomial.Chebyshev([1, 2, 3]) + >>> np.polynomial.set_default_printstyle('unicode') + >>> print(p) + 1.0 + 2.0·x + 3.0·x² + >>> print(c) + 1.0 + 2.0·T₁(x) + 3.0·T₂(x) + >>> np.polynomial.set_default_printstyle('ascii') + >>> print(p) + 1.0 + 2.0 x + 3.0 x**2 + >>> print(c) + 1.0 + 2.0 T_1(x) + 3.0 T_2(x) + >>> # Formatting supersedes all class/package-level defaults + >>> print(f"{p:unicode}") + 1.0 + 2.0·x + 3.0·x² + """ + if style not in ('unicode', 'ascii'): + raise ValueError( + f"Unsupported format string '{style}'. Valid options are 'ascii' " + f"and 'unicode'" + ) + _use_unicode = True + if style == 'ascii': + _use_unicode = False + from ._polybase import ABCPolyBase + ABCPolyBase._use_unicode = _use_unicode + + +from numpy._pytesttester import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__init__.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c5dccfe16dee8889508150ecfe963297f24a5fd0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/__init__.pyi @@ -0,0 +1,24 @@ +from typing import Final, Literal + +from .polynomial import Polynomial +from .chebyshev import Chebyshev +from .legendre import Legendre +from .hermite import Hermite +from .hermite_e import HermiteE +from .laguerre import Laguerre +from . import polynomial, chebyshev, legendre, hermite, hermite_e, laguerre + +__all__ = [ + "set_default_printstyle", + "polynomial", "Polynomial", + "chebyshev", "Chebyshev", + "legendre", "Legendre", + "hermite", "Hermite", + "hermite_e", "HermiteE", + "laguerre", "Laguerre", +] + +def set_default_printstyle(style: Literal["ascii", "unicode"]) -> None: ... + +from numpy._pytesttester import PytestTester as _PytestTester +test: Final[_PytestTester] diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4f6f23f4242a56fd68a9341c54e45cb0042995a Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/_polybase.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/_polybase.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87b7e35b12599d724058df2bfd06d3e2d4983f1d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/_polybase.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/chebyshev.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/chebyshev.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..088e9aec34332e6e0129ae434f1eda8c7fe03ee7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/chebyshev.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a85e34fbf7a467deb3f4d7e24ec13fc75dce106 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite_e.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite_e.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73f3ce0cde9a34ea9427f1d4665b09f698b074da Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/hermite_e.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/laguerre.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/laguerre.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87ed0dcf6c5b01933bfa28530dba30881428a4cf Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/laguerre.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/legendre.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/legendre.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24c379da002d1286018a15af2b37150ae89eb749 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/legendre.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polynomial.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polynomial.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43e13a37c6c3d9afd3d7ee630bc0fd67e10551d7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polynomial.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polyutils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polyutils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..750ae7cc6026fdb68644fd5a6115960726473b90 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/__pycache__/polyutils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.py b/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.py new file mode 100644 index 0000000000000000000000000000000000000000..1c3d16c6efd7af25ef0cdfc32083802ecce2a92f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.py @@ -0,0 +1,1197 @@ +""" +Abstract base class for the various polynomial Classes. + +The ABCPolyBase class provides the methods needed to implement the common API +for the various polynomial classes. It operates as a mixin, but uses the +abc module from the stdlib, hence it is only available for Python >= 2.6. + +""" +import os +import abc +import numbers +from typing import Callable + +import numpy as np +from . import polyutils as pu + +__all__ = ['ABCPolyBase'] + +class ABCPolyBase(abc.ABC): + """An abstract base class for immutable series classes. + + ABCPolyBase provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' along with the + methods listed below. + + Parameters + ---------- + coef : array_like + Series coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*P_0(x) + 2*P_1(x) + 3*P_2(x)``, where + ``P_i`` is the basis polynomials of degree ``i``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is the derived class domain. + window : (2,) array_like, optional + Window, see domain for its use. The default value is the + derived class window. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + Attributes + ---------- + coef : (N,) ndarray + Series coefficients in order of increasing degree. + domain : (2,) ndarray + Domain that is mapped to window. + window : (2,) ndarray + Window that domain is mapped to. + symbol : str + Symbol representing the independent variable. + + Class Attributes + ---------------- + maxpower : int + Maximum power allowed, i.e., the largest number ``n`` such that + ``p(x)**n`` is allowed. This is to limit runaway polynomial size. + domain : (2,) ndarray + Default domain of the class. + window : (2,) ndarray + Default window of the class. + + """ + + # Not hashable + __hash__ = None + + # Opt out of numpy ufuncs and Python ops with ndarray subclasses. + __array_ufunc__ = None + + # Limit runaway size. T_n^m has degree n*m + maxpower = 100 + + # Unicode character mappings for improved __str__ + _superscript_mapping = str.maketrans({ + "0": "⁰", + "1": "¹", + "2": "²", + "3": "³", + "4": "⁴", + "5": "⁵", + "6": "⁶", + "7": "⁷", + "8": "⁸", + "9": "⁹" + }) + _subscript_mapping = str.maketrans({ + "0": "₀", + "1": "₁", + "2": "₂", + "3": "₃", + "4": "₄", + "5": "₅", + "6": "₆", + "7": "₇", + "8": "₈", + "9": "₉" + }) + # Some fonts don't support full unicode character ranges necessary for + # the full set of superscripts and subscripts, including common/default + # fonts in Windows shells/terminals. Therefore, default to ascii-only + # printing on windows. + _use_unicode = not os.name == 'nt' + + @property + def symbol(self): + return self._symbol + + @property + @abc.abstractmethod + def domain(self): + pass + + @property + @abc.abstractmethod + def window(self): + pass + + @property + @abc.abstractmethod + def basis_name(self): + pass + + @staticmethod + @abc.abstractmethod + def _add(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _sub(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _mul(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _div(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _pow(c, pow, maxpower=None): + pass + + @staticmethod + @abc.abstractmethod + def _val(x, c): + pass + + @staticmethod + @abc.abstractmethod + def _int(c, m, k, lbnd, scl): + pass + + @staticmethod + @abc.abstractmethod + def _der(c, m, scl): + pass + + @staticmethod + @abc.abstractmethod + def _fit(x, y, deg, rcond, full): + pass + + @staticmethod + @abc.abstractmethod + def _line(off, scl): + pass + + @staticmethod + @abc.abstractmethod + def _roots(c): + pass + + @staticmethod + @abc.abstractmethod + def _fromroots(r): + pass + + def has_samecoef(self, other): + """Check if coefficients match. + + Parameters + ---------- + other : class instance + The other class must have the ``coef`` attribute. + + Returns + ------- + bool : boolean + True if the coefficients are the same, False otherwise. + + """ + if len(self.coef) != len(other.coef): + return False + elif not np.all(self.coef == other.coef): + return False + else: + return True + + def has_samedomain(self, other): + """Check if domains match. + + Parameters + ---------- + other : class instance + The other class must have the ``domain`` attribute. + + Returns + ------- + bool : boolean + True if the domains are the same, False otherwise. + + """ + return np.all(self.domain == other.domain) + + def has_samewindow(self, other): + """Check if windows match. + + Parameters + ---------- + other : class instance + The other class must have the ``window`` attribute. + + Returns + ------- + bool : boolean + True if the windows are the same, False otherwise. + + """ + return np.all(self.window == other.window) + + def has_sametype(self, other): + """Check if types match. + + Parameters + ---------- + other : object + Class instance. + + Returns + ------- + bool : boolean + True if other is same class as self + + """ + return isinstance(other, self.__class__) + + def _get_coefficients(self, other): + """Interpret other as polynomial coefficients. + + The `other` argument is checked to see if it is of the same + class as self with identical domain and window. If so, + return its coefficients, otherwise return `other`. + + Parameters + ---------- + other : anything + Object to be checked. + + Returns + ------- + coef + The coefficients of`other` if it is a compatible instance, + of ABCPolyBase, otherwise `other`. + + Raises + ------ + TypeError + When `other` is an incompatible instance of ABCPolyBase. + + """ + if isinstance(other, ABCPolyBase): + if not isinstance(other, self.__class__): + raise TypeError("Polynomial types differ") + elif not np.all(self.domain == other.domain): + raise TypeError("Domains differ") + elif not np.all(self.window == other.window): + raise TypeError("Windows differ") + elif self.symbol != other.symbol: + raise ValueError("Polynomial symbols differ") + return other.coef + return other + + def __init__(self, coef, domain=None, window=None, symbol='x'): + [coef] = pu.as_series([coef], trim=False) + self.coef = coef + + if domain is not None: + [domain] = pu.as_series([domain], trim=False) + if len(domain) != 2: + raise ValueError("Domain has wrong number of elements.") + self.domain = domain + + if window is not None: + [window] = pu.as_series([window], trim=False) + if len(window) != 2: + raise ValueError("Window has wrong number of elements.") + self.window = window + + # Validation for symbol + try: + if not symbol.isidentifier(): + raise ValueError( + "Symbol string must be a valid Python identifier" + ) + # If a user passes in something other than a string, the above + # results in an AttributeError. Catch this and raise a more + # informative exception + except AttributeError: + raise TypeError("Symbol must be a non-empty string") + + self._symbol = symbol + + def __repr__(self): + coef = repr(self.coef)[6:-1] + domain = repr(self.domain)[6:-1] + window = repr(self.window)[6:-1] + name = self.__class__.__name__ + return (f"{name}({coef}, domain={domain}, window={window}, " + f"symbol='{self.symbol}')") + + def __format__(self, fmt_str): + if fmt_str == '': + return self.__str__() + if fmt_str not in ('ascii', 'unicode'): + raise ValueError( + f"Unsupported format string '{fmt_str}' passed to " + f"{self.__class__}.__format__. Valid options are " + f"'ascii' and 'unicode'" + ) + if fmt_str == 'ascii': + return self._generate_string(self._str_term_ascii) + return self._generate_string(self._str_term_unicode) + + def __str__(self): + if self._use_unicode: + return self._generate_string(self._str_term_unicode) + return self._generate_string(self._str_term_ascii) + + def _generate_string(self, term_method): + """ + Generate the full string representation of the polynomial, using + ``term_method`` to generate each polynomial term. + """ + # Get configuration for line breaks + linewidth = np.get_printoptions().get('linewidth', 75) + if linewidth < 1: + linewidth = 1 + out = pu.format_float(self.coef[0]) + + off, scale = self.mapparms() + + scaled_symbol, needs_parens = self._format_term(pu.format_float, + off, scale) + if needs_parens: + scaled_symbol = '(' + scaled_symbol + ')' + + for i, coef in enumerate(self.coef[1:]): + out += " " + power = str(i + 1) + # Polynomial coefficient + # The coefficient array can be an object array with elements that + # will raise a TypeError with >= 0 (e.g. strings or Python + # complex). In this case, represent the coefficient as-is. + try: + if coef >= 0: + next_term = "+ " + pu.format_float(coef, parens=True) + else: + next_term = "- " + pu.format_float(-coef, parens=True) + except TypeError: + next_term = f"+ {coef}" + # Polynomial term + next_term += term_method(power, scaled_symbol) + # Length of the current line with next term added + line_len = len(out.split('\n')[-1]) + len(next_term) + # If not the last term in the polynomial, it will be two + # characters longer due to the +/- with the next term + if i < len(self.coef[1:]) - 1: + line_len += 2 + # Handle linebreaking + if line_len >= linewidth: + next_term = next_term.replace(" ", "\n", 1) + out += next_term + return out + + @classmethod + def _str_term_unicode(cls, i, arg_str): + """ + String representation of single polynomial term using unicode + characters for superscripts and subscripts. + """ + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis_name, or override " + "_str_term_unicode(cls, i, arg_str)" + ) + return (f"·{cls.basis_name}{i.translate(cls._subscript_mapping)}" + f"({arg_str})") + + @classmethod + def _str_term_ascii(cls, i, arg_str): + """ + String representation of a single polynomial term using ** and _ to + represent superscripts and subscripts, respectively. + """ + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis_name, or override " + "_str_term_ascii(cls, i, arg_str)" + ) + return f" {cls.basis_name}_{i}({arg_str})" + + @classmethod + def _repr_latex_term(cls, i, arg_str, needs_parens): + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis name, or override " + "_repr_latex_term(i, arg_str, needs_parens)") + # since we always add parens, we don't care if the expression needs them + return f"{{{cls.basis_name}}}_{{{i}}}({arg_str})" + + @staticmethod + def _repr_latex_scalar(x, parens=False): + # TODO: we're stuck with disabling math formatting until we handle + # exponents in this function + return r'\text{{{}}}'.format(pu.format_float(x, parens=parens)) + + def _format_term(self, scalar_format: Callable, off: float, scale: float): + """ Format a single term in the expansion """ + if off == 0 and scale == 1: + term = self.symbol + needs_parens = False + elif scale == 1: + term = f"{scalar_format(off)} + {self.symbol}" + needs_parens = True + elif off == 0: + term = f"{scalar_format(scale)}{self.symbol}" + needs_parens = True + else: + term = ( + f"{scalar_format(off)} + " + f"{scalar_format(scale)}{self.symbol}" + ) + needs_parens = True + return term, needs_parens + + def _repr_latex_(self): + # get the scaled argument string to the basis functions + off, scale = self.mapparms() + term, needs_parens = self._format_term(self._repr_latex_scalar, + off, scale) + + mute = r"\color{{LightGray}}{{{}}}".format + + parts = [] + for i, c in enumerate(self.coef): + # prevent duplication of + and - signs + if i == 0: + coef_str = f"{self._repr_latex_scalar(c)}" + elif not isinstance(c, numbers.Real): + coef_str = f" + ({self._repr_latex_scalar(c)})" + elif c >= 0: + coef_str = f" + {self._repr_latex_scalar(c, parens=True)}" + else: + coef_str = f" - {self._repr_latex_scalar(-c, parens=True)}" + + # produce the string for the term + term_str = self._repr_latex_term(i, term, needs_parens) + if term_str == '1': + part = coef_str + else: + part = rf"{coef_str}\,{term_str}" + + if c == 0: + part = mute(part) + + parts.append(part) + + if parts: + body = ''.join(parts) + else: + # in case somehow there are no coefficients at all + body = '0' + + return rf"${self.symbol} \mapsto {body}$" + + + + # Pickle and copy + + def __getstate__(self): + ret = self.__dict__.copy() + ret['coef'] = self.coef.copy() + ret['domain'] = self.domain.copy() + ret['window'] = self.window.copy() + ret['symbol'] = self.symbol + return ret + + def __setstate__(self, dict): + self.__dict__ = dict + + # Call + + def __call__(self, arg): + arg = pu.mapdomain(arg, self.domain, self.window) + return self._val(arg, self.coef) + + def __iter__(self): + return iter(self.coef) + + def __len__(self): + return len(self.coef) + + # Numeric properties. + + def __neg__(self): + return self.__class__( + -self.coef, self.domain, self.window, self.symbol + ) + + def __pos__(self): + return self + + def __add__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._add(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __sub__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._sub(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __mul__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._mul(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __truediv__(self, other): + # there is no true divide if the rhs is not a Number, although it + # could return the first n elements of an infinite series. + # It is hard to see where n would come from, though. + if not isinstance(other, numbers.Number) or isinstance(other, bool): + raise TypeError( + f"unsupported types for true division: " + f"'{type(self)}', '{type(other)}'" + ) + return self.__floordiv__(other) + + def __floordiv__(self, other): + res = self.__divmod__(other) + if res is NotImplemented: + return res + return res[0] + + def __mod__(self, other): + res = self.__divmod__(other) + if res is NotImplemented: + return res + return res[1] + + def __divmod__(self, other): + othercoef = self._get_coefficients(other) + try: + quo, rem = self._div(self.coef, othercoef) + except ZeroDivisionError: + raise + except Exception: + return NotImplemented + quo = self.__class__(quo, self.domain, self.window, self.symbol) + rem = self.__class__(rem, self.domain, self.window, self.symbol) + return quo, rem + + def __pow__(self, other): + coef = self._pow(self.coef, other, maxpower=self.maxpower) + res = self.__class__(coef, self.domain, self.window, self.symbol) + return res + + def __radd__(self, other): + try: + coef = self._add(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rsub__(self, other): + try: + coef = self._sub(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rmul__(self, other): + try: + coef = self._mul(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rdiv__(self, other): + # set to __floordiv__ /. + return self.__rfloordiv__(other) + + def __rtruediv__(self, other): + # An instance of ABCPolyBase is not considered a + # Number. + return NotImplemented + + def __rfloordiv__(self, other): + res = self.__rdivmod__(other) + if res is NotImplemented: + return res + return res[0] + + def __rmod__(self, other): + res = self.__rdivmod__(other) + if res is NotImplemented: + return res + return res[1] + + def __rdivmod__(self, other): + try: + quo, rem = self._div(other, self.coef) + except ZeroDivisionError: + raise + except Exception: + return NotImplemented + quo = self.__class__(quo, self.domain, self.window, self.symbol) + rem = self.__class__(rem, self.domain, self.window, self.symbol) + return quo, rem + + def __eq__(self, other): + res = (isinstance(other, self.__class__) and + np.all(self.domain == other.domain) and + np.all(self.window == other.window) and + (self.coef.shape == other.coef.shape) and + np.all(self.coef == other.coef) and + (self.symbol == other.symbol)) + return res + + def __ne__(self, other): + return not self.__eq__(other) + + # + # Extra methods. + # + + def copy(self): + """Return a copy. + + Returns + ------- + new_series : series + Copy of self. + + """ + return self.__class__(self.coef, self.domain, self.window, self.symbol) + + def degree(self): + """The degree of the series. + + Returns + ------- + degree : int + Degree of the series, one less than the number of coefficients. + + Examples + -------- + + Create a polynomial object for ``1 + 7*x + 4*x**2``: + + >>> poly = np.polynomial.Polynomial([1, 7, 4]) + >>> print(poly) + 1.0 + 7.0·x + 4.0·x² + >>> poly.degree() + 2 + + Note that this method does not check for non-zero coefficients. + You must trim the polynomial to remove any trailing zeroes: + + >>> poly = np.polynomial.Polynomial([1, 7, 0]) + >>> print(poly) + 1.0 + 7.0·x + 0.0·x² + >>> poly.degree() + 2 + >>> poly.trim().degree() + 1 + + """ + return len(self) - 1 + + def cutdeg(self, deg): + """Truncate series to the given degree. + + Reduce the degree of the series to `deg` by discarding the + high order terms. If `deg` is greater than the current degree a + copy of the current series is returned. This can be useful in least + squares where the coefficients of the high degree terms may be very + small. + + Parameters + ---------- + deg : non-negative int + The series is reduced to degree `deg` by discarding the high + order terms. The value of `deg` must be a non-negative integer. + + Returns + ------- + new_series : series + New instance of series with reduced degree. + + """ + return self.truncate(deg + 1) + + def trim(self, tol=0): + """Remove trailing coefficients + + Remove trailing coefficients until a coefficient is reached whose + absolute value greater than `tol` or the beginning of the series is + reached. If all the coefficients would be removed the series is set + to ``[0]``. A new series instance is returned with the new + coefficients. The current instance remains unchanged. + + Parameters + ---------- + tol : non-negative number. + All trailing coefficients less than `tol` will be removed. + + Returns + ------- + new_series : series + New instance of series with trimmed coefficients. + + """ + coef = pu.trimcoef(self.coef, tol) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def truncate(self, size): + """Truncate series to length `size`. + + Reduce the series to length `size` by discarding the high + degree terms. The value of `size` must be a positive integer. This + can be useful in least squares where the coefficients of the + high degree terms may be very small. + + Parameters + ---------- + size : positive int + The series is reduced to length `size` by discarding the high + degree terms. The value of `size` must be a positive integer. + + Returns + ------- + new_series : series + New instance of series with truncated coefficients. + + """ + isize = int(size) + if isize != size or isize < 1: + raise ValueError("size must be a positive integer") + if isize >= len(self.coef): + coef = self.coef + else: + coef = self.coef[:isize] + return self.__class__(coef, self.domain, self.window, self.symbol) + + def convert(self, domain=None, kind=None, window=None): + """Convert series to a different kind and/or domain and/or window. + + Parameters + ---------- + domain : array_like, optional + The domain of the converted series. If the value is None, + the default domain of `kind` is used. + kind : class, optional + The polynomial series type class to which the current instance + should be converted. If kind is None, then the class of the + current instance is used. + window : array_like, optional + The window of the converted series. If the value is None, + the default window of `kind` is used. + + Returns + ------- + new_series : series + The returned class can be of different type than the current + instance and/or have a different domain and/or different + window. + + Notes + ----- + Conversion between domains and class types can result in + numerically ill defined series. + + """ + if kind is None: + kind = self.__class__ + if domain is None: + domain = kind.domain + if window is None: + window = kind.window + return self(kind.identity(domain, window=window, symbol=self.symbol)) + + def mapparms(self): + """Return the mapping parameters. + + The returned values define a linear map ``off + scl*x`` that is + applied to the input arguments before the series is evaluated. The + map depends on the ``domain`` and ``window``; if the current + ``domain`` is equal to the ``window`` the resulting map is the + identity. If the coefficients of the series instance are to be + used by themselves outside this class, then the linear function + must be substituted for the ``x`` in the standard representation of + the base polynomials. + + Returns + ------- + off, scl : float or complex + The mapping function is defined by ``off + scl*x``. + + Notes + ----- + If the current domain is the interval ``[l1, r1]`` and the window + is ``[l2, r2]``, then the linear mapping function ``L`` is + defined by the equations:: + + L(l1) = l2 + L(r1) = r2 + + """ + return pu.mapparms(self.domain, self.window) + + def integ(self, m=1, k=[], lbnd=None): + """Integrate. + + Return a series instance that is the definite integral of the + current series. + + Parameters + ---------- + m : non-negative int + The number of integrations to perform. + k : array_like + Integration constants. The first constant is applied to the + first integration, the second to the second, and so on. The + list of values must less than or equal to `m` in length and any + missing values are set to zero. + lbnd : Scalar + The lower bound of the definite integral. + + Returns + ------- + new_series : series + A new series representing the integral. The domain is the same + as the domain of the integrated series. + + """ + off, scl = self.mapparms() + if lbnd is None: + lbnd = 0 + else: + lbnd = off + scl*lbnd + coef = self._int(self.coef, m, k, lbnd, 1./scl) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def deriv(self, m=1): + """Differentiate. + + Return a series instance of that is the derivative of the current + series. + + Parameters + ---------- + m : non-negative int + Find the derivative of order `m`. + + Returns + ------- + new_series : series + A new series representing the derivative. The domain is the same + as the domain of the differentiated series. + + """ + off, scl = self.mapparms() + coef = self._der(self.coef, m, scl) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def roots(self): + """Return the roots of the series polynomial. + + Compute the roots for the series. Note that the accuracy of the + roots decreases the further outside the `domain` they lie. + + Returns + ------- + roots : ndarray + Array containing the roots of the series. + + """ + roots = self._roots(self.coef) + return pu.mapdomain(roots, self.window, self.domain) + + def linspace(self, n=100, domain=None): + """Return x, y values at equally spaced points in domain. + + Returns the x, y values at `n` linearly spaced points across the + domain. Here y is the value of the polynomial at the points x. By + default the domain is the same as that of the series instance. + This method is intended mostly as a plotting aid. + + Parameters + ---------- + n : int, optional + Number of point pairs to return. The default value is 100. + domain : {None, array_like}, optional + If not None, the specified domain is used instead of that of + the calling instance. It should be of the form ``[beg,end]``. + The default is None which case the class domain is used. + + Returns + ------- + x, y : ndarray + x is equal to linspace(self.domain[0], self.domain[1], n) and + y is the series evaluated at element of x. + + """ + if domain is None: + domain = self.domain + x = np.linspace(domain[0], domain[1], n) + y = self(x) + return x, y + + @classmethod + def fit(cls, x, y, deg, domain=None, rcond=None, full=False, w=None, + window=None, symbol='x'): + """Least squares fit to data. + + Return a series instance that is the least squares fit to the data + `y` sampled at `x`. The domain of the returned instance can be + specified and this will often result in a superior fit with less + chance of ill conditioning. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) + y-coordinates of the M sample points ``(x[i], y[i])``. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + domain : {None, [beg, end], []}, optional + Domain to use for the returned series. If ``None``, + then a minimal domain that covers the points `x` is chosen. If + ``[]`` the class domain is used. The default value was the + class domain in NumPy 1.4 and ``None`` in later versions. + The ``[]`` option was added in numpy 1.5.0. + rcond : float, optional + Relative condition number of the fit. Singular values smaller + than this relative to the largest singular value will be + ignored. The default value is ``len(x)*eps``, where eps is the + relative precision of the float type, about 2e-16 in most + cases. + full : bool, optional + Switch determining nature of return value. When it is False + (the default) just the coefficients are returned, when True + diagnostic information from the singular value decomposition is + also returned. + w : array_like, shape (M,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have + the same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + window : {[beg, end]}, optional + Window to use for the returned series. The default + value is the default class domain + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + A series that represents the least squares fit to the data and + has the domain and window specified in the call. If the + coefficients for the unscaled and unshifted basis polynomials are + of interest, do ``new_series.convert().coef``. + + [resid, rank, sv, rcond] : list + These values are only returned if ``full == True`` + + - resid -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - sv -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `linalg.lstsq`. + + """ + if domain is None: + domain = pu.getdomain(x) + if domain[0] == domain[1]: + domain[0] -= 1 + domain[1] += 1 + elif type(domain) is list and len(domain) == 0: + domain = cls.domain + + if window is None: + window = cls.window + + xnew = pu.mapdomain(x, domain, window) + res = cls._fit(xnew, y, deg, w=w, rcond=rcond, full=full) + if full: + [coef, status] = res + return ( + cls(coef, domain=domain, window=window, symbol=symbol), status + ) + else: + coef = res + return cls(coef, domain=domain, window=window, symbol=symbol) + + @classmethod + def fromroots(cls, roots, domain=[], window=None, symbol='x'): + """Return series instance that has the specified roots. + + Returns a series representing the product + ``(x - r[0])*(x - r[1])*...*(x - r[n-1])``, where ``r`` is a + list of roots. + + Parameters + ---------- + roots : array_like + List of roots. + domain : {[], None, array_like}, optional + Domain for the resulting series. If None the domain is the + interval from the smallest root to the largest. If [] the + domain is the class domain. The default is []. + window : {None, array_like}, optional + Window for the returned series. If None the class window is + used. The default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + Series with the specified roots. + + """ + [roots] = pu.as_series([roots], trim=False) + if domain is None: + domain = pu.getdomain(roots) + elif type(domain) is list and len(domain) == 0: + domain = cls.domain + + if window is None: + window = cls.window + + deg = len(roots) + off, scl = pu.mapparms(domain, window) + rnew = off + scl*roots + coef = cls._fromroots(rnew) / scl**deg + return cls(coef, domain=domain, window=window, symbol=symbol) + + @classmethod + def identity(cls, domain=None, window=None, symbol='x'): + """Identity function. + + If ``p`` is the returned series, then ``p(x) == x`` for all + values of x. + + Parameters + ---------- + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + Series of representing the identity. + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + off, scl = pu.mapparms(window, domain) + coef = cls._line(off, scl) + return cls(coef, domain, window, symbol) + + @classmethod + def basis(cls, deg, domain=None, window=None, symbol='x'): + """Series basis polynomial of degree `deg`. + + Returns the series representing the basis polynomial of degree `deg`. + + Parameters + ---------- + deg : int + Degree of the basis polynomial for the series. Must be >= 0. + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + A series with the coefficient of the `deg` term set to one and + all others zero. + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + ideg = int(deg) + + if ideg != deg or ideg < 0: + raise ValueError("deg must be non-negative integer") + return cls([0]*ideg + [1], domain, window, symbol) + + @classmethod + def cast(cls, series, domain=None, window=None): + """Convert series to series of this class. + + The `series` is expected to be an instance of some polynomial + series of one of the types supported by by the numpy.polynomial + module, but could be some other class that supports the convert + method. + + Parameters + ---------- + series : series + The series instance to be converted. + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + + Returns + ------- + new_series : series + A series of the same kind as the calling class and equal to + `series` when evaluated. + + See Also + -------- + convert : similar instance method + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + return series.convert(domain, cls, window) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.pyi new file mode 100644 index 0000000000000000000000000000000000000000..ca7ca628d5140c7584ef42a92fb633625ca8a657 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/_polybase.pyi @@ -0,0 +1,287 @@ +import abc +import decimal +import numbers +from collections.abc import Iterator, Mapping, Sequence +from typing import ( + Any, + ClassVar, + Final, + Generic, + Literal, + SupportsIndex, + TypeAlias, + TypeGuard, + overload, +) + +import numpy as np +import numpy.typing as npt +from numpy._typing import ( + _FloatLike_co, + _NumberLike_co, + + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, +) + +from ._polytypes import ( + _AnyInt, + _CoefLike_co, + + _Array2, + _Tuple2, + + _Series, + _CoefSeries, + + _SeriesLikeInt_co, + _SeriesLikeCoef_co, + + _ArrayLikeCoefObject_co, + _ArrayLikeCoef_co, +) + +from typing_extensions import LiteralString, TypeVar + + +__all__: Final[Sequence[str]] = ("ABCPolyBase",) + + +_NameCo = TypeVar("_NameCo", bound=LiteralString | None, covariant=True, default=LiteralString | None) +_Self = TypeVar("_Self") +_Other = TypeVar("_Other", bound=ABCPolyBase) + +_AnyOther: TypeAlias = ABCPolyBase | _CoefLike_co | _SeriesLikeCoef_co +_Hundred: TypeAlias = Literal[100] + + +class ABCPolyBase(Generic[_NameCo], metaclass=abc.ABCMeta): + __hash__: ClassVar[None] # type: ignore[assignment] + __array_ufunc__: ClassVar[None] + + maxpower: ClassVar[_Hundred] + _superscript_mapping: ClassVar[Mapping[int, str]] + _subscript_mapping: ClassVar[Mapping[int, str]] + _use_unicode: ClassVar[bool] + + basis_name: _NameCo + coef: _CoefSeries + domain: _Array2[np.inexact[Any] | np.object_] + window: _Array2[np.inexact[Any] | np.object_] + + _symbol: LiteralString + @property + def symbol(self, /) -> LiteralString: ... + + def __init__( + self, + /, + coef: _SeriesLikeCoef_co, + domain: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> None: ... + + @overload + def __call__(self, /, arg: _Other) -> _Other: ... + # TODO: Once `_ShapeType@ndarray` is covariant and bounded (see #26081), + # additionally include 0-d arrays as input types with scalar return type. + @overload + def __call__( + self, + /, + arg: _FloatLike_co | decimal.Decimal | numbers.Real | np.object_, + ) -> np.float64 | np.complex128: ... + @overload + def __call__( + self, + /, + arg: _NumberLike_co | numbers.Complex, + ) -> np.complex128: ... + @overload + def __call__(self, /, arg: _ArrayLikeFloat_co) -> ( + npt.NDArray[np.float64] + | npt.NDArray[np.complex128] + | npt.NDArray[np.object_] + ): ... + @overload + def __call__( + self, + /, + arg: _ArrayLikeComplex_co, + ) -> npt.NDArray[np.complex128] | npt.NDArray[np.object_]: ... + @overload + def __call__( + self, + /, + arg: _ArrayLikeCoefObject_co, + ) -> npt.NDArray[np.object_]: ... + + def __format__(self, fmt_str: str, /) -> str: ... + def __eq__(self, x: object, /) -> bool: ... + def __ne__(self, x: object, /) -> bool: ... + def __neg__(self: _Self, /) -> _Self: ... + def __pos__(self: _Self, /) -> _Self: ... + def __add__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __sub__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __mul__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __truediv__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __floordiv__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __mod__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __divmod__(self: _Self, x: _AnyOther, /) -> _Tuple2[_Self]: ... + def __pow__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __radd__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rsub__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rmul__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rtruediv__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rfloordiv__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rmod__(self: _Self, x: _AnyOther, /) -> _Self: ... + def __rdivmod__(self: _Self, x: _AnyOther, /) -> _Tuple2[_Self]: ... + def __len__(self, /) -> int: ... + def __iter__(self, /) -> Iterator[np.inexact[Any] | object]: ... + def __getstate__(self, /) -> dict[str, Any]: ... + def __setstate__(self, dict: dict[str, Any], /) -> None: ... + + def has_samecoef(self, /, other: ABCPolyBase) -> bool: ... + def has_samedomain(self, /, other: ABCPolyBase) -> bool: ... + def has_samewindow(self, /, other: ABCPolyBase) -> bool: ... + @overload + def has_sametype(self: _Self, /, other: ABCPolyBase) -> TypeGuard[_Self]: ... + @overload + def has_sametype(self, /, other: object) -> Literal[False]: ... + + def copy(self: _Self, /) -> _Self: ... + def degree(self, /) -> int: ... + def cutdeg(self: _Self, /) -> _Self: ... + def trim(self: _Self, /, tol: _FloatLike_co = ...) -> _Self: ... + def truncate(self: _Self, /, size: _AnyInt) -> _Self: ... + + @overload + def convert( + self, + domain: None | _SeriesLikeCoef_co, + kind: type[_Other], + /, + window: None | _SeriesLikeCoef_co = ..., + ) -> _Other: ... + @overload + def convert( + self, + /, + domain: None | _SeriesLikeCoef_co = ..., + *, + kind: type[_Other], + window: None | _SeriesLikeCoef_co = ..., + ) -> _Other: ... + @overload + def convert( + self: _Self, + /, + domain: None | _SeriesLikeCoef_co = ..., + kind: None | type[_Self] = ..., + window: None | _SeriesLikeCoef_co = ..., + ) -> _Self: ... + + def mapparms(self, /) -> _Tuple2[Any]: ... + + def integ( + self: _Self, /, + m: SupportsIndex = ..., + k: _CoefLike_co | _SeriesLikeCoef_co = ..., + lbnd: None | _CoefLike_co = ..., + ) -> _Self: ... + + def deriv(self: _Self, /, m: SupportsIndex = ...) -> _Self: ... + + def roots(self, /) -> _CoefSeries: ... + + def linspace( + self, /, + n: SupportsIndex = ..., + domain: None | _SeriesLikeCoef_co = ..., + ) -> _Tuple2[_Series[np.float64 | np.complex128]]: ... + + @overload + @classmethod + def fit( + cls: type[_Self], /, + x: _SeriesLikeCoef_co, + y: _SeriesLikeCoef_co, + deg: int | _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co = ..., + rcond: _FloatLike_co = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> _Self: ... + @overload + @classmethod + def fit( + cls: type[_Self], /, + x: _SeriesLikeCoef_co, + y: _SeriesLikeCoef_co, + deg: int | _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co = ..., + rcond: _FloatLike_co = ..., + *, + full: Literal[True], + w: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> tuple[_Self, Sequence[np.inexact[Any] | np.int32]]: ... + @overload + @classmethod + def fit( + cls: type[_Self], + x: _SeriesLikeCoef_co, + y: _SeriesLikeCoef_co, + deg: int | _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co, + rcond: _FloatLike_co, + full: Literal[True], /, + w: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> tuple[_Self, Sequence[np.inexact[Any] | np.int32]]: ... + + @classmethod + def fromroots( + cls: type[_Self], /, + roots: _ArrayLikeCoef_co, + domain: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> _Self: ... + + @classmethod + def identity( + cls: type[_Self], /, + domain: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> _Self: ... + + @classmethod + def basis( + cls: type[_Self], /, + deg: _AnyInt, + domain: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + symbol: str = ..., + ) -> _Self: ... + + @classmethod + def cast( + cls: type[_Self], /, + series: ABCPolyBase, + domain: None | _SeriesLikeCoef_co = ..., + window: None | _SeriesLikeCoef_co = ..., + ) -> _Self: ... + + @classmethod + def _str_term_unicode(cls, /, i: str, arg_str: str) -> str: ... + @staticmethod + def _str_term_ascii(i: str, arg_str: str) -> str: ... + @staticmethod + def _repr_latex_term(i: str, arg_str: str, needs_parens: bool) -> str: ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/_polytypes.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/_polytypes.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b0794eb61831d396c339d54f34dc43c1554657b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/_polytypes.pyi @@ -0,0 +1,888 @@ +from collections.abc import Callable, Sequence +from typing import ( + Any, + Literal, + NoReturn, + Protocol, + SupportsIndex, + SupportsInt, + TypeAlias, + overload, + type_check_only, +) + +import numpy as np +import numpy.typing as npt +from numpy._typing import ( + # array-likes + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeNumber_co, + _ArrayLikeObject_co, + _NestedSequence, + _SupportsArray, + + # scalar-likes + _IntLike_co, + _FloatLike_co, + _ComplexLike_co, + _NumberLike_co, +) + +from typing_extensions import LiteralString, TypeVar + + +_T = TypeVar("_T") +_T_contra = TypeVar("_T_contra", contravariant=True) +_Self = TypeVar("_Self") +_SCT = TypeVar("_SCT", bound=np.number[Any] | np.bool | np.object_) + +# compatible with e.g. int, float, complex, Decimal, Fraction, and ABCPolyBase +@type_check_only +class _SupportsCoefOps(Protocol[_T_contra]): + def __eq__(self, x: object, /) -> bool: ... + def __ne__(self, x: object, /) -> bool: ... + + def __neg__(self: _Self, /) -> _Self: ... + def __pos__(self: _Self, /) -> _Self: ... + + def __add__(self: _Self, x: _T_contra, /) -> _Self: ... + def __sub__(self: _Self, x: _T_contra, /) -> _Self: ... + def __mul__(self: _Self, x: _T_contra, /) -> _Self: ... + def __pow__(self: _Self, x: _T_contra, /) -> _Self | float: ... + + def __radd__(self: _Self, x: _T_contra, /) -> _Self: ... + def __rsub__(self: _Self, x: _T_contra, /) -> _Self: ... + def __rmul__(self: _Self, x: _T_contra, /) -> _Self: ... + +_Series: TypeAlias = np.ndarray[tuple[int], np.dtype[_SCT]] + +_FloatSeries: TypeAlias = _Series[np.floating[Any]] +_ComplexSeries: TypeAlias = _Series[np.complexfloating[Any, Any]] +_ObjectSeries: TypeAlias = _Series[np.object_] +_CoefSeries: TypeAlias = _Series[np.inexact[Any] | np.object_] + +_FloatArray: TypeAlias = npt.NDArray[np.floating[Any]] +_ComplexArray: TypeAlias = npt.NDArray[np.complexfloating[Any, Any]] +_ObjectArray: TypeAlias = npt.NDArray[np.object_] +_CoefArray: TypeAlias = npt.NDArray[np.inexact[Any] | np.object_] + +_Tuple2: TypeAlias = tuple[_T, _T] +_Array1: TypeAlias = np.ndarray[tuple[Literal[1]], np.dtype[_SCT]] +_Array2: TypeAlias = np.ndarray[tuple[Literal[2]], np.dtype[_SCT]] + +_AnyInt: TypeAlias = SupportsInt | SupportsIndex + +_CoefObjectLike_co: TypeAlias = np.object_ | _SupportsCoefOps[Any] +_CoefLike_co: TypeAlias = _NumberLike_co | _CoefObjectLike_co + +# The term "series" is used here to refer to 1-d arrays of numeric scalars. +_SeriesLikeBool_co: TypeAlias = ( + _SupportsArray[np.dtype[np.bool]] + | Sequence[bool | np.bool] +) +_SeriesLikeInt_co: TypeAlias = ( + _SupportsArray[np.dtype[np.integer[Any] | np.bool]] + | Sequence[_IntLike_co] +) +_SeriesLikeFloat_co: TypeAlias = ( + _SupportsArray[np.dtype[np.floating[Any] | np.integer[Any] | np.bool]] + | Sequence[_FloatLike_co] +) +_SeriesLikeComplex_co: TypeAlias = ( + _SupportsArray[np.dtype[np.inexact[Any] | np.integer[Any] | np.bool]] + | Sequence[_ComplexLike_co] +) +_SeriesLikeObject_co: TypeAlias = ( + _SupportsArray[np.dtype[np.object_]] + | Sequence[_CoefObjectLike_co] +) +_SeriesLikeCoef_co: TypeAlias = ( + _SupportsArray[np.dtype[np.number[Any] | np.bool | np.object_]] + | Sequence[_CoefLike_co] +) + +_ArrayLikeCoefObject_co: TypeAlias = ( + _CoefObjectLike_co + | _SeriesLikeObject_co + | _NestedSequence[_SeriesLikeObject_co] +) +_ArrayLikeCoef_co: TypeAlias = ( + npt.NDArray[np.number[Any] | np.bool | np.object_] + | _ArrayLikeNumber_co + | _ArrayLikeCoefObject_co +) + +_Name_co = TypeVar("_Name_co", bound=LiteralString, covariant=True, default=LiteralString) + +@type_check_only +class _Named(Protocol[_Name_co]): + @property + def __name__(self, /) -> _Name_co: ... + +_Line: TypeAlias = np.ndarray[tuple[Literal[1, 2]], np.dtype[_SCT]] + +@type_check_only +class _FuncLine(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__(self, /, off: _SCT, scl: _SCT) -> _Line[_SCT]: ... + @overload + def __call__(self, /, off: int, scl: int) -> _Line[np.int_] : ... + @overload + def __call__(self, /, off: float, scl: float) -> _Line[np.float64]: ... + @overload + def __call__( + self, + /, + off: complex, + scl: complex, + ) -> _Line[np.complex128]: ... + @overload + def __call__( + self, + /, + off: _SupportsCoefOps[Any], + scl: _SupportsCoefOps[Any], + ) -> _Line[np.object_]: ... + +@type_check_only +class _FuncFromRoots(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__(self, /, roots: _SeriesLikeFloat_co) -> _FloatSeries: ... + @overload + def __call__(self, /, roots: _SeriesLikeComplex_co) -> _ComplexSeries: ... + @overload + def __call__(self, /, roots: _SeriesLikeCoef_co) -> _ObjectSeries: ... + +@type_check_only +class _FuncBinOp(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c1: _SeriesLikeBool_co, + c2: _SeriesLikeBool_co, + ) -> NoReturn: ... + @overload + def __call__( + self, + /, + c1: _SeriesLikeFloat_co, + c2: _SeriesLikeFloat_co, + ) -> _FloatSeries: ... + @overload + def __call__( + self, + /, + c1: _SeriesLikeComplex_co, + c2: _SeriesLikeComplex_co, + ) -> _ComplexSeries: ... + @overload + def __call__( + self, + /, + c1: _SeriesLikeCoef_co, + c2: _SeriesLikeCoef_co, + ) -> _ObjectSeries: ... + +@type_check_only +class _FuncUnOp(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__(self, /, c: _SeriesLikeFloat_co) -> _FloatSeries: ... + @overload + def __call__(self, /, c: _SeriesLikeComplex_co) -> _ComplexSeries: ... + @overload + def __call__(self, /, c: _SeriesLikeCoef_co) -> _ObjectSeries: ... + +@type_check_only +class _FuncPoly2Ortho(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__(self, /, pol: _SeriesLikeFloat_co) -> _FloatSeries: ... + @overload + def __call__(self, /, pol: _SeriesLikeComplex_co) -> _ComplexSeries: ... + @overload + def __call__(self, /, pol: _SeriesLikeCoef_co) -> _ObjectSeries: ... + +@type_check_only +class _FuncPow(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _SeriesLikeFloat_co, + pow: _IntLike_co, + maxpower: None | _IntLike_co = ..., + ) -> _FloatSeries: ... + @overload + def __call__( + self, + /, + c: _SeriesLikeComplex_co, + pow: _IntLike_co, + maxpower: None | _IntLike_co = ..., + ) -> _ComplexSeries: ... + @overload + def __call__( + self, + /, + c: _SeriesLikeCoef_co, + pow: _IntLike_co, + maxpower: None | _IntLike_co = ..., + ) -> _ObjectSeries: ... + +@type_check_only +class _FuncDer(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _ArrayLikeFloat_co, + m: SupportsIndex = ..., + scl: _FloatLike_co = ..., + axis: SupportsIndex = ..., + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + c: _ArrayLikeComplex_co, + m: SupportsIndex = ..., + scl: _ComplexLike_co = ..., + axis: SupportsIndex = ..., + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + c: _ArrayLikeCoef_co, + m: SupportsIndex = ..., + scl: _CoefLike_co = ..., + axis: SupportsIndex = ..., + ) -> _ObjectArray: ... + +@type_check_only +class _FuncInteg(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _ArrayLikeFloat_co, + m: SupportsIndex = ..., + k: _FloatLike_co | _SeriesLikeFloat_co = ..., + lbnd: _FloatLike_co = ..., + scl: _FloatLike_co = ..., + axis: SupportsIndex = ..., + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + c: _ArrayLikeComplex_co, + m: SupportsIndex = ..., + k: _ComplexLike_co | _SeriesLikeComplex_co = ..., + lbnd: _ComplexLike_co = ..., + scl: _ComplexLike_co = ..., + axis: SupportsIndex = ..., + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + c: _ArrayLikeCoef_co, + m: SupportsIndex = ..., + k: _CoefLike_co | _SeriesLikeCoef_co = ..., + lbnd: _CoefLike_co = ..., + scl: _CoefLike_co = ..., + axis: SupportsIndex = ..., + ) -> _ObjectArray: ... + +@type_check_only +class _FuncValFromRoots(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _FloatLike_co, + r: _FloatLike_co, + tensor: bool = ..., + ) -> np.floating[Any]: ... + @overload + def __call__( + self, + /, + x: _NumberLike_co, + r: _NumberLike_co, + tensor: bool = ..., + ) -> np.complexfloating[Any, Any]: ... + @overload + def __call__( + self, + /, + x: _FloatLike_co | _ArrayLikeFloat_co, + r: _ArrayLikeFloat_co, + tensor: bool = ..., + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _NumberLike_co | _ArrayLikeComplex_co, + r: _ArrayLikeComplex_co, + tensor: bool = ..., + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _CoefLike_co | _ArrayLikeCoef_co, + r: _ArrayLikeCoef_co, + tensor: bool = ..., + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: _CoefLike_co, + r: _CoefLike_co, + tensor: bool = ..., + ) -> _SupportsCoefOps[Any]: ... + +@type_check_only +class _FuncVal(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _FloatLike_co, + c: _SeriesLikeFloat_co, + tensor: bool = ..., + ) -> np.floating[Any]: ... + @overload + def __call__( + self, + /, + x: _NumberLike_co, + c: _SeriesLikeComplex_co, + tensor: bool = ..., + ) -> np.complexfloating[Any, Any]: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + c: _ArrayLikeFloat_co, + tensor: bool = ..., + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + c: _ArrayLikeComplex_co, + tensor: bool = ..., + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + c: _ArrayLikeCoef_co, + tensor: bool = ..., + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: _CoefLike_co, + c: _SeriesLikeObject_co, + tensor: bool = ..., + ) -> _SupportsCoefOps[Any]: ... + +@type_check_only +class _FuncVal2D(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _FloatLike_co, + y: _FloatLike_co, + c: _SeriesLikeFloat_co, + ) -> np.floating[Any]: ... + @overload + def __call__( + self, + /, + x: _NumberLike_co, + y: _NumberLike_co, + c: _SeriesLikeComplex_co, + ) -> np.complexfloating[Any, Any]: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + c: _ArrayLikeFloat_co, + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + c: _ArrayLikeComplex_co, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + y: _ArrayLikeCoef_co, + c: _ArrayLikeCoef_co, + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: _CoefLike_co, + y: _CoefLike_co, + c: _SeriesLikeCoef_co, + ) -> _SupportsCoefOps[Any]: ... + +@type_check_only +class _FuncVal3D(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _FloatLike_co, + y: _FloatLike_co, + z: _FloatLike_co, + c: _SeriesLikeFloat_co + ) -> np.floating[Any]: ... + @overload + def __call__( + self, + /, + x: _NumberLike_co, + y: _NumberLike_co, + z: _NumberLike_co, + c: _SeriesLikeComplex_co, + ) -> np.complexfloating[Any, Any]: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + z: _ArrayLikeFloat_co, + c: _ArrayLikeFloat_co, + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + z: _ArrayLikeComplex_co, + c: _ArrayLikeComplex_co, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + y: _ArrayLikeCoef_co, + z: _ArrayLikeCoef_co, + c: _ArrayLikeCoef_co, + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: _CoefLike_co, + y: _CoefLike_co, + z: _CoefLike_co, + c: _SeriesLikeCoef_co, + ) -> _SupportsCoefOps[Any]: ... + +_AnyValF: TypeAlias = Callable[ + [npt.ArrayLike, npt.ArrayLike, bool], + _CoefArray, +] + +@type_check_only +class _FuncValND(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + val_f: _AnyValF, + c: _SeriesLikeFloat_co, + /, + *args: _FloatLike_co, + ) -> np.floating[Any]: ... + @overload + def __call__( + self, + val_f: _AnyValF, + c: _SeriesLikeComplex_co, + /, + *args: _NumberLike_co, + ) -> np.complexfloating[Any, Any]: ... + @overload + def __call__( + self, + val_f: _AnyValF, + c: _ArrayLikeFloat_co, + /, + *args: _ArrayLikeFloat_co, + ) -> _FloatArray: ... + @overload + def __call__( + self, + val_f: _AnyValF, + c: _ArrayLikeComplex_co, + /, + *args: _ArrayLikeComplex_co, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + val_f: _AnyValF, + c: _SeriesLikeObject_co, + /, + *args: _CoefObjectLike_co, + ) -> _SupportsCoefOps[Any]: ... + @overload + def __call__( + self, + val_f: _AnyValF, + c: _ArrayLikeCoef_co, + /, + *args: _ArrayLikeCoef_co, + ) -> _ObjectArray: ... + +@type_check_only +class _FuncVander(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + deg: SupportsIndex, + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + deg: SupportsIndex, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + deg: SupportsIndex, + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: npt.ArrayLike, + deg: SupportsIndex, + ) -> _CoefArray: ... + +_AnyDegrees: TypeAlias = Sequence[SupportsIndex] + +@type_check_only +class _FuncVander2D(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: _AnyDegrees, + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: _AnyDegrees, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + y: _ArrayLikeCoef_co, + deg: _AnyDegrees, + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: npt.ArrayLike, + y: npt.ArrayLike, + deg: _AnyDegrees, + ) -> _CoefArray: ... + +@type_check_only +class _FuncVander3D(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _ArrayLikeFloat_co, + y: _ArrayLikeFloat_co, + z: _ArrayLikeFloat_co, + deg: _AnyDegrees, + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeComplex_co, + y: _ArrayLikeComplex_co, + z: _ArrayLikeComplex_co, + deg: _AnyDegrees, + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + x: _ArrayLikeCoef_co, + y: _ArrayLikeCoef_co, + z: _ArrayLikeCoef_co, + deg: _AnyDegrees, + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + x: npt.ArrayLike, + y: npt.ArrayLike, + z: npt.ArrayLike, + deg: _AnyDegrees, + ) -> _CoefArray: ... + +# keep in sync with the broadest overload of `._FuncVander` +_AnyFuncVander: TypeAlias = Callable[ + [npt.ArrayLike, SupportsIndex], + _CoefArray, +] + +@type_check_only +class _FuncVanderND(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + vander_fs: Sequence[_AnyFuncVander], + points: Sequence[_ArrayLikeFloat_co], + degrees: Sequence[SupportsIndex], + ) -> _FloatArray: ... + @overload + def __call__( + self, + /, + vander_fs: Sequence[_AnyFuncVander], + points: Sequence[_ArrayLikeComplex_co], + degrees: Sequence[SupportsIndex], + ) -> _ComplexArray: ... + @overload + def __call__( + self, + /, + vander_fs: Sequence[_AnyFuncVander], + points: Sequence[ + _ArrayLikeObject_co | _ArrayLikeComplex_co, + ], + degrees: Sequence[SupportsIndex], + ) -> _ObjectArray: ... + @overload + def __call__( + self, + /, + vander_fs: Sequence[_AnyFuncVander], + points: Sequence[npt.ArrayLike], + degrees: Sequence[SupportsIndex], + ) -> _CoefArray: ... + +_FullFitResult: TypeAlias = Sequence[np.inexact[Any] | np.int32] + +@type_check_only +class _FuncFit(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + x: _SeriesLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeFloat_co = ..., + ) -> _FloatArray: ... + @overload + def __call__( + self, + x: _SeriesLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float, + full: Literal[True], + /, + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_FloatArray, _FullFitResult]: ... + @overload + def __call__( + self, + /, + x: _SeriesLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + *, + full: Literal[True], + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_FloatArray, _FullFitResult]: ... + + @overload + def __call__( + self, + /, + x: _SeriesLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeFloat_co = ..., + ) -> _ComplexArray: ... + @overload + def __call__( + self, + x: _SeriesLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float, + full: Literal[True], + /, + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_ComplexArray, _FullFitResult]: ... + @overload + def __call__( + self, + /, + x: _SeriesLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + *, + full: Literal[True], + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_ComplexArray, _FullFitResult]: ... + + @overload + def __call__( + self, + /, + x: _SeriesLikeComplex_co, + y: _ArrayLikeCoef_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeFloat_co = ..., + ) -> _ObjectArray: ... + @overload + def __call__( + self, + x: _SeriesLikeComplex_co, + y: _ArrayLikeCoef_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float, + full: Literal[True], + /, + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_ObjectArray, _FullFitResult]: ... + @overload + def __call__( + self, + /, + x: _SeriesLikeComplex_co, + y: _ArrayLikeCoef_co, + deg: int | _SeriesLikeInt_co, + rcond: None | float = ..., + *, + full: Literal[True], + w: None | _SeriesLikeFloat_co = ..., + ) -> tuple[_ObjectArray, _FullFitResult]: ... + +@type_check_only +class _FuncRoots(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _SeriesLikeFloat_co, + ) -> _Series[np.float64]: ... + @overload + def __call__( + self, + /, + c: _SeriesLikeComplex_co, + ) -> _Series[np.complex128]: ... + @overload + def __call__(self, /, c: _SeriesLikeCoef_co) -> _ObjectSeries: ... + + +_Companion: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_SCT]] + +@type_check_only +class _FuncCompanion(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _SeriesLikeFloat_co, + ) -> _Companion[np.float64]: ... + @overload + def __call__( + self, + /, + c: _SeriesLikeComplex_co, + ) -> _Companion[np.complex128]: ... + @overload + def __call__(self, /, c: _SeriesLikeCoef_co) -> _Companion[np.object_]: ... + +@type_check_only +class _FuncGauss(_Named[_Name_co], Protocol[_Name_co]): + def __call__( + self, + /, + deg: SupportsIndex, + ) -> _Tuple2[_Series[np.float64]]: ... + +@type_check_only +class _FuncWeight(_Named[_Name_co], Protocol[_Name_co]): + @overload + def __call__( + self, + /, + c: _ArrayLikeFloat_co, + ) -> npt.NDArray[np.float64]: ... + @overload + def __call__( + self, + /, + c: _ArrayLikeComplex_co, + ) -> npt.NDArray[np.complex128]: ... + @overload + def __call__(self, /, c: _ArrayLikeCoef_co) -> _ObjectArray: ... + +@type_check_only +class _FuncPts(_Named[_Name_co], Protocol[_Name_co]): + def __call__(self, /, npts: _AnyInt) -> _Series[np.float64]: ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.py b/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.py new file mode 100644 index 0000000000000000000000000000000000000000..837847e45110a9cf5cf202c496c68c5e437c4e67 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.py @@ -0,0 +1,2003 @@ +""" +==================================================== +Chebyshev Series (:mod:`numpy.polynomial.chebyshev`) +==================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Chebyshev series, including a `Chebyshev` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- + +.. autosummary:: + :toctree: generated/ + + Chebyshev + + +Constants +--------- + +.. autosummary:: + :toctree: generated/ + + chebdomain + chebzero + chebone + chebx + +Arithmetic +---------- + +.. autosummary:: + :toctree: generated/ + + chebadd + chebsub + chebmulx + chebmul + chebdiv + chebpow + chebval + chebval2d + chebval3d + chebgrid2d + chebgrid3d + +Calculus +-------- + +.. autosummary:: + :toctree: generated/ + + chebder + chebint + +Misc Functions +-------------- + +.. autosummary:: + :toctree: generated/ + + chebfromroots + chebroots + chebvander + chebvander2d + chebvander3d + chebgauss + chebweight + chebcompanion + chebfit + chebpts1 + chebpts2 + chebtrim + chebline + cheb2poly + poly2cheb + chebinterpolate + +See also +-------- +`numpy.polynomial` + +Notes +----- +The implementations of multiplication, division, integration, and +differentiation use the algebraic identities [1]_: + +.. math:: + T_n(x) = \\frac{z^n + z^{-n}}{2} \\\\ + z\\frac{dx}{dz} = \\frac{z - z^{-1}}{2}. + +where + +.. math:: x = \\frac{z + z^{-1}}{2}. + +These identities allow a Chebyshev series to be expressed as a finite, +symmetric Laurent series. In this module, this sort of Laurent series +is referred to as a "z-series." + +References +---------- +.. [1] A. T. Benjamin, et al., "Combinatorial Trigonometry with Chebyshev + Polynomials," *Journal of Statistical Planning and Inference 14*, 2008 + (https://web.archive.org/web/20080221202153/https://www.math.hmc.edu/~benjamin/papers/CombTrig.pdf, pg. 4) + +""" +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'chebzero', 'chebone', 'chebx', 'chebdomain', 'chebline', 'chebadd', + 'chebsub', 'chebmulx', 'chebmul', 'chebdiv', 'chebpow', 'chebval', + 'chebder', 'chebint', 'cheb2poly', 'poly2cheb', 'chebfromroots', + 'chebvander', 'chebfit', 'chebtrim', 'chebroots', 'chebpts1', + 'chebpts2', 'Chebyshev', 'chebval2d', 'chebval3d', 'chebgrid2d', + 'chebgrid3d', 'chebvander2d', 'chebvander3d', 'chebcompanion', + 'chebgauss', 'chebweight', 'chebinterpolate'] + +chebtrim = pu.trimcoef + +# +# A collection of functions for manipulating z-series. These are private +# functions and do minimal error checking. +# + +def _cseries_to_zseries(c): + """Convert Chebyshev series to z-series. + + Convert a Chebyshev series to the equivalent z-series. The result is + never an empty array. The dtype of the return is the same as that of + the input. No checks are run on the arguments as this routine is for + internal use. + + Parameters + ---------- + c : 1-D ndarray + Chebyshev coefficients, ordered from low to high + + Returns + ------- + zs : 1-D ndarray + Odd length symmetric z-series, ordered from low to high. + + """ + n = c.size + zs = np.zeros(2*n-1, dtype=c.dtype) + zs[n-1:] = c/2 + return zs + zs[::-1] + + +def _zseries_to_cseries(zs): + """Convert z-series to a Chebyshev series. + + Convert a z series to the equivalent Chebyshev series. The result is + never an empty array. The dtype of the return is the same as that of + the input. No checks are run on the arguments as this routine is for + internal use. + + Parameters + ---------- + zs : 1-D ndarray + Odd length symmetric z-series, ordered from low to high. + + Returns + ------- + c : 1-D ndarray + Chebyshev coefficients, ordered from low to high. + + """ + n = (zs.size + 1)//2 + c = zs[n-1:].copy() + c[1:n] *= 2 + return c + + +def _zseries_mul(z1, z2): + """Multiply two z-series. + + Multiply two z-series to produce a z-series. + + Parameters + ---------- + z1, z2 : 1-D ndarray + The arrays must be 1-D but this is not checked. + + Returns + ------- + product : 1-D ndarray + The product z-series. + + Notes + ----- + This is simply convolution. If symmetric/anti-symmetric z-series are + denoted by S/A then the following rules apply: + + S*S, A*A -> S + S*A, A*S -> A + + """ + return np.convolve(z1, z2) + + +def _zseries_div(z1, z2): + """Divide the first z-series by the second. + + Divide `z1` by `z2` and return the quotient and remainder as z-series. + Warning: this implementation only applies when both z1 and z2 have the + same symmetry, which is sufficient for present purposes. + + Parameters + ---------- + z1, z2 : 1-D ndarray + The arrays must be 1-D and have the same symmetry, but this is not + checked. + + Returns + ------- + + (quotient, remainder) : 1-D ndarrays + Quotient and remainder as z-series. + + Notes + ----- + This is not the same as polynomial division on account of the desired form + of the remainder. If symmetric/anti-symmetric z-series are denoted by S/A + then the following rules apply: + + S/S -> S,S + A/A -> S,A + + The restriction to types of the same symmetry could be fixed but seems like + unneeded generality. There is no natural form for the remainder in the case + where there is no symmetry. + + """ + z1 = z1.copy() + z2 = z2.copy() + lc1 = len(z1) + lc2 = len(z2) + if lc2 == 1: + z1 /= z2 + return z1, z1[:1]*0 + elif lc1 < lc2: + return z1[:1]*0, z1 + else: + dlen = lc1 - lc2 + scl = z2[0] + z2 /= scl + quo = np.empty(dlen + 1, dtype=z1.dtype) + i = 0 + j = dlen + while i < j: + r = z1[i] + quo[i] = z1[i] + quo[dlen - i] = r + tmp = r*z2 + z1[i:i+lc2] -= tmp + z1[j:j+lc2] -= tmp + i += 1 + j -= 1 + r = z1[i] + quo[i] = r + tmp = r*z2 + z1[i:i+lc2] -= tmp + quo /= scl + rem = z1[i+1:i-1+lc2].copy() + return quo, rem + + +def _zseries_der(zs): + """Differentiate a z-series. + + The derivative is with respect to x, not z. This is achieved using the + chain rule and the value of dx/dz given in the module notes. + + Parameters + ---------- + zs : z-series + The z-series to differentiate. + + Returns + ------- + derivative : z-series + The derivative + + Notes + ----- + The zseries for x (ns) has been multiplied by two in order to avoid + using floats that are incompatible with Decimal and likely other + specialized scalar types. This scaling has been compensated by + multiplying the value of zs by two also so that the two cancels in the + division. + + """ + n = len(zs)//2 + ns = np.array([-1, 0, 1], dtype=zs.dtype) + zs *= np.arange(-n, n+1)*2 + d, r = _zseries_div(zs, ns) + return d + + +def _zseries_int(zs): + """Integrate a z-series. + + The integral is with respect to x, not z. This is achieved by a change + of variable using dx/dz given in the module notes. + + Parameters + ---------- + zs : z-series + The z-series to integrate + + Returns + ------- + integral : z-series + The indefinite integral + + Notes + ----- + The zseries for x (ns) has been multiplied by two in order to avoid + using floats that are incompatible with Decimal and likely other + specialized scalar types. This scaling has been compensated by + dividing the resulting zs by two. + + """ + n = 1 + len(zs)//2 + ns = np.array([-1, 0, 1], dtype=zs.dtype) + zs = _zseries_mul(zs, ns) + div = np.arange(-n, n+1)*2 + zs[:n] /= div[:n] + zs[n+1:] /= div[n+1:] + zs[n] = 0 + return zs + +# +# Chebyshev series functions +# + + +def poly2cheb(pol): + """ + Convert a polynomial to a Chebyshev series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Chebyshev series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Chebyshev + series. + + See Also + -------- + cheb2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> p = P.Polynomial(range(4)) + >>> p + Polynomial([0., 1., 2., 3.], domain=[-1., 1.], window=[-1., 1.], symbol='x') + >>> c = p.convert(kind=P.Chebyshev) + >>> c + Chebyshev([1. , 3.25, 1. , 0.75], domain=[-1., 1.], window=[-1., ... + >>> P.chebyshev.poly2cheb(range(4)) + array([1. , 3.25, 1. , 0.75]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = chebadd(chebmulx(res), pol[i]) + return res + + +def cheb2poly(c): + """ + Convert a Chebyshev series to a polynomial. + + Convert an array representing the coefficients of a Chebyshev series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Chebyshev series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2cheb + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> c = P.Chebyshev(range(4)) + >>> c + Chebyshev([0., 1., 2., 3.], domain=[-1., 1.], window=[-1., 1.], symbol='x') + >>> p = c.convert(kind=P.Polynomial) + >>> p + Polynomial([-2., -8., 4., 12.], domain=[-1., 1.], window=[-1., 1.], ... + >>> P.chebyshev.cheb2poly(range(4)) + array([-2., -8., 4., 12.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n < 3: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1) + c1 = polyadd(tmp, polymulx(c1)*2) + return polyadd(c0, polymulx(c1)) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Chebyshev default domain. +chebdomain = np.array([-1., 1.]) + +# Chebyshev coefficients representing zero. +chebzero = np.array([0]) + +# Chebyshev coefficients representing one. +chebone = np.array([1]) + +# Chebyshev coefficients representing the identity x. +chebx = np.array([0, 1]) + + +def chebline(off, scl): + """ + Chebyshev series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Chebyshev series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebline(3,2) + array([3, 2]) + >>> C.chebval(-3, C.chebline(3,2)) # should be -3 + -3.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def chebfromroots(roots): + """ + Generate a Chebyshev series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Chebyshev form, where the :math:`r_n` are the roots specified in + `roots`. If a zero has multiplicity n, then it must appear in `roots` + n times. For instance, if 2 is a root of multiplicity three and 3 is a + root of multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. + The roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * T_1(x) + ... + c_n * T_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Chebyshev form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebfromroots((-1,0,1)) # x^3 - x relative to the standard basis + array([ 0. , -0.25, 0. , 0.25]) + >>> j = complex(0,1) + >>> C.chebfromroots((-j,j)) # x^2 + 1 relative to the standard basis + array([1.5+0.j, 0. +0.j, 0.5+0.j]) + + """ + return pu._fromroots(chebline, chebmul, roots) + + +def chebadd(c1, c2): + """ + Add one Chebyshev series to another. + + Returns the sum of two Chebyshev series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Chebyshev series of their sum. + + See Also + -------- + chebsub, chebmulx, chebmul, chebdiv, chebpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Chebyshev series + is a Chebyshev series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebadd(c1,c2) + array([4., 4., 4.]) + + """ + return pu._add(c1, c2) + + +def chebsub(c1, c2): + """ + Subtract one Chebyshev series from another. + + Returns the difference of two Chebyshev series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Chebyshev series coefficients representing their difference. + + See Also + -------- + chebadd, chebmulx, chebmul, chebdiv, chebpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Chebyshev + series is a Chebyshev series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebsub(c1,c2) + array([-2., 0., 2.]) + >>> C.chebsub(c2,c1) # -C.chebsub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def chebmulx(c): + """Multiply a Chebyshev series by x. + + Multiply the polynomial `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + chebadd, chebsub, chebmul, chebdiv, chebpow + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> C.chebmulx([1,2,3]) + array([1. , 2.5, 1. , 1.5]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + if len(c) > 1: + tmp = c[1:]/2 + prd[2:] = tmp + prd[0:-2] += tmp + return prd + + +def chebmul(c1, c2): + """ + Multiply one Chebyshev series by another. + + Returns the product of two Chebyshev series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Chebyshev series coefficients representing their product. + + See Also + -------- + chebadd, chebsub, chebmulx, chebdiv, chebpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Chebyshev polynomial basis set. Thus, to express + the product as a C-series, it is typically necessary to "reproject" + the product onto said basis set, which typically produces + "unintuitive live" (but correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebmul(c1,c2) # multiplication requires "reprojection" + array([ 6.5, 12. , 12. , 4. , 1.5]) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + z1 = _cseries_to_zseries(c1) + z2 = _cseries_to_zseries(c2) + prd = _zseries_mul(z1, z2) + ret = _zseries_to_cseries(prd) + return pu.trimseq(ret) + + +def chebdiv(c1, c2): + """ + Divide one Chebyshev series by another. + + Returns the quotient-with-remainder of two Chebyshev series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Chebyshev series coefficients representing the quotient and + remainder. + + See Also + -------- + chebadd, chebsub, chebmulx, chebmul, chebpow + + Notes + ----- + In general, the (polynomial) division of one C-series by another + results in quotient and remainder terms that are not in the Chebyshev + polynomial basis set. Thus, to express these results as C-series, it + is typically necessary to "reproject" the results onto said basis + set, which typically produces "unintuitive" (but correct) results; + see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebdiv(c1,c2) # quotient "intuitive," remainder not + (array([3.]), array([-8., -4.])) + >>> c2 = (0,1,2,3) + >>> C.chebdiv(c2,c1) # neither "intuitive" + (array([0., 2.]), array([-2., -4.])) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError # FIXME: add message with details to exception + + # note: this is more efficient than `pu._div(chebmul, c1, c2)` + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + z1 = _cseries_to_zseries(c1) + z2 = _cseries_to_zseries(c2) + quo, rem = _zseries_div(z1, z2) + quo = pu.trimseq(_zseries_to_cseries(quo)) + rem = pu.trimseq(_zseries_to_cseries(rem)) + return quo, rem + + +def chebpow(c, pow, maxpower=16): + """Raise a Chebyshev series to a power. + + Returns the Chebyshev series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``T_0 + 2*T_1 + 3*T_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Chebyshev series of power. + + See Also + -------- + chebadd, chebsub, chebmulx, chebmul, chebdiv + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> C.chebpow([1, 2, 3, 4], 2) + array([15.5, 22. , 16. , ..., 12.5, 12. , 8. ]) + + """ + # note: this is more efficient than `pu._pow(chebmul, c1, c2)`, as it + # avoids converting between z and c series repeatedly + + # c is a trimmed copy + [c] = pu.as_series([c]) + power = int(pow) + if power != pow or power < 0: + raise ValueError("Power must be a non-negative integer.") + elif maxpower is not None and power > maxpower: + raise ValueError("Power is too large") + elif power == 0: + return np.array([1], dtype=c.dtype) + elif power == 1: + return c + else: + # This can be made more efficient by using powers of two + # in the usual way. + zs = _cseries_to_zseries(c) + prd = zs + for i in range(2, power + 1): + prd = np.convolve(prd, zs) + return _zseries_to_cseries(prd) + + +def chebder(c, m=1, scl=1, axis=0): + """ + Differentiate a Chebyshev series. + + Returns the Chebyshev series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*T_0 + 2*T_1 + 3*T_2`` + while [[1,2],[1,2]] represents ``1*T_0(x)*T_0(y) + 1*T_1(x)*T_0(y) + + 2*T_0(x)*T_1(y) + 2*T_1(x)*T_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Chebyshev series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Chebyshev series of the derivative. + + See Also + -------- + chebint + + Notes + ----- + In general, the result of differentiating a C-series needs to be + "reprojected" onto the C-series basis set. Thus, typically, the + result of this function is "unintuitive," albeit correct; see Examples + section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c = (1,2,3,4) + >>> C.chebder(c) + array([14., 12., 24.]) + >>> C.chebder(c,3) + array([96.]) + >>> C.chebder(c,scl=-1) + array([-14., -12., -24.]) + >>> C.chebder(c,2,-1) + array([12., 96.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 2, -1): + der[j - 1] = (2*j)*c[j] + c[j - 2] += (j*c[j])/(j - 2) + if n > 1: + der[1] = 4*c[2] + der[0] = c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def chebint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Chebyshev series. + + Returns the Chebyshev series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``T_0 + 2*T_1 + 3*T_2`` while [[1,2],[1,2]] + represents ``1*T_0(x)*T_0(y) + 1*T_1(x)*T_0(y) + 2*T_0(x)*T_1(y) + + 2*T_1(x)*T_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Chebyshev series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at zero + is the first value in the list, the value of the second integral + at zero is the second value, etc. If ``k == []`` (the default), + all constants are set to zero. If ``m == 1``, a single scalar can + be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + C-series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 1``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + chebder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a`- perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c = (1,2,3) + >>> C.chebint(c) + array([ 0.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,3) + array([ 0.03125 , -0.1875 , 0.04166667, -0.05208333, 0.01041667, # may vary + 0.00625 ]) + >>> C.chebint(c, k=3) + array([ 3.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,lbnd=-2) + array([ 8.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,scl=-2) + array([-1., 1., -1., -1.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + if n > 1: + tmp[2] = c[1]/4 + for j in range(2, n): + tmp[j + 1] = c[j]/(2*(j + 1)) + tmp[j - 1] -= c[j]/(2*(j - 1)) + tmp[0] += k[i] - chebval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def chebval(x, c, tensor=True): + """ + Evaluate a Chebyshev series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * T_0(x) + c_1 * T_1(x) + ... + c_n * T_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + chebval2d, chebgrid2d, chebval3d, chebgrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + x2 = 2*x + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + c0 = c[-i] - c1 + c1 = tmp + c1*x2 + return c0 + c1*x + + +def chebval2d(x, y, c): + """ + Evaluate a 2-D Chebyshev series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * T_i(x) * T_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than 2 the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points formed + from pairs of corresponding values from `x` and `y`. + + See Also + -------- + chebval, chebgrid2d, chebval3d, chebgrid3d + """ + return pu._valnd(chebval, c, x, y) + + +def chebgrid2d(x, y, c): + """ + Evaluate a 2-D Chebyshev series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * T_i(a) * T_j(b), + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + chebval, chebval2d, chebval3d, chebgrid3d + """ + return pu._gridnd(chebval, c, x, y) + + +def chebval3d(x, y, z, c): + """ + Evaluate a 3-D Chebyshev series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * T_i(x) * T_j(y) * T_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + ``(x, y, z)``, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + chebval, chebval2d, chebgrid2d, chebgrid3d + """ + return pu._valnd(chebval, c, x, y, z) + + +def chebgrid3d(x, y, z, c): + """ + Evaluate a 3-D Chebyshev series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * T_i(a) * T_j(b) * T_k(c) + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + chebval, chebval2d, chebgrid2d, chebval3d + """ + return pu._gridnd(chebval, c, x, y, z) + + +def chebvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = T_i(x), + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Chebyshev polynomial. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + matrix ``V = chebvander(x, n)``, then ``np.dot(V, c)`` and + ``chebval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Chebyshev series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Chebyshev polynomial. The dtype will be the same as + the converted `x`. + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + # Use forward recursion to generate the entries. + v[0] = x*0 + 1 + if ideg > 0: + x2 = 2*x + v[1] = x + for i in range(2, ideg + 1): + v[i] = v[i-1]*x2 - v[i-2] + return np.moveaxis(v, 0, -1) + + +def chebvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = T_i(x) * T_j(y), + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the degrees of + the Chebyshev polynomials. + + If ``V = chebvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``chebval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Chebyshev + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + chebvander, chebvander3d, chebval2d, chebval3d + """ + return pu._vander_nd_flat((chebvander, chebvander), (x, y), deg) + + +def chebvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = T_i(x)*T_j(y)*T_k(z), + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the degrees of the Chebyshev polynomials. + + If ``V = chebvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``chebval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Chebyshev + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + chebvander, chebvander3d, chebval2d, chebval3d + """ + return pu._vander_nd_flat((chebvander, chebvander, chebvander), (x, y, z), deg) + + +def chebfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Chebyshev series to data. + + Return the coefficients of a Chebyshev series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * T_1(x) + ... + c_n * T_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer, + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is ``len(x)*eps``, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Chebyshev coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + chebval : Evaluates a Chebyshev series. + chebvander : Vandermonde matrix of Chebyshev series. + chebweight : Chebyshev weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Chebyshev series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where :math:`w_j` are the weights. This problem is solved by setting up + as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `~exceptions.RankWarning` will be issued. This means that + the coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Chebyshev series are usually better conditioned than fits + using power series, but much can depend on the distribution of the + sample points and the smoothness of the data. If the quality of the fit + is inadequate splines may be a good alternative. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + + """ + return pu._fit(chebvander, x, y, deg, rcond, full, w) + + +def chebcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is a Chebyshev basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.array([1.] + [np.sqrt(.5)]*(n-1)) + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[0] = np.sqrt(.5) + top[1:] = 1/2 + bot[...] = top + mat[:, -1] -= (c[:-1]/c[-1])*(scl/scl[-1])*.5 + return mat + + +def chebroots(c): + """ + Compute the roots of a Chebyshev series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * T_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Chebyshev series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> import numpy.polynomial.chebyshev as cheb + >>> cheb.chebroots((-1, 1,-1, 1)) # T3 - T2 + T1 - T0 has real roots + array([ -5.00000000e-01, 2.60860684e-17, 1.00000000e+00]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = chebcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def chebinterpolate(func, deg, args=()): + """Interpolate a function at the Chebyshev points of the first kind. + + Returns the Chebyshev series that interpolates `func` at the Chebyshev + points of the first kind in the interval [-1, 1]. The interpolating + series tends to a minmax approximation to `func` with increasing `deg` + if the function is continuous in the interval. + + Parameters + ---------- + func : function + The function to be approximated. It must be a function of a single + variable of the form ``f(x, a, b, c...)``, where ``a, b, c...`` are + extra arguments passed in the `args` parameter. + deg : int + Degree of the interpolating polynomial + args : tuple, optional + Extra arguments to be used in the function call. Default is no extra + arguments. + + Returns + ------- + coef : ndarray, shape (deg + 1,) + Chebyshev coefficients of the interpolating series ordered from low to + high. + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebinterpolate(lambda x: np.tanh(x) + 0.5, 8) + array([ 5.00000000e-01, 8.11675684e-01, -9.86864911e-17, + -5.42457905e-02, -2.71387850e-16, 4.51658839e-03, + 2.46716228e-17, -3.79694221e-04, -3.26899002e-16]) + + Notes + ----- + The Chebyshev polynomials used in the interpolation are orthogonal when + sampled at the Chebyshev points of the first kind. If it is desired to + constrain some of the coefficients they can simply be set to the desired + value after the interpolation, no new interpolation or fit is needed. This + is especially useful if it is known apriori that some of coefficients are + zero. For instance, if the function is even then the coefficients of the + terms of odd degree in the result can be set to zero. + + """ + deg = np.asarray(deg) + + # check arguments. + if deg.ndim > 0 or deg.dtype.kind not in 'iu' or deg.size == 0: + raise TypeError("deg must be an int") + if deg < 0: + raise ValueError("expected deg >= 0") + + order = deg + 1 + xcheb = chebpts1(order) + yfunc = func(xcheb, *args) + m = chebvander(xcheb, deg) + c = np.dot(m.T, yfunc) + c[0] /= order + c[1:] /= 0.5*order + + return c + + +def chebgauss(deg): + """ + Gauss-Chebyshev quadrature. + + Computes the sample points and weights for Gauss-Chebyshev quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-1, 1]` with + the weight function :math:`f(x) = 1/\\sqrt{1 - x^2}`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + The results have only been tested up to degree 100, higher degrees may + be problematic. For Gauss-Chebyshev there are closed form solutions for + the sample points and weights. If n = `deg`, then + + .. math:: x_i = \\cos(\\pi (2 i - 1) / (2 n)) + + .. math:: w_i = \\pi / n + + """ + ideg = pu._as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + x = np.cos(np.pi * np.arange(1, 2*ideg, 2) / (2.0*ideg)) + w = np.ones(ideg)*(np.pi/ideg) + + return x, w + + +def chebweight(x): + """ + The weight function of the Chebyshev polynomials. + + The weight function is :math:`1/\\sqrt{1 - x^2}` and the interval of + integration is :math:`[-1, 1]`. The Chebyshev polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + """ + w = 1./(np.sqrt(1. + x) * np.sqrt(1. - x)) + return w + + +def chebpts1(npts): + """ + Chebyshev points of the first kind. + + The Chebyshev points of the first kind are the points ``cos(x)``, + where ``x = [pi*(k + .5)/npts for k in range(npts)]``. + + Parameters + ---------- + npts : int + Number of sample points desired. + + Returns + ------- + pts : ndarray + The Chebyshev points of the first kind. + + See Also + -------- + chebpts2 + """ + _npts = int(npts) + if _npts != npts: + raise ValueError("npts must be integer") + if _npts < 1: + raise ValueError("npts must be >= 1") + + x = 0.5 * np.pi / _npts * np.arange(-_npts+1, _npts+1, 2) + return np.sin(x) + + +def chebpts2(npts): + """ + Chebyshev points of the second kind. + + The Chebyshev points of the second kind are the points ``cos(x)``, + where ``x = [pi*k/(npts - 1) for k in range(npts)]`` sorted in ascending + order. + + Parameters + ---------- + npts : int + Number of sample points desired. + + Returns + ------- + pts : ndarray + The Chebyshev points of the second kind. + """ + _npts = int(npts) + if _npts != npts: + raise ValueError("npts must be integer") + if _npts < 2: + raise ValueError("npts must be >= 2") + + x = np.linspace(-np.pi, 0, _npts) + return np.cos(x) + + +# +# Chebyshev series class +# + +class Chebyshev(ABCPolyBase): + """A Chebyshev series class. + + The Chebyshev class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + Chebyshev coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*T_0(x) + 2*T_1(x) + 3*T_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(chebadd) + _sub = staticmethod(chebsub) + _mul = staticmethod(chebmul) + _div = staticmethod(chebdiv) + _pow = staticmethod(chebpow) + _val = staticmethod(chebval) + _int = staticmethod(chebint) + _der = staticmethod(chebder) + _fit = staticmethod(chebfit) + _line = staticmethod(chebline) + _roots = staticmethod(chebroots) + _fromroots = staticmethod(chebfromroots) + + @classmethod + def interpolate(cls, func, deg, domain=None, args=()): + """Interpolate a function at the Chebyshev points of the first kind. + + Returns the series that interpolates `func` at the Chebyshev points of + the first kind scaled and shifted to the `domain`. The resulting series + tends to a minmax approximation of `func` when the function is + continuous in the domain. + + Parameters + ---------- + func : function + The function to be interpolated. It must be a function of a single + variable of the form ``f(x, a, b, c...)``, where ``a, b, c...`` are + extra arguments passed in the `args` parameter. + deg : int + Degree of the interpolating polynomial. + domain : {None, [beg, end]}, optional + Domain over which `func` is interpolated. The default is None, in + which case the domain is [-1, 1]. + args : tuple, optional + Extra arguments to be used in the function call. Default is no + extra arguments. + + Returns + ------- + polynomial : Chebyshev instance + Interpolating Chebyshev instance. + + Notes + ----- + See `numpy.polynomial.chebinterpolate` for more details. + + """ + if domain is None: + domain = cls.domain + xfunc = lambda x: func(pu.mapdomain(x, cls.window, domain), *args) + coef = chebinterpolate(xfunc, deg) + return cls(coef, domain=domain) + + # Virtual properties + domain = np.array(chebdomain) + window = np.array(chebdomain) + basis_name = 'T' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.pyi new file mode 100644 index 0000000000000000000000000000000000000000..067af81d635d75511469f6cd130d774f00391be6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/chebyshev.pyi @@ -0,0 +1,192 @@ +from collections.abc import Callable, Iterable +from typing import ( + Any, + Concatenate, + Final, + Literal as L, + TypeVar, + overload, +) + +import numpy as np +import numpy.typing as npt +from numpy._typing import _IntLike_co + +from ._polybase import ABCPolyBase +from ._polytypes import ( + _SeriesLikeCoef_co, + _Array1, + _Series, + _Array2, + _CoefSeries, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncGauss, + _FuncInteg, + _FuncLine, + _FuncPoly2Ortho, + _FuncPow, + _FuncPts, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVal2D, + _FuncVal3D, + _FuncValFromRoots, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncWeight, +) +from .polyutils import trimcoef as chebtrim + +__all__ = [ + "chebzero", + "chebone", + "chebx", + "chebdomain", + "chebline", + "chebadd", + "chebsub", + "chebmulx", + "chebmul", + "chebdiv", + "chebpow", + "chebval", + "chebder", + "chebint", + "cheb2poly", + "poly2cheb", + "chebfromroots", + "chebvander", + "chebfit", + "chebtrim", + "chebroots", + "chebpts1", + "chebpts2", + "Chebyshev", + "chebval2d", + "chebval3d", + "chebgrid2d", + "chebgrid3d", + "chebvander2d", + "chebvander3d", + "chebcompanion", + "chebgauss", + "chebweight", + "chebinterpolate", +] + +_SCT = TypeVar("_SCT", bound=np.number[Any] | np.object_) +def _cseries_to_zseries(c: npt.NDArray[_SCT]) -> _Series[_SCT]: ... +def _zseries_to_cseries(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ... +def _zseries_mul( + z1: npt.NDArray[_SCT], + z2: npt.NDArray[_SCT], +) -> _Series[_SCT]: ... +def _zseries_div( + z1: npt.NDArray[_SCT], + z2: npt.NDArray[_SCT], +) -> _Series[_SCT]: ... +def _zseries_der(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ... +def _zseries_int(zs: npt.NDArray[_SCT]) -> _Series[_SCT]: ... + +poly2cheb: _FuncPoly2Ortho[L["poly2cheb"]] +cheb2poly: _FuncUnOp[L["cheb2poly"]] + +chebdomain: Final[_Array2[np.float64]] +chebzero: Final[_Array1[np.int_]] +chebone: Final[_Array1[np.int_]] +chebx: Final[_Array2[np.int_]] + +chebline: _FuncLine[L["chebline"]] +chebfromroots: _FuncFromRoots[L["chebfromroots"]] +chebadd: _FuncBinOp[L["chebadd"]] +chebsub: _FuncBinOp[L["chebsub"]] +chebmulx: _FuncUnOp[L["chebmulx"]] +chebmul: _FuncBinOp[L["chebmul"]] +chebdiv: _FuncBinOp[L["chebdiv"]] +chebpow: _FuncPow[L["chebpow"]] +chebder: _FuncDer[L["chebder"]] +chebint: _FuncInteg[L["chebint"]] +chebval: _FuncVal[L["chebval"]] +chebval2d: _FuncVal2D[L["chebval2d"]] +chebval3d: _FuncVal3D[L["chebval3d"]] +chebvalfromroots: _FuncValFromRoots[L["chebvalfromroots"]] +chebgrid2d: _FuncVal2D[L["chebgrid2d"]] +chebgrid3d: _FuncVal3D[L["chebgrid3d"]] +chebvander: _FuncVander[L["chebvander"]] +chebvander2d: _FuncVander2D[L["chebvander2d"]] +chebvander3d: _FuncVander3D[L["chebvander3d"]] +chebfit: _FuncFit[L["chebfit"]] +chebcompanion: _FuncCompanion[L["chebcompanion"]] +chebroots: _FuncRoots[L["chebroots"]] +chebgauss: _FuncGauss[L["chebgauss"]] +chebweight: _FuncWeight[L["chebweight"]] +chebpts1: _FuncPts[L["chebpts1"]] +chebpts2: _FuncPts[L["chebpts2"]] + +# keep in sync with `Chebyshev.interpolate` +_RT = TypeVar("_RT", bound=np.number[Any] | np.bool | np.object_) +@overload +def chebinterpolate( + func: np.ufunc, + deg: _IntLike_co, + args: tuple[()] = ..., +) -> npt.NDArray[np.float64 | np.complex128 | np.object_]: ... +@overload +def chebinterpolate( + func: Callable[[npt.NDArray[np.float64]], _RT], + deg: _IntLike_co, + args: tuple[()] = ..., +) -> npt.NDArray[_RT]: ... +@overload +def chebinterpolate( + func: Callable[Concatenate[npt.NDArray[np.float64], ...], _RT], + deg: _IntLike_co, + args: Iterable[Any], +) -> npt.NDArray[_RT]: ... + +_Self = TypeVar("_Self", bound=object) + +class Chebyshev(ABCPolyBase[L["T"]]): + @overload + @classmethod + def interpolate( + cls: type[_Self], + /, + func: Callable[[npt.NDArray[np.float64]], _CoefSeries], + deg: _IntLike_co, + domain: None | _SeriesLikeCoef_co = ..., + args: tuple[()] = ..., + ) -> _Self: ... + @overload + @classmethod + def interpolate( + cls: type[_Self], + /, + func: Callable[ + Concatenate[npt.NDArray[np.float64], ...], + _CoefSeries, + ], + deg: _IntLike_co, + domain: None | _SeriesLikeCoef_co = ..., + *, + args: Iterable[Any], + ) -> _Self: ... + @overload + @classmethod + def interpolate( + cls: type[_Self], + func: Callable[ + Concatenate[npt.NDArray[np.float64], ...], + _CoefSeries, + ], + deg: _IntLike_co, + domain: None | _SeriesLikeCoef_co, + args: Iterable[Any], + /, + ) -> _Self: ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.py b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.py new file mode 100644 index 0000000000000000000000000000000000000000..24e51dca7fa55c83dfa467013440e160b260d9d9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.py @@ -0,0 +1,1740 @@ +""" +============================================================== +Hermite Series, "Physicists" (:mod:`numpy.polynomial.hermite`) +============================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Hermite series, including a `Hermite` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Hermite + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + hermdomain + hermzero + hermone + hermx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + hermadd + hermsub + hermmulx + hermmul + hermdiv + hermpow + hermval + hermval2d + hermval3d + hermgrid2d + hermgrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + hermder + hermint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + hermfromroots + hermroots + hermvander + hermvander2d + hermvander3d + hermgauss + hermweight + hermcompanion + hermfit + hermtrim + hermline + herm2poly + poly2herm + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'hermzero', 'hermone', 'hermx', 'hermdomain', 'hermline', 'hermadd', + 'hermsub', 'hermmulx', 'hermmul', 'hermdiv', 'hermpow', 'hermval', + 'hermder', 'hermint', 'herm2poly', 'poly2herm', 'hermfromroots', + 'hermvander', 'hermfit', 'hermtrim', 'hermroots', 'Hermite', + 'hermval2d', 'hermval3d', 'hermgrid2d', 'hermgrid3d', 'hermvander2d', + 'hermvander3d', 'hermcompanion', 'hermgauss', 'hermweight'] + +hermtrim = pu.trimcoef + + +def poly2herm(pol): + """ + poly2herm(pol) + + Convert a polynomial to a Hermite series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Hermite series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Hermite + series. + + See Also + -------- + herm2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite import poly2herm + >>> poly2herm(np.arange(4)) + array([1. , 2.75 , 0.5 , 0.375]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = hermadd(hermmulx(res), pol[i]) + return res + + +def herm2poly(c): + """ + Convert a Hermite series to a polynomial. + + Convert an array representing the coefficients of a Hermite series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Hermite series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2herm + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite import herm2poly + >>> herm2poly([ 1. , 2.75 , 0.5 , 0.375]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + if n == 2: + c[1] *= 2 + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1*(2*(i - 1))) + c1 = polyadd(tmp, polymulx(c1)*2) + return polyadd(c0, polymulx(c1)*2) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Hermite +hermdomain = np.array([-1., 1.]) + +# Hermite coefficients representing zero. +hermzero = np.array([0]) + +# Hermite coefficients representing one. +hermone = np.array([1]) + +# Hermite coefficients representing the identity x. +hermx = np.array([0, 1/2]) + + +def hermline(off, scl): + """ + Hermite series whose graph is a straight line. + + + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Hermite series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial.hermite import hermline, hermval + >>> hermval(0,hermline(3, 2)) + 3.0 + >>> hermval(1,hermline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off, scl/2]) + else: + return np.array([off]) + + +def hermfromroots(roots): + """ + Generate a Hermite series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Hermite form, where the :math:`r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * H_1(x) + ... + c_n * H_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Hermite form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> from numpy.polynomial.hermite import hermfromroots, hermval + >>> coef = hermfromroots((-1, 0, 1)) + >>> hermval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = hermfromroots((-1j, 1j)) + >>> hermval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(hermline, hermmul, roots) + + +def hermadd(c1, c2): + """ + Add one Hermite series to another. + + Returns the sum of two Hermite series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Hermite series of their sum. + + See Also + -------- + hermsub, hermmulx, hermmul, hermdiv, hermpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Hermite series + is a Hermite series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite import hermadd + >>> hermadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + """ + return pu._add(c1, c2) + + +def hermsub(c1, c2): + """ + Subtract one Hermite series from another. + + Returns the difference of two Hermite series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their difference. + + See Also + -------- + hermadd, hermmulx, hermmul, hermdiv, hermpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Hermite + series is a Hermite series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite import hermsub + >>> hermsub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def hermmulx(c): + """Multiply a Hermite series by x. + + Multiply the Hermite series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + hermadd, hermsub, hermmul, hermdiv, hermpow + + Notes + ----- + The multiplication uses the recursion relationship for Hermite + polynomials in the form + + .. math:: + + xP_i(x) = (P_{i + 1}(x)/2 + i*P_{i - 1}(x)) + + Examples + -------- + >>> from numpy.polynomial.hermite import hermmulx + >>> hermmulx([1, 2, 3]) + array([2. , 6.5, 1. , 1.5]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0]/2 + for i in range(1, len(c)): + prd[i + 1] = c[i]/2 + prd[i - 1] += c[i]*i + return prd + + +def hermmul(c1, c2): + """ + Multiply one Hermite series by another. + + Returns the product of two Hermite series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their product. + + See Also + -------- + hermadd, hermsub, hermmulx, hermdiv, hermpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Hermite polynomial basis set. Thus, to express + the product as a Hermite series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermmul + >>> hermmul([1, 2, 3], [0, 1, 2]) + array([52., 29., 52., 7., 6.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = hermsub(c[-i]*xs, c1*(2*(nd - 1))) + c1 = hermadd(tmp, hermmulx(c1)*2) + return hermadd(c0, hermmulx(c1)*2) + + +def hermdiv(c1, c2): + """ + Divide one Hermite series by another. + + Returns the quotient-with-remainder of two Hermite series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Hermite series coefficients representing the quotient and + remainder. + + See Also + -------- + hermadd, hermsub, hermmulx, hermmul, hermpow + + Notes + ----- + In general, the (polynomial) division of one Hermite series by another + results in quotient and remainder terms that are not in the Hermite + polynomial basis set. Thus, to express these results as a Hermite + series, it is necessary to "reproject" the results onto the Hermite + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermdiv + >>> hermdiv([ 52., 29., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> hermdiv([ 54., 31., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([2., 2.])) + >>> hermdiv([ 53., 30., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 1.])) + + """ + return pu._div(hermmul, c1, c2) + + +def hermpow(c, pow, maxpower=16): + """Raise a Hermite series to a power. + + Returns the Hermite series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Hermite series of power. + + See Also + -------- + hermadd, hermsub, hermmulx, hermmul, hermdiv + + Examples + -------- + >>> from numpy.polynomial.hermite import hermpow + >>> hermpow([1, 2, 3], 2) + array([81., 52., 82., 12., 9.]) + + """ + return pu._pow(hermmul, c, pow, maxpower) + + +def hermder(c, m=1, scl=1, axis=0): + """ + Differentiate a Hermite series. + + Returns the Hermite series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*H_0 + 2*H_1 + 3*H_2`` + while [[1,2],[1,2]] represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + + 2*H_0(x)*H_1(y) + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite series coefficients. If `c` is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Hermite series of the derivative. + + See Also + -------- + hermint + + Notes + ----- + In general, the result of differentiating a Hermite series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermder + >>> hermder([ 1. , 0.5, 0.5, 0.5]) + array([1., 2., 3.]) + >>> hermder([-0.5, 1./2., 1./8., 1./12., 1./16.], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 0, -1): + der[j - 1] = (2*j)*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Hermite series. + + Returns the Hermite series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``H_0 + 2*H_1 + 3*H_2`` while [[1,2],[1,2]] + represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + 2*H_0(x)*H_1(y) + + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + Hermite series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + hermder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermint + >>> hermint([1,2,3]) # integrate once, value 0 at 0. + array([1. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], m=2) # integrate twice, value & deriv 0 at 0 + array([-0.5 , 0.5 , 0.125 , 0.08333333, 0.0625 ]) # may vary + >>> hermint([1,2,3], k=1) # integrate once, value 1 at 0. + array([2. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], lbnd=-1) # integrate once, value 0 at -1 + array([-2. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], m=2, k=[1,2], lbnd=-1) + array([ 1.66666667, -0.5 , 0.125 , 0.08333333, 0.0625 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0]/2 + for j in range(1, n): + tmp[j + 1] = c[j]/(2*(j + 1)) + tmp[0] += k[i] - hermval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermval(x, c, tensor=True): + """ + Evaluate an Hermite series at points x. + + If `c` is of length ``n + 1``, this function returns the value: + + .. math:: p(x) = c_0 * H_0(x) + c_1 * H_1(x) + ... + c_n * H_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + hermval2d, hermgrid2d, hermval3d, hermgrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermval + >>> coef = [1,2,3] + >>> hermval(1, coef) + 11.0 + >>> hermval([[1,2],[3,4]], coef) + array([[ 11., 51.], + [115., 203.]]) + + """ + c = np.array(c, ndmin=1, copy=None) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + x2 = x*2 + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - c1*(2*(nd - 1)) + c1 = tmp + c1*x2 + return c0 + c1*x2 + + +def hermval2d(x, y, c): + """ + Evaluate a 2-D Hermite series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * H_i(x) * H_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + hermval, hermgrid2d, hermval3d, hermgrid3d + + Examples + -------- + >>> from numpy.polynomial.hermite import hermval2d + >>> x = [1, 2] + >>> y = [4, 5] + >>> c = [[1, 2, 3], [4, 5, 6]] + >>> hermval2d(x, y, c) + array([1035., 2883.]) + + """ + return pu._valnd(hermval, c, x, y) + + +def hermgrid2d(x, y, c): + """ + Evaluate a 2-D Hermite series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * H_i(a) * H_j(b) + + where the points ``(a, b)`` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermval, hermval2d, hermval3d, hermgrid3d + + Examples + -------- + >>> from numpy.polynomial.hermite import hermgrid2d + >>> x = [1, 2, 3] + >>> y = [4, 5] + >>> c = [[1, 2, 3], [4, 5, 6]] + >>> hermgrid2d(x, y, c) + array([[1035., 1599.], + [1867., 2883.], + [2699., 4167.]]) + + """ + return pu._gridnd(hermval, c, x, y) + + +def hermval3d(x, y, z, c): + """ + Evaluate a 3-D Hermite series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * H_i(x) * H_j(y) * H_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + ``(x, y, z)``, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + hermval, hermval2d, hermgrid2d, hermgrid3d + + Examples + -------- + >>> from numpy.polynomial.hermite import hermval3d + >>> x = [1, 2] + >>> y = [4, 5] + >>> z = [6, 7] + >>> c = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]] + >>> hermval3d(x, y, z, c) + array([ 40077., 120131.]) + + """ + return pu._valnd(hermval, c, x, y, z) + + +def hermgrid3d(x, y, z, c): + """ + Evaluate a 3-D Hermite series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * H_i(a) * H_j(b) * H_k(c) + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermval, hermval2d, hermgrid2d, hermval3d + + Examples + -------- + >>> from numpy.polynomial.hermite import hermgrid3d + >>> x = [1, 2] + >>> y = [4, 5] + >>> z = [6, 7] + >>> c = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]] + >>> hermgrid3d(x, y, z, c) + array([[[ 40077., 54117.], + [ 49293., 66561.]], + [[ 72375., 97719.], + [ 88975., 120131.]]]) + + """ + return pu._gridnd(hermval, c, x, y, z) + + +def hermvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = H_i(x), + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Hermite polynomial. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + array ``V = hermvander(x, n)``, then ``np.dot(V, c)`` and + ``hermval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Hermite series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Hermite polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite import hermvander + >>> x = np.array([-1, 0, 1]) + >>> hermvander(x, 3) + array([[ 1., -2., 2., 4.], + [ 1., 0., -2., -0.], + [ 1., 2., 2., -4.]]) + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + x2 = x*2 + v[1] = x2 + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x2 - v[i-2]*(2*(i - 1))) + return np.moveaxis(v, 0, -1) + + +def hermvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = H_i(x) * H_j(y), + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the degrees of + the Hermite polynomials. + + If ``V = hermvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``hermval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Hermite + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + hermvander, hermvander3d, hermval2d, hermval3d + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite import hermvander2d + >>> x = np.array([-1, 0, 1]) + >>> y = np.array([-1, 0, 1]) + >>> hermvander2d(x, y, [2, 2]) + array([[ 1., -2., 2., -2., 4., -4., 2., -4., 4.], + [ 1., 0., -2., 0., 0., -0., -2., -0., 4.], + [ 1., 2., 2., 2., 4., 4., 2., 4., 4.]]) + + """ + return pu._vander_nd_flat((hermvander, hermvander), (x, y), deg) + + +def hermvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = H_i(x)*H_j(y)*H_k(z), + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the degrees of the Hermite polynomials. + + If ``V = hermvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``hermval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Hermite + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + hermvander, hermvander3d, hermval2d, hermval3d + + Examples + -------- + >>> from numpy.polynomial.hermite import hermvander3d + >>> x = np.array([-1, 0, 1]) + >>> y = np.array([-1, 0, 1]) + >>> z = np.array([-1, 0, 1]) + >>> hermvander3d(x, y, z, [0, 1, 2]) + array([[ 1., -2., 2., -2., 4., -4.], + [ 1., 0., -2., 0., 0., -0.], + [ 1., 2., 2., 2., 4., 4.]]) + + """ + return pu._vander_nd_flat((hermvander, hermvander, hermvander), (x, y, z), deg) + + +def hermfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Hermite series to data. + + Return the coefficients of a Hermite series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * H_1(x) + ... + c_n * H_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Hermite coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.polynomial.polyfit + numpy.polynomial.hermite_e.hermefit + hermval : Evaluates a Hermite series. + hermvander : Vandermonde matrix of Hermite series. + hermweight : Hermite weight function + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Hermite series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `~exceptions.RankWarning` will be issued. This means that + the coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Hermite series are probably most useful when the data can be + approximated by ``sqrt(w(x)) * p(x)``, where ``w(x)`` is the Hermite + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `hermweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite import hermfit, hermval + >>> x = np.linspace(-10, 10) + >>> rng = np.random.default_rng() + >>> err = rng.normal(scale=1./10, size=len(x)) + >>> y = hermval(x, [1, 2, 3]) + err + >>> hermfit(x, y, 2) + array([1.02294967, 2.00016403, 2.99994614]) # may vary + + """ + return pu._fit(hermvander, x, y, deg, rcond, full, w) + + +def hermcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an Hermite basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + + Examples + -------- + >>> from numpy.polynomial.hermite import hermcompanion + >>> hermcompanion([1, 0, 1]) + array([[0. , 0.35355339], + [0.70710678, 0. ]]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-.5*c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.hstack((1., 1./np.sqrt(2.*np.arange(n - 1, 0, -1)))) + scl = np.multiply.accumulate(scl)[::-1] + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.sqrt(.5*np.arange(1, n)) + bot[...] = top + mat[:, -1] -= scl*c[:-1]/(2.0*c[-1]) + return mat + + +def hermroots(c): + """ + Compute the roots of a Hermite series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * H_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Hermite series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermroots, hermfromroots + >>> coef = hermfromroots([-1, 0, 1]) + >>> coef + array([0. , 0.25 , 0. , 0.125]) + >>> hermroots(coef) + array([-1.00000000e+00, -1.38777878e-17, 1.00000000e+00]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-.5*c[0]/c[1]]) + + # rotated companion matrix reduces error + m = hermcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def _normed_hermite_n(x, n): + """ + Evaluate a normalized Hermite polynomial. + + Compute the value of the normalized Hermite polynomial of degree ``n`` + at the points ``x``. + + + Parameters + ---------- + x : ndarray of double. + Points at which to evaluate the function + n : int + Degree of the normalized Hermite function to be evaluated. + + Returns + ------- + values : ndarray + The shape of the return value is described above. + + Notes + ----- + This function is needed for finding the Gauss points and integration + weights for high degrees. The values of the standard Hermite functions + overflow when n >= 207. + + """ + if n == 0: + return np.full(x.shape, 1/np.sqrt(np.sqrt(np.pi))) + + c0 = 0. + c1 = 1./np.sqrt(np.sqrt(np.pi)) + nd = float(n) + for i in range(n - 1): + tmp = c0 + c0 = -c1*np.sqrt((nd - 1.)/nd) + c1 = tmp + c1*x*np.sqrt(2./nd) + nd = nd - 1.0 + return c0 + c1*x*np.sqrt(2) + + +def hermgauss(deg): + """ + Gauss-Hermite quadrature. + + Computes the sample points and weights for Gauss-Hermite quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-\\inf, \\inf]` + with the weight function :math:`f(x) = \\exp(-x^2)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (H'_n(x_k) * H_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`H_n`, and then scaling the results to get + the right value when integrating 1. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermgauss + >>> hermgauss(2) + (array([-0.70710678, 0.70710678]), array([0.88622693, 0.88622693])) + + """ + ideg = pu._as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1], dtype=np.float64) + m = hermcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = _normed_hermite_n(x, ideg) + df = _normed_hermite_n(x, ideg - 1) * np.sqrt(2*ideg) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = _normed_hermite_n(x, ideg - 1) + fm /= np.abs(fm).max() + w = 1/(fm * fm) + + # for Hermite we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= np.sqrt(np.pi) / w.sum() + + return x, w + + +def hermweight(x): + """ + Weight function of the Hermite polynomials. + + The weight function is :math:`\\exp(-x^2)` and the interval of + integration is :math:`[-\\inf, \\inf]`. the Hermite polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite import hermweight + >>> x = np.arange(-2, 2) + >>> hermweight(x) + array([0.01831564, 0.36787944, 1. , 0.36787944]) + + """ + w = np.exp(-x**2) + return w + + +# +# Hermite series class +# + +class Hermite(ABCPolyBase): + """An Hermite series class. + + The Hermite class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + Hermite coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*H_0(x) + 2*H_1(x) + 3*H_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(hermadd) + _sub = staticmethod(hermsub) + _mul = staticmethod(hermmul) + _div = staticmethod(hermdiv) + _pow = staticmethod(hermpow) + _val = staticmethod(hermval) + _int = staticmethod(hermint) + _der = staticmethod(hermder) + _fit = staticmethod(hermfit) + _line = staticmethod(hermline) + _roots = staticmethod(hermroots) + _fromroots = staticmethod(hermfromroots) + + # Virtual properties + domain = np.array(hermdomain) + window = np.array(hermdomain) + basis_name = 'H' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.pyi new file mode 100644 index 0000000000000000000000000000000000000000..07db43d0c0006601781cd24ee3269ae2f32a0445 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite.pyi @@ -0,0 +1,106 @@ +from typing import Any, Final, Literal as L, TypeVar + +import numpy as np + +from ._polybase import ABCPolyBase +from ._polytypes import ( + _Array1, + _Array2, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncGauss, + _FuncInteg, + _FuncLine, + _FuncPoly2Ortho, + _FuncPow, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVal2D, + _FuncVal3D, + _FuncValFromRoots, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncWeight, +) +from .polyutils import trimcoef as hermtrim + +__all__ = [ + "hermzero", + "hermone", + "hermx", + "hermdomain", + "hermline", + "hermadd", + "hermsub", + "hermmulx", + "hermmul", + "hermdiv", + "hermpow", + "hermval", + "hermder", + "hermint", + "herm2poly", + "poly2herm", + "hermfromroots", + "hermvander", + "hermfit", + "hermtrim", + "hermroots", + "Hermite", + "hermval2d", + "hermval3d", + "hermgrid2d", + "hermgrid3d", + "hermvander2d", + "hermvander3d", + "hermcompanion", + "hermgauss", + "hermweight", +] + +poly2herm: _FuncPoly2Ortho[L["poly2herm"]] +herm2poly: _FuncUnOp[L["herm2poly"]] + +hermdomain: Final[_Array2[np.float64]] +hermzero: Final[_Array1[np.int_]] +hermone: Final[_Array1[np.int_]] +hermx: Final[_Array2[np.int_]] + +hermline: _FuncLine[L["hermline"]] +hermfromroots: _FuncFromRoots[L["hermfromroots"]] +hermadd: _FuncBinOp[L["hermadd"]] +hermsub: _FuncBinOp[L["hermsub"]] +hermmulx: _FuncUnOp[L["hermmulx"]] +hermmul: _FuncBinOp[L["hermmul"]] +hermdiv: _FuncBinOp[L["hermdiv"]] +hermpow: _FuncPow[L["hermpow"]] +hermder: _FuncDer[L["hermder"]] +hermint: _FuncInteg[L["hermint"]] +hermval: _FuncVal[L["hermval"]] +hermval2d: _FuncVal2D[L["hermval2d"]] +hermval3d: _FuncVal3D[L["hermval3d"]] +hermvalfromroots: _FuncValFromRoots[L["hermvalfromroots"]] +hermgrid2d: _FuncVal2D[L["hermgrid2d"]] +hermgrid3d: _FuncVal3D[L["hermgrid3d"]] +hermvander: _FuncVander[L["hermvander"]] +hermvander2d: _FuncVander2D[L["hermvander2d"]] +hermvander3d: _FuncVander3D[L["hermvander3d"]] +hermfit: _FuncFit[L["hermfit"]] +hermcompanion: _FuncCompanion[L["hermcompanion"]] +hermroots: _FuncRoots[L["hermroots"]] + +_ND = TypeVar("_ND", bound=Any) +def _normed_hermite_n( + x: np.ndarray[_ND, np.dtype[np.float64]], + n: int | np.intp, +) -> np.ndarray[_ND, np.dtype[np.float64]]: ... + +hermgauss: _FuncGauss[L["hermgauss"]] +hermweight: _FuncWeight[L["hermweight"]] + +class Hermite(ABCPolyBase[L["H"]]): ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.py b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.py new file mode 100644 index 0000000000000000000000000000000000000000..c820760ef75c1db162b0a6e0897c88ba18582464 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.py @@ -0,0 +1,1642 @@ +""" +=================================================================== +HermiteE Series, "Probabilists" (:mod:`numpy.polynomial.hermite_e`) +=================================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Hermite_e series, including a `HermiteE` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + HermiteE + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + hermedomain + hermezero + hermeone + hermex + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + hermeadd + hermesub + hermemulx + hermemul + hermediv + hermepow + hermeval + hermeval2d + hermeval3d + hermegrid2d + hermegrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + hermeder + hermeint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + hermefromroots + hermeroots + hermevander + hermevander2d + hermevander3d + hermegauss + hermeweight + hermecompanion + hermefit + hermetrim + hermeline + herme2poly + poly2herme + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'hermezero', 'hermeone', 'hermex', 'hermedomain', 'hermeline', + 'hermeadd', 'hermesub', 'hermemulx', 'hermemul', 'hermediv', + 'hermepow', 'hermeval', 'hermeder', 'hermeint', 'herme2poly', + 'poly2herme', 'hermefromroots', 'hermevander', 'hermefit', 'hermetrim', + 'hermeroots', 'HermiteE', 'hermeval2d', 'hermeval3d', 'hermegrid2d', + 'hermegrid3d', 'hermevander2d', 'hermevander3d', 'hermecompanion', + 'hermegauss', 'hermeweight'] + +hermetrim = pu.trimcoef + + +def poly2herme(pol): + """ + poly2herme(pol) + + Convert a polynomial to a Hermite series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Hermite series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Hermite + series. + + See Also + -------- + herme2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite_e import poly2herme + >>> poly2herme(np.arange(4)) + array([ 2., 10., 2., 3.]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = hermeadd(hermemulx(res), pol[i]) + return res + + +def herme2poly(c): + """ + Convert a Hermite series to a polynomial. + + Convert an array representing the coefficients of a Hermite series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Hermite series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2herme + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import herme2poly + >>> herme2poly([ 2., 10., 2., 3.]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + if n == 2: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1*(i - 1)) + c1 = polyadd(tmp, polymulx(c1)) + return polyadd(c0, polymulx(c1)) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Hermite +hermedomain = np.array([-1., 1.]) + +# Hermite coefficients representing zero. +hermezero = np.array([0]) + +# Hermite coefficients representing one. +hermeone = np.array([1]) + +# Hermite coefficients representing the identity x. +hermex = np.array([0, 1]) + + +def hermeline(off, scl): + """ + Hermite series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Hermite series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeline + >>> from numpy.polynomial.hermite_e import hermeline, hermeval + >>> hermeval(0,hermeline(3, 2)) + 3.0 + >>> hermeval(1,hermeline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def hermefromroots(roots): + """ + Generate a HermiteE series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in HermiteE form, where the :math:`r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * He_1(x) + ... + c_n * He_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in HermiteE form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.chebyshev.chebfromroots + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermefromroots, hermeval + >>> coef = hermefromroots((-1, 0, 1)) + >>> hermeval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = hermefromroots((-1j, 1j)) + >>> hermeval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(hermeline, hermemul, roots) + + +def hermeadd(c1, c2): + """ + Add one Hermite series to another. + + Returns the sum of two Hermite series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Hermite series of their sum. + + See Also + -------- + hermesub, hermemulx, hermemul, hermediv, hermepow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Hermite series + is a Hermite series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeadd + >>> hermeadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + """ + return pu._add(c1, c2) + + +def hermesub(c1, c2): + """ + Subtract one Hermite series from another. + + Returns the difference of two Hermite series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their difference. + + See Also + -------- + hermeadd, hermemulx, hermemul, hermediv, hermepow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Hermite + series is a Hermite series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermesub + >>> hermesub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def hermemulx(c): + """Multiply a Hermite series by x. + + Multiply the Hermite series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + hermeadd, hermesub, hermemul, hermediv, hermepow + + Notes + ----- + The multiplication uses the recursion relationship for Hermite + polynomials in the form + + .. math:: + + xP_i(x) = (P_{i + 1}(x) + iP_{i - 1}(x))) + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermemulx + >>> hermemulx([1, 2, 3]) + array([2., 7., 2., 3.]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + for i in range(1, len(c)): + prd[i + 1] = c[i] + prd[i - 1] += c[i]*i + return prd + + +def hermemul(c1, c2): + """ + Multiply one Hermite series by another. + + Returns the product of two Hermite series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their product. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermediv, hermepow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Hermite polynomial basis set. Thus, to express + the product as a Hermite series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermemul + >>> hermemul([1, 2, 3], [0, 1, 2]) + array([14., 15., 28., 7., 6.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = hermesub(c[-i]*xs, c1*(nd - 1)) + c1 = hermeadd(tmp, hermemulx(c1)) + return hermeadd(c0, hermemulx(c1)) + + +def hermediv(c1, c2): + """ + Divide one Hermite series by another. + + Returns the quotient-with-remainder of two Hermite series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Hermite series coefficients representing the quotient and + remainder. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermemul, hermepow + + Notes + ----- + In general, the (polynomial) division of one Hermite series by another + results in quotient and remainder terms that are not in the Hermite + polynomial basis set. Thus, to express these results as a Hermite + series, it is necessary to "reproject" the results onto the Hermite + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermediv + >>> hermediv([ 14., 15., 28., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> hermediv([ 15., 17., 28., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 2.])) + + """ + return pu._div(hermemul, c1, c2) + + +def hermepow(c, pow, maxpower=16): + """Raise a Hermite series to a power. + + Returns the Hermite series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Hermite series of power. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermemul, hermediv + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermepow + >>> hermepow([1, 2, 3], 2) + array([23., 28., 46., 12., 9.]) + + """ + return pu._pow(hermemul, c, pow, maxpower) + + +def hermeder(c, m=1, scl=1, axis=0): + """ + Differentiate a Hermite_e series. + + Returns the series coefficients `c` differentiated `m` times along + `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*He_0 + 2*He_1 + 3*He_2`` + while [[1,2],[1,2]] represents ``1*He_0(x)*He_0(y) + 1*He_1(x)*He_0(y) + + 2*He_0(x)*He_1(y) + 2*He_1(x)*He_1(y)`` if axis=0 is ``x`` and axis=1 + is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite_e series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Hermite series of the derivative. + + See Also + -------- + hermeint + + Notes + ----- + In general, the result of differentiating a Hermite series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeder + >>> hermeder([ 1., 1., 1., 1.]) + array([1., 2., 3.]) + >>> hermeder([-0.25, 1., 1./2., 1./3., 1./4 ], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + return c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 0, -1): + der[j - 1] = j*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermeint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Hermite_e series. + + Returns the Hermite_e series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``H_0 + 2*H_1 + 3*H_2`` while [[1,2],[1,2]] + represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + 2*H_0(x)*H_1(y) + + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite_e series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + Hermite_e series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + hermeder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeint + >>> hermeint([1, 2, 3]) # integrate once, value 0 at 0. + array([1., 1., 1., 1.]) + >>> hermeint([1, 2, 3], m=2) # integrate twice, value & deriv 0 at 0 + array([-0.25 , 1. , 0.5 , 0.33333333, 0.25 ]) # may vary + >>> hermeint([1, 2, 3], k=1) # integrate once, value 1 at 0. + array([2., 1., 1., 1.]) + >>> hermeint([1, 2, 3], lbnd=-1) # integrate once, value 0 at -1 + array([-1., 1., 1., 1.]) + >>> hermeint([1, 2, 3], m=2, k=[1, 2], lbnd=-1) + array([ 1.83333333, 0. , 0.5 , 0.33333333, 0.25 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + for j in range(1, n): + tmp[j + 1] = c[j]/(j + 1) + tmp[0] += k[i] - hermeval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermeval(x, c, tensor=True): + """ + Evaluate an HermiteE series at points x. + + If `c` is of length ``n + 1``, this function returns the value: + + .. math:: p(x) = c_0 * He_0(x) + c_1 * He_1(x) + ... + c_n * He_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + hermeval2d, hermegrid2d, hermeval3d, hermegrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeval + >>> coef = [1,2,3] + >>> hermeval(1, coef) + 3.0 + >>> hermeval([[1,2],[3,4]], coef) + array([[ 3., 14.], + [31., 54.]]) + + """ + c = np.array(c, ndmin=1, copy=None) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - c1*(nd - 1) + c1 = tmp + c1*x + return c0 + c1*x + + +def hermeval2d(x, y, c): + """ + Evaluate a 2-D HermiteE series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * He_i(x) * He_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + hermeval, hermegrid2d, hermeval3d, hermegrid3d + """ + return pu._valnd(hermeval, c, x, y) + + +def hermegrid2d(x, y, c): + """ + Evaluate a 2-D HermiteE series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * H_i(a) * H_j(b) + + where the points ``(a, b)`` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermeval, hermeval2d, hermeval3d, hermegrid3d + """ + return pu._gridnd(hermeval, c, x, y) + + +def hermeval3d(x, y, z, c): + """ + Evaluate a 3-D Hermite_e series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * He_i(x) * He_j(y) * He_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + hermeval, hermeval2d, hermegrid2d, hermegrid3d + """ + return pu._valnd(hermeval, c, x, y, z) + + +def hermegrid3d(x, y, z, c): + """ + Evaluate a 3-D HermiteE series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * He_i(a) * He_j(b) * He_k(c) + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermeval, hermeval2d, hermegrid2d, hermeval3d + """ + return pu._gridnd(hermeval, c, x, y, z) + + +def hermevander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = He_i(x), + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the degree of the HermiteE polynomial. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + array ``V = hermevander(x, n)``, then ``np.dot(V, c)`` and + ``hermeval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of HermiteE series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding HermiteE polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite_e import hermevander + >>> x = np.array([-1, 0, 1]) + >>> hermevander(x, 3) + array([[ 1., -1., 0., 2.], + [ 1., 0., -1., -0.], + [ 1., 1., 0., -2.]]) + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x - v[i-2]*(i - 1)) + return np.moveaxis(v, 0, -1) + + +def hermevander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = He_i(x) * He_j(y), + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the degrees of + the HermiteE polynomials. + + If ``V = hermevander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``hermeval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D HermiteE + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + hermevander, hermevander3d, hermeval2d, hermeval3d + """ + return pu._vander_nd_flat((hermevander, hermevander), (x, y), deg) + + +def hermevander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then Hehe pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = He_i(x)*He_j(y)*He_k(z), + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the degrees of the HermiteE polynomials. + + If ``V = hermevander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``hermeval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D HermiteE + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + hermevander, hermevander3d, hermeval2d, hermeval3d + """ + return pu._vander_nd_flat((hermevander, hermevander, hermevander), (x, y, z), deg) + + +def hermefit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Hermite series to data. + + Return the coefficients of a HermiteE series of degree `deg` that is + the least squares fit to the data values `y` given at points `x`. If + `y` is 1-D the returned coefficients will also be 1-D. If `y` is 2-D + multiple fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * He_1(x) + ... + c_n * He_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Hermite coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full = False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.polynomial.polyfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.laguerre.lagfit + hermeval : Evaluates a Hermite series. + hermevander : pseudo Vandermonde matrix of Hermite series. + hermeweight : HermiteE weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the HermiteE series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the pseudo Vandermonde matrix of `x`, the elements of `c` + are the coefficients to be solved for, and the elements of `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `~exceptions.RankWarning` will be issued. This means that + the coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using HermiteE series are probably most useful when the data can + be approximated by ``sqrt(w(x)) * p(x)``, where ``w(x)`` is the HermiteE + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `hermeweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.hermite_e import hermefit, hermeval + >>> x = np.linspace(-10, 10) + >>> rng = np.random.default_rng() + >>> err = rng.normal(scale=1./10, size=len(x)) + >>> y = hermeval(x, [1, 2, 3]) + err + >>> hermefit(x, y, 2) + array([1.02284196, 2.00032805, 2.99978457]) # may vary + + """ + return pu._fit(hermevander, x, y, deg, rcond, full, w) + + +def hermecompanion(c): + """ + Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an HermiteE basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of HermiteE series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.hstack((1., 1./np.sqrt(np.arange(n - 1, 0, -1)))) + scl = np.multiply.accumulate(scl)[::-1] + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.sqrt(np.arange(1, n)) + bot[...] = top + mat[:, -1] -= scl*c[:-1]/c[-1] + return mat + + +def hermeroots(c): + """ + Compute the roots of a HermiteE series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * He_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.chebyshev.chebroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The HermiteE series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeroots, hermefromroots + >>> coef = hermefromroots([-1, 0, 1]) + >>> coef + array([0., 2., 0., 1.]) + >>> hermeroots(coef) + array([-1., 0., 1.]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = hermecompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def _normed_hermite_e_n(x, n): + """ + Evaluate a normalized HermiteE polynomial. + + Compute the value of the normalized HermiteE polynomial of degree ``n`` + at the points ``x``. + + + Parameters + ---------- + x : ndarray of double. + Points at which to evaluate the function + n : int + Degree of the normalized HermiteE function to be evaluated. + + Returns + ------- + values : ndarray + The shape of the return value is described above. + + Notes + ----- + This function is needed for finding the Gauss points and integration + weights for high degrees. The values of the standard HermiteE functions + overflow when n >= 207. + + """ + if n == 0: + return np.full(x.shape, 1/np.sqrt(np.sqrt(2*np.pi))) + + c0 = 0. + c1 = 1./np.sqrt(np.sqrt(2*np.pi)) + nd = float(n) + for i in range(n - 1): + tmp = c0 + c0 = -c1*np.sqrt((nd - 1.)/nd) + c1 = tmp + c1*x*np.sqrt(1./nd) + nd = nd - 1.0 + return c0 + c1*x + + +def hermegauss(deg): + """ + Gauss-HermiteE quadrature. + + Computes the sample points and weights for Gauss-HermiteE quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-\\inf, \\inf]` + with the weight function :math:`f(x) = \\exp(-x^2/2)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (He'_n(x_k) * He_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`He_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = hermecompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = _normed_hermite_e_n(x, ideg) + df = _normed_hermite_e_n(x, ideg - 1) * np.sqrt(ideg) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = _normed_hermite_e_n(x, ideg - 1) + fm /= np.abs(fm).max() + w = 1/(fm * fm) + + # for Hermite_e we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= np.sqrt(2*np.pi) / w.sum() + + return x, w + + +def hermeweight(x): + """Weight function of the Hermite_e polynomials. + + The weight function is :math:`\\exp(-x^2/2)` and the interval of + integration is :math:`[-\\inf, \\inf]`. the HermiteE polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + """ + w = np.exp(-.5*x**2) + return w + + +# +# HermiteE series class +# + +class HermiteE(ABCPolyBase): + """An HermiteE series class. + + The HermiteE class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + HermiteE coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*He_0(x) + 2*He_1(X) + 3*He_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(hermeadd) + _sub = staticmethod(hermesub) + _mul = staticmethod(hermemul) + _div = staticmethod(hermediv) + _pow = staticmethod(hermepow) + _val = staticmethod(hermeval) + _int = staticmethod(hermeint) + _der = staticmethod(hermeder) + _fit = staticmethod(hermefit) + _line = staticmethod(hermeline) + _roots = staticmethod(hermeroots) + _fromroots = staticmethod(hermefromroots) + + # Virtual properties + domain = np.array(hermedomain) + window = np.array(hermedomain) + basis_name = 'He' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.pyi new file mode 100644 index 0000000000000000000000000000000000000000..94ad7248f268b9d4e4de1685063187c94db25fd7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/hermite_e.pyi @@ -0,0 +1,106 @@ +from typing import Any, Final, Literal as L, TypeVar + +import numpy as np + +from ._polybase import ABCPolyBase +from ._polytypes import ( + _Array1, + _Array2, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncGauss, + _FuncInteg, + _FuncLine, + _FuncPoly2Ortho, + _FuncPow, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVal2D, + _FuncVal3D, + _FuncValFromRoots, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncWeight, +) +from .polyutils import trimcoef as hermetrim + +__all__ = [ + "hermezero", + "hermeone", + "hermex", + "hermedomain", + "hermeline", + "hermeadd", + "hermesub", + "hermemulx", + "hermemul", + "hermediv", + "hermepow", + "hermeval", + "hermeder", + "hermeint", + "herme2poly", + "poly2herme", + "hermefromroots", + "hermevander", + "hermefit", + "hermetrim", + "hermeroots", + "HermiteE", + "hermeval2d", + "hermeval3d", + "hermegrid2d", + "hermegrid3d", + "hermevander2d", + "hermevander3d", + "hermecompanion", + "hermegauss", + "hermeweight", +] + +poly2herme: _FuncPoly2Ortho[L["poly2herme"]] +herme2poly: _FuncUnOp[L["herme2poly"]] + +hermedomain: Final[_Array2[np.float64]] +hermezero: Final[_Array1[np.int_]] +hermeone: Final[_Array1[np.int_]] +hermex: Final[_Array2[np.int_]] + +hermeline: _FuncLine[L["hermeline"]] +hermefromroots: _FuncFromRoots[L["hermefromroots"]] +hermeadd: _FuncBinOp[L["hermeadd"]] +hermesub: _FuncBinOp[L["hermesub"]] +hermemulx: _FuncUnOp[L["hermemulx"]] +hermemul: _FuncBinOp[L["hermemul"]] +hermediv: _FuncBinOp[L["hermediv"]] +hermepow: _FuncPow[L["hermepow"]] +hermeder: _FuncDer[L["hermeder"]] +hermeint: _FuncInteg[L["hermeint"]] +hermeval: _FuncVal[L["hermeval"]] +hermeval2d: _FuncVal2D[L["hermeval2d"]] +hermeval3d: _FuncVal3D[L["hermeval3d"]] +hermevalfromroots: _FuncValFromRoots[L["hermevalfromroots"]] +hermegrid2d: _FuncVal2D[L["hermegrid2d"]] +hermegrid3d: _FuncVal3D[L["hermegrid3d"]] +hermevander: _FuncVander[L["hermevander"]] +hermevander2d: _FuncVander2D[L["hermevander2d"]] +hermevander3d: _FuncVander3D[L["hermevander3d"]] +hermefit: _FuncFit[L["hermefit"]] +hermecompanion: _FuncCompanion[L["hermecompanion"]] +hermeroots: _FuncRoots[L["hermeroots"]] + +_ND = TypeVar("_ND", bound=Any) +def _normed_hermite_e_n( + x: np.ndarray[_ND, np.dtype[np.float64]], + n: int | np.intp, +) -> np.ndarray[_ND, np.dtype[np.float64]]: ... + +hermegauss: _FuncGauss[L["hermegauss"]] +hermeweight: _FuncWeight[L["hermeweight"]] + +class HermiteE(ABCPolyBase[L["He"]]): ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.py b/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.py new file mode 100644 index 0000000000000000000000000000000000000000..b2cc5817c30cb892f58f1c366746b5967670d2ad --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.py @@ -0,0 +1,1675 @@ +""" +================================================== +Laguerre Series (:mod:`numpy.polynomial.laguerre`) +================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Laguerre series, including a `Laguerre` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Laguerre + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + lagdomain + lagzero + lagone + lagx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + lagadd + lagsub + lagmulx + lagmul + lagdiv + lagpow + lagval + lagval2d + lagval3d + laggrid2d + laggrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + lagder + lagint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + lagfromroots + lagroots + lagvander + lagvander2d + lagvander3d + laggauss + lagweight + lagcompanion + lagfit + lagtrim + lagline + lag2poly + poly2lag + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'lagzero', 'lagone', 'lagx', 'lagdomain', 'lagline', 'lagadd', + 'lagsub', 'lagmulx', 'lagmul', 'lagdiv', 'lagpow', 'lagval', 'lagder', + 'lagint', 'lag2poly', 'poly2lag', 'lagfromroots', 'lagvander', + 'lagfit', 'lagtrim', 'lagroots', 'Laguerre', 'lagval2d', 'lagval3d', + 'laggrid2d', 'laggrid3d', 'lagvander2d', 'lagvander3d', 'lagcompanion', + 'laggauss', 'lagweight'] + +lagtrim = pu.trimcoef + + +def poly2lag(pol): + """ + poly2lag(pol) + + Convert a polynomial to a Laguerre series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Laguerre series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Laguerre + series. + + See Also + -------- + lag2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.laguerre import poly2lag + >>> poly2lag(np.arange(4)) + array([ 23., -63., 58., -18.]) + + """ + [pol] = pu.as_series([pol]) + res = 0 + for p in pol[::-1]: + res = lagadd(lagmulx(res), p) + return res + + +def lag2poly(c): + """ + Convert a Laguerre series to a polynomial. + + Convert an array representing the coefficients of a Laguerre series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Laguerre series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2lag + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lag2poly + >>> lag2poly([ 23., -63., 58., -18.]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], (c1*(i - 1))/i) + c1 = polyadd(tmp, polysub((2*i - 1)*c1, polymulx(c1))/i) + return polyadd(c0, polysub(c1, polymulx(c1))) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Laguerre +lagdomain = np.array([0., 1.]) + +# Laguerre coefficients representing zero. +lagzero = np.array([0]) + +# Laguerre coefficients representing one. +lagone = np.array([1]) + +# Laguerre coefficients representing the identity x. +lagx = np.array([1, -1]) + + +def lagline(off, scl): + """ + Laguerre series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Laguerre series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagline, lagval + >>> lagval(0,lagline(3, 2)) + 3.0 + >>> lagval(1,lagline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off + scl, -scl]) + else: + return np.array([off]) + + +def lagfromroots(roots): + """ + Generate a Laguerre series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Laguerre form, where the :math:`r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Laguerre form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagfromroots, lagval + >>> coef = lagfromroots((-1, 0, 1)) + >>> lagval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = lagfromroots((-1j, 1j)) + >>> lagval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(lagline, lagmul, roots) + + +def lagadd(c1, c2): + """ + Add one Laguerre series to another. + + Returns the sum of two Laguerre series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Laguerre series of their sum. + + See Also + -------- + lagsub, lagmulx, lagmul, lagdiv, lagpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Laguerre series + is a Laguerre series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagadd + >>> lagadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + """ + return pu._add(c1, c2) + + +def lagsub(c1, c2): + """ + Subtract one Laguerre series from another. + + Returns the difference of two Laguerre series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Laguerre series coefficients representing their difference. + + See Also + -------- + lagadd, lagmulx, lagmul, lagdiv, lagpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Laguerre + series is a Laguerre series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagsub + >>> lagsub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def lagmulx(c): + """Multiply a Laguerre series by x. + + Multiply the Laguerre series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + lagadd, lagsub, lagmul, lagdiv, lagpow + + Notes + ----- + The multiplication uses the recursion relationship for Laguerre + polynomials in the form + + .. math:: + + xP_i(x) = (-(i + 1)*P_{i + 1}(x) + (2i + 1)P_{i}(x) - iP_{i - 1}(x)) + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagmulx + >>> lagmulx([1, 2, 3]) + array([-1., -1., 11., -9.]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0] + prd[1] = -c[0] + for i in range(1, len(c)): + prd[i + 1] = -c[i]*(i + 1) + prd[i] += c[i]*(2*i + 1) + prd[i - 1] -= c[i]*i + return prd + + +def lagmul(c1, c2): + """ + Multiply one Laguerre series by another. + + Returns the product of two Laguerre series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Laguerre series coefficients representing their product. + + See Also + -------- + lagadd, lagsub, lagmulx, lagdiv, lagpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Laguerre polynomial basis set. Thus, to express + the product as a Laguerre series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagmul + >>> lagmul([1, 2, 3], [0, 1, 2]) + array([ 8., -13., 38., -51., 36.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = lagsub(c[-i]*xs, (c1*(nd - 1))/nd) + c1 = lagadd(tmp, lagsub((2*nd - 1)*c1, lagmulx(c1))/nd) + return lagadd(c0, lagsub(c1, lagmulx(c1))) + + +def lagdiv(c1, c2): + """ + Divide one Laguerre series by another. + + Returns the quotient-with-remainder of two Laguerre series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Laguerre series coefficients representing the quotient and + remainder. + + See Also + -------- + lagadd, lagsub, lagmulx, lagmul, lagpow + + Notes + ----- + In general, the (polynomial) division of one Laguerre series by another + results in quotient and remainder terms that are not in the Laguerre + polynomial basis set. Thus, to express these results as a Laguerre + series, it is necessary to "reproject" the results onto the Laguerre + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagdiv + >>> lagdiv([ 8., -13., 38., -51., 36.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> lagdiv([ 9., -12., 38., -51., 36.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 1.])) + + """ + return pu._div(lagmul, c1, c2) + + +def lagpow(c, pow, maxpower=16): + """Raise a Laguerre series to a power. + + Returns the Laguerre series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Laguerre series of power. + + See Also + -------- + lagadd, lagsub, lagmulx, lagmul, lagdiv + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagpow + >>> lagpow([1, 2, 3], 2) + array([ 14., -16., 56., -72., 54.]) + + """ + return pu._pow(lagmul, c, pow, maxpower) + + +def lagder(c, m=1, scl=1, axis=0): + """ + Differentiate a Laguerre series. + + Returns the Laguerre series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*L_0 + 2*L_1 + 3*L_2`` + while [[1,2],[1,2]] represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + + 2*L_0(x)*L_1(y) + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Laguerre series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Laguerre series of the derivative. + + See Also + -------- + lagint + + Notes + ----- + In general, the result of differentiating a Laguerre series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagder + >>> lagder([ 1., 1., 1., -3.]) + array([1., 2., 3.]) + >>> lagder([ 1., 0., 0., -4., 3.], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 1, -1): + der[j - 1] = -c[j] + c[j - 1] += c[j] + der[0] = -c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def lagint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Laguerre series. + + Returns the Laguerre series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``L_0 + 2*L_1 + 3*L_2`` while [[1,2],[1,2]] + represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + 2*L_0(x)*L_1(y) + + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + + Parameters + ---------- + c : array_like + Array of Laguerre series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + Laguerre series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + lagder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagint + >>> lagint([1,2,3]) + array([ 1., 1., 1., -3.]) + >>> lagint([1,2,3], m=2) + array([ 1., 0., 0., -4., 3.]) + >>> lagint([1,2,3], k=1) + array([ 2., 1., 1., -3.]) + >>> lagint([1,2,3], lbnd=-1) + array([11.5, 1. , 1. , -3. ]) + >>> lagint([1,2], m=2, k=[1,2], lbnd=-1) + array([ 11.16666667, -5. , -3. , 2. ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0] + tmp[1] = -c[0] + for j in range(1, n): + tmp[j] += c[j] + tmp[j + 1] = -c[j] + tmp[0] += k[i] - lagval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def lagval(x, c, tensor=True): + """ + Evaluate a Laguerre series at points x. + + If `c` is of length ``n + 1``, this function returns the value: + + .. math:: p(x) = c_0 * L_0(x) + c_1 * L_1(x) + ... + c_n * L_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + lagval2d, laggrid2d, lagval3d, laggrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagval + >>> coef = [1, 2, 3] + >>> lagval(1, coef) + -0.5 + >>> lagval([[1, 2],[3, 4]], coef) + array([[-0.5, -4. ], + [-4.5, -2. ]]) + + """ + c = np.array(c, ndmin=1, copy=None) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - (c1*(nd - 1))/nd + c1 = tmp + (c1*((2*nd - 1) - x))/nd + return c0 + c1*(1 - x) + + +def lagval2d(x, y, c): + """ + Evaluate a 2-D Laguerre series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * L_i(x) * L_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + lagval, laggrid2d, lagval3d, laggrid3d + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagval2d + >>> c = [[1, 2],[3, 4]] + >>> lagval2d(1, 1, c) + 1.0 + """ + return pu._valnd(lagval, c, x, y) + + +def laggrid2d(x, y, c): + """ + Evaluate a 2-D Laguerre series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * L_i(a) * L_j(b) + + where the points ``(a, b)`` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + lagval, lagval2d, lagval3d, laggrid3d + + Examples + -------- + >>> from numpy.polynomial.laguerre import laggrid2d + >>> c = [[1, 2], [3, 4]] + >>> laggrid2d([0, 1], [0, 1], c) + array([[10., 4.], + [ 3., 1.]]) + + """ + return pu._gridnd(lagval, c, x, y) + + +def lagval3d(x, y, z, c): + """ + Evaluate a 3-D Laguerre series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * L_i(x) * L_j(y) * L_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + ``(x, y, z)``, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + lagval, lagval2d, laggrid2d, laggrid3d + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagval3d + >>> c = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> lagval3d(1, 1, 2, c) + -1.0 + + """ + return pu._valnd(lagval, c, x, y, z) + + +def laggrid3d(x, y, z, c): + """ + Evaluate a 3-D Laguerre series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * L_i(a) * L_j(b) * L_k(c) + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + lagval, lagval2d, laggrid2d, lagval3d + + Examples + -------- + >>> from numpy.polynomial.laguerre import laggrid3d + >>> c = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> laggrid3d([0, 1], [0, 1], [2, 4], c) + array([[[ -4., -44.], + [ -2., -18.]], + [[ -2., -14.], + [ -1., -5.]]]) + + """ + return pu._gridnd(lagval, c, x, y, z) + + +def lagvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = L_i(x) + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Laguerre polynomial. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + array ``V = lagvander(x, n)``, then ``np.dot(V, c)`` and + ``lagval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Laguerre series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Laguerre polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.laguerre import lagvander + >>> x = np.array([0, 1, 2]) + >>> lagvander(x, 3) + array([[ 1. , 1. , 1. , 1. ], + [ 1. , 0. , -0.5 , -0.66666667], + [ 1. , -1. , -1. , -0.33333333]]) + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = 1 - x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*(2*i - 1 - x) - v[i-2]*(i - 1))/i + return np.moveaxis(v, 0, -1) + + +def lagvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = L_i(x) * L_j(y), + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the degrees of + the Laguerre polynomials. + + If ``V = lagvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``lagval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Laguerre + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + lagvander, lagvander3d, lagval2d, lagval3d + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.laguerre import lagvander2d + >>> x = np.array([0]) + >>> y = np.array([2]) + >>> lagvander2d(x, y, [2, 1]) + array([[ 1., -1., 1., -1., 1., -1.]]) + + """ + return pu._vander_nd_flat((lagvander, lagvander), (x, y), deg) + + +def lagvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = L_i(x)*L_j(y)*L_k(z), + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the degrees of the Laguerre polynomials. + + If ``V = lagvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``lagval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Laguerre + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + lagvander, lagvander3d, lagval2d, lagval3d + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.laguerre import lagvander3d + >>> x = np.array([0]) + >>> y = np.array([2]) + >>> z = np.array([0]) + >>> lagvander3d(x, y, z, [2, 1, 3]) + array([[ 1., 1., 1., 1., -1., -1., -1., -1., 1., 1., 1., 1., -1., + -1., -1., -1., 1., 1., 1., 1., -1., -1., -1., -1.]]) + + """ + return pu._vander_nd_flat((lagvander, lagvander, lagvander), (x, y, z), deg) + + +def lagfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Laguerre series to data. + + Return the coefficients of a Laguerre series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x), + + where ``n`` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Laguerre coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column *k* of `y` are in column + *k*. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.legendre.legfit + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + lagval : Evaluates a Laguerre series. + lagvander : pseudo Vandermonde matrix of Laguerre series. + lagweight : Laguerre weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Laguerre series ``p`` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where ``V`` is the weighted pseudo Vandermonde matrix of `x`, ``c`` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of ``V``. + + If some of the singular values of `V` are so small that they are + neglected, then a `~exceptions.RankWarning` will be issued. This means that + the coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Laguerre series are probably most useful when the data can + be approximated by ``sqrt(w(x)) * p(x)``, where ``w(x)`` is the Laguerre + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `lagweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.laguerre import lagfit, lagval + >>> x = np.linspace(0, 10) + >>> rng = np.random.default_rng() + >>> err = rng.normal(scale=1./10, size=len(x)) + >>> y = lagval(x, [1, 2, 3]) + err + >>> lagfit(x, y, 2) + array([1.00578369, 1.99417356, 2.99827656]) # may vary + + """ + return pu._fit(lagvander, x, y, deg, rcond, full, w) + + +def lagcompanion(c): + """ + Return the companion matrix of c. + + The usual companion matrix of the Laguerre polynomials is already + symmetric when `c` is a basis Laguerre polynomial, so no scaling is + applied. + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Companion matrix of dimensions (deg, deg). + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagcompanion + >>> lagcompanion([1, 2, 3]) + array([[ 1. , -0.33333333], + [-1. , 4.33333333]]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[1 + c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + top = mat.reshape(-1)[1::n+1] + mid = mat.reshape(-1)[0::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = -np.arange(1, n) + mid[...] = 2.*np.arange(n) + 1. + bot[...] = top + mat[:, -1] += (c[:-1]/c[-1])*n + return mat + + +def lagroots(c): + """ + Compute the roots of a Laguerre series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * L_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Laguerre series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagroots, lagfromroots + >>> coef = lagfromroots([0, 1, 2]) + >>> coef + array([ 2., -8., 12., -6.]) + >>> lagroots(coef) + array([-4.4408921e-16, 1.0000000e+00, 2.0000000e+00]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([1 + c[0]/c[1]]) + + # rotated companion matrix reduces error + m = lagcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def laggauss(deg): + """ + Gauss-Laguerre quadrature. + + Computes the sample points and weights for Gauss-Laguerre quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[0, \\inf]` + with the weight function :math:`f(x) = \\exp(-x)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + The results have only been tested up to degree 100 higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (L'_n(x_k) * L_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`L_n`, and then scaling the results to get + the right value when integrating 1. + + Examples + -------- + >>> from numpy.polynomial.laguerre import laggauss + >>> laggauss(2) + (array([0.58578644, 3.41421356]), array([0.85355339, 0.14644661])) + + """ + ideg = pu._as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = lagcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = lagval(x, c) + df = lagval(x, lagder(c)) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = lagval(x, c[1:]) + fm /= np.abs(fm).max() + df /= np.abs(df).max() + w = 1/(fm * df) + + # scale w to get the right value, 1 in this case + w /= w.sum() + + return x, w + + +def lagweight(x): + """Weight function of the Laguerre polynomials. + + The weight function is :math:`exp(-x)` and the interval of integration + is :math:`[0, \\inf]`. The Laguerre polynomials are orthogonal, but not + normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagweight + >>> x = np.array([0, 1, 2]) + >>> lagweight(x) + array([1. , 0.36787944, 0.13533528]) + + """ + w = np.exp(-x) + return w + +# +# Laguerre series class +# + +class Laguerre(ABCPolyBase): + """A Laguerre series class. + + The Laguerre class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + Laguerre coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*L_0(x) + 2*L_1(X) + 3*L_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [0., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [0., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(lagadd) + _sub = staticmethod(lagsub) + _mul = staticmethod(lagmul) + _div = staticmethod(lagdiv) + _pow = staticmethod(lagpow) + _val = staticmethod(lagval) + _int = staticmethod(lagint) + _der = staticmethod(lagder) + _fit = staticmethod(lagfit) + _line = staticmethod(lagline) + _roots = staticmethod(lagroots) + _fromroots = staticmethod(lagfromroots) + + # Virtual properties + domain = np.array(lagdomain) + window = np.array(lagdomain) + basis_name = 'L' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.pyi new file mode 100644 index 0000000000000000000000000000000000000000..ee81157957482006cde90445fa73cf4223723d5f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/laguerre.pyi @@ -0,0 +1,100 @@ +from typing import Final, Literal as L + +import numpy as np + +from ._polybase import ABCPolyBase +from ._polytypes import ( + _Array1, + _Array2, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncGauss, + _FuncInteg, + _FuncLine, + _FuncPoly2Ortho, + _FuncPow, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVal2D, + _FuncVal3D, + _FuncValFromRoots, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncWeight, +) +from .polyutils import trimcoef as lagtrim + +__all__ = [ + "lagzero", + "lagone", + "lagx", + "lagdomain", + "lagline", + "lagadd", + "lagsub", + "lagmulx", + "lagmul", + "lagdiv", + "lagpow", + "lagval", + "lagder", + "lagint", + "lag2poly", + "poly2lag", + "lagfromroots", + "lagvander", + "lagfit", + "lagtrim", + "lagroots", + "Laguerre", + "lagval2d", + "lagval3d", + "laggrid2d", + "laggrid3d", + "lagvander2d", + "lagvander3d", + "lagcompanion", + "laggauss", + "lagweight", +] + +poly2lag: _FuncPoly2Ortho[L["poly2lag"]] +lag2poly: _FuncUnOp[L["lag2poly"]] + +lagdomain: Final[_Array2[np.float64]] +lagzero: Final[_Array1[np.int_]] +lagone: Final[_Array1[np.int_]] +lagx: Final[_Array2[np.int_]] + +lagline: _FuncLine[L["lagline"]] +lagfromroots: _FuncFromRoots[L["lagfromroots"]] +lagadd: _FuncBinOp[L["lagadd"]] +lagsub: _FuncBinOp[L["lagsub"]] +lagmulx: _FuncUnOp[L["lagmulx"]] +lagmul: _FuncBinOp[L["lagmul"]] +lagdiv: _FuncBinOp[L["lagdiv"]] +lagpow: _FuncPow[L["lagpow"]] +lagder: _FuncDer[L["lagder"]] +lagint: _FuncInteg[L["lagint"]] +lagval: _FuncVal[L["lagval"]] +lagval2d: _FuncVal2D[L["lagval2d"]] +lagval3d: _FuncVal3D[L["lagval3d"]] +lagvalfromroots: _FuncValFromRoots[L["lagvalfromroots"]] +laggrid2d: _FuncVal2D[L["laggrid2d"]] +laggrid3d: _FuncVal3D[L["laggrid3d"]] +lagvander: _FuncVander[L["lagvander"]] +lagvander2d: _FuncVander2D[L["lagvander2d"]] +lagvander3d: _FuncVander3D[L["lagvander3d"]] +lagfit: _FuncFit[L["lagfit"]] +lagcompanion: _FuncCompanion[L["lagcompanion"]] +lagroots: _FuncRoots[L["lagroots"]] +laggauss: _FuncGauss[L["laggauss"]] +lagweight: _FuncWeight[L["lagweight"]] + + +class Laguerre(ABCPolyBase[L["L"]]): ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.py b/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.py new file mode 100644 index 0000000000000000000000000000000000000000..c2cd3fbfe76021c908b0e5a004f68617c1da6d7f --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.py @@ -0,0 +1,1605 @@ +""" +================================================== +Legendre Series (:mod:`numpy.polynomial.legendre`) +================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Legendre series, including a `Legendre` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Legendre + +Constants +--------- + +.. autosummary:: + :toctree: generated/ + + legdomain + legzero + legone + legx + +Arithmetic +---------- + +.. autosummary:: + :toctree: generated/ + + legadd + legsub + legmulx + legmul + legdiv + legpow + legval + legval2d + legval3d + leggrid2d + leggrid3d + +Calculus +-------- + +.. autosummary:: + :toctree: generated/ + + legder + legint + +Misc Functions +-------------- + +.. autosummary:: + :toctree: generated/ + + legfromroots + legroots + legvander + legvander2d + legvander3d + leggauss + legweight + legcompanion + legfit + legtrim + legline + leg2poly + poly2leg + +See also +-------- +numpy.polynomial + +""" +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'legzero', 'legone', 'legx', 'legdomain', 'legline', 'legadd', + 'legsub', 'legmulx', 'legmul', 'legdiv', 'legpow', 'legval', 'legder', + 'legint', 'leg2poly', 'poly2leg', 'legfromroots', 'legvander', + 'legfit', 'legtrim', 'legroots', 'Legendre', 'legval2d', 'legval3d', + 'leggrid2d', 'leggrid3d', 'legvander2d', 'legvander3d', 'legcompanion', + 'leggauss', 'legweight'] + +legtrim = pu.trimcoef + + +def poly2leg(pol): + """ + Convert a polynomial to a Legendre series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Legendre series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Legendre + series. + + See Also + -------- + leg2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> import numpy as np + >>> from numpy import polynomial as P + >>> p = P.Polynomial(np.arange(4)) + >>> p + Polynomial([0., 1., 2., 3.], domain=[-1., 1.], window=[-1., 1.], ... + >>> c = P.Legendre(P.legendre.poly2leg(p.coef)) + >>> c + Legendre([ 1. , 3.25, 1. , 0.75], domain=[-1, 1], window=[-1, 1]) # may vary + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = legadd(legmulx(res), pol[i]) + return res + + +def leg2poly(c): + """ + Convert a Legendre series to a polynomial. + + Convert an array representing the coefficients of a Legendre series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Legendre series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2leg + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> c = P.Legendre(range(4)) + >>> c + Legendre([0., 1., 2., 3.], domain=[-1., 1.], window=[-1., 1.], symbol='x') + >>> p = c.convert(kind=P.Polynomial) + >>> p + Polynomial([-1. , -3.5, 3. , 7.5], domain=[-1., 1.], window=[-1., ... + >>> P.legendre.leg2poly(range(4)) + array([-1. , -3.5, 3. , 7.5]) + + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n < 3: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], (c1*(i - 1))/i) + c1 = polyadd(tmp, (polymulx(c1)*(2*i - 1))/i) + return polyadd(c0, polymulx(c1)) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Legendre +legdomain = np.array([-1., 1.]) + +# Legendre coefficients representing zero. +legzero = np.array([0]) + +# Legendre coefficients representing one. +legone = np.array([1]) + +# Legendre coefficients representing the identity x. +legx = np.array([0, 1]) + + +def legline(off, scl): + """ + Legendre series whose graph is a straight line. + + + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Legendre series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> import numpy.polynomial.legendre as L + >>> L.legline(3,2) + array([3, 2]) + >>> L.legval(-3, L.legline(3,2)) # should be -3 + -3.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def legfromroots(roots): + """ + Generate a Legendre series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Legendre form, where the :math:`r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Legendre form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> import numpy.polynomial.legendre as L + >>> L.legfromroots((-1,0,1)) # x^3 - x relative to the standard basis + array([ 0. , -0.4, 0. , 0.4]) + >>> j = complex(0,1) + >>> L.legfromroots((-j,j)) # x^2 + 1 relative to the standard basis + array([ 1.33333333+0.j, 0.00000000+0.j, 0.66666667+0.j]) # may vary + + """ + return pu._fromroots(legline, legmul, roots) + + +def legadd(c1, c2): + """ + Add one Legendre series to another. + + Returns the sum of two Legendre series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Legendre series of their sum. + + See Also + -------- + legsub, legmulx, legmul, legdiv, legpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Legendre series + is a Legendre series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legadd(c1,c2) + array([4., 4., 4.]) + + """ + return pu._add(c1, c2) + + +def legsub(c1, c2): + """ + Subtract one Legendre series from another. + + Returns the difference of two Legendre series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Legendre series coefficients representing their difference. + + See Also + -------- + legadd, legmulx, legmul, legdiv, legpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Legendre + series is a Legendre series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legsub(c1,c2) + array([-2., 0., 2.]) + >>> L.legsub(c2,c1) # -C.legsub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def legmulx(c): + """Multiply a Legendre series by x. + + Multiply the Legendre series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + legadd, legsub, legmul, legdiv, legpow + + Notes + ----- + The multiplication uses the recursion relationship for Legendre + polynomials in the form + + .. math:: + + xP_i(x) = ((i + 1)*P_{i + 1}(x) + i*P_{i - 1}(x))/(2i + 1) + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> L.legmulx([1,2,3]) + array([ 0.66666667, 2.2, 1.33333333, 1.8]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + for i in range(1, len(c)): + j = i + 1 + k = i - 1 + s = i + j + prd[j] = (c[i]*j)/s + prd[k] += (c[i]*i)/s + return prd + + +def legmul(c1, c2): + """ + Multiply one Legendre series by another. + + Returns the product of two Legendre series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Legendre series coefficients representing their product. + + See Also + -------- + legadd, legsub, legmulx, legdiv, legpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Legendre polynomial basis set. Thus, to express + the product as a Legendre series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2) + >>> L.legmul(c1,c2) # multiplication requires "reprojection" + array([ 4.33333333, 10.4 , 11.66666667, 3.6 ]) # may vary + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = legsub(c[-i]*xs, (c1*(nd - 1))/nd) + c1 = legadd(tmp, (legmulx(c1)*(2*nd - 1))/nd) + return legadd(c0, legmulx(c1)) + + +def legdiv(c1, c2): + """ + Divide one Legendre series by another. + + Returns the quotient-with-remainder of two Legendre series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + quo, rem : ndarrays + Of Legendre series coefficients representing the quotient and + remainder. + + See Also + -------- + legadd, legsub, legmulx, legmul, legpow + + Notes + ----- + In general, the (polynomial) division of one Legendre series by another + results in quotient and remainder terms that are not in the Legendre + polynomial basis set. Thus, to express these results as a Legendre + series, it is necessary to "reproject" the results onto the Legendre + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legdiv(c1,c2) # quotient "intuitive," remainder not + (array([3.]), array([-8., -4.])) + >>> c2 = (0,1,2,3) + >>> L.legdiv(c2,c1) # neither "intuitive" + (array([-0.07407407, 1.66666667]), array([-1.03703704, -2.51851852])) # may vary + + """ + return pu._div(legmul, c1, c2) + + +def legpow(c, pow, maxpower=16): + """Raise a Legendre series to a power. + + Returns the Legendre series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Legendre series of power. + + See Also + -------- + legadd, legsub, legmulx, legmul, legdiv + + """ + return pu._pow(legmul, c, pow, maxpower) + + +def legder(c, m=1, scl=1, axis=0): + """ + Differentiate a Legendre series. + + Returns the Legendre series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*L_0 + 2*L_1 + 3*L_2`` + while [[1,2],[1,2]] represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + + 2*L_0(x)*L_1(y) + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Legendre series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Legendre series of the derivative. + + See Also + -------- + legint + + Notes + ----- + In general, the result of differentiating a Legendre series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c = (1,2,3,4) + >>> L.legder(c) + array([ 6., 9., 20.]) + >>> L.legder(c, 3) + array([60.]) + >>> L.legder(c, scl=-1) + array([ -6., -9., -20.]) + >>> L.legder(c, 2,-1) + array([ 9., 60.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 2, -1): + der[j - 1] = (2*j - 1)*c[j] + c[j - 2] += c[j] + if n > 1: + der[1] = 3*c[2] + der[0] = c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def legint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Legendre series. + + Returns the Legendre series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``L_0 + 2*L_1 + 3*L_2`` while [[1,2],[1,2]] + represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + 2*L_0(x)*L_1(y) + + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Legendre series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + Legendre series coefficient array of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + legder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c = (1,2,3) + >>> L.legint(c) + array([ 0.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, 3) + array([ 1.66666667e-02, -1.78571429e-02, 4.76190476e-02, # may vary + -1.73472348e-18, 1.90476190e-02, 9.52380952e-03]) + >>> L.legint(c, k=3) + array([ 3.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, lbnd=-2) + array([ 7.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, scl=2) + array([ 0.66666667, 0.8 , 1.33333333, 1.2 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + if n > 1: + tmp[2] = c[1]/3 + for j in range(2, n): + t = c[j]/(2*j + 1) + tmp[j + 1] = t + tmp[j - 1] -= t + tmp[0] += k[i] - legval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def legval(x, c, tensor=True): + """ + Evaluate a Legendre series at points x. + + If `c` is of length ``n + 1``, this function returns the value: + + .. math:: p(x) = c_0 * L_0(x) + c_1 * L_1(x) + ... + c_n * L_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + legval2d, leggrid2d, legval3d, leggrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + """ + c = np.array(c, ndmin=1, copy=None) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - (c1*(nd - 1))/nd + c1 = tmp + (c1*x*(2*nd - 1))/nd + return c0 + c1*x + + +def legval2d(x, y, c): + """ + Evaluate a 2-D Legendre series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * L_i(x) * L_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Legendre series at points formed + from pairs of corresponding values from `x` and `y`. + + See Also + -------- + legval, leggrid2d, legval3d, leggrid3d + """ + return pu._valnd(legval, c, x, y) + + +def leggrid2d(x, y, c): + """ + Evaluate a 2-D Legendre series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * L_i(a) * L_j(b) + + where the points ``(a, b)`` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + legval, legval2d, legval3d, leggrid3d + """ + return pu._gridnd(legval, c, x, y) + + +def legval3d(x, y, z, c): + """ + Evaluate a 3-D Legendre series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * L_i(x) * L_j(y) * L_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + ``(x, y, z)``, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + legval, legval2d, leggrid2d, leggrid3d + """ + return pu._valnd(legval, c, x, y, z) + + +def leggrid3d(x, y, z, c): + """ + Evaluate a 3-D Legendre series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * L_i(a) * L_j(b) * L_k(c) + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + legval, legval2d, leggrid2d, legval3d + """ + return pu._gridnd(legval, c, x, y, z) + + +def legvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = L_i(x) + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Legendre polynomial. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + array ``V = legvander(x, n)``, then ``np.dot(V, c)`` and + ``legval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Legendre series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Legendre polynomial. The dtype will be the same as + the converted `x`. + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + # Use forward recursion to generate the entries. This is not as accurate + # as reverse recursion in this application but it is more efficient. + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x*(2*i - 1) - v[i-2]*(i - 1))/i + return np.moveaxis(v, 0, -1) + + +def legvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = L_i(x) * L_j(y), + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the degrees of + the Legendre polynomials. + + If ``V = legvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``legval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Legendre + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + legvander, legvander3d, legval2d, legval3d + """ + return pu._vander_nd_flat((legvander, legvander), (x, y), deg) + + +def legvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = L_i(x)*L_j(y)*L_k(z), + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the degrees of the Legendre polynomials. + + If ``V = legvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``legval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Legendre + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + legvander, legvander3d, legval2d, legval3d + """ + return pu._vander_nd_flat((legvander, legvander, legvander), (x, y, z), deg) + + +def legfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Legendre series to data. + + Return the coefficients of a Legendre series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Legendre coefficients ordered from low to high. If `y` was + 2-D, the coefficients for the data in column k of `y` are in + column `k`. If `deg` is specified as a list, coefficients for + terms not included in the fit are set equal to zero in the + returned `coef`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + legval : Evaluates a Legendre series. + legvander : Vandermonde matrix of Legendre series. + legweight : Legendre weight function (= 1). + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Legendre series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where :math:`w_j` are the weights. This problem is solved by setting up + as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `~exceptions.RankWarning` will be issued. This means that + the coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Legendre series are usually better conditioned than fits + using power series, but much can depend on the distribution of the + sample points and the smoothness of the data. If the quality of the fit + is inadequate splines may be a good alternative. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + + """ + return pu._fit(legvander, x, y, deg, rcond, full, w) + + +def legcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an Legendre basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = 1./np.sqrt(2*np.arange(n) + 1) + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.arange(1, n)*scl[:n-1]*scl[1:n] + bot[...] = top + mat[:, -1] -= (c[:-1]/c[-1])*(scl/scl[-1])*(n/(2*n - 1)) + return mat + + +def legroots(c): + """ + Compute the roots of a Legendre series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * L_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such values. + Roots with multiplicity greater than 1 will also show larger errors as + the value of the series near such points is relatively insensitive to + errors in the roots. Isolated roots near the origin can be improved by + a few iterations of Newton's method. + + The Legendre series basis polynomials aren't powers of ``x`` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> import numpy.polynomial.legendre as leg + >>> leg.legroots((1, 2, 3, 4)) # 4L_3 + 3L_2 + 2L_1 + 1L_0, all real roots + array([-0.85099543, -0.11407192, 0.51506735]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = legcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def leggauss(deg): + """ + Gauss-Legendre quadrature. + + Computes the sample points and weights for Gauss-Legendre quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-1, 1]` with + the weight function :math:`f(x) = 1`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (L'_n(x_k) * L_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`L_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = legcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = legval(x, c) + df = legval(x, legder(c)) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = legval(x, c[1:]) + fm /= np.abs(fm).max() + df /= np.abs(df).max() + w = 1/(fm * df) + + # for Legendre we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= 2. / w.sum() + + return x, w + + +def legweight(x): + """ + Weight function of the Legendre polynomials. + + The weight function is :math:`1` and the interval of integration is + :math:`[-1, 1]`. The Legendre polynomials are orthogonal, but not + normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + """ + w = x*0.0 + 1.0 + return w + +# +# Legendre series class +# + +class Legendre(ABCPolyBase): + """A Legendre series class. + + The Legendre class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + Legendre coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*P_0(x) + 2*P_1(x) + 3*P_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(legadd) + _sub = staticmethod(legsub) + _mul = staticmethod(legmul) + _div = staticmethod(legdiv) + _pow = staticmethod(legpow) + _val = staticmethod(legval) + _int = staticmethod(legint) + _der = staticmethod(legder) + _fit = staticmethod(legfit) + _line = staticmethod(legline) + _roots = staticmethod(legroots) + _fromroots = staticmethod(legfromroots) + + # Virtual properties + domain = np.array(legdomain) + window = np.array(legdomain) + basis_name = 'P' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.pyi new file mode 100644 index 0000000000000000000000000000000000000000..d81f3e6f54a4f72fd2cbc341f0efaa973aa3195a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/legendre.pyi @@ -0,0 +1,99 @@ +from typing import Final, Literal as L + +import numpy as np + +from ._polybase import ABCPolyBase +from ._polytypes import ( + _Array1, + _Array2, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncGauss, + _FuncInteg, + _FuncLine, + _FuncPoly2Ortho, + _FuncPow, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVal2D, + _FuncVal3D, + _FuncValFromRoots, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncWeight, +) +from .polyutils import trimcoef as legtrim + +__all__ = [ + "legzero", + "legone", + "legx", + "legdomain", + "legline", + "legadd", + "legsub", + "legmulx", + "legmul", + "legdiv", + "legpow", + "legval", + "legder", + "legint", + "leg2poly", + "poly2leg", + "legfromroots", + "legvander", + "legfit", + "legtrim", + "legroots", + "Legendre", + "legval2d", + "legval3d", + "leggrid2d", + "leggrid3d", + "legvander2d", + "legvander3d", + "legcompanion", + "leggauss", + "legweight", +] + +poly2leg: _FuncPoly2Ortho[L["poly2leg"]] +leg2poly: _FuncUnOp[L["leg2poly"]] + +legdomain: Final[_Array2[np.float64]] +legzero: Final[_Array1[np.int_]] +legone: Final[_Array1[np.int_]] +legx: Final[_Array2[np.int_]] + +legline: _FuncLine[L["legline"]] +legfromroots: _FuncFromRoots[L["legfromroots"]] +legadd: _FuncBinOp[L["legadd"]] +legsub: _FuncBinOp[L["legsub"]] +legmulx: _FuncUnOp[L["legmulx"]] +legmul: _FuncBinOp[L["legmul"]] +legdiv: _FuncBinOp[L["legdiv"]] +legpow: _FuncPow[L["legpow"]] +legder: _FuncDer[L["legder"]] +legint: _FuncInteg[L["legint"]] +legval: _FuncVal[L["legval"]] +legval2d: _FuncVal2D[L["legval2d"]] +legval3d: _FuncVal3D[L["legval3d"]] +legvalfromroots: _FuncValFromRoots[L["legvalfromroots"]] +leggrid2d: _FuncVal2D[L["leggrid2d"]] +leggrid3d: _FuncVal3D[L["leggrid3d"]] +legvander: _FuncVander[L["legvander"]] +legvander2d: _FuncVander2D[L["legvander2d"]] +legvander3d: _FuncVander3D[L["legvander3d"]] +legfit: _FuncFit[L["legfit"]] +legcompanion: _FuncCompanion[L["legcompanion"]] +legroots: _FuncRoots[L["legroots"]] +leggauss: _FuncGauss[L["leggauss"]] +legweight: _FuncWeight[L["legweight"]] + +class Legendre(ABCPolyBase[L["P"]]): ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.py b/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.py new file mode 100644 index 0000000000000000000000000000000000000000..86ea3a5d1d6e030929bc9de2f4744983a2a0417e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.py @@ -0,0 +1,1617 @@ +""" +================================================= +Power Series (:mod:`numpy.polynomial.polynomial`) +================================================= + +This module provides a number of objects (mostly functions) useful for +dealing with polynomials, including a `Polynomial` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with polynomial objects is in +the docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Polynomial + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + polydomain + polyzero + polyone + polyx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + polyadd + polysub + polymulx + polymul + polydiv + polypow + polyval + polyval2d + polyval3d + polygrid2d + polygrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + polyder + polyint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + polyfromroots + polyroots + polyvalfromroots + polyvander + polyvander2d + polyvander3d + polycompanion + polyfit + polytrim + polyline + +See Also +-------- +`numpy.polynomial` + +""" +__all__ = [ + 'polyzero', 'polyone', 'polyx', 'polydomain', 'polyline', 'polyadd', + 'polysub', 'polymulx', 'polymul', 'polydiv', 'polypow', 'polyval', + 'polyvalfromroots', 'polyder', 'polyint', 'polyfromroots', 'polyvander', + 'polyfit', 'polytrim', 'polyroots', 'Polynomial', 'polyval2d', 'polyval3d', + 'polygrid2d', 'polygrid3d', 'polyvander2d', 'polyvander3d', + 'polycompanion'] + +import numpy as np +import numpy.linalg as la +from numpy.lib.array_utils import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +polytrim = pu.trimcoef + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Polynomial default domain. +polydomain = np.array([-1., 1.]) + +# Polynomial coefficients representing zero. +polyzero = np.array([0]) + +# Polynomial coefficients representing one. +polyone = np.array([1]) + +# Polynomial coefficients representing the identity x. +polyx = np.array([0, 1]) + +# +# Polynomial series functions +# + + +def polyline(off, scl): + """ + Returns an array representing a linear polynomial. + + Parameters + ---------- + off, scl : scalars + The "y-intercept" and "slope" of the line, respectively. + + Returns + ------- + y : ndarray + This module's representation of the linear polynomial ``off + + scl*x``. + + See Also + -------- + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polyline(1, -1) + array([ 1, -1]) + >>> P.polyval(1, P.polyline(1, -1)) # should be 0 + 0.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def polyfromroots(roots): + """ + Generate a monic polynomial with given roots. + + Return the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + where the :math:`r_n` are the roots specified in `roots`. If a zero has + multiplicity n, then it must appear in `roots` n times. For instance, + if 2 is a root of multiplicity three and 3 is a root of multiplicity 2, + then `roots` looks something like [2, 2, 2, 3, 3]. The roots can appear + in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * x + ... + x^n + + The coefficient of the last term is 1 for monic polynomials in this + form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of the polynomial's coefficients If all the roots are + real, then `out` is also real, otherwise it is complex. (see + Examples below). + + See Also + -------- + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Notes + ----- + The coefficients are determined by multiplying together linear factors + of the form ``(x - r_i)``, i.e. + + .. math:: p(x) = (x - r_0) (x - r_1) ... (x - r_n) + + where ``n == len(roots) - 1``; note that this implies that ``1`` is always + returned for :math:`a_n`. + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polyfromroots((-1,0,1)) # x(x - 1)(x + 1) = x^3 - x + array([ 0., -1., 0., 1.]) + >>> j = complex(0,1) + >>> P.polyfromroots((-j,j)) # complex returned, though values are real + array([1.+0.j, 0.+0.j, 1.+0.j]) + + """ + return pu._fromroots(polyline, polymul, roots) + + +def polyadd(c1, c2): + """ + Add one polynomial to another. + + Returns the sum of two polynomials `c1` + `c2`. The arguments are + sequences of coefficients from lowest order term to highest, i.e., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to high. + + Returns + ------- + out : ndarray + The coefficient array representing their sum. + + See Also + -------- + polysub, polymulx, polymul, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1, 2, 3) + >>> c2 = (3, 2, 1) + >>> sum = P.polyadd(c1,c2); sum + array([4., 4., 4.]) + >>> P.polyval(2, sum) # 4 + 4(2) + 4(2**2) + 28.0 + + """ + return pu._add(c1, c2) + + +def polysub(c1, c2): + """ + Subtract one polynomial from another. + + Returns the difference of two polynomials `c1` - `c2`. The arguments + are sequences of coefficients from lowest order term to highest, i.e., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of coefficients representing their difference. + + See Also + -------- + polyadd, polymulx, polymul, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1, 2, 3) + >>> c2 = (3, 2, 1) + >>> P.polysub(c1,c2) + array([-2., 0., 2.]) + >>> P.polysub(c2, c1) # -P.polysub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def polymulx(c): + """Multiply a polynomial by x. + + Multiply the polynomial `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + polyadd, polysub, polymul, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1, 2, 3) + >>> P.polymulx(c) + array([0., 1., 2., 3.]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1:] = c + return prd + + +def polymul(c1, c2): + """ + Multiply one polynomial by another. + + Returns the product of two polynomials `c1` * `c2`. The arguments are + sequences of coefficients, from lowest order term to highest, e.g., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2.`` + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of coefficients representing a polynomial, relative to the + "standard" basis, and ordered from lowest order term to highest. + + Returns + ------- + out : ndarray + Of the coefficients of their product. + + See Also + -------- + polyadd, polysub, polymulx, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1, 2, 3) + >>> c2 = (3, 2, 1) + >>> P.polymul(c1, c2) + array([ 3., 8., 14., 8., 3.]) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + ret = np.convolve(c1, c2) + return pu.trimseq(ret) + + +def polydiv(c1, c2): + """ + Divide one polynomial by another. + + Returns the quotient-with-remainder of two polynomials `c1` / `c2`. + The arguments are sequences of coefficients, from lowest order term + to highest, e.g., [1,2,3] represents ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to high. + + Returns + ------- + [quo, rem] : ndarrays + Of coefficient series representing the quotient and remainder. + + See Also + -------- + polyadd, polysub, polymulx, polymul, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1, 2, 3) + >>> c2 = (3, 2, 1) + >>> P.polydiv(c1, c2) + (array([3.]), array([-8., -4.])) + >>> P.polydiv(c2, c1) + (array([ 0.33333333]), array([ 2.66666667, 1.33333333])) # may vary + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError # FIXME: add message with details to exception + + # note: this is more efficient than `pu._div(polymul, c1, c2)` + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + dlen = lc1 - lc2 + scl = c2[-1] + c2 = c2[:-1]/scl + i = dlen + j = lc1 - 1 + while i >= 0: + c1[i:j] -= c2*c1[j] + i -= 1 + j -= 1 + return c1[j+1:]/scl, pu.trimseq(c1[:j+1]) + + +def polypow(c, pow, maxpower=None): + """Raise a polynomial to a power. + + Returns the polynomial `c` raised to the power `pow`. The argument + `c` is a sequence of coefficients ordered from low to high. i.e., + [1,2,3] is the series ``1 + 2*x + 3*x**2.`` + + Parameters + ---------- + c : array_like + 1-D array of array of series coefficients ordered from low to + high degree. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Power series of power. + + See Also + -------- + polyadd, polysub, polymulx, polymul, polydiv + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polypow([1, 2, 3], 2) + array([ 1., 4., 10., 12., 9.]) + + """ + # note: this is more efficient than `pu._pow(polymul, c1, c2)`, as it + # avoids calling `as_series` repeatedly + return pu._pow(np.convolve, c, pow, maxpower) + + +def polyder(c, m=1, scl=1, axis=0): + """ + Differentiate a polynomial. + + Returns the polynomial coefficients `c` differentiated `m` times along + `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The + argument `c` is an array of coefficients from low to high degree along + each axis, e.g., [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2`` + while [[1,2],[1,2]] represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is + ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of polynomial coefficients. If c is multidimensional the + different axis correspond to different variables with the degree + in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change + of variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + Returns + ------- + der : ndarray + Polynomial coefficients of the derivative. + + See Also + -------- + polyint + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1, 2, 3, 4) + >>> P.polyder(c) # (d/dx)(c) + array([ 2., 6., 12.]) + >>> P.polyder(c, 3) # (d**3/dx**3)(c) + array([24.]) + >>> P.polyder(c, scl=-1) # (d/d(-x))(c) + array([ -2., -6., -12.]) + >>> P.polyder(c, 2, -1) # (d**2/d(-x)**2)(c) + array([ 6., 24.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype fails with NA + c = c + 0.0 + cdt = c.dtype + cnt = pu._as_int(m, "the order of derivation") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=cdt) + for j in range(n, 0, -1): + der[j - 1] = j*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def polyint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a polynomial. + + Returns the polynomial coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients, from low to high degree along each axis, e.g., [1,2,3] + represents the polynomial ``1 + 2*x + 3*x**2`` while [[1,2],[1,2]] + represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients, ordered from low to high. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at zero + is the first value in the list, the value of the second integral + at zero is the second value, etc. If ``k == []`` (the default), + all constants are set to zero. If ``m == 1``, a single scalar can + be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + Returns + ------- + S : ndarray + Coefficient array of the integral. + + Raises + ------ + ValueError + If ``m < 1``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + polyder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. Why + is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1, 2, 3) + >>> P.polyint(c) # should return array([0, 1, 1, 1]) + array([0., 1., 1., 1.]) + >>> P.polyint(c, 3) # should return array([0, 0, 0, 1/6, 1/12, 1/20]) + array([ 0. , 0. , 0. , 0.16666667, 0.08333333, # may vary + 0.05 ]) + >>> P.polyint(c, k=3) # should return array([3, 1, 1, 1]) + array([3., 1., 1., 1.]) + >>> P.polyint(c,lbnd=-2) # should return array([6, 1, 1, 1]) + array([6., 1., 1., 1.]) + >>> P.polyint(c,scl=-2) # should return array([0, -2, -2, -2]) + array([ 0., -2., -2., -2.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype doesn't preserve mask attribute. + c = c + 0.0 + cdt = c.dtype + if not np.iterable(k): + k = [k] + cnt = pu._as_int(m, "the order of integration") + iaxis = pu._as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + k = list(k) + [0]*(cnt - len(k)) + c = np.moveaxis(c, iaxis, 0) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=cdt) + tmp[0] = c[0]*0 + tmp[1] = c[0] + for j in range(1, n): + tmp[j + 1] = c[j]/(j + 1) + tmp[0] += k[i] - polyval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def polyval(x, c, tensor=True): + """ + Evaluate a polynomial at points x. + + If `c` is of length ``n + 1``, this function returns the value + + .. math:: p(x) = c_0 + c_1 * x + ... + c_n * x^n + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + Returns + ------- + values : ndarray, compatible object + The shape of the returned array is described above. + + See Also + -------- + polyval2d, polygrid2d, polyval3d, polygrid3d + + Notes + ----- + The evaluation uses Horner's method. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial.polynomial import polyval + >>> polyval(1, [1,2,3]) + 6.0 + >>> a = np.arange(4).reshape(2,2) + >>> a + array([[0, 1], + [2, 3]]) + >>> polyval(a, [1, 2, 3]) + array([[ 1., 6.], + [17., 34.]]) + >>> coef = np.arange(4).reshape(2, 2) # multidimensional coefficients + >>> coef + array([[0, 1], + [2, 3]]) + >>> polyval([1, 2], coef, tensor=True) + array([[2., 4.], + [4., 7.]]) + >>> polyval([1, 2], coef, tensor=False) + array([2., 7.]) + + """ + c = np.array(c, ndmin=1, copy=None) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype fails with NA + c = c + 0.0 + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + c0 = c[-1] + x*0 + for i in range(2, len(c) + 1): + c0 = c[-i] + c0*x + return c0 + + +def polyvalfromroots(x, r, tensor=True): + """ + Evaluate a polynomial specified by its roots at points x. + + If `r` is of length ``N``, this function returns the value + + .. math:: p(x) = \\prod_{n=1}^{N} (x - r_n) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `r`. + + If `r` is a 1-D array, then ``p(x)`` will have the same shape as `x`. If `r` + is multidimensional, then the shape of the result depends on the value of + `tensor`. If `tensor` is ``True`` the shape will be r.shape[1:] + x.shape; + that is, each polynomial is evaluated at every value of `x`. If `tensor` is + ``False``, the shape will be r.shape[1:]; that is, each polynomial is + evaluated only for the corresponding broadcast value of `x`. Note that + scalars have shape (,). + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `r`. + r : array_like + Array of roots. If `r` is multidimensional the first index is the + root index, while the remaining indices enumerate multiple + polynomials. For instance, in the two dimensional case the roots + of each polynomial may be thought of as stored in the columns of `r`. + tensor : boolean, optional + If True, the shape of the roots array is extended with ones on the + right, one for each dimension of `x`. Scalars have dimension 0 for this + action. The result is that every column of coefficients in `r` is + evaluated for every element of `x`. If False, `x` is broadcast over the + columns of `r` for the evaluation. This keyword is useful when `r` is + multidimensional. The default value is True. + + Returns + ------- + values : ndarray, compatible object + The shape of the returned array is described above. + + See Also + -------- + polyroots, polyfromroots, polyval + + Examples + -------- + >>> from numpy.polynomial.polynomial import polyvalfromroots + >>> polyvalfromroots(1, [1, 2, 3]) + 0.0 + >>> a = np.arange(4).reshape(2, 2) + >>> a + array([[0, 1], + [2, 3]]) + >>> polyvalfromroots(a, [-1, 0, 1]) + array([[-0., 0.], + [ 6., 24.]]) + >>> r = np.arange(-2, 2).reshape(2,2) # multidimensional coefficients + >>> r # each column of r defines one polynomial + array([[-2, -1], + [ 0, 1]]) + >>> b = [-2, 1] + >>> polyvalfromroots(b, r, tensor=True) + array([[-0., 3.], + [ 3., 0.]]) + >>> polyvalfromroots(b, r, tensor=False) + array([-0., 0.]) + + """ + r = np.array(r, ndmin=1, copy=None) + if r.dtype.char in '?bBhHiIlLqQpP': + r = r.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray): + if tensor: + r = r.reshape(r.shape + (1,)*x.ndim) + elif x.ndim >= r.ndim: + raise ValueError("x.ndim must be < r.ndim when tensor == False") + return np.prod(x - r, axis=0) + + +def polyval2d(x, y, c): + """ + Evaluate a 2-D polynomial at points (x, y). + + This function returns the value + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * x^i * y^j + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points ``(x, y)``, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + polyval, polygrid2d, polyval3d, polygrid3d + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = ((1, 2, 3), (4, 5, 6)) + >>> P.polyval2d(1, 1, c) + 21.0 + + """ + return pu._valnd(polyval, c, x, y) + + +def polygrid2d(x, y, c): + """ + Evaluate a 2-D polynomial on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * a^i * b^j + + where the points ``(a, b)`` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + polyval, polyval2d, polyval3d, polygrid3d + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = ((1, 2, 3), (4, 5, 6)) + >>> P.polygrid2d([0, 1], [0, 1], c) + array([[ 1., 6.], + [ 5., 21.]]) + + """ + return pu._gridnd(polyval, c, x, y) + + +def polyval3d(x, y, z, c): + """ + Evaluate a 3-D polynomial at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * x^i * y^j * z^k + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + ``(x, y, z)``, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + polyval, polyval2d, polygrid2d, polygrid3d + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = ((1, 2, 3), (4, 5, 6), (7, 8, 9)) + >>> P.polyval3d(1, 1, 1, c) + 45.0 + + """ + return pu._valnd(polyval, c, x, y, z) + + +def polygrid3d(x, y, z, c): + """ + Evaluate a 3-D polynomial on the Cartesian product of x, y and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * a^i * b^j * c^k + + where the points ``(a, b, c)`` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`, `y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + polyval, polyval2d, polygrid2d, polyval3d + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = ((1, 2, 3), (4, 5, 6), (7, 8, 9)) + >>> P.polygrid3d([0, 1], [0, 1], [0, 1], c) + array([[ 1., 13.], + [ 6., 51.]]) + + """ + return pu._gridnd(polyval, c, x, y, z) + + +def polyvander(x, deg): + """Vandermonde matrix of given degree. + + Returns the Vandermonde matrix of degree `deg` and sample points + `x`. The Vandermonde matrix is defined by + + .. math:: V[..., i] = x^i, + + where ``0 <= i <= deg``. The leading indices of `V` index the elements of + `x` and the last index is the power of `x`. + + If `c` is a 1-D array of coefficients of length ``n + 1`` and `V` is the + matrix ``V = polyvander(x, n)``, then ``np.dot(V, c)`` and + ``polyval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of polynomials of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray. + The Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where the last index is the power of `x`. + The dtype will be the same as the converted `x`. + + See Also + -------- + polyvander2d, polyvander3d + + Examples + -------- + The Vandermonde matrix of degree ``deg = 5`` and sample points + ``x = [-1, 2, 3]`` contains the element-wise powers of `x` + from 0 to 5 as its columns. + + >>> from numpy.polynomial import polynomial as P + >>> x, deg = [-1, 2, 3], 5 + >>> P.polyvander(x=x, deg=deg) + array([[ 1., -1., 1., -1., 1., -1.], + [ 1., 2., 4., 8., 16., 32.], + [ 1., 3., 9., 27., 81., 243.]]) + + """ + ideg = pu._as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=None, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = v[i-1]*x + return np.moveaxis(v, 0, -1) + + +def polyvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y)``. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = x^i * y^j, + + where ``0 <= i <= deg[0]`` and ``0 <= j <= deg[1]``. The leading indices of + `V` index the points ``(x, y)`` and the last index encodes the powers of + `x` and `y`. + + If ``V = polyvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``polyval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D polynomials + of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg([1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + polyvander, polyvander3d, polyval2d, polyval3d + + Examples + -------- + >>> import numpy as np + + The 2-D pseudo-Vandermonde matrix of degree ``[1, 2]`` and sample + points ``x = [-1, 2]`` and ``y = [1, 3]`` is as follows: + + >>> from numpy.polynomial import polynomial as P + >>> x = np.array([-1, 2]) + >>> y = np.array([1, 3]) + >>> m, n = 1, 2 + >>> deg = np.array([m, n]) + >>> V = P.polyvander2d(x=x, y=y, deg=deg) + >>> V + array([[ 1., 1., 1., -1., -1., -1.], + [ 1., 3., 9., 2., 6., 18.]]) + + We can verify the columns for any ``0 <= i <= m`` and ``0 <= j <= n``: + + >>> i, j = 0, 1 + >>> V[:, (deg[1]+1)*i + j] == x**i * y**j + array([ True, True]) + + The (1D) Vandermonde matrix of sample points ``x`` and degree ``m`` is a + special case of the (2D) pseudo-Vandermonde matrix with ``y`` points all + zero and degree ``[m, 0]``. + + >>> P.polyvander2d(x=x, y=0*x, deg=(m, 0)) == P.polyvander(x=x, deg=m) + array([[ True, True], + [ True, True]]) + + """ + return pu._vander_nd_flat((polyvander, polyvander), (x, y), deg) + + +def polyvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points ``(x, y, z)``. If `l`, `m`, `n` are the given degrees in `x`, `y`, `z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = x^i * y^j * z^k, + + where ``0 <= i <= l``, ``0 <= j <= m``, and ``0 <= j <= n``. The leading + indices of `V` index the points ``(x, y, z)`` and the last index encodes + the powers of `x`, `y`, and `z`. + + If ``V = polyvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``polyval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D polynomials + of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg([1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + polyvander, polyvander3d, polyval2d, polyval3d + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial import polynomial as P + >>> x = np.asarray([-1, 2, 1]) + >>> y = np.asarray([1, -2, -3]) + >>> z = np.asarray([2, 2, 5]) + >>> l, m, n = [2, 2, 1] + >>> deg = [l, m, n] + >>> V = P.polyvander3d(x=x, y=y, z=z, deg=deg) + >>> V + array([[ 1., 2., 1., 2., 1., 2., -1., -2., -1., + -2., -1., -2., 1., 2., 1., 2., 1., 2.], + [ 1., 2., -2., -4., 4., 8., 2., 4., -4., + -8., 8., 16., 4., 8., -8., -16., 16., 32.], + [ 1., 5., -3., -15., 9., 45., 1., 5., -3., + -15., 9., 45., 1., 5., -3., -15., 9., 45.]]) + + We can verify the columns for any ``0 <= i <= l``, ``0 <= j <= m``, + and ``0 <= k <= n`` + + >>> i, j, k = 2, 1, 0 + >>> V[:, (m+1)*(n+1)*i + (n+1)*j + k] == x**i * y**j * z**k + array([ True, True, True]) + + """ + return pu._vander_nd_flat((polyvander, polyvander, polyvander), (x, y, z), deg) + + +def polyfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least-squares fit of a polynomial to data. + + Return the coefficients of a polynomial of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * x + ... + c_n * x^n, + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (`M`,) + x-coordinates of the `M` sample (data) points ``(x[i], y[i])``. + y : array_like, shape (`M`,) or (`M`, `K`) + y-coordinates of the sample points. Several sets of sample points + sharing the same x-coordinates can be (independently) fit with one + call to `polyfit` by passing in for `y` a 2-D array that contains + one data set per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller + than `rcond`, relative to the largest singular value, will be + ignored. The default value is ``len(x)*eps``, where `eps` is the + relative precision of the platform's float type, about 2e-16 in + most cases. + full : bool, optional + Switch determining the nature of the return value. When ``False`` + (the default) just the coefficients are returned; when ``True``, + diagnostic information from the singular value decomposition (used + to solve the fit's matrix equation) is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (`deg` + 1,) or (`deg` + 1, `K`) + Polynomial coefficients ordered from low to high. If `y` was 2-D, + the coefficients in column `k` of `coef` represent the polynomial + fit to the data in `y`'s `k`-th column. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Raises + ------ + RankWarning + Raised if the matrix in the least-squares fit is rank deficient. + The warning is only raised if ``full == False``. The warnings can + be turned off by: + + >>> import warnings + >>> warnings.simplefilter('ignore', np.exceptions.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + polyval : Evaluates a polynomial. + polyvander : Vandermonde matrix for powers. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the polynomial `p` that minimizes + the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) over-determined matrix equation: + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected (and `full` == ``False``), a `~exceptions.RankWarning` will be + raised. This means that the coefficient values may be poorly determined. + Fitting to a lower order polynomial will usually get rid of the warning + (but may not be what you want, of course; if you have independent + reason(s) for choosing the degree which isn't working, you may have to: + a) reconsider those reasons, and/or b) reconsider the quality of your + data). The `rcond` parameter can also be set to a value smaller than + its default, but the resulting fit may be spurious and have large + contributions from roundoff error. + + Polynomial fits using double precision tend to "fail" at about + (polynomial) degree 20. Fits using Chebyshev or Legendre series are + generally better conditioned, but much can still depend on the + distribution of the sample points and the smoothness of the data. If + the quality of the fit is inadequate, splines may be a good + alternative. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial import polynomial as P + >>> x = np.linspace(-1,1,51) # x "data": [-1, -0.96, ..., 0.96, 1] + >>> rng = np.random.default_rng() + >>> err = rng.normal(size=len(x)) + >>> y = x**3 - x + err # x^3 - x + Gaussian noise + >>> c, stats = P.polyfit(x,y,3,full=True) + >>> c # c[0], c[1] approx. -1, c[2] should be approx. 0, c[3] approx. 1 + array([ 0.23111996, -1.02785049, -0.2241444 , 1.08405657]) # may vary + >>> stats # note the large SSR, explaining the rather poor results + [array([48.312088]), # may vary + 4, + array([1.38446749, 1.32119158, 0.50443316, 0.28853036]), + 1.1324274851176597e-14] + + Same thing without the added noise + + >>> y = x**3 - x + >>> c, stats = P.polyfit(x,y,3,full=True) + >>> c # c[0], c[1] ~= -1, c[2] should be "very close to 0", c[3] ~= 1 + array([-6.73496154e-17, -1.00000000e+00, 0.00000000e+00, 1.00000000e+00]) + >>> stats # note the minuscule SSR + [array([8.79579319e-31]), + np.int32(4), + array([1.38446749, 1.32119158, 0.50443316, 0.28853036]), + 1.1324274851176597e-14] + + """ + return pu._fit(polyvander, x, y, deg, rcond, full, w) + + +def polycompanion(c): + """ + Return the companion matrix of c. + + The companion matrix for power series cannot be made symmetric by + scaling the basis, so this function differs from those for the + orthogonal polynomials. + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Companion matrix of dimensions (deg, deg). + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1, 2, 3) + >>> P.polycompanion(c) + array([[ 0. , -0.33333333], + [ 1. , -0.66666667]]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + bot = mat.reshape(-1)[n::n+1] + bot[...] = 1 + mat[:, -1] -= c[:-1]/c[-1] + return mat + + +def polyroots(c): + """ + Compute the roots of a polynomial. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * x^i. + + Parameters + ---------- + c : 1-D array_like + 1-D array of polynomial coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the polynomial. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the power series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + Examples + -------- + >>> import numpy.polynomial.polynomial as poly + >>> poly.polyroots(poly.polyfromroots((-1,0,1))) + array([-1., 0., 1.]) + >>> poly.polyroots(poly.polyfromroots((-1,0,1))).dtype + dtype('float64') + >>> j = complex(0,1) + >>> poly.polyroots(poly.polyfromroots((-j,0,j))) + array([ 0.00000000e+00+0.j, 0.00000000e+00+1.j, 2.77555756e-17-1.j]) # may vary + + """ # noqa: E501 + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = polycompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +# +# polynomial class +# + +class Polynomial(ABCPolyBase): + """A power series class. + + The Polynomial class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed below. + + Parameters + ---------- + coef : array_like + Polynomial coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` give ``1 + 2*x + 3*x**2``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1., 1.]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1., 1.]. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(polyadd) + _sub = staticmethod(polysub) + _mul = staticmethod(polymul) + _div = staticmethod(polydiv) + _pow = staticmethod(polypow) + _val = staticmethod(polyval) + _int = staticmethod(polyint) + _der = staticmethod(polyder) + _fit = staticmethod(polyfit) + _line = staticmethod(polyline) + _roots = staticmethod(polyroots) + _fromroots = staticmethod(polyfromroots) + + # Virtual properties + domain = np.array(polydomain) + window = np.array(polydomain) + basis_name = None + + @classmethod + def _str_term_unicode(cls, i, arg_str): + if i == '1': + return f"·{arg_str}" + else: + return f"·{arg_str}{i.translate(cls._superscript_mapping)}" + + @staticmethod + def _str_term_ascii(i, arg_str): + if i == '1': + return f" {arg_str}" + else: + return f" {arg_str}**{i}" + + @staticmethod + def _repr_latex_term(i, arg_str, needs_parens): + if needs_parens: + arg_str = rf"\left({arg_str}\right)" + if i == 0: + return '1' + elif i == 1: + return arg_str + else: + return f"{arg_str}^{{{i}}}" diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.pyi new file mode 100644 index 0000000000000000000000000000000000000000..89a8b57185f3e326f8891e71ab2b47f48cd908e9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/polynomial.pyi @@ -0,0 +1,87 @@ +from typing import Final, Literal as L + +import numpy as np +from ._polybase import ABCPolyBase +from ._polytypes import ( + _Array1, + _Array2, + _FuncVal2D, + _FuncVal3D, + _FuncBinOp, + _FuncCompanion, + _FuncDer, + _FuncFit, + _FuncFromRoots, + _FuncInteg, + _FuncLine, + _FuncPow, + _FuncRoots, + _FuncUnOp, + _FuncVal, + _FuncVander, + _FuncVander2D, + _FuncVander3D, + _FuncValFromRoots, +) +from .polyutils import trimcoef as polytrim + +__all__ = [ + "polyzero", + "polyone", + "polyx", + "polydomain", + "polyline", + "polyadd", + "polysub", + "polymulx", + "polymul", + "polydiv", + "polypow", + "polyval", + "polyvalfromroots", + "polyder", + "polyint", + "polyfromroots", + "polyvander", + "polyfit", + "polytrim", + "polyroots", + "Polynomial", + "polyval2d", + "polyval3d", + "polygrid2d", + "polygrid3d", + "polyvander2d", + "polyvander3d", + "polycompanion", +] + +polydomain: Final[_Array2[np.float64]] +polyzero: Final[_Array1[np.int_]] +polyone: Final[_Array1[np.int_]] +polyx: Final[_Array2[np.int_]] + +polyline: _FuncLine[L["Polyline"]] +polyfromroots: _FuncFromRoots[L["polyfromroots"]] +polyadd: _FuncBinOp[L["polyadd"]] +polysub: _FuncBinOp[L["polysub"]] +polymulx: _FuncUnOp[L["polymulx"]] +polymul: _FuncBinOp[L["polymul"]] +polydiv: _FuncBinOp[L["polydiv"]] +polypow: _FuncPow[L["polypow"]] +polyder: _FuncDer[L["polyder"]] +polyint: _FuncInteg[L["polyint"]] +polyval: _FuncVal[L["polyval"]] +polyval2d: _FuncVal2D[L["polyval2d"]] +polyval3d: _FuncVal3D[L["polyval3d"]] +polyvalfromroots: _FuncValFromRoots[L["polyvalfromroots"]] +polygrid2d: _FuncVal2D[L["polygrid2d"]] +polygrid3d: _FuncVal3D[L["polygrid3d"]] +polyvander: _FuncVander[L["polyvander"]] +polyvander2d: _FuncVander2D[L["polyvander2d"]] +polyvander3d: _FuncVander3D[L["polyvander3d"]] +polyfit: _FuncFit[L["polyfit"]] +polycompanion: _FuncCompanion[L["polycompanion"]] +polyroots: _FuncRoots[L["polyroots"]] + +class Polynomial(ABCPolyBase[None]): ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.py b/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.py new file mode 100644 index 0000000000000000000000000000000000000000..1a6813b786c9bdd7eaa7961b5c50a5b187f7837a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.py @@ -0,0 +1,757 @@ +""" +Utility classes and functions for the polynomial modules. + +This module provides: error and warning objects; a polynomial base class; +and some routines used in both the `polynomial` and `chebyshev` modules. + +Functions +--------- + +.. autosummary:: + :toctree: generated/ + + as_series convert list of array_likes into 1-D arrays of common type. + trimseq remove trailing zeros. + trimcoef remove small trailing coefficients. + getdomain return the domain appropriate for a given set of abscissae. + mapdomain maps points between domains. + mapparms parameters of the linear map between domains. + +""" +import operator +import functools +import warnings + +import numpy as np + +from numpy._core.multiarray import dragon4_positional, dragon4_scientific +from numpy.exceptions import RankWarning + +__all__ = [ + 'as_series', 'trimseq', 'trimcoef', 'getdomain', 'mapdomain', 'mapparms', + 'format_float'] + +# +# Helper functions to convert inputs to 1-D arrays +# +def trimseq(seq): + """Remove small Poly series coefficients. + + Parameters + ---------- + seq : sequence + Sequence of Poly series coefficients. + + Returns + ------- + series : sequence + Subsequence with trailing zeros removed. If the resulting sequence + would be empty, return the first element. The returned sequence may + or may not be a view. + + Notes + ----- + Do not lose the type info if the sequence contains unknown objects. + + """ + if len(seq) == 0 or seq[-1] != 0: + return seq + else: + for i in range(len(seq) - 1, -1, -1): + if seq[i] != 0: + break + return seq[:i+1] + + +def as_series(alist, trim=True): + """ + Return argument as a list of 1-d arrays. + + The returned list contains array(s) of dtype double, complex double, or + object. A 1-d argument of shape ``(N,)`` is parsed into ``N`` arrays of + size one; a 2-d argument of shape ``(M,N)`` is parsed into ``M`` arrays + of size ``N`` (i.e., is "parsed by row"); and a higher dimensional array + raises a Value Error if it is not first reshaped into either a 1-d or 2-d + array. + + Parameters + ---------- + alist : array_like + A 1- or 2-d array_like + trim : boolean, optional + When True, trailing zeros are removed from the inputs. + When False, the inputs are passed through intact. + + Returns + ------- + [a1, a2,...] : list of 1-D arrays + A copy of the input data as a list of 1-d arrays. + + Raises + ------ + ValueError + Raised when `as_series` cannot convert its input to 1-d arrays, or at + least one of the resulting arrays is empty. + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial import polyutils as pu + >>> a = np.arange(4) + >>> pu.as_series(a) + [array([0.]), array([1.]), array([2.]), array([3.])] + >>> b = np.arange(6).reshape((2,3)) + >>> pu.as_series(b) + [array([0., 1., 2.]), array([3., 4., 5.])] + + >>> pu.as_series((1, np.arange(3), np.arange(2, dtype=np.float16))) + [array([1.]), array([0., 1., 2.]), array([0., 1.])] + + >>> pu.as_series([2, [1.1, 0.]]) + [array([2.]), array([1.1])] + + >>> pu.as_series([2, [1.1, 0.]], trim=False) + [array([2.]), array([1.1, 0. ])] + + """ + arrays = [np.array(a, ndmin=1, copy=None) for a in alist] + for a in arrays: + if a.size == 0: + raise ValueError("Coefficient array is empty") + if any(a.ndim != 1 for a in arrays): + raise ValueError("Coefficient array is not 1-d") + if trim: + arrays = [trimseq(a) for a in arrays] + + if any(a.dtype == np.dtype(object) for a in arrays): + ret = [] + for a in arrays: + if a.dtype != np.dtype(object): + tmp = np.empty(len(a), dtype=np.dtype(object)) + tmp[:] = a[:] + ret.append(tmp) + else: + ret.append(a.copy()) + else: + try: + dtype = np.common_type(*arrays) + except Exception as e: + raise ValueError("Coefficient arrays have no common type") from e + ret = [np.array(a, copy=True, dtype=dtype) for a in arrays] + return ret + + +def trimcoef(c, tol=0): + """ + Remove "small" "trailing" coefficients from a polynomial. + + "Small" means "small in absolute value" and is controlled by the + parameter `tol`; "trailing" means highest order coefficient(s), e.g., in + ``[0, 1, 1, 0, 0]`` (which represents ``0 + x + x**2 + 0*x**3 + 0*x**4``) + both the 3-rd and 4-th order coefficients would be "trimmed." + + Parameters + ---------- + c : array_like + 1-d array of coefficients, ordered from lowest order to highest. + tol : number, optional + Trailing (i.e., highest order) elements with absolute value less + than or equal to `tol` (default value is zero) are removed. + + Returns + ------- + trimmed : ndarray + 1-d array with trailing zeros removed. If the resulting series + would be empty, a series containing a single zero is returned. + + Raises + ------ + ValueError + If `tol` < 0 + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> pu.trimcoef((0,0,3,0,5,0,0)) + array([0., 0., 3., 0., 5.]) + >>> pu.trimcoef((0,0,1e-3,0,1e-5,0,0),1e-3) # item == tol is trimmed + array([0.]) + >>> i = complex(0,1) # works for complex + >>> pu.trimcoef((3e-4,1e-3*(1-i),5e-4,2e-5*(1+i)), 1e-3) + array([0.0003+0.j , 0.001 -0.001j]) + + """ + if tol < 0: + raise ValueError("tol must be non-negative") + + [c] = as_series([c]) + [ind] = np.nonzero(np.abs(c) > tol) + if len(ind) == 0: + return c[:1]*0 + else: + return c[:ind[-1] + 1].copy() + +def getdomain(x): + """ + Return a domain suitable for given abscissae. + + Find a domain suitable for a polynomial or Chebyshev series + defined at the values supplied. + + Parameters + ---------- + x : array_like + 1-d array of abscissae whose domain will be determined. + + Returns + ------- + domain : ndarray + 1-d array containing two values. If the inputs are complex, then + the two returned points are the lower left and upper right corners + of the smallest rectangle (aligned with the axes) in the complex + plane containing the points `x`. If the inputs are real, then the + two points are the ends of the smallest interval containing the + points `x`. + + See Also + -------- + mapparms, mapdomain + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial import polyutils as pu + >>> points = np.arange(4)**2 - 5; points + array([-5, -4, -1, 4]) + >>> pu.getdomain(points) + array([-5., 4.]) + >>> c = np.exp(complex(0,1)*np.pi*np.arange(12)/6) # unit circle + >>> pu.getdomain(c) + array([-1.-1.j, 1.+1.j]) + + """ + [x] = as_series([x], trim=False) + if x.dtype.char in np.typecodes['Complex']: + rmin, rmax = x.real.min(), x.real.max() + imin, imax = x.imag.min(), x.imag.max() + return np.array((complex(rmin, imin), complex(rmax, imax))) + else: + return np.array((x.min(), x.max())) + +def mapparms(old, new): + """ + Linear map parameters between domains. + + Return the parameters of the linear map ``offset + scale*x`` that maps + `old` to `new` such that ``old[i] -> new[i]``, ``i = 0, 1``. + + Parameters + ---------- + old, new : array_like + Domains. Each domain must (successfully) convert to a 1-d array + containing precisely two values. + + Returns + ------- + offset, scale : scalars + The map ``L(x) = offset + scale*x`` maps the first domain to the + second. + + See Also + -------- + getdomain, mapdomain + + Notes + ----- + Also works for complex numbers, and thus can be used to calculate the + parameters required to map any line in the complex plane to any other + line therein. + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> pu.mapparms((-1,1),(-1,1)) + (0.0, 1.0) + >>> pu.mapparms((1,-1),(-1,1)) + (-0.0, -1.0) + >>> i = complex(0,1) + >>> pu.mapparms((-i,-1),(1,i)) + ((1+1j), (1-0j)) + + """ + oldlen = old[1] - old[0] + newlen = new[1] - new[0] + off = (old[1]*new[0] - old[0]*new[1])/oldlen + scl = newlen/oldlen + return off, scl + +def mapdomain(x, old, new): + """ + Apply linear map to input points. + + The linear map ``offset + scale*x`` that maps the domain `old` to + the domain `new` is applied to the points `x`. + + Parameters + ---------- + x : array_like + Points to be mapped. If `x` is a subtype of ndarray the subtype + will be preserved. + old, new : array_like + The two domains that determine the map. Each must (successfully) + convert to 1-d arrays containing precisely two values. + + Returns + ------- + x_out : ndarray + Array of points of the same shape as `x`, after application of the + linear map between the two domains. + + See Also + -------- + getdomain, mapparms + + Notes + ----- + Effectively, this implements: + + .. math:: + x\\_out = new[0] + m(x - old[0]) + + where + + .. math:: + m = \\frac{new[1]-new[0]}{old[1]-old[0]} + + Examples + -------- + >>> import numpy as np + >>> from numpy.polynomial import polyutils as pu + >>> old_domain = (-1,1) + >>> new_domain = (0,2*np.pi) + >>> x = np.linspace(-1,1,6); x + array([-1. , -0.6, -0.2, 0.2, 0.6, 1. ]) + >>> x_out = pu.mapdomain(x, old_domain, new_domain); x_out + array([ 0. , 1.25663706, 2.51327412, 3.76991118, 5.02654825, # may vary + 6.28318531]) + >>> x - pu.mapdomain(x_out, new_domain, old_domain) + array([0., 0., 0., 0., 0., 0.]) + + Also works for complex numbers (and thus can be used to map any line in + the complex plane to any other line therein). + + >>> i = complex(0,1) + >>> old = (-1 - i, 1 + i) + >>> new = (-1 + i, 1 - i) + >>> z = np.linspace(old[0], old[1], 6); z + array([-1. -1.j , -0.6-0.6j, -0.2-0.2j, 0.2+0.2j, 0.6+0.6j, 1. +1.j ]) + >>> new_z = pu.mapdomain(z, old, new); new_z + array([-1.0+1.j , -0.6+0.6j, -0.2+0.2j, 0.2-0.2j, 0.6-0.6j, 1.0-1.j ]) # may vary + + """ + if type(x) not in (int, float, complex) and not isinstance(x, np.generic): + x = np.asanyarray(x) + off, scl = mapparms(old, new) + return off + scl*x + + +def _nth_slice(i, ndim): + sl = [np.newaxis] * ndim + sl[i] = slice(None) + return tuple(sl) + + +def _vander_nd(vander_fs, points, degrees): + r""" + A generalization of the Vandermonde matrix for N dimensions + + The result is built by combining the results of 1d Vandermonde matrices, + + .. math:: + W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{V_k(x_k)[i_0, \ldots, i_M, j_k]} + + where + + .. math:: + N &= \texttt{len(points)} = \texttt{len(degrees)} = \texttt{len(vander\_fs)} \\ + M &= \texttt{points[k].ndim} \\ + V_k &= \texttt{vander\_fs[k]} \\ + x_k &= \texttt{points[k]} \\ + 0 \le j_k &\le \texttt{degrees[k]} + + Expanding the one-dimensional :math:`V_k` functions gives: + + .. math:: + W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{B_{k, j_k}(x_k[i_0, \ldots, i_M])} + + where :math:`B_{k,m}` is the m'th basis of the polynomial construction used along + dimension :math:`k`. For a regular polynomial, :math:`B_{k, m}(x) = P_m(x) = x^m`. + + Parameters + ---------- + vander_fs : Sequence[function(array_like, int) -> ndarray] + The 1d vander function to use for each axis, such as ``polyvander`` + points : Sequence[array_like] + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + This must be the same length as `vander_fs`. + degrees : Sequence[int] + The maximum degree (inclusive) to use for each axis. + This must be the same length as `vander_fs`. + + Returns + ------- + vander_nd : ndarray + An array of shape ``points[0].shape + tuple(d + 1 for d in degrees)``. + """ + n_dims = len(vander_fs) + if n_dims != len(points): + raise ValueError( + f"Expected {n_dims} dimensions of sample points, got {len(points)}") + if n_dims != len(degrees): + raise ValueError( + f"Expected {n_dims} dimensions of degrees, got {len(degrees)}") + if n_dims == 0: + raise ValueError("Unable to guess a dtype or shape when no points are given") + + # convert to the same shape and type + points = tuple(np.asarray(tuple(points)) + 0.0) + + # produce the vandermonde matrix for each dimension, placing the last + # axis of each in an independent trailing axis of the output + vander_arrays = ( + vander_fs[i](points[i], degrees[i])[(...,) + _nth_slice(i, n_dims)] + for i in range(n_dims) + ) + + # we checked this wasn't empty already, so no `initial` needed + return functools.reduce(operator.mul, vander_arrays) + + +def _vander_nd_flat(vander_fs, points, degrees): + """ + Like `_vander_nd`, but flattens the last ``len(degrees)`` axes into a single axis + + Used to implement the public ``vanderd`` functions. + """ + v = _vander_nd(vander_fs, points, degrees) + return v.reshape(v.shape[:-len(degrees)] + (-1,)) + + +def _fromroots(line_f, mul_f, roots): + """ + Helper function used to implement the ``fromroots`` functions. + + Parameters + ---------- + line_f : function(float, float) -> ndarray + The ``line`` function, such as ``polyline`` + mul_f : function(array_like, array_like) -> ndarray + The ``mul`` function, such as ``polymul`` + roots + See the ``fromroots`` functions for more detail + """ + if len(roots) == 0: + return np.ones(1) + else: + [roots] = as_series([roots], trim=False) + roots.sort() + p = [line_f(-r, 1) for r in roots] + n = len(p) + while n > 1: + m, r = divmod(n, 2) + tmp = [mul_f(p[i], p[i+m]) for i in range(m)] + if r: + tmp[0] = mul_f(tmp[0], p[-1]) + p = tmp + n = m + return p[0] + + +def _valnd(val_f, c, *args): + """ + Helper function used to implement the ``vald`` functions. + + Parameters + ---------- + val_f : function(array_like, array_like, tensor: bool) -> array_like + The ``val`` function, such as ``polyval`` + c, args + See the ``vald`` functions for more detail + """ + args = [np.asanyarray(a) for a in args] + shape0 = args[0].shape + if not all(a.shape == shape0 for a in args[1:]): + if len(args) == 3: + raise ValueError('x, y, z are incompatible') + elif len(args) == 2: + raise ValueError('x, y are incompatible') + else: + raise ValueError('ordinates are incompatible') + it = iter(args) + x0 = next(it) + + # use tensor on only the first + c = val_f(x0, c) + for xi in it: + c = val_f(xi, c, tensor=False) + return c + + +def _gridnd(val_f, c, *args): + """ + Helper function used to implement the ``gridd`` functions. + + Parameters + ---------- + val_f : function(array_like, array_like, tensor: bool) -> array_like + The ``val`` function, such as ``polyval`` + c, args + See the ``gridd`` functions for more detail + """ + for xi in args: + c = val_f(xi, c) + return c + + +def _div(mul_f, c1, c2): + """ + Helper function used to implement the ``div`` functions. + + Implementation uses repeated subtraction of c2 multiplied by the nth basis. + For some polynomial types, a more efficient approach may be possible. + + Parameters + ---------- + mul_f : function(array_like, array_like) -> array_like + The ``mul`` function, such as ``polymul`` + c1, c2 + See the ``div`` functions for more detail + """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError # FIXME: add message with details to exception + + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + quo = np.empty(lc1 - lc2 + 1, dtype=c1.dtype) + rem = c1 + for i in range(lc1 - lc2, - 1, -1): + p = mul_f([0]*i + [1], c2) + q = rem[-1]/p[-1] + rem = rem[:-1] - q*p[:-1] + quo[i] = q + return quo, trimseq(rem) + + +def _add(c1, c2): + """ Helper function used to implement the ``add`` functions. """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if len(c1) > len(c2): + c1[:c2.size] += c2 + ret = c1 + else: + c2[:c1.size] += c1 + ret = c2 + return trimseq(ret) + + +def _sub(c1, c2): + """ Helper function used to implement the ``sub`` functions. """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if len(c1) > len(c2): + c1[:c2.size] -= c2 + ret = c1 + else: + c2 = -c2 + c2[:c1.size] += c1 + ret = c2 + return trimseq(ret) + + +def _fit(vander_f, x, y, deg, rcond=None, full=False, w=None): + """ + Helper function used to implement the ``fit`` functions. + + Parameters + ---------- + vander_f : function(array_like, int) -> ndarray + The 1d vander function, such as ``polyvander`` + c1, c2 + See the ``fit`` functions for more detail + """ + x = np.asarray(x) + 0.0 + y = np.asarray(y) + 0.0 + deg = np.asarray(deg) + + # check arguments. + if deg.ndim > 1 or deg.dtype.kind not in 'iu' or deg.size == 0: + raise TypeError("deg must be an int or non-empty 1-D array of int") + if deg.min() < 0: + raise ValueError("expected deg >= 0") + if x.ndim != 1: + raise TypeError("expected 1D vector for x") + if x.size == 0: + raise TypeError("expected non-empty vector for x") + if y.ndim < 1 or y.ndim > 2: + raise TypeError("expected 1D or 2D array for y") + if len(x) != len(y): + raise TypeError("expected x and y to have same length") + + if deg.ndim == 0: + lmax = deg + order = lmax + 1 + van = vander_f(x, lmax) + else: + deg = np.sort(deg) + lmax = deg[-1] + order = len(deg) + van = vander_f(x, lmax)[:, deg] + + # set up the least squares matrices in transposed form + lhs = van.T + rhs = y.T + if w is not None: + w = np.asarray(w) + 0.0 + if w.ndim != 1: + raise TypeError("expected 1D vector for w") + if len(x) != len(w): + raise TypeError("expected x and w to have same length") + # apply weights. Don't use inplace operations as they + # can cause problems with NA. + lhs = lhs * w + rhs = rhs * w + + # set rcond + if rcond is None: + rcond = len(x)*np.finfo(x.dtype).eps + + # Determine the norms of the design matrix columns. + if issubclass(lhs.dtype.type, np.complexfloating): + scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1)) + else: + scl = np.sqrt(np.square(lhs).sum(1)) + scl[scl == 0] = 1 + + # Solve the least squares problem. + c, resids, rank, s = np.linalg.lstsq(lhs.T/scl, rhs.T, rcond) + c = (c.T/scl).T + + # Expand c to include non-fitted coefficients which are set to zero + if deg.ndim > 0: + if c.ndim == 2: + cc = np.zeros((lmax+1, c.shape[1]), dtype=c.dtype) + else: + cc = np.zeros(lmax+1, dtype=c.dtype) + cc[deg] = c + c = cc + + # warn on rank reduction + if rank != order and not full: + msg = "The fit may be poorly conditioned" + warnings.warn(msg, RankWarning, stacklevel=2) + + if full: + return c, [resids, rank, s, rcond] + else: + return c + + +def _pow(mul_f, c, pow, maxpower): + """ + Helper function used to implement the ``pow`` functions. + + Parameters + ---------- + mul_f : function(array_like, array_like) -> ndarray + The ``mul`` function, such as ``polymul`` + c : array_like + 1-D array of array of series coefficients + pow, maxpower + See the ``pow`` functions for more detail + """ + # c is a trimmed copy + [c] = as_series([c]) + power = int(pow) + if power != pow or power < 0: + raise ValueError("Power must be a non-negative integer.") + elif maxpower is not None and power > maxpower: + raise ValueError("Power is too large") + elif power == 0: + return np.array([1], dtype=c.dtype) + elif power == 1: + return c + else: + # This can be made more efficient by using powers of two + # in the usual way. + prd = c + for i in range(2, power + 1): + prd = mul_f(prd, c) + return prd + + +def _as_int(x, desc): + """ + Like `operator.index`, but emits a custom exception when passed an + incorrect type + + Parameters + ---------- + x : int-like + Value to interpret as an integer + desc : str + description to include in any error message + + Raises + ------ + TypeError : if x is a float or non-numeric + """ + try: + return operator.index(x) + except TypeError as e: + raise TypeError(f"{desc} must be an integer, received {x}") from e + + +def format_float(x, parens=False): + if not np.issubdtype(type(x), np.floating): + return str(x) + + opts = np.get_printoptions() + + if np.isnan(x): + return opts['nanstr'] + elif np.isinf(x): + return opts['infstr'] + + exp_format = False + if x != 0: + a = np.abs(x) + if a >= 1.e8 or a < 10**min(0, -(opts['precision']-1)//2): + exp_format = True + + trim, unique = '0', True + if opts['floatmode'] == 'fixed': + trim, unique = 'k', False + + if exp_format: + s = dragon4_scientific(x, precision=opts['precision'], + unique=unique, trim=trim, + sign=opts['sign'] == '+') + if parens: + s = '(' + s + ')' + else: + s = dragon4_positional(x, precision=opts['precision'], + fractional=True, + unique=unique, trim=trim, + sign=opts['sign'] == '+') + return s diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.pyi b/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9299b23975b1ff9c59d36c9e6e804e06d415cf4b --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/polyutils.pyi @@ -0,0 +1,431 @@ +from collections.abc import Callable, Iterable, Sequence +from typing import ( + Any, + Final, + Literal, + SupportsIndex, + TypeAlias, + TypeVar, + overload, +) + +import numpy as np +import numpy.typing as npt +from numpy._typing import ( + _FloatLike_co, + _NumberLike_co, + + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, +) + +from ._polytypes import ( + _AnyInt, + _CoefLike_co, + + _Array2, + _Tuple2, + + _FloatSeries, + _CoefSeries, + _ComplexSeries, + _ObjectSeries, + + _ComplexArray, + _FloatArray, + _CoefArray, + _ObjectArray, + + _SeriesLikeInt_co, + _SeriesLikeFloat_co, + _SeriesLikeComplex_co, + _SeriesLikeCoef_co, + + _ArrayLikeCoef_co, + + _FuncBinOp, + _FuncValND, + _FuncVanderND, +) + +__all__: Final[Sequence[str]] = [ + "as_series", + "format_float", + "getdomain", + "mapdomain", + "mapparms", + "trimcoef", + "trimseq", +] + +_AnyLineF: TypeAlias = Callable[ + [_CoefLike_co, _CoefLike_co], + _CoefArray, +] +_AnyMulF: TypeAlias = Callable[ + [npt.ArrayLike, npt.ArrayLike], + _CoefArray, +] +_AnyVanderF: TypeAlias = Callable[ + [npt.ArrayLike, SupportsIndex], + _CoefArray, +] + +@overload +def as_series( + alist: npt.NDArray[np.integer[Any]] | _FloatArray, + trim: bool = ..., +) -> list[_FloatSeries]: ... +@overload +def as_series( + alist: _ComplexArray, + trim: bool = ..., +) -> list[_ComplexSeries]: ... +@overload +def as_series( + alist: _ObjectArray, + trim: bool = ..., +) -> list[_ObjectSeries]: ... +@overload +def as_series( # type: ignore[overload-overlap] + alist: Iterable[_FloatArray | npt.NDArray[np.integer[Any]]], + trim: bool = ..., +) -> list[_FloatSeries]: ... +@overload +def as_series( + alist: Iterable[_ComplexArray], + trim: bool = ..., +) -> list[_ComplexSeries]: ... +@overload +def as_series( + alist: Iterable[_ObjectArray], + trim: bool = ..., +) -> list[_ObjectSeries]: ... +@overload +def as_series( # type: ignore[overload-overlap] + alist: Iterable[_SeriesLikeFloat_co | float], + trim: bool = ..., +) -> list[_FloatSeries]: ... +@overload +def as_series( + alist: Iterable[_SeriesLikeComplex_co | complex], + trim: bool = ..., +) -> list[_ComplexSeries]: ... +@overload +def as_series( + alist: Iterable[_SeriesLikeCoef_co | object], + trim: bool = ..., +) -> list[_ObjectSeries]: ... + +_T_seq = TypeVar("_T_seq", bound=_CoefArray | Sequence[_CoefLike_co]) +def trimseq(seq: _T_seq) -> _T_seq: ... + +@overload +def trimcoef( # type: ignore[overload-overlap] + c: npt.NDArray[np.integer[Any]] | _FloatArray, + tol: _FloatLike_co = ..., +) -> _FloatSeries: ... +@overload +def trimcoef( + c: _ComplexArray, + tol: _FloatLike_co = ..., +) -> _ComplexSeries: ... +@overload +def trimcoef( + c: _ObjectArray, + tol: _FloatLike_co = ..., +) -> _ObjectSeries: ... +@overload +def trimcoef( # type: ignore[overload-overlap] + c: _SeriesLikeFloat_co | float, + tol: _FloatLike_co = ..., +) -> _FloatSeries: ... +@overload +def trimcoef( + c: _SeriesLikeComplex_co | complex, + tol: _FloatLike_co = ..., +) -> _ComplexSeries: ... +@overload +def trimcoef( + c: _SeriesLikeCoef_co | object, + tol: _FloatLike_co = ..., +) -> _ObjectSeries: ... + +@overload +def getdomain( # type: ignore[overload-overlap] + x: _FloatArray | npt.NDArray[np.integer[Any]], +) -> _Array2[np.float64]: ... +@overload +def getdomain( + x: _ComplexArray, +) -> _Array2[np.complex128]: ... +@overload +def getdomain( + x: _ObjectArray, +) -> _Array2[np.object_]: ... +@overload +def getdomain( # type: ignore[overload-overlap] + x: _SeriesLikeFloat_co | float, +) -> _Array2[np.float64]: ... +@overload +def getdomain( + x: _SeriesLikeComplex_co | complex, +) -> _Array2[np.complex128]: ... +@overload +def getdomain( + x: _SeriesLikeCoef_co | object, +) -> _Array2[np.object_]: ... + +@overload +def mapparms( # type: ignore[overload-overlap] + old: npt.NDArray[np.floating[Any] | np.integer[Any]], + new: npt.NDArray[np.floating[Any] | np.integer[Any]], +) -> _Tuple2[np.floating[Any]]: ... +@overload +def mapparms( + old: npt.NDArray[np.number[Any]], + new: npt.NDArray[np.number[Any]], +) -> _Tuple2[np.complexfloating[Any, Any]]: ... +@overload +def mapparms( + old: npt.NDArray[np.object_ | np.number[Any]], + new: npt.NDArray[np.object_ | np.number[Any]], +) -> _Tuple2[object]: ... +@overload +def mapparms( # type: ignore[overload-overlap] + old: Sequence[float], + new: Sequence[float], +) -> _Tuple2[float]: ... +@overload +def mapparms( + old: Sequence[complex], + new: Sequence[complex], +) -> _Tuple2[complex]: ... +@overload +def mapparms( + old: _SeriesLikeFloat_co, + new: _SeriesLikeFloat_co, +) -> _Tuple2[np.floating[Any]]: ... +@overload +def mapparms( + old: _SeriesLikeComplex_co, + new: _SeriesLikeComplex_co, +) -> _Tuple2[np.complexfloating[Any, Any]]: ... +@overload +def mapparms( + old: _SeriesLikeCoef_co, + new: _SeriesLikeCoef_co, +) -> _Tuple2[object]: ... + +@overload +def mapdomain( # type: ignore[overload-overlap] + x: _FloatLike_co, + old: _SeriesLikeFloat_co, + new: _SeriesLikeFloat_co, +) -> np.floating[Any]: ... +@overload +def mapdomain( + x: _NumberLike_co, + old: _SeriesLikeComplex_co, + new: _SeriesLikeComplex_co, +) -> np.complexfloating[Any, Any]: ... +@overload +def mapdomain( # type: ignore[overload-overlap] + x: npt.NDArray[np.floating[Any] | np.integer[Any]], + old: npt.NDArray[np.floating[Any] | np.integer[Any]], + new: npt.NDArray[np.floating[Any] | np.integer[Any]], +) -> _FloatSeries: ... +@overload +def mapdomain( + x: npt.NDArray[np.number[Any]], + old: npt.NDArray[np.number[Any]], + new: npt.NDArray[np.number[Any]], +) -> _ComplexSeries: ... +@overload +def mapdomain( + x: npt.NDArray[np.object_ | np.number[Any]], + old: npt.NDArray[np.object_ | np.number[Any]], + new: npt.NDArray[np.object_ | np.number[Any]], +) -> _ObjectSeries: ... +@overload +def mapdomain( # type: ignore[overload-overlap] + x: _SeriesLikeFloat_co, + old: _SeriesLikeFloat_co, + new: _SeriesLikeFloat_co, +) -> _FloatSeries: ... +@overload +def mapdomain( + x: _SeriesLikeComplex_co, + old: _SeriesLikeComplex_co, + new: _SeriesLikeComplex_co, +) -> _ComplexSeries: ... +@overload +def mapdomain( + x: _SeriesLikeCoef_co, + old:_SeriesLikeCoef_co, + new: _SeriesLikeCoef_co, +) -> _ObjectSeries: ... +@overload +def mapdomain( + x: _CoefLike_co, + old: _SeriesLikeCoef_co, + new: _SeriesLikeCoef_co, +) -> object: ... + +def _nth_slice( + i: SupportsIndex, + ndim: SupportsIndex, +) -> tuple[None | slice, ...]: ... + +_vander_nd: _FuncVanderND[Literal["_vander_nd"]] +_vander_nd_flat: _FuncVanderND[Literal["_vander_nd_flat"]] + +# keep in sync with `._polytypes._FuncFromRoots` +@overload +def _fromroots( # type: ignore[overload-overlap] + line_f: _AnyLineF, + mul_f: _AnyMulF, + roots: _SeriesLikeFloat_co, +) -> _FloatSeries: ... +@overload +def _fromroots( + line_f: _AnyLineF, + mul_f: _AnyMulF, + roots: _SeriesLikeComplex_co, +) -> _ComplexSeries: ... +@overload +def _fromroots( + line_f: _AnyLineF, + mul_f: _AnyMulF, + roots: _SeriesLikeCoef_co, +) -> _ObjectSeries: ... +@overload +def _fromroots( + line_f: _AnyLineF, + mul_f: _AnyMulF, + roots: _SeriesLikeCoef_co, +) -> _CoefSeries: ... + +_valnd: _FuncValND[Literal["_valnd"]] +_gridnd: _FuncValND[Literal["_gridnd"]] + +# keep in sync with `_polytypes._FuncBinOp` +@overload +def _div( # type: ignore[overload-overlap] + mul_f: _AnyMulF, + c1: _SeriesLikeFloat_co, + c2: _SeriesLikeFloat_co, +) -> _Tuple2[_FloatSeries]: ... +@overload +def _div( + mul_f: _AnyMulF, + c1: _SeriesLikeComplex_co, + c2: _SeriesLikeComplex_co, +) -> _Tuple2[_ComplexSeries]: ... +@overload +def _div( + mul_f: _AnyMulF, + c1: _SeriesLikeCoef_co, + c2: _SeriesLikeCoef_co, +) -> _Tuple2[_ObjectSeries]: ... +@overload +def _div( + mul_f: _AnyMulF, + c1: _SeriesLikeCoef_co, + c2: _SeriesLikeCoef_co, +) -> _Tuple2[_CoefSeries]: ... + +_add: Final[_FuncBinOp] +_sub: Final[_FuncBinOp] + +# keep in sync with `_polytypes._FuncPow` +@overload +def _pow( # type: ignore[overload-overlap] + mul_f: _AnyMulF, + c: _SeriesLikeFloat_co, + pow: _AnyInt, + maxpower: None | _AnyInt = ..., +) -> _FloatSeries: ... +@overload +def _pow( + mul_f: _AnyMulF, + c: _SeriesLikeComplex_co, + pow: _AnyInt, + maxpower: None | _AnyInt = ..., +) -> _ComplexSeries: ... +@overload +def _pow( + mul_f: _AnyMulF, + c: _SeriesLikeCoef_co, + pow: _AnyInt, + maxpower: None | _AnyInt = ..., +) -> _ObjectSeries: ... +@overload +def _pow( + mul_f: _AnyMulF, + c: _SeriesLikeCoef_co, + pow: _AnyInt, + maxpower: None | _AnyInt = ..., +) -> _CoefSeries: ... + +# keep in sync with `_polytypes._FuncFit` +@overload +def _fit( # type: ignore[overload-overlap] + vander_f: _AnyVanderF, + x: _SeriesLikeFloat_co, + y: _ArrayLikeFloat_co, + deg: _SeriesLikeInt_co, + domain: None | _SeriesLikeFloat_co = ..., + rcond: None | _FloatLike_co = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeFloat_co = ..., +) -> _FloatArray: ... +@overload +def _fit( + vander_f: _AnyVanderF, + x: _SeriesLikeComplex_co, + y: _ArrayLikeComplex_co, + deg: _SeriesLikeInt_co, + domain: None | _SeriesLikeComplex_co = ..., + rcond: None | _FloatLike_co = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeComplex_co = ..., +) -> _ComplexArray: ... +@overload +def _fit( + vander_f: _AnyVanderF, + x: _SeriesLikeCoef_co, + y: _ArrayLikeCoef_co, + deg: _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co = ..., + rcond: None | _FloatLike_co = ..., + full: Literal[False] = ..., + w: None | _SeriesLikeCoef_co = ..., +) -> _CoefArray: ... +@overload +def _fit( + vander_f: _AnyVanderF, + x: _SeriesLikeCoef_co, + y: _SeriesLikeCoef_co, + deg: _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co, + rcond: None | _FloatLike_co , + full: Literal[True], + /, + w: None | _SeriesLikeCoef_co = ..., +) -> tuple[_CoefSeries, Sequence[np.inexact[Any] | np.int32]]: ... +@overload +def _fit( + vander_f: _AnyVanderF, + x: _SeriesLikeCoef_co, + y: _SeriesLikeCoef_co, + deg: _SeriesLikeInt_co, + domain: None | _SeriesLikeCoef_co = ..., + rcond: None | _FloatLike_co = ..., + *, + full: Literal[True], + w: None | _SeriesLikeCoef_co = ..., +) -> tuple[_CoefSeries, Sequence[np.inexact[Any] | np.int32]]: ... + +def _as_int(x: SupportsIndex, desc: str) -> int: ... +def format_float(x: _FloatLike_co, parens: bool = ...) -> str: ... diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__init__.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2654e1fffd2d2e0261c465b862aec348427817f7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_chebyshev.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_chebyshev.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f3b9ee410f97e3c985800a5b8eeff3359588d9b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_chebyshev.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30043ed36f6a545743860f7005654af228b43772 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df570f0b29373105f93e94c98fe4819e1296d642 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite_e.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite_e.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20c3c2438278e65682644243724059ecfdf75b2b Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_hermite_e.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff59d9e968bd7f2cd7bac7e727e7e33c818a66a3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8e2ba76fc460e9cf42baef842bbb47c95bb61dd Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f079a6e40bbac7fa6e5216f0504bfb75829ed9c Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f321dfd502d78871a8b1d4184cb2fdb4146480ac Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ed70f7e9ead622141c76043d2cdf5a2165d3f7d Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-310.pyc b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7979c8944d4386b6c04f432c974b79efb0cc938 Binary files /dev/null and b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_chebyshev.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_chebyshev.py new file mode 100644 index 0000000000000000000000000000000000000000..2f54bebfdb27d54f436378e4ab6d6c8f2426dd90 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_chebyshev.py @@ -0,0 +1,619 @@ +"""Tests for chebyshev module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.chebyshev as cheb +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + + +def trim(x): + return cheb.chebtrim(x, tol=1e-6) + +T0 = [1] +T1 = [0, 1] +T2 = [-1, 0, 2] +T3 = [0, -3, 0, 4] +T4 = [1, 0, -8, 0, 8] +T5 = [0, 5, 0, -20, 0, 16] +T6 = [-1, 0, 18, 0, -48, 0, 32] +T7 = [0, -7, 0, 56, 0, -112, 0, 64] +T8 = [1, 0, -32, 0, 160, 0, -256, 0, 128] +T9 = [0, 9, 0, -120, 0, 432, 0, -576, 0, 256] + +Tlist = [T0, T1, T2, T3, T4, T5, T6, T7, T8, T9] + + +class TestPrivate: + + def test__cseries_to_zseries(self): + for i in range(5): + inp = np.array([2] + [1]*i, np.double) + tgt = np.array([.5]*i + [2] + [.5]*i, np.double) + res = cheb._cseries_to_zseries(inp) + assert_equal(res, tgt) + + def test__zseries_to_cseries(self): + for i in range(5): + inp = np.array([.5]*i + [2] + [.5]*i, np.double) + tgt = np.array([2] + [1]*i, np.double) + res = cheb._zseries_to_cseries(inp) + assert_equal(res, tgt) + + +class TestConstants: + + def test_chebdomain(self): + assert_equal(cheb.chebdomain, [-1, 1]) + + def test_chebzero(self): + assert_equal(cheb.chebzero, [0]) + + def test_chebone(self): + assert_equal(cheb.chebone, [1]) + + def test_chebx(self): + assert_equal(cheb.chebx, [0, 1]) + + +class TestArithmetic: + + def test_chebadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = cheb.chebadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = cheb.chebsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebmulx(self): + assert_equal(cheb.chebmulx([0]), [0]) + assert_equal(cheb.chebmulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [.5, 0, .5] + assert_equal(cheb.chebmulx(ser), tgt) + + def test_chebmul(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(i + j + 1) + tgt[i + j] += .5 + tgt[abs(i - j)] += .5 + res = cheb.chebmul([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = cheb.chebadd(ci, cj) + quo, rem = cheb.chebdiv(tgt, ci) + res = cheb.chebadd(cheb.chebmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(cheb.chebmul, [c]*j, np.array([1])) + res = cheb.chebpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2.5, 2., 1.5]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_chebval(self): + #check empty input + assert_equal(cheb.chebval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Tlist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = cheb.chebval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(cheb.chebval(x, [1]).shape, dims) + assert_equal(cheb.chebval(x, [1, 0]).shape, dims) + assert_equal(cheb.chebval(x, [1, 0, 0]).shape, dims) + + def test_chebval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, cheb.chebval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = cheb.chebval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_chebval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, cheb.chebval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = cheb.chebval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_chebgrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = cheb.chebgrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebgrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_chebgrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = cheb.chebgrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebgrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_chebint(self): + # check exceptions + assert_raises(TypeError, cheb.chebint, [0], .5) + assert_raises(ValueError, cheb.chebint, [0], -1) + assert_raises(ValueError, cheb.chebint, [0], 1, [0, 0]) + assert_raises(ValueError, cheb.chebint, [0], lbnd=[0]) + assert_raises(ValueError, cheb.chebint, [0], scl=[0]) + assert_raises(TypeError, cheb.chebint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = cheb.chebint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i]) + res = cheb.cheb2poly(chebint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(cheb.chebval(-1, chebint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i], scl=2) + res = cheb.cheb2poly(chebint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1) + res = cheb.chebint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k]) + res = cheb.chebint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k], lbnd=-1) + res = cheb.chebint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k], scl=2) + res = cheb.chebint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([cheb.chebint(c) for c in c2d.T]).T + res = cheb.chebint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebint(c) for c in c2d]) + res = cheb.chebint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebint(c, k=3) for c in c2d]) + res = cheb.chebint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_chebder(self): + # check exceptions + assert_raises(TypeError, cheb.chebder, [0], .5) + assert_raises(ValueError, cheb.chebder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = cheb.chebder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = cheb.chebder(cheb.chebint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = cheb.chebder(cheb.chebint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([cheb.chebder(c) for c in c2d.T]).T + res = cheb.chebder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebder(c) for c in c2d]) + res = cheb.chebder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_chebvander(self): + # check for 1d x + x = np.arange(3) + v = cheb.chebvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], cheb.chebval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = cheb.chebvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], cheb.chebval(x, coef)) + + def test_chebvander2d(self): + # also tests chebval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = cheb.chebvander2d(x1, x2, [1, 2]) + tgt = cheb.chebval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = cheb.chebvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_chebvander3d(self): + # also tests chebval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = cheb.chebvander3d(x1, x2, x3, [1, 2, 3]) + tgt = cheb.chebval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = cheb.chebvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_chebfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, cheb.chebfit, [1], [1], -1) + assert_raises(TypeError, cheb.chebfit, [[1]], [1], 0) + assert_raises(TypeError, cheb.chebfit, [], [1], 0) + assert_raises(TypeError, cheb.chebfit, [1], [[[1]]], 0) + assert_raises(TypeError, cheb.chebfit, [1, 2], [1], 0) + assert_raises(TypeError, cheb.chebfit, [1], [1, 2], 0) + assert_raises(TypeError, cheb.chebfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, cheb.chebfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, cheb.chebfit, [1], [1], [-1,]) + assert_raises(ValueError, cheb.chebfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, cheb.chebfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = cheb.chebfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(cheb.chebval(x, coef3), y) + coef3 = cheb.chebfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(cheb.chebval(x, coef3), y) + # + coef4 = cheb.chebfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + coef4 = cheb.chebfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = cheb.chebfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + # + coef2d = cheb.chebfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = cheb.chebfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = cheb.chebfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = cheb.chebfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = cheb.chebfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = cheb.chebfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(cheb.chebfit(x, x, 1), [0, 1]) + assert_almost_equal(cheb.chebfit(x, x, [0, 1]), [0, 1]) + # test fitting only even polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = cheb.chebfit(x, y, 4) + assert_almost_equal(cheb.chebval(x, coef1), y) + coef2 = cheb.chebfit(x, y, [0, 2, 4]) + assert_almost_equal(cheb.chebval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestInterpolate: + + def f(self, x): + return x * (x - 1) * (x - 2) + + def test_raises(self): + assert_raises(ValueError, cheb.chebinterpolate, self.f, -1) + assert_raises(TypeError, cheb.chebinterpolate, self.f, 10.) + + def test_dimensions(self): + for deg in range(1, 5): + assert_(cheb.chebinterpolate(self.f, deg).shape == (deg + 1,)) + + def test_approximation(self): + + def powx(x, p): + return x**p + + x = np.linspace(-1, 1, 10) + for deg in range(0, 10): + for p in range(0, deg + 1): + c = cheb.chebinterpolate(powx, deg, (p,)) + assert_almost_equal(cheb.chebval(x, c), powx(x, p), decimal=12) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, cheb.chebcompanion, []) + assert_raises(ValueError, cheb.chebcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(cheb.chebcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(cheb.chebcompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = cheb.chebgauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = cheb.chebvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.pi + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_chebfromroots(self): + res = cheb.chebfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + tgt = [0]*i + [1] + res = cheb.chebfromroots(roots)*2**(i-1) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebroots(self): + assert_almost_equal(cheb.chebroots([1]), []) + assert_almost_equal(cheb.chebroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = cheb.chebroots(cheb.chebfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, cheb.chebtrim, coef, -1) + + # Test results + assert_equal(cheb.chebtrim(coef), coef[:-1]) + assert_equal(cheb.chebtrim(coef, 1), coef[:-3]) + assert_equal(cheb.chebtrim(coef, 2), [0]) + + def test_chebline(self): + assert_equal(cheb.chebline(3, 4), [3, 4]) + + def test_cheb2poly(self): + for i in range(10): + assert_almost_equal(cheb.cheb2poly([0]*i + [1]), Tlist[i]) + + def test_poly2cheb(self): + for i in range(10): + assert_almost_equal(cheb.poly2cheb(Tlist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-1, 1, 11)[1:-1] + tgt = 1./(np.sqrt(1 + x) * np.sqrt(1 - x)) + res = cheb.chebweight(x) + assert_almost_equal(res, tgt) + + def test_chebpts1(self): + #test exceptions + assert_raises(ValueError, cheb.chebpts1, 1.5) + assert_raises(ValueError, cheb.chebpts1, 0) + + #test points + tgt = [0] + assert_almost_equal(cheb.chebpts1(1), tgt) + tgt = [-0.70710678118654746, 0.70710678118654746] + assert_almost_equal(cheb.chebpts1(2), tgt) + tgt = [-0.86602540378443871, 0, 0.86602540378443871] + assert_almost_equal(cheb.chebpts1(3), tgt) + tgt = [-0.9238795325, -0.3826834323, 0.3826834323, 0.9238795325] + assert_almost_equal(cheb.chebpts1(4), tgt) + + def test_chebpts2(self): + #test exceptions + assert_raises(ValueError, cheb.chebpts2, 1.5) + assert_raises(ValueError, cheb.chebpts2, 1) + + #test points + tgt = [-1, 1] + assert_almost_equal(cheb.chebpts2(2), tgt) + tgt = [-1, 0, 1] + assert_almost_equal(cheb.chebpts2(3), tgt) + tgt = [-1, -0.5, .5, 1] + assert_almost_equal(cheb.chebpts2(4), tgt) + tgt = [-1.0, -0.707106781187, 0, 0.707106781187, 1.0] + assert_almost_equal(cheb.chebpts2(5), tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_classes.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_classes.py new file mode 100644 index 0000000000000000000000000000000000000000..75672a148524d8887663b986ec5d9e6c13d1193a --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_classes.py @@ -0,0 +1,607 @@ +"""Test inter-conversion of different polynomial classes. + +This tests the convert and cast methods of all the polynomial classes. + +""" +import operator as op +from numbers import Number + +import pytest +import numpy as np +from numpy.polynomial import ( + Polynomial, Legendre, Chebyshev, Laguerre, Hermite, HermiteE) +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) +from numpy.exceptions import RankWarning + +# +# fixtures +# + +classes = ( + Polynomial, Legendre, Chebyshev, Laguerre, + Hermite, HermiteE + ) +classids = tuple(cls.__name__ for cls in classes) + +@pytest.fixture(params=classes, ids=classids) +def Poly(request): + return request.param + +# +# helper functions +# +random = np.random.random + + +def assert_poly_almost_equal(p1, p2, msg=""): + try: + assert_(np.all(p1.domain == p2.domain)) + assert_(np.all(p1.window == p2.window)) + assert_almost_equal(p1.coef, p2.coef) + except AssertionError: + msg = f"Result: {p1}\nTarget: {p2}" + raise AssertionError(msg) + + +# +# Test conversion methods that depend on combinations of two classes. +# + +Poly1 = Poly +Poly2 = Poly + + +def test_conversion(Poly1, Poly2): + x = np.linspace(0, 1, 10) + coef = random((3,)) + + d1 = Poly1.domain + random((2,))*.25 + w1 = Poly1.window + random((2,))*.25 + p1 = Poly1(coef, domain=d1, window=w1) + + d2 = Poly2.domain + random((2,))*.25 + w2 = Poly2.window + random((2,))*.25 + p2 = p1.convert(kind=Poly2, domain=d2, window=w2) + + assert_almost_equal(p2.domain, d2) + assert_almost_equal(p2.window, w2) + assert_almost_equal(p2(x), p1(x)) + + +def test_cast(Poly1, Poly2): + x = np.linspace(0, 1, 10) + coef = random((3,)) + + d1 = Poly1.domain + random((2,))*.25 + w1 = Poly1.window + random((2,))*.25 + p1 = Poly1(coef, domain=d1, window=w1) + + d2 = Poly2.domain + random((2,))*.25 + w2 = Poly2.window + random((2,))*.25 + p2 = Poly2.cast(p1, domain=d2, window=w2) + + assert_almost_equal(p2.domain, d2) + assert_almost_equal(p2.window, w2) + assert_almost_equal(p2(x), p1(x)) + + +# +# test methods that depend on one class +# + + +def test_identity(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + x = np.linspace(d[0], d[1], 11) + p = Poly.identity(domain=d, window=w) + assert_equal(p.domain, d) + assert_equal(p.window, w) + assert_almost_equal(p(x), x) + + +def test_basis(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly.basis(5, domain=d, window=w) + assert_equal(p.domain, d) + assert_equal(p.window, w) + assert_equal(p.coef, [0]*5 + [1]) + + +def test_fromroots(Poly): + # check that requested roots are zeros of a polynomial + # of correct degree, domain, and window. + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + r = random((5,)) + p1 = Poly.fromroots(r, domain=d, window=w) + assert_equal(p1.degree(), len(r)) + assert_equal(p1.domain, d) + assert_equal(p1.window, w) + assert_almost_equal(p1(r), 0) + + # check that polynomial is monic + pdom = Polynomial.domain + pwin = Polynomial.window + p2 = Polynomial.cast(p1, domain=pdom, window=pwin) + assert_almost_equal(p2.coef[-1], 1) + + +def test_bad_conditioned_fit(Poly): + + x = [0., 0., 1.] + y = [1., 2., 3.] + + # check RankWarning is raised + with pytest.warns(RankWarning) as record: + Poly.fit(x, y, 2) + assert record[0].message.args[0] == "The fit may be poorly conditioned" + + +def test_fit(Poly): + + def f(x): + return x*(x - 1)*(x - 2) + x = np.linspace(0, 3) + y = f(x) + + # check default value of domain and window + p = Poly.fit(x, y, 3) + assert_almost_equal(p.domain, [0, 3]) + assert_almost_equal(p(x), y) + assert_equal(p.degree(), 3) + + # check with given domains and window + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly.fit(x, y, 3, domain=d, window=w) + assert_almost_equal(p(x), y) + assert_almost_equal(p.domain, d) + assert_almost_equal(p.window, w) + p = Poly.fit(x, y, [0, 1, 2, 3], domain=d, window=w) + assert_almost_equal(p(x), y) + assert_almost_equal(p.domain, d) + assert_almost_equal(p.window, w) + + # check with class domain default + p = Poly.fit(x, y, 3, []) + assert_equal(p.domain, Poly.domain) + assert_equal(p.window, Poly.window) + p = Poly.fit(x, y, [0, 1, 2, 3], []) + assert_equal(p.domain, Poly.domain) + assert_equal(p.window, Poly.window) + + # check that fit accepts weights. + w = np.zeros_like(x) + z = y + random(y.shape)*.25 + w[::2] = 1 + p1 = Poly.fit(x[::2], z[::2], 3) + p2 = Poly.fit(x, z, 3, w=w) + p3 = Poly.fit(x, z, [0, 1, 2, 3], w=w) + assert_almost_equal(p1(x), p2(x)) + assert_almost_equal(p2(x), p3(x)) + + +def test_equal(Poly): + p1 = Poly([1, 2, 3], domain=[0, 1], window=[2, 3]) + p2 = Poly([1, 1, 1], domain=[0, 1], window=[2, 3]) + p3 = Poly([1, 2, 3], domain=[1, 2], window=[2, 3]) + p4 = Poly([1, 2, 3], domain=[0, 1], window=[1, 2]) + assert_(p1 == p1) + assert_(not p1 == p2) + assert_(not p1 == p3) + assert_(not p1 == p4) + + +def test_not_equal(Poly): + p1 = Poly([1, 2, 3], domain=[0, 1], window=[2, 3]) + p2 = Poly([1, 1, 1], domain=[0, 1], window=[2, 3]) + p3 = Poly([1, 2, 3], domain=[1, 2], window=[2, 3]) + p4 = Poly([1, 2, 3], domain=[0, 1], window=[1, 2]) + assert_(not p1 != p1) + assert_(p1 != p2) + assert_(p1 != p3) + assert_(p1 != p4) + + +def test_add(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 + p2 + assert_poly_almost_equal(p2 + p1, p3) + assert_poly_almost_equal(p1 + c2, p3) + assert_poly_almost_equal(c2 + p1, p3) + assert_poly_almost_equal(p1 + tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) + p1, p3) + assert_poly_almost_equal(p1 + np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) + p1, p3) + assert_raises(TypeError, op.add, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.add, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.add, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.add, p1, Polynomial([0])) + + +def test_sub(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 - p2 + assert_poly_almost_equal(p2 - p1, -p3) + assert_poly_almost_equal(p1 - c2, p3) + assert_poly_almost_equal(c2 - p1, -p3) + assert_poly_almost_equal(p1 - tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) - p1, -p3) + assert_poly_almost_equal(p1 - np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) - p1, -p3) + assert_raises(TypeError, op.sub, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.sub, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.sub, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.sub, p1, Polynomial([0])) + + +def test_mul(Poly): + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 * p2 + assert_poly_almost_equal(p2 * p1, p3) + assert_poly_almost_equal(p1 * c2, p3) + assert_poly_almost_equal(c2 * p1, p3) + assert_poly_almost_equal(p1 * tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) * p1, p3) + assert_poly_almost_equal(p1 * np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) * p1, p3) + assert_poly_almost_equal(p1 * 2, p1 * Poly([2])) + assert_poly_almost_equal(2 * p1, p1 * Poly([2])) + assert_raises(TypeError, op.mul, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.mul, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.mul, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.mul, p1, Polynomial([0])) + + +def test_floordiv(Poly): + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + assert_poly_almost_equal(p4 // p2, p1) + assert_poly_almost_equal(p4 // c2, p1) + assert_poly_almost_equal(c4 // p2, p1) + assert_poly_almost_equal(p4 // tuple(c2), p1) + assert_poly_almost_equal(tuple(c4) // p2, p1) + assert_poly_almost_equal(p4 // np.array(c2), p1) + assert_poly_almost_equal(np.array(c4) // p2, p1) + assert_poly_almost_equal(2 // p2, Poly([0])) + assert_poly_almost_equal(p2 // 2, 0.5*p2) + assert_raises( + TypeError, op.floordiv, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises( + TypeError, op.floordiv, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.floordiv, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.floordiv, p1, Polynomial([0])) + + +def test_truediv(Poly): + # true division is valid only if the denominator is a Number and + # not a python bool. + p1 = Poly([1,2,3]) + p2 = p1 * 5 + + for stype in np.ScalarType: + if not issubclass(stype, Number) or issubclass(stype, bool): + continue + s = stype(5) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for stype in (int, float): + s = stype(5) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for stype in [complex]: + s = stype(5, 0) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for s in [tuple(), list(), dict(), bool(), np.array([1])]: + assert_raises(TypeError, op.truediv, p2, s) + assert_raises(TypeError, op.truediv, s, p2) + for ptype in classes: + assert_raises(TypeError, op.truediv, p2, ptype(1)) + + +def test_mod(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + assert_poly_almost_equal(p4 % p2, p3) + assert_poly_almost_equal(p4 % c2, p3) + assert_poly_almost_equal(c4 % p2, p3) + assert_poly_almost_equal(p4 % tuple(c2), p3) + assert_poly_almost_equal(tuple(c4) % p2, p3) + assert_poly_almost_equal(p4 % np.array(c2), p3) + assert_poly_almost_equal(np.array(c4) % p2, p3) + assert_poly_almost_equal(2 % p2, Poly([2])) + assert_poly_almost_equal(p2 % 2, Poly([0])) + assert_raises(TypeError, op.mod, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.mod, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.mod, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.mod, p1, Polynomial([0])) + + +def test_divmod(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + quo, rem = divmod(p4, p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, c2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(c4, p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, tuple(c2)) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(tuple(c4), p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, np.array(c2)) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(np.array(c4), p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p2, 2) + assert_poly_almost_equal(quo, 0.5*p2) + assert_poly_almost_equal(rem, Poly([0])) + quo, rem = divmod(2, p2) + assert_poly_almost_equal(quo, Poly([0])) + assert_poly_almost_equal(rem, Poly([2])) + assert_raises(TypeError, divmod, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, divmod, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, divmod, p1, Chebyshev([0])) + else: + assert_raises(TypeError, divmod, p1, Polynomial([0])) + + +def test_roots(Poly): + d = Poly.domain * 1.25 + .25 + w = Poly.window + tgt = np.linspace(d[0], d[1], 5) + res = np.sort(Poly.fromroots(tgt, domain=d, window=w).roots()) + assert_almost_equal(res, tgt) + # default domain and window + res = np.sort(Poly.fromroots(tgt).roots()) + assert_almost_equal(res, tgt) + + +def test_degree(Poly): + p = Poly.basis(5) + assert_equal(p.degree(), 5) + + +def test_copy(Poly): + p1 = Poly.basis(5) + p2 = p1.copy() + assert_(p1 == p2) + assert_(p1 is not p2) + assert_(p1.coef is not p2.coef) + assert_(p1.domain is not p2.domain) + assert_(p1.window is not p2.window) + + +def test_integ(Poly): + P = Polynomial + # Check defaults + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ()) + p2 = P.cast(p0.integ(2)) + assert_poly_almost_equal(p1, P([0, 2, 3, 4])) + assert_poly_almost_equal(p2, P([0, 0, 1, 1, 1])) + # Check with k + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ(k=1)) + p2 = P.cast(p0.integ(2, k=[1, 1])) + assert_poly_almost_equal(p1, P([1, 2, 3, 4])) + assert_poly_almost_equal(p2, P([1, 1, 1, 1, 1])) + # Check with lbnd + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ(lbnd=1)) + p2 = P.cast(p0.integ(2, lbnd=1)) + assert_poly_almost_equal(p1, P([-9, 2, 3, 4])) + assert_poly_almost_equal(p2, P([6, -9, 1, 1, 1])) + # Check scaling + d = 2*Poly.domain + p0 = Poly.cast(P([1*2, 2*3, 3*4]), domain=d) + p1 = P.cast(p0.integ()) + p2 = P.cast(p0.integ(2)) + assert_poly_almost_equal(p1, P([0, 2, 3, 4])) + assert_poly_almost_equal(p2, P([0, 0, 1, 1, 1])) + + +def test_deriv(Poly): + # Check that the derivative is the inverse of integration. It is + # assumes that the integration has been checked elsewhere. + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p1 = Poly([1, 2, 3], domain=d, window=w) + p2 = p1.integ(2, k=[1, 2]) + p3 = p1.integ(1, k=[1]) + assert_almost_equal(p2.deriv(1).coef, p3.coef) + assert_almost_equal(p2.deriv(2).coef, p1.coef) + # default domain and window + p1 = Poly([1, 2, 3]) + p2 = p1.integ(2, k=[1, 2]) + p3 = p1.integ(1, k=[1]) + assert_almost_equal(p2.deriv(1).coef, p3.coef) + assert_almost_equal(p2.deriv(2).coef, p1.coef) + + +def test_linspace(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly([1, 2, 3], domain=d, window=w) + # check default domain + xtgt = np.linspace(d[0], d[1], 20) + ytgt = p(xtgt) + xres, yres = p.linspace(20) + assert_almost_equal(xres, xtgt) + assert_almost_equal(yres, ytgt) + # check specified domain + xtgt = np.linspace(0, 2, 20) + ytgt = p(xtgt) + xres, yres = p.linspace(20, domain=[0, 2]) + assert_almost_equal(xres, xtgt) + assert_almost_equal(yres, ytgt) + + +def test_pow(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + tgt = Poly([1], domain=d, window=w) + tst = Poly([1, 2, 3], domain=d, window=w) + for i in range(5): + assert_poly_almost_equal(tst**i, tgt) + tgt = tgt * tst + # default domain and window + tgt = Poly([1]) + tst = Poly([1, 2, 3]) + for i in range(5): + assert_poly_almost_equal(tst**i, tgt) + tgt = tgt * tst + # check error for invalid powers + assert_raises(ValueError, op.pow, tgt, 1.5) + assert_raises(ValueError, op.pow, tgt, -1) + + +def test_call(Poly): + P = Polynomial + d = Poly.domain + x = np.linspace(d[0], d[1], 11) + + # Check defaults + p = Poly.cast(P([1, 2, 3])) + tgt = 1 + x*(2 + 3*x) + res = p(x) + assert_almost_equal(res, tgt) + + +def test_call_with_list(Poly): + p = Poly([1, 2, 3]) + x = [-1, 0, 2] + res = p(x) + assert_equal(res, p(np.array(x))) + + +def test_cutdeg(Poly): + p = Poly([1, 2, 3]) + assert_raises(ValueError, p.cutdeg, .5) + assert_raises(ValueError, p.cutdeg, -1) + assert_equal(len(p.cutdeg(3)), 3) + assert_equal(len(p.cutdeg(2)), 3) + assert_equal(len(p.cutdeg(1)), 2) + assert_equal(len(p.cutdeg(0)), 1) + + +def test_truncate(Poly): + p = Poly([1, 2, 3]) + assert_raises(ValueError, p.truncate, .5) + assert_raises(ValueError, p.truncate, 0) + assert_equal(len(p.truncate(4)), 3) + assert_equal(len(p.truncate(3)), 3) + assert_equal(len(p.truncate(2)), 2) + assert_equal(len(p.truncate(1)), 1) + + +def test_trim(Poly): + c = [1, 1e-6, 1e-12, 0] + p = Poly(c) + assert_equal(p.trim().coef, c[:3]) + assert_equal(p.trim(1e-10).coef, c[:2]) + assert_equal(p.trim(1e-5).coef, c[:1]) + + +def test_mapparms(Poly): + # check with defaults. Should be identity. + d = Poly.domain + w = Poly.window + p = Poly([1], domain=d, window=w) + assert_almost_equal([0, 1], p.mapparms()) + # + w = 2*d + 1 + p = Poly([1], domain=d, window=w) + assert_almost_equal([1, 2], p.mapparms()) + + +def test_ufunc_override(Poly): + p = Poly([1, 2, 3]) + x = np.ones(3) + assert_raises(TypeError, np.add, p, x) + assert_raises(TypeError, np.add, x, p) + + +# +# Test class method that only exists for some classes +# + + +class TestInterpolate: + + def f(self, x): + return x * (x - 1) * (x - 2) + + def test_raises(self): + assert_raises(ValueError, Chebyshev.interpolate, self.f, -1) + assert_raises(TypeError, Chebyshev.interpolate, self.f, 10.) + + def test_dimensions(self): + for deg in range(1, 5): + assert_(Chebyshev.interpolate(self.f, deg).degree() == deg) + + def test_approximation(self): + + def powx(x, p): + return x**p + + x = np.linspace(0, 2, 10) + for deg in range(0, 10): + for t in range(0, deg + 1): + p = Chebyshev.interpolate(powx, deg, domain=[0, 2], args=(t,)) + assert_almost_equal(p(x), powx(x, t), decimal=11) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite.py new file mode 100644 index 0000000000000000000000000000000000000000..2188800853f2f5e9a98d2d7087893a7cf11440ef --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite.py @@ -0,0 +1,555 @@ +"""Tests for hermite module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.hermite as herm +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +H0 = np.array([1]) +H1 = np.array([0, 2]) +H2 = np.array([-2, 0, 4]) +H3 = np.array([0, -12, 0, 8]) +H4 = np.array([12, 0, -48, 0, 16]) +H5 = np.array([0, 120, 0, -160, 0, 32]) +H6 = np.array([-120, 0, 720, 0, -480, 0, 64]) +H7 = np.array([0, -1680, 0, 3360, 0, -1344, 0, 128]) +H8 = np.array([1680, 0, -13440, 0, 13440, 0, -3584, 0, 256]) +H9 = np.array([0, 30240, 0, -80640, 0, 48384, 0, -9216, 0, 512]) + +Hlist = [H0, H1, H2, H3, H4, H5, H6, H7, H8, H9] + + +def trim(x): + return herm.hermtrim(x, tol=1e-6) + + +class TestConstants: + + def test_hermdomain(self): + assert_equal(herm.hermdomain, [-1, 1]) + + def test_hermzero(self): + assert_equal(herm.hermzero, [0]) + + def test_hermone(self): + assert_equal(herm.hermone, [1]) + + def test_hermx(self): + assert_equal(herm.hermx, [0, .5]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_hermadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = herm.hermadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = herm.hermsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermmulx(self): + assert_equal(herm.hermmulx([0]), [0]) + assert_equal(herm.hermmulx([1]), [0, .5]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i, 0, .5] + assert_equal(herm.hermmulx(ser), tgt) + + def test_hermmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = herm.hermval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = herm.hermval(self.x, pol2) + pol3 = herm.hermmul(pol1, pol2) + val3 = herm.hermval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_hermdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = herm.hermadd(ci, cj) + quo, rem = herm.hermdiv(tgt, ci) + res = herm.hermadd(herm.hermmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(herm.hermmul, [c]*j, np.array([1])) + res = herm.hermpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2.5, 1., .75]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_hermval(self): + #check empty input + assert_equal(herm.hermval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Hlist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = herm.hermval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(herm.hermval(x, [1]).shape, dims) + assert_equal(herm.hermval(x, [1, 0]).shape, dims) + assert_equal(herm.hermval(x, [1, 0, 0]).shape, dims) + + def test_hermval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herm.hermval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = herm.hermval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_hermval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herm.hermval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = herm.hermval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_hermgrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = herm.hermgrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermgrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_hermgrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = herm.hermgrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermgrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_hermint(self): + # check exceptions + assert_raises(TypeError, herm.hermint, [0], .5) + assert_raises(ValueError, herm.hermint, [0], -1) + assert_raises(ValueError, herm.hermint, [0], 1, [0, 0]) + assert_raises(ValueError, herm.hermint, [0], lbnd=[0]) + assert_raises(ValueError, herm.hermint, [0], scl=[0]) + assert_raises(TypeError, herm.hermint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = herm.hermint([0], m=i, k=k) + assert_almost_equal(res, [0, .5]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i]) + res = herm.herm2poly(hermint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(herm.hermval(-1, hermint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i], scl=2) + res = herm.herm2poly(hermint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1) + res = herm.hermint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k]) + res = herm.hermint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k], lbnd=-1) + res = herm.hermint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k], scl=2) + res = herm.hermint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herm.hermint(c) for c in c2d.T]).T + res = herm.hermint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermint(c) for c in c2d]) + res = herm.hermint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermint(c, k=3) for c in c2d]) + res = herm.hermint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_hermder(self): + # check exceptions + assert_raises(TypeError, herm.hermder, [0], .5) + assert_raises(ValueError, herm.hermder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = herm.hermder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herm.hermder(herm.hermint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herm.hermder(herm.hermint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herm.hermder(c) for c in c2d.T]).T + res = herm.hermder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermder(c) for c in c2d]) + res = herm.hermder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_hermvander(self): + # check for 1d x + x = np.arange(3) + v = herm.hermvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herm.hermval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = herm.hermvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herm.hermval(x, coef)) + + def test_hermvander2d(self): + # also tests hermval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = herm.hermvander2d(x1, x2, [1, 2]) + tgt = herm.hermval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herm.hermvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_hermvander3d(self): + # also tests hermval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = herm.hermvander3d(x1, x2, x3, [1, 2, 3]) + tgt = herm.hermval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herm.hermvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_hermfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, herm.hermfit, [1], [1], -1) + assert_raises(TypeError, herm.hermfit, [[1]], [1], 0) + assert_raises(TypeError, herm.hermfit, [], [1], 0) + assert_raises(TypeError, herm.hermfit, [1], [[[1]]], 0) + assert_raises(TypeError, herm.hermfit, [1, 2], [1], 0) + assert_raises(TypeError, herm.hermfit, [1], [1, 2], 0) + assert_raises(TypeError, herm.hermfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, herm.hermfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, herm.hermfit, [1], [1], [-1,]) + assert_raises(ValueError, herm.hermfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, herm.hermfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = herm.hermfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(herm.hermval(x, coef3), y) + coef3 = herm.hermfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(herm.hermval(x, coef3), y) + # + coef4 = herm.hermfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + coef4 = herm.hermfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = herm.hermfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + # + coef2d = herm.hermfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = herm.hermfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = herm.hermfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = herm.hermfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = herm.hermfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = herm.hermfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(herm.hermfit(x, x, 1), [0, .5]) + assert_almost_equal(herm.hermfit(x, x, [0, 1]), [0, .5]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = herm.hermfit(x, y, 4) + assert_almost_equal(herm.hermval(x, coef1), y) + coef2 = herm.hermfit(x, y, [0, 2, 4]) + assert_almost_equal(herm.hermval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, herm.hermcompanion, []) + assert_raises(ValueError, herm.hermcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(herm.hermcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(herm.hermcompanion([1, 2])[0, 0] == -.25) + + +class TestGauss: + + def test_100(self): + x, w = herm.hermgauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = herm.hermvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.sqrt(np.pi) + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_hermfromroots(self): + res = herm.hermfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = herm.hermfromroots(roots) + res = herm.hermval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(herm.herm2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_hermroots(self): + assert_almost_equal(herm.hermroots([1]), []) + assert_almost_equal(herm.hermroots([1, 1]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = herm.hermroots(herm.hermfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, herm.hermtrim, coef, -1) + + # Test results + assert_equal(herm.hermtrim(coef), coef[:-1]) + assert_equal(herm.hermtrim(coef, 1), coef[:-3]) + assert_equal(herm.hermtrim(coef, 2), [0]) + + def test_hermline(self): + assert_equal(herm.hermline(3, 4), [3, 2]) + + def test_herm2poly(self): + for i in range(10): + assert_almost_equal(herm.herm2poly([0]*i + [1]), Hlist[i]) + + def test_poly2herm(self): + for i in range(10): + assert_almost_equal(herm.poly2herm(Hlist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-5, 5, 11) + tgt = np.exp(-x**2) + res = herm.hermweight(x) + assert_almost_equal(res, tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite_e.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite_e.py new file mode 100644 index 0000000000000000000000000000000000000000..2d262a3306222bd79f682b09763b0bd2b90ba8fe --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_hermite_e.py @@ -0,0 +1,556 @@ +"""Tests for hermite_e module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.hermite_e as herme +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +He0 = np.array([1]) +He1 = np.array([0, 1]) +He2 = np.array([-1, 0, 1]) +He3 = np.array([0, -3, 0, 1]) +He4 = np.array([3, 0, -6, 0, 1]) +He5 = np.array([0, 15, 0, -10, 0, 1]) +He6 = np.array([-15, 0, 45, 0, -15, 0, 1]) +He7 = np.array([0, -105, 0, 105, 0, -21, 0, 1]) +He8 = np.array([105, 0, -420, 0, 210, 0, -28, 0, 1]) +He9 = np.array([0, 945, 0, -1260, 0, 378, 0, -36, 0, 1]) + +Helist = [He0, He1, He2, He3, He4, He5, He6, He7, He8, He9] + + +def trim(x): + return herme.hermetrim(x, tol=1e-6) + + +class TestConstants: + + def test_hermedomain(self): + assert_equal(herme.hermedomain, [-1, 1]) + + def test_hermezero(self): + assert_equal(herme.hermezero, [0]) + + def test_hermeone(self): + assert_equal(herme.hermeone, [1]) + + def test_hermex(self): + assert_equal(herme.hermex, [0, 1]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_hermeadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = herme.hermeadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermesub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = herme.hermesub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermemulx(self): + assert_equal(herme.hermemulx([0]), [0]) + assert_equal(herme.hermemulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i, 0, 1] + assert_equal(herme.hermemulx(ser), tgt) + + def test_hermemul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = herme.hermeval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = herme.hermeval(self.x, pol2) + pol3 = herme.hermemul(pol1, pol2) + val3 = herme.hermeval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_hermediv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = herme.hermeadd(ci, cj) + quo, rem = herme.hermediv(tgt, ci) + res = herme.hermeadd(herme.hermemul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermepow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(herme.hermemul, [c]*j, np.array([1])) + res = herme.hermepow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([4., 2., 3.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_hermeval(self): + #check empty input + assert_equal(herme.hermeval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Helist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = herme.hermeval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(herme.hermeval(x, [1]).shape, dims) + assert_equal(herme.hermeval(x, [1, 0]).shape, dims) + assert_equal(herme.hermeval(x, [1, 0, 0]).shape, dims) + + def test_hermeval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herme.hermeval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = herme.hermeval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermeval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_hermeval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herme.hermeval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = herme.hermeval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermeval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_hermegrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = herme.hermegrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermegrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_hermegrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = herme.hermegrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermegrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_hermeint(self): + # check exceptions + assert_raises(TypeError, herme.hermeint, [0], .5) + assert_raises(ValueError, herme.hermeint, [0], -1) + assert_raises(ValueError, herme.hermeint, [0], 1, [0, 0]) + assert_raises(ValueError, herme.hermeint, [0], lbnd=[0]) + assert_raises(ValueError, herme.hermeint, [0], scl=[0]) + assert_raises(TypeError, herme.hermeint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = herme.hermeint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i]) + res = herme.herme2poly(hermeint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i], lbnd=-1) + assert_almost_equal(herme.hermeval(-1, hermeint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i], scl=2) + res = herme.herme2poly(hermeint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1) + res = herme.hermeint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k]) + res = herme.hermeint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k], lbnd=-1) + res = herme.hermeint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k], scl=2) + res = herme.hermeint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermeint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herme.hermeint(c) for c in c2d.T]).T + res = herme.hermeint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeint(c) for c in c2d]) + res = herme.hermeint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeint(c, k=3) for c in c2d]) + res = herme.hermeint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_hermeder(self): + # check exceptions + assert_raises(TypeError, herme.hermeder, [0], .5) + assert_raises(ValueError, herme.hermeder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = herme.hermeder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herme.hermeder(herme.hermeint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herme.hermeder( + herme.hermeint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermeder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herme.hermeder(c) for c in c2d.T]).T + res = herme.hermeder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeder(c) for c in c2d]) + res = herme.hermeder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_hermevander(self): + # check for 1d x + x = np.arange(3) + v = herme.hermevander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herme.hermeval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = herme.hermevander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herme.hermeval(x, coef)) + + def test_hermevander2d(self): + # also tests hermeval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = herme.hermevander2d(x1, x2, [1, 2]) + tgt = herme.hermeval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herme.hermevander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_hermevander3d(self): + # also tests hermeval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = herme.hermevander3d(x1, x2, x3, [1, 2, 3]) + tgt = herme.hermeval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herme.hermevander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_hermefit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, herme.hermefit, [1], [1], -1) + assert_raises(TypeError, herme.hermefit, [[1]], [1], 0) + assert_raises(TypeError, herme.hermefit, [], [1], 0) + assert_raises(TypeError, herme.hermefit, [1], [[[1]]], 0) + assert_raises(TypeError, herme.hermefit, [1, 2], [1], 0) + assert_raises(TypeError, herme.hermefit, [1], [1, 2], 0) + assert_raises(TypeError, herme.hermefit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, herme.hermefit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, herme.hermefit, [1], [1], [-1,]) + assert_raises(ValueError, herme.hermefit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, herme.hermefit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = herme.hermefit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(herme.hermeval(x, coef3), y) + coef3 = herme.hermefit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(herme.hermeval(x, coef3), y) + # + coef4 = herme.hermefit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + coef4 = herme.hermefit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = herme.hermefit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + # + coef2d = herme.hermefit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = herme.hermefit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = herme.hermefit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = herme.hermefit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = herme.hermefit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = herme.hermefit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(herme.hermefit(x, x, 1), [0, 1]) + assert_almost_equal(herme.hermefit(x, x, [0, 1]), [0, 1]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = herme.hermefit(x, y, 4) + assert_almost_equal(herme.hermeval(x, coef1), y) + coef2 = herme.hermefit(x, y, [0, 2, 4]) + assert_almost_equal(herme.hermeval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, herme.hermecompanion, []) + assert_raises(ValueError, herme.hermecompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(herme.hermecompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(herme.hermecompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = herme.hermegauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = herme.hermevander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.sqrt(2*np.pi) + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_hermefromroots(self): + res = herme.hermefromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = herme.hermefromroots(roots) + res = herme.hermeval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(herme.herme2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_hermeroots(self): + assert_almost_equal(herme.hermeroots([1]), []) + assert_almost_equal(herme.hermeroots([1, 1]), [-1]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = herme.hermeroots(herme.hermefromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermetrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, herme.hermetrim, coef, -1) + + # Test results + assert_equal(herme.hermetrim(coef), coef[:-1]) + assert_equal(herme.hermetrim(coef, 1), coef[:-3]) + assert_equal(herme.hermetrim(coef, 2), [0]) + + def test_hermeline(self): + assert_equal(herme.hermeline(3, 4), [3, 4]) + + def test_herme2poly(self): + for i in range(10): + assert_almost_equal(herme.herme2poly([0]*i + [1]), Helist[i]) + + def test_poly2herme(self): + for i in range(10): + assert_almost_equal(herme.poly2herme(Helist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-5, 5, 11) + tgt = np.exp(-.5*x**2) + res = herme.hermeweight(x) + assert_almost_equal(res, tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_laguerre.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_laguerre.py new file mode 100644 index 0000000000000000000000000000000000000000..49f7c7e115bec499a04f58c38d803d3e8be1247e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_laguerre.py @@ -0,0 +1,537 @@ +"""Tests for laguerre module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.laguerre as lag +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +L0 = np.array([1])/1 +L1 = np.array([1, -1])/1 +L2 = np.array([2, -4, 1])/2 +L3 = np.array([6, -18, 9, -1])/6 +L4 = np.array([24, -96, 72, -16, 1])/24 +L5 = np.array([120, -600, 600, -200, 25, -1])/120 +L6 = np.array([720, -4320, 5400, -2400, 450, -36, 1])/720 + +Llist = [L0, L1, L2, L3, L4, L5, L6] + + +def trim(x): + return lag.lagtrim(x, tol=1e-6) + + +class TestConstants: + + def test_lagdomain(self): + assert_equal(lag.lagdomain, [0, 1]) + + def test_lagzero(self): + assert_equal(lag.lagzero, [0]) + + def test_lagone(self): + assert_equal(lag.lagone, [1]) + + def test_lagx(self): + assert_equal(lag.lagx, [1, -1]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_lagadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = lag.lagadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = lag.lagsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagmulx(self): + assert_equal(lag.lagmulx([0]), [0]) + assert_equal(lag.lagmulx([1]), [1, -1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [-i, 2*i + 1, -(i + 1)] + assert_almost_equal(lag.lagmulx(ser), tgt) + + def test_lagmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = lag.lagval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = lag.lagval(self.x, pol2) + pol3 = lag.lagmul(pol1, pol2) + val3 = lag.lagval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_lagdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = lag.lagadd(ci, cj) + quo, rem = lag.lagdiv(tgt, ci) + res = lag.lagadd(lag.lagmul(quo, ci), rem) + assert_almost_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(lag.lagmul, [c]*j, np.array([1])) + res = lag.lagpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([9., -14., 6.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_lagval(self): + #check empty input + assert_equal(lag.lagval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Llist] + for i in range(7): + msg = f"At i={i}" + tgt = y[i] + res = lag.lagval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(lag.lagval(x, [1]).shape, dims) + assert_equal(lag.lagval(x, [1, 0]).shape, dims) + assert_equal(lag.lagval(x, [1, 0, 0]).shape, dims) + + def test_lagval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, lag.lagval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = lag.lagval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.lagval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_lagval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, lag.lagval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = lag.lagval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.lagval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_laggrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = lag.laggrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.laggrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_laggrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = lag.laggrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.laggrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_lagint(self): + # check exceptions + assert_raises(TypeError, lag.lagint, [0], .5) + assert_raises(ValueError, lag.lagint, [0], -1) + assert_raises(ValueError, lag.lagint, [0], 1, [0, 0]) + assert_raises(ValueError, lag.lagint, [0], lbnd=[0]) + assert_raises(ValueError, lag.lagint, [0], scl=[0]) + assert_raises(TypeError, lag.lagint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = lag.lagint([0], m=i, k=k) + assert_almost_equal(res, [1, -1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i]) + res = lag.lag2poly(lagint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(lag.lagval(-1, lagint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i], scl=2) + res = lag.lag2poly(lagint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1) + res = lag.lagint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k]) + res = lag.lagint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k], lbnd=-1) + res = lag.lagint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k], scl=2) + res = lag.lagint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([lag.lagint(c) for c in c2d.T]).T + res = lag.lagint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagint(c) for c in c2d]) + res = lag.lagint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagint(c, k=3) for c in c2d]) + res = lag.lagint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_lagder(self): + # check exceptions + assert_raises(TypeError, lag.lagder, [0], .5) + assert_raises(ValueError, lag.lagder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = lag.lagder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = lag.lagder(lag.lagint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = lag.lagder(lag.lagint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([lag.lagder(c) for c in c2d.T]).T + res = lag.lagder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagder(c) for c in c2d]) + res = lag.lagder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_lagvander(self): + # check for 1d x + x = np.arange(3) + v = lag.lagvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], lag.lagval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = lag.lagvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], lag.lagval(x, coef)) + + def test_lagvander2d(self): + # also tests lagval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = lag.lagvander2d(x1, x2, [1, 2]) + tgt = lag.lagval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = lag.lagvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_lagvander3d(self): + # also tests lagval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = lag.lagvander3d(x1, x2, x3, [1, 2, 3]) + tgt = lag.lagval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = lag.lagvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_lagfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + # Test exceptions + assert_raises(ValueError, lag.lagfit, [1], [1], -1) + assert_raises(TypeError, lag.lagfit, [[1]], [1], 0) + assert_raises(TypeError, lag.lagfit, [], [1], 0) + assert_raises(TypeError, lag.lagfit, [1], [[[1]]], 0) + assert_raises(TypeError, lag.lagfit, [1, 2], [1], 0) + assert_raises(TypeError, lag.lagfit, [1], [1, 2], 0) + assert_raises(TypeError, lag.lagfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, lag.lagfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, lag.lagfit, [1], [1], [-1,]) + assert_raises(ValueError, lag.lagfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, lag.lagfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = lag.lagfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(lag.lagval(x, coef3), y) + coef3 = lag.lagfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(lag.lagval(x, coef3), y) + # + coef4 = lag.lagfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(lag.lagval(x, coef4), y) + coef4 = lag.lagfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(lag.lagval(x, coef4), y) + # + coef2d = lag.lagfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = lag.lagfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = lag.lagfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = lag.lagfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = lag.lagfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = lag.lagfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(lag.lagfit(x, x, 1), [1, -1]) + assert_almost_equal(lag.lagfit(x, x, [0, 1]), [1, -1]) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, lag.lagcompanion, []) + assert_raises(ValueError, lag.lagcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(lag.lagcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(lag.lagcompanion([1, 2])[0, 0] == 1.5) + + +class TestGauss: + + def test_100(self): + x, w = lag.laggauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = lag.lagvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = 1.0 + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_lagfromroots(self): + res = lag.lagfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = lag.lagfromroots(roots) + res = lag.lagval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(lag.lag2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_lagroots(self): + assert_almost_equal(lag.lagroots([1]), []) + assert_almost_equal(lag.lagroots([0, 1]), [1]) + for i in range(2, 5): + tgt = np.linspace(0, 3, i) + res = lag.lagroots(lag.lagfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, lag.lagtrim, coef, -1) + + # Test results + assert_equal(lag.lagtrim(coef), coef[:-1]) + assert_equal(lag.lagtrim(coef, 1), coef[:-3]) + assert_equal(lag.lagtrim(coef, 2), [0]) + + def test_lagline(self): + assert_equal(lag.lagline(3, 4), [7, -4]) + + def test_lag2poly(self): + for i in range(7): + assert_almost_equal(lag.lag2poly([0]*i + [1]), Llist[i]) + + def test_poly2lag(self): + for i in range(7): + assert_almost_equal(lag.poly2lag(Llist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(0, 10, 11) + tgt = np.exp(-x) + res = lag.lagweight(x) + assert_almost_equal(res, tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_legendre.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_legendre.py new file mode 100644 index 0000000000000000000000000000000000000000..9f1c9733a91121e208d7037f8e93b27f0cdbf9bb --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_legendre.py @@ -0,0 +1,568 @@ +"""Tests for legendre module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.legendre as leg +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +L0 = np.array([1]) +L1 = np.array([0, 1]) +L2 = np.array([-1, 0, 3])/2 +L3 = np.array([0, -3, 0, 5])/2 +L4 = np.array([3, 0, -30, 0, 35])/8 +L5 = np.array([0, 15, 0, -70, 0, 63])/8 +L6 = np.array([-5, 0, 105, 0, -315, 0, 231])/16 +L7 = np.array([0, -35, 0, 315, 0, -693, 0, 429])/16 +L8 = np.array([35, 0, -1260, 0, 6930, 0, -12012, 0, 6435])/128 +L9 = np.array([0, 315, 0, -4620, 0, 18018, 0, -25740, 0, 12155])/128 + +Llist = [L0, L1, L2, L3, L4, L5, L6, L7, L8, L9] + + +def trim(x): + return leg.legtrim(x, tol=1e-6) + + +class TestConstants: + + def test_legdomain(self): + assert_equal(leg.legdomain, [-1, 1]) + + def test_legzero(self): + assert_equal(leg.legzero, [0]) + + def test_legone(self): + assert_equal(leg.legone, [1]) + + def test_legx(self): + assert_equal(leg.legx, [0, 1]) + + +class TestArithmetic: + x = np.linspace(-1, 1, 100) + + def test_legadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = leg.legadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = leg.legsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legmulx(self): + assert_equal(leg.legmulx([0]), [0]) + assert_equal(leg.legmulx([1]), [0, 1]) + for i in range(1, 5): + tmp = 2*i + 1 + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i/tmp, 0, (i + 1)/tmp] + assert_equal(leg.legmulx(ser), tgt) + + def test_legmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = leg.legval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = leg.legval(self.x, pol2) + pol3 = leg.legmul(pol1, pol2) + val3 = leg.legval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_legdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = leg.legadd(ci, cj) + quo, rem = leg.legdiv(tgt, ci) + res = leg.legadd(leg.legmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(leg.legmul, [c]*j, np.array([1])) + res = leg.legpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2., 2., 2.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_legval(self): + #check empty input + assert_equal(leg.legval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Llist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = leg.legval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(leg.legval(x, [1]).shape, dims) + assert_equal(leg.legval(x, [1, 0]).shape, dims) + assert_equal(leg.legval(x, [1, 0, 0]).shape, dims) + + def test_legval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, leg.legval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = leg.legval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.legval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_legval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, leg.legval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = leg.legval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.legval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_leggrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = leg.leggrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.leggrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_leggrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = leg.leggrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.leggrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_legint(self): + # check exceptions + assert_raises(TypeError, leg.legint, [0], .5) + assert_raises(ValueError, leg.legint, [0], -1) + assert_raises(ValueError, leg.legint, [0], 1, [0, 0]) + assert_raises(ValueError, leg.legint, [0], lbnd=[0]) + assert_raises(ValueError, leg.legint, [0], scl=[0]) + assert_raises(TypeError, leg.legint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = leg.legint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i]) + res = leg.leg2poly(legint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(leg.legval(-1, legint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i], scl=2) + res = leg.leg2poly(legint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1) + res = leg.legint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k]) + res = leg.legint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k], lbnd=-1) + res = leg.legint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k], scl=2) + res = leg.legint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([leg.legint(c) for c in c2d.T]).T + res = leg.legint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legint(c) for c in c2d]) + res = leg.legint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legint(c, k=3) for c in c2d]) + res = leg.legint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + def test_legint_zerointord(self): + assert_equal(leg.legint((1, 2, 3), 0), (1, 2, 3)) + + +class TestDerivative: + + def test_legder(self): + # check exceptions + assert_raises(TypeError, leg.legder, [0], .5) + assert_raises(ValueError, leg.legder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = leg.legder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = leg.legder(leg.legint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = leg.legder(leg.legint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([leg.legder(c) for c in c2d.T]).T + res = leg.legder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legder(c) for c in c2d]) + res = leg.legder(c2d, axis=1) + assert_almost_equal(res, tgt) + + def test_legder_orderhigherthancoeff(self): + c = (1, 2, 3, 4) + assert_equal(leg.legder(c, 4), [0]) + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_legvander(self): + # check for 1d x + x = np.arange(3) + v = leg.legvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], leg.legval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = leg.legvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], leg.legval(x, coef)) + + def test_legvander2d(self): + # also tests polyval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = leg.legvander2d(x1, x2, [1, 2]) + tgt = leg.legval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = leg.legvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_legvander3d(self): + # also tests polyval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = leg.legvander3d(x1, x2, x3, [1, 2, 3]) + tgt = leg.legval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = leg.legvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + def test_legvander_negdeg(self): + assert_raises(ValueError, leg.legvander, (1, 2, 3), -1) + + +class TestFitting: + + def test_legfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, leg.legfit, [1], [1], -1) + assert_raises(TypeError, leg.legfit, [[1]], [1], 0) + assert_raises(TypeError, leg.legfit, [], [1], 0) + assert_raises(TypeError, leg.legfit, [1], [[[1]]], 0) + assert_raises(TypeError, leg.legfit, [1, 2], [1], 0) + assert_raises(TypeError, leg.legfit, [1], [1, 2], 0) + assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, leg.legfit, [1], [1], [-1,]) + assert_raises(ValueError, leg.legfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, leg.legfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = leg.legfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(leg.legval(x, coef3), y) + coef3 = leg.legfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(leg.legval(x, coef3), y) + # + coef4 = leg.legfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + coef4 = leg.legfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = leg.legfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + # + coef2d = leg.legfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = leg.legfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = leg.legfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = leg.legfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = leg.legfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = leg.legfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(leg.legfit(x, x, 1), [0, 1]) + assert_almost_equal(leg.legfit(x, x, [0, 1]), [0, 1]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = leg.legfit(x, y, 4) + assert_almost_equal(leg.legval(x, coef1), y) + coef2 = leg.legfit(x, y, [0, 2, 4]) + assert_almost_equal(leg.legval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, leg.legcompanion, []) + assert_raises(ValueError, leg.legcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(leg.legcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(leg.legcompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = leg.leggauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = leg.legvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = 2.0 + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_legfromroots(self): + res = leg.legfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = leg.legfromroots(roots) + res = leg.legval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(leg.leg2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_legroots(self): + assert_almost_equal(leg.legroots([1]), []) + assert_almost_equal(leg.legroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = leg.legroots(leg.legfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, leg.legtrim, coef, -1) + + # Test results + assert_equal(leg.legtrim(coef), coef[:-1]) + assert_equal(leg.legtrim(coef, 1), coef[:-3]) + assert_equal(leg.legtrim(coef, 2), [0]) + + def test_legline(self): + assert_equal(leg.legline(3, 4), [3, 4]) + + def test_legline_zeroscl(self): + assert_equal(leg.legline(3, 0), [3]) + + def test_leg2poly(self): + for i in range(10): + assert_almost_equal(leg.leg2poly([0]*i + [1]), Llist[i]) + + def test_poly2leg(self): + for i in range(10): + assert_almost_equal(leg.poly2leg(Llist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-1, 1, 11) + tgt = 1. + res = leg.legweight(x) + assert_almost_equal(res, tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polynomial.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polynomial.py new file mode 100644 index 0000000000000000000000000000000000000000..d36b07dbd9536b4c1bd1f3129ae7ccaa2a320ed3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polynomial.py @@ -0,0 +1,647 @@ +"""Tests for polynomial module. + +""" +from functools import reduce +from fractions import Fraction +import numpy as np +import numpy.polynomial.polynomial as poly +import numpy.polynomial.polyutils as pu +import pickle +from copy import deepcopy +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + assert_array_equal, assert_raises_regex, assert_warns) + + +def trim(x): + return poly.polytrim(x, tol=1e-6) + +T0 = [1] +T1 = [0, 1] +T2 = [-1, 0, 2] +T3 = [0, -3, 0, 4] +T4 = [1, 0, -8, 0, 8] +T5 = [0, 5, 0, -20, 0, 16] +T6 = [-1, 0, 18, 0, -48, 0, 32] +T7 = [0, -7, 0, 56, 0, -112, 0, 64] +T8 = [1, 0, -32, 0, 160, 0, -256, 0, 128] +T9 = [0, 9, 0, -120, 0, 432, 0, -576, 0, 256] + +Tlist = [T0, T1, T2, T3, T4, T5, T6, T7, T8, T9] + + +class TestConstants: + + def test_polydomain(self): + assert_equal(poly.polydomain, [-1, 1]) + + def test_polyzero(self): + assert_equal(poly.polyzero, [0]) + + def test_polyone(self): + assert_equal(poly.polyone, [1]) + + def test_polyx(self): + assert_equal(poly.polyx, [0, 1]) + + def test_copy(self): + x = poly.Polynomial([1, 2, 3]) + y = deepcopy(x) + assert_equal(x, y) + + def test_pickle(self): + x = poly.Polynomial([1, 2, 3]) + y = pickle.loads(pickle.dumps(x)) + assert_equal(x, y) + +class TestArithmetic: + + def test_polyadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = poly.polyadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polysub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = poly.polysub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polymulx(self): + assert_equal(poly.polymulx([0]), [0]) + assert_equal(poly.polymulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i + 1) + [1] + assert_equal(poly.polymulx(ser), tgt) + + def test_polymul(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(i + j + 1) + tgt[i + j] += 1 + res = poly.polymul([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polydiv(self): + # check zero division + assert_raises(ZeroDivisionError, poly.polydiv, [1], [0]) + + # check scalar division + quo, rem = poly.polydiv([2], [2]) + assert_equal((quo, rem), (1, 0)) + quo, rem = poly.polydiv([2, 2], [2]) + assert_equal((quo, rem), ((1, 1), 0)) + + # check rest. + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1, 2] + cj = [0]*j + [1, 2] + tgt = poly.polyadd(ci, cj) + quo, rem = poly.polydiv(tgt, ci) + res = poly.polyadd(poly.polymul(quo, ci), rem) + assert_equal(res, tgt, err_msg=msg) + + def test_polypow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(poly.polymul, [c]*j, np.array([1])) + res = poly.polypow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + +class TestFraction: + + def test_Fraction(self): + # assert we can use Polynomials with coefficients of object dtype + f = Fraction(2, 3) + one = Fraction(1, 1) + zero = Fraction(0, 1) + p = poly.Polynomial([f, f], domain=[zero, one], window=[zero, one]) + + x = 2 * p + p ** 2 + assert_equal(x.coef, np.array([Fraction(16, 9), Fraction(20, 9), + Fraction(4, 9)], dtype=object)) + assert_equal(p.domain, [zero, one]) + assert_equal(p.coef.dtype, np.dtypes.ObjectDType()) + assert_(isinstance(p(f), Fraction)) + assert_equal(p(f), Fraction(10, 9)) + p_deriv = poly.Polynomial([Fraction(2, 3)], domain=[zero, one], + window=[zero, one]) + assert_equal(p.deriv(), p_deriv) + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([1., 2., 3.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = poly.polyval(x, [1., 2., 3.]) + + def test_polyval(self): + #check empty input + assert_equal(poly.polyval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [x**i for i in range(5)] + for i in range(5): + tgt = y[i] + res = poly.polyval(x, [0]*i + [1]) + assert_almost_equal(res, tgt) + tgt = x*(x**2 - 1) + res = poly.polyval(x, [0, -1, 0, 1]) + assert_almost_equal(res, tgt) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(poly.polyval(x, [1]).shape, dims) + assert_equal(poly.polyval(x, [1, 0]).shape, dims) + assert_equal(poly.polyval(x, [1, 0, 0]).shape, dims) + + #check masked arrays are processed correctly + mask = [False, True, False] + mx = np.ma.array([1, 2, 3], mask=mask) + res = np.polyval([7, 5, 3], mx) + assert_array_equal(res.mask, mask) + + #check subtypes of ndarray are preserved + class C(np.ndarray): + pass + + cx = np.array([1, 2, 3]).view(C) + assert_equal(type(np.polyval([2, 3, 4], cx)), C) + + def test_polyvalfromroots(self): + # check exception for broadcasting x values over root array with + # too few dimensions + assert_raises(ValueError, poly.polyvalfromroots, + [1], [1], tensor=False) + + # check empty input + assert_equal(poly.polyvalfromroots([], [1]).size, 0) + assert_(poly.polyvalfromroots([], [1]).shape == (0,)) + + # check empty input + multidimensional roots + assert_equal(poly.polyvalfromroots([], [[1] * 5]).size, 0) + assert_(poly.polyvalfromroots([], [[1] * 5]).shape == (5, 0)) + + # check scalar input + assert_equal(poly.polyvalfromroots(1, 1), 0) + assert_(poly.polyvalfromroots(1, np.ones((3, 3))).shape == (3,)) + + # check normal input) + x = np.linspace(-1, 1) + y = [x**i for i in range(5)] + for i in range(1, 5): + tgt = y[i] + res = poly.polyvalfromroots(x, [0]*i) + assert_almost_equal(res, tgt) + tgt = x*(x - 1)*(x + 1) + res = poly.polyvalfromroots(x, [-1, 0, 1]) + assert_almost_equal(res, tgt) + + # check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(poly.polyvalfromroots(x, [1]).shape, dims) + assert_equal(poly.polyvalfromroots(x, [1, 0]).shape, dims) + assert_equal(poly.polyvalfromroots(x, [1, 0, 0]).shape, dims) + + # check compatibility with factorization + ptest = [15, 2, -16, -2, 1] + r = poly.polyroots(ptest) + x = np.linspace(-1, 1) + assert_almost_equal(poly.polyval(x, ptest), + poly.polyvalfromroots(x, r)) + + # check multidimensional arrays of roots and values + # check tensor=False + rshape = (3, 5) + x = np.arange(-3, 2) + r = np.random.randint(-5, 5, size=rshape) + res = poly.polyvalfromroots(x, r, tensor=False) + tgt = np.empty(r.shape[1:]) + for ii in range(tgt.size): + tgt[ii] = poly.polyvalfromroots(x[ii], r[:, ii]) + assert_equal(res, tgt) + + # check tensor=True + x = np.vstack([x, 2*x]) + res = poly.polyvalfromroots(x, r, tensor=True) + tgt = np.empty(r.shape[1:] + x.shape) + for ii in range(r.shape[1]): + for jj in range(x.shape[0]): + tgt[ii, jj, :] = poly.polyvalfromroots(x[jj], r[:, ii]) + assert_equal(res, tgt) + + def test_polyval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises_regex(ValueError, 'incompatible', + poly.polyval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = poly.polyval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polyval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_polyval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises_regex(ValueError, 'incompatible', + poly.polyval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = poly.polyval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polyval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_polygrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = poly.polygrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polygrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_polygrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = poly.polygrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polygrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_polyint(self): + # check exceptions + assert_raises(TypeError, poly.polyint, [0], .5) + assert_raises(ValueError, poly.polyint, [0], -1) + assert_raises(ValueError, poly.polyint, [0], 1, [0, 0]) + assert_raises(ValueError, poly.polyint, [0], lbnd=[0]) + assert_raises(ValueError, poly.polyint, [0], scl=[0]) + assert_raises(TypeError, poly.polyint, [0], axis=.5) + assert_raises(TypeError, poly.polyint, [1, 1], 1.) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = poly.polyint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + res = poly.polyint(pol, m=1, k=[i]) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + res = poly.polyint(pol, m=1, k=[i], lbnd=-1) + assert_almost_equal(poly.polyval(-1, res), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + res = poly.polyint(pol, m=1, k=[i], scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1) + res = poly.polyint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k]) + res = poly.polyint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k], lbnd=-1) + res = poly.polyint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k], scl=2) + res = poly.polyint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([poly.polyint(c) for c in c2d.T]).T + res = poly.polyint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyint(c) for c in c2d]) + res = poly.polyint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyint(c, k=3) for c in c2d]) + res = poly.polyint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_polyder(self): + # check exceptions + assert_raises(TypeError, poly.polyder, [0], .5) + assert_raises(ValueError, poly.polyder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = poly.polyder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = poly.polyder(poly.polyint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = poly.polyder(poly.polyint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([poly.polyder(c) for c in c2d.T]).T + res = poly.polyder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyder(c) for c in c2d]) + res = poly.polyder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_polyvander(self): + # check for 1d x + x = np.arange(3) + v = poly.polyvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], poly.polyval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = poly.polyvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], poly.polyval(x, coef)) + + def test_polyvander2d(self): + # also tests polyval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = poly.polyvander2d(x1, x2, [1, 2]) + tgt = poly.polyval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = poly.polyvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_polyvander3d(self): + # also tests polyval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = poly.polyvander3d(x1, x2, x3, [1, 2, 3]) + tgt = poly.polyval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = poly.polyvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + def test_polyvandernegdeg(self): + x = np.arange(3) + assert_raises(ValueError, poly.polyvander, x, -1) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, poly.polycompanion, []) + assert_raises(ValueError, poly.polycompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(poly.polycompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(poly.polycompanion([1, 2])[0, 0] == -.5) + + +class TestMisc: + + def test_polyfromroots(self): + res = poly.polyfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + tgt = Tlist[i] + res = poly.polyfromroots(roots)*2**(i-1) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyroots(self): + assert_almost_equal(poly.polyroots([1]), []) + assert_almost_equal(poly.polyroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = poly.polyroots(poly.polyfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, poly.polyfit, [1], [1], -1) + assert_raises(TypeError, poly.polyfit, [[1]], [1], 0) + assert_raises(TypeError, poly.polyfit, [], [1], 0) + assert_raises(TypeError, poly.polyfit, [1], [[[1]]], 0) + assert_raises(TypeError, poly.polyfit, [1, 2], [1], 0) + assert_raises(TypeError, poly.polyfit, [1], [1, 2], 0) + assert_raises(TypeError, poly.polyfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, poly.polyfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, poly.polyfit, [1], [1], [-1,]) + assert_raises(ValueError, poly.polyfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, poly.polyfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = poly.polyfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(poly.polyval(x, coef3), y) + coef3 = poly.polyfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(poly.polyval(x, coef3), y) + # + coef4 = poly.polyfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(poly.polyval(x, coef4), y) + coef4 = poly.polyfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(poly.polyval(x, coef4), y) + # + coef2d = poly.polyfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = poly.polyfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + yw[0::2] = 0 + wcoef3 = poly.polyfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = poly.polyfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = poly.polyfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = poly.polyfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(poly.polyfit(x, x, 1), [0, 1]) + assert_almost_equal(poly.polyfit(x, x, [0, 1]), [0, 1]) + # test fitting only even Polyendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = poly.polyfit(x, y, 4) + assert_almost_equal(poly.polyval(x, coef1), y) + coef2 = poly.polyfit(x, y, [0, 2, 4]) + assert_almost_equal(poly.polyval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + def test_polytrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, poly.polytrim, coef, -1) + + # Test results + assert_equal(poly.polytrim(coef), coef[:-1]) + assert_equal(poly.polytrim(coef, 1), coef[:-3]) + assert_equal(poly.polytrim(coef, 2), [0]) + + def test_polyline(self): + assert_equal(poly.polyline(3, 4), [3, 4]) + + def test_polyline_zero(self): + assert_equal(poly.polyline(3, 0), [3]) + + def test_fit_degenerate_domain(self): + p = poly.Polynomial.fit([1], [2], deg=0) + assert_equal(p.coef, [2.]) + p = poly.Polynomial.fit([1, 1], [2, 2.1], deg=0) + assert_almost_equal(p.coef, [2.05]) + with assert_warns(pu.RankWarning): + p = poly.Polynomial.fit([1, 1], [2, 2.1], deg=1) + + def test_result_type(self): + w = np.array([-1, 1], dtype=np.float32) + p = np.polynomial.Polynomial(w, domain=w, window=w) + v = p(2) + assert_equal(v.dtype, np.float32) + + arr = np.polydiv(1, np.float32(1)) + assert_equal(arr[0].dtype, np.float64) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polyutils.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polyutils.py new file mode 100644 index 0000000000000000000000000000000000000000..e5143ed5c3e4a1651c67b5260cef47112c5ea071 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_polyutils.py @@ -0,0 +1,125 @@ +"""Tests for polyutils module. + +""" +import numpy as np +import numpy.polynomial.polyutils as pu +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + + +class TestMisc: + + def test_trimseq(self): + tgt = [1] + for num_trailing_zeros in range(5): + res = pu.trimseq([1] + [0] * num_trailing_zeros) + assert_equal(res, tgt) + + def test_trimseq_empty_input(self): + for empty_seq in [[], np.array([], dtype=np.int32)]: + assert_equal(pu.trimseq(empty_seq), empty_seq) + + def test_as_series(self): + # check exceptions + assert_raises(ValueError, pu.as_series, [[]]) + assert_raises(ValueError, pu.as_series, [[[1, 2]]]) + assert_raises(ValueError, pu.as_series, [[1], ['a']]) + # check common types + types = ['i', 'd', 'O'] + for i in range(len(types)): + for j in range(i): + ci = np.ones(1, types[i]) + cj = np.ones(1, types[j]) + [resi, resj] = pu.as_series([ci, cj]) + assert_(resi.dtype.char == resj.dtype.char) + assert_(resj.dtype.char == types[i]) + + def test_trimcoef(self): + coef = [2, -1, 1, 0] + # Test exceptions + assert_raises(ValueError, pu.trimcoef, coef, -1) + # Test results + assert_equal(pu.trimcoef(coef), coef[:-1]) + assert_equal(pu.trimcoef(coef, 1), coef[:-3]) + assert_equal(pu.trimcoef(coef, 2), [0]) + + def test_vander_nd_exception(self): + # n_dims != len(points) + assert_raises(ValueError, pu._vander_nd, (), (1, 2, 3), [90]) + # n_dims != len(degrees) + assert_raises(ValueError, pu._vander_nd, (), (), [90.65]) + # n_dims == 0 + assert_raises(ValueError, pu._vander_nd, (), (), []) + + def test_div_zerodiv(self): + # c2[-1] == 0 + assert_raises(ZeroDivisionError, pu._div, pu._div, (1, 2, 3), [0]) + + def test_pow_too_large(self): + # power > maxpower + assert_raises(ValueError, pu._pow, (), [1, 2, 3], 5, 4) + +class TestDomain: + + def test_getdomain(self): + # test for real values + x = [1, 10, 3, -1] + tgt = [-1, 10] + res = pu.getdomain(x) + assert_almost_equal(res, tgt) + + # test for complex values + x = [1 + 1j, 1 - 1j, 0, 2] + tgt = [-1j, 2 + 1j] + res = pu.getdomain(x) + assert_almost_equal(res, tgt) + + def test_mapdomain(self): + # test for real values + dom1 = [0, 4] + dom2 = [1, 3] + tgt = dom2 + res = pu.mapdomain(dom1, dom1, dom2) + assert_almost_equal(res, tgt) + + # test for complex values + dom1 = [0 - 1j, 2 + 1j] + dom2 = [-2, 2] + tgt = dom2 + x = dom1 + res = pu.mapdomain(x, dom1, dom2) + assert_almost_equal(res, tgt) + + # test for multidimensional arrays + dom1 = [0, 4] + dom2 = [1, 3] + tgt = np.array([dom2, dom2]) + x = np.array([dom1, dom1]) + res = pu.mapdomain(x, dom1, dom2) + assert_almost_equal(res, tgt) + + # test that subtypes are preserved. + class MyNDArray(np.ndarray): + pass + + dom1 = [0, 4] + dom2 = [1, 3] + x = np.array([dom1, dom1]).view(MyNDArray) + res = pu.mapdomain(x, dom1, dom2) + assert_(isinstance(res, MyNDArray)) + + def test_mapparms(self): + # test for real values + dom1 = [0, 4] + dom2 = [1, 3] + tgt = [1, .5] + res = pu. mapparms(dom1, dom2) + assert_almost_equal(res, tgt) + + # test for complex values + dom1 = [0 - 1j, 2 + 1j] + dom2 = [-2, 2] + tgt = [-1 + 1j, 1 - 1j] + res = pu.mapparms(dom1, dom2) + assert_almost_equal(res, tgt) diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_printing.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_printing.py new file mode 100644 index 0000000000000000000000000000000000000000..6651f6cd92056f94d19f62cd818eeed642df2b2e --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_printing.py @@ -0,0 +1,552 @@ +from math import nan, inf +import pytest +from numpy._core import array, arange, printoptions +import numpy.polynomial as poly +from numpy.testing import assert_equal, assert_ + +# For testing polynomial printing with object arrays +from fractions import Fraction +from decimal import Decimal + + +class TestStrUnicodeSuperSubscripts: + + @pytest.fixture(scope='class', autouse=True) + def use_unicode(self): + poly.set_default_printstyle('unicode') + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·x + 3.0·x²"), + ([-1, 0, 3, -1], "-1.0 + 0.0·x + 3.0·x² - 1.0·x³"), + (arange(12), ("0.0 + 1.0·x + 2.0·x² + 3.0·x³ + 4.0·x⁴ + 5.0·x⁵ + " + "6.0·x⁶ + 7.0·x⁷ +\n8.0·x⁸ + 9.0·x⁹ + 10.0·x¹⁰ + " + "11.0·x¹¹")), + )) + def test_polynomial_str(self, inp, tgt): + p = poly.Polynomial(inp) + res = str(p) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·T₁(x) + 3.0·T₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·T₁(x) + 3.0·T₂(x) - 1.0·T₃(x)"), + (arange(12), ("0.0 + 1.0·T₁(x) + 2.0·T₂(x) + 3.0·T₃(x) + 4.0·T₄(x) + " + "5.0·T₅(x) +\n6.0·T₆(x) + 7.0·T₇(x) + 8.0·T₈(x) + " + "9.0·T₉(x) + 10.0·T₁₀(x) + 11.0·T₁₁(x)")), + )) + def test_chebyshev_str(self, inp, tgt): + res = str(poly.Chebyshev(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·P₁(x) + 3.0·P₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·P₁(x) + 3.0·P₂(x) - 1.0·P₃(x)"), + (arange(12), ("0.0 + 1.0·P₁(x) + 2.0·P₂(x) + 3.0·P₃(x) + 4.0·P₄(x) + " + "5.0·P₅(x) +\n6.0·P₆(x) + 7.0·P₇(x) + 8.0·P₈(x) + " + "9.0·P₉(x) + 10.0·P₁₀(x) + 11.0·P₁₁(x)")), + )) + def test_legendre_str(self, inp, tgt): + res = str(poly.Legendre(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·H₁(x) + 3.0·H₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·H₁(x) + 3.0·H₂(x) - 1.0·H₃(x)"), + (arange(12), ("0.0 + 1.0·H₁(x) + 2.0·H₂(x) + 3.0·H₃(x) + 4.0·H₄(x) + " + "5.0·H₅(x) +\n6.0·H₆(x) + 7.0·H₇(x) + 8.0·H₈(x) + " + "9.0·H₉(x) + 10.0·H₁₀(x) + 11.0·H₁₁(x)")), + )) + def test_hermite_str(self, inp, tgt): + res = str(poly.Hermite(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·He₁(x) + 3.0·He₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·He₁(x) + 3.0·He₂(x) - 1.0·He₃(x)"), + (arange(12), ("0.0 + 1.0·He₁(x) + 2.0·He₂(x) + 3.0·He₃(x) + " + "4.0·He₄(x) + 5.0·He₅(x) +\n6.0·He₆(x) + 7.0·He₇(x) + " + "8.0·He₈(x) + 9.0·He₉(x) + 10.0·He₁₀(x) +\n" + "11.0·He₁₁(x)")), + )) + def test_hermiteE_str(self, inp, tgt): + res = str(poly.HermiteE(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·L₁(x) + 3.0·L₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·L₁(x) + 3.0·L₂(x) - 1.0·L₃(x)"), + (arange(12), ("0.0 + 1.0·L₁(x) + 2.0·L₂(x) + 3.0·L₃(x) + 4.0·L₄(x) + " + "5.0·L₅(x) +\n6.0·L₆(x) + 7.0·L₇(x) + 8.0·L₈(x) + " + "9.0·L₉(x) + 10.0·L₁₀(x) + 11.0·L₁₁(x)")), + )) + def test_laguerre_str(self, inp, tgt): + res = str(poly.Laguerre(inp)) + assert_equal(res, tgt) + + def test_polynomial_str_domains(self): + res = str(poly.Polynomial([0, 1])) + tgt = '0.0 + 1.0·x' + assert_equal(res, tgt) + + res = str(poly.Polynomial([0, 1], domain=[1, 2])) + tgt = '0.0 + 1.0·(-3.0 + 2.0x)' + assert_equal(res, tgt) + +class TestStrAscii: + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 x + 3.0 x**2"), + ([-1, 0, 3, -1], "-1.0 + 0.0 x + 3.0 x**2 - 1.0 x**3"), + (arange(12), ("0.0 + 1.0 x + 2.0 x**2 + 3.0 x**3 + 4.0 x**4 + " + "5.0 x**5 + 6.0 x**6 +\n7.0 x**7 + 8.0 x**8 + " + "9.0 x**9 + 10.0 x**10 + 11.0 x**11")), + )) + def test_polynomial_str(self, inp, tgt): + res = str(poly.Polynomial(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 T_1(x) + 3.0 T_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 T_1(x) + 3.0 T_2(x) - 1.0 T_3(x)"), + (arange(12), ("0.0 + 1.0 T_1(x) + 2.0 T_2(x) + 3.0 T_3(x) + " + "4.0 T_4(x) + 5.0 T_5(x) +\n6.0 T_6(x) + 7.0 T_7(x) + " + "8.0 T_8(x) + 9.0 T_9(x) + 10.0 T_10(x) +\n" + "11.0 T_11(x)")), + )) + def test_chebyshev_str(self, inp, tgt): + res = str(poly.Chebyshev(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 P_1(x) + 3.0 P_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 P_1(x) + 3.0 P_2(x) - 1.0 P_3(x)"), + (arange(12), ("0.0 + 1.0 P_1(x) + 2.0 P_2(x) + 3.0 P_3(x) + " + "4.0 P_4(x) + 5.0 P_5(x) +\n6.0 P_6(x) + 7.0 P_7(x) + " + "8.0 P_8(x) + 9.0 P_9(x) + 10.0 P_10(x) +\n" + "11.0 P_11(x)")), + )) + def test_legendre_str(self, inp, tgt): + res = str(poly.Legendre(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 H_1(x) + 3.0 H_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 H_1(x) + 3.0 H_2(x) - 1.0 H_3(x)"), + (arange(12), ("0.0 + 1.0 H_1(x) + 2.0 H_2(x) + 3.0 H_3(x) + " + "4.0 H_4(x) + 5.0 H_5(x) +\n6.0 H_6(x) + 7.0 H_7(x) + " + "8.0 H_8(x) + 9.0 H_9(x) + 10.0 H_10(x) +\n" + "11.0 H_11(x)")), + )) + def test_hermite_str(self, inp, tgt): + res = str(poly.Hermite(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 He_1(x) + 3.0 He_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 He_1(x) + 3.0 He_2(x) - 1.0 He_3(x)"), + (arange(12), ("0.0 + 1.0 He_1(x) + 2.0 He_2(x) + 3.0 He_3(x) + " + "4.0 He_4(x) +\n5.0 He_5(x) + 6.0 He_6(x) + " + "7.0 He_7(x) + 8.0 He_8(x) + 9.0 He_9(x) +\n" + "10.0 He_10(x) + 11.0 He_11(x)")), + )) + def test_hermiteE_str(self, inp, tgt): + res = str(poly.HermiteE(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 L_1(x) + 3.0 L_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 L_1(x) + 3.0 L_2(x) - 1.0 L_3(x)"), + (arange(12), ("0.0 + 1.0 L_1(x) + 2.0 L_2(x) + 3.0 L_3(x) + " + "4.0 L_4(x) + 5.0 L_5(x) +\n6.0 L_6(x) + 7.0 L_7(x) + " + "8.0 L_8(x) + 9.0 L_9(x) + 10.0 L_10(x) +\n" + "11.0 L_11(x)")), + )) + def test_laguerre_str(self, inp, tgt): + res = str(poly.Laguerre(inp)) + assert_equal(res, tgt) + + def test_polynomial_str_domains(self): + res = str(poly.Polynomial([0, 1])) + tgt = '0.0 + 1.0 x' + assert_equal(res, tgt) + + res = str(poly.Polynomial([0, 1], domain=[1, 2])) + tgt = '0.0 + 1.0 (-3.0 + 2.0x)' + assert_equal(res, tgt) + +class TestLinebreaking: + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + def test_single_line_one_less(self): + # With 'ascii' style, len(str(p)) is default linewidth - 1 (i.e. 74) + p = poly.Polynomial([12345678, 12345678, 12345678, 12345678, 123]) + assert_equal(len(str(p)), 74) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 + 123.0 x**4' + )) + + def test_num_chars_is_linewidth(self): + # len(str(p)) == default linewidth == 75 + p = poly.Polynomial([12345678, 12345678, 12345678, 12345678, 1234]) + assert_equal(len(str(p)), 75) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 +\n1234.0 x**4' + )) + + def test_first_linebreak_multiline_one_less_than_linewidth(self): + # Multiline str where len(first_line) + len(next_term) == lw - 1 == 74 + p = poly.Polynomial( + [12345678, 12345678, 12345678, 12345678, 1, 12345678] + ) + assert_equal(len(str(p).split('\n')[0]), 74) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 + 1.0 x**4 +\n12345678.0 x**5' + )) + + def test_first_linebreak_multiline_on_linewidth(self): + # First line is one character longer than previous test + p = poly.Polynomial( + [12345678, 12345678, 12345678, 12345678.12, 1, 12345678] + ) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.12 x**3 +\n1.0 x**4 + 12345678.0 x**5' + )) + + @pytest.mark.parametrize(('lw', 'tgt'), ( + (75, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 + 40000.0 x**4 + ' + '500000.0 x**5 +\n600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 + ' + '900.0 x**9')), + (45, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 +\n40000.0 x**4 + ' + '500000.0 x**5 +\n600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 +\n' + '900.0 x**9')), + (132, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 + 40000.0 x**4 + ' + '500000.0 x**5 + 600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 + ' + '900.0 x**9')), + )) + def test_linewidth_printoption(self, lw, tgt): + p = poly.Polynomial( + [0, 10, 200, 3000, 40000, 500000, 600000, 70000, 8000, 900] + ) + with printoptions(linewidth=lw): + assert_equal(str(p), tgt) + for line in str(p).split('\n'): + assert_(len(line) < lw) + + +def test_set_default_printoptions(): + p = poly.Polynomial([1, 2, 3]) + c = poly.Chebyshev([1, 2, 3]) + poly.set_default_printstyle('ascii') + assert_equal(str(p), "1.0 + 2.0 x + 3.0 x**2") + assert_equal(str(c), "1.0 + 2.0 T_1(x) + 3.0 T_2(x)") + poly.set_default_printstyle('unicode') + assert_equal(str(p), "1.0 + 2.0·x + 3.0·x²") + assert_equal(str(c), "1.0 + 2.0·T₁(x) + 3.0·T₂(x)") + with pytest.raises(ValueError): + poly.set_default_printstyle('invalid_input') + + +def test_complex_coefficients(): + """Test both numpy and built-in complex.""" + coefs = [0+1j, 1+1j, -2+2j, 3+0j] + # numpy complex + p1 = poly.Polynomial(coefs) + # Python complex + p2 = poly.Polynomial(array(coefs, dtype=object)) + poly.set_default_printstyle('unicode') + assert_equal(str(p1), "1j + (1+1j)·x - (2-2j)·x² + (3+0j)·x³") + assert_equal(str(p2), "1j + (1+1j)·x + (-2+2j)·x² + (3+0j)·x³") + poly.set_default_printstyle('ascii') + assert_equal(str(p1), "1j + (1+1j) x - (2-2j) x**2 + (3+0j) x**3") + assert_equal(str(p2), "1j + (1+1j) x + (-2+2j) x**2 + (3+0j) x**3") + + +@pytest.mark.parametrize(('coefs', 'tgt'), ( + (array([Fraction(1, 2), Fraction(3, 4)], dtype=object), ( + "1/2 + 3/4·x" + )), + (array([1, 2, Fraction(5, 7)], dtype=object), ( + "1 + 2·x + 5/7·x²" + )), + (array([Decimal('1.00'), Decimal('2.2'), 3], dtype=object), ( + "1.00 + 2.2·x + 3·x²" + )), +)) +def test_numeric_object_coefficients(coefs, tgt): + p = poly.Polynomial(coefs) + poly.set_default_printstyle('unicode') + assert_equal(str(p), tgt) + + +@pytest.mark.parametrize(('coefs', 'tgt'), ( + (array([1, 2, 'f'], dtype=object), '1 + 2·x + f·x²'), + (array([1, 2, [3, 4]], dtype=object), '1 + 2·x + [3, 4]·x²'), +)) +def test_nonnumeric_object_coefficients(coefs, tgt): + """ + Test coef fallback for object arrays of non-numeric coefficients. + """ + p = poly.Polynomial(coefs) + poly.set_default_printstyle('unicode') + assert_equal(str(p), tgt) + + +class TestFormat: + def test_format_unicode(self): + poly.set_default_printstyle('ascii') + p = poly.Polynomial([1, 2, 0, -1]) + assert_equal(format(p, 'unicode'), "1.0 + 2.0·x + 0.0·x² - 1.0·x³") + + def test_format_ascii(self): + poly.set_default_printstyle('unicode') + p = poly.Polynomial([1, 2, 0, -1]) + assert_equal( + format(p, 'ascii'), "1.0 + 2.0 x + 0.0 x**2 - 1.0 x**3" + ) + + def test_empty_formatstr(self): + poly.set_default_printstyle('ascii') + p = poly.Polynomial([1, 2, 3]) + assert_equal(format(p), "1.0 + 2.0 x + 3.0 x**2") + assert_equal(f"{p}", "1.0 + 2.0 x + 3.0 x**2") + + def test_bad_formatstr(self): + p = poly.Polynomial([1, 2, 0, -1]) + with pytest.raises(ValueError): + format(p, '.2f') + + +@pytest.mark.parametrize(('poly', 'tgt'), ( + (poly.Polynomial, '1.0 + 2.0·z + 3.0·z²'), + (poly.Chebyshev, '1.0 + 2.0·T₁(z) + 3.0·T₂(z)'), + (poly.Hermite, '1.0 + 2.0·H₁(z) + 3.0·H₂(z)'), + (poly.HermiteE, '1.0 + 2.0·He₁(z) + 3.0·He₂(z)'), + (poly.Laguerre, '1.0 + 2.0·L₁(z) + 3.0·L₂(z)'), + (poly.Legendre, '1.0 + 2.0·P₁(z) + 3.0·P₂(z)'), +)) +def test_symbol(poly, tgt): + p = poly([1, 2, 3], symbol='z') + assert_equal(f"{p:unicode}", tgt) + + +class TestRepr: + def test_polynomial_repr(self): + res = repr(poly.Polynomial([0, 1])) + tgt = ( + "Polynomial([0., 1.], domain=[-1., 1.], window=[-1., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_chebyshev_repr(self): + res = repr(poly.Chebyshev([0, 1])) + tgt = ( + "Chebyshev([0., 1.], domain=[-1., 1.], window=[-1., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_legendre_repr(self): + res = repr(poly.Legendre([0, 1])) + tgt = ( + "Legendre([0., 1.], domain=[-1., 1.], window=[-1., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_hermite_repr(self): + res = repr(poly.Hermite([0, 1])) + tgt = ( + "Hermite([0., 1.], domain=[-1., 1.], window=[-1., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_hermiteE_repr(self): + res = repr(poly.HermiteE([0, 1])) + tgt = ( + "HermiteE([0., 1.], domain=[-1., 1.], window=[-1., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_laguerre_repr(self): + res = repr(poly.Laguerre([0, 1])) + tgt = ( + "Laguerre([0., 1.], domain=[0., 1.], window=[0., 1.], " + "symbol='x')" + ) + assert_equal(res, tgt) + + +class TestLatexRepr: + """Test the latex repr used by Jupyter""" + + @staticmethod + def as_latex(obj): + # right now we ignore the formatting of scalars in our tests, since + # it makes them too verbose. Ideally, the formatting of scalars will + # be fixed such that tests below continue to pass + obj._repr_latex_scalar = lambda x, parens=False: str(x) + try: + return obj._repr_latex_() + finally: + del obj._repr_latex_scalar + + def test_simple_polynomial(self): + # default input + p = poly.Polynomial([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,x + 3.0\,x^{2}$') + + # translated input + p = poly.Polynomial([1, 2, 3], domain=[-2, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(1.0 + x\right) + 3.0\,\left(1.0 + x\right)^{2}$') + + # scaled input + p = poly.Polynomial([1, 2, 3], domain=[-0.5, 0.5]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(2.0x\right) + 3.0\,\left(2.0x\right)^{2}$') + + # affine input + p = poly.Polynomial([1, 2, 3], domain=[-1, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(1.0 + 2.0x\right) + 3.0\,\left(1.0 + 2.0x\right)^{2}$') + + def test_basis_func(self): + p = poly.Chebyshev([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{T}_{0}(x) + 2.0\,{T}_{1}(x) + 3.0\,{T}_{2}(x)$') + # affine input - check no surplus parens are added + p = poly.Chebyshev([1, 2, 3], domain=[-1, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{T}_{0}(1.0 + 2.0x) + 2.0\,{T}_{1}(1.0 + 2.0x) + 3.0\,{T}_{2}(1.0 + 2.0x)$') + + def test_multichar_basis_func(self): + p = poly.HermiteE([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{He}_{0}(x) + 2.0\,{He}_{1}(x) + 3.0\,{He}_{2}(x)$') + + def test_symbol_basic(self): + # default input + p = poly.Polynomial([1, 2, 3], symbol='z') + assert_equal(self.as_latex(p), + r'$z \mapsto 1.0 + 2.0\,z + 3.0\,z^{2}$') + + # translated input + p = poly.Polynomial([1, 2, 3], domain=[-2, 0], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(1.0 + z\right) + 3.0\,' + r'\left(1.0 + z\right)^{2}$' + ), + ) + + # scaled input + p = poly.Polynomial([1, 2, 3], domain=[-0.5, 0.5], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(2.0z\right) + 3.0\,' + r'\left(2.0z\right)^{2}$' + ), + ) + + # affine input + p = poly.Polynomial([1, 2, 3], domain=[-1, 0], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(1.0 + 2.0z\right) + 3.0\,' + r'\left(1.0 + 2.0z\right)^{2}$' + ), + ) + + def test_numeric_object_coefficients(self): + coefs = array([Fraction(1, 2), Fraction(1)]) + p = poly.Polynomial(coefs) + assert_equal(self.as_latex(p), '$x \\mapsto 1/2 + 1\\,x$') + +SWITCH_TO_EXP = ( + '1.0 + (1.0e-01) x + (1.0e-02) x**2', + '1.2 + (1.2e-01) x + (1.2e-02) x**2', + '1.23 + 0.12 x + (1.23e-02) x**2 + (1.23e-03) x**3', + '1.235 + 0.123 x + (1.235e-02) x**2 + (1.235e-03) x**3', + '1.2346 + 0.1235 x + 0.0123 x**2 + (1.2346e-03) x**3 + (1.2346e-04) x**4', + '1.23457 + 0.12346 x + 0.01235 x**2 + (1.23457e-03) x**3 + ' + '(1.23457e-04) x**4', + '1.234568 + 0.123457 x + 0.012346 x**2 + 0.001235 x**3 + ' + '(1.234568e-04) x**4 + (1.234568e-05) x**5', + '1.2345679 + 0.1234568 x + 0.0123457 x**2 + 0.0012346 x**3 + ' + '(1.2345679e-04) x**4 + (1.2345679e-05) x**5') + +class TestPrintOptions: + """ + Test the output is properly configured via printoptions. + The exponential notation is enabled automatically when the values + are too small or too large. + """ + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + def test_str(self): + p = poly.Polynomial([1/2, 1/7, 1/7*10**8, 1/7*10**9]) + assert_equal(str(p), '0.5 + 0.14285714 x + 14285714.28571429 x**2 ' + '+ (1.42857143e+08) x**3') + + with printoptions(precision=3): + assert_equal(str(p), '0.5 + 0.143 x + 14285714.286 x**2 ' + '+ (1.429e+08) x**3') + + def test_latex(self): + p = poly.Polynomial([1/2, 1/7, 1/7*10**8, 1/7*10**9]) + assert_equal(p._repr_latex_(), + r'$x \mapsto \text{0.5} + \text{0.14285714}\,x + ' + r'\text{14285714.28571429}\,x^{2} + ' + r'\text{(1.42857143e+08)}\,x^{3}$') + + with printoptions(precision=3): + assert_equal(p._repr_latex_(), + r'$x \mapsto \text{0.5} + \text{0.143}\,x + ' + r'\text{14285714.286}\,x^{2} + \text{(1.429e+08)}\,x^{3}$') + + def test_fixed(self): + p = poly.Polynomial([1/2]) + assert_equal(str(p), '0.5') + + with printoptions(floatmode='fixed'): + assert_equal(str(p), '0.50000000') + + with printoptions(floatmode='fixed', precision=4): + assert_equal(str(p), '0.5000') + + def test_switch_to_exp(self): + for i, s in enumerate(SWITCH_TO_EXP): + with printoptions(precision=i): + p = poly.Polynomial([1.23456789*10**-i + for i in range(i//2+3)]) + assert str(p).replace('\n', ' ') == s + + def test_non_finite(self): + p = poly.Polynomial([nan, inf]) + assert str(p) == 'nan + inf x' + assert p._repr_latex_() == r'$x \mapsto \text{nan} + \text{inf}\,x$' + with printoptions(nanstr='NAN', infstr='INF'): + assert str(p) == 'NAN + INF x' + assert p._repr_latex_() == \ + r'$x \mapsto \text{NAN} + \text{INF}\,x$' diff --git a/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_symbol.py b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_symbol.py new file mode 100644 index 0000000000000000000000000000000000000000..f985533f9fe8c639f224daead98e31dc6f798cc4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/numpy/polynomial/tests/test_symbol.py @@ -0,0 +1,216 @@ +""" +Tests related to the ``symbol`` attribute of the ABCPolyBase class. +""" + +import pytest +import numpy.polynomial as poly +from numpy._core import array +from numpy.testing import assert_equal, assert_raises, assert_ + + +class TestInit: + """ + Test polynomial creation with symbol kwarg. + """ + c = [1, 2, 3] + + def test_default_symbol(self): + p = poly.Polynomial(self.c) + assert_equal(p.symbol, 'x') + + @pytest.mark.parametrize(('bad_input', 'exception'), ( + ('', ValueError), + ('3', ValueError), + (None, TypeError), + (1, TypeError), + )) + def test_symbol_bad_input(self, bad_input, exception): + with pytest.raises(exception): + p = poly.Polynomial(self.c, symbol=bad_input) + + @pytest.mark.parametrize('symbol', ( + 'x', + 'x_1', + 'A', + 'xyz', + 'β', + )) + def test_valid_symbols(self, symbol): + """ + Values for symbol that should pass input validation. + """ + p = poly.Polynomial(self.c, symbol=symbol) + assert_equal(p.symbol, symbol) + + def test_property(self): + """ + 'symbol' attribute is read only. + """ + p = poly.Polynomial(self.c, symbol='x') + with pytest.raises(AttributeError): + p.symbol = 'z' + + def test_change_symbol(self): + p = poly.Polynomial(self.c, symbol='y') + # Create new polynomial from p with different symbol + pt = poly.Polynomial(p.coef, symbol='t') + assert_equal(pt.symbol, 't') + + +class TestUnaryOperators: + p = poly.Polynomial([1, 2, 3], symbol='z') + + def test_neg(self): + n = -self.p + assert_equal(n.symbol, 'z') + + def test_scalarmul(self): + out = self.p * 10 + assert_equal(out.symbol, 'z') + + def test_rscalarmul(self): + out = 10 * self.p + assert_equal(out.symbol, 'z') + + def test_pow(self): + out = self.p ** 3 + assert_equal(out.symbol, 'z') + + +@pytest.mark.parametrize( + 'rhs', + ( + poly.Polynomial([4, 5, 6], symbol='z'), + array([4, 5, 6]), + ), +) +class TestBinaryOperatorsSameSymbol: + """ + Ensure symbol is preserved for numeric operations on polynomials with + the same symbol + """ + p = poly.Polynomial([1, 2, 3], symbol='z') + + def test_add(self, rhs): + out = self.p + rhs + assert_equal(out.symbol, 'z') + + def test_sub(self, rhs): + out = self.p - rhs + assert_equal(out.symbol, 'z') + + def test_polymul(self, rhs): + out = self.p * rhs + assert_equal(out.symbol, 'z') + + def test_divmod(self, rhs): + for out in divmod(self.p, rhs): + assert_equal(out.symbol, 'z') + + def test_radd(self, rhs): + out = rhs + self.p + assert_equal(out.symbol, 'z') + + def test_rsub(self, rhs): + out = rhs - self.p + assert_equal(out.symbol, 'z') + + def test_rmul(self, rhs): + out = rhs * self.p + assert_equal(out.symbol, 'z') + + def test_rdivmod(self, rhs): + for out in divmod(rhs, self.p): + assert_equal(out.symbol, 'z') + + +class TestBinaryOperatorsDifferentSymbol: + p = poly.Polynomial([1, 2, 3], symbol='x') + other = poly.Polynomial([4, 5, 6], symbol='y') + ops = (p.__add__, p.__sub__, p.__mul__, p.__floordiv__, p.__mod__) + + @pytest.mark.parametrize('f', ops) + def test_binops_fails(self, f): + assert_raises(ValueError, f, self.other) + + +class TestEquality: + p = poly.Polynomial([1, 2, 3], symbol='x') + + def test_eq(self): + other = poly.Polynomial([1, 2, 3], symbol='x') + assert_(self.p == other) + + def test_neq(self): + other = poly.Polynomial([1, 2, 3], symbol='y') + assert_(not self.p == other) + + +class TestExtraMethods: + """ + Test other methods for manipulating/creating polynomial objects. + """ + p = poly.Polynomial([1, 2, 3, 0], symbol='z') + + def test_copy(self): + other = self.p.copy() + assert_equal(other.symbol, 'z') + + def test_trim(self): + other = self.p.trim() + assert_equal(other.symbol, 'z') + + def test_truncate(self): + other = self.p.truncate(2) + assert_equal(other.symbol, 'z') + + @pytest.mark.parametrize('kwarg', ( + {'domain': [-10, 10]}, + {'window': [-10, 10]}, + {'kind': poly.Chebyshev}, + )) + def test_convert(self, kwarg): + other = self.p.convert(**kwarg) + assert_equal(other.symbol, 'z') + + def test_integ(self): + other = self.p.integ() + assert_equal(other.symbol, 'z') + + def test_deriv(self): + other = self.p.deriv() + assert_equal(other.symbol, 'z') + + +def test_composition(): + p = poly.Polynomial([3, 2, 1], symbol="t") + q = poly.Polynomial([5, 1, 0, -1], symbol="λ_1") + r = p(q) + assert r.symbol == "λ_1" + + +# +# Class methods that result in new polynomial class instances +# + + +def test_fit(): + x, y = (range(10),)*2 + p = poly.Polynomial.fit(x, y, deg=1, symbol='z') + assert_equal(p.symbol, 'z') + + +def test_froomroots(): + roots = [-2, 2] + p = poly.Polynomial.fromroots(roots, symbol='z') + assert_equal(p.symbol, 'z') + + +def test_identity(): + p = poly.Polynomial.identity(domain=[-1, 1], window=[5, 20], symbol='z') + assert_equal(p.symbol, 'z') + + +def test_basis(): + p = poly.Polynomial.basis(3, symbol='z') + assert_equal(p.symbol, 'z') diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0d8d05f17dde9634f260d0bdb1dfa5e4e09ac40f --- /dev/null +++ b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/METADATA @@ -0,0 +1,138 @@ +Metadata-Version: 2.2 +Name: nvidia-cusparselt-cu12 +Version: 0.7.1 +Summary: NVIDIA cuSPARSELt +Home-page: https://developer.nvidia.com/cusparselt +Author: NVIDIA Corporation +Author-email: cuda_installer@nvidia.com +License: NVIDIA Proprietary Software +Keywords: cuda,nvidia,machine learning,high-performance computing +Classifier: Topic :: Scientific/Engineering +Classifier: Environment :: GPU :: NVIDIA CUDA +Classifier: Environment :: GPU :: NVIDIA CUDA :: 12 +Description-Content-Type: text/x-rst +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: description +Dynamic: description-content-type +Dynamic: home-page +Dynamic: keywords +Dynamic: license +Dynamic: summary + +################################################################################### +cuSPARSELt: A High-Performance CUDA Library for Sparse Matrix-Matrix Multiplication +################################################################################### + +**NVIDIA cuSPARSELt** is a high-performance CUDA library dedicated to general matrix-matrix operations in which at least one operand is a sparse matrix: + +.. math:: + + D = Activation(\alpha op(A) \cdot op(B) + \beta op(C) + bias) \cdot scale + +where :math:`op(A)/op(B)` refers to in-place operations such as transpose/non-transpose, and :math:`alpha, beta, scale` are scalars. + +The *cuSPARSELt APIs* allow flexibility in the algorithm/operation selection, epilogue, and matrix characteristics, including memory layout, alignment, and data types. + +**Download:** `developer.nvidia.com/cusparselt/downloads `_ + +**Provide Feedback:** `Math-Libs-Feedback@nvidia.com `_ + +**Examples**: +`cuSPARSELt Example 1 `_, +`cuSPARSELt Example 2 `_ + +**Blog post**: + +- `Exploiting NVIDIA Ampere Structured Sparsity with cuSPARSELt `_ +- `Structured Sparsity in the NVIDIA Ampere Architecture and Applications in Search Engines `__ +- `Making the Most of Structured Sparsity in the NVIDIA Ampere Architecture `__ + +================================================================================ +Key Features +================================================================================ + +* *NVIDIA Sparse MMA tensor core* support +* Mixed-precision computation support: + + +--------------+----------------+-----------------+-------------+ + | Input A/B | Input C | Output D | Compute | + +==============+================+=================+=============+ + | `FP32` | `FP32` | `FP32` | `FP32` | + +--------------+----------------+-----------------+-------------+ + | `FP16` | `FP16` | `FP16` | `FP32` | + + + + +-------------+ + | | | | `FP16` | + +--------------+----------------+-----------------+-------------+ + | `BF16` | `BF16` | `BF16` | `FP32` | + +--------------+----------------+-----------------+-------------+ + | `INT8` | `INT8` | `INT8` | `INT32` | + + +----------------+-----------------+ + + | | `INT32` | `INT32` | | + + +----------------+-----------------+ + + | | `FP16` | `FP16` | | + + +----------------+-----------------+ + + | | `BF16` | `BF16` | | + +--------------+----------------+-----------------+-------------+ + | `E4M3` | `FP16` | `E4M3` | `FP32` | + + +----------------+-----------------+ + + | | `BF16` | `E4M3` | | + + +----------------+-----------------+ + + | | `FP16` | `FP16` | | + + +----------------+-----------------+ + + | | `BF16` | `BF16` | | + + +----------------+-----------------+ + + | | `FP32` | `FP32` | | + +--------------+----------------+-----------------+-------------+ + | `E5M2` | `FP16` | `E5M2` | `FP32` | + + +----------------+-----------------+ + + | | `BF16` | `E5M2` | | + + +----------------+-----------------+ + + | | `FP16` | `FP16` | | + + +----------------+-----------------+ + + | | `BF16` | `BF16` | | + + +----------------+-----------------+ + + | | `FP32` | `FP32` | | + +--------------+----------------+-----------------+-------------+ + +* Matrix pruning and compression functionalities +* Activation functions, bias vector, and output scaling +* Batched computation (multiple matrices in a single run) +* GEMM Split-K mode +* Auto-tuning functionality (see `cusparseLtMatmulSearch()`) +* NVTX ranging and Logging functionalities + +================================================================================ +Support +================================================================================ + +* *Supported SM Architectures*: `SM 8.0`, `SM 8.6`, `SM 8.9`, `SM 9.0`, `SM 10.0`, `SM 12.0` +* *Supported CPU architectures and operating systems*: + ++------------+--------------------+ +| OS | CPU archs | ++============+====================+ +| `Windows` | `x86_64` | ++------------+--------------------+ +| `Linux` | `x86_64`, `Arm64` | ++------------+--------------------+ + + +================================================================================ +Documentation +================================================================================ + +Please refer to https://docs.nvidia.com/cuda/cusparselt/index.html for the cuSPARSELt documentation. + +================================================================================ +Installation +================================================================================ + +The cuSPARSELt wheel can be installed as follows: + +.. code-block:: bash + + pip install nvidia-cusparselt-cuXX + +where XX is the CUDA major version (currently CUDA 12 only is supported). diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..4561b2445a29c59ddae3b0260347836c3d957839 --- /dev/null +++ b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/RECORD @@ -0,0 +1,9 @@ +nvidia/cusparselt/LICENSE.txt,sha256=6NFYiFpoG5Xsem_AbdjUpSmJ83TLE4DIpMj7J_09XV4,17948 +nvidia/cusparselt/include/cusparseLt.h,sha256=dFgNMQTtWOFwjS7nRvZcKp4VV_kbHRWNLZPBWR4RjDg,17876 +nvidia/cusparselt/lib/libcusparseLt.so.0,sha256=s9haAxHq39POyGsCkcmktXacui8kip5-gWf78_-pmGM,451980513 +nvidia_cusparselt_cu12-0.7.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +nvidia_cusparselt_cu12-0.7.1.dist-info/METADATA,sha256=smTOppUbcLUqP9L_qoj5EJ6vN5Yzqo1uCDLsyS3fb7o,6974 +nvidia_cusparselt_cu12-0.7.1.dist-info/RECORD,, +nvidia_cusparselt_cu12-0.7.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia_cusparselt_cu12-0.7.1.dist-info/WHEEL,sha256=q0csfyxC0CgL_VbODy1dAROUKOBO6j0wA7STVVTEknI,108 +nvidia_cusparselt_cu12-0.7.1.dist-info/top_level.txt,sha256=ofIC-dbK0s22jN55sgfVpbhHWT63ZdJLWeSb6K_1-BI,18 diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..b37526f046b54acc36a14784ff0ee77425c6ffa3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (75.8.0) +Root-Is-Purelib: true +Tag: py3-none-manylinux2014_x86_64 + diff --git a/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..9b2aa4bee1c5613e5c54c1332a435b1c2db4c6cf --- /dev/null +++ b/venv/lib/python3.10/site-packages/nvidia_cusparselt_cu12-0.7.1.dist-info/top_level.txt @@ -0,0 +1 @@ +nvidia/cusparselt diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..3d5b2613e8b60616b55580f66be43c8b237d3f20 --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/METADATA @@ -0,0 +1,175 @@ +Metadata-Version: 2.4 +Name: pillow +Version: 12.0.0 +Summary: Python Imaging Library (fork) +Author-email: "Jeffrey A. Clark" +License-Expression: MIT-CMU +Project-URL: Changelog, https://github.com/python-pillow/Pillow/releases +Project-URL: Documentation, https://pillow.readthedocs.io +Project-URL: Funding, https://tidelift.com/subscription/pkg/pypi-pillow?utm_source=pypi-pillow&utm_medium=pypi +Project-URL: Homepage, https://python-pillow.github.io +Project-URL: Mastodon, https://fosstodon.org/@pillow +Project-URL: Release notes, https://pillow.readthedocs.io/en/stable/releasenotes/index.html +Project-URL: Source, https://github.com/python-pillow/Pillow +Keywords: Imaging +Classifier: Development Status :: 6 - Mature +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Multimedia :: Graphics +Classifier: Topic :: Multimedia :: Graphics :: Capture :: Digital Camera +Classifier: Topic :: Multimedia :: Graphics :: Capture :: Screen Capture +Classifier: Topic :: Multimedia :: Graphics :: Graphics Conversion +Classifier: Topic :: Multimedia :: Graphics :: Viewers +Classifier: Typing :: Typed +Requires-Python: >=3.10 +Description-Content-Type: text/markdown +License-File: LICENSE +Provides-Extra: docs +Requires-Dist: furo; extra == "docs" +Requires-Dist: olefile; extra == "docs" +Requires-Dist: sphinx>=8.2; extra == "docs" +Requires-Dist: sphinx-autobuild; extra == "docs" +Requires-Dist: sphinx-copybutton; extra == "docs" +Requires-Dist: sphinx-inline-tabs; extra == "docs" +Requires-Dist: sphinxext-opengraph; extra == "docs" +Provides-Extra: fpx +Requires-Dist: olefile; extra == "fpx" +Provides-Extra: mic +Requires-Dist: olefile; extra == "mic" +Provides-Extra: test-arrow +Requires-Dist: arro3-compute; extra == "test-arrow" +Requires-Dist: arro3-core; extra == "test-arrow" +Requires-Dist: nanoarrow; extra == "test-arrow" +Requires-Dist: pyarrow; extra == "test-arrow" +Provides-Extra: tests +Requires-Dist: check-manifest; extra == "tests" +Requires-Dist: coverage>=7.4.2; extra == "tests" +Requires-Dist: defusedxml; extra == "tests" +Requires-Dist: markdown2; extra == "tests" +Requires-Dist: olefile; extra == "tests" +Requires-Dist: packaging; extra == "tests" +Requires-Dist: pyroma>=5; extra == "tests" +Requires-Dist: pytest; extra == "tests" +Requires-Dist: pytest-cov; extra == "tests" +Requires-Dist: pytest-timeout; extra == "tests" +Requires-Dist: pytest-xdist; extra == "tests" +Requires-Dist: trove-classifiers>=2024.10.12; extra == "tests" +Provides-Extra: xmp +Requires-Dist: defusedxml; extra == "xmp" +Dynamic: license-file + +

+ Pillow logo +

+ +# Pillow + +## Python Imaging Library (Fork) + +Pillow is the friendly PIL fork by [Jeffrey A. Clark and +contributors](https://github.com/python-pillow/Pillow/graphs/contributors). +PIL is the Python Imaging Library by Fredrik Lundh and contributors. +As of 2019, Pillow development is +[supported by Tidelift](https://tidelift.com/subscription/pkg/pypi-pillow?utm_source=pypi-pillow&utm_medium=readme&utm_campaign=enterprise). + + + + + + + + + + + + + + + + + + +
docs + Documentation Status +
tests + GitHub Actions build status (Lint) + GitHub Actions build status (Test Linux and macOS) + GitHub Actions build status (Test Windows) + GitHub Actions build status (Test MinGW) + GitHub Actions build status (Test Docker) + GitHub Actions build status (Wheels) + Code coverage + Fuzzing Status +
package + Zenodo + Tidelift + Newest PyPI version + Number of PyPI downloads + OpenSSF Best Practices +
social + Join the chat at https://gitter.im/python-pillow/Pillow + Follow on https://fosstodon.org/@pillow +
+ +## Overview + +The Python Imaging Library adds image processing capabilities to your Python interpreter. + +This library provides extensive file format support, an efficient internal representation, and fairly powerful image processing capabilities. + +The core image library is designed for fast access to data stored in a few basic pixel formats. It should provide a solid foundation for a general image processing tool. + +## More information + +- [Documentation](https://pillow.readthedocs.io/) + - [Installation](https://pillow.readthedocs.io/en/latest/installation/basic-installation.html) + - [Handbook](https://pillow.readthedocs.io/en/latest/handbook/index.html) +- [Contribute](https://github.com/python-pillow/Pillow/blob/main/.github/CONTRIBUTING.md) + - [Issues](https://github.com/python-pillow/Pillow/issues) + - [Pull requests](https://github.com/python-pillow/Pillow/pulls) +- [Release notes](https://pillow.readthedocs.io/en/stable/releasenotes/index.html) +- [Changelog](https://github.com/python-pillow/Pillow/releases) + - [Pre-fork](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst#pre-fork) + +## Report a vulnerability + +To report a security vulnerability, please follow the procedure described in the [Tidelift security policy](https://tidelift.com/docs/security). diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..1d75369f5c134831a6342871c4c4812be2ffa767 --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/RECORD @@ -0,0 +1,236 @@ +PIL/AvifImagePlugin.py,sha256=5IiDMvMZQXLnS3t25XJjlwgNWmeVSNaGfReWAp-V5lo,8994 +PIL/BdfFontFile.py,sha256=PhlZfIRmEfmorbhZZeSM5eebGo1Ei7fL-lR9XlfTZZA,3285 +PIL/BlpImagePlugin.py,sha256=Ub4vVKBEniiNBEgNizxScEpO1VKbC1w6iecWUU7T-Vs,16533 +PIL/BmpImagePlugin.py,sha256=-SNdj2godmaKYAc08dEng6z3mRPbYYHezjveIR5e-tU,19855 +PIL/BufrStubImagePlugin.py,sha256=JSqDhkPNPnFw0Qcz-gQJl-D_iSCFdtcLvPynshKJ4WM,1730 +PIL/ContainerIO.py,sha256=wkBqL2GDAb5fh3wrtfTGUfqioJipCl-lg2GxbjQrTZw,4604 +PIL/CurImagePlugin.py,sha256=-WEsgwQbA9rQzXB0HG0LK1V_qbuwHosPZ0T2IjfN8r0,1791 +PIL/DcxImagePlugin.py,sha256=DhqsmW7MjmnUSTGZ-Skv9hz1XeX3XoQQoAl9GWLAEEY,2145 +PIL/DdsImagePlugin.py,sha256=fjdfZK_eQtUp_-bjoRmt-5wgOT5GTmvg6aI-itch4mo,18906 +PIL/EpsImagePlugin.py,sha256=Q91Of8yr6VY12picGSU6k6HvgU9FgnQJvWrJQryiLnU,16552 +PIL/ExifTags.py,sha256=zW6kVikCosiyoCo7J7R62evD3hoxjKPchnVh8po7CZc,9931 +PIL/FitsImagePlugin.py,sha256=-oDJnAH113CK5qPvwz9lL81fkV1gla_tNfqLcq8zKgo,4644 +PIL/FliImagePlugin.py,sha256=4zxH8IXBX9DGi6dJRM6Y5NMdbA1d99x696mcGZHxHzI,4929 +PIL/FontFile.py,sha256=St7MxO5Q-oakCLWn3ZrgrtaT3wSsmAarxm8AU-G8Moc,3577 +PIL/FpxImagePlugin.py,sha256=aXfg0YdvNeJhxqh-f-f22D1NobQ8tSVCj-tpLE2PKfE,7293 +PIL/FtexImagePlugin.py,sha256=v2I5YkdfNA3iW35JzKnWry9v6Rgvr0oezGVOuArREac,3535 +PIL/GbrImagePlugin.py,sha256=ADLgy4hlBcC_3Rr2HsnIuXPoUvof1ddkhbEo5Yw8OMQ,2979 +PIL/GdImageFile.py,sha256=LP4Uxv3Y2ivGZIyOVuGJarDDVS7zK6F1Q6SNl4wyGuQ,2788 +PIL/GifImagePlugin.py,sha256=VNTEgDJRP6OIze8JVt0EXY1gMv3Xx90oECxawggCFAE,42213 +PIL/GimpGradientFile.py,sha256=gqqUkDbKVFCtBxt5VAhPS0HtLZDYFI6KWEaUhhTNNE8,3982 +PIL/GimpPaletteFile.py,sha256=hIHQ9LJ5ri0hy1e_vZYeD-n67UWdhEDlKc4vDxgaUdg,1860 +PIL/GribStubImagePlugin.py,sha256=I-_ZlKsSKANo7adUTnIx7pTUhQt-0B60DacLDOVm_3E,1759 +PIL/Hdf5StubImagePlugin.py,sha256=OuEQijGqVwTTSG4dB2vAyQzmN-NYT22tiuZHFH0Q0Sw,1741 +PIL/IcnsImagePlugin.py,sha256=dr_p68k2ECoONrw3Dqw3ISig39uXo39YY3nTfshUNHw,12405 +PIL/IcoImagePlugin.py,sha256=QCo29Toh08UX8vEcdCAaIeuidSolbPiZlCnQ4rUu2SQ,12491 +PIL/ImImagePlugin.py,sha256=wo5OL2PAcQW2MwRkJnS-N16toZzXWL95jx9FBM7l9ok,11567 +PIL/Image.py,sha256=2naP8UMkmSyU64EFc5crz4t02nV4f8ptFKfZdvbtBHQ,148085 +PIL/ImageChops.py,sha256=GEjlymcoDtA5OOeIxQVIX96BD-s6AXhb7TmSLYn2tUg,7946 +PIL/ImageCms.py,sha256=IuCm3gXKpb5Eu1kn-TB8cD9XJLZEa8fpkEjzVqAIKNk,40676 +PIL/ImageColor.py,sha256=IGA9C2umeED_EzS2Cvj6KsU0VutC9RstWIYPe8uDsVk,9441 +PIL/ImageDraw.py,sha256=FMn0AK_gxxJBYB8afOGF2FUP2KJPFvUf3UZp_KrJz7A,36287 +PIL/ImageDraw2.py,sha256=pdVMW7bVw3KwhXvRZh28Md4y-2xFfuo5fHcDnaYqVK4,7227 +PIL/ImageEnhance.py,sha256=4Elhz_lyyxLmx0GkSHrwOAmNJ2TkqVQPHejzGihZUMI,3627 +PIL/ImageFile.py,sha256=m6Se6q-5zsmnE7Bezp13q-H5F2yt5d4tO7YD06FGHx4,29600 +PIL/ImageFilter.py,sha256=MO1MBrbXDiX2IAGESdGm_0087bwmSZ_14ecAj28ojCY,18729 +PIL/ImageFont.py,sha256=2PGC3YI127GKrXYg4zP7_Tul2KCQ3c4ajU8LONVO9bc,63101 +PIL/ImageGrab.py,sha256=I9PHpsQf2VyNX4T8QL-8awFNotyAzB1mGxTt_I5FbTE,6471 +PIL/ImageMath.py,sha256=RQl6cRXGuszba4KwtbIudin_8U65shpWrajr9gTn1rw,10369 +PIL/ImageMode.py,sha256=aaZVHAiCEanOA2K1jN3DlW3NPKa8Dm5nIXTXErzyFms,2395 +PIL/ImageMorph.py,sha256=dobO2v2w7c8SjH7stFc3TP6HtU9O7JAWQz5Nu6mBxYg,8562 +PIL/ImageOps.py,sha256=bIcQFK_MtovfNSYTcOesp4So9OgsGrwt3cGsB7xlGRM,25567 +PIL/ImagePalette.py,sha256=M5tYUgadWR7mxUEByyVl7IV9QFFzAGiKKmAhCZtdG0w,9009 +PIL/ImagePath.py,sha256=5yUG5XCUil1KKTTA_8PgGhcmg-mnue-GK0FwTBlhjw4,371 +PIL/ImageQt.py,sha256=PTt5TPyngWL-Vuvx_bwnH17EOBe3tE7l4huVmvGQP5Y,6684 +PIL/ImageSequence.py,sha256=Mphgkr79scmYBgmi9ZguhDfVwHvpLSX5uZVHDZlrn0I,2253 +PIL/ImageShow.py,sha256=Ju0_Db2B4_n3yKJV9sDsF7_HAgciEdXlq6I1Eiw1YTo,10106 +PIL/ImageStat.py,sha256=FVTiYWGCciPW1QD61b7DYZlcDqR0dS6hsLjq-gcKcG4,5495 +PIL/ImageText.py,sha256=rkdTrW6pQCquXFOTu_0OoBfvCYiC9zQG__8JjGwnPYE,12103 +PIL/ImageTk.py,sha256=b5SntckGXs0ECsI2MmdJg3CSX6AtELsWh0Ohxu41u_k,8132 +PIL/ImageTransform.py,sha256=-qek7P3lzLddcXt9cWt5w_L11JGp2yY3AJtOfmJAkDc,3916 +PIL/ImageWin.py,sha256=LT05w8_vTfRrC3n9S9pM0TNbXrzZLEJHlCJil7Xv80k,8085 +PIL/ImtImagePlugin.py,sha256=SL5IrsHcblltxtX4v_HVFhYnR6haJ0AOd2NHhZKMImY,2665 +PIL/IptcImagePlugin.py,sha256=cOFy4epsqpMOWNgQ3Gj_dOrt2TPjbO0gjCvp_f1mUxk,6444 +PIL/Jpeg2KImagePlugin.py,sha256=IabyXVrNchWV9oOUU79eNjGSudq5tlvnihFimZH4VAA,13932 +PIL/JpegImagePlugin.py,sha256=ZMvTMZTxi2UHu87NXauQmhLC3tWEMxa2CgUhYcFq7yw,31318 +PIL/JpegPresets.py,sha256=lnqWHo4DLIHIulcdHp0NJ7CWexHt8T3w51kIKlLfkIA,12379 +PIL/McIdasImagePlugin.py,sha256=baOIkD-CIIeCgBFTf8kos928PKBuCUqYYa38u3WES_8,1877 +PIL/MicImagePlugin.py,sha256=aoIwkWVyr_X-dPvB6ldZOJF3a9kd_OeuEW3say5Y0QM,2564 +PIL/MpegImagePlugin.py,sha256=g7BZd93kWpFi41SG_wKFoi0yEPsioI4kj45b2F-3Vrw,2010 +PIL/MpoImagePlugin.py,sha256=S45qt7OcY7rBjYlwEk0nUmEj5IOu5z8KVLo066V1RBE,6722 +PIL/MspImagePlugin.py,sha256=oxk_MLUDvzJ4JDuOZCHkmqOPXniG42PHOyNGwe60slY,5892 +PIL/PSDraw.py,sha256=KMBGj3vXaFpblaIcA9KjFFTpdal41AQggY-UgzqoMkQ,6918 +PIL/PaletteFile.py,sha256=suDdAL6VMljXw4oEn1vhTt4DQ4vbpIHGd3A4oxOgE6s,1216 +PIL/PalmImagePlugin.py,sha256=WJ1b8I1xTSAXYDJhIpkVFCLu2LlpbiBD5d1Hr-m2l08,8748 +PIL/PcdImagePlugin.py,sha256=-gnMUqQH0R-aljsd3nZS9eBI1j75ijWD_HZfadE3RsQ,1774 +PIL/PcfFontFile.py,sha256=DqcyydQgP2vtiPFzj57KYHLuF2v-0oMTB-VkgYYHKhE,7223 +PIL/PcxImagePlugin.py,sha256=1xAq6CdH34cOsOgTPi4Wu2SKQCdQiTLVyqaMkYQZUP4,6245 +PIL/PdfImagePlugin.py,sha256=6lZLoQMVbAE-x1ESrv6PgGSyM9Ueck7e6E6ps-YQ-vI,9321 +PIL/PdfParser.py,sha256=Hr3ImLDSIKwUF6OrQ1GjlAnGi6ZpGVLWhGfKhqQ_DRM,37996 +PIL/PixarImagePlugin.py,sha256=l_4GwBd0mATnIXYJbwmmODU2vP7wewLu6BRviHCB2EI,1758 +PIL/PngImagePlugin.py,sha256=jGtbaGMrt9x0i9c500Nh1ofQ4M6w23Wlu67oTuwYEIA,51144 +PIL/PpmImagePlugin.py,sha256=vb5SP0IjQPzDRDE8jSPtcJv9K3Rh1LczAlt0Pg26i90,12391 +PIL/PsdImagePlugin.py,sha256=ImnNRG4VANs2GATXVEB5Q-yy1Jskc6XRVRtZYi2fALg,8685 +PIL/QoiImagePlugin.py,sha256=RPO63QsgHAsyPpcxh7ymeMYlnjVu5gT5ELolkvJt0vc,8572 +PIL/SgiImagePlugin.py,sha256=3Ql89s8vycNWjcxJwMw28iksV9Yj2xWoKBQ6c5DHXBg,6389 +PIL/SpiderImagePlugin.py,sha256=Bsg6pfZMctas1xYx__oL-ZZseUReZdnLy5a-aKEJhpE,10249 +PIL/SunImagePlugin.py,sha256=Hdxkhk0pxpBGxYhPJfCDLwsYcO1KjxjtplNMFYibIvk,4589 +PIL/TarIO.py,sha256=BqYUChCBb9F7Sh-uZ86iz1Dtoy2D0obNwGm65z1rdc0,1442 +PIL/TgaImagePlugin.py,sha256=2vDsFTcBUBHw1V80wpVv4tgpLDbPr6yVHi6Fvaqf0HY,6980 +PIL/TiffImagePlugin.py,sha256=cIQ48x3zmm5PFSG01wqweC8DJUhJxrX1R62c8Edw1Jg,85002 +PIL/TiffTags.py,sha256=cMmOVPxiq8Yt99J9DEQz4tXu8IZvsJFCSL3zJDGw3fM,17251 +PIL/WalImageFile.py,sha256=4o52MngMxr9dlMmCyIXu-11-QpEN6-MRcJnEfyjdc4M,5687 +PIL/WebPImagePlugin.py,sha256=h8hosK6SWJ5tAuSFFCboKTJ_dQCFthCGT9ooYq6TVCk,10054 +PIL/WmfImagePlugin.py,sha256=y1z3RPYozRQY8AOEs-iark--cv835yF9xENm7b0GNXo,5244 +PIL/XVThumbImagePlugin.py,sha256=cJSapkBasFt11O6XYXxqcyA-njxA5BD3wHhNj6VC7Fk,2115 +PIL/XbmImagePlugin.py,sha256=Fd6GVDEo73nyFICA3Z3w4LjkwoZWvhHB6rKCm5yVrho,2669 +PIL/XpmImagePlugin.py,sha256=jtUKavJCYwIAsJaJwSx8vJsx1oTbCywfDxePENmA93w,4400 +PIL/__init__.py,sha256=Q4KOEpR7S_Xsj30fvOsvR94xEpX4KUsVeUwaVP1fU80,2031 +PIL/__main__.py,sha256=Lpj4vef8mI7jA1sRCUAoVYaeePD_Uc898xF5c7XLx1A,133 +PIL/__pycache__/AvifImagePlugin.cpython-310.pyc,, +PIL/__pycache__/BdfFontFile.cpython-310.pyc,, +PIL/__pycache__/BlpImagePlugin.cpython-310.pyc,, +PIL/__pycache__/BmpImagePlugin.cpython-310.pyc,, +PIL/__pycache__/BufrStubImagePlugin.cpython-310.pyc,, +PIL/__pycache__/ContainerIO.cpython-310.pyc,, +PIL/__pycache__/CurImagePlugin.cpython-310.pyc,, +PIL/__pycache__/DcxImagePlugin.cpython-310.pyc,, +PIL/__pycache__/DdsImagePlugin.cpython-310.pyc,, +PIL/__pycache__/EpsImagePlugin.cpython-310.pyc,, +PIL/__pycache__/ExifTags.cpython-310.pyc,, +PIL/__pycache__/FitsImagePlugin.cpython-310.pyc,, +PIL/__pycache__/FliImagePlugin.cpython-310.pyc,, +PIL/__pycache__/FontFile.cpython-310.pyc,, +PIL/__pycache__/FpxImagePlugin.cpython-310.pyc,, +PIL/__pycache__/FtexImagePlugin.cpython-310.pyc,, +PIL/__pycache__/GbrImagePlugin.cpython-310.pyc,, +PIL/__pycache__/GdImageFile.cpython-310.pyc,, +PIL/__pycache__/GifImagePlugin.cpython-310.pyc,, +PIL/__pycache__/GimpGradientFile.cpython-310.pyc,, +PIL/__pycache__/GimpPaletteFile.cpython-310.pyc,, +PIL/__pycache__/GribStubImagePlugin.cpython-310.pyc,, +PIL/__pycache__/Hdf5StubImagePlugin.cpython-310.pyc,, +PIL/__pycache__/IcnsImagePlugin.cpython-310.pyc,, +PIL/__pycache__/IcoImagePlugin.cpython-310.pyc,, +PIL/__pycache__/ImImagePlugin.cpython-310.pyc,, +PIL/__pycache__/Image.cpython-310.pyc,, +PIL/__pycache__/ImageChops.cpython-310.pyc,, +PIL/__pycache__/ImageCms.cpython-310.pyc,, +PIL/__pycache__/ImageColor.cpython-310.pyc,, +PIL/__pycache__/ImageDraw.cpython-310.pyc,, +PIL/__pycache__/ImageDraw2.cpython-310.pyc,, +PIL/__pycache__/ImageEnhance.cpython-310.pyc,, +PIL/__pycache__/ImageFile.cpython-310.pyc,, +PIL/__pycache__/ImageFilter.cpython-310.pyc,, +PIL/__pycache__/ImageFont.cpython-310.pyc,, +PIL/__pycache__/ImageGrab.cpython-310.pyc,, +PIL/__pycache__/ImageMath.cpython-310.pyc,, +PIL/__pycache__/ImageMode.cpython-310.pyc,, +PIL/__pycache__/ImageMorph.cpython-310.pyc,, +PIL/__pycache__/ImageOps.cpython-310.pyc,, +PIL/__pycache__/ImagePalette.cpython-310.pyc,, +PIL/__pycache__/ImagePath.cpython-310.pyc,, +PIL/__pycache__/ImageQt.cpython-310.pyc,, +PIL/__pycache__/ImageSequence.cpython-310.pyc,, +PIL/__pycache__/ImageShow.cpython-310.pyc,, +PIL/__pycache__/ImageStat.cpython-310.pyc,, +PIL/__pycache__/ImageText.cpython-310.pyc,, +PIL/__pycache__/ImageTk.cpython-310.pyc,, +PIL/__pycache__/ImageTransform.cpython-310.pyc,, +PIL/__pycache__/ImageWin.cpython-310.pyc,, +PIL/__pycache__/ImtImagePlugin.cpython-310.pyc,, +PIL/__pycache__/IptcImagePlugin.cpython-310.pyc,, +PIL/__pycache__/Jpeg2KImagePlugin.cpython-310.pyc,, +PIL/__pycache__/JpegImagePlugin.cpython-310.pyc,, +PIL/__pycache__/JpegPresets.cpython-310.pyc,, +PIL/__pycache__/McIdasImagePlugin.cpython-310.pyc,, +PIL/__pycache__/MicImagePlugin.cpython-310.pyc,, +PIL/__pycache__/MpegImagePlugin.cpython-310.pyc,, +PIL/__pycache__/MpoImagePlugin.cpython-310.pyc,, +PIL/__pycache__/MspImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PSDraw.cpython-310.pyc,, +PIL/__pycache__/PaletteFile.cpython-310.pyc,, +PIL/__pycache__/PalmImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PcdImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PcfFontFile.cpython-310.pyc,, +PIL/__pycache__/PcxImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PdfImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PdfParser.cpython-310.pyc,, +PIL/__pycache__/PixarImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PngImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PpmImagePlugin.cpython-310.pyc,, +PIL/__pycache__/PsdImagePlugin.cpython-310.pyc,, +PIL/__pycache__/QoiImagePlugin.cpython-310.pyc,, +PIL/__pycache__/SgiImagePlugin.cpython-310.pyc,, +PIL/__pycache__/SpiderImagePlugin.cpython-310.pyc,, +PIL/__pycache__/SunImagePlugin.cpython-310.pyc,, +PIL/__pycache__/TarIO.cpython-310.pyc,, +PIL/__pycache__/TgaImagePlugin.cpython-310.pyc,, +PIL/__pycache__/TiffImagePlugin.cpython-310.pyc,, +PIL/__pycache__/TiffTags.cpython-310.pyc,, +PIL/__pycache__/WalImageFile.cpython-310.pyc,, +PIL/__pycache__/WebPImagePlugin.cpython-310.pyc,, +PIL/__pycache__/WmfImagePlugin.cpython-310.pyc,, +PIL/__pycache__/XVThumbImagePlugin.cpython-310.pyc,, +PIL/__pycache__/XbmImagePlugin.cpython-310.pyc,, +PIL/__pycache__/XpmImagePlugin.cpython-310.pyc,, +PIL/__pycache__/__init__.cpython-310.pyc,, +PIL/__pycache__/__main__.cpython-310.pyc,, +PIL/__pycache__/_binary.cpython-310.pyc,, +PIL/__pycache__/_deprecate.cpython-310.pyc,, +PIL/__pycache__/_tkinter_finder.cpython-310.pyc,, +PIL/__pycache__/_typing.cpython-310.pyc,, +PIL/__pycache__/_util.cpython-310.pyc,, +PIL/__pycache__/_version.cpython-310.pyc,, +PIL/__pycache__/features.cpython-310.pyc,, +PIL/__pycache__/report.cpython-310.pyc,, +PIL/_avif.cpython-310-x86_64-linux-gnu.so,sha256=hPbJfk_-phRx-TB1I918qbJg6aI7_m-lV_pp5zkwuQA,87889 +PIL/_avif.pyi,sha256=3fBxcSppJr6EOEcUojvflG3Eegg7lv2Qp0dNQQILrP4,63 +PIL/_binary.py,sha256=pcM6AL04GxgmGeLfcH1V1BZHENwIrQH0uxhJ7r0HIL0,2550 +PIL/_deprecate.py,sha256=2t747uUfRLL0rYxIEvykqyDaB_1b0eDSnwgpTh5O5fI,1970 +PIL/_imaging.cpython-310-x86_64-linux-gnu.so,sha256=HU0KrYAaFpYYHz4D6fDK3WEX0YXfefCp1MZpCjaVZDk,3320681 +PIL/_imaging.pyi,sha256=StMbXUZS32AegATP1sUHfs5P05A3TD_BiQKsDHQBW40,868 +PIL/_imagingcms.cpython-310-x86_64-linux-gnu.so,sha256=P8zeod3ZI31LJXSX9QPh1N_P0OZHzQV_ynbLbejNcUc,157841 +PIL/_imagingcms.pyi,sha256=ZZ8iIoi6EHWLvgAdfm1hPD5CQmxi75LiJl5x8yGxYoU,4433 +PIL/_imagingft.cpython-310-x86_64-linux-gnu.so,sha256=RYNDE8kjaeyAKOnzgY8_894-ogteJM2Z2OEZu16Vrq8,314777 +PIL/_imagingft.pyi,sha256=cYySzvcKBCiHPBsvttMie9AdfUcEsqZR-3256YQtz2Q,1833 +PIL/_imagingmath.cpython-310-x86_64-linux-gnu.so,sha256=tKr1VDBoDhJz02NT1FfUNcalxbrTjpKFKr97CTAO4Cc,162480 +PIL/_imagingmath.pyi,sha256=3fBxcSppJr6EOEcUojvflG3Eegg7lv2Qp0dNQQILrP4,63 +PIL/_imagingmorph.cpython-310-x86_64-linux-gnu.so,sha256=jIU7JR_UbZj4tonWFHeuyL-mLMg84KptN-G9wo0dSSA,36480 +PIL/_imagingmorph.pyi,sha256=3fBxcSppJr6EOEcUojvflG3Eegg7lv2Qp0dNQQILrP4,63 +PIL/_imagingtk.cpython-310-x86_64-linux-gnu.so,sha256=M01LHJMcp7W9c2_L5ARqfo5asWhxT8UmiisLeXwVovM,45920 +PIL/_imagingtk.pyi,sha256=3fBxcSppJr6EOEcUojvflG3Eegg7lv2Qp0dNQQILrP4,63 +PIL/_tkinter_finder.py,sha256=GIZ4stmFhUosmHKSrdxcjStiocDNfyJn7RBie2SWxU0,538 +PIL/_typing.py,sha256=2z33ZUp9aQnkSqXzNR3Zn7l04d2W-oAj1OiZhiyFF68,919 +PIL/_util.py,sha256=fxhWdrLARyc2PsMgN3m9_U1dY3oUKbV7mkoHcXgoeeA,684 +PIL/_version.py,sha256=OYLJ24lJylQ6NFpzq3LMRSnDyAJmp1pnFQ_c3PMZR44,87 +PIL/_webp.cpython-310-x86_64-linux-gnu.so,sha256=uFrT8o5JP3sCsciW1gH9Jb-LG_hftx78tHk9QzyXSrw,108849 +PIL/_webp.pyi,sha256=3fBxcSppJr6EOEcUojvflG3Eegg7lv2Qp0dNQQILrP4,63 +PIL/features.py,sha256=FPkEhjtBaRSqpkgHNYduwxiFtycu4NjZKwEMWxtemPU,10775 +PIL/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +PIL/report.py,sha256=4JY6-IU7sH1RKuRbOvy1fUt0dAoi79FX4tYJN3p1DT0,100 +pillow-12.0.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pillow-12.0.0.dist-info/METADATA,sha256=rWIEUr-lPL-ilNfae2bataCGpU7eDOF4vXTjgOvENMg,8808 +pillow-12.0.0.dist-info/RECORD,, +pillow-12.0.0.dist-info/WHEEL,sha256=Obtqci3x5vy5ZivY2BiOH9GHO8-xQ0d4HFhTOtXMNhw,152 +pillow-12.0.0.dist-info/licenses/LICENSE,sha256=MBeL96_5-NyCr-01CGzTeKkGTnf8tDgEhfOLXaM3cFI,68061 +pillow-12.0.0.dist-info/top_level.txt,sha256=riZqrk-hyZqh5f1Z0Zwii3dKfxEsByhu9cU9IODF-NY,4 +pillow-12.0.0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1 +pillow.libs/libXau-154567c4.so.6.0.0,sha256=BUhNJL94y47QMWnxywZyBNgpy3ryHeiCBADSnRFeQyA,22081 +pillow.libs/libavif-01e67780.so.16.3.0,sha256=xCXlA8_rggzsCA_EsWKFzgqPBmT1ihpZCo3iMyvD1P0,5142057 +pillow.libs/libbrotlicommon-c55a5f7a.so.1.1.0,sha256=HaLbMm3YehX759wgF7ZU0kVwhdgX4ukfvQNKytoarw8,144425 +pillow.libs/libbrotlidec-2ced2f3a.so.1.1.0,sha256=BOwekVTiRipkYusnBXmzGGhiPsBwBI6DDWPLkvxAbRE,62337 +pillow.libs/libfreetype-5bb46249.so.6.20.4,sha256=Tg7wTbDEPfL0YvLbadZ40tNwy4fKxjhuxuCsaDHzCrw,1463609 +pillow.libs/libharfbuzz-525aa570.so.0.61210.0,sha256=HyXdr1Fa9eEK7Q3VvhhzDw8DA-kp-BjzV6Ql5oX-lbc,941033 +pillow.libs/libjpeg-a41b0190.so.62.4.0,sha256=nMxHb3xGeNHn_aC68P7_-TP__RdUvJooSl_7JEFa1mU,836273 +pillow.libs/liblcms2-cc10e42f.so.2.0.17,sha256=5JMjEDVKMwxcinhbQl6qhRLaezAiQFYEPSz-KultHe0,519073 +pillow.libs/liblzma-64b7ab39.so.5.8.1,sha256=hN2B2RPEM6wOgvER_g43fNjbNQ_SsrenX2wlAfHW-nA,266369 +pillow.libs/libopenjp2-94e588ba.so.2.5.4,sha256=5ye2mwzPYSo447VzscBGO_ZWXyBgFEPFRY45hlXLIw0,585849 +pillow.libs/libpng16-00127801.so.16.50.0,sha256=KIB4XEQDJkGL2nMIrURzNgiSxMwl87FIyB0em53ffx8,278001 +pillow.libs/libsharpyuv-95d8a097.so.0.1.2,sha256=EtR2hzr_XVKswxcXrVFfDixKeUd1TMu56F0oHwH3Els,46113 +pillow.libs/libtiff-295fd75c.so.6.2.0,sha256=fkRD5oAmt0kNw9iL_CfKz0aqsJ3AgnDyVs1U0CRkEZU,754729 +pillow.libs/libwebp-d8b9687f.so.7.2.0,sha256=k-gbdtoXnzmBItYhvYxdhV11ff08TsQgWSMUUHmMzPs,731209 +pillow.libs/libwebpdemux-747f2b49.so.2.0.17,sha256=jsiJz7rjNfyn9TOqFT2OPydmTYax0iNBCCBQg9st9vw,30217 +pillow.libs/libwebpmux-7f11e5ce.so.3.1.2,sha256=Pes9BQ-MFyCzlQH9n07pWNtt7O0afgVlBsMJ-6kuU_o,58617 +pillow.libs/libxcb-64009ff3.so.1.1.0,sha256=t0N-0WuuesRJgEn9FOENG9HD59FdDl6rHS6tQqg6SdE,251425 +pillow.libs/libzstd-761a17b6.so.1.5.7,sha256=jKEGQObGqZaFbxpkGZX8jzstJVcHhJPh4SShnGgZjB8,1800497 diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..12e67ca29ff180ca4133d9ec4769226edb266ce4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: false +Tag: cp310-cp310-manylinux_2_27_x86_64 +Tag: cp310-cp310-manylinux_2_28_x86_64 + diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a5ad83b9b79dbb611fd7f68dce93043625c305be --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/licenses/LICENSE @@ -0,0 +1,1523 @@ +The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh and contributors + +Pillow is the friendly PIL fork. It is + + Copyright © 2010 by Jeffrey A. Clark and contributors + +Like PIL, Pillow is licensed under the open source MIT-CMU License: + +By obtaining, using, and/or copying this software and/or its associated +documentation, you agree that you have read, understood, and will comply +with the following terms and conditions: + +Permission to use, copy, modify and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appears in all copies, and that +both that copyright notice and this permission notice appear in supporting +documentation, and that the name of Secret Labs AB or the author not be +used in advertising or publicity pertaining to distribution of the software +without specific, written prior permission. + +SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS +SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. +IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, +INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE +OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. + + +---- + +AOM + +Copyright (c) 2016, Alliance for Open Media. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. + + +---- + +BROTLI + +Copyright (c) 2009, 2010, 2013-2016 by the Brotli Authors. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. + + +---- + +BZIP2 + + +-------------------------------------------------------------------------- + +This program, "bzip2", the associated library "libbzip2", and all +documentation, are copyright (C) 1996-2019 Julian R Seward. All +rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. The origin of this software must not be misrepresented; you must + not claim that you wrote the original software. If you use this + software in a product, an acknowledgment in the product + documentation would be appreciated but is not required. + +3. Altered source versions must be plainly marked as such, and must + not be misrepresented as being the original software. + +4. The name of the author may not be used to endorse or promote + products derived from this software without specific prior written + permission. + +THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS +OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +Julian Seward, jseward@acm.org +bzip2/libbzip2 version 1.0.8 of 13 July 2019 + +-------------------------------------------------------------------------- + + +---- + +DAV1D + +Copyright © 2018-2019, VideoLAN and dav1d authors +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +---- + +FREETYPE2 + +The FreeType 2 font engine is copyrighted work and cannot be used +legally without a software license. In order to make this project +usable to a vast majority of developers, we distribute it under two +mutually exclusive open-source licenses. + +This means that *you* must choose *one* of the two licenses described +below, then obey all its terms and conditions when using FreeType 2 in +any of your projects or products. + + - The FreeType License, found in the file `docs/FTL.TXT`, which is + similar to the original BSD license *with* an advertising clause + that forces you to explicitly cite the FreeType project in your + product's documentation. All details are in the license file. + This license is suited to products which don't use the GNU General + Public License. + + Note that this license is compatible to the GNU General Public + License version 3, but not version 2. + + - The GNU General Public License version 2, found in + `docs/GPLv2.TXT` (any later version can be used also), for + programs which already use the GPL. Note that the FTL is + incompatible with GPLv2 due to its advertisement clause. + +The contributed BDF and PCF drivers come with a license similar to +that of the X Window System. It is compatible to the above two +licenses (see files `src/bdf/README` and `src/pcf/README`). The same +holds for the source code files `src/base/fthash.c` and +`include/freetype/internal/fthash.h`; they were part of the BDF driver +in earlier FreeType versions. + +The gzip module uses the zlib license (see `src/gzip/zlib.h`) which +too is compatible to the above two licenses. + +The files `src/autofit/ft-hb.c` and `src/autofit/ft-hb.h` contain code +taken almost verbatim from the HarfBuzz file `hb-ft.cc`, which uses +the 'Old MIT' license, compatible to the above two licenses. + +The MD5 checksum support (only used for debugging in development +builds) is in the public domain. + +-------------------------------------------------------------------------- + + The FreeType Project LICENSE + ---------------------------- + + 2006-Jan-27 + + Copyright 1996-2002, 2006 by + David Turner, Robert Wilhelm, and Werner Lemberg + + + +Introduction +============ + + The FreeType Project is distributed in several archive packages; + some of them may contain, in addition to the FreeType font engine, + various tools and contributions which rely on, or relate to, the + FreeType Project. + + This license applies to all files found in such packages, and + which do not fall under their own explicit license. The license + affects thus the FreeType font engine, the test programs, + documentation and makefiles, at the very least. + + This license was inspired by the BSD, Artistic, and IJG + (Independent JPEG Group) licenses, which all encourage inclusion + and use of free software in commercial and freeware products + alike. As a consequence, its main points are that: + + o We don't promise that this software works. However, we will be + interested in any kind of bug reports. (`as is' distribution) + + o You can use this software for whatever you want, in parts or + full form, without having to pay us. (`royalty-free' usage) + + o You may not pretend that you wrote this software. If you use + it, or only parts of it, in a program, you must acknowledge + somewhere in your documentation that you have used the + FreeType code. (`credits') + + We specifically permit and encourage the inclusion of this + software, with or without modifications, in commercial products. + We disclaim all warranties covering The FreeType Project and + assume no liability related to The FreeType Project. + + + Finally, many people asked us for a preferred form for a + credit/disclaimer to use in compliance with this license. We thus + encourage you to use the following text: + + """ + Portions of this software are copyright © The FreeType + Project (www.freetype.org). All rights reserved. + """ + + Please replace with the value from the FreeType version you + actually use. + + +Legal Terms +=========== + +0. Definitions +-------------- + + Throughout this license, the terms `package', `FreeType Project', + and `FreeType archive' refer to the set of files originally + distributed by the authors (David Turner, Robert Wilhelm, and + Werner Lemberg) as the `FreeType Project', be they named as alpha, + beta or final release. + + `You' refers to the licensee, or person using the project, where + `using' is a generic term including compiling the project's source + code as well as linking it to form a `program' or `executable'. + This program is referred to as `a program using the FreeType + engine'. + + This license applies to all files distributed in the original + FreeType Project, including all source code, binaries and + documentation, unless otherwise stated in the file in its + original, unmodified form as distributed in the original archive. + If you are unsure whether or not a particular file is covered by + this license, you must contact us to verify this. + + The FreeType Project is copyright (C) 1996-2000 by David Turner, + Robert Wilhelm, and Werner Lemberg. All rights reserved except as + specified below. + +1. No Warranty +-------------- + + THE FREETYPE PROJECT IS PROVIDED `AS IS' WITHOUT WARRANTY OF ANY + KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + PURPOSE. IN NO EVENT WILL ANY OF THE AUTHORS OR COPYRIGHT HOLDERS + BE LIABLE FOR ANY DAMAGES CAUSED BY THE USE OR THE INABILITY TO + USE, OF THE FREETYPE PROJECT. + +2. Redistribution +----------------- + + This license grants a worldwide, royalty-free, perpetual and + irrevocable right and license to use, execute, perform, compile, + display, copy, create derivative works of, distribute and + sublicense the FreeType Project (in both source and object code + forms) and derivative works thereof for any purpose; and to + authorize others to exercise some or all of the rights granted + herein, subject to the following conditions: + + o Redistribution of source code must retain this license file + (`FTL.TXT') unaltered; any additions, deletions or changes to + the original files must be clearly indicated in accompanying + documentation. The copyright notices of the unaltered, + original files must be preserved in all copies of source + files. + + o Redistribution in binary form must provide a disclaimer that + states that the software is based in part of the work of the + FreeType Team, in the distribution documentation. We also + encourage you to put an URL to the FreeType web page in your + documentation, though this isn't mandatory. + + These conditions apply to any software derived from or based on + the FreeType Project, not just the unmodified files. If you use + our work, you must acknowledge us. However, no fee need be paid + to us. + +3. Advertising +-------------- + + Neither the FreeType authors and contributors nor you shall use + the name of the other for commercial, advertising, or promotional + purposes without specific prior written permission. + + We suggest, but do not require, that you use one or more of the + following phrases to refer to this software in your documentation + or advertising materials: `FreeType Project', `FreeType Engine', + `FreeType library', or `FreeType Distribution'. + + As you have not signed this license, you are not required to + accept it. However, as the FreeType Project is copyrighted + material, only this license, or another one contracted with the + authors, grants you the right to use, distribute, and modify it. + Therefore, by using, distributing, or modifying the FreeType + Project, you indicate that you understand and accept all the terms + of this license. + +4. Contacts +----------- + + There are two mailing lists related to FreeType: + + o freetype@nongnu.org + + Discusses general use and applications of FreeType, as well as + future and wanted additions to the library and distribution. + If you are looking for support, start in this list if you + haven't found anything to help you in the documentation. + + o freetype-devel@nongnu.org + + Discusses bugs, as well as engine internals, design issues, + specific licenses, porting, etc. + + Our home page can be found at + + https://www.freetype.org + + +--- end of FTL.TXT --- + +The following license details are part of `src/bdf/README`: + +``` +License +******* + +Copyright (C) 2001-2002 by Francesco Zappa Nardelli + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +*** Portions of the driver (that is, bdflib.c and bdf.h): + +Copyright 2000 Computing Research Labs, New Mexico State University +Copyright 2001-2002, 2011 Francesco Zappa Nardelli + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the "Software"), +to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, +and/or sell copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +THE COMPUTING RESEARCH LAB OR NEW MEXICO STATE UNIVERSITY BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT +OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR +THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + +Credits +******* + +This driver is based on excellent Mark Leisher's bdf library. If you +find something good in this driver you should probably thank him, not +me. +``` + +The following license details are part of `src/pcf/README`: + +``` +License +******* + +Copyright (C) 2000 by Francesco Zappa Nardelli + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + +Credits +******* + +Keith Packard wrote the pcf driver found in XFree86. His work is at +the same time the specification and the sample implementation of the +PCF format. Undoubtedly, this driver is inspired from his work. +``` + + +---- + +HARFBUZZ + +HarfBuzz is licensed under the so-called "Old MIT" license. Details follow. +For parts of HarfBuzz that are licensed under different licenses see individual +files names COPYING in subdirectories where applicable. + +Copyright © 2010-2022 Google, Inc. +Copyright © 2015-2020 Ebrahim Byagowi +Copyright © 2019,2020 Facebook, Inc. +Copyright © 2012,2015 Mozilla Foundation +Copyright © 2011 Codethink Limited +Copyright © 2008,2010 Nokia Corporation and/or its subsidiary(-ies) +Copyright © 2009 Keith Stribley +Copyright © 2011 Martin Hosken and SIL International +Copyright © 2007 Chris Wilson +Copyright © 2005,2006,2020,2021,2022,2023 Behdad Esfahbod +Copyright © 2004,2007,2008,2009,2010,2013,2021,2022,2023 Red Hat, Inc. +Copyright © 1998-2005 David Turner and Werner Lemberg +Copyright © 2016 Igalia S.L. +Copyright © 2022 Matthias Clasen +Copyright © 2018,2021 Khaled Hosny +Copyright © 2018,2019,2020 Adobe, Inc +Copyright © 2013-2015 Alexei Podtelezhnikov + +For full copyright notices consult the individual files in the package. + + +Permission is hereby granted, without written agreement and without +license or royalty fees, to use, copy, modify, and distribute this +software and its documentation for any purpose, provided that the +above copyright notice and the following two paragraphs appear in +all copies of this software. + +IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE TO ANY PARTY FOR +DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES +ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN +IF THE COPYRIGHT HOLDER HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH +DAMAGE. + +THE COPYRIGHT HOLDER SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, +BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS FOR A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS +ON AN "AS IS" BASIS, AND THE COPYRIGHT HOLDER HAS NO OBLIGATION TO +PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS. + + +---- + +LCMS2 + +Little CMS +Copyright (c) 1998-2020 Marti Maria Saguer + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + +---- + +LIBAVIF + +Copyright 2019 Joe Drago. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +------------------------------------------------------------------------------ + +Files: src/obu.c + +Copyright © 2018-2019, VideoLAN and dav1d authors +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +------------------------------------------------------------------------------ + +Files: third_party/iccjpeg/* + +In plain English: + +1. We don't promise that this software works. (But if you find any bugs, + please let us know!) +2. You can use this software for whatever you want. You don't have to pay us. +3. You may not pretend that you wrote this software. If you use it in a + program, you must acknowledge somewhere in your documentation that + you've used the IJG code. + +In legalese: + +The authors make NO WARRANTY or representation, either express or implied, +with respect to this software, its quality, accuracy, merchantability, or +fitness for a particular purpose. This software is provided "AS IS", and you, +its user, assume the entire risk as to its quality and accuracy. + +This software is copyright (C) 1991-2013, Thomas G. Lane, Guido Vollbeding. +All Rights Reserved except as specified below. + +Permission is hereby granted to use, copy, modify, and distribute this +software (or portions thereof) for any purpose, without fee, subject to these +conditions: +(1) If any part of the source code for this software is distributed, then this +README file must be included, with this copyright and no-warranty notice +unaltered; and any additions, deletions, or changes to the original files +must be clearly indicated in accompanying documentation. +(2) If only executable code is distributed, then the accompanying +documentation must state that "this software is based in part on the work of +the Independent JPEG Group". +(3) Permission for use of this software is granted only if the user accepts +full responsibility for any undesirable consequences; the authors accept +NO LIABILITY for damages of any kind. + +These conditions apply to any software derived from or based on the IJG code, +not just to the unmodified library. If you use our work, you ought to +acknowledge us. + +Permission is NOT granted for the use of any IJG author's name or company name +in advertising or publicity relating to this software or products derived from +it. This software may be referred to only as "the Independent JPEG Group's +software". + +We specifically permit and encourage the use of this software as the basis of +commercial products, provided that all warranty or liability claims are +assumed by the product vendor. + + +The Unix configuration script "configure" was produced with GNU Autoconf. +It is copyright by the Free Software Foundation but is freely distributable. +The same holds for its supporting scripts (config.guess, config.sub, +ltmain.sh). Another support script, install-sh, is copyright by X Consortium +but is also freely distributable. + +The IJG distribution formerly included code to read and write GIF files. +To avoid entanglement with the Unisys LZW patent, GIF reading support has +been removed altogether, and the GIF writer has been simplified to produce +"uncompressed GIFs". This technique does not use the LZW algorithm; the +resulting GIF files are larger than usual, but are readable by all standard +GIF decoders. + +We are required to state that + "The Graphics Interchange Format(c) is the Copyright property of + CompuServe Incorporated. GIF(sm) is a Service Mark property of + CompuServe Incorporated." + +------------------------------------------------------------------------------ + +Files: contrib/gdk-pixbuf/* + +Copyright 2020 Emmanuel Gil Peyrot. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this +list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, +this list of conditions and the following disclaimer in the documentation +and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +------------------------------------------------------------------------------ + +Files: android_jni/gradlew* + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. + +------------------------------------------------------------------------------ + +Files: third_party/libyuv/* + +Copyright 2011 The LibYuv Project Authors. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + + * Neither the name of Google nor the names of its contributors may + be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +---- + +LIBJPEG + +1. We don't promise that this software works. (But if you find any bugs, + please let us know!) +2. You can use this software for whatever you want. You don't have to pay us. +3. You may not pretend that you wrote this software. If you use it in a + program, you must acknowledge somewhere in your documentation that + you've used the IJG code. + +In legalese: + +The authors make NO WARRANTY or representation, either express or implied, +with respect to this software, its quality, accuracy, merchantability, or +fitness for a particular purpose. This software is provided "AS IS", and you, +its user, assume the entire risk as to its quality and accuracy. + +This software is copyright (C) 1991-2020, Thomas G. Lane, Guido Vollbeding. +All Rights Reserved except as specified below. + +Permission is hereby granted to use, copy, modify, and distribute this +software (or portions thereof) for any purpose, without fee, subject to these +conditions: +(1) If any part of the source code for this software is distributed, then this +README file must be included, with this copyright and no-warranty notice +unaltered; and any additions, deletions, or changes to the original files +must be clearly indicated in accompanying documentation. +(2) If only executable code is distributed, then the accompanying +documentation must state that "this software is based in part on the work of +the Independent JPEG Group". +(3) Permission for use of this software is granted only if the user accepts +full responsibility for any undesirable consequences; the authors accept +NO LIABILITY for damages of any kind. + +These conditions apply to any software derived from or based on the IJG code, +not just to the unmodified library. If you use our work, you ought to +acknowledge us. + +Permission is NOT granted for the use of any IJG author's name or company name +in advertising or publicity relating to this software or products derived from +it. This software may be referred to only as "the Independent JPEG Group's +software". + +We specifically permit and encourage the use of this software as the basis of +commercial products, provided that all warranty or liability claims are +assumed by the product vendor. + + +---- + +LIBLZMA + +XZ Utils Licensing +================== + + Different licenses apply to different files in this package. Here + is a rough summary of which licenses apply to which parts of this + package (but check the individual files to be sure!): + + - liblzma is in the public domain. + + - xz, xzdec, and lzmadec command line tools are in the public + domain unless GNU getopt_long had to be compiled and linked + in from the lib directory. The getopt_long code is under + GNU LGPLv2.1+. + + - The scripts to grep, diff, and view compressed files have been + adapted from gzip. These scripts and their documentation are + under GNU GPLv2+. + + - All the documentation in the doc directory and most of the + XZ Utils specific documentation files in other directories + are in the public domain. + + - Translated messages are in the public domain. + + - The build system contains public domain files, and files that + are under GNU GPLv2+ or GNU GPLv3+. None of these files end up + in the binaries being built. + + - Test files and test code in the tests directory, and debugging + utilities in the debug directory are in the public domain. + + - The extra directory may contain public domain files, and files + that are under various free software licenses. + + You can do whatever you want with the files that have been put into + the public domain. If you find public domain legally problematic, + take the previous sentence as a license grant. If you still find + the lack of copyright legally problematic, you have too many + lawyers. + + As usual, this software is provided "as is", without any warranty. + + If you copy significant amounts of public domain code from XZ Utils + into your project, acknowledging this somewhere in your software is + polite (especially if it is proprietary, non-free software), but + naturally it is not legally required. Here is an example of a good + notice to put into "about box" or into documentation: + + This software includes code from XZ Utils . + + The following license texts are included in the following files: + - COPYING.LGPLv2.1: GNU Lesser General Public License version 2.1 + - COPYING.GPLv2: GNU General Public License version 2 + - COPYING.GPLv3: GNU General Public License version 3 + + Note that the toolchain (compiler, linker etc.) may add some code + pieces that are copyrighted. Thus, it is possible that e.g. liblzma + binary wouldn't actually be in the public domain in its entirety + even though it contains no copyrighted code from the XZ Utils source + package. + + If you have questions, don't hesitate to ask the author(s) for more + information. + + +---- + +LIBPNG + +COPYRIGHT NOTICE, DISCLAIMER, and LICENSE +========================================= + +PNG Reference Library License version 2 +--------------------------------------- + + * Copyright (c) 1995-2022 The PNG Reference Library Authors. + * Copyright (c) 2018-2022 Cosmin Truta. + * Copyright (c) 2000-2002, 2004, 2006-2018 Glenn Randers-Pehrson. + * Copyright (c) 1996-1997 Andreas Dilger. + * Copyright (c) 1995-1996 Guy Eric Schalnat, Group 42, Inc. + +The software is supplied "as is", without warranty of any kind, +express or implied, including, without limitation, the warranties +of merchantability, fitness for a particular purpose, title, and +non-infringement. In no event shall the Copyright owners, or +anyone distributing the software, be liable for any damages or +other liability, whether in contract, tort or otherwise, arising +from, out of, or in connection with the software, or the use or +other dealings in the software, even if advised of the possibility +of such damage. + +Permission is hereby granted to use, copy, modify, and distribute +this software, or portions hereof, for any purpose, without fee, +subject to the following restrictions: + + 1. The origin of this software must not be misrepresented; you + must not claim that you wrote the original software. If you + use this software in a product, an acknowledgment in the product + documentation would be appreciated, but is not required. + + 2. Altered source versions must be plainly marked as such, and must + not be misrepresented as being the original software. + + 3. This Copyright notice may not be removed or altered from any + source or altered source distribution. + + +PNG Reference Library License version 1 (for libpng 0.5 through 1.6.35) +----------------------------------------------------------------------- + +libpng versions 1.0.7, July 1, 2000, through 1.6.35, July 15, 2018 are +Copyright (c) 2000-2002, 2004, 2006-2018 Glenn Randers-Pehrson, are +derived from libpng-1.0.6, and are distributed according to the same +disclaimer and license as libpng-1.0.6 with the following individuals +added to the list of Contributing Authors: + + Simon-Pierre Cadieux + Eric S. Raymond + Mans Rullgard + Cosmin Truta + Gilles Vollant + James Yu + Mandar Sahastrabuddhe + Google Inc. + Vadim Barkov + +and with the following additions to the disclaimer: + + There is no warranty against interference with your enjoyment of + the library or against infringement. There is no warranty that our + efforts or the library will fulfill any of your particular purposes + or needs. This library is provided with all faults, and the entire + risk of satisfactory quality, performance, accuracy, and effort is + with the user. + +Some files in the "contrib" directory and some configure-generated +files that are distributed with libpng have other copyright owners, and +are released under other open source licenses. + +libpng versions 0.97, January 1998, through 1.0.6, March 20, 2000, are +Copyright (c) 1998-2000 Glenn Randers-Pehrson, are derived from +libpng-0.96, and are distributed according to the same disclaimer and +license as libpng-0.96, with the following individuals added to the +list of Contributing Authors: + + Tom Lane + Glenn Randers-Pehrson + Willem van Schaik + +libpng versions 0.89, June 1996, through 0.96, May 1997, are +Copyright (c) 1996-1997 Andreas Dilger, are derived from libpng-0.88, +and are distributed according to the same disclaimer and license as +libpng-0.88, with the following individuals added to the list of +Contributing Authors: + + John Bowler + Kevin Bracey + Sam Bushell + Magnus Holmgren + Greg Roelofs + Tom Tanner + +Some files in the "scripts" directory have other copyright owners, +but are released under this license. + +libpng versions 0.5, May 1995, through 0.88, January 1996, are +Copyright (c) 1995-1996 Guy Eric Schalnat, Group 42, Inc. + +For the purposes of this copyright and license, "Contributing Authors" +is defined as the following set of individuals: + + Andreas Dilger + Dave Martindale + Guy Eric Schalnat + Paul Schmidt + Tim Wegner + +The PNG Reference Library is supplied "AS IS". The Contributing +Authors and Group 42, Inc. disclaim all warranties, expressed or +implied, including, without limitation, the warranties of +merchantability and of fitness for any purpose. The Contributing +Authors and Group 42, Inc. assume no liability for direct, indirect, +incidental, special, exemplary, or consequential damages, which may +result from the use of the PNG Reference Library, even if advised of +the possibility of such damage. + +Permission is hereby granted to use, copy, modify, and distribute this +source code, or portions hereof, for any purpose, without fee, subject +to the following restrictions: + + 1. The origin of this source code must not be misrepresented. + + 2. Altered versions must be plainly marked as such and must not + be misrepresented as being the original source. + + 3. This Copyright notice may not be removed or altered from any + source or altered source distribution. + +The Contributing Authors and Group 42, Inc. specifically permit, +without fee, and encourage the use of this source code as a component +to supporting the PNG file format in commercial products. If you use +this source code in a product, acknowledgment is not required but would +be appreciated. + + +---- + +LIBTIFF + +Copyright (c) 1988-1997 Sam Leffler +Copyright (c) 1991-1997 Silicon Graphics, Inc. + +Permission to use, copy, modify, distribute, and sell this software and +its documentation for any purpose is hereby granted without fee, provided +that (i) the above copyright notices and this permission notice appear in +all copies of the software and related documentation, and (ii) the names of +Sam Leffler and Silicon Graphics may not be used in any advertising or +publicity relating to the software without the specific, prior written +permission of Sam Leffler and Silicon Graphics. + +THE SOFTWARE IS PROVIDED "AS-IS" AND WITHOUT WARRANTY OF ANY KIND, +EXPRESS, IMPLIED OR OTHERWISE, INCLUDING WITHOUT LIMITATION, ANY +WARRANTY OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. + +IN NO EVENT SHALL SAM LEFFLER OR SILICON GRAPHICS BE LIABLE FOR +ANY SPECIAL, INCIDENTAL, INDIRECT OR CONSEQUENTIAL DAMAGES OF ANY KIND, +OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, +WHETHER OR NOT ADVISED OF THE POSSIBILITY OF DAMAGE, AND ON ANY THEORY OF +LIABILITY, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE +OF THIS SOFTWARE. + + +---- + +LIBWEBP + +Copyright (c) 2010, Google Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + + * Neither the name of Google nor the names of its contributors may + be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +---- + +LIBYUV + +Copyright 2011 The LibYuv Project Authors. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in + the documentation and/or other materials provided with the + distribution. + + * Neither the name of Google nor the names of its contributors may + be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +---- + +OPENJPEG + +* + * The copyright in this software is being made available under the 2-clauses + * BSD License, included below. This software may be subject to other third + * party and contributor rights, including patent rights, and no such rights + * are granted under this license. + * + * Copyright (c) 2002-2014, Universite catholique de Louvain (UCL), Belgium + * Copyright (c) 2002-2014, Professor Benoit Macq + * Copyright (c) 2003-2014, Antonin Descampe + * Copyright (c) 2003-2009, Francois-Olivier Devaux + * Copyright (c) 2005, Herve Drolon, FreeImage Team + * Copyright (c) 2002-2003, Yannick Verschueren + * Copyright (c) 2001-2003, David Janssens + * Copyright (c) 2011-2012, Centre National d'Etudes Spatiales (CNES), France + * Copyright (c) 2012, CS Systemes d'Information, France + * + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * 1. Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS `AS IS' + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + */ + + +---- + +RAQM + +The MIT License (MIT) + +Copyright © 2015 Information Technology Authority (ITA) +Copyright © 2016 Khaled Hosny + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + +---- + +XAU + +Copyright 1988, 1993, 1994, 1998 The Open Group + +Permission to use, copy, modify, distribute, and sell this software and its +documentation for any purpose is hereby granted without fee, provided that +the above copyright notice appear in all copies and that both that +copyright notice and this permission notice appear in supporting +documentation. + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +OPEN GROUP BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN +AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +Except as contained in this notice, the name of The Open Group shall not be +used in advertising or otherwise to promote the sale, use or other dealings +in this Software without prior written authorization from The Open Group. + + +---- + +XCB + +Copyright (C) 2001-2006 Bart Massey, Jamey Sharp, and Josh Triplett. +All Rights Reserved. + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, +sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, +subject to the following conditions: + +The above copyright notice and this permission notice shall +be included in all copies or substantial portions of the +Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY +KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE +WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR +PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS +BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +OTHER DEALINGS IN THE SOFTWARE. + +Except as contained in this notice, the names of the authors +or their institutions shall not be used in advertising or +otherwise to promote the sale, use or other dealings in this +Software without prior written authorization from the +authors. + + +---- + +XDMCP + +Copyright 1989, 1998 The Open Group + +Permission to use, copy, modify, distribute, and sell this software and its +documentation for any purpose is hereby granted without fee, provided that +the above copyright notice appear in all copies and that both that +copyright notice and this permission notice appear in supporting +documentation. + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +OPEN GROUP BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN +AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +Except as contained in this notice, the name of The Open Group shall not be +used in advertising or otherwise to promote the sale, use or other dealings +in this Software without prior written authorization from The Open Group. + +Author: Keith Packard, MIT X Consortium + + +---- + +ZLIB + + (C) 1995-2017 Jean-loup Gailly and Mark Adler + + This software is provided 'as-is', without any express or implied + warranty. In no event will the authors be held liable for any damages + arising from the use of this software. + + Permission is granted to anyone to use this software for any purpose, + including commercial applications, and to alter it and redistribute it + freely, subject to the following restrictions: + + 1. The origin of this software must not be misrepresented; you must not + claim that you wrote the original software. If you use this software + in a product, an acknowledgment in the product documentation would be + appreciated but is not required. + 2. Altered source versions must be plainly marked as such, and must not be + misrepresented as being the original software. + 3. This notice may not be removed or altered from any source distribution. + + Jean-loup Gailly Mark Adler + jloup@gzip.org madler@alumni.caltech.edu + +If you use the zlib library in a product, we would appreciate *not* receiving +lengthy legal documents to sign. The sources are provided for free but without +warranty of any kind. The library has been entirely written by Jean-loup +Gailly and Mark Adler; it does not include third-party code. + +If you redistribute modified sources, we would appreciate that you include in +the file ChangeLog history information documenting your changes. Please read +the FAQ for more information on the distribution of modified source versions. + + +---- + +ZSTD + +BSD License + +For Zstandard software + +Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name Facebook, nor Meta, nor the names of its contributors may + be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..b338169ce0c740c335bfe82912227ae8637bd492 --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/top_level.txt @@ -0,0 +1 @@ +PIL diff --git a/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/zip-safe b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/zip-safe new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/venv/lib/python3.10/site-packages/pillow-12.0.0.dist-info/zip-safe @@ -0,0 +1 @@ + diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0dd40dc641c0966fad791870ee6bb0f5aa21d131 --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/METADATA @@ -0,0 +1,443 @@ +Metadata-Version: 2.4 +Name: propcache +Version: 0.4.1 +Summary: Accelerated property cache +Home-page: https://github.com/aio-libs/propcache +Author: Andrew Svetlov +Author-email: andrew.svetlov@gmail.com +Maintainer: aiohttp team +Maintainer-email: team@aiohttp.org +License: Apache-2.0 +Project-URL: Chat: Matrix, https://matrix.to/#/#aio-libs:matrix.org +Project-URL: Chat: Matrix Space, https://matrix.to/#/#aio-libs-space:matrix.org +Project-URL: CI: GitHub Workflows, https://github.com/aio-libs/propcache/actions?query=branch:master +Project-URL: Code of Conduct, https://github.com/aio-libs/.github/blob/master/CODE_OF_CONDUCT.md +Project-URL: Coverage: codecov, https://codecov.io/github/aio-libs/propcache +Project-URL: Docs: Changelog, https://propcache.readthedocs.io/en/latest/changes/ +Project-URL: Docs: RTD, https://propcache.readthedocs.io +Project-URL: GitHub: issues, https://github.com/aio-libs/propcache/issues +Project-URL: GitHub: repo, https://github.com/aio-libs/propcache +Keywords: cython,cext,propcache +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Programming Language :: Cython +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Topic :: Internet :: WWW/HTTP +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Requires-Python: >=3.9 +Description-Content-Type: text/x-rst +License-File: LICENSE +License-File: NOTICE +Dynamic: license-file + +propcache +========= + +The module provides a fast implementation of cached properties for Python 3.9+. + +.. image:: https://github.com/aio-libs/propcache/actions/workflows/ci-cd.yml/badge.svg + :target: https://github.com/aio-libs/propcache/actions?query=workflow%3ACI + :align: right + +.. image:: https://codecov.io/gh/aio-libs/propcache/branch/master/graph/badge.svg + :target: https://codecov.io/gh/aio-libs/propcache + +.. image:: https://badge.fury.io/py/propcache.svg + :target: https://badge.fury.io/py/propcache + + +.. image:: https://readthedocs.org/projects/propcache/badge/?version=latest + :target: https://propcache.readthedocs.io + + +.. image:: https://img.shields.io/pypi/pyversions/propcache.svg + :target: https://pypi.python.org/pypi/propcache + +.. image:: https://img.shields.io/matrix/aio-libs:matrix.org?label=Discuss%20on%20Matrix%20at%20%23aio-libs%3Amatrix.org&logo=matrix&server_fqdn=matrix.org&style=flat + :target: https://matrix.to/#/%23aio-libs:matrix.org + :alt: Matrix Room — #aio-libs:matrix.org + +.. image:: https://img.shields.io/matrix/aio-libs-space:matrix.org?label=Discuss%20on%20Matrix%20at%20%23aio-libs-space%3Amatrix.org&logo=matrix&server_fqdn=matrix.org&style=flat + :target: https://matrix.to/#/%23aio-libs-space:matrix.org + :alt: Matrix Space — #aio-libs-space:matrix.org + +Introduction +------------ + +The API is designed to be nearly identical to the built-in ``functools.cached_property`` class, +except for the additional ``under_cached_property`` class which uses ``self._cache`` +instead of ``self.__dict__`` to store the cached values and prevents ``__set__`` from being called. + +For full documentation please read https://propcache.readthedocs.io. + +Installation +------------ + +:: + + $ pip install propcache + +The library is Python 3 only! + +PyPI contains binary wheels for Linux, Windows and MacOS. If you want to install +``propcache`` on another operating system where wheels are not provided, +the the tarball will be used to compile the library from +the source code. It requires a C compiler and and Python headers installed. + +To skip the compilation you must explicitly opt-in by using a PEP 517 +configuration setting ``pure-python``, or setting the ``PROPCACHE_NO_EXTENSIONS`` +environment variable to a non-empty value, e.g.: + +.. code-block:: console + + $ pip install propcache --config-settings=pure-python=false + +Please note that the pure-Python (uncompiled) version is much slower. However, +PyPy always uses a pure-Python implementation, and, as such, it is unaffected +by this variable. + + +API documentation +------------------ + +The documentation is located at https://propcache.readthedocs.io. + +Source code +----------- + +The project is hosted on GitHub_ + +Please file an issue on the `bug tracker +`_ if you have found a bug +or have some suggestion in order to improve the library. + +Discussion list +--------------- + +*aio-libs* google group: https://groups.google.com/forum/#!forum/aio-libs + +Feel free to post your questions and ideas here. + + +Authors and License +------------------- + +The ``propcache`` package is derived from ``yarl`` which is written by Andrew Svetlov. + +It's *Apache 2* licensed and freely available. + + +.. _GitHub: https://github.com/aio-libs/propcache + +========= +Changelog +========= + +.. + You should *NOT* be adding new change log entries to this file, this + file is managed by towncrier. You *may* edit previous change logs to + fix problems like typo corrections or such. + To add a new change log entry, please see + https://pip.pypa.io/en/latest/development/#adding-a-news-entry + we named the news folder "changes". + + WARNING: Don't drop the next directive! + +.. towncrier release notes start + +0.4.1 +===== + +*(2025-10-08)* + + +Bug fixes +--------- + +- Fixed reference leak caused by ``Py_INCREF`` because Cython has its own reference counter systems -- by `@Vizonex `__. + + *Related issues and pull requests on GitHub:* + `#162 `__. + + +Contributor-facing changes +-------------------------- + +- Fixes the default value for the ``os`` + parameter in ``reusable-build-wheel.yml`` + to be ``ubuntu-latest`` instead of + ``ubuntu``. + + *Related issues and pull requests on GitHub:* + `#155 `__. + + +---- + + +0.4.0 +===== + +*(2025-10-04)* + + +Features +-------- + +- Optimized propcache by replacing sentinel ``object`` for checking if + the ``object`` is ``NULL`` and changed ``dict`` API for + Python C-API -- by `@Vizonex `__. + + *Related issues and pull requests on GitHub:* + `#121 `__. + + +Contributor-facing changes +-------------------------- + +- Builds have been added for arm64 Windows + wheels and the ``reusable-build-wheel.yml`` + workflow has been modified to allow for + an OS value (``windows-11-arm``) which + does not include the ``-latest`` postfix + -- by `@finnagin `__. + + *Related issues and pull requests on GitHub:* + `#133 `__. + +- Added CI for CPython 3.14 -- by `@kumaraditya303 `__. + + *Related issues and pull requests on GitHub:* + `#140 `__. + + +---- + + +0.3.2 +===== + +*(2025-06-09)* + + +Improved documentation +---------------------- + +- Fixed incorrect decorator usage in the ``~propcache.api.under_cached_property`` example code -- by `@meanmail `__. + + *Related issues and pull requests on GitHub:* + `#109 `__. + + +Packaging updates and notes for downstreams +------------------------------------------- + +- Updated to use Cython 3.1 universally across the build path -- by `@lysnikolaou `__. + + *Related issues and pull requests on GitHub:* + `#117 `__. + +- Made Cython line tracing opt-in via the ``with-cython-tracing`` build config setting -- by `@bdraco `__. + + Previously, line tracing was enabled by default in ``pyproject.toml``, which caused build issues for some users and made wheels nearly twice as slow. + + Now line tracing is only enabled when explicitly requested via ``pip install . --config-setting=with-cython-tracing=true`` or by setting the ``PROPCACHE_CYTHON_TRACING`` environment variable. + + *Related issues and pull requests on GitHub:* + `#118 `__. + + +---- + + +0.3.1 +===== + +*(2025-03-25)* + + +Bug fixes +--------- + +- Improved typing annotations, fixing some type errors under correct usage + and improving typing robustness generally -- by `@Dreamsorcerer `__. + + *Related issues and pull requests on GitHub:* + `#103 `__. + + +---- + + +0.3.0 +===== + +*(2025-02-20)* + + +Features +-------- + +- Implemented support for the free-threaded build of CPython 3.13 -- by `@lysnikolaou `__. + + *Related issues and pull requests on GitHub:* + `#84 `__. + + +Packaging updates and notes for downstreams +------------------------------------------- + +- Started building wheels for the free-threaded build of CPython 3.13 -- by `@lysnikolaou `__. + + *Related issues and pull requests on GitHub:* + `#84 `__. + + +Contributor-facing changes +-------------------------- + +- GitHub Actions CI/CD is now configured to manage caching pip-ecosystem + dependencies using `re-actors/cache-python-deps`_ -- an action by + `@webknjaz `__ that takes into account ABI stability and the exact + version of Python runtime. + + .. _`re-actors/cache-python-deps`: + https://github.com/marketplace/actions/cache-python-deps + + *Related issues and pull requests on GitHub:* + `#93 `__. + + +---- + + +0.2.1 +===== + +*(2024-12-01)* + + +Bug fixes +--------- + +- Stopped implicitly allowing the use of Cython pre-release versions when + building the distribution package -- by `@ajsanchezsanz `__ and + `@markgreene74 `__. + + *Related commits on GitHub:* + `64df0a6 `__. + +- Fixed ``wrapped`` and ``func`` not being accessible in the Cython versions of ``propcache.api.cached_property`` and ``propcache.api.under_cached_property`` decorators -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#72 `__. + + +Removals and backward incompatible breaking changes +--------------------------------------------------- + +- Removed support for Python 3.8 as it has reached end of life -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#57 `__. + + +Packaging updates and notes for downstreams +------------------------------------------- + +- Stopped implicitly allowing the use of Cython pre-release versions when + building the distribution package -- by `@ajsanchezsanz `__ and + `@markgreene74 `__. + + *Related commits on GitHub:* + `64df0a6 `__. + + +---- + + +0.2.0 +===== + +*(2024-10-07)* + + +Bug fixes +--------- + +- Fixed loading the C-extensions on Python 3.8 -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#26 `__. + + +Features +-------- + +- Improved typing for the ``propcache.api.under_cached_property`` decorator -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#38 `__. + + +Improved documentation +---------------------- + +- Added API documentation for the ``propcache.api.cached_property`` and ``propcache.api.under_cached_property`` decorators -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#16 `__. + + +Packaging updates and notes for downstreams +------------------------------------------- + +- Moved ``propcache.api.under_cached_property`` and ``propcache.api.cached_property`` to `propcache.api` -- by `@bdraco `__. + + Both decorators remain importable from the top-level package, however importing from `propcache.api` is now the recommended way to use them. + + *Related issues and pull requests on GitHub:* + `#19 `__, `#24 `__, `#32 `__. + +- Converted project to use a src layout -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#22 `__, `#29 `__, `#37 `__. + + +---- + + +0.1.0 +===== + +*(2024-10-03)* + + +Features +-------- + +- Added ``armv7l`` wheels -- by `@bdraco `__. + + *Related issues and pull requests on GitHub:* + `#5 `__. + + +---- + + +0.0.0 +===== + +*(2024-10-02)* + + +- Initial release. diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..6e4dd85a161884f087767245a8785880dd227b29 --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/RECORD @@ -0,0 +1,19 @@ +propcache-0.4.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +propcache-0.4.1.dist-info/METADATA,sha256=_CPzxSszbVM1zQU_R9kf9jFpSeZa0fHj38zVKE9Nbs4,13745 +propcache-0.4.1.dist-info/RECORD,, +propcache-0.4.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +propcache-0.4.1.dist-info/WHEEL,sha256=DZl4yYurviXJJQsilHH1qAzCRDsdz--oNtDZ-hUrZUk,190 +propcache-0.4.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358 +propcache-0.4.1.dist-info/licenses/NOTICE,sha256=VtasbIEFwKUTBMIdsGDjYa-ajqCvmnXCOcKLXRNpODg,609 +propcache-0.4.1.dist-info/top_level.txt,sha256=pVF_GbqSAITPMiX27kfU3QP9-ufhRvkADmudDxWdF3w,10 +propcache/__init__.py,sha256=8kebeGvYn7s-ow1AFmK0A4EvonZMpyM7Lkzs2Ktia3Y,965 +propcache/__pycache__/__init__.cpython-310.pyc,, +propcache/__pycache__/_helpers.cpython-310.pyc,, +propcache/__pycache__/_helpers_py.cpython-310.pyc,, +propcache/__pycache__/api.cpython-310.pyc,, +propcache/_helpers.py,sha256=68SQm6kETN8Mnt9Ol26LJYgHgmB0mKy1tp92888zN4k,1553 +propcache/_helpers_c.cpython-310-x86_64-linux-gnu.so,sha256=cp-xCTiqZ6AZoq3ObGW6e7X3uY0cjoakONTjJXhMbdc,642864 +propcache/_helpers_c.pyx,sha256=kcJa1U5lh54TPCqAeZ0cVB7URcb3I8ZbJieOrkNhLQE,3265 +propcache/_helpers_py.py,sha256=Wixs2zWA-FBU-j4zLPyBUU24FEfPhKk-UunFSp9q95U,1909 +propcache/api.py,sha256=wvgB-ypkkI5uf72VVYl2NFGc_TnzUQA2CxC7dTlL5ak,179 +propcache/py.typed,sha256=ay5OMO475PlcZ_Fbun9maHW7Y6MBTk0UXL4ztHx3Iug,14 diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..6d6bf82d678c683effaf22b612f4dea70dd9933d --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/WHEEL @@ -0,0 +1,7 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: false +Tag: cp310-cp310-manylinux_2_17_x86_64 +Tag: cp310-cp310-manylinux2014_x86_64 +Tag: cp310-cp310-manylinux_2_28_x86_64 + diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..d645695673349e3947e8e5ae42332d0ac3164cd7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/NOTICE b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..fa53b2b138df881c4c95239d0e4bede831b36ab5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/licenses/NOTICE @@ -0,0 +1,13 @@ + Copyright 2016-2021, Andrew Svetlov and aio-libs team + + 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. diff --git a/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..8c9accf6226df7e4011a41ac5d6014223685cfed --- /dev/null +++ b/venv/lib/python3.10/site-packages/propcache-0.4.1.dist-info/top_level.txt @@ -0,0 +1 @@ +propcache diff --git a/venv/lib/python3.10/site-packages/safetensors/__init__.py b/venv/lib/python3.10/site-packages/safetensors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0aea2154eb6b993974d3e2874f2d3ad09e21d9c0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/__init__.py @@ -0,0 +1,10 @@ +# Re-export this +from ._safetensors_rust import ( # noqa: F401 + SafetensorError, + __version__, + deserialize, + safe_open, + _safe_open_handle, + serialize, + serialize_file, +) diff --git a/venv/lib/python3.10/site-packages/safetensors/__init__.pyi b/venv/lib/python3.10/site-packages/safetensors/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a89e0655a1522224f44de863e9276d0562b27162 --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/__init__.pyi @@ -0,0 +1,164 @@ +# Generated content DO NOT EDIT +@staticmethod +def deserialize(bytes): + """ + Opens a safetensors lazily and returns tensors as asked + + Args: + data (`bytes`): + The byte content of a file + + Returns: + (`List[str, Dict[str, Dict[str, any]]]`): + The deserialized content is like: + [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)] + """ + pass + +@staticmethod +def serialize(tensor_dict, metadata=None): + """ + Serializes raw data. + + Args: + tensor_dict (`Dict[str, Dict[Any]]`): + The tensor dict is like: + {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} + metadata (`Dict[str, str]`, *optional*): + The optional purely text annotations + + Returns: + (`bytes`): + The serialized content. + """ + pass + +@staticmethod +def serialize_file(tensor_dict, filename, metadata=None): + """ + Serializes raw data into file. + + Args: + tensor_dict (`Dict[str, Dict[Any]]`): + The tensor dict is like: + {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} + filename (`str`, or `os.PathLike`): + The name of the file to write into. + metadata (`Dict[str, str]`, *optional*): + The optional purely text annotations + + Returns: + (`NoneType`): + On success return None + """ + pass + +class safe_open: + """ + Opens a safetensors lazily and returns tensors as asked + + Args: + filename (`str`, or `os.PathLike`): + The filename to open + + framework (`str`): + The framework you want you tensors in. Supported values: + `pt`, `tf`, `flax`, `numpy`. + + device (`str`, defaults to `"cpu"`): + The device on which you want the tensors. + """ + def __init__(self, filename, framework, device=...): + pass + + def __enter__(self): + """ + Start the context manager + """ + pass + + def __exit__(self, _exc_type, _exc_value, _traceback): + """ + Exits the context manager + """ + pass + + def get_slice(self, name): + """ + Returns a full slice view object + + Args: + name (`str`): + The name of the tensor you want + + Returns: + (`PySafeSlice`): + A dummy object you can slice into to get a real tensor + Example: + ```python + from safetensors import safe_open + + with safe_open("model.safetensors", framework="pt", device=0) as f: + tensor_part = f.get_slice("embedding")[:, ::8] + + ``` + """ + pass + + def get_tensor(self, name): + """ + Returns a full tensor + + Args: + name (`str`): + The name of the tensor you want + + Returns: + (`Tensor`): + The tensor in the framework you opened the file for. + + Example: + ```python + from safetensors import safe_open + + with safe_open("model.safetensors", framework="pt", device=0) as f: + tensor = f.get_tensor("embedding") + + ``` + """ + pass + + def keys(self): + """ + Returns the names of the tensors in the file. + + Returns: + (`List[str]`): + The name of the tensors contained in that file + """ + pass + + def metadata(self): + """ + Return the special non tensor information in the header + + Returns: + (`Dict[str, str]`): + The freeform metadata. + """ + pass + + def offset_keys(self): + """ + Returns the names of the tensors in the file, ordered by offset. + + Returns: + (`List[str]`): + The name of the tensors contained in that file + """ + pass + +class SafetensorError(Exception): + """ + Custom Python Exception for Safetensor errors. + """ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81c2e26ae2804e048b44bce1573258e579ffbbb0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/flax.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/flax.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2db0d6b05867fb8727b1ac082cd4c78a8bbdba66 Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/flax.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/mlx.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/mlx.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08a455d76e39201d0f036609c93b73f51f05d260 Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/mlx.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/numpy.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/numpy.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01455eaa22e5c4f5adb931d23f0eed279fe1783b Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/numpy.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/paddle.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/paddle.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd5e9bc3703ad4a18540211e86f153e0a313e37a Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/paddle.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/tensorflow.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/tensorflow.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c366a87c59231d6a001ccb4dde6f99f4c757119 Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/tensorflow.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/__pycache__/torch.cpython-310.pyc b/venv/lib/python3.10/site-packages/safetensors/__pycache__/torch.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fb2426eca0616750d38081e956f88dc7c1db972 Binary files /dev/null and b/venv/lib/python3.10/site-packages/safetensors/__pycache__/torch.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/safetensors/flax.py b/venv/lib/python3.10/site-packages/safetensors/flax.py new file mode 100644 index 0000000000000000000000000000000000000000..a6471507dc44bff6c4a75e5b75a24241a693dad5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/flax.py @@ -0,0 +1,138 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np + +import jax.numpy as jnp +from jax import Array +from safetensors import numpy, safe_open + + +def save(tensors: Dict[str, Array], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, Array]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.flax import save + from jax import numpy as jnp + + tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _jnp2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, Array], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, Array]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.flax import save_file + from jax import numpy as jnp + + tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _jnp2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, Array]: + """ + Loads a safetensors file into flax format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, Array]`: dictionary that contains name as key, value as `Array` on cpu + + Example: + + ```python + from safetensors.flax import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2jnp(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, Array]: + """ + Loads a safetensors file into flax format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + + Returns: + `Dict[str, Array]`: dictionary that contains name as key, value as `Array` + + Example: + + ```python + from safetensors.flax import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="flax") as f: + for k in f.offset_keys(): + result[k] = f.get_tensor(k) + return result + + +def _np2jnp(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, Array]: + for k, v in numpy_dict.items(): + numpy_dict[k] = jnp.array(v) + return numpy_dict + + +def _jnp2np(jnp_dict: Dict[str, Array]) -> Dict[str, np.array]: + for k, v in jnp_dict.items(): + jnp_dict[k] = np.asarray(v) + return jnp_dict diff --git a/venv/lib/python3.10/site-packages/safetensors/mlx.py b/venv/lib/python3.10/site-packages/safetensors/mlx.py new file mode 100644 index 0000000000000000000000000000000000000000..985b73b98ef19afc2104ed1b2464fdeaf829d3bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/mlx.py @@ -0,0 +1,140 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np + +import mlx.core as mx +from safetensors import numpy, safe_open + + +def save( + tensors: Dict[str, mx.array], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, mx.array]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.mlx import save + import mlx.core as mx + + tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _mx2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, mx.array], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, mx.array]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.mlx import save_file + import mlx.core as mx + + tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _mx2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, mx.array]: + """ + Loads a safetensors file into MLX format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array` + + Example: + + ```python + from safetensors.mlx import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2mx(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, mx.array]: + """ + Loads a safetensors file into MLX format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + + Returns: + `Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array` + + Example: + + ```python + from safetensors.flax import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="mlx") as f: + for k in f.offset_keys(): + result[k] = f.get_tensor(k) + return result + + +def _np2mx(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, mx.array]: + for k, v in numpy_dict.items(): + numpy_dict[k] = mx.array(v) + return numpy_dict + + +def _mx2np(mx_dict: Dict[str, mx.array]) -> Dict[str, np.array]: + new_dict = {} + for k, v in mx_dict.items(): + new_dict[k] = np.asarray(v) + return new_dict diff --git a/venv/lib/python3.10/site-packages/safetensors/numpy.py b/venv/lib/python3.10/site-packages/safetensors/numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..282f287186a835cc7218b88f17eb237ec93add7f --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/numpy.py @@ -0,0 +1,186 @@ +import os +import sys +from typing import Dict, Optional, Union + +import numpy as np + +from safetensors import deserialize, safe_open, serialize, serialize_file + + +def _tobytes(tensor: np.ndarray) -> bytes: + if not _is_little_endian(tensor): + tensor = tensor.byteswap(inplace=False) + return tensor.tobytes() + + +def save( + tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensor_dict (`Dict[str, np.ndarray]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.numpy import save + import numpy as np + + tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + flattened = { + k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} + for k, v in tensor_dict.items() + } + serialized = serialize(flattened, metadata=metadata) + result = bytes(serialized) + return result + + +def save_file( + tensor_dict: Dict[str, np.ndarray], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensor_dict (`Dict[str, np.ndarray]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.numpy import save_file + import numpy as np + + tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + flattened = { + k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} + for k, v in tensor_dict.items() + } + serialize_file(flattened, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, np.ndarray]: + """ + Loads a safetensors file into numpy format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` on cpu + + Example: + + ```python + from safetensors.numpy import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = deserialize(data) + return _view2np(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, np.ndarray]: + """ + Loads a safetensors file into numpy format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + + Returns: + `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` + + Example: + + ```python + from safetensors.numpy import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="np") as f: + for k in f.offset_keys(): + result[k] = f.get_tensor(k) + return result + + +_TYPES = { + "F64": np.float64, + "F32": np.float32, + "F16": np.float16, + "I64": np.int64, + "U64": np.uint64, + "I32": np.int32, + "U32": np.uint32, + "I16": np.int16, + "U16": np.uint16, + "I8": np.int8, + "U8": np.uint8, + "BOOL": bool, +} + + +def _getdtype(dtype_str: str) -> np.dtype: + return _TYPES[dtype_str] + + +def _view2np(safeview) -> Dict[str, np.ndarray]: + result = {} + for k, v in safeview: + dtype = _getdtype(v["dtype"]) + arr = np.frombuffer(v["data"], dtype=dtype).reshape(v["shape"]) + result[k] = arr + return result + + +def _is_little_endian(tensor: np.ndarray) -> bool: + byteorder = tensor.dtype.byteorder + if byteorder == "=": + if sys.byteorder == "little": + return True + else: + return False + elif byteorder == "|": + return True + elif byteorder == "<": + return True + elif byteorder == ">": + return False + raise ValueError(f"Unexpected byte order {byteorder}") diff --git a/venv/lib/python3.10/site-packages/safetensors/paddle.py b/venv/lib/python3.10/site-packages/safetensors/paddle.py new file mode 100644 index 0000000000000000000000000000000000000000..5ec459e3ac45ea56583376ecbffd77734bf82381 --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/paddle.py @@ -0,0 +1,144 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np + +import paddle +from safetensors import numpy + + +def save( + tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, paddle.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.paddle import save + import paddle + + tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _paddle2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, paddle.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, paddle.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.paddle import save_file + import paddle + + tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _paddle2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]: + """ + Loads a safetensors file into paddle format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` on cpu + + Example: + + ```python + from safetensors.paddle import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2paddle(flat, device) + + +def load_file( + filename: Union[str, os.PathLike], device="cpu" +) -> Dict[str, paddle.Tensor]: + """ + Loads a safetensors file into paddle format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + device (`Union[Dict[str, any], str]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular paddle device locations + + Returns: + `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` + + Example: + + ```python + from safetensors.paddle import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + flat = numpy.load_file(filename) + output = _np2paddle(flat, device) + return output + + +def _np2paddle( + numpy_dict: Dict[str, np.ndarray], device: str = "cpu" +) -> Dict[str, paddle.Tensor]: + for k, v in numpy_dict.items(): + numpy_dict[k] = paddle.to_tensor(v, place=device) + return numpy_dict + + +def _paddle2np(paddle_dict: Dict[str, paddle.Tensor]) -> Dict[str, np.array]: + for k, v in paddle_dict.items(): + paddle_dict[k] = v.detach().cpu().numpy() + return paddle_dict diff --git a/venv/lib/python3.10/site-packages/safetensors/py.typed b/venv/lib/python3.10/site-packages/safetensors/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/safetensors/tensorflow.py b/venv/lib/python3.10/site-packages/safetensors/tensorflow.py new file mode 100644 index 0000000000000000000000000000000000000000..c2bfc936ebae5350c0599a79c8ec3841fd6cda0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/tensorflow.py @@ -0,0 +1,139 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np +import tensorflow as tf + +from safetensors import numpy, safe_open + + +def save( + tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, tf.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.tensorflow import save + import tensorflow as tf + + tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _tf2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, tf.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, tf.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.tensorflow import save_file + import tensorflow as tf + + tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _tf2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, tf.Tensor]: + """ + Loads a safetensors file into tensorflow format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` on cpu + + Example: + + ```python + from safetensors.tensorflow import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2tf(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, tf.Tensor]: + """ + Loads a safetensors file into tensorflow format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + + Returns: + `Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` + + Example: + + ```python + from safetensors.tensorflow import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="tf") as f: + for k in f.offset_keys(): + result[k] = f.get_tensor(k) + return result + + +def _np2tf(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, tf.Tensor]: + for k, v in numpy_dict.items(): + numpy_dict[k] = tf.convert_to_tensor(v) + return numpy_dict + + +def _tf2np(tf_dict: Dict[str, tf.Tensor]) -> Dict[str, np.array]: + for k, v in tf_dict.items(): + tf_dict[k] = v.numpy() + return tf_dict diff --git a/venv/lib/python3.10/site-packages/safetensors/torch.py b/venv/lib/python3.10/site-packages/safetensors/torch.py new file mode 100644 index 0000000000000000000000000000000000000000..f2285e0f0d10dbdc31fef92496ccfe859a64aa4e --- /dev/null +++ b/venv/lib/python3.10/site-packages/safetensors/torch.py @@ -0,0 +1,592 @@ +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List, Optional, Set, Tuple, Union +from packaging.version import Version + +import torch + +from safetensors import deserialize, safe_open, serialize, serialize_file + + +def storage_ptr(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().data_ptr() + except Exception: + # Fallback for torch==1.10 + try: + return tensor.storage().data_ptr() + except NotImplementedError: + # Fallback for meta storage + return 0 + + +def _end_ptr(tensor: torch.Tensor) -> int: + if tensor.nelement(): + stop = tensor.view(-1)[-1].data_ptr() + _SIZE[tensor.dtype] + else: + stop = tensor.data_ptr() + return stop + + +def storage_size(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().nbytes() + except AttributeError: + # Fallback for torch==1.10 + try: + return tensor.storage().size() * _SIZE[tensor.dtype] + except NotImplementedError: + # Fallback for meta storage + # On torch >=2.0 this is the tensor size + return tensor.nelement() * _SIZE[tensor.dtype] + + +def _filter_shared_not_shared( + tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor] +) -> List[Set[str]]: + filtered_tensors = [] + for shared in tensors: + if len(shared) < 2: + filtered_tensors.append(shared) + continue + + areas = [] + for name in shared: + tensor = state_dict[name] + areas.append((tensor.data_ptr(), _end_ptr(tensor), name)) + areas.sort() + + _, last_stop, last_name = areas[0] + filtered_tensors.append({last_name}) + for start, stop, name in areas[1:]: + if start >= last_stop: + filtered_tensors.append({name}) + else: + filtered_tensors[-1].add(name) + last_stop = stop + + return filtered_tensors + + +def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]: + tensors = defaultdict(set) + for k, v in state_dict.items(): + if ( + v.device != torch.device("meta") + and storage_ptr(v) != 0 + and storage_size(v) != 0 + ): + # Need to add device as key because of multiple GPU. + tensors[(v.device, storage_ptr(v), storage_size(v))].add(k) + tensors = list(sorted(tensors.values())) + tensors = _filter_shared_not_shared(tensors, state_dict) + return tensors + + +def _is_complete(tensor: torch.Tensor) -> bool: + return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[ + tensor.dtype + ] == storage_size(tensor) + + +def _remove_duplicate_names( + state_dict: Dict[str, torch.Tensor], + *, + preferred_names: Optional[List[str]] = None, + discard_names: Optional[List[str]] = None, +) -> Dict[str, List[str]]: + if preferred_names is None: + preferred_names = [] + preferred_names = set(preferred_names) + if discard_names is None: + discard_names = [] + discard_names = set(discard_names) + + shareds = _find_shared_tensors(state_dict) + to_remove = defaultdict(list) + for shared in shareds: + complete_names = set( + [name for name in shared if _is_complete(state_dict[name])] + ) + if not complete_names: + raise RuntimeError( + "Error while trying to find names to remove to save state dict, but found no suitable name to keep" + f" for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model" + " since you could be storing much more memory than needed. Please refer to" + " https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an" + " issue." + ) + + keep_name = sorted(list(complete_names))[0] + + # Mechanism to preferentially select keys to keep + # coming from the on-disk file to allow + # loading models saved with a different choice + # of keep_name + preferred = complete_names.difference(discard_names) + if preferred: + keep_name = sorted(list(preferred))[0] + + if preferred_names: + preferred = preferred_names.intersection(complete_names) + if preferred: + keep_name = sorted(list(preferred))[0] + for name in sorted(shared): + if name != keep_name: + to_remove[keep_name].append(name) + return to_remove + + +def save_model( + model: torch.nn.Module, + filename: str, + metadata: Optional[Dict[str, str]] = None, + force_contiguous: bool = True, +): + """ + Saves a given torch model to specified filename. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to save on disk. + filename (`str`): + The filename location to save the file + metadata (`Dict[str, str]`, *optional*): + Extra information to save along with the file. + Some metadata will be added for each dropped tensors. + This information will not be enough to recover the entire + shared structure but might help understanding things + force_contiguous (`boolean`, *optional*, defaults to True): + Forcing the state_dict to be saved as contiguous tensors. + This has no effect on the correctness of the model, but it + could potentially change performance if the layout of the tensor + was chosen specifically for that reason. + """ + state_dict = model.state_dict() + to_removes = _remove_duplicate_names(state_dict) + + for kept_name, to_remove_group in to_removes.items(): + for to_remove in to_remove_group: + if metadata is None: + metadata = {} + + if to_remove not in metadata: + # Do not override user data + metadata[to_remove] = kept_name + del state_dict[to_remove] + if force_contiguous: + state_dict = {k: v.contiguous() for k, v in state_dict.items()} + try: + save_file(state_dict, filename, metadata=metadata) + except ValueError as e: + msg = str(e) + msg += " Or use save_model(..., force_contiguous=True), read the docs for potential caveats." + raise ValueError(msg) + + +def load_model( + model: torch.nn.Module, + filename: Union[str, os.PathLike], + strict: bool = True, + device: Union[str, int] = "cpu", +) -> Tuple[List[str], List[str]]: + """ + Loads a given filename onto a torch model. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to load onto. + filename (`str`, or `os.PathLike`): + The filename location to load the file from. + strict (`bool`, *optional*, defaults to True): + Whether to fail if you're missing keys or having unexpected ones. + When false, the function simply returns missing and unexpected names. + device (`Union[str, int]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular torch device locations. + + Returns: + `(missing, unexpected): (List[str], List[str])` + `missing` are names in the model which were not modified during loading + `unexpected` are names that are on the file, but weren't used during + the load. + """ + state_dict = load_file(filename, device=device) + model_state_dict = model.state_dict() + to_removes = _remove_duplicate_names( + model_state_dict, preferred_names=state_dict.keys() + ) + + reverse_to_remove = {} + for key, to_remove_group in to_removes.items(): + for to_remove in to_remove_group: + reverse_to_remove[to_remove] = key + + # We iterate on the model, so we'll add keys we find missing + # here + missing = set() + # We start with all keys on disk declared as unexpected, we'll + # slowly remove them when we find them + unexpected = set(state_dict.keys()) + # Some keys can be invalid too. + invalid = set() + + for k, mv in model_state_dict.items(): + actual_k = reverse_to_remove.get(k, None) + if actual_k is not None: + look_k = actual_k + else: + look_k = k + v = state_dict.get(look_k, None) + if v is None: + missing.add(k) + else: + # We can actually check for the shapes while we're at it. + # For the device, it's trickier given torch's internals + # There might be some Meta device for faster initiation + if v.dtype != mv.dtype or v.shape != mv.shape: + invalid.add(k) + if actual_k is None: + unexpected.remove(k) + + missing = set(missing) + unexpected = set(unexpected) + if strict and (missing or unexpected or invalid): + missing_keys = ", ".join([f'"{k}"' for k in sorted(missing)]) + unexpected_keys = ", ".join([f'"{k}"' for k in sorted(unexpected)]) + invalid_keys = ", ".join([f'"{k}"' for k in sorted(invalid)]) + error = f"Error(s) in loading state_dict for {model.__class__.__name__}:" + if missing: + error += f"\n Missing key(s) in state_dict: {missing_keys}" + if unexpected: + error += f"\n Unexpected key(s) in state_dict: {unexpected_keys}" + if invalid: + error += f"\n Invalid key(s) in state_dict: {invalid_keys}, mismatched dtypes or shape." + del state_dict + raise RuntimeError(error) + + torch_missing, torch_unexpected = model.load_state_dict(state_dict, strict=False) + # Sanity check that the work we've done matches + # Pytorch internal loading. + torch_missing = set(torch_missing) + torch_unexpected = set(torch_unexpected) + for to_remove_group in to_removes.values(): + for to_remove in to_remove_group: + if to_remove not in torch_missing: + torch_unexpected.add(to_remove) + else: + torch_missing.remove(to_remove) + assert torch_missing == missing, f"{torch_missing} != {missing}" + assert torch_unexpected == unexpected, f"{torch_unexpected} != {unexpected}" + return missing, unexpected + + +def save( + tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.torch import save + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + serialized = serialize(_flatten(tensors), metadata=metadata) + result = bytes(serialized) + return result + + +def save_file( + tensors: Dict[str, torch.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +): + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.torch import save_file + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + serialize_file(_flatten(tensors), filename, metadata=metadata) + + +def load_file( + filename: Union[str, os.PathLike], device: Union[str, int] = "cpu" +) -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format. + + Args: + filename (`str`, or `os.PathLike`): + The name of the file which contains the tensors + device (`Union[str, int]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular torch device locations. + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` + + Example: + + ```python + from safetensors.torch import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="pt", device=device) as f: + for k in f.offset_keys(): + result[k] = f.get_tensor(k) + return result + + +def load(data: bytes) -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` on cpu + + Example: + + ```python + from safetensors.torch import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = deserialize(data) + return _view2torch(flat) + + +# torch.float8 formats require 2.1; we do not support these dtypes on earlier versions +_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None) +_float8_e5m2 = getattr(torch, "float8_e5m2", None) +_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None) +_float4_e2m1_x2 = getattr(torch, "float4_e2m1fn_x2", None) + +_SIZE = { + torch.int64: 8, + torch.float32: 4, + torch.int32: 4, + torch.bfloat16: 2, + torch.float16: 2, + torch.int16: 2, + torch.uint8: 1, + torch.int8: 1, + torch.bool: 1, + torch.float64: 8, + _float8_e4m3fn: 1, + _float8_e5m2: 1, + _float8_e8m0: 1, + _float4_e2m1_x2: 1, +} +if Version(torch.__version__) >= Version("2.3.0"): + _SIZE.update( + { + torch.uint64: 8, + torch.uint32: 4, + torch.uint16: 2, + } + ) + +_TYPES = { + "F64": torch.float64, + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I64": torch.int64, + "I32": torch.int32, + "I16": torch.int16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, + "F8_E4M3": _float8_e4m3fn, + "F8_E5M2": _float8_e5m2, +} +if Version(torch.__version__) >= Version("2.3.0"): + _TYPES.update( + { + "U64": torch.uint64, + "U32": torch.uint32, + "U16": torch.uint16, + } + ) + + +def _getdtype(dtype_str: str) -> torch.dtype: + return _TYPES[dtype_str] + + +def _view2torch(safeview) -> Dict[str, torch.Tensor]: + result = {} + for k, v in safeview: + dtype = _getdtype(v["dtype"]) + if len(v["data"]) == 0: + # Workaround because frombuffer doesn't accept zero-size tensors + assert any(x == 0 for x in v["shape"]) + arr = torch.empty(v["shape"], dtype=dtype) + else: + arr = torch.frombuffer(v["data"], dtype=dtype).reshape(v["shape"]) + if sys.byteorder == "big": + arr = torch.from_numpy(arr.numpy().byteswap(inplace=False)) + result[k] = arr + + return result + + +def _tobytes(tensor: torch.Tensor, name: str) -> bytes: + if tensor.layout != torch.strided: + raise ValueError( + f"You are trying to save a sparse tensor: `{name}` which this library does not support." + " You can make it a dense tensor before saving with `.to_dense()` but be aware this might" + " make a much larger file than needed." + ) + + if not tensor.is_contiguous(): + raise ValueError( + f"You are trying to save a non contiguous tensor: `{name}` which is not allowed. It either means you" + " are trying to save tensors which are reference of each other in which case it's recommended to save" + " only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to" + " pack it before saving." + ) + if tensor.device.type != "cpu": + # Moving tensor to cpu before saving + tensor = tensor.to("cpu") + + import ctypes + + import numpy as np + + # When shape is empty (scalar), np.prod returns a float + # we need a int for the following calculations + length = int(np.prod(tensor.shape).item()) + bytes_per_item = _SIZE[tensor.dtype] + + total_bytes = length * bytes_per_item + + ptr = tensor.data_ptr() + if ptr == 0: + return b"" + newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte)) + data = np.ctypeslib.as_array(newptr, (total_bytes,)) # no internal copy + if sys.byteorder == "big": + NPDTYPES = { + torch.int64: np.int64, + torch.float32: np.float32, + torch.int32: np.int32, + # XXX: This is ok because both have the same width + torch.bfloat16: np.float16, + torch.float16: np.float16, + torch.int16: np.int16, + torch.uint8: np.uint8, + torch.int8: np.int8, + torch.bool: bool, + torch.float64: np.float64, + # XXX: This is ok because both have the same width and byteswap is a no-op anyway + _float8_e4m3fn: np.uint8, + _float8_e5m2: np.uint8, + } + npdtype = NPDTYPES[tensor.dtype] + # Not in place as that would potentially modify a live running model + data = data.view(npdtype).byteswap(inplace=False) + return data.tobytes() + + +def _flatten(tensors: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, Any]]: + if not isinstance(tensors, dict): + raise ValueError( + f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}" + ) + + invalid_tensors = [] + for k, v in tensors.items(): + if not isinstance(v, torch.Tensor): + raise ValueError( + f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}" + ) + + if v.layout != torch.strided: + invalid_tensors.append(k) + if invalid_tensors: + raise ValueError( + f"You are trying to save a sparse tensors: `{invalid_tensors}` which this library does not support." + " You can make it a dense tensor before saving with `.to_dense()` but be aware this might" + " make a much larger file than needed." + ) + + shared_pointers = _find_shared_tensors(tensors) + failing = [] + for names in shared_pointers: + if len(names) > 1: + failing.append(names) + + if failing: + raise RuntimeError( + f""" + Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: {failing}. + A potential way to correctly save your model is to use `save_model`. + More information at https://huggingface.co/docs/safetensors/torch_shared_tensors + """ + ) + + return { + k: { + "dtype": str(v.dtype).split(".")[-1], + "shape": v.shape, + "data": _tobytes(v, k), + } + for k, v in tensors.items() + } diff --git a/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/METADATA b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..e125c86dfb85b77d4bc592c77ce06d031a28ff03 --- /dev/null +++ b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/METADATA @@ -0,0 +1,210 @@ +Metadata-Version: 2.4 +Name: tokenizers +Version: 0.22.1 +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Requires-Dist: huggingface-hub>=0.16.4,<2.0 +Requires-Dist: pytest ; extra == 'testing' +Requires-Dist: pytest-asyncio ; extra == 'testing' +Requires-Dist: requests ; extra == 'testing' +Requires-Dist: numpy ; extra == 'testing' +Requires-Dist: datasets ; extra == 'testing' +Requires-Dist: black==22.3 ; extra == 'testing' +Requires-Dist: ruff ; extra == 'testing' +Requires-Dist: sphinx ; extra == 'docs' +Requires-Dist: sphinx-rtd-theme ; extra == 'docs' +Requires-Dist: setuptools-rust ; extra == 'docs' +Requires-Dist: tokenizers[testing] ; extra == 'dev' +Provides-Extra: testing +Provides-Extra: docs +Provides-Extra: dev +Keywords: NLP,tokenizer,BPE,transformer,deep learning +Author-email: Nicolas Patry , Anthony Moi +Requires-Python: >=3.9 +Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM +Project-URL: Homepage, https://github.com/huggingface/tokenizers +Project-URL: Source, https://github.com/huggingface/tokenizers + +

+
+ +
+

+

+ + Build + + + GitHub + +

+
+ +# Tokenizers + +Provides an implementation of today's most used tokenizers, with a focus on performance and +versatility. + +Bindings over the [Rust](https://github.com/huggingface/tokenizers/tree/master/tokenizers) implementation. +If you are interested in the High-level design, you can go check it there. + +Otherwise, let's dive in! + +## Main features: + + - Train new vocabularies and tokenize using 4 pre-made tokenizers (Bert WordPiece and the 3 + most common BPE versions). + - Extremely fast (both training and tokenization), thanks to the Rust implementation. Takes + less than 20 seconds to tokenize a GB of text on a server's CPU. + - Easy to use, but also extremely versatile. + - Designed for research and production. + - Normalization comes with alignments tracking. It's always possible to get the part of the + original sentence that corresponds to a given token. + - Does all the pre-processing: Truncate, Pad, add the special tokens your model needs. + +### Installation + +#### With pip: + +```bash +pip install tokenizers +``` + +#### From sources: + +To use this method, you need to have the Rust installed: + +```bash +# Install with: +curl https://sh.rustup.rs -sSf | sh -s -- -y +export PATH="$HOME/.cargo/bin:$PATH" +``` + +Once Rust is installed, you can compile doing the following + +```bash +git clone https://github.com/huggingface/tokenizers +cd tokenizers/bindings/python + +# Create a virtual env (you can use yours as well) +python -m venv .env +source .env/bin/activate + +# Install `tokenizers` in the current virtual env +pip install -e . +``` + +### Load a pretrained tokenizer from the Hub + +```python +from tokenizers import Tokenizer + +tokenizer = Tokenizer.from_pretrained("bert-base-cased") +``` + +### Using the provided Tokenizers + +We provide some pre-build tokenizers to cover the most common cases. You can easily load one of +these using some `vocab.json` and `merges.txt` files: + +```python +from tokenizers import CharBPETokenizer + +# Initialize a tokenizer +vocab = "./path/to/vocab.json" +merges = "./path/to/merges.txt" +tokenizer = CharBPETokenizer(vocab, merges) + +# And then encode: +encoded = tokenizer.encode("I can feel the magic, can you?") +print(encoded.ids) +print(encoded.tokens) +``` + +And you can train them just as simply: + +```python +from tokenizers import CharBPETokenizer + +# Initialize a tokenizer +tokenizer = CharBPETokenizer() + +# Then train it! +tokenizer.train([ "./path/to/files/1.txt", "./path/to/files/2.txt" ]) + +# Now, let's use it: +encoded = tokenizer.encode("I can feel the magic, can you?") + +# And finally save it somewhere +tokenizer.save("./path/to/directory/my-bpe.tokenizer.json") +``` + +#### Provided Tokenizers + + - `CharBPETokenizer`: The original BPE + - `ByteLevelBPETokenizer`: The byte level version of the BPE + - `SentencePieceBPETokenizer`: A BPE implementation compatible with the one used by SentencePiece + - `BertWordPieceTokenizer`: The famous Bert tokenizer, using WordPiece + +All of these can be used and trained as explained above! + +### Build your own + +Whenever these provided tokenizers don't give you enough freedom, you can build your own tokenizer, +by putting all the different parts you need together. +You can check how we implemented the [provided tokenizers](https://github.com/huggingface/tokenizers/tree/master/bindings/python/py_src/tokenizers/implementations) and adapt them easily to your own needs. + +#### Building a byte-level BPE + +Here is an example showing how to build your own byte-level BPE by putting all the different pieces +together, and then saving it to a single file: + +```python +from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors + +# Initialize a tokenizer +tokenizer = Tokenizer(models.BPE()) + +# Customize pre-tokenization and decoding +tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True) +tokenizer.decoder = decoders.ByteLevel() +tokenizer.post_processor = processors.ByteLevel(trim_offsets=True) + +# And then train +trainer = trainers.BpeTrainer( + vocab_size=20000, + min_frequency=2, + initial_alphabet=pre_tokenizers.ByteLevel.alphabet() +) +tokenizer.train([ + "./path/to/dataset/1.txt", + "./path/to/dataset/2.txt", + "./path/to/dataset/3.txt" +], trainer=trainer) + +# And Save it +tokenizer.save("byte-level-bpe.tokenizer.json", pretty=True) +``` + +Now, when you want to use this tokenizer, this is as simple as: + +```python +from tokenizers import Tokenizer + +tokenizer = Tokenizer.from_file("byte-level-bpe.tokenizer.json") + +encoded = tokenizer.encode("I can feel the magic, can you?") +``` + diff --git a/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/RECORD b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..7e1f24bc60f5f9fdf52b49e1194a69d7c0019d2e --- /dev/null +++ b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/RECORD @@ -0,0 +1,46 @@ +tokenizers-0.22.1.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +tokenizers-0.22.1.dist-info/METADATA,sha256=3ckvBh_0IvsF0Z2ljpFhxQmNoUPDiqI3eb-KW0jOdso,6779 +tokenizers-0.22.1.dist-info/RECORD,, +tokenizers-0.22.1.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +tokenizers-0.22.1.dist-info/WHEEL,sha256=EmbG9zyqShfWQD8iunhabgzlMZGQwxFV_zZkaywLDn0,127 +tokenizers/__init__.py,sha256=ZE5ZagUvobBScrHBQdEobhx4wqM0bsq9F9aLYkBNjYQ,2615 +tokenizers/__init__.pyi,sha256=4To1kfbT82HE2tSszJJwSKUmy7m3Y5dw6Oqzy_-0Iao,45154 +tokenizers/__pycache__/__init__.cpython-310.pyc,, +tokenizers/decoders/__init__.py,sha256=hfwM6CFUDvlMGGL4-xsaaYz81K9P5rQI5ZL5UHWK8Y4,372 +tokenizers/decoders/__init__.pyi,sha256=72hidxCIWgV-dbJxFe9KReGs2YOXAdjq_kJhJI_mgLY,7395 +tokenizers/decoders/__pycache__/__init__.cpython-310.pyc,, +tokenizers/implementations/__init__.py,sha256=VzAsplaIo7rl4AFO8Miu7ig7MfZjvonwVblZw01zR6M,310 +tokenizers/implementations/__pycache__/__init__.cpython-310.pyc,, +tokenizers/implementations/__pycache__/base_tokenizer.cpython-310.pyc,, +tokenizers/implementations/__pycache__/bert_wordpiece.cpython-310.pyc,, +tokenizers/implementations/__pycache__/byte_level_bpe.cpython-310.pyc,, +tokenizers/implementations/__pycache__/char_level_bpe.cpython-310.pyc,, +tokenizers/implementations/__pycache__/sentencepiece_bpe.cpython-310.pyc,, +tokenizers/implementations/__pycache__/sentencepiece_unigram.cpython-310.pyc,, +tokenizers/implementations/base_tokenizer.py,sha256=HzK6Nm36LxJaKTZrIM1Cx_Ld9gNjeJdB9E29M0aYGBI,15791 +tokenizers/implementations/bert_wordpiece.py,sha256=sKCum0FKPYdSgJFJN8LDerVBoTDRSqyqSdrcm-lvQqI,5520 +tokenizers/implementations/byte_level_bpe.py,sha256=iBepM_z1s5Ky7zFDVrYLc3L5byYrIouk7-k0JGuF10s,4272 +tokenizers/implementations/char_level_bpe.py,sha256=Nag_HFq8Rvcucqi8MhV1-0xtoR0C7FjHOecFVURL7ss,5449 +tokenizers/implementations/sentencepiece_bpe.py,sha256=c08fKf6i92E2RsKgsxy7LzZfYX8-MACHSRG8U_I5ytY,3721 +tokenizers/implementations/sentencepiece_unigram.py,sha256=SYiVXL8ZtqLXKpuqwnwmrfxgGotu8yAkOu7dLztEXIo,7580 +tokenizers/models/__init__.py,sha256=eJZ4HTAQZpxnKILNylWaTFqxXy-Ba6OKswWN47feeV8,176 +tokenizers/models/__init__.pyi,sha256=clPTwiyjz7FlVdEuwo_3Wa_TmQrbZhW0SGmnNylepnY,16929 +tokenizers/models/__pycache__/__init__.cpython-310.pyc,, +tokenizers/normalizers/__init__.py,sha256=_06w4cqRItveEgIddYaLMScgkSOkIAMIzYCesb5AA4U,841 +tokenizers/normalizers/__init__.pyi,sha256=lSFqDb_lPZBfRxEG99EcFEaU1HlnIhIQUu7zZIyP4AY,20898 +tokenizers/normalizers/__pycache__/__init__.cpython-310.pyc,, +tokenizers/pre_tokenizers/__init__.py,sha256=KV9-EsAykGENUUzkGWCbv4n6YM6hYa1hfnY-gzBpMNE,598 +tokenizers/pre_tokenizers/__init__.pyi,sha256=n6BFClhxm8y7miCC0lJd7oVeo8oj3kPg2tU9ObV4PGU,26556 +tokenizers/pre_tokenizers/__pycache__/__init__.cpython-310.pyc,, +tokenizers/processors/__init__.py,sha256=xM2DEKwKtHIumHsszM8AMkq-AlaqvBZFXWgLU8SNhOY,307 +tokenizers/processors/__init__.pyi,sha256=hx767ZY8SHhxb_hiXPRxm-f_KcoR4XDx7vfK2c0lR-Q,11357 +tokenizers/processors/__pycache__/__init__.cpython-310.pyc,, +tokenizers/tokenizers.abi3.so,sha256=-T5w9yJa7pkS35XUNSagc57jj1Nvn4dayt95bQDaX10,10376544 +tokenizers/tools/__init__.py,sha256=xG8caB9OHC8cbB01S5vYV14HZxhO6eWbLehsb70ppio,55 +tokenizers/tools/__pycache__/__init__.cpython-310.pyc,, +tokenizers/tools/__pycache__/visualizer.cpython-310.pyc,, +tokenizers/tools/visualizer-styles.css,sha256=zAydq1oGWD8QEll4-eyL8Llw0B1sty_hpIE3tYxL02k,4850 +tokenizers/tools/visualizer.py,sha256=zEXELCLSxXL_RgQgCTjVQrc3G_ptuDwr7deaU95b3dA,14625 +tokenizers/trainers/__init__.py,sha256=UTu22AGcp76IvpW45xLRbJWET04NxPW6NfCb2YYz0EM,248 +tokenizers/trainers/__init__.pyi,sha256=OwdiVOlMXhU5hOq7a5TYYG1vw3fk8nTqH88tVr05NZ0,5860 +tokenizers/trainers/__pycache__/__init__.cpython-310.pyc,, diff --git a/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/WHEEL b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..3ed6a04ed82d3ad089b0cc2485a62e34ce27e9b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/tokenizers-0.22.1.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: maturin (1.9.4) +Root-Is-Purelib: false +Tag: cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64 diff --git a/venv/lib/python3.10/site-packages/torch/_C.cpython-310-x86_64-linux-gnu.so b/venv/lib/python3.10/site-packages/torch/_C.cpython-310-x86_64-linux-gnu.so new file mode 100644 index 0000000000000000000000000000000000000000..5e6b897e794661345e4884974cf39b8925ee2d17 Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/_C.cpython-310-x86_64-linux-gnu.so differ diff --git a/venv/lib/python3.10/site-packages/torch/_VF.py b/venv/lib/python3.10/site-packages/torch/_VF.py new file mode 100644 index 0000000000000000000000000000000000000000..94166b51f1786593b584629744adc24036e8b1d7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_VF.py @@ -0,0 +1,31 @@ +""" +This makes the functions in torch._C._VariableFunctions available as + torch._VF. +without mypy being able to find them. + +A subset of those functions are mapped to ATen functions in +torch/jit/_builtins.py + +See https://github.com/pytorch/pytorch/issues/21478 for the reason for +introducing torch._VF + +""" + +import sys +import types + +import torch + + +class VFModule(types.ModuleType): + vf: types.ModuleType + + def __init__(self, name: str): + super().__init__(name) + self.vf = torch._C._VariableFunctions + + def __getattr__(self, name: str) -> object: + return getattr(self.vf, name) + + +sys.modules[__name__] = VFModule(__name__) diff --git a/venv/lib/python3.10/site-packages/torch/_VF.pyi b/venv/lib/python3.10/site-packages/torch/_VF.pyi new file mode 100644 index 0000000000000000000000000000000000000000..b86f038e60cf4f79735913961b7cf4ee9469287a --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_VF.pyi @@ -0,0 +1,33541 @@ +# @generated by tools/pyi/gen_pyi.py from torch/_C/_VariableFunctions.pyi.in +# mypy: disable-error-code="type-arg" +# mypy: allow-untyped-defs +# ruff: noqa: F401,PYI054 + +from collections.abc import Sequence +from types import EllipsisType +from typing import Any, Callable, Literal, overload, TypeVar + +import torch +from torch import ( + contiguous_format, + Generator, + inf, + memory_format, + strided, + SymInt, + Tensor, +) +from torch._prims_common import DeviceLikeType +from torch.types import ( + _bool, + _complex, + _device, + _dtype, + _float, + _int, + _layout, + _qscheme, + _size, + Device, + Number, +) + +__all__ = [ + "__and__", + "__lshift__", + "__or__", + "__rshift__", + "__xor__", + "_adaptive_avg_pool2d", + "_adaptive_avg_pool3d", + "_add_batch_dim", + "_add_relu", + "_add_relu_", + "_addmm_activation", + "_aminmax", + "_amp_foreach_non_finite_check_and_unscale_", + "_amp_update_scale_", + "_assert_async", + "_assert_scalar", + "_assert_tensor_metadata", + "_batch_norm_impl_index", + "_cast_Byte", + "_cast_Char", + "_cast_Double", + "_cast_Float", + "_cast_Half", + "_cast_Int", + "_cast_Long", + "_cast_Short", + "_choose_qparams_per_tensor", + "_chunk_cat", + "_coalesce", + "_compute_linear_combination", + "_conj", + "_conj_copy", + "_conj_physical", + "_convert_indices_from_coo_to_csr", + "_convert_indices_from_csr_to_coo", + "_convert_weight_to_int4pack", + "_convert_weight_to_int4pack_for_cpu", + "_convolution", + "_convolution_mode", + "_copy_from", + "_copy_from_and_resize", + "_cslt_compress", + "_cslt_sparse_mm", + "_cslt_sparse_mm_search", + "_ctc_loss", + "_cudnn_ctc_loss", + "_cudnn_init_dropout_state", + "_cudnn_rnn", + "_cudnn_rnn_flatten_weight", + "_cufft_clear_plan_cache", + "_cufft_get_plan_cache_max_size", + "_cufft_get_plan_cache_size", + "_cufft_set_plan_cache_max_size", + "_cummax_helper", + "_cummin_helper", + "_debug_has_internal_overlap", + "_dim_arange", + "_dirichlet_grad", + "_disable_functionalization", + "_dyn_quant_matmul_4bit", + "_dyn_quant_pack_4bit_weight", + "_efficientzerotensor", + "_embedding_bag", + "_embedding_bag_forward_only", + "_empty_affine_quantized", + "_empty_per_channel_affine_quantized", + "_enable_functionalization", + "_euclidean_dist", + "_fake_quantize_learnable_per_channel_affine", + "_fake_quantize_learnable_per_tensor_affine", + "_fake_quantize_per_tensor_affine_cachemask_tensor_qparams", + "_fft_c2c", + "_fft_c2r", + "_fft_r2c", + "_fill_mem_eff_dropout_mask_", + "_foobar", + "_foreach_abs", + "_foreach_abs_", + "_foreach_acos", + "_foreach_acos_", + "_foreach_add", + "_foreach_add_", + "_foreach_addcdiv", + "_foreach_addcdiv_", + "_foreach_addcmul", + "_foreach_addcmul_", + "_foreach_asin", + "_foreach_asin_", + "_foreach_atan", + "_foreach_atan_", + "_foreach_ceil", + "_foreach_ceil_", + "_foreach_clamp_max", + "_foreach_clamp_max_", + "_foreach_clamp_min", + "_foreach_clamp_min_", + "_foreach_copy_", + "_foreach_cos", + "_foreach_cos_", + "_foreach_cosh", + "_foreach_cosh_", + "_foreach_div", + "_foreach_div_", + "_foreach_erf", + "_foreach_erf_", + "_foreach_erfc", + "_foreach_erfc_", + "_foreach_exp", + "_foreach_exp_", + "_foreach_expm1", + "_foreach_expm1_", + "_foreach_floor", + "_foreach_floor_", + "_foreach_frac", + "_foreach_frac_", + "_foreach_lerp", + "_foreach_lerp_", + "_foreach_lgamma", + "_foreach_lgamma_", + "_foreach_log", + "_foreach_log10", + "_foreach_log10_", + "_foreach_log1p", + "_foreach_log1p_", + "_foreach_log2", + "_foreach_log2_", + "_foreach_log_", + "_foreach_max", + "_foreach_maximum", + "_foreach_maximum_", + "_foreach_minimum", + "_foreach_minimum_", + "_foreach_mul", + "_foreach_mul_", + "_foreach_neg", + "_foreach_neg_", + "_foreach_norm", + "_foreach_pow", + "_foreach_pow_", + "_foreach_reciprocal", + "_foreach_reciprocal_", + "_foreach_round", + "_foreach_round_", + "_foreach_rsqrt", + "_foreach_rsqrt_", + "_foreach_sigmoid", + "_foreach_sigmoid_", + "_foreach_sign", + "_foreach_sign_", + "_foreach_sin", + "_foreach_sin_", + "_foreach_sinh", + "_foreach_sinh_", + "_foreach_sqrt", + "_foreach_sqrt_", + "_foreach_sub", + "_foreach_sub_", + "_foreach_tan", + "_foreach_tan_", + "_foreach_tanh", + "_foreach_tanh_", + "_foreach_trunc", + "_foreach_trunc_", + "_foreach_zero_", + "_from_functional_tensor", + "_functional_assert_async", + "_functional_assert_scalar", + "_functional_sym_constrain_range", + "_functional_sym_constrain_range_for_size", + "_functionalize_apply_view_metas", + "_functionalize_are_all_mutations_hidden_from_autograd", + "_functionalize_are_all_mutations_under_no_grad_or_inference_mode", + "_functionalize_commit_update", + "_functionalize_has_metadata_mutation", + "_functionalize_inductor_storage_resized_counter", + "_functionalize_is_symbolic", + "_functionalize_mark_mutation_hidden_from_autograd", + "_functionalize_mark_storage_changed", + "_functionalize_mutation_counter", + "_functionalize_replace", + "_functionalize_storage_changed_counter", + "_functionalize_sync", + "_functionalize_unsafe_set", + "_functionalize_was_inductor_storage_resized", + "_functionalize_was_storage_changed", + "_fused_adagrad_", + "_fused_adam_", + "_fused_adamw_", + "_fused_dropout", + "_fused_moving_avg_obs_fq_helper", + "_fused_rms_norm", + "_fused_sdp_choice", + "_fused_sgd_", + "_fw_primal_copy", + "_grid_sampler_2d_cpu_fallback", + "_grouped_mm", + "_has_compatible_shallow_copy_type", + "_histogramdd_bin_edges", + "_histogramdd_from_bin_cts", + "_histogramdd_from_bin_tensors", + "_index_put_impl_", + "_indices_copy", + "_int_mm", + "_is_all_true", + "_is_any_true", + "_is_functional_tensor", + "_is_functional_tensor_base", + "_is_zerotensor", + "_lazy_clone", + "_linalg_check_errors", + "_linalg_det", + "_linalg_eigh", + "_linalg_slogdet", + "_linalg_solve_ex", + "_linalg_svd", + "_log_softmax", + "_log_softmax_backward_data", + "_logcumsumexp", + "_lstm_mps", + "_lu_with_info", + "_make_dep_token", + "_make_dual", + "_make_dual_copy", + "_make_per_channel_quantized_tensor", + "_make_per_tensor_quantized_tensor", + "_masked_scale", + "_masked_softmax", + "_mixed_dtypes_linear", + "_mkldnn_reshape", + "_mkldnn_transpose", + "_mkldnn_transpose_", + "_mps_convolution", + "_mps_convolution_transpose", + "_native_batch_norm_legit", + "_native_batch_norm_legit_no_training", + "_native_multi_head_attention", + "_neg_view", + "_neg_view_copy", + "_nested_compute_contiguous_strides_offsets", + "_nested_from_padded", + "_nested_from_padded_and_nested_example", + "_nested_from_padded_tensor", + "_nested_get_jagged_dummy", + "_nested_get_lengths", + "_nested_get_max_seqlen", + "_nested_get_min_seqlen", + "_nested_get_offsets", + "_nested_get_ragged_idx", + "_nested_get_values", + "_nested_get_values_copy", + "_nested_tensor_from_mask", + "_nested_tensor_from_mask_left_aligned", + "_nested_tensor_from_tensor_list", + "_nested_tensor_softmax_with_shape", + "_nested_view_from_buffer", + "_nested_view_from_buffer_copy", + "_nested_view_from_jagged", + "_nested_view_from_jagged_copy", + "_nnpack_available", + "_nnpack_spatial_convolution", + "_pack_padded_sequence", + "_pad_packed_sequence", + "_pin_memory", + "_prelu_kernel", + "_print", + "_propagate_xla_data", + "_remove_batch_dim", + "_reshape_alias_copy", + "_reshape_from_tensor", + "_resize_output_", + "_rowwise_prune", + "_safe_softmax", + "_sample_dirichlet", + "_saturate_weight_to_fp16", + "_scaled_dot_product_attention_math", + "_scaled_dot_product_attention_math_for_mps", + "_scaled_dot_product_cudnn_attention", + "_scaled_dot_product_efficient_attention", + "_scaled_dot_product_flash_attention", + "_scaled_dot_product_flash_attention_for_cpu", + "_scaled_grouped_mm", + "_scaled_mm", + "_shape_as_tensor", + "_sobol_engine_draw", + "_sobol_engine_ff_", + "_sobol_engine_initialize_state_", + "_sobol_engine_scramble_", + "_softmax", + "_softmax_backward_data", + "_sparse_broadcast_to", + "_sparse_broadcast_to_copy", + "_sparse_csr_prod", + "_sparse_csr_sum", + "_sparse_log_softmax_backward_data", + "_sparse_semi_structured_addmm", + "_sparse_semi_structured_apply", + "_sparse_semi_structured_apply_dense", + "_sparse_semi_structured_linear", + "_sparse_semi_structured_mm", + "_sparse_semi_structured_tile", + "_sparse_softmax_backward_data", + "_sparse_sparse_matmul", + "_sparse_sum", + "_stack", + "_standard_gamma", + "_standard_gamma_grad", + "_sync", + "_test_autograd_multiple_dispatch", + "_test_autograd_multiple_dispatch_view", + "_test_autograd_multiple_dispatch_view_copy", + "_test_check_tensor", + "_test_functorch_fallback", + "_test_parallel_materialize", + "_test_serialization_subcmul", + "_to_cpu", + "_to_functional_tensor", + "_to_sparse_semi_structured", + "_transform_bias_rescale_qkv", + "_transformer_encoder_layer_fwd", + "_trilinear", + "_triton_multi_head_attention", + "_triton_scaled_dot_attention", + "_unique", + "_unique2", + "_unpack_dual", + "_unsafe_index", + "_unsafe_index_put", + "_unsafe_masked_index", + "_unsafe_masked_index_put_accumulate", + "_use_cudnn_ctc_loss", + "_use_cudnn_rnn_flatten_weight", + "_validate_compressed_sparse_indices", + "_validate_sparse_bsc_tensor_args", + "_validate_sparse_bsr_tensor_args", + "_validate_sparse_compressed_tensor_args", + "_validate_sparse_coo_tensor_args", + "_validate_sparse_csc_tensor_args", + "_validate_sparse_csr_tensor_args", + "_values_copy", + "_weight_int4pack_mm", + "_weight_int4pack_mm_for_cpu", + "_weight_int4pack_mm_with_scales_and_zeros", + "_weight_int8pack_mm", + "_weight_norm", + "_weight_norm_interface", + "_wrapped_linear_prepack", + "_wrapped_quantized_linear_prepacked", + "abs", + "abs_", + "absolute", + "acos", + "acos_", + "acosh", + "acosh_", + "adaptive_avg_pool1d", + "adaptive_max_pool1d", + "add", + "addbmm", + "addcdiv", + "addcmul", + "addmm", + "addmv", + "addmv_", + "addr", + "adjoint", + "affine_grid_generator", + "alias_copy", + "all", + "allclose", + "alpha_dropout", + "alpha_dropout_", + "amax", + "amin", + "aminmax", + "angle", + "any", + "arange", + "arccos", + "arccos_", + "arccosh", + "arccosh_", + "arcsin", + "arcsin_", + "arcsinh", + "arcsinh_", + "arctan", + "arctan2", + "arctan_", + "arctanh", + "arctanh_", + "argmax", + "argmin", + "argsort", + "argwhere", + "as_strided", + "as_strided_", + "as_strided_copy", + "as_strided_scatter", + "as_tensor", + "asarray", + "asin", + "asin_", + "asinh", + "asinh_", + "atan", + "atan2", + "atan_", + "atanh", + "atanh_", + "avg_pool1d", + "baddbmm", + "bartlett_window", + "batch_norm", + "batch_norm_backward_elemt", + "batch_norm_backward_reduce", + "batch_norm_elemt", + "batch_norm_gather_stats", + "batch_norm_gather_stats_with_counts", + "batch_norm_stats", + "batch_norm_update_stats", + "bernoulli", + "bilinear", + "binary_cross_entropy_with_logits", + "bincount", + "binomial", + "bitwise_and", + "bitwise_left_shift", + "bitwise_not", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "blackman_window", + "bmm", + "broadcast_to", + "bucketize", + "can_cast", + "cat", + "ccol_indices_copy", + "ceil", + "ceil_", + "celu", + "celu_", + "channel_shuffle", + "cholesky", + "cholesky_inverse", + "cholesky_solve", + "choose_qparams_optimized", + "chunk", + "clamp", + "clamp_", + "clamp_max", + "clamp_max_", + "clamp_min", + "clamp_min_", + "clip", + "clip_", + "clone", + "col_indices_copy", + "column_stack", + "combinations", + "complex", + "concat", + "concatenate", + "conj", + "conj_physical", + "conj_physical_", + "constant_pad_nd", + "conv1d", + "conv2d", + "conv3d", + "conv_tbc", + "conv_transpose1d", + "conv_transpose2d", + "conv_transpose3d", + "convolution", + "copysign", + "corrcoef", + "cos", + "cos_", + "cosh", + "cosh_", + "cosine_embedding_loss", + "cosine_similarity", + "count_nonzero", + "cov", + "cross", + "crow_indices_copy", + "ctc_loss", + "cudnn_affine_grid_generator", + "cudnn_batch_norm", + "cudnn_convolution", + "cudnn_convolution_add_relu", + "cudnn_convolution_relu", + "cudnn_convolution_transpose", + "cudnn_grid_sampler", + "cudnn_is_acceptable", + "cummax", + "cummin", + "cumprod", + "cumsum", + "cumulative_trapezoid", + "deg2rad", + "deg2rad_", + "dequantize", + "det", + "detach", + "detach_", + "detach_copy", + "diag", + "diag_embed", + "diagflat", + "diagonal", + "diagonal_copy", + "diagonal_scatter", + "diff", + "digamma", + "dist", + "div", + "divide", + "dot", + "dropout", + "dropout_", + "dsmm", + "dsplit", + "dstack", + "embedding", + "embedding_bag", + "embedding_renorm_", + "empty", + "empty_like", + "empty_permuted", + "empty_quantized", + "empty_strided", + "eq", + "equal", + "erf", + "erf_", + "erfc", + "erfc_", + "erfinv", + "exp", + "exp2", + "exp2_", + "exp_", + "expand_copy", + "expm1", + "expm1_", + "eye", + "fake_quantize_per_channel_affine", + "fake_quantize_per_tensor_affine", + "fbgemm_linear_fp16_weight", + "fbgemm_linear_fp16_weight_fp32_activation", + "fbgemm_linear_int8_weight", + "fbgemm_linear_int8_weight_fp32_activation", + "fbgemm_linear_quantize_weight", + "fbgemm_pack_gemm_matrix_fp16", + "fbgemm_pack_quantized_matrix", + "feature_alpha_dropout", + "feature_alpha_dropout_", + "feature_dropout", + "feature_dropout_", + "fill", + "fill_", + "fix", + "fix_", + "flatten", + "flip", + "fliplr", + "flipud", + "float_power", + "floor", + "floor_", + "floor_divide", + "fmax", + "fmin", + "fmod", + "frac", + "frac_", + "frexp", + "frobenius_norm", + "from_file", + "from_numpy", + "frombuffer", + "full", + "full_like", + "fused_moving_avg_obs_fake_quant", + "gather", + "gcd", + "gcd_", + "ge", + "geqrf", + "ger", + "get_default_dtype", + "get_num_interop_threads", + "get_num_threads", + "gradient", + "greater", + "greater_equal", + "grid_sampler", + "grid_sampler_2d", + "grid_sampler_3d", + "group_norm", + "gru", + "gru_cell", + "gt", + "hamming_window", + "hann_window", + "hardshrink", + "hash_tensor", + "heaviside", + "hinge_embedding_loss", + "histc", + "histogram", + "histogramdd", + "hsmm", + "hsplit", + "hspmm", + "hstack", + "hypot", + "i0", + "i0_", + "igamma", + "igammac", + "imag", + "index_add", + "index_copy", + "index_fill", + "index_put", + "index_put_", + "index_reduce", + "index_select", + "indices_copy", + "init_num_threads", + "inner", + "instance_norm", + "int_repr", + "inverse", + "is_complex", + "is_conj", + "is_distributed", + "is_floating_point", + "is_grad_enabled", + "is_inference", + "is_inference_mode_enabled", + "is_neg", + "is_nonzero", + "is_same_size", + "is_signed", + "is_vulkan_available", + "isclose", + "isfinite", + "isin", + "isinf", + "isnan", + "isneginf", + "isposinf", + "isreal", + "istft", + "kaiser_window", + "kl_div", + "kron", + "kthvalue", + "layer_norm", + "lcm", + "lcm_", + "ldexp", + "ldexp_", + "le", + "lerp", + "less", + "less_equal", + "lgamma", + "linspace", + "log", + "log10", + "log10_", + "log1p", + "log1p_", + "log2", + "log2_", + "log_", + "log_softmax", + "logaddexp", + "logaddexp2", + "logcumsumexp", + "logdet", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "logit", + "logit_", + "logspace", + "logsumexp", + "lstm", + "lstm_cell", + "lt", + "lu_solve", + "lu_unpack", + "margin_ranking_loss", + "masked_fill", + "masked_scatter", + "masked_select", + "matmul", + "matrix_exp", + "matrix_power", + "max", + "max_pool1d", + "max_pool1d_with_indices", + "max_pool2d", + "max_pool3d", + "maximum", + "mean", + "median", + "min", + "minimum", + "miopen_batch_norm", + "miopen_convolution", + "miopen_convolution_add_relu", + "miopen_convolution_relu", + "miopen_convolution_transpose", + "miopen_depthwise_convolution", + "miopen_rnn", + "mkldnn_adaptive_avg_pool2d", + "mkldnn_convolution", + "mkldnn_linear_backward_weights", + "mkldnn_max_pool2d", + "mkldnn_max_pool3d", + "mkldnn_rnn_layer", + "mm", + "mode", + "moveaxis", + "movedim", + "msort", + "mul", + "multinomial", + "multiply", + "mv", + "mvlgamma", + "nan_to_num", + "nan_to_num_", + "nanmean", + "nanmedian", + "nanquantile", + "nansum", + "narrow", + "narrow_copy", + "native_batch_norm", + "native_channel_shuffle", + "native_dropout", + "native_group_norm", + "native_layer_norm", + "native_norm", + "ne", + "neg", + "neg_", + "negative", + "negative_", + "nextafter", + "nonzero", + "nonzero_static", + "norm_except_dim", + "normal", + "not_equal", + "nuclear_norm", + "numel", + "ones", + "ones_like", + "orgqr", + "ormqr", + "outer", + "pairwise_distance", + "pdist", + "permute", + "permute_copy", + "pinverse", + "pixel_shuffle", + "pixel_unshuffle", + "poisson", + "poisson_nll_loss", + "polar", + "polygamma", + "positive", + "pow", + "prelu", + "prod", + "promote_types", + "put", + "q_per_channel_axis", + "q_per_channel_scales", + "q_per_channel_zero_points", + "q_scale", + "q_zero_point", + "qr", + "quantile", + "quantize_per_channel", + "quantize_per_tensor", + "quantize_per_tensor_dynamic", + "quantized_batch_norm", + "quantized_gru_cell", + "quantized_lstm_cell", + "quantized_max_pool1d", + "quantized_max_pool2d", + "quantized_max_pool3d", + "quantized_rnn_relu_cell", + "quantized_rnn_tanh_cell", + "rad2deg", + "rad2deg_", + "rand", + "rand_like", + "randint", + "randint_like", + "randn", + "randn_like", + "randperm", + "range", + "ravel", + "real", + "reciprocal", + "reciprocal_", + "relu", + "relu_", + "remainder", + "renorm", + "repeat_interleave", + "reshape", + "resize_as_", + "resize_as_sparse_", + "resolve_conj", + "resolve_neg", + "result_type", + "rms_norm", + "rnn_relu", + "rnn_relu_cell", + "rnn_tanh", + "rnn_tanh_cell", + "roll", + "rot90", + "round", + "round_", + "row_indices_copy", + "row_stack", + "rrelu", + "rrelu_", + "rsqrt", + "rsqrt_", + "rsub", + "saddmm", + "scalar_tensor", + "scatter", + "scatter_add", + "scatter_reduce", + "searchsorted", + "segment_reduce", + "select", + "select_copy", + "select_scatter", + "selu", + "selu_", + "set_flush_denormal", + "set_num_interop_threads", + "set_num_threads", + "sgn", + "sigmoid", + "sigmoid_", + "sign", + "signbit", + "sin", + "sin_", + "sinc", + "sinc_", + "sinh", + "sinh_", + "slice_copy", + "slice_inverse", + "slice_scatter", + "slogdet", + "smm", + "softmax", + "sort", + "sparse_bsc_tensor", + "sparse_bsr_tensor", + "sparse_compressed_tensor", + "sparse_coo_tensor", + "sparse_csc_tensor", + "sparse_csr_tensor", + "split_copy", + "split_with_sizes", + "split_with_sizes_copy", + "spmm", + "sqrt", + "sqrt_", + "square", + "square_", + "squeeze", + "squeeze_copy", + "sspaddmm", + "stack", + "std", + "std_mean", + "sub", + "subtract", + "sum", + "svd", + "swapaxes", + "swapdims", + "sym_constrain_range", + "sym_constrain_range_for_size", + "t", + "t_copy", + "take", + "take_along_dim", + "tan", + "tan_", + "tanh", + "tanh_", + "tensor", + "tensor_split", + "threshold", + "threshold_", + "tile", + "topk", + "trace", + "transpose", + "transpose_copy", + "trapezoid", + "trapz", + "triangular_solve", + "tril", + "tril_indices", + "triplet_margin_loss", + "triu", + "triu_indices", + "true_divide", + "trunc", + "trunc_", + "unbind", + "unbind_copy", + "unflatten", + "unfold_copy", + "unique_dim", + "unsafe_chunk", + "unsafe_split", + "unsafe_split_with_sizes", + "unsqueeze", + "unsqueeze_copy", + "values_copy", + "vander", + "var", + "var_mean", + "vdot", + "view_as_complex", + "view_as_complex_copy", + "view_as_real", + "view_as_real_copy", + "view_copy", + "vsplit", + "vstack", + "where", + "xlogy", + "xlogy_", + "zero_", + "zeros", + "zeros_like", +] + +@overload +def __and__(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def __and__(input: Tensor, other: Number | _complex) -> Tensor: ... +@overload +def __lshift__(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def __lshift__(input: Tensor, other: Number | _complex) -> Tensor: ... +@overload +def __or__(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def __or__(input: Tensor, other: Number | _complex) -> Tensor: ... +@overload +def __rshift__(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def __rshift__(input: Tensor, other: Number | _complex) -> Tensor: ... +@overload +def __xor__(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def __xor__(input: Tensor, other: Number | _complex) -> Tensor: ... +def _adaptive_avg_pool2d( + input: Tensor, + output_size: _int | SymInt | Sequence[_int | SymInt], +) -> Tensor: ... +def _adaptive_avg_pool3d( + input: Tensor, + output_size: _int | SymInt | Sequence[_int | SymInt], +) -> Tensor: ... +def _add_batch_dim(input: Tensor, batch_dim: _int, level: _int) -> Tensor: ... +@overload +def _add_relu( + input: Tensor, + other: Tensor, + *, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def _add_relu( + input: Tensor, + other: Number | _complex, + alpha: Number | _complex = 1, +) -> Tensor: ... +@overload +def _add_relu_( + input: Tensor, + other: Tensor, + *, + alpha: Number | _complex = 1, +) -> Tensor: ... +@overload +def _add_relu_( + input: Tensor, + other: Number | _complex, + alpha: Number | _complex = 1, +) -> Tensor: ... +def _addmm_activation( + input: Tensor, + mat1: Tensor, + mat2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + use_gelu: _bool = False, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def _aminmax(input: Tensor) -> tuple[Tensor, Tensor]: ... +@overload +def _aminmax( + input: Tensor, + dim: _int, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: ... +def _amp_foreach_non_finite_check_and_unscale_( + self: tuple[Tensor, ...] | list[Tensor] | None, + found_inf: Tensor, + inv_scale: Tensor, +) -> None: ... +def _amp_update_scale_( + input: Tensor, + growth_tracker: Tensor, + found_inf: Tensor, + scale_growth_factor: _float, + scale_backoff_factor: _float, + growth_interval: _int, +) -> Tensor: ... +@overload +def _assert_async(input: Tensor) -> None: + r""" + _assert_async(tensor) -> void + + Asynchronously assert that the contents of tensor are nonzero. For CPU tensors, + this is equivalent to ``assert tensor`` or ``assert tensor.is_nonzero()``; for + CUDA tensors, we DO NOT synchronize and you may only find out the assertion + failed at a later CUDA kernel launch. Asynchronous assertion can be helpful for + testing invariants in CUDA tensors without giving up performance. This function + is NOT intended to be used for regular error checking, as it will trash your CUDA + context if the assert fails (forcing you to restart your PyTorch process.) + + Args: + tensor (Tensor): a one element tensor to test to see if it is nonzero. Zero + elements (including False for boolean tensors) cause an assertion failure + to be raised. + """ + +@overload +def _assert_async(input: Tensor, assert_msg: str) -> None: + r""" + _assert_async(tensor) -> void + + Asynchronously assert that the contents of tensor are nonzero. For CPU tensors, + this is equivalent to ``assert tensor`` or ``assert tensor.is_nonzero()``; for + CUDA tensors, we DO NOT synchronize and you may only find out the assertion + failed at a later CUDA kernel launch. Asynchronous assertion can be helpful for + testing invariants in CUDA tensors without giving up performance. This function + is NOT intended to be used for regular error checking, as it will trash your CUDA + context if the assert fails (forcing you to restart your PyTorch process.) + + Args: + tensor (Tensor): a one element tensor to test to see if it is nonzero. Zero + elements (including False for boolean tensors) cause an assertion failure + to be raised. + """ + +def _assert_scalar(self: Number | _complex, assert_msg: str) -> None: ... +def _assert_tensor_metadata( + a: Tensor, + size: Sequence[_int | SymInt] | None = None, + stride: Sequence[_int | SymInt] | None = None, + dtype: _dtype | None = None, + *, + device: DeviceLikeType | None = None, + layout: _layout | None = None, +) -> None: ... +def _batch_norm_impl_index( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + training: _bool, + momentum: _float, + eps: _float, + cudnn_enabled: _bool, +) -> tuple[Tensor, Tensor, Tensor, Tensor, _int]: ... +def _cast_Byte(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Char(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Double(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Float(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Half(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Int(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Long(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _cast_Short(input: Tensor, non_blocking: _bool = False) -> Tensor: ... +def _choose_qparams_per_tensor( + input: Tensor, + reduce_range: _bool = False, +) -> tuple[_float, _int]: ... +def _chunk_cat( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int, + num_chunks: _int, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _coalesce(input: Tensor) -> Tensor: ... +def _compute_linear_combination( + input: Tensor, + coefficients: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _conj(input: Tensor) -> Tensor: ... +def _conj_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: ... +def _conj_physical(input: Tensor) -> Tensor: ... +def _convert_indices_from_coo_to_csr( + input: Tensor, + size: _int, + *, + out_int32: _bool = False, + out: Tensor | None = None, +) -> Tensor: ... +def _convert_indices_from_csr_to_coo( + crow_indices: Tensor, + col_indices: Tensor, + *, + out_int32: _bool = False, + transpose: _bool = False, + out: Tensor | None = None, +) -> Tensor: ... +def _convert_weight_to_int4pack(input: Tensor, innerKTiles: _int) -> Tensor: ... +def _convert_weight_to_int4pack_for_cpu( + input: Tensor, + innerKTiles: _int, +) -> Tensor: ... +@overload +def _convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + transposed: _bool, + output_padding: _size, + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, + cudnn_enabled: _bool, +) -> Tensor: ... +@overload +def _convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + transposed: _bool, + output_padding: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, + cudnn_enabled: _bool, + allow_tf32: _bool, +) -> Tensor: ... +def _convolution_mode( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: str, + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def _copy_from( + input: Tensor, + dst: Tensor, + non_blocking: _bool = False, +) -> Tensor: ... +def _copy_from_and_resize(input: Tensor, dst: Tensor) -> Tensor: ... +def _cslt_compress(input: Tensor) -> Tensor: ... +def _cslt_sparse_mm( + compressed_A: Tensor, + dense_B: Tensor, + bias: Tensor | None = None, + alpha: Tensor | None = None, + out_dtype: _dtype | None = None, + transpose_result: _bool = False, + alg_id: _int = 0, + split_k: _int = 1, + split_k_mode: _int = -1, +) -> Tensor: ... +def _cslt_sparse_mm_search( + compressed_A: Tensor, + dense_B: Tensor, + bias: Tensor | None = None, + alpha: Tensor | None = None, + out_dtype: _dtype | None = None, + transpose_result: _bool = False, +) -> _int: ... +@overload +def _ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: _size, + target_lengths: _size, + blank: _int = 0, + zero_infinity: _bool = False, +) -> tuple[Tensor, Tensor]: ... +@overload +def _ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: Tensor, + target_lengths: Tensor, + blank: _int = 0, + zero_infinity: _bool = False, +) -> tuple[Tensor, Tensor]: ... +@overload +def _cudnn_ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: _size, + target_lengths: _size, + blank: _int, + deterministic: _bool, + zero_infinity: _bool, +) -> tuple[Tensor, Tensor]: ... +@overload +def _cudnn_ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: Tensor, + target_lengths: Tensor, + blank: _int, + deterministic: _bool, + zero_infinity: _bool, +) -> tuple[Tensor, Tensor]: ... +def _cudnn_init_dropout_state( + dropout: _float, + train: _bool, + dropout_seed: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +def _cudnn_rnn( + input: Tensor, + weight: tuple[Tensor, ...] | list[Tensor] | None, + weight_stride0: _int, + weight_buf: Tensor | None, + hx: Tensor, + cx: Tensor | None, + mode: _int, + hidden_size: _int | SymInt, + proj_size: _int | SymInt, + num_layers: _int, + batch_first: _bool, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_sizes: Sequence[_int | SymInt], + dropout_state: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: ... +def _cudnn_rnn_flatten_weight( + weight_arr: tuple[Tensor, ...] | list[Tensor] | None, + weight_stride0: _int, + input_size: _int | SymInt, + mode: _int, + hidden_size: _int | SymInt, + proj_size: _int | SymInt, + num_layers: _int, + batch_first: _bool, + bidirectional: _bool, +) -> Tensor: ... +def _cufft_clear_plan_cache(device_index: _int) -> None: ... +def _cufft_get_plan_cache_max_size(device_index: _int) -> _int: ... +def _cufft_get_plan_cache_size(device_index: _int) -> _int: ... +def _cufft_set_plan_cache_max_size( + device_index: _int, + max_size: _int, +) -> None: ... +def _cummax_helper( + input: Tensor, + values: Tensor, + indices: Tensor, + dim: _int, +) -> None: ... +def _cummin_helper( + input: Tensor, + values: Tensor, + indices: Tensor, + dim: _int, +) -> None: ... +def _debug_has_internal_overlap(input: Tensor) -> _int: ... +def _dim_arange(like: Tensor, dim: _int) -> Tensor: ... +def _dirichlet_grad(x: Tensor, alpha: Tensor, total: Tensor) -> Tensor: ... +def _disable_functionalization(): ... +def _dyn_quant_matmul_4bit( + inp: Tensor, + packed_weights: Tensor, + block_size: _int, + in_features: _int, + out_features: _int, +) -> Tensor: ... +def _dyn_quant_pack_4bit_weight( + weights: Tensor, + scales_zeros: Tensor, + bias: Tensor | None, + block_size: _int, + in_features: _int, + out_features: _int, +) -> Tensor: ... +@overload +def _efficientzerotensor( + size: Sequence[_int | SymInt], + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +@overload +def _efficientzerotensor( + *size: _int | SymInt, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +def _embedding_bag( + weight: Tensor, + indices: Tensor, + offsets: Tensor, + scale_grad_by_freq: _bool = False, + mode: _int = 0, + sparse: _bool = False, + per_sample_weights: Tensor | None = None, + include_last_offset: _bool = False, + padding_idx: _int = -1, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +def _embedding_bag_forward_only( + weight: Tensor, + indices: Tensor, + offsets: Tensor, + scale_grad_by_freq: _bool = False, + mode: _int = 0, + sparse: _bool = False, + per_sample_weights: Tensor | None = None, + include_last_offset: _bool = False, + padding_idx: _int = -1, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +@overload +def _empty_affine_quantized( + size: Sequence[_int | SymInt], + *, + scale: _float = 1, + zero_point: _int = 0, + memory_format: memory_format | None = contiguous_format, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +@overload +def _empty_affine_quantized( + *size: _int | SymInt, + scale: _float = 1, + zero_point: _int = 0, + memory_format: memory_format | None = contiguous_format, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +@overload +def _empty_per_channel_affine_quantized( + size: Sequence[_int | SymInt], + *, + scales: Tensor, + zero_points: Tensor, + axis: _int, + memory_format: memory_format | None = contiguous_format, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +@overload +def _empty_per_channel_affine_quantized( + *size: _int | SymInt, + scales: Tensor, + zero_points: Tensor, + axis: _int, + memory_format: memory_format | None = contiguous_format, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +def _enable_functionalization(*, reapply_views: _bool = False) -> None: ... +def _euclidean_dist(x1: Tensor, x2: Tensor) -> Tensor: ... +def _fake_quantize_learnable_per_channel_affine( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + axis: _int, + quant_min: _int, + quant_max: _int, + grad_factor: _float = 1.0, +) -> Tensor: ... +def _fake_quantize_learnable_per_tensor_affine( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + quant_min: _int, + quant_max: _int, + grad_factor: _float = 1.0, +) -> Tensor: ... +def _fake_quantize_per_tensor_affine_cachemask_tensor_qparams( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + fake_quant_enabled: Tensor, + quant_min: _int, + quant_max: _int, +) -> torch.return_types._fake_quantize_per_tensor_affine_cachemask_tensor_qparams: # fmt: skip + ... +def _fft_c2c( + input: Tensor, + dim: Sequence[_int | SymInt], + normalization: _int, + forward: _bool, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _fft_c2r( + input: Tensor, + dim: _size, + normalization: _int, + last_dim_size: _int | SymInt, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _fft_r2c( + input: Tensor, + dim: _size, + normalization: _int, + onesided: _bool, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _fill_mem_eff_dropout_mask_( + input: Tensor, + dropout_p: _float, + seed: _int, + offset: _int, +) -> Tensor: ... +def _foobar( + input: Tensor, + arg1: _bool = True, + arg2: _bool = True, + *, + arg3: _bool = True, +) -> Tensor: ... +def _foreach_abs( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_abs(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.abs` to each Tensor of the input list. + """ + +def _foreach_abs_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_abs_(self: List[Tensor]) -> None + + Apply :func:`torch.abs` to each Tensor of the input list. + """ + +def _foreach_acos( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_acos(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.acos` to each Tensor of the input list. + """ + +def _foreach_acos_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_acos_(self: List[Tensor]) -> None + + Apply :func:`torch.acos` to each Tensor of the input list. + """ + +@overload +def _foreach_add( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_add( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, + *, + alpha: Number | _complex = 1, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_add( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, + *, + alpha: Number | _complex = 1, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_add( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_add_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_add_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, + *, + alpha: Number | _complex = 1, +) -> None: ... +@overload +def _foreach_add_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, + *, + alpha: Number | _complex = 1, +) -> None: ... +@overload +def _foreach_add_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_addcdiv( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcdiv( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Tensor, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcdiv( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + value: Number | _complex = 1, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcdiv_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_addcdiv_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Tensor, +) -> None: ... +@overload +def _foreach_addcdiv_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + value: Number | _complex = 1, +) -> None: ... +@overload +def _foreach_addcmul( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcmul( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Tensor, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcmul( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + value: Number | _complex = 1, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_addcmul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_addcmul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Tensor, +) -> None: ... +@overload +def _foreach_addcmul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensor1: tuple[Tensor, ...] | list[Tensor] | None, + tensor2: tuple[Tensor, ...] | list[Tensor] | None, + value: Number | _complex = 1, +) -> None: ... +def _foreach_asin( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_asin(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.asin` to each Tensor of the input list. + """ + +def _foreach_asin_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_asin_(self: List[Tensor]) -> None + + Apply :func:`torch.asin` to each Tensor of the input list. + """ + +def _foreach_atan( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_atan(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.atan` to each Tensor of the input list. + """ + +def _foreach_atan_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_atan_(self: List[Tensor]) -> None + + Apply :func:`torch.atan` to each Tensor of the input list. + """ + +def _foreach_ceil( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_ceil(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.ceil` to each Tensor of the input list. + """ + +def _foreach_ceil_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_ceil_(self: List[Tensor]) -> None + + Apply :func:`torch.ceil` to each Tensor of the input list. + """ + +@overload +def _foreach_clamp_max( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_max( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_max( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_max_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_clamp_max_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_clamp_max_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +@overload +def _foreach_clamp_min( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_min( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_min( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_clamp_min_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_clamp_min_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_clamp_min_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +def _foreach_copy_( + self: tuple[Tensor, ...] | list[Tensor] | None, + src: tuple[Tensor, ...] | list[Tensor] | None, + non_blocking: _bool = False, +) -> None: ... +def _foreach_cos( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_cos(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.cos` to each Tensor of the input list. + """ + +def _foreach_cos_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_cos_(self: List[Tensor]) -> None + + Apply :func:`torch.cos` to each Tensor of the input list. + """ + +def _foreach_cosh( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_cosh(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.cosh` to each Tensor of the input list. + """ + +def _foreach_cosh_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_cosh_(self: List[Tensor]) -> None + + Apply :func:`torch.cosh` to each Tensor of the input list. + """ + +@overload +def _foreach_div( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_div( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_div( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_div( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_div_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_div_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, +) -> None: ... +@overload +def _foreach_div_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_div_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +def _foreach_erf( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_erf(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.erf` to each Tensor of the input list. + """ + +def _foreach_erf_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_erf_(self: List[Tensor]) -> None + + Apply :func:`torch.erf` to each Tensor of the input list. + """ + +def _foreach_erfc( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_erfc(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.erfc` to each Tensor of the input list. + """ + +def _foreach_erfc_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_erfc_(self: List[Tensor]) -> None + + Apply :func:`torch.erfc` to each Tensor of the input list. + """ + +def _foreach_exp( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_exp(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.exp` to each Tensor of the input list. + """ + +def _foreach_exp_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_exp_(self: List[Tensor]) -> None + + Apply :func:`torch.exp` to each Tensor of the input list. + """ + +def _foreach_expm1( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_expm1(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.expm1` to each Tensor of the input list. + """ + +def _foreach_expm1_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_expm1_(self: List[Tensor]) -> None + + Apply :func:`torch.expm1` to each Tensor of the input list. + """ + +def _foreach_floor( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_floor(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.floor` to each Tensor of the input list. + """ + +def _foreach_floor_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_floor_(self: List[Tensor]) -> None + + Apply :func:`torch.floor` to each Tensor of the input list. + """ + +def _foreach_frac( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_frac(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.frac` to each Tensor of the input list. + """ + +def _foreach_frac_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_frac_(self: List[Tensor]) -> None + + Apply :func:`torch.frac` to each Tensor of the input list. + """ + +@overload +def _foreach_lerp( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weight: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_lerp( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weight: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_lerp( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weights: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_lerp_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weight: Number | _complex, +) -> None: ... +@overload +def _foreach_lerp_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weight: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_lerp_( + self: tuple[Tensor, ...] | list[Tensor] | None, + tensors1: tuple[Tensor, ...] | list[Tensor] | None, + weights: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +def _foreach_lgamma( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_lgamma(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.lgamma` to each Tensor of the input list. + """ + +def _foreach_lgamma_( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: + r""" + _foreach_lgamma_(self: List[Tensor]) -> None + + Apply :func:`torch.lgamma` to each Tensor of the input list. + """ + +def _foreach_log( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_log(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.log` to each Tensor of the input list. + """ + +def _foreach_log10( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_log10(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.log10` to each Tensor of the input list. + """ + +def _foreach_log10_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_log10_(self: List[Tensor]) -> None + + Apply :func:`torch.log10` to each Tensor of the input list. + """ + +def _foreach_log1p( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_log1p(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.log1p` to each Tensor of the input list. + """ + +def _foreach_log1p_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_log1p_(self: List[Tensor]) -> None + + Apply :func:`torch.log1p` to each Tensor of the input list. + """ + +def _foreach_log2( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_log2(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.log2` to each Tensor of the input list. + """ + +def _foreach_log2_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_log2_(self: List[Tensor]) -> None + + Apply :func:`torch.log2` to each Tensor of the input list. + """ + +def _foreach_log_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_log_(self: List[Tensor]) -> None + + Apply :func:`torch.log` to each Tensor of the input list. + """ + +def _foreach_max( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_maximum( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_maximum( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_maximum( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_maximum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_maximum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_maximum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +@overload +def _foreach_minimum( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_minimum( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_minimum( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_minimum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_minimum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_minimum_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +@overload +def _foreach_mul( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_mul( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_mul( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_mul( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_mul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_mul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: Tensor, +) -> None: ... +@overload +def _foreach_mul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +@overload +def _foreach_mul_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +def _foreach_neg( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_neg(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.neg` to each Tensor of the input list. + """ + +def _foreach_neg_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_neg_(self: List[Tensor]) -> None + + Apply :func:`torch.neg` to each Tensor of the input list. + """ + +def _foreach_norm( + self: tuple[Tensor, ...] | list[Tensor] | None, + ord: Number | _complex = 2, + dtype: _dtype | None = None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_pow( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_pow( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_pow( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_pow( + self: Number | _complex, + exponent: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_pow_( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_pow_( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: Number | _complex, +) -> None: ... +@overload +def _foreach_pow_( + self: tuple[Tensor, ...] | list[Tensor] | None, + exponent: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: ... +def _foreach_reciprocal( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_reciprocal(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.reciprocal` to each Tensor of the input list. + """ + +def _foreach_reciprocal_( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: + r""" + _foreach_reciprocal_(self: List[Tensor]) -> None + + Apply :func:`torch.reciprocal` to each Tensor of the input list. + """ + +def _foreach_round( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_round(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.round` to each Tensor of the input list. + """ + +def _foreach_round_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_round_(self: List[Tensor]) -> None + + Apply :func:`torch.round` to each Tensor of the input list. + """ + +def _foreach_rsqrt( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +def _foreach_rsqrt_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: ... +def _foreach_sigmoid( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_sigmoid(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.sigmoid` to each Tensor of the input list. + """ + +def _foreach_sigmoid_( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> None: + r""" + _foreach_sigmoid_(self: List[Tensor]) -> None + + Apply :func:`torch.sigmoid` to each Tensor of the input list. + """ + +def _foreach_sign( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +def _foreach_sign_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: ... +def _foreach_sin( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_sin(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.sin` to each Tensor of the input list. + """ + +def _foreach_sin_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_sin_(self: List[Tensor]) -> None + + Apply :func:`torch.sin` to each Tensor of the input list. + """ + +def _foreach_sinh( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_sinh(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.sinh` to each Tensor of the input list. + """ + +def _foreach_sinh_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_sinh_(self: List[Tensor]) -> None + + Apply :func:`torch.sinh` to each Tensor of the input list. + """ + +def _foreach_sqrt( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_sqrt(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.sqrt` to each Tensor of the input list. + """ + +def _foreach_sqrt_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_sqrt_(self: List[Tensor]) -> None + + Apply :func:`torch.sqrt` to each Tensor of the input list. + """ + +@overload +def _foreach_sub( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_sub( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, + *, + alpha: Number | _complex = 1, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_sub( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> tuple[Tensor, ...]: ... +@overload +def _foreach_sub_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalars: Sequence[Number | _complex], +) -> None: ... +@overload +def _foreach_sub_( + self: tuple[Tensor, ...] | list[Tensor] | None, + other: tuple[Tensor, ...] | list[Tensor] | None, + *, + alpha: Number | _complex = 1, +) -> None: ... +@overload +def _foreach_sub_( + self: tuple[Tensor, ...] | list[Tensor] | None, + scalar: Number | _complex, +) -> None: ... +def _foreach_tan( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_tan(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.tan` to each Tensor of the input list. + """ + +def _foreach_tan_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_tan_(self: List[Tensor]) -> None + + Apply :func:`torch.tan` to each Tensor of the input list. + """ + +def _foreach_tanh( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_tanh(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.tanh` to each Tensor of the input list. + """ + +def _foreach_tanh_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_tanh_(self: List[Tensor]) -> None + + Apply :func:`torch.tanh` to each Tensor of the input list. + """ + +def _foreach_trunc( + self: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + _foreach_trunc(self: List[Tensor]) -> List[Tensor] + + Apply :func:`torch.trunc` to each Tensor of the input list. + """ + +def _foreach_trunc_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_trunc_(self: List[Tensor]) -> None + + Apply :func:`torch.trunc` to each Tensor of the input list. + """ + +def _foreach_zero_(self: tuple[Tensor, ...] | list[Tensor] | None) -> None: + r""" + _foreach_zero_(self: List[Tensor]) -> None + + Apply :func:`torch.zero` to each Tensor of the input list. + """ + +def _from_functional_tensor(t: Tensor) -> Tensor: ... +def _functional_assert_async( + input: Tensor, + assert_msg: str, + dep_token: Tensor, +) -> Tensor: ... +def _functional_assert_scalar( + self: Number | _complex, + assert_msg: str, + dep_token: Tensor, +) -> Tensor: ... +def _functional_sym_constrain_range( + size: Number | _complex, + min: _int | None, + max: _int | None, + dep_token: Tensor, +) -> Tensor: ... +def _functional_sym_constrain_range_for_size( + size: Number | _complex, + min: _int | None, + max: _int | None, + dep_token: Tensor, +) -> Tensor: ... +def _functionalize_apply_view_metas(tensor: Tensor, base: Tensor) -> Tensor: ... +def _functionalize_are_all_mutations_hidden_from_autograd( + t: Tensor, +) -> _bool: ... +def _functionalize_are_all_mutations_under_no_grad_or_inference_mode( + t: Tensor, +) -> _bool: ... +def _functionalize_commit_update(t: Tensor) -> None: ... +def _functionalize_has_metadata_mutation(tensor: Tensor) -> _bool: ... +def _functionalize_inductor_storage_resized_counter(t: Tensor) -> _int: ... +def _functionalize_is_symbolic(tensor: Tensor) -> _bool: ... +def _functionalize_mark_mutation_hidden_from_autograd(t: Tensor) -> None: ... +def _functionalize_mark_storage_changed(tensor: Tensor) -> _bool: ... +def _functionalize_mutation_counter(t: Tensor) -> _int: ... +def _functionalize_replace(self_: Tensor, other: Tensor) -> None: ... +def _functionalize_storage_changed_counter(t: Tensor) -> _int: ... +def _functionalize_sync(t: Tensor) -> None: ... +def _functionalize_unsafe_set(dst: Tensor, src: Tensor) -> None: ... +def _functionalize_was_inductor_storage_resized(t: Tensor) -> _bool: ... +def _functionalize_was_storage_changed(tensor: Tensor) -> _bool: ... +@overload +def _fused_adagrad_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + state_sums: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: Tensor, + lr_decay: _float, + weight_decay: _float, + eps: _float, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_adagrad_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + state_sums: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: _float, + lr_decay: _float, + weight_decay: _float, + eps: _float, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_adam_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + exp_avgs: tuple[Tensor, ...] | list[Tensor] | None, + exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + max_exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: Tensor, + beta1: _float, + beta2: _float, + weight_decay: _float, + eps: _float, + amsgrad: _bool, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_adam_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + exp_avgs: tuple[Tensor, ...] | list[Tensor] | None, + exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + max_exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: _float, + beta1: _float, + beta2: _float, + weight_decay: _float, + eps: _float, + amsgrad: _bool, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_adamw_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + exp_avgs: tuple[Tensor, ...] | list[Tensor] | None, + exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + max_exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: Tensor, + beta1: _float, + beta2: _float, + weight_decay: _float, + eps: _float, + amsgrad: _bool, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_adamw_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + exp_avgs: tuple[Tensor, ...] | list[Tensor] | None, + exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + max_exp_avg_sqs: tuple[Tensor, ...] | list[Tensor] | None, + state_steps: tuple[Tensor, ...] | list[Tensor] | None, + *, + lr: _float, + beta1: _float, + beta2: _float, + weight_decay: _float, + eps: _float, + amsgrad: _bool, + maximize: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +def _fused_dropout( + input: Tensor, + p: _float, + generator: Generator | None = None, +) -> tuple[Tensor, Tensor]: ... +def _fused_moving_avg_obs_fq_helper( + input: Tensor, + observer_on: Tensor, + fake_quant_on: Tensor, + running_min: Tensor, + running_max: Tensor, + scale: Tensor, + zero_point: Tensor, + averaging_const: _float, + quant_min: _int, + quant_max: _int, + ch_axis: _int, + per_row_fake_quant: _bool = False, + symmetric_quant: _bool = False, +) -> torch.return_types._fused_moving_avg_obs_fq_helper: ... +def _fused_rms_norm( + input: Tensor, + normalized_shape: _size, + weight: Tensor | None, + eps: _float | None, +) -> tuple[Tensor, Tensor]: ... +def _fused_sdp_choice( + query: Tensor, + key: Tensor, + value: Tensor, + attn_mask: Tensor | None = None, + dropout_p: _float = 0.0, + is_causal: _bool = False, + *, + scale: _float | None = None, + enable_gqa: _bool = False, +) -> _int: ... +@overload +def _fused_sgd_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + momentum_buffer_list: tuple[Tensor, ...] | list[Tensor] | None, + *, + weight_decay: _float, + momentum: _float, + lr: Tensor, + dampening: _float, + nesterov: _bool, + maximize: _bool, + is_first_step: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +@overload +def _fused_sgd_( + self: tuple[Tensor, ...] | list[Tensor] | None, + grads: tuple[Tensor, ...] | list[Tensor] | None, + momentum_buffer_list: tuple[Tensor, ...] | list[Tensor] | None, + *, + weight_decay: _float, + momentum: _float, + lr: _float, + dampening: _float, + nesterov: _bool, + maximize: _bool, + is_first_step: _bool, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, +) -> None: ... +def _fw_primal_copy( + input: Tensor, + level: _int, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _grid_sampler_2d_cpu_fallback( + input: Tensor, + grid: Tensor, + interpolation_mode: _int, + padding_mode: _int, + align_corners: _bool, +) -> Tensor: ... +def _grouped_mm( + input: Tensor, + mat2: Tensor, + offs: Tensor | None = None, + bias: Tensor | None = None, + out_dtype: _dtype | None = None, +) -> Tensor: ... +def _has_compatible_shallow_copy_type( + input: Tensor, + from_: Tensor, +) -> _bool: ... +def _histogramdd_bin_edges( + input: Tensor, + bins: _size, + *, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, +) -> tuple[Tensor, ...]: ... +def _histogramdd_from_bin_cts( + input: Tensor, + bins: _size, + *, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, +) -> Tensor: ... +def _histogramdd_from_bin_tensors( + input: Tensor, + bins: tuple[Tensor, ...] | list[Tensor] | None, + *, + weight: Tensor | None = None, + density: _bool = False, +) -> Tensor: ... +def _index_put_impl_( + input: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + values: Tensor, + accumulate: _bool = False, + unsafe: _bool = False, +) -> Tensor: ... +def _indices_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: ... +def _int_mm( + input: Tensor, + mat2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _is_all_true(input: Tensor) -> Tensor: ... +def _is_any_true(input: Tensor) -> Tensor: ... +def _is_functional_tensor(t: Tensor) -> _bool: ... +def _is_functional_tensor_base(t: Tensor) -> _bool: ... +def _is_zerotensor(input: Tensor) -> _bool: ... +def _lazy_clone(input: Tensor) -> Tensor: ... +def _linalg_check_errors( + info: Tensor, + api_name: str, + *, + is_matrix: _bool, +) -> None: ... +def _linalg_det( + A: Tensor, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types._linalg_det: ... +def _linalg_eigh( + A: Tensor, + UPLO: str = "L", + compute_v: _bool = True, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types._linalg_eigh: ... +def _linalg_slogdet( + A: Tensor, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types._linalg_slogdet: ... +def _linalg_solve_ex( + A: Tensor, + B: Tensor, + *, + left: _bool = True, + check_errors: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types._linalg_solve_ex: ... +def _linalg_svd( + A: Tensor, + full_matrices: _bool = False, + compute_uv: _bool = True, + *, + driver: str | None = None, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types._linalg_svd: ... +def _log_softmax( + input: Tensor, + dim: _int, + half_to_float: _bool, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _log_softmax_backward_data( + grad_output: Tensor, + output: Tensor, + dim: _int, + input_dtype: _dtype, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _logcumsumexp( + input: Tensor, + dim: _int, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _lstm_mps( + input: Tensor, + hx: tuple[Tensor, ...] | list[Tensor] | None, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_first: _bool, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: ... +def _lu_with_info( + input: Tensor, + pivot: _bool = True, + check_errors: _bool = True, +) -> torch.return_types._lu_with_info: ... +def _make_dep_token( + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +def _make_dual(primal: Tensor, tangent: Tensor, level: _int) -> Tensor: ... +def _make_dual_copy( + primal: Tensor, + tangent: Tensor, + level: _int, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _make_per_channel_quantized_tensor( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + axis: _int, +) -> Tensor: ... +def _make_per_tensor_quantized_tensor( + input: Tensor, + scale: _float, + zero_point: _int, +) -> Tensor: ... +def _masked_scale(input: Tensor, mask: Tensor, scale: _float) -> Tensor: ... +def _masked_softmax( + input: Tensor, + mask: Tensor, + dim: _int | None = None, + mask_type: _int | None = None, +) -> Tensor: ... +def _mixed_dtypes_linear( + input: Tensor, + weight: Tensor, + scale: Tensor, + *, + bias: Tensor | None = None, + activation: str | None = None, +) -> Tensor: ... +def _mkldnn_reshape(input: Tensor, shape: _size) -> Tensor: ... +def _mkldnn_transpose(input: Tensor, dim0: _int, dim1: _int) -> Tensor: ... +def _mkldnn_transpose_(input: Tensor, dim0: _int, dim1: _int) -> Tensor: ... +def _mps_convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def _mps_convolution_transpose( + input: Tensor, + weight: Tensor, + padding: Sequence[_int | SymInt], + output_padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +@overload +def _native_batch_norm_legit( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor, + running_var: Tensor, + training: _bool, + momentum: _float, + eps: _float, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> tuple[Tensor, Tensor, Tensor]: ... +@overload +def _native_batch_norm_legit( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + training: _bool, + momentum: _float, + eps: _float, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> tuple[Tensor, Tensor, Tensor]: ... +def _native_batch_norm_legit_no_training( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor, + running_var: Tensor, + momentum: _float, + eps: _float, +) -> tuple[Tensor, Tensor, Tensor]: ... +def _native_multi_head_attention( + query: Tensor, + key: Tensor, + value: Tensor, + embed_dim: _int, + num_head: _int, + qkv_weight: Tensor, + qkv_bias: Tensor, + proj_weight: Tensor, + proj_bias: Tensor, + mask: Tensor | None = None, + need_weights: _bool = True, + average_attn_weights: _bool = True, + mask_type: _int | None = None, +) -> tuple[Tensor, Tensor]: ... +def _neg_view(input: Tensor) -> Tensor: ... +def _neg_view_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: ... +def _nested_compute_contiguous_strides_offsets( + nested_size: Tensor, +) -> tuple[Tensor, Tensor]: ... +def _nested_from_padded( + padded: Tensor, + cpu_nested_shape_example: Tensor, + fuse_transform_0213: _bool = False, +) -> Tensor: ... +def _nested_from_padded_and_nested_example( + padded: Tensor, + nt_example: Tensor, +) -> Tensor: ... +def _nested_from_padded_tensor( + padded: Tensor, + offsets: Tensor, + dummy: Tensor, + ragged_idx: _int = 1, + min_seqlen: Tensor | None = None, + max_seqlen: Tensor | None = None, + sum_S: _int | SymInt | None = None, +) -> Tensor: ... +def _nested_get_jagged_dummy(any: Tensor) -> Tensor: ... +def _nested_get_lengths(input: Tensor) -> Tensor: ... +def _nested_get_max_seqlen(input: Tensor) -> Tensor: ... +def _nested_get_min_seqlen(input: Tensor) -> Tensor: ... +def _nested_get_offsets(input: Tensor) -> Tensor: ... +def _nested_get_ragged_idx(input: Tensor) -> _int: ... +def _nested_get_values(input: Tensor) -> Tensor: ... +def _nested_get_values_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _nested_tensor_from_mask( + t: Tensor, + mask: Tensor, + mask_check: _bool = True, +) -> Tensor: ... +def _nested_tensor_from_mask_left_aligned(t: Tensor, mask: Tensor) -> _bool: ... +def _nested_tensor_from_tensor_list( + list: tuple[Tensor, ...] | list[Tensor] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = None, +) -> Tensor: ... +def _nested_tensor_softmax_with_shape( + input: Tensor, + query: Tensor, +) -> Tensor: ... +def _nested_view_from_buffer( + input: Tensor, + nested_size: Tensor, + nested_strides: Tensor, + offsets: Tensor, +) -> Tensor: ... +def _nested_view_from_buffer_copy( + input: Tensor, + nested_size: Tensor, + nested_strides: Tensor, + offsets: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _nested_view_from_jagged( + input: Tensor, + offsets: Tensor, + dummy: Tensor, + lengths: Tensor | None = None, + ragged_idx: _int = 1, + min_seqlen: Tensor | None = None, + max_seqlen: Tensor | None = None, +) -> Tensor: ... +def _nested_view_from_jagged_copy( + input: Tensor, + offsets: Tensor, + dummy: Tensor, + lengths: Tensor | None = None, + ragged_idx: _int = 1, + min_seqlen: Tensor | None = None, + max_seqlen: Tensor | None = None, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _nnpack_available() -> _bool: ... +def _nnpack_spatial_convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: _int | SymInt | Sequence[_int | SymInt], + stride: _int | SymInt | Sequence[_int | SymInt] = 1, +) -> Tensor: ... +def _pack_padded_sequence( + input: Tensor, + lengths: Tensor, + batch_first: _bool, +) -> tuple[Tensor, Tensor]: ... +def _pad_packed_sequence( + data: Tensor, + batch_sizes: Tensor, + batch_first: _bool, + padding_value: Number | _complex, + total_length: _int, +) -> tuple[Tensor, Tensor]: ... +def _pin_memory( + input: Tensor, + device: DeviceLikeType | None = None, +) -> Tensor: ... +def _prelu_kernel(input: Tensor, weight: Tensor) -> Tensor: ... +def _print(s: str) -> None: ... +def _propagate_xla_data(input: Tensor, output: Tensor) -> None: ... +def _remove_batch_dim( + input: Tensor, + level: _int, + batch_size: _int | SymInt, + out_dim: _int, +) -> Tensor: ... +def _reshape_alias_copy( + input: Tensor, + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + *, + out: Tensor | None = None, +) -> Tensor: ... +def _reshape_from_tensor(input: Tensor, shape: Tensor) -> Tensor: ... +def _resize_output_( + input: Tensor, + size: Sequence[_int | SymInt], + device: DeviceLikeType | None, +) -> Tensor: ... +def _rowwise_prune( + weight: Tensor, + mask: Tensor, + compressed_indices_dtype: _dtype, +) -> tuple[Tensor, Tensor]: ... +def _safe_softmax( + input: Tensor, + dim: _int, + dtype: _dtype | None = None, +) -> Tensor: ... +def _sample_dirichlet( + input: Tensor, + generator: Generator | None = None, +) -> Tensor: ... +def _saturate_weight_to_fp16(weight: Tensor) -> Tensor: ... +def _scaled_dot_product_attention_math( + query: Tensor, + key: Tensor, + value: Tensor, + attn_mask: Tensor | None = None, + dropout_p: _float = 0.0, + is_causal: _bool = False, + dropout_mask: Tensor | None = None, + *, + scale: _float | None = None, + enable_gqa: _bool = False, +) -> tuple[Tensor, Tensor]: ... +def _scaled_dot_product_attention_math_for_mps( + query: Tensor, + key: Tensor, + value: Tensor, + attn_mask: Tensor | None = None, + dropout_p: _float = 0.0, + is_causal: _bool = False, + dropout_mask: Tensor | None = None, + *, + scale: _float | None = None, +) -> tuple[Tensor, Tensor]: ... +def _scaled_dot_product_cudnn_attention( + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Tensor | None, + compute_log_sumexp: _bool, + dropout_p: _float = 0.0, + is_causal: _bool = False, + return_debug_mask: _bool = False, + *, + scale: _float | None = None, +) -> torch.return_types._scaled_dot_product_cudnn_attention: ... +def _scaled_dot_product_efficient_attention( + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Tensor | None, + compute_log_sumexp: _bool, + dropout_p: _float = 0.0, + is_causal: _bool = False, + *, + scale: _float | None = None, +) -> torch.return_types._scaled_dot_product_efficient_attention: ... +def _scaled_dot_product_flash_attention( + query: Tensor, + key: Tensor, + value: Tensor, + dropout_p: _float = 0.0, + is_causal: _bool = False, + return_debug_mask: _bool = False, + *, + scale: _float | None = None, +) -> torch.return_types._scaled_dot_product_flash_attention: ... +def _scaled_dot_product_flash_attention_for_cpu( + query: Tensor, + key: Tensor, + value: Tensor, + dropout_p: _float = 0.0, + is_causal: _bool = False, + *, + attn_mask: Tensor | None = None, + scale: _float | None = None, +) -> torch.return_types._scaled_dot_product_flash_attention_for_cpu: ... +def _scaled_grouped_mm( + input: Tensor, + mat2: Tensor, + scale_a: Tensor, + scale_b: Tensor, + offs: Tensor | None = None, + bias: Tensor | None = None, + scale_result: Tensor | None = None, + out_dtype: _dtype | None = None, + use_fast_accum: _bool = False, +) -> Tensor: ... +def _scaled_mm( + input: Tensor, + mat2: Tensor, + scale_a: Tensor, + scale_b: Tensor, + bias: Tensor | None = None, + scale_result: Tensor | None = None, + out_dtype: _dtype | None = None, + use_fast_accum: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _shape_as_tensor(input: Tensor) -> Tensor: ... +def _sobol_engine_draw( + quasi: Tensor, + n: _int, + sobolstate: Tensor, + dimension: _int, + num_generated: _int, + dtype: _dtype | None, +) -> tuple[Tensor, Tensor]: ... +def _sobol_engine_ff_( + input: Tensor, + n: _int, + sobolstate: Tensor, + dimension: _int, + num_generated: _int, +) -> Tensor: ... +def _sobol_engine_initialize_state_( + input: Tensor, + dimension: _int, +) -> Tensor: ... +def _sobol_engine_scramble_( + input: Tensor, + ltm: Tensor, + dimension: _int, +) -> Tensor: ... +def _softmax( + input: Tensor, + dim: _int, + half_to_float: _bool, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _softmax_backward_data( + grad_output: Tensor, + output: Tensor, + dim: _int, + input_dtype: _dtype, + *, + grad_input: Tensor | None = None, +) -> Tensor: ... +def _sparse_broadcast_to(input: Tensor, size: _size) -> Tensor: ... +def _sparse_broadcast_to_copy( + input: Tensor, + size: _size, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _sparse_csr_prod( + input: Tensor, + dim: _int | _size, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, +) -> Tensor: ... +def _sparse_csr_sum( + input: Tensor, + dim: _int | _size, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, +) -> Tensor: ... +def _sparse_log_softmax_backward_data( + grad_output: Tensor, + output: Tensor, + dim: _int, + input: Tensor, +) -> Tensor: ... +def _sparse_semi_structured_addmm( + input: Tensor, + mat1: Tensor, + mat1_meta: Tensor, + mat2: Tensor, + *, + alpha: Number | _complex = 1, + beta: Number | _complex = 1, + out_dtype: _dtype | None = None, +) -> Tensor: ... +def _sparse_semi_structured_apply( + input: Tensor, + thread_masks: Tensor, +) -> tuple[Tensor, Tensor]: ... +def _sparse_semi_structured_apply_dense( + input: Tensor, + thread_masks: Tensor, +) -> Tensor: ... +def _sparse_semi_structured_linear( + input: Tensor, + weight: Tensor, + meta: Tensor, + *, + bias: Tensor | None = None, + activation: str | None = None, + out_dtype: _dtype | None = None, +) -> Tensor: ... +def _sparse_semi_structured_mm( + mat1: Tensor, + mat1_meta: Tensor, + mat2: Tensor, + *, + out_dtype: _dtype | None = None, +) -> Tensor: ... +def _sparse_semi_structured_tile( + input: Tensor, + algorithm: str = "", + use_cutlass: _bool = True, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: ... +def _sparse_softmax_backward_data( + grad_output: Tensor, + output: Tensor, + dim: _int, + input: Tensor, +) -> Tensor: ... +def _sparse_sparse_matmul(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def _sparse_sum(input: Tensor) -> Tensor: ... +@overload +def _sparse_sum(input: Tensor, *, dtype: _dtype) -> Tensor: ... +@overload +def _sparse_sum(input: Tensor, dim: _int | _size) -> Tensor: ... +@overload +def _sparse_sum( + input: Tensor, + dim: _int | _size, + *, + dtype: _dtype, +) -> Tensor: ... +def _stack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _standard_gamma( + input: Tensor, + generator: Generator | None = None, +) -> Tensor: ... +def _standard_gamma_grad(input: Tensor, output: Tensor) -> Tensor: ... +def _sync(t: Tensor) -> None: ... +@overload +def _test_autograd_multiple_dispatch(input: Tensor) -> Tensor: ... +@overload +def _test_autograd_multiple_dispatch(input: Tensor, b: _bool) -> Tensor: ... +def _test_autograd_multiple_dispatch_view(input: Tensor) -> Tensor: ... +def _test_autograd_multiple_dispatch_view_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def _test_check_tensor(input: Tensor) -> Tensor: ... +def _test_functorch_fallback(input: Tensor, other: Tensor) -> Tensor: ... +def _test_parallel_materialize( + input: Tensor, + num_parallel: _int, + skip_first: _bool = False, +) -> Tensor: ... +def _test_serialization_subcmul( + input: Tensor, + other: Tensor, + alpha: Number | _complex = 1, +) -> Tensor: ... +def _to_cpu( + tensors: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: ... +def _to_functional_tensor(t: Tensor) -> Tensor: ... +def _to_sparse_semi_structured(dense: Tensor) -> tuple[Tensor, Tensor]: ... +def _transform_bias_rescale_qkv( + qkv: Tensor, + qkv_bias: Tensor, + num_heads: _int, +) -> tuple[Tensor, Tensor, Tensor]: ... +def _transformer_encoder_layer_fwd( + src: Tensor, + embed_dim: _int, + num_heads: _int, + qkv_weight: Tensor, + qkv_bias: Tensor, + proj_weight: Tensor, + proj_bias: Tensor, + use_gelu: _bool, + norm_first: _bool, + eps: _float, + norm_weight_1: Tensor, + norm_bias_1: Tensor, + norm_weight_2: Tensor, + norm_bias_2: Tensor, + ffn_weight_1: Tensor, + ffn_bias_1: Tensor, + ffn_weight_2: Tensor, + ffn_bias_2: Tensor, + mask: Tensor | None = None, + mask_type: _int | None = None, +) -> Tensor: ... +def _trilinear( + i1: Tensor, + i2: Tensor, + i3: Tensor, + expand1: _size, + expand2: _size, + expand3: _size, + sumdim: _size, + unroll_dim: _int = 1, +) -> Tensor: ... +def _triton_multi_head_attention( + query: Tensor, + key: Tensor, + value: Tensor, + embed_dim: _int, + num_head: _int, + qkv_weight: Tensor, + qkv_bias: Tensor, + proj_weight: Tensor, + proj_bias: Tensor, + mask: Tensor | None = None, +) -> Tensor: ... +def _triton_scaled_dot_attention( + q: Tensor, + k: Tensor, + v: Tensor, + dropout_p: _float = 0.0, +) -> Tensor: ... +def _unique( + input: Tensor, + sorted: _bool = True, + return_inverse: _bool = False, +) -> tuple[Tensor, Tensor]: ... +def _unique2( + input: Tensor, + sorted: _bool = True, + return_inverse: _bool = False, + return_counts: _bool = False, +) -> tuple[Tensor, Tensor, Tensor]: ... +def _unpack_dual( + dual: Tensor, + level: _int, +) -> torch.return_types._unpack_dual: ... +def _unsafe_index( + input: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, +) -> Tensor: ... +def _unsafe_index_put( + input: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + values: Tensor, + accumulate: _bool = False, +) -> Tensor: ... +def _unsafe_masked_index( + input: Tensor, + mask: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + fill: Number | _complex, +) -> Tensor: ... +def _unsafe_masked_index_put_accumulate( + input: Tensor, + mask: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + values: Tensor, +) -> Tensor: ... +@overload +def _use_cudnn_ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: Tensor, + target_lengths: Tensor, + blank: _int, +) -> _bool: ... +@overload +def _use_cudnn_ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: _size, + target_lengths: _size, + blank: _int, +) -> _bool: ... +def _use_cudnn_rnn_flatten_weight() -> _bool: ... +def _validate_compressed_sparse_indices( + is_crow: _bool, + compressed_idx: Tensor, + plain_idx: Tensor, + cdim: _int, + dim: _int, + nnz: _int, +) -> None: ... +def _validate_sparse_bsc_tensor_args( + ccol_indices: Tensor, + row_indices: Tensor, + values: Tensor, + size: _size, + check_pinning: _bool | None = None, +) -> None: ... +def _validate_sparse_bsr_tensor_args( + crow_indices: Tensor, + col_indices: Tensor, + values: Tensor, + size: _size, + check_pinning: _bool | None = None, +) -> None: ... +def _validate_sparse_compressed_tensor_args( + compressed_indices: Tensor, + plain_indices: Tensor, + values: Tensor, + size: _size, + layout: _layout, + check_pinning: _bool | None = None, +) -> None: ... +def _validate_sparse_coo_tensor_args( + indices: Tensor, + values: Tensor, + size: _size, + is_coalesced: _bool | None = None, + check_pinning: _bool | None = None, +) -> None: ... +def _validate_sparse_csc_tensor_args( + ccol_indices: Tensor, + row_indices: Tensor, + values: Tensor, + size: _size, + check_pinning: _bool | None = None, +) -> None: ... +def _validate_sparse_csr_tensor_args( + crow_indices: Tensor, + col_indices: Tensor, + values: Tensor, + size: _size, + check_pinning: _bool | None = None, +) -> None: ... +def _values_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: ... +def _weight_int4pack_mm( + input: Tensor, + mat2: Tensor, + qGroupSize: _int, + qScaleAndZeros: Tensor, +) -> Tensor: ... +def _weight_int4pack_mm_for_cpu( + input: Tensor, + mat2: Tensor, + qGroupSize: _int, + qScaleAndZeros: Tensor, +) -> Tensor: ... +def _weight_int4pack_mm_with_scales_and_zeros( + input: Tensor, + mat2: Tensor, + qGroupSize: _int, + qScale: Tensor, + qZeros: Tensor, +) -> Tensor: ... +def _weight_int8pack_mm( + input: Tensor, + mat2: Tensor, + scales: Tensor, +) -> Tensor: ... +def _weight_norm(v: Tensor, g: Tensor, dim: _int = 0) -> Tensor: ... +def _weight_norm_interface( + v: Tensor, + g: Tensor, + dim: _int = 0, +) -> tuple[Tensor, Tensor]: ... +def _wrapped_linear_prepack( + weight: Tensor, + weight_scale: Tensor, + weight_zero_point: Tensor, + bias: Tensor, +) -> Tensor: ... +def _wrapped_quantized_linear_prepacked( + input: Tensor, + input_scale: Tensor, + input_zero_point: Tensor, + packed_weight: Tensor, + output_scale: Tensor, + output_zero_point: Tensor, + out_channel: _int, +) -> Tensor: ... +def abs(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + abs(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the absolute value of each element in :attr:`input`. + + .. math:: + \text{out}_{i} = |\text{input}_{i}| + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.abs(torch.tensor([-1, -2, 3])) + tensor([ 1, 2, 3]) + """ + +def abs_(input: Tensor) -> Tensor: ... +def absolute(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + absolute(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.abs` + """ + +def acos(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + acos(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the inverse cosine of each element in :attr:`input`. + + .. math:: + \text{out}_{i} = \cos^{-1}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.3348, -0.5889, 0.2005, -0.1584]) + >>> torch.acos(a) + tensor([ 1.2294, 2.2004, 1.3690, 1.7298]) + """ + +def acos_(input: Tensor) -> Tensor: ... +def acosh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + acosh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the inverse hyperbolic cosine of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \cosh^{-1}(\text{input}_{i}) + + Note: + The domain of the inverse hyperbolic cosine is `[1, inf)` and values outside this range + will be mapped to ``NaN``, except for `+ INF` for which the output is mapped to `+ INF`. + + Args: + input (Tensor): the input tensor. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4).uniform_(1, 2) + >>> a + tensor([ 1.3192, 1.9915, 1.9674, 1.7151 ]) + >>> torch.acosh(a) + tensor([ 0.7791, 1.3120, 1.2979, 1.1341 ]) + """ + +def acosh_(input: Tensor) -> Tensor: ... +def adaptive_avg_pool1d(input: Tensor, output_size: _int | _size) -> Tensor: ... +def adaptive_max_pool1d( + input: Tensor, + output_size: _int | _size, +) -> tuple[Tensor, Tensor]: ... +@overload +def add( + input: Tensor | Number | _complex, + other: Tensor | Number | _complex, + *, + alpha: Number | _complex | None = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + add(input, other, *, alpha=1, out=None) -> Tensor + + Adds :attr:`other`, scaled by :attr:`alpha`, to :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to add to :attr:`input`. + + Keyword arguments: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Examples:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.0202, 1.0985, 1.3506, -0.6056]) + >>> torch.add(a, 20) + tensor([ 20.0202, 21.0985, 21.3506, 19.3944]) + + >>> b = torch.randn(4) + >>> b + tensor([-0.9732, -0.3497, 0.6245, 0.4022]) + >>> c = torch.randn(4, 1) + >>> c + tensor([[ 0.3743], + [-1.7724], + [-0.5811], + [-0.8017]]) + >>> torch.add(b, c, alpha=10) + tensor([[ 2.7695, 3.3930, 4.3672, 4.1450], + [-18.6971, -18.0736, -17.0994, -17.3216], + [ -6.7845, -6.1610, -5.1868, -5.4090], + [ -8.9902, -8.3667, -7.3925, -7.6147]]) + """ + +@overload +def add(self: Tensor, alpha: Number | _complex, other: Tensor) -> Tensor: + r""" + add(input, other, *, alpha=1, out=None) -> Tensor + + Adds :attr:`other`, scaled by :attr:`alpha`, to :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to add to :attr:`input`. + + Keyword arguments: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Examples:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.0202, 1.0985, 1.3506, -0.6056]) + >>> torch.add(a, 20) + tensor([ 20.0202, 21.0985, 21.3506, 19.3944]) + + >>> b = torch.randn(4) + >>> b + tensor([-0.9732, -0.3497, 0.6245, 0.4022]) + >>> c = torch.randn(4, 1) + >>> c + tensor([[ 0.3743], + [-1.7724], + [-0.5811], + [-0.8017]]) + >>> torch.add(b, c, alpha=10) + tensor([[ 2.7695, 3.3930, 4.3672, 4.1450], + [-18.6971, -18.0736, -17.0994, -17.3216], + [ -6.7845, -6.1610, -5.1868, -5.4090], + [ -8.9902, -8.3667, -7.3925, -7.6147]]) + """ + +@overload +def add( + self: Tensor, + alpha: Number | _complex, + other: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + add(input, other, *, alpha=1, out=None) -> Tensor + + Adds :attr:`other`, scaled by :attr:`alpha`, to :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to add to :attr:`input`. + + Keyword arguments: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Examples:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.0202, 1.0985, 1.3506, -0.6056]) + >>> torch.add(a, 20) + tensor([ 20.0202, 21.0985, 21.3506, 19.3944]) + + >>> b = torch.randn(4) + >>> b + tensor([-0.9732, -0.3497, 0.6245, 0.4022]) + >>> c = torch.randn(4, 1) + >>> c + tensor([[ 0.3743], + [-1.7724], + [-0.5811], + [-0.8017]]) + >>> torch.add(b, c, alpha=10) + tensor([[ 2.7695, 3.3930, 4.3672, 4.1450], + [-18.6971, -18.0736, -17.0994, -17.3216], + [ -6.7845, -6.1610, -5.1868, -5.4090], + [ -8.9902, -8.3667, -7.3925, -7.6147]]) + """ + +@overload +def addbmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + batch1: Tensor, + batch2: Tensor, +) -> Tensor: + r""" + addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored + in :attr:`batch1` and :attr:`batch2`, + with a reduced add step (all matrix multiplications get accumulated + along the first dimension). + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the + same number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + .. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` + must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) + """ + +@overload +def addbmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + batch1: Tensor, + batch2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored + in :attr:`batch1` and :attr:`batch2`, + with a reduced add step (all matrix multiplications get accumulated + along the first dimension). + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the + same number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + .. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` + must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) + """ + +@overload +def addbmm( + input: Tensor, + batch1: Tensor, + batch2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored + in :attr:`batch1` and :attr:`batch2`, + with a reduced add step (all matrix multiplications get accumulated + along the first dimension). + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the + same number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + .. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` + must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) + """ + +@overload +def addbmm( + beta: Number | _complex, + self: Tensor, + batch1: Tensor, + batch2: Tensor, +) -> Tensor: + r""" + addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored + in :attr:`batch1` and :attr:`batch2`, + with a reduced add step (all matrix multiplications get accumulated + along the first dimension). + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the + same number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + .. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` + must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) + """ + +@overload +def addbmm( + beta: Number | _complex, + self: Tensor, + batch1: Tensor, + batch2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored + in :attr:`batch1` and :attr:`batch2`, + with a reduced add step (all matrix multiplications get accumulated + along the first dimension). + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the + same number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + .. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` + must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) + """ + +@overload +def addcdiv( + self: Tensor, + value: Number | _complex, + tensor1: Tensor, + tensor2: Tensor, +) -> Tensor: + r""" + addcdiv(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise division of :attr:`tensor1` by :attr:`tensor2`, + multiplies the result by the scalar :attr:`value` and adds it to :attr:`input`. + + .. warning:: + Integer division with addcdiv is no longer supported, and in a future + release addcdiv will perform a true division of tensor1 and tensor2. + The historic addcdiv behavior can be implemented as + (input + value * torch.trunc(tensor1 / tensor2)).to(input.dtype) + for integer inputs and as (input + value * tensor1 / tensor2) for float inputs. + The future addcdiv behavior is just the latter implementation: + (input + value * tensor1 / tensor2), for all dtypes. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} + + + The shapes of :attr:`input`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the numerator tensor + tensor2 (Tensor): the denominator tensor + + Keyword args: + value (Number, optional): multiplier for :math:`\text{tensor1} / \text{tensor2}` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcdiv(t, t1, t2, value=0.1) + tensor([[-0.2312, -3.6496, 0.1312], + [-1.0428, 3.4292, -0.1030], + [-0.5369, -0.9829, 0.0430]]) + """ + +@overload +def addcdiv( + self: Tensor, + value: Number | _complex, + tensor1: Tensor, + tensor2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addcdiv(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise division of :attr:`tensor1` by :attr:`tensor2`, + multiplies the result by the scalar :attr:`value` and adds it to :attr:`input`. + + .. warning:: + Integer division with addcdiv is no longer supported, and in a future + release addcdiv will perform a true division of tensor1 and tensor2. + The historic addcdiv behavior can be implemented as + (input + value * torch.trunc(tensor1 / tensor2)).to(input.dtype) + for integer inputs and as (input + value * tensor1 / tensor2) for float inputs. + The future addcdiv behavior is just the latter implementation: + (input + value * tensor1 / tensor2), for all dtypes. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} + + + The shapes of :attr:`input`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the numerator tensor + tensor2 (Tensor): the denominator tensor + + Keyword args: + value (Number, optional): multiplier for :math:`\text{tensor1} / \text{tensor2}` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcdiv(t, t1, t2, value=0.1) + tensor([[-0.2312, -3.6496, 0.1312], + [-1.0428, 3.4292, -0.1030], + [-0.5369, -0.9829, 0.0430]]) + """ + +@overload +def addcdiv( + input: Tensor, + tensor1: Tensor, + tensor2: Tensor, + *, + value: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addcdiv(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise division of :attr:`tensor1` by :attr:`tensor2`, + multiplies the result by the scalar :attr:`value` and adds it to :attr:`input`. + + .. warning:: + Integer division with addcdiv is no longer supported, and in a future + release addcdiv will perform a true division of tensor1 and tensor2. + The historic addcdiv behavior can be implemented as + (input + value * torch.trunc(tensor1 / tensor2)).to(input.dtype) + for integer inputs and as (input + value * tensor1 / tensor2) for float inputs. + The future addcdiv behavior is just the latter implementation: + (input + value * tensor1 / tensor2), for all dtypes. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} + + + The shapes of :attr:`input`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the numerator tensor + tensor2 (Tensor): the denominator tensor + + Keyword args: + value (Number, optional): multiplier for :math:`\text{tensor1} / \text{tensor2}` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcdiv(t, t1, t2, value=0.1) + tensor([[-0.2312, -3.6496, 0.1312], + [-1.0428, 3.4292, -0.1030], + [-0.5369, -0.9829, 0.0430]]) + """ + +@overload +def addcmul( + self: Tensor, + value: Number | _complex, + tensor1: Tensor, + tensor2: Tensor, +) -> Tensor: + r""" + addcmul(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise multiplication of :attr:`tensor1` + by :attr:`tensor2`, multiplies the result by the scalar :attr:`value` + and adds it to :attr:`input`. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \text{tensor1}_i \times \text{tensor2}_i + + The shapes of :attr:`tensor`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the tensor to be multiplied + tensor2 (Tensor): the tensor to be multiplied + + Keyword args: + value (Number, optional): multiplier for :math:`tensor1 .* tensor2` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcmul(t, t1, t2, value=0.1) + tensor([[-0.8635, -0.6391, 1.6174], + [-0.7617, -0.5879, 1.7388], + [-0.8353, -0.6249, 1.6511]]) + """ + +@overload +def addcmul( + self: Tensor, + value: Number | _complex, + tensor1: Tensor, + tensor2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addcmul(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise multiplication of :attr:`tensor1` + by :attr:`tensor2`, multiplies the result by the scalar :attr:`value` + and adds it to :attr:`input`. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \text{tensor1}_i \times \text{tensor2}_i + + The shapes of :attr:`tensor`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the tensor to be multiplied + tensor2 (Tensor): the tensor to be multiplied + + Keyword args: + value (Number, optional): multiplier for :math:`tensor1 .* tensor2` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcmul(t, t1, t2, value=0.1) + tensor([[-0.8635, -0.6391, 1.6174], + [-0.7617, -0.5879, 1.7388], + [-0.8353, -0.6249, 1.6511]]) + """ + +@overload +def addcmul( + input: Tensor, + tensor1: Tensor, + tensor2: Tensor, + *, + value: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addcmul(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + + Performs the element-wise multiplication of :attr:`tensor1` + by :attr:`tensor2`, multiplies the result by the scalar :attr:`value` + and adds it to :attr:`input`. + + .. math:: + \text{out}_i = \text{input}_i + \text{value} \times \text{tensor1}_i \times \text{tensor2}_i + + The shapes of :attr:`tensor`, :attr:`tensor1`, and :attr:`tensor2` must be + :ref:`broadcastable `. + + For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be + a real number, otherwise an integer. + + Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the tensor to be multiplied + tensor2 (Tensor): the tensor to be multiplied + + Keyword args: + value (Number, optional): multiplier for :math:`tensor1 .* tensor2` + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcmul(t, t1, t2, value=0.1) + tensor([[-0.8635, -0.6391, 1.6174], + [-0.7617, -0.5879, 1.7388], + [-0.8353, -0.6249, 1.6511]]) + """ + +@overload +def addmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat1: Tensor, + mat2: Tensor, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat1: Tensor, + mat2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmm( + input: Tensor, + mat1: Tensor, + mat2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmm( + input: Tensor, + mat1: Tensor, + mat2: Tensor, + out_dtype: _dtype, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmm( + beta: Number | _complex, + self: Tensor, + mat1: Tensor, + mat2: Tensor, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmm( + beta: Number | _complex, + self: Tensor, + mat1: Tensor, + mat2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. + The matrix :attr:`input` is added to the final result. + + If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a :math:`(n \times p)` tensor + and :attr:`out` will be a :math:`(n \times p)` tensor. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operation has support for arguments with :ref:`sparse layouts`. If + :attr:`input` is sparse the result will have the same layout and if :attr:`out` + is provided it must have the same layout as :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) + """ + +@overload +def addmv( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat: Tensor, + vec: Tensor, +) -> Tensor: + r""" + addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`mat` and + the vector :attr:`vec`. + The vector :attr:`input` is added to the final result. + + If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a 1-D tensor of size `n` and + :attr:`out` will be 1-D tensor of size `n`. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) + """ + +@overload +def addmv( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat: Tensor, + vec: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`mat` and + the vector :attr:`vec`. + The vector :attr:`input` is added to the final result. + + If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a 1-D tensor of size `n` and + :attr:`out` will be 1-D tensor of size `n`. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) + """ + +@overload +def addmv( + input: Tensor, + mat: Tensor, + vec: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`mat` and + the vector :attr:`vec`. + The vector :attr:`input` is added to the final result. + + If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a 1-D tensor of size `n` and + :attr:`out` will be 1-D tensor of size `n`. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) + """ + +@overload +def addmv( + beta: Number | _complex, + self: Tensor, + mat: Tensor, + vec: Tensor, +) -> Tensor: + r""" + addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`mat` and + the vector :attr:`vec`. + The vector :attr:`input` is added to the final result. + + If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a 1-D tensor of size `n` and + :attr:`out` will be 1-D tensor of size `n`. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) + """ + +@overload +def addmv( + beta: Number | _complex, + self: Tensor, + mat: Tensor, + vec: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`mat` and + the vector :attr:`vec`. + The vector :attr:`input` is added to the final result. + + If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a 1-D tensor of size `n` and + :attr:`out` will be 1-D tensor of size `n`. + + :attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between + :attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) + """ + +@overload +def addmv_( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat: Tensor, + vec: Tensor, +) -> Tensor: ... +@overload +def addmv_( + input: Tensor, + mat: Tensor, + vec: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, +) -> Tensor: ... +@overload +def addmv_( + beta: Number | _complex, + self: Tensor, + mat: Tensor, + vec: Tensor, +) -> Tensor: ... +@overload +def addr( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + vec1: Tensor, + vec2: Tensor, +) -> Tensor: + r""" + addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` + and adds it to the matrix :attr:`input`. + + Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the + outer product between :attr:`vec1` and :attr:`vec2` and the added matrix + :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector + of size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a matrix of size + :math:`(n \times m)` and :attr:`out` will be a matrix of size + :math:`(n \times m)`. + + Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{vec1} \otimes \text{vec2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) + """ + +@overload +def addr( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + vec1: Tensor, + vec2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` + and adds it to the matrix :attr:`input`. + + Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the + outer product between :attr:`vec1` and :attr:`vec2` and the added matrix + :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector + of size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a matrix of size + :math:`(n \times m)` and :attr:`out` will be a matrix of size + :math:`(n \times m)`. + + Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{vec1} \otimes \text{vec2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) + """ + +@overload +def addr( + input: Tensor, + vec1: Tensor, + vec2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` + and adds it to the matrix :attr:`input`. + + Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the + outer product between :attr:`vec1` and :attr:`vec2` and the added matrix + :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector + of size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a matrix of size + :math:`(n \times m)` and :attr:`out` will be a matrix of size + :math:`(n \times m)`. + + Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{vec1} \otimes \text{vec2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) + """ + +@overload +def addr( + beta: Number | _complex, + self: Tensor, + vec1: Tensor, + vec2: Tensor, +) -> Tensor: + r""" + addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` + and adds it to the matrix :attr:`input`. + + Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the + outer product between :attr:`vec1` and :attr:`vec2` and the added matrix + :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector + of size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a matrix of size + :math:`(n \times m)` and :attr:`out` will be a matrix of size + :math:`(n \times m)`. + + Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{vec1} \otimes \text{vec2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) + """ + +@overload +def addr( + beta: Number | _complex, + self: Tensor, + vec1: Tensor, + vec2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + + Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` + and adds it to the matrix :attr:`input`. + + Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the + outer product between :attr:`vec1` and :attr:`vec2` and the added matrix + :attr:`input` respectively. + + .. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector + of size `m`, then :attr:`input` must be + :ref:`broadcastable ` with a matrix of size + :math:`(n \times m)` and :attr:`out` will be a matrix of size + :math:`(n \times m)`. + + Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{vec1} \otimes \text{vec2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) + """ + +def adjoint(input: Tensor) -> Tensor: + r""" + adjoint(input: Tensor) -> Tensor + Returns a view of the tensor conjugated and with the last two dimensions transposed. + + ``x.adjoint()`` is equivalent to ``x.transpose(-2, -1).conj()`` for complex tensors and + to ``x.transpose(-2, -1)`` for real tensors. + + Args: + {input} + + Example:: + + >>> x = torch.arange(4, dtype=torch.float) + >>> A = torch.complex(x, x).reshape(2, 2) + >>> A + tensor([[0.+0.j, 1.+1.j], + [2.+2.j, 3.+3.j]]) + >>> A.adjoint() + tensor([[0.-0.j, 2.-2.j], + [1.-1.j, 3.-3.j]]) + >>> (A.adjoint() == A.mH).all() + tensor(True) + """ + +def affine_grid_generator( + theta: Tensor, + size: Sequence[_int | SymInt], + align_corners: _bool, +) -> Tensor: ... +def alias_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.alias`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def all(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + all(input: Tensor, *, out=None) -> Tensor + + Tests if all elements in :attr:`input` evaluate to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.all(a) + tensor(False, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.all(a) + tensor(False) + + .. function:: all(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if all elements in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(4, 2).bool() + >>> a + tensor([[True, True], + [True, False], + [True, True], + [True, True]], dtype=torch.bool) + >>> torch.all(a, dim=1) + tensor([ True, False, True, True], dtype=torch.bool) + >>> torch.all(a, dim=0) + tensor([ True, False], dtype=torch.bool) + """ + +@overload +def all( + input: Tensor, + dim: _size | None = None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + all(input: Tensor, *, out=None) -> Tensor + + Tests if all elements in :attr:`input` evaluate to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.all(a) + tensor(False, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.all(a) + tensor(False) + + .. function:: all(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if all elements in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(4, 2).bool() + >>> a + tensor([[True, True], + [True, False], + [True, True], + [True, True]], dtype=torch.bool) + >>> torch.all(a, dim=1) + tensor([ True, False, True, True], dtype=torch.bool) + >>> torch.all(a, dim=0) + tensor([ True, False], dtype=torch.bool) + """ + +@overload +def all( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + all(input: Tensor, *, out=None) -> Tensor + + Tests if all elements in :attr:`input` evaluate to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.all(a) + tensor(False, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.all(a) + tensor(False) + + .. function:: all(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if all elements in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(4, 2).bool() + >>> a + tensor([[True, True], + [True, False], + [True, True], + [True, True]], dtype=torch.bool) + >>> torch.all(a, dim=1) + tensor([ True, False, True, True], dtype=torch.bool) + >>> torch.all(a, dim=0) + tensor([ True, False], dtype=torch.bool) + """ + +@overload +def all( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + all(input: Tensor, *, out=None) -> Tensor + + Tests if all elements in :attr:`input` evaluate to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.all(a) + tensor(False, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.all(a) + tensor(False) + + .. function:: all(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if all elements in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(4, 2).bool() + >>> a + tensor([[True, True], + [True, False], + [True, True], + [True, True]], dtype=torch.bool) + >>> torch.all(a, dim=1) + tensor([ True, False, True, True], dtype=torch.bool) + >>> torch.all(a, dim=0) + tensor([ True, False], dtype=torch.bool) + """ + +def allclose( + input: Tensor, + other: Tensor, + rtol: _float = 1e-05, + atol: _float = 1e-08, + equal_nan: _bool = False, +) -> _bool: + r""" + allclose(input: Tensor, other: Tensor, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) -> bool + + This function checks if :attr:`input` and :attr:`other` satisfy the condition: + + .. math:: + \lvert \text{input}_i - \text{other}_i \rvert \leq \texttt{atol} + \texttt{rtol} \times \lvert \text{other}_i \rvert + + elementwise, for all elements of :attr:`input` and :attr:`other`. The behaviour of this function is analogous to + `numpy.allclose `_ + + Args: + input (Tensor): first tensor to compare + other (Tensor): second tensor to compare + atol (float, optional): absolute tolerance. Default: 1e-08 + rtol (float, optional): relative tolerance. Default: 1e-05 + equal_nan (bool, optional): if ``True``, then two ``NaN`` s will be considered equal. Default: ``False`` + + Example:: + + >>> torch.allclose(torch.tensor([10000., 1e-07]), torch.tensor([10000.1, 1e-08])) + False + >>> torch.allclose(torch.tensor([10000., 1e-08]), torch.tensor([10000.1, 1e-09])) + True + >>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')])) + False + >>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]), equal_nan=True) + True + """ + +def alpha_dropout(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def alpha_dropout_(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def amax( + input: Tensor, + dim: _int | _size = (), + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + amax(input, dim, keepdim=False, *, out=None) -> Tensor + + Returns the maximum value of each slice of the :attr:`input` tensor in the given + dimension(s) :attr:`dim`. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.8177, 1.4878, -0.2491, 0.9130], + [-0.7158, 1.1775, 2.0992, 0.4817], + [-0.0053, 0.0164, -1.3738, -0.0507], + [ 1.9700, 1.1106, -1.0318, -1.0816]]) + >>> torch.amax(a, 1) + tensor([1.4878, 2.0992, 0.0164, 1.9700]) + """ + +def amin( + input: Tensor, + dim: _int | _size = (), + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + amin(input, dim, keepdim=False, *, out=None) -> Tensor + + Returns the minimum value of each slice of the :attr:`input` tensor in the given + dimension(s) :attr:`dim`. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.6451, -0.4866, 0.2987, -1.3312], + [-0.5744, 1.2980, 1.8397, -0.2713], + [ 0.9128, 0.9214, -1.7268, -0.2995], + [ 0.9023, 0.4853, 0.9075, -1.6165]]) + >>> torch.amin(a, 1) + tensor([-1.3312, -0.5744, -1.7268, -1.6165]) + """ + +def aminmax( + input: Tensor, + *, + dim: _int | None = None, + keepdim: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.aminmax: + r""" + aminmax(input, *, dim=None, keepdim=False, out=None) -> (Tensor min, Tensor max) + + Computes the minimum and maximum values of the :attr:`input` tensor. + + Args: + input (Tensor): + The input tensor + + Keyword Args: + dim (Optional[int]): + The dimension along which to compute the values. If `None`, + computes the values over the entire :attr:`input` tensor. + Default is `None`. + keepdim (bool): + If `True`, the reduced dimensions will be kept in the output + tensor as dimensions with size 1 for broadcasting, otherwise + they will be removed, as if calling (:func:`torch.squeeze`). + Default is `False`. + out (Optional[Tuple[Tensor, Tensor]]): + Optional tensors on which to write the result. Must have the same + shape and dtype as the expected output. + Default is `None`. + + Returns: + A named tuple `(min, max)` containing the minimum and maximum values. + + Raises: + RuntimeError + If any of the dimensions to compute the values over has size 0. + + .. note:: + NaN values are propagated to the output if at least one value is NaN. + + .. seealso:: + :func:`torch.amin` computes just the minimum value + :func:`torch.amax` computes just the maximum value + + Example:: + + >>> torch.aminmax(torch.tensor([1, -3, 5])) + torch.return_types.aminmax( + min=tensor(-3), + max=tensor(5)) + + >>> # aminmax propagates NaNs + >>> torch.aminmax(torch.tensor([1, -3, 5, torch.nan])) + torch.return_types.aminmax( + min=tensor(nan), + max=tensor(nan)) + + >>> t = torch.arange(10).view(2, 5) + >>> t + tensor([[0, 1, 2, 3, 4], + [5, 6, 7, 8, 9]]) + >>> t.aminmax(dim=0, keepdim=True) + torch.return_types.aminmax( + min=tensor([[0, 1, 2, 3, 4]]), + max=tensor([[5, 6, 7, 8, 9]])) + """ + +def angle(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + angle(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the element-wise angle (in radians) of the given :attr:`input` tensor. + + .. math:: + \text{out}_{i} = angle(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + .. note:: Starting in PyTorch 1.8, angle returns pi for negative real numbers, + zero for non-negative real numbers, and propagates NaNs. Previously + the function would return zero for all real numbers and not propagate + floating-point NaNs. + + Example:: + + >>> torch.angle(torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]))*180/3.14159 + tensor([ 135., 135, -45]) + """ + +@overload +def any(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + any(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Tests if any element in :attr:`input` evaluates to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.any(a) + tensor(True, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.any(a) + tensor(True) + + .. function:: any(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if any element in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 2) < 0 + >>> a + tensor([[ True, True], + [False, True], + [ True, True], + [False, False]]) + >>> torch.any(a, 1) + tensor([ True, True, True, False]) + >>> torch.any(a, 0) + tensor([True, True]) + """ + +@overload +def any( + input: Tensor, + dim: _size | None = None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + any(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Tests if any element in :attr:`input` evaluates to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.any(a) + tensor(True, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.any(a) + tensor(True) + + .. function:: any(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if any element in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 2) < 0 + >>> a + tensor([[ True, True], + [False, True], + [ True, True], + [False, False]]) + >>> torch.any(a, 1) + tensor([ True, True, True, False]) + >>> torch.any(a, 0) + tensor([True, True]) + """ + +@overload +def any( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + any(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Tests if any element in :attr:`input` evaluates to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.any(a) + tensor(True, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.any(a) + tensor(True) + + .. function:: any(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if any element in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 2) < 0 + >>> a + tensor([[ True, True], + [False, True], + [ True, True], + [False, False]]) + >>> torch.any(a, 1) + tensor([ True, True, True, False]) + >>> torch.any(a, 0) + tensor([True, True]) + """ + +@overload +def any( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + any(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Tests if any element in :attr:`input` evaluates to `True`. + + .. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.any(a) + tensor(True, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.any(a) + tensor(True) + + .. function:: any(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + + For each row of :attr:`input` in the given dimension :attr:`dim`, + returns `True` if any element in the row evaluate to `True` and `False` otherwise. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 2) < 0 + >>> a + tensor([[ True, True], + [False, True], + [ True, True], + [False, False]]) + >>> torch.any(a, 1) + tensor([ True, True, True, False]) + >>> torch.any(a, 0) + tensor([True, True]) + """ + +@overload +def arange( + start: Number, + end: Number, + step: Number, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +@overload +def arange( + start: Number, + end: Number, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +@overload +def arange( + end: Number, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +@overload +def arange( + end: Number | _complex, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +@overload +def arange( + start: Number | _complex, + end: Number | _complex, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +@overload +def arange( + start: Number | _complex, + end: Number | _complex, + step: Number | _complex = 1, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` + with values from the interval ``[start, end)`` taken with common difference + :attr:`step` beginning from `start`. + + Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), + the results may be affected by floating-point rounding behavior. Some values in the sequence + might not be exactly representable in certain floating-point formats, which can lead to + repeated values or unexpected rounding. For precise sequences, it is recommended to use + integer dtypes instead of floating-point dtypes. + + Note that non-integer :attr:`step` is subject to floating point rounding errors when + comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` + in such cases. + + .. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} + + Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) + """ + +def arccos(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arccos(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.acos`. + """ + +def arccos_(input: Tensor) -> Tensor: ... +def arccosh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arccosh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.acosh`. + """ + +def arccosh_(input: Tensor) -> Tensor: ... +def arcsin(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arcsin(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.asin`. + """ + +def arcsin_(input: Tensor) -> Tensor: ... +def arcsinh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arcsinh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.asinh`. + """ + +def arcsinh_(input: Tensor) -> Tensor: ... +def arctan(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arctan(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.atan`. + """ + +def arctan2( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + arctan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + Alias for :func:`torch.atan2`. + """ + +def arctan_(input: Tensor) -> Tensor: ... +def arctanh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + arctanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Alias for :func:`torch.atanh`. + """ + +def arctanh_(input: Tensor) -> Tensor: ... +def argmax( + input: Tensor, + dim: _int | None = None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + argmax(input) -> LongTensor + + Returns the indices of the maximum value of all elements in the :attr:`input` tensor. + + This is the second value returned by :meth:`torch.max`. See its + documentation for the exact semantics of this method. + + .. note:: If there are multiple maximal values then the indices of the first maximal value are returned. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], + [-0.7401, -0.8805, -0.3402, -1.1936], + [ 0.4907, -1.3948, -1.0691, -0.3132], + [-1.6092, 0.5419, -0.2993, 0.3195]]) + >>> torch.argmax(a) + tensor(0) + + .. function:: argmax(input, dim, keepdim=False) -> LongTensor + :noindex: + + Returns the indices of the maximum values of a tensor across a dimension. + + This is the second value returned by :meth:`torch.max`. See its + documentation for the exact semantics of this method. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, the argmax of the flattened input is returned. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], + [-0.7401, -0.8805, -0.3402, -1.1936], + [ 0.4907, -1.3948, -1.0691, -0.3132], + [-1.6092, 0.5419, -0.2993, 0.3195]]) + >>> torch.argmax(a, dim=1) + tensor([ 0, 2, 0, 1]) + """ + +def argmin( + input: Tensor, + dim: _int | None = None, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + argmin(input, dim=None, keepdim=False) -> LongTensor + + Returns the indices of the minimum value(s) of the flattened tensor or along a dimension + + This is the second value returned by :meth:`torch.min`. See its + documentation for the exact semantics of this method. + + .. note:: If there are multiple minimal values then the indices of the first minimal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, the argmin of the flattened input is returned. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.1139, 0.2254, -0.1381, 0.3687], + [ 1.0100, -1.1975, -0.0102, -0.4732], + [-0.9240, 0.1207, -0.7506, -1.0213], + [ 1.7809, -1.2960, 0.9384, 0.1438]]) + >>> torch.argmin(a) + tensor(13) + >>> torch.argmin(a, dim=1) + tensor([ 2, 1, 3, 1]) + >>> torch.argmin(a, dim=1, keepdim=True) + tensor([[2], + [1], + [3], + [1]]) + """ + +@overload +def argsort( + input: Tensor, + *, + stable: _bool, + dim: _int = -1, + descending: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + argsort(input, dim=-1, descending=False, *, stable=False) -> Tensor + + Returns the indices that sort a tensor along a given dimension in ascending + order by value. + + This is the second value returned by :meth:`torch.sort`. See its documentation + for the exact semantics of this method. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. If ``False``, the relative order of values + which compare equal is not guaranteed. ``True`` is slower. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): controls the relative order of equivalent elements + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0785, 1.5267, -0.8521, 0.4065], + [ 0.1598, 0.0788, -0.0745, -1.2700], + [ 1.2208, 1.0722, -0.7064, 1.2564], + [ 0.0669, -0.2318, -0.8229, -0.9280]]) + + + >>> torch.argsort(a, dim=1) + tensor([[2, 0, 3, 1], + [3, 2, 1, 0], + [2, 1, 0, 3], + [3, 2, 1, 0]]) + """ + +@overload +def argsort( + input: Tensor, + dim: _int = -1, + descending: _bool = False, +) -> Tensor: + r""" + argsort(input, dim=-1, descending=False, *, stable=False) -> Tensor + + Returns the indices that sort a tensor along a given dimension in ascending + order by value. + + This is the second value returned by :meth:`torch.sort`. See its documentation + for the exact semantics of this method. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. If ``False``, the relative order of values + which compare equal is not guaranteed. ``True`` is slower. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): controls the relative order of equivalent elements + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0785, 1.5267, -0.8521, 0.4065], + [ 0.1598, 0.0788, -0.0745, -1.2700], + [ 1.2208, 1.0722, -0.7064, 1.2564], + [ 0.0669, -0.2318, -0.8229, -0.9280]]) + + + >>> torch.argsort(a, dim=1) + tensor([[2, 0, 3, 1], + [3, 2, 1, 0], + [2, 1, 0, 3], + [3, 2, 1, 0]]) + """ + +@overload +def argsort( + input: Tensor, + dim: str | EllipsisType | None, + descending: _bool = False, +) -> Tensor: + r""" + argsort(input, dim=-1, descending=False, *, stable=False) -> Tensor + + Returns the indices that sort a tensor along a given dimension in ascending + order by value. + + This is the second value returned by :meth:`torch.sort`. See its documentation + for the exact semantics of this method. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. If ``False``, the relative order of values + which compare equal is not guaranteed. ``True`` is slower. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): controls the relative order of equivalent elements + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0785, 1.5267, -0.8521, 0.4065], + [ 0.1598, 0.0788, -0.0745, -1.2700], + [ 1.2208, 1.0722, -0.7064, 1.2564], + [ 0.0669, -0.2318, -0.8229, -0.9280]]) + + + >>> torch.argsort(a, dim=1) + tensor([[2, 0, 3, 1], + [3, 2, 1, 0], + [2, 1, 0, 3], + [3, 2, 1, 0]]) + """ + +def argwhere(input: Tensor) -> Tensor: + r""" + argwhere(input) -> Tensor + + Returns a tensor containing the indices of all non-zero elements of + :attr:`input`. Each row in the result contains the indices of a non-zero + element in :attr:`input`. The result is sorted lexicographically, with + the last index changing the fastest (C-style). + + If :attr:`input` has :math:`n` dimensions, then the resulting indices tensor + :attr:`out` is of size :math:`(z \times n)`, where :math:`z` is the total number of + non-zero elements in the :attr:`input` tensor. + + .. note:: + This function is similar to NumPy's `argwhere`. + + When :attr:`input` is on CUDA, this function causes host-device synchronization. + + Args: + {input} + + Example:: + + >>> t = torch.tensor([1, 0, 1]) + >>> torch.argwhere(t) + tensor([[0], + [2]]) + >>> t = torch.tensor([[1, 0, 1], [0, 1, 1]]) + >>> torch.argwhere(t) + tensor([[0, 0], + [0, 2], + [1, 1], + [1, 2]]) + """ + +def as_strided( + input: Tensor, + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + storage_offset: _int | SymInt | None = None, +) -> Tensor: + r""" + as_strided(input, size, stride, storage_offset=None) -> Tensor + + Create a view of an existing `torch.Tensor` :attr:`input` with specified + :attr:`size`, :attr:`stride` and :attr:`storage_offset`. + + .. warning:: + Prefer using other view functions, like :meth:`torch.Tensor.view` or + :meth:`torch.Tensor.expand`, to setting a view's strides manually with + `as_strided`, as this function will throw an error on non-standard Pytorch + backends (that do not have a concept of stride) and the result will depend + on the current layout in memory. The constructed view must only refer to + elements within the Tensor's storage or a runtime error will be thrown. + If the generated view is "overlapped" (with multiple indices referring to + the same element in memory), the behavior of inplace operations on this view + is undefined (and might not throw runtime errors). + + Args: + input (Tensor): the input tensor. + size (tuple or ints): the shape of the output tensor + stride (tuple or ints): the stride of the output tensor + storage_offset (int, optional): the offset in the underlying storage of the output tensor. + If ``None``, the storage_offset of the output tensor will match the input tensor. + + Example:: + + >>> x = torch.randn(3, 3) + >>> x + tensor([[ 0.9039, 0.6291, 1.0795], + [ 0.1586, 2.1939, -0.4900], + [-0.1909, -0.7503, 1.9355]]) + >>> t = torch.as_strided(x, (2, 2), (1, 2)) + >>> t + tensor([[0.9039, 1.0795], + [0.6291, 0.1586]]) + >>> t = torch.as_strided(x, (2, 2), (1, 2), 1) + tensor([[0.6291, 0.1586], + [1.0795, 2.1939]]) + """ + +def as_strided_( + input: Tensor, + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + storage_offset: _int | SymInt | None = None, +) -> Tensor: ... +def as_strided_copy( + input: Tensor, + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + storage_offset: _int | SymInt | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.as_strided`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def as_strided_scatter( + input: Tensor, + src: Tensor, + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + storage_offset: _int | SymInt | None = None, +) -> Tensor: + r""" + as_strided_scatter(input, src, size, stride, storage_offset=None) -> Tensor + + Embeds the values of the :attr:`src` tensor into :attr:`input` along + the elements corresponding to the result of calling + input.as_strided(size, stride, storage_offset). + + This function returns a tensor with fresh storage; it does not + return a view. + + Args: + input (Tensor): the input tensor. + size (tuple or ints): the shape of the output tensor + stride (tuple or ints): the stride of the output tensor + storage_offset (int, optional): the offset in the underlying storage of the output tensor + + .. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + `torch.as_strided(input, size, stride, storage_offset)` + + Example:: + + >>> a = torch.arange(4).reshape(2, 2) + 1 + >>> a + tensor([[1, 2], + [3, 4]]) + >>> b = torch.zeros(3, 3) + >>> b + tensor([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]]) + >>> torch.as_strided_scatter(b, a, (2, 2), (1, 2)) + tensor([[1., 3., 2.], + [4., 0., 0.], + [0., 0., 0.]]) + """ + +def as_tensor( + data: Any, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, +) -> Tensor: + r""" + as_tensor(data: Any, dtype: Optional[dtype] = None, device: Optional[DeviceLikeType]) -> Tensor + + Converts :attr:`data` into a tensor, sharing data and preserving autograd + history if possible. + + If :attr:`data` is already a tensor with the requested dtype and device + then :attr:`data` itself is returned, but if :attr:`data` is a + tensor with a different dtype or device then it's copied as if using + `data.to(dtype=dtype, device=device)`. + + If :attr:`data` is a NumPy array (an ndarray) with the same dtype and device then a + tensor is constructed using :func:`torch.from_numpy`. + + If :attr:`data` is a CuPy array, the returned tensor will be located on the same device as the CuPy array unless + specifically overwritten by :attr:`device` or a default device. The device of the CuPy array is inferred from the + pointer of the array using `cudaPointerGetAttributes` unless :attr:`device` is provided with an explicit device index. + + .. seealso:: + + :func:`torch.tensor` never shares its data and creates a new "leaf tensor" (see :doc:`/notes/autograd`). + + + Args: + data (array_like): Initial data for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, infers data type from :attr:`data`. + device (:class:`torch.device`, optional): the device of the constructed tensor. If None and data is a tensor + then the device of data is used. If None and data is not a tensor then + the result tensor is constructed on the current device. + + + Example:: + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.as_tensor(a) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.as_tensor(a, device=torch.device('cuda')) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([1, 2, 3]) + """ + +def asarray( + obj: Any, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + copy: _bool | None = None, + requires_grad: _bool = False, +) -> Tensor: + r""" + asarray(obj: Any, *, dtype: Optional[dtype], device: Optional[DeviceLikeType], copy: Optional[bool] = None, requires_grad: bool = False) -> Tensor # noqa: B950 + + Converts :attr:`obj` to a tensor. + + :attr:`obj` can be one of: + + 1. a tensor + 2. a NumPy array or a NumPy scalar + 3. a DLPack capsule + 4. an object that implements Python's buffer protocol + 5. a scalar + 6. a sequence of scalars + + When :attr:`obj` is a tensor, NumPy array, or DLPack capsule the returned tensor will, + by default, not require a gradient, have the same datatype as :attr:`obj`, be on the + same device, and share memory with it. These properties can be controlled with the + :attr:`dtype`, :attr:`device`, :attr:`copy`, and :attr:`requires_grad` keyword arguments. + If the returned tensor is of a different datatype, on a different device, or a copy is + requested then it will not share its memory with :attr:`obj`. If :attr:`requires_grad` + is ``True`` then the returned tensor will require a gradient, and if :attr:`obj` is + also a tensor with an autograd history then the returned tensor will have the same history. + + When :attr:`obj` is not a tensor, NumPy array, or DLPack capsule but implements Python's + buffer protocol then the buffer is interpreted as an array of bytes grouped according to + the size of the datatype passed to the :attr:`dtype` keyword argument. (If no datatype is + passed then the default floating point datatype is used, instead.) The returned tensor + will have the specified datatype (or default floating point datatype if none is specified) + and, by default, be on the CPU device and share memory with the buffer. + + When :attr:`obj` is a NumPy scalar, the returned tensor will be a 0-dimensional tensor on + the CPU and that doesn't share its memory (i.e. ``copy=True``). By default datatype will + be the PyTorch datatype corresponding to the NumPy's scalar's datatype. + + When :attr:`obj` is none of the above but a scalar, or a sequence of scalars then the + returned tensor will, by default, infer its datatype from the scalar values, be on the + current default device, and not share its memory. + + .. seealso:: + + :func:`torch.tensor` creates a tensor that always copies the data from the input object. + :func:`torch.from_numpy` creates a tensor that always shares memory from NumPy arrays. + :func:`torch.frombuffer` creates a tensor that always shares memory from objects that + implement the buffer protocol. + :func:`torch.from_dlpack` creates a tensor that always shares memory from + DLPack capsules. + + Args: + obj (object): a tensor, NumPy array, DLPack Capsule, object that implements Python's + buffer protocol, scalar, or sequence of scalars. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the datatype of the returned tensor. + Default: ``None``, which causes the datatype of the returned tensor to be + inferred from :attr:`obj`. + copy (bool, optional): controls whether the returned tensor shares memory with :attr:`obj`. + Default: ``None``, which causes the returned tensor to share memory with :attr:`obj` + whenever possible. If ``True`` then the returned tensor does not share its memory. + If ``False`` then the returned tensor shares its memory with :attr:`obj` and an + error is thrown if it cannot. + device (:class:`torch.device`, optional): the device of the returned tensor. + Default: ``None``, which causes the device of :attr:`obj` to be used. Or, if + :attr:`obj` is a Python sequence, the current default device will be used. + requires_grad (bool, optional): whether the returned tensor requires grad. + Default: ``False``, which causes the returned tensor not to require a gradient. + If ``True``, then the returned tensor will require a gradient, and if :attr:`obj` + is also a tensor with an autograd history then the returned tensor will have + the same history. + + Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> # Shares memory with tensor 'a' + >>> b = torch.asarray(a) + >>> a.data_ptr() == b.data_ptr() + True + >>> # Forces memory copy + >>> c = torch.asarray(a, copy=True) + >>> a.data_ptr() == c.data_ptr() + False + + >>> a = torch.tensor([1., 2., 3.], requires_grad=True) + >>> b = a + 2 + >>> b + tensor([3., 4., 5.], grad_fn=) + >>> # Shares memory with tensor 'b', with no grad + >>> c = torch.asarray(b) + >>> c + tensor([3., 4., 5.]) + >>> # Shares memory with tensor 'b', retaining autograd history + >>> d = torch.asarray(b, requires_grad=True) + >>> d + tensor([3., 4., 5.], grad_fn=) + + >>> array = numpy.array([1, 2, 3]) + >>> # Shares memory with array 'array' + >>> t1 = torch.asarray(array) + >>> array.__array_interface__['data'][0] == t1.data_ptr() + True + >>> # Copies memory due to dtype mismatch + >>> t2 = torch.asarray(array, dtype=torch.float32) + >>> array.__array_interface__['data'][0] == t2.data_ptr() + False + + >>> scalar = numpy.float64(0.5) + >>> torch.asarray(scalar) + tensor(0.5000, dtype=torch.float64) + """ + +def asin(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + asin(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the arcsine of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \sin^{-1}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.5962, 1.4985, -0.4396, 1.4525]) + >>> torch.asin(a) + tensor([-0.6387, nan, -0.4552, nan]) + """ + +def asin_(input: Tensor) -> Tensor: ... +def asinh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + asinh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the inverse hyperbolic sine of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \sinh^{-1}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.1606, -1.4267, -1.0899, -1.0250 ]) + >>> torch.asinh(a) + tensor([ 0.1599, -1.1534, -0.9435, -0.8990 ]) + """ + +def asinh_(input: Tensor) -> Tensor: ... +def atan(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + atan(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the arctangent of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \tan^{-1}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.2341, 0.2539, -0.6256, -0.6448]) + >>> torch.atan(a) + tensor([ 0.2299, 0.2487, -0.5591, -0.5727]) + """ + +def atan2( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + atan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + + Element-wise arctangent of :math:`\text{input}_{i} / \text{other}_{i}` + with consideration of the quadrant. Returns a new tensor with the signed angles + in radians between vector :math:`(\text{other}_{i}, \text{input}_{i})` + and vector :math:`(1, 0)`. (Note that :math:`\text{other}_{i}`, the second + parameter, is the x-coordinate, while :math:`\text{input}_{i}`, the first + parameter, is the y-coordinate.) + + The shapes of ``input`` and ``other`` must be + :ref:`broadcastable `. + + Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.9041, 0.0196, -0.3108, -2.4423]) + >>> torch.atan2(a, torch.randn(4)) + tensor([ 0.9833, 0.0811, -1.9743, -1.4151]) + """ + +def atan_(input: Tensor) -> Tensor: ... +def atanh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + atanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the inverse hyperbolic tangent of the elements of :attr:`input`. + + Note: + The domain of the inverse hyperbolic tangent is `(-1, 1)` and values outside this range + will be mapped to ``NaN``, except for the values `1` and `-1` for which the output is + mapped to `+/-INF` respectively. + + .. math:: + \text{out}_{i} = \tanh^{-1}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4).uniform_(-1, 1) + >>> a + tensor([ -0.9385, 0.2968, -0.8591, -0.1871 ]) + >>> torch.atanh(a) + tensor([ -1.7253, 0.3060, -1.2899, -0.1893 ]) + """ + +def atanh_(input: Tensor) -> Tensor: ... +def avg_pool1d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + ceil_mode: _bool = False, + count_include_pad: _bool = True, +) -> Tensor: ... +@overload +def baddbmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + batch1: Tensor, + batch2: Tensor, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def baddbmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + batch1: Tensor, + batch2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def baddbmm( + input: Tensor, + batch1: Tensor, + batch2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def baddbmm( + input: Tensor, + batch1: Tensor, + batch2: Tensor, + out_dtype: _dtype, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def baddbmm( + beta: Number | _complex, + self: Tensor, + batch1: Tensor, + batch2: Tensor, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def baddbmm( + beta: Number | _complex, + self: Tensor, + batch1: Tensor, + batch2: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices in :attr:`batch1` + and :attr:`batch2`. + :attr:`input` is added to the final result. + + :attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same + number of matrices. + + If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a + :math:`(b \times m \times p)` tensor, then :attr:`input` must be + :ref:`broadcastable ` with a + :math:`(b \times n \times p)` tensor and :attr:`out` will be a + :math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the + same as the scaling factors used in :meth:`torch.addbmm`. + + .. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + + If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in + it will not be propagated. + + For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and + :attr:`alpha` must be real numbers, otherwise they should be integers. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{batch1} \mathbin{@} \text{batch2}` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + + Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) + """ + +@overload +def bartlett_window( + window_length: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + bartlett_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Bartlett window function. + + .. math:: + w[n] = 1 - \left| \frac{2n}{N-1} - 1 \right| = \begin{cases} + \frac{2n}{N - 1} & \text{if } 0 \leq n \leq \frac{N - 1}{2} \\ + 2 - \frac{2n}{N - 1} & \text{if } \frac{N - 1}{2} < n < N \\ + \end{cases}, + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.bartlett_window(L, periodic=True)`` equal to + ``torch.bartlett_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +@overload +def bartlett_window( + window_length: _int, + periodic: _bool, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + bartlett_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Bartlett window function. + + .. math:: + w[n] = 1 - \left| \frac{2n}{N-1} - 1 \right| = \begin{cases} + \frac{2n}{N - 1} & \text{if } 0 \leq n \leq \frac{N - 1}{2} \\ + 2 - \frac{2n}{N - 1} & \text{if } \frac{N - 1}{2} < n < N \\ + \end{cases}, + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.bartlett_window(L, periodic=True)`` equal to + ``torch.bartlett_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +def batch_norm( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + training: _bool, + momentum: _float, + eps: _float, + cudnn_enabled: _bool, +) -> Tensor: ... +def batch_norm_backward_elemt( + grad_out: Tensor, + input: Tensor, + mean: Tensor, + invstd: Tensor, + weight: Tensor | None, + sum_dy: Tensor, + sum_dy_xmu: Tensor, + count: Tensor, +) -> Tensor: ... +def batch_norm_backward_reduce( + grad_out: Tensor, + input: Tensor, + mean: Tensor, + invstd: Tensor, + weight: Tensor | None, + input_g: _bool, + weight_g: _bool, + bias_g: _bool, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +def batch_norm_elemt( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + mean: Tensor, + invstd: Tensor, + eps: _float, + *, + out: Tensor | None = None, +) -> Tensor: ... +def batch_norm_gather_stats( + input: Tensor, + mean: Tensor, + invstd: Tensor, + running_mean: Tensor | None, + running_var: Tensor | None, + momentum: _float, + eps: _float, + count: _int, +) -> tuple[Tensor, Tensor]: ... +def batch_norm_gather_stats_with_counts( + input: Tensor, + mean: Tensor, + invstd: Tensor, + running_mean: Tensor | None, + running_var: Tensor | None, + momentum: _float, + eps: _float, + counts: Tensor, +) -> tuple[Tensor, Tensor]: ... +def batch_norm_stats(input: Tensor, eps: _float) -> tuple[Tensor, Tensor]: ... +def batch_norm_update_stats( + input: Tensor, + running_mean: Tensor | None, + running_var: Tensor | None, + momentum: _float, +) -> tuple[Tensor, Tensor]: ... +@overload +def bernoulli( + input: Tensor, + *, + generator: Generator | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + bernoulli(input: Tensor, *, generator: Optional[Generator], out: Optional[Tensor]) -> Tensor + + Draws binary random numbers (0 or 1) from a Bernoulli distribution. + + The :attr:`input` tensor should be a tensor containing probabilities + to be used for drawing the binary random number. + Hence, all values in :attr:`input` have to be in the range: + :math:`0 \leq \text{input}_i \leq 1`. + + The :math:`\text{i}^{th}` element of the output tensor will draw a + value :math:`1` according to the :math:`\text{i}^{th}` probability value given + in :attr:`input`. + + .. math:: + \text{out}_{i} \sim \mathrm{Bernoulli}(p = \text{input}_{i}) + + The returned :attr:`out` tensor only has values 0 or 1 and is of the same + shape as :attr:`input`. + + :attr:`out` can have integral ``dtype``, but :attr:`input` must have floating + point ``dtype``. + + Args: + input (Tensor): the input tensor of probability values for the Bernoulli distribution + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.empty(3, 3).uniform_(0, 1) # generate a uniform random matrix with range [0, 1] + >>> a + tensor([[ 0.1737, 0.0950, 0.3609], + [ 0.7148, 0.0289, 0.2676], + [ 0.9456, 0.8937, 0.7202]]) + >>> torch.bernoulli(a) + tensor([[ 1., 0., 0.], + [ 0., 0., 0.], + [ 1., 1., 1.]]) + + >>> a = torch.ones(3, 3) # probability of drawing "1" is 1 + >>> torch.bernoulli(a) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.]]) + >>> a = torch.zeros(3, 3) # probability of drawing "1" is 0 + >>> torch.bernoulli(a) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.], + [ 0., 0., 0.]]) + """ + +@overload +def bernoulli( + input: Tensor, + p: _float, + *, + generator: Generator | None = None, +) -> Tensor: + r""" + bernoulli(input: Tensor, *, generator: Optional[Generator], out: Optional[Tensor]) -> Tensor + + Draws binary random numbers (0 or 1) from a Bernoulli distribution. + + The :attr:`input` tensor should be a tensor containing probabilities + to be used for drawing the binary random number. + Hence, all values in :attr:`input` have to be in the range: + :math:`0 \leq \text{input}_i \leq 1`. + + The :math:`\text{i}^{th}` element of the output tensor will draw a + value :math:`1` according to the :math:`\text{i}^{th}` probability value given + in :attr:`input`. + + .. math:: + \text{out}_{i} \sim \mathrm{Bernoulli}(p = \text{input}_{i}) + + The returned :attr:`out` tensor only has values 0 or 1 and is of the same + shape as :attr:`input`. + + :attr:`out` can have integral ``dtype``, but :attr:`input` must have floating + point ``dtype``. + + Args: + input (Tensor): the input tensor of probability values for the Bernoulli distribution + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.empty(3, 3).uniform_(0, 1) # generate a uniform random matrix with range [0, 1] + >>> a + tensor([[ 0.1737, 0.0950, 0.3609], + [ 0.7148, 0.0289, 0.2676], + [ 0.9456, 0.8937, 0.7202]]) + >>> torch.bernoulli(a) + tensor([[ 1., 0., 0.], + [ 0., 0., 0.], + [ 1., 1., 1.]]) + + >>> a = torch.ones(3, 3) # probability of drawing "1" is 1 + >>> torch.bernoulli(a) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.]]) + >>> a = torch.zeros(3, 3) # probability of drawing "1" is 0 + >>> torch.bernoulli(a) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.], + [ 0., 0., 0.]]) + """ + +def bilinear( + input1: Tensor, + input2: Tensor, + weight: Tensor, + bias: Tensor | None = None, +) -> Tensor: ... +def binary_cross_entropy_with_logits( + input: Tensor, + target: Tensor, + weight: Tensor | None = None, + pos_weight: Tensor | None = None, + reduction: _int = 1, +) -> Tensor: ... +def bincount( + input: Tensor, + weights: Tensor | None = None, + minlength: _int | SymInt = 0, +) -> Tensor: + r""" + bincount(input, weights=None, minlength=0) -> Tensor + + Count the frequency of each value in an array of non-negative ints. + + The number of bins (size 1) is one larger than the largest value in + :attr:`input` unless :attr:`input` is empty, in which case the result is a + tensor of size 0. If :attr:`minlength` is specified, the number of bins is at least + :attr:`minlength` and if :attr:`input` is empty, then the result is tensor of size + :attr:`minlength` filled with zeros. If ``n`` is the value at position ``i``, + ``out[n] += weights[i]`` if :attr:`weights` is specified else + ``out[n] += 1``. + + Note: + This operation may produce nondeterministic gradients when given tensors on a CUDA device. See :doc:`/notes/randomness` for more information. + + Arguments: + input (Tensor): 1-d int tensor + weights (Tensor): optional, weight for each value in the input tensor. + Should be of same size as input tensor. + minlength (int): optional, minimum number of bins. Should be non-negative. + + Returns: + output (Tensor): a tensor of shape ``Size([max(input) + 1])`` if + :attr:`input` is non-empty, else ``Size(0)`` + + Example:: + + >>> input = torch.randint(0, 8, (5,), dtype=torch.int64) + >>> weights = torch.linspace(0, 1, steps=5) + >>> input, weights + (tensor([4, 3, 6, 3, 4]), + tensor([ 0.0000, 0.2500, 0.5000, 0.7500, 1.0000]) + + >>> torch.bincount(input) + tensor([0, 0, 0, 2, 2, 0, 1]) + + >>> input.bincount(weights) + tensor([0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 0.0000, 0.5000]) + """ + +def binomial( + count: Tensor, + prob: Tensor, + generator: Generator | None = None, +) -> Tensor: ... +@overload +def bitwise_and( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_and(input, other, *, out=None) -> Tensor + + Computes the bitwise AND of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical AND. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_and(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([1, 0, 3], dtype=torch.int8) + >>> torch.bitwise_and(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ False, True, False]) + """ + +@overload +def bitwise_and(self: Number | _complex, other: Tensor) -> Tensor: + r""" + bitwise_and(input, other, *, out=None) -> Tensor + + Computes the bitwise AND of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical AND. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_and(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([1, 0, 3], dtype=torch.int8) + >>> torch.bitwise_and(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ False, True, False]) + """ + +@overload +def bitwise_and( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_and(input, other, *, out=None) -> Tensor + + Computes the bitwise AND of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical AND. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_and(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([1, 0, 3], dtype=torch.int8) + >>> torch.bitwise_and(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ False, True, False]) + """ + +@overload +def bitwise_left_shift( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_left_shift(input, other, *, out=None) -> Tensor + + Computes the left arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i << \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_left_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 24], dtype=torch.int8) + """ + +@overload +def bitwise_left_shift(self: Number | _complex, other: Tensor) -> Tensor: + r""" + bitwise_left_shift(input, other, *, out=None) -> Tensor + + Computes the left arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i << \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_left_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 24], dtype=torch.int8) + """ + +@overload +def bitwise_left_shift( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_left_shift(input, other, *, out=None) -> Tensor + + Computes the left arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i << \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_left_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 24], dtype=torch.int8) + """ + +def bitwise_not(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + bitwise_not(input, *, out=None) -> Tensor + + Computes the bitwise NOT of the given input tensor. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical NOT. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_not(torch.tensor([-1, -2, 3], dtype=torch.int8)) + tensor([ 0, 1, -4], dtype=torch.int8) + """ + +@overload +def bitwise_or( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_or(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the bitwise OR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical OR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_or(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -2, 3], dtype=torch.int8) + >>> torch.bitwise_or(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, True, False]) + """ + +@overload +def bitwise_or(self: Number | _complex, other: Tensor) -> Tensor: + r""" + bitwise_or(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the bitwise OR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical OR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_or(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -2, 3], dtype=torch.int8) + >>> torch.bitwise_or(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, True, False]) + """ + +@overload +def bitwise_or( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_or(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the bitwise OR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical OR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_or(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -2, 3], dtype=torch.int8) + >>> torch.bitwise_or(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, True, False]) + """ + +@overload +def bitwise_right_shift( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_right_shift(input, other, *, out=None) -> Tensor + + Computes the right arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + In any case, if the value of the right operand is negative or is greater + or equal to the number of bits in the promoted left operand, the behavior is undefined. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i >> \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_right_shift(torch.tensor([-2, -7, 31], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -7, 3], dtype=torch.int8) + """ + +@overload +def bitwise_right_shift(self: Number | _complex, other: Tensor) -> Tensor: + r""" + bitwise_right_shift(input, other, *, out=None) -> Tensor + + Computes the right arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + In any case, if the value of the right operand is negative or is greater + or equal to the number of bits in the promoted left operand, the behavior is undefined. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i >> \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_right_shift(torch.tensor([-2, -7, 31], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -7, 3], dtype=torch.int8) + """ + +@overload +def bitwise_right_shift( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_right_shift(input, other, *, out=None) -> Tensor + + Computes the right arithmetic shift of :attr:`input` by :attr:`other` bits. + The input tensor must be of integral type. This operator supports + :ref:`broadcasting to a common shape ` and + :ref:`type promotion `. + In any case, if the value of the right operand is negative or is greater + or equal to the number of bits in the promoted left operand, the behavior is undefined. + + The operation applied is: + + .. math:: + \text{out}_i = \text{input}_i >> \text{other}_i + + Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_right_shift(torch.tensor([-2, -7, 31], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -7, 3], dtype=torch.int8) + """ + +@overload +def bitwise_xor( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_xor(input, other, *, out=None) -> Tensor + + Computes the bitwise XOR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical XOR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_xor(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 0], dtype=torch.int8) + >>> torch.bitwise_xor(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, False, False]) + """ + +@overload +def bitwise_xor(self: Number | _complex, other: Tensor) -> Tensor: + r""" + bitwise_xor(input, other, *, out=None) -> Tensor + + Computes the bitwise XOR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical XOR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_xor(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 0], dtype=torch.int8) + >>> torch.bitwise_xor(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, False, False]) + """ + +@overload +def bitwise_xor( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bitwise_xor(input, other, *, out=None) -> Tensor + + Computes the bitwise XOR of :attr:`input` and :attr:`other`. The input tensor must be of + integral or Boolean types. For bool tensors, it computes the logical XOR. + + Args: + input: the first input tensor + other: the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.bitwise_xor(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 0], dtype=torch.int8) + >>> torch.bitwise_xor(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, False, False]) + """ + +@overload +def blackman_window( + window_length: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + blackman_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Blackman window function. + + .. math:: + w[n] = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{N - 1} \right) + 0.08 \cos \left( \frac{4 \pi n}{N - 1} \right) + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.blackman_window(L, periodic=True)`` equal to + ``torch.blackman_window(L + 1, periodic=False)[:-1]``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +@overload +def blackman_window( + window_length: _int, + periodic: _bool, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + blackman_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Blackman window function. + + .. math:: + w[n] = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{N - 1} \right) + 0.08 \cos \left( \frac{4 \pi n}{N - 1} \right) + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.blackman_window(L, periodic=True)`` equal to + ``torch.blackman_window(L + 1, periodic=False)[:-1]``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +@overload +def bmm( + input: Tensor, + mat2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bmm(input, mat2, out_dtype=None, *, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored in :attr:`input` + and :attr:`mat2`. + + :attr:`input` and :attr:`mat2` must be 3-D tensors each containing + the same number of matrices. + + If :attr:`input` is a :math:`(b \times n \times m)` tensor, :attr:`mat2` is a + :math:`(b \times m \times p)` tensor, :attr:`out` will be a + :math:`(b \times n \times p)` tensor. + + .. math:: + \text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + .. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + + Args: + input (Tensor): the first batch of matrices to be multiplied + mat2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword Args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> input = torch.randn(10, 3, 4) + >>> mat2 = torch.randn(10, 4, 5) + >>> res = torch.bmm(input, mat2) + >>> res.size() + torch.Size([10, 3, 5]) + """ + +@overload +def bmm( + input: Tensor, + mat2: Tensor, + out_dtype: _dtype, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + bmm(input, mat2, out_dtype=None, *, out=None) -> Tensor + + Performs a batch matrix-matrix product of matrices stored in :attr:`input` + and :attr:`mat2`. + + :attr:`input` and :attr:`mat2` must be 3-D tensors each containing + the same number of matrices. + + If :attr:`input` is a :math:`(b \times n \times m)` tensor, :attr:`mat2` is a + :math:`(b \times m \times p)` tensor, :attr:`out` will be a + :math:`(b \times n \times p)` tensor. + + .. math:: + \text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + .. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + + Args: + input (Tensor): the first batch of matrices to be multiplied + mat2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword Args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> input = torch.randn(10, 3, 4) + >>> mat2 = torch.randn(10, 4, 5) + >>> res = torch.bmm(input, mat2) + >>> res.size() + torch.Size([10, 3, 5]) + """ + +def broadcast_to(input: Tensor, size: Sequence[_int | SymInt]) -> Tensor: + r""" + broadcast_to(input, shape) -> Tensor + + Broadcasts :attr:`input` to the shape :attr:`\shape`. + Equivalent to calling ``input.expand(shape)``. See :meth:`~Tensor.expand` for details. + + Args: + input (Tensor): the input tensor. + shape (list, tuple, or :class:`torch.Size`): the new shape. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> torch.broadcast_to(x, (3, 3)) + tensor([[1, 2, 3], + [1, 2, 3], + [1, 2, 3]]) + """ + +@overload +def bucketize( + input: Tensor, + boundaries: Tensor, + *, + out_int32: _bool = False, + right: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + bucketize(input, boundaries, *, out_int32=False, right=False, out=None) -> Tensor + + Returns the indices of the buckets to which each value in the :attr:`input` belongs, where the + boundaries of the buckets are set by :attr:`boundaries`. Return a new tensor with the same size + as :attr:`input`. If :attr:`right` is False (default), then the left boundary is open. Note that + this behavior is opposite the behavior of + `numpy.digitize `_. + More formally, the returned index satisfies the following rules: + + .. list-table:: + :widths: 15 85 + :header-rows: 1 + + * - :attr:`right` + - *returned index satisfies* + * - False + - ``boundaries[i-1] < input[m][n]...[l][x] <= boundaries[i]`` + * - True + - ``boundaries[i-1] <= input[m][n]...[l][x] < boundaries[i]`` + + Args: + input (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + boundaries (Tensor): 1-D tensor, must contain a strictly increasing sequence, or the return value is undefined. + + Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): determines the behavior for values in :attr:`boundaries`. See the table above. + out (Tensor, optional): the output tensor, must be the same size as :attr:`input` if provided. + + + Example:: + + >>> boundaries = torch.tensor([1, 3, 5, 7, 9]) + >>> boundaries + tensor([1, 3, 5, 7, 9]) + >>> v = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> v + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.bucketize(v, boundaries) + tensor([[1, 3, 4], + [1, 3, 4]]) + >>> torch.bucketize(v, boundaries, right=True) + tensor([[2, 3, 5], + [2, 3, 5]]) + """ + +@overload +def bucketize( + self: Number | _complex, + boundaries: Tensor, + *, + out_int32: _bool = False, + right: _bool = False, +) -> Tensor: + r""" + bucketize(input, boundaries, *, out_int32=False, right=False, out=None) -> Tensor + + Returns the indices of the buckets to which each value in the :attr:`input` belongs, where the + boundaries of the buckets are set by :attr:`boundaries`. Return a new tensor with the same size + as :attr:`input`. If :attr:`right` is False (default), then the left boundary is open. Note that + this behavior is opposite the behavior of + `numpy.digitize `_. + More formally, the returned index satisfies the following rules: + + .. list-table:: + :widths: 15 85 + :header-rows: 1 + + * - :attr:`right` + - *returned index satisfies* + * - False + - ``boundaries[i-1] < input[m][n]...[l][x] <= boundaries[i]`` + * - True + - ``boundaries[i-1] <= input[m][n]...[l][x] < boundaries[i]`` + + Args: + input (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + boundaries (Tensor): 1-D tensor, must contain a strictly increasing sequence, or the return value is undefined. + + Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): determines the behavior for values in :attr:`boundaries`. See the table above. + out (Tensor, optional): the output tensor, must be the same size as :attr:`input` if provided. + + + Example:: + + >>> boundaries = torch.tensor([1, 3, 5, 7, 9]) + >>> boundaries + tensor([1, 3, 5, 7, 9]) + >>> v = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> v + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.bucketize(v, boundaries) + tensor([[1, 3, 4], + [1, 3, 4]]) + >>> torch.bucketize(v, boundaries, right=True) + tensor([[2, 3, 5], + [2, 3, 5]]) + """ + +def can_cast(from_: _dtype, to: _dtype) -> _bool: + r""" + can_cast(from_, to) -> bool + + Determines if a type conversion is allowed under PyTorch casting rules + described in the type promotion :ref:`documentation `. + + Args: + from\_ (dtype): The original :class:`torch.dtype`. + to (dtype): The target :class:`torch.dtype`. + + Example:: + + >>> torch.can_cast(torch.double, torch.float) + True + >>> torch.can_cast(torch.float, torch.int) + False + """ + +@overload +def cat( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cat(tensors, dim=0, *, out=None) -> Tensor + + Concatenates the given sequence of tensors in :attr:`tensors` in the given dimension. + All tensors must either have the same shape (except in the concatenating + dimension) or be a 1-D empty tensor with size ``(0,)``. + + :func:`torch.cat` can be seen as an inverse operation for :func:`torch.split` + and :func:`torch.chunk`. + + :func:`torch.cat` can be best understood via examples. + + .. seealso:: + + :func:`torch.stack` concatenates the given sequence along a new dimension. + + Args: + tensors (sequence of Tensors): Non-empty tensors provided must have the same shape, + except in the cat dimension. + + dim (int, optional): the dimension over which the tensors are concatenated + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 0) + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 1) + tensor([[ 0.6580, -1.0969, -0.4614, 0.6580, -1.0969, -0.4614, 0.6580, + -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497, -0.1034, -0.5790, 0.1497, -0.1034, + -0.5790, 0.1497]]) + """ + +@overload +def cat( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: str | EllipsisType | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cat(tensors, dim=0, *, out=None) -> Tensor + + Concatenates the given sequence of tensors in :attr:`tensors` in the given dimension. + All tensors must either have the same shape (except in the concatenating + dimension) or be a 1-D empty tensor with size ``(0,)``. + + :func:`torch.cat` can be seen as an inverse operation for :func:`torch.split` + and :func:`torch.chunk`. + + :func:`torch.cat` can be best understood via examples. + + .. seealso:: + + :func:`torch.stack` concatenates the given sequence along a new dimension. + + Args: + tensors (sequence of Tensors): Non-empty tensors provided must have the same shape, + except in the cat dimension. + + dim (int, optional): the dimension over which the tensors are concatenated + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 0) + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 1) + tensor([[ 0.6580, -1.0969, -0.4614, 0.6580, -1.0969, -0.4614, 0.6580, + -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497, -0.1034, -0.5790, 0.1497, -0.1034, + -0.5790, 0.1497]]) + """ + +def ccol_indices_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +def ceil(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + ceil(input, *, out=None) -> Tensor + + Returns a new tensor with the ceil of the elements of :attr:`input`, + the smallest integer greater than or equal to each element. + + For integer inputs, follows the array-api convention of returning a + copy of the input tensor. + + .. math:: + \text{out}_{i} = \left\lceil \text{input}_{i} \right\rceil + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.6341, -1.4208, -1.0900, 0.5826]) + >>> torch.ceil(a) + tensor([-0., -1., -1., 1.]) + """ + +def ceil_(input: Tensor) -> Tensor: ... +def celu(input: Tensor, alpha: Number | _complex = 1.0) -> Tensor: ... +def celu_(input: Tensor, alpha: Number | _complex = 1.0) -> Tensor: ... +def channel_shuffle(input: Tensor, groups: _int | SymInt) -> Tensor: ... +def cholesky( + input: Tensor, + upper: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cholesky(input, upper=False, *, out=None) -> Tensor + + Computes the Cholesky decomposition of a symmetric positive-definite + matrix :math:`A` or for batches of symmetric positive-definite matrices. + + If :attr:`upper` is ``True``, the returned matrix ``U`` is upper-triangular, and + the decomposition has the form: + + .. math:: + + A = U^TU + + If :attr:`upper` is ``False``, the returned matrix ``L`` is lower-triangular, and + the decomposition has the form: + + .. math:: + + A = LL^T + + If :attr:`upper` is ``True``, and :math:`A` is a batch of symmetric positive-definite + matrices, then the returned tensor will be composed of upper-triangular Cholesky factors + of each of the individual matrices. Similarly, when :attr:`upper` is ``False``, the returned + tensor will be composed of lower-triangular Cholesky factors of each of the individual + matrices. + + .. warning:: + + :func:`torch.cholesky` is deprecated in favor of :func:`torch.linalg.cholesky` + and will be removed in a future PyTorch release. + + ``L = torch.cholesky(A)`` should be replaced with + + .. code:: python + + L = torch.linalg.cholesky(A) + + ``U = torch.cholesky(A, upper=True)`` should be replaced with + + .. code:: python + + U = torch.linalg.cholesky(A).mH + + This transform will produce equivalent results for all valid (symmetric positive definite) inputs. + + Args: + input (Tensor): the input tensor :math:`A` of size :math:`(*, n, n)` where `*` is zero or more + batch dimensions consisting of symmetric positive-definite matrices. + upper (bool, optional): flag that indicates whether to return a + upper or lower triangular matrix. Default: ``False`` + + Keyword args: + out (Tensor, optional): the output matrix + + Example:: + + >>> a = torch.randn(3, 3) + >>> a = a @ a.mT + 1e-3 # make symmetric positive-definite + >>> l = torch.cholesky(a) + >>> a + tensor([[ 2.4112, -0.7486, 1.4551], + [-0.7486, 1.3544, 0.1294], + [ 1.4551, 0.1294, 1.6724]]) + >>> l + tensor([[ 1.5528, 0.0000, 0.0000], + [-0.4821, 1.0592, 0.0000], + [ 0.9371, 0.5487, 0.7023]]) + >>> l @ l.mT + tensor([[ 2.4112, -0.7486, 1.4551], + [-0.7486, 1.3544, 0.1294], + [ 1.4551, 0.1294, 1.6724]]) + >>> a = torch.randn(3, 2, 2) # Example for batched input + >>> a = a @ a.mT + 1e-03 # make symmetric positive-definite + >>> l = torch.cholesky(a) + >>> z = l @ l.mT + >>> torch.dist(z, a) + tensor(2.3842e-07) + """ + +def cholesky_inverse( + input: Tensor, + upper: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cholesky_inverse(L, upper=False, *, out=None) -> Tensor + + Computes the inverse of a complex Hermitian or real symmetric + positive-definite matrix given its Cholesky decomposition. + + Let :math:`A` be a complex Hermitian or real symmetric positive-definite matrix, + and :math:`L` its Cholesky decomposition such that: + + .. math:: + + A = LL^{\text{H}} + + where :math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, + and the transpose when :math:`L` is real-valued. + + Computes the inverse matrix :math:`A^{-1}`. + + Supports input of float, double, cfloat and cdouble dtypes. + Also supports batches of matrices, and if :math:`A` is a batch of matrices + then the output has the same batch dimensions. + + Args: + L (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of lower or upper triangular Cholesky decompositions of + symmetric or Hermitian positive-definite matrices. + upper (bool, optional): flag that indicates whether :math:`L` is lower triangular + or upper triangular. Default: ``False`` + + Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + + Example:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.T + torch.eye(3) * 1e-3 # Creates a symmetric positive-definite matrix + >>> L = torch.linalg.cholesky(A) # Extract Cholesky decomposition + >>> torch.cholesky_inverse(L) + tensor([[ 1.9314, 1.2251, -0.0889], + [ 1.2251, 2.4439, 0.2122], + [-0.0889, 0.2122, 0.1412]]) + >>> A.inverse() + tensor([[ 1.9314, 1.2251, -0.0889], + [ 1.2251, 2.4439, 0.2122], + [-0.0889, 0.2122, 0.1412]]) + + >>> A = torch.randn(3, 2, 2, dtype=torch.complex64) + >>> A = A @ A.mH + torch.eye(2) * 1e-3 # Batch of Hermitian positive-definite matrices + >>> L = torch.linalg.cholesky(A) + >>> torch.dist(torch.inverse(A), torch.cholesky_inverse(L)) + tensor(5.6358e-7) + """ + +def cholesky_solve( + input: Tensor, + input2: Tensor, + upper: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cholesky_solve(B, L, upper=False, *, out=None) -> Tensor + + Computes the solution of a system of linear equations with complex Hermitian + or real symmetric positive-definite lhs given its Cholesky decomposition. + + Let :math:`A` be a complex Hermitian or real symmetric positive-definite matrix, + and :math:`L` its Cholesky decomposition such that: + + .. math:: + + A = LL^{\text{H}} + + where :math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, + and the transpose when :math:`L` is real-valued. + + Returns the solution :math:`X` of the following linear system: + + .. math:: + + AX = B + + Supports inputs of float, double, cfloat and cdouble dtypes. + Also supports batches of matrices, and if :math:`A` or :math:`B` is a batch of matrices + then the output has the same batch dimensions. + + Args: + B (Tensor): right-hand side tensor of shape `(*, n, k)` + where :math:`*` is zero or more batch dimensions + L (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of lower or upper triangular Cholesky decompositions of + symmetric or Hermitian positive-definite matrices. + upper (bool, optional): flag that indicates whether :math:`L` is lower triangular + or upper triangular. Default: ``False``. + + Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + + Example:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.T + torch.eye(3) * 1e-3 # Creates a symmetric positive-definite matrix + >>> L = torch.linalg.cholesky(A) # Extract Cholesky decomposition + >>> B = torch.randn(3, 2) + >>> torch.cholesky_solve(B, L) + tensor([[ -8.1625, 19.6097], + [ -5.8398, 14.2387], + [ -4.3771, 10.4173]]) + >>> A.inverse() @ B + tensor([[ -8.1626, 19.6097], + [ -5.8398, 14.2387], + [ -4.3771, 10.4173]]) + + >>> A = torch.randn(3, 2, 2, dtype=torch.complex64) + >>> A = A @ A.mH + torch.eye(2) * 1e-3 # Batch of Hermitian positive-definite matrices + >>> L = torch.linalg.cholesky(A) + >>> B = torch.randn(2, 1, dtype=torch.complex64) + >>> X = torch.cholesky_solve(B, L) + >>> torch.dist(X, A.inverse() @ B) + tensor(1.6881e-5) + """ + +def choose_qparams_optimized( + input: Tensor, + numel: _int, + n_bins: _int, + ratio: _float, + bit_width: _int, +) -> tuple[Tensor, Tensor]: ... +def chunk(input: Tensor, chunks: _int, dim: _int = 0) -> tuple[Tensor, ...]: + r""" + chunk(input: Tensor, chunks: int, dim: int = 0) -> Tuple[Tensor, ...] + + Attempts to split a tensor into the specified number of chunks. Each chunk is a view of + the input tensor. + + + .. note:: + + This function may return fewer than the specified number of chunks! + + .. seealso:: + + :func:`torch.tensor_split` a function that always returns exactly the specified number of chunks + + If the tensor size along the given dimension :attr:`dim` is divisible by :attr:`chunks`, + all returned chunks will be the same size. + If the tensor size along the given dimension :attr:`dim` is not divisible by :attr:`chunks`, + all returned chunks will be the same size, except the last one. + If such division is not possible, this function may return fewer + than the specified number of chunks. + + Arguments: + input (Tensor): the tensor to split + chunks (int): number of chunks to return + dim (int): dimension along which to split the tensor + + Example: + >>> torch.arange(11).chunk(6) + (tensor([0, 1]), + tensor([2, 3]), + tensor([4, 5]), + tensor([6, 7]), + tensor([8, 9]), + tensor([10])) + >>> torch.arange(12).chunk(6) + (tensor([0, 1]), + tensor([2, 3]), + tensor([4, 5]), + tensor([6, 7]), + tensor([8, 9]), + tensor([10, 11])) + >>> torch.arange(13).chunk(6) + (tensor([0, 1, 2]), + tensor([3, 4, 5]), + tensor([6, 7, 8]), + tensor([ 9, 10, 11]), + tensor([12])) + """ + +@overload +def clamp( + input: Tensor, + min: Tensor | None = None, + max: Tensor | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + clamp(input, min=None, max=None, *, out=None) -> Tensor + + Clamps all elements in :attr:`input` into the range `[` :attr:`min`, :attr:`max` `]`. + Letting min_value and max_value be :attr:`min` and :attr:`max`, respectively, this returns: + + .. math:: + y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i) + + If :attr:`min` is ``None``, there is no lower bound. + Or, if :attr:`max` is ``None`` there is no upper bound. + + + .. note:: + If :attr:`min` is greater than :attr:`max` :func:`torch.clamp(..., min, max) ` + sets all elements in :attr:`input` to the value of :attr:`max`. + + Args: + input (Tensor): the input tensor. + min (Number or Tensor, optional): lower-bound of the range to be clamped to + max (Number or Tensor, optional): upper-bound of the range to be clamped to + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-1.7120, 0.1734, -0.0478, -0.0922]) + >>> torch.clamp(a, min=-0.5, max=0.5) + tensor([-0.5000, 0.1734, -0.0478, -0.0922]) + + >>> min = torch.linspace(-1, 1, steps=4) + >>> torch.clamp(a, min=min) + tensor([-1.0000, 0.1734, 0.3333, 1.0000]) + """ + +@overload +def clamp( + input: Tensor, + min: Number | _complex | None = None, + max: Number | _complex | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + clamp(input, min=None, max=None, *, out=None) -> Tensor + + Clamps all elements in :attr:`input` into the range `[` :attr:`min`, :attr:`max` `]`. + Letting min_value and max_value be :attr:`min` and :attr:`max`, respectively, this returns: + + .. math:: + y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i) + + If :attr:`min` is ``None``, there is no lower bound. + Or, if :attr:`max` is ``None`` there is no upper bound. + + + .. note:: + If :attr:`min` is greater than :attr:`max` :func:`torch.clamp(..., min, max) ` + sets all elements in :attr:`input` to the value of :attr:`max`. + + Args: + input (Tensor): the input tensor. + min (Number or Tensor, optional): lower-bound of the range to be clamped to + max (Number or Tensor, optional): upper-bound of the range to be clamped to + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-1.7120, 0.1734, -0.0478, -0.0922]) + >>> torch.clamp(a, min=-0.5, max=0.5) + tensor([-0.5000, 0.1734, -0.0478, -0.0922]) + + >>> min = torch.linspace(-1, 1, steps=4) + >>> torch.clamp(a, min=min) + tensor([-1.0000, 0.1734, 0.3333, 1.0000]) + """ + +@overload +def clamp_( + input: Tensor, + min: Tensor | None = None, + max: Tensor | None = None, +) -> Tensor: ... +@overload +def clamp_( + input: Tensor, + min: Number | _complex | None = None, + max: Number | _complex | None = None, +) -> Tensor: ... +@overload +def clamp_max( + input: Tensor, + max: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def clamp_max( + input: Tensor, + max: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def clamp_max_(input: Tensor, max: Tensor) -> Tensor: ... +@overload +def clamp_max_(input: Tensor, max: Number | _complex) -> Tensor: ... +@overload +def clamp_min( + input: Tensor, + min: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def clamp_min( + input: Tensor, + min: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def clamp_min_(input: Tensor, min: Tensor) -> Tensor: ... +@overload +def clamp_min_(input: Tensor, min: Number | _complex) -> Tensor: ... +@overload +def clip( + input: Tensor, + min: Tensor | None = None, + max: Tensor | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + clip(input, min=None, max=None, *, out=None) -> Tensor + + Alias for :func:`torch.clamp`. + """ + +@overload +def clip( + input: Tensor, + min: Number | _complex | None = None, + max: Number | _complex | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + clip(input, min=None, max=None, *, out=None) -> Tensor + + Alias for :func:`torch.clamp`. + """ + +@overload +def clip_( + input: Tensor, + min: Tensor | None = None, + max: Tensor | None = None, +) -> Tensor: ... +@overload +def clip_( + input: Tensor, + min: Number | _complex | None = None, + max: Number | _complex | None = None, +) -> Tensor: ... +def clone( + input: Tensor, + *, + memory_format: memory_format | None = None, +) -> Tensor: + r""" + clone(input, *, memory_format=torch.preserve_format) -> Tensor + + Returns a copy of :attr:`input`. + + .. note:: + + This function is differentiable, so gradients will flow back from the + result of this operation to :attr:`input`. To create a tensor without an + autograd relationship to :attr:`input` see :meth:`~Tensor.detach`. + + In addition, when ``torch.preserve_format`` is used: + If the input tensor is dense (i.e., non-overlapping strided), + its memory format (including strides) is retained. + Otherwise (e.g., a non-dense view like a stepped slice), + the output is converted to the dense (contiguous) format. + + Args: + input (Tensor): the input tensor. + + Keyword args: + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned tensor. Default: ``torch.preserve_format``. + """ + +def col_indices_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.col_indices`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def column_stack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + column_stack(tensors, *, out=None) -> Tensor + + Creates a new tensor by horizontally stacking the tensors in :attr:`tensors`. + + Equivalent to ``torch.hstack(tensors)``, except each zero or one dimensional tensor ``t`` + in :attr:`tensors` is first reshaped into a ``(t.numel(), 1)`` column before being stacked horizontally. + + Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.column_stack((a, b)) + tensor([[1, 4], + [2, 5], + [3, 6]]) + >>> a = torch.arange(5) + >>> b = torch.arange(10).reshape(5, 2) + >>> torch.column_stack((a, b, b)) + tensor([[0, 0, 1, 0, 1], + [1, 2, 3, 2, 3], + [2, 4, 5, 4, 5], + [3, 6, 7, 6, 7], + [4, 8, 9, 8, 9]]) + """ + +def combinations( + input: Tensor, + r: _int = 2, + with_replacement: _bool = False, +) -> Tensor: + r""" + combinations(input: Tensor, r: int = 2, with_replacement: bool = False) -> seq + + Compute combinations of length :math:`r` of the given tensor. The behavior is similar to + python's `itertools.combinations` when `with_replacement` is set to `False`, and + `itertools.combinations_with_replacement` when `with_replacement` is set to `True`. + + Arguments: + input (Tensor): 1D vector. + r (int, optional): number of elements to combine + with_replacement (bool, optional): whether to allow duplication in combination + + Returns: + Tensor: A tensor equivalent to converting all the input tensors into lists, do + `itertools.combinations` or `itertools.combinations_with_replacement` on these + lists, and finally convert the resulting list into tensor. + + Example:: + + >>> a = [1, 2, 3] + >>> list(itertools.combinations(a, r=2)) + [(1, 2), (1, 3), (2, 3)] + >>> list(itertools.combinations(a, r=3)) + [(1, 2, 3)] + >>> list(itertools.combinations_with_replacement(a, r=2)) + [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)] + >>> tensor_a = torch.tensor(a) + >>> torch.combinations(tensor_a) + tensor([[1, 2], + [1, 3], + [2, 3]]) + >>> torch.combinations(tensor_a, r=3) + tensor([[1, 2, 3]]) + >>> torch.combinations(tensor_a, with_replacement=True) + tensor([[1, 1], + [1, 2], + [1, 3], + [2, 2], + [2, 3], + [3, 3]]) + """ + +def complex( + real: Tensor, + imag: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + complex(real, imag, *, out=None) -> Tensor + + Constructs a complex tensor with its real part equal to :attr:`real` and its + imaginary part equal to :attr:`imag`. + + Args: + real (Tensor): The real part of the complex tensor. Must be half, float or double. + imag (Tensor): The imaginary part of the complex tensor. Must be same dtype + as :attr:`real`. + + Keyword args: + out (Tensor): If the inputs are ``torch.float32``, must be + ``torch.complex64``. If the inputs are ``torch.float64``, must be + ``torch.complex128``. + + Example:: + + >>> real = torch.tensor([1, 2], dtype=torch.float32) + >>> imag = torch.tensor([3, 4], dtype=torch.float32) + >>> z = torch.complex(real, imag) + >>> z + tensor([(1.+3.j), (2.+4.j)]) + >>> z.dtype + torch.complex64 + """ + +@overload +def concat( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + concat(tensors, dim=0, *, out=None) -> Tensor + + Alias of :func:`torch.cat`. + """ + +@overload +def concat( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: str | EllipsisType | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + concat(tensors, dim=0, *, out=None) -> Tensor + + Alias of :func:`torch.cat`. + """ + +@overload +def concatenate( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + concatenate(tensors, axis=0, out=None) -> Tensor + + Alias of :func:`torch.cat`. + """ + +@overload +def concatenate( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: str | EllipsisType | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + concatenate(tensors, axis=0, out=None) -> Tensor + + Alias of :func:`torch.cat`. + """ + +def conj(input: Tensor) -> Tensor: + r""" + conj(input) -> Tensor + + Returns a view of :attr:`input` with a flipped conjugate bit. If :attr:`input` has a non-complex dtype, + this function just returns :attr:`input`. + + .. note:: + :func:`torch.conj` performs a lazy conjugation, but the actual conjugated tensor can be materialized + at any time using :func:`torch.resolve_conj`. + + .. warning:: In the future, :func:`torch.conj` may return a non-writeable view for an :attr:`input` of + non-complex dtype. It's recommended that programs not modify the tensor returned by :func:`torch.conj_physical` + when :attr:`input` is of non-complex dtype to be compatible with this change. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> x.is_conj() + False + >>> y = torch.conj(x) + >>> y.is_conj() + True + """ + +def conj_physical(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + conj_physical(input, *, out=None) -> Tensor + + Computes the element-wise conjugate of the given :attr:`input` tensor. + If :attr:`input` has a non-complex dtype, this function just returns :attr:`input`. + + .. note:: + This performs the conjugate operation regardless of the fact conjugate bit is set or not. + + .. warning:: In the future, :func:`torch.conj_physical` may return a non-writeable view for an :attr:`input` of + non-complex dtype. It's recommended that programs not modify the tensor returned by :func:`torch.conj_physical` + when :attr:`input` is of non-complex dtype to be compatible with this change. + + .. math:: + \text{out}_{i} = conj(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.conj_physical(torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j])) + tensor([-1 - 1j, -2 - 2j, 3 + 3j]) + """ + +def conj_physical_(input: Tensor) -> Tensor: ... +def constant_pad_nd( + input: Tensor, + pad: Sequence[_int | SymInt], + value: Number | _complex = 0, +) -> Tensor: ... +@overload +def conv1d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +@overload +def conv1d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: str = "valid", + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +@overload +def conv2d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +@overload +def conv2d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: str = "valid", + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +@overload +def conv3d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +@overload +def conv3d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: str = "valid", + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, + groups: _int | SymInt = 1, +) -> Tensor: ... +def conv_tbc( + input: Tensor, + weight: Tensor, + bias: Tensor, + pad: _int = 0, +) -> Tensor: ... +def conv_transpose1d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + output_padding: _int | SymInt | Sequence[_int | SymInt] = 0, + groups: _int | SymInt = 1, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, +) -> Tensor: ... +def conv_transpose2d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + output_padding: _int | SymInt | Sequence[_int | SymInt] = 0, + groups: _int | SymInt = 1, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, +) -> Tensor: ... +def conv_transpose3d( + input: Tensor, + weight: Tensor, + bias: Tensor | None = None, + stride: _int | SymInt | Sequence[_int | SymInt] = 1, + padding: _int | SymInt | Sequence[_int | SymInt] = 0, + output_padding: _int | SymInt | Sequence[_int | SymInt] = 0, + groups: _int | SymInt = 1, + dilation: _int | SymInt | Sequence[_int | SymInt] = 1, +) -> Tensor: ... +def convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + transposed: _bool, + output_padding: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +@overload +def copysign( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + copysign(input, other, *, out=None) -> Tensor + + Create a new floating-point tensor with the magnitude of :attr:`input` and the sign of :attr:`other`, elementwise. + + .. math:: + \text{out}_{i} = \begin{cases} + -|\text{input}_{i}| & \text{if } \text{other}_{i} \leq -0.0 \\ + |\text{input}_{i}| & \text{if } \text{other}_{i} \geq 0.0 \\ + \end{cases} + + + Supports :ref:`broadcasting to a common shape `, + and integer and float inputs. + + Args: + input (Tensor): magnitudes. + other (Tensor or Number): contains value(s) whose signbit(s) are + applied to the magnitudes in :attr:`input`. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(5) + >>> a + tensor([-1.2557, -0.0026, -0.5387, 0.4740, -0.9244]) + >>> torch.copysign(a, 1) + tensor([1.2557, 0.0026, 0.5387, 0.4740, 0.9244]) + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.7079, 0.2778, -1.0249, 0.5719], + [-0.0059, -0.2600, -0.4475, -1.3948], + [ 0.3667, -0.9567, -2.5757, -0.1751], + [ 0.2046, -0.0742, 0.2998, -0.1054]]) + >>> b = torch.randn(4) + tensor([ 0.2373, 0.3120, 0.3190, -1.1128]) + >>> torch.copysign(a, b) + tensor([[ 0.7079, 0.2778, 1.0249, -0.5719], + [ 0.0059, 0.2600, 0.4475, -1.3948], + [ 0.3667, 0.9567, 2.5757, -0.1751], + [ 0.2046, 0.0742, 0.2998, -0.1054]]) + >>> a = torch.tensor([1.]) + >>> b = torch.tensor([-0.]) + >>> torch.copysign(a, b) + tensor([-1.]) + + .. note:: + copysign handles signed zeros. If the other argument has a negative zero (-0), + the corresponding output value will be negative. + """ + +@overload +def copysign( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + copysign(input, other, *, out=None) -> Tensor + + Create a new floating-point tensor with the magnitude of :attr:`input` and the sign of :attr:`other`, elementwise. + + .. math:: + \text{out}_{i} = \begin{cases} + -|\text{input}_{i}| & \text{if } \text{other}_{i} \leq -0.0 \\ + |\text{input}_{i}| & \text{if } \text{other}_{i} \geq 0.0 \\ + \end{cases} + + + Supports :ref:`broadcasting to a common shape `, + and integer and float inputs. + + Args: + input (Tensor): magnitudes. + other (Tensor or Number): contains value(s) whose signbit(s) are + applied to the magnitudes in :attr:`input`. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(5) + >>> a + tensor([-1.2557, -0.0026, -0.5387, 0.4740, -0.9244]) + >>> torch.copysign(a, 1) + tensor([1.2557, 0.0026, 0.5387, 0.4740, 0.9244]) + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.7079, 0.2778, -1.0249, 0.5719], + [-0.0059, -0.2600, -0.4475, -1.3948], + [ 0.3667, -0.9567, -2.5757, -0.1751], + [ 0.2046, -0.0742, 0.2998, -0.1054]]) + >>> b = torch.randn(4) + tensor([ 0.2373, 0.3120, 0.3190, -1.1128]) + >>> torch.copysign(a, b) + tensor([[ 0.7079, 0.2778, 1.0249, -0.5719], + [ 0.0059, 0.2600, 0.4475, -1.3948], + [ 0.3667, 0.9567, 2.5757, -0.1751], + [ 0.2046, 0.0742, 0.2998, -0.1054]]) + >>> a = torch.tensor([1.]) + >>> b = torch.tensor([-0.]) + >>> torch.copysign(a, b) + tensor([-1.]) + + .. note:: + copysign handles signed zeros. If the other argument has a negative zero (-0), + the corresponding output value will be negative. + """ + +def corrcoef(input: Tensor) -> Tensor: + r""" + corrcoef(input) -> Tensor + + Estimates the Pearson product-moment correlation coefficient matrix of the variables given by the :attr:`input` matrix, + where rows are the variables and columns are the observations. + + .. note:: + + The correlation coefficient matrix R is computed using the covariance matrix C as given by + :math:`R_{ij} = \frac{ C_{ij} } { \sqrt{ C_{ii} * C_{jj} } }` + + .. note:: + + Due to floating point rounding, the resulting array may not be Hermitian and its diagonal elements may not be 1. + The real and imaginary values are clipped to the interval [-1, 1] in an attempt to improve this situation. + + Args: + input (Tensor): A 2D matrix containing multiple variables and observations, or a + Scalar or 1D vector representing a single variable. + + Returns: + (Tensor) The correlation coefficient matrix of the variables. + + .. seealso:: + + :func:`torch.cov` covariance matrix. + + Example:: + + >>> x = torch.tensor([[0, 1, 2], [2, 1, 0]]) + >>> torch.corrcoef(x) + tensor([[ 1., -1.], + [-1., 1.]]) + >>> x = torch.randn(2, 4) + >>> x + tensor([[-0.2678, -0.0908, -0.3766, 0.2780], + [-0.5812, 0.1535, 0.2387, 0.2350]]) + >>> torch.corrcoef(x) + tensor([[1.0000, 0.3582], + [0.3582, 1.0000]]) + >>> torch.corrcoef(x[0]) + tensor(1.) + """ + +def cos(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + cos(input, *, out=None) -> Tensor + + Returns a new tensor with the cosine of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \cos(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 1.4309, 1.2706, -0.8562, 0.9796]) + >>> torch.cos(a) + tensor([ 0.1395, 0.2957, 0.6553, 0.5574]) + """ + +def cos_(input: Tensor) -> Tensor: ... +def cosh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + cosh(input, *, out=None) -> Tensor + + Returns a new tensor with the hyperbolic cosine of the elements of + :attr:`input`. + + .. math:: + \text{out}_{i} = \cosh(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.1632, 1.1835, -0.6979, -0.7325]) + >>> torch.cosh(a) + tensor([ 1.0133, 1.7860, 1.2536, 1.2805]) + + .. note:: + When :attr:`input` is on the CPU, the implementation of torch.cosh may use + the Sleef library, which rounds very large results to infinity or negative + infinity. See `here `_ for details. + """ + +def cosh_(input: Tensor) -> Tensor: ... +def cosine_embedding_loss( + input1: Tensor, + input2: Tensor, + target: Tensor, + margin: _float = 0.0, + reduction: _int = 1, +) -> Tensor: ... +def cosine_similarity( + x1: Tensor, + x2: Tensor, + dim: _int = 1, + eps: _float = 1e-08, +) -> Tensor: ... +@overload +def count_nonzero(input: Tensor, dim: _int | None = None) -> Tensor: + r""" + count_nonzero(input, dim=None) -> Tensor + + Counts the number of non-zero values in the tensor :attr:`input` along the given :attr:`dim`. + If no dim is specified then all non-zeros in the tensor are counted. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): Dim or tuple of dims along which to count non-zeros. + + Example:: + + >>> x = torch.zeros(3,3) + >>> x[torch.randn(3,3) > 0.5] = 1 + >>> x + tensor([[0., 1., 1.], + [0., 0., 0.], + [0., 0., 1.]]) + >>> torch.count_nonzero(x) + tensor(3) + >>> torch.count_nonzero(x, dim=0) + tensor([0, 1, 2]) + """ + +@overload +def count_nonzero(input: Tensor, dim: _size) -> Tensor: + r""" + count_nonzero(input, dim=None) -> Tensor + + Counts the number of non-zero values in the tensor :attr:`input` along the given :attr:`dim`. + If no dim is specified then all non-zeros in the tensor are counted. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): Dim or tuple of dims along which to count non-zeros. + + Example:: + + >>> x = torch.zeros(3,3) + >>> x[torch.randn(3,3) > 0.5] = 1 + >>> x + tensor([[0., 1., 1.], + [0., 0., 0.], + [0., 0., 1.]]) + >>> torch.count_nonzero(x) + tensor(3) + >>> torch.count_nonzero(x, dim=0) + tensor([0, 1, 2]) + """ + +def cov( + input: Tensor, + *, + correction: _int = 1, + fweights: Tensor | None = None, + aweights: Tensor | None = None, +) -> Tensor: + r""" + cov(input, *, correction=1, fweights=None, aweights=None) -> Tensor + + Estimates the covariance matrix of the variables given by the :attr:`input` matrix, where rows are + the variables and columns are the observations. + + A covariance matrix is a square matrix giving the covariance of each pair of variables. The diagonal contains + the variance of each variable (covariance of a variable with itself). By definition, if :attr:`input` represents + a single variable (Scalar or 1D) then its variance is returned. + + The sample covariance of the variables :math:`x` and :math:`y` is given by: + + .. math:: + \text{cov}(x,y) = \frac{\sum^{N}_{i = 1}(x_{i} - \bar{x})(y_{i} - \bar{y})}{\max(0,~N~-~\delta N)} + + where :math:`\bar{x}` and :math:`\bar{y}` are the simple means of the :math:`x` and :math:`y` respectively, and + :math:`\delta N` is the :attr:`correction`. + + If :attr:`fweights` and/or :attr:`aweights` are provided, the weighted covariance + is calculated, which is given by: + + .. math:: + \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}w_i(x_{i} - \mu_x^*)(y_{i} - \mu_y^*)} + {\max(0,~\sum^{N}_{i = 1}w_i~-~\frac{\sum^{N}_{i = 1}w_ia_i}{\sum^{N}_{i = 1}w_i}~\delta N)} + + where :math:`w` denotes :attr:`fweights` or :attr:`aweights` (``f`` and ``a`` for brevity) based on whichever is + provided, or :math:`w = f \times a` if both are provided, and + :math:`\mu_x^* = \frac{\sum^{N}_{i = 1}w_ix_{i} }{\sum^{N}_{i = 1}w_i}` is the weighted mean of the variable. If not + provided, ``f`` and/or ``a`` can be seen as a :math:`\mathbb{1}` vector of appropriate size. + + Args: + input (Tensor): A 2D matrix containing multiple variables and observations, or a + Scalar or 1D vector representing a single variable. + + Keyword Args: + correction (int, optional): difference between the sample size and sample degrees of freedom. + Defaults to Bessel's correction, ``correction = 1`` which returns the unbiased estimate, + even if both :attr:`fweights` and :attr:`aweights` are specified. ``correction = 0`` + will return the simple average. Defaults to ``1``. + fweights (tensor, optional): A Scalar or 1D tensor of observation vector frequencies representing the number of + times each observation should be repeated. Its numel must equal the number of columns of :attr:`input`. + Must have integral dtype. Ignored if ``None``. Defaults to ``None``. + aweights (tensor, optional): A Scalar or 1D array of observation vector weights. + These relative weights are typically large for observations considered "important" and smaller for + observations considered less "important". Its numel must equal the number of columns of :attr:`input`. + Must have floating point dtype. Ignored if ``None``. Defaults to ``None``. + + Returns: + (Tensor) The covariance matrix of the variables. + + .. seealso:: + + :func:`torch.corrcoef` normalized covariance matrix. + + Example:: + + >>> x = torch.tensor([[0, 2], [1, 1], [2, 0]]).T + >>> x + tensor([[0, 1, 2], + [2, 1, 0]]) + >>> torch.cov(x) + tensor([[ 1., -1.], + [-1., 1.]]) + >>> torch.cov(x, correction=0) + tensor([[ 0.6667, -0.6667], + [-0.6667, 0.6667]]) + >>> fw = torch.randint(1, 10, (3,)) + >>> fw + tensor([1, 6, 9]) + >>> aw = torch.rand(3) + >>> aw + tensor([0.4282, 0.0255, 0.4144]) + >>> torch.cov(x, fweights=fw, aweights=aw) + tensor([[ 0.4169, -0.4169], + [-0.4169, 0.4169]]) + """ + +def cross( + input: Tensor, + other: Tensor, + dim: _int | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + cross(input, other, dim=None, *, out=None) -> Tensor + + + Returns the cross product of vectors in dimension :attr:`dim` of :attr:`input` + and :attr:`other`. + + Supports input of float, double, cfloat and cdouble dtypes. Also supports batches + of vectors, for which it computes the product along the dimension :attr:`dim`. + In this case, the output has the same batch dimensions as the inputs. + + .. warning:: + If :attr:`dim` is not given, it defaults to the first dimension found + with the size 3. Note that this might be unexpected. + + This behavior is deprecated and will be changed to match that of :func:`torch.linalg.cross` + in a future release. + + .. seealso:: + :func:`torch.linalg.cross` which has dim=-1 as default. + + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + dim (int, optional): the dimension to take the cross-product in. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4, 3) + >>> a + tensor([[-0.3956, 1.1455, 1.6895], + [-0.5849, 1.3672, 0.3599], + [-1.1626, 0.7180, -0.0521], + [-0.1339, 0.9902, -2.0225]]) + >>> b = torch.randn(4, 3) + >>> b + tensor([[-0.0257, -1.4725, -1.2251], + [-1.1479, -0.7005, -1.9757], + [-1.3904, 0.3726, -1.1836], + [-0.9688, -0.7153, 0.2159]]) + >>> torch.cross(a, b, dim=1) + tensor([[ 1.0844, -0.5281, 0.6120], + [-2.4490, -1.5687, 1.9792], + [-0.8304, -1.3037, 0.5650], + [-1.2329, 1.9883, 1.0551]]) + >>> torch.cross(a, b) + tensor([[ 1.0844, -0.5281, 0.6120], + [-2.4490, -1.5687, 1.9792], + [-0.8304, -1.3037, 0.5650], + [-1.2329, 1.9883, 1.0551]]) + """ + +def crow_indices_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.crow_indices`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: _size, + target_lengths: _size, + blank: _int = 0, + reduction: _int = 1, + zero_infinity: _bool = False, +) -> Tensor: ... +@overload +def ctc_loss( + log_probs: Tensor, + targets: Tensor, + input_lengths: Tensor, + target_lengths: Tensor, + blank: _int = 0, + reduction: _int = 1, + zero_infinity: _bool = False, +) -> Tensor: ... +def cudnn_affine_grid_generator( + theta: Tensor, + N: _int, + C: _int, + H: _int, + W: _int, +) -> Tensor: ... +def cudnn_batch_norm( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + training: _bool, + exponential_average_factor: _float, + epsilon: _float, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +def cudnn_convolution( + input: Tensor, + weight: Tensor, + padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, + allow_tf32: _bool, + *, + out: Tensor | None = None, +) -> Tensor: ... +def cudnn_convolution_add_relu( + input: Tensor, + weight: Tensor, + z: Tensor, + alpha: Number | _complex | None, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def cudnn_convolution_relu( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def cudnn_convolution_transpose( + input: Tensor, + weight: Tensor, + padding: Sequence[_int | SymInt], + output_padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, + allow_tf32: _bool, +) -> Tensor: ... +def cudnn_grid_sampler(input: Tensor, grid: Tensor) -> Tensor: ... +def cudnn_is_acceptable(input: Tensor) -> _bool: ... +@overload +def cummax( + input: Tensor, + dim: _int, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.cummax: + r""" + cummax(input, dim, *, out=None) -> (Tensor, LongTensor) + Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative maximum of + elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index + location of each maximum value found in the dimension :attr:`dim`. + + .. math:: + y_i = max(x_1, x_2, x_3, \dots, x_i) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.3449, -1.5447, 0.0685, -1.5104, -1.1706, 0.2259, 1.4696, -1.3284, + 1.9946, -0.8209]) + >>> torch.cummax(a, dim=0) + torch.return_types.cummax( + values=tensor([-0.3449, -0.3449, 0.0685, 0.0685, 0.0685, 0.2259, 1.4696, 1.4696, + 1.9946, 1.9946]), + indices=tensor([0, 0, 2, 2, 2, 5, 6, 6, 8, 8])) + """ + +@overload +def cummax( + input: Tensor, + dim: str | EllipsisType | None, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.cummax: + r""" + cummax(input, dim, *, out=None) -> (Tensor, LongTensor) + Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative maximum of + elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index + location of each maximum value found in the dimension :attr:`dim`. + + .. math:: + y_i = max(x_1, x_2, x_3, \dots, x_i) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.3449, -1.5447, 0.0685, -1.5104, -1.1706, 0.2259, 1.4696, -1.3284, + 1.9946, -0.8209]) + >>> torch.cummax(a, dim=0) + torch.return_types.cummax( + values=tensor([-0.3449, -0.3449, 0.0685, 0.0685, 0.0685, 0.2259, 1.4696, 1.4696, + 1.9946, 1.9946]), + indices=tensor([0, 0, 2, 2, 2, 5, 6, 6, 8, 8])) + """ + +@overload +def cummin( + input: Tensor, + dim: _int, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.cummin: + r""" + cummin(input, dim, *, out=None) -> (Tensor, LongTensor) + Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative minimum of + elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index + location of each maximum value found in the dimension :attr:`dim`. + + .. math:: + y_i = min(x_1, x_2, x_3, \dots, x_i) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.2284, -0.6628, 0.0975, 0.2680, -1.3298, -0.4220, -0.3885, 1.1762, + 0.9165, 1.6684]) + >>> torch.cummin(a, dim=0) + torch.return_types.cummin( + values=tensor([-0.2284, -0.6628, -0.6628, -0.6628, -1.3298, -1.3298, -1.3298, -1.3298, + -1.3298, -1.3298]), + indices=tensor([0, 1, 1, 1, 4, 4, 4, 4, 4, 4])) + """ + +@overload +def cummin( + input: Tensor, + dim: str | EllipsisType | None, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.cummin: + r""" + cummin(input, dim, *, out=None) -> (Tensor, LongTensor) + Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative minimum of + elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index + location of each maximum value found in the dimension :attr:`dim`. + + .. math:: + y_i = min(x_1, x_2, x_3, \dots, x_i) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.2284, -0.6628, 0.0975, 0.2680, -1.3298, -0.4220, -0.3885, 1.1762, + 0.9165, 1.6684]) + >>> torch.cummin(a, dim=0) + torch.return_types.cummin( + values=tensor([-0.2284, -0.6628, -0.6628, -0.6628, -1.3298, -1.3298, -1.3298, -1.3298, + -1.3298, -1.3298]), + indices=tensor([0, 1, 1, 1, 4, 4, 4, 4, 4, 4])) + """ + +@overload +def cumprod( + input: Tensor, + dim: _int, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + cumprod(input, dim, *, dtype=None, out=None) -> Tensor + + Returns the cumulative product of elements of :attr:`input` in the dimension + :attr:`dim`. + + For example, if :attr:`input` is a vector of size N, the result will also be + a vector of size N, with elements. + + .. math:: + y_i = x_1 \times x_2\times x_3\times \dots \times x_i + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([ 0.6001, 0.2069, -0.1919, 0.9792, 0.6727, 1.0062, 0.4126, + -0.2129, -0.4206, 0.1968]) + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0158, -0.0065, + 0.0014, -0.0006, -0.0001]) + + >>> a[5] = 0.0 + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0000, -0.0000, + 0.0000, -0.0000, -0.0000]) + """ + +@overload +def cumprod( + input: Tensor, + dim: str | EllipsisType | None, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + cumprod(input, dim, *, dtype=None, out=None) -> Tensor + + Returns the cumulative product of elements of :attr:`input` in the dimension + :attr:`dim`. + + For example, if :attr:`input` is a vector of size N, the result will also be + a vector of size N, with elements. + + .. math:: + y_i = x_1 \times x_2\times x_3\times \dots \times x_i + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(10) + >>> a + tensor([ 0.6001, 0.2069, -0.1919, 0.9792, 0.6727, 1.0062, 0.4126, + -0.2129, -0.4206, 0.1968]) + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0158, -0.0065, + 0.0014, -0.0006, -0.0001]) + + >>> a[5] = 0.0 + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0000, -0.0000, + 0.0000, -0.0000, -0.0000]) + """ + +@overload +def cumsum( + input: Tensor, + dim: _int, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + cumsum(input, dim, *, dtype=None, out=None) -> Tensor + + Returns the cumulative sum of elements of :attr:`input` in the dimension + :attr:`dim`. + + For example, if :attr:`input` is a vector of size N, the result will also be + a vector of size N, with elements. + + .. math:: + y_i = x_1 + x_2 + x_3 + \dots + x_i + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randint(1, 20, (10,)) + >>> a + tensor([13, 7, 3, 10, 13, 3, 15, 10, 9, 10]) + >>> torch.cumsum(a, dim=0) + tensor([13, 20, 23, 33, 46, 49, 64, 74, 83, 93]) + """ + +@overload +def cumsum( + input: Tensor, + dim: str | EllipsisType | None, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + cumsum(input, dim, *, dtype=None, out=None) -> Tensor + + Returns the cumulative sum of elements of :attr:`input` in the dimension + :attr:`dim`. + + For example, if :attr:`input` is a vector of size N, the result will also be + a vector of size N, with elements. + + .. math:: + y_i = x_1 + x_2 + x_3 + \dots + x_i + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randint(1, 20, (10,)) + >>> a + tensor([13, 7, 3, 10, 13, 3, 15, 10, 9, 10]) + >>> torch.cumsum(a, dim=0) + tensor([13, 20, 23, 33, 46, 49, 64, 74, 83, 93]) + """ + +@overload +def cumulative_trapezoid(y: Tensor, x: Tensor, *, dim: _int = -1) -> Tensor: + r""" + cumulative_trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + + Cumulatively computes the `trapezoidal rule `_ + along :attr:`dim`. By default the spacing between elements is assumed to be 1, but + :attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be + used to specify arbitrary spacing along :attr:`dim`. + + For more details, please read :func:`torch.trapezoid`. The difference between :func:`torch.trapezoid` + and this function is that, :func:`torch.trapezoid` returns a value for each integration, + where as this function returns a cumulative value for every spacing within the integration. This + is analogous to how `.sum` returns a value and `.cumsum` returns a cumulative sum. + + Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + + Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + + Examples:: + + >>> # Cumulatively computes the trapezoidal rule in 1D, spacing is implicitly 1. + >>> y = torch.tensor([1, 5, 10]) + >>> torch.cumulative_trapezoid(y) + tensor([3., 10.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> (1 + 5) / 2 + 3.0 + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Cumulatively computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.cumulative_trapezoid(y, dx=2) + tensor([6., 21.]) + + >>> # Cumulatively computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([6., 28.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> ((3 - 1) * (1 + 5)) / 2 + 6.0 + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.cumulative_trapezoid(y) + tensor([[ 0.5, 2.], + [ 3.5, 8.], + [ 6.5, 14.]]) + + >>> # Cumulatively computes the trapezoidal rule for each column of the matrix + >>> torch.cumulative_trapezoid(y, dim=0) + tensor([[ 1.5, 2.5, 3.5], + [ 6.0, 8.0, 10.0]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[2., 5.], + [2., 5.], + [2., 5.]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[1., 2.], + [2., 4.], + [3., 6.]]) + """ + +@overload +def cumulative_trapezoid( + y: Tensor, + *, + dx: Number | _complex = 1, + dim: _int = -1, +) -> Tensor: + r""" + cumulative_trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + + Cumulatively computes the `trapezoidal rule `_ + along :attr:`dim`. By default the spacing between elements is assumed to be 1, but + :attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be + used to specify arbitrary spacing along :attr:`dim`. + + For more details, please read :func:`torch.trapezoid`. The difference between :func:`torch.trapezoid` + and this function is that, :func:`torch.trapezoid` returns a value for each integration, + where as this function returns a cumulative value for every spacing within the integration. This + is analogous to how `.sum` returns a value and `.cumsum` returns a cumulative sum. + + Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + + Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + + Examples:: + + >>> # Cumulatively computes the trapezoidal rule in 1D, spacing is implicitly 1. + >>> y = torch.tensor([1, 5, 10]) + >>> torch.cumulative_trapezoid(y) + tensor([3., 10.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> (1 + 5) / 2 + 3.0 + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Cumulatively computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.cumulative_trapezoid(y, dx=2) + tensor([6., 21.]) + + >>> # Cumulatively computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([6., 28.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> ((3 - 1) * (1 + 5)) / 2 + 6.0 + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.cumulative_trapezoid(y) + tensor([[ 0.5, 2.], + [ 3.5, 8.], + [ 6.5, 14.]]) + + >>> # Cumulatively computes the trapezoidal rule for each column of the matrix + >>> torch.cumulative_trapezoid(y, dim=0) + tensor([[ 1.5, 2.5, 3.5], + [ 6.0, 8.0, 10.0]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[2., 5.], + [2., 5.], + [2., 5.]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[1., 2.], + [2., 4.], + [3., 6.]]) + """ + +def deg2rad(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + deg2rad(input, *, out=None) -> Tensor + + Returns a new tensor with each of the elements of :attr:`input` + converted from angles in degrees to radians. + + Args: + input (Tensor): the input tensor. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([[180.0, -180.0], [360.0, -360.0], [90.0, -90.0]]) + >>> torch.deg2rad(a) + tensor([[ 3.1416, -3.1416], + [ 6.2832, -6.2832], + [ 1.5708, -1.5708]]) + """ + +def deg2rad_(input: Tensor) -> Tensor: ... +@overload +def dequantize(input: Tensor) -> Tensor: + r""" + dequantize(tensor) -> Tensor + + Returns an fp32 Tensor by dequantizing a quantized Tensor + + Args: + tensor (Tensor): A quantized Tensor + + .. function:: dequantize(tensors) -> sequence of Tensors + :noindex: + + Given a list of quantized Tensors, dequantize them and return a list of fp32 Tensors + + Args: + tensors (sequence of Tensors): A list of quantized Tensors + """ + +@overload +def dequantize( + tensors: tuple[Tensor, ...] | list[Tensor] | None, +) -> tuple[Tensor, ...]: + r""" + dequantize(tensor) -> Tensor + + Returns an fp32 Tensor by dequantizing a quantized Tensor + + Args: + tensor (Tensor): A quantized Tensor + + .. function:: dequantize(tensors) -> sequence of Tensors + :noindex: + + Given a list of quantized Tensors, dequantize them and return a list of fp32 Tensors + + Args: + tensors (sequence of Tensors): A list of quantized Tensors + """ + +def det(input: Tensor) -> Tensor: + r""" + det(input) -> Tensor + + Alias for :func:`torch.linalg.det` + """ + +def detach(input: Tensor) -> Tensor: ... +def detach_(input: Tensor) -> Tensor: ... +def detach_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.detach`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def diag( + input: Tensor, + diagonal: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + diag(input, diagonal=0, *, out=None) -> Tensor + + - If :attr:`input` is a vector (1-D tensor), then returns a 2-D square tensor + with the elements of :attr:`input` as the diagonal. + - If :attr:`input` is a matrix (2-D tensor), then returns a 1-D tensor with + the diagonal elements of :attr:`input`. + + The argument :attr:`diagonal` controls which diagonal to consider: + + - If :attr:`diagonal` = 0, it is the main diagonal. + - If :attr:`diagonal` > 0, it is above the main diagonal. + - If :attr:`diagonal` < 0, it is below the main diagonal. + + Args: + input (Tensor): the input tensor. + diagonal (int, optional): the diagonal to consider + + Keyword args: + out (Tensor, optional): the output tensor. + + .. seealso:: + + :func:`torch.diagonal` always returns the diagonal of its input. + + :func:`torch.diagflat` always constructs a tensor with diagonal elements + specified by the input. + + Examples: + + Get the square matrix where the input vector is the diagonal:: + + >>> a = torch.randn(3) + >>> a + tensor([ 0.5950,-0.0872, 2.3298]) + >>> torch.diag(a) + tensor([[ 0.5950, 0.0000, 0.0000], + [ 0.0000,-0.0872, 0.0000], + [ 0.0000, 0.0000, 2.3298]]) + >>> torch.diag(a, 1) + tensor([[ 0.0000, 0.5950, 0.0000, 0.0000], + [ 0.0000, 0.0000,-0.0872, 0.0000], + [ 0.0000, 0.0000, 0.0000, 2.3298], + [ 0.0000, 0.0000, 0.0000, 0.0000]]) + + Get the k-th diagonal of a given matrix:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-0.4264, 0.0255,-0.1064], + [ 0.8795,-0.2429, 0.1374], + [ 0.1029,-0.6482,-1.6300]]) + >>> torch.diag(a, 0) + tensor([-0.4264,-0.2429,-1.6300]) + >>> torch.diag(a, 1) + tensor([ 0.0255, 0.1374]) + """ + +def diag_embed( + input: Tensor, + offset: _int = 0, + dim1: _int = -2, + dim2: _int = -1, +) -> Tensor: + r""" + diag_embed(input, offset=0, dim1=-2, dim2=-1) -> Tensor + + Creates a tensor whose diagonals of certain 2D planes (specified by + :attr:`dim1` and :attr:`dim2`) are filled by :attr:`input`. + To facilitate creating batched diagonal matrices, the 2D planes formed by + the last two dimensions of the returned tensor are chosen by default. + + The argument :attr:`offset` controls which diagonal to consider: + + - If :attr:`offset` = 0, it is the main diagonal. + - If :attr:`offset` > 0, it is above the main diagonal. + - If :attr:`offset` < 0, it is below the main diagonal. + + The size of the new matrix will be calculated to make the specified diagonal + of the size of the last input dimension. + Note that for :attr:`offset` other than :math:`0`, the order of :attr:`dim1` + and :attr:`dim2` matters. Exchanging them is equivalent to changing the + sign of :attr:`offset`. + + Applying :meth:`torch.diagonal` to the output of this function with + the same arguments yields a matrix identical to input. However, + :meth:`torch.diagonal` has different default dimensions, so those + need to be explicitly specified. + + Args: + input (Tensor): the input tensor. Must be at least 1-dimensional. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: -2. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: -1. + + Example:: + + >>> a = torch.randn(2, 3) + >>> torch.diag_embed(a) + tensor([[[ 1.5410, 0.0000, 0.0000], + [ 0.0000, -0.2934, 0.0000], + [ 0.0000, 0.0000, -2.1788]], + + [[ 0.5684, 0.0000, 0.0000], + [ 0.0000, -1.0845, 0.0000], + [ 0.0000, 0.0000, -1.3986]]]) + + >>> torch.diag_embed(a, offset=1, dim1=0, dim2=2) + tensor([[[ 0.0000, 1.5410, 0.0000, 0.0000], + [ 0.0000, 0.5684, 0.0000, 0.0000]], + + [[ 0.0000, 0.0000, -0.2934, 0.0000], + [ 0.0000, 0.0000, -1.0845, 0.0000]], + + [[ 0.0000, 0.0000, 0.0000, -2.1788], + [ 0.0000, 0.0000, 0.0000, -1.3986]], + + [[ 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.0000, 0.0000, 0.0000, 0.0000]]]) + """ + +def diagflat(input: Tensor, offset: _int = 0) -> Tensor: + r""" + diagflat(input, offset=0) -> Tensor + + - If :attr:`input` is a vector (1-D tensor), then returns a 2-D square tensor + with the elements of :attr:`input` as the diagonal. + - If :attr:`input` is a tensor with more than one dimension, then returns a + 2-D tensor with diagonal elements equal to a flattened :attr:`input`. + + The argument :attr:`offset` controls which diagonal to consider: + + - If :attr:`offset` = 0, it is the main diagonal. + - If :attr:`offset` > 0, it is above the main diagonal. + - If :attr:`offset` < 0, it is below the main diagonal. + + Args: + input (Tensor): the input tensor. + offset (int, optional): the diagonal to consider. Default: 0 (main + diagonal). + + Examples:: + + >>> a = torch.randn(3) + >>> a + tensor([-0.2956, -0.9068, 0.1695]) + >>> torch.diagflat(a) + tensor([[-0.2956, 0.0000, 0.0000], + [ 0.0000, -0.9068, 0.0000], + [ 0.0000, 0.0000, 0.1695]]) + >>> torch.diagflat(a, 1) + tensor([[ 0.0000, -0.2956, 0.0000, 0.0000], + [ 0.0000, 0.0000, -0.9068, 0.0000], + [ 0.0000, 0.0000, 0.0000, 0.1695], + [ 0.0000, 0.0000, 0.0000, 0.0000]]) + + >>> a = torch.randn(2, 2) + >>> a + tensor([[ 0.2094, -0.3018], + [-0.1516, 1.9342]]) + >>> torch.diagflat(a) + tensor([[ 0.2094, 0.0000, 0.0000, 0.0000], + [ 0.0000, -0.3018, 0.0000, 0.0000], + [ 0.0000, 0.0000, -0.1516, 0.0000], + [ 0.0000, 0.0000, 0.0000, 1.9342]]) + """ + +@overload +def diagonal( + input: Tensor, + offset: _int = 0, + dim1: _int = 0, + dim2: _int = 1, +) -> Tensor: + r""" + diagonal(input, offset=0, dim1=0, dim2=1) -> Tensor + + Returns a partial view of :attr:`input` with the its diagonal elements + with respect to :attr:`dim1` and :attr:`dim2` appended as a dimension + at the end of the shape. + + The argument :attr:`offset` controls which diagonal to consider: + + - If :attr:`offset` = 0, it is the main diagonal. + - If :attr:`offset` > 0, it is above the main diagonal. + - If :attr:`offset` < 0, it is below the main diagonal. + + Applying :meth:`torch.diag_embed` to the output of this function with + the same arguments yields a diagonal matrix with the diagonal entries + of the input. However, :meth:`torch.diag_embed` has different default + dimensions, so those need to be explicitly specified. + + Args: + input (Tensor): the input tensor. Must be at least 2-dimensional. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: 0. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: 1. + + .. note:: To take a batch diagonal, pass in dim1=-2, dim2=-1. + + Examples:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-1.0854, 1.1431, -0.1752], + [ 0.8536, -0.0905, 0.0360], + [ 0.6927, -0.3735, -0.4945]]) + + + >>> torch.diagonal(a) + tensor([-1.0854, -0.0905, -0.4945]) + + + >>> torch.diagonal(a, 1) + tensor([ 1.1431, 0.0360]) + + >>> b = torch.randn(2, 5) + >>> b + tensor([[-1.7948, -1.2731, -0.3181, 2.0200, -1.6745], + [ 1.8262, -1.5049, 0.4114, 1.0704, -1.2607]]) + + >>> torch.diagonal(b, 1, 1, 0) + tensor([1.8262]) + + >>> x = torch.randn(2, 5, 4, 2) + >>> torch.diagonal(x, offset=-1, dim1=1, dim2=2) + tensor([[[-1.2631, 0.3755, -1.5977, -1.8172], + [-1.1065, 1.0401, -0.2235, -0.7938]], + + [[-1.7325, -0.3081, 0.6166, 0.2335], + [ 1.0500, 0.7336, -0.3836, -1.1015]]]) + """ + +@overload +def diagonal( + input: Tensor, + *, + outdim: str | EllipsisType | None, + dim1: str | EllipsisType | None, + dim2: str | EllipsisType | None, + offset: _int = 0, +) -> Tensor: + r""" + diagonal(input, offset=0, dim1=0, dim2=1) -> Tensor + + Returns a partial view of :attr:`input` with the its diagonal elements + with respect to :attr:`dim1` and :attr:`dim2` appended as a dimension + at the end of the shape. + + The argument :attr:`offset` controls which diagonal to consider: + + - If :attr:`offset` = 0, it is the main diagonal. + - If :attr:`offset` > 0, it is above the main diagonal. + - If :attr:`offset` < 0, it is below the main diagonal. + + Applying :meth:`torch.diag_embed` to the output of this function with + the same arguments yields a diagonal matrix with the diagonal entries + of the input. However, :meth:`torch.diag_embed` has different default + dimensions, so those need to be explicitly specified. + + Args: + input (Tensor): the input tensor. Must be at least 2-dimensional. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: 0. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: 1. + + .. note:: To take a batch diagonal, pass in dim1=-2, dim2=-1. + + Examples:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-1.0854, 1.1431, -0.1752], + [ 0.8536, -0.0905, 0.0360], + [ 0.6927, -0.3735, -0.4945]]) + + + >>> torch.diagonal(a) + tensor([-1.0854, -0.0905, -0.4945]) + + + >>> torch.diagonal(a, 1) + tensor([ 1.1431, 0.0360]) + + >>> b = torch.randn(2, 5) + >>> b + tensor([[-1.7948, -1.2731, -0.3181, 2.0200, -1.6745], + [ 1.8262, -1.5049, 0.4114, 1.0704, -1.2607]]) + + >>> torch.diagonal(b, 1, 1, 0) + tensor([1.8262]) + + >>> x = torch.randn(2, 5, 4, 2) + >>> torch.diagonal(x, offset=-1, dim1=1, dim2=2) + tensor([[[-1.2631, 0.3755, -1.5977, -1.8172], + [-1.1065, 1.0401, -0.2235, -0.7938]], + + [[-1.7325, -0.3081, 0.6166, 0.2335], + [ 1.0500, 0.7336, -0.3836, -1.1015]]]) + """ + +def diagonal_copy( + input: Tensor, + offset: _int = 0, + dim1: _int = 0, + dim2: _int = 1, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.diagonal`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def diagonal_scatter( + input: Tensor, + src: Tensor, + offset: _int = 0, + dim1: _int = 0, + dim2: _int = 1, +) -> Tensor: + r""" + diagonal_scatter(input, src, offset=0, dim1=0, dim2=1) -> Tensor + + Embeds the values of the :attr:`src` tensor into :attr:`input` along + the diagonal elements of :attr:`input`, with respect to :attr:`dim1` + and :attr:`dim2`. + + This function returns a tensor with fresh storage; it does not + return a view. + + The argument :attr:`offset` controls which diagonal to consider: + + - If :attr:`offset` = 0, it is the main diagonal. + - If :attr:`offset` > 0, it is above the main diagonal. + - If :attr:`offset` < 0, it is below the main diagonal. + + Args: + input (Tensor): the input tensor. Must be at least 2-dimensional. + src (Tensor): the tensor to embed into :attr:`input`. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: 0. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: 1. + + .. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + ``torch.diagonal(input, offset, dim1, dim2)`` + + Examples:: + + >>> a = torch.zeros(3, 3) + >>> a + tensor([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]]) + + >>> torch.diagonal_scatter(a, torch.ones(3), 0) + tensor([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + + >>> torch.diagonal_scatter(a, torch.ones(2), 1) + tensor([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) + """ + +def diff( + input: Tensor, + n: _int = 1, + dim: _int = -1, + prepend: Tensor | None = None, + append: Tensor | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + diff(input, n=1, dim=-1, prepend=None, append=None) -> Tensor + + Computes the n-th forward difference along the given dimension. + + The first-order differences are given by `out[i] = input[i + 1] - input[i]`. Higher-order + differences are calculated by using :func:`torch.diff` recursively. + + Args: + input (Tensor): the tensor to compute the differences on + n (int, optional): the number of times to recursively compute the difference + dim (int, optional): the dimension to compute the difference along. + Default is the last dimension. + prepend, append (Tensor, optional): values to prepend or append to + :attr:`input` along :attr:`dim` before computing the difference. + Their dimensions must be equivalent to that of input, and their shapes + must match input's shape except on :attr:`dim`. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([1, 3, 2]) + >>> torch.diff(a) + tensor([ 2, -1]) + >>> b = torch.tensor([4, 5]) + >>> torch.diff(a, append=b) + tensor([ 2, -1, 2, 1]) + >>> c = torch.tensor([[1, 2, 3], [3, 4, 5]]) + >>> torch.diff(c, dim=0) + tensor([[2, 2, 2]]) + >>> torch.diff(c, dim=1) + tensor([[1, 1], + [1, 1]]) + """ + +def digamma(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + digamma(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.digamma`. + """ + +def dist(input: Tensor, other: Tensor, p: Number | _complex = 2) -> Tensor: + r""" + dist(input, other, p=2) -> Tensor + + Returns the p-norm of (:attr:`input` - :attr:`other`) + + The shapes of :attr:`input` and :attr:`other` must be + :ref:`broadcastable `. + + Args: + input (Tensor): the input tensor. + other (Tensor): the Right-hand-side input tensor + p (float, optional): the norm to be computed + + Example:: + + >>> x = torch.randn(4) + >>> x + tensor([-1.5393, -0.8675, 0.5916, 1.6321]) + >>> y = torch.randn(4) + >>> y + tensor([ 0.0967, -1.0511, 0.6295, 0.8360]) + >>> torch.dist(x, y, 3.5) + tensor(1.6727) + >>> torch.dist(x, y, 3) + tensor(1.6973) + >>> torch.dist(x, y, 0) + tensor(4.) + >>> torch.dist(x, y, 1) + tensor(2.6537) + """ + +def div( + input: Tensor | Number, + other: Tensor | Number, + *, + rounding_mode: str | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + div(input, other, *, rounding_mode=None, out=None) -> Tensor + + Divides each element of the input ``input`` by the corresponding element of + :attr:`other`. + + .. math:: + \text{out}_i = \frac{\text{input}_i}{\text{other}_i} + + .. note:: + By default, this performs a "true" division like Python 3. + See the :attr:`rounding_mode` argument for floor division. + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + Always promotes integer types to the default scalar type. + + Args: + input (Tensor): the dividend + other (Tensor or Number): the divisor + + Keyword args: + rounding_mode (str, optional): Type of rounding applied to the result: + + * None - default behavior. Performs no rounding and, if both :attr:`input` and + :attr:`other` are integer types, promotes the inputs to the default scalar type. + Equivalent to true division in Python (the ``/`` operator) and NumPy's ``np.true_divide``. + * ``"trunc"`` - rounds the results of the division towards zero. + Equivalent to C-style integer division. + * ``"floor"`` - rounds the results of the division down. + Equivalent to floor division in Python (the ``//`` operator) and NumPy's ``np.floor_divide``. + + out (Tensor, optional): the output tensor. + + Examples:: + + >>> x = torch.tensor([ 0.3810, 1.2774, -0.2972, -0.3719, 0.4637]) + >>> torch.div(x, 0.5) + tensor([ 0.7620, 2.5548, -0.5944, -0.7438, 0.9274]) + + >>> a = torch.tensor([[-0.3711, -1.9353, -0.4605, -0.2917], + ... [ 0.1815, -1.0111, 0.9805, -1.5923], + ... [ 0.1062, 1.4581, 0.7759, -1.2344], + ... [-0.1830, -0.0313, 1.1908, -1.4757]]) + >>> b = torch.tensor([ 0.8032, 0.2930, -0.8113, -0.2308]) + >>> torch.div(a, b) + tensor([[-0.4620, -6.6051, 0.5676, 1.2639], + [ 0.2260, -3.4509, -1.2086, 6.8990], + [ 0.1322, 4.9764, -0.9564, 5.3484], + [-0.2278, -0.1068, -1.4678, 6.3938]]) + + >>> torch.div(a, b, rounding_mode='trunc') + tensor([[-0., -6., 0., 1.], + [ 0., -3., -1., 6.], + [ 0., 4., -0., 5.], + [-0., -0., -1., 6.]]) + + >>> torch.div(a, b, rounding_mode='floor') + tensor([[-1., -7., 0., 1.], + [ 0., -4., -2., 6.], + [ 0., 4., -1., 5.], + [-1., -1., -2., 6.]]) + """ + +@overload +def divide( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + divide(input, other, *, rounding_mode=None, out=None) -> Tensor + + Alias for :func:`torch.div`. + """ + +@overload +def divide( + input: Tensor, + other: Tensor, + *, + rounding_mode: str | None, + out: Tensor | None = None, +) -> Tensor: + r""" + divide(input, other, *, rounding_mode=None, out=None) -> Tensor + + Alias for :func:`torch.div`. + """ + +@overload +def divide( + input: Tensor, + other: Number | _complex, + *, + rounding_mode: str | None, +) -> Tensor: + r""" + divide(input, other, *, rounding_mode=None, out=None) -> Tensor + + Alias for :func:`torch.div`. + """ + +@overload +def divide(input: Tensor, other: Number | _complex) -> Tensor: + r""" + divide(input, other, *, rounding_mode=None, out=None) -> Tensor + + Alias for :func:`torch.div`. + """ + +def dot( + input: Tensor, + tensor: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + dot(input, tensor, *, out=None) -> Tensor + + Computes the dot product of two 1D tensors. + + .. note:: + + Unlike NumPy's dot, torch.dot intentionally only supports computing the dot product + of two 1D tensors with the same number of elements. + + Args: + input (Tensor): first tensor in the dot product, must be 1D. + tensor (Tensor): second tensor in the dot product, must be 1D. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.dot(torch.tensor([2, 3]), torch.tensor([2, 1])) + tensor(7) + + >>> t1, t2 = torch.tensor([0, 1]), torch.tensor([2, 3]) + >>> torch.dot(t1, t2) + tensor(3) + """ + +def dropout(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def dropout_(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def dsmm(input: Tensor, mat2: Tensor) -> Tensor: ... +@overload +def dsplit(input: Tensor, sections: _int) -> tuple[Tensor, ...]: + r""" + dsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with three or more dimensions, into multiple tensors + depthwise according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=2) + (the split dimension is 2), except that if :attr:`indices_or_sections` is an integer + it must evenly divide the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.dsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(2, 2, 4) + >>> t + tensor([[[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.]], + [[ 8., 9., 10., 11.], + [12., 13., 14., 15.]]]) + >>> torch.dsplit(t, 2) + (tensor([[[ 0., 1.], + [ 4., 5.]], + [[ 8., 9.], + [12., 13.]]]), + tensor([[[ 2., 3.], + [ 6., 7.]], + [[10., 11.], + [14., 15.]]])) + + >>> torch.dsplit(t, [3, 6]) + (tensor([[[ 0., 1., 2.], + [ 4., 5., 6.]], + [[ 8., 9., 10.], + [12., 13., 14.]]]), + tensor([[[ 3.], + [ 7.]], + [[11.], + [15.]]]), + tensor([], size=(2, 2, 0))) + """ + +@overload +def dsplit(input: Tensor, indices: _size) -> tuple[Tensor, ...]: + r""" + dsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with three or more dimensions, into multiple tensors + depthwise according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=2) + (the split dimension is 2), except that if :attr:`indices_or_sections` is an integer + it must evenly divide the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.dsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(2, 2, 4) + >>> t + tensor([[[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.]], + [[ 8., 9., 10., 11.], + [12., 13., 14., 15.]]]) + >>> torch.dsplit(t, 2) + (tensor([[[ 0., 1.], + [ 4., 5.]], + [[ 8., 9.], + [12., 13.]]]), + tensor([[[ 2., 3.], + [ 6., 7.]], + [[10., 11.], + [14., 15.]]])) + + >>> torch.dsplit(t, [3, 6]) + (tensor([[[ 0., 1., 2.], + [ 4., 5., 6.]], + [[ 8., 9., 10.], + [12., 13., 14.]]]), + tensor([[[ 3.], + [ 7.]], + [[11.], + [15.]]]), + tensor([], size=(2, 2, 0))) + """ + +def dstack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + dstack(tensors, *, out=None) -> Tensor + + Stack tensors in sequence depthwise (along third axis). + + This is equivalent to concatenation along the third axis after 1-D and 2-D tensors have been reshaped by :func:`torch.atleast_3d`. + + Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.dstack((a,b)) + tensor([[[1, 4], + [2, 5], + [3, 6]]]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.dstack((a,b)) + tensor([[[1, 4]], + [[2, 5]], + [[3, 6]]]) + """ + +def embedding( + weight: Tensor, + indices: Tensor, + padding_idx: _int | SymInt = -1, + scale_grad_by_freq: _bool = False, + sparse: _bool = False, +) -> Tensor: ... +@overload +def embedding_bag( + weight: Tensor, + indices: Tensor, + offsets: Tensor, + scale_grad_by_freq: _bool, + mode: _int, + sparse: _bool, + per_sample_weights: Tensor | None, + include_last_offset: _bool, + padding_idx: _int | None, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +@overload +def embedding_bag( + weight: Tensor, + indices: Tensor, + offsets: Tensor, + scale_grad_by_freq: _bool = False, + mode: _int = 0, + sparse: _bool = False, + per_sample_weights: Tensor | None = None, + include_last_offset: _bool = False, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +def embedding_renorm_( + input: Tensor, + indices: Tensor, + max_norm: _float, + norm_type: _float, +) -> Tensor: ... +@overload +def empty( + size: Sequence[_int | SymInt], + *, + memory_format: memory_format | None = None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, memory_format=torch.contiguous_format) -> Tensor + + Returns a tensor filled with uninitialized data. The shape of the tensor is + defined by the variable argument :attr:`size`. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + + Example:: + + >>> torch.empty((2,3), dtype=torch.int64) + tensor([[ 9.4064e+13, 2.8000e+01, 9.3493e+13], + [ 7.5751e+18, 7.1428e+18, 7.5955e+18]]) + """ + +@overload +def empty( + *size: _int | SymInt, + memory_format: memory_format | None = None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, memory_format=torch.contiguous_format) -> Tensor + + Returns a tensor filled with uninitialized data. The shape of the tensor is + defined by the variable argument :attr:`size`. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + + Example:: + + >>> torch.empty((2,3), dtype=torch.int64) + tensor([[ 9.4064e+13, 2.8000e+01, 9.3493e+13], + [ 7.5751e+18, 7.1428e+18, 7.5955e+18]]) + """ + +@overload +def empty( + size: _size, + *, + names: Sequence[str | EllipsisType | None] | None, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, memory_format=torch.contiguous_format) -> Tensor + + Returns a tensor filled with uninitialized data. The shape of the tensor is + defined by the variable argument :attr:`size`. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + + Example:: + + >>> torch.empty((2,3), dtype=torch.int64) + tensor([[ 9.4064e+13, 2.8000e+01, 9.3493e+13], + [ 7.5751e+18, 7.1428e+18, 7.5955e+18]]) + """ + +@overload +def empty( + *size: _int, + names: Sequence[str | EllipsisType | None] | None, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, memory_format=torch.contiguous_format) -> Tensor + + Returns a tensor filled with uninitialized data. The shape of the tensor is + defined by the variable argument :attr:`size`. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + + Example:: + + >>> torch.empty((2,3), dtype=torch.int64) + tensor([[ 9.4064e+13, 2.8000e+01, 9.3493e+13], + [ 7.5751e+18, 7.1428e+18, 7.5955e+18]]) + """ + +def empty_like( + input: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns an uninitialized tensor with the same size as :attr:`input`. + ``torch.empty_like(input)`` is equivalent to + ``torch.empty(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + When ``torch.preserve_format`` is used: + If the input tensor is dense (i.e., non-overlapping strided), + its memory format (including strides) is retained. + Otherwise (e.g., a non-dense view like a stepped slice), + the output is converted to the dense format. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + + Example:: + + >>> a=torch.empty((2,3), dtype=torch.int32, device = 'cuda') + >>> torch.empty_like(a) + tensor([[0, 0, 0], + [0, 0, 0]], device='cuda:0', dtype=torch.int32) + """ + +def empty_permuted( + size: Sequence[_int | SymInt], + physical_layout: _size, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty_permuted(size, physical_layout, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Creates an uninitialized, non-overlapping and dense tensor with the + specified :attr:`size`, with :attr:`physical_layout` specifying how the + dimensions are physically laid out in memory (each logical dimension is listed + from outermost to innermost). :attr:`physical_layout` is a generalization + of NCHW/NHWC notation: if each dimension is assigned a number according to + what order they occur in size (N=0, C=1, H=2, W=3), then NCHW is ``(0, 1, 2, 3)`` + while NHWC is ``(0, 2, 3, 1)``. Equivalently, the strides of the output + tensor ``t`` are such that ``t.stride(physical_layout[i]) == contiguous_strides[i]`` + (notably, this function is *not* equivalent to ``torch.empty(size).permute(physical_layout)``). + + Unlike :func:`torch.empty_strided`, this is guaranteed to produce a dense + tensor with no overlaps. If possible, prefer using this function over + :func:`torch.empty_strided` or manual use of :func:`torch.as_strided`. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (tuple of int): the shape of the output tensor + physical_layout (tuple of int): the ordering of dimensions physically in memory + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Examples: + + >>> torch.empty((2, 3, 5, 7)).stride() + (105, 35, 7, 1) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 1, 2, 3)).stride() + (105, 35, 7, 1) + >>> torch.empty((2, 3, 5, 7), memory_format=torch.channels_last).stride() + (105, 1, 21, 3) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 2, 3, 1)).stride() + (105, 1, 21, 3) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 2, 3, 1)).dim_order() + (0, 2, 3, 1) + """ + +def empty_quantized( + size: _size, + qtensor: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +def empty_strided( + size: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + empty_strided(size, stride, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Creates a tensor with the specified :attr:`size` and :attr:`stride` and filled with undefined data. + + .. warning:: + If the constructed tensor is "overlapped" (with multiple indices referring to the same element + in memory) its behavior is undefined. + + .. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + Args: + size (tuple of int): the shape of the output tensor + stride (tuple of int): the strides of the output tensor + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> a = torch.empty_strided((2, 3), (1, 2)) + >>> a + tensor([[8.9683e-44, 4.4842e-44, 5.1239e+07], + [0.0000e+00, 0.0000e+00, 3.0705e-41]]) + >>> a.stride() + (1, 2) + >>> a.size() + torch.Size([2, 3]) + """ + +@overload +def eq( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + eq(input, other, *, out=None) -> Tensor + + Computes element-wise equality + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.eq(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[ True, False], + [False, True]]) + """ + +@overload +def eq( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + eq(input, other, *, out=None) -> Tensor + + Computes element-wise equality + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.eq(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[ True, False], + [False, True]]) + """ + +def equal(input: Tensor, other: Tensor) -> _bool: + r""" + equal(input, other) -> bool + + ``True`` if two tensors have the same size and elements, ``False`` otherwise. + + .. note:: + + Tensors containing NaNs are never equal to each other. Additionally, this function does not + differentiate between the data types of the tensors during comparison. For more thorough tensor checks, + use :meth:`torch.testing.assert_close`. + + Example:: + + >>> torch.equal(torch.tensor([1, 2]), torch.tensor([1, 2])) + True + >>> torch.equal(torch.tensor([3, torch.nan]), torch.tensor([3, torch.nan])) + False + >>> torch.equal(torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([1, 2, 3], dtype=torch.float32)) + True + """ + +def erf(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + erf(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.erf`. + """ + +def erf_(input: Tensor) -> Tensor: ... +def erfc(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + erfc(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.erfc`. + """ + +def erfc_(input: Tensor) -> Tensor: ... +def erfinv(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + erfinv(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.erfinv`. + """ + +def exp(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + exp(input, *, out=None) -> Tensor + + Returns a new tensor with the exponential of the elements + of the input tensor :attr:`input`. + + .. math:: + y_{i} = e^{x_{i}} + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.exp(torch.tensor([0, math.log(2.)])) + tensor([ 1., 2.]) + """ + +def exp2(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + exp2(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.exp2`. + """ + +def exp2_(input: Tensor) -> Tensor: ... +def exp_(input: Tensor) -> Tensor: ... +def expand_copy( + input: Tensor, + size: Sequence[_int | SymInt], + *, + implicit: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.Tensor.expand`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def expm1(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + expm1(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.expm1`. + """ + +def expm1_(input: Tensor) -> Tensor: ... +@overload +def eye( + n: _int | SymInt, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + eye(n, m=None, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 2-D tensor with ones on the diagonal and zeros elsewhere. + + Args: + n (int): the number of rows + m (int, optional): the number of columns with default being :attr:`n` + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 2-D tensor with ones on the diagonal and zeros elsewhere + + Example:: + + >>> torch.eye(3) + tensor([[ 1., 0., 0.], + [ 0., 1., 0.], + [ 0., 0., 1.]]) + """ + +@overload +def eye( + n: _int | SymInt, + m: _int | SymInt, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + eye(n, m=None, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 2-D tensor with ones on the diagonal and zeros elsewhere. + + Args: + n (int): the number of rows + m (int, optional): the number of columns with default being :attr:`n` + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 2-D tensor with ones on the diagonal and zeros elsewhere + + Example:: + + >>> torch.eye(3) + tensor([[ 1., 0., 0.], + [ 0., 1., 0.], + [ 0., 0., 1.]]) + """ + +def fake_quantize_per_channel_affine( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + axis: _int, + quant_min: _int, + quant_max: _int, +) -> Tensor: + r""" + fake_quantize_per_channel_affine(input, scale, zero_point, axis, quant_min, quant_max) -> Tensor + + Returns a new tensor with the data in :attr:`input` fake quantized per channel using :attr:`scale`, + :attr:`zero_point`, :attr:`quant_min` and :attr:`quant_max`, across the channel specified by :attr:`axis`. + + .. math:: + \text{output} = ( + min( + \text{quant\_max}, + max( + \text{quant\_min}, + \text{std::nearby\_int}(\text{input} / \text{scale}) + \text{zero\_point} + ) + ) - \text{zero\_point} + ) \times \text{scale} + + Args: + input (Tensor): the input value(s), in ``torch.float32`` + scale (Tensor): quantization scale, per channel in ``torch.float32`` + zero_point (Tensor): quantization zero_point, per channel in ``torch.int32`` or ``torch.half`` or ``torch.float32`` + axis (int32): channel axis + quant_min (int64): lower bound of the quantized domain + quant_max (int64): upper bound of the quantized domain + + Returns: + Tensor: A newly fake_quantized per channel ``torch.float32`` tensor + + Example:: + + >>> x = torch.randn(2, 2, 2) + >>> x + tensor([[[-0.2525, -0.0466], + [ 0.3491, -0.2168]], + + [[-0.5906, 1.6258], + [ 0.6444, -0.0542]]]) + >>> scales = (torch.randn(2) + 1) * 0.05 + >>> scales + tensor([0.0475, 0.0486]) + >>> zero_points = torch.zeros(2).to(torch.int32) + >>> zero_points + tensor([0, 0]) + >>> torch.fake_quantize_per_channel_affine(x, scales, zero_points, 1, 0, 255) + tensor([[[0.0000, 0.0000], + [0.3405, 0.0000]], + + [[0.0000, 1.6134], + [0.6323, 0.0000]]]) + """ + +@overload +def fake_quantize_per_tensor_affine( + input: Tensor, + scale: _float, + zero_point: _int, + quant_min: _int, + quant_max: _int, +) -> Tensor: + r""" + fake_quantize_per_tensor_affine(input, scale, zero_point, quant_min, quant_max) -> Tensor + + Returns a new tensor with the data in :attr:`input` fake quantized using :attr:`scale`, + :attr:`zero_point`, :attr:`quant_min` and :attr:`quant_max`. + + .. math:: + \text{output} = ( + min( + \text{quant\_max}, + max( + \text{quant\_min}, + \text{std::nearby\_int}(\text{input} / \text{scale}) + \text{zero\_point} + ) + ) - \text{zero\_point} + ) \times \text{scale} + + Args: + input (Tensor): the input value(s), ``torch.float32`` tensor + scale (double scalar or ``float32`` Tensor): quantization scale + zero_point (int64 scalar or ``int32`` Tensor): quantization zero_point + quant_min (int64): lower bound of the quantized domain + quant_max (int64): upper bound of the quantized domain + + Returns: + Tensor: A newly fake_quantized ``torch.float32`` tensor + + Example:: + + >>> x = torch.randn(4) + >>> x + tensor([ 0.0552, 0.9730, 0.3973, -1.0780]) + >>> torch.fake_quantize_per_tensor_affine(x, 0.1, 0, 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) + >>> torch.fake_quantize_per_tensor_affine(x, torch.tensor(0.1), torch.tensor(0), 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) + """ + +@overload +def fake_quantize_per_tensor_affine( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + quant_min: _int, + quant_max: _int, +) -> Tensor: + r""" + fake_quantize_per_tensor_affine(input, scale, zero_point, quant_min, quant_max) -> Tensor + + Returns a new tensor with the data in :attr:`input` fake quantized using :attr:`scale`, + :attr:`zero_point`, :attr:`quant_min` and :attr:`quant_max`. + + .. math:: + \text{output} = ( + min( + \text{quant\_max}, + max( + \text{quant\_min}, + \text{std::nearby\_int}(\text{input} / \text{scale}) + \text{zero\_point} + ) + ) - \text{zero\_point} + ) \times \text{scale} + + Args: + input (Tensor): the input value(s), ``torch.float32`` tensor + scale (double scalar or ``float32`` Tensor): quantization scale + zero_point (int64 scalar or ``int32`` Tensor): quantization zero_point + quant_min (int64): lower bound of the quantized domain + quant_max (int64): upper bound of the quantized domain + + Returns: + Tensor: A newly fake_quantized ``torch.float32`` tensor + + Example:: + + >>> x = torch.randn(4) + >>> x + tensor([ 0.0552, 0.9730, 0.3973, -1.0780]) + >>> torch.fake_quantize_per_tensor_affine(x, 0.1, 0, 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) + >>> torch.fake_quantize_per_tensor_affine(x, torch.tensor(0.1), torch.tensor(0), 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) + """ + +@overload +def fbgemm_linear_fp16_weight( + input: Tensor, + packed_weight: Tensor, + bias: Tensor, +) -> Tensor: ... +@overload +def fbgemm_linear_fp16_weight( + input: Tensor, + packed_weight: Tensor, + bias: Tensor, + output: Tensor, +) -> Tensor: ... +@overload +def fbgemm_linear_fp16_weight_fp32_activation( + input: Tensor, + packed_weight: Tensor, + bias: Tensor | None, +) -> Tensor: ... +@overload +def fbgemm_linear_fp16_weight_fp32_activation( + input: Tensor, + packed_weight: Tensor, + bias: Tensor | None, + output: Tensor, +) -> Tensor: ... +def fbgemm_linear_int8_weight( + input: Tensor, + weight: Tensor, + packed: Tensor, + col_offsets: Tensor, + weight_scale: Number | _complex, + weight_zero_point: Number | _complex, + bias: Tensor, +) -> Tensor: ... +def fbgemm_linear_int8_weight_fp32_activation( + input: Tensor, + weight: Tensor, + packed: Tensor, + col_offsets: Tensor, + weight_scale: Number | _complex, + weight_zero_point: Number | _complex, + bias: Tensor, +) -> Tensor: ... +def fbgemm_linear_quantize_weight( + input: Tensor, +) -> tuple[Tensor, Tensor, _float, _int]: ... +def fbgemm_pack_gemm_matrix_fp16(input: Tensor) -> Tensor: ... +@overload +def fbgemm_pack_quantized_matrix(input: Tensor) -> Tensor: ... +@overload +def fbgemm_pack_quantized_matrix(input: Tensor, K: _int, N: _int) -> Tensor: ... +def feature_alpha_dropout(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def feature_alpha_dropout_( + input: Tensor, + p: _float, + train: _bool, +) -> Tensor: ... +def feature_dropout(input: Tensor, p: _float, train: _bool) -> Tensor: ... +def feature_dropout_(input: Tensor, p: _float, train: _bool) -> Tensor: ... +@overload +def fill(input: Tensor, value: Tensor) -> Tensor: ... +@overload +def fill(input: Tensor, value: Number | _complex) -> Tensor: ... +@overload +def fill_(input: Tensor, value: Tensor) -> Tensor: ... +@overload +def fill_(input: Tensor, value: Number | _complex) -> Tensor: ... +def fix(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + fix(input, *, out=None) -> Tensor + + Alias for :func:`torch.trunc` + """ + +def fix_(input: Tensor) -> Tensor: ... +@overload +def flatten( + input: Tensor, + start_dim: _int = 0, + end_dim: _int = -1, +) -> Tensor: + r""" + flatten(input, start_dim=0, end_dim=-1) -> Tensor + + Flattens :attr:`input` by reshaping it into a one-dimensional tensor. If :attr:`start_dim` or :attr:`end_dim` + are passed, only dimensions starting with :attr:`start_dim` and ending with :attr:`end_dim` are flattened. + The order of elements in :attr:`input` is unchanged. + + Unlike NumPy's flatten, which always copies input's data, this function may return the original object, a view, + or copy. If no dimensions are flattened, then the original object :attr:`input` is returned. Otherwise, if input can + be viewed as the flattened shape, then that view is returned. Finally, only if the input cannot be viewed as the + flattened shape is input's data copied. See :meth:`torch.Tensor.view` for details on when a view will be returned. + + .. note:: + Flattening a zero-dimensional tensor will return a one-dimensional view. + + Args: + input (Tensor): the input tensor. + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + + Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.flatten(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + >>> torch.flatten(t, start_dim=1) + tensor([[1, 2, 3, 4], + [5, 6, 7, 8]]) + """ + +@overload +def flatten( + input: Tensor, + start_dim: _int, + end_dim: _int, + out_dim: str | EllipsisType | None, +) -> Tensor: + r""" + flatten(input, start_dim=0, end_dim=-1) -> Tensor + + Flattens :attr:`input` by reshaping it into a one-dimensional tensor. If :attr:`start_dim` or :attr:`end_dim` + are passed, only dimensions starting with :attr:`start_dim` and ending with :attr:`end_dim` are flattened. + The order of elements in :attr:`input` is unchanged. + + Unlike NumPy's flatten, which always copies input's data, this function may return the original object, a view, + or copy. If no dimensions are flattened, then the original object :attr:`input` is returned. Otherwise, if input can + be viewed as the flattened shape, then that view is returned. Finally, only if the input cannot be viewed as the + flattened shape is input's data copied. See :meth:`torch.Tensor.view` for details on when a view will be returned. + + .. note:: + Flattening a zero-dimensional tensor will return a one-dimensional view. + + Args: + input (Tensor): the input tensor. + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + + Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.flatten(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + >>> torch.flatten(t, start_dim=1) + tensor([[1, 2, 3, 4], + [5, 6, 7, 8]]) + """ + +@overload +def flatten( + input: Tensor, + start_dim: str | EllipsisType | None, + end_dim: str | EllipsisType | None, + out_dim: str | EllipsisType | None, +) -> Tensor: + r""" + flatten(input, start_dim=0, end_dim=-1) -> Tensor + + Flattens :attr:`input` by reshaping it into a one-dimensional tensor. If :attr:`start_dim` or :attr:`end_dim` + are passed, only dimensions starting with :attr:`start_dim` and ending with :attr:`end_dim` are flattened. + The order of elements in :attr:`input` is unchanged. + + Unlike NumPy's flatten, which always copies input's data, this function may return the original object, a view, + or copy. If no dimensions are flattened, then the original object :attr:`input` is returned. Otherwise, if input can + be viewed as the flattened shape, then that view is returned. Finally, only if the input cannot be viewed as the + flattened shape is input's data copied. See :meth:`torch.Tensor.view` for details on when a view will be returned. + + .. note:: + Flattening a zero-dimensional tensor will return a one-dimensional view. + + Args: + input (Tensor): the input tensor. + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + + Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.flatten(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + >>> torch.flatten(t, start_dim=1) + tensor([[1, 2, 3, 4], + [5, 6, 7, 8]]) + """ + +@overload +def flatten( + input: Tensor, + dims: Sequence[str | EllipsisType | None], + out_dim: str | EllipsisType | None, +) -> Tensor: + r""" + flatten(input, start_dim=0, end_dim=-1) -> Tensor + + Flattens :attr:`input` by reshaping it into a one-dimensional tensor. If :attr:`start_dim` or :attr:`end_dim` + are passed, only dimensions starting with :attr:`start_dim` and ending with :attr:`end_dim` are flattened. + The order of elements in :attr:`input` is unchanged. + + Unlike NumPy's flatten, which always copies input's data, this function may return the original object, a view, + or copy. If no dimensions are flattened, then the original object :attr:`input` is returned. Otherwise, if input can + be viewed as the flattened shape, then that view is returned. Finally, only if the input cannot be viewed as the + flattened shape is input's data copied. See :meth:`torch.Tensor.view` for details on when a view will be returned. + + .. note:: + Flattening a zero-dimensional tensor will return a one-dimensional view. + + Args: + input (Tensor): the input tensor. + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + + Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.flatten(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + >>> torch.flatten(t, start_dim=1) + tensor([[1, 2, 3, 4], + [5, 6, 7, 8]]) + """ + +def flip(input: Tensor, dims: _size) -> Tensor: + r""" + flip(input, dims) -> Tensor + + Reverse the order of an n-D tensor along given axis in dims. + + .. note:: + `torch.flip` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.flip`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.flip` is expected to be slower than `np.flip`. + + Args: + input (Tensor): the input tensor. + dims (a list or tuple): axis to flip on + + Example:: + + >>> x = torch.arange(8).view(2, 2, 2) + >>> x + tensor([[[ 0, 1], + [ 2, 3]], + + [[ 4, 5], + [ 6, 7]]]) + >>> torch.flip(x, [0, 1]) + tensor([[[ 6, 7], + [ 4, 5]], + + [[ 2, 3], + [ 0, 1]]]) + """ + +def fliplr(input: Tensor) -> Tensor: + r""" + fliplr(input) -> Tensor + + Flip tensor in the left/right direction, returning a new tensor. + + Flip the entries in each row in the left/right direction. + Columns are preserved, but appear in a different order than before. + + Note: + Requires the tensor to be at least 2-D. + + .. note:: + `torch.fliplr` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.fliplr`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.fliplr` is expected to be slower than `np.fliplr`. + + Args: + input (Tensor): Must be at least 2-dimensional. + + Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.fliplr(x) + tensor([[1, 0], + [3, 2]]) + """ + +def flipud(input: Tensor) -> Tensor: + r""" + flipud(input) -> Tensor + + Flip tensor in the up/down direction, returning a new tensor. + + Flip the entries in each column in the up/down direction. + Rows are preserved, but appear in a different order than before. + + Note: + Requires the tensor to be at least 1-D. + + .. note:: + `torch.flipud` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.flipud`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.flipud` is expected to be slower than `np.flipud`. + + Args: + input (Tensor): Must be at least 1-dimensional. + + Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.flipud(x) + tensor([[2, 3], + [0, 1]]) + """ + +@overload +def float_power( + input: Tensor, + exponent: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + float_power(input, exponent, *, out=None) -> Tensor + + Raises :attr:`input` to the power of :attr:`exponent`, elementwise, in double precision. + If neither input is complex returns a ``torch.float64`` tensor, + and if one or more inputs is complex returns a ``torch.complex128`` tensor. + + .. note:: + This function always computes in double precision, unlike :func:`torch.pow`, + which implements more typical :ref:`type promotion `. + This is useful when the computation needs to be performed in a wider or more precise dtype, + or the results of the computation may contain fractional values not representable in the input dtypes, + like when an integer base is raised to a negative integer exponent. + + Args: + input (Tensor or Number): the base value(s) + exponent (Tensor or Number): the exponent value(s) + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randint(10, (4,)) + >>> a + tensor([6, 4, 7, 1]) + >>> torch.float_power(a, 2) + tensor([36., 16., 49., 1.], dtype=torch.float64) + + >>> a = torch.arange(1, 5) + >>> a + tensor([ 1, 2, 3, 4]) + >>> exp = torch.tensor([2, -3, 4, -5]) + >>> exp + tensor([ 2, -3, 4, -5]) + >>> torch.float_power(a, exp) + tensor([1.0000e+00, 1.2500e-01, 8.1000e+01, 9.7656e-04], dtype=torch.float64) + """ + +@overload +def float_power( + self: Number | _complex, + exponent: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + float_power(input, exponent, *, out=None) -> Tensor + + Raises :attr:`input` to the power of :attr:`exponent`, elementwise, in double precision. + If neither input is complex returns a ``torch.float64`` tensor, + and if one or more inputs is complex returns a ``torch.complex128`` tensor. + + .. note:: + This function always computes in double precision, unlike :func:`torch.pow`, + which implements more typical :ref:`type promotion `. + This is useful when the computation needs to be performed in a wider or more precise dtype, + or the results of the computation may contain fractional values not representable in the input dtypes, + like when an integer base is raised to a negative integer exponent. + + Args: + input (Tensor or Number): the base value(s) + exponent (Tensor or Number): the exponent value(s) + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randint(10, (4,)) + >>> a + tensor([6, 4, 7, 1]) + >>> torch.float_power(a, 2) + tensor([36., 16., 49., 1.], dtype=torch.float64) + + >>> a = torch.arange(1, 5) + >>> a + tensor([ 1, 2, 3, 4]) + >>> exp = torch.tensor([2, -3, 4, -5]) + >>> exp + tensor([ 2, -3, 4, -5]) + >>> torch.float_power(a, exp) + tensor([1.0000e+00, 1.2500e-01, 8.1000e+01, 9.7656e-04], dtype=torch.float64) + """ + +@overload +def float_power( + input: Tensor, + exponent: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + float_power(input, exponent, *, out=None) -> Tensor + + Raises :attr:`input` to the power of :attr:`exponent`, elementwise, in double precision. + If neither input is complex returns a ``torch.float64`` tensor, + and if one or more inputs is complex returns a ``torch.complex128`` tensor. + + .. note:: + This function always computes in double precision, unlike :func:`torch.pow`, + which implements more typical :ref:`type promotion `. + This is useful when the computation needs to be performed in a wider or more precise dtype, + or the results of the computation may contain fractional values not representable in the input dtypes, + like when an integer base is raised to a negative integer exponent. + + Args: + input (Tensor or Number): the base value(s) + exponent (Tensor or Number): the exponent value(s) + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randint(10, (4,)) + >>> a + tensor([6, 4, 7, 1]) + >>> torch.float_power(a, 2) + tensor([36., 16., 49., 1.], dtype=torch.float64) + + >>> a = torch.arange(1, 5) + >>> a + tensor([ 1, 2, 3, 4]) + >>> exp = torch.tensor([2, -3, 4, -5]) + >>> exp + tensor([ 2, -3, 4, -5]) + >>> torch.float_power(a, exp) + tensor([1.0000e+00, 1.2500e-01, 8.1000e+01, 9.7656e-04], dtype=torch.float64) + """ + +def floor(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + floor(input, *, out=None) -> Tensor + + Returns a new tensor with the floor of the elements of :attr:`input`, + the largest integer less than or equal to each element. + + For integer inputs, follows the array-api convention of returning a + copy of the input tensor. + + .. math:: + \text{out}_{i} = \left\lfloor \text{input}_{i} \right\rfloor + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.8166, 1.5308, -0.2530, -0.2091]) + >>> torch.floor(a) + tensor([-1., 1., -1., -1.]) + """ + +def floor_(input: Tensor) -> Tensor: ... +def floor_divide( + input: Tensor | Number, + other: Tensor | Number, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + floor_divide(input, other, *, out=None) -> Tensor + + .. note:: + + Before PyTorch 1.13 :func:`torch.floor_divide` incorrectly performed + truncation division. To restore the previous behavior use + :func:`torch.div` with ``rounding_mode='trunc'``. + + Computes :attr:`input` divided by :attr:`other`, elementwise, and floors + the result. + + .. math:: + \text{{out}}_i = \text{floor} \left( \frac{{\text{{input}}_i}}{{\text{{other}}_i}} \right) + + + + Supports broadcasting to a common shape, type promotion, and integer and float inputs. + + Args: + input (Tensor or Number): the dividend + other (Tensor or Number): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([4.0, 3.0]) + >>> b = torch.tensor([2.0, 2.0]) + >>> torch.floor_divide(a, b) + tensor([2.0, 1.0]) + >>> torch.floor_divide(a, 1.4) + tensor([2.0, 2.0]) + """ + +def fmax( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + fmax(input, other, *, out=None) -> Tensor + + Computes the element-wise maximum of :attr:`input` and :attr:`other`. + + This is like :func:`torch.maximum` except it handles NaNs differently: + if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the maximum. + Only if both elements are NaN is NaN propagated. + + This function is a wrapper around C++'s ``std::fmax`` and is similar to NumPy's ``fmax`` function. + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and floating-point inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([9.7, float('nan'), 3.1, float('nan')]) + >>> b = torch.tensor([-2.2, 0.5, float('nan'), float('nan')]) + >>> torch.fmax(a, b) + tensor([9.7000, 0.5000, 3.1000, nan]) + """ + +def fmin( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + fmin(input, other, *, out=None) -> Tensor + + Computes the element-wise minimum of :attr:`input` and :attr:`other`. + + This is like :func:`torch.minimum` except it handles NaNs differently: + if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the minimum. + Only if both elements are NaN is NaN propagated. + + This function is a wrapper around C++'s ``std::fmin`` and is similar to NumPy's ``fmin`` function. + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and floating-point inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([2.2, float('nan'), 2.1, float('nan')]) + >>> b = torch.tensor([-9.3, 0.1, float('nan'), float('nan')]) + >>> torch.fmin(a, b) + tensor([-9.3000, 0.1000, 2.1000, nan]) + """ + +@overload +def fmod( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + fmod(input, other, *, out=None) -> Tensor + + Applies C++'s `std::fmod `_ entrywise. + The result has the same sign as the dividend :attr:`input` and its absolute value + is less than that of :attr:`other`. + + This function may be defined in terms of :func:`torch.div` as + + .. code:: python + + torch.fmod(a, b) == a - a.div(b, rounding_mode="trunc") * b + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and float inputs. + + .. note:: + + When the divisor is zero, returns ``NaN`` for floating point dtypes + on both CPU and GPU; raises ``RuntimeError`` for integer division by + zero on CPU; Integer division by zero on GPU may return any value. + + .. note:: + + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + + .. seealso:: + + :func:`torch.remainder` which implements Python's modulus operator. + This one is defined using division rounding down the result. + + Args: + input (Tensor): the dividend + other (Tensor or Scalar): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.fmod(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([-1., -0., -1., 1., 0., 1.]) + >>> torch.fmod(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([1.0000, 0.5000, 0.0000, 1.0000, 0.5000]) + """ + +@overload +def fmod( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + fmod(input, other, *, out=None) -> Tensor + + Applies C++'s `std::fmod `_ entrywise. + The result has the same sign as the dividend :attr:`input` and its absolute value + is less than that of :attr:`other`. + + This function may be defined in terms of :func:`torch.div` as + + .. code:: python + + torch.fmod(a, b) == a - a.div(b, rounding_mode="trunc") * b + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and float inputs. + + .. note:: + + When the divisor is zero, returns ``NaN`` for floating point dtypes + on both CPU and GPU; raises ``RuntimeError`` for integer division by + zero on CPU; Integer division by zero on GPU may return any value. + + .. note:: + + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + + .. seealso:: + + :func:`torch.remainder` which implements Python's modulus operator. + This one is defined using division rounding down the result. + + Args: + input (Tensor): the dividend + other (Tensor or Scalar): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.fmod(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([-1., -0., -1., 1., 0., 1.]) + >>> torch.fmod(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([1.0000, 0.5000, 0.0000, 1.0000, 0.5000]) + """ + +def frac(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + frac(input, *, out=None) -> Tensor + + Computes the fractional portion of each element in :attr:`input`. + + .. math:: + \text{out}_{i} = \text{input}_{i} - \left\lfloor |\text{input}_{i}| \right\rfloor * \operatorname{sgn}(\text{input}_{i}) + + Example:: + + >>> torch.frac(torch.tensor([1, 2.5, -3.2])) + tensor([ 0.0000, 0.5000, -0.2000]) + """ + +def frac_(input: Tensor) -> Tensor: ... +def frexp( + input: Tensor, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.frexp: + r""" + frexp(input, *, out=None) -> (Tensor mantissa, Tensor exponent) + + Decomposes :attr:`input` into mantissa and exponent tensors + such that :math:`\text{input} = \text{mantissa} \times 2^{\text{exponent}}`. + + The range of mantissa is the open interval (-1, 1). + + Supports float inputs. + + Args: + input (Tensor): the input tensor + + + Keyword args: + out (tuple, optional): the output tensors + + Example:: + + >>> x = torch.arange(9.) + >>> mantissa, exponent = torch.frexp(x) + >>> mantissa + tensor([0.0000, 0.5000, 0.5000, 0.7500, 0.5000, 0.6250, 0.7500, 0.8750, 0.5000]) + >>> exponent + tensor([0, 1, 2, 2, 3, 3, 3, 3, 4], dtype=torch.int32) + >>> torch.ldexp(mantissa, exponent) + tensor([0., 1., 2., 3., 4., 5., 6., 7., 8.]) + """ + +def frobenius_norm( + input: Tensor, + dim: _int | _size, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: ... +def from_file( + filename: str, + shared: _bool | None = None, + size: _int | None = 0, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + from_file(filename, shared=None, size=0, *, dtype=None, layout=None, device=None, pin_memory=False) + + Creates a CPU tensor with a storage backed by a memory-mapped file. + + If ``shared`` is True, then memory is shared between processes. All changes are written to the file. + If ``shared`` is False, then changes to the tensor do not affect the file. + + ``size`` is the number of elements in the Tensor. If ``shared`` is ``False``, then the file must contain + at least ``size * sizeof(dtype)`` bytes. If ``shared`` is ``True`` the file will be created if needed. + + .. note:: + Only CPU tensors can be mapped to files. + + .. note:: + For now, tensors with storages backed by a memory-mapped file cannot be created in pinned memory. + + + Args: + filename (str): file name to map + shared (bool): whether to share memory (whether ``MAP_SHARED`` or ``MAP_PRIVATE`` is passed to the + underlying `mmap(2) call `_) + size (int): number of elements in the tensor + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> t = torch.randn(2, 5, dtype=torch.float64) + >>> t.numpy().tofile('storage.pt') + >>> t_mapped = torch.from_file('storage.pt', shared=False, size=10, dtype=torch.float64) + """ + +def from_numpy(ndarray) -> Tensor: + r""" + from_numpy(ndarray) -> Tensor + + Creates a :class:`Tensor` from a :class:`numpy.ndarray`. + + The returned tensor and :attr:`ndarray` share the same memory. Modifications to + the tensor will be reflected in the :attr:`ndarray` and vice versa. The returned + tensor is not resizable. + + It currently accepts :attr:`ndarray` with dtypes of ``numpy.float64``, + ``numpy.float32``, ``numpy.float16``, ``numpy.complex64``, ``numpy.complex128``, + ``numpy.int64``, ``numpy.int32``, ``numpy.int16``, ``numpy.int8``, ``numpy.uint8``, + and ``bool``. + + .. warning:: + Writing to a tensor created from a read-only NumPy array is not supported and will result in undefined behavior. + + Example:: + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.from_numpy(a) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) + """ + +def frombuffer( + buffer: Any, + *, + dtype: _dtype, + count: int = -1, + offset: int = 0, + requires_grad: _bool = False, +) -> Tensor: + r""" + frombuffer(buffer, *, dtype, count=-1, offset=0, requires_grad=False) -> Tensor + + Creates a 1-dimensional :class:`Tensor` from an object that implements + the Python buffer protocol. + + Skips the first :attr:`offset` bytes in the buffer, and interprets the rest of + the raw bytes as a 1-dimensional tensor of type :attr:`dtype` with :attr:`count` + elements. + + Note that either of the following must be true: + + 1. :attr:`count` is a positive non-zero number, and the total number of bytes + in the buffer is more than :attr:`offset` plus :attr:`count` times the size + (in bytes) of :attr:`dtype`. + + 2. :attr:`count` is negative, and the length (number of bytes) of the buffer + subtracted by the :attr:`offset` is a multiple of the size (in bytes) of + :attr:`dtype`. + + The returned tensor and buffer share the same memory. Modifications to + the tensor will be reflected in the buffer and vice versa. The returned + tensor is not resizable. + + .. note:: + This function increments the reference count for the object that + owns the shared memory. Therefore, such memory will not be deallocated + before the returned tensor goes out of scope. + + .. warning:: + This function's behavior is undefined when passed an object implementing + the buffer protocol whose data is not on the CPU. Doing so is likely to + cause a segmentation fault. + + .. warning:: + This function does not try to infer the :attr:`dtype` (hence, it is not + optional). Passing a different :attr:`dtype` than its source may result + in unexpected behavior. + + Args: + buffer (object): a Python object that exposes the buffer interface. + + Keyword args: + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + count (int, optional): the number of desired elements to be read. + If negative, all the elements (until the end of the buffer) will be + read. Default: -1. + offset (int, optional): the number of bytes to skip at the start of + the buffer. Default: 0. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> import array + >>> a = array.array('i', [1, 2, 3]) + >>> t = torch.frombuffer(a, dtype=torch.int32) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) + + >>> # Interprets the signed char bytes as 32-bit integers. + >>> # Each 4 signed char elements will be interpreted as + >>> # 1 signed 32-bit integer. + >>> import array + >>> a = array.array('b', [-1, 0, 0, 0]) + >>> torch.frombuffer(a, dtype=torch.int32) + tensor([255], dtype=torch.int32) + """ + +@overload +def full( + size: _size, + fill_value: Number | _complex, + *, + out: Tensor | None = None, + layout: _layout = strided, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a tensor of size :attr:`size` filled with :attr:`fill_value`. The + tensor's dtype is inferred from :attr:`fill_value`. + + Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + fill_value (Scalar): the value to fill the output tensor with. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.full((2, 3), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + """ + +@overload +def full( + size: _size, + fill_value: Number | _complex, + *, + names: list[str | None], + layout: _layout = strided, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a tensor of size :attr:`size` filled with :attr:`fill_value`. The + tensor's dtype is inferred from :attr:`fill_value`. + + Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + fill_value (Scalar): the value to fill the output tensor with. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.full((2, 3), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + """ + +@overload +def full( + size: Sequence[_int | SymInt], + fill_value: Number | _complex, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a tensor of size :attr:`size` filled with :attr:`fill_value`. The + tensor's dtype is inferred from :attr:`fill_value`. + + Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + fill_value (Scalar): the value to fill the output tensor with. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.full((2, 3), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + """ + +@overload +def full( + size: _size, + fill_value: Number | _complex, + *, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a tensor of size :attr:`size` filled with :attr:`fill_value`. The + tensor's dtype is inferred from :attr:`fill_value`. + + Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + fill_value (Scalar): the value to fill the output tensor with. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.full((2, 3), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + """ + +def full_like( + input: Tensor, + fill_value: Number | _complex, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + full_like(input, fill_value, \*, dtype=None, layout=torch.strided, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same size as :attr:`input` filled with :attr:`fill_value`. + ``torch.full_like(input, fill_value)`` is equivalent to + ``torch.full(input.size(), fill_value, dtype=input.dtype, layout=input.layout, device=input.device)``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + fill_value: the number to fill the output tensor with. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + + Example:: + + >>> x = torch.ones(2, 3) + >>> torch.full_like(x, 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + >>> torch.full_like(x, 7) + tensor([[7., 7., 7.], + [7., 7., 7.]]) + >>> torch.full_like(x, 0.5, dtype=torch.int32) + tensor([[0, 0, 0], + [0, 0, 0]], dtype=torch.int32) + >>> y = torch.randn(3, 4, dtype=torch.float64) + >>> torch.full_like(y, -1.0) + tensor([[-1., -1., -1., -1.], + [-1., -1., -1., -1.], + [-1., -1., -1., -1.]], dtype=torch.float64) + """ + +def fused_moving_avg_obs_fake_quant( + input: Tensor, + observer_on: Tensor, + fake_quant_on: Tensor, + running_min: Tensor, + running_max: Tensor, + scale: Tensor, + zero_point: Tensor, + averaging_const: _float, + quant_min: _int, + quant_max: _int, + ch_axis: _int, + per_row_fake_quant: _bool = False, + symmetric_quant: _bool = False, +) -> Tensor: ... +@overload +def gather( + input: Tensor, + dim: _int, + index: Tensor, + *, + sparse_grad: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + gather(input, dim, index, *, sparse_grad=False, out=None) -> Tensor + + Gathers values along an axis specified by `dim`. + + For a 3-D tensor the output is specified by:: + + out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 + out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 + out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 + + :attr:`input` and :attr:`index` must have the same number of dimensions. + It is also required that ``index.size(d) <= input.size(d)`` for all + dimensions ``d != dim``. :attr:`out` will have the same shape as :attr:`index`. + Note that ``input`` and ``index`` do not broadcast against each other. + When :attr:`index` is empty, we always return an empty output with the same shape + without further error checking. + + Args: + input (Tensor): the source tensor + dim (int): the axis along which to index + index (LongTensor): the indices of elements to gather + + Keyword arguments: + sparse_grad (bool, optional): If ``True``, gradient w.r.t. :attr:`input` will be a sparse tensor. + out (Tensor, optional): the destination tensor + + Example:: + + >>> t = torch.tensor([[1, 2], [3, 4]]) + >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) + tensor([[ 1, 1], + [ 4, 3]]) + """ + +@overload +def gather( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + *, + sparse_grad: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + gather(input, dim, index, *, sparse_grad=False, out=None) -> Tensor + + Gathers values along an axis specified by `dim`. + + For a 3-D tensor the output is specified by:: + + out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 + out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 + out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 + + :attr:`input` and :attr:`index` must have the same number of dimensions. + It is also required that ``index.size(d) <= input.size(d)`` for all + dimensions ``d != dim``. :attr:`out` will have the same shape as :attr:`index`. + Note that ``input`` and ``index`` do not broadcast against each other. + When :attr:`index` is empty, we always return an empty output with the same shape + without further error checking. + + Args: + input (Tensor): the source tensor + dim (int): the axis along which to index + index (LongTensor): the indices of elements to gather + + Keyword arguments: + sparse_grad (bool, optional): If ``True``, gradient w.r.t. :attr:`input` will be a sparse tensor. + out (Tensor, optional): the destination tensor + + Example:: + + >>> t = torch.tensor([[1, 2], [3, 4]]) + >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) + tensor([[ 1, 1], + [ 4, 3]]) + """ + +def gcd( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + gcd(input, other, *, out=None) -> Tensor + + Computes the element-wise greatest common divisor (GCD) of :attr:`input` and :attr:`other`. + + Both :attr:`input` and :attr:`other` must have integer types. + + .. note:: + This defines :math:`gcd(0, 0) = 0`. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([5, 10, 15]) + >>> b = torch.tensor([3, 4, 5]) + >>> torch.gcd(a, b) + tensor([1, 2, 5]) + >>> c = torch.tensor([3]) + >>> torch.gcd(a, c) + tensor([1, 1, 3]) + """ + +def gcd_(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def ge( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ge(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \geq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is greater than or equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.ge(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, True], [False, True]]) + """ + +@overload +def ge( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ge(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \geq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is greater than or equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.ge(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, True], [False, True]]) + """ + +def geqrf( + input: Tensor, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.geqrf: + r""" + geqrf(input, *, out=None) -> (Tensor, Tensor) + + This is a low-level function for calling LAPACK's geqrf directly. This function + returns a namedtuple (a, tau) as defined in `LAPACK documentation for geqrf`_ . + + Computes a QR decomposition of :attr:`input`. + Both `Q` and `R` matrices are stored in the same output tensor `a`. + The elements of `R` are stored on and above the diagonal. + Elementary reflectors (or Householder vectors) implicitly defining matrix `Q` + are stored below the diagonal. + The results of this function can be used together with :func:`torch.linalg.householder_product` + to obtain the `Q` matrix or + with :func:`torch.ormqr`, which uses an implicit representation of the `Q` matrix, + for an efficient matrix-matrix multiplication. + + See `LAPACK documentation for geqrf`_ for further details. + + .. note:: + See also :func:`torch.linalg.qr`, which computes Q and R matrices, and :func:`torch.linalg.lstsq` + with the ``driver="gels"`` option for a function that can solve matrix equations using a QR decomposition. + + Args: + input (Tensor): the input matrix + + Keyword args: + out (tuple, optional): the output tuple of (Tensor, Tensor). Ignored if `None`. Default: `None`. + + .. _LAPACK documentation for geqrf: + http://www.netlib.org/lapack/explore-html/df/dc5/group__variants_g_ecomputational_ga3766ea903391b5cf9008132f7440ec7b.html + """ + +def ger( + input: Tensor, + vec2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ger(input, vec2, *, out=None) -> Tensor + + Alias of :func:`torch.outer`. + + .. warning:: + This function is deprecated and will be removed in a future PyTorch release. + Use :func:`torch.outer` instead. + """ + +def get_default_dtype() -> _dtype: + r""" + get_default_dtype() -> torch.dtype + + Get the current default floating point :class:`torch.dtype`. + + Example:: + + >>> torch.get_default_dtype() # initial default for floating point is torch.float32 + torch.float32 + >>> torch.set_default_dtype(torch.float64) + >>> torch.get_default_dtype() # default is now changed to torch.float64 + torch.float64 + """ + +def get_num_interop_threads() -> _int: + r""" + get_num_interop_threads() -> int + + Returns the number of threads used for inter-op parallelism on CPU + (e.g. in JIT interpreter) + """ + +def get_num_threads() -> _int: + r""" + get_num_threads() -> int + + Returns the number of threads used for parallelizing CPU operations + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: Number | _complex | None = None, + dim: _int | None = None, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: Sequence[Number | _complex], + dim: _int | None = None, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: Sequence[Number | _complex], + dim: _size, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int | None = None, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: Number | _complex, + dim: _size, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + spacing: tuple[Tensor, ...] | list[Tensor] | None, + dim: _size, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def gradient( + input: Tensor, + *, + dim: _size, + edge_order: _int = 1, +) -> tuple[Tensor, ...]: + r""" + gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + + Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in + one or more dimensions using the `second-order accurate central differences method + `_ and + either first or second order estimates at the boundaries. + + The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not + specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates + to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional + :attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and + :math:`g(1, 2, 3)\ == input[1, 2, 3]`. + + When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. + This is detailed in the "Keyword Arguments" section below. + + The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is + accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be + improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative + is estimated using `Taylor's theorem with remainder `_. + Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring + it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + + .. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + + Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + + .. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + + .. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + + The value of each partial derivative at the boundary points is computed differently. See edge_order below. + + Args: + input (``Tensor``): the tensor that represents the values of the function + + Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + + Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + """ + +@overload +def greater( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + greater(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.gt`. + """ + +@overload +def greater( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + greater(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.gt`. + """ + +@overload +def greater_equal( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + greater_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.ge`. + """ + +@overload +def greater_equal( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + greater_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.ge`. + """ + +def grid_sampler( + input: Tensor, + grid: Tensor, + interpolation_mode: _int, + padding_mode: _int, + align_corners: _bool, +) -> Tensor: ... +def grid_sampler_2d( + input: Tensor, + grid: Tensor, + interpolation_mode: _int, + padding_mode: _int, + align_corners: _bool, +) -> Tensor: ... +def grid_sampler_3d( + input: Tensor, + grid: Tensor, + interpolation_mode: _int, + padding_mode: _int, + align_corners: _bool, +) -> Tensor: ... +def group_norm( + input: Tensor, + num_groups: _int, + weight: Tensor | None = None, + bias: Tensor | None = None, + eps: _float = 1e-05, + cudnn_enabled: _bool = True, +) -> Tensor: ... +@overload +def gru( + data: Tensor, + batch_sizes: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, +) -> tuple[Tensor, Tensor]: ... +@overload +def gru( + input: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_first: _bool, +) -> tuple[Tensor, Tensor]: ... +def gru_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor | None = None, + b_hh: Tensor | None = None, +) -> Tensor: ... +@overload +def gt( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + gt(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} > \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is greater than :attr:`other` and False elsewhere + + Example:: + + >>> torch.gt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [False, False]]) + """ + +@overload +def gt( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + gt(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} > \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is greater than :attr:`other` and False elsewhere + + Example:: + + >>> torch.gt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [False, False]]) + """ + +@overload +def hamming_window( + window_length: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hamming_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False) -> Tensor + + Hamming window function. + + .. math:: + w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hamming_window(L, periodic=True)`` equal to + ``torch.hamming_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + .. note:: + This is a generalized version of :meth:`torch.hann_window`. + + Arguments: + window_length (int): the size of returned window + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic and alpha specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, float beta, *, dtype=None, layout=None, \ + device=None, pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic, alpha and beta specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + beta (float): The coefficient :math:`\beta` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + """ + +@overload +def hamming_window( + window_length: _int, + periodic: _bool, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hamming_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False) -> Tensor + + Hamming window function. + + .. math:: + w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hamming_window(L, periodic=True)`` equal to + ``torch.hamming_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + .. note:: + This is a generalized version of :meth:`torch.hann_window`. + + Arguments: + window_length (int): the size of returned window + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic and alpha specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, float beta, *, dtype=None, layout=None, \ + device=None, pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic, alpha and beta specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + beta (float): The coefficient :math:`\beta` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + """ + +@overload +def hamming_window( + window_length: _int, + periodic: _bool, + alpha: _float, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hamming_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False) -> Tensor + + Hamming window function. + + .. math:: + w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hamming_window(L, periodic=True)`` equal to + ``torch.hamming_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + .. note:: + This is a generalized version of :meth:`torch.hann_window`. + + Arguments: + window_length (int): the size of returned window + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic and alpha specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, float beta, *, dtype=None, layout=None, \ + device=None, pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic, alpha and beta specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + beta (float): The coefficient :math:`\beta` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + """ + +@overload +def hamming_window( + window_length: _int, + periodic: _bool, + alpha: _float, + beta: _float, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hamming_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False) -> Tensor + + Hamming window function. + + .. math:: + w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hamming_window(L, periodic=True)`` equal to + ``torch.hamming_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + .. note:: + This is a generalized version of :meth:`torch.hann_window`. + + Arguments: + window_length (int): the size of returned window + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic and alpha specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + + .. function:: hamming_window(window_length, periodic, float alpha, float beta, *, dtype=None, layout=None, \ + device=None, pin_memory=False, requires_grad=False) -> Tensor + :noindex: + + Hamming window function with periodic, alpha and beta specified. + + Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + beta (float): The coefficient :math:`\beta` in the equation above + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window. + """ + +@overload +def hann_window( + window_length: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hann_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Hann window function. + + .. math:: + w[n] = \frac{1}{2}\ \left[1 - \cos \left( \frac{2 \pi n}{N - 1} \right)\right] = + \sin^2 \left( \frac{\pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hann_window(L, periodic=True)`` equal to + ``torch.hann_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +@overload +def hann_window( + window_length: _int, + periodic: _bool, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + hann_window(window_length, periodic=True, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Hann window function. + + .. math:: + w[n] = \frac{1}{2}\ \left[1 - \cos \left( \frac{2 \pi n}{N - 1} \right)\right] = + \sin^2 \left( \frac{\pi n}{N - 1} \right), + + where :math:`N` is the full window size. + + The input :attr:`window_length` is a positive integer controlling the + returned window size. :attr:`periodic` flag determines whether the returned + window trims off the last duplicate value from the symmetric window and is + ready to be used as a periodic window with functions like + :meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in + above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have + ``torch.hann_window(L, periodic=True)`` equal to + ``torch.hann_window(L + 1, periodic=False)[:-1])``. + + .. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + + Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Returns: + Tensor: A 1-D tensor of size :math:`(\text{window\_length},)` containing the window + """ + +def hardshrink( + input: Tensor, + lambd: Number | _complex = 0.5, + *, + out: Tensor | None = None, +) -> Tensor: ... +def hash_tensor( + input: Tensor, + dim: _int | _size = (), + *, + keepdim: _bool = False, + mode: _int = 0, + out: Tensor | None = None, +) -> Tensor: + r""" + hash_tensor(input, *, mode=0) -> Tensor + + Returns a hash of all elements in the :attr:`input` tensor. + + Currently only mode=0 (reduction via xor) is supported. The output will always + be of type ``torch.uint64``. The elements of ``input`` are upcasted to their + 64 bit float / integer equivalent and bitcasted to ``torch.uint64`` before + reduction via xor. + + Args: + input (Tensor): the input tensor. + + Keyword Args: + mode (int) : The hash to use. Default: 0 (xor_reduction) + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.1918, -1.1813, 0.3373]]) + >>> torch.hash_tensor(a) + tensor(13822780554648485888, dtype=torch.uint64) + + .. function:: hash_tensor(input, dim, *, keepdim=False, mode=0) -> Tensor + :noindex: + + Returns the hash of each row of the :attr:`input` tensor in the given + dimension :attr:`dim` given by mode. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword Args: + mode (int) : The hash to use. Default: 0 (xor_reduction) + + Example:: + + >>> a = torch.randn(2, 4) + >>> a + tensor([[ 0.1317, -0.5554, -1.4724, -1.1391], + [ 0.0778, -0.6070, 0.6375, 0.1798]]) + >>> torch.hash_tensor(a, 1) + tensor([9233691267014066176, 9255993250844508160], dtype=torch.uint64) + """ + +def heaviside( + input: Tensor, + values: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + heaviside(input, values, *, out=None) -> Tensor + + Computes the Heaviside step function for each element in :attr:`input`. + The Heaviside step function is defined as: + + .. math:: + \text{{heaviside}}(input, values) = \begin{cases} + 0, & \text{if input < 0}\\ + values, & \text{if input == 0}\\ + 1, & \text{if input > 0} + \end{cases} + + + Args: + input (Tensor): the input tensor. + values (Tensor): The values to use where :attr:`input` is zero. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> input = torch.tensor([-1.5, 0, 2.0]) + >>> values = torch.tensor([0.5]) + >>> torch.heaviside(input, values) + tensor([0.0000, 0.5000, 1.0000]) + >>> values = torch.tensor([1.2, -2.0, 3.5]) + >>> torch.heaviside(input, values) + tensor([0., -2., 1.]) + """ + +def hinge_embedding_loss( + input: Tensor, + target: Tensor, + margin: _float = 1.0, + reduction: _int = 1, +) -> Tensor: ... +def histc( + input: Tensor, + bins: _int = 100, + min: Number | _complex = 0, + max: Number | _complex = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + histc(input, bins=100, min=0, max=0, *, out=None) -> Tensor + + Computes the histogram of a tensor. + + The elements are sorted into equal width bins between :attr:`min` and + :attr:`max`. If :attr:`min` and :attr:`max` are both zero, the minimum and + maximum values of the data are used. + + Elements lower than min and higher than max and ``NaN`` elements are ignored. + + Args: + input (Tensor): the input tensor. + bins (int): number of histogram bins + min (Scalar): lower end of the range (inclusive) + max (Scalar): upper end of the range (inclusive) + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: Histogram represented as a tensor + + Example:: + + >>> torch.histc(torch.tensor([1., 2, 1]), bins=4, min=0, max=3) + tensor([ 0., 2., 1., 0.]) + """ + +@overload +def histogram( + input: Tensor, + bins: Tensor, + *, + weight: Tensor | None = None, + density: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.histogram: + r""" + histogram(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor) + + Computes a histogram of the values in a tensor. + + :attr:`bins` can be an integer or a 1D tensor. + + If :attr:`bins` is an int, it specifies the number of equal-width bins. + By default, the lower and upper range of the bins is determined by the + minimum and maximum elements of the input tensor. The :attr:`range` + argument can be provided to specify a range for the bins. + + If :attr:`bins` is a 1D tensor, it specifies the sequence of bin edges + including the rightmost edge. It should contain at least 2 elements + and its elements should be increasing. + + Args: + input (Tensor): the input tensor. + bins: int or 1D Tensor. If int, defines the number of equal-width bins. If tensor, + defines the sequence of bin edges including the rightmost edge. + + Keyword args: + range (tuple of float): Defines the range of the bins. + weight (Tensor): If provided, weight should have the same shape as input. Each value in + input contributes its associated weight towards its bin's result. + density (bool): If False, the result will contain the count (or total weight) in each bin. + If True, the result is the value of the probability density function over the bins, + normalized such that the integral over the range of the bins is 1. + out (Tensor, optional): the output tensor. (tuple, optional): The result tuple of two output tensors (hist, bin_edges). + + Returns: + hist (Tensor): 1D Tensor containing the values of the histogram. + bin_edges(Tensor): 1D Tensor containing the edges of the histogram bins. + + Example:: + + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.])) + (tensor([ 0., 5., 2., 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.]), density=True) + (tensor([ 0., 0.9524, 0.3810, 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) + """ + +@overload +def histogram( + input: Tensor, + bins: _int = 100, + *, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.histogram: + r""" + histogram(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor) + + Computes a histogram of the values in a tensor. + + :attr:`bins` can be an integer or a 1D tensor. + + If :attr:`bins` is an int, it specifies the number of equal-width bins. + By default, the lower and upper range of the bins is determined by the + minimum and maximum elements of the input tensor. The :attr:`range` + argument can be provided to specify a range for the bins. + + If :attr:`bins` is a 1D tensor, it specifies the sequence of bin edges + including the rightmost edge. It should contain at least 2 elements + and its elements should be increasing. + + Args: + input (Tensor): the input tensor. + bins: int or 1D Tensor. If int, defines the number of equal-width bins. If tensor, + defines the sequence of bin edges including the rightmost edge. + + Keyword args: + range (tuple of float): Defines the range of the bins. + weight (Tensor): If provided, weight should have the same shape as input. Each value in + input contributes its associated weight towards its bin's result. + density (bool): If False, the result will contain the count (or total weight) in each bin. + If True, the result is the value of the probability density function over the bins, + normalized such that the integral over the range of the bins is 1. + out (Tensor, optional): the output tensor. (tuple, optional): The result tuple of two output tensors (hist, bin_edges). + + Returns: + hist (Tensor): 1D Tensor containing the values of the histogram. + bin_edges(Tensor): 1D Tensor containing the edges of the histogram bins. + + Example:: + + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.])) + (tensor([ 0., 5., 2., 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.]), density=True) + (tensor([ 0., 0.9524, 0.3810, 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) + """ + +@overload +def histogramdd( + input: Tensor, + bins: _int, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, +) -> torch.return_types.histogramdd: + r""" + histogramdd(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor[]) + + Computes a multi-dimensional histogram of the values in a tensor. + + Interprets the elements of an input tensor whose innermost dimension has size N + as a collection of N-dimensional points. Maps each of the points into a set of + N-dimensional bins and returns the number of points (or total weight) in each bin. + + :attr:`input` must be a tensor with at least 2 dimensions. + If input has shape (M, N), each of its M rows defines a point in N-dimensional space. + If input has three or more dimensions, all but the last dimension are flattened. + + Each dimension is independently associated with its own strictly increasing sequence + of bin edges. Bin edges may be specified explicitly by passing a sequence of 1D + tensors. Alternatively, bin edges may be constructed automatically by passing a + sequence of integers specifying the number of equal-width bins in each dimension. + + For each N-dimensional point in input: + - Each of its coordinates is binned independently among the bin edges + corresponding to its dimension + - Binning results are combined to identify the N-dimensional bin (if any) + into which the point falls + - If the point falls into a bin, the bin's count (or total weight) is incremented + - Points which do not fall into any bin do not contribute to the output + + :attr:`bins` can be a sequence of N 1D tensors, a sequence of N ints, or a single int. + + If :attr:`bins` is a sequence of N 1D tensors, it explicitly specifies the N sequences + of bin edges. Each 1D tensor should contain a strictly increasing sequence with at + least one element. A sequence of K bin edges defines K-1 bins, explicitly specifying + the left and right edges of all bins. Every bin is inclusive of its left edge. Only + the rightmost bin is inclusive of its right edge. + + If :attr:`bins` is a sequence of N ints, it specifies the number of equal-width bins + in each dimension. By default, the leftmost and rightmost bin edges in each dimension + are determined by the minimum and maximum elements of the input tensor in the + corresponding dimension. The :attr:`range` argument can be provided to manually + specify the leftmost and rightmost bin edges in each dimension. + + If :attr:`bins` is an int, it specifies the number of equal-width bins for all dimensions. + + .. note:: + See also :func:`torch.histogram`, which specifically computes 1D histograms. + While :func:`torch.histogramdd` infers the dimensionality of its bins and + binned values from the shape of :attr:`input`, :func:`torch.histogram` + accepts and flattens :attr:`input` of any shape. + + Args: + input (Tensor): the input tensor. + bins: Tensor[], int[], or int. + If Tensor[], defines the sequences of bin edges. + If int[], defines the number of equal-width bins in each dimension. + If int, defines the number of equal-width bins for all dimensions. + Keyword args: + range (sequence of float): Defines the leftmost and rightmost bin edges + in each dimension. + weight (Tensor): By default, each value in the input has weight 1. If a weight + tensor is passed, each N-dimensional coordinate in input + contributes its associated weight towards its bin's result. + The weight tensor should have the same shape as the :attr:`input` + tensor excluding its innermost dimension N. + density (bool): If False (default), the result will contain the count (or total weight) + in each bin. If True, each count (weight) is divided by the total count + (total weight), then divided by the volume of its associated bin. + Returns: + hist (Tensor): N-dimensional Tensor containing the values of the histogram. + bin_edges(Tensor[]): sequence of N 1D Tensors containing the bin edges. + + Example:: + + >>> torch.histogramdd(torch.tensor([[0., 1.], [1., 0.], [2., 0.], [2., 2.]]), bins=[3, 3], + ... weight=torch.tensor([1., 2., 4., 8.])) + torch.return_types.histogramdd( + hist=tensor([[0., 1., 0.], + [2., 0., 0.], + [4., 0., 8.]]), + bin_edges=(tensor([0.0000, 0.6667, 1.3333, 2.0000]), + tensor([0.0000, 0.6667, 1.3333, 2.0000]))) + + >>> torch.histogramdd(torch.tensor([[0., 0.], [1., 1.], [2., 2.]]), bins=[2, 2], + ... range=[0., 1., 0., 1.], density=True) + torch.return_types.histogramdd( + hist=tensor([[2., 0.], + [0., 2.]]), + bin_edges=(tensor([0.0000, 0.5000, 1.0000]), + tensor([0.0000, 0.5000, 1.0000]))) + """ + +@overload +def histogramdd( + input: Tensor, + bins: _size, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, +) -> torch.return_types.histogramdd: + r""" + histogramdd(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor[]) + + Computes a multi-dimensional histogram of the values in a tensor. + + Interprets the elements of an input tensor whose innermost dimension has size N + as a collection of N-dimensional points. Maps each of the points into a set of + N-dimensional bins and returns the number of points (or total weight) in each bin. + + :attr:`input` must be a tensor with at least 2 dimensions. + If input has shape (M, N), each of its M rows defines a point in N-dimensional space. + If input has three or more dimensions, all but the last dimension are flattened. + + Each dimension is independently associated with its own strictly increasing sequence + of bin edges. Bin edges may be specified explicitly by passing a sequence of 1D + tensors. Alternatively, bin edges may be constructed automatically by passing a + sequence of integers specifying the number of equal-width bins in each dimension. + + For each N-dimensional point in input: + - Each of its coordinates is binned independently among the bin edges + corresponding to its dimension + - Binning results are combined to identify the N-dimensional bin (if any) + into which the point falls + - If the point falls into a bin, the bin's count (or total weight) is incremented + - Points which do not fall into any bin do not contribute to the output + + :attr:`bins` can be a sequence of N 1D tensors, a sequence of N ints, or a single int. + + If :attr:`bins` is a sequence of N 1D tensors, it explicitly specifies the N sequences + of bin edges. Each 1D tensor should contain a strictly increasing sequence with at + least one element. A sequence of K bin edges defines K-1 bins, explicitly specifying + the left and right edges of all bins. Every bin is inclusive of its left edge. Only + the rightmost bin is inclusive of its right edge. + + If :attr:`bins` is a sequence of N ints, it specifies the number of equal-width bins + in each dimension. By default, the leftmost and rightmost bin edges in each dimension + are determined by the minimum and maximum elements of the input tensor in the + corresponding dimension. The :attr:`range` argument can be provided to manually + specify the leftmost and rightmost bin edges in each dimension. + + If :attr:`bins` is an int, it specifies the number of equal-width bins for all dimensions. + + .. note:: + See also :func:`torch.histogram`, which specifically computes 1D histograms. + While :func:`torch.histogramdd` infers the dimensionality of its bins and + binned values from the shape of :attr:`input`, :func:`torch.histogram` + accepts and flattens :attr:`input` of any shape. + + Args: + input (Tensor): the input tensor. + bins: Tensor[], int[], or int. + If Tensor[], defines the sequences of bin edges. + If int[], defines the number of equal-width bins in each dimension. + If int, defines the number of equal-width bins for all dimensions. + Keyword args: + range (sequence of float): Defines the leftmost and rightmost bin edges + in each dimension. + weight (Tensor): By default, each value in the input has weight 1. If a weight + tensor is passed, each N-dimensional coordinate in input + contributes its associated weight towards its bin's result. + The weight tensor should have the same shape as the :attr:`input` + tensor excluding its innermost dimension N. + density (bool): If False (default), the result will contain the count (or total weight) + in each bin. If True, each count (weight) is divided by the total count + (total weight), then divided by the volume of its associated bin. + Returns: + hist (Tensor): N-dimensional Tensor containing the values of the histogram. + bin_edges(Tensor[]): sequence of N 1D Tensors containing the bin edges. + + Example:: + + >>> torch.histogramdd(torch.tensor([[0., 1.], [1., 0.], [2., 0.], [2., 2.]]), bins=[3, 3], + ... weight=torch.tensor([1., 2., 4., 8.])) + torch.return_types.histogramdd( + hist=tensor([[0., 1., 0.], + [2., 0., 0.], + [4., 0., 8.]]), + bin_edges=(tensor([0.0000, 0.6667, 1.3333, 2.0000]), + tensor([0.0000, 0.6667, 1.3333, 2.0000]))) + + >>> torch.histogramdd(torch.tensor([[0., 0.], [1., 1.], [2., 2.]]), bins=[2, 2], + ... range=[0., 1., 0., 1.], density=True) + torch.return_types.histogramdd( + hist=tensor([[2., 0.], + [0., 2.]]), + bin_edges=(tensor([0.0000, 0.5000, 1.0000]), + tensor([0.0000, 0.5000, 1.0000]))) + """ + +@overload +def histogramdd( + input: Tensor, + bins: tuple[Tensor, ...] | list[Tensor] | None, + range: Sequence[_float] | None = None, + weight: Tensor | None = None, + density: _bool = False, +) -> torch.return_types.histogramdd: + r""" + histogramdd(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor[]) + + Computes a multi-dimensional histogram of the values in a tensor. + + Interprets the elements of an input tensor whose innermost dimension has size N + as a collection of N-dimensional points. Maps each of the points into a set of + N-dimensional bins and returns the number of points (or total weight) in each bin. + + :attr:`input` must be a tensor with at least 2 dimensions. + If input has shape (M, N), each of its M rows defines a point in N-dimensional space. + If input has three or more dimensions, all but the last dimension are flattened. + + Each dimension is independently associated with its own strictly increasing sequence + of bin edges. Bin edges may be specified explicitly by passing a sequence of 1D + tensors. Alternatively, bin edges may be constructed automatically by passing a + sequence of integers specifying the number of equal-width bins in each dimension. + + For each N-dimensional point in input: + - Each of its coordinates is binned independently among the bin edges + corresponding to its dimension + - Binning results are combined to identify the N-dimensional bin (if any) + into which the point falls + - If the point falls into a bin, the bin's count (or total weight) is incremented + - Points which do not fall into any bin do not contribute to the output + + :attr:`bins` can be a sequence of N 1D tensors, a sequence of N ints, or a single int. + + If :attr:`bins` is a sequence of N 1D tensors, it explicitly specifies the N sequences + of bin edges. Each 1D tensor should contain a strictly increasing sequence with at + least one element. A sequence of K bin edges defines K-1 bins, explicitly specifying + the left and right edges of all bins. Every bin is inclusive of its left edge. Only + the rightmost bin is inclusive of its right edge. + + If :attr:`bins` is a sequence of N ints, it specifies the number of equal-width bins + in each dimension. By default, the leftmost and rightmost bin edges in each dimension + are determined by the minimum and maximum elements of the input tensor in the + corresponding dimension. The :attr:`range` argument can be provided to manually + specify the leftmost and rightmost bin edges in each dimension. + + If :attr:`bins` is an int, it specifies the number of equal-width bins for all dimensions. + + .. note:: + See also :func:`torch.histogram`, which specifically computes 1D histograms. + While :func:`torch.histogramdd` infers the dimensionality of its bins and + binned values from the shape of :attr:`input`, :func:`torch.histogram` + accepts and flattens :attr:`input` of any shape. + + Args: + input (Tensor): the input tensor. + bins: Tensor[], int[], or int. + If Tensor[], defines the sequences of bin edges. + If int[], defines the number of equal-width bins in each dimension. + If int, defines the number of equal-width bins for all dimensions. + Keyword args: + range (sequence of float): Defines the leftmost and rightmost bin edges + in each dimension. + weight (Tensor): By default, each value in the input has weight 1. If a weight + tensor is passed, each N-dimensional coordinate in input + contributes its associated weight towards its bin's result. + The weight tensor should have the same shape as the :attr:`input` + tensor excluding its innermost dimension N. + density (bool): If False (default), the result will contain the count (or total weight) + in each bin. If True, each count (weight) is divided by the total count + (total weight), then divided by the volume of its associated bin. + Returns: + hist (Tensor): N-dimensional Tensor containing the values of the histogram. + bin_edges(Tensor[]): sequence of N 1D Tensors containing the bin edges. + + Example:: + + >>> torch.histogramdd(torch.tensor([[0., 1.], [1., 0.], [2., 0.], [2., 2.]]), bins=[3, 3], + ... weight=torch.tensor([1., 2., 4., 8.])) + torch.return_types.histogramdd( + hist=tensor([[0., 1., 0.], + [2., 0., 0.], + [4., 0., 8.]]), + bin_edges=(tensor([0.0000, 0.6667, 1.3333, 2.0000]), + tensor([0.0000, 0.6667, 1.3333, 2.0000]))) + + >>> torch.histogramdd(torch.tensor([[0., 0.], [1., 1.], [2., 2.]]), bins=[2, 2], + ... range=[0., 1., 0., 1.], density=True) + torch.return_types.histogramdd( + hist=tensor([[2., 0.], + [0., 2.]]), + bin_edges=(tensor([0.0000, 0.5000, 1.0000]), + tensor([0.0000, 0.5000, 1.0000]))) + """ + +def hsmm(input: Tensor, mat2: Tensor) -> Tensor: ... +@overload +def hsplit(input: Tensor, sections: _int) -> tuple[Tensor, ...]: + r""" + hsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with one or more dimensions, into multiple tensors + horizontally according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + If :attr:`input` is one dimensional this is equivalent to calling + torch.tensor_split(input, indices_or_sections, dim=0) (the split dimension is + zero), and if :attr:`input` has two or more dimensions it's equivalent to calling + torch.tensor_split(input, indices_or_sections, dim=1) (the split dimension is 1), + except that if :attr:`indices_or_sections` is an integer it must evenly divide + the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.hsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.hsplit(t, 2) + (tensor([[ 0., 1.], + [ 4., 5.], + [ 8., 9.], + [12., 13.]]), + tensor([[ 2., 3.], + [ 6., 7.], + [10., 11.], + [14., 15.]])) + >>> torch.hsplit(t, [3, 6]) + (tensor([[ 0., 1., 2.], + [ 4., 5., 6.], + [ 8., 9., 10.], + [12., 13., 14.]]), + tensor([[ 3.], + [ 7.], + [11.], + [15.]]), + tensor([], size=(4, 0))) + """ + +@overload +def hsplit(input: Tensor, indices: _size) -> tuple[Tensor, ...]: + r""" + hsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with one or more dimensions, into multiple tensors + horizontally according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + If :attr:`input` is one dimensional this is equivalent to calling + torch.tensor_split(input, indices_or_sections, dim=0) (the split dimension is + zero), and if :attr:`input` has two or more dimensions it's equivalent to calling + torch.tensor_split(input, indices_or_sections, dim=1) (the split dimension is 1), + except that if :attr:`indices_or_sections` is an integer it must evenly divide + the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.hsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.hsplit(t, 2) + (tensor([[ 0., 1.], + [ 4., 5.], + [ 8., 9.], + [12., 13.]]), + tensor([[ 2., 3.], + [ 6., 7.], + [10., 11.], + [14., 15.]])) + >>> torch.hsplit(t, [3, 6]) + (tensor([[ 0., 1., 2.], + [ 4., 5., 6.], + [ 8., 9., 10.], + [12., 13., 14.]]), + tensor([[ 3.], + [ 7.], + [11.], + [15.]]), + tensor([], size=(4, 0))) + """ + +def hspmm( + mat1: Tensor, + mat2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + hspmm(mat1, mat2, *, out=None) -> Tensor + + Performs a matrix multiplication of a :ref:`sparse COO matrix + ` :attr:`mat1` and a strided matrix :attr:`mat2`. The + result is a (1 + 1)-dimensional :ref:`hybrid COO matrix + `. + + Args: + mat1 (Tensor): the first sparse matrix to be matrix multiplied + mat2 (Tensor): the second strided matrix to be matrix multiplied + + Keyword args: + out (Tensor, optional): the output tensor. + """ + +def hstack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + hstack(tensors, *, out=None) -> Tensor + + Stack tensors in sequence horizontally (column wise). + + This is equivalent to concatenation along the first axis for 1-D tensors, and along the second axis for all other tensors. + + Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.hstack((a,b)) + tensor([1, 2, 3, 4, 5, 6]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.hstack((a,b)) + tensor([[1, 4], + [2, 5], + [3, 6]]) + """ + +def hypot( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + hypot(input, other, *, out=None) -> Tensor + + Given the legs of a right triangle, return its hypotenuse. + + .. math:: + \text{out}_{i} = \sqrt{\text{input}_{i}^{2} + \text{other}_{i}^{2}} + + The shapes of ``input`` and ``other`` must be + :ref:`broadcastable `. + + Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.hypot(torch.tensor([4.0]), torch.tensor([3.0, 4.0, 5.0])) + tensor([5.0000, 5.6569, 6.4031]) + """ + +def i0(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + i0(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.i0`. + """ + +def i0_(input: Tensor) -> Tensor: ... +def igamma( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + igamma(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.special.gammainc`. + """ + +def igammac( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + igammac(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.special.gammaincc`. + """ + +def imag(input: Tensor) -> Tensor: + r""" + imag(input) -> Tensor + + Returns a new tensor containing imaginary values of the :attr:`self` tensor. + The returned tensor and :attr:`self` share the same underlying storage. + + .. warning:: + :func:`imag` is only supported for tensors with complex dtypes. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.imag + tensor([ 0.3553, -0.7896, -0.0633, -0.8119]) + """ + +@overload +def index_add( + input: Tensor, + dim: _int, + index: Tensor, + source: Tensor, + *, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + index_add(input: Tensor, dim: int, index: Tensor, source: Tensor, *, alpha: Union[Number, _complex] = 1, out: Optional[Tensor]) -> Tensor # noqa: B950 + + See :meth:`~Tensor.index_add_` for function description. + """ + +@overload +def index_add( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + source: Tensor, + *, + alpha: Number | _complex = 1, +) -> Tensor: + r""" + index_add(input: Tensor, dim: int, index: Tensor, source: Tensor, *, alpha: Union[Number, _complex] = 1, out: Optional[Tensor]) -> Tensor # noqa: B950 + + See :meth:`~Tensor.index_add_` for function description. + """ + +@overload +def index_copy( + input: Tensor, + dim: _int, + index: Tensor, + source: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor, *, out: Optional[Tensor]) -> Tensor + + See :meth:`~Tensor.index_add_` for function description. + """ + +@overload +def index_copy( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + source: Tensor, +) -> Tensor: + r""" + index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor, *, out: Optional[Tensor]) -> Tensor + + See :meth:`~Tensor.index_add_` for function description. + """ + +@overload +def index_fill( + input: Tensor, + dim: _int, + index: Tensor, + value: Tensor, +) -> Tensor: ... +@overload +def index_fill( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + value: Tensor, +) -> Tensor: ... +@overload +def index_fill( + input: Tensor, + dim: _int, + index: Tensor, + value: Number | _complex, +) -> Tensor: ... +@overload +def index_fill( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + value: Number | _complex, +) -> Tensor: ... +def index_put( + input: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + values: Tensor, + accumulate: _bool = False, +) -> Tensor: ... +def index_put_( + input: Tensor, + indices: tuple[Tensor, ...] | list[Tensor] | None, + values: Tensor, + accumulate: _bool = False, +) -> Tensor: ... +def index_reduce( + input: Tensor, + dim: _int, + index: Tensor, + source: Tensor, + reduce: str, + *, + include_self: _bool = True, + out: Tensor | None = None, +) -> Tensor: + r""" + index_reduce(input: Tensor, dim: int, index: Tensor, source: Tensor, reduce: str, *, include_self: bool = True, out: Optional[Tensor]) -> Tensor # noqa: B950 + + See :meth:`~Tensor.index_reduce_` for function description. + """ + +@overload +def index_select( + input: Tensor, + dim: _int, + index: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + index_select(input, dim, index, *, out=None) -> Tensor + + Returns a new tensor which indexes the :attr:`input` tensor along dimension + :attr:`dim` using the entries in :attr:`index` which is a `LongTensor`. + + The returned tensor has the same number of dimensions as the original tensor + (:attr:`input`). The :attr:`dim`\ th dimension has the same size as the length + of :attr:`index`; other dimensions have the same size as in the original tensor. + + .. note:: The returned tensor does **not** use the same storage as the original + tensor. If :attr:`out` has a different shape than expected, we + silently change it to the correct shape, reallocating the underlying + storage if necessary. + + Args: + input (Tensor): the input tensor. + dim (int): the dimension in which we index + index (IntTensor or LongTensor): the 1-D tensor containing the indices to index + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(3, 4) + >>> x + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> indices = torch.tensor([0, 2]) + >>> torch.index_select(x, 0, indices) + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> torch.index_select(x, 1, indices) + tensor([[ 0.1427, -0.5414], + [-0.4664, -0.1228], + [-1.1734, 0.7230]]) + """ + +@overload +def index_select( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + index_select(input, dim, index, *, out=None) -> Tensor + + Returns a new tensor which indexes the :attr:`input` tensor along dimension + :attr:`dim` using the entries in :attr:`index` which is a `LongTensor`. + + The returned tensor has the same number of dimensions as the original tensor + (:attr:`input`). The :attr:`dim`\ th dimension has the same size as the length + of :attr:`index`; other dimensions have the same size as in the original tensor. + + .. note:: The returned tensor does **not** use the same storage as the original + tensor. If :attr:`out` has a different shape than expected, we + silently change it to the correct shape, reallocating the underlying + storage if necessary. + + Args: + input (Tensor): the input tensor. + dim (int): the dimension in which we index + index (IntTensor or LongTensor): the 1-D tensor containing the indices to index + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(3, 4) + >>> x + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> indices = torch.tensor([0, 2]) + >>> torch.index_select(x, 0, indices) + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> torch.index_select(x, 1, indices) + tensor([[ 0.1427, -0.5414], + [-0.4664, -0.1228], + [-1.1734, 0.7230]]) + """ + +def indices_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.indices`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def init_num_threads() -> None: ... +def inner( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + inner(input, other, *, out=None) -> Tensor + + Computes the dot product for 1D tensors. For higher dimensions, sums the product + of elements from :attr:`input` and :attr:`other` along their last dimension. + + .. note:: + + If either :attr:`input` or :attr:`other` is a scalar, the result is equivalent + to `torch.mul(input, other)`. + + If both :attr:`input` and :attr:`other` are non-scalars, the size of their last + dimension must match and the result is equivalent to `torch.tensordot(input, + other, dims=([-1], [-1]))` + + Args: + input (Tensor): First input tensor + other (Tensor): Second input tensor + + Keyword args: + out (Tensor, optional): Optional output tensor to write result into. The output + shape is `input.shape[:-1] + other.shape[:-1]`. + + Example:: + + # Dot product + >>> torch.inner(torch.tensor([1, 2, 3]), torch.tensor([0, 2, 1])) + tensor(7) + + # Multidimensional input tensors + >>> a = torch.randn(2, 3) + >>> a + tensor([[0.8173, 1.0874, 1.1784], + [0.3279, 0.1234, 2.7894]]) + >>> b = torch.randn(2, 4, 3) + >>> b + tensor([[[-0.4682, -0.7159, 0.1506], + [ 0.4034, -0.3657, 1.0387], + [ 0.9892, -0.6684, 0.1774], + [ 0.9482, 1.3261, 0.3917]], + + [[ 0.4537, 0.7493, 1.1724], + [ 0.2291, 0.5749, -0.2267], + [-0.7920, 0.3607, -0.3701], + [ 1.3666, -0.5850, -1.7242]]]) + >>> torch.inner(a, b) + tensor([[[-0.9837, 1.1560, 0.2907, 2.6785], + [ 2.5671, 0.5452, -0.6912, -1.5509]], + + [[ 0.1782, 2.9843, 0.7366, 1.5672], + [ 3.5115, -0.4864, -1.2476, -4.4337]]]) + + # Scalar input + >>> torch.inner(a, torch.tensor(2)) + tensor([[1.6347, 2.1748, 2.3567], + [0.6558, 0.2469, 5.5787]]) + """ + +def instance_norm( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + use_input_stats: _bool, + momentum: _float, + eps: _float, + cudnn_enabled: _bool, +) -> Tensor: ... +def int_repr(input: Tensor) -> Tensor: ... +def inverse(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + inverse(input, *, out=None) -> Tensor + + Alias for :func:`torch.linalg.inv` + """ + +def is_complex(input: Tensor) -> _bool: + r""" + is_complex(input: Tensor) -> bool + + Returns True if the data type of :attr:`input` is a complex data type i.e., + one of ``torch.complex64``, and ``torch.complex128``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.complex64)) + True + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.complex128)) + True + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.int32)) + False + >>> torch.is_complex(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)) + False + """ + +def is_conj(input: Tensor) -> _bool: + r""" + is_conj(input) -> (bool) + + Returns True if the :attr:`input` is a conjugated tensor, i.e. its conjugate bit is set to `True`. + + Args: + input (Tensor): the input tensor. + """ + +def is_distributed(input: Tensor) -> _bool: ... +def is_floating_point(input: Tensor) -> _bool: + r""" + is_floating_point(input: Tensor) -> bool + + Returns True if the data type of :attr:`input` is a floating point data type i.e., + one of ``torch.float64``, ``torch.float32``, ``torch.float16``, and ``torch.bfloat16``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> torch.is_floating_point(torch.tensor([1.0, 2.0, 3.0])) + True + >>> torch.is_floating_point(torch.tensor([1, 2, 3], dtype=torch.int32)) + False + >>> torch.is_floating_point(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)) + True + >>> torch.is_floating_point(torch.tensor([1, 2, 3], dtype=torch.complex64)) + False + """ + +def is_grad_enabled() -> _bool: + r""" + is_grad_enabled() -> (bool) + + Returns True if grad mode is currently enabled. + """ + +def is_inference(input: Tensor) -> _bool: + r""" + is_inference(input) -> (bool) + + Returns True if :attr:`input` is an inference tensor. + + A non-view tensor is an inference tensor if and only if it was + allocated during inference mode. A view tensor is an inference + tensor if and only if the tensor it is a view of is an inference tensor. + + For details on inference mode please see + `Inference Mode `_. + + Args: + input (Tensor): the input tensor. + """ + +def is_inference_mode_enabled() -> _bool: + r""" + is_inference_mode_enabled() -> (bool) + + Returns True if inference mode is currently enabled. + """ + +def is_neg(input: Tensor) -> _bool: ... +def is_nonzero(input: Tensor) -> _bool: + r""" + is_nonzero(input) -> (bool) + + Returns True if the :attr:`input` is a single element tensor which is not equal to zero + after type conversions. + i.e. not equal to ``torch.tensor([0.])`` or ``torch.tensor([0])`` or + ``torch.tensor([False])``. + Throws a ``RuntimeError`` if ``torch.numel() != 1`` (even in case + of sparse tensors). + + Args: + input (Tensor): the input tensor. + + Examples:: + + >>> torch.is_nonzero(torch.tensor([0.])) + False + >>> torch.is_nonzero(torch.tensor([1.5])) + True + >>> torch.is_nonzero(torch.tensor([False])) + False + >>> torch.is_nonzero(torch.tensor([3])) + True + >>> torch.is_nonzero(torch.tensor([1, 3, 5])) + Traceback (most recent call last): + ... + RuntimeError: Boolean value of Tensor with more than one value is ambiguous + >>> torch.is_nonzero(torch.tensor([])) + Traceback (most recent call last): + ... + RuntimeError: Boolean value of Tensor with no values is ambiguous + """ + +def is_same_size(input: Tensor, other: Tensor) -> _bool: ... +def is_signed(input: Tensor) -> _bool: ... +def is_vulkan_available() -> _bool: ... +def isclose( + input: Tensor, + other: Tensor, + rtol: _float = 1e-05, + atol: _float = 1e-08, + equal_nan: _bool = False, +) -> Tensor: + r""" + isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False) -> Tensor + + Returns a new tensor with boolean elements representing if each element of + :attr:`input` is "close" to the corresponding element of :attr:`other`. + Closeness is defined as: + + .. math:: + \lvert \text{input}_i - \text{other}_i \rvert \leq \texttt{rtol} \times \lvert \text{other}_i \rvert + \texttt{atol} + + + where :attr:`input` and :attr:`other` are finite. Where :attr:`input` + and/or :attr:`other` are nonfinite they are close if and only if + they are equal, with NaNs being considered equal to each other when + :attr:`equal_nan` is True. + + Args: + input (Tensor): first tensor to compare + other (Tensor): second tensor to compare + rtol (float, optional): relative tolerance. Default: 1e-05 + atol (float, optional): absolute tolerance. Default: 1e-08 + equal_nan (bool, optional): if ``True``, then two ``NaN`` s will be considered equal. Default: ``False`` + + Examples:: + + >>> torch.isclose(torch.tensor((1., 2, 3)), torch.tensor((1 + 1e-10, 3, 4))) + tensor([ True, False, False]) + >>> torch.isclose(torch.tensor((float('inf'), 4)), torch.tensor((float('inf'), 6)), rtol=.5) + tensor([True, True]) + """ + +def isfinite(input: Tensor) -> Tensor: + r""" + isfinite(input) -> Tensor + + Returns a new tensor with boolean elements representing if each element is `finite` or not. + + Real values are finite when they are not NaN, negative infinity, or infinity. + Complex values are finite when both their real and imaginary parts are finite. + + Args: + input (Tensor): the input tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is finite and False elsewhere + + Example:: + + >>> torch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')])) + tensor([True, False, True, False, False]) + """ + +@overload +def isin( + elements: Tensor, + test_elements: Tensor, + *, + assume_unique: _bool = False, + invert: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + isin(elements, test_elements, *, assume_unique=False, invert=False) -> Tensor + + Tests if each element of :attr:`elements` is in :attr:`test_elements`. Returns + a boolean tensor of the same shape as :attr:`elements` that is True for elements + in :attr:`test_elements` and False otherwise. + + .. note:: + One of :attr:`elements` or :attr:`test_elements` can be a scalar, but not both. + + Args: + elements (Tensor or Scalar): Input elements + test_elements (Tensor or Scalar): Values against which to test for each input element + assume_unique (bool, optional): If True, assumes both :attr:`elements` and + :attr:`test_elements` contain unique elements, which can speed up the + calculation. Default: False + invert (bool, optional): If True, inverts the boolean return tensor, resulting in True + values for elements *not* in :attr:`test_elements`. Default: False + + Returns: + A boolean tensor of the same shape as :attr:`elements` that is True for elements in + :attr:`test_elements` and False otherwise + + Example: + >>> torch.isin(torch.tensor([[1, 2], [3, 4]]), torch.tensor([2, 3])) + tensor([[False, True], + [ True, False]]) + """ + +@overload +def isin( + element: Number | _complex, + test_elements: Tensor, + *, + assume_unique: _bool = False, + invert: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + isin(elements, test_elements, *, assume_unique=False, invert=False) -> Tensor + + Tests if each element of :attr:`elements` is in :attr:`test_elements`. Returns + a boolean tensor of the same shape as :attr:`elements` that is True for elements + in :attr:`test_elements` and False otherwise. + + .. note:: + One of :attr:`elements` or :attr:`test_elements` can be a scalar, but not both. + + Args: + elements (Tensor or Scalar): Input elements + test_elements (Tensor or Scalar): Values against which to test for each input element + assume_unique (bool, optional): If True, assumes both :attr:`elements` and + :attr:`test_elements` contain unique elements, which can speed up the + calculation. Default: False + invert (bool, optional): If True, inverts the boolean return tensor, resulting in True + values for elements *not* in :attr:`test_elements`. Default: False + + Returns: + A boolean tensor of the same shape as :attr:`elements` that is True for elements in + :attr:`test_elements` and False otherwise + + Example: + >>> torch.isin(torch.tensor([[1, 2], [3, 4]]), torch.tensor([2, 3])) + tensor([[False, True], + [ True, False]]) + """ + +@overload +def isin( + elements: Tensor, + test_element: Number | _complex, + *, + assume_unique: _bool = False, + invert: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + isin(elements, test_elements, *, assume_unique=False, invert=False) -> Tensor + + Tests if each element of :attr:`elements` is in :attr:`test_elements`. Returns + a boolean tensor of the same shape as :attr:`elements` that is True for elements + in :attr:`test_elements` and False otherwise. + + .. note:: + One of :attr:`elements` or :attr:`test_elements` can be a scalar, but not both. + + Args: + elements (Tensor or Scalar): Input elements + test_elements (Tensor or Scalar): Values against which to test for each input element + assume_unique (bool, optional): If True, assumes both :attr:`elements` and + :attr:`test_elements` contain unique elements, which can speed up the + calculation. Default: False + invert (bool, optional): If True, inverts the boolean return tensor, resulting in True + values for elements *not* in :attr:`test_elements`. Default: False + + Returns: + A boolean tensor of the same shape as :attr:`elements` that is True for elements in + :attr:`test_elements` and False otherwise + + Example: + >>> torch.isin(torch.tensor([[1, 2], [3, 4]]), torch.tensor([2, 3])) + tensor([[False, True], + [ True, False]]) + """ + +def isinf(input: Tensor) -> Tensor: + r""" + isinf(input) -> Tensor + + Tests if each element of :attr:`input` is infinite + (positive or negative infinity) or not. + + .. note:: + Complex values are infinite when their real or imaginary part is + infinite. + + Args: + input (Tensor): the input tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is infinite and False elsewhere + + Example:: + + >>> torch.isinf(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')])) + tensor([False, True, False, True, False]) + """ + +def isnan(input: Tensor) -> Tensor: + r""" + isnan(input) -> Tensor + + Returns a new tensor with boolean elements representing if each element of :attr:`input` + is NaN or not. Complex values are considered NaN when either their real + and/or imaginary part is NaN. + + Arguments: + input (Tensor): the input tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is NaN and False elsewhere + + Example:: + + >>> torch.isnan(torch.tensor([1, float('nan'), 2])) + tensor([False, True, False]) + """ + +def isneginf(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + isneginf(input, *, out=None) -> Tensor + Tests if each element of :attr:`input` is negative infinity or not. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([-float('inf'), float('inf'), 1.2]) + >>> torch.isneginf(a) + tensor([ True, False, False]) + """ + +def isposinf(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + isposinf(input, *, out=None) -> Tensor + Tests if each element of :attr:`input` is positive infinity or not. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([-float('inf'), float('inf'), 1.2]) + >>> torch.isposinf(a) + tensor([False, True, False]) + """ + +def isreal(input: Tensor) -> Tensor: + r""" + isreal(input) -> Tensor + + Returns a new tensor with boolean elements representing if each element of :attr:`input` is real-valued or not. + All real-valued types are considered real. Complex values are considered real when their imaginary part is 0. + + Arguments: + input (Tensor): the input tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is real and False elsewhere + + Example:: + + >>> torch.isreal(torch.tensor([1, 1+1j, 2+0j])) + tensor([True, False, True]) + """ + +def istft( + input: Tensor, + n_fft: _int, + hop_length: _int | None = None, + win_length: _int | None = None, + window: Tensor | None = None, + center: _bool = True, + normalized: _bool = False, + onesided: _bool | None = None, + length: _int | None = None, + return_complex: _bool = False, +) -> Tensor: ... +@overload +def kaiser_window( + window_length: _int, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + kaiser_window(window_length, periodic=True, beta=12.0, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Computes the Kaiser window with window length :attr:`window_length` and shape parameter :attr:`beta`. + + Let I_0 be the zeroth order modified Bessel function of the first kind (see :func:`torch.i0`) and + ``N = L - 1`` if :attr:`periodic` is False and ``L`` if :attr:`periodic` is True, + where ``L`` is the :attr:`window_length`. This function computes: + + .. math:: + out_i = I_0 \left( \beta \sqrt{1 - \left( {\frac{i - N/2}{N/2}} \right) ^2 } \right) / I_0( \beta ) + + Calling ``torch.kaiser_window(L, B, periodic=True)`` is equivalent to calling + ``torch.kaiser_window(L + 1, B, periodic=False)[:-1])``. + The :attr:`periodic` argument is intended as a helpful shorthand + to produce a periodic window as input to functions like :func:`torch.stft`. + + .. note:: + If :attr:`window_length` is one, then the returned window is a single element tensor containing a one. + + + Args: + window_length (int): length of the window. + periodic (bool, optional): If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta (float, optional): shape parameter for the window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + """ + +@overload +def kaiser_window( + window_length: _int, + periodic: _bool, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + kaiser_window(window_length, periodic=True, beta=12.0, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Computes the Kaiser window with window length :attr:`window_length` and shape parameter :attr:`beta`. + + Let I_0 be the zeroth order modified Bessel function of the first kind (see :func:`torch.i0`) and + ``N = L - 1`` if :attr:`periodic` is False and ``L`` if :attr:`periodic` is True, + where ``L`` is the :attr:`window_length`. This function computes: + + .. math:: + out_i = I_0 \left( \beta \sqrt{1 - \left( {\frac{i - N/2}{N/2}} \right) ^2 } \right) / I_0( \beta ) + + Calling ``torch.kaiser_window(L, B, periodic=True)`` is equivalent to calling + ``torch.kaiser_window(L + 1, B, periodic=False)[:-1])``. + The :attr:`periodic` argument is intended as a helpful shorthand + to produce a periodic window as input to functions like :func:`torch.stft`. + + .. note:: + If :attr:`window_length` is one, then the returned window is a single element tensor containing a one. + + + Args: + window_length (int): length of the window. + periodic (bool, optional): If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta (float, optional): shape parameter for the window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + """ + +@overload +def kaiser_window( + window_length: _int, + periodic: _bool, + beta: _float, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + kaiser_window(window_length, periodic=True, beta=12.0, *, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Computes the Kaiser window with window length :attr:`window_length` and shape parameter :attr:`beta`. + + Let I_0 be the zeroth order modified Bessel function of the first kind (see :func:`torch.i0`) and + ``N = L - 1`` if :attr:`periodic` is False and ``L`` if :attr:`periodic` is True, + where ``L`` is the :attr:`window_length`. This function computes: + + .. math:: + out_i = I_0 \left( \beta \sqrt{1 - \left( {\frac{i - N/2}{N/2}} \right) ^2 } \right) / I_0( \beta ) + + Calling ``torch.kaiser_window(L, B, periodic=True)`` is equivalent to calling + ``torch.kaiser_window(L + 1, B, periodic=False)[:-1])``. + The :attr:`periodic` argument is intended as a helpful shorthand + to produce a periodic window as input to functions like :func:`torch.stft`. + + .. note:: + If :attr:`window_length` is one, then the returned window is a single element tensor containing a one. + + + Args: + window_length (int): length of the window. + periodic (bool, optional): If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta (float, optional): shape parameter for the window. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + """ + +def kl_div( + input: Tensor, + target: Tensor, + reduction: _int = 1, + *, + log_target: _bool = False, +) -> Tensor: ... +def kron( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + kron(input, other, *, out=None) -> Tensor + + Computes the Kronecker product, denoted by :math:`\otimes`, of :attr:`input` and :attr:`other`. + + If :attr:`input` is a :math:`(a_0 \times a_1 \times \dots \times a_n)` tensor and :attr:`other` is a + :math:`(b_0 \times b_1 \times \dots \times b_n)` tensor, the result will be a + :math:`(a_0*b_0 \times a_1*b_1 \times \dots \times a_n*b_n)` tensor with the following entries: + + .. math:: + (\text{input} \otimes \text{other})_{k_0, k_1, \dots, k_n} = + \text{input}_{i_0, i_1, \dots, i_n} * \text{other}_{j_0, j_1, \dots, j_n}, + + where :math:`k_t = i_t * b_t + j_t` for :math:`0 \leq t \leq n`. + If one tensor has fewer dimensions than the other it is unsqueezed until it has the same number of dimensions. + + Supports real-valued and complex-valued inputs. + + .. note:: + This function generalizes the typical definition of the Kronecker product for two matrices to two tensors, + as described above. When :attr:`input` is a :math:`(m \times n)` matrix and :attr:`other` is a + :math:`(p \times q)` matrix, the result will be a :math:`(p*m \times q*n)` block matrix: + + .. math:: + \mathbf{A} \otimes \mathbf{B}=\begin{bmatrix} + a_{11} \mathbf{B} & \cdots & a_{1 n} \mathbf{B} \\ + \vdots & \ddots & \vdots \\ + a_{m 1} \mathbf{B} & \cdots & a_{m n} \mathbf{B} \end{bmatrix} + + where :attr:`input` is :math:`\mathbf{A}` and :attr:`other` is :math:`\mathbf{B}`. + + Arguments: + input (Tensor) + other (Tensor) + + Keyword args: + out (Tensor, optional): The output tensor. Ignored if ``None``. Default: ``None`` + + Examples:: + + >>> mat1 = torch.eye(2) + >>> mat2 = torch.ones(2, 2) + >>> torch.kron(mat1, mat2) + tensor([[1., 1., 0., 0.], + [1., 1., 0., 0.], + [0., 0., 1., 1.], + [0., 0., 1., 1.]]) + + >>> mat1 = torch.eye(2) + >>> mat2 = torch.arange(1, 5).reshape(2, 2) + >>> torch.kron(mat1, mat2) + tensor([[1., 2., 0., 0.], + [3., 4., 0., 0.], + [0., 0., 1., 2.], + [0., 0., 3., 4.]]) + """ + +@overload +def kthvalue( + input: Tensor, + k: _int | SymInt, + dim: _int = -1, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.kthvalue: + r""" + kthvalue(input, k, dim=None, keepdim=False, *, out=None) -> (Tensor, LongTensor) + + Returns a namedtuple ``(values, indices)`` where ``values`` is the :attr:`k` th + smallest element of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each element found. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`keepdim` is ``True``, both the :attr:`values` and :attr:`indices` tensors + are the same size as :attr:`input`, except in the dimension :attr:`dim` where + they are of size 1. Otherwise, :attr:`dim` is squeezed + (see :func:`torch.squeeze`), resulting in both the :attr:`values` and + :attr:`indices` tensors having 1 fewer dimension than the :attr:`input` tensor. + + .. note:: + When :attr:`input` is a CUDA tensor and there are multiple valid + :attr:`k` th values, this function may nondeterministically return + :attr:`indices` for any of them. + + Args: + input (Tensor): the input tensor. + k (int): k for the k-th smallest element + dim (int, optional): the dimension to find the kth value along + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the output tuple of (Tensor, LongTensor) + can be optionally given to be used as output buffers + + Example:: + + >>> x = torch.arange(1., 6.) + >>> x + tensor([ 1., 2., 3., 4., 5.]) + >>> torch.kthvalue(x, 4) + torch.return_types.kthvalue(values=tensor(4.), indices=tensor(3)) + + >>> x=torch.arange(1.,7.).resize_(2,3) + >>> x + tensor([[ 1., 2., 3.], + [ 4., 5., 6.]]) + >>> torch.kthvalue(x, 2, 0, True) + torch.return_types.kthvalue(values=tensor([[4., 5., 6.]]), indices=tensor([[1, 1, 1]])) + """ + +@overload +def kthvalue( + input: Tensor, + k: _int | SymInt, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.kthvalue: + r""" + kthvalue(input, k, dim=None, keepdim=False, *, out=None) -> (Tensor, LongTensor) + + Returns a namedtuple ``(values, indices)`` where ``values`` is the :attr:`k` th + smallest element of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each element found. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`keepdim` is ``True``, both the :attr:`values` and :attr:`indices` tensors + are the same size as :attr:`input`, except in the dimension :attr:`dim` where + they are of size 1. Otherwise, :attr:`dim` is squeezed + (see :func:`torch.squeeze`), resulting in both the :attr:`values` and + :attr:`indices` tensors having 1 fewer dimension than the :attr:`input` tensor. + + .. note:: + When :attr:`input` is a CUDA tensor and there are multiple valid + :attr:`k` th values, this function may nondeterministically return + :attr:`indices` for any of them. + + Args: + input (Tensor): the input tensor. + k (int): k for the k-th smallest element + dim (int, optional): the dimension to find the kth value along + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the output tuple of (Tensor, LongTensor) + can be optionally given to be used as output buffers + + Example:: + + >>> x = torch.arange(1., 6.) + >>> x + tensor([ 1., 2., 3., 4., 5.]) + >>> torch.kthvalue(x, 4) + torch.return_types.kthvalue(values=tensor(4.), indices=tensor(3)) + + >>> x=torch.arange(1.,7.).resize_(2,3) + >>> x + tensor([[ 1., 2., 3.], + [ 4., 5., 6.]]) + >>> torch.kthvalue(x, 2, 0, True) + torch.return_types.kthvalue(values=tensor([[4., 5., 6.]]), indices=tensor([[1, 1, 1]])) + """ + +def layer_norm( + input: Tensor, + normalized_shape: Sequence[_int | SymInt], + weight: Tensor | None = None, + bias: Tensor | None = None, + eps: _float = 1e-05, + cudnn_enable: _bool = True, +) -> Tensor: ... +def lcm( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lcm(input, other, *, out=None) -> Tensor + + Computes the element-wise least common multiple (LCM) of :attr:`input` and :attr:`other`. + + Both :attr:`input` and :attr:`other` must have integer types. + + .. note:: + This defines :math:`lcm(0, 0) = 0` and :math:`lcm(0, a) = 0`. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([5, 10, 15]) + >>> b = torch.tensor([3, 4, 5]) + >>> torch.lcm(a, b) + tensor([15, 20, 15]) + >>> c = torch.tensor([3]) + >>> torch.lcm(a, c) + tensor([15, 30, 15]) + """ + +def lcm_(input: Tensor, other: Tensor) -> Tensor: ... +def ldexp( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ldexp(input, other, *, out=None) -> Tensor + + Multiplies :attr:`input` by 2 ** :attr:`other`. + + .. math:: + \text{{out}}_i = \text{{input}}_i * 2^\text{{other}}_i + + + Typically this function is used to construct floating point numbers by multiplying + mantissas in :attr:`input` with integral powers of two created from the exponents + in :attr:`other`. + + Args: + input (Tensor): the input tensor. + other (Tensor): a tensor of exponents, typically integers. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.ldexp(torch.tensor([1.]), torch.tensor([1])) + tensor([2.]) + >>> torch.ldexp(torch.tensor([1.0]), torch.tensor([1, 2, 3, 4])) + tensor([ 2., 4., 8., 16.]) + """ + +def ldexp_(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def le( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + le(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \leq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or Scalar): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is less than or equal to + :attr:`other` and False elsewhere + + Example:: + + >>> torch.le(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, False], [True, True]]) + """ + +@overload +def le( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + le(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \leq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or Scalar): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is less than or equal to + :attr:`other` and False elsewhere + + Example:: + + >>> torch.le(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, False], [True, True]]) + """ + +@overload +def lerp( + input: Tensor, + end: Tensor, + weight: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lerp(input, end, weight, *, out=None) + + Does a linear interpolation of two tensors :attr:`start` (given by :attr:`input`) and :attr:`end` based + on a scalar or tensor :attr:`weight` and returns the resulting :attr:`out` tensor. + + .. math:: + \text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i) + + The shapes of :attr:`start` and :attr:`end` must be + :ref:`broadcastable `. If :attr:`weight` is a tensor, then + the shapes of :attr:`weight`, :attr:`start`, and :attr:`end` must be :ref:`broadcastable `. + + Args: + input (Tensor): the tensor with the starting points + end (Tensor): the tensor with the ending points + weight (float or tensor): the weight for the interpolation formula + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> start = torch.arange(1., 5.) + >>> end = torch.empty(4).fill_(10) + >>> start + tensor([ 1., 2., 3., 4.]) + >>> end + tensor([ 10., 10., 10., 10.]) + >>> torch.lerp(start, end, 0.5) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) + >>> torch.lerp(start, end, torch.full_like(start, 0.5)) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) + """ + +@overload +def lerp( + input: Tensor, + end: Tensor, + weight: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lerp(input, end, weight, *, out=None) + + Does a linear interpolation of two tensors :attr:`start` (given by :attr:`input`) and :attr:`end` based + on a scalar or tensor :attr:`weight` and returns the resulting :attr:`out` tensor. + + .. math:: + \text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i) + + The shapes of :attr:`start` and :attr:`end` must be + :ref:`broadcastable `. If :attr:`weight` is a tensor, then + the shapes of :attr:`weight`, :attr:`start`, and :attr:`end` must be :ref:`broadcastable `. + + Args: + input (Tensor): the tensor with the starting points + end (Tensor): the tensor with the ending points + weight (float or tensor): the weight for the interpolation formula + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> start = torch.arange(1., 5.) + >>> end = torch.empty(4).fill_(10) + >>> start + tensor([ 1., 2., 3., 4.]) + >>> end + tensor([ 10., 10., 10., 10.]) + >>> torch.lerp(start, end, 0.5) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) + >>> torch.lerp(start, end, torch.full_like(start, 0.5)) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) + """ + +@overload +def less( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + less(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.lt`. + """ + +@overload +def less( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + less(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.lt`. + """ + +@overload +def less_equal( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + less_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.le`. + """ + +@overload +def less_equal( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + less_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.le`. + """ + +def lgamma(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + lgamma(input, *, out=None) -> Tensor + + Computes the natural logarithm of the absolute value of the gamma function on :attr:`input`. + + .. math:: + \text{out}_{i} = \ln |\Gamma(\text{input}_{i})| + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.arange(0.5, 2, 0.5) + >>> torch.lgamma(a) + tensor([ 0.5724, 0.0000, -0.1208]) + """ + +@overload +def linspace( + start: Number, + end: Number, + steps: _int | None = None, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + + .. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) + + + From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + + Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) + """ + +@overload +def linspace( + start: Tensor, + end: Tensor, + steps: _int, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + + .. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) + + + From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + + Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) + """ + +@overload +def linspace( + start: Number | _complex, + end: Tensor, + steps: _int, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + + .. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) + + + From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + + Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) + """ + +@overload +def linspace( + start: Tensor, + end: Number | _complex, + steps: _int, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + + .. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) + + + From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + + Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) + """ + +@overload +def linspace( + start: Number | _complex, + end: Number | _complex, + steps: _int, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + + .. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) + + + From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + + Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) + """ + +def log(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + log(input, *, out=None) -> Tensor + + Returns a new tensor with the natural logarithm of the elements + of :attr:`input`. + + .. math:: + y_{i} = \log_{e} (x_{i}) + + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(5) * 5 + >>> a + tensor([4.7767, 4.3234, 1.2156, 0.2411, 4.5739]) + >>> torch.log(a) + tensor([ 1.5637, 1.4640, 0.1952, -1.4226, 1.5204]) + """ + +def log10(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + log10(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the logarithm to the base 10 of the elements + of :attr:`input`. + + .. math:: + y_{i} = \log_{10} (x_{i}) + + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(5) + >>> a + tensor([ 0.5224, 0.9354, 0.7257, 0.1301, 0.2251]) + + + >>> torch.log10(a) + tensor([-0.2820, -0.0290, -0.1392, -0.8857, -0.6476]) + """ + +def log10_(input: Tensor) -> Tensor: ... +def log1p(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + log1p(input, *, out=None) -> Tensor + + Returns a new tensor with the natural logarithm of (1 + :attr:`input`). + + .. math:: + y_i = \log_{e} (x_i + 1) + + .. note:: This function is more accurate than :func:`torch.log` for small + values of :attr:`input` + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(5) + >>> a + tensor([-1.0090, -0.9923, 1.0249, -0.5372, 0.2492]) + >>> torch.log1p(a) + tensor([ nan, -4.8653, 0.7055, -0.7705, 0.2225]) + """ + +def log1p_(input: Tensor) -> Tensor: ... +def log2(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + log2(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the logarithm to the base 2 of the elements + of :attr:`input`. + + .. math:: + y_{i} = \log_{2} (x_{i}) + + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.rand(5) + >>> a + tensor([ 0.8419, 0.8003, 0.9971, 0.5287, 0.0490]) + + + >>> torch.log2(a) + tensor([-0.2483, -0.3213, -0.0042, -0.9196, -4.3504]) + """ + +def log2_(input: Tensor) -> Tensor: ... +def log_(input: Tensor) -> Tensor: ... +@overload +def log_softmax( + input: Tensor, + dim: _int, + dtype: _dtype | None = None, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def log_softmax( + input: Tensor, + dim: str | EllipsisType | None, + *, + dtype: _dtype | None = None, +) -> Tensor: ... +def logaddexp( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logaddexp(input, other, *, out=None) -> Tensor + + Logarithm of the sum of exponentiations of the inputs. + + Calculates pointwise :math:`\log\left(e^x + e^y\right)`. This function is useful + in statistics where the calculated probabilities of events may be so small as to + exceed the range of normal floating point numbers. In such cases the logarithm + of the calculated probability is stored. This function allows adding + probabilities stored in such a fashion. + + This op should be disambiguated with :func:`torch.logsumexp` which performs a + reduction on a single tensor. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.logaddexp(torch.tensor([-1.0]), torch.tensor([-1.0, -2, -3])) + tensor([-0.3069, -0.6867, -0.8731]) + >>> torch.logaddexp(torch.tensor([-100.0, -200, -300]), torch.tensor([-1.0, -2, -3])) + tensor([-1., -2., -3.]) + >>> torch.logaddexp(torch.tensor([1.0, 2000, 30000]), torch.tensor([-1.0, -2, -3])) + tensor([1.1269e+00, 2.0000e+03, 3.0000e+04]) + """ + +def logaddexp2( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logaddexp2(input, other, *, out=None) -> Tensor + + Logarithm of the sum of exponentiations of the inputs in base-2. + + Calculates pointwise :math:`\log_2\left(2^x + 2^y\right)`. See + :func:`torch.logaddexp` for more details. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword arguments: + out (Tensor, optional): the output tensor. + """ + +@overload +def logcumsumexp( + input: Tensor, + dim: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logcumsumexp(input, dim, *, out=None) -> Tensor + Returns the logarithm of the cumulative summation of the exponentiation of + elements of :attr:`input` in the dimension :attr:`dim`. + + For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{logcumsumexp}(x)_{ij} = \log \sum\limits_{k=0}^{j} \exp(x_{ik}) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(10) + >>> torch.logcumsumexp(a, dim=0) + tensor([-0.42296738, -0.04462666, 0.86278635, 0.94622083, 1.05277811, + 1.39202815, 1.83525007, 1.84492621, 2.06084887, 2.06844475])) + """ + +@overload +def logcumsumexp( + input: Tensor, + dim: str | EllipsisType | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logcumsumexp(input, dim, *, out=None) -> Tensor + Returns the logarithm of the cumulative summation of the exponentiation of + elements of :attr:`input` in the dimension :attr:`dim`. + + For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{logcumsumexp}(x)_{ij} = \log \sum\limits_{k=0}^{j} \exp(x_{ik}) + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to do the operation over + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(10) + >>> torch.logcumsumexp(a, dim=0) + tensor([-0.42296738, -0.04462666, 0.86278635, 0.94622083, 1.05277811, + 1.39202815, 1.83525007, 1.84492621, 2.06084887, 2.06844475])) + """ + +def logdet(input: Tensor) -> Tensor: + r""" + logdet(input) -> Tensor + + Calculates log determinant of a square matrix or batches of square matrices. + + It returns ``-inf`` if the input has a determinant of zero, and ``NaN`` if it has + a negative determinant. + + .. note:: + Backward through :meth:`logdet` internally uses SVD results when :attr:`input` + is not invertible. In this case, double backward through :meth:`logdet` will + be unstable in when :attr:`input` doesn't have distinct singular values. See + :func:`torch.linalg.svd` for details. + + .. seealso:: + + :func:`torch.linalg.slogdet` computes the sign (resp. angle) and natural logarithm of the + absolute value of the determinant of real-valued (resp. complex) square matrices. + + Arguments: + input (Tensor): the input tensor of size ``(*, n, n)`` where ``*`` is zero or more + batch dimensions. + + Example:: + + >>> A = torch.randn(3, 3) + >>> torch.det(A) + tensor(0.2611) + >>> torch.logdet(A) + tensor(-1.3430) + >>> A + tensor([[[ 0.9254, -0.6213], + [-0.5787, 1.6843]], + + [[ 0.3242, -0.9665], + [ 0.4539, -0.0887]], + + [[ 1.1336, -0.4025], + [-0.7089, 0.9032]]]) + >>> A.det() + tensor([1.1990, 0.4099, 0.7386]) + >>> A.det().log() + tensor([ 0.1815, -0.8917, -0.3031]) + """ + +def logical_and( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logical_and(input, other, *, out=None) -> Tensor + + Computes the element-wise logical AND of the given input tensors. Zeros are treated as ``False`` and nonzeros are + treated as ``True``. + + Args: + input (Tensor): the input tensor. + other (Tensor): the tensor to compute AND with + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.logical_and(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([ True, False, False]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_and(a, b) + tensor([False, False, True, False]) + >>> torch.logical_and(a.double(), b.double()) + tensor([False, False, True, False]) + >>> torch.logical_and(a.double(), b) + tensor([False, False, True, False]) + >>> torch.logical_and(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([False, False, True, False]) + """ + +def logical_not(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + logical_not(input, *, out=None) -> Tensor + + Computes the element-wise logical NOT of the given input tensor. If not specified, the output tensor will have the bool + dtype. If the input tensor is not a bool tensor, zeros are treated as ``False`` and non-zeros are treated as ``True``. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.logical_not(torch.tensor([True, False])) + tensor([False, True]) + >>> torch.logical_not(torch.tensor([0, 1, -10], dtype=torch.int8)) + tensor([ True, False, False]) + >>> torch.logical_not(torch.tensor([0., 1.5, -10.], dtype=torch.double)) + tensor([ True, False, False]) + >>> torch.logical_not(torch.tensor([0., 1., -10.], dtype=torch.double), out=torch.empty(3, dtype=torch.int16)) + tensor([1, 0, 0], dtype=torch.int16) + """ + +def logical_or( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logical_or(input, other, *, out=None) -> Tensor + + Computes the element-wise logical OR of the given input tensors. Zeros are treated as ``False`` and nonzeros are + treated as ``True``. + + Args: + input (Tensor): the input tensor. + other (Tensor): the tensor to compute OR with + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.logical_or(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([ True, False, True]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_or(a, b) + tensor([ True, True, True, False]) + >>> torch.logical_or(a.double(), b.double()) + tensor([ True, True, True, False]) + >>> torch.logical_or(a.double(), b) + tensor([ True, True, True, False]) + >>> torch.logical_or(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([ True, True, True, False]) + """ + +def logical_xor( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logical_xor(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + + Computes the element-wise logical XOR of the given input tensors. Zeros are treated as ``False`` and nonzeros are + treated as ``True``. + + Args: + input (Tensor): the input tensor. + other (Tensor): the tensor to compute XOR with + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.logical_xor(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([False, False, True]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_xor(a, b) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a.double(), b.double()) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a.double(), b) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([ True, True, False, False]) + """ + +def logit( + input: Tensor, + eps: _float | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logit(input, eps=None, *, out=None) -> Tensor + + Alias for :func:`torch.special.logit`. + """ + +def logit_(input: Tensor, eps: _float | None = None) -> Tensor: ... +@overload +def logspace( + start: Number, + end: Number, + steps: _int | None = None, + base: _float = 10.0, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + logspace(start, end, steps, base=10.0, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to + :math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale + with base :attr:`base`. That is, the values are: + + .. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) + + + + From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) + """ + +@overload +def logspace( + start: Tensor, + end: Tensor, + steps: _int, + base: _float = 10.0, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + logspace(start, end, steps, base=10.0, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to + :math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale + with base :attr:`base`. That is, the values are: + + .. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) + + + + From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) + """ + +@overload +def logspace( + start: Number | _complex, + end: Tensor, + steps: _int, + base: _float = 10.0, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + logspace(start, end, steps, base=10.0, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to + :math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale + with base :attr:`base`. That is, the values are: + + .. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) + + + + From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) + """ + +@overload +def logspace( + start: Tensor, + end: Number | _complex, + steps: _int, + base: _float = 10.0, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + logspace(start, end, steps, base=10.0, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to + :math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale + with base :attr:`base`. That is, the values are: + + .. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) + + + + From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) + """ + +@overload +def logspace( + start: Number | _complex, + end: Number | _complex, + steps: _int, + base: _float = 10.0, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + logspace(start, end, steps, base=10.0, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + + Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly + spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to + :math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale + with base :attr:`base`. That is, the values are: + + .. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) + + + + From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + + Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) + """ + +@overload +def logsumexp( + input: Tensor, + dim: _int | _size, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logsumexp(input, dim, keepdim=False, *, out=None) + + Returns the log of summed exponentials of each row of the :attr:`input` + tensor in the given dimension :attr:`dim`. The computation is numerically + stabilized. + + For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{logsumexp}(x)_{i} = \log \sum_j \exp(x_{ij}) + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints): the dimension or dimensions to reduce. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(3, 3) + >>> torch.logsumexp(a, 1) + tensor([1.4907, 1.0593, 1.5696]) + >>> torch.dist(torch.logsumexp(a, 1), torch.log(torch.sum(torch.exp(a), 1))) + tensor(1.6859e-07) + """ + +@overload +def logsumexp( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + logsumexp(input, dim, keepdim=False, *, out=None) + + Returns the log of summed exponentials of each row of the :attr:`input` + tensor in the given dimension :attr:`dim`. The computation is numerically + stabilized. + + For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{logsumexp}(x)_{i} = \log \sum_j \exp(x_{ij}) + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints): the dimension or dimensions to reduce. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(3, 3) + >>> torch.logsumexp(a, 1) + tensor([1.4907, 1.0593, 1.5696]) + >>> torch.dist(torch.logsumexp(a, 1), torch.log(torch.sum(torch.exp(a), 1))) + tensor(1.6859e-07) + """ + +@overload +def lstm( + data: Tensor, + batch_sizes: Tensor, + hx: tuple[Tensor, ...] | list[Tensor] | None, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, +) -> tuple[Tensor, Tensor, Tensor]: ... +@overload +def lstm( + input: Tensor, + hx: tuple[Tensor, ...] | list[Tensor] | None, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_first: _bool, +) -> tuple[Tensor, Tensor, Tensor]: ... +def lstm_cell( + input: Tensor, + hx: tuple[Tensor, ...] | list[Tensor] | None, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor | None = None, + b_hh: Tensor | None = None, +) -> tuple[Tensor, Tensor]: ... +@overload +def lt( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lt(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} < \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is less than :attr:`other` and False elsewhere + + Example:: + + >>> torch.lt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, False], [True, False]]) + """ + +@overload +def lt( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lt(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} < \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is less than :attr:`other` and False elsewhere + + Example:: + + >>> torch.lt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, False], [True, False]]) + """ + +def lu_solve( + input: Tensor, + LU_data: Tensor, + LU_pivots: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + lu_solve(b, LU_data, LU_pivots, *, out=None) -> Tensor + + Returns the LU solve of the linear system :math:`Ax = b` using the partially pivoted + LU factorization of A from :func:`~linalg.lu_factor`. + + This function supports ``float``, ``double``, ``cfloat`` and ``cdouble`` dtypes for :attr:`input`. + + .. warning:: + + :func:`torch.lu_solve` is deprecated in favor of :func:`torch.linalg.lu_solve`. + :func:`torch.lu_solve` will be removed in a future PyTorch release. + ``X = torch.lu_solve(B, LU, pivots)`` should be replaced with + + .. code:: python + + X = linalg.lu_solve(LU, pivots, B) + + Arguments: + b (Tensor): the RHS tensor of size :math:`(*, m, k)`, where :math:`*` + is zero or more batch dimensions. + LU_data (Tensor): the pivoted LU factorization of A from :meth:`~linalg.lu_factor` of size :math:`(*, m, m)`, + where :math:`*` is zero or more batch dimensions. + LU_pivots (IntTensor): the pivots of the LU factorization from :meth:`~linalg.lu_factor` of size :math:`(*, m)`, + where :math:`*` is zero or more batch dimensions. + The batch dimensions of :attr:`LU_pivots` must be equal to the batch dimensions of + :attr:`LU_data`. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> A = torch.randn(2, 3, 3) + >>> b = torch.randn(2, 3, 1) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> x = torch.lu_solve(b, LU, pivots) + >>> torch.dist(A @ x, b) + tensor(1.00000e-07 * + 2.8312) + """ + +def lu_unpack( + LU_data: Tensor, + LU_pivots: Tensor, + unpack_data: _bool = True, + unpack_pivots: _bool = True, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.lu_unpack: + r""" + lu_unpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True, *, out=None) -> (Tensor, Tensor, Tensor) + + Unpacks the LU decomposition returned by :func:`~linalg.lu_factor` into the `P, L, U` matrices. + + .. seealso:: + + :func:`~linalg.lu` returns the matrices from the LU decomposition. Its gradient formula is more efficient + than that of doing :func:`~linalg.lu_factor` followed by :func:`~linalg.lu_unpack`. + + Args: + LU_data (Tensor): the packed LU factorization data + LU_pivots (Tensor): the packed LU factorization pivots + unpack_data (bool): flag indicating if the data should be unpacked. + If ``False``, then the returned ``L`` and ``U`` are empty tensors. + Default: ``True`` + unpack_pivots (bool): flag indicating if the pivots should be unpacked into a permutation matrix ``P``. + If ``False``, then the returned ``P`` is an empty tensor. + Default: ``True`` + + Keyword args: + out (tuple, optional): output tuple of three tensors. Ignored if `None`. + + Returns: + A namedtuple ``(P, L, U)`` + + Examples:: + + >>> A = torch.randn(2, 3, 3) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> P, L, U = torch.lu_unpack(LU, pivots) + >>> # We can recover A from the factorization + >>> A_ = P @ L @ U + >>> torch.allclose(A, A_) + True + + >>> # LU factorization of a rectangular matrix: + >>> A = torch.randn(2, 3, 2) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> P, L, U = torch.lu_unpack(LU, pivots) + >>> # P, L, U are the same as returned by linalg.lu + >>> P_, L_, U_ = torch.linalg.lu(A) + >>> torch.allclose(P, P_) and torch.allclose(L, L_) and torch.allclose(U, U_) + True + """ + +def margin_ranking_loss( + input1: Tensor, + input2: Tensor, + target: Tensor, + margin: _float = 0.0, + reduction: _int = 1, +) -> Tensor: ... +@overload +def masked_fill(input: Tensor, mask: Tensor, value: Tensor) -> Tensor: ... +@overload +def masked_fill( + input: Tensor, + mask: Tensor, + value: Number | _complex, +) -> Tensor: ... +def masked_scatter(input: Tensor, mask: Tensor, source: Tensor) -> Tensor: ... +def masked_select( + input: Tensor, + mask: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + masked_select(input, mask, *, out=None) -> Tensor + + Returns a new 1-D tensor which indexes the :attr:`input` tensor according to + the boolean mask :attr:`mask` which is a `BoolTensor`. + + The shapes of the :attr:`mask` tensor and the :attr:`input` tensor don't need + to match, but they must be :ref:`broadcastable `. + + .. note:: The returned tensor does **not** use the same storage + as the original tensor + + Args: + input (Tensor): the input tensor. + mask (BoolTensor): the tensor containing the binary mask to index with + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(3, 4) + >>> x + tensor([[ 0.3552, -2.3825, -0.8297, 0.3477], + [-1.2035, 1.2252, 0.5002, 0.6248], + [ 0.1307, -2.0608, 0.1244, 2.0139]]) + >>> mask = x.ge(0.5) + >>> mask + tensor([[False, False, False, False], + [False, True, True, True], + [False, False, False, True]]) + >>> torch.masked_select(x, mask) + tensor([ 1.2252, 0.5002, 0.6248, 2.0139]) + """ + +def matmul( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + matmul(input, other, *, out=None) -> Tensor + + Matrix product of two tensors. + + The behavior depends on the dimensionality of the tensors as follows: + + - If both tensors are 1-dimensional, the dot product (scalar) is returned. + - If both arguments are 2-dimensional, the matrix-matrix product is returned. + - If the first argument is 1-dimensional and the second argument is 2-dimensional, + a 1 is prepended to its dimension for the purpose of the matrix multiply. + After the matrix multiply, the prepended dimension is removed. + - If the first argument is 2-dimensional and the second argument is 1-dimensional, + the matrix-vector product is returned. + - If both arguments are at least 1-dimensional and at least one argument is + N-dimensional (where N > 2), then a batched matrix multiply is returned. If the first + argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the + batched matrix multiply and removed after. If the second argument is 1-dimensional, a + 1 is appended to its dimension for the purpose of the batched matrix multiply and removed after. + + The first N-2 dimensions of each argument, the batch dimensions, are + :ref:`broadcast ` (and thus must be broadcastable). + The last 2, the matrix dimensions, are handled as in the matrix-matrix product. + + For example, if :attr:`input` is a + :math:`(j \times 1 \times n \times m)` tensor and :attr:`other` is a :math:`(k \times m \times p)` + tensor, the batch dimensions are :math:`(j \times 1)` and :math:`(k)`, + and the matrix dimensions are :math:`(n \times m)` and :math:`(m \times p)`. + :attr:`out` will be a :math:`(j \times k \times n \times p)` tensor. + + This operation has support for arguments with :ref:`sparse layouts`. In particular the + matrix-matrix (both arguments 2-dimensional) supports sparse arguments with the same restrictions + as :func:`torch.mm` + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + .. note:: + + The 1-dimensional dot product version of this function does not support an :attr:`out` parameter. + + Arguments: + input (Tensor): the first tensor to be multiplied + other (Tensor): the second tensor to be multiplied + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> # vector x vector + >>> tensor1 = torch.randn(3) + >>> tensor2 = torch.randn(3) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([]) + >>> # matrix x vector + >>> tensor1 = torch.randn(3, 4) + >>> tensor2 = torch.randn(4) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([3]) + >>> # batched matrix x broadcasted vector + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(4) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3]) + >>> # batched matrix x batched matrix + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(10, 4, 5) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3, 5]) + >>> # batched matrix x broadcasted matrix + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(4, 5) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3, 5]) + """ + +def matrix_exp(input: Tensor) -> Tensor: + r""" + matrix_exp(A) -> Tensor + + Alias for :func:`torch.linalg.matrix_exp`. + """ + +def matrix_power( + input: Tensor, + n: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + matrix_power(input, n, *, out=None) -> Tensor + + Alias for :func:`torch.linalg.matrix_power` + """ + +@overload +def max(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + max(input, *, out=None) -> Tensor + + Returns the maximum value of all elements in the ``input`` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6763, 0.7445, -2.2369]]) + >>> torch.max(a) + tensor(0.7445) + + .. function:: max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the maximum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each maximum value found + (argmax). + + If ``keepdim`` is ``True``, the output tensors are of the same size + as ``input`` except in the dimension ``dim`` where they are of size 1. + Otherwise, ``dim`` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than ``input``. + + .. note:: If there are multiple maximal values in a reduced row then + the indices of the first maximal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (max, max_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-1.2360, -0.2942, -0.1222, 0.8475], + [ 1.1949, -1.1127, -2.2379, -0.6702], + [ 1.5717, -0.9207, 0.1297, -1.8768], + [-0.6172, 1.0036, -0.6060, -0.2432]]) + >>> torch.max(a, 1) + torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1])) + >>> a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + >>> a.max(dim=1, keepdim=True) + torch.return_types.max( + values=tensor([[2.], [4.]]), + indices=tensor([[1], [1]])) + >>> a.max(dim=1, keepdim=False) + torch.return_types.max( + values=tensor([2., 4.]), + indices=tensor([1, 1])) + + .. function:: max(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.maximum`. + """ + +@overload +def max( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + max(input, *, out=None) -> Tensor + + Returns the maximum value of all elements in the ``input`` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6763, 0.7445, -2.2369]]) + >>> torch.max(a) + tensor(0.7445) + + .. function:: max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the maximum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each maximum value found + (argmax). + + If ``keepdim`` is ``True``, the output tensors are of the same size + as ``input`` except in the dimension ``dim`` where they are of size 1. + Otherwise, ``dim`` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than ``input``. + + .. note:: If there are multiple maximal values in a reduced row then + the indices of the first maximal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (max, max_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-1.2360, -0.2942, -0.1222, 0.8475], + [ 1.1949, -1.1127, -2.2379, -0.6702], + [ 1.5717, -0.9207, 0.1297, -1.8768], + [-0.6172, 1.0036, -0.6060, -0.2432]]) + >>> torch.max(a, 1) + torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1])) + >>> a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + >>> a.max(dim=1, keepdim=True) + torch.return_types.max( + values=tensor([[2.], [4.]]), + indices=tensor([[1], [1]])) + >>> a.max(dim=1, keepdim=False) + torch.return_types.max( + values=tensor([2., 4.]), + indices=tensor([1, 1])) + + .. function:: max(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.maximum`. + """ + +@overload +def max( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.max: + r""" + max(input, *, out=None) -> Tensor + + Returns the maximum value of all elements in the ``input`` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6763, 0.7445, -2.2369]]) + >>> torch.max(a) + tensor(0.7445) + + .. function:: max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the maximum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each maximum value found + (argmax). + + If ``keepdim`` is ``True``, the output tensors are of the same size + as ``input`` except in the dimension ``dim`` where they are of size 1. + Otherwise, ``dim`` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than ``input``. + + .. note:: If there are multiple maximal values in a reduced row then + the indices of the first maximal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (max, max_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-1.2360, -0.2942, -0.1222, 0.8475], + [ 1.1949, -1.1127, -2.2379, -0.6702], + [ 1.5717, -0.9207, 0.1297, -1.8768], + [-0.6172, 1.0036, -0.6060, -0.2432]]) + >>> torch.max(a, 1) + torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1])) + >>> a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + >>> a.max(dim=1, keepdim=True) + torch.return_types.max( + values=tensor([[2.], [4.]]), + indices=tensor([[1], [1]])) + >>> a.max(dim=1, keepdim=False) + torch.return_types.max( + values=tensor([2., 4.]), + indices=tensor([1, 1])) + + .. function:: max(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.maximum`. + """ + +@overload +def max( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.max: + r""" + max(input, *, out=None) -> Tensor + + Returns the maximum value of all elements in the ``input`` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6763, 0.7445, -2.2369]]) + >>> torch.max(a) + tensor(0.7445) + + .. function:: max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the maximum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each maximum value found + (argmax). + + If ``keepdim`` is ``True``, the output tensors are of the same size + as ``input`` except in the dimension ``dim`` where they are of size 1. + Otherwise, ``dim`` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than ``input``. + + .. note:: If there are multiple maximal values in a reduced row then + the indices of the first maximal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (max, max_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-1.2360, -0.2942, -0.1222, 0.8475], + [ 1.1949, -1.1127, -2.2379, -0.6702], + [ 1.5717, -0.9207, 0.1297, -1.8768], + [-0.6172, 1.0036, -0.6060, -0.2432]]) + >>> torch.max(a, 1) + torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1])) + >>> a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + >>> a.max(dim=1, keepdim=True) + torch.return_types.max( + values=tensor([[2.], [4.]]), + indices=tensor([[1], [1]])) + >>> a.max(dim=1, keepdim=False) + torch.return_types.max( + values=tensor([2., 4.]), + indices=tensor([1, 1])) + + .. function:: max(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.maximum`. + """ + +def max_pool1d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def max_pool1d_with_indices( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> tuple[Tensor, Tensor]: ... +def max_pool2d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def max_pool3d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def maximum( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + maximum(input, other, *, out=None) -> Tensor + + Computes the element-wise maximum of :attr:`input` and :attr:`other`. + + .. note:: + If one of the elements being compared is a NaN, then that element is returned. + :func:`maximum` is not supported for tensors with complex dtypes. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor((1, 2, -1)) + >>> b = torch.tensor((3, 0, 4)) + >>> torch.maximum(a, b) + tensor([3, 2, 4]) + """ + +@overload +def mean( + input: Tensor, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + mean(input, *, dtype=None) -> Tensor + + .. note:: + If the `input` tensor is empty, ``torch.mean()`` returns ``nan``. + This behavior is consistent with NumPy and follows the definition + that the mean over an empty set is undefined. + + + Returns the mean value of all elements in the :attr:`input` tensor. Input must be floating point or complex. + + Args: + input (Tensor): + the input tensor, either of floating point or complex dtype + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.2294, -0.5481, 1.3288]]) + >>> torch.mean(a) + tensor(0.3367) + + .. function:: mean(input, dim, keepdim=False, *, dtype=None, out=None) -> Tensor + :noindex: + + Returns the mean value of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + .. seealso:: + + :func:`torch.nanmean` computes the mean value of `non-NaN` elements. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.3841, 0.6320, 0.4254, -0.7384], + [-0.9644, 1.0131, -0.6549, -1.4279], + [-0.2951, -1.3350, -0.7694, 0.5600], + [ 1.0842, -0.9580, 0.3623, 0.2343]]) + >>> torch.mean(a, 1) + tensor([-0.0163, -0.5085, -0.4599, 0.1807]) + >>> torch.mean(a, 1, True) + tensor([[-0.0163], + [-0.5085], + [-0.4599], + [ 0.1807]]) + """ + +@overload +def mean( + input: Tensor, + dim: _int | _size | None, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + mean(input, *, dtype=None) -> Tensor + + .. note:: + If the `input` tensor is empty, ``torch.mean()`` returns ``nan``. + This behavior is consistent with NumPy and follows the definition + that the mean over an empty set is undefined. + + + Returns the mean value of all elements in the :attr:`input` tensor. Input must be floating point or complex. + + Args: + input (Tensor): + the input tensor, either of floating point or complex dtype + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.2294, -0.5481, 1.3288]]) + >>> torch.mean(a) + tensor(0.3367) + + .. function:: mean(input, dim, keepdim=False, *, dtype=None, out=None) -> Tensor + :noindex: + + Returns the mean value of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + .. seealso:: + + :func:`torch.nanmean` computes the mean value of `non-NaN` elements. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.3841, 0.6320, 0.4254, -0.7384], + [-0.9644, 1.0131, -0.6549, -1.4279], + [-0.2951, -1.3350, -0.7694, 0.5600], + [ 1.0842, -0.9580, 0.3623, 0.2343]]) + >>> torch.mean(a, 1) + tensor([-0.0163, -0.5085, -0.4599, 0.1807]) + >>> torch.mean(a, 1, True) + tensor([[-0.0163], + [-0.5085], + [-0.4599], + [ 0.1807]]) + """ + +@overload +def mean( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + mean(input, *, dtype=None) -> Tensor + + .. note:: + If the `input` tensor is empty, ``torch.mean()`` returns ``nan``. + This behavior is consistent with NumPy and follows the definition + that the mean over an empty set is undefined. + + + Returns the mean value of all elements in the :attr:`input` tensor. Input must be floating point or complex. + + Args: + input (Tensor): + the input tensor, either of floating point or complex dtype + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.2294, -0.5481, 1.3288]]) + >>> torch.mean(a) + tensor(0.3367) + + .. function:: mean(input, dim, keepdim=False, *, dtype=None, out=None) -> Tensor + :noindex: + + Returns the mean value of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + .. seealso:: + + :func:`torch.nanmean` computes the mean value of `non-NaN` elements. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.3841, 0.6320, 0.4254, -0.7384], + [-0.9644, 1.0131, -0.6549, -1.4279], + [-0.2951, -1.3350, -0.7694, 0.5600], + [ 1.0842, -0.9580, 0.3623, 0.2343]]) + >>> torch.mean(a, 1) + tensor([-0.0163, -0.5085, -0.4599, 0.1807]) + >>> torch.mean(a, 1, True) + tensor([[-0.0163], + [-0.5085], + [-0.4599], + [ 0.1807]]) + """ + +@overload +def median(input: Tensor) -> Tensor: + r""" + median(input) -> Tensor + + Returns the median of the values in :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements. In this case the lower of the two medians is returned. To + compute the mean of both medians, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + This function produces deterministic (sub)gradients unlike ``median(dim=0)`` + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.5219, -1.5212, 0.2202]]) + >>> torch.median(a) + tensor(0.2202) + + .. function:: median(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, and ``indices`` contains the index of the median values found in the dimension :attr:`dim`. + + By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + + If :attr:`keepdim` is ``True``, the output tensors are of the same size + as :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the outputs tensor having 1 fewer dimension than :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements in the dimension :attr:`dim`. In this case the lower of the + two medians is returned. To compute the mean of both medians in + :attr:`input`, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + ``indices`` does not necessarily contain the first occurrence of each + median value found, unless it is unique. + The exact implementation details are device-specific. + Do not expect the same result when run on CPU and GPU in general. + For the same reason do not expect the gradients to be deterministic. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.randn(4, 5) + >>> a + tensor([[ 0.2505, -0.3982, -0.9948, 0.3518, -1.3131], + [ 0.3180, -0.6993, 1.0436, 0.0438, 0.2270], + [-0.2751, 0.7303, 0.2192, 0.3321, 0.2488], + [ 1.0778, -1.9510, 0.7048, 0.4742, -0.7125]]) + >>> torch.median(a, 1) + torch.return_types.median(values=tensor([-0.3982, 0.2270, 0.2488, 0.4742]), indices=tensor([1, 4, 4, 3])) + """ + +@overload +def median( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.median: + r""" + median(input) -> Tensor + + Returns the median of the values in :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements. In this case the lower of the two medians is returned. To + compute the mean of both medians, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + This function produces deterministic (sub)gradients unlike ``median(dim=0)`` + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.5219, -1.5212, 0.2202]]) + >>> torch.median(a) + tensor(0.2202) + + .. function:: median(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, and ``indices`` contains the index of the median values found in the dimension :attr:`dim`. + + By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + + If :attr:`keepdim` is ``True``, the output tensors are of the same size + as :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the outputs tensor having 1 fewer dimension than :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements in the dimension :attr:`dim`. In this case the lower of the + two medians is returned. To compute the mean of both medians in + :attr:`input`, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + ``indices`` does not necessarily contain the first occurrence of each + median value found, unless it is unique. + The exact implementation details are device-specific. + Do not expect the same result when run on CPU and GPU in general. + For the same reason do not expect the gradients to be deterministic. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.randn(4, 5) + >>> a + tensor([[ 0.2505, -0.3982, -0.9948, 0.3518, -1.3131], + [ 0.3180, -0.6993, 1.0436, 0.0438, 0.2270], + [-0.2751, 0.7303, 0.2192, 0.3321, 0.2488], + [ 1.0778, -1.9510, 0.7048, 0.4742, -0.7125]]) + >>> torch.median(a, 1) + torch.return_types.median(values=tensor([-0.3982, 0.2270, 0.2488, 0.4742]), indices=tensor([1, 4, 4, 3])) + """ + +@overload +def median( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.median: + r""" + median(input) -> Tensor + + Returns the median of the values in :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements. In this case the lower of the two medians is returned. To + compute the mean of both medians, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + This function produces deterministic (sub)gradients unlike ``median(dim=0)`` + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.5219, -1.5212, 0.2202]]) + >>> torch.median(a) + tensor(0.2202) + + .. function:: median(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, and ``indices`` contains the index of the median values found in the dimension :attr:`dim`. + + By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + + If :attr:`keepdim` is ``True``, the output tensors are of the same size + as :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the outputs tensor having 1 fewer dimension than :attr:`input`. + + .. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements in the dimension :attr:`dim`. In this case the lower of the + two medians is returned. To compute the mean of both medians in + :attr:`input`, use :func:`torch.quantile` with ``q=0.5`` instead. + + .. warning:: + ``indices`` does not necessarily contain the first occurrence of each + median value found, unless it is unique. + The exact implementation details are device-specific. + Do not expect the same result when run on CPU and GPU in general. + For the same reason do not expect the gradients to be deterministic. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.randn(4, 5) + >>> a + tensor([[ 0.2505, -0.3982, -0.9948, 0.3518, -1.3131], + [ 0.3180, -0.6993, 1.0436, 0.0438, 0.2270], + [-0.2751, 0.7303, 0.2192, 0.3321, 0.2488], + [ 1.0778, -1.9510, 0.7048, 0.4742, -0.7125]]) + >>> torch.median(a, 1) + torch.return_types.median(values=tensor([-0.3982, 0.2270, 0.2488, 0.4742]), indices=tensor([1, 4, 4, 3])) + """ + +@overload +def min(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + min(input, *, out=None) -> Tensor + + Returns the minimum value of all elements in the :attr:`input` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6750, 1.0857, 1.7197]]) + >>> torch.min(a) + tensor(0.6750) + + .. function:: min(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the minimum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each minimum value found + (argmin). + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: If there are multiple minimal values in a reduced row then + the indices of the first minimal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the tuple of two output tensors (min, min_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.6248, 1.1334, -1.1899, -0.2803], + [-1.4644, -0.2635, -0.3651, 0.6134], + [ 0.2457, 0.0384, 1.0128, 0.7015], + [-0.1153, 2.9849, 2.1458, 0.5788]]) + >>> torch.min(a, 1) + torch.return_types.min(values=tensor([-1.1899, -1.4644, 0.0384, -0.1153]), indices=tensor([2, 0, 1, 0])) + + .. function:: min(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.minimum`. + """ + +@overload +def min( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + min(input, *, out=None) -> Tensor + + Returns the minimum value of all elements in the :attr:`input` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6750, 1.0857, 1.7197]]) + >>> torch.min(a) + tensor(0.6750) + + .. function:: min(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the minimum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each minimum value found + (argmin). + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: If there are multiple minimal values in a reduced row then + the indices of the first minimal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the tuple of two output tensors (min, min_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.6248, 1.1334, -1.1899, -0.2803], + [-1.4644, -0.2635, -0.3651, 0.6134], + [ 0.2457, 0.0384, 1.0128, 0.7015], + [-0.1153, 2.9849, 2.1458, 0.5788]]) + >>> torch.min(a, 1) + torch.return_types.min(values=tensor([-1.1899, -1.4644, 0.0384, -0.1153]), indices=tensor([2, 0, 1, 0])) + + .. function:: min(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.minimum`. + """ + +@overload +def min( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.min: + r""" + min(input, *, out=None) -> Tensor + + Returns the minimum value of all elements in the :attr:`input` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6750, 1.0857, 1.7197]]) + >>> torch.min(a) + tensor(0.6750) + + .. function:: min(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the minimum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each minimum value found + (argmin). + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: If there are multiple minimal values in a reduced row then + the indices of the first minimal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the tuple of two output tensors (min, min_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.6248, 1.1334, -1.1899, -0.2803], + [-1.4644, -0.2635, -0.3651, 0.6134], + [ 0.2457, 0.0384, 1.0128, 0.7015], + [-0.1153, 2.9849, 2.1458, 0.5788]]) + >>> torch.min(a, 1) + torch.return_types.min(values=tensor([-1.1899, -1.4644, 0.0384, -0.1153]), indices=tensor([2, 0, 1, 0])) + + .. function:: min(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.minimum`. + """ + +@overload +def min( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.min: + r""" + min(input, *, out=None) -> Tensor + + Returns the minimum value of all elements in the :attr:`input` tensor. + + .. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6750, 1.0857, 1.7197]]) + >>> torch.min(a) + tensor(0.6750) + + .. function:: min(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` is the minimum + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`. And ``indices`` is the index location of each minimum value found + (argmin). + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: If there are multiple minimal values in a reduced row then + the indices of the first minimal value are returned. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the tuple of two output tensors (min, min_indices) + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.6248, 1.1334, -1.1899, -0.2803], + [-1.4644, -0.2635, -0.3651, 0.6134], + [ 0.2457, 0.0384, 1.0128, 0.7015], + [-0.1153, 2.9849, 2.1458, 0.5788]]) + >>> torch.min(a, 1) + torch.return_types.min(values=tensor([-1.1899, -1.4644, 0.0384, -0.1153]), indices=tensor([2, 0, 1, 0])) + + .. function:: min(input, other, *, out=None) -> Tensor + :noindex: + + See :func:`torch.minimum`. + """ + +def minimum( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + minimum(input, other, *, out=None) -> Tensor + + Computes the element-wise minimum of :attr:`input` and :attr:`other`. + + .. note:: + If one of the elements being compared is a NaN, then that element is returned. + :func:`minimum` is not supported for tensors with complex dtypes. + + Args: + input (Tensor): the input tensor. + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor((1, 2, -1)) + >>> b = torch.tensor((3, 0, 4)) + >>> torch.minimum(a, b) + tensor([1, 0, -1]) + """ + +def miopen_batch_norm( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + training: _bool, + exponential_average_factor: _float, + epsilon: _float, +) -> tuple[Tensor, Tensor, Tensor]: ... +def miopen_convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, +) -> Tensor: ... +def miopen_convolution_add_relu( + input: Tensor, + weight: Tensor, + z: Tensor, + alpha: Number | _complex | None, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def miopen_convolution_relu( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + stride: Sequence[_int | SymInt], + padding: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def miopen_convolution_transpose( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: Sequence[_int | SymInt], + output_padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, +) -> Tensor: ... +def miopen_depthwise_convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, + benchmark: _bool, + deterministic: _bool, +) -> Tensor: ... +def miopen_rnn( + input: Tensor, + weight: tuple[Tensor, ...] | list[Tensor] | None, + weight_stride0: _int, + hx: Tensor, + cx: Tensor | None, + mode: _int, + hidden_size: _int, + num_layers: _int, + batch_first: _bool, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_sizes: _size, + dropout_state: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: ... +def mkldnn_adaptive_avg_pool2d( + input: Tensor, + output_size: _int | _size, + *, + out: Tensor | None = None, +) -> Tensor: ... +def mkldnn_convolution( + input: Tensor, + weight: Tensor, + bias: Tensor | None, + padding: Sequence[_int | SymInt], + stride: Sequence[_int | SymInt], + dilation: Sequence[_int | SymInt], + groups: _int | SymInt, +) -> Tensor: ... +def mkldnn_linear_backward_weights( + grad_output: Tensor, + input: Tensor, + weight: Tensor, + bias_defined: _bool, +) -> tuple[Tensor, Tensor]: ... +def mkldnn_max_pool2d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def mkldnn_max_pool3d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def mkldnn_rnn_layer( + input: Tensor, + weight0: Tensor, + weight1: Tensor, + weight2: Tensor, + weight3: Tensor, + hx_: Tensor, + cx_: Tensor, + reverse: _bool, + batch_sizes: _size, + mode: _int, + hidden_size: _int, + num_layers: _int, + has_biases: _bool, + bidirectional: _bool, + batch_first: _bool, + train: _bool, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: ... +@overload +def mm(input: Tensor, mat2: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + mm(input, mat2, out_dtype=None, *, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`input` and :attr:`mat2`. + + If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, :attr:`out` will be a :math:`(n \times p)` tensor. + + .. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + + Supports strided and sparse 2-D tensors as inputs, autograd with + respect to strided inputs. + + This operation has support for arguments with :ref:`sparse layouts`. + If :attr:`out` is provided its layout will be used. Otherwise, the result + layout will be deduced from that of :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.mm(mat1, mat2) + tensor([[ 0.4851, 0.5037, -0.3633], + [-0.0760, -3.6705, 2.4784]]) + """ + +@overload +def mm( + input: Tensor, + mat2: Tensor, + out_dtype: _dtype, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + mm(input, mat2, out_dtype=None, *, out=None) -> Tensor + + Performs a matrix multiplication of the matrices :attr:`input` and :attr:`mat2`. + + If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`mat2` is a + :math:`(m \times p)` tensor, :attr:`out` will be a :math:`(n \times p)` tensor. + + .. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + + Supports strided and sparse 2-D tensors as inputs, autograd with + respect to strided inputs. + + This operation has support for arguments with :ref:`sparse layouts`. + If :attr:`out` is provided its layout will be used. Otherwise, the result + layout will be deduced from that of :attr:`input`. + + + .. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request. + + This operator supports :ref:`TensorFloat32`. + + On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision` for backward. + + Args: + input (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.mm(mat1, mat2) + tensor([[ 0.4851, 0.5037, -0.3633], + [-0.0760, -3.6705, 2.4784]]) + """ + +@overload +def mode( + input: Tensor, + dim: _int = -1, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.mode: + r""" + mode(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + + Returns a namedtuple ``(values, indices)`` where ``values`` is the mode + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`, i.e. a value which appears most often + in that row, and ``indices`` is the index location of each mode value found. + + By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: This function is not defined for ``torch.cuda.Tensor`` yet. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> b = torch.tensor([[0, 0, 0, 2, 0, 0, 2], + ... [0, 3, 0, 0, 2, 0, 1], + ... [2, 2, 2, 0, 0, 0, 3], + ... [2, 2, 3, 0, 1, 1, 0], + ... [1, 1, 0, 0, 2, 0, 2]]) + >>> torch.mode(b, 0) + torch.return_types.mode( + values=tensor([0, 2, 0, 0, 0, 0, 2]), + indices=tensor([1, 3, 4, 4, 2, 4, 4])) + """ + +@overload +def mode( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.mode: + r""" + mode(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + + Returns a namedtuple ``(values, indices)`` where ``values`` is the mode + value of each row of the :attr:`input` tensor in the given dimension + :attr:`dim`, i.e. a value which appears most often + in that row, and ``indices`` is the index location of each mode value found. + + By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + + If :attr:`keepdim` is ``True``, the output tensors are of the same size as + :attr:`input` except in the dimension :attr:`dim` where they are of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting + in the output tensors having 1 fewer dimension than :attr:`input`. + + .. note:: This function is not defined for ``torch.cuda.Tensor`` yet. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + + Example:: + + >>> b = torch.tensor([[0, 0, 0, 2, 0, 0, 2], + ... [0, 3, 0, 0, 2, 0, 1], + ... [2, 2, 2, 0, 0, 0, 3], + ... [2, 2, 3, 0, 1, 1, 0], + ... [1, 1, 0, 0, 2, 0, 2]]) + >>> torch.mode(b, 0) + torch.return_types.mode( + values=tensor([0, 2, 0, 0, 0, 0, 2]), + indices=tensor([1, 3, 4, 4, 2, 4, 4])) + """ + +@overload +def moveaxis(input: Tensor, source: _int, destination: _int) -> Tensor: + r""" + moveaxis(input, source, destination) -> Tensor + + Alias for :func:`torch.movedim`. + + This function is equivalent to NumPy's moveaxis function. + + Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.moveaxis(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.moveaxis(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.moveaxis(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.moveaxis(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) + """ + +@overload +def moveaxis(input: Tensor, source: _size, destination: _size) -> Tensor: + r""" + moveaxis(input, source, destination) -> Tensor + + Alias for :func:`torch.movedim`. + + This function is equivalent to NumPy's moveaxis function. + + Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.moveaxis(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.moveaxis(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.moveaxis(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.moveaxis(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) + """ + +@overload +def movedim(input: Tensor, source: _int, destination: _int) -> Tensor: + r""" + movedim(input, source, destination) -> Tensor + + Moves the dimension(s) of :attr:`input` at the position(s) in :attr:`source` + to the position(s) in :attr:`destination`. + + Other dimensions of :attr:`input` that are not explicitly moved remain in + their original order and appear at the positions not specified in :attr:`destination`. + + Args: + input (Tensor): the input tensor. + source (int or tuple of ints): Original positions of the dims to move. These must be unique. + destination (int or tuple of ints): Destination positions for each of the original dims. These must also be unique. + + Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.movedim(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.movedim(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.movedim(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.movedim(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) + """ + +@overload +def movedim(input: Tensor, source: _size, destination: _size) -> Tensor: + r""" + movedim(input, source, destination) -> Tensor + + Moves the dimension(s) of :attr:`input` at the position(s) in :attr:`source` + to the position(s) in :attr:`destination`. + + Other dimensions of :attr:`input` that are not explicitly moved remain in + their original order and appear at the positions not specified in :attr:`destination`. + + Args: + input (Tensor): the input tensor. + source (int or tuple of ints): Original positions of the dims to move. These must be unique. + destination (int or tuple of ints): Destination positions for each of the original dims. These must also be unique. + + Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.movedim(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.movedim(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.movedim(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.movedim(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) + """ + +def msort(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + msort(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Sorts the elements of the :attr:`input` tensor along its first dimension + in ascending order by value. + + .. note:: `torch.msort(t)` is equivalent to `torch.sort(t, dim=0)[0]`. + See also :func:`torch.sort`. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.randn(3, 4) + >>> t + tensor([[-0.1321, 0.4370, -1.2631, -1.1289], + [-2.0527, -1.1250, 0.2275, 0.3077], + [-0.0881, -0.1259, -0.5495, 1.0284]]) + >>> torch.msort(t) + tensor([[-2.0527, -1.1250, -1.2631, -1.1289], + [-0.1321, -0.1259, -0.5495, 0.3077], + [-0.0881, 0.4370, 0.2275, 1.0284]]) + """ + +def mul( + input: Tensor | Number | _complex, + other: Tensor | Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + mul(input, other, *, out=None) -> Tensor + + Multiplies :attr:`input` by :attr:`other`. + + + .. math:: + \text{out}_i = \text{input}_i \times \text{other}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to multiply input by. + + Keyword args: + out (Tensor, optional): the output tensor. + + Examples:: + + >>> a = torch.randn(3) + >>> a + tensor([ 0.2015, -0.4255, 2.6087]) + >>> torch.mul(a, 100) + tensor([ 20.1494, -42.5491, 260.8663]) + + >>> b = torch.randn(4, 1) + >>> b + tensor([[ 1.1207], + [-0.3137], + [ 0.0700], + [ 0.8378]]) + >>> c = torch.randn(1, 4) + >>> c + tensor([[ 0.5146, 0.1216, -0.5244, 2.2382]]) + >>> torch.mul(b, c) + tensor([[ 0.5767, 0.1363, -0.5877, 2.5083], + [-0.1614, -0.0382, 0.1645, -0.7021], + [ 0.0360, 0.0085, -0.0367, 0.1567], + [ 0.4312, 0.1019, -0.4394, 1.8753]]) + """ + +def multinomial( + input: Tensor, + num_samples: _int | SymInt, + replacement: _bool = False, + *, + generator: Generator | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + multinomial(input, num_samples, replacement=False, *, generator=None, out=None) -> LongTensor + + Returns a tensor where each row contains :attr:`num_samples` indices sampled + from the multinomial (a stricter definition would be multivariate, + refer to :class:`torch.distributions.multinomial.Multinomial` for more details) + probability distribution located in the corresponding row + of tensor :attr:`input`. + + .. note:: + The rows of :attr:`input` do not need to sum to one (in which case we use + the values as weights), but must be non-negative, finite and have + a non-zero sum. + + Indices are ordered from left to right according to when each was sampled + (first samples are placed in first column). + + If :attr:`input` is a vector, :attr:`out` is a vector of size :attr:`num_samples`. + + If :attr:`input` is a matrix with `m` rows, :attr:`out` is an matrix of shape + :math:`(m \times \text{num\_samples})`. + + If replacement is ``True``, samples are drawn with replacement. + + If not, they are drawn without replacement, which means that when a + sample index is drawn for a row, it cannot be drawn again for that row. + + .. note:: + When drawn without replacement, :attr:`num_samples` must be lower than + number of non-zero elements in :attr:`input` (or the min number of non-zero + elements in each row of :attr:`input` if it is a matrix). + + Args: + input (Tensor): the input tensor containing probabilities + num_samples (int): number of samples to draw + replacement (bool, optional): whether to draw with replacement or not + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights + >>> torch.multinomial(weights, 2) + tensor([1, 2]) + >>> torch.multinomial(weights, 5) # ERROR! + RuntimeError: cannot sample n_sample > prob_dist.size(-1) samples without replacement + >>> torch.multinomial(weights, 4, replacement=True) + tensor([ 2, 1, 1, 1]) + """ + +@overload +def multiply( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + multiply(input, other, *, out=None) + + Alias for :func:`torch.mul`. + """ + +@overload +def multiply(input: Tensor, other: Number | _complex) -> Tensor: + r""" + multiply(input, other, *, out=None) + + Alias for :func:`torch.mul`. + """ + +def mv(input: Tensor, vec: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + mv(input, vec, *, out=None) -> Tensor + + Performs a matrix-vector product of the matrix :attr:`input` and the vector + :attr:`vec`. + + If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of + size :math:`m`, :attr:`out` will be 1-D of size :math:`n`. + + .. note:: This function does not :ref:`broadcast `. + + Args: + input (Tensor): matrix to be multiplied + vec (Tensor): vector to be multiplied + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.mv(mat, vec) + tensor([ 1.0404, -0.6361]) + """ + +def mvlgamma( + input: Tensor, + p: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + mvlgamma(input, p, *, out=None) -> Tensor + + Alias for :func:`torch.special.multigammaln`. + """ + +def nan_to_num( + input: Tensor, + nan: _float | None = None, + posinf: _float | None = None, + neginf: _float | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None) -> Tensor + + Replaces :literal:`NaN`, positive infinity, and negative infinity values in :attr:`input` + with the values specified by :attr:`nan`, :attr:`posinf`, and :attr:`neginf`, respectively. + By default, :literal:`NaN`\ s are replaced with zero, positive infinity is replaced with the + greatest finite value representable by :attr:`input`'s dtype, and negative infinity + is replaced with the least finite value representable by :attr:`input`'s dtype. + + Args: + input (Tensor): the input tensor. + nan (Number, optional): the value to replace :literal:`NaN`\s with. Default is zero. + posinf (Number, optional): if a Number, the value to replace positive infinity values with. + If None, positive infinity values are replaced with the greatest finite value representable by :attr:`input`'s dtype. + Default is None. + neginf (Number, optional): if a Number, the value to replace negative infinity values with. + If None, negative infinity values are replaced with the lowest finite value representable by :attr:`input`'s dtype. + Default is None. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.tensor([float('nan'), float('inf'), -float('inf'), 3.14]) + >>> torch.nan_to_num(x) + tensor([ 0.0000e+00, 3.4028e+38, -3.4028e+38, 3.1400e+00]) + >>> torch.nan_to_num(x, nan=2.0) + tensor([ 2.0000e+00, 3.4028e+38, -3.4028e+38, 3.1400e+00]) + >>> torch.nan_to_num(x, nan=2.0, posinf=1.0) + tensor([ 2.0000e+00, 1.0000e+00, -3.4028e+38, 3.1400e+00]) + """ + +def nan_to_num_( + input: Tensor, + nan: _float | None = None, + posinf: _float | None = None, + neginf: _float | None = None, +) -> Tensor: ... +def nanmean( + input: Tensor, + dim: _int | _size | None = None, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + nanmean(input, dim=None, keepdim=False, *, dtype=None, out=None) -> Tensor + + Computes the mean of all `non-NaN` elements along the specified dimensions. + Input must be floating point or complex. + + This function is identical to :func:`torch.mean` when there are no `NaN` values + in the :attr:`input` tensor. In the presence of `NaN`, :func:`torch.mean` will + propagate the `NaN` to the output whereas :func:`torch.nanmean` will ignore the + `NaN` values (`torch.nanmean(a)` is equivalent to `torch.mean(a[~a.isnan()])`). + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor, either of floating point or complex dtype + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + out (Tensor, optional): the output tensor. + + .. seealso:: + + :func:`torch.mean` computes the mean value, propagating `NaN`. + + Example:: + + >>> x = torch.tensor([[torch.nan, 1, 2], [1, 2, 3]]) + >>> x.mean() + tensor(nan) + >>> x.nanmean() + tensor(1.8000) + >>> x.mean(dim=0) + tensor([ nan, 1.5000, 2.5000]) + >>> x.nanmean(dim=0) + tensor([1.0000, 1.5000, 2.5000]) + + # If all elements in the reduced dimensions are NaN then the result is NaN + >>> torch.tensor([torch.nan]).nanmean() + tensor(nan) + """ + +@overload +def nanmedian(input: Tensor) -> Tensor: + r""" + nanmedian(input) -> Tensor + + Returns the median of the values in :attr:`input`, ignoring ``NaN`` values. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in :attr:`input`. + When :attr:`input` has one or more ``NaN`` values, :func:`torch.median` will always return ``NaN``, + while this function will return the median of the non-``NaN`` elements in :attr:`input`. + If all the elements in :attr:`input` are ``NaN`` it will also return ``NaN``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.tensor([1, float('nan'), 3, 2]) + >>> a.median() + tensor(nan) + >>> a.nanmedian() + tensor(2.) + + .. function:: nanmedian(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, ignoring ``NaN`` values, and ``indices`` contains the index of the median values + found in the dimension :attr:`dim`. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in a reduced row. When a reduced row has + one or more ``NaN`` values, :func:`torch.median` will always reduce it to ``NaN``, while this function will reduce it to the + median of the non-``NaN`` elements. If all the elements in a reduced row are ``NaN`` then it will be reduced to ``NaN``, too. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.tensor([[2, 3, 1], [float('nan'), 1, float('nan')]]) + >>> a + tensor([[2., 3., 1.], + [nan, 1., nan]]) + >>> a.median(0) + torch.return_types.median(values=tensor([nan, 1., nan]), indices=tensor([1, 1, 1])) + >>> a.nanmedian(0) + torch.return_types.nanmedian(values=tensor([2., 1., 1.]), indices=tensor([0, 1, 0])) + """ + +@overload +def nanmedian( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.nanmedian: + r""" + nanmedian(input) -> Tensor + + Returns the median of the values in :attr:`input`, ignoring ``NaN`` values. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in :attr:`input`. + When :attr:`input` has one or more ``NaN`` values, :func:`torch.median` will always return ``NaN``, + while this function will return the median of the non-``NaN`` elements in :attr:`input`. + If all the elements in :attr:`input` are ``NaN`` it will also return ``NaN``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.tensor([1, float('nan'), 3, 2]) + >>> a.median() + tensor(nan) + >>> a.nanmedian() + tensor(2.) + + .. function:: nanmedian(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, ignoring ``NaN`` values, and ``indices`` contains the index of the median values + found in the dimension :attr:`dim`. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in a reduced row. When a reduced row has + one or more ``NaN`` values, :func:`torch.median` will always reduce it to ``NaN``, while this function will reduce it to the + median of the non-``NaN`` elements. If all the elements in a reduced row are ``NaN`` then it will be reduced to ``NaN``, too. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.tensor([[2, 3, 1], [float('nan'), 1, float('nan')]]) + >>> a + tensor([[2., 3., 1.], + [nan, 1., nan]]) + >>> a.median(0) + torch.return_types.median(values=tensor([nan, 1., nan]), indices=tensor([1, 1, 1])) + >>> a.nanmedian(0) + torch.return_types.nanmedian(values=tensor([2., 1., 1.]), indices=tensor([0, 1, 0])) + """ + +@overload +def nanmedian( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.nanmedian: + r""" + nanmedian(input) -> Tensor + + Returns the median of the values in :attr:`input`, ignoring ``NaN`` values. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in :attr:`input`. + When :attr:`input` has one or more ``NaN`` values, :func:`torch.median` will always return ``NaN``, + while this function will return the median of the non-``NaN`` elements in :attr:`input`. + If all the elements in :attr:`input` are ``NaN`` it will also return ``NaN``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.tensor([1, float('nan'), 3, 2]) + >>> a.median() + tensor(nan) + >>> a.nanmedian() + tensor(2.) + + .. function:: nanmedian(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + + Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` + in the dimension :attr:`dim`, ignoring ``NaN`` values, and ``indices`` contains the index of the median values + found in the dimension :attr:`dim`. + + This function is identical to :func:`torch.median` when there are no ``NaN`` values in a reduced row. When a reduced row has + one or more ``NaN`` values, :func:`torch.median` will always reduce it to ``NaN``, while this function will reduce it to the + median of the non-``NaN`` elements. If all the elements in a reduced row are ``NaN`` then it will be reduced to ``NaN``, too. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + + Example:: + + >>> a = torch.tensor([[2, 3, 1], [float('nan'), 1, float('nan')]]) + >>> a + tensor([[2., 3., 1.], + [nan, 1., nan]]) + >>> a.median(0) + torch.return_types.median(values=tensor([nan, 1., nan]), indices=tensor([1, 1, 1])) + >>> a.nanmedian(0) + torch.return_types.nanmedian(values=tensor([2., 1., 1.]), indices=tensor([0, 1, 0])) + """ + +@overload +def nanquantile( + input: Tensor, + q: Tensor, + dim: _int | None = None, + keepdim: _bool = False, + *, + interpolation: str = "linear", + out: Tensor | None = None, +) -> Tensor: + r""" + nanquantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + + This is a variant of :func:`torch.quantile` that "ignores" ``NaN`` values, + computing the quantiles :attr:`q` as if ``NaN`` values in :attr:`input` did + not exist. If all values in a reduced row are ``NaN`` then the quantiles for + that reduction will be ``NaN``. See the documentation for :func:`torch.quantile`. + + Args: + input (Tensor): the input tensor. + q (float or Tensor): a scalar or 1D tensor of quantile values in the range [0, 1] + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword arguments: + interpolation (str): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.tensor([float('nan'), 1, 2]) + >>> t.quantile(0.5) + tensor(nan) + >>> t.nanquantile(0.5) + tensor(1.5000) + >>> t = torch.tensor([[float('nan'), float('nan')], [1, 2]]) + >>> t + tensor([[nan, nan], + [1., 2.]]) + >>> t.nanquantile(0.5, dim=0) + tensor([1., 2.]) + >>> t.nanquantile(0.5, dim=1) + tensor([ nan, 1.5000]) + """ + +@overload +def nanquantile( + input: Tensor, + q: _float, + dim: _int | None = None, + keepdim: _bool = False, + *, + interpolation: str = "linear", + out: Tensor | None = None, +) -> Tensor: + r""" + nanquantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + + This is a variant of :func:`torch.quantile` that "ignores" ``NaN`` values, + computing the quantiles :attr:`q` as if ``NaN`` values in :attr:`input` did + not exist. If all values in a reduced row are ``NaN`` then the quantiles for + that reduction will be ``NaN``. See the documentation for :func:`torch.quantile`. + + Args: + input (Tensor): the input tensor. + q (float or Tensor): a scalar or 1D tensor of quantile values in the range [0, 1] + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword arguments: + interpolation (str): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.tensor([float('nan'), 1, 2]) + >>> t.quantile(0.5) + tensor(nan) + >>> t.nanquantile(0.5) + tensor(1.5000) + >>> t = torch.tensor([[float('nan'), float('nan')], [1, 2]]) + >>> t + tensor([[nan, nan], + [1., 2.]]) + >>> t.nanquantile(0.5, dim=0) + tensor([1., 2.]) + >>> t.nanquantile(0.5, dim=1) + tensor([ nan, 1.5000]) + """ + +def nansum( + input: Tensor, + dim: _int | _size | None = None, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + nansum(input, *, dtype=None) -> Tensor + + Returns the sum of all elements, treating Not a Numbers (NaNs) as zero. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.tensor([1., 2., float('nan'), 4.]) + >>> torch.nansum(a) + tensor(7.) + + .. function:: nansum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the sum of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`, treating Not a Numbers (NaNs) as zero. + If :attr:`dim` is a list of dimensions, reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> torch.nansum(torch.tensor([1., float("nan")])) + tensor(1.) + >>> a = torch.tensor([[1, 2], [3., float("nan")]]) + >>> torch.nansum(a) + tensor(6.) + >>> torch.nansum(a, dim=0) + tensor([4., 2.]) + >>> torch.nansum(a, dim=1) + tensor([3., 3.]) + """ + +@overload +def narrow( + input: Tensor, + dim: _int, + start: Tensor, + length: _int | SymInt, +) -> Tensor: + r""" + narrow(input, dim, start, length) -> Tensor + + Returns a new tensor that is a narrowed version of :attr:`input` tensor. The + dimension :attr:`dim` is input from :attr:`start` to ``start + length``. The + returned tensor and :attr:`input` tensor share the same underlying storage. + + Args: + input (Tensor): the tensor to narrow + dim (int): the dimension along which to narrow + start (int or Tensor): index of the element to start the narrowed dimension + from. Can be negative, which means indexing from the end of `dim`. If + `Tensor`, it must be an 0-dim integral `Tensor` (bools not allowed) + length (int): length of the narrowed dimension, must be weakly positive + + Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> torch.narrow(x, 0, 0, 2) + tensor([[ 1, 2, 3], + [ 4, 5, 6]]) + >>> torch.narrow(x, 1, 1, 2) + tensor([[ 2, 3], + [ 5, 6], + [ 8, 9]]) + >>> torch.narrow(x, -1, torch.tensor(-1), 1) + tensor([[3], + [6], + [9]]) + """ + +@overload +def narrow( + input: Tensor, + dim: _int, + start: _int | SymInt, + length: _int | SymInt, +) -> Tensor: + r""" + narrow(input, dim, start, length) -> Tensor + + Returns a new tensor that is a narrowed version of :attr:`input` tensor. The + dimension :attr:`dim` is input from :attr:`start` to ``start + length``. The + returned tensor and :attr:`input` tensor share the same underlying storage. + + Args: + input (Tensor): the tensor to narrow + dim (int): the dimension along which to narrow + start (int or Tensor): index of the element to start the narrowed dimension + from. Can be negative, which means indexing from the end of `dim`. If + `Tensor`, it must be an 0-dim integral `Tensor` (bools not allowed) + length (int): length of the narrowed dimension, must be weakly positive + + Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> torch.narrow(x, 0, 0, 2) + tensor([[ 1, 2, 3], + [ 4, 5, 6]]) + >>> torch.narrow(x, 1, 1, 2) + tensor([[ 2, 3], + [ 5, 6], + [ 8, 9]]) + >>> torch.narrow(x, -1, torch.tensor(-1), 1) + tensor([[3], + [6], + [9]]) + """ + +def narrow_copy( + input: Tensor, + dim: _int, + start: _int | SymInt, + length: _int | SymInt, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + narrow_copy(input, dim, start, length, *, out=None) -> Tensor + + Same as :meth:`Tensor.narrow` except this returns a copy rather + than shared storage. This is primarily for sparse tensors, which + do not have a shared-storage narrow method. + + Args: + input (Tensor): the tensor to narrow + dim (int): the dimension along which to narrow + start (int): index of the element to start the narrowed dimension from. Can + be negative, which means indexing from the end of `dim` + length (int): length of the narrowed dimension, must be weakly positive + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> torch.narrow_copy(x, 0, 0, 2) + tensor([[ 1, 2, 3], + [ 4, 5, 6]]) + >>> torch.narrow_copy(x, 1, 1, 2) + tensor([[ 2, 3], + [ 5, 6], + [ 8, 9]]) + >>> s = torch.arange(16).reshape(2, 2, 2, 2).to_sparse(2) + >>> torch.narrow_copy(s, 0, 0, 1) + tensor(indices=tensor([[0, 0], + [0, 1]]), + values=tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]), + size=(1, 2, 2, 2), nnz=2, layout=torch.sparse_coo) + + .. seealso:: + + :func:`torch.narrow` for a non copy variant + """ + +def native_batch_norm( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + running_mean: Tensor | None, + running_var: Tensor | None, + training: _bool, + momentum: _float, + eps: _float, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> tuple[Tensor, Tensor, Tensor]: ... +def native_channel_shuffle(input: Tensor, groups: _int | SymInt) -> Tensor: ... +def native_dropout( + input: Tensor, + p: _float, + train: _bool | None, +) -> tuple[Tensor, Tensor]: ... +def native_group_norm( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + N: _int | SymInt, + C: _int | SymInt, + HxW: _int | SymInt, + group: _int, + eps: _float, +) -> tuple[Tensor, Tensor, Tensor]: ... +def native_layer_norm( + input: Tensor, + normalized_shape: Sequence[_int | SymInt], + weight: Tensor | None, + bias: Tensor | None, + eps: _float, +) -> tuple[Tensor, Tensor, Tensor]: ... +@overload +def native_norm( + input: Tensor, + p: Number | _complex | None, + dim: _int | _size, + keepdim: _bool, + dtype: _dtype | None, +) -> Tensor: ... +@overload +def native_norm(input: Tensor, p: Number | _complex = 2) -> Tensor: ... +@overload +def ne( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ne(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \neq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is not equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.ne(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [True, False]]) + """ + +@overload +def ne( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ne(input, other, *, out=None) -> Tensor + + Computes :math:`\text{input} \neq \text{other}` element-wise. + + + The second argument can be a number or a tensor whose shape is + :ref:`broadcastable ` with the first argument. + + Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + A boolean tensor that is True where :attr:`input` is not equal to :attr:`other` and False elsewhere + + Example:: + + >>> torch.ne(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [True, False]]) + """ + +def neg(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + neg(input, *, out=None) -> Tensor + + Returns a new tensor with the negative of the elements of :attr:`input`. + + .. math:: + \text{out} = -1 \times \text{input} + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(5) + >>> a + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) + >>> torch.neg(a) + tensor([-0.0090, 0.2262, 0.0682, 0.2866, -0.3940]) + """ + +def neg_(input: Tensor) -> Tensor: ... +def negative(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + negative(input, *, out=None) -> Tensor + + Alias for :func:`torch.neg` + """ + +def negative_(input: Tensor) -> Tensor: ... +def nextafter( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + nextafter(input, other, *, out=None) -> Tensor + + Return the next floating-point value after :attr:`input` towards :attr:`other`, elementwise. + + The shapes of ``input`` and ``other`` must be + :ref:`broadcastable `. + + Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> eps = torch.finfo(torch.float32).eps + >>> torch.nextafter(torch.tensor([1.0, 2.0]), torch.tensor([2.0, 1.0])) == torch.tensor([eps + 1, 2 - eps]) + tensor([True, True]) + """ + +@overload +def nonzero( + input: Tensor, + *, + as_tuple: Literal[False] = False, + out: Tensor | None = None, +) -> Tensor: + r""" + nonzero(input, *, out=None, as_tuple=False) -> LongTensor or tuple of LongTensors + + .. note:: + :func:`torch.nonzero(..., as_tuple=False) ` (default) returns a + 2-D tensor where each row is the index for a nonzero value. + + :func:`torch.nonzero(..., as_tuple=True) ` returns a tuple of 1-D + index tensors, allowing for advanced indexing, so ``x[x.nonzero(as_tuple=True)]`` + gives all nonzero values of tensor ``x``. Of the returned tuple, each index tensor + contains nonzero indices for a certain dimension. + + See below for more details on the two behaviors. + + When :attr:`input` is on CUDA, :func:`torch.nonzero() ` causes + host-device synchronization. + + **When** :attr:`as_tuple` **is** ``False`` **(default)**: + + Returns a tensor containing the indices of all non-zero elements of + :attr:`input`. Each row in the result contains the indices of a non-zero + element in :attr:`input`. The result is sorted lexicographically, with + the last index changing the fastest (C-style). + + If :attr:`input` has :math:`n` dimensions, then the resulting indices tensor + :attr:`out` is of size :math:`(z \times n)`, where :math:`z` is the total number of + non-zero elements in the :attr:`input` tensor. + + **When** :attr:`as_tuple` **is** ``True``: + + Returns a tuple of 1-D tensors, one for each dimension in :attr:`input`, + each containing the indices (in that dimension) of all non-zero elements of + :attr:`input` . + + If :attr:`input` has :math:`n` dimensions, then the resulting tuple contains :math:`n` + tensors of size :math:`z`, where :math:`z` is the total number of + non-zero elements in the :attr:`input` tensor. + + As a special case, when :attr:`input` has zero dimensions and a nonzero scalar + value, it is treated as a one-dimensional tensor with one element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (LongTensor, optional): the output tensor containing indices + + Returns: + LongTensor or tuple of LongTensor: If :attr:`as_tuple` is ``False``, the output + tensor containing indices. If :attr:`as_tuple` is ``True``, one 1-D tensor for + each dimension, containing the indices of each nonzero element along that + dimension. + + Example:: + + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1])) + tensor([[ 0], + [ 1], + [ 2], + [ 4]]) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]])) + tensor([[ 0, 0], + [ 1, 1], + [ 2, 2], + [ 3, 3]]) + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]), as_tuple=True) + (tensor([0, 1, 2, 4]),) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]]), as_tuple=True) + (tensor([0, 1, 2, 3]), tensor([0, 1, 2, 3])) + >>> torch.nonzero(torch.tensor(5), as_tuple=True) + (tensor([0]),) + """ + +@overload +def nonzero( + input: Tensor, + *, + as_tuple: Literal[True], +) -> tuple[Tensor, ...]: + r""" + nonzero(input, *, out=None, as_tuple=False) -> LongTensor or tuple of LongTensors + + .. note:: + :func:`torch.nonzero(..., as_tuple=False) ` (default) returns a + 2-D tensor where each row is the index for a nonzero value. + + :func:`torch.nonzero(..., as_tuple=True) ` returns a tuple of 1-D + index tensors, allowing for advanced indexing, so ``x[x.nonzero(as_tuple=True)]`` + gives all nonzero values of tensor ``x``. Of the returned tuple, each index tensor + contains nonzero indices for a certain dimension. + + See below for more details on the two behaviors. + + When :attr:`input` is on CUDA, :func:`torch.nonzero() ` causes + host-device synchronization. + + **When** :attr:`as_tuple` **is** ``False`` **(default)**: + + Returns a tensor containing the indices of all non-zero elements of + :attr:`input`. Each row in the result contains the indices of a non-zero + element in :attr:`input`. The result is sorted lexicographically, with + the last index changing the fastest (C-style). + + If :attr:`input` has :math:`n` dimensions, then the resulting indices tensor + :attr:`out` is of size :math:`(z \times n)`, where :math:`z` is the total number of + non-zero elements in the :attr:`input` tensor. + + **When** :attr:`as_tuple` **is** ``True``: + + Returns a tuple of 1-D tensors, one for each dimension in :attr:`input`, + each containing the indices (in that dimension) of all non-zero elements of + :attr:`input` . + + If :attr:`input` has :math:`n` dimensions, then the resulting tuple contains :math:`n` + tensors of size :math:`z`, where :math:`z` is the total number of + non-zero elements in the :attr:`input` tensor. + + As a special case, when :attr:`input` has zero dimensions and a nonzero scalar + value, it is treated as a one-dimensional tensor with one element. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (LongTensor, optional): the output tensor containing indices + + Returns: + LongTensor or tuple of LongTensor: If :attr:`as_tuple` is ``False``, the output + tensor containing indices. If :attr:`as_tuple` is ``True``, one 1-D tensor for + each dimension, containing the indices of each nonzero element along that + dimension. + + Example:: + + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1])) + tensor([[ 0], + [ 1], + [ 2], + [ 4]]) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]])) + tensor([[ 0, 0], + [ 1, 1], + [ 2, 2], + [ 3, 3]]) + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]), as_tuple=True) + (tensor([0, 1, 2, 4]),) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]]), as_tuple=True) + (tensor([0, 1, 2, 3]), tensor([0, 1, 2, 3])) + >>> torch.nonzero(torch.tensor(5), as_tuple=True) + (tensor([0]),) + """ + +def nonzero_static( + input: Tensor, + *, + size: _int | SymInt, + fill_value: _int = -1, + out: Tensor | None = None, +) -> Tensor: ... +def norm_except_dim(v: Tensor, pow: _int = 2, dim: _int = 0) -> Tensor: ... +@overload +def normal( + mean: Tensor, + std: Tensor, + *, + generator: Generator | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + normal(mean, std, *, generator=None, out=None) -> Tensor + + Returns a tensor of random numbers drawn from separate normal distributions + whose mean and standard deviation are given. + + The :attr:`mean` is a tensor with the mean of + each output element's normal distribution + + The :attr:`std` is a tensor with the standard deviation of + each output element's normal distribution + + The shapes of :attr:`mean` and :attr:`std` don't need to match, but the + total number of elements in each tensor need to be the same. + + .. note:: When the shapes do not match, the shape of :attr:`mean` + is used as the shape for the returned output tensor + + .. note:: When :attr:`std` is a CUDA tensor, this function synchronizes + its device with the CPU. + + Args: + mean (Tensor): the tensor of per-element means + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1)) + tensor([ 1.0425, 3.5672, 2.7969, 4.2925, 4.7229, 6.2134, + 8.0505, 8.1408, 9.0563, 10.0566]) + + .. function:: normal(mean=0.0, std, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means are shared among all drawn + elements. + + Args: + mean (float, optional): the mean for all distributions + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=0.5, std=torch.arange(1., 6.)) + tensor([-1.2793, -1.0732, -2.0687, 5.1177, -1.2303]) + + .. function:: normal(mean, std=1.0, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the standard deviations are shared among + all drawn elements. + + Args: + mean (Tensor): the tensor of per-element means + std (float, optional): the standard deviation for all distributions + + Keyword args: + out (Tensor, optional): the output tensor + + Example:: + + >>> torch.normal(mean=torch.arange(1., 6.)) + tensor([ 1.1552, 2.6148, 2.6535, 5.8318, 4.2361]) + + .. function:: normal(mean, std, size, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means and standard deviations are shared + among all drawn elements. The resulting tensor has size given by :attr:`size`. + + Args: + mean (float): the mean for all distributions + std (float): the standard deviation for all distributions + size (int...): a sequence of integers defining the shape of the output tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(2, 3, size=(1, 4)) + tensor([[-1.3987, -1.9544, 3.6048, 0.7909]]) + """ + +@overload +def normal( + mean: Tensor, + std: _float = 1, + *, + generator: Generator | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + normal(mean, std, *, generator=None, out=None) -> Tensor + + Returns a tensor of random numbers drawn from separate normal distributions + whose mean and standard deviation are given. + + The :attr:`mean` is a tensor with the mean of + each output element's normal distribution + + The :attr:`std` is a tensor with the standard deviation of + each output element's normal distribution + + The shapes of :attr:`mean` and :attr:`std` don't need to match, but the + total number of elements in each tensor need to be the same. + + .. note:: When the shapes do not match, the shape of :attr:`mean` + is used as the shape for the returned output tensor + + .. note:: When :attr:`std` is a CUDA tensor, this function synchronizes + its device with the CPU. + + Args: + mean (Tensor): the tensor of per-element means + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1)) + tensor([ 1.0425, 3.5672, 2.7969, 4.2925, 4.7229, 6.2134, + 8.0505, 8.1408, 9.0563, 10.0566]) + + .. function:: normal(mean=0.0, std, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means are shared among all drawn + elements. + + Args: + mean (float, optional): the mean for all distributions + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=0.5, std=torch.arange(1., 6.)) + tensor([-1.2793, -1.0732, -2.0687, 5.1177, -1.2303]) + + .. function:: normal(mean, std=1.0, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the standard deviations are shared among + all drawn elements. + + Args: + mean (Tensor): the tensor of per-element means + std (float, optional): the standard deviation for all distributions + + Keyword args: + out (Tensor, optional): the output tensor + + Example:: + + >>> torch.normal(mean=torch.arange(1., 6.)) + tensor([ 1.1552, 2.6148, 2.6535, 5.8318, 4.2361]) + + .. function:: normal(mean, std, size, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means and standard deviations are shared + among all drawn elements. The resulting tensor has size given by :attr:`size`. + + Args: + mean (float): the mean for all distributions + std (float): the standard deviation for all distributions + size (int...): a sequence of integers defining the shape of the output tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(2, 3, size=(1, 4)) + tensor([[-1.3987, -1.9544, 3.6048, 0.7909]]) + """ + +@overload +def normal( + mean: _float, + std: Tensor, + *, + generator: Generator | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + normal(mean, std, *, generator=None, out=None) -> Tensor + + Returns a tensor of random numbers drawn from separate normal distributions + whose mean and standard deviation are given. + + The :attr:`mean` is a tensor with the mean of + each output element's normal distribution + + The :attr:`std` is a tensor with the standard deviation of + each output element's normal distribution + + The shapes of :attr:`mean` and :attr:`std` don't need to match, but the + total number of elements in each tensor need to be the same. + + .. note:: When the shapes do not match, the shape of :attr:`mean` + is used as the shape for the returned output tensor + + .. note:: When :attr:`std` is a CUDA tensor, this function synchronizes + its device with the CPU. + + Args: + mean (Tensor): the tensor of per-element means + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1)) + tensor([ 1.0425, 3.5672, 2.7969, 4.2925, 4.7229, 6.2134, + 8.0505, 8.1408, 9.0563, 10.0566]) + + .. function:: normal(mean=0.0, std, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means are shared among all drawn + elements. + + Args: + mean (float, optional): the mean for all distributions + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=0.5, std=torch.arange(1., 6.)) + tensor([-1.2793, -1.0732, -2.0687, 5.1177, -1.2303]) + + .. function:: normal(mean, std=1.0, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the standard deviations are shared among + all drawn elements. + + Args: + mean (Tensor): the tensor of per-element means + std (float, optional): the standard deviation for all distributions + + Keyword args: + out (Tensor, optional): the output tensor + + Example:: + + >>> torch.normal(mean=torch.arange(1., 6.)) + tensor([ 1.1552, 2.6148, 2.6535, 5.8318, 4.2361]) + + .. function:: normal(mean, std, size, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means and standard deviations are shared + among all drawn elements. The resulting tensor has size given by :attr:`size`. + + Args: + mean (float): the mean for all distributions + std (float): the standard deviation for all distributions + size (int...): a sequence of integers defining the shape of the output tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(2, 3, size=(1, 4)) + tensor([[-1.3987, -1.9544, 3.6048, 0.7909]]) + """ + +@overload +def normal( + mean: _float, + std: _float, + size: Sequence[_int | SymInt], + *, + generator: Generator | None = None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + normal(mean, std, *, generator=None, out=None) -> Tensor + + Returns a tensor of random numbers drawn from separate normal distributions + whose mean and standard deviation are given. + + The :attr:`mean` is a tensor with the mean of + each output element's normal distribution + + The :attr:`std` is a tensor with the standard deviation of + each output element's normal distribution + + The shapes of :attr:`mean` and :attr:`std` don't need to match, but the + total number of elements in each tensor need to be the same. + + .. note:: When the shapes do not match, the shape of :attr:`mean` + is used as the shape for the returned output tensor + + .. note:: When :attr:`std` is a CUDA tensor, this function synchronizes + its device with the CPU. + + Args: + mean (Tensor): the tensor of per-element means + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1)) + tensor([ 1.0425, 3.5672, 2.7969, 4.2925, 4.7229, 6.2134, + 8.0505, 8.1408, 9.0563, 10.0566]) + + .. function:: normal(mean=0.0, std, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means are shared among all drawn + elements. + + Args: + mean (float, optional): the mean for all distributions + std (Tensor): the tensor of per-element standard deviations + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(mean=0.5, std=torch.arange(1., 6.)) + tensor([-1.2793, -1.0732, -2.0687, 5.1177, -1.2303]) + + .. function:: normal(mean, std=1.0, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the standard deviations are shared among + all drawn elements. + + Args: + mean (Tensor): the tensor of per-element means + std (float, optional): the standard deviation for all distributions + + Keyword args: + out (Tensor, optional): the output tensor + + Example:: + + >>> torch.normal(mean=torch.arange(1., 6.)) + tensor([ 1.1552, 2.6148, 2.6535, 5.8318, 4.2361]) + + .. function:: normal(mean, std, size, *, out=None) -> Tensor + :noindex: + + Similar to the function above, but the means and standard deviations are shared + among all drawn elements. The resulting tensor has size given by :attr:`size`. + + Args: + mean (float): the mean for all distributions + std (float): the standard deviation for all distributions + size (int...): a sequence of integers defining the shape of the output tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.normal(2, 3, size=(1, 4)) + tensor([[-1.3987, -1.9544, 3.6048, 0.7909]]) + """ + +@overload +def not_equal( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + not_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.ne`. + """ + +@overload +def not_equal( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + not_equal(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.ne`. + """ + +@overload +def nuclear_norm( + input: Tensor, + dim: _int | _size, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: ... +@overload +def nuclear_norm( + input: Tensor, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: ... +def numel(self: Tensor) -> _int: + r""" + numel(input: Tensor) -> int + + Returns the total number of elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> a = torch.randn(1, 2, 3, 4, 5) + >>> torch.numel(a) + 120 + >>> a = torch.zeros(4,4) + >>> torch.numel(a) + 16 + """ + +@overload +def ones( + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `1`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.ones(2, 3) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + + >>> torch.ones(5) + tensor([ 1., 1., 1., 1., 1.]) + """ + +@overload +def ones( + *size: _int | SymInt, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `1`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.ones(2, 3) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + + >>> torch.ones(5) + tensor([ 1., 1., 1., 1., 1.]) + """ + +@overload +def ones( + size: _size, + *, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `1`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.ones(2, 3) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + + >>> torch.ones(5) + tensor([ 1., 1., 1., 1., 1.]) + """ + +@overload +def ones( + *size: _int, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `1`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword arguments: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.ones(2, 3) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + + >>> torch.ones(5) + tensor([ 1., 1., 1., 1., 1.]) + """ + +def ones_like( + input: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + ones_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor filled with the scalar value `1`, with the same size as + :attr:`input`. ``torch.ones_like(input)`` is equivalent to + ``torch.ones(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + + .. warning:: + As of 0.4, this function does not support an :attr:`out` keyword. As an alternative, + the old ``torch.ones_like(input, out=output)`` is equivalent to + ``torch.ones(input.size(), out=output)``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + + Keyword arguments: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + + Example:: + + >>> input = torch.empty(2, 3) + >>> torch.ones_like(input) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + """ + +def orgqr( + input: Tensor, + input2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + orgqr(input, tau) -> Tensor + + Alias for :func:`torch.linalg.householder_product`. + """ + +def ormqr( + input: Tensor, + input2: Tensor, + input3: Tensor, + left: _bool = True, + transpose: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + ormqr(input, tau, other, left=True, transpose=False, *, out=None) -> Tensor + + Computes the matrix-matrix multiplication of a product of Householder matrices with a general matrix. + + Multiplies a :math:`m \times n` matrix `C` (given by :attr:`other`) with a matrix `Q`, + where `Q` is represented using Householder reflectors `(input, tau)`. + See `Representation of Orthogonal or Unitary Matrices`_ for further details. + + If :attr:`left` is `True` then `op(Q)` times `C` is computed, otherwise the result is `C` times `op(Q)`. + When :attr:`left` is `True`, the implicit matrix `Q` has size :math:`m \times m`. + It has size :math:`n \times n` otherwise. + If :attr:`transpose` is `True` then `op` is the conjugate transpose operation, otherwise it's a no-op. + + Supports inputs of float, double, cfloat and cdouble dtypes. + Also supports batched inputs, and, if the input is batched, the output is batched with the same dimensions. + + .. seealso:: + :func:`torch.geqrf` can be used to form the Householder representation `(input, tau)` of matrix `Q` + from the QR decomposition. + + .. note:: + This function supports backward but it is only fast when ``(input, tau)`` do not require gradients + and/or ``tau.size(-1)`` is very small. + `` + + Args: + input (Tensor): tensor of shape `(*, mn, k)` where `*` is zero or more batch dimensions + and `mn` equals to `m` or `n` depending on the :attr:`left`. + tau (Tensor): tensor of shape `(*, min(mn, k))` where `*` is zero or more batch dimensions. + other (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + left (bool): controls the order of multiplication. + transpose (bool): controls whether the matrix `Q` is conjugate transposed or not. + + Keyword args: + out (Tensor, optional): the output Tensor. Ignored if `None`. Default: `None`. + + .. _Representation of Orthogonal or Unitary Matrices: + https://www.netlib.org/lapack/lug/node128.html + """ + +def outer( + input: Tensor, + vec2: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + outer(input, vec2, *, out=None) -> Tensor + + Outer product of :attr:`input` and :attr:`vec2`. + If :attr:`input` is a vector of size :math:`n` and :attr:`vec2` is a vector of + size :math:`m`, then :attr:`out` must be a matrix of size :math:`(n \times m)`. + + .. note:: This function does not :ref:`broadcast `. + + Args: + input (Tensor): 1-D input vector + vec2 (Tensor): 1-D input vector + + Keyword args: + out (Tensor, optional): optional output matrix + + Example:: + + >>> v1 = torch.arange(1., 5.) + >>> v2 = torch.arange(1., 4.) + >>> torch.outer(v1, v2) + tensor([[ 1., 2., 3.], + [ 2., 4., 6.], + [ 3., 6., 9.], + [ 4., 8., 12.]]) + """ + +def pairwise_distance( + x1: Tensor, + x2: Tensor, + p: _float = 2, + eps: _float = 1e-06, + keepdim: _bool = False, +) -> Tensor: ... +def pdist(input: Tensor, p: _float = 2) -> Tensor: ... +def permute(input: Tensor, dims: _size) -> Tensor: + r""" + permute(input, dims) -> Tensor + + Returns a view of the original tensor :attr:`input` with its dimensions permuted. + + Args: + input (Tensor): the input tensor. + dims (tuple of int): The desired ordering of dimensions + + Example: + >>> x = torch.randn(2, 3, 5) + >>> x.size() + torch.Size([2, 3, 5]) + >>> torch.permute(x, (2, 0, 1)).size() + torch.Size([5, 2, 3]) + """ + +def permute_copy( + input: Tensor, + dims: _size, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.permute`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def pinverse(input: Tensor, rcond: _float = 1e-15) -> Tensor: + r""" + pinverse(input, rcond=1e-15) -> Tensor + + Alias for :func:`torch.linalg.pinv` + """ + +def pixel_shuffle(input: Tensor, upscale_factor: _int) -> Tensor: ... +def pixel_unshuffle(input: Tensor, downscale_factor: _int) -> Tensor: ... +def poisson(input: Tensor, generator: Generator | None = None) -> Tensor: + r""" + poisson(input, generator=None) -> Tensor + + Returns a tensor of the same size as :attr:`input` with each element + sampled from a Poisson distribution with rate parameter given by the corresponding + element in :attr:`input` i.e., + + .. math:: + \text{out}_i \sim \text{Poisson}(\text{input}_i) + + :attr:`input` must be non-negative. + + Args: + input (Tensor): the input tensor containing the rates of the Poisson distribution + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + + Example:: + + >>> rates = torch.rand(4, 4) * 5 # rate parameter between 0 and 5 + >>> torch.poisson(rates) + tensor([[9., 1., 3., 5.], + [8., 6., 6., 0.], + [0., 4., 5., 3.], + [2., 1., 4., 2.]]) + """ + +def poisson_nll_loss( + input: Tensor, + target: Tensor, + log_input: _bool, + full: _bool, + eps: _float, + reduction: _int, +) -> Tensor: ... +def polar( + abs: Tensor, + angle: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + polar(abs, angle, *, out=None) -> Tensor + + Constructs a complex tensor whose elements are Cartesian coordinates + corresponding to the polar coordinates with absolute value :attr:`abs` and angle + :attr:`angle`. + + .. math:: + \text{out} = \text{abs} \cdot \cos(\text{angle}) + \text{abs} \cdot \sin(\text{angle}) \cdot j + + .. note:: + `torch.polar` is similar to + `std::polar `_ + and does not compute the polar decomposition + of a complex tensor like Python's `cmath.polar` and SciPy's `linalg.polar` do. + The behavior of this function is undefined if `abs` is negative or NaN, or if `angle` is + infinite. + + + Args: + abs (Tensor): The absolute value the complex tensor. Must be float or double. + angle (Tensor): The angle of the complex tensor. Must be same dtype as + :attr:`abs`. + + Keyword args: + out (Tensor): If the inputs are ``torch.float32``, must be + ``torch.complex64``. If the inputs are ``torch.float64``, must be + ``torch.complex128``. + + Example:: + + >>> import numpy as np + >>> abs = torch.tensor([1, 2], dtype=torch.float64) + >>> angle = torch.tensor([np.pi / 2, 5 * np.pi / 4], dtype=torch.float64) + >>> z = torch.polar(abs, angle) + >>> z + tensor([(0.0000+1.0000j), (-1.4142-1.4142j)], dtype=torch.complex128) + """ + +def polygamma( + n: _int, + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + polygamma(n, input, *, out=None) -> Tensor + + Alias for :func:`torch.special.polygamma`. + """ + +def positive(input: Tensor) -> Tensor: + r""" + positive(input) -> Tensor + + Returns :attr:`input`. + Throws a runtime error if :attr:`input` is a bool tensor. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> t = torch.randn(5) + >>> t + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) + >>> torch.positive(t) + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) + """ + +@overload +def pow( + input: Tensor, + exponent: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + pow(input, exponent, *, out=None) -> Tensor + + Takes the power of each element in :attr:`input` with :attr:`exponent` and + returns a tensor with the result. + + :attr:`exponent` can be either a single ``float`` number or a `Tensor` + with the same number of elements as :attr:`input`. + + When :attr:`exponent` is a scalar value, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ \text{exponent} + + When :attr:`exponent` is a tensor, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ {\text{exponent}_i} + + When :attr:`exponent` is a tensor, the shapes of :attr:`input` + and :attr:`exponent` must be :ref:`broadcastable `. + + Args: + input (Tensor): the input tensor. + exponent (float or tensor): the exponent value + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.4331, 1.2475, 0.6834, -0.2791]) + >>> torch.pow(a, 2) + tensor([ 0.1875, 1.5561, 0.4670, 0.0779]) + >>> exp = torch.arange(1., 5.) + + >>> a = torch.arange(1., 5.) + >>> a + tensor([ 1., 2., 3., 4.]) + >>> exp + tensor([ 1., 2., 3., 4.]) + >>> torch.pow(a, exp) + tensor([ 1., 4., 27., 256.]) + + .. function:: pow(self, exponent, *, out=None) -> Tensor + :noindex: + + :attr:`self` is a scalar ``float`` value, and :attr:`exponent` is a tensor. + The returned tensor :attr:`out` is of the same shape as :attr:`exponent` + + The operation applied is: + + .. math:: + \text{out}_i = \text{self} ^ {\text{exponent}_i} + + Args: + self (float): the scalar base value for the power operation + exponent (Tensor): the exponent tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> exp = torch.arange(1., 5.) + >>> base = 2 + >>> torch.pow(base, exp) + tensor([ 2., 4., 8., 16.]) + """ + +@overload +def pow( + self: Number | _complex, + exponent: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + pow(input, exponent, *, out=None) -> Tensor + + Takes the power of each element in :attr:`input` with :attr:`exponent` and + returns a tensor with the result. + + :attr:`exponent` can be either a single ``float`` number or a `Tensor` + with the same number of elements as :attr:`input`. + + When :attr:`exponent` is a scalar value, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ \text{exponent} + + When :attr:`exponent` is a tensor, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ {\text{exponent}_i} + + When :attr:`exponent` is a tensor, the shapes of :attr:`input` + and :attr:`exponent` must be :ref:`broadcastable `. + + Args: + input (Tensor): the input tensor. + exponent (float or tensor): the exponent value + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.4331, 1.2475, 0.6834, -0.2791]) + >>> torch.pow(a, 2) + tensor([ 0.1875, 1.5561, 0.4670, 0.0779]) + >>> exp = torch.arange(1., 5.) + + >>> a = torch.arange(1., 5.) + >>> a + tensor([ 1., 2., 3., 4.]) + >>> exp + tensor([ 1., 2., 3., 4.]) + >>> torch.pow(a, exp) + tensor([ 1., 4., 27., 256.]) + + .. function:: pow(self, exponent, *, out=None) -> Tensor + :noindex: + + :attr:`self` is a scalar ``float`` value, and :attr:`exponent` is a tensor. + The returned tensor :attr:`out` is of the same shape as :attr:`exponent` + + The operation applied is: + + .. math:: + \text{out}_i = \text{self} ^ {\text{exponent}_i} + + Args: + self (float): the scalar base value for the power operation + exponent (Tensor): the exponent tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> exp = torch.arange(1., 5.) + >>> base = 2 + >>> torch.pow(base, exp) + tensor([ 2., 4., 8., 16.]) + """ + +@overload +def pow( + input: Tensor, + exponent: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + pow(input, exponent, *, out=None) -> Tensor + + Takes the power of each element in :attr:`input` with :attr:`exponent` and + returns a tensor with the result. + + :attr:`exponent` can be either a single ``float`` number or a `Tensor` + with the same number of elements as :attr:`input`. + + When :attr:`exponent` is a scalar value, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ \text{exponent} + + When :attr:`exponent` is a tensor, the operation applied is: + + .. math:: + \text{out}_i = x_i ^ {\text{exponent}_i} + + When :attr:`exponent` is a tensor, the shapes of :attr:`input` + and :attr:`exponent` must be :ref:`broadcastable `. + + Args: + input (Tensor): the input tensor. + exponent (float or tensor): the exponent value + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.4331, 1.2475, 0.6834, -0.2791]) + >>> torch.pow(a, 2) + tensor([ 0.1875, 1.5561, 0.4670, 0.0779]) + >>> exp = torch.arange(1., 5.) + + >>> a = torch.arange(1., 5.) + >>> a + tensor([ 1., 2., 3., 4.]) + >>> exp + tensor([ 1., 2., 3., 4.]) + >>> torch.pow(a, exp) + tensor([ 1., 4., 27., 256.]) + + .. function:: pow(self, exponent, *, out=None) -> Tensor + :noindex: + + :attr:`self` is a scalar ``float`` value, and :attr:`exponent` is a tensor. + The returned tensor :attr:`out` is of the same shape as :attr:`exponent` + + The operation applied is: + + .. math:: + \text{out}_i = \text{self} ^ {\text{exponent}_i} + + Args: + self (float): the scalar base value for the power operation + exponent (Tensor): the exponent tensor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> exp = torch.arange(1., 5.) + >>> base = 2 + >>> torch.pow(base, exp) + tensor([ 2., 4., 8., 16.]) + """ + +def prelu(input: Tensor, weight: Tensor) -> Tensor: ... +@overload +def prod(input: Tensor, *, dtype: _dtype | None = None) -> Tensor: + r""" + prod(input: Tensor, *, dtype: Optional[_dtype]) -> Tensor + + Returns the product of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[-0.8020, 0.5428, -1.5854]]) + >>> torch.prod(a) + tensor(0.6902) + + .. function:: prod(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the product of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensor having 1 fewer dimension than :attr:`input`. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 2) + >>> a + tensor([[ 0.5261, -0.3837], + [ 1.1857, -0.2498], + [-1.1646, 0.0705], + [ 1.1131, -1.0629]]) + >>> torch.prod(a, 1) + tensor([-0.2018, -0.2962, -0.0821, -1.1831]) + """ + +@overload +def prod( + input: Tensor, + dim: _int, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + prod(input: Tensor, *, dtype: Optional[_dtype]) -> Tensor + + Returns the product of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[-0.8020, 0.5428, -1.5854]]) + >>> torch.prod(a) + tensor(0.6902) + + .. function:: prod(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the product of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensor having 1 fewer dimension than :attr:`input`. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 2) + >>> a + tensor([[ 0.5261, -0.3837], + [ 1.1857, -0.2498], + [-1.1646, 0.0705], + [ 1.1131, -1.0629]]) + >>> torch.prod(a, 1) + tensor([-0.2018, -0.2962, -0.0821, -1.1831]) + """ + +@overload +def prod( + input: Tensor, + dim: str | EllipsisType | None, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + prod(input: Tensor, *, dtype: Optional[_dtype]) -> Tensor + + Returns the product of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[-0.8020, 0.5428, -1.5854]]) + >>> torch.prod(a) + tensor(0.6902) + + .. function:: prod(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the product of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in + the output tensor having 1 fewer dimension than :attr:`input`. + + Args: + input (Tensor): the input tensor. + + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 2) + >>> a + tensor([[ 0.5261, -0.3837], + [ 1.1857, -0.2498], + [-1.1646, 0.0705], + [ 1.1131, -1.0629]]) + >>> torch.prod(a, 1) + tensor([-0.2018, -0.2962, -0.0821, -1.1831]) + """ + +def promote_types(type1: _dtype, type2: _dtype) -> _dtype: + r""" + promote_types(type1, type2) -> dtype + + Returns the :class:`torch.dtype` with the smallest size and scalar kind that is + not smaller nor of lower kind than either `type1` or `type2`. See type promotion + :ref:`documentation ` for more information on the type + promotion logic. + + Args: + type1 (:class:`torch.dtype`) + type2 (:class:`torch.dtype`) + + Example:: + + >>> torch.promote_types(torch.int32, torch.float32) + torch.float32 + >>> torch.promote_types(torch.uint8, torch.long) + torch.long + """ + +def put( + input: Tensor, + index: Tensor, + source: Tensor, + accumulate: _bool = False, +) -> Tensor: ... +def q_per_channel_axis(input: Tensor) -> _int: ... +def q_per_channel_scales(input: Tensor) -> Tensor: ... +def q_per_channel_zero_points(input: Tensor) -> Tensor: ... +def q_scale(input: Tensor) -> _float: ... +def q_zero_point(input: Tensor) -> _int: ... +def qr( + input: Tensor, + some: _bool = True, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.qr: + r""" + qr(input: Tensor, some: bool = True, *, out: Union[Tensor, Tuple[Tensor, ...], List[Tensor], None]) -> (Tensor, Tensor) + + Computes the QR decomposition of a matrix or a batch of matrices :attr:`input`, + and returns a namedtuple (Q, R) of tensors such that :math:`\text{input} = Q R` + with :math:`Q` being an orthogonal matrix or batch of orthogonal matrices and + :math:`R` being an upper triangular matrix or batch of upper triangular matrices. + + If :attr:`some` is ``True``, then this function returns the thin (reduced) QR factorization. + Otherwise, if :attr:`some` is ``False``, this function returns the complete QR factorization. + + .. warning:: + + :func:`torch.qr` is deprecated in favor of :func:`torch.linalg.qr` + and will be removed in a future PyTorch release. The boolean parameter :attr:`some` has been + replaced with a string parameter :attr:`mode`. + + ``Q, R = torch.qr(A)`` should be replaced with + + .. code:: python + + Q, R = torch.linalg.qr(A) + + ``Q, R = torch.qr(A, some=False)`` should be replaced with + + .. code:: python + + Q, R = torch.linalg.qr(A, mode="complete") + + .. warning:: + If you plan to backpropagate through QR, note that the current backward implementation + is only well-defined when the first :math:`\min(input.size(-1), input.size(-2))` + columns of :attr:`input` are linearly independent. + This behavior will probably change once QR supports pivoting. + + .. note:: This function uses LAPACK for CPU inputs and MAGMA for CUDA inputs, + and may produce different (valid) decompositions on different device types + or different platforms. + + Args: + input (Tensor): the input tensor of size :math:`(*, m, n)` where `*` is zero or more + batch dimensions consisting of matrices of dimension :math:`m \times n`. + some (bool, optional): Set to ``True`` for reduced QR decomposition and ``False`` for + complete QR decomposition. If `k = min(m, n)` then: + + * ``some=True`` : returns `(Q, R)` with dimensions (m, k), (k, n) (default) + + * ``'some=False'``: returns `(Q, R)` with dimensions (m, m), (m, n) + + Keyword args: + out (tuple, optional): tuple of `Q` and `R` tensors. + The dimensions of `Q` and `R` are detailed in the description of :attr:`some` above. + + Example:: + + >>> a = torch.tensor([[12., -51, 4], [6, 167, -68], [-4, 24, -41]]) + >>> q, r = torch.qr(a) + >>> q + tensor([[-0.8571, 0.3943, 0.3314], + [-0.4286, -0.9029, -0.0343], + [ 0.2857, -0.1714, 0.9429]]) + >>> r + tensor([[ -14.0000, -21.0000, 14.0000], + [ 0.0000, -175.0000, 70.0000], + [ 0.0000, 0.0000, -35.0000]]) + >>> torch.mm(q, r).round() + tensor([[ 12., -51., 4.], + [ 6., 167., -68.], + [ -4., 24., -41.]]) + >>> torch.mm(q.t(), q).round() + tensor([[ 1., 0., 0.], + [ 0., 1., -0.], + [ 0., -0., 1.]]) + >>> a = torch.randn(3, 4, 5) + >>> q, r = torch.qr(a, some=False) + >>> torch.allclose(torch.matmul(q, r), a) + True + >>> torch.allclose(torch.matmul(q.mT, q), torch.eye(5)) + True + """ + +@overload +def quantile( + input: Tensor, + q: Tensor, + dim: _int | None = None, + keepdim: _bool = False, + *, + interpolation: str = "linear", + out: Tensor | None = None, +) -> Tensor: + r""" + quantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + + Computes the q-th quantiles of each row of the :attr:`input` tensor along the dimension :attr:`dim`. + + To compute the quantile, we map q in [0, 1] to the range of indices [0, n] to find the location + of the quantile in the sorted input. If the quantile lies between two data points ``a < b`` with + indices ``i`` and ``j`` in the sorted order, result is computed according to the given + :attr:`interpolation` method as follows: + + - ``linear``: ``a + (b - a) * fraction``, where ``fraction`` is the fractional part of the computed quantile index. + - ``lower``: ``a``. + - ``higher``: ``b``. + - ``nearest``: ``a`` or ``b``, whichever's index is closer to the computed quantile index (rounding down for .5 fractions). + - ``midpoint``: ``(a + b) / 2``. + + If :attr:`q` is a 1D tensor, the first dimension of the output represents the quantiles and has size + equal to the size of :attr:`q`, the remaining dimensions are what remains from the reduction. + + .. note:: + By default :attr:`dim` is ``None`` resulting in the :attr:`input` tensor being flattened before computation. + + Args: + input (Tensor): the input tensor. + q (float or Tensor): a scalar or 1D tensor of values in the range [0, 1]. + + dim (int, optional): the dimension to reduce. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword arguments: + interpolation (str, optional): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(2, 3) + >>> a + tensor([[ 0.0795, -1.2117, 0.9765], + [ 1.1707, 0.6706, 0.4884]]) + >>> q = torch.tensor([0.25, 0.5, 0.75]) + >>> torch.quantile(a, q, dim=1, keepdim=True) + tensor([[[-0.5661], + [ 0.5795]], + + [[ 0.0795], + [ 0.6706]], + + [[ 0.5280], + [ 0.9206]]]) + >>> torch.quantile(a, q, dim=1, keepdim=True).shape + torch.Size([3, 2, 1]) + >>> a = torch.arange(4.) + >>> a + tensor([0., 1., 2., 3.]) + >>> torch.quantile(a, 0.6, interpolation='linear') + tensor(1.8000) + >>> torch.quantile(a, 0.6, interpolation='lower') + tensor(1.) + >>> torch.quantile(a, 0.6, interpolation='higher') + tensor(2.) + >>> torch.quantile(a, 0.6, interpolation='midpoint') + tensor(1.5000) + >>> torch.quantile(a, 0.6, interpolation='nearest') + tensor(2.) + >>> torch.quantile(a, 0.4, interpolation='nearest') + tensor(1.) + """ + +@overload +def quantile( + input: Tensor, + q: _float, + dim: _int | None = None, + keepdim: _bool = False, + *, + interpolation: str = "linear", + out: Tensor | None = None, +) -> Tensor: + r""" + quantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + + Computes the q-th quantiles of each row of the :attr:`input` tensor along the dimension :attr:`dim`. + + To compute the quantile, we map q in [0, 1] to the range of indices [0, n] to find the location + of the quantile in the sorted input. If the quantile lies between two data points ``a < b`` with + indices ``i`` and ``j`` in the sorted order, result is computed according to the given + :attr:`interpolation` method as follows: + + - ``linear``: ``a + (b - a) * fraction``, where ``fraction`` is the fractional part of the computed quantile index. + - ``lower``: ``a``. + - ``higher``: ``b``. + - ``nearest``: ``a`` or ``b``, whichever's index is closer to the computed quantile index (rounding down for .5 fractions). + - ``midpoint``: ``(a + b) / 2``. + + If :attr:`q` is a 1D tensor, the first dimension of the output represents the quantiles and has size + equal to the size of :attr:`q`, the remaining dimensions are what remains from the reduction. + + .. note:: + By default :attr:`dim` is ``None`` resulting in the :attr:`input` tensor being flattened before computation. + + Args: + input (Tensor): the input tensor. + q (float or Tensor): a scalar or 1D tensor of values in the range [0, 1]. + + dim (int, optional): the dimension to reduce. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword arguments: + interpolation (str, optional): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(2, 3) + >>> a + tensor([[ 0.0795, -1.2117, 0.9765], + [ 1.1707, 0.6706, 0.4884]]) + >>> q = torch.tensor([0.25, 0.5, 0.75]) + >>> torch.quantile(a, q, dim=1, keepdim=True) + tensor([[[-0.5661], + [ 0.5795]], + + [[ 0.0795], + [ 0.6706]], + + [[ 0.5280], + [ 0.9206]]]) + >>> torch.quantile(a, q, dim=1, keepdim=True).shape + torch.Size([3, 2, 1]) + >>> a = torch.arange(4.) + >>> a + tensor([0., 1., 2., 3.]) + >>> torch.quantile(a, 0.6, interpolation='linear') + tensor(1.8000) + >>> torch.quantile(a, 0.6, interpolation='lower') + tensor(1.) + >>> torch.quantile(a, 0.6, interpolation='higher') + tensor(2.) + >>> torch.quantile(a, 0.6, interpolation='midpoint') + tensor(1.5000) + >>> torch.quantile(a, 0.6, interpolation='nearest') + tensor(2.) + >>> torch.quantile(a, 0.4, interpolation='nearest') + tensor(1.) + """ + +def quantize_per_channel( + input: Tensor, + scales: Tensor, + zero_points: Tensor, + axis: _int, + dtype: _dtype, +) -> Tensor: + r""" + quantize_per_channel(input, scales, zero_points, axis, dtype) -> Tensor + + Converts a float tensor to a per-channel quantized tensor with given scales and zero points. + + Arguments: + input (Tensor): float tensor to quantize + scales (Tensor): float 1D tensor of scales to use, size should match ``input.size(axis)`` + zero_points (int): integer 1D tensor of offset to use, size should match ``input.size(axis)`` + axis (int): dimension on which apply per-channel quantization + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + + Returns: + Tensor: A newly quantized tensor + + Example:: + + >>> x = torch.tensor([[-1.0, 0.0], [1.0, 2.0]]) + >>> torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8) + tensor([[-1., 0.], + [ 1., 2.]], size=(2, 2), dtype=torch.quint8, + quantization_scheme=torch.per_channel_affine, + scale=tensor([0.1000, 0.0100], dtype=torch.float64), + zero_point=tensor([10, 0]), axis=0) + >>> torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8).int_repr() + tensor([[ 0, 10], + [100, 200]], dtype=torch.uint8) + """ + +@overload +def quantize_per_tensor( + input: Tensor, + scale: Tensor, + zero_point: Tensor, + dtype: _dtype, +) -> Tensor: + r""" + quantize_per_tensor(input, scale, zero_point, dtype) -> Tensor + + Converts a float tensor to a quantized tensor with given scale and zero point. + + Arguments: + input (Tensor): float tensor or list of tensors to quantize + scale (float or Tensor): scale to apply in quantization formula + zero_point (int or Tensor): offset in integer value that maps to float zero + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + + Returns: + Tensor: A newly quantized tensor or list of quantized tensors. + + Example:: + + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8).int_repr() + tensor([ 0, 10, 20, 30], dtype=torch.uint8) + >>> torch.quantize_per_tensor([torch.tensor([-1.0, 0.0]), torch.tensor([-2.0, 2.0])], + >>> torch.tensor([0.1, 0.2]), torch.tensor([10, 20]), torch.quint8) + (tensor([-1., 0.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10), + tensor([-2., 2.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=20)) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.tensor(0.1), torch.tensor(10), torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.10, zero_point=10) + """ + +@overload +def quantize_per_tensor( + input: Tensor, + scale: _float, + zero_point: _int, + dtype: _dtype, +) -> Tensor: + r""" + quantize_per_tensor(input, scale, zero_point, dtype) -> Tensor + + Converts a float tensor to a quantized tensor with given scale and zero point. + + Arguments: + input (Tensor): float tensor or list of tensors to quantize + scale (float or Tensor): scale to apply in quantization formula + zero_point (int or Tensor): offset in integer value that maps to float zero + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + + Returns: + Tensor: A newly quantized tensor or list of quantized tensors. + + Example:: + + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8).int_repr() + tensor([ 0, 10, 20, 30], dtype=torch.uint8) + >>> torch.quantize_per_tensor([torch.tensor([-1.0, 0.0]), torch.tensor([-2.0, 2.0])], + >>> torch.tensor([0.1, 0.2]), torch.tensor([10, 20]), torch.quint8) + (tensor([-1., 0.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10), + tensor([-2., 2.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=20)) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.tensor(0.1), torch.tensor(10), torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.10, zero_point=10) + """ + +@overload +def quantize_per_tensor( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + scales: Tensor, + zero_points: Tensor, + dtype: _dtype, +) -> tuple[Tensor, ...]: + r""" + quantize_per_tensor(input, scale, zero_point, dtype) -> Tensor + + Converts a float tensor to a quantized tensor with given scale and zero point. + + Arguments: + input (Tensor): float tensor or list of tensors to quantize + scale (float or Tensor): scale to apply in quantization formula + zero_point (int or Tensor): offset in integer value that maps to float zero + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + + Returns: + Tensor: A newly quantized tensor or list of quantized tensors. + + Example:: + + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8).int_repr() + tensor([ 0, 10, 20, 30], dtype=torch.uint8) + >>> torch.quantize_per_tensor([torch.tensor([-1.0, 0.0]), torch.tensor([-2.0, 2.0])], + >>> torch.tensor([0.1, 0.2]), torch.tensor([10, 20]), torch.quint8) + (tensor([-1., 0.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10), + tensor([-2., 2.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=20)) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.tensor(0.1), torch.tensor(10), torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.10, zero_point=10) + """ + +def quantize_per_tensor_dynamic( + input: Tensor, + dtype: _dtype, + reduce_range: _bool, +) -> Tensor: + r""" + quantize_per_tensor_dynamic(input, dtype, reduce_range) -> Tensor + + Converts a float tensor to a quantized tensor with scale and zero_point calculated + dynamically based on the input. + + Arguments: + input (Tensor): float tensor or list of tensors to quantize + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8`` + reduce_range (bool): a flag to indicate whether to reduce the range of quantized + data by 1 bit, it's required to avoid instruction overflow for some hardwares + + Returns: + Tensor: A newly (dynamically) quantized tensor + + Example:: + + >>> t = torch.quantize_per_tensor_dynamic(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.quint8, False) + >>> print(t) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.011764705882352941, + zero_point=85) + >>> t.int_repr() + tensor([ 0, 85, 170, 255], dtype=torch.uint8) + """ + +def quantized_batch_norm( + input: Tensor, + weight: Tensor | None, + bias: Tensor | None, + mean: Tensor, + var: Tensor, + eps: _float, + output_scale: _float, + output_zero_point: _int, +) -> Tensor: + r""" + quantized_batch_norm(input, weight=None, bias=None, mean, var, eps, output_scale, output_zero_point) -> Tensor + + Applies batch normalization on a 4D (NCHW) quantized tensor. + + .. math:: + + y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta + + Arguments: + input (Tensor): quantized tensor + weight (Tensor): float tensor that corresponds to the gamma, size C + bias (Tensor): float tensor that corresponds to the beta, size C + mean (Tensor): float mean value in batch normalization, size C + var (Tensor): float tensor for variance, size C + eps (float): a value added to the denominator for numerical stability. + output_scale (float): output quantized tensor scale + output_zero_point (int): output quantized tensor zero_point + + Returns: + Tensor: A quantized tensor with batch normalization applied. + + Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2, 2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_batch_norm(qx, torch.ones(2), torch.zeros(2), torch.rand(2), torch.rand(2), 0.00001, 0.2, 2) + tensor([[[[-0.2000, -0.2000], + [ 1.6000, -0.2000]], + + [[-0.4000, -0.4000], + [-0.4000, 0.6000]]], + + + [[[-0.2000, -0.2000], + [-0.2000, -0.2000]], + + [[ 0.6000, -0.4000], + [ 0.6000, -0.4000]]]], size=(2, 2, 2, 2), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=2) + """ + +def quantized_gru_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor, + b_hh: Tensor, + packed_ih: Tensor, + packed_hh: Tensor, + col_offsets_ih: Tensor, + col_offsets_hh: Tensor, + scale_ih: Number | _complex, + scale_hh: Number | _complex, + zero_point_ih: Number | _complex, + zero_point_hh: Number | _complex, +) -> Tensor: ... +def quantized_lstm_cell( + input: Tensor, + hx: tuple[Tensor, ...] | list[Tensor] | None, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor, + b_hh: Tensor, + packed_ih: Tensor, + packed_hh: Tensor, + col_offsets_ih: Tensor, + col_offsets_hh: Tensor, + scale_ih: Number | _complex, + scale_hh: Number | _complex, + zero_point_ih: Number | _complex, + zero_point_hh: Number | _complex, +) -> tuple[Tensor, Tensor]: ... +def quantized_max_pool1d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: + r""" + quantized_max_pool1d(input, kernel_size, stride=[], padding=0, dilation=1, ceil_mode=False) -> Tensor + + Applies a 1D max pooling over an input quantized tensor composed of several input planes. + + Arguments: + input (Tensor): quantized tensor + kernel_size (list of int): the size of the sliding window + stride (``list of int``, optional): the stride of the sliding window + padding (``list of int``, optional): padding to be added on both sides, must be >= 0 and <= kernel_size / 2 + dilation (``list of int``, optional): The stride between elements within a sliding window, must be > 0. Default 1 + ceil_mode (bool, optional): If True, will use ceil instead of floor to compute the output shape. + Defaults to False. + + + Returns: + Tensor: A quantized tensor with max_pool1d applied. + + Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_max_pool1d(qx, [2]) + tensor([[0.0000], + [1.5000]], size=(2, 1), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=1.5, zero_point=3) + """ + +def quantized_max_pool2d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: + r""" + quantized_max_pool2d(input, kernel_size, stride=[], padding=0, dilation=1, ceil_mode=False) -> Tensor + + Applies a 2D max pooling over an input quantized tensor composed of several input planes. + + Arguments: + input (Tensor): quantized tensor + kernel_size (``list of int``): the size of the sliding window + stride (``list of int``, optional): the stride of the sliding window + padding (``list of int``, optional): padding to be added on both sides, must be >= 0 and <= kernel_size / 2 + dilation (``list of int``, optional): The stride between elements within a sliding window, must be > 0. Default 1 + ceil_mode (bool, optional): If True, will use ceil instead of floor to compute the output shape. + Defaults to False. + + + Returns: + Tensor: A quantized tensor with max_pool2d applied. + + Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2, 2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_max_pool2d(qx, [2,2]) + tensor([[[[1.5000]], + + [[1.5000]]], + + + [[[0.0000]], + + [[0.0000]]]], size=(2, 2, 1, 1), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=1.5, zero_point=3) + """ + +def quantized_max_pool3d( + input: Tensor, + kernel_size: _int | _size, + stride: _int | _size = (), + padding: _int | _size = 0, + dilation: _int | _size = 1, + ceil_mode: _bool = False, +) -> Tensor: ... +def quantized_rnn_relu_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor, + b_hh: Tensor, + packed_ih: Tensor, + packed_hh: Tensor, + col_offsets_ih: Tensor, + col_offsets_hh: Tensor, + scale_ih: Number | _complex, + scale_hh: Number | _complex, + zero_point_ih: Number | _complex, + zero_point_hh: Number | _complex, +) -> Tensor: ... +def quantized_rnn_tanh_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor, + b_hh: Tensor, + packed_ih: Tensor, + packed_hh: Tensor, + col_offsets_ih: Tensor, + col_offsets_hh: Tensor, + scale_ih: Number | _complex, + scale_hh: Number | _complex, + zero_point_ih: Number | _complex, + zero_point_hh: Number | _complex, +) -> Tensor: ... +def rad2deg(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + rad2deg(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with each of the elements of :attr:`input` + converted from angles in radians to degrees. + + Args: + input (Tensor): the input tensor. + + Keyword arguments: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([[3.142, -3.142], [6.283, -6.283], [1.570, -1.570]]) + >>> torch.rad2deg(a) + tensor([[ 180.0233, -180.0233], + [ 359.9894, -359.9894], + [ 89.9544, -89.9544]]) + """ + +def rad2deg_(input: Tensor) -> Tensor: ... +@overload +def rand( + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + *size: _int | SymInt, + generator: Generator | None, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + *size: _int | SymInt, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + *size: _int | SymInt, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + size: Sequence[_int | SymInt], + *, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +@overload +def rand( + *size: _int | SymInt, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a tensor filled with random numbers from a uniform distribution + on the interval :math:`[0, 1)` + + The shape of the tensor is defined by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) + """ + +def rand_like( + input: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + rand_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same size as :attr:`input` that is filled with + random numbers from a uniform distribution on the interval :math:`[0, 1)`. + ``torch.rand_like(input)`` is equivalent to + ``torch.rand(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + """ + +@overload +def randint( + low: _int, + high: _int, + size: _size, + *, + generator: Generator | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint( + high: _int, + size: _size, + *, + generator: Generator | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint( + high: _int | SymInt, + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint( + high: _int | SymInt, + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint( + low: _int | SymInt, + high: _int | SymInt, + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint( + low: _int | SymInt, + high: _int | SymInt, + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with random integers generated uniformly + between :attr:`low` (inclusive) and :attr:`high` (exclusive). + + The shape of the tensor is defined by the variable argument :attr:`size`. + + .. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + """ + +@overload +def randint_like( + input: Tensor, + low: _int | SymInt, + high: _int | SymInt, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint_like(input, low=0, high, \*, dtype=None, layout=torch.strided, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same shape as Tensor :attr:`input` filled with + random integers generated uniformly between :attr:`low` (inclusive) and + :attr:`high` (exclusive). + + .. note: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + """ + +@overload +def randint_like( + input: Tensor, + high: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint_like(input, low=0, high, \*, dtype=None, layout=torch.strided, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same shape as Tensor :attr:`input` filled with + random integers generated uniformly between :attr:`low` (inclusive) and + :attr:`high` (exclusive). + + .. note: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + """ + +@overload +def randint_like( + input: Tensor, + high: _int | SymInt, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randint_like(input, low=0, high, \*, dtype=None, layout=torch.strided, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same shape as Tensor :attr:`input` filled with + random integers generated uniformly between :attr:`low` (inclusive) and + :attr:`high` (exclusive). + + .. note: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + """ + +@overload +def randn( + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + *size: _int | SymInt, + generator: Generator | None, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + size: Sequence[_int | SymInt], + *, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + *size: _int | SymInt, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + *size: _int | SymInt, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + size: Sequence[_int | SymInt], + *, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +@overload +def randn( + *size: _int | SymInt, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + + Returns a tensor filled with random numbers from a normal distribution + with mean `0` and variance `1` (also called the standard normal + distribution). + + .. math:: + \text{out}_{i} \sim \mathcal{N}(0, 1) + + For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and + unit variance as + + .. math:: + \text{out}_{i} \sim \mathcal{CN}(0, 1) + + This is equivalent to separately sampling the real :math:`(\operatorname{Re})` and imaginary + :math:`(\operatorname{Im})` part of :math:`\text{out}_i` as + + .. math:: + \operatorname{Re}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}),\quad + \operatorname{Im}(\text{out}_{i}) \sim \mathcal{N}(0, \frac{1}{2}) + + The shape of the tensor is defined by the variable argument :attr:`size`. + + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + + .. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution + """ + +def randn_like( + input: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randn_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor with the same size as :attr:`input` that is filled with + random numbers from a normal distribution with mean 0 and variance 1. Please refer to :func:`torch.randn` for the + sampling process of complex dtypes. ``torch.randn_like(input)`` is equivalent to + ``torch.randn(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + """ + +@overload +def randperm( + n: _int | SymInt, + *, + generator: Generator | None, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randperm(n, *, generator=None, out=None, dtype=torch.int64,layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a random permutation of integers from ``0`` to ``n - 1``. + + Args: + n (int): the upper bound (exclusive) + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randperm(4) + tensor([2, 1, 0, 3]) + """ + +@overload +def randperm( + n: _int | SymInt, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + randperm(n, *, generator=None, out=None, dtype=torch.int64,layout=torch.strided, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Returns a random permutation of integers from ``0`` to ``n - 1``. + + Args: + n (int): the upper bound (exclusive) + + Keyword args: + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: ``torch.int64``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + Example:: + + >>> torch.randperm(4) + tensor([2, 1, 0, 3]) + """ + +def range( + start: Number, + end: Number, + step: Number = 1, + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + range(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a 1-D tensor of size :math:`\left\lfloor \frac{\text{end} - \text{start}}{\text{step}} \right\rfloor + 1` + with values from :attr:`start` to :attr:`end` with step :attr:`step`. Step is + the gap between two values in the tensor. + + .. math:: + \text{out}_{i+1} = \text{out}_i + \text{step}. + + .. warning:: + This function is deprecated and will be removed in a future release because its behavior is inconsistent with + Python's range builtin. Instead, use :func:`torch.arange`, which produces values in [start, end). + + Args: + start (float, optional): the starting value for the set of points. Default: ``0``. + end (float): the ending value for the set of points + step (float, optional): the gap between each pair of adjacent points. Default: ``1``. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `step` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.range(1, 4) + tensor([ 1., 2., 3., 4.]) + >>> torch.range(1, 4, 0.5) + tensor([ 1.0000, 1.5000, 2.0000, 2.5000, 3.0000, 3.5000, 4.0000]) + """ + +def ravel(input: Tensor) -> Tensor: + r""" + ravel(input) -> Tensor + + Return a contiguous flattened tensor. A copy is made only if needed. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.ravel(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + """ + +def real(input: Tensor) -> Tensor: + r""" + real(input) -> Tensor + + Returns a new tensor containing real values of the :attr:`self` tensor. + The returned tensor and :attr:`self` share the same underlying storage. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.real + tensor([ 0.3100, -0.5445, -1.6492, -0.0638]) + """ + +def reciprocal(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + reciprocal(input, *, out=None) -> Tensor + + Returns a new tensor with the reciprocal of the elements of :attr:`input` + + .. math:: + \text{out}_{i} = \frac{1}{\text{input}_{i}} + + .. note:: + Unlike NumPy's reciprocal, torch.reciprocal supports integral inputs. Integral + inputs to reciprocal are automatically :ref:`promoted ` to + the default scalar type. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.4595, -2.1219, -1.4314, 0.7298]) + >>> torch.reciprocal(a) + tensor([-2.1763, -0.4713, -0.6986, 1.3702]) + """ + +def reciprocal_(input: Tensor) -> Tensor: ... +def relu(input: Tensor) -> Tensor: ... +def relu_(input: Tensor) -> Tensor: ... +@overload +def remainder( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + remainder(input, other, *, out=None) -> Tensor + + Computes + `Python's modulus operation `_ + entrywise. The result has the same sign as the divisor :attr:`other` and its absolute value + is less than that of :attr:`other`. + + It may also be defined in terms of :func:`torch.div` as + + .. code:: python + + torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and float inputs. + + .. note:: + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + See :func:`torch.fmod` for how division by zero is handled. + + .. seealso:: + + :func:`torch.fmod` which implements C++'s `std::fmod `_. + This one is defined in terms of division rounding towards zero. + + Args: + input (Tensor or Scalar): the dividend + other (Tensor or Scalar): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.remainder(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([ 1., 0., 1., 1., 0., 1.]) + >>> torch.remainder(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([ -0.5000, -1.0000, 0.0000, -0.5000, -1.0000 ]) + """ + +@overload +def remainder(self: Number | _complex, other: Tensor) -> Tensor: + r""" + remainder(input, other, *, out=None) -> Tensor + + Computes + `Python's modulus operation `_ + entrywise. The result has the same sign as the divisor :attr:`other` and its absolute value + is less than that of :attr:`other`. + + It may also be defined in terms of :func:`torch.div` as + + .. code:: python + + torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and float inputs. + + .. note:: + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + See :func:`torch.fmod` for how division by zero is handled. + + .. seealso:: + + :func:`torch.fmod` which implements C++'s `std::fmod `_. + This one is defined in terms of division rounding towards zero. + + Args: + input (Tensor or Scalar): the dividend + other (Tensor or Scalar): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.remainder(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([ 1., 0., 1., 1., 0., 1.]) + >>> torch.remainder(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([ -0.5000, -1.0000, 0.0000, -0.5000, -1.0000 ]) + """ + +@overload +def remainder( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + remainder(input, other, *, out=None) -> Tensor + + Computes + `Python's modulus operation `_ + entrywise. The result has the same sign as the divisor :attr:`other` and its absolute value + is less than that of :attr:`other`. + + It may also be defined in terms of :func:`torch.div` as + + .. code:: python + + torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer and float inputs. + + .. note:: + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + See :func:`torch.fmod` for how division by zero is handled. + + .. seealso:: + + :func:`torch.fmod` which implements C++'s `std::fmod `_. + This one is defined in terms of division rounding towards zero. + + Args: + input (Tensor or Scalar): the dividend + other (Tensor or Scalar): the divisor + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.remainder(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([ 1., 0., 1., 1., 0., 1.]) + >>> torch.remainder(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([ -0.5000, -1.0000, 0.0000, -0.5000, -1.0000 ]) + """ + +def renorm( + input: Tensor, + p: Number | _complex, + dim: _int, + maxnorm: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + renorm(input, p, dim, maxnorm, *, out=None) -> Tensor + + Returns a tensor where each sub-tensor of :attr:`input` along dimension + :attr:`dim` is normalized such that the `p`-norm of the sub-tensor is lower + than the value :attr:`maxnorm` + + .. note:: If the norm of a row is lower than `maxnorm`, the row is unchanged + + Args: + input (Tensor): the input tensor. + p (float): the power for the norm computation + dim (int): the dimension to slice over to get the sub-tensors + maxnorm (float): the maximum norm to keep each sub-tensor under + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.ones(3, 3) + >>> x[1].fill_(2) + tensor([ 2., 2., 2.]) + >>> x[2].fill_(3) + tensor([ 3., 3., 3.]) + >>> x + tensor([[ 1., 1., 1.], + [ 2., 2., 2.], + [ 3., 3., 3.]]) + >>> torch.renorm(x, 1, 0, 5) + tensor([[ 1.0000, 1.0000, 1.0000], + [ 1.6667, 1.6667, 1.6667], + [ 1.6667, 1.6667, 1.6667]]) + """ + +@overload +def repeat_interleave( + input: Tensor, + repeats: Tensor, + dim: _int | None = None, + *, + output_size: _int | SymInt | None = None, +) -> Tensor: + r""" + repeat_interleave(input, repeats, dim=None, *, output_size=None) -> Tensor + + Repeat elements of a tensor. + + .. warning:: + + This is different from :meth:`torch.Tensor.repeat` but similar to ``numpy.repeat``. + + Args: + input (Tensor): the input tensor. + repeats (Tensor or int): The number of repetitions for each element. + repeats is broadcasted to fit the shape of the given axis. + dim (int, optional): The dimension along which to repeat values. + By default, use the flattened input array, and return a flat output + array. + + Keyword args: + output_size (int, optional): Total output size for the given axis + ( e.g. sum of repeats). If given, it will avoid stream synchronization + needed to calculate output shape of the tensor. + + Returns: + Tensor: Repeated tensor which has the same shape as input, except along the given axis. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.repeat_interleave(2) + tensor([1, 1, 2, 2, 3, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.repeat_interleave(y, 2) + tensor([1, 1, 2, 2, 3, 3, 4, 4]) + >>> torch.repeat_interleave(y, 3, dim=1) + tensor([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0) + tensor([[1, 2], + [3, 4], + [3, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0, output_size=3) + tensor([[1, 2], + [3, 4], + [3, 4]]) + + If the `repeats` is `tensor([n1, n2, n3, ...])`, then the output will be + `tensor([0, 0, ..., 1, 1, ..., 2, 2, ..., ...])` where `0` appears `n1` times, + `1` appears `n2` times, `2` appears `n3` times, etc. + + .. function:: repeat_interleave(repeats, *) -> Tensor + :noindex: + + Repeats 0 repeats[0] times, 1 repeats[1] times, 2 repeats[2] times, etc. + + Args: + repeats (Tensor): The number of repetitions for each element. + + Returns: + Tensor: Repeated tensor of size `sum(repeats)`. + + Example:: + + >>> torch.repeat_interleave(torch.tensor([1, 2, 3])) + tensor([0, 1, 1, 2, 2, 2]) + """ + +@overload +def repeat_interleave( + repeats: Tensor, + *, + output_size: _int | SymInt | None = None, +) -> Tensor: + r""" + repeat_interleave(input, repeats, dim=None, *, output_size=None) -> Tensor + + Repeat elements of a tensor. + + .. warning:: + + This is different from :meth:`torch.Tensor.repeat` but similar to ``numpy.repeat``. + + Args: + input (Tensor): the input tensor. + repeats (Tensor or int): The number of repetitions for each element. + repeats is broadcasted to fit the shape of the given axis. + dim (int, optional): The dimension along which to repeat values. + By default, use the flattened input array, and return a flat output + array. + + Keyword args: + output_size (int, optional): Total output size for the given axis + ( e.g. sum of repeats). If given, it will avoid stream synchronization + needed to calculate output shape of the tensor. + + Returns: + Tensor: Repeated tensor which has the same shape as input, except along the given axis. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.repeat_interleave(2) + tensor([1, 1, 2, 2, 3, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.repeat_interleave(y, 2) + tensor([1, 1, 2, 2, 3, 3, 4, 4]) + >>> torch.repeat_interleave(y, 3, dim=1) + tensor([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0) + tensor([[1, 2], + [3, 4], + [3, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0, output_size=3) + tensor([[1, 2], + [3, 4], + [3, 4]]) + + If the `repeats` is `tensor([n1, n2, n3, ...])`, then the output will be + `tensor([0, 0, ..., 1, 1, ..., 2, 2, ..., ...])` where `0` appears `n1` times, + `1` appears `n2` times, `2` appears `n3` times, etc. + + .. function:: repeat_interleave(repeats, *) -> Tensor + :noindex: + + Repeats 0 repeats[0] times, 1 repeats[1] times, 2 repeats[2] times, etc. + + Args: + repeats (Tensor): The number of repetitions for each element. + + Returns: + Tensor: Repeated tensor of size `sum(repeats)`. + + Example:: + + >>> torch.repeat_interleave(torch.tensor([1, 2, 3])) + tensor([0, 1, 1, 2, 2, 2]) + """ + +@overload +def repeat_interleave( + input: Tensor, + repeats: _int | SymInt, + dim: _int | None = None, + *, + output_size: _int | SymInt | None = None, +) -> Tensor: + r""" + repeat_interleave(input, repeats, dim=None, *, output_size=None) -> Tensor + + Repeat elements of a tensor. + + .. warning:: + + This is different from :meth:`torch.Tensor.repeat` but similar to ``numpy.repeat``. + + Args: + input (Tensor): the input tensor. + repeats (Tensor or int): The number of repetitions for each element. + repeats is broadcasted to fit the shape of the given axis. + dim (int, optional): The dimension along which to repeat values. + By default, use the flattened input array, and return a flat output + array. + + Keyword args: + output_size (int, optional): Total output size for the given axis + ( e.g. sum of repeats). If given, it will avoid stream synchronization + needed to calculate output shape of the tensor. + + Returns: + Tensor: Repeated tensor which has the same shape as input, except along the given axis. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.repeat_interleave(2) + tensor([1, 1, 2, 2, 3, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.repeat_interleave(y, 2) + tensor([1, 1, 2, 2, 3, 3, 4, 4]) + >>> torch.repeat_interleave(y, 3, dim=1) + tensor([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0) + tensor([[1, 2], + [3, 4], + [3, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0, output_size=3) + tensor([[1, 2], + [3, 4], + [3, 4]]) + + If the `repeats` is `tensor([n1, n2, n3, ...])`, then the output will be + `tensor([0, 0, ..., 1, 1, ..., 2, 2, ..., ...])` where `0` appears `n1` times, + `1` appears `n2` times, `2` appears `n3` times, etc. + + .. function:: repeat_interleave(repeats, *) -> Tensor + :noindex: + + Repeats 0 repeats[0] times, 1 repeats[1] times, 2 repeats[2] times, etc. + + Args: + repeats (Tensor): The number of repetitions for each element. + + Returns: + Tensor: Repeated tensor of size `sum(repeats)`. + + Example:: + + >>> torch.repeat_interleave(torch.tensor([1, 2, 3])) + tensor([0, 1, 1, 2, 2, 2]) + """ + +def reshape(input: Tensor, shape: Sequence[_int | SymInt]) -> Tensor: + r""" + reshape(input, shape) -> Tensor + + Returns a tensor with the same data and number of elements as :attr:`input`, + but with the specified shape. When possible, the returned tensor will be a view + of :attr:`input`. Otherwise, it will be a copy. Contiguous inputs and inputs + with compatible strides can be reshaped without copying, but you should not + depend on the copying vs. viewing behavior. + + See :meth:`torch.Tensor.view` on when it is possible to return a view. + + A single dimension may be -1, in which case it's inferred from the remaining + dimensions and the number of elements in :attr:`input`. + + Args: + input (Tensor): the tensor to be reshaped + shape (tuple of int): the new shape + + Example:: + + >>> a = torch.arange(4.) + >>> torch.reshape(a, (2, 2)) + tensor([[ 0., 1.], + [ 2., 3.]]) + >>> b = torch.tensor([[0, 1], [2, 3]]) + >>> torch.reshape(b, (-1,)) + tensor([ 0, 1, 2, 3]) + """ + +def resize_as_( + input: Tensor, + the_template: Tensor, + *, + memory_format: memory_format | None = None, +) -> Tensor: ... +def resize_as_sparse_(input: Tensor, the_template: Tensor) -> Tensor: ... +def resolve_conj(input: Tensor) -> Tensor: + r""" + resolve_conj(input) -> Tensor + + Returns a new tensor with materialized conjugation if :attr:`input`'s conjugate bit is set to `True`, + else returns :attr:`input`. The output tensor will always have its conjugate bit set to `False`. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> y = x.conj() + >>> y.is_conj() + True + >>> z = y.resolve_conj() + >>> z + tensor([-1 - 1j, -2 - 2j, 3 + 3j]) + >>> z.is_conj() + False + """ + +def resolve_neg(input: Tensor) -> Tensor: + r""" + resolve_neg(input) -> Tensor + + Returns a new tensor with materialized negation if :attr:`input`'s negative bit is set to `True`, + else returns :attr:`input`. The output tensor will always have its negative bit set to `False`. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> y = x.conj() + >>> z = y.imag + >>> z.is_neg() + True + >>> out = z.resolve_neg() + >>> out + tensor([-1., -2., 3.]) + >>> out.is_neg() + False + """ + +@overload +def result_type(tensor: Tensor, other: Tensor) -> _dtype: + r""" + result_type(tensor1, tensor2) -> dtype + + Returns the :class:`torch.dtype` that would result from performing an arithmetic + operation on the provided input tensors. See type promotion :ref:`documentation ` + for more information on the type promotion logic. + + Args: + tensor1 (Tensor or Number): an input tensor or number + tensor2 (Tensor or Number): an input tensor or number + + Example:: + + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.int), 1.0) + torch.float32 + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.uint8), torch.tensor(1)) + torch.uint8 + """ + +@overload +def result_type(scalar: Number | _complex, tensor: Tensor) -> _dtype: + r""" + result_type(tensor1, tensor2) -> dtype + + Returns the :class:`torch.dtype` that would result from performing an arithmetic + operation on the provided input tensors. See type promotion :ref:`documentation ` + for more information on the type promotion logic. + + Args: + tensor1 (Tensor or Number): an input tensor or number + tensor2 (Tensor or Number): an input tensor or number + + Example:: + + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.int), 1.0) + torch.float32 + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.uint8), torch.tensor(1)) + torch.uint8 + """ + +@overload +def result_type(tensor: Tensor, other: Number | _complex) -> _dtype: + r""" + result_type(tensor1, tensor2) -> dtype + + Returns the :class:`torch.dtype` that would result from performing an arithmetic + operation on the provided input tensors. See type promotion :ref:`documentation ` + for more information on the type promotion logic. + + Args: + tensor1 (Tensor or Number): an input tensor or number + tensor2 (Tensor or Number): an input tensor or number + + Example:: + + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.int), 1.0) + torch.float32 + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.uint8), torch.tensor(1)) + torch.uint8 + """ + +@overload +def result_type( + scalar1: Number | _complex, + scalar2: Number | _complex, +) -> _dtype: + r""" + result_type(tensor1, tensor2) -> dtype + + Returns the :class:`torch.dtype` that would result from performing an arithmetic + operation on the provided input tensors. See type promotion :ref:`documentation ` + for more information on the type promotion logic. + + Args: + tensor1 (Tensor or Number): an input tensor or number + tensor2 (Tensor or Number): an input tensor or number + + Example:: + + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.int), 1.0) + torch.float32 + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.uint8), torch.tensor(1)) + torch.uint8 + """ + +def rms_norm( + input: Tensor, + normalized_shape: Sequence[_int | SymInt], + weight: Tensor | None = None, + eps: _float | None = None, +) -> Tensor: ... +@overload +def rnn_relu( + data: Tensor, + batch_sizes: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, +) -> tuple[Tensor, Tensor]: ... +@overload +def rnn_relu( + input: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_first: _bool, +) -> tuple[Tensor, Tensor]: ... +def rnn_relu_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor | None = None, + b_hh: Tensor | None = None, +) -> Tensor: ... +@overload +def rnn_tanh( + data: Tensor, + batch_sizes: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, +) -> tuple[Tensor, Tensor]: ... +@overload +def rnn_tanh( + input: Tensor, + hx: Tensor, + params: tuple[Tensor, ...] | list[Tensor] | None, + has_biases: _bool, + num_layers: _int, + dropout: _float, + train: _bool, + bidirectional: _bool, + batch_first: _bool, +) -> tuple[Tensor, Tensor]: ... +def rnn_tanh_cell( + input: Tensor, + hx: Tensor, + w_ih: Tensor, + w_hh: Tensor, + b_ih: Tensor | None = None, + b_hh: Tensor | None = None, +) -> Tensor: ... +def roll( + input: Tensor, + shifts: _int | SymInt | Sequence[_int | SymInt], + dims: _int | _size = (), +) -> Tensor: + r""" + roll(input, shifts, dims=None) -> Tensor + + Roll the tensor :attr:`input` along the given dimension(s). Elements that are + shifted beyond the last position are re-introduced at the first position. If + :attr:`dims` is `None`, the tensor will be flattened before rolling and then + restored to the original shape. + + Args: + input (Tensor): the input tensor. + shifts (int or tuple of ints): The number of places by which the elements + of the tensor are shifted. If shifts is a tuple, dims must be a tuple of + the same size, and each dimension will be rolled by the corresponding + value + dims (int or tuple of ints): Axis along which to roll + + Example:: + + >>> x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]).view(4, 2) + >>> x + tensor([[1, 2], + [3, 4], + [5, 6], + [7, 8]]) + >>> torch.roll(x, 1) + tensor([[8, 1], + [2, 3], + [4, 5], + [6, 7]]) + >>> torch.roll(x, 1, 0) + tensor([[7, 8], + [1, 2], + [3, 4], + [5, 6]]) + >>> torch.roll(x, -1, 0) + tensor([[3, 4], + [5, 6], + [7, 8], + [1, 2]]) + >>> torch.roll(x, shifts=(2, 1), dims=(0, 1)) + tensor([[6, 5], + [8, 7], + [2, 1], + [4, 3]]) + """ + +def rot90(input: Tensor, k: _int = 1, dims: _size = (0, 1)) -> Tensor: + r""" + rot90(input, k=1, dims=(0, 1)) -> Tensor + + Rotate an n-D tensor by 90 degrees in the plane specified by dims axis. + Rotation direction is from the first towards the second axis if k > 0, and from the second towards the first for k < 0. + + Args: + input (Tensor): the input tensor. + k (int): number of times to rotate. Default value is 1 + dims (a list or tuple): axis to rotate. Default value is [0, 1] + + Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.rot90(x, 1, [0, 1]) + tensor([[1, 3], + [0, 2]]) + + >>> x = torch.arange(8).view(2, 2, 2) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.rot90(x, 1, [1, 2]) + tensor([[[1, 3], + [0, 2]], + + [[5, 7], + [4, 6]]]) + """ + +@overload +def round(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + round(input, *, decimals=0, out=None) -> Tensor + + Rounds elements of :attr:`input` to the nearest integer. + + For integer inputs, follows the array-api convention of returning a + copy of the input tensor. + The return type of output is same as that of input's dtype. + + .. note:: + This function implements the "round half to even" to + break ties when a number is equidistant from two + integers (e.g. `round(2.5)` is 2). + + When the :attr:\`decimals\` argument is specified the + algorithm used is similar to NumPy's `around`. This + algorithm is fast but inexact and it can easily + overflow for low precision dtypes. + Eg. `round(tensor([10000], dtype=torch.float16), decimals=3)` is `inf`. + + .. seealso:: + :func:`torch.ceil`, which rounds up. + :func:`torch.floor`, which rounds down. + :func:`torch.trunc`, which rounds towards zero. + + Args: + input (Tensor): the input tensor. + decimals (int): Number of decimal places to round to (default: 0). + If decimals is negative, it specifies the number of positions + to the left of the decimal point. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.round(torch.tensor((4.7, -2.3, 9.1, -7.7))) + tensor([ 5., -2., 9., -8.]) + + >>> # Values equidistant from two integers are rounded towards the + >>> # the nearest even value (zero is treated as even) + >>> torch.round(torch.tensor([-0.5, 0.5, 1.5, 2.5])) + tensor([-0., 0., 2., 2.]) + + >>> # A positive decimals argument rounds to the to that decimal place + >>> torch.round(torch.tensor([0.1234567]), decimals=3) + tensor([0.1230]) + + >>> # A negative decimals argument rounds to the left of the decimal + >>> torch.round(torch.tensor([1200.1234567]), decimals=-3) + tensor([1000.]) + """ + +@overload +def round( + input: Tensor, + *, + decimals: _int, + out: Tensor | None = None, +) -> Tensor: + r""" + round(input, *, decimals=0, out=None) -> Tensor + + Rounds elements of :attr:`input` to the nearest integer. + + For integer inputs, follows the array-api convention of returning a + copy of the input tensor. + The return type of output is same as that of input's dtype. + + .. note:: + This function implements the "round half to even" to + break ties when a number is equidistant from two + integers (e.g. `round(2.5)` is 2). + + When the :attr:\`decimals\` argument is specified the + algorithm used is similar to NumPy's `around`. This + algorithm is fast but inexact and it can easily + overflow for low precision dtypes. + Eg. `round(tensor([10000], dtype=torch.float16), decimals=3)` is `inf`. + + .. seealso:: + :func:`torch.ceil`, which rounds up. + :func:`torch.floor`, which rounds down. + :func:`torch.trunc`, which rounds towards zero. + + Args: + input (Tensor): the input tensor. + decimals (int): Number of decimal places to round to (default: 0). + If decimals is negative, it specifies the number of positions + to the left of the decimal point. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> torch.round(torch.tensor((4.7, -2.3, 9.1, -7.7))) + tensor([ 5., -2., 9., -8.]) + + >>> # Values equidistant from two integers are rounded towards the + >>> # the nearest even value (zero is treated as even) + >>> torch.round(torch.tensor([-0.5, 0.5, 1.5, 2.5])) + tensor([-0., 0., 2., 2.]) + + >>> # A positive decimals argument rounds to the to that decimal place + >>> torch.round(torch.tensor([0.1234567]), decimals=3) + tensor([0.1230]) + + >>> # A negative decimals argument rounds to the left of the decimal + >>> torch.round(torch.tensor([1200.1234567]), decimals=-3) + tensor([1000.]) + """ + +@overload +def round_(input: Tensor) -> Tensor: ... +@overload +def round_(input: Tensor, *, decimals: _int) -> Tensor: ... +def row_indices_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: ... +def row_stack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + row_stack(tensors, *, out=None) -> Tensor + + Alias of :func:`torch.vstack`. + """ + +def rrelu( + input: Tensor, + lower: Number | _complex = 0.125, + upper: Number | _complex = 0.3333333333333333, + training: _bool = False, + generator: Generator | None = None, +) -> Tensor: ... +def rrelu_( + input: Tensor, + lower: Number | _complex = 0.125, + upper: Number | _complex = 0.3333333333333333, + training: _bool = False, + generator: Generator | None = None, +) -> Tensor: ... +def rsqrt(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + rsqrt(input, *, out=None) -> Tensor + + Returns a new tensor with the reciprocal of the square-root of each of + the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \frac{1}{\sqrt{\text{input}_{i}}} + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.0370, 0.2970, 1.5420, -0.9105]) + >>> torch.rsqrt(a) + tensor([ nan, 1.8351, 0.8053, nan]) + """ + +def rsqrt_(input: Tensor) -> Tensor: ... +@overload +def rsub( + input: Tensor, + other: Tensor, + *, + alpha: Number | _complex = 1, +) -> Tensor: ... +@overload +def rsub( + input: Tensor, + other: Number | _complex, + alpha: Number | _complex = 1, +) -> Tensor: ... +def saddmm( + input: Tensor, + mat1: Tensor, + mat2: Tensor, + *, + beta: Number = 1, + alpha: Number = 1, + out: Tensor | None = None, +) -> Tensor: ... +def scalar_tensor( + s: Number | _complex, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: ... +@overload +def scatter( + input: Tensor, + dim: _int, + index: Tensor, + src: Tensor, + *, + reduce: str, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter( + input: Tensor, + dim: _int, + index: Tensor, + src: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter( + input: Tensor, + dim: _int, + index: Tensor, + value: Number | _complex, + *, + reduce: str, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + src: Tensor, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter( + input: Tensor, + dim: _int, + index: Tensor, + value: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + value: Number | _complex, +) -> Tensor: + r""" + scatter(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_` + """ + +@overload +def scatter_add( + input: Tensor, + dim: _int, + index: Tensor, + src: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter_add(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_add_` + """ + +@overload +def scatter_add( + input: Tensor, + dim: str | EllipsisType | None, + index: Tensor, + src: Tensor, +) -> Tensor: + r""" + scatter_add(input, dim, index, src) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_add_` + """ + +def scatter_reduce( + input: Tensor, + dim: _int, + index: Tensor, + src: Tensor, + reduce: str, + *, + include_self: _bool = True, + out: Tensor | None = None, +) -> Tensor: + r""" + scatter_reduce(input, dim, index, src, reduce, *, include_self=True) -> Tensor + + Out-of-place version of :meth:`torch.Tensor.scatter_reduce_` + """ + +@overload +def searchsorted( + sorted_sequence: Tensor, + input: Tensor, + *, + out_int32: _bool = False, + right: _bool = False, + side: str | None = None, + sorter: Tensor | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + searchsorted(sorted_sequence, values, *, out_int32=False, right=False, side=None, out=None, sorter=None) -> Tensor + + Find the indices from the *innermost* dimension of :attr:`sorted_sequence` such that, if the + corresponding values in :attr:`values` were inserted before the indices, when sorted, the order + of the corresponding *innermost* dimension within :attr:`sorted_sequence` would be preserved. + Return a new tensor with the same size as :attr:`values`. More formally, + the returned index satisfies the following rules: + + .. list-table:: + :widths: 12 10 78 + :header-rows: 1 + + * - :attr:`sorted_sequence` + - :attr:`right` + - *returned index satisfies* + * - 1-D + - False + - ``sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]`` + * - 1-D + - True + - ``sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]`` + * - N-D + - False + - ``sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]`` + * - N-D + - True + - ``sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]`` + + Args: + sorted_sequence (Tensor): N-D or 1-D tensor, containing monotonically increasing sequence on the *innermost* + dimension unless :attr:`sorter` is provided, in which case the sequence does not + need to be sorted + values (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + + Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): if False, return the first suitable location that is found. If True, return the + last such index. If no suitable index found, return 0 for non-numerical value + (eg. nan, inf) or the size of *innermost* dimension within :attr:`sorted_sequence` + (one pass the last index of the *innermost* dimension). In other words, if False, + gets the lower bound index for each value in :attr:`values` on the corresponding + *innermost* dimension of the :attr:`sorted_sequence`. If True, gets the upper + bound index instead. Default value is False. :attr:`side` does the same and is + preferred. It will error if :attr:`side` is set to "left" while this is True. + side (str, optional): the same as :attr:`right` but preferred. "left" corresponds to False for :attr:`right` + and "right" corresponds to True for :attr:`right`. It will error if this is set to + "left" while :attr:`right` is True. Default value is None. + out (Tensor, optional): the output tensor, must be the same size as :attr:`values` if provided. + sorter (LongTensor, optional): if provided, a tensor matching the shape of the unsorted + :attr:`sorted_sequence` containing a sequence of indices that sort it in the + ascending order on the innermost dimension + + + Example:: + + >>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]) + >>> sorted_sequence + tensor([[ 1, 3, 5, 7, 9], + [ 2, 4, 6, 8, 10]]) + >>> values = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> values + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.searchsorted(sorted_sequence, values) + tensor([[1, 3, 4], + [1, 2, 4]]) + >>> torch.searchsorted(sorted_sequence, values, side='right') + tensor([[2, 3, 5], + [1, 3, 4]]) + + >>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9]) + >>> sorted_sequence_1d + tensor([1, 3, 5, 7, 9]) + >>> torch.searchsorted(sorted_sequence_1d, values) + tensor([[1, 3, 4], + [1, 3, 4]]) + """ + +@overload +def searchsorted( + sorted_sequence: Tensor, + self: Number | _complex, + *, + out_int32: _bool = False, + right: _bool = False, + side: str | None = None, + sorter: Tensor | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + searchsorted(sorted_sequence, values, *, out_int32=False, right=False, side=None, out=None, sorter=None) -> Tensor + + Find the indices from the *innermost* dimension of :attr:`sorted_sequence` such that, if the + corresponding values in :attr:`values` were inserted before the indices, when sorted, the order + of the corresponding *innermost* dimension within :attr:`sorted_sequence` would be preserved. + Return a new tensor with the same size as :attr:`values`. More formally, + the returned index satisfies the following rules: + + .. list-table:: + :widths: 12 10 78 + :header-rows: 1 + + * - :attr:`sorted_sequence` + - :attr:`right` + - *returned index satisfies* + * - 1-D + - False + - ``sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]`` + * - 1-D + - True + - ``sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]`` + * - N-D + - False + - ``sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]`` + * - N-D + - True + - ``sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]`` + + Args: + sorted_sequence (Tensor): N-D or 1-D tensor, containing monotonically increasing sequence on the *innermost* + dimension unless :attr:`sorter` is provided, in which case the sequence does not + need to be sorted + values (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + + Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): if False, return the first suitable location that is found. If True, return the + last such index. If no suitable index found, return 0 for non-numerical value + (eg. nan, inf) or the size of *innermost* dimension within :attr:`sorted_sequence` + (one pass the last index of the *innermost* dimension). In other words, if False, + gets the lower bound index for each value in :attr:`values` on the corresponding + *innermost* dimension of the :attr:`sorted_sequence`. If True, gets the upper + bound index instead. Default value is False. :attr:`side` does the same and is + preferred. It will error if :attr:`side` is set to "left" while this is True. + side (str, optional): the same as :attr:`right` but preferred. "left" corresponds to False for :attr:`right` + and "right" corresponds to True for :attr:`right`. It will error if this is set to + "left" while :attr:`right` is True. Default value is None. + out (Tensor, optional): the output tensor, must be the same size as :attr:`values` if provided. + sorter (LongTensor, optional): if provided, a tensor matching the shape of the unsorted + :attr:`sorted_sequence` containing a sequence of indices that sort it in the + ascending order on the innermost dimension + + + Example:: + + >>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]) + >>> sorted_sequence + tensor([[ 1, 3, 5, 7, 9], + [ 2, 4, 6, 8, 10]]) + >>> values = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> values + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.searchsorted(sorted_sequence, values) + tensor([[1, 3, 4], + [1, 2, 4]]) + >>> torch.searchsorted(sorted_sequence, values, side='right') + tensor([[2, 3, 5], + [1, 3, 4]]) + + >>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9]) + >>> sorted_sequence_1d + tensor([1, 3, 5, 7, 9]) + >>> torch.searchsorted(sorted_sequence_1d, values) + tensor([[1, 3, 4], + [1, 3, 4]]) + """ + +def segment_reduce( + data: Tensor, + reduce: str, + *, + lengths: Tensor | None = None, + indices: Tensor | None = None, + offsets: Tensor | None = None, + axis: _int = 0, + unsafe: _bool = False, + initial: Number | _complex | None = None, +) -> Tensor: + r""" + segment_reduce(data: Tensor, reduce: str, *, lengths: Tensor | None = None, indices: Tensor | None = None, offsets: Tensor | None = None, axis: _int = 0, unsafe: _bool = False, initial: Number | _complex | None = None) -> Tensor # noqa: B950 + + Perform a segment reduction operation on the input tensor along the specified axis. + + Args: + data (Tensor): The input tensor on which the segment reduction operation will be performed. + reduce (str): The type of reduction operation. Supported values are ``sum``, ``mean``, ``max``, ``min``, ``prod``. + + Keyword args: + lengths (Tensor, optional): Length of each segment. Default: ``None``. + offsets (Tensor, optional): Offset of each segment. Default: ``None``. + axis (int, optional): The axis perform reduction. Default: ``0``. + unsafe (bool, optional): Skip validation If `True`. Default: ``False``. + initial (Number, optional): The initial value for the reduction operation. Default: ``None``. + + Example:: + + >>> data = torch.tensor([[1, 2, 3, 4],[5, 6, 7, 8],[9, 10, 11, 12]], dtype=torch.float32, device='cuda') + >>> lengths = torch.tensor([2, 1], device='cuda') + >>> torch.segment_reduce(data, 'max', lengths=lengths) + tensor([[ 5., 6., 7., 8.], + [ 9., 10., 11., 12.]], device='cuda:0') + """ + +@overload +def select(input: Tensor, dim: _int, index: _int | SymInt) -> Tensor: + r""" + select(input, dim, index) -> Tensor + + Slices the :attr:`input` tensor along the selected dimension at the given index. + This function returns a view of the original tensor with the given dimension removed. + + .. note:: If :attr:`input` is a sparse tensor and returning a view of + the tensor is not possible, a RuntimeError exception is + raised. In this is the case, consider using + :func:`torch.select_copy` function. + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to slice + index (int): the index to select with + + .. note:: + + :meth:`select` is equivalent to slicing. For example, + ``tensor.select(0, index)`` is equivalent to ``tensor[index]`` and + ``tensor.select(2, index)`` is equivalent to ``tensor[:,:,index]``. + """ + +@overload +def select( + input: Tensor, + dim: str | EllipsisType | None, + index: _int, +) -> Tensor: + r""" + select(input, dim, index) -> Tensor + + Slices the :attr:`input` tensor along the selected dimension at the given index. + This function returns a view of the original tensor with the given dimension removed. + + .. note:: If :attr:`input` is a sparse tensor and returning a view of + the tensor is not possible, a RuntimeError exception is + raised. In this is the case, consider using + :func:`torch.select_copy` function. + + Args: + input (Tensor): the input tensor. + dim (int): the dimension to slice + index (int): the index to select with + + .. note:: + + :meth:`select` is equivalent to slicing. For example, + ``tensor.select(0, index)`` is equivalent to ``tensor[index]`` and + ``tensor.select(2, index)`` is equivalent to ``tensor[:,:,index]``. + """ + +def select_copy( + input: Tensor, + dim: _int, + index: _int | SymInt, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.select`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def select_scatter( + input: Tensor, + src: Tensor, + dim: _int, + index: _int | SymInt, +) -> Tensor: + r""" + select_scatter(input, src, dim, index) -> Tensor + + Embeds the values of the :attr:`src` tensor into :attr:`input` at the given index. + This function returns a tensor with fresh storage; it does not create a view. + + + Args: + input (Tensor): the input tensor. + src (Tensor): The tensor to embed into :attr:`input` + dim (int): the dimension to insert the slice into. + index (int): the index to select with + + .. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + ``torch.select(input, dim, index)`` + + Example:: + + >>> a = torch.zeros(2, 2) + >>> b = torch.ones(2) + >>> a.select_scatter(b, 0, 0) + tensor([[1., 1.], + [0., 0.]]) + """ + +def selu(input: Tensor) -> Tensor: ... +def selu_(input: Tensor) -> Tensor: ... +def set_flush_denormal(mode: _bool) -> _bool: + r""" + set_flush_denormal(mode) -> bool + + Disables denormal floating numbers on CPU. + + Returns ``True`` if your system supports flushing denormal numbers and it + successfully configures flush denormal mode. :meth:`~torch.set_flush_denormal` + is supported on x86 architectures supporting SSE3 and AArch64 architecture. + + Args: + mode (bool): Controls whether to enable flush denormal mode or not + + Example:: + + >>> torch.set_flush_denormal(True) + True + >>> torch.tensor([1e-323], dtype=torch.float64) + tensor([ 0.], dtype=torch.float64) + >>> torch.set_flush_denormal(False) + True + >>> torch.tensor([1e-323], dtype=torch.float64) + tensor(9.88131e-324 * + [ 1.0000], dtype=torch.float64) + """ + +def set_num_interop_threads(num: _int) -> None: + r""" + set_num_interop_threads(int) + + Sets the number of threads used for interop parallelism + (e.g. in JIT interpreter) on CPU. + + .. warning:: + Can only be called once and before any inter-op parallel work + is started (e.g. JIT execution). + """ + +def set_num_threads(num: _int) -> None: + r""" + set_num_threads(int) + + Sets the number of threads used for intraop parallelism on CPU. + + .. warning:: + To ensure that the correct number of threads is used, set_num_threads + must be called before running eager, JIT or autograd code. + """ + +def sgn(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sgn(input, *, out=None) -> Tensor + + This function is an extension of torch.sign() to complex tensors. + It computes a new tensor whose elements have + the same angles as the corresponding elements of :attr:`input` and + absolute values (i.e. magnitudes) of one for complex tensors and + is equivalent to torch.sign() for non-complex tensors. + + .. math:: + \text{out}_{i} = \begin{cases} + 0 & |\text{{input}}_i| == 0 \\ + \frac{{\text{{input}}_i}}{|{\text{{input}}_i}|} & \text{otherwise} + \end{cases} + + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.tensor([3+4j, 7-24j, 0, 1+2j]) + >>> t.sgn() + tensor([0.6000+0.8000j, 0.2800-0.9600j, 0.0000+0.0000j, 0.4472+0.8944j]) + """ + +def sigmoid(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sigmoid(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.expit`. + """ + +def sigmoid_(input: Tensor) -> Tensor: ... +def sign(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sign(input, *, out=None) -> Tensor + + Returns a new tensor with the signs of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \operatorname{sgn}(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([0.7, -1.2, 0., 2.3]) + >>> a + tensor([ 0.7000, -1.2000, 0.0000, 2.3000]) + >>> torch.sign(a) + tensor([ 1., -1., 0., 1.]) + """ + +def signbit(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + signbit(input, *, out=None) -> Tensor + + Tests if each element of :attr:`input` has its sign bit set or not. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([0.7, -1.2, 0., 2.3]) + >>> torch.signbit(a) + tensor([ False, True, False, False]) + >>> a = torch.tensor([-0.0, 0.0]) + >>> torch.signbit(a) + tensor([ True, False]) + + .. note:: + signbit handles signed zeros, so negative zero (-0) returns True. + """ + +def sin(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sin(input, *, out=None) -> Tensor + + Returns a new tensor with the sine of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \sin(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.5461, 0.1347, -2.7266, -0.2746]) + >>> torch.sin(a) + tensor([-0.5194, 0.1343, -0.4032, -0.2711]) + """ + +def sin_(input: Tensor) -> Tensor: ... +def sinc(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sinc(input, *, out=None) -> Tensor + + Alias for :func:`torch.special.sinc`. + """ + +def sinc_(input: Tensor) -> Tensor: ... +def sinh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sinh(input, *, out=None) -> Tensor + + Returns a new tensor with the hyperbolic sine of the elements of + :attr:`input`. + + .. math:: + \text{out}_{i} = \sinh(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.5380, -0.8632, -0.1265, 0.9399]) + >>> torch.sinh(a) + tensor([ 0.5644, -0.9744, -0.1268, 1.0845]) + + .. note:: + When :attr:`input` is on the CPU, the implementation of torch.sinh may use + the Sleef library, which rounds very large results to infinity or negative + infinity. See `here `_ for details. + """ + +def sinh_(input: Tensor) -> Tensor: ... +def slice_copy( + input: Tensor, + dim: _int = 0, + start: _int | SymInt | None = None, + end: _int | SymInt | None = None, + step: _int | SymInt = 1, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.slice`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def slice_inverse( + input: Tensor, + src: Tensor, + dim: _int = 0, + start: _int | SymInt | None = None, + end: _int | SymInt | None = None, + step: _int | SymInt = 1, +) -> Tensor: ... +def slice_scatter( + input: Tensor, + src: Tensor, + dim: _int = 0, + start: _int | SymInt | None = None, + end: _int | SymInt | None = None, + step: _int | SymInt = 1, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + slice_scatter(input, src, dim=0, start=None, end=None, step=1) -> Tensor + + Embeds the values of the :attr:`src` tensor into :attr:`input` at the given + dimension. + This function returns a tensor with fresh storage; it does not create a view. + + + Args: + input (Tensor): the input tensor. + src (Tensor): The tensor to embed into :attr:`input` + dim (int): the dimension to insert the slice into + start (Optional[int]): the start index of where to insert the slice + end (Optional[int]): the end index of where to insert the slice + step (int): the how many elements to skip in + + Example:: + + >>> a = torch.zeros(8, 8) + >>> b = torch.ones(2, 8) + >>> a.slice_scatter(b, start=6) + tensor([[0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1., 1., 1.]]) + + >>> b = torch.ones(8, 2) + >>> a.slice_scatter(b, dim=1, start=2, end=6, step=2) + tensor([[0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.]]) + """ + +def slogdet( + input: Tensor, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.slogdet: + r""" + slogdet(input) -> (Tensor, Tensor) + + Alias for :func:`torch.linalg.slogdet` + """ + +def smm(input: Tensor, mat2: Tensor) -> Tensor: + r""" + smm(input, mat) -> Tensor + + Performs a matrix multiplication of the sparse matrix :attr:`input` + with the dense matrix :attr:`mat`. + + Args: + input (Tensor): a sparse matrix to be matrix multiplied + mat (Tensor): a dense matrix to be matrix multiplied + """ + +@overload +def softmax( + input: Tensor, + dim: _int, + dtype: _dtype | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + softmax(input, dim, *, dtype=None) -> Tensor + + Alias for :func:`torch.nn.functional.softmax`. + """ + +@overload +def softmax( + input: Tensor, + dim: str | EllipsisType | None, + *, + dtype: _dtype | None = None, +) -> Tensor: + r""" + softmax(input, dim, *, dtype=None) -> Tensor + + Alias for :func:`torch.nn.functional.softmax`. + """ + +@overload +def sort( + input: Tensor, + *, + stable: _bool | None, + dim: _int = -1, + descending: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.sort: + r""" + sort(input, dim=-1, descending=False, *, stable=False, out=None) -> (Tensor, LongTensor) + + Sorts the elements of the :attr:`input` tensor along a given dimension + in ascending order by value. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`descending` is ``True`` then the elements are sorted in descending + order by value. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. + + A namedtuple of (values, indices) is returned, where the `values` are the + sorted values and `indices` are the indices of the elements in the original + `input` tensor. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): makes the sorting routine stable, which guarantees that the order + of equivalent elements is preserved. + out (tuple, optional): the output tuple of (`Tensor`, `LongTensor`) that can + be optionally given to be used as output buffers + + Example:: + + >>> x = torch.randn(3, 4) + >>> sorted, indices = torch.sort(x) + >>> sorted + tensor([[-0.2162, 0.0608, 0.6719, 2.3332], + [-0.5793, 0.0061, 0.6058, 0.9497], + [-0.5071, 0.3343, 0.9553, 1.0960]]) + >>> indices + tensor([[ 1, 0, 2, 3], + [ 3, 1, 0, 2], + [ 0, 3, 1, 2]]) + + >>> sorted, indices = torch.sort(x, 0) + >>> sorted + tensor([[-0.5071, -0.2162, 0.6719, -0.5793], + [ 0.0608, 0.0061, 0.9497, 0.3343], + [ 0.6058, 0.9553, 1.0960, 2.3332]]) + >>> indices + tensor([[ 2, 0, 0, 1], + [ 0, 1, 1, 2], + [ 1, 2, 2, 0]]) + >>> x = torch.tensor([0, 1] * 9) + >>> x.sort() + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 2, 16, 4, 6, 14, 8, 0, 10, 12, 9, 17, 15, 13, 11, 7, 5, 3, 1])) + >>> x.sort(stable=True) + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 0, 2, 4, 6, 8, 10, 12, 14, 16, 1, 3, 5, 7, 9, 11, 13, 15, 17])) + """ + +@overload +def sort( + input: Tensor, + dim: _int = -1, + descending: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.sort: + r""" + sort(input, dim=-1, descending=False, *, stable=False, out=None) -> (Tensor, LongTensor) + + Sorts the elements of the :attr:`input` tensor along a given dimension + in ascending order by value. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`descending` is ``True`` then the elements are sorted in descending + order by value. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. + + A namedtuple of (values, indices) is returned, where the `values` are the + sorted values and `indices` are the indices of the elements in the original + `input` tensor. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): makes the sorting routine stable, which guarantees that the order + of equivalent elements is preserved. + out (tuple, optional): the output tuple of (`Tensor`, `LongTensor`) that can + be optionally given to be used as output buffers + + Example:: + + >>> x = torch.randn(3, 4) + >>> sorted, indices = torch.sort(x) + >>> sorted + tensor([[-0.2162, 0.0608, 0.6719, 2.3332], + [-0.5793, 0.0061, 0.6058, 0.9497], + [-0.5071, 0.3343, 0.9553, 1.0960]]) + >>> indices + tensor([[ 1, 0, 2, 3], + [ 3, 1, 0, 2], + [ 0, 3, 1, 2]]) + + >>> sorted, indices = torch.sort(x, 0) + >>> sorted + tensor([[-0.5071, -0.2162, 0.6719, -0.5793], + [ 0.0608, 0.0061, 0.9497, 0.3343], + [ 0.6058, 0.9553, 1.0960, 2.3332]]) + >>> indices + tensor([[ 2, 0, 0, 1], + [ 0, 1, 1, 2], + [ 1, 2, 2, 0]]) + >>> x = torch.tensor([0, 1] * 9) + >>> x.sort() + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 2, 16, 4, 6, 14, 8, 0, 10, 12, 9, 17, 15, 13, 11, 7, 5, 3, 1])) + >>> x.sort(stable=True) + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 0, 2, 4, 6, 8, 10, 12, 14, 16, 1, 3, 5, 7, 9, 11, 13, 15, 17])) + """ + +@overload +def sort( + input: Tensor, + *, + stable: _bool | None, + dim: str | EllipsisType | None, + descending: _bool = False, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.sort: + r""" + sort(input, dim=-1, descending=False, *, stable=False, out=None) -> (Tensor, LongTensor) + + Sorts the elements of the :attr:`input` tensor along a given dimension + in ascending order by value. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`descending` is ``True`` then the elements are sorted in descending + order by value. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. + + A namedtuple of (values, indices) is returned, where the `values` are the + sorted values and `indices` are the indices of the elements in the original + `input` tensor. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): makes the sorting routine stable, which guarantees that the order + of equivalent elements is preserved. + out (tuple, optional): the output tuple of (`Tensor`, `LongTensor`) that can + be optionally given to be used as output buffers + + Example:: + + >>> x = torch.randn(3, 4) + >>> sorted, indices = torch.sort(x) + >>> sorted + tensor([[-0.2162, 0.0608, 0.6719, 2.3332], + [-0.5793, 0.0061, 0.6058, 0.9497], + [-0.5071, 0.3343, 0.9553, 1.0960]]) + >>> indices + tensor([[ 1, 0, 2, 3], + [ 3, 1, 0, 2], + [ 0, 3, 1, 2]]) + + >>> sorted, indices = torch.sort(x, 0) + >>> sorted + tensor([[-0.5071, -0.2162, 0.6719, -0.5793], + [ 0.0608, 0.0061, 0.9497, 0.3343], + [ 0.6058, 0.9553, 1.0960, 2.3332]]) + >>> indices + tensor([[ 2, 0, 0, 1], + [ 0, 1, 1, 2], + [ 1, 2, 2, 0]]) + >>> x = torch.tensor([0, 1] * 9) + >>> x.sort() + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 2, 16, 4, 6, 14, 8, 0, 10, 12, 9, 17, 15, 13, 11, 7, 5, 3, 1])) + >>> x.sort(stable=True) + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 0, 2, 4, 6, 8, 10, 12, 14, 16, 1, 3, 5, 7, 9, 11, 13, 15, 17])) + """ + +@overload +def sort( + input: Tensor, + dim: str | EllipsisType | None, + descending: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.sort: + r""" + sort(input, dim=-1, descending=False, *, stable=False, out=None) -> (Tensor, LongTensor) + + Sorts the elements of the :attr:`input` tensor along a given dimension + in ascending order by value. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`descending` is ``True`` then the elements are sorted in descending + order by value. + + If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving + the order of equivalent elements. + + A namedtuple of (values, indices) is returned, where the `values` are the + sorted values and `indices` are the indices of the elements in the original + `input` tensor. + + Args: + input (Tensor): the input tensor. + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + + Keyword args: + stable (bool, optional): makes the sorting routine stable, which guarantees that the order + of equivalent elements is preserved. + out (tuple, optional): the output tuple of (`Tensor`, `LongTensor`) that can + be optionally given to be used as output buffers + + Example:: + + >>> x = torch.randn(3, 4) + >>> sorted, indices = torch.sort(x) + >>> sorted + tensor([[-0.2162, 0.0608, 0.6719, 2.3332], + [-0.5793, 0.0061, 0.6058, 0.9497], + [-0.5071, 0.3343, 0.9553, 1.0960]]) + >>> indices + tensor([[ 1, 0, 2, 3], + [ 3, 1, 0, 2], + [ 0, 3, 1, 2]]) + + >>> sorted, indices = torch.sort(x, 0) + >>> sorted + tensor([[-0.5071, -0.2162, 0.6719, -0.5793], + [ 0.0608, 0.0061, 0.9497, 0.3343], + [ 0.6058, 0.9553, 1.0960, 2.3332]]) + >>> indices + tensor([[ 2, 0, 0, 1], + [ 0, 1, 1, 2], + [ 1, 2, 2, 0]]) + >>> x = torch.tensor([0, 1] * 9) + >>> x.sort() + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 2, 16, 4, 6, 14, 8, 0, 10, 12, 9, 17, 15, 13, 11, 7, 5, 3, 1])) + >>> x.sort(stable=True) + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 0, 2, 4, 6, 8, 10, 12, 14, 16, 1, 3, 5, 7, 9, 11, 13, 15, 17])) + """ + +def sparse_bsc_tensor( + ccol_indices: Tensor | list, + row_indices: Tensor | list, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, +) -> Tensor: + r""" + sparse_bsc_tensor(ccol_indices, row_indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + + Constructs a :ref:`sparse tensor in BSC (Block Compressed Sparse + Column)) ` with specified 2-dimensional blocks at the + given :attr:`ccol_indices` and :attr:`row_indices`. Sparse matrix + multiplication operations in BSC format are typically faster than that + for sparse tensors in COO format. Make you have a look at :ref:`the + note on the data type of the indices `. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + ccol_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, ncolblocks + 1)``. The last element of each + batch is the number of non-zeros. This tensor encodes the + index in values and row_indices depending on where the given + column starts. Each successive number in the tensor subtracted + by the number before it denotes the number of elements in a + given column. + row_indices (array_like): Row block coordinates of each block in + values. (B+1)-dimensional tensor with the same length + as values. + values (array_list): Initial blocks for the tensor. Can be a list, + tuple, NumPy ``ndarray``, and other types that + represents a (1 + 2 + K)-dimensional tensor where ``K`` is the + number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` If not provided, the size will be + inferred as the minimum size big enough to hold all non-zero + blocks. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + + Example:: + + >>> ccol_indices = [0, 1, 2] + >>> row_indices = [0, 1] + >>> values = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> torch.sparse_bsc_tensor(torch.tensor(ccol_indices, dtype=torch.int64), + ... torch.tensor(row_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(ccol_indices=tensor([0, 1, 2]), + row_indices=tensor([0, 1]), + values=tensor([[[1., 2.], + [3., 4.]], + [[5., 6.], + [7., 8.]]]), size=(2, 2), nnz=2, dtype=torch.float64, + layout=torch.sparse_bsc) + """ + +def sparse_bsr_tensor( + crow_indices: Tensor | list, + col_indices: Tensor | list, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, +) -> Tensor: + r""" + sparse_bsr_tensor(crow_indices, col_indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + + Constructs a :ref:`sparse tensor in BSR (Block Compressed Sparse Row)) + ` with specified 2-dimensional blocks at the given + :attr:`crow_indices` and :attr:`col_indices`. Sparse matrix + multiplication operations in BSR format are typically faster than that + for sparse tensors in COO format. Make you have a look at :ref:`the + note on the data type of the indices `. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + crow_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, nrowblocks + 1)``. The last element of each + batch is the number of non-zeros. This tensor encodes the + block index in values and col_indices depending on where the + given row block starts. Each successive number in the tensor + subtracted by the number before it denotes the number of + blocks in a given row. + col_indices (array_like): Column block coordinates of each block + in values. (B+1)-dimensional tensor with the same length as + values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1 + 2 + K)-dimensional tensor where ``K`` is the + number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` where ``blocksize == + values.shape[1:3]``. If not provided, the size will be + inferred as the minimum size big enough to hold all non-zero + blocks. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + + Example:: + + >>> crow_indices = [0, 1, 2] + >>> col_indices = [0, 1] + >>> values = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> torch.sparse_bsr_tensor(torch.tensor(crow_indices, dtype=torch.int64), + ... torch.tensor(col_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(crow_indices=tensor([0, 1, 2]), + col_indices=tensor([0, 1]), + values=tensor([[[1., 2.], + [3., 4.]], + [[5., 6.], + [7., 8.]]]), size=(2, 2), nnz=2, dtype=torch.float64, + layout=torch.sparse_bsr) + """ + +def sparse_compressed_tensor( + compressed_indices: Tensor | list, + plain_indices: Tensor | list, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, +) -> Tensor: + r""" + sparse_compressed_tensor(compressed_indices, plain_indices, values, size=None, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + + Constructs a :ref:`sparse tensor in Compressed Sparse format - CSR, + CSC, BSR, or BSC - ` with specified values at + the given :attr:`compressed_indices` and :attr:`plain_indices`. Sparse + matrix multiplication operations in Compressed Sparse format are + typically faster than that for sparse tensors in COO format. Make you + have a look at :ref:`the note on the data type of the indices + `. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + compressed_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, compressed_dim_size + 1)``. The last element of + each batch is the number of non-zero elements or blocks. This + tensor encodes the index in ``values`` and ``plain_indices`` + depending on where the given compressed dimension (row or + column) starts. Each successive number in the tensor + subtracted by the number before it denotes the number of + elements or blocks in a given compressed dimension. + plain_indices (array_like): Plain dimension (column or row) + coordinates of each element or block in values. (B+1)-dimensional + tensor with the same length as values. + + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types. that + represents a (1+K)-dimensional (for CSR and CSC layouts) or + (1+2+K)-dimensional tensor (for BSR and BSC layouts) where + ``K`` is the number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` where ``blocksize[0] == + blocksize[1] == 1`` for CSR and CSC formats. If not provided, + the size will be inferred as the minimum size big enough to + hold all non-zero elements or blocks. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + layout (:class:`torch.layout`, required): the desired layout of + returned tensor: :attr:`torch.sparse_csr`, + :attr:`torch.sparse_csc`, :attr:`torch.sparse_bsr`, or + :attr:`torch.sparse_bsc`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + + Example:: + + >>> compressed_indices = [0, 2, 4] + >>> plain_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_compressed_tensor(torch.tensor(compressed_indices, dtype=torch.int64), + ... torch.tensor(plain_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double, layout=torch.sparse_csr) + tensor(crow_indices=tensor([0, 2, 4]), + col_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csr) + """ + +def sparse_coo_tensor( + indices: Tensor, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, + is_coalesced: _bool | None = None, +) -> Tensor: + r""" + sparse_coo_tensor(indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None, is_coalesced=None) -> Tensor + + Constructs a :ref:`sparse tensor in COO(rdinate) format + ` with specified values at the given + :attr:`indices`. + + .. note:: + + This function returns an :ref:`uncoalesced tensor + ` when :attr:`is_coalesced` is + unspecified or ``None``. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + indices (array_like): Initial data for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. Will be cast to a :class:`torch.LongTensor` + internally. The indices are the coordinates of the non-zero values in the matrix, and thus + should be two-dimensional where the first dimension is the number of tensor dimensions and + the second dimension is the number of non-zero values. + values (array_like): Initial values for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. + size (list, tuple, or :class:`torch.Size`, optional): Size of the sparse tensor. If not + provided the size will be inferred as the minimum size big enough to hold all non-zero + elements. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if None, infers data type from :attr:`values`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if None, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + is_coalesced (bool, optional): When``True``, the caller is + responsible for providing tensor indices that correspond to a + coalesced tensor. If the :attr:`check_invariants` flag is + False, no error will be raised if the prerequisites are not + met and this will lead to silently incorrect results. To force + coalescion please use :meth:`coalesce` on the resulting + Tensor. + Default: None: except for trivial cases (e.g. nnz < 2) the + resulting Tensor has is_coalesced set to ``False```. + + Example:: + + >>> i = torch.tensor([[0, 1, 1], + ... [2, 0, 2]]) + >>> v = torch.tensor([3, 4, 5], dtype=torch.float32) + >>> torch.sparse_coo_tensor(i, v, [2, 4]) + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + size=(2, 4), nnz=3, layout=torch.sparse_coo) + + >>> torch.sparse_coo_tensor(i, v) # Shape inference + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + size=(2, 3), nnz=3, layout=torch.sparse_coo) + + >>> torch.sparse_coo_tensor(i, v, [2, 4], + ... dtype=torch.float64, + ... device=torch.device('cuda:0')) + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + device='cuda:0', size=(2, 4), nnz=3, dtype=torch.float64, + layout=torch.sparse_coo) + + # Create an empty sparse tensor with the following invariants: + # 1. sparse_dim + dense_dim = len(SparseTensor.shape) + # 2. SparseTensor._indices().shape = (sparse_dim, nnz) + # 3. SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:]) + # + # For instance, to create an empty sparse tensor with nnz = 0, dense_dim = 0 and + # sparse_dim = 1 (hence indices is a 2D tensor of shape = (1, 0)) + >>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), [], [1]) + tensor(indices=tensor([], size=(1, 0)), + values=tensor([], size=(0,)), + size=(1,), nnz=0, layout=torch.sparse_coo) + + # and to create an empty sparse tensor with nnz = 0, dense_dim = 1 and + # sparse_dim = 1 + >>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), torch.empty([0, 2]), [1, 2]) + tensor(indices=tensor([], size=(1, 0)), + values=tensor([], size=(0, 2)), + size=(1, 2), nnz=0, layout=torch.sparse_coo) + + .. _torch.sparse: https://pytorch.org/docs/stable/sparse.html + """ + +def sparse_csc_tensor( + ccol_indices: Tensor | list, + row_indices: Tensor | list, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, +) -> Tensor: + r""" + sparse_csc_tensor(ccol_indices, row_indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + + Constructs a :ref:`sparse tensor in CSC (Compressed Sparse Column) + ` with specified values at the given + :attr:`ccol_indices` and :attr:`row_indices`. Sparse matrix + multiplication operations in CSC format are typically faster than that + for sparse tensors in COO format. Make you have a look at :ref:`the + note on the data type of the indices `. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + ccol_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, ncols + 1)``. The last element of each batch + is the number of non-zeros. This tensor encodes the index in + values and row_indices depending on where the given column + starts. Each successive number in the tensor subtracted by the + number before it denotes the number of elements in a given + column. + row_indices (array_like): Row coordinates of each element in + values. (B+1)-dimensional tensor with the same length as + values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1+K)-dimensional tensor where ``K`` is the number + of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows, ncols, *densesize)``. If + not provided, the size will be inferred as the minimum size + big enough to hold all non-zero elements. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + + Example:: + + >>> ccol_indices = [0, 2, 4] + >>> row_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_csc_tensor(torch.tensor(ccol_indices, dtype=torch.int64), + ... torch.tensor(row_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(ccol_indices=tensor([0, 2, 4]), + row_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csc) + """ + +def sparse_csr_tensor( + crow_indices: Tensor | list, + col_indices: Tensor | list, + values: Tensor | list, + size: _size | None = None, + *, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + check_invariants: _bool | None = None, +) -> Tensor: + r""" + sparse_csr_tensor(crow_indices, col_indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + + Constructs a :ref:`sparse tensor in CSR (Compressed Sparse Row) ` with specified + values at the given :attr:`crow_indices` and :attr:`col_indices`. Sparse matrix multiplication operations + in CSR format are typically faster than that for sparse tensors in COO format. Make you have a look + at :ref:`the note on the data type of the indices `. + + .. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor. + + Args: + crow_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, nrows + 1)``. The last element of each batch + is the number of non-zeros. This tensor encodes the index in + values and col_indices depending on where the given row + starts. Each successive number in the tensor subtracted by the + number before it denotes the number of elements in a given + row. + col_indices (array_like): Column coordinates of each element in + values. (B+1)-dimensional tensor with the same length + as values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1+K)-dimensional tensor where ``K`` is the number + of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows, ncols, *densesize)``. If + not provided, the size will be inferred as the minimum size + big enough to hold all non-zero elements. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. + + Example:: + + >>> crow_indices = [0, 2, 4] + >>> col_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_csr_tensor(torch.tensor(crow_indices, dtype=torch.int64), + ... torch.tensor(col_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(crow_indices=tensor([0, 2, 4]), + col_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csr) + """ + +def split_copy( + input: Tensor, + split_size: _int | SymInt, + dim: _int = 0, + *, + out: tuple[Tensor, ...] | list[Tensor] | None = None, +) -> None: + r""" + Performs the same operation as :func:`torch.split`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def split_with_sizes( + input: Tensor, + split_sizes: Sequence[_int | SymInt], + dim: _int = 0, +) -> tuple[Tensor, ...]: ... +def split_with_sizes_copy( + input: Tensor, + split_sizes: Sequence[_int | SymInt], + dim: _int = 0, + *, + out: tuple[Tensor, ...] | list[Tensor] | None = None, +) -> None: + r""" + Performs the same operation as :func:`torch.split_with_sizes`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def spmm(input: Tensor, mat2: Tensor) -> Tensor: ... +def sqrt(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + sqrt(input, *, out=None) -> Tensor + + Returns a new tensor with the square-root of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \sqrt{\text{input}_{i}} + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-2.0755, 1.0226, 0.0831, 0.4806]) + >>> torch.sqrt(a) + tensor([ nan, 1.0112, 0.2883, 0.6933]) + """ + +def sqrt_(input: Tensor) -> Tensor: ... +def square(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + square(input: Tensor, *, out: Optional[Tensor]) -> Tensor + + Returns a new tensor with the square of the elements of :attr:`input`. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-2.0755, 1.0226, 0.0831, 0.4806]) + >>> torch.square(a) + tensor([ 4.3077, 1.0457, 0.0069, 0.2310]) + """ + +def square_(input: Tensor) -> Tensor: ... +@overload +def squeeze(input: Tensor) -> Tensor: + r""" + squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor + + Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed. + + For example, if `input` is of shape: + :math:`(A \times 1 \times B \times C \times 1 \times D)` then the `input.squeeze()` + will be of shape: :math:`(A \times B \times C \times D)`. + + When :attr:`dim` is given, a squeeze operation is done only in the given + dimension(s). If `input` is of shape: :math:`(A \times 1 \times B)`, + ``squeeze(input, 0)`` leaves the tensor unchanged, but ``squeeze(input, 1)`` + will squeeze the tensor to the shape :math:`(A \times B)`. + + .. note:: The returned tensor shares the storage with the input tensor, + so changing the contents of one will change the contents of the other. + + .. warning:: If the tensor has a batch dimension of size 1, then `squeeze(input)` + will also remove the batch dimension, which can lead to unexpected + errors. Consider specifying only the dims you wish to be squeezed. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): if given, the input will be squeezed + only in the specified dimensions. + + .. versionchanged:: 2.0 + :attr:`dim` now accepts tuples of dimensions. + + Example:: + + >>> x = torch.zeros(2, 1, 2, 1, 2) + >>> x.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x) + >>> y.size() + torch.Size([2, 2, 2]) + >>> y = torch.squeeze(x, 0) + >>> y.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x, 1) + >>> y.size() + torch.Size([2, 2, 1, 2]) + >>> y = torch.squeeze(x, (1, 2, 3)) + torch.Size([2, 2, 2]) + """ + +@overload +def squeeze(input: Tensor, dim: _int) -> Tensor: + r""" + squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor + + Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed. + + For example, if `input` is of shape: + :math:`(A \times 1 \times B \times C \times 1 \times D)` then the `input.squeeze()` + will be of shape: :math:`(A \times B \times C \times D)`. + + When :attr:`dim` is given, a squeeze operation is done only in the given + dimension(s). If `input` is of shape: :math:`(A \times 1 \times B)`, + ``squeeze(input, 0)`` leaves the tensor unchanged, but ``squeeze(input, 1)`` + will squeeze the tensor to the shape :math:`(A \times B)`. + + .. note:: The returned tensor shares the storage with the input tensor, + so changing the contents of one will change the contents of the other. + + .. warning:: If the tensor has a batch dimension of size 1, then `squeeze(input)` + will also remove the batch dimension, which can lead to unexpected + errors. Consider specifying only the dims you wish to be squeezed. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): if given, the input will be squeezed + only in the specified dimensions. + + .. versionchanged:: 2.0 + :attr:`dim` now accepts tuples of dimensions. + + Example:: + + >>> x = torch.zeros(2, 1, 2, 1, 2) + >>> x.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x) + >>> y.size() + torch.Size([2, 2, 2]) + >>> y = torch.squeeze(x, 0) + >>> y.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x, 1) + >>> y.size() + torch.Size([2, 2, 1, 2]) + >>> y = torch.squeeze(x, (1, 2, 3)) + torch.Size([2, 2, 2]) + """ + +@overload +def squeeze(input: Tensor, dim: _size) -> Tensor: + r""" + squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor + + Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed. + + For example, if `input` is of shape: + :math:`(A \times 1 \times B \times C \times 1 \times D)` then the `input.squeeze()` + will be of shape: :math:`(A \times B \times C \times D)`. + + When :attr:`dim` is given, a squeeze operation is done only in the given + dimension(s). If `input` is of shape: :math:`(A \times 1 \times B)`, + ``squeeze(input, 0)`` leaves the tensor unchanged, but ``squeeze(input, 1)`` + will squeeze the tensor to the shape :math:`(A \times B)`. + + .. note:: The returned tensor shares the storage with the input tensor, + so changing the contents of one will change the contents of the other. + + .. warning:: If the tensor has a batch dimension of size 1, then `squeeze(input)` + will also remove the batch dimension, which can lead to unexpected + errors. Consider specifying only the dims you wish to be squeezed. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): if given, the input will be squeezed + only in the specified dimensions. + + .. versionchanged:: 2.0 + :attr:`dim` now accepts tuples of dimensions. + + Example:: + + >>> x = torch.zeros(2, 1, 2, 1, 2) + >>> x.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x) + >>> y.size() + torch.Size([2, 2, 2]) + >>> y = torch.squeeze(x, 0) + >>> y.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x, 1) + >>> y.size() + torch.Size([2, 2, 1, 2]) + >>> y = torch.squeeze(x, (1, 2, 3)) + torch.Size([2, 2, 2]) + """ + +@overload +def squeeze(input: Tensor, dim: str | EllipsisType | None) -> Tensor: + r""" + squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor + + Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed. + + For example, if `input` is of shape: + :math:`(A \times 1 \times B \times C \times 1 \times D)` then the `input.squeeze()` + will be of shape: :math:`(A \times B \times C \times D)`. + + When :attr:`dim` is given, a squeeze operation is done only in the given + dimension(s). If `input` is of shape: :math:`(A \times 1 \times B)`, + ``squeeze(input, 0)`` leaves the tensor unchanged, but ``squeeze(input, 1)`` + will squeeze the tensor to the shape :math:`(A \times B)`. + + .. note:: The returned tensor shares the storage with the input tensor, + so changing the contents of one will change the contents of the other. + + .. warning:: If the tensor has a batch dimension of size 1, then `squeeze(input)` + will also remove the batch dimension, which can lead to unexpected + errors. Consider specifying only the dims you wish to be squeezed. + + Args: + input (Tensor): the input tensor. + dim (int or tuple of ints, optional): if given, the input will be squeezed + only in the specified dimensions. + + .. versionchanged:: 2.0 + :attr:`dim` now accepts tuples of dimensions. + + Example:: + + >>> x = torch.zeros(2, 1, 2, 1, 2) + >>> x.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x) + >>> y.size() + torch.Size([2, 2, 2]) + >>> y = torch.squeeze(x, 0) + >>> y.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x, 1) + >>> y.size() + torch.Size([2, 2, 1, 2]) + >>> y = torch.squeeze(x, (1, 2, 3)) + torch.Size([2, 2, 2]) + """ + +@overload +def squeeze_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.squeeze`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def squeeze_copy( + input: Tensor, + dim: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.squeeze`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def squeeze_copy( + input: Tensor, + dim: _size, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.squeeze`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def sspaddmm( + beta: Number | _complex, + self: Tensor, + alpha: Number | _complex, + mat1: Tensor, + mat2: Tensor, +) -> Tensor: + r""" + sspaddmm(input, mat1, mat2, *, beta=1, alpha=1, out=None) -> Tensor + + Matrix multiplies a sparse tensor :attr:`mat1` with a dense tensor + :attr:`mat2`, then adds the sparse tensor :attr:`input` to the result. + + Note: This function is equivalent to :func:`torch.addmm`, except + :attr:`input` and :attr:`mat1` are sparse. + + Args: + input (Tensor): a sparse matrix to be added + mat1 (Tensor): a sparse matrix to be matrix multiplied + mat2 (Tensor): a dense matrix to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + """ + +@overload +def sspaddmm( + input: Tensor, + mat1: Tensor, + mat2: Tensor, + *, + beta: Number | _complex = 1, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + sspaddmm(input, mat1, mat2, *, beta=1, alpha=1, out=None) -> Tensor + + Matrix multiplies a sparse tensor :attr:`mat1` with a dense tensor + :attr:`mat2`, then adds the sparse tensor :attr:`input` to the result. + + Note: This function is equivalent to :func:`torch.addmm`, except + :attr:`input` and :attr:`mat1` are sparse. + + Args: + input (Tensor): a sparse matrix to be added + mat1 (Tensor): a sparse matrix to be matrix multiplied + mat2 (Tensor): a dense matrix to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + """ + +@overload +def sspaddmm( + beta: Number | _complex, + self: Tensor, + mat1: Tensor, + mat2: Tensor, +) -> Tensor: + r""" + sspaddmm(input, mat1, mat2, *, beta=1, alpha=1, out=None) -> Tensor + + Matrix multiplies a sparse tensor :attr:`mat1` with a dense tensor + :attr:`mat2`, then adds the sparse tensor :attr:`input` to the result. + + Note: This function is equivalent to :func:`torch.addmm`, except + :attr:`input` and :attr:`mat1` are sparse. + + Args: + input (Tensor): a sparse matrix to be added + mat1 (Tensor): a sparse matrix to be matrix multiplied + mat2 (Tensor): a dense matrix to be matrix multiplied + + Keyword args: + beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): the output tensor. + """ + +def stack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + dim: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + stack(tensors, dim=0, *, out=None) -> Tensor + + Concatenates a sequence of tensors along a new dimension. + + All tensors need to be of the same size. + + .. seealso:: + + :func:`torch.cat` concatenates the given sequence along an existing dimension. + + Arguments: + tensors (sequence of Tensors): sequence of tensors to concatenate + dim (int, optional): dimension to insert. Has to be between 0 and the number + of dimensions of concatenated tensors (inclusive). Default: 0 + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]]) + >>> torch.stack((x, x)) # same as torch.stack((x, x), dim=0) + tensor([[[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]], + + [[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]]]) + >>> torch.stack((x, x)).size() + torch.Size([2, 2, 3]) + >>> torch.stack((x, x), dim=1) + tensor([[[ 0.3367, 0.1288, 0.2345], + [ 0.3367, 0.1288, 0.2345]], + + [[ 0.2303, -1.1229, -0.1863], + [ 0.2303, -1.1229, -0.1863]]]) + >>> torch.stack((x, x), dim=2) + tensor([[[ 0.3367, 0.3367], + [ 0.1288, 0.1288], + [ 0.2345, 0.2345]], + + [[ 0.2303, 0.2303], + [-1.1229, -1.1229], + [-0.1863, -0.1863]]]) + >>> torch.stack((x, x), dim=-1) + tensor([[[ 0.3367, 0.3367], + [ 0.1288, 0.1288], + [ 0.2345, 0.2345]], + + [[ 0.2303, 0.2303], + [-1.1229, -1.1229], + [-0.1863, -0.1863]]]) + """ + +@overload +def std( + input: Tensor, + dim: _int | _size | None, + unbiased: _bool = True, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the standard deviation over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std( + input: Tensor, + dim: _int | _size | None = None, + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the standard deviation over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std(input: Tensor, unbiased: _bool = True) -> Tensor: + r""" + std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the standard deviation over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the standard deviation over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + unbiased: _bool = True, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the standard deviation over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std_mean( + input: Tensor, + dim: _int | _size | None, + unbiased: _bool = True, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the standard deviation and mean over the dimensions specified by + :attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or + ``None`` to reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (std, mean) containing the standard deviation and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std_mean( + input: Tensor, + dim: _int | _size | None = None, + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the standard deviation and mean over the dimensions specified by + :attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or + ``None`` to reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (std, mean) containing the standard deviation and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std_mean( + input: Tensor, + unbiased: _bool = True, +) -> tuple[Tensor, Tensor]: + r""" + std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the standard deviation and mean over the dimensions specified by + :attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or + ``None`` to reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (std, mean) containing the standard deviation and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std_mean( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the standard deviation and mean over the dimensions specified by + :attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or + ``None`` to reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (std, mean) containing the standard deviation and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def std_mean( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + unbiased: _bool = True, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the standard deviation and mean over the dimensions specified by + :attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or + ``None`` to reduce over all dimensions. + + The standard deviation (:math:`\sigma`) is calculated as + + .. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (std, mean) containing the standard deviation and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def sub( + input: Tensor | Number | _complex, + other: Tensor | Number | _complex, + *, + alpha: Number | _complex | None = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + sub(input, other, *, alpha=1, out=None) -> Tensor + + Subtracts :attr:`other`, scaled by :attr:`alpha`, from :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to subtract from :attr:`input`. + + Keyword args: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor((1, 2)) + >>> b = torch.tensor((0, 1)) + >>> torch.sub(a, b, alpha=2) + tensor([1, 0]) + """ + +@overload +def sub(self: Tensor, alpha: Number | _complex, other: Tensor) -> Tensor: + r""" + sub(input, other, *, alpha=1, out=None) -> Tensor + + Subtracts :attr:`other`, scaled by :attr:`alpha`, from :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to subtract from :attr:`input`. + + Keyword args: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor((1, 2)) + >>> b = torch.tensor((0, 1)) + >>> torch.sub(a, b, alpha=2) + tensor([1, 0]) + """ + +@overload +def sub( + self: Tensor, + alpha: Number | _complex, + other: Tensor, + *, + out: Tensor, +) -> Tensor: + r""" + sub(input, other, *, alpha=1, out=None) -> Tensor + + Subtracts :attr:`other`, scaled by :attr:`alpha`, from :attr:`input`. + + .. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i + + + Supports :ref:`broadcasting to a common shape `, + :ref:`type promotion `, and integer, float, and complex inputs. + + Args: + input (Tensor): the input tensor. + other (Tensor or Number): the tensor or number to subtract from :attr:`input`. + + Keyword args: + alpha (Number): the multiplier for :attr:`other`. + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor((1, 2)) + >>> b = torch.tensor((0, 1)) + >>> torch.sub(a, b, alpha=2) + tensor([1, 0]) + """ + +@overload +def subtract( + input: Tensor, + other: Tensor, + *, + alpha: Number | _complex = 1, + out: Tensor | None = None, +) -> Tensor: + r""" + subtract(input, other, *, alpha=1, out=None) -> Tensor + + Alias for :func:`torch.sub`. + """ + +@overload +def subtract( + input: Tensor, + other: Number | _complex, + alpha: Number | _complex = 1, +) -> Tensor: + r""" + subtract(input, other, *, alpha=1, out=None) -> Tensor + + Alias for :func:`torch.sub`. + """ + +@overload +def sum(input: Tensor, *, dtype: _dtype | None = None) -> Tensor: + r""" + sum(input, *, dtype=None) -> Tensor + + Returns the sum of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + .. note:: Use the `dtype` argument if you need the result in a specific tensor type. + Otherwise, the result type may be automatically promoted (e.g., from `torch.int32` to `torch.int64`). + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.1133, -0.9567, 0.2958]]) + >>> torch.sum(a) + tensor(-0.5475) + + .. function:: sum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the sum of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0569, -0.2475, 0.0737, -0.3429], + [-0.2993, 0.9138, 0.9337, -1.6864], + [ 0.1132, 0.7892, -0.1003, 0.5688], + [ 0.3637, -0.9906, -0.4752, -1.5197]]) + >>> torch.sum(a, 1) + tensor([-0.4598, -0.1381, 1.3708, -2.6217]) + >>> b = torch.arange(4 * 5 * 6).view(4, 5, 6) + >>> torch.sum(b, (2, 1)) + tensor([ 435., 1335., 2235., 3135.]) + """ + +@overload +def sum( + input: Tensor, + dim: _int | _size | None, + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + sum(input, *, dtype=None) -> Tensor + + Returns the sum of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + .. note:: Use the `dtype` argument if you need the result in a specific tensor type. + Otherwise, the result type may be automatically promoted (e.g., from `torch.int32` to `torch.int64`). + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.1133, -0.9567, 0.2958]]) + >>> torch.sum(a) + tensor(-0.5475) + + .. function:: sum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the sum of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0569, -0.2475, 0.0737, -0.3429], + [-0.2993, 0.9138, 0.9337, -1.6864], + [ 0.1132, 0.7892, -0.1003, 0.5688], + [ 0.3637, -0.9906, -0.4752, -1.5197]]) + >>> torch.sum(a, 1) + tensor([-0.4598, -0.1381, 1.3708, -2.6217]) + >>> b = torch.arange(4 * 5 * 6).view(4, 5, 6) + >>> torch.sum(b, (2, 1)) + tensor([ 435., 1335., 2235., 3135.]) + """ + +@overload +def sum( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + keepdim: _bool = False, + *, + dtype: _dtype | None = None, + out: Tensor | None = None, +) -> Tensor: + r""" + sum(input, *, dtype=None) -> Tensor + + Returns the sum of all elements in the :attr:`input` tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + .. note:: Use the `dtype` argument if you need the result in a specific tensor type. + Otherwise, the result type may be automatically promoted (e.g., from `torch.int32` to `torch.int64`). + + Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.1133, -0.9567, 0.2958]]) + >>> torch.sum(a) + tensor(-0.5475) + + .. function:: sum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + + Returns the sum of each row of the :attr:`input` tensor in the given + dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + + Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0569, -0.2475, 0.0737, -0.3429], + [-0.2993, 0.9138, 0.9337, -1.6864], + [ 0.1132, 0.7892, -0.1003, 0.5688], + [ 0.3637, -0.9906, -0.4752, -1.5197]]) + >>> torch.sum(a, 1) + tensor([-0.4598, -0.1381, 1.3708, -2.6217]) + >>> b = torch.arange(4 * 5 * 6).view(4, 5, 6) + >>> torch.sum(b, (2, 1)) + tensor([ 435., 1335., 2235., 3135.]) + """ + +def svd( + input: Tensor, + some: _bool = True, + compute_uv: _bool = True, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.svd: + r""" + svd(input, some=True, compute_uv=True, *, out=None) -> (Tensor, Tensor, Tensor) + + Computes the singular value decomposition of either a matrix or batch of + matrices :attr:`input`. The singular value decomposition is represented as a + namedtuple `(U, S, V)`, such that :attr:`input` :math:`= U \text{diag}(S) V^{\text{H}}`. + where :math:`V^{\text{H}}` is the transpose of `V` for real inputs, + and the conjugate transpose of `V` for complex inputs. + If :attr:`input` is a batch of matrices, then `U`, `S`, and `V` are also + batched with the same batch dimensions as :attr:`input`. + + If :attr:`some` is `True` (default), the method returns the reduced singular + value decomposition. In this case, if the last two dimensions of :attr:`input` are + `m` and `n`, then the returned `U` and `V` matrices will contain only + `min(n, m)` orthonormal columns. + + If :attr:`compute_uv` is `False`, the returned `U` and `V` will be + zero-filled matrices of shape `(m, m)` and `(n, n)` + respectively, and the same device as :attr:`input`. The argument :attr:`some` + has no effect when :attr:`compute_uv` is `False`. + + Supports :attr:`input` of float, double, cfloat and cdouble data types. + The dtypes of `U` and `V` are the same as :attr:`input`'s. `S` will + always be real-valued, even if :attr:`input` is complex. + + .. warning:: + + :func:`torch.svd` is deprecated in favor of :func:`torch.linalg.svd` + and will be removed in a future PyTorch release. + + ``U, S, V = torch.svd(A, some=some, compute_uv=True)`` (default) should be replaced with + + .. code:: python + + U, S, Vh = torch.linalg.svd(A, full_matrices=not some) + V = Vh.mH + + ``_, S, _ = torch.svd(A, some=some, compute_uv=False)`` should be replaced with + + .. code:: python + + S = torch.linalg.svdvals(A) + + .. note:: Differences with :func:`torch.linalg.svd`: + + * :attr:`some` is the opposite of + :func:`torch.linalg.svd`'s :attr:`full_matrices`. Note that + default value for both is `True`, so the default behavior is + effectively the opposite. + * :func:`torch.svd` returns `V`, whereas :func:`torch.linalg.svd` returns + `Vh`, that is, :math:`V^{\text{H}}`. + * If :attr:`compute_uv` is `False`, :func:`torch.svd` returns zero-filled + tensors for `U` and `Vh`, whereas :func:`torch.linalg.svd` returns + empty tensors. + + .. note:: The singular values are returned in descending order. If :attr:`input` is a batch of matrices, + then the singular values of each matrix in the batch are returned in descending order. + + .. note:: The `S` tensor can only be used to compute gradients if :attr:`compute_uv` is `True`. + + .. note:: When :attr:`some` is `False`, the gradients on `U[..., :, min(m, n):]` + and `V[..., :, min(m, n):]` will be ignored in the backward pass, as those vectors + can be arbitrary bases of the corresponding subspaces. + + .. note:: The implementation of :func:`torch.linalg.svd` on CPU uses LAPACK's routine `?gesdd` + (a divide-and-conquer algorithm) instead of `?gesvd` for speed. Analogously, + on GPU, it uses cuSOLVER's routines `gesvdj` and `gesvdjBatched` on CUDA 10.1.243 + and later, and MAGMA's routine `gesdd` on earlier versions of CUDA. + + .. note:: The returned `U` will not be contiguous. The matrix (or batch of matrices) will + be represented as a column-major matrix (i.e. Fortran-contiguous). + + .. warning:: The gradients with respect to `U` and `V` will only be finite when the input does not + have zero nor repeated singular values. + + .. warning:: If the distance between any two singular values is close to zero, the gradients with respect to + `U` and `V` will be numerically unstable, as they depends on + :math:`\frac{1}{\min_{i \neq j} \sigma_i^2 - \sigma_j^2}`. The same happens when the matrix + has small singular values, as these gradients also depend on `S^{-1}`. + + .. warning:: For complex-valued :attr:`input` the singular value decomposition is not unique, + as `U` and `V` may be multiplied by an arbitrary phase factor :math:`e^{i \phi}` on every column. + The same happens when :attr:`input` has repeated singular values, where one may multiply + the columns of the spanning subspace in `U` and `V` by a rotation matrix + and `the resulting vectors will span the same subspace`_. + Different platforms, like NumPy, or inputs on different device types, + may produce different `U` and `V` tensors. + + Args: + input (Tensor): the input tensor of size `(*, m, n)` where `*` is zero or more + batch dimensions consisting of `(m, n)` matrices. + some (bool, optional): controls whether to compute the reduced or full decomposition, and + consequently, the shape of returned `U` and `V`. Default: `True`. + compute_uv (bool, optional): controls whether to compute `U` and `V`. Default: `True`. + + Keyword args: + out (tuple, optional): the output tuple of tensors + + Example:: + + >>> a = torch.randn(5, 3) + >>> a + tensor([[ 0.2364, -0.7752, 0.6372], + [ 1.7201, 0.7394, -0.0504], + [-0.3371, -1.0584, 0.5296], + [ 0.3550, -0.4022, 1.5569], + [ 0.2445, -0.0158, 1.1414]]) + >>> u, s, v = torch.svd(a) + >>> u + tensor([[ 0.4027, 0.0287, 0.5434], + [-0.1946, 0.8833, 0.3679], + [ 0.4296, -0.2890, 0.5261], + [ 0.6604, 0.2717, -0.2618], + [ 0.4234, 0.2481, -0.4733]]) + >>> s + tensor([2.3289, 2.0315, 0.7806]) + >>> v + tensor([[-0.0199, 0.8766, 0.4809], + [-0.5080, 0.4054, -0.7600], + [ 0.8611, 0.2594, -0.4373]]) + >>> torch.dist(a, torch.mm(torch.mm(u, torch.diag(s)), v.t())) + tensor(8.6531e-07) + >>> a_big = torch.randn(7, 5, 3) + >>> u, s, v = torch.svd(a_big) + >>> torch.dist(a_big, torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.mT)) + tensor(2.6503e-06) + + .. _the resulting vectors will span the same subspace: + (https://en.wikipedia.org/wiki/Singular_value_decomposition#Singular_values,_singular_vectors,_and_their_relation_to_the_SVD) + """ + +def swapaxes(input: Tensor, axis0: _int, axis1: _int) -> Tensor: + r""" + swapaxes(input, axis0, axis1) -> Tensor + + Alias for :func:`torch.transpose`. + + This function is equivalent to NumPy's swapaxes function. + + Examples:: + + >>> x = torch.tensor([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.swapaxes(x, 0, 1) + tensor([[[0, 1], + [4, 5]], + + [[2, 3], + [6, 7]]]) + >>> torch.swapaxes(x, 0, 2) + tensor([[[0, 4], + [2, 6]], + + [[1, 5], + [3, 7]]]) + """ + +def swapdims(input: Tensor, dim0: _int, dim1: _int) -> Tensor: + r""" + swapdims(input, dim0, dim1) -> Tensor + + Alias for :func:`torch.transpose`. + + This function is equivalent to NumPy's swapaxes function. + + Examples:: + + >>> x = torch.tensor([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.swapdims(x, 0, 1) + tensor([[[0, 1], + [4, 5]], + + [[2, 3], + [6, 7]]]) + >>> torch.swapdims(x, 0, 2) + tensor([[[0, 4], + [2, 6]], + + [[1, 5], + [3, 7]]]) + """ + +def sym_constrain_range( + size: Number | _complex, + *, + min: _int | None = None, + max: _int | None = None, +) -> None: ... +def sym_constrain_range_for_size( + size: Number | _complex, + *, + min: _int | None = None, + max: _int | None = None, +) -> None: ... +def t(input: Tensor) -> Tensor: + r""" + t(input) -> Tensor + + Expects :attr:`input` to be <= 2-D tensor and transposes dimensions 0 + and 1. + + 0-D and 1-D tensors are returned as is. When input is a 2-D tensor this + is equivalent to ``transpose(input, 0, 1)``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x = torch.randn(()) + >>> x + tensor(0.1995) + >>> torch.t(x) + tensor(0.1995) + >>> x = torch.randn(3) + >>> x + tensor([ 2.4320, -0.4608, 0.7702]) + >>> torch.t(x) + tensor([ 2.4320, -0.4608, 0.7702]) + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.4875, 0.9158, -0.5872], + [ 0.3938, -0.6929, 0.6932]]) + >>> torch.t(x) + tensor([[ 0.4875, 0.3938], + [ 0.9158, -0.6929], + [-0.5872, 0.6932]]) + + See also :func:`torch.transpose`. + """ + +def t_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.t`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def take( + input: Tensor, + index: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + take(input, index) -> Tensor + + Returns a new tensor with the elements of :attr:`input` at the given indices. + The input tensor is treated as if it were viewed as a 1-D tensor. The result + takes the same shape as the indices. + + Args: + input (Tensor): the input tensor. + index (LongTensor): the indices into tensor + + Example:: + + >>> src = torch.tensor([[4, 3, 5], + ... [6, 7, 8]]) + >>> torch.take(src, torch.tensor([0, 2, 5])) + tensor([ 4, 5, 8]) + """ + +def take_along_dim( + input: Tensor, + indices: Tensor, + dim: _int | None = None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + take_along_dim(input, indices, dim=None, *, out=None) -> Tensor + + Selects values from :attr:`input` at the 1-dimensional indices from :attr:`indices` along the given :attr:`dim`. + + If :attr:`dim` is None, the input array is treated as if it has been flattened to 1d. + + Functions that return indices along a dimension, like :func:`torch.argmax` and :func:`torch.argsort`, + are designed to work with this function. See the examples below. + + .. note:: + This function is similar to NumPy's `take_along_axis`. + See also :func:`torch.gather`. + + Args: + input (Tensor): the input tensor. + indices (LongTensor): the indices into :attr:`input`. Must have long dtype. + dim (int, optional): dimension to select along. Default: 0 + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> t = torch.tensor([[10, 30, 20], [60, 40, 50]]) + >>> max_idx = torch.argmax(t) + >>> torch.take_along_dim(t, max_idx) + tensor([60]) + >>> sorted_idx = torch.argsort(t, dim=1) + >>> torch.take_along_dim(t, sorted_idx, dim=1) + tensor([[10, 20, 30], + [40, 50, 60]]) + """ + +def tan(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + tan(input, *, out=None) -> Tensor + + Returns a new tensor with the tangent of the elements of :attr:`input`. + + .. math:: + \text{out}_{i} = \tan(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-1.2027, -1.7687, 0.4412, -1.3856]) + >>> torch.tan(a) + tensor([-2.5930, 4.9859, 0.4722, -5.3366]) + """ + +def tan_(input: Tensor) -> Tensor: ... +def tanh(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + tanh(input, *, out=None) -> Tensor + + Returns a new tensor with the hyperbolic tangent of the elements + of :attr:`input`. + + .. math:: + \text{out}_{i} = \tanh(\text{input}_{i}) + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.8986, -0.7279, 1.1745, 0.2611]) + >>> torch.tanh(a) + tensor([ 0.7156, -0.6218, 0.8257, 0.2553]) + """ + +def tanh_(input: Tensor) -> Tensor: ... +def tensor( + data: Any, + dtype: _dtype | None = None, + device: DeviceLikeType | None = None, + requires_grad: _bool = False, + pin_memory: _bool = False, +) -> Tensor: + r""" + tensor(data, *, dtype=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + + Constructs a tensor with no autograd history (also known as a "leaf tensor", see :doc:`/notes/autograd`) by copying :attr:`data`. + + .. warning:: + + When working with tensors prefer using :func:`torch.Tensor.clone`, + :func:`torch.Tensor.detach`, and :func:`torch.Tensor.requires_grad_` for + readability. Letting `t` be a tensor, ``torch.tensor(t)`` is equivalent to + ``t.detach().clone()``, and ``torch.tensor(t, requires_grad=True)`` + is equivalent to ``t.detach().clone().requires_grad_(True)``. + + .. seealso:: + + :func:`torch.as_tensor` preserves autograd history and avoids copies where possible. + :func:`torch.from_numpy` creates a tensor that shares storage with a NumPy array. + + Args: + data (array_like): Initial data for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, infers data type from :attr:`data`. + device (:class:`torch.device`, optional): the device of the constructed tensor. If None and data is a tensor + then the device of data is used. If None and data is not a tensor then + the result tensor is constructed on the current device. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + + + Example:: + + >>> torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]]) + tensor([[ 0.1000, 1.2000], + [ 2.2000, 3.1000], + [ 4.9000, 5.2000]]) + + >>> torch.tensor([0, 1]) # Type inference on data + tensor([ 0, 1]) + + >>> torch.tensor([[0.11111, 0.222222, 0.3333333]], + ... dtype=torch.float64, + ... device=torch.device('cuda:0')) # creates a double tensor on a CUDA device + tensor([[ 0.1111, 0.2222, 0.3333]], dtype=torch.float64, device='cuda:0') + + >>> torch.tensor(3.14159) # Create a zero-dimensional (scalar) tensor + tensor(3.1416) + + >>> torch.tensor([]) # Create an empty tensor (of size (0,)) + tensor([]) + """ + +@overload +def tensor_split( + input: Tensor, + tensor_indices_or_sections: Tensor, + dim: _int = 0, +) -> tuple[Tensor, ...]: + r""" + tensor_split(input, indices_or_sections, dim=0) -> List of Tensors + + Splits a tensor into multiple sub-tensors, all of which are views of :attr:`input`, + along dimension :attr:`dim` according to the indices or number of sections specified + by :attr:`indices_or_sections`. This function is based on NumPy's + :func:`numpy.array_split`. + + Args: + input (Tensor): the tensor to split + indices_or_sections (Tensor, int or list or tuple of ints): + If :attr:`indices_or_sections` is an integer ``n`` or a zero dimensional long tensor + with value ``n``, :attr:`input` is split into ``n`` sections along dimension :attr:`dim`. + If :attr:`input` is divisible by ``n`` along dimension :attr:`dim`, each + section will be of equal size, :code:`input.size(dim) / n`. If :attr:`input` + is not divisible by ``n``, the sizes of the first :code:`int(input.size(dim) % n)` + sections will have size :code:`int(input.size(dim) / n) + 1`, and the rest will + have size :code:`int(input.size(dim) / n)`. + + If :attr:`indices_or_sections` is a list or tuple of ints, or a one-dimensional long + tensor, then :attr:`input` is split along dimension :attr:`dim` at each of the indices + in the list, tuple or tensor. For instance, :code:`indices_or_sections=[2, 3]` and :code:`dim=0` + would result in the tensors :code:`input[:2]`, :code:`input[2:3]`, and :code:`input[3:]`. + + If :attr:`indices_or_sections` is a tensor, it must be a zero-dimensional or one-dimensional + long tensor on the CPU. + + dim (int, optional): dimension along which to split the tensor. Default: ``0`` + + Example:: + + >>> x = torch.arange(8) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7])) + + >>> x = torch.arange(7) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6])) + >>> torch.tensor_split(x, (1, 6)) + (tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6])) + + >>> x = torch.arange(14).reshape(2, 7) + >>> x + tensor([[ 0, 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12, 13]]) + >>> torch.tensor_split(x, 3, dim=1) + (tensor([[0, 1, 2], + [7, 8, 9]]), + tensor([[ 3, 4], + [10, 11]]), + tensor([[ 5, 6], + [12, 13]])) + >>> torch.tensor_split(x, (1, 6), dim=1) + (tensor([[0], + [7]]), + tensor([[ 1, 2, 3, 4, 5], + [ 8, 9, 10, 11, 12]]), + tensor([[ 6], + [13]])) + """ + +@overload +def tensor_split( + input: Tensor, + sections: _int | SymInt, + dim: _int = 0, +) -> tuple[Tensor, ...]: + r""" + tensor_split(input, indices_or_sections, dim=0) -> List of Tensors + + Splits a tensor into multiple sub-tensors, all of which are views of :attr:`input`, + along dimension :attr:`dim` according to the indices or number of sections specified + by :attr:`indices_or_sections`. This function is based on NumPy's + :func:`numpy.array_split`. + + Args: + input (Tensor): the tensor to split + indices_or_sections (Tensor, int or list or tuple of ints): + If :attr:`indices_or_sections` is an integer ``n`` or a zero dimensional long tensor + with value ``n``, :attr:`input` is split into ``n`` sections along dimension :attr:`dim`. + If :attr:`input` is divisible by ``n`` along dimension :attr:`dim`, each + section will be of equal size, :code:`input.size(dim) / n`. If :attr:`input` + is not divisible by ``n``, the sizes of the first :code:`int(input.size(dim) % n)` + sections will have size :code:`int(input.size(dim) / n) + 1`, and the rest will + have size :code:`int(input.size(dim) / n)`. + + If :attr:`indices_or_sections` is a list or tuple of ints, or a one-dimensional long + tensor, then :attr:`input` is split along dimension :attr:`dim` at each of the indices + in the list, tuple or tensor. For instance, :code:`indices_or_sections=[2, 3]` and :code:`dim=0` + would result in the tensors :code:`input[:2]`, :code:`input[2:3]`, and :code:`input[3:]`. + + If :attr:`indices_or_sections` is a tensor, it must be a zero-dimensional or one-dimensional + long tensor on the CPU. + + dim (int, optional): dimension along which to split the tensor. Default: ``0`` + + Example:: + + >>> x = torch.arange(8) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7])) + + >>> x = torch.arange(7) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6])) + >>> torch.tensor_split(x, (1, 6)) + (tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6])) + + >>> x = torch.arange(14).reshape(2, 7) + >>> x + tensor([[ 0, 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12, 13]]) + >>> torch.tensor_split(x, 3, dim=1) + (tensor([[0, 1, 2], + [7, 8, 9]]), + tensor([[ 3, 4], + [10, 11]]), + tensor([[ 5, 6], + [12, 13]])) + >>> torch.tensor_split(x, (1, 6), dim=1) + (tensor([[0], + [7]]), + tensor([[ 1, 2, 3, 4, 5], + [ 8, 9, 10, 11, 12]]), + tensor([[ 6], + [13]])) + """ + +@overload +def tensor_split( + input: Tensor, + indices: Sequence[_int | SymInt], + dim: _int = 0, +) -> tuple[Tensor, ...]: + r""" + tensor_split(input, indices_or_sections, dim=0) -> List of Tensors + + Splits a tensor into multiple sub-tensors, all of which are views of :attr:`input`, + along dimension :attr:`dim` according to the indices or number of sections specified + by :attr:`indices_or_sections`. This function is based on NumPy's + :func:`numpy.array_split`. + + Args: + input (Tensor): the tensor to split + indices_or_sections (Tensor, int or list or tuple of ints): + If :attr:`indices_or_sections` is an integer ``n`` or a zero dimensional long tensor + with value ``n``, :attr:`input` is split into ``n`` sections along dimension :attr:`dim`. + If :attr:`input` is divisible by ``n`` along dimension :attr:`dim`, each + section will be of equal size, :code:`input.size(dim) / n`. If :attr:`input` + is not divisible by ``n``, the sizes of the first :code:`int(input.size(dim) % n)` + sections will have size :code:`int(input.size(dim) / n) + 1`, and the rest will + have size :code:`int(input.size(dim) / n)`. + + If :attr:`indices_or_sections` is a list or tuple of ints, or a one-dimensional long + tensor, then :attr:`input` is split along dimension :attr:`dim` at each of the indices + in the list, tuple or tensor. For instance, :code:`indices_or_sections=[2, 3]` and :code:`dim=0` + would result in the tensors :code:`input[:2]`, :code:`input[2:3]`, and :code:`input[3:]`. + + If :attr:`indices_or_sections` is a tensor, it must be a zero-dimensional or one-dimensional + long tensor on the CPU. + + dim (int, optional): dimension along which to split the tensor. Default: ``0`` + + Example:: + + >>> x = torch.arange(8) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7])) + + >>> x = torch.arange(7) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6])) + >>> torch.tensor_split(x, (1, 6)) + (tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6])) + + >>> x = torch.arange(14).reshape(2, 7) + >>> x + tensor([[ 0, 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12, 13]]) + >>> torch.tensor_split(x, 3, dim=1) + (tensor([[0, 1, 2], + [7, 8, 9]]), + tensor([[ 3, 4], + [10, 11]]), + tensor([[ 5, 6], + [12, 13]])) + >>> torch.tensor_split(x, (1, 6), dim=1) + (tensor([[0], + [7]]), + tensor([[ 1, 2, 3, 4, 5], + [ 8, 9, 10, 11, 12]]), + tensor([[ 6], + [13]])) + """ + +def threshold( + input: Tensor, + threshold: Number | _complex, + value: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: ... +def threshold_( + input: Tensor, + threshold: Number | _complex, + value: Number | _complex, +) -> Tensor: ... +def tile(input: Tensor, dims: Sequence[_int | SymInt]) -> Tensor: + r""" + tile(input, dims) -> Tensor + + Constructs a tensor by repeating the elements of :attr:`input`. + The :attr:`dims` argument specifies the number of repetitions + in each dimension. + + If :attr:`dims` specifies fewer dimensions than :attr:`input` has, then + ones are prepended to :attr:`dims` until all dimensions are specified. + For example, if :attr:`input` has shape (8, 6, 4, 2) and :attr:`dims` + is (2, 2), then :attr:`dims` is treated as (1, 1, 2, 2). + + Analogously, if :attr:`input` has fewer dimensions than :attr:`dims` + specifies, then :attr:`input` is treated as if it were unsqueezed at + dimension zero until it has as many dimensions as :attr:`dims` specifies. + For example, if :attr:`input` has shape (4, 2) and :attr:`dims` + is (3, 3, 2, 2), then :attr:`input` is treated as if it had the + shape (1, 1, 4, 2). + + .. note:: + + This function is similar to NumPy's tile function. + + Args: + input (Tensor): the tensor whose elements to repeat. + dims (tuple): the number of repetitions per dimension. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.tile((2,)) + tensor([1, 2, 3, 1, 2, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.tile(y, (2, 2)) + tensor([[1, 2, 1, 2], + [3, 4, 3, 4], + [1, 2, 1, 2], + [3, 4, 3, 4]]) + """ + +def topk( + input: Tensor, + k: _int | SymInt, + dim: _int = -1, + largest: _bool = True, + sorted: _bool = True, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.topk: + r""" + topk(input, k, dim=None, largest=True, sorted=True, *, out=None) -> (Tensor, LongTensor) + + Returns the :attr:`k` largest elements of the given :attr:`input` tensor along + a given dimension. + + If :attr:`dim` is not given, the last dimension of the `input` is chosen. + + If :attr:`largest` is ``False`` then the `k` smallest elements are returned. + + A namedtuple of `(values, indices)` is returned with the `values` and + `indices` of the largest `k` elements of each row of the `input` tensor in the + given dimension `dim`. + + The boolean option :attr:`sorted` if ``True``, will make sure that the returned + `k` elements are themselves sorted + + .. note:: + When using `torch.topk`, the indices of tied elements are not guaranteed to be stable + and may vary across different invocations. + + Args: + input (Tensor): the input tensor. + k (int): the k in "top-k" + dim (int, optional): the dimension to sort along + largest (bool, optional): controls whether to return largest or + smallest elements + sorted (bool, optional): controls whether to return the elements + in sorted order + + Keyword args: + out (tuple, optional): the output tuple of (Tensor, LongTensor) that can be + optionally given to be used as output buffers + + Example:: + + >>> x = torch.arange(1., 6.) + >>> x + tensor([ 1., 2., 3., 4., 5.]) + >>> torch.topk(x, 3) + torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2])) + """ + +def trace(input: Tensor) -> Tensor: + r""" + trace(input) -> Tensor + + Returns the sum of the elements of the diagonal of the input 2-D matrix. + + Example:: + + >>> x = torch.arange(1., 10.).view(3, 3) + >>> x + tensor([[ 1., 2., 3.], + [ 4., 5., 6.], + [ 7., 8., 9.]]) + >>> torch.trace(x) + tensor(15.) + """ + +@overload +def transpose(input: Tensor, dim0: _int, dim1: _int) -> Tensor: + r""" + transpose(input, dim0, dim1) -> Tensor + + Returns a tensor that is a transposed version of :attr:`input`. + The given dimensions :attr:`dim0` and :attr:`dim1` are swapped. + + If :attr:`input` is a strided tensor then the resulting :attr:`out` + tensor shares its underlying storage with the :attr:`input` tensor, so + changing the content of one would change the content of the other. + + If :attr:`input` is a :ref:`sparse tensor ` then the + resulting :attr:`out` tensor *does not* share the underlying storage + with the :attr:`input` tensor. + + If :attr:`input` is a :ref:`sparse tensor ` with compressed + layout (SparseCSR, SparseBSR, SparseCSC or SparseBSC) the arguments + :attr:`dim0` and :attr:`dim1` must be both batch dimensions, or must + both be sparse dimensions. The batch dimensions of a sparse tensor are the + dimensions preceding the sparse dimensions. + + .. note:: + Transpositions which interchange the sparse dimensions of a `SparseCSR` + or `SparseCSC` layout tensor will result in the layout changing between + the two options. Transposition of the sparse dimensions of a ` SparseBSR` + or `SparseBSC` layout tensor will likewise generate a result with the + opposite layout. + + + Args: + input (Tensor): the input tensor. + dim0 (int): the first dimension to be transposed + dim1 (int): the second dimension to be transposed + + Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 1.0028, -0.9893, 0.5809], + [-0.1669, 0.7299, 0.4942]]) + >>> torch.transpose(x, 0, 1) + tensor([[ 1.0028, -0.1669], + [-0.9893, 0.7299], + [ 0.5809, 0.4942]]) + + See also :func:`torch.t`. + """ + +@overload +def transpose( + input: Tensor, + dim0: str | EllipsisType | None, + dim1: str | EllipsisType | None, +) -> Tensor: + r""" + transpose(input, dim0, dim1) -> Tensor + + Returns a tensor that is a transposed version of :attr:`input`. + The given dimensions :attr:`dim0` and :attr:`dim1` are swapped. + + If :attr:`input` is a strided tensor then the resulting :attr:`out` + tensor shares its underlying storage with the :attr:`input` tensor, so + changing the content of one would change the content of the other. + + If :attr:`input` is a :ref:`sparse tensor ` then the + resulting :attr:`out` tensor *does not* share the underlying storage + with the :attr:`input` tensor. + + If :attr:`input` is a :ref:`sparse tensor ` with compressed + layout (SparseCSR, SparseBSR, SparseCSC or SparseBSC) the arguments + :attr:`dim0` and :attr:`dim1` must be both batch dimensions, or must + both be sparse dimensions. The batch dimensions of a sparse tensor are the + dimensions preceding the sparse dimensions. + + .. note:: + Transpositions which interchange the sparse dimensions of a `SparseCSR` + or `SparseCSC` layout tensor will result in the layout changing between + the two options. Transposition of the sparse dimensions of a ` SparseBSR` + or `SparseBSC` layout tensor will likewise generate a result with the + opposite layout. + + + Args: + input (Tensor): the input tensor. + dim0 (int): the first dimension to be transposed + dim1 (int): the second dimension to be transposed + + Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 1.0028, -0.9893, 0.5809], + [-0.1669, 0.7299, 0.4942]]) + >>> torch.transpose(x, 0, 1) + tensor([[ 1.0028, -0.1669], + [-0.9893, 0.7299], + [ 0.5809, 0.4942]]) + + See also :func:`torch.t`. + """ + +def transpose_copy( + input: Tensor, + dim0: _int, + dim1: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.transpose`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def trapezoid(y: Tensor, x: Tensor, *, dim: _int = -1) -> Tensor: + r""" + trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + + Computes the `trapezoidal rule `_ along + :attr:`dim`. By default the spacing between elements is assumed to be 1, but + :attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be + used to specify arbitrary spacing along :attr:`dim`. Only one of :attr:`x` or :attr:`dx` should be specified. + + + Assuming :attr:`y` is a one-dimensional tensor with elements :math:`{y_0, y_1, ..., y_n}`, + the default computation is + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{1}{2} (y_i + y_{i-1}) + \end{aligned} + + When :attr:`dx` is specified the computation becomes + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{\Delta x}{2} (y_i + y_{i-1}) + \end{aligned} + + effectively multiplying the result by :attr:`dx`. When :attr:`x` is specified, + assuming :attr:`x` is also a one-dimensional tensor with + elements :math:`{x_0, x_1, ..., x_n}`, the computation becomes + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{(x_i - x_{i-1})}{2} (y_i + y_{i-1}) + \end{aligned} + + When :attr:`x` and :attr:`y` have the same size, the computation is as described above and no broadcasting is needed. + The broadcasting behavior of this function is as follows when their sizes are different. For both :attr:`x` + and :attr:`y`, the function computes the difference between consecutive elements along + dimension :attr:`dim`. This effectively creates two tensors, `x_diff` and `y_diff`, that have + the same shape as the original tensors except their lengths along the dimension :attr:`dim` is reduced by 1. + After that, those two tensors are broadcast together to compute final output as part of the trapezoidal rule. + See the examples below for details. + + .. note:: + The trapezoidal rule is a technique for approximating the definite integral of a function + by averaging its left and right Riemann sums. The approximation becomes more accurate as + the resolution of the partition increases. + + Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + + Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + + Examples:: + + >>> # Computes the trapezoidal rule in 1D, spacing is implicitly 1 + >>> y = torch.tensor([1, 5, 10]) + >>> torch.trapezoid(y) + tensor(10.5) + + >>> # Computes the same trapezoidal rule directly to verify + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.trapezoid(y, dx=2) + 21.0 + + >>> # Computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + 28.5 + + >>> # Computes the same trapezoidal rule directly to verify + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.trapezoid(y) + tensor([ 2., 8., 14.]) + + >>> # Computes the trapezoidal rule for each column of the matrix + >>> torch.trapezoid(y, dim=0) + tensor([ 6., 8., 10.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + array([5., 5., 5.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.trapezoid(y, x) + array([2., 4., 6.]) + """ + +@overload +def trapezoid( + y: Tensor, + *, + dx: Number | _complex = 1, + dim: _int = -1, +) -> Tensor: + r""" + trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + + Computes the `trapezoidal rule `_ along + :attr:`dim`. By default the spacing between elements is assumed to be 1, but + :attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be + used to specify arbitrary spacing along :attr:`dim`. Only one of :attr:`x` or :attr:`dx` should be specified. + + + Assuming :attr:`y` is a one-dimensional tensor with elements :math:`{y_0, y_1, ..., y_n}`, + the default computation is + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{1}{2} (y_i + y_{i-1}) + \end{aligned} + + When :attr:`dx` is specified the computation becomes + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{\Delta x}{2} (y_i + y_{i-1}) + \end{aligned} + + effectively multiplying the result by :attr:`dx`. When :attr:`x` is specified, + assuming :attr:`x` is also a one-dimensional tensor with + elements :math:`{x_0, x_1, ..., x_n}`, the computation becomes + + .. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{(x_i - x_{i-1})}{2} (y_i + y_{i-1}) + \end{aligned} + + When :attr:`x` and :attr:`y` have the same size, the computation is as described above and no broadcasting is needed. + The broadcasting behavior of this function is as follows when their sizes are different. For both :attr:`x` + and :attr:`y`, the function computes the difference between consecutive elements along + dimension :attr:`dim`. This effectively creates two tensors, `x_diff` and `y_diff`, that have + the same shape as the original tensors except their lengths along the dimension :attr:`dim` is reduced by 1. + After that, those two tensors are broadcast together to compute final output as part of the trapezoidal rule. + See the examples below for details. + + .. note:: + The trapezoidal rule is a technique for approximating the definite integral of a function + by averaging its left and right Riemann sums. The approximation becomes more accurate as + the resolution of the partition increases. + + Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + + Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + + Examples:: + + >>> # Computes the trapezoidal rule in 1D, spacing is implicitly 1 + >>> y = torch.tensor([1, 5, 10]) + >>> torch.trapezoid(y) + tensor(10.5) + + >>> # Computes the same trapezoidal rule directly to verify + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.trapezoid(y, dx=2) + 21.0 + + >>> # Computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + 28.5 + + >>> # Computes the same trapezoidal rule directly to verify + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.trapezoid(y) + tensor([ 2., 8., 14.]) + + >>> # Computes the trapezoidal rule for each column of the matrix + >>> torch.trapezoid(y, dim=0) + tensor([ 6., 8., 10.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + array([5., 5., 5.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.trapezoid(y, x) + array([2., 4., 6.]) + """ + +@overload +def trapz(y: Tensor, *, dx: _float = 1, dim: _int = -1) -> Tensor: + r""" + trapz(y, x, *, dim=-1) -> Tensor + + Alias for :func:`torch.trapezoid`. + """ + +@overload +def trapz(y: Tensor, x: Tensor, *, dim: _int = -1) -> Tensor: + r""" + trapz(y, x, *, dim=-1) -> Tensor + + Alias for :func:`torch.trapezoid`. + """ + +def triangular_solve( + input: Tensor, + A: Tensor, + upper: _bool = True, + transpose: _bool = False, + unitriangular: _bool = False, + *, + out: Tensor | tuple[Tensor, ...] | list[Tensor] | None = None, +) -> torch.return_types.triangular_solve: + r""" + triangular_solve(b, A, upper=True, transpose=False, unitriangular=False, *, out=None) -> (Tensor, Tensor) + + Solves a system of equations with a square upper or lower triangular invertible matrix :math:`A` + and multiple right-hand sides :math:`b`. + + In symbols, it solves :math:`AX = b` and assumes :math:`A` is square upper-triangular + (or lower-triangular if :attr:`upper`\ `= False`) and does not have zeros on the diagonal. + + `torch.triangular_solve(b, A)` can take in 2D inputs `b, A` or inputs that are + batches of 2D matrices. If the inputs are batches, then returns + batched outputs `X` + + If the diagonal of :attr:`A` contains zeros or elements that are very close to zero and + :attr:`unitriangular`\ `= False` (default) or if the input matrix is badly conditioned, + the result may contain `NaN` s. + + Supports input of float, double, cfloat and cdouble data types. + + .. warning:: + + :func:`torch.triangular_solve` is deprecated in favor of :func:`torch.linalg.solve_triangular` + and will be removed in a future PyTorch release. + :func:`torch.linalg.solve_triangular` has its arguments reversed and does not return a + copy of one of the inputs. + + ``X = torch.triangular_solve(B, A).solution`` should be replaced with + + .. code:: python + + X = torch.linalg.solve_triangular(A, B) + + Args: + b (Tensor): multiple right-hand sides of size :math:`(*, m, k)` where + :math:`*` is zero of more batch dimensions + A (Tensor): the input triangular coefficient matrix of size :math:`(*, m, m)` + where :math:`*` is zero or more batch dimensions + upper (bool, optional): whether :math:`A` is upper or lower triangular. Default: ``True``. + transpose (bool, optional): solves `op(A)X = b` where `op(A) = A^T` if this flag is ``True``, + and `op(A) = A` if it is ``False``. Default: ``False``. + unitriangular (bool, optional): whether :math:`A` is unit triangular. + If True, the diagonal elements of :math:`A` are assumed to be + 1 and not referenced from :math:`A`. Default: ``False``. + + Keyword args: + out ((Tensor, Tensor), optional): tuple of two tensors to write + the output to. Ignored if `None`. Default: `None`. + + Returns: + A namedtuple `(solution, cloned_coefficient)` where `cloned_coefficient` + is a clone of :math:`A` and `solution` is the solution :math:`X` to :math:`AX = b` + (or whatever variant of the system of equations, depending on the keyword arguments.) + + Examples:: + + >>> A = torch.randn(2, 2).triu() + >>> A + tensor([[ 1.1527, -1.0753], + [ 0.0000, 0.7986]]) + >>> b = torch.randn(2, 3) + >>> b + tensor([[-0.0210, 2.3513, -1.5492], + [ 1.5429, 0.7403, -1.0243]]) + >>> torch.triangular_solve(b, A) + torch.return_types.triangular_solve( + solution=tensor([[ 1.7841, 2.9046, -2.5405], + [ 1.9320, 0.9270, -1.2826]]), + cloned_coefficient=tensor([[ 1.1527, -1.0753], + [ 0.0000, 0.7986]])) + """ + +def tril( + input: Tensor, + diagonal: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + tril(input, diagonal=0, *, out=None) -> Tensor + + Returns the lower triangular part of the matrix (2-D tensor) or batch of matrices + :attr:`input`, the other elements of the result tensor :attr:`out` are set to 0. + + The lower triangular part of the matrix is defined as the elements on and + below the diagonal. + + The argument :attr:`diagonal` controls which diagonal to consider. If + :attr:`diagonal` = 0, all elements on and below the main diagonal are + retained. A positive value includes just as many diagonals above the main + diagonal, and similarly a negative value excludes just as many diagonals below + the main diagonal. The main diagonal are the set of indices + :math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` where + :math:`d_{1}, d_{2}` are the dimensions of the matrix. + + Args: + input (Tensor): the input tensor. + diagonal (int, optional): the diagonal to consider + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-1.0813, -0.8619, 0.7105], + [ 0.0935, 0.1380, 2.2112], + [-0.3409, -0.9828, 0.0289]]) + >>> torch.tril(a) + tensor([[-1.0813, 0.0000, 0.0000], + [ 0.0935, 0.1380, 0.0000], + [-0.3409, -0.9828, 0.0289]]) + + >>> b = torch.randn(4, 6) + >>> b + tensor([[ 1.2219, 0.5653, -0.2521, -0.2345, 1.2544, 0.3461], + [ 0.4785, -0.4477, 0.6049, 0.6368, 0.8775, 0.7145], + [ 1.1502, 3.2716, -1.1243, -0.5413, 0.3615, 0.6864], + [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024, 0.0978]]) + >>> torch.tril(b, diagonal=1) + tensor([[ 1.2219, 0.5653, 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.4785, -0.4477, 0.6049, 0.0000, 0.0000, 0.0000], + [ 1.1502, 3.2716, -1.1243, -0.5413, 0.0000, 0.0000], + [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024, 0.0000]]) + >>> torch.tril(b, diagonal=-1) + tensor([[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.4785, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [ 1.1502, 3.2716, 0.0000, 0.0000, 0.0000, 0.0000], + [-0.0614, -0.7344, -1.3164, 0.0000, 0.0000, 0.0000]]) + """ + +def tril_indices( + row: _int, + col: _int, + offset: _int = 0, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + tril_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) -> Tensor + + Returns the indices of the lower triangular part of a :attr:`row`-by- + :attr:`col` matrix in a 2-by-N Tensor, where the first row contains row + coordinates of all indices and the second row contains column coordinates. + Indices are ordered based on rows and then columns. + + The lower triangular part of the matrix is defined as the elements on and + below the diagonal. + + The argument :attr:`offset` controls which diagonal to consider. If + :attr:`offset` = 0, all elements on and below the main diagonal are + retained. A positive value includes just as many diagonals above the main + diagonal, and similarly a negative value excludes just as many diagonals below + the main diagonal. The main diagonal are the set of indices + :math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` + where :math:`d_{1}, d_{2}` are the dimensions of the matrix. + + .. note:: + When running on CUDA, ``row * col`` must be less than :math:`2^{59}` to + prevent overflow during calculation. + + Args: + row (``int``): number of rows in the 2-D matrix. + col (``int``): number of columns in the 2-D matrix. + offset (``int``): diagonal offset from the main diagonal. + Default: if not provided, 0. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor, + only support ``torch.int``, ``torch.long``. Default: if ``None``, ``torch.long``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + layout (:class:`torch.layout`, optional): currently only support ``torch.strided``. + + Example:: + + >>> a = torch.tril_indices(3, 3) + >>> a + tensor([[0, 1, 1, 2, 2, 2], + [0, 0, 1, 0, 1, 2]]) + + >>> a = torch.tril_indices(4, 3, -1) + >>> a + tensor([[1, 2, 2, 3, 3, 3], + [0, 0, 1, 0, 1, 2]]) + + >>> a = torch.tril_indices(4, 3, 1) + >>> a + tensor([[0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], + [0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2]]) + """ + +def triplet_margin_loss( + anchor: Tensor, + positive: Tensor, + negative: Tensor, + margin: _float = 1.0, + p: _float = 2, + eps: _float = 1e-06, + swap: _bool = False, + reduction: _int = 1, +) -> Tensor: ... +def triu( + input: Tensor, + diagonal: _int = 0, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + triu(input, diagonal=0, *, out=None) -> Tensor + + Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices + :attr:`input`, the other elements of the result tensor :attr:`out` are set to 0. + + The upper triangular part of the matrix is defined as the elements on and + above the diagonal. + + The argument :attr:`diagonal` controls which diagonal to consider. If + :attr:`diagonal` = 0, all elements on and above the main diagonal are + retained. A positive value excludes just as many diagonals above the main + diagonal, and similarly a negative value includes just as many diagonals below + the main diagonal. The main diagonal are the set of indices + :math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` where + :math:`d_{1}, d_{2}` are the dimensions of the matrix. + + Args: + input (Tensor): the input tensor. + diagonal (int, optional): the diagonal to consider + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.2072, -1.0680, 0.6602], + [ 0.3480, -0.5211, -0.4573]]) + >>> torch.triu(a) + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.0000, -1.0680, 0.6602], + [ 0.0000, 0.0000, -0.4573]]) + >>> torch.triu(a, diagonal=1) + tensor([[ 0.0000, 0.5207, 2.0049], + [ 0.0000, 0.0000, 0.6602], + [ 0.0000, 0.0000, 0.0000]]) + >>> torch.triu(a, diagonal=-1) + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.2072, -1.0680, 0.6602], + [ 0.0000, -0.5211, -0.4573]]) + + >>> b = torch.randn(4, 6) + >>> b + tensor([[ 0.5876, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [-0.2447, 0.9556, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.4333, 0.3146, 0.6576, -1.0432, 0.9348, -0.4410], + [-0.9888, 1.0679, -1.3337, -1.6556, 0.4798, 0.2830]]) + >>> torch.triu(b, diagonal=1) + tensor([[ 0.0000, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [ 0.0000, 0.0000, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.0000, 0.0000, 0.0000, -1.0432, 0.9348, -0.4410], + [ 0.0000, 0.0000, 0.0000, 0.0000, 0.4798, 0.2830]]) + >>> torch.triu(b, diagonal=-1) + tensor([[ 0.5876, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [-0.2447, 0.9556, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.0000, 0.3146, 0.6576, -1.0432, 0.9348, -0.4410], + [ 0.0000, 0.0000, -1.3337, -1.6556, 0.4798, 0.2830]]) + """ + +def triu_indices( + row: _int, + col: _int, + offset: _int = 0, + *, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + triu_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) -> Tensor + + Returns the indices of the upper triangular part of a :attr:`row` by + :attr:`col` matrix in a 2-by-N Tensor, where the first row contains row + coordinates of all indices and the second row contains column coordinates. + Indices are ordered based on rows and then columns. + + The upper triangular part of the matrix is defined as the elements on and + above the diagonal. + + The argument :attr:`offset` controls which diagonal to consider. If + :attr:`offset` = 0, all elements on and above the main diagonal are + retained. A positive value excludes just as many diagonals above the main + diagonal, and similarly a negative value includes just as many diagonals below + the main diagonal. The main diagonal are the set of indices + :math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` + where :math:`d_{1}, d_{2}` are the dimensions of the matrix. + + .. note:: + When running on CUDA, ``row * col`` must be less than :math:`2^{59}` to + prevent overflow during calculation. + + Args: + row (``int``): number of rows in the 2-D matrix. + col (``int``): number of columns in the 2-D matrix. + offset (``int``): diagonal offset from the main diagonal. + Default: if not provided, 0. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor, + only support ``torch.int``, ``torch.long``. Default: if ``None``, ``torch.long``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + layout (:class:`torch.layout`, optional): currently only support ``torch.strided``. + + Example:: + + >>> a = torch.triu_indices(3, 3) + >>> a + tensor([[0, 0, 0, 1, 1, 2], + [0, 1, 2, 1, 2, 2]]) + + >>> a = torch.triu_indices(4, 3, -1) + >>> a + tensor([[0, 0, 0, 1, 1, 1, 2, 2, 3], + [0, 1, 2, 0, 1, 2, 1, 2, 2]]) + + >>> a = torch.triu_indices(4, 3, 1) + >>> a + tensor([[0, 0, 1], + [1, 2, 2]]) + """ + +def true_divide( + input: Tensor | Number, + other: Tensor | Number, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + true_divide(dividend, divisor, *, out) -> Tensor + + Alias for :func:`torch.div` with ``rounding_mode=None``. + """ + +def trunc(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + trunc(input, *, out=None) -> Tensor + + Returns a new tensor with the truncated integer values of + the elements of :attr:`input`. + + For integer inputs, follows the array-api convention of returning a + copy of the input tensor. + + Args: + input (Tensor): the input tensor. + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 3.4742, 0.5466, -0.8008, -0.9079]) + >>> torch.trunc(a) + tensor([ 3., 0., -0., -0.]) + """ + +def trunc_(input: Tensor) -> Tensor: ... +@overload +def unbind(input: Tensor, dim: _int = 0) -> tuple[Tensor, ...]: + r""" + unbind(input, dim=0) -> seq + + Removes a tensor dimension. + + Returns a tuple of all slices along a given dimension, already without it. + + Arguments: + input (Tensor): the tensor to unbind + dim (int): dimension to remove + + Example:: + + >>> torch.unbind(torch.tensor([[1, 2, 3], + >>> [4, 5, 6], + >>> [7, 8, 9]])) + (tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9])) + """ + +@overload +def unbind( + input: Tensor, + dim: str | EllipsisType | None, +) -> tuple[Tensor, ...]: + r""" + unbind(input, dim=0) -> seq + + Removes a tensor dimension. + + Returns a tuple of all slices along a given dimension, already without it. + + Arguments: + input (Tensor): the tensor to unbind + dim (int): dimension to remove + + Example:: + + >>> torch.unbind(torch.tensor([[1, 2, 3], + >>> [4, 5, 6], + >>> [7, 8, 9]])) + (tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9])) + """ + +def unbind_copy( + input: Tensor, + dim: _int = 0, + *, + out: tuple[Tensor, ...] | list[Tensor] | None = None, +) -> None: + r""" + Performs the same operation as :func:`torch.unbind`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def unflatten( + input: Tensor, + dim: str | EllipsisType | None, + sizes: Sequence[_int | SymInt], + names: Sequence[str | EllipsisType | None], +) -> Tensor: + r""" + unflatten(input, dim, sizes) -> Tensor + + Expands a dimension of the input tensor over multiple dimensions. + + .. seealso:: + + :func:`torch.flatten` the inverse of this function. It coalesces several dimensions into one. + + Args: + input (Tensor): the input tensor. + dim (int): Dimension to be unflattened, specified as an index into + ``input.shape``. + sizes (Tuple[int]): New shape of the unflattened dimension. + One of its elements can be `-1` in which case the corresponding output + dimension is inferred. Otherwise, the product of ``sizes`` *must* + equal ``input.shape[dim]``. + + Returns: + A View of input with the specified dimension unflattened. + + Examples:: + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (2, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (-1, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(5, 12, 3), -2, (2, 2, 3, 1, 1)).shape + torch.Size([5, 2, 2, 3, 1, 1, 3]) + """ + +@overload +def unflatten( + input: Tensor, + dim: _int, + sizes: Sequence[_int | SymInt], +) -> Tensor: + r""" + unflatten(input, dim, sizes) -> Tensor + + Expands a dimension of the input tensor over multiple dimensions. + + .. seealso:: + + :func:`torch.flatten` the inverse of this function. It coalesces several dimensions into one. + + Args: + input (Tensor): the input tensor. + dim (int): Dimension to be unflattened, specified as an index into + ``input.shape``. + sizes (Tuple[int]): New shape of the unflattened dimension. + One of its elements can be `-1` in which case the corresponding output + dimension is inferred. Otherwise, the product of ``sizes`` *must* + equal ``input.shape[dim]``. + + Returns: + A View of input with the specified dimension unflattened. + + Examples:: + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (2, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (-1, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(5, 12, 3), -2, (2, 2, 3, 1, 1)).shape + torch.Size([5, 2, 2, 3, 1, 1, 3]) + """ + +def unfold_copy( + input: Tensor, + dimension: _int, + size: _int, + step: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.unfold`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def unique_dim( + input: Tensor, + dim: _int, + sorted: _bool = True, + return_inverse: _bool = False, + return_counts: _bool = False, +) -> tuple[Tensor, Tensor, Tensor]: ... +def unsafe_chunk( + input: Tensor, + chunks: _int, + dim: _int = 0, +) -> tuple[Tensor, ...]: + r""" + unsafe_chunk(input, chunks, dim=0) -> List of Tensors + + Works like :func:`torch.chunk` but without enforcing the autograd restrictions + on inplace modification of the outputs. + + .. warning:: + This function is safe to use as long as only the input, or only the outputs + are modified inplace after calling this function. It is user's + responsibility to ensure that is the case. If both the input and one or more + of the outputs are modified inplace, gradients computed by autograd will be + silently incorrect. + """ + +def unsafe_split( + input: Tensor, + split_size: _int | SymInt, + dim: _int = 0, +) -> tuple[Tensor, ...]: + r""" + unsafe_split(tensor, split_size_or_sections, dim=0) -> List of Tensors + + Works like :func:`torch.split` but without enforcing the autograd restrictions + on inplace modification of the outputs. + + .. warning:: + This function is safe to use as long as only the input, or only the outputs + are modified inplace after calling this function. It is user's + responsibility to ensure that is the case. If both the input and one or more + of the outputs are modified inplace, gradients computed by autograd will be + silently incorrect. + """ + +def unsafe_split_with_sizes( + input: Tensor, + split_sizes: Sequence[_int | SymInt], + dim: _int = 0, +) -> tuple[Tensor, ...]: ... +def unsqueeze(input: Tensor, dim: _int) -> Tensor: + r""" + unsqueeze(input, dim) -> Tensor + + Returns a new tensor with a dimension of size one inserted at the + specified position. + + The returned tensor shares the same underlying data with this tensor. + + A :attr:`dim` value within the range ``[-input.dim() - 1, input.dim() + 1)`` + can be used. Negative :attr:`dim` will correspond to :meth:`unsqueeze` + applied at :attr:`dim` = ``dim + input.dim() + 1``. + + Args: + input (Tensor): the input tensor. + dim (int): the index at which to insert the singleton dimension + + Example:: + + >>> x = torch.tensor([1, 2, 3, 4]) + >>> torch.unsqueeze(x, 0) + tensor([[ 1, 2, 3, 4]]) + >>> torch.unsqueeze(x, 1) + tensor([[ 1], + [ 2], + [ 3], + [ 4]]) + """ + +def unsqueeze_copy( + input: Tensor, + dim: _int, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.unsqueeze`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def values_copy(input: Tensor, *, out: Tensor | None = None) -> Tensor: + r""" + Performs the same operation as :func:`torch.values`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def vander( + x: Tensor, + N: _int | None = None, + increasing: _bool = False, +) -> Tensor: + r""" + vander(x, N=None, increasing=False) -> Tensor + + Generates a Vandermonde matrix. + + The columns of the output matrix are elementwise powers of the input vector :math:`x^{(N-1)}, x^{(N-2)}, ..., x^0`. + If increasing is True, the order of the columns is reversed :math:`x^0, x^1, ..., x^{(N-1)}`. Such a + matrix with a geometric progression in each row is named for Alexandre-Theophile Vandermonde. + + Arguments: + x (Tensor): 1-D input tensor. + N (int, optional): Number of columns in the output. If N is not specified, + a square array is returned :math:`(N = len(x))`. + increasing (bool, optional): Order of the powers of the columns. If True, + the powers increase from left to right, if False (the default) they are reversed. + + Returns: + Tensor: Vandermonde matrix. If increasing is False, the first column is :math:`x^{(N-1)}`, + the second :math:`x^{(N-2)}` and so forth. If increasing is True, the columns + are :math:`x^0, x^1, ..., x^{(N-1)}`. + + Example:: + + >>> x = torch.tensor([1, 2, 3, 5]) + >>> torch.vander(x) + tensor([[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]]) + >>> torch.vander(x, N=3) + tensor([[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]]) + >>> torch.vander(x, N=3, increasing=True) + tensor([[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]]) + """ + +@overload +def var( + input: Tensor, + dim: _int | _size | None, + unbiased: _bool = True, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` + can be a single dimension, list of dimensions, or ``None`` to reduce over all + dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var( + input: Tensor, + dim: _int | _size | None = None, + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` + can be a single dimension, list of dimensions, or ``None`` to reduce over all + dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var(input: Tensor, unbiased: _bool = True) -> Tensor: + r""" + var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` + can be a single dimension, list of dimensions, or ``None`` to reduce over all + dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, + out: Tensor | None = None, +) -> Tensor: + r""" + var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` + can be a single dimension, list of dimensions, or ``None`` to reduce over all + dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + unbiased: _bool = True, + keepdim: _bool = False, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + + Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` + can be a single dimension, list of dimensions, or ``None`` to reduce over all + dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var_mean( + input: Tensor, + dim: _int | _size | None, + unbiased: _bool = True, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the variance and mean over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (var, mean) containing the variance and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var_mean( + input: Tensor, + dim: _int | _size | None = None, + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the variance and mean over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (var, mean) containing the variance and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var_mean( + input: Tensor, + unbiased: _bool = True, +) -> tuple[Tensor, Tensor]: + r""" + var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the variance and mean over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (var, mean) containing the variance and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var_mean( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + *, + correction: Number | _complex | None = None, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the variance and mean over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (var, mean) containing the variance and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +@overload +def var_mean( + input: Tensor, + dim: Sequence[str | EllipsisType | None], + unbiased: _bool = True, + keepdim: _bool = False, +) -> tuple[Tensor, Tensor]: + r""" + var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + + Calculates the variance and mean over the dimensions specified by :attr:`dim`. + :attr:`dim` can be a single dimension, list of dimensions, or ``None`` to + reduce over all dimensions. + + The variance (:math:`\sigma^2`) is calculated as + + .. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + + where :math:`x` is the sample set of elements, :math:`\bar{x}` is the + sample mean, :math:`N` is the number of samples and :math:`\delta N` is + the :attr:`correction`. + + + + If :attr:`keepdim` is ``True``, the output tensor is of the same size + as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. + Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the + output tensor having 1 (or ``len(dim)``) fewer dimension(s). + + + Args: + input (Tensor): the input tensor. + + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. + + + Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. + + out (Tensor, optional): the output tensor. + + Returns: + A tuple (var, mean) containing the variance and mean. + + Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + + .. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + """ + +def vdot( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + vdot(input, other, *, out=None) -> Tensor + + Computes the dot product of two 1D vectors along a dimension. + + In symbols, this function computes + + .. math:: + + \sum_{i=1}^n \overline{x_i}y_i. + + where :math:`\overline{x_i}` denotes the conjugate for complex + vectors, and it is the identity for real vectors. + + .. note:: + + Unlike NumPy's vdot, torch.vdot intentionally only supports computing the dot product + of two 1D tensors with the same number of elements. + + .. seealso:: + + :func:`torch.linalg.vecdot` computes the dot product of two batches of vectors along a dimension. + + Args: + input (Tensor): first tensor in the dot product, must be 1D. Its conjugate is used if it's complex. + other (Tensor): second tensor in the dot product, must be 1D. + + Keyword args: + + .. note:: out (Tensor, optional): the output tensor. + + + Example:: + + >>> torch.vdot(torch.tensor([2, 3]), torch.tensor([2, 1])) + tensor(7) + >>> a = torch.tensor((1 +2j, 3 - 1j)) + >>> b = torch.tensor((2 +1j, 4 - 0j)) + >>> torch.vdot(a, b) + tensor([16.+1.j]) + >>> torch.vdot(b, a) + tensor([16.-1.j]) + """ + +def view_as_complex(input: Tensor) -> Tensor: + r""" + view_as_complex(input) -> Tensor + + Returns a view of :attr:`input` as a complex tensor. For an input complex + tensor of :attr:`size` :math:`m1, m2, \dots, mi, 2`, this function returns a + new complex tensor of :attr:`size` :math:`m1, m2, \dots, mi` where the last + dimension of the input tensor is expected to represent the real and imaginary + components of complex numbers. + + .. warning:: + :func:`view_as_complex` is only supported for tensors with + :class:`torch.dtype` ``torch.float64`` and ``torch.float32``. The input is + expected to have the last dimension of :attr:`size` 2. In addition, the + tensor must have a `stride` of 1 for its last dimension. The strides of all + other dimensions must be even numbers. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x=torch.randn(4, 2) + >>> x + tensor([[ 1.6116, -0.5772], + [-1.4606, -0.9120], + [ 0.0786, -1.7497], + [-0.6561, -1.6623]]) + >>> torch.view_as_complex(x) + tensor([(1.6116-0.5772j), (-1.4606-0.9120j), (0.0786-1.7497j), (-0.6561-1.6623j)]) + """ + +def view_as_complex_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.view_as_complex`, but all output tensors + are freshly created instead of aliasing the input. + """ + +def view_as_real(input: Tensor) -> Tensor: + r""" + view_as_real(input) -> Tensor + + Returns a view of :attr:`input` as a real tensor. For an input complex tensor of + :attr:`size` :math:`m1, m2, \dots, mi`, this function returns a new + real tensor of size :math:`m1, m2, \dots, mi, 2`, where the last dimension of size 2 + represents the real and imaginary components of complex numbers. + + .. warning:: + :func:`view_as_real` is only supported for tensors with ``complex dtypes``. + + Args: + input (Tensor): the input tensor. + + Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.4737-0.3839j), (-0.2098-0.6699j), (0.3470-0.9451j), (-0.5174-1.3136j)]) + >>> torch.view_as_real(x) + tensor([[ 0.4737, -0.3839], + [-0.2098, -0.6699], + [ 0.3470, -0.9451], + [-0.5174, -1.3136]]) + """ + +def view_as_real_copy( + input: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.view_as_real`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def view_copy( + input: Tensor, + dtype: _dtype, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.view`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def view_copy( + input: Tensor, + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, +) -> Tensor: + r""" + Performs the same operation as :func:`torch.view`, but all output tensors + are freshly created instead of aliasing the input. + """ + +@overload +def vsplit(input: Tensor, sections: _int) -> tuple[Tensor, ...]: + r""" + vsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with two or more dimensions, into multiple tensors + vertically according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=0) + (the split dimension is 0), except that if :attr:`indices_or_sections` is an integer + it must evenly divide the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.vsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.vsplit(t, 2) + (tensor([[0., 1., 2., 3.], + [4., 5., 6., 7.]]), + tensor([[ 8., 9., 10., 11.], + [12., 13., 14., 15.]])) + >>> torch.vsplit(t, [3, 6]) + (tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]]), + tensor([[12., 13., 14., 15.]]), + tensor([], size=(0, 4))) + """ + +@overload +def vsplit(input: Tensor, indices: _size) -> tuple[Tensor, ...]: + r""" + vsplit(input, indices_or_sections) -> List of Tensors + + Splits :attr:`input`, a tensor with two or more dimensions, into multiple tensors + vertically according to :attr:`indices_or_sections`. Each split is a view of + :attr:`input`. + + This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=0) + (the split dimension is 0), except that if :attr:`indices_or_sections` is an integer + it must evenly divide the split dimension or a runtime error will be thrown. + + This function is based on NumPy's :func:`numpy.vsplit`. + + Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + + Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.vsplit(t, 2) + (tensor([[0., 1., 2., 3.], + [4., 5., 6., 7.]]), + tensor([[ 8., 9., 10., 11.], + [12., 13., 14., 15.]])) + >>> torch.vsplit(t, [3, 6]) + (tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]]), + tensor([[12., 13., 14., 15.]]), + tensor([], size=(0, 4))) + """ + +def vstack( + tensors: tuple[Tensor, ...] | list[Tensor] | None, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + vstack(tensors, *, out=None) -> Tensor + + Stack tensors in sequence vertically (row wise). + + This is equivalent to concatenation along the first axis after all 1-D tensors have been reshaped by :func:`torch.atleast_2d`. + + Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + + Keyword args: + out (Tensor, optional): the output tensor. + + Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.vstack((a,b)) + tensor([[1, 2, 3], + [4, 5, 6]]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.vstack((a,b)) + tensor([[1], + [2], + [3], + [4], + [5], + [6]]) + """ + +@overload +def where(condition: Tensor) -> tuple[Tensor, ...]: + r""" + where(condition, input, other, *, out=None) -> Tensor + + Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + + The operation is defined as: + + .. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} + + .. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + + Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + + Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + + .. function:: where(condition) -> tuple of LongTensor + :noindex: + + ``torch.where(condition)`` is identical to + ``torch.nonzero(condition, as_tuple=True)``. + + .. note:: + See also :func:`torch.nonzero`. + """ + +@overload +def where( + condition: Tensor, + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + where(condition, input, other, *, out=None) -> Tensor + + Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + + The operation is defined as: + + .. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} + + .. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + + Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + + Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + + .. function:: where(condition) -> tuple of LongTensor + :noindex: + + ``torch.where(condition)`` is identical to + ``torch.nonzero(condition, as_tuple=True)``. + + .. note:: + See also :func:`torch.nonzero`. + """ + +@overload +def where( + condition: Tensor, + self: Number | _complex, + other: Tensor, +) -> Tensor: + r""" + where(condition, input, other, *, out=None) -> Tensor + + Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + + The operation is defined as: + + .. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} + + .. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + + Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + + Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + + .. function:: where(condition) -> tuple of LongTensor + :noindex: + + ``torch.where(condition)`` is identical to + ``torch.nonzero(condition, as_tuple=True)``. + + .. note:: + See also :func:`torch.nonzero`. + """ + +@overload +def where( + condition: Tensor, + input: Tensor, + other: Number | _complex, +) -> Tensor: + r""" + where(condition, input, other, *, out=None) -> Tensor + + Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + + The operation is defined as: + + .. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} + + .. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + + Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + + Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + + .. function:: where(condition) -> tuple of LongTensor + :noindex: + + ``torch.where(condition)`` is identical to + ``torch.nonzero(condition, as_tuple=True)``. + + .. note:: + See also :func:`torch.nonzero`. + """ + +@overload +def where( + condition: Tensor, + self: Number | _complex, + other: Number | _complex, +) -> Tensor: + r""" + where(condition, input, other, *, out=None) -> Tensor + + Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + + The operation is defined as: + + .. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} + + .. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + + Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + + Keyword args: + out (Tensor, optional): the output tensor. + + Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + + Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + + .. function:: where(condition) -> tuple of LongTensor + :noindex: + + ``torch.where(condition)`` is identical to + ``torch.nonzero(condition, as_tuple=True)``. + + .. note:: + See also :func:`torch.nonzero`. + """ + +@overload +def xlogy( + input: Tensor, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + xlogy(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.special.xlogy`. + """ + +@overload +def xlogy( + self: Number | _complex, + other: Tensor, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + xlogy(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.special.xlogy`. + """ + +@overload +def xlogy( + input: Tensor, + other: Number | _complex, + *, + out: Tensor | None = None, +) -> Tensor: + r""" + xlogy(input, other, *, out=None) -> Tensor + + Alias for :func:`torch.special.xlogy`. + """ + +@overload +def xlogy_(input: Tensor, other: Tensor) -> Tensor: ... +@overload +def xlogy_(input: Tensor, other: Number | _complex) -> Tensor: ... +def zero_(input: Tensor) -> Tensor: ... +@overload +def zeros( + size: Sequence[_int | SymInt], + *, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + zeros(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `0`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.zeros(2, 3) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + + >>> torch.zeros(5) + tensor([ 0., 0., 0., 0., 0.]) + """ + +@overload +def zeros( + *size: _int | SymInt, + out: Tensor | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + zeros(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `0`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.zeros(2, 3) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + + >>> torch.zeros(5) + tensor([ 0., 0., 0., 0., 0.]) + """ + +@overload +def zeros( + size: _size, + *, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + zeros(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `0`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.zeros(2, 3) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + + >>> torch.zeros(5) + tensor([ 0., 0., 0., 0., 0.]) + """ + +@overload +def zeros( + *size: _int, + names: Sequence[str | EllipsisType | None] | None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + zeros(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + + Returns a tensor filled with the scalar value `0`, with the shape defined + by the variable argument :attr:`size`. + + Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + + Keyword args: + out (Tensor, optional): the output tensor. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + + Example:: + + >>> torch.zeros(2, 3) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + + >>> torch.zeros(5) + tensor([ 0., 0., 0., 0., 0.]) + """ + +def zeros_like( + input: Tensor, + *, + memory_format: memory_format | None = None, + dtype: _dtype | None = None, + layout: _layout | None = None, + device: DeviceLikeType | None = None, + pin_memory: _bool | None = False, + requires_grad: _bool | None = False, +) -> Tensor: + r""" + zeros_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + + Returns a tensor filled with the scalar value `0`, with the same size as + :attr:`input`. ``torch.zeros_like(input)`` is equivalent to + ``torch.zeros(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + + .. warning:: + As of 0.4, this function does not support an :attr:`out` keyword. As an alternative, + the old ``torch.zeros_like(input, out=output)`` is equivalent to + ``torch.zeros(input.size(), out=output)``. + + Args: + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + + Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. + + Example:: + + >>> input = torch.empty(2, 3) + >>> torch.zeros_like(input) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + """ diff --git a/venv/lib/python3.10/site-packages/torch/__config__.py b/venv/lib/python3.10/site-packages/torch/__config__.py new file mode 100644 index 0000000000000000000000000000000000000000..1187fab3713996932d4f3bad71bac243c6baff35 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/__config__.py @@ -0,0 +1,22 @@ +import torch + + +def show() -> str: + """ + Return a human-readable string with descriptions of the + configuration of PyTorch. + """ + return torch._C._show_config() + + +# TODO: In principle, we could provide more structured version/config +# information here. For now only CXX_FLAGS is exposed, as Timer +# uses them. +def _cxx_flags() -> str: + """Returns the CXX_FLAGS used when building PyTorch.""" + return torch._C._cxx_flags() + + +def parallel_info() -> str: + r"""Returns detailed string with parallelization settings""" + return torch._C._parallel_info() diff --git a/venv/lib/python3.10/site-packages/torch/__future__.py b/venv/lib/python3.10/site-packages/torch/__future__.py new file mode 100644 index 0000000000000000000000000000000000000000..f172ee3c8fe223aa316667f37f356e5b6658d20e --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/__future__.py @@ -0,0 +1,75 @@ +_overwrite_module_params_on_conversion: bool = False +_swap_module_params_on_conversion: bool = False + + +def set_overwrite_module_params_on_conversion(value: bool) -> None: + """ + Sets whether to assign new tensors to the parameters instead of changing the + existing parameters in-place when converting an ``nn.Module``. + + When enabled, the following methods will assign new parameters to the module: + + #. ``module.{device}()`` (e.g. :meth:`nn.Module.cuda()`) for moving a module between devices + #. ``module.{dtype}()`` (e.g. :meth:`nn.Module.float()`) for converting a module to a different dtype + #. :meth:`nn.Module.to` + #. :meth:`nn.Module.to_empty` + + Args: + value (bool): Whether to assign new tensors or not. + + """ + global _overwrite_module_params_on_conversion + _overwrite_module_params_on_conversion = value + + +def get_overwrite_module_params_on_conversion() -> bool: + """ + Returns whether to assign new tensors to the parameters instead of changing the + existing parameters in-place when converting an :class:`torch.nn.Module`. Defaults to ``False``. + + See :func:`~torch.__future__.set_overwrite_module_params_on_conversion` for more information. + """ + return _overwrite_module_params_on_conversion + + +def set_swap_module_params_on_conversion(value: bool) -> None: + """ + Sets whether to use :func:`~torch.utils.swap_tensors` instead of setting ``.data`` to + change the existing parameters in-place when converting an ``nn.Module`` and instead + of ``param.copy_(state_dict[key])`` when loading a state dict into an ``nn.Module``. + + .. note:: + This function takes precedence over :func:`~torch.__future__.get_overwrite_module_params_on_conversion` + + When enabled, the following methods will swap the existing parameters in-place: + + #. ``module.{device}()`` (e.g. :meth:`nn.Module.cuda()`) for moving a module between devices + #. ``module.{dtype}()`` (e.g. :meth:`nn.Module.float()`) for converting a module to a different dtype + #. :meth:`nn.Module.to` + #. :meth:`nn.Module.to_empty` + #. :meth:`nn.Module.load_state_dict` + + The semantics for :meth:`~nn.Module.load_state_dict` when this is set are as follows: + + #. For each parameter/buffer, its corresponding ``state_dict['key']`` is transformed via + :meth:`~torch.Tensor.module_load` (i.e. ``res = param.module_load(state_dict['key'])``) + #. If necessary, ``res`` will be wrapped in an :class:`~nn.Parameter` + #. The parameter/buffer in the module will be swapped via :func:`~torch.utils.swap_tensors` + with ``res`` + + Args: + value (bool): Whether to use :func:`~torch.utils.swap_tensors` or not. + + """ + global _swap_module_params_on_conversion + _swap_module_params_on_conversion = value + + +def get_swap_module_params_on_conversion() -> bool: + """ + Returns whether to use :func:`~torch.utils.swap_tensors` instead of setting .data to + change the existing parameters in-place when converting an ``nn.Module``. Defaults to ``False``. + + See :func:`~torch.__future__.set_swap_module_params_on_conversion` for more information. + """ + return _swap_module_params_on_conversion diff --git a/venv/lib/python3.10/site-packages/torch/__init__.py b/venv/lib/python3.10/site-packages/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0625ad60bfff6b0f0f3be06306f7bd937e5031ac --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/__init__.py @@ -0,0 +1,2878 @@ +""" +The torch package contains data structures for multi-dimensional +tensors and defines mathematical operations over these tensors. +Additionally, it provides many utilities for efficient serialization of +Tensors and arbitrary types, and other useful utilities. + +It has a CUDA counterpart, that enables you to run your tensor computations +on an NVIDIA GPU with compute capability >= 3.0. +""" + +# mypy: allow-untyped-defs + +import builtins +import ctypes +import functools +import glob +import importlib +import inspect +import math +import os +import platform +import sys +import textwrap +import threading +from typing import ( + Any as _Any, + Callable as _Callable, + get_origin as _get_origin, + Optional as _Optional, + overload as _overload, + TYPE_CHECKING, + TypeVar as _TypeVar, + Union as _Union, +) +from typing_extensions import ParamSpec as _ParamSpec, TypeIs as _TypeIs + + +# As a bunch of torch.packages internally still have this check +# we need to keep this. @todo: Remove tests that rely on this check as +# they are likely stale. +def _running_with_deploy() -> builtins.bool: + return False + + +from torch._utils import ( + _functionalize_sync as _sync, + _import_dotted_name, + classproperty, +) +from torch._utils_internal import ( + get_file_path, + prepare_multiprocessing_environment, + profiler_allow_cudagraph_cupti_lazy_reinit_cuda12, + USE_GLOBAL_DEPS, + USE_RTLD_GLOBAL_WITH_LIBTORCH, +) +from torch.torch_version import __version__ as __version__ + + +if TYPE_CHECKING: + from torch.types import Device, IntLikeType + + +__all__ = [ + "BoolStorage", + "BoolTensor", + "ByteStorage", + "ByteTensor", + "CharStorage", + "CharTensor", + "DoubleStorage", + "DoubleTensor", + "FloatStorage", + "FloatTensor", + "GradScaler", + "IntStorage", + "IntTensor", + "LongStorage", + "LongTensor", + "ShortStorage", + "ShortTensor", + "SymBool", + "SymFloat", + "SymInt", + "Tensor", + "TypedStorage", + "UntypedStorage", + "are_deterministic_algorithms_enabled", + "autocast", + "chunk", + "compile", + "cond", + "enable_grad", + "export", + "get_default_device", + "get_deterministic_debug_mode", + "get_device_module", + "get_float32_matmul_precision", + "get_rng_state", + "inference_mode", + "initial_seed", + "is_deterministic_algorithms_warn_only_enabled", + "is_storage", + "is_tensor", + "is_warn_always_enabled", + "load", + "lobpcg", + "manual_seed", + "matmul", + "no_grad", + "rand", + "randn", + "save", + "seed", + "set_default_device", + "set_default_tensor_type", + "set_deterministic_debug_mode", + "set_float32_matmul_precision", + "set_printoptions", + "set_rng_state", + "set_warn_always", + "split", + "stack", + "sym_float", + "sym_fresh_size", + "sym_int", + "sym_ite", + "sym_max", + "sym_min", + "sym_not", + "sym_sum", + "typename", + "unravel_index", + "use_deterministic_algorithms", + "vmap", +] + +# Please keep this list sorted +assert __all__ == sorted(__all__) + +################################################################################ +# Load the extension module +################################################################################ + +# If PyTorch was built against the ROCm runtime wheels, then there will be +# a _rocm_init module and it will define an initialize() function which can +# prepare ROCm for use. See general documentation on ROCm runtime wheels: +# https://github.com/ROCm/TheRock/blob/main/docs/packaging/python_packaging.md +# Since this module is only ever added to the wheel if built for such a +# deployment, it is always safe to attempt. +try: + from . import _rocm_init # type: ignore[attr-defined] +except ImportError: + pass +else: + _rocm_init.initialize() + del _rocm_init + + +if sys.platform == "win32": + + def _load_dll_libraries() -> None: + import sysconfig + + from torch.version import cuda as cuda_version + + pfiles_path = os.getenv("ProgramFiles", r"C:\Program Files") + py_dll_path = os.path.join(sys.exec_prefix, "Library", "bin") + th_dll_path = os.path.join(os.path.dirname(__file__), "lib") + usebase_path = os.path.join( + sysconfig.get_config_var("userbase"), "Library", "bin" + ) + py_root_bin_path = os.path.join(sys.exec_prefix, "bin") + + # When users create a virtualenv that inherits the base environment, + # we will need to add the corresponding library directory into + # DLL search directories. Otherwise, it will rely on `PATH` which + # is dependent on user settings. + if sys.exec_prefix != sys.base_exec_prefix: + base_py_dll_path = os.path.join(sys.base_exec_prefix, "Library", "bin") + else: + base_py_dll_path = "" + + dll_paths = [ + p + for p in ( + th_dll_path, + py_dll_path, + base_py_dll_path, + usebase_path, + py_root_bin_path, + ) + if os.path.exists(p) + ] + + if not builtins.any( + os.path.exists(os.path.join(p, "nvToolsExt64_1.dll")) for p in dll_paths + ): + nvtoolsext_dll_path = os.path.join( + os.getenv( + "NVTOOLSEXT_PATH", + os.path.join(pfiles_path, "NVIDIA Corporation", "NvToolsExt"), + ), + "bin", + "x64", + ) + else: + nvtoolsext_dll_path = "" + + if cuda_version and builtins.all( + not glob.glob(os.path.join(p, "cudart64*.dll")) for p in dll_paths + ): + cuda_version_1 = cuda_version.replace(".", "_") + cuda_path_var = "CUDA_PATH_V" + cuda_version_1 + default_path = os.path.join( + pfiles_path, "NVIDIA GPU Computing Toolkit", "CUDA", f"v{cuda_version}" + ) + cuda_path = os.path.join(os.getenv(cuda_path_var, default_path), "bin") + else: + cuda_path = "" + + dll_paths.extend( + p for p in (nvtoolsext_dll_path, cuda_path) if os.path.exists(p) + ) + + kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) + with_load_library_flags = hasattr(kernel32, "AddDllDirectory") + prev_error_mode = kernel32.SetErrorMode(0x0001) + + kernel32.LoadLibraryW.restype = ctypes.c_void_p + if with_load_library_flags: + kernel32.LoadLibraryExW.restype = ctypes.c_void_p + + for dll_path in dll_paths: + os.add_dll_directory(dll_path) + + try: + ctypes.CDLL("vcruntime140.dll") + ctypes.CDLL("msvcp140.dll") + if platform.machine() != "ARM64": + ctypes.CDLL("vcruntime140_1.dll") + except OSError: + print( + textwrap.dedent( + """ + Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure. + It can be downloaded at https://aka.ms/vs/17/release/vc_redist.x64.exe + """ + ).strip() + ) + + dlls = glob.glob(os.path.join(th_dll_path, "*.dll")) + path_patched = False + for dll in dlls: + is_loaded = False + if with_load_library_flags: + res = kernel32.LoadLibraryExW(dll, None, 0x00001100) + last_error = ctypes.get_last_error() + if res is None and last_error != 126: + err = ctypes.WinError(last_error) + err.strerror += ( + f' Error loading "{dll}" or one of its dependencies.' + ) + raise err + elif res is not None: + is_loaded = True + if not is_loaded: + if not path_patched: + os.environ["PATH"] = ";".join(dll_paths + [os.environ["PATH"]]) + path_patched = True + res = kernel32.LoadLibraryW(dll) + if res is None: + err = ctypes.WinError(ctypes.get_last_error()) + err.strerror += ( + f' Error loading "{dll}" or one of its dependencies.' + ) + raise err + + kernel32.SetErrorMode(prev_error_mode) + + _load_dll_libraries() + del _load_dll_libraries + + +def _get_cuda_dep_paths(path: str, lib_folder: str, lib_name: str) -> list[str]: + # Libraries can either be in + # path/nvidia/lib_folder/lib or + # path/nvidia/cuXX/lib (since CUDA 13.0) or + # path/lib_folder/lib + from torch.version import cuda as cuda_version + + nvidia_lib_paths = glob.glob( + os.path.join(path, "nvidia", lib_folder, "lib", lib_name) + ) + if cuda_version is not None: + maj_cuda_version = cuda_version.split(".")[0] + nvidia_lib_paths += glob.glob( + os.path.join(path, "nvidia", f"cu{maj_cuda_version}", "lib", lib_name) + ) + lib_paths = glob.glob(os.path.join(path, lib_folder, "lib", lib_name)) + + return nvidia_lib_paths + lib_paths + + +def _preload_cuda_deps(lib_folder: str, lib_name: str, required: bool = True) -> None: # type: ignore[valid-type] + """Preloads cuda deps if they could not be found otherwise.""" + # Should only be called on Linux if default path resolution have failed + assert platform.system() == "Linux", "Should only be called on Linux" + + lib_path = None + for path in sys.path: + candidate_lib_paths = _get_cuda_dep_paths(path, lib_folder, lib_name) + if candidate_lib_paths: + lib_path = candidate_lib_paths[0] + break + if not lib_path and required: + raise ValueError(f"{lib_name} not found in the system path {sys.path}") + if lib_path: + ctypes.CDLL(lib_path) + + +# See Note [Global dependencies] +def _load_global_deps() -> None: + if platform.system() == "Windows": + return + + # Determine the file extension based on the platform + lib_ext = ".dylib" if platform.system() == "Darwin" else ".so" + lib_name = f"libtorch_global_deps{lib_ext}" + here = os.path.abspath(__file__) + global_deps_lib_path = os.path.join(os.path.dirname(here), "lib", lib_name) + + try: + ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL) + # Workaround slim-wheel CUDA dependency bugs in cusparse and cudnn by preloading nvjitlink + # and nvrtc. In CUDA-12.4+ cusparse depends on nvjitlink, but does not have rpath when + # shipped as wheel, which results in OS picking wrong/older version of nvjitlink library + # if `LD_LIBRARY_PATH` is defined, see https://github.com/pytorch/pytorch/issues/138460 + # Similar issue exist in cudnn that dynamically loads nvrtc, unaware of its relative path. + # See https://github.com/pytorch/pytorch/issues/145580 + try: + with open("/proc/self/maps") as f: + _maps = f.read() + + # libtorch_global_deps.so always depends in cudart, check if its installed and loaded + if "libcudart.so" not in _maps: + return + # If all above-mentioned conditions are met, preload nvrtc and nvjitlink + _preload_cuda_deps("cuda_nvrtc", "libnvrtc.so.*[0-9]") + _preload_cuda_deps("cuda_nvrtc", "libnvrtc-builtins.so.*[0-9]") + _preload_cuda_deps("nvjitlink", "libnvJitLink.so.*[0-9]") + except Exception: + pass + + except OSError as err: + # Can only happen for wheel with cuda libs as PYPI deps + # As PyTorch is not purelib, but nvidia-*-cu12 is + cuda_libs: dict[str, str] = { + "cublas": "libcublas.so.*[0-9]", + "cudnn": "libcudnn.so.*[0-9]", + "cuda_nvrtc": "libnvrtc.so.*[0-9]", + "cuda_runtime": "libcudart.so.*[0-9]", + "cuda_cupti": "libcupti.so.*[0-9]", + "cufft": "libcufft.so.*[0-9]", + "curand": "libcurand.so.*[0-9]", + "nvjitlink": "libnvJitLink.so.*[0-9]", + "cusparse": "libcusparse.so.*[0-9]", + "cusparselt": "libcusparseLt.so.*[0-9]", + "cusolver": "libcusolver.so.*[0-9]", + "nccl": "libnccl.so.*[0-9]", + "nvshmem": "libnvshmem_host.so.*[0-9]", + "cufile": "libcufile.so.*[0-9]", + } + + is_cuda_lib_err = [ + lib for lib in cuda_libs.values() if lib.split(".")[0] in err.args[0] + ] + if not is_cuda_lib_err: + raise err + for lib_folder, lib_name in cuda_libs.items(): + _preload_cuda_deps(lib_folder, lib_name) + + # libnvToolsExt is Optional Dependency + _preload_cuda_deps("nvtx", "libnvToolsExt.so.*[0-9]", required=False) + ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL) + + +if (USE_RTLD_GLOBAL_WITH_LIBTORCH or os.getenv("TORCH_USE_RTLD_GLOBAL")) and ( + platform.system() != "Windows" +): + # Do it the hard way. You might want to load libtorch with RTLD_GLOBAL in a + # few circumstances: + # + # 1. You're in a build environment (e.g., fbcode) where + # libtorch_global_deps is not available, but you still need + # to get mkl to link in with RTLD_GLOBAL or it will just + # not work. + # + # 2. You're trying to run PyTorch under UBSAN and you need + # to ensure that only one copy of libtorch is loaded, so + # vptr checks work properly + # + # If you're using this setting, you must verify that all the libraries + # you load consistently use the same libstdc++, or you may have + # mysterious segfaults. + # + old_flags = sys.getdlopenflags() + sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_LAZY) + + from torch._C import * # noqa: F403 + + sys.setdlopenflags(old_flags) + del old_flags + +else: + # Easy way. You want this most of the time, because it will prevent + # C++ symbols from libtorch clobbering C++ symbols from other + # libraries, leading to mysterious segfaults. + # + # If building in an environment where libtorch_global_deps isn't available + # like parts of fbsource, but where RTLD_GLOBAL causes segfaults, you will + # want USE_RTLD_GLOBAL_WITH_LIBTORCH = False and USE_GLOBAL_DEPS = False + # + # See Note [Global dependencies] + if USE_GLOBAL_DEPS: + _load_global_deps() + from torch._C import * # noqa: F403 + + +class SymInt: + """ + Like an int (including magic methods), but redirects all operations on the + wrapped node. This is used in particular to symbolically record operations + in the symbolic shape workflow. + """ + + def __init__(self, node): + # This field MUST be named node; C++ binding code assumes that this + # class has a field named node that stores SymNode + self.node = node + + def __bool__(self): + return builtins.bool(self != 0) + + def __int__(self): + return self.node.int_() + + def __index__(self): + return self.node.int_() + + # Magic methods installed by torch.fx.experimental.sym_node + + def __round__(self, ndigits=None): + return self + + def __truediv__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(self).__float_truediv__(other) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + return self.__int_truediv__(other) + + def __rtruediv__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(self).__rfloat_truediv__(other) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + return self.__rint_truediv__(other) + + def __floordiv__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(math.floor(sym_float(self) / other)) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + return self.__int_floordiv__(other) + + def __rfloordiv__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(math.floor(other / sym_float(self))) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + return self.__rint_floordiv__(other) + + # nb: complex is impossible to handle correctly lol, with + # negative base and integral float need to diverge semantics and + # just always return complex. Neener neener pretend this problem + # doesn't exist + def __pow__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(self).__pow__(other) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + # Guards! This guard is necessary because we need to know it to + # determine the output type of this operation + if other >= 0: + return self.__pow_by_natural__(other) + else: + # Mercifully, when the exponent is negative, Python just promotes + # to doubles and does a float pow: + # + # if (Py_SIZE(b) < 0 && c == NULL) { + # /* if exponent is negative and there's no modulus: + # return a float. This works because we know + # that this calls float_pow() which converts its + # arguments to double. */ + # Py_DECREF(a); + # Py_DECREF(b); + # return PyFloat_Type.tp_as_number->nb_power(v, w, x); + # } + return sym_float(self).__pow__(sym_float(other)) + + def __rpow__(self, other): + if isinstance(other, (builtins.float, SymFloat)): + return sym_float(self).__rpow__(other) + if not isinstance(other, (builtins.int, SymInt)): + return NotImplemented + if self >= 0: # self is exponent + return self.__rpow_by_natural__(other) + else: + return sym_float(self).__rpow__(sym_float(other)) + + def __eq__(self, other: object) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __lt__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __gt__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __le__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __ge__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __add__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __radd__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __rmul__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __mod__(self, other: "IntLikeType") -> "SymInt": + raise TypeError("type stub not overridden") + + def __mul__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __pow_by_natural__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __rpow_by_natural__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __int_truediv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __rint_truediv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __int_floordiv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __rint_floordiv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __sym_max__(self, other): + raise TypeError("type stub not overridden") + + def __sym_min__(self, other): + raise TypeError("type stub not overridden") + + def __sym_float__(self): + raise TypeError("type stub not overridden") + + def __neg__(self): + raise TypeError("type stub not overridden") + + def __sub__(self, other: "IntLikeType") -> "SymInt": + raise TypeError("type stub not overridden") + + def __rsub__(self, other: "IntLikeType") -> "SymInt": + raise TypeError("type stub not overridden") + + def __and__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __or__(self, other) -> "SymInt": + raise TypeError("type stub not overridden") + + def __repr__(self): + return self.node._graph_repr() + + def _sympy_(self): + return self.node.expr + + def __hash__(self) -> builtins.int: + if self.node.is_nested_int(): + return hash(self.node.nested_int()) + else: + # We could support constant SymInts as well, but not doing it for now + raise TypeError("unhashable type: non-nested SymInt") + # TODO: Force specialization + # This can't be done because the TypeError here is load bearing + # for einops + # https://github.com/arogozhnikov/einops/blob/6181e1e95dc58c00a3143c1726da1c6ee0463164/einops/einops.py#L237 + # return hash(builtins.int(self)) + + def as_integer_ratio(self) -> tuple["SymInt", builtins.int]: + """Represent this int as an exact integer ratio""" + return self, 1 + + def bit_length(self) -> builtins.int: + # TODO: A more relaxed guard is possible here, where you guard to + # allow all integer quantities which would result in the same bit + # length. We can also just make a dedicated Sympy function for + # computing this quantity and represent it symbolically. + return builtins.int(self).bit_length() + + def conjugate(self) -> "SymInt": + return self + + +class SymFloat: + """ + Like a float (including magic methods), but redirects all operations on the + wrapped node. This is used in particular to symbolically record operations + in the symbolic shape workflow. + """ + + def __init__(self, node): + # This field MUST be named node; C++ binding code assumes that this + # class has a field named node that stores SymNode + self.node = node + + def __truediv__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + return self.__float_truediv__(sym_float(other)) + + def __rtruediv__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + return self.__rfloat_truediv__(sym_float(other)) + + def __floordiv__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + return sym_float(math.floor(self / sym_float(other))) + + def __rfloordiv__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + return sym_float(math.floor(sym_float(other) / self)) + + def __bool__(self): + return self.node.bool_() + + def __float__(self): + return self.node.guard_float("", 0) + + # Symbolic power does NOT work with negative base, this is to avoid + # potential complex outputs + def __pow__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + torch._check(self >= 0) + return self.__float_pow__(other) + + def __rpow__(self, other): + if not isinstance(other, (builtins.int, builtins.float, SymInt, SymFloat)): + return NotImplemented + torch._check(other >= 0) + return self.__rfloat_pow__(other) + + # Magic methods installed by torch.fx.experimental.sym_node + + def __eq__(self, other: object) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __lt__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __gt__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __le__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __ge__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __float_pow__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __rfloat_pow__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __float_truediv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __rfloat_truediv__(self, other) -> "SymFloat": + raise TypeError("type stub not overridden") + + def __trunc__(self): + raise TypeError("type stub not overridden") + + def __sym_max__(self, other): + raise TypeError("type stub not overridden") + + def __sym_min__(self, other): + raise TypeError("type stub not overridden") + + def __sym_int__(self): + raise TypeError("type stub not overridden") + + def is_integer(self): + """Return True if the float is an integer.""" + raise TypeError("type stub not overridden") + + def as_integer_ratio(self) -> tuple[builtins.int, builtins.int]: + """Represent this float as an exact integer ratio""" + return builtins.float(self).as_integer_ratio() + + def __repr__(self): + return self.node._graph_repr() + + def _sympy_(self): + return self.node.expr + + def __hash__(self): + return hash(builtins.float(self)) + + def conjugate(self) -> "SymFloat": + """Returns the complex conjugate of the float.""" + return self + + def hex(self) -> str: + """Returns the hexadecimal representation of the float.""" + return self.node.guard_float("", 0).hex() + + +class SymBool: + """ + Like a bool (including magic methods), but redirects all operations on the + wrapped node. This is used in particular to symbolically record operations + in the symbolic shape workflow. + + Unlike regular bools, regular boolean operators will force extra guards instead + of symbolically evaluate. Use the bitwise operators instead to handle this. + """ + + def __init__(self, node): + # This field MUST be named node; C++ binding code assumes that this + # class has a field named node that stores SymNode + self.node = node + + def __bool__(self): + return self.node.bool_() + + def __int__(self): + return builtins.int(self.node.bool_()) + + # Magic methods installed by torch.fx.experimental.sym_node + def __and__(self, other) -> "SymBool": + raise TypeError("type stub not overridden") + + def __or__(self, other) -> "SymBool": + raise TypeError("type stub not overridden") + + # We very carefully define __sym_not__, and not a number of other + # plausible alternatives: + # + # - We do not override __not__ because this is not a real magic + # method; you cannot override the meaning of the not builtin in + # Python. We use the name 'sym_not' to clarify that in user code you + # cannot use the builtin not or operator.not_ or operator.__not__ and + # hit this magic method; you must use our custom sym_not operator. + # + # - We do not override the __invert__ method because SymBool is + # meant to be usable in situations where bool is expected. However, + # bitwise negation ~a does the wrong thing with booleans (because + # bool is a subclass of int, so ~1 = -2 which is not falseish.) + # This would be a giant footgun, so we get around it by defining + # our own operator. Note that bitwise and/or do the right thing, + # so we reuse the conventional operators there for readability. + # + def __sym_not__(self) -> "SymBool": + raise TypeError("type stub not overridden") + + def __sym_ite__(self, then_val, else_val): + raise TypeError("type stub not overridden") + + def __eq__(self, other) -> builtins.bool: + raise TypeError("type stub not overridden") + + def __repr__(self): + return self.node._graph_repr() + + def _sympy_(self): + return self.node.expr + + def __hash__(self): + if self.node.is_constant(): + return hash(self.node.bool_()) + else: + # Force specialization + return hash(builtins.bool(self)) + + +def sym_not(a): + r"""SymInt-aware utility for logical negation. + + Args: + a (SymBool or bool): Object to negate + """ + import sympy + + if overrides.has_torch_function_unary(a): + return overrides.handle_torch_function(sym_not, (a,), a) + if hasattr(a, "__sym_not__"): + return a.__sym_not__() + if isinstance(a, sympy.Basic): + return ~a # type: ignore[operator] + return not a + + +def sym_float(a): + r"""SymInt-aware utility for float casting. + + Args: + a (SymInt, SymFloat, or object): Object to cast + """ + if overrides.has_torch_function_unary(a): + return overrides.handle_torch_function(sym_float, (a,), a) + if isinstance(a, SymFloat): + return a + elif hasattr(a, "__sym_float__"): + return a.__sym_float__() + return builtins.float(a) # type: ignore[operator] + + +def sym_int(a): + r"""SymInt-aware utility for int casting. + + Args: + a (SymInt, SymFloat, or object): Object to cast + """ + if overrides.has_torch_function_unary(a): + return overrides.handle_torch_function(sym_int, (a,), a) + if isinstance(a, SymInt): + return a + elif isinstance(a, SymFloat): + return math.trunc(a) + return builtins.int(a) # type: ignore[operator] + + +def sym_max(a, b): + """ + SymInt-aware utility for max which avoids branching on a < b. + Unlike builtins.max(), this only works for int/float, and it always + promotes to float if any argument is float (unlike builtins.max, which + will faithfully preserve the type of the input argument). + """ + if overrides.has_torch_function((a, b)): + return overrides.handle_torch_function(sym_max, (a, b), a, b) + if isinstance(a, (SymInt, SymFloat)): + return a.__sym_max__(b) + elif isinstance(b, (SymInt, SymFloat)): + # Due to promotion semantics, this is operator is commutative: + # max(1, 1.0) === max(1.0, 1) === 1.0 + return b.__sym_max__(a) + # TODO: Probably can make bool work too, just lazy + + all_types, float_types = __all_and_float_types() + + assert isinstance(a, all_types), type(a) + assert isinstance(b, all_types), type(b) + if isinstance(a, float_types) or isinstance(b, float_types): + return builtins.float(builtins.max(a, b)) # type: ignore[call-overload] + else: + return builtins.max(a, b) # type: ignore[call-overload] + + +def __all_and_float_types() -> tuple[tuple[type, ...], tuple[type, ...]]: + try: + import numpy as np + + all_types: tuple[type, ...] = ( + np.integer, + np.floating, + builtins.int, + builtins.float, + ) + float_types: tuple[type, ...] = (np.floating, builtins.float) + except ModuleNotFoundError: + all_types = (builtins.int, builtins.float) + float_types = (builtins.float,) + + return all_types, float_types + + +def sym_min(a, b): + """SymInt-aware utility for min().""" + if overrides.has_torch_function((a, b)): + return overrides.handle_torch_function(sym_min, (a, b), a, b) + if isinstance(a, (SymInt, SymFloat)): + return a.__sym_min__(b) + elif isinstance(b, (SymInt, SymFloat)): + return b.__sym_min__(a) + + all_types, float_types = __all_and_float_types() + + assert isinstance(a, all_types), type(a) + assert isinstance(b, all_types), type(b) + if isinstance(a, float_types) or isinstance(b, float_types): + return builtins.float(builtins.min(a, b)) # type: ignore[call-overload] + else: + return builtins.min(a, b) # type: ignore[call-overload] + + +def sym_sum(args): + """ + N-ary add which is faster to compute for long lists than iterated binary + addition. Only does something special for integers. + """ + if overrides.has_torch_function(args): + return overrides.handle_torch_function(sym_sum, args, args) + + found = None + for a in args: + if not isinstance(a, (SymInt, builtins.int)): + return builtins.sum(args) + if isinstance(a, SymInt): + found = a.node + if found is None: + return builtins.sum(args) + + from torch.fx.experimental.sym_node import to_node, wrap_node + + return wrap_node(found.sym_sum(tuple(to_node(found, a) for a in args))) + + +# Drop in replacement for math.sqrt, math.sin, math.cos etc +def _get_sym_math_fn(name): + def fn(a): + if overrides.has_torch_function_unary(a): + return overrides.handle_torch_function(fn, (a,), a) + if isinstance(a, SymInt): + a = torch.sym_float(a) + if hasattr(a, f"__sym_{name}__"): + return getattr(a, f"__sym_{name}__")() + return getattr(math, name)(a) + + return fn + + +__fn, __name, __sym_name = None, "", "" +for __name in ( + "sqrt", + "cos", + "cosh", + "sin", + "sinh", + "tan", + "tanh", + "asin", + "acos", + "atan", + "log2", +): + __sym_name = f"_sym_{__name}" + __fn = _get_sym_math_fn(__name) + __fn.__qualname__ = __fn.__name__ = __sym_name + globals()[__sym_name] = __fn + + +del __fn, __name, __sym_name, _get_sym_math_fn + +# Adding temporary shortcut +sym_sqrt = globals()["_sym_sqrt"] +__all__.append("sym_sqrt") + + +def sym_ite(b, t, f): + """SymInt-aware utility for ternary operator (``t if b else f``.)""" + if overrides.has_torch_function((b, t, f)): + return overrides.handle_torch_function(sym_ite, (b, t, f), b, t, f) + assert isinstance(b, (SymBool, builtins.bool)) and type(t) == type(f) + if isinstance(b, SymBool): + return b.__sym_ite__(t, f) + return t if b else f + + +# Create a fresh unbacked int, from an (possibly unbacked int) expression. +def sym_fresh_size(expr): + return torch.tensor(expr).item() + + +# Check to see if we can load C extensions, and if not provide some guidance +# on what the problem might be. +try: + # _initExtension is chosen (arbitrarily) as a sentinel. + from torch._C import _initExtension +except ImportError: + import torch._C as _C_for_compiled_check + + # The __file__ check only works for Python 3.7 and above. + if _C_for_compiled_check.__file__ is None: + raise ImportError( + textwrap.dedent( + """ + Failed to load PyTorch C extensions: + It appears that PyTorch has loaded the `torch/_C` folder + of the PyTorch repository rather than the C extensions which + are expected in the `torch._C` namespace. This can occur when + using the `install` workflow. e.g. + $ python -m pip install --no-build-isolation -v . && python -c "import torch" + + This error can generally be solved using the `develop` workflow + $ python -m pip install --no-build-isolation -v -e . && python -c "import torch" # This should succeed + or by running Python from a different directory. + """ + ).strip() + ) from None + raise # If __file__ is not None the cause is unknown, so just re-raise. + +# The torch._C submodule is already loaded via `from torch._C import *` above +# Make an explicit reference to the _C submodule to appease linters +from torch import _C as _C + + +__name, __obj = "", None +for __name in dir(_C): + if __name[0] != "_" and not __name.endswith("Base"): + __all__.append(__name) + __obj = getattr(_C, __name) + if callable(__obj) or inspect.isclass(__obj): + if __obj.__module__ != __name__: # "torch" + # TODO: fix their module from C++ side + if __name not in { + "DisableTorchFunctionSubclass", + "DisableTorchFunction", + "Generator", + }: + __obj.__module__ = __name__ # "torch" + elif __name == "TensorBase": + # issue 109438 / pr 109940. Prevent TensorBase from being copied into torch. + delattr(sys.modules[__name__], __name) + +del __name, __obj + +if not TYPE_CHECKING: + # issue 38137 and python issue 43367. Submodules of a C extension are + # non-standard, and attributes of those submodules cannot be pickled since + # pickle expect to be able to import them as "from _C.sub import attr" + # which fails with "_C is not a package + def _import_extension_to_sys_modules(module, memo=None): + if memo is None: + memo = set() + if module in memo: + return + memo.add(module) + module_name = module.__name__ + for name in dir(module): + member = getattr(module, name) + member_name = getattr(member, "__name__", "") + if inspect.ismodule(member) and member_name.startswith(module_name): + sys.modules.setdefault(member_name, member) + # Recurse for submodules (e.g., `_C._dynamo.eval_frame`) + _import_extension_to_sys_modules(member, memo) + + _import_extension_to_sys_modules(_C) + del _import_extension_to_sys_modules + +################################################################################ +# Define basic utilities +################################################################################ + + +def typename(obj: _Any, /) -> str: + """ + String representation of the type of an object. + + This function returns a fully qualified string representation of an object's type. + Args: + obj (object): The object whose type to represent + Returns: + str: the type of the object `o` + Example: + >>> x = torch.tensor([1, 2, 3]) + >>> torch.typename(x) + 'torch.LongTensor' + >>> torch.typename(torch.nn.Parameter) + 'torch.nn.parameter.Parameter' + """ + if isinstance(obj, torch.Tensor): + return obj.type() + + module = getattr(obj, "__module__", "") or "" + qualname = "" + + if hasattr(obj, "__qualname__"): + qualname = obj.__qualname__ + elif hasattr(obj, "__name__"): + qualname = obj.__name__ + else: + module = obj.__class__.__module__ or "" + qualname = obj.__class__.__qualname__ + + if module in {"", "builtins"}: + return qualname + return f"{module}.{qualname}" + + +def is_tensor(obj: _Any, /) -> _TypeIs["torch.Tensor"]: + r"""Returns True if `obj` is a PyTorch tensor. + + Note that this function is simply doing ``isinstance(obj, Tensor)``. + Using that ``isinstance`` check is better for type checking with mypy, + and more explicit - so it's recommended to use that instead of + ``is_tensor``. + + Args: + obj (object): Object to test + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> torch.is_tensor(x) + True + + """ + return isinstance(obj, torch.Tensor) + + +def is_storage(obj: _Any, /) -> _TypeIs[_Union["TypedStorage", "UntypedStorage"]]: + r"""Returns True if `obj` is a PyTorch storage object. + + Args: + obj (Object): Object to test + """ + return type(obj) in _storage_classes + + +_GLOBAL_DEVICE_CONTEXT = threading.local() + + +def get_default_device() -> "torch.device": + r"""Gets the default ``torch.Tensor`` to be allocated on ``device``""" + global _GLOBAL_DEVICE_CONTEXT + + from torch.overrides import _get_current_function_mode_stack + from torch.utils._device import DeviceContext + + def _get_device_with_index(device): + if device.index is not None: + return device + else: + # TODO: Call like get_device_index() method corresponding to + # each device type + return torch.tensor([]).device + + # Get device from any active DeviceContext. + device_mode = next( + filter( + lambda mode: isinstance(mode, DeviceContext), + reversed(_get_current_function_mode_stack()), + ), + None, + ) + if device_mode: + device = device_mode.device + return _get_device_with_index(device) + + if hasattr(_GLOBAL_DEVICE_CONTEXT, "device_context"): + device = _GLOBAL_DEVICE_CONTEXT.device_context.device + return _get_device_with_index(device) + else: + return torch.device("cpu") + + +def set_default_device(device: "Device") -> None: + """Sets the default ``torch.Tensor`` to be allocated on ``device``. This + does not affect factory function calls which are called with an explicit + ``device`` argument. Factory calls will be performed as if they + were passed ``device`` as an argument. + + To only temporarily change the default device instead of setting it + globally, use ``with torch.device(device):`` instead. + + The default device is initially ``cpu``. If you set the default tensor + device to another device (e.g., ``cuda``) without a device index, tensors + will be allocated on whatever the current device for the device type, + even after :func:`torch.cuda.set_device` is called. + + .. warning:: + + This function imposes a slight performance cost on every Python + call to the torch API (not just factory functions). If this + is causing problems for you, please comment on + https://github.com/pytorch/pytorch/issues/92701 + + .. note:: + + This doesn't affect functions that create tensors that share the same memory as the input, like: + :func:`torch.from_numpy` and :func:`torch.frombuffer` + + Args: + device (device or string): the device to set as default + + Example:: + + >>> # xdoctest: +SKIP("requires cuda, changes global state") + >>> torch.get_default_device() + device(type='cpu') + >>> torch.set_default_device('cuda') # current device is 0 + >>> torch.get_default_device() + device(type='cuda', index=0) + >>> torch.set_default_device('cuda') + >>> torch.cuda.set_device('cuda:1') # current device is 1 + >>> torch.get_default_device() + device(type='cuda', index=1) + >>> torch.set_default_device('cuda:1') + >>> torch.get_default_device() + device(type='cuda', index=1) + + """ + global _GLOBAL_DEVICE_CONTEXT + if hasattr(_GLOBAL_DEVICE_CONTEXT, "device_context"): + device_context = _GLOBAL_DEVICE_CONTEXT.device_context + if device_context is not None: + device_context.__exit__(None, None, None) + + if device is None: + device_context = None + else: + from torch.utils._device import DeviceContext + + device_context = DeviceContext(device) + device_context.__enter__() + _GLOBAL_DEVICE_CONTEXT.device_context = device_context + + +def set_default_tensor_type(t: _Union[type["torch.Tensor"], str], /) -> None: + r""" + .. warning:: + + This function is deprecated as of PyTorch 2.1, please use :func:`torch.set_default_dtype()` and + :func:`torch.set_default_device()` as alternatives. + + Sets the default ``torch.Tensor`` type to floating point tensor type + ``t``. This type will also be used as default floating point type for + type inference in :func:`torch.tensor`. + + The default floating point tensor type is initially ``torch.FloatTensor``. + + Args: + t (type or string): the floating point tensor type or its name + + Example:: + + >>> # xdoctest: +SKIP("Other tests may have changed the default type. Can we reset it?") + >>> torch.tensor([1.2, 3]).dtype # initial default for floating point is torch.float32 + torch.float32 + >>> torch.set_default_tensor_type(torch.DoubleTensor) + >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor + torch.float64 + + """ + if isinstance(t, str): + t = _import_dotted_name(t) + _C._set_default_tensor_type(t) + + +def set_default_dtype(d: "torch.dtype", /) -> None: + r""" + + Sets the default floating point dtype to :attr:`d`. Supports floating point dtype + as inputs. Other dtypes will cause torch to raise an exception. + + When PyTorch is initialized its default floating point dtype is torch.float32, + and the intent of set_default_dtype(torch.float64) is to facilitate NumPy-like + type inference. The default floating point dtype is used to: + + 1. Implicitly determine the default complex dtype. When the default floating type is float16, + the default complex dtype is complex32. For float32, the default complex dtype is complex64. + For float64, it is complex128. For bfloat16, an exception will be raised because + there is no corresponding complex type for bfloat16. + 2. Infer the dtype for tensors constructed using Python floats or complex Python + numbers. See examples below. + 3. Determine the result of type promotion between bool and integer tensors and + Python floats and complex Python numbers. + + Args: + d (:class:`torch.dtype`): the floating point dtype to make the default. + + Example: + >>> # xdoctest: +SKIP("Other tests may have changed the default type. Can we reset it?") + >>> # initial default for floating point is torch.float32 + >>> # Python floats are interpreted as float32 + >>> torch.tensor([1.2, 3]).dtype + torch.float32 + >>> # initial default for floating point is torch.complex64 + >>> # Complex Python numbers are interpreted as complex64 + >>> torch.tensor([1.2, 3j]).dtype + torch.complex64 + + >>> torch.set_default_dtype(torch.float64) + >>> # Python floats are now interpreted as float64 + >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor + torch.float64 + >>> # Complex Python numbers are now interpreted as complex128 + >>> torch.tensor([1.2, 3j]).dtype # a new complex tensor + torch.complex128 + + >>> torch.set_default_dtype(torch.float16) + >>> # Python floats are now interpreted as float16 + >>> torch.tensor([1.2, 3]).dtype # a new floating point tensor + torch.float16 + >>> # Complex Python numbers are now interpreted as complex128 + >>> torch.tensor([1.2, 3j]).dtype # a new complex tensor + torch.complex32 + + """ + _C._set_default_dtype(d) + + +def use_deterministic_algorithms( + mode: builtins.bool, + *, + warn_only: builtins.bool = False, +) -> None: + r"""Sets whether PyTorch operations must use "deterministic" + algorithms. That is, algorithms which, given the same input, and when + run on the same software and hardware, always produce the same output. + When enabled, operations will use deterministic algorithms when available, + and if only nondeterministic algorithms are available they will throw a + :class:`RuntimeError` when called. + + .. note:: This setting alone is not always enough to make an application + reproducible. Refer to :ref:`reproducibility` for more information. + + .. note:: :func:`torch.set_deterministic_debug_mode` offers an alternative + interface for this feature. + + The following normally-nondeterministic operations will act + deterministically when ``mode=True``: + + * :class:`torch.nn.Conv1d` when called on CUDA tensor + * :class:`torch.nn.Conv2d` when called on CUDA tensor + * :class:`torch.nn.Conv3d` when called on CUDA tensor + * :class:`torch.nn.ConvTranspose1d` when called on CUDA tensor + * :class:`torch.nn.ConvTranspose2d` when called on CUDA tensor + * :class:`torch.nn.ConvTranspose3d` when called on CUDA tensor + * :class:`torch.nn.ReplicationPad1d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.ReplicationPad2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.ReplicationPad3d` when attempting to differentiate a CUDA tensor + * :func:`torch.bmm` when called on sparse-dense CUDA tensors + * :func:`torch.Tensor.__getitem__` when attempting to differentiate a CPU tensor + and the index is a list of tensors + * :func:`torch.Tensor.index_put` with ``accumulate=False`` + * :func:`torch.Tensor.index_put` with ``accumulate=True`` when called on a CPU + tensor + * :func:`torch.Tensor.put_` with ``accumulate=True`` when called on a CPU + tensor + * :func:`torch.Tensor.scatter_add_` when called on a CUDA tensor + * :func:`torch.gather` when called on a CUDA tensor that requires grad + * :func:`torch.index_add` when called on CUDA tensor + * :func:`torch.index_select` when attempting to differentiate a CUDA tensor + * :func:`torch.repeat_interleave` when attempting to differentiate a CUDA tensor + * :func:`torch.Tensor.index_copy` when called on a CPU or CUDA tensor + * :func:`torch.Tensor.scatter` when `src` type is Tensor and called on CUDA tensor + * :func:`torch.Tensor.scatter_reduce` when ``reduce='sum'`` or ``reduce='mean'`` and called on CUDA tensor + + The following normally-nondeterministic operations will throw a + :class:`RuntimeError` when ``mode=True``: + + * :class:`torch.nn.AvgPool3d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.AdaptiveAvgPool2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.AdaptiveAvgPool3d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.MaxPool3d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.AdaptiveMaxPool2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.FractionalMaxPool2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.FractionalMaxPool3d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.MaxUnpool1d` + * :class:`torch.nn.MaxUnpool2d` + * :class:`torch.nn.MaxUnpool3d` + * :func:`torch.nn.functional.interpolate` when attempting to differentiate a CUDA tensor + and one of the following modes is used: + + - ``linear`` + - ``bilinear`` + - ``bicubic`` + - ``trilinear`` + + * :class:`torch.nn.ReflectionPad1d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.ReflectionPad2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.ReflectionPad3d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.NLLLoss` when called on a CUDA tensor + * :class:`torch.nn.CTCLoss` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.EmbeddingBag` when attempting to differentiate a CUDA tensor when + ``mode='max'`` + * :func:`torch.Tensor.put_` when ``accumulate=False`` + * :func:`torch.Tensor.put_` when ``accumulate=True`` and called on a CUDA tensor + * :func:`torch.histc` when called on a CUDA tensor + * :func:`torch.bincount` when called on a CUDA tensor and ``weights`` + tensor is given + * :func:`torch.kthvalue` with called on a CUDA tensor + * :func:`torch.median` with indices output when called on a CUDA tensor + * :func:`torch.nn.functional.grid_sample` when attempting to differentiate a CUDA tensor + * :func:`torch.cumsum` when called on a CUDA tensor when dtype is floating point or complex + * :func:`torch.Tensor.scatter_reduce` when ``reduce='prod'`` and called on CUDA tensor + * :func:`torch.Tensor.resize_` when called with a quantized tensor + + In addition, several operations fill uninitialized memory when this setting + is turned on and when + :attr:`torch.utils.deterministic.fill_uninitialized_memory` is turned on. + See the documentation for that attribute for more information. + + A handful of CUDA operations are nondeterministic if the CUDA version is + 10.2 or greater, unless the environment variable ``CUBLAS_WORKSPACE_CONFIG=:4096:8`` + or ``CUBLAS_WORKSPACE_CONFIG=:16:8`` is set. See the CUDA documentation for more + details: ``_ + If one of these environment variable configurations is not set, a :class:`RuntimeError` + will be raised from these operations when called with CUDA tensors: + + * :func:`torch.mm` + * :func:`torch.mv` + * :func:`torch.bmm` + + Note that deterministic operations tend to have worse performance than + nondeterministic operations. + + .. note:: + + This flag does not detect or prevent nondeterministic behavior caused + by calling an inplace operation on a tensor with an internal memory + overlap or by giving such a tensor as the :attr:`out` argument for an + operation. In these cases, multiple writes of different data may target + a single memory location, and the order of writes is not guaranteed. + + Args: + mode (:class:`bool`): If True, makes potentially nondeterministic + operations switch to a deterministic algorithm or throw a runtime + error. If False, allows nondeterministic operations. + + Keyword args: + warn_only (:class:`bool`, optional): If True, operations that do not + have a deterministic implementation will throw a warning instead of + an error. Default: ``False`` + + Example:: + + >>> # xdoctest: +SKIP + >>> torch.use_deterministic_algorithms(True) + + # Forward mode nondeterministic error + >>> torch.randn(10, device='cuda').kthvalue(1) + ... + RuntimeError: kthvalue CUDA does not have a deterministic implementation... + + # Backward mode nondeterministic error + >>> torch.nn.AvgPool3d(1)(torch.randn(3, 4, 5, 6, requires_grad=True).cuda()).sum().backward() + ... + RuntimeError: avg_pool3d_backward_cuda does not have a deterministic implementation... + """ + _C._set_deterministic_algorithms(mode, warn_only=warn_only) + + +def are_deterministic_algorithms_enabled() -> builtins.bool: + r"""Returns True if the global deterministic flag is turned on. Refer to + :func:`torch.use_deterministic_algorithms` documentation for more details. + """ + return _C._get_deterministic_algorithms() + + +def is_deterministic_algorithms_warn_only_enabled() -> builtins.bool: + r"""Returns True if the global deterministic flag is set to warn only. + Refer to :func:`torch.use_deterministic_algorithms` documentation for more + details. + """ + return _C._get_deterministic_algorithms_warn_only() + + +def set_deterministic_debug_mode(debug_mode: _Union[builtins.int, str]) -> None: + r"""Sets the debug mode for deterministic operations. + + .. note:: This is an alternative interface for + :func:`torch.use_deterministic_algorithms`. Refer to that function's + documentation for details about affected operations. + + Args: + debug_mode(str or int): If "default" or 0, don't error or warn on + nondeterministic operations. If "warn" or 1, warn on + nondeterministic operations. If "error" or 2, error on + nondeterministic operations. + """ + + # NOTE: builtins.int is used here because int in this scope resolves + # to torch.int + if not isinstance(debug_mode, (builtins.int, str)): + raise TypeError(f"debug_mode must be str or int, but got {type(debug_mode)}") + + if isinstance(debug_mode, str): + if debug_mode == "default": + debug_mode = 0 + elif debug_mode == "warn": + debug_mode = 1 + elif debug_mode == "error": + debug_mode = 2 + else: + raise RuntimeError( + "invalid value of debug_mode, expected one of `default`, " + f"`warn`, `error`, but got {debug_mode}" + ) + + if debug_mode == 0: + _C._set_deterministic_algorithms(False) + elif debug_mode == 1: + _C._set_deterministic_algorithms(True, warn_only=True) + elif debug_mode == 2: + _C._set_deterministic_algorithms(True) + else: + raise RuntimeError( + f"invalid value of debug_mode, expected 0, 1, or 2, but got {debug_mode}" + ) + + +def get_deterministic_debug_mode() -> builtins.int: + r"""Returns the current value of the debug mode for deterministic + operations. Refer to :func:`torch.set_deterministic_debug_mode` + documentation for more details. + """ + + if _C._get_deterministic_algorithms(): + if _C._get_deterministic_algorithms_warn_only(): + return 1 + else: + return 2 + else: + return 0 + + +def get_float32_matmul_precision() -> str: + r"""Returns the current value of float32 matrix multiplication precision. Refer to + :func:`torch.set_float32_matmul_precision` documentation for more details. + """ + return _C._get_float32_matmul_precision() + + +def set_float32_matmul_precision(precision: str) -> None: + r"""Sets the internal precision of float32 matrix multiplications. + + Running float32 matrix multiplications in lower precision may significantly increase + performance, and in some programs the loss of precision has a negligible impact. + + Supports three settings: + + * "highest", float32 matrix multiplications use the float32 datatype (24 mantissa + bits with 23 bits explicitly stored) for internal computations. + * "high", float32 matrix multiplications either use the TensorFloat32 datatype (10 + mantissa bits explicitly stored) or treat each float32 number as the sum of two bfloat16 numbers + (approximately 16 mantissa bits with 14 bits explicitly stored), if the appropriate fast matrix multiplication + algorithms are available. Otherwise float32 matrix multiplications are computed + as if the precision is "highest". See below for more information on the bfloat16 + approach. + * "medium", float32 matrix multiplications use the bfloat16 datatype (8 mantissa + bits with 7 bits explicitly stored) for internal computations, if a fast matrix multiplication algorithm + using that datatype internally is available. Otherwise float32 + matrix multiplications are computed as if the precision is "high". + + When using "high" precision, float32 multiplications may use a bfloat16-based algorithm + that is more complicated than simply truncating to some smaller number mantissa bits + (e.g. 10 for TensorFloat32, 7 for bfloat16 explicitly stored). Refer to [Henry2019]_ for a complete + description of this algorithm. To briefly explain here, the first step is to realize + that we can perfectly encode a single float32 number as the sum of three bfloat16 + numbers (because float32 has 23 mantissa bits while bfloat16 has 7 explicitly stored, and both have the + same number of exponent bits). This means that the product of two float32 numbers can + be exactly given by the sum of nine products of bfloat16 numbers. We can then trade + accuracy for speed by dropping some of these products. The "high" precision algorithm + specifically keeps only the three most significant products, which conveniently excludes + all of the products involving the last 8 mantissa bits of either input. This means that + we can represent our inputs as the sum of two bfloat16 numbers rather than three. + Because bfloat16 fused-multiply-add (FMA) instructions are typically >10x faster than + float32 ones, it's faster to do three multiplications and 2 additions with bfloat16 + precision than it is to do a single multiplication with float32 precision. + + .. [Henry2019] http://arxiv.org/abs/1904.06376 + + .. note:: + + This does not change the output dtype of float32 matrix multiplications, + it controls how the internal computation of the matrix multiplication is performed. + + .. note:: + + This does not change the precision of convolution operations. Other flags, + like `torch.backends.cudnn.allow_tf32`, may control the precision of convolution + operations. + + .. note:: + + This flag currently only affects one native device type: CUDA. + If "high" or "medium" are set then the TensorFloat32 datatype will be used + when computing float32 matrix multiplications, equivalent to setting + `torch.backends.cuda.matmul.allow_tf32 = True`. When "highest" (the default) + is set then the float32 datatype is used for internal computations, equivalent + to setting `torch.backends.cuda.matmul.allow_tf32 = False`. + + Args: + precision(str): can be set to "highest" (default), "high", or "medium" (see above). + + """ + _C._set_float32_matmul_precision(precision) + + +def set_warn_always(b: builtins.bool, /) -> None: + r"""When this flag is False (default) then some PyTorch warnings may only + appear once per process. This helps avoid excessive warning information. + Setting it to True causes these warnings to always appear, which may be + helpful when debugging. + + Args: + b (:class:`bool`): If True, force warnings to always be emitted + If False, set to the default behaviour + """ + _C._set_warnAlways(b) + + +def is_warn_always_enabled() -> builtins.bool: + r"""Returns True if the global warn_always flag is turned on. Refer to + :func:`torch.set_warn_always` documentation for more details. + """ + return _C._get_warnAlways() + + +################################################################################ +# Define error checking functions +################################################################################ + +# These error checking functions must be kept consistent with their C++ +# equivalents. Their C++ equivalents are mentioned where applicable. + + +def _check_with( + error_type, + cond: _Union[builtins.bool, SymBool], + message: _Callable[[], str], +): # noqa: F811 + if not isinstance(cond, (builtins.bool, SymBool)): + raise TypeError(f"cond must be a bool, but got {type(cond)}") + + from torch.fx.experimental.symbolic_shapes import expect_true + + if expect_true(cond): + return + + # error_type must be a subclass of Exception and not subclass of Warning + assert issubclass(error_type, Exception) and not issubclass(error_type, Warning) + + if message is None: + message_evaluated = ( + "Expected cond to be True, but got False. (Could this error " + "message be improved? If so, please report an enhancement request " + "to PyTorch.)" + ) + + else: + if not callable(message): + raise TypeError("message must be a callable") + + message_evaluated = str(message()) + + raise error_type(message_evaluated) + + +def _check(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``RuntimeError`` + + C++ equivalent: ``TORCH_CHECK`` + + Args: + cond (:class:`bool`): If False, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_with(RuntimeError, cond, message) + + +def _check_is_size(i, message=None, *, max=None): + """Checks that a given integer is a valid size (i.e., is non-negative). + You should use this over ``_check(i >= 0)`` because it can prevent + ``GuardOnDataDependentSymNode`` exceptions by opting yourself into alternate + semantics for ``guard_size_oblivious`` tests that treat values 0 and 1 + equivalently to all other values. + + When max is not None, this specifies an upper bound equivalent to + ``_check(i <= max)``. This bound is also subject to alternate semantics: + in ``guard_size_oblivious`` tests, we assume that a constant max bound is + treated equivalently to all other values. Symbolic max bounds are not yet + supported. + + NB: Do NOT use this in contexts where a -1 size would be valid (indicating + to infer the size from context, or if you should wrap-around or truncate). + Only use this if the only valid value is an honest to goodness size. + """ + # This is responsible for the expect_true + _check(i >= 0, message) + from torch.fx.experimental.symbolic_shapes import _advise_is_size + + _advise_is_size(i) + + if max is not None: + _check(i <= max, message) + + from torch.fx.experimental.symbolic_shapes import _advise_is_bounded + + _advise_is_bounded(i, max) + + +def _check_index(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``IndexError`` + + C++ equivalent: ``TORCH_CHECK_INDEX`` + + Args: + cond (:class:`bool`): If False, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_with(IndexError, cond, message) + + +def _check_value(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``ValueError`` + + C++ equivalent: ``TORCH_CHECK_VALUE`` + + Args: + cond (:class:`bool`): If False, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_with(ValueError, cond, message) + + +def _check_type(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``TypeError`` + + C++ equivalent: ``TORCH_CHECK_TYPE`` + + Args: + cond (:class:`bool`): If False, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_with(TypeError, cond, message) + + +def _check_not_implemented(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``NotImplementedError`` + + C++ equivalent: ``TORCH_CHECK_NOT_IMPLEMENTED`` + + Args: + cond (:class:`bool`): If False, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_with(NotImplementedError, cond, message) + + +def _check_tensor_all_with(error_type, cond, message=None): # noqa: F811 + if not is_tensor(cond): + raise TypeError(f"cond must be a tensor, but got {type(cond)}") + + if not cond.dtype == torch.bool: + raise TypeError(f"cond tensor must have dtype torch.bool, but got {cond.dtype}") + + _check_with(error_type, cond._is_all_true().item(), message) # type: ignore[arg-type] + + +# C++ equivalent: `TORCH_CHECK_TENSOR_ALL` +def _check_tensor_all(cond, message=None): # noqa: F811 + r"""Throws error containing an optional message if the specified condition + is False. + + Error type: ``RuntimeError`` + + C++ equivalent: ``TORCH_CHECK_TENSOR_ALL`` + + Args: + cond (:class:`torch.Tensor`): Tensor of dtype ``torch.bool``. If any + element is ``False``, throw error + + message (Callable, optional): Callable that returns either a string or + an object that has a ``__str__()`` method to be used as the error + message. Default: ``None`` + """ + _check_tensor_all_with(RuntimeError, cond, message) + + +################################################################################ +# Define numeric constants +################################################################################ + +# For Python Array API (https://data-apis.org/array-api/latest/API_specification/constants.html) and +# NumPy consistency (https://numpy.org/devdocs/reference/constants.html) +from math import e, inf, nan, pi + + +newaxis: None = None + +__all__.extend(["e", "pi", "nan", "inf", "newaxis"]) + +################################################################################ +# Define Storage and Tensor classes +################################################################################ + +from torch._tensor import Tensor # usort: skip + +# needs to be after torch.Tensor is defined to avoid circular dependencies +from torch import storage as storage # usort: skip +from torch.storage import ( + _LegacyStorage, + _StorageBase, + _warn_typed_storage_removal, + TypedStorage, + UntypedStorage, +) + + +# NOTE: New Storage classes should never be added. When adding a new +# dtype, use torch.storage.TypedStorage directly. +class ByteStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.uint8 + + +class DoubleStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.double + + +class FloatStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.float + + +class HalfStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.half + + +class LongStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.long + + +class IntStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.int + + +class ShortStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.short + + +class CharStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.int8 + + +class BoolStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.bool + + +class BFloat16Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.bfloat16 + + +class ComplexDoubleStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.cdouble + + +class ComplexFloatStorage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.cfloat + + +class QUInt8Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.quint8 + + +class QInt8Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.qint8 + + +class QInt32Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.qint32 + + +class QUInt4x2Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.quint4x2 + + +class QUInt2x4Storage(_LegacyStorage): + @classproperty + def dtype(self): + _warn_typed_storage_removal(stacklevel=3) + return self._dtype + + @classproperty + def _dtype(self): + return torch.quint2x4 + + +_storage_classes: set[type[_Union[TypedStorage, UntypedStorage]]] = { + UntypedStorage, + DoubleStorage, + FloatStorage, + LongStorage, + IntStorage, + ShortStorage, + CharStorage, + ByteStorage, + HalfStorage, + BoolStorage, + QUInt8Storage, + QInt8Storage, + QInt32Storage, + BFloat16Storage, + ComplexFloatStorage, + ComplexDoubleStorage, + QUInt4x2Storage, + QUInt2x4Storage, + TypedStorage, +} + +# The _tensor_classes set is initialized by the call to initialize_python_bindings. +_tensor_classes: set[type["torch.Tensor"]] = set() + +# If you edit these imports, please update torch/__init__.py.in as well +from torch import amp as amp, random as random, serialization as serialization +from torch._tensor_str import set_printoptions +from torch.amp import autocast, GradScaler +from torch.random import get_rng_state, initial_seed, manual_seed, seed, set_rng_state +from torch.serialization import load, save + + +################################################################################ +# Initialize extension +################################################################################ + + +# Shared memory manager needs to know the exact location of manager executable +def _manager_path(): + if platform.system() == "Windows": + return b"" + path = get_file_path("torch", "bin", "torch_shm_manager") + prepare_multiprocessing_environment(get_file_path("torch")) + if not os.path.exists(path): + raise RuntimeError("Unable to find torch_shm_manager at " + path) + return path.encode("utf-8") + + +_C._initExtension(_manager_path()) + +del _manager_path + +# Appease the type checker: it can't deal with direct setting of globals(). +# Note that we will see "too many" functions when reexporting this way; there +# is not a good way to fix this problem. Perhaps, try to redesign VariableFunctions +# so that this import is good enough +if TYPE_CHECKING: + # Some type signatures pulled in from _VariableFunctions here clash with + # signatures already imported. For now these clashes are ignored; see + # PR #43339 for details. + from torch._C._VariableFunctions import * # type: ignore[assignment, misc] # noqa: F403 + + # Fixup segment_reduce visibility + _segment_reduce = segment_reduce + del segment_reduce # noqa: F821 + +# Ops not to be exposed in `torch` namespace, +# mostly helper ops. +PRIVATE_OPS = ("unique_dim",) + +__name, __obj = "", None +for __name in dir(_C._VariableFunctions): + if __name.startswith("__") or __name in PRIVATE_OPS: + continue + __obj = getattr(_C._VariableFunctions, __name) + __obj.__module__ = __name__ # "torch" + # Hide some APIs that should not be public + if __name == "segment_reduce": + # TODO: Once the undocumented FC window is passed, remove the line below + globals()[__name] = __obj + __name = "_" + __name + globals()[__name] = __obj + if not __name.startswith("_"): + __all__.append(__name) + +del __name, __obj + +################################################################################ +# Add torch.dtype instances to the public API +################################################################################ + +import torch + + +__all__.extend( + name for name in dir(torch) if isinstance(getattr(torch, name), torch.dtype) +) + +################################################################################ +# Import TorchDynamo's lazy APIs to avoid circular dependencies +################################################################################ + +# needs to be before from torch.functional import * to avoid circular dependencies +from torch._compile import _disable_dynamo # usort: skip + +################################################################################ +# Import interface functions defined in Python +################################################################################ + +# needs to be after the above ATen bindings so we can overwrite from Python side +from torch import _VF as _VF, functional as functional # usort: skip +from torch.functional import * # usort: skip # noqa: F403 + +################################################################################ +# Remove unnecessary members +################################################################################ + +del _StorageBase +del _LegacyStorage + +################################################################################ +# Define _assert +################################################################################ + + +# needs to be before the submodule imports to avoid circular dependencies +def _assert(condition, message): + r"""A wrapper around Python's assert which is symbolically traceable.""" + if type(condition) is not torch.Tensor and overrides.has_torch_function( + (condition,) + ): + return overrides.handle_torch_function( + _assert, (condition,), condition, message + ) + assert condition, message + + +################################################################################ +# Import most common subpackages +################################################################################ + +# Use the redundant form so that type checkers know that these are a part of +# the public API. The "regular" import lines are there solely for the runtime +# side effect of adding to the imported module's members for other users. + +# needs to be before import torch.nn as nn to avoid circular dependencies +from torch.autograd import ( # usort: skip + enable_grad as enable_grad, + inference_mode as inference_mode, + no_grad as no_grad, + set_grad_enabled as set_grad_enabled, +) + +from torch import ( + __config__ as __config__, + __future__ as __future__, + _awaits as _awaits, + accelerator as accelerator, + autograd as autograd, + backends as backends, + cpu as cpu, + cuda as cuda, + distributed as distributed, + distributions as distributions, + fft as fft, + futures as futures, + hub as hub, + jit as jit, + linalg as linalg, + mps as mps, + mtia as mtia, + multiprocessing as multiprocessing, + nested as nested, + nn as nn, + optim as optim, + overrides as overrides, + profiler as profiler, + sparse as sparse, + special as special, + testing as testing, + types as types, + utils as utils, + version as version, + xpu as xpu, +) +from torch.signal import windows as windows + + +# Quantized, sparse, AO, etc. should be last to get imported, as nothing +# is expected to depend on them. +from torch import ao as ao # usort: skip + +# nn.quant* depends on ao -- so should be after those. +import torch.nn.intrinsic +import torch.nn.qat +import torch.nn.quantizable +import torch.nn.quantized + + +_C._init_names(list(_storage_classes)) + +# attach docstrings to torch and tensor functions +from torch import _size_docs, _storage_docs, _tensor_docs, _torch_docs + + +del _torch_docs, _tensor_docs, _storage_docs, _size_docs + + +def compiled_with_cxx11_abi() -> builtins.bool: + r"""Returns whether PyTorch was built with _GLIBCXX_USE_CXX11_ABI=1""" + return True + + +from torch import _library as _library, _ops as _ops + + +# Import the ops and classes "namespace" +from torch._ops import ops as ops # usort: skip +from torch._classes import classes as classes # usort: skip + +sys.modules.setdefault(f"{__name__}.ops", ops) +sys.modules.setdefault(f"{__name__}.classes", classes) + +# quantization depends on torch.fx and torch.ops +# Import quantization +from torch import quantization as quantization # usort: skip + +# Import the quasi random sampler +from torch import quasirandom as quasirandom # usort: skip + +# If you are seeing this, it means that this call site was not checked if +# the memory format could be preserved, and it was switched to old default +# behaviour of contiguous +legacy_contiguous_format = contiguous_format # defined by _C._initExtension() + +# Register fork handler to initialize OpenMP in child processes (see gh-28389) +from torch.multiprocessing._atfork import register_after_fork + + +register_after_fork(torch.get_num_threads) +del register_after_fork + +# Import tools that require fully imported torch (for applying +# torch.jit.script as a decorator, for instance): +from torch._lobpcg import lobpcg as lobpcg + + +# These were previously defined in native_functions.yaml and appeared on the +# `torch` namespace, but we moved them to c10 dispatch to facilitate custom +# class usage. We add these lines here to preserve backward compatibility. +quantized_lstm = ops.aten.quantized_lstm +quantized_gru = ops.aten.quantized_gru + +# Import experimental masked operations support. See +# [RFC-0016](https://github.com/pytorch/rfcs/pull/27) for more +# information. +from torch import masked as masked + +# Import removed ops with error message about removal +from torch._linalg_utils import ( # type: ignore[misc] + _symeig as symeig, + eig, + lstsq, + matrix_rank, + solve, +) +from torch.utils.dlpack import from_dlpack, to_dlpack + + +class _TorchCompileInductorWrapper: + compiler_name = "inductor" + + def __init__(self, mode, options, dynamic): + from torch._inductor.compiler_bisector import CompilerBisector + + self.config: dict[str, _Any] = {} + self.dynamic = dynamic + self.apply_mode(mode) + self.apply_options(options) + self.apply_options(CompilerBisector.get_config_change("inductor")) + + cuda_version = None + if hasattr(torch, "version"): + from torch.torch_version import TorchVersion + + cuda_version = TorchVersion(getattr(torch.version, "cuda", "0.0")) + + if self.config.get("triton.cudagraphs", False) and ( + (cuda_version and cuda_version < "12.6") + or not profiler_allow_cudagraph_cupti_lazy_reinit_cuda12() + ): + os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1" + # FIXME: CUDA Graph does not work well with CUPTI teardown. + # 1) crashes on 1st lazy CUPTI re-init after teardown (CUDA 11) + # 2) crashes on 2nd non-lazy CUPTI re-init after teardown (CUDA 12) + # Workaround: turn off CUPTI teardown when using CUDA Graphs. + os.environ["TEARDOWN_CUPTI"] = "0" + + def __eq__(self, other): + return ( + isinstance(other, _TorchCompileInductorWrapper) + and self.config == other.config + and self.dynamic == other.dynamic + ) + + def apply_mode(self, mode: _Optional[str]): + if mode and mode != "default": + from torch._inductor import list_mode_options + + self.apply_options(list_mode_options(mode, self.dynamic)) + + def apply_options(self, options: _Optional[dict[str, _Any]]): + if not options: + return + + from torch._inductor import config + + current_config: dict[str, _Any] = config.get_config_copy() + + for key, val in options.items(): + attr_name = key.replace("-", "_") + if attr_name not in current_config: + raise RuntimeError( + f"Unexpected optimization option {key}, known options are {list(current_config.keys())}" + ) + attr_type = config.get_type(attr_name) # type: ignore[attr-defined] + # Subscriptable generic types don't support isinstance so skip the type + # check. There doesn't seem to be a good way of checking membership without + # 3rd party libraries. + if _get_origin(attr_type) is None: + if not isinstance(val, attr_type): + val_type_str = type(val).__name__ + expected_type_str = type(current_config[attr_name]).__name__ + raise RuntimeError( + f"Unexpected type of attr {key}, got {val_type_str} should be {expected_type_str}" + ) + self.config[attr_name] = val + + def __call__(self, model_, inputs_): + from torch._inductor.compile_fx import compile_fx + + return compile_fx(model_, inputs_, config_patches=self.config) + + def get_compiler_config(self): + from torch._inductor.compile_fx import get_patched_config_dict + + return get_patched_config_dict(config_patches=self.config) + + def reset(self): + from torch._inductor import config + + if "triton.cudagraphs" in self.config or config.triton.cudagraphs: + if self.config.get("triton.cudagraphs", True): + from torch._inductor.cudagraph_trees import reset_cudagraph_trees + + reset_cudagraph_trees() + + +class _TorchCompileWrapper: + def __init__(self, backend, mode, options, dynamic): + from torch._dynamo.backends.registry import lookup_backend + + if isinstance(backend, str): + self.compiler_name = backend + elif hasattr(backend, "__name__"): + self.compiler_name = backend.__name__ + else: + self.compiler_name = str(backend) + self.dynamic = dynamic + self.compiler_fn = lookup_backend(backend) + self.kwargs = {} + # only pass the args if they non-empty + if mode and mode != "default": + self.kwargs["mode"] = mode + if options: + self.kwargs["options"] = options + + def __eq__(self, other): + return ( + isinstance(other, _TorchCompileWrapper) + and self.compiler_fn == other.compiler_fn + and self.kwargs == other.kwargs + and self.dynamic == other.dynamic + ) + + def __call__(self, model_, inputs_): + return self.compiler_fn(model_, inputs_, **self.kwargs) + + def reset(self): + if hasattr(self.compiler_fn, "reset"): + self.compiler_fn.reset() + + +_InputT = _ParamSpec("_InputT") +_RetT = _TypeVar("_RetT") + + +@_overload +def compile( + model: _Callable[_InputT, _RetT], + *, + fullgraph: builtins.bool = False, + dynamic: _Optional[builtins.bool] = None, + backend: _Union[str, _Callable] = "inductor", + mode: _Union[str, None] = None, + options: _Optional[ + dict[str, _Union[str, builtins.int, builtins.bool, _Callable]] + ] = None, + disable: builtins.bool = False, +) -> _Callable[_InputT, _RetT]: ... + + +@_overload +def compile( + model: None = None, + *, + fullgraph: builtins.bool = False, + dynamic: _Optional[builtins.bool] = None, + backend: _Union[str, _Callable] = "inductor", + mode: _Union[str, None] = None, + options: _Optional[ + dict[str, _Union[str, builtins.int, builtins.bool, _Callable]] + ] = None, + disable: builtins.bool = False, +) -> _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]]: ... + + +def compile( + model: _Optional[_Callable[_InputT, _RetT]] = None, + *, + fullgraph: builtins.bool = False, + dynamic: _Optional[builtins.bool] = None, + backend: _Union[str, _Callable] = "inductor", + mode: _Union[str, None] = None, + options: _Optional[ + dict[str, _Union[str, builtins.int, builtins.bool, _Callable]] + ] = None, + disable: builtins.bool = False, +) -> _Union[ + _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]], + _Callable[_InputT, _RetT], +]: + """ + Optimizes given model/function using TorchDynamo and specified backend. + If you are compiling an :class:`torch.nn.Module`, you can also use :meth:`torch.nn.Module.compile` + to compile the module inplace without changing its structure. + + Concretely, for every frame executed within the compiled region, we will attempt + to compile it and cache the compiled result on the code object for future + use. A single frame may be compiled multiple times if previous compiled + results are not applicable for subsequent calls (this is called a "guard + failure), you can use TORCH_LOGS=guards to debug these situations. + Multiple compiled results can be associated with a frame up to + ``torch._dynamo.config.recompile_limit``, which defaults to 8; at which + point we will fall back to eager. Note that compile caches are per + *code object*, not frame; if you dynamically create multiple copies of a + function, they will all share the same code cache. + + Args: + model (Callable or None): Module/function to optimize + fullgraph (bool): If False (default), torch.compile attempts to discover compilable regions + in the function that it will optimize. If True, then we require that the entire function be + capturable into a single graph. If this is not possible (that is, if there are graph breaks), + then this will raise an error. + dynamic (bool or None): Use dynamic shape tracing. When this is True, we will up-front attempt + to generate a kernel that is as dynamic as possible to avoid recompilations when + sizes change. This may not always work as some operations/optimizations will + force specialization; use TORCH_LOGS=dynamic to debug overspecialization. + When this is False, we will NEVER generate dynamic kernels, we will always specialize. + By default (None), we automatically detect if dynamism has occurred and compile a more + dynamic kernel upon recompile. + backend (str or Callable): backend to be used + + - "inductor" is the default backend, which is a good balance between performance and overhead + + - Non experimental in-tree backends can be seen with `torch._dynamo.list_backends()` + + - Experimental or debug in-tree backends can be seen with `torch._dynamo.list_backends(None)` + + - To register an out-of-tree custom backend: + https://pytorch.org/docs/main/torch.compiler_custom_backends.html#registering-custom-backends + mode (str): Can be either "default", "reduce-overhead", "max-autotune" or "max-autotune-no-cudagraphs" + + - "default" is the default mode, which is a good balance between performance and overhead + + - "reduce-overhead" is a mode that reduces the overhead of python with CUDA graphs, + useful for small batches. Reduction of overhead can come at the cost of more memory + usage, as we will cache the workspace memory required for the invocation so that we + do not have to reallocate it on subsequent runs. Reduction of overhead is not guaranteed + to work; today, we only reduce overhead for CUDA only graphs which do not mutate inputs. + There are other circumstances where CUDA graphs are not applicable; use TORCH_LOG=perf_hints + to debug. + + - "max-autotune" is a mode that leverages Triton or template based matrix multiplications + on supported devices and Triton based convolutions on GPU. + It enables CUDA graphs by default on GPU. + + - "max-autotune-no-cudagraphs" is a mode similar to "max-autotune" but without CUDA graphs + + - To see the exact configs that each mode sets you can call `torch._inductor.list_mode_options()` + + options (dict): A dictionary of options to pass to the backend. Some notable ones to try out are + + - `epilogue_fusion` which fuses pointwise ops into templates. Requires `max_autotune` to also be set + + - `max_autotune` which will profile to pick the best matmul configuration + + - `fallback_random` which is useful when debugging accuracy issues + + - `shape_padding` which pads matrix shapes to better align loads on GPUs especially for tensor cores + + - `triton.cudagraphs` which will reduce the overhead of python with CUDA graphs + + - `trace.enabled` which is the most useful debugging flag to turn on + + - `trace.graph_diagram` which will show you a picture of your graph after fusion + + - `guard_filter_fn` that controls which dynamo guards are saved with compilations. + This is an unsafe feature and there is no backward compatibility guarantee provided + for dynamo guards as data types. + For stable helper functions to use, see the documentations in `torch.compiler`, for example: + - `torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe` + - `torch.compiler.skip_guard_on_all_nn_modules_unsafe` + - `torch.compiler.keep_tensor_guards_unsafe` + + - For inductor you can see the full list of configs that it supports by calling `torch._inductor.list_options()` + disable (bool): Turn torch.compile() into a no-op for testing + + Example:: + + @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True) + def foo(x): + return torch.sin(x) + torch.cos(x) + + """ + import sysconfig + + _C._log_api_usage_once("torch.compile") + if sys.version_info >= (3, 14): + raise RuntimeError("torch.compile is not supported on Python 3.14+") + elif sysconfig.get_config_var("Py_GIL_DISABLED") == 1 and sys.version_info < ( + 3, + 13, + 3, + ): + raise RuntimeError( + "torch.compile is not supported on Python < 3.13.3 built with GIL disabled. " + "Please use Python 3.13.3+." + ) + + # Decorator mode + if model is None: + + def fn(model: _Callable[_InputT, _RetT]) -> _Callable[_InputT, _RetT]: + if model is None: + raise RuntimeError("Model can't be None") + return compile( + model, + fullgraph=fullgraph, + dynamic=dynamic, + backend=backend, + mode=mode, + options=options, + disable=disable, + ) + + return fn + + if mode is not None and options is not None: + raise RuntimeError( + "Either mode or options can be specified, but both can't be specified at the same time." + ) + if mode is None and options is None: + mode = "default" + + from torch._inductor.compiler_bisector import CompilerBisector + + if bisect_backend := CompilerBisector.get_backend(): + backend = bisect_backend + + guard_filter_fn = None + if options and isinstance(options, dict): + guard_filter_fn = options.pop("guard_filter_fn", None) + + if backend == "inductor": + backend = _TorchCompileInductorWrapper(mode, options, dynamic) + else: + backend = _TorchCompileWrapper(backend, mode, options, dynamic) + + return torch._dynamo.optimize( + backend=backend, + nopython=fullgraph, + dynamic=dynamic, + disable=disable, + guard_filter_fn=guard_filter_fn, + )(model) # type: ignore[return-value] + + +def _register_device_module(device_type, module): + r"""Register an external runtime module of the specific :attr:`device_type` + supported by torch. + + After the :attr:`module` is registered correctly, the user can refer + the external runtime module as part of torch with attribute torch.xxx. + """ + # Make sure the device_type represent a supported device type for torch. + device_type = torch.device(device_type).type + m = sys.modules[__name__] + if hasattr(m, device_type): + raise RuntimeError( + f"The runtime module of '{device_type}' has already " + f"been registered with '{getattr(m, device_type)}'" + ) + setattr(m, device_type, module) + torch_module_name = ".".join([__name__, device_type]) + sys.modules[torch_module_name] = module + + +from torch import ( + export as export, + func as func, + library as library, + return_types as return_types, +) +from torch._higher_order_ops import cond as cond, while_loop as while_loop +from torch.func import vmap as vmap + + +if not TYPE_CHECKING: + from torch import _meta_registrations + +# Enable CUDA Sanitizer +if "TORCH_CUDA_SANITIZER" in os.environ: + import torch.cuda._sanitizer as csan + + csan.enable_cuda_sanitizer() + +# Populate magic methods on SymInt and SymFloat +import torch.fx.experimental.sym_node +from torch import fx as fx + + +# Register MPS specific decomps +torch.backends.mps._init() + +from torch import compiler as compiler + + +class _TritonLibrary: + lib = torch.library.Library("triton", "DEF") + ops_table: dict[tuple[str, str], _Callable] = {} + + @classmethod + def registerOp(cls, op_key, full_schema, op_impl, dispatch_key): + if (op_key, dispatch_key) not in cls.ops_table: + cls.lib.define(full_schema) + cls.lib.impl("triton::" + op_key, op_impl, dispatch_key) + cls.ops_table[(op_key, dispatch_key)] = op_impl + + return cls.ops_table[(op_key, dispatch_key)] + + +# Deprecated attributes +_deprecated_attrs = { + "has_mps": torch.backends.mps.is_built, + "has_cuda": torch.backends.cuda.is_built, + "has_cudnn": torch.backends.cudnn.is_available, + "has_mkldnn": torch.backends.mkldnn.is_available, +} + +if TYPE_CHECKING: + # Import the following modules during type checking to enable code intelligence features, + # such as auto-completion in tools like pylance, even when these modules are not explicitly + # imported in user code. + from torch import ( + _dynamo as _dynamo, + _inductor as _inductor, + _subclasses as _subclasses, + onnx as onnx, + ) + +else: + _lazy_modules = { + "_dynamo", + "_inductor", + "_export", + # ONNX must be imported after _dynamo, _ops, _subclasses, fx, func and jit + "onnx", + } + + def __getattr__(name): + # Deprecated attrs + replacement = _deprecated_attrs.get(name) + if replacement is not None: + import warnings + + warnings.warn( + f"'{name}' is deprecated, please use '{replacement.__module__}.{replacement.__name__}()'", + stacklevel=2, + ) + return replacement() + + # Lazy modules + if name in _lazy_modules: + return importlib.import_module(f".{name}", __name__) + + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +@functools.cache +def get_device_module(device: _Optional[_Union[torch.device, str]] = None): + """ + Returns the module associated with a given device(e.g., torch.device('cuda'), "mtia:0", "xpu", ...). + If no device is given, return the module for the current accelerator or CPU if none is present. + """ + if isinstance(device, torch.device): + device_module_name = device.type + elif isinstance(device, str): + device_module_name = torch.device(device).type + elif device is None: + # Using default accelerator type. If no accelerator is available, it automatically returns CPU device. + device_module_name = torch._C._get_accelerator().type + else: + raise RuntimeError( + f"Invalid value of device '{device}', expect torch.device, str, or None" + ) + device_module = getattr(torch, device_module_name, None) + if device_module is None: + raise RuntimeError( + f"Device '{device_module_name}' does not have a corresponding module registered as 'torch.{device_module_name}'." + ) + return device_module + + +def _constrain_as_size( + symbol, + min: _Optional[builtins.int] = None, + max: _Optional[builtins.int] = None, +): + """ + This indicates that a given int is size-like, and can be used in any context where a size is expected. + You will typically use this when reading out integers from Tensors, e.g., max.item() or lengths.tolist() + which then need to be used as tensor constructors. Providing these assertions to PyTorch can help resolve + GuardOnDataDependentSymNode errors upon export, since we cannot guard on unbacked SymInts. + + This function has unusual semantics in some circumstances in framework + code, we will treat this int as >= 2 (when we do a size-oblivious guard). + This makes it easier to use the unbacked int in size contexts, + as we will often attempt to guard on a size being zero/one + (e.g., when computing the contiguity of a tensor, or testing if + broadcasting can occur), which will not work on unbacked SymInts. + However, if we conservatively assume that the size is not zero/one, we will + end up with a graph that will still work even if the size is zero/one. + + For more details, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit + ``` + """ + torch.sym_constrain_range_for_size(symbol, min=min, max=max) + + +from torch import _logging + + +_logging._init_logs() + + +def _import_device_backends(): + """ + Leverage the Python plugin mechanism to load out-of-the-tree device extensions. + See this RFC: https://github.com/pytorch/pytorch/issues/122468 + """ + from importlib.metadata import entry_points + + group_name = "torch.backends" + if sys.version_info < (3, 10): + backend_extensions = entry_points().get(group_name, ()) + else: + backend_extensions = entry_points(group=group_name) + + for backend_extension in backend_extensions: + try: + # Load the extension + entrypoint = backend_extension.load() + # Call the entrypoint + entrypoint() + except Exception as err: + raise RuntimeError( + f"Failed to load the backend extension: {backend_extension.name}. " + f"You can disable extension auto-loading with TORCH_DEVICE_BACKEND_AUTOLOAD=0." + ) from err + + +def _is_device_backend_autoload_enabled() -> builtins.bool: + """ + Whether autoloading out-of-the-tree device extensions is enabled. + The switch depends on the value of the environment variable + `TORCH_DEVICE_BACKEND_AUTOLOAD`. + + Returns: + bool: Whether to enable autoloading the extensions. Enabled by default. + + Examples: + >>> torch._is_device_backend_autoload_enabled() + True + """ + # enabled by default + return os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1") == "1" + + +def _as_tensor_fullprec(t): + """ + Like torch.as_tensor, but when given Python data types it will keep + them in full precision. Used for calling convention for Dynamo. + """ + ty = type(t) + if ty is builtins.float: + return torch.as_tensor(t, dtype=torch.float64) + elif ty is builtins.int: + return torch.as_tensor(t, dtype=torch.int64) + else: + return torch.as_tensor(t) + + +# `_import_device_backends` should be kept at the end to ensure +# all the other functions in this module that may be accessed by +# an autoloaded backend are defined +if _is_device_backend_autoload_enabled(): + _import_device_backends() diff --git a/venv/lib/python3.10/site-packages/torch/_appdirs.py b/venv/lib/python3.10/site-packages/torch/_appdirs.py new file mode 100644 index 0000000000000000000000000000000000000000..291963f6f6f627cb69081e4fc764d7cdc0c02df5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_appdirs.py @@ -0,0 +1,666 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# Copyright (c) 2005-2010 ActiveState Software Inc. +# Copyright (c) 2013 Eddy Petrișor + +# flake8: noqa + +""" +This file is directly from +https://github.com/ActiveState/appdirs/blob/3fe6a83776843a46f20c2e5587afcffe05e03b39/appdirs.py + +The license of https://github.com/ActiveState/appdirs copied below: + + +# This is the MIT license + +Copyright (c) 2010 ActiveState Software Inc. + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +""" + +"""Utilities for determining application-specific dirs. + +See for details and usage. +""" +# Dev Notes: +# - Windows "Known Folders": https://learn.microsoft.com/en-us/windows/win32/shell/csidl +# - macOS File System Programming Guide: https://developer.apple.com/library/archive/documentation/FileManagement/Conceptual/FileSystemProgrammingGuide/Introduction/Introduction.html +# - XDG spec for Un*x: https://standards.freedesktop.org/basedir-spec/basedir-spec-latest.html + +__version__ = "1.4.4" +__version_info__ = tuple(int(segment) for segment in __version__.split(".")) + + +import os +import sys + + +unicode = str + +if sys.platform.startswith("java"): + import platform + + os_name = platform.java_ver()[3][0] + if os_name.startswith("Windows"): # "Windows XP", "Windows 7", etc. + system = "win32" + elif os_name.startswith("Mac"): # "Mac OS X", etc. + system = "darwin" + else: # "Linux", "SunOS", "FreeBSD", etc. + # Setting this to "linux2" is not ideal, but only Windows or Mac + # are actually checked for and the rest of the module expects + # *sys.platform* style strings. + system = "linux2" +else: + system = sys.platform + + +def user_data_dir(appname=None, appauthor=None, version=None, roaming=False): + r"""Return full path to the user-specific data dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "roaming" (boolean, default False) can be set True to use the Windows + roaming appdata directory. That means that for users on a Windows + network setup for roaming profiles, this user data will be + sync'd on login. See + + for a discussion of issues. + + Typical user data directories are: + Mac OS X: ~/Library/Application Support/ + Unix: ~/.local/share/ # or in $XDG_DATA_HOME, if defined + Win XP (not roaming): C:\Documents and Settings\\Application Data\\ + Win XP (roaming): C:\Documents and Settings\\Local Settings\Application Data\\ + Win 7 (not roaming): C:\Users\\AppData\Local\\ + Win 7 (roaming): C:\Users\\AppData\Roaming\\ + + For Unix, we follow the XDG spec and support $XDG_DATA_HOME. + That means, by default "~/.local/share/". + """ + if system == "win32": + if appauthor is None: + appauthor = appname + const = roaming and "CSIDL_APPDATA" or "CSIDL_LOCAL_APPDATA" + path = os.path.normpath(_get_win_folder(const)) + if appname: + if appauthor is not False: + path = os.path.join(path, appauthor, appname) + else: + path = os.path.join(path, appname) + elif system == "darwin": + path = os.path.expanduser("~/Library/Application Support/") + if appname: + path = os.path.join(path, appname) + else: + path = os.getenv("XDG_DATA_HOME", os.path.expanduser("~/.local/share")) + if appname: + path = os.path.join(path, appname) + if appname and version: + path = os.path.join(path, version) + return path + + +def site_data_dir(appname=None, appauthor=None, version=None, multipath=False): + r"""Return full path to the user-shared data dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "multipath" is an optional parameter only applicable to *nix + which indicates that the entire list of data dirs should be + returned. By default, the first item from XDG_DATA_DIRS is + returned, or '/usr/local/share/', + if XDG_DATA_DIRS is not set + + Typical site data directories are: + Mac OS X: /Library/Application Support/ + Unix: /usr/local/share/ or /usr/share/ + Win XP: C:\Documents and Settings\All Users\Application Data\\ + Vista: (Fail! "C:\ProgramData" is a hidden *system* directory on Vista.) + Win 7: C:\ProgramData\\ # Hidden, but writeable on Win 7. + + For Unix, this is using the $XDG_DATA_DIRS[0] default. + + WARNING: Do not use this on Windows. See the Vista-Fail note above for why. + """ + if system == "win32": + if appauthor is None: + appauthor = appname + path = os.path.normpath(_get_win_folder("CSIDL_COMMON_APPDATA")) + if appname: + if appauthor is not False: + path = os.path.join(path, appauthor, appname) + else: + path = os.path.join(path, appname) + elif system == "darwin": + path = os.path.expanduser("/Library/Application Support") + if appname: + path = os.path.join(path, appname) + else: + # XDG default for $XDG_DATA_DIRS + # only first, if multipath is False + path = os.getenv( + "XDG_DATA_DIRS", os.pathsep.join(["/usr/local/share", "/usr/share"]) + ) + pathlist = [ + os.path.expanduser(x.rstrip(os.sep)) for x in path.split(os.pathsep) + ] + if appname: + if version: + appname = os.path.join(appname, version) + pathlist = [os.sep.join([x, appname]) for x in pathlist] + + if multipath: + path = os.pathsep.join(pathlist) + else: + path = pathlist[0] + return path + + if appname and version: + path = os.path.join(path, version) + return path + + +def user_config_dir(appname=None, appauthor=None, version=None, roaming=False): + r"""Return full path to the user-specific config dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "roaming" (boolean, default False) can be set True to use the Windows + roaming appdata directory. That means that for users on a Windows + network setup for roaming profiles, this user data will be + sync'd on login. See + + for a discussion of issues. + + Typical user config directories are: + Mac OS X: ~/Library/Preferences/ + Unix: ~/.config/ # or in $XDG_CONFIG_HOME, if defined + Win *: same as user_data_dir + + For Unix, we follow the XDG spec and support $XDG_CONFIG_HOME. + That means, by default "~/.config/". + """ + if system == "win32": + path = user_data_dir(appname, appauthor, None, roaming) + elif system == "darwin": + path = os.path.expanduser("~/Library/Preferences/") + if appname: + path = os.path.join(path, appname) + else: + path = os.getenv("XDG_CONFIG_HOME", os.path.expanduser("~/.config")) + if appname: + path = os.path.join(path, appname) + if appname and version: + path = os.path.join(path, version) + return path + + +def site_config_dir(appname=None, appauthor=None, version=None, multipath=False): + r"""Return full path to the user-shared data dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "multipath" is an optional parameter only applicable to *nix + which indicates that the entire list of config dirs should be + returned. By default, the first item from XDG_CONFIG_DIRS is + returned, or '/etc/xdg/', if XDG_CONFIG_DIRS is not set + + Typical site config directories are: + Mac OS X: same as site_data_dir + Unix: /etc/xdg/ or $XDG_CONFIG_DIRS[i]/ for each value in + $XDG_CONFIG_DIRS + Win *: same as site_data_dir + Vista: (Fail! "C:\ProgramData" is a hidden *system* directory on Vista.) + + For Unix, this is using the $XDG_CONFIG_DIRS[0] default, if multipath=False + + WARNING: Do not use this on Windows. See the Vista-Fail note above for why. + """ + if system == "win32": + path = site_data_dir(appname, appauthor) + if appname and version: + path = os.path.join(path, version) + elif system == "darwin": + path = os.path.expanduser("/Library/Preferences") + if appname: + path = os.path.join(path, appname) + else: + # XDG default for $XDG_CONFIG_DIRS + # only first, if multipath is False + path = os.getenv("XDG_CONFIG_DIRS", "/etc/xdg") + pathlist = [ + os.path.expanduser(x.rstrip(os.sep)) for x in path.split(os.pathsep) + ] + if appname: + if version: + appname = os.path.join(appname, version) + pathlist = [os.sep.join([x, appname]) for x in pathlist] + + if multipath: + path = os.pathsep.join(pathlist) + else: + path = pathlist[0] + return path + + +def user_cache_dir(appname=None, appauthor=None, version=None, opinion=True): + r"""Return full path to the user-specific cache dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "opinion" (boolean) can be False to disable the appending of + "Cache" to the base app data dir for Windows. See + discussion below. + + Typical user cache directories are: + Mac OS X: ~/Library/Caches/ + Unix: ~/.cache/ (XDG default) + Win XP: C:\Documents and Settings\\Local Settings\Application Data\\\Cache + Vista: C:\Users\\AppData\Local\\\Cache + + On Windows the only suggestion in the MSDN docs is that local settings go in + the `CSIDL_LOCAL_APPDATA` directory. This is identical to the non-roaming + app data dir (the default returned by `user_data_dir` above). Apps typically + put cache data somewhere *under* the given dir here. Some examples: + ...\Mozilla\Firefox\Profiles\\Cache + ...\Acme\SuperApp\Cache\1.0 + OPINION: This function appends "Cache" to the `CSIDL_LOCAL_APPDATA` value. + This can be disabled with the `opinion=False` option. + """ + if system == "win32": + if appauthor is None: + appauthor = appname + path = os.path.normpath(_get_win_folder("CSIDL_LOCAL_APPDATA")) + if appname: + if appauthor is not False: + path = os.path.join(path, appauthor, appname) + else: + path = os.path.join(path, appname) + if opinion: + path = os.path.join(path, "Cache") + elif system == "darwin": + path = os.path.expanduser("~/Library/Caches") + if appname: + path = os.path.join(path, appname) + else: + path = os.getenv("XDG_CACHE_HOME", os.path.expanduser("~/.cache")) + if appname: + path = os.path.join(path, appname) + if appname and version: + path = os.path.join(path, version) + return path + + +def user_state_dir(appname=None, appauthor=None, version=None, roaming=False): + r"""Return full path to the user-specific state dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "roaming" (boolean, default False) can be set True to use the Windows + roaming appdata directory. That means that for users on a Windows + network setup for roaming profiles, this user data will be + sync'd on login. See + + for a discussion of issues. + + Typical user state directories are: + Mac OS X: same as user_data_dir + Unix: ~/.local/state/ # or in $XDG_STATE_HOME, if defined + Win *: same as user_data_dir + + For Unix, we follow this Debian proposal + to extend the XDG spec and support $XDG_STATE_HOME. + + That means, by default "~/.local/state/". + """ + if system in ["win32", "darwin"]: + path = user_data_dir(appname, appauthor, None, roaming) + else: + path = os.getenv("XDG_STATE_HOME", os.path.expanduser("~/.local/state")) + if appname: + path = os.path.join(path, appname) + if appname and version: + path = os.path.join(path, version) + return path + + +def user_log_dir(appname=None, appauthor=None, version=None, opinion=True): + r"""Return full path to the user-specific log dir for this application. + + "appname" is the name of application. + If None, just the system directory is returned. + "appauthor" (only used on Windows) is the name of the + appauthor or distributing body for this application. Typically + it is the owning company name. This falls back to appname. You may + pass False to disable it. + "version" is an optional version path element to append to the + path. You might want to use this if you want multiple versions + of your app to be able to run independently. If used, this + would typically be ".". + Only applied when appname is present. + "opinion" (boolean) can be False to disable the appending of + "Logs" to the base app data dir for Windows, and "log" to the + base cache dir for Unix. See discussion below. + + Typical user log directories are: + Mac OS X: ~/Library/Logs/ + Unix: ~/.cache//log # or under $XDG_CACHE_HOME if defined + Win XP: C:\Documents and Settings\\Local Settings\Application Data\\\Logs + Vista: C:\Users\\AppData\Local\\\Logs + + On Windows the only suggestion in the MSDN docs is that local settings + go in the `CSIDL_LOCAL_APPDATA` directory. (Note: I'm interested in + examples of what some windows apps use for a logs dir.) + + OPINION: This function appends "Logs" to the `CSIDL_LOCAL_APPDATA` + value for Windows and appends "log" to the user cache dir for Unix. + This can be disabled with the `opinion=False` option. + """ + if system == "darwin": + path = os.path.join(os.path.expanduser("~/Library/Logs"), appname) + elif system == "win32": + path = user_data_dir(appname, appauthor, version) + version = False + if opinion: + path = os.path.join(path, "Logs") + else: + path = user_cache_dir(appname, appauthor, version) + version = False + if opinion: + path = os.path.join(path, "log") + if appname and version: + path = os.path.join(path, version) + return path + + +class AppDirs(object): + """Convenience wrapper for getting application dirs.""" + + def __init__( + self, appname=None, appauthor=None, version=None, roaming=False, multipath=False + ): + self.appname = appname + self.appauthor = appauthor + self.version = version + self.roaming = roaming + self.multipath = multipath + + @property + def user_data_dir(self): + return user_data_dir( + self.appname, self.appauthor, version=self.version, roaming=self.roaming + ) + + @property + def site_data_dir(self): + return site_data_dir( + self.appname, self.appauthor, version=self.version, multipath=self.multipath + ) + + @property + def user_config_dir(self): + return user_config_dir( + self.appname, self.appauthor, version=self.version, roaming=self.roaming + ) + + @property + def site_config_dir(self): + return site_config_dir( + self.appname, self.appauthor, version=self.version, multipath=self.multipath + ) + + @property + def user_cache_dir(self): + return user_cache_dir(self.appname, self.appauthor, version=self.version) + + @property + def user_state_dir(self): + return user_state_dir(self.appname, self.appauthor, version=self.version) + + @property + def user_log_dir(self): + return user_log_dir(self.appname, self.appauthor, version=self.version) + + +# ---- internal support stuff + + +def _get_win_folder_from_registry(csidl_name): + """This is a fallback technique at best. I'm not sure if using the + registry for this guarantees us the correct answer for all CSIDL_* + names. + """ + import winreg as _winreg + + shell_folder_name = { + "CSIDL_APPDATA": "AppData", + "CSIDL_COMMON_APPDATA": "Common AppData", + "CSIDL_LOCAL_APPDATA": "Local AppData", + }[csidl_name] + + key = _winreg.OpenKey( + _winreg.HKEY_CURRENT_USER, + r"Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders", + ) + dir, _type = _winreg.QueryValueEx(key, shell_folder_name) + return dir + + +def _get_win_folder_with_pywin32(csidl_name): + from win32com.shell import shell, shellcon + + dir = shell.SHGetFolderPath(0, getattr(shellcon, csidl_name), 0, 0) + # Try to make this a unicode path because SHGetFolderPath does + # not return unicode strings when there is unicode data in the + # path. + try: + dir = unicode(dir) + + # Downgrade to short path name if have highbit chars. See + # . + has_high_char = False + for c in dir: + if ord(c) > 255: + has_high_char = True + break + if has_high_char: + try: + import win32api + + dir = win32api.GetShortPathName(dir) + except ImportError: + pass + except UnicodeError: + pass + return dir + + +def _get_win_folder_with_ctypes(csidl_name): + import ctypes + + csidl_const = { + "CSIDL_APPDATA": 26, + "CSIDL_COMMON_APPDATA": 35, + "CSIDL_LOCAL_APPDATA": 28, + }[csidl_name] + + buf = ctypes.create_unicode_buffer(1024) + ctypes.windll.shell32.SHGetFolderPathW(None, csidl_const, None, 0, buf) + + # Downgrade to short path name if have highbit chars. See + # . + has_high_char = False + for c in buf: + if ord(c) > 255: + has_high_char = True + break + if has_high_char: + buf2 = ctypes.create_unicode_buffer(1024) + if ctypes.windll.kernel32.GetShortPathNameW(buf.value, buf2, 1024): + buf = buf2 + + return buf.value + + +def _get_win_folder_with_jna(csidl_name): + import array + + from com.sun import jna + from com.sun.jna.platform import win32 + + buf_size = win32.WinDef.MAX_PATH * 2 + buf = array.zeros("c", buf_size) + shell = win32.Shell32.INSTANCE + shell.SHGetFolderPath( + None, + getattr(win32.ShlObj, csidl_name), + None, + win32.ShlObj.SHGFP_TYPE_CURRENT, + buf, + ) + dir = jna.Native.toString(buf.tostring()).rstrip("\0") + + # Downgrade to short path name if have highbit chars. See + # . + has_high_char = False + for c in dir: + if ord(c) > 255: + has_high_char = True + break + if has_high_char: + buf = array.zeros("c", buf_size) + kernel = win32.Kernel32.INSTANCE + if kernel.GetShortPathName(dir, buf, buf_size): + dir = jna.Native.toString(buf.tostring()).rstrip("\0") + + return dir + + +if system == "win32": + try: + import win32com.shell + + _get_win_folder = _get_win_folder_with_pywin32 + except ImportError: + try: + from ctypes import windll + + _get_win_folder = _get_win_folder_with_ctypes + except ImportError: + try: + import com.sun.jna + + _get_win_folder = _get_win_folder_with_jna + except ImportError: + _get_win_folder = _get_win_folder_from_registry + + +# ---- self test code + +if __name__ == "__main__": + appname = "MyApp" + appauthor = "MyCompany" + + props = ( + "user_data_dir", + "user_config_dir", + "user_cache_dir", + "user_state_dir", + "user_log_dir", + "site_data_dir", + "site_config_dir", + ) + + print(f"-- app dirs {__version__} --") + + print("-- app dirs (with optional 'version')") + dirs = AppDirs(appname, appauthor, version="1.0") + for prop in props: + print(f"{prop}: {getattr(dirs, prop)}") + + print("\n-- app dirs (without optional 'version')") + dirs = AppDirs(appname, appauthor) + for prop in props: + print(f"{prop}: {getattr(dirs, prop)}") + + print("\n-- app dirs (without optional 'appauthor')") + dirs = AppDirs(appname) + for prop in props: + print(f"{prop}: {getattr(dirs, prop)}") + + print("\n-- app dirs (with disabled 'appauthor')") + dirs = AppDirs(appname, appauthor=False) + for prop in props: + print(f"{prop}: {getattr(dirs, prop)}") diff --git a/venv/lib/python3.10/site-packages/torch/_classes.py b/venv/lib/python3.10/site-packages/torch/_classes.py new file mode 100644 index 0000000000000000000000000000000000000000..a811c7c30be612bf0a89e5c8a5473e9d54d02ba3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_classes.py @@ -0,0 +1,56 @@ +import types +from typing import Any + +import torch._C + + +class _ClassNamespace(types.ModuleType): + def __init__(self, name: str) -> None: + super().__init__("torch.classes" + name) + self.name = name + + def __getattr__(self, attr: str) -> Any: + proxy = torch._C._get_custom_class_python_wrapper(self.name, attr) + if proxy is None: + raise RuntimeError(f"Class {self.name}.{attr} not registered!") + return proxy + + +class _Classes(types.ModuleType): + __file__ = "_classes.py" + + def __init__(self) -> None: + super().__init__("torch.classes") + + def __getattr__(self, name: str) -> _ClassNamespace: + namespace = _ClassNamespace(name) + setattr(self, name, namespace) + return namespace + + @property + def loaded_libraries(self) -> Any: + return torch.ops.loaded_libraries + + def load_library(self, path: str) -> None: + """ + Loads a shared library from the given path into the current process. + + The library being loaded may run global initialization code to register + custom classes with the PyTorch JIT runtime. This allows dynamically + loading custom classes. For this, you should compile your class + and the static registration code into a shared library object, and then + call ``torch.classes.load_library('path/to/libcustom.so')`` to load the + shared object. + + After the library is loaded, it is added to the + ``torch.classes.loaded_libraries`` attribute, a set that may be inspected + for the paths of all libraries loaded using this function. + + Args: + path (str): A path to a shared library to load. + """ + torch.ops.load_library(path) + + +# The classes "namespace" +classes = _Classes() diff --git a/venv/lib/python3.10/site-packages/torch/_compile.py b/venv/lib/python3.10/site-packages/torch/_compile.py new file mode 100644 index 0000000000000000000000000000000000000000..33855b44b70572f7f465ca6d2e483075ffa82be9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_compile.py @@ -0,0 +1,59 @@ +""" +APIs related to torch.compile which lazily import torch._dynamo to avoid +circular dependencies. +""" + +import functools +from typing import Callable, Literal, Optional, overload, TypeVar, Union +from typing_extensions import ParamSpec + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + + +@overload +def _disable_dynamo( + fn: Callable[_P, _T], recursive: bool = True +) -> Callable[_P, _T]: ... + + +@overload +def _disable_dynamo( + fn: Literal[None] = None, recursive: bool = True +) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ... + + +def _disable_dynamo( + fn: Optional[Callable[_P, _T]] = None, recursive: bool = True +) -> Union[Callable[_P, _T], Callable[[Callable[_P, _T]], Callable[_P, _T]]]: + """ + This API should be only used inside torch, external users should still use + torch._dynamo.disable. The main goal of this API is to avoid circular + imports issues that is common while using _dynamo.disable inside torch + itself. + + This API avoids it by lazily importing torch._dynamo from the import time to + the invocation of the decorated function. + """ + if fn is not None: + + @functools.wraps(fn) + def inner(*args: _P.args, **kwargs: _P.kwargs) -> _T: + # cache this on the first invocation to avoid adding too much overhead. + disable_fn = getattr(fn, "__dynamo_disable", None) + if disable_fn is None: + import torch._dynamo + + # We can safely turn off functools.wraps here because the inner + # already wraps fn in the outer scope. + disable_fn = torch._dynamo.disable(fn, recursive, wrapping=False) + fn.__dynamo_disable = disable_fn # type: ignore[attr-defined] + + return disable_fn(*args, **kwargs) + + return inner + else: + # decorator usage like @_disable_dynamo(recursive=False). The resulting + # object expects the original decorated function as the arg. + return functools.partial(_disable_dynamo, recursive=recursive) diff --git a/venv/lib/python3.10/site-packages/torch/_custom_ops.py b/venv/lib/python3.10/site-packages/torch/_custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..5203da640fa58e9ab7bf6785eb26ae1186059bee --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_custom_ops.py @@ -0,0 +1,326 @@ +# mypy: allow-untyped-defs +import inspect + +from torch._custom_op.impl import ( + _custom_op_with_schema, + _find_custom_op, + infer_schema, + parse_qualname, + validate_namespace, +) +from torch.library import get_ctx + + +__all__ = [ + "custom_op", + "impl", + "impl_abstract", + "get_ctx", + "impl_save_for_backward", + "impl_backward", +] + + +def custom_op(qualname, func_or_schema=None): + r"""Register a new custom operator + + In PyTorch, defining an op (short for "operator") is a two step-process: + - we need to define the op (by providing an operator name and schema) + - we need to implement behavior for how the operator interacts with + various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc. + + This entrypoint defines the custom operator (the first step) + you must then perform the second step by calling various + ``impl_*`` APIs. + + This API may be used as a decorator (see examples). + + For a detailed guide on custom ops, please see + https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk + + Arguments: + qualname (str): Should be a string that looks like + "namespace::operator_name". Operators in PyTorch need a namespace to + avoid name collisions; a given operator may only be created once. + If you are writing a Python library, we recommend the namespace to + be the name of your top-level module. + func_or_schema (Union[Callable, str]): Each PyTorch operator needs a + schema that tells PyTorch the types of the inputs/outputs. + If this is a Callable, we will automatically infer the schema from + the type annotations on the function (see examples). Otherwise, + if you don't want to use type annotations, you may provide us the + schema string. + + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> import torch + >>> import numpy as np + >>> from torch import Tensor + >>> + >>> # Step 1: define the custom op. + >>> # We need to provide the API a "prototype function" + >>> # (a function that returns NotImplementedError), from which + >>> # we will infer the types of the inputs and outputs. + >>> @torch._custom_ops.custom_op("mylibrary::numpy_sin") + >>> def numpy_sin(x: Tensor) -> Tensor: + >>> raise NotImplementedError + >>> + >>> # The custom op is now accessible via the torch.ops module: + >>> torch.ops.mylibrary.numpy_sin + >>> + >>> # Step 2: Register an implementation for various PyTorch subsystems + >>> + >>> # Register an implementation for CPU tensors + >>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cpu") + >>> def numpy_sin_impl_cpu(x): + >>> return torch.from_numpy(np.sin(x.numpy())) + >>> + >>> # Register an implementation for CUDA tensors + >>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cuda") + >>> def numpy_sin_impl_cuda(x): + >>> return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device) + >>> + >>> x = torch.randn(3) + >>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cpu + >>> + >>> x_cuda = x.cuda() + >>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cuda + + """ + ns, name = parse_qualname(qualname) + validate_namespace(ns) + + def inner(func): + if not inspect.isfunction(func): + raise ValueError( + f"custom_op(...)(func): Expected `func` to be a Python " + f"function, got: {type(func)}" + ) + + if func.__name__ != name: + raise ValueError( + f"custom_op(qualname='{qualname}', ...)(func): expected `func` " + f"to have name '{name}' but got '{func.__name__}'. " + f"Please either change the name of `func` or the qualname that " + f"is passed to `custom_op`" + ) + + schema = infer_schema(func, mutates_args=()) + _custom_op_with_schema(qualname, schema) + return func + + if func_or_schema is None: + return inner + if isinstance(func_or_schema, str): + _custom_op_with_schema(qualname, func_or_schema) + else: + return inner(func_or_schema) + + +def impl(qualname, *, device_types=("cpu", "cuda"), func=None): + r"""Register an implementation for a device type for this custom op. + + If the op is passed multiple Tensor inputs with different device + types, it will dispatch to the registered implementation for the highest + priority device type among those present. + The supported device types, in order of priority, are {'cuda', 'cpu'}. + + This API may be used as a decorator (see examples). + + For a detailed guide on custom ops, please see + https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk + + Arguments: + device_types (str or Iterable[str]): the device type(s) to register the function for. + + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> import torch + >>> import numpy as np + >>> from torch import Tensor + >>> + >>> # Step 1: define the custom op. + >>> # We need to provide the API a "prototype function" + >>> # (a function that returns NotImplementedError), from which + >>> # we will infer the types of the inputs and outputs. + >>> @torch._custom_ops.custom_op("mylibrary::numpy_cos") + >>> def numpy_cos(x: Tensor) -> Tensor: + >>> raise NotImplementedError + >>> + >>> # The custom op is now accessible via the torch.ops module: + >>> torch.ops.mylibrary.numpy_cos + >>> + >>> # Step 2: Register an implementation for various PyTorch subsystems + >>> + >>> # Register an implementation for CPU tensors + >>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cpu") + >>> def numpy_cos_impl_cpu(x): + >>> return torch.from_numpy(np.cos(x.numpy())) + >>> + >>> # Register an implementation for CUDA tensors + >>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cuda") + >>> def numpy_cos_impl_cuda(x): + >>> return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device) + >>> + >>> x = torch.randn(3) + >>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cpu + >>> + >>> x_cuda = x.cuda() + >>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cuda + + """ + + def inner(func): + custom_op = _find_custom_op(qualname, also_check_torch_library=True) + custom_op.impl(device_types, _stacklevel=3)(func) + return func + + if func is None: + return inner + return inner(func) + + +def impl_abstract(qualname, *, func=None): + r"""Register an abstract implementation for this operator. + + An "abstract implementation" specifies the behavior of this operator on + Tensors that carry no data. Given some input Tensors with certain properties + (sizes/strides/storage_offset/device), it specifies what the properties of + the output Tensors are. + + The abstract implementation has the same signature as the operator. + It is run for both FakeTensors and meta tensors. To write an abstract + implementation, assume that all Tensor inputs to the operator are + regular CPU/CUDA/Meta tensors, but they do not have storage, and + you are trying to return regular CPU/CUDA/Meta tensor(s) as output. + The abstract implementation must consist of only PyTorch operations + (and may not directly access the storage or data of any input or + intermediate Tensors). + + This API may be used as a decorator (see examples). + + For a detailed guide on custom ops, please see + https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk + + Examples:: + >>> import numpy as np + >>> from torch import Tensor + >>> + >>> # Example 1: an operator without data-dependent output shape + >>> @torch._custom_ops.custom_op("mylibrary::custom_linear") + >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor: + >>> raise NotImplementedError + >>> + >>> @torch._custom_ops.impl_abstract("mylibrary::custom_linear") + >>> def custom_linear_abstract(x, weight): + >>> assert x.dim() == 2 + >>> assert weight.dim() == 2 + >>> assert bias.dim() == 1 + >>> assert x.shape[1] == weight.shape[1] + >>> assert weight.shape[0] == bias.shape[0] + >>> assert x.device == weight.device + >>> + >>> return (x @ weight.t()) + bias + >>> + >>> # Example 2: an operator with data-dependent output shape + >>> @torch._custom_ops.custom_op('mylibrary::custom_nonzero') + >>> def custom_nonzero(x: Tensor) -> Tensor: + >>> ... + >>> + >>> @torch._custom_ops.impl_abstract("mylibrary::custom_nonzero") + >>> def custom_nonzero_abstract(x): + >>> # Number of nonzero-elements is data-dependent. + >>> # Since we cannot peek at the data in an abstract impl, + >>> # we use the ctx object to construct a new symint that + >>> # represents the data-dependent size. + >>> ctx = torch._custom_ops.get_ctx() + >>> nnz = ctx.create_unbacked_symint() + >>> shape = [x.dim(), nnz] + >>> result = x.new_empty(shape, dtype=torch.long) + >>> return result + >>> + >>> @torch._custom_ops.impl("mylibrary::custom_nonzero") + >>> def custom_nonzero_impl(x): + >>> x_np = to_numpy(x) + >>> res = np.stack(np.nonzero(x_np), axis=1) + >>> # unbacked symbolic ints in PyTorch must be >= 2, so we + >>> # constrain the range to at least 2 + >>> if res.shape[0] <= 1: + >>> raise RuntimeError("not supported") + >>> return torch.tensor(res, device=x.device) + + """ + import torch.library + + return torch.library.register_fake(qualname, func, _stacklevel=2) + + +def impl_save_for_backward(qualname, *, func=None): + r"""Register a function that tells us what to save for backward. + + Please see :func:`impl_backward` for more details. + """ + + def inner(func): + custom_op = _find_custom_op(qualname, also_check_torch_library=True) + custom_op.impl_save_for_backward(_stacklevel=3)(func) + return func + + if func is None: + return inner + return inner(func) + + +def impl_backward(qualname, output_differentiability=None, *, func=None): + r"""Registers a backward formula for an operator. + + In order for an operator to work with autograd, you need to register + a backward formula. There are two pieces to this: + 1. You must give us a function to specify what to save for backward. + Call this the "save for backward" function. + 2. You must give us a function that computes gradients. Call this the + "backward" function. + + Use `impl_save_for_backward` to define a "save for backward" function + that specifies what gets saved for backward. The function should accept + two arguments ``(inputs, output)`` and return the quantities to be saved + for backward. + + During runtime, when you call the operator in a forwards pass, PyTorch + will invoke the "save for backward" function with the inputs and output + of the operator. + + Use `impl_backward` to define the "backward" function. The backward + function must accept ``(ctx, saved, *grads)``: + - ``ctx`` is a context object where we may provide information + - ``saved`` is exactly what gets returned from the "save for backward" + function + - ``grads`` is one or more gradients. The number of gradients matches + the number of outputs of the operator. + + The backward function must return a dict that maps the name of + an input to the operator to its corresponding gradient. All inputs that + were declared to be Tensors in the operator definition must be accounted + for in the dict. The gradient may be a Tensor or None. + + For a detailed guide on custom ops, please see + https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk + + """ + + def inner(func): + custom_op = _find_custom_op(qualname, also_check_torch_library=True) + custom_op.impl_backward(output_differentiability, _stacklevel=3)(func) + return func + + if func is None: + return inner + return inner(func) + + +def _destroy(qualname): + """De-registers a custom op. For testing purposes only""" + custom_op = _find_custom_op(qualname) + custom_op._destroy() diff --git a/venv/lib/python3.10/site-packages/torch/_environment.py b/venv/lib/python3.10/site-packages/torch/_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..65cbd5d35ad5164f66f37a7d75ac29e4df5f5cfa --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_environment.py @@ -0,0 +1,2 @@ +def is_fbcode() -> bool: + return False diff --git a/venv/lib/python3.10/site-packages/torch/_guards.py b/venv/lib/python3.10/site-packages/torch/_guards.py new file mode 100644 index 0000000000000000000000000000000000000000..f6f053ea064cb7257002119c6329b71eeb2f824c --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_guards.py @@ -0,0 +1,1180 @@ +from __future__ import annotations + +import contextlib +import dataclasses +import enum +import functools +import logging +import re +import threading +import traceback +import unittest.mock +import weakref +from abc import abstractmethod +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Generic, + NamedTuple, + Optional, + TYPE_CHECKING, + TypeVar, + Union, +) + +import torch +from torch.utils import _pytree as pytree +from torch.utils._backport_slots import dataclass_slots +from torch.utils._traceback import CapturedTraceback, format_frame +from torch.utils.weak import WeakTensorKeyDictionary + + +log = logging.getLogger(__name__) + + +if TYPE_CHECKING: + from collections.abc import Generator, Iterator + from types import CodeType + + import sympy + + from torch._dynamo.backends.distributed import DDPOptimizerContext + from torch._dynamo.codegen import PyCodegen + from torch._functorch._aot_autograd.schemas import ViewAndMutationMeta + from torch._subclasses.fake_tensor import FakeTensorMode + + +""" +torch._guards is the definitional source of truth for general purpose guard structures. + +An important thing to keep in mind here is the preservation of layering. There should be no dynamo notions, +and no guard installation notions here. +""" + +COMPILE_ID_PATTERN = re.compile(r"^(?P\d+)/(?P\d+)$") +CA_COMPILE_ID_PATTERN = re.compile( + r"^!(?P\d+)(?:/(?P\d+)/(?P\d+))?$" +) + +# [Note: Updating CompiledId] +# +# CompiledId represents a unique program-level identifier, and we want to keep that +# property as the codebase evolves. This property is relied on even outside of the pytorch +# repo, e.g. tlparse or other internal tooling. The in-memory format can be freely changed, +# as those dependencies only consume the string serialization. +# +# The string form should be: +# 1. Program-level uid: CompileId can uniquely identify a compiled graph. +# 2. Storage efficient: This object is logged in nearly every entry. We should elide symbols when possible. +# 3. Compact: The string form is directly displayed by some tools. Special symbols are okay. + + +# TODO: mark as kw_only=True once we drop support for str: + # NOTE: Keep this in sync with both from_string and the tlparse repo + if self.compiled_autograd_id is not None: + assert (self.frame_id is None) == (self.frame_compile_id is None) + frame_str = "" + if self.frame_id is not None: + frame_str = f"/{self.frame_id}/{self.frame_compile_id}" + + return f"!{self.compiled_autograd_id}{frame_str}" + else: + assert self.frame_id is not None and self.frame_compile_id is not None + return f"{self.frame_id}/{self.frame_compile_id}" + + @classmethod + def from_string(cls, compile_id: Optional[str]) -> Optional[CompileId]: + """ + Factory method that creates a CompileId from its string representation. + Keep this in sync with the __str__ method. + """ + if compile_id is None: + return None + try: + for pattern in (COMPILE_ID_PATTERN, CA_COMPILE_ID_PATTERN): + if match := pattern.match(compile_id): + groups = match.groupdict() + for k, v in groups.items(): + if v is not None: + groups[k] = int(v) + return cls(**groups) # type: ignore[arg-type] + else: + raise ValueError + + except Exception as e: + raise ValueError(f"Invalid compile_id '{compile_id}'") from e + + +class TraceId(NamedTuple): + compile_id: CompileId + # This starts off as 0, and every time we restart analysis it goes + # up by one + attempt: int + + def __str__(self) -> str: + # Keep this in sync with tlparse repo + if self.attempt == 0: + return str(self.compile_id) + else: + return f"{self.compile_id}_{self.attempt}" + + +class GuardSource(enum.Enum): + LOCAL = 0 + GLOBAL = 1 + LOCAL_SPECIALIZED_NN_MODULE = 2 + GLOBAL_SPECIALIZED_NN_MODULE = 3 + CONSTANT = 4 + RANDOM_VALUE = 5 + SHAPE_ENV = 6 + LOCAL_FSDP_MODULE = 7 + GLOBAL_FSDP_MODULE = 8 + BACKWARD_STATE = 9 + EPHEMERAL = 10 + SYNTHETIC_LOCAL = 11 + LOCAL_UNSPECIALIZED_NN_MODULE = 12 + GLOBAL_UNSPECIALIZED_NN_MODULE = 13 + LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE = 14 + GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULE = 15 + + def is_fsdp_module(self) -> bool: + return self in (GuardSource.GLOBAL_FSDP_MODULE, GuardSource.LOCAL_FSDP_MODULE) + + def is_specialized_nn_module(self) -> bool: + import torch._dynamo.config as config + + if config._unsafe_skip_fsdp_module_guards: + return ( + self + in ( + GuardSource.GLOBAL_SPECIALIZED_NN_MODULE, + GuardSource.LOCAL_SPECIALIZED_NN_MODULE, + ) + or self.is_fsdp_module() + ) + return self in ( + GuardSource.GLOBAL_SPECIALIZED_NN_MODULE, + GuardSource.LOCAL_SPECIALIZED_NN_MODULE, + ) + + def is_unspecialized_nn_module(self) -> bool: + return self in ( + GuardSource.GLOBAL_UNSPECIALIZED_NN_MODULE, + GuardSource.LOCAL_UNSPECIALIZED_NN_MODULE, + GuardSource.GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULE, + GuardSource.LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE, + ) + + def is_unspecialized_builtin_nn_module(self) -> bool: + return self in ( + GuardSource.GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULE, + GuardSource.LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE, + ) + + def is_local(self) -> bool: + return self in ( + GuardSource.LOCAL, + GuardSource.LOCAL_SPECIALIZED_NN_MODULE, + GuardSource.LOCAL_FSDP_MODULE, + GuardSource.LOCAL_UNSPECIALIZED_NN_MODULE, + GuardSource.LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE, + ) + + +""" +Base class for a "GuardBuilder" role. + +The GuardBuilderBase role is to represent a scope within which to build a guard. The name is a little +confusing, as its not a builder, but for the sake of avoiding a lot of renames and keeping the original reference +to torchdynamo's GuardBuilder. + +Note: create_fn is invoked with a GuardBuilderBase and a Guard. A GuardBuilder is chosen based +on GuardSource's select function. + +There is value in keeping this GuardBuilderBase empty to keep layering clean. +""" + + +class GuardBuilderBase: + pass + + +@dataclasses.dataclass(frozen=True) +class SLoc: + framework_loc: Optional[Union[traceback.FrameSummary, str]] + maybe_user_loc: Optional[str] + + def __str__(self) -> str: + floc = ( + self.framework_loc + if isinstance(self.framework_loc, str) + else format_frame(self.framework_loc) + ) + if self.maybe_user_loc is not None: + return f"{self.maybe_user_loc} ({floc})" + else: + return f"({floc})" + + +class ShapeGuard(NamedTuple): + expr: sympy.logic.boolalg.Boolean + sloc: SLoc + size_oblivious: bool + + +@dataclass_slots +@dataclasses.dataclass +class Guard: + # originating_source is the source that called the make_guard method to + # construct this guard object. The property name specifies what exactly it + # is the guard is guarding on. The meaning of the name is dependent on the + # create_fn; you must look at the use-site inside create_fn to know what + # name means. + # + # That being said, although you might think this is just a "name", name is + # usually an arbitrary Python expression that will be evaluated with all + # globals (and locals, if you create a LOCAL guard) to extract the Python + # object that we want to perform guard tests on. This evaluation + # typically happens in GuardBuilder.eval. In these cases, name is + # typically produced by originating_source.name() (not to be confused with + # GuardSource - the property source). + # + # Occasionally, name is not a valid Python expression; sometimes + # it is meaningless. Example create_fns that are like this include + # GRAD_MODE and SHAPE_ENV. + originating_source: Source + create_fn: Callable[[GuardBuilderBase, Guard], None] + + # Export only. These values are written to at time of guard check_fn creation. + guard_types: Optional[list[str]] = None + code_list: Optional[list[str]] = None + obj_weakref: Optional[object] = None + guarded_class_weakref: Optional[weakref.ReferenceType[Any]] = None + + stack: Optional[CapturedTraceback] = None + user_stack: Optional[traceback.StackSummary] = None + _hash: Optional[int] = None + _unserializable: bool = False + + def __hash__(self) -> int: + if self._hash is None: + self._hash = hash((self.name, self.source, id(self.create_fn))) + return self._hash + + def sort_key(self) -> tuple[bool, int, int, str, int]: + # Put the duplicate input guards at the end. The duplicate guards have + # two sources while guard.name only considers one source. + + is_duplicate_input = ( + isinstance(self.create_fn, functools.partial) + and self.create_fn.func is torch._dynamo.guards.GuardBuilder.DUPLICATE_INPUT + ) + return ( + is_duplicate_input, + self.source.value if self.source else -1, + len(self.name), + self.name, + self.inner_create_fn().__code__.co_firstlineno, + ) + + def __lt__(self, other: Guard) -> bool: + return self.sort_key() < other.sort_key() + + def inner_create_fn(self) -> Callable[[GuardBuilderBase, Guard], Any]: + if isinstance(self.create_fn, functools.partial): + return self.create_fn.func + else: + return self.create_fn + + @property + def name(self) -> str: + return self.originating_source.name() + + @property + def source(self) -> GuardSource: + return self.originating_source.guard_source() + + @staticmethod + def weakref_to_str(obj_weakref: object) -> str: + """ + This is a workaround of a Python weakref bug. + + `obj_weakref` is instance returned by `weakref.ref`, + `str(obj_weakref)` is buggy if the original obj overrides __getattr__, e.g: + + class MyConfig(dict): + def __getattr__(self, x): + return self[x] + + obj = MyConfig(offset=5) + obj_weakref = weakref.ref(obj) + str(obj_weakref) # raise error: KeyError: '__name__' + """ + if isinstance(obj_weakref, weakref.ReferenceType): + obj = obj_weakref() + if obj is not None: + return f"" + else: + return f"" + else: + return str(obj_weakref) + + def __repr__(self) -> str: + s = f""" + {self.source.name.lower() if self.source else ""} {repr(self.name)} {self.inner_create_fn().__name__} + {{ + 'guard_types': {self.guard_types}, + 'code': {self.code_list}, + 'obj_weakref': {self.weakref_to_str(self.obj_weakref)} + 'guarded_class': {self.guarded_class_weakref} + }} + """ + return s + + def __str__(self) -> str: + output = f"Name: {repr(self.name)}\n" + source = self.source.name.lower() if self.source else "" + output += f" Source: {source}\n" + output += f" Create Function: {self.inner_create_fn().__name__}\n" + output += f" Guard Types: {self.guard_types}\n" + output += f" Code List: {self.code_list}\n" + output += f" Object Weakref: {self.weakref_to_str(self.obj_weakref)}\n" + output += f" Guarded Class Weakref: {self.guarded_class_weakref}\n" + return output + + def create(self, builder: GuardBuilderBase) -> Any: + try: + return self.create_fn(builder, self) + except Exception: + log.exception("Error while creating guard:\n%s", str(self).rstrip()) + if self.stack: + log.error("Created at:\n%s", "".join(self.stack.format()[-4:]).rstrip()) + raise + + def is_specialized_nn_module(self) -> bool: + return self.source.is_specialized_nn_module() + + def is_fsdp_module(self) -> bool: + return self.source.is_fsdp_module() + + def is_local(self) -> bool: + return self.source.is_local() + + def create_fn_name(self) -> str: + if isinstance(self.create_fn, functools.partial): + create_fn = self.create_fn.func # type: ignore[attr-defined] + else: + create_fn = self.create_fn + return create_fn.__name__ + + def set_export_info( + self, + guard_type: str, + guarded_class: Optional[weakref.ReferenceType[Any]], + code_list: list[str], + obj_weakref: object, + ) -> None: + if not self.guard_types: + self.guard_types = [] + + self.guard_types.append(guard_type) + + assert self.guarded_class_weakref in ( + guarded_class, + None, + ), "Guarded class id must be identical, or None" + self.guarded_class_weakref = guarded_class + + if not self.code_list: + self.code_list = code_list + else: + self.code_list.extend(code_list) + + # Some objects are ephemeral, e.g., list[slice(1, 2)]. If we have + # multiple guards on the same object, the weakref can die between the + # invocation of set_export_info calls. So a dead weakref is also + # acceptable. + assert ( + self.obj_weakref in (obj_weakref, None) + or callable(self.obj_weakref) + and self.obj_weakref() is None + ), "Guarded object must be identical, None or ephemeral (dead weakref)" + self.obj_weakref = obj_weakref + + +T = TypeVar("T") + +""" +Parent structure for guard env expressions. +A GuardEnvExpr can have any subtype. +Note: All subtypes must be handled exhaustively in +torch._dynamo.guards._parse_guard_env_guards to avoid a RuntimeError. +""" + + +@dataclasses.dataclass(frozen=True) +class GuardEnvExpr: + pass + + +""" +A class representing a pair of duplicate inputs. +input_pos_a and input_pos_b are input positions we have deduped. +""" + + +@dataclasses.dataclass(frozen=True) +class DuplicateInputs(GuardEnvExpr): + input_source_a: Source + input_source_b: Source + + def __post_init__(self) -> None: + assert self.input_source_a != self.input_source_b + + +""" +A class representing storage overlap relations among inputs that aliases the same storage. + +Given that a set of tensors alias the same storage, this guard checks whether they actually +have overlapping storages. + +While non_overlapping_sources represent input tensors that definitely don't have any storage +overlapping with any other input, overlapping_sources represent tensors that either: + +1. Do overlap some other input tensor +2. Might not overlap some other input tensor, but we are not sure +""" + + +@dataclasses.dataclass(frozen=True) +class StorageOverlap(GuardEnvExpr): + overlapping_sources: list[Source] + non_overlapping_sources: list[Source] + + +""" +Checkpointable is an interface for driving state snapshotting, left purposely vague for now. + +copy_graphstate() -> T, a somewhat legacy name, is expected to emit a snapshot of any type that +can also be taken in at restore_graphstate(T) calls. + +When to snapshot, is, at the moment, an implementation detail of upstream callers. Checkpointable +does not provide any guarantees around consistency, idempotency, or safety of calling its APIs, yet. + +In the future, it will have a closer coupling to a generic Checkpoint management system. +""" + + +class Checkpointable(Generic[T]): + @abstractmethod + def copy_graphstate(self) -> T: ... + + @abstractmethod + def restore_graphstate(self, state: T) -> None: ... + + +class GuardsCheckpointState: + """ + The GuardCheckpointState - it is the T of Checkpointable[T] for GuardsContext + """ + + dynamo_guards: set[Guard] = set() + + def __init__(self, dynamo_guards: set[Guard]) -> None: + self.dynamo_guards = dynamo_guards + + def diff(self, other: GuardsCheckpointState) -> Optional[set[Guard]]: + """ + Produces a delta against another GuardsCheckpointState. + + Returns None if no delta is found, otherwise, return a set() of mismatched + Guard type objects. + """ + r = self.dynamo_guards.difference(other.dynamo_guards) + if len(r) == 0: + return None + return r + + def __eq__(self, other: object) -> bool: + if not isinstance(other, GuardsCheckpointState): + return False + return self.diff(other) is None + + +class ModuleContextCheckpointState: + nn_modules: dict[str, torch.nn.Module] = {} + + def __init__(self, nn_modules: dict[str, torch.nn.Module]) -> None: + self.nn_modules = nn_modules + + def diff(self, other: ModuleContextCheckpointState) -> Optional[set[str]]: + """ + Produces a delta against another ModuleContextCheckpointState. + + Returns None if no delta is found, otherwise, return a set() of mismatched + module key names. + """ + r = set(self.nn_modules.keys()).difference(set(other.nn_modules.keys())) + if len(r) == 0: + return None + return r + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ModuleContextCheckpointState): + return False + return self.diff(other) is None + + +class ModuleContext(Checkpointable[ModuleContextCheckpointState]): + def __init__(self) -> None: + self.nn_modules: dict[str, Any] = {} + + def copy_graphstate(self) -> ModuleContextCheckpointState: + return ModuleContextCheckpointState(dict(self.nn_modules)) + + def restore_graphstate(self, state: ModuleContextCheckpointState) -> None: + assert isinstance(state, ModuleContextCheckpointState) + self.nn_modules = state.nn_modules + + +class GlobalContextCheckpointState: + global_state: dict[str, tuple[Callable, Any]] = {} + + def __init__(self, global_states: dict[str, tuple[Callable, Any]]) -> None: + self.global_state = global_states + + def diff(self, other: GlobalContextCheckpointState) -> Optional[set[str]]: + """ + Produces a delta against another GlobalContextCheckpointState. + + Returns None if no delta is found, otherwise, return a set() of mismatched + global key names. + """ + r = set(self.global_state.keys()).difference(set(other.global_state.keys())) + if len(r) == 0: + return None + return r + + def __eq__(self, other: object) -> bool: + if not isinstance(other, GlobalContextCheckpointState): + return False + return self.diff(other) is None + + +class GlobalContext(Checkpointable[GlobalContextCheckpointState]): + """ + This keeps track of the global torch state during tracing of a function. + For example, torch.is_grad_enabled. + """ + + _supported_global_states = { + "grad_enabled", + "autocast_enabled", + "autocast_cpu_enabled", + "autocast_gpu_dtype", + "autocast_cpu_dtype", + "autocast_cache_enabled", + } + + def __init__(self) -> None: + self.global_state: dict[str, tuple[Callable, Any]] = {} + + def copy_graphstate(self) -> GlobalContextCheckpointState: + return GlobalContextCheckpointState(self.global_state) + + def restore_graphstate(self, state: GlobalContextCheckpointState) -> None: + assert isinstance(state, GlobalContextCheckpointState) + self.global_state = state.global_state + assert ( + len(self.global_state) == len(self._supported_global_states) + and set(self.global_state.keys()) == self._supported_global_states + ), "Global state mismatch" + for func, args in self.global_state.values(): + func(args) + + +# Like a Set[Guard] but will record the user stack on all guards at the +# time they were installed at their destination +class GuardsSet: + def __init__(self, inner: Optional[set[Guard]] = None) -> None: + if inner is None: + inner = set() + self.inner = inner + + def __iter__(self) -> Iterator[Guard]: + return iter(self.inner) + + def __len__(self) -> int: + return len(self.inner) + + # Subtraction along with bool is typically used to determine the delta of + # added guards between checkpoints for higher order ops + def __sub__(self, other: GuardsSet) -> GuardsSet: + return GuardsSet(self.inner - other.inner) + + def __bool__(self) -> bool: + return bool(self.inner) + + def add( + self, guard: Guard, *, collect_debug_stack: bool = True, skip: int = 0 + ) -> None: + if guard in self.inner: + return + if collect_debug_stack: + if guard.stack is None: + guard.stack = CapturedTraceback.extract(skip=1 + skip) + if guard.user_stack is None: + guard.user_stack = TracingContext.extract_stack() + self.inner.add(guard) + + def update(self, *others: set[Guard]) -> None: + for o in others: + for g in o: + self.add(g, skip=1) + + def remove_guards_with_source(self, source: Source) -> None: + """Delete all guards that contains a given source""" + from ._dynamo.source import is_from_source + + self.inner = { + g for g in self.inner if not is_from_source(g.originating_source, source) + } + + +""" +A GuardsContext is a checkpointable representation of all the guards in the current tracing +context. It's lifecycle is bound 1:1 to the tracing context, and it should never be instantiated +directly outside of it. For passing around internal state representations of this object, +prefer to extract them with copy_graphstate to produce a GuardsCheckpointState. +""" + + +class GuardsContext(Checkpointable[GuardsCheckpointState]): + def __init__(self) -> None: + self.dynamo_guards: GuardsSet = GuardsSet() + self.aotautograd_guards: list[GuardEnvExpr] = [] + + def copy_graphstate(self) -> GuardsCheckpointState: + return GuardsCheckpointState(set(self.dynamo_guards.inner)) + + def restore_graphstate(self, state: GuardsCheckpointState) -> None: + # NB: "steals" the passed in state + assert isinstance(state, GuardsCheckpointState) + self.dynamo_guards = GuardsSet(state.dynamo_guards) + + +class HopSubgraphCache: + @abstractmethod + def add_dynamo_installed_submodule(self, fn_id: int, identifier: str) -> None: ... + + @abstractmethod + def get_dynamo_installed_submodules(self, fn_id: int) -> list[str]: ... + + @abstractmethod + def add_autograd_key_entry(self, identifier: str, key: Callable) -> None: ... + + @abstractmethod + def get_autograd_key_entry(self, identifier: str) -> Optional[Callable]: ... + + @abstractmethod + def add_proxy_dispatch_entry(self, identifier: str, key: Callable) -> None: ... + + @abstractmethod + def get_proxy_dispatch_entry(self, identifier: str) -> Optional[Callable]: ... + + @abstractmethod + def add_lazy_bwd_entry( + self, + identifier: str, + tangent_metadata: tuple[object], + gmod: torch.fx.GraphModule, + ) -> int: ... + + @abstractmethod + def get_lazy_bwd_entry( + self, identifier: str, tangent_metadata: tuple[object] + ) -> tuple[Optional[torch.fx.GraphModule], Optional[int]]: ... + + +class InvokeSubgraphCache(HopSubgraphCache): + def __init__(self) -> None: + self.autograd_cache: dict[str, Callable] = {} + self.proxy_dispatch_cache: dict[str, Callable] = {} + self.dynamo_installed_submodules: dict[int, list[str]] = defaultdict(list) + self.lazy_bwd_cache: dict[ + str, dict[tuple[object], tuple[torch.fx.GraphModule, int]] + ] = defaultdict(dict) + + def add_dynamo_installed_submodule(self, fn_id: int, identifier: str) -> None: + self.dynamo_installed_submodules[fn_id].append(identifier) + + def get_dynamo_installed_submodules(self, fn_id: int) -> list[str]: + return self.dynamo_installed_submodules.get(fn_id, []) + + def add_autograd_key_entry(self, identifier: str, key: Callable) -> None: + self.autograd_cache[identifier] = key + + def get_autograd_key_entry(self, identifier: str) -> Optional[Callable]: + return self.autograd_cache.get(identifier, None) + + def add_proxy_dispatch_entry(self, identifier: str, key: Callable) -> None: + self.proxy_dispatch_cache[identifier] = key + + def get_proxy_dispatch_entry(self, identifier: str) -> Optional[Callable]: + return self.proxy_dispatch_cache.get(identifier, None) + + def add_lazy_bwd_entry( + self, + identifier: str, + tangent_metadata: tuple[object], + gmod: torch.fx.GraphModule, + ) -> int: + # Save the number of existing graph modules in the dictionary to get the suffix + num_gmods = len(self.lazy_bwd_cache[identifier]) + self.lazy_bwd_cache[identifier][tangent_metadata] = (gmod, num_gmods) + return num_gmods + + def get_lazy_bwd_entry( + self, identifier: str, tangent_metadata: tuple[object] + ) -> tuple[Optional[torch.fx.GraphModule], Optional[int]]: + if identifier not in self.lazy_bwd_cache: + return (None, None) + + return self.lazy_bwd_cache[identifier].get(tangent_metadata, (None, None)) + + +class HopDispatchSetCache: + def __init__(self) -> None: + # Delayed import to avoid circular dependency + from torch._higher_order_ops.invoke_subgraph import invoke_subgraph + + self.hop_cache_map = {invoke_subgraph: InvokeSubgraphCache()} + + def get_cache( + self, op: torch._ops.HigherOrderOperator + ) -> Optional[HopSubgraphCache]: + if op not in self.hop_cache_map: + return None + return self.hop_cache_map[op] # type: ignore[index] + + +_TLS = threading.local() + +""" +TracingContext is the source of truth for all currently accumulated information +needed to trace. Its lifecycle is kept 1:1 when using TorchDynamo, but other systems +are open to managing their own TracingContext with that in mind. + +The purpose of TracingContext is not to be a dumping ground, or god object, but rather to avoid +having to plumb complex subsystems across multiple verticals. + +Ex: A common example is guard accumulation between dynamo, shape_env, aot_autograd, and inductor. +Accessing the current tracing context via +TracingContext.get() allows users to accumulate their own guards for processing, without needing to know how +to plumb objects back up to where frame interpretation happened. + +Note that you can end up with multiple TracingContext for a single compilation +of a frame, as we reset the TracingContext whenever we restart analysis. +CompileContext is a more overarching context that encompasses multiple restarts. +""" + + +class CompileContext: + @staticmethod + def get() -> CompileContext: + assert _TLS.compile_context is not None + return _TLS.compile_context + + @staticmethod + def try_get() -> Optional[CompileContext]: + return getattr(_TLS, "compile_context", None) + + def __init__(self, compile_id: Optional[CompileId]) -> None: + assert compile_id is None or isinstance(compile_id, CompileId) + self.compile_id: Optional[CompileId] = compile_id + self.attempt = 0 + # Verbose ShapeEnv guards produced. + self.shape_env_guards: list[str] = [] + + @staticmethod + def current_compile_id() -> Optional[CompileId]: + self = CompileContext.try_get() + if self is None: + return None + return self.compile_id + + @staticmethod + def current_trace_id() -> Optional[TraceId]: + self = CompileContext.try_get() + if self is None: + return None + if self.compile_id is None: + return None + return TraceId(self.compile_id, self.attempt) + + +class TracingContext: + """ + Provides the currently installed TracingContext, or None. + + Note that it is a staticmethod, and invocations outside of `with tracing()` (see below), are valid but + will return None. + """ + + @staticmethod + def try_get() -> Optional[TracingContext]: + return getattr(_TLS, "tracing_context", None) + + @staticmethod + def get() -> TracingContext: + if ctx := TracingContext.try_get(): + return ctx + raise RuntimeError( + "TracingContext.get() must be called within an ongoing trace." + ) + + def __init__(self, fake_mode: Optional[FakeTensorMode]) -> None: + self.guards_context = GuardsContext() + self.module_context = ModuleContext() + self.global_context = GlobalContext() + self.previously_inlined_functions: dict[Any, Any] = dict() + self.previously_cleaned_instructions: dict[Any, Any] = dict() + self.fake_mode: Optional[FakeTensorMode] = fake_mode + self.frame_summary_stack: list[traceback.FrameSummary] = [] + # This is morally part of frame_summary_stack, but it is kept separate + # for clarity. As we process a frame, this variable gets updated + # to keep track of what line we are in the function. We make a + # function call, this gets cleared and the frame location is pushed + # to frame_summary_stack (prepping this variable for the inner frame's + # progress) + self.loc_in_frame: Optional[tuple[str, int, str]] = None + # this is only set after aot_autograd + self.fw_metadata: Optional[ViewAndMutationMeta] = None + # this is only set when the DDPOptimizer is used + self.ddp_optimizer_ctx: Optional[DDPOptimizerContext] = None + # this is only set after aot_autograd + self.aot_graph_name: Optional[list[str]] = None + self.params_flat: Optional[list[Any]] = None + self.params_flat_unwrap_subclasses: Optional[list[Any]] = None + self.params_unwrapped_to_flat_index: Optional[list[Any]] = None + # this is for extended return calling convention from backend + # compiler to aot_autograd + # Per output, what the compiler specified stride of the output is, + # or None if no stride is known. This is always the HINT, it + # is never a SymInt (it would be better if it was a SymInt, but + # I can't conveniently get this from Inductor atm. Also, be + # careful not to accidentally induce guards on the SymInt if + # you ever do change this in aot_autograd.py; you should check + # on permutations preferentially.) + self.output_strides: Optional[list[Optional[tuple[int, ...]]]] = None + # When this is True, whenever we encounter an int in Dynamo tracing, + # we will (1) force unspec it and (2) force it as a size-like unbacked + # integer. This is currently used when processing certain lists of + # ints that are known to be size-like and may have 0/1 entries that we + # must not specialize on. + self.force_unspec_int_unbacked_size_like = False + # See note [Tensor Fakification and Symbol Caching] + self.tensor_to_context = WeakTensorKeyDictionary() + + # If this true, Aot Autograd will return output Fake Tensors with appropriate + # meta on the first invocation + # see note: [Returning Fake Tensors on First AOT Autograd Call] + self.fakify_first_call = False + self.hop_dispatch_set_cache = HopDispatchSetCache() + # list of code objects for inlined functions + self.traced_code: list[CodeType] = [] + + def clear(self) -> None: + # Look at the note in output_graph.py in function `save_global_state` + # for the context on clearing global context. + self.global_context.global_state = {} + self.previously_inlined_functions.clear() + self.previously_cleaned_instructions.clear() + + @staticmethod + @contextmanager + def patch(**kwargs: Any) -> Generator[None, None, None]: + prior = {} + ctx = TracingContext.get() + + for key in kwargs.keys(): + # KeyError on invalid entry + prior[key] = getattr(ctx, key) + for key, val in kwargs.items(): + setattr(ctx, key, val) + try: + yield + finally: + for key, val in prior.items(): + setattr(ctx, key, val) + + @staticmethod + def extract_stack() -> traceback.StackSummary: + self = TracingContext.try_get() + if self is None: + return traceback.StackSummary() + stack = self.frame_summary_stack + if self.loc_in_frame is not None: + stack = stack + [self._populate_loc_in_frame_summary()] + return traceback.StackSummary.from_list(stack) + + def _populate_loc_in_frame_summary(self) -> traceback.FrameSummary: + assert self.loc_in_frame is not None + filename, lineno, frame_name = self.loc_in_frame + return traceback.FrameSummary(filename, lineno, frame_name, lookup_line=False) + + # Call this when you want to call into some code that isn't necessarily + # associated with the current frame state + @staticmethod + @contextlib.contextmanager + def clear_frame() -> Generator[None, None, None]: + tc = TracingContext.get() + with ( + unittest.mock.patch.object(tc, "frame_summary_stack", []), + unittest.mock.patch.object(tc, "loc_in_frame", None), + ): + try: + yield + except Exception as e: + # Prevent real_stack from getting attached + # + # The invariant is that if an Exception as real_stack, we've + # appropriately attached a user stack and we no longer need to + # attach anything. Because we cannot conveniently interpose + # when an exception is thrown, we instead interpose everywhere + # we set what the user stack is set (using the context + # manager). However, our compiler stack does "tail calls" + # (when it calls into user compiler), at which point the + # parent exception frames would incorrectly attach an + # incorrect frame. + # + # However, if, somehow, someone raised an exception with this + # scope that had a stack (for example, because they are + # restoring the user stack state appropriately as they process + # node by node), we should respect it. Thus, we cannot + # unconditionally set None. + if not hasattr(e, "real_stack"): + e.real_stack = None # type: ignore[attr-defined] + raise + + @staticmethod + @contextlib.contextmanager + def current_frame( + frame_summary: Optional[traceback.FrameSummary], + ) -> Generator[None, None, None]: + # frame_summary can be None to solely take advantage of real_stack + # attachment to thrown exceptions + tc = TracingContext.get() + if frame_summary is not None: + tc.frame_summary_stack.append(frame_summary) + old = tc.loc_in_frame + tc.loc_in_frame = None + try: + yield + except Exception as e: + if not hasattr(e, "real_stack"): + e.real_stack = tc.extract_stack() # type: ignore[attr-defined] + raise + finally: + if frame_summary is not None: + tc.frame_summary_stack.pop() + tc.loc_in_frame = old + + @staticmethod + @contextlib.contextmanager + def report_output_strides() -> Generator[ + Optional[list[Optional[tuple[int, ...]]]], None, None + ]: + tc = TracingContext.try_get() + if tc is None: + yield None + return + old_output_strides = tc.output_strides + tc.output_strides = [] + try: + yield tc.output_strides + finally: + tc.output_strides = old_output_strides + + @staticmethod + def set_current_loc(filename: str, lineno: int, frame_name: str) -> None: + # Save the current location in the frame. Lazily generate the + # framesummary. + TracingContext.get().loc_in_frame = (filename, lineno, frame_name) + + @staticmethod + def get_traced_code() -> Optional[list[CodeType]]: + tc = TracingContext.try_get() + if tc is None: + return None + return tc.traced_code + + +@contextmanager +def compile_context( + context: Optional[CompileContext], +) -> Generator[Optional[CompileContext], None, None]: + old_context = getattr(_TLS, "compile_context", None) + _TLS.compile_context = context + try: + yield context + finally: + _TLS.compile_context = old_context + + +@contextmanager +def tracing( + context: Optional[TracingContext], +) -> Generator[Optional[TracingContext], None, None]: + """ + This function installs the passed in tracing context as a dynamic scoped + global variable. + + Calls to TracingContext.get() while not under a `with tracing()` context + will return None. + """ + old_context = getattr(_TLS, "tracing_context", None) + _TLS.tracing_context = context + try: + yield context + except Exception as e: + if not hasattr(e, "real_stack") and context is not None: + e.real_stack = context.extract_stack() # type: ignore[attr-defined] + raise + finally: + if ( + context is not None + and context.fake_mode is not None + and context.fake_mode.shape_env is not None + ): + context.fake_mode.shape_env.cleanup() + _TLS.tracing_context = old_context + + +# Subclasses can be found in torch/_dynamo/source.py +# TODO(voz): Consider a toplevel torch/_source.py +@dataclasses.dataclass(frozen=True) +class Source: + def is_dict_key(self) -> bool: + return False + + def is_ephemeral(self) -> bool: + return False + + def reconstruct(self, codegen: PyCodegen) -> None: + raise NotImplementedError + + def guard_source(self) -> GuardSource: + raise NotImplementedError + + def name(self) -> str: + raise NotImplementedError + + def make_guard(self, fn: Callable[..., Any]) -> Guard: + if self.guard_source() is GuardSource.CONSTANT: + raise NotImplementedError + return Guard(self, fn) + + def is_specialized_nn_module(self) -> bool: + return self.guard_source().is_specialized_nn_module() + + def subguards_allowed(self) -> bool: + """True if you can guard on attributes of this""" + return self.guard_source() != GuardSource.SYNTHETIC_LOCAL + + +# Subclasses can be found in torch/_dynamo/source.py +@dataclasses.dataclass(frozen=True) +class ChainedSource(Source): + base: Source + + def is_dict_key(self) -> bool: + # Recurse until you either hit a ConstDictKey or a Source + return self.base.is_dict_key() + + def is_ephemeral(self) -> bool: + return self.base.is_ephemeral() + + def get_base(self) -> Source: + current: Source = self + while isinstance(current, ChainedSource): + current = current.base + return current + + +def detect_fake_mode(inputs: Any = None) -> Optional[FakeTensorMode]: + """ + Attempts to "detect" what the current fake mode is. If there is one ambiently + available from TracingContext, we preferentially use that. Otherwise, we + heuristically detect the fake mode via the following sources, in order of + priority: + + - Currently active fake mode on stack + - Fake mode associated with passed in tensors (inputs does not + have to be flattened) + """ + from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode + + fake_modes = [] + + if context := TracingContext.try_get(): + fake_mode = context.fake_mode + if fake_mode is not None: + fake_modes.append((fake_mode, "tracing context", 0)) + + from torch.utils._python_dispatch import _get_current_dispatch_mode_stack + + for i, m in enumerate(reversed(_get_current_dispatch_mode_stack())): + if isinstance(m, FakeTensorMode): + fake_modes.append((m, "active fake mode", i)) + + flat_inputs = pytree.tree_leaves(inputs) + for i, flat_input in enumerate(flat_inputs): + if isinstance(flat_input, FakeTensor): + fake_modes.append((flat_input.fake_mode, "fake tensor input", i)) + + if fake_modes: + fake_mode, desc1, i1 = fake_modes[0] + for m, desc2, i2 in fake_modes[1:]: + assert fake_mode is m, ( + f"fake mode ({fake_mode}) from {desc1} {i1} doesn't match mode ({m}) from {desc2} {i2}\n\n" + f"fake mode from {desc1} {i1} allocated at:\n{fake_mode.stack}\n" + f"fake mode from {desc2} {i2} allocated at:\n{m.stack}" + ) + return fake_mode + else: + return None + + +def active_fake_mode() -> Optional[FakeTensorMode]: + """ + Inspects the dispatch mode stack for an active fake mode and returns it. + Returns None if no fake mode is active. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.utils._python_dispatch import _get_current_dispatch_mode_stack + + for _, m in enumerate(reversed(_get_current_dispatch_mode_stack())): + if isinstance(m, FakeTensorMode): + return m + + return None diff --git a/venv/lib/python3.10/site-packages/torch/_jit_internal.py b/venv/lib/python3.10/site-packages/torch/_jit_internal.py new file mode 100644 index 0000000000000000000000000000000000000000..be6d23bbbc53cb451be92ddbc3bfa7af5c436ada --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_jit_internal.py @@ -0,0 +1,1552 @@ +# mypy: allow-untyped-defs +""" +The weak_script annotation needs to be here instead of inside torch/jit/ so it +can be used in other places in torch/ (namely torch.nn) without running into +circular dependency problems +""" + +import ast +import builtins +import collections +import contextlib +import enum +import inspect +import io +import pickle +import sys +import textwrap +import threading +import types +import typing +import warnings +import weakref +from typing import ( # noqa: UP035, F401 # (Dict, List, Tuple) imported by torch.jit.annotations + Any, + Callable, + Dict, + Final, + ForwardRef, + get_args, + get_origin, + List, + Optional, + Tuple, + TypeVar, + Union, +) +from typing_extensions import ParamSpec + +import torch + +# This is needed. `torch._jit_internal` is imported before `torch.distributed.__init__`. +# Explicitly ask to import `torch.distributed.__init__` first. +# Otherwise, "AttributeError: module 'torch' has no attribute 'distributed'" is raised. +import torch.distributed.rpc +import torch.package._mangling as package_mangling +from torch._awaits import _Await +from torch._C import _Await as CAwait, Future as CFuture +from torch._sources import fake_range, get_source_lines_and_file, parse_def +from torch.futures import Future + + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +IS_PY310_PLUS: Final[bool] = sys.version_info >= (3, 10) + +BuiltinUnionType: Union[type, tuple[type, ...]] +if sys.version_info >= (3, 10): + # NOTE: IS_PY310_PLUS doesn't work with mypy. + # cf. https://mypy.readthedocs.io/en/stable/common_issues.html#python-version-and-system-platform-checks + BuiltinUnionType = types.UnionType +else: + BuiltinUnionType = () # trick: this makes isinstance short circuit. + +LockType: type +try: + import _thread + + LockType = _thread.LockType +except ImportError: + import _dummy_thread # type: ignore[import-not-found] + + LockType = _dummy_thread.LockType + +# Wrapper functions that can call either of 2 functions depending on a boolean +# argument +boolean_dispatched: "weakref.WeakKeyDictionary[Callable, dict[str, Callable]]" = ( + weakref.WeakKeyDictionary() +) # noqa: T484 + + +FAKE_FILENAME_PREFIX = "__torch_jit_dataclass" + + +def is_final(ann) -> bool: + return ( + hasattr(ann, "__module__") + and ann.__module__ in {"typing", "typing_extensions"} + and (get_origin(ann) is Final or isinstance(ann, type(Final))) + ) + + +# allows BroadcastingList instance to be subscriptable +class BroadcastingListCls: + def __getitem__(self, types): + return + + +# mypy doesn't support parameters on types, so we have to explicitly type each +# list size +BroadcastingList1 = BroadcastingListCls() +for i in range(2, 7): + globals()[f"BroadcastingList{i}"] = BroadcastingList1 + + +def is_scripting() -> bool: + r""" + Function that returns True when in compilation and False otherwise. This + is useful especially with the @unused decorator to leave code in your + model that is not yet TorchScript compatible. + .. testcode:: + + import torch + + @torch.jit.unused + def unsupported_linear_op(x): + return x + + def linear(x): + if torch.jit.is_scripting(): + return torch.linear(x) + else: + return unsupported_linear_op(x) + """ + return False + + +# Retrieves a fully-qualified name (module hierarchy + classname) for a given obj. +def _qualified_name(obj, mangle_name=True) -> str: + # This special case allows us to override the qualified name on a type. + # It's currently used in conjunction with tracing, where we create a + # fake module to filter only supported attributes. However, since this + # new type is defined as a local class, we need a mechanism to override + # its qualname so it appears correctly in the TorchScript system. This, + # we set '_jit_override_qualname' with the original traced module's + # qualified name, which is picked up here + if hasattr(obj, "_jit_override_qualname"): + return obj._jit_override_qualname + # short-circuit in cases where the object already has a known qualified name + if isinstance(obj, torch._C.ScriptFunction): + return obj.qualified_name + + if getattr(obj, "__name__", None): + name = obj.__name__ + # Enum classes do not have `__name__` attr, instead they have `name`. + elif isinstance(obj, enum.Enum): + name = obj.name + else: + raise RuntimeError("Could not get name of python class object") + + if name == "": + name = "_lambda" # make name a valid identifier + + module_name = obj.__module__ + + # If the module is actually a torchbind module, then we should short circuit + if module_name == "torch._classes": + return obj.qualified_name + + # The Python docs are very clear that `__module__` can be None, but I can't + # figure out when it actually would be. + if module_name is None: + raise RuntimeError( + f"Could not get qualified name for class '{name}': " + "__module__ can't be None." + ) + + # if getattr(sys.modules[module_name], name) is not obj: + # raise RuntimeError(f"Could not get qualified name for class '{name}': " + # f"the attr {name} on module {module_name} is not the class") + + # torch.package and TorchScript have separate mangling schemes to avoid + # name collisions from multiple packages. To avoid them interfering with + # each other, normalize the package managing here. + if package_mangling.is_mangled(module_name): + module_name = module_name.replace("<", "_") + module_name = module_name.replace(">", "_") + + # The PythonExceptionValue C++ class in torch/csrc/jit/python/python_sugared_value.h + # does not need mangle the python class name. + if mangle_name: + # __main__ is a builtin module, so rewrite it to "__torch__". + if module_name == "__main__": + module_name = "__torch__" + else: + # Everything else gets a "__torch__" prefix to avoid name collisions + # with the names of user values. + module_name = "__torch__." + module_name + + if "." in name: + raise RuntimeError( + f"Could not get qualified name for class '{name}': " + f"'{name}' is not a valid identifier" + ) + + return module_name + "." + name + + +class SourceLoader: + def __init__(self): + self.content = {} + + def cache(self, fn, source): + self.content[fn] = source + + def get_source(self, fn): + return self.content.get(fn) + + +loader = SourceLoader() + + +def createResolutionCallbackFromEnv(lookup_base): + """ + Creates a resolution callback that will look up qualified names in an + environment, starting with `lookup_base` for the base of any qualified + names, then proceeding down the lookup chain with the resolved object. + + You should not use this directly, it should only be used from the other + createResolutionCallbackFrom* functions. + """ + + def lookupInModule(qualified_name, module): + if "." in qualified_name: + base, remaining_pieces = qualified_name.split(".", maxsplit=1) + module_value = getattr(module, base) + return lookupInModule(remaining_pieces, module_value) + else: + return getattr(module, qualified_name) + + def parseNestedExpr(expr, module) -> tuple[Any, int]: + i = 0 + while i < len(expr) and expr[i] not in (",", "[", "]"): + i += 1 + + # Special case logic for the empty Tuple as a subscript (used + # in the type annotation `Tuple[()]`) + if expr[:i] == "()": + return (), i + + base = lookupInModule(expr[:i].strip(), module) + assert base is not None, f"Unresolvable type {expr[:i]}" + if i == len(expr) or expr[i] != "[": + return base, i + + assert expr[i] == "[" + parts = [] + while expr[i] != "]": + part_len = 0 + i += 1 + part, part_len = parseNestedExpr(expr[i:], module) + parts.append(part) + i += part_len + if len(parts) > 1: + return base[tuple(parts)], i + 1 + else: + return base[parts[0]], i + 1 + + def parseExpr(expr, module): + try: + value, len_parsed = parseNestedExpr(expr, module) + assert len_parsed == len(expr), ( + "whole expression was not parsed, falling back to c++ parser" + ) + return value + except Exception: + """ + The python resolver fails in several cases in known unit tests, and is intended + to fall back gracefully to the c++ resolver in general. For example, python 2 style + annotations which are frequent in our unit tests often fail with types e.g. int not + resolvable from the calling frame. + """ + return None + + return lambda expr: parseExpr(expr, lookup_base) + + +def createResolutionCallbackFromFrame(frames_up: int = 0): + """ + Creates a function which, given a string variable name, + returns the value of the variable in the scope of the caller of + the function which called createResolutionCallbackFromFrame (by default). + + This is used to enable access in-scope Python variables inside + TorchScript fragments. + + frames_up is number of additional frames to go up on the stack. + The default value is 0, which correspond to the frame of the caller + of createResolutionCallbackFromFrame. Also for example, if frames_up is set + to 1, then the frame of the caller's caller of createResolutionCallbackFromFrame + will be taken. + + For example, the following program prints 2:: + + def bar(): + cb = createResolutionCallbackFromFrame(1) + print(cb("foo")) + + + def baz(): + foo = 2 + bar() + + + baz() + """ + frame = inspect.currentframe() + i = 0 + while i < frames_up + 1: + assert frame is not None + frame = frame.f_back + i += 1 + + assert frame is not None + f_locals = frame.f_locals + f_globals = frame.f_globals + + class env: + def __getattr__(self, key): + if key in f_locals: + return f_locals[key] + elif key in f_globals: + return f_globals[key] + elif key in dir(builtins): + return getattr(builtins, key) + + return createResolutionCallbackFromEnv(env()) + + +def get_closure(fn): + """ + Get a dictionary of closed over variables from a function + """ + captures = {} + captures.update(fn.__globals__) + + for index, captured_name in enumerate(fn.__code__.co_freevars): + captures[captured_name] = fn.__closure__[index].cell_contents + + return captures + + +# [local resolution in python] +# Depending on where a variable is defined, and where it is used, we may +# or may not be able to recover its value when recursively compiling a +# script function. Remember in the general case, a module or function is +# first defined and then later scripted. This means we do not have a +# chance to capture the active frames when the function is defined. Hence any +# name resolution has to happen later on the created closure. The way +# python captures type annotations restricts what we can recover. The +# follow example illustrates the different cases: +# +# class MyGlobalClass: +# ... +# def my_local_scope(): +# @torch.jit.script +# class MyClass: +# ... +# @torch.jit.script +# class MyClassUsedAsVar: +# ... +# def eg(x: MyClass, y: MyGlobalClass): +# a_local_capture : Foo +# return MyClassUsedAsVar(x) +# +# MyGlobalClass is defined in the __globals__ dictionary of function +# 'eg', so it is always recoverable. my_local_scope introduces a new local +# variable scope in the function. Classes defined here are only visible as +# local variables. For the case of MyClassUsedAsVar, it is captured +# because it is used as a variable inside the body of the function, and we +# can resolve it using the captures returned from `get_closure`. However, +# the type annotations are not captured by the closure. In Python +# 3.0--3.9, the _value_ of MyClass and MyGlobalClass will be available as +# annotations on `eg``, but starting in Python 4.0, they will represented as +# strings and no longer present. Furthermore, since the body of `eg` does +# not reference those names, they do not appear in the list of closed over +# variables. In Python 2.x, type annotations are in comments, leading to a +# similar situation where their definitions are not available. We anticipate +# that most users will not run into this issue because their modules and +# functions will be defined at a global scope like MyGlobalClass. In cases +# where they are not, it is possible to work around issues by declaring the +# values global in the function. +# In Python 3.9 declaring class as global will make it invisible to +# `inspect.getsource`, see https://bugs.python.org/issue42666 . +# This could be worked around by manually adding it to `global()` dictionary. + + +def createResolutionCallbackFromClosure(fn): + """ + Create a resolutionCallback by introspecting the function instead of + looking up the stack for the enclosing scope + """ + closure = get_closure(fn) + + class closure_lookup: + # This is a class since `closure` is a dict and it's easier in + # `env_helper` if everything just works with `getattr` calls + def __getattr__(self, key): + if key in closure: + return closure[key] + elif hasattr(typing, key): + return getattr(typing, key) + elif hasattr(builtins, key): + return getattr(builtins, key) + return None + + return createResolutionCallbackFromEnv(closure_lookup()) + + +def can_compile_class(cls) -> bool: + # If any of the functions on a type don't have a code object, this type can't + # be compiled and is probably a builtin / bound from C + if is_ignored_fn(cls): + return False + + # Ignore the following list of built-in classes. + ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception) + if issubclass(cls, ignored_builtin_classes): + return False + + names = cls.__dict__ + fns = [ + getattr(cls, name) + for name in names + if inspect.isroutine(getattr(cls, name, None)) + ] + has_code = [hasattr(fn, "__code__") for fn in fns] + return all(has_code) + + +def get_callable_argument_names(fn) -> list[str]: + """ + Gets names of all POSITIONAL_OR_KEYWORD arguments for callable `fn`. + Returns an empty list when other types of arguments are present. + + This is used by `torch.jit.trace` to assign meaningful argument names to + traced functions and modules. + + Args: + fn: A callable. + Returns: + Argument names: List[str] + """ + # inspect.signature may fail, give up in that case. + try: + callable_signature = inspect.signature(fn) + except Exception: + return [] + + argument_names = [] + for name, param in callable_signature.parameters.items(): + # All four other types of arguments do not map to individual values + # with a keyword as name. + if not param.kind == param.POSITIONAL_OR_KEYWORD: + continue + + argument_names.append(name) + + return argument_names + + +def get_annotation_str(annotation): + """ + Convert an AST node containing a type annotation to the string present in the source + that represents the same annotation. + """ + if isinstance(annotation, ast.Name): + return annotation.id + elif isinstance(annotation, ast.Attribute): + return ".".join([get_annotation_str(annotation.value), annotation.attr]) + elif isinstance(annotation, ast.Subscript): + # In Python3.9+ subscript indices are not wrapped in ast.Index + subscript_slice = annotation.slice + return f"{get_annotation_str(annotation.value)}[{get_annotation_str(subscript_slice)}]" + elif isinstance(annotation, ast.Tuple): + return ",".join([get_annotation_str(elt) for elt in annotation.elts]) + elif isinstance(annotation, ast.Constant): + return f"{annotation.value}" + + # If an AST node is not handled here, it's probably handled in ScriptTypeParser. + return None + + +def get_type_hint_captures(fn): + """ + Get a dictionary containing type resolution mappings necessary to resolve types + for the literal annotations on 'fn'. These are not considered to be closed-over by fn + and must be obtained separately (e.g. using this function). + + Args: + fn: A callable. + Returns: + A Dict[str, Any] containing a mapping from the literal annotations used on + fn to the Python objects they refer to. + """ + # First, try to get the source of the function. We'll need to parse it to find the actual string names + # that were used to annotate the types, since inspect.signature() will only return the class object that + # the annotation refers to, not the string name. If we can't get the source, simply return an empty dict. + # This may happen in cases where the function is synthesized dynamically at runtime. + src = loader.get_source(fn) + if src is None: + try: + src = inspect.getsource(fn) + except OSError as e: + raise OSError( + f"Failed to get source for {fn} using inspect.getsource" + ) from e + + # Gather a dictionary of parameter name -> type, skipping any parameters whose annotated + # types are strings. These are only understood by TorchScript in the context of a type annotation + # that refers to a class in its own definition, but trying to include a mapping for this in the result + # function would cause infinite recursion because the class is currently being compiled. + # In addition, there is logic in ScriptTypeParser to handle this. + signature = inspect.signature(fn) + name_to_type = { + name: parameter.annotation + for name, parameter in signature.parameters.items() + if parameter.annotation is not inspect.Parameter.empty + and not isinstance(parameter.annotation, str) + } + + # Then, get the literal type annotations from the function declaration + # by source inspection. This accounts for the case in which aliases are used + # to annotate the arguments (e.g device_t = torch.device, and then d: device_t). + # frontend.py cannot be used here because it includes _jit_internal, so use ast instead. + a = ast.parse(textwrap.dedent(src)) + if len(a.body) != 1 or not isinstance(a.body[0], ast.FunctionDef): + raise RuntimeError(f"Expected {fn} to be a function") + f = a.body[0] + + # Prepare a dictionary of source annotation -> type, which will be the final result of this function, + # by using the parsed AST (f) to reconstruct source annotations as strings for each parameter and mapping + # them to the type object corresponding to the annotation via name_to_type using the parameter name. + annotation_to_type = {} + + for arg in f.args.args: + # Get the source type annotation string for this argument if possible. + arg_annotation_str = ( + get_annotation_str(arg.annotation) if arg.annotation else None + ) + + # If the argument has no annotation or get_annotation_str cannot convert it to a string, + # arg_annotation_str will be None. Skip this arg; ScriptTypeParser will probably handle + # this in the latter case. + if arg_annotation_str is None: + continue + + # Insert {arg_annotation_str: type} into annotation_to_type if possible. One reason arg_name may not + # be present in name_to_type is that the annotation itself is a string and not a type object + # (common for self-refential annotations in classes). Once again, let ScriptTypeParser handle this. + arg_name = arg.arg + if arg_name in name_to_type: + annotation_to_type[arg_annotation_str] = name_to_type[arg_name] + + # If there is a valid return annotation, include it in annotation_to_type. As with argument annotations, + # the literal annotation has to be convertible to a string by get_annotation_str, and the actual type + # of the annotation cannot be a string. + literal_return_annotation = get_annotation_str(f.returns) + valid_literal_annotation = literal_return_annotation is not None + return_annotation = signature.return_annotation + valid_return_annotation_type = ( + return_annotation is not inspect.Parameter.empty + and not isinstance(return_annotation, str) + ) + if valid_literal_annotation and valid_return_annotation_type: + annotation_to_type[literal_return_annotation] = return_annotation + + return annotation_to_type + + +def createResolutionCallbackForClassMethods(cls): + """ + This looks at all the methods defined in a class and pulls their closed-over + variables into a dictionary and uses that to resolve variables. + """ + # cls is a type here, so `ismethod` is false since the methods on the type + # aren't bound to anything, so Python treats them as regular functions + fns = [ + getattr(cls, name) + for name in cls.__dict__ + if inspect.isroutine(getattr(cls, name)) + ] + # Skip built-ins, as they do not have global scope nor type hints + # Needed to support `enum.Enum` derived classes in Python-3.11 + # That adds `_new_member_` property which is an alias to `__new__` + fns = [fn for fn in fns if not inspect.isbuiltin(fn) and hasattr(fn, "__globals__")] + captures = {} + + for fn in fns: + captures.update(get_closure(fn)) + captures.update(get_type_hint_captures(fn)) + + def lookup_in_class(key): + if key in captures: + return captures[key] + else: + return getattr(builtins, key, None) + + return lookup_in_class + + +def boolean_dispatch( + arg_name, + arg_index, + default, + if_true, + if_false, + module_name, + func_name, +): + """ + Dispatches to either of 2 script functions based on a boolean argument. + In TorchScript, the boolean argument must be constant so that the correct + function to use can be determined at compile time. + """ + + def fn(*args, **kwargs): + dispatch_flag = default + if arg_name in kwargs: + dispatch_flag = kwargs[arg_name] + elif arg_index < len(args): + dispatch_flag = args[arg_index] + + if dispatch_flag: + return if_true(*args, **kwargs) + else: + return if_false(*args, **kwargs) + + if if_true.__doc__ is None and if_false.__doc__ is not None: + doc = if_false.__doc__ + if_true.__doc__ = doc + elif if_false.__doc__ is None and if_true.__doc__ is not None: + doc = if_true.__doc__ + if_false.__doc__ = doc + elif if_false.__doc__ is None and if_true.__doc__ is None: + # neither function has a docstring + doc = None + else: + raise RuntimeError("only one function can have a docstring") + fn.__doc__ = doc + + if module_name is not None: + fn.__module__ = module_name + if func_name is not None: + fn.__name__ = func_name + + boolean_dispatched[fn] = { + "if_true": if_true, + "if_false": if_false, + "index": arg_index, + "default": default, + "arg_name": arg_name, + } + return fn + + +class FunctionModifiers: + """ + Used to denote the behavior of a function in TorchScript. See export() and + ignore() for details. + """ + + UNUSED = "unused (ignored and replaced with raising of an exception)" + IGNORE = "ignore (leave as a call to Python, cannot be torch.jit.save'd)" + EXPORT = "export (compile this function even if nothing calls it)" + DEFAULT = "default (compile if called from a exported function / forward)" + COPY_TO_SCRIPT_WRAPPER = ( + "if this method is not scripted, copy the python method onto the scripted model" + ) + _DROP = "_drop (function is fully ignored, declaration can be unscriptable)" + + +def export(fn: Callable[_P, _R]) -> Callable[_P, _R]: + """ + This decorator indicates that a method on an ``nn.Module`` is used as an entry point into a + :class:`ScriptModule` and should be compiled. + + ``forward`` implicitly is assumed to be an entry point, so it does not need this decorator. + Functions and methods called from ``forward`` are compiled as they are seen + by the compiler, so they do not need this decorator either. + + Example (using ``@torch.jit.export`` on a method): + + .. testcode:: + + import torch + import torch.nn as nn + + class MyModule(nn.Module): + def implicitly_compiled_method(self, x): + return x + 99 + + # `forward` is implicitly decorated with `@torch.jit.export`, + # so adding it here would have no effect + def forward(self, x): + return x + 10 + + @torch.jit.export + def another_forward(self, x): + # When the compiler sees this call, it will compile + # `implicitly_compiled_method` + return self.implicitly_compiled_method(x) + + def unused_method(self, x): + return x - 20 + + # `m` will contain compiled methods: + # `forward` + # `another_forward` + # `implicitly_compiled_method` + # `unused_method` will not be compiled since it was not called from + # any compiled methods and wasn't decorated with `@torch.jit.export` + m = torch.jit.script(MyModule()) + """ + fn._torchscript_modifier = FunctionModifiers.EXPORT # type:ignore[attr-defined] + return fn + + +def unused(fn: Callable[_P, _R]) -> Callable[_P, _R]: + """ + This decorator indicates to the compiler that a function or method should + be ignored and replaced with the raising of an exception. This allows you + to leave code in your model that is not yet TorchScript compatible and still + export your model. + + Example (using ``@torch.jit.unused`` on a method):: + + import torch + import torch.nn as nn + + + class MyModule(nn.Module): + def __init__(self, use_memory_efficient): + super().__init__() + self.use_memory_efficient = use_memory_efficient + + @torch.jit.unused + def memory_efficient(self, x): + import pdb + + pdb.set_trace() + return x + 10 + + def forward(self, x): + # Use not-yet-scriptable memory efficient mode + if self.use_memory_efficient: + return self.memory_efficient(x) + else: + return x + 10 + + + m = torch.jit.script(MyModule(use_memory_efficient=False)) + m.save("m.pt") + + m = torch.jit.script(MyModule(use_memory_efficient=True)) + # exception raised + m(torch.rand(100)) + """ + if isinstance(fn, property): + prop = fn + setattr( # noqa: B010 + prop.fget, "_torchscript_modifier", FunctionModifiers.UNUSED + ) + + if prop.fset: + setattr( # noqa: B010 + prop.fset, "_torchscript_modifier", FunctionModifiers.UNUSED + ) + + return prop + + fn._torchscript_modifier = FunctionModifiers.UNUSED # type: ignore[attr-defined] + return fn + + +# No op context manager from python side +class _IgnoreContextManager(contextlib.AbstractContextManager): + def __init__(self, **kwargs): + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + pass + + +def ignore(drop=False, **kwargs): + """ + This decorator indicates to the compiler that a function or method should + be ignored and left as a Python function. This allows you to leave code in + your model that is not yet TorchScript compatible. If called from TorchScript, + ignored functions will dispatch the call to the Python interpreter. Models with ignored + functions cannot be exported; use :func:`@torch.jit.unused ` instead. + + Example (using ``@torch.jit.ignore`` on a method):: + + import torch + import torch.nn as nn + + + class MyModule(nn.Module): + @torch.jit.ignore + def debugger(self, x): + import pdb + + pdb.set_trace() + + def forward(self, x): + x += 10 + # The compiler would normally try to compile `debugger`, + # but since it is `@ignore`d, it will be left as a call + # to Python + self.debugger(x) + return x + + + m = torch.jit.script(MyModule()) + + # Error! The call `debugger` cannot be saved since it calls into Python + m.save("m.pt") + + Example (using ``@torch.jit.ignore(drop=True)`` on a method): + + .. testcode:: + + import torch + import torch.nn as nn + + class MyModule(nn.Module): + @torch.jit.ignore(drop=True) + def training_method(self, x): + import pdb + pdb.set_trace() + + def forward(self, x): + if self.training: + self.training_method(x) + return x + + m = torch.jit.script(MyModule()) + + # This is OK since `training_method` is not saved, the call is replaced + # with a `raise`. + m.save("m.pt") + + .. testcleanup:: + + import os + os.remove('m.pt') + """ + + if callable(drop): + # used without any args, so drop is actually a function + # @torch.jit.ignore + # def fn(...): + fn = drop + fn._torchscript_modifier = FunctionModifiers.IGNORE + return fn + + if not isinstance(drop, bool): + raise RuntimeError( + f"Argument to @torch.jit.ignore must be a bool or a function but got {drop}" + ) + + # for backwards compat + drop_on_export = kwargs.pop("drop_on_export", None) + if drop_on_export: + warnings.warn( + "ignore(drop_on_export=True) has been deprecated. TorchScript will now drop the function " + "call on compilation. Use torch.jit.unused now. {}", + category=FutureWarning, + ) + + drop = drop_on_export + elif drop: + warnings.warn( + "ignore(True) has been deprecated. TorchScript will now drop the function " + "call on compilation. Use torch.jit.unused now. {}", + category=FutureWarning, + ) + + def decorator(fn): + if drop: + fn._torchscript_modifier = FunctionModifiers.UNUSED + else: + fn._torchscript_modifier = FunctionModifiers.IGNORE + return fn + + return decorator + + +def _drop(fn: Callable[_P, _R]) -> Callable[_P, _R]: + fn._torchscript_modifier = FunctionModifiers._DROP # type: ignore[attr-defined] + return fn + + +def _copy_to_script_wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]: + fn._torchscript_modifier = FunctionModifiers.COPY_TO_SCRIPT_WRAPPER # type: ignore[attr-defined] + return fn + + +def module_has_exports(mod): + for name in dir(mod): + if hasattr(mod, name): + item = getattr(mod, name) + if callable(item): + if get_torchscript_modifier(item) is FunctionModifiers.EXPORT: + return True + return False + + +# WARNING: should_drop is currently being used by our JIT code coverage plug-in to mark JIT'd code as covered. If you +# rename this function, please update references in tools/coverage_plugins_package/src/coverage_plugins/jit_plugin.py to +# allow JIT'd code to still be covered. +def should_drop(fn) -> bool: + attr = get_torchscript_modifier(fn) + if attr is None: + return False + return attr is FunctionModifiers.UNUSED or attr is FunctionModifiers._DROP + + +def is_ignored_fn(fn) -> bool: + mod = get_torchscript_modifier(fn) + return ( + mod is FunctionModifiers.UNUSED + or mod is FunctionModifiers.IGNORE + or mod is FunctionModifiers._DROP + ) + + +def _is_drop_fn(fn) -> bool: + mod = get_torchscript_modifier(fn) + return mod is FunctionModifiers._DROP + + +def is_static_fn(cls, fn) -> bool: + return isinstance(inspect.getattr_static(cls, fn, default=None), staticmethod) + + +def get_static_fn(cls, fn): + return inspect.getattr_static(cls, fn).__func__ + + +def get_torchscript_modifier(fn): + if not callable(fn): + return None + if hasattr(fn, "__func__"): + fn = fn.__func__ + return getattr(fn, "_torchscript_modifier", FunctionModifiers.DEFAULT) + + +def copy_torchscript_modifier(orig, new) -> None: + attr = get_torchscript_modifier(orig) + if attr is None: + return + new._torchscript_modifier = attr + + +# overloading registration +# overloads get registered in this file, and compiled in torch/jit/__init__.py +# so that they can be imported in nn/functional.py without an import cycle + +# qualified_name => list[overload_functions] +_overloaded_fns: dict[str, list[Callable]] = {} # noqa: T484 + + +_OVERLOAD_EXAMPLE = """ +Example usage of overload function: +@torch.jit._overload +def my_function(x: type0) -> type0: # decl 1 + pass + +@torch.jit._overload +def my_function(x: type1) -> type1: # decl 2 + pass + +def my_function(x): # implementation + if isinstance(x, type0): + return x + elif isinstance(x, type1): + return x +""" + + +def get_overload_no_implementation_error_message(kind, obj): + sourcelines, file_lineno, filename = get_source_lines_and_file(obj) + return ( + f'Implementation for the {kind} "{_qualified_name(obj)}" is missing. Please make ' + f"sure a definition is provided and defined after all overload declarations.\n" + f'File "{filename}", line {file_lineno}:\n' + + "".join(sourcelines) + + "\n" + + _OVERLOAD_EXAMPLE + ) + + +def _check_overload_body(func): + try: + parsed_def = parse_def(func) + except OSError: + # Parsing the function definition can raise an OSError if source is unavailable. + # Since this is just an initial check, just raise a warning if this is the case. + warnings.warn( + f"Unable to retrieve source for @torch.jit._overload function: {func}." + ) + return + + body = parsed_def.ast.body[0].body + + def is_pass(x): + return isinstance(x, ast.Pass) + + def is_ellipsis(x): + return ( + isinstance(x, ast.Expr) + and isinstance(x.value, ast.Constant) + and x.value.value is Ellipsis + ) + + if len(body) != 1 or not (is_pass(body[0]) or is_ellipsis(body[0])): + msg = ( + "Only `pass` statement or `...` can be the body of overload declaration:\n" + ) + msg += "\n".join(parsed_def.source.split("\n")[:3]) + msg += " <- Expecting `pass` or `...` here!\n" + _OVERLOAD_EXAMPLE + raise RuntimeError(msg) + + +def _overload(func): + _check_overload_body(func) + qual_name = _qualified_name(func) + global _overloaded_fns + fn_overload_list = _overloaded_fns.get(qual_name) + if fn_overload_list is None: + fn_overload_list = [] + _overloaded_fns[qual_name] = fn_overload_list + fn_overload_list.append(func) + return func + + +def _get_fn_overloads(qual_name): + return _overloaded_fns.get(qual_name) + + +def _clear_fn_overloads(qual_name) -> None: + del _overloaded_fns[qual_name] + + +def get_class_name_lineno(method) -> tuple[str, int]: + current_frame = inspect.currentframe() + + # one for the get_class_name call, one for _overload_method call + for _ in range(2): + assert ( + current_frame is not None + ) # assert current frame is not an Optional[FrameType] + current_frame = current_frame.f_back + + assert current_frame is not None # same here + class_name = current_frame.f_code.co_name + line_no = current_frame.f_code.co_firstlineno + return class_name, line_no + + +# At the point the decorator is applied to class methods the method +# has no reference to its owning class. _qualified_name would not include +# the class it is defined in, so any methods with the same name in the same file +# would have the same _qualified_name, even if they were defined in different +# classes. This problem only exists in python 2. +# We get around this problem by looking at the stack frame and identifying +# the class name, and throwing an error whenever overloads are used +# when modules of the same name are in the same file + +# qualified_name => class name => list[overload_functions] +_overloaded_methods: dict[str, dict[str, list[Callable]]] = {} # noqa: T484 + + +# (qualified_name, class name) => class_fileno +_overloaded_method_class_fileno: dict[tuple[str, str], int] = {} + + +def _overload_method(func): + _check_overload_body(func) + qual_name = _qualified_name(func) + global _overloaded_methods + class_name_map = _overloaded_methods.get(qual_name, None) + if class_name_map is None: + class_name_map = {} + _overloaded_methods[qual_name] = class_name_map + + class_name, line_no = get_class_name_lineno(func) + method_overloads = class_name_map.get(class_name, None) + if method_overloads is None: + method_overloads = [] + class_name_map[class_name] = method_overloads + _overloaded_method_class_fileno[(qual_name, class_name)] = line_no + else: + existing_lineno = _overloaded_method_class_fileno[(qual_name, class_name)] + if existing_lineno != line_no: + raise RuntimeError( + "Cannot currently overload the same method name in two different" + " classes with the same name in the same module" + ) + + method_overloads.append(func) + return func + + +def _get_overloaded_methods(method, mod_class): + # TODO: __name__ not set for submodules in recursive script + if not hasattr(method, "__name__"): + return None + qual_name = _qualified_name(method) + class_name_map = _overloaded_methods.get(qual_name, None) + if class_name_map is None: + return None + overloads = class_name_map.get(mod_class.__name__, None) + if overloads is None: + return None + + method_line_no = get_source_lines_and_file(method)[1] + mod_class_fileno = get_source_lines_and_file(mod_class)[1] + mod_end_fileno = mod_class_fileno + len(get_source_lines_and_file(mod_class)[0]) + if not (method_line_no >= mod_class_fileno and method_line_no <= mod_end_fileno): + raise AssertionError( + "Overloads are not usable when a module is redeclared within the same file: " + + str(method) + ) + return overloads + + +def is_tuple(ann) -> bool: + # Check for typing.Tuple missing args (but `tuple` is fine) + if ann is typing.Tuple: # noqa: UP006 + raise_error_container_parameter_missing("Tuple") + + # For some reason Python 3.7 violates the Type[A, B].__origin__ == Type rule + if not hasattr(ann, "__module__"): + return False + + ann_origin = get_origin(ann) + return ann.__module__ in ("builtins", "typing") and ann_origin is tuple + + +def is_list(ann) -> bool: + # Check for typing.List missing args (but `list` is fine) + if ann is typing.List: # noqa: UP006 + raise_error_container_parameter_missing("List") + + if not hasattr(ann, "__module__"): + return False + + ann_origin = get_origin(ann) + return ann.__module__ in ("builtins", "typing") and ann_origin is list + + +def is_dict(ann) -> bool: + # Check for typing.Dict missing args (but `dict` is fine) + if ann is typing.Dict: # noqa: UP006 + raise_error_container_parameter_missing("Dict") + + if not hasattr(ann, "__module__"): + return False + + ann_origin = get_origin(ann) + return ann.__module__ in ("builtins", "typing") and ann_origin is dict + + +def is_union(ann): + if ann is Union: + raise_error_container_parameter_missing("Union") + + return isinstance(ann, BuiltinUnionType) or ( + hasattr(ann, "__module__") + and ann.__module__ == "typing" + and (get_origin(ann) is Union) + ) + + +def is_optional(ann): + if ann is Optional: + raise_error_container_parameter_missing("Optional") + + def is_optional_as_optional(ann): + return ( + hasattr(ann, "__module__") + and ann.__module__ == "typing" + and (get_origin(ann) is Optional) + ) + + def is_union_as_optional(ann): + ann_args = get_args(ann) + return len(ann_args) == 2 and (None in ann_args or type(None) in ann_args) + + return is_optional_as_optional(ann) or (is_union(ann) and is_union_as_optional(ann)) + + +def is_future(ann) -> bool: + if ann is Future: + raise RuntimeError( + "Attempted to use Future without a " + "contained type. Please add a contained type, e.g. " + "Future[int]" + ) + return get_origin(ann) is Future + + +def is_await(ann) -> bool: + if ann is _Await: + return True + return get_origin(ann) is _Await + + +if torch.distributed.rpc.is_available(): + from torch._C._distributed_rpc import PyRRef + from torch.distributed.rpc import RRef + + def is_rref(ann) -> bool: + if ann is RRef: + raise RuntimeError( + "Attempted to use RRef without a " + "contained type. Please add a contained type, e.g. " + "RRef[int]" + ) + return get_origin(ann) is RRef + + def is_rref_instance(obj) -> bool: + return isinstance(obj, PyRRef) + +else: + + def is_rref_instance(obj) -> bool: + # If the RPC module doesn't exist then RRefs don't exist either. + return False + + +def _try_get_dispatched_fn(fn): + if not callable(fn): + return None + return boolean_dispatched.get(fn) + + +def _get_named_tuple_properties( + obj, + loc: Optional[torch._C._jit_tree_views.SourceRange] = None, + rcb=None, +): + if loc is None: + loc = fake_range() + + assert issubclass(obj, tuple) and hasattr(obj, "_fields") + if hasattr(obj, "_field_defaults"): + defaults = [ + obj._field_defaults[field] + for field in obj._fields + if field in obj._field_defaults + ] + else: + defaults = [] + # In 3.10 recommended way to get annotations is to call `inspect.get_annotations` function + # Also, annotations from base class are not inherited so they need to be queried explicitly + if sys.version_info[:2] < (3, 10): + obj_annotations = getattr(obj, "__annotations__", {}) + else: + obj_annotations = inspect.get_annotations(obj) + if len(obj_annotations) == 0 and hasattr(obj, "__base__"): + obj_annotations = inspect.get_annotations(obj.__base__) + + annotations = [] + for field in obj._fields: + if field in obj_annotations: + field_type = obj_annotations[field] + # [Note: ForwardRef annotations in NamedTuple attributes] + # NamedTuple types are slightly different from normal types. + # + # Normally, annotations are evaluated like this (during jit.script): + # 1. Load strings of python code into c++ and parse. + # 2. Get annotations as strings + # 3. Use the PythonResolver's resolution callback (rcb) to convert + # the string into a python object + # 4. We call into annotations.py:ann_to_type to convert python obj + # from step 3 into a type that torchscript understands. + # + # NamedTuples are more complicated, because it has sub-types. + # Normally, once we have the NamedTuple type object from #3, + # we can just look at the annotation literal values and use + # ann_to_type directly on them. + # + # But sometimes, users will annotate with string literals, e.g. + # x: 'int' + # This also happens with PEP563 (from __forward__ import annotations) + # + # These annotations appear in the annotation dict as ForwardRef('int'). + # + # Then, we need to convert the string into a python object. This + # requires having local context for custom objects or imported types. + # rcb() is what gives us this. So, we plumb rcb through the stack so + # it can be used in this context for the if block below. + # + # FAQ: + # - Why do we need this special handling for NamedTuple but string + # annotations work fine for normal types? Normally, we parse the + # string directly and then call rcb() directly from C++. + # - Why not use ForwardRef._evaluate? For that, we need globals() + # and locals() for the local context where the NamedTuple was defined. + # rcb is what lets us look up into these. So, basically rcb does the + # hard work for us. + if isinstance(field_type, ForwardRef) and rcb is not None: + rcb_type = rcb(field_type.__forward_arg__) + # rcb returns None if it can't find anything. + if rcb_type is None: + raise ValueError( + f"Unknown type annotation: '{field_type}' in NamedTuple {obj.__name__}." + f" Likely due to partial support for ForwardRef parameters in NamedTuples, see #95858." + f" Issue occurred at {loc.highlight()}" + ) + field_type = rcb_type + the_type = torch.jit.annotations.ann_to_type(field_type, loc, rcb) + annotations.append(the_type) + else: + annotations.append(torch._C.TensorType.getInferred()) + return type(obj).__name__, obj._fields, annotations, defaults + + +def _create_named_tuple( + t, + unqual_name: str, + field_names: list[str], + defaults: tuple[Any, ...], +): + TupleType = collections.namedtuple(unqual_name, field_names, defaults=defaults) # type: ignore[call-arg, no-redef, misc] + return TupleType(*t) + + +@contextlib.contextmanager +def _disable_emit_hooks(): + hooks = torch._C._jit_get_emit_hooks() + torch._C._jit_set_emit_hooks(None, None) + try: + yield + finally: + torch._C._jit_set_emit_hooks(hooks[0], hooks[1]) + + +def _disable_emit_hooks_decorator(_DecoratorContextManager) -> None: # noqa: F811 + # noqa: F841 + def __enter__(self) -> None: + self.hooks = torch._C._jit_get_emit_hooks() + torch._C._jit_set_emit_hooks(None, None) + + def __exit__(self, *args) -> None: + torch._C._jit_set_emit_hooks(self.hooks[0], self.hooks[1]) + + +def _is_exception(obj) -> bool: + if not inspect.isclass(obj): + return False + return issubclass(obj, Exception) + + +def raise_error_container_parameter_missing(target_type) -> None: + if target_type.endswith("ict"): + raise RuntimeError( + f"Attempted to use {target_type} without " + "contained types. Please add contained type, e.g. " + f"{target_type}[int, int]" + ) + raise RuntimeError( + f"Attempted to use {target_type} without a " + "contained type. Please add a contained type, e.g. " + f"{target_type}[int]" + ) + + +_RAW_TYPE_NAME_MAPPING = { + dict: "dict", + list: "list", + tuple: "tuple", + typing.Dict: "Dict", # noqa: UP006 + typing.List: "List", # noqa: UP006 + typing.Optional: "Optional", + typing.Tuple: "Tuple", # noqa: UP006 +} + + +def check_args_exist(target_type) -> None: + if name := _RAW_TYPE_NAME_MAPPING.get(target_type): + raise_error_container_parameter_missing(name) + + +def check_empty_containers(obj) -> None: + if obj == [] or obj == {} or obj == (): + warnings.warn( + "The inner type of a container is lost when " + "calling torch.jit.isinstance in eager mode. For " + "example, List[int] would become list and " + "therefore falsely return True for List[float] or" + " List[str]." + ) + + +# supports List/Dict/Tuple and Optional types +# TODO support future +def container_checker(obj, target_type) -> bool: + origin_type = get_origin(target_type) + check_args_exist(target_type) + if origin_type is None: + return False + elif origin_type is list or origin_type is typing.List: # noqa: UP006 + check_empty_containers(obj) + if not isinstance(obj, list): + return False + arg_type = get_args(target_type)[0] + arg_origin = get_origin(arg_type) + for el in obj: + # check if nested container, ex: List[List[str]] + if arg_origin: # processes nested container, ex: List[List[str]] + if not container_checker(el, arg_type): + return False + elif not isinstance(el, arg_type): + return False + return True + elif origin_type is typing.Dict or origin_type is dict: # noqa: UP006 + check_empty_containers(obj) + if not isinstance(obj, dict): + return False + key_type = get_args(target_type)[0] + val_type = get_args(target_type)[1] + for key, val in obj.items(): + # check if keys are of right type + if not isinstance(key, key_type): + return False + val_origin = get_origin(val_type) + if val_origin: + if not container_checker(val, val_type): + return False + elif not isinstance(val, val_type): + return False + return True + elif origin_type is typing.Tuple or origin_type is tuple: # noqa: UP006 + check_empty_containers(obj) + if not isinstance(obj, tuple): + return False + arg_types = get_args(target_type) + if len(obj) != len(arg_types): + return False + for el, el_type in zip(obj, arg_types): + el_origin = get_origin(el_type) + if el_origin: + if not container_checker(el, el_type): + return False + elif not isinstance(el, el_type): + return False + return True + elif origin_type is Union or issubclass( + origin_type, BuiltinUnionType + ): # also handles Optional + if obj is None: # check before recursion because None is always fine + return True + inner_types = get_args(target_type) + for t in inner_types: + t_origin = get_origin(t) + if t_origin: + return container_checker(obj, t) + elif isinstance(obj, t): + return True + return False + + +def _isinstance(obj, target_type) -> bool: + if isinstance(target_type, collections.abc.Container): + if not isinstance(target_type, tuple): + raise RuntimeError( + "The second argument to " + "`torch.jit.isinstance` must be a type " + "or a tuple of types" + ) + for t_type in target_type: + if _isinstance(obj, t_type): + return True + return False + + origin_type = get_origin(target_type) + if origin_type: + return container_checker(obj, target_type) + + # Check to handle non-typed optional origin returns as none instead + # of as optional in 3.7-3.8 + check_args_exist(target_type) + + # handle non-containers + return isinstance(obj, target_type) + + +class _TensorExtractor(pickle.Pickler): + def __init__(self, *args, tensors: list[torch.Tensor], **kwargs): + super().__init__(*args, **kwargs) + self.tensors = tensors + + def persistent_id(self, obj): + if isinstance(obj, torch.Tensor): + self.tensors.append(obj) + return "" + # Since we just want to extract tensors, we don't mind if an object is + # unpicklable if it doesn't contain tensors, as we can just ignore/skip + # it. To play it safe, we only do so for common objects that we're sure + # don't contain tensors. Feel free to add new types here. Note also that + # even if a type isn't listed here this won't block users, since they + # can just add a __getstate__ or __reduce__ method to their class. + if isinstance(obj, LockType): + return "" + # Futures and RRefs don't technically contain a value, they just offer + # the means to access a value. + if isinstance(obj, CFuture) or is_rref_instance(obj): + return "" + if isinstance(obj, CAwait): + return "" + if isinstance(obj, torch.cuda.Event): + return "" + if isinstance(obj, threading.Thread): + return "" + return None + + +def _extract_tensors(obj): + r""" + This function is exclusively called from C++. + See ``torch/csrc/jit/python/python_ivalue.h``. + + It extracts the tensors contained in the given object, through pickling. + """ + tensors: list[torch.Tensor] = [] + extractor = _TensorExtractor(io.BytesIO(), protocol=-1, tensors=tensors) + extractor.dump(obj) + return tensors + + +def _get_model_id(obj) -> Optional[str]: + if isinstance(obj, torch.jit.ScriptModule): + return str(obj._c._type()) + elif isinstance(obj, torch.jit.ScriptFunction): + return obj.qualified_name + else: + return None + + +# In Python-3.11+ typed enums (i.e. IntEnum for example) retain number of base class methods in subclass +# that were previously dropped. To preserve the behavior, explicitly drop them there + +if sys.version_info >= (3, 11): + _drop(enum.Enum.__new__) + _drop(enum.Enum.__format__) + _drop(enum.Enum.__repr__) + _drop(enum.Enum.__str__) diff --git a/venv/lib/python3.10/site-packages/torch/_linalg_utils.py b/venv/lib/python3.10/site-packages/torch/_linalg_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..43c8b65767e00295df20e556b9ecac4e63eac8bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_linalg_utils.py @@ -0,0 +1,150 @@ +# mypy: allow-untyped-defs +"""Various linear algebra utility methods for internal use.""" + +from typing import Optional + +import torch +from torch import Tensor + + +def is_sparse(A): + """Check if tensor A is a sparse COO tensor. All other sparse storage formats (CSR, CSC, etc...) will return False.""" + if isinstance(A, torch.Tensor): + return A.layout == torch.sparse_coo + + error_str = "expected Tensor" + if not torch.jit.is_scripting(): + error_str += f" but got {type(A)}" + raise TypeError(error_str) + + +def get_floating_dtype(A): + """Return the floating point dtype of tensor A. + + Integer types map to float32. + """ + dtype = A.dtype + if dtype in (torch.float16, torch.float32, torch.float64): + return dtype + return torch.float32 + + +def matmul(A: Optional[Tensor], B: Tensor) -> Tensor: + """Multiply two matrices. + + If A is None, return B. A can be sparse or dense. B is always + dense. + """ + if A is None: + return B + if is_sparse(A): + return torch.sparse.mm(A, B) + return torch.matmul(A, B) + + +def bform(X: Tensor, A: Optional[Tensor], Y: Tensor) -> Tensor: + """Return bilinear form of matrices: :math:`X^T A Y`.""" + return matmul(X.mT, matmul(A, Y)) + + +def qform(A: Optional[Tensor], S: Tensor): + """Return quadratic form :math:`S^T A S`.""" + return bform(S, A, S) + + +def basis(A): + """Return orthogonal basis of A columns.""" + return torch.linalg.qr(A).Q + + +def symeig(A: Tensor, largest: Optional[bool] = False) -> tuple[Tensor, Tensor]: + """Return eigenpairs of A with specified ordering.""" + if largest is None: + largest = False + E, Z = torch.linalg.eigh(A, UPLO="U") + # assuming that E is ordered + if largest: + E = torch.flip(E, dims=(-1,)) + Z = torch.flip(Z, dims=(-1,)) + return E, Z + + +# These functions were deprecated and removed +# This nice error message can be removed in version 1.13+ +def matrix_rank(input, tol=None, symmetric=False, *, out=None) -> Tensor: + raise RuntimeError( + "This function was deprecated since version 1.9 and is now removed.\n" + "Please use the `torch.linalg.matrix_rank` function instead. " + "The parameter 'symmetric' was renamed in `torch.linalg.matrix_rank()` to 'hermitian'." + ) + + +def solve(input: Tensor, A: Tensor, *, out=None) -> tuple[Tensor, Tensor]: + raise RuntimeError( + "This function was deprecated since version 1.9 and is now removed. " + "`torch.solve` is deprecated in favor of `torch.linalg.solve`. " + "`torch.linalg.solve` has its arguments reversed and does not return the LU factorization.\n\n" + "To get the LU factorization see `torch.lu`, which can be used with `torch.lu_solve` or `torch.lu_unpack`.\n" + "X = torch.solve(B, A).solution " + "should be replaced with:\n" + "X = torch.linalg.solve(A, B)" + ) + + +def lstsq(input: Tensor, A: Tensor, *, out=None) -> tuple[Tensor, Tensor]: + raise RuntimeError( + "This function was deprecated since version 1.9 and is now removed. " + "`torch.lstsq` is deprecated in favor of `torch.linalg.lstsq`.\n" + "`torch.linalg.lstsq` has reversed arguments and does not return the QR decomposition in " + "the returned tuple (although it returns other information about the problem).\n\n" + "To get the QR decomposition consider using `torch.linalg.qr`.\n\n" + "The returned solution in `torch.lstsq` stored the residuals of the solution in the " + "last m - n columns of the returned value whenever m > n. In torch.linalg.lstsq, " + "the residuals are in the field 'residuals' of the returned named tuple.\n\n" + "The unpacking of the solution, as in\n" + "X, _ = torch.lstsq(B, A).solution[:A.size(1)]\n" + "should be replaced with:\n" + "X = torch.linalg.lstsq(A, B).solution" + ) + + +def _symeig( + input, + eigenvectors=False, + upper=True, + *, + out=None, +) -> tuple[Tensor, Tensor]: + raise RuntimeError( + "This function was deprecated since version 1.9 and is now removed. " + "The default behavior has changed from using the upper triangular portion of the matrix by default " + "to using the lower triangular portion.\n\n" + "L, _ = torch.symeig(A, upper=upper) " + "should be replaced with:\n" + "L = torch.linalg.eigvalsh(A, UPLO='U' if upper else 'L')\n\n" + "and\n\n" + "L, V = torch.symeig(A, eigenvectors=True) " + "should be replaced with:\n" + "L, V = torch.linalg.eigh(A, UPLO='U' if upper else 'L')" + ) + + +def eig( + self: Tensor, + eigenvectors: bool = False, + *, + e=None, + v=None, +) -> tuple[Tensor, Tensor]: + raise RuntimeError( + "This function was deprecated since version 1.9 and is now removed. " + "`torch.linalg.eig` returns complex tensors of dtype `cfloat` or `cdouble` rather than real tensors " + "mimicking complex tensors.\n\n" + "L, _ = torch.eig(A) " + "should be replaced with:\n" + "L_complex = torch.linalg.eigvals(A)\n\n" + "and\n\n" + "L, V = torch.eig(A, eigenvectors=True) " + "should be replaced with:\n" + "L_complex, V_complex = torch.linalg.eig(A)" + ) diff --git a/venv/lib/python3.10/site-packages/torch/_lobpcg.py b/venv/lib/python3.10/site-packages/torch/_lobpcg.py new file mode 100644 index 0000000000000000000000000000000000000000..a3f57411b8f548dcf6757fc1d1721c0fefae1e9e --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_lobpcg.py @@ -0,0 +1,1157 @@ +# mypy: allow-untyped-defs +"""Locally Optimal Block Preconditioned Conjugate Gradient methods.""" +# Author: Pearu Peterson +# Created: February 2020 + +from typing import Optional + +import torch +from torch import _linalg_utils as _utils, Tensor +from torch.overrides import handle_torch_function, has_torch_function + + +__all__ = ["lobpcg"] + + +def _symeig_backward_complete_eigenspace(D_grad, U_grad, A, D, U): + # compute F, such that F_ij = (d_j - d_i)^{-1} for i != j, F_ii = 0 + F = D.unsqueeze(-2) - D.unsqueeze(-1) + F.diagonal(dim1=-2, dim2=-1).fill_(float("inf")) + F.pow_(-1) + + # A.grad = U (D.grad + (U^T U.grad * F)) U^T + Ut = U.mT.contiguous() + res = torch.matmul( + U, torch.matmul(torch.diag_embed(D_grad) + torch.matmul(Ut, U_grad) * F, Ut) + ) + + return res + + +def _polynomial_coefficients_given_roots(roots): + """ + Given the `roots` of a polynomial, find the polynomial's coefficients. + + If roots = (r_1, ..., r_n), then the method returns + coefficients (a_0, a_1, ..., a_n (== 1)) so that + p(x) = (x - r_1) * ... * (x - r_n) + = x^n + a_{n-1} * x^{n-1} + ... a_1 * x_1 + a_0 + + Note: for better performance requires writing a low-level kernel + """ + poly_order = roots.shape[-1] + poly_coeffs_shape = list(roots.shape) + # we assume p(x) = x^n + a_{n-1} * x^{n-1} + ... + a_1 * x + a_0, + # so poly_coeffs = {a_0, ..., a_n, a_{n+1}(== 1)}, + # but we insert one extra coefficient to enable better vectorization below + poly_coeffs_shape[-1] += 2 + poly_coeffs = roots.new_zeros(poly_coeffs_shape) + poly_coeffs[..., 0] = 1 + poly_coeffs[..., -1] = 1 + + # perform the Horner's rule + for i in range(1, poly_order + 1): + # note that it is computationally hard to compute backward for this method, + # because then given the coefficients it would require finding the roots and/or + # calculating the sensitivity based on the Vieta's theorem. + # So the code below tries to circumvent the explicit root finding by series + # of operations on memory copies imitating the Horner's method. + # The memory copies are required to construct nodes in the computational graph + # by exploiting the explicit (not in-place, separate node for each step) + # recursion of the Horner's method. + # Needs more memory, O(... * k^2), but with only O(... * k^2) complexity. + poly_coeffs_new = poly_coeffs.clone() if roots.requires_grad else poly_coeffs + out = poly_coeffs_new.narrow(-1, poly_order - i, i + 1) + out -= roots.narrow(-1, i - 1, 1) * poly_coeffs.narrow( + -1, poly_order - i + 1, i + 1 + ) + poly_coeffs = poly_coeffs_new + + return poly_coeffs.narrow(-1, 1, poly_order + 1) + + +def _polynomial_value(poly, x, zero_power, transition): + """ + A generic method for computing poly(x) using the Horner's rule. + + Args: + poly (Tensor): the (possibly batched) 1D Tensor representing + polynomial coefficients such that + poly[..., i] = (a_{i_0}, ..., a{i_n} (==1)), and + poly(x) = poly[..., 0] * zero_power + ... + poly[..., n] * x^n + + x (Tensor): the value (possible batched) to evaluate the polynomial `poly` at. + + zero_power (Tensor): the representation of `x^0`. It is application-specific. + + transition (Callable): the function that accepts some intermediate result `int_val`, + the `x` and a specific polynomial coefficient + `poly[..., k]` for some iteration `k`. + It basically performs one iteration of the Horner's rule + defined as `x * int_val + poly[..., k] * zero_power`. + Note that `zero_power` is not a parameter, + because the step `+ poly[..., k] * zero_power` depends on `x`, + whether it is a vector, a matrix, or something else, so this + functionality is delegated to the user. + """ + + res = zero_power.clone() + for k in range(poly.size(-1) - 2, -1, -1): + res = transition(res, x, poly[..., k]) + return res + + +def _matrix_polynomial_value(poly, x, zero_power=None): + """ + Evaluates `poly(x)` for the (batched) matrix input `x`. + Check out `_polynomial_value` function for more details. + """ + + # matrix-aware Horner's rule iteration + def transition(curr_poly_val, x, poly_coeff): + res = x.matmul(curr_poly_val) + res.diagonal(dim1=-2, dim2=-1).add_(poly_coeff.unsqueeze(-1)) + return res + + if zero_power is None: + zero_power = torch.eye( + x.size(-1), x.size(-1), dtype=x.dtype, device=x.device + ).view(*([1] * len(list(x.shape[:-2]))), x.size(-1), x.size(-1)) + + return _polynomial_value(poly, x, zero_power, transition) + + +def _vector_polynomial_value(poly, x, zero_power=None): + """ + Evaluates `poly(x)` for the (batched) vector input `x`. + Check out `_polynomial_value` function for more details. + """ + + # vector-aware Horner's rule iteration + def transition(curr_poly_val, x, poly_coeff): + res = torch.addcmul(poly_coeff.unsqueeze(-1), x, curr_poly_val) + return res + + if zero_power is None: + zero_power = x.new_ones(1).expand(x.shape) + + return _polynomial_value(poly, x, zero_power, transition) + + +def _symeig_backward_partial_eigenspace(D_grad, U_grad, A, D, U, largest): + # compute a projection operator onto an orthogonal subspace spanned by the + # columns of U defined as (I - UU^T) + Ut = U.mT.contiguous() + proj_U_ortho = -U.matmul(Ut) + proj_U_ortho.diagonal(dim1=-2, dim2=-1).add_(1) + + # compute U_ortho, a basis for the orthogonal complement to the span(U), + # by projecting a random [..., m, m - k] matrix onto the subspace spanned + # by the columns of U. + # + # fix generator for determinism + gen = torch.Generator(A.device) + + # orthogonal complement to the span(U) + U_ortho = proj_U_ortho.matmul( + torch.randn( + (*A.shape[:-1], A.size(-1) - D.size(-1)), + dtype=A.dtype, + device=A.device, + generator=gen, + ) + ) + U_ortho_t = U_ortho.mT.contiguous() + + # compute the coefficients of the characteristic polynomial of the tensor D. + # Note that D is diagonal, so the diagonal elements are exactly the roots + # of the characteristic polynomial. + chr_poly_D = _polynomial_coefficients_given_roots(D) + + # the code below finds the explicit solution to the Sylvester equation + # U_ortho^T A U_ortho dX - dX D = -U_ortho^T A U + # and incorporates it into the whole gradient stored in the `res` variable. + # + # Equivalent to the following naive implementation: + # res = A.new_zeros(A.shape) + # p_res = A.new_zeros(*A.shape[:-1], D.size(-1)) + # for k in range(1, chr_poly_D.size(-1)): + # p_res.zero_() + # for i in range(0, k): + # p_res += (A.matrix_power(k - 1 - i) @ U_grad) * D.pow(i).unsqueeze(-2) + # res -= chr_poly_D[k] * (U_ortho @ poly_D_at_A.inverse() @ U_ortho_t @ p_res @ U.t()) + # + # Note that dX is a differential, so the gradient contribution comes from the backward sensitivity + # Tr(f(U_grad, D_grad, A, U, D)^T dX) = Tr(g(U_grad, A, U, D)^T dA) for some functions f and g, + # and we need to compute g(U_grad, A, U, D) + # + # The naive implementation is based on the paper + # Hu, Qingxi, and Daizhan Cheng. + # "The polynomial solution to the Sylvester matrix equation." + # Applied mathematics letters 19.9 (2006): 859-864. + # + # We can modify the computation of `p_res` from above in a more efficient way + # p_res = U_grad * (chr_poly_D[1] * D.pow(0) + ... + chr_poly_D[k] * D.pow(k)).unsqueeze(-2) + # + A U_grad * (chr_poly_D[2] * D.pow(0) + ... + chr_poly_D[k] * D.pow(k - 1)).unsqueeze(-2) + # + ... + # + A.matrix_power(k - 1) U_grad * chr_poly_D[k] + # Note that this saves us from redundant matrix products with A (elimination of matrix_power) + U_grad_projected = U_grad + series_acc = U_grad_projected.new_zeros(U_grad_projected.shape) + for k in range(1, chr_poly_D.size(-1)): + poly_D = _vector_polynomial_value(chr_poly_D[..., k:], D) + series_acc += U_grad_projected * poly_D.unsqueeze(-2) + U_grad_projected = A.matmul(U_grad_projected) + + # compute chr_poly_D(A) which essentially is: + # + # chr_poly_D_at_A = A.new_zeros(A.shape) + # for k in range(chr_poly_D.size(-1)): + # chr_poly_D_at_A += chr_poly_D[k] * A.matrix_power(k) + # + # Note, however, for better performance we use the Horner's rule + chr_poly_D_at_A = _matrix_polynomial_value(chr_poly_D, A) + + # compute the action of `chr_poly_D_at_A` restricted to U_ortho_t + chr_poly_D_at_A_to_U_ortho = torch.matmul( + U_ortho_t, torch.matmul(chr_poly_D_at_A, U_ortho) + ) + # we need to invert 'chr_poly_D_at_A_to_U_ortho`, for that we compute its + # Cholesky decomposition and then use `torch.cholesky_solve` for better stability. + # Cholesky decomposition requires the input to be positive-definite. + # Note that `chr_poly_D_at_A_to_U_ortho` is positive-definite if + # 1. `largest` == False, or + # 2. `largest` == True and `k` is even + # under the assumption that `A` has distinct eigenvalues. + # + # check if `chr_poly_D_at_A_to_U_ortho` is positive-definite or negative-definite + chr_poly_D_at_A_to_U_ortho_sign = -1 if (largest and (k % 2 == 1)) else +1 + chr_poly_D_at_A_to_U_ortho_L = torch.linalg.cholesky( + chr_poly_D_at_A_to_U_ortho_sign * chr_poly_D_at_A_to_U_ortho + ) + + # compute the gradient part in span(U) + res = _symeig_backward_complete_eigenspace(D_grad, U_grad, A, D, U) + + # incorporate the Sylvester equation solution into the full gradient + # it resides in span(U_ortho) + res -= U_ortho.matmul( + chr_poly_D_at_A_to_U_ortho_sign + * torch.cholesky_solve( + U_ortho_t.matmul(series_acc), chr_poly_D_at_A_to_U_ortho_L + ) + ).matmul(Ut) + + return res + + +def _symeig_backward(D_grad, U_grad, A, D, U, largest): + # if `U` is square, then the columns of `U` is a complete eigenspace + if U.size(-1) == U.size(-2): + return _symeig_backward_complete_eigenspace(D_grad, U_grad, A, D, U) + else: + return _symeig_backward_partial_eigenspace(D_grad, U_grad, A, D, U, largest) + + +class LOBPCGAutogradFunction(torch.autograd.Function): + @staticmethod + def forward( # type: ignore[override] + ctx, + A: Tensor, + k: Optional[int] = None, + B: Optional[Tensor] = None, + X: Optional[Tensor] = None, + n: Optional[int] = None, + iK: Optional[Tensor] = None, + niter: Optional[int] = None, + tol: Optional[float] = None, + largest: Optional[bool] = None, + method: Optional[str] = None, + tracker: None = None, + ortho_iparams: Optional[dict[str, int]] = None, + ortho_fparams: Optional[dict[str, float]] = None, + ortho_bparams: Optional[dict[str, bool]] = None, + ) -> tuple[Tensor, Tensor]: + # makes sure that input is contiguous for efficiency. + # Note: autograd does not support dense gradients for sparse input yet. + A = A.contiguous() if (not A.is_sparse) else A + if B is not None: + B = B.contiguous() if (not B.is_sparse) else B + + D, U = _lobpcg( + A, + k, + B, + X, + n, + iK, + niter, + tol, + largest, + method, + tracker, + ortho_iparams, + ortho_fparams, + ortho_bparams, + ) + + ctx.save_for_backward(A, B, D, U) + ctx.largest = largest + + return D, U + + @staticmethod + def backward(ctx, D_grad, U_grad): + A_grad = B_grad = None + grads = [None] * 14 + + A, B, D, U = ctx.saved_tensors + largest = ctx.largest + + # lobpcg.backward has some limitations. Checks for unsupported input + if A.is_sparse or (B is not None and B.is_sparse and ctx.needs_input_grad[2]): + raise ValueError( + "lobpcg.backward does not support sparse input yet." + "Note that lobpcg.forward does though." + ) + if ( + A.dtype in (torch.complex64, torch.complex128) + or B is not None + and B.dtype in (torch.complex64, torch.complex128) + ): + raise ValueError( + "lobpcg.backward does not support complex input yet." + "Note that lobpcg.forward does though." + ) + if B is not None: + raise ValueError( + "lobpcg.backward does not support backward with B != I yet." + ) + + if largest is None: + largest = True + + # symeig backward + if B is None: + A_grad = _symeig_backward(D_grad, U_grad, A, D, U, largest) + + # A has index 0 + grads[0] = A_grad + # B has index 2 + grads[2] = B_grad + return tuple(grads) + + +def lobpcg( + A: Tensor, + k: Optional[int] = None, + B: Optional[Tensor] = None, + X: Optional[Tensor] = None, + n: Optional[int] = None, + iK: Optional[Tensor] = None, + niter: Optional[int] = None, + tol: Optional[float] = None, + largest: Optional[bool] = None, + method: Optional[str] = None, + tracker: None = None, + ortho_iparams: Optional[dict[str, int]] = None, + ortho_fparams: Optional[dict[str, float]] = None, + ortho_bparams: Optional[dict[str, bool]] = None, +) -> tuple[Tensor, Tensor]: + """Find the k largest (or smallest) eigenvalues and the corresponding + eigenvectors of a symmetric positive definite generalized + eigenvalue problem using matrix-free LOBPCG methods. + + This function is a front-end to the following LOBPCG algorithms + selectable via `method` argument: + + `method="basic"` - the LOBPCG method introduced by Andrew + Knyazev, see [Knyazev2001]. A less robust method, may fail when + Cholesky is applied to singular input. + + `method="ortho"` - the LOBPCG method with orthogonal basis + selection [StathopoulosEtal2002]. A robust method. + + Supported inputs are dense, sparse, and batches of dense matrices. + + .. note:: In general, the basic method spends least time per + iteration. However, the robust methods converge much faster and + are more stable. So, the usage of the basic method is generally + not recommended but there exist cases where the usage of the + basic method may be preferred. + + .. warning:: The backward method does not support sparse and complex inputs. + It works only when `B` is not provided (i.e. `B == None`). + We are actively working on extensions, and the details of + the algorithms are going to be published promptly. + + .. warning:: While it is assumed that `A` is symmetric, `A.grad` is not. + To make sure that `A.grad` is symmetric, so that `A - t * A.grad` is symmetric + in first-order optimization routines, prior to running `lobpcg` + we do the following symmetrization map: `A -> (A + A.t()) / 2`. + The map is performed only when the `A` requires gradients. + + .. warning:: LOBPCG algorithm is not applicable when the number of `A`'s rows + is smaller than 3x the number of requested eigenpairs `n`. + + Args: + + A (Tensor): the input tensor of size :math:`(*, m, m)` + + k (integer, optional): the number of requested + eigenpairs. Default is the number of :math:`X` + columns (when specified) or `1`. + + B (Tensor, optional): the input tensor of size :math:`(*, m, + m)`. When not specified, `B` is interpreted as + identity matrix. + + X (tensor, optional): the input tensor of size :math:`(*, m, n)` + where `k <= n <= m`. When specified, it is used as + initial approximation of eigenvectors. X must be a + dense tensor. + + n (integer, optional): if :math:`X` is not specified then `n` + specifies the size of the generated random + approximation of eigenvectors. Default value for `n` + is `k`. If :math:`X` is specified, any provided value of `n` is + ignored and `n` is automatically set to the number of + columns in :math:`X`. + + iK (tensor, optional): the input tensor of size :math:`(*, m, + m)`. When specified, it will be used as preconditioner. + + niter (int, optional): maximum number of iterations. When + reached, the iteration process is hard-stopped and + the current approximation of eigenpairs is returned. + For infinite iteration but until convergence criteria + is met, use `-1`. + + tol (float, optional): residual tolerance for stopping + criterion. Default is `feps ** 0.5` where `feps` is + smallest non-zero floating-point number of the given + input tensor `A` data type. + + largest (bool, optional): when True, solve the eigenproblem for + the largest eigenvalues. Otherwise, solve the + eigenproblem for smallest eigenvalues. Default is + `True`. + + method (str, optional): select LOBPCG method. See the + description of the function above. Default is + "ortho". + + tracker (callable, optional) : a function for tracing the + iteration process. When specified, it is called at + each iteration step with LOBPCG instance as an + argument. The LOBPCG instance holds the full state of + the iteration process in the following attributes: + + `iparams`, `fparams`, `bparams` - dictionaries of + integer, float, and boolean valued input + parameters, respectively + + `ivars`, `fvars`, `bvars`, `tvars` - dictionaries + of integer, float, boolean, and Tensor valued + iteration variables, respectively. + + `A`, `B`, `iK` - input Tensor arguments. + + `E`, `X`, `S`, `R` - iteration Tensor variables. + + For instance: + + `ivars["istep"]` - the current iteration step + `X` - the current approximation of eigenvectors + `E` - the current approximation of eigenvalues + `R` - the current residual + `ivars["converged_count"]` - the current number of converged eigenpairs + `tvars["rerr"]` - the current state of convergence criteria + + Note that when `tracker` stores Tensor objects from + the LOBPCG instance, it must make copies of these. + + If `tracker` sets `bvars["force_stop"] = True`, the + iteration process will be hard-stopped. + + ortho_iparams, ortho_fparams, ortho_bparams (dict, optional): + various parameters to LOBPCG algorithm when using + `method="ortho"`. + + Returns: + + E (Tensor): tensor of eigenvalues of size :math:`(*, k)` + + X (Tensor): tensor of eigenvectors of size :math:`(*, m, k)` + + References: + + [Knyazev2001] Andrew V. Knyazev. (2001) Toward the Optimal + Preconditioned Eigensolver: Locally Optimal Block Preconditioned + Conjugate Gradient Method. SIAM J. Sci. Comput., 23(2), + 517-541. (25 pages) + https://epubs.siam.org/doi/abs/10.1137/S1064827500366124 + + [StathopoulosEtal2002] Andreas Stathopoulos and Kesheng + Wu. (2002) A Block Orthogonalization Procedure with Constant + Synchronization Requirements. SIAM J. Sci. Comput., 23(6), + 2165-2182. (18 pages) + https://epubs.siam.org/doi/10.1137/S1064827500370883 + + [DuerschEtal2018] Jed A. Duersch, Meiyue Shao, Chao Yang, Ming + Gu. (2018) A Robust and Efficient Implementation of LOBPCG. + SIAM J. Sci. Comput., 40(5), C655-C676. (22 pages) + https://arxiv.org/abs/1704.07458 + + """ + + if not torch.jit.is_scripting(): + tensor_ops = (A, B, X, iK) + if not set(map(type, tensor_ops)).issubset( + (torch.Tensor, type(None)) + ) and has_torch_function(tensor_ops): + return handle_torch_function( + lobpcg, + tensor_ops, + A, + k=k, + B=B, + X=X, + n=n, + iK=iK, + niter=niter, + tol=tol, + largest=largest, + method=method, + tracker=tracker, + ortho_iparams=ortho_iparams, + ortho_fparams=ortho_fparams, + ortho_bparams=ortho_bparams, + ) + + if not torch._jit_internal.is_scripting(): + if A.requires_grad or (B is not None and B.requires_grad): + # While it is expected that `A` is symmetric, + # the `A_grad` might be not. Therefore we perform the trick below, + # so that `A_grad` becomes symmetric. + # The symmetrization is important for first-order optimization methods, + # so that (A - alpha * A_grad) is still a symmetric matrix. + # Same holds for `B`. + A_sym = (A + A.mT) / 2 + B_sym = (B + B.mT) / 2 if (B is not None) else None + + return LOBPCGAutogradFunction.apply( + A_sym, + k, + B_sym, + X, + n, + iK, + niter, + tol, + largest, + method, + tracker, + ortho_iparams, + ortho_fparams, + ortho_bparams, + ) + else: + if A.requires_grad or (B is not None and B.requires_grad): + raise RuntimeError( + "Script and require grads is not supported atm." + "If you just want to do the forward, use .detach()" + "on A and B before calling into lobpcg" + ) + + return _lobpcg( + A, + k, + B, + X, + n, + iK, + niter, + tol, + largest, + method, + tracker, + ortho_iparams, + ortho_fparams, + ortho_bparams, + ) + + +def _lobpcg( + A: Tensor, + k: Optional[int] = None, + B: Optional[Tensor] = None, + X: Optional[Tensor] = None, + n: Optional[int] = None, + iK: Optional[Tensor] = None, + niter: Optional[int] = None, + tol: Optional[float] = None, + largest: Optional[bool] = None, + method: Optional[str] = None, + tracker: None = None, + ortho_iparams: Optional[dict[str, int]] = None, + ortho_fparams: Optional[dict[str, float]] = None, + ortho_bparams: Optional[dict[str, bool]] = None, +) -> tuple[Tensor, Tensor]: + # A must be square: + assert A.shape[-2] == A.shape[-1], A.shape + if B is not None: + # A and B must have the same shapes: + assert A.shape == B.shape, (A.shape, B.shape) + + dtype = _utils.get_floating_dtype(A) + device = A.device + if tol is None: + feps = {torch.float32: 1.2e-07, torch.float64: 2.23e-16}[dtype] + tol = feps**0.5 + + m = A.shape[-1] + k = (1 if X is None else X.shape[-1]) if k is None else k + n = (k if n is None else n) if X is None else X.shape[-1] + + if m < 3 * n: + raise ValueError( + f"LPBPCG algorithm is not applicable when the number of A rows (={m})" + f" is smaller than 3 x the number of requested eigenpairs (={n})" + ) + + method = "ortho" if method is None else method + + iparams = { + "m": m, + "n": n, + "k": k, + "niter": 1000 if niter is None else niter, + } + + fparams = { + "tol": tol, + } + + bparams = {"largest": True if largest is None else largest} + + if method == "ortho": + if ortho_iparams is not None: + iparams.update(ortho_iparams) + if ortho_fparams is not None: + fparams.update(ortho_fparams) + if ortho_bparams is not None: + bparams.update(ortho_bparams) + iparams["ortho_i_max"] = iparams.get("ortho_i_max", 3) + iparams["ortho_j_max"] = iparams.get("ortho_j_max", 3) + fparams["ortho_tol"] = fparams.get("ortho_tol", tol) + fparams["ortho_tol_drop"] = fparams.get("ortho_tol_drop", tol) + fparams["ortho_tol_replace"] = fparams.get("ortho_tol_replace", tol) + bparams["ortho_use_drop"] = bparams.get("ortho_use_drop", False) + + if not torch.jit.is_scripting(): + LOBPCG.call_tracker = LOBPCG_call_tracker # type: ignore[method-assign] + + if len(A.shape) > 2: + N = int(torch.prod(torch.tensor(A.shape[:-2]))) + bA = A.reshape((N,) + A.shape[-2:]) + bB = B.reshape((N,) + A.shape[-2:]) if B is not None else None + bX = X.reshape((N,) + X.shape[-2:]) if X is not None else None + bE = torch.empty((N, k), dtype=dtype, device=device) + bXret = torch.empty((N, m, k), dtype=dtype, device=device) + + for i in range(N): + A_ = bA[i] + B_ = bB[i] if bB is not None else None + X_ = ( + torch.randn((m, n), dtype=dtype, device=device) if bX is None else bX[i] + ) + assert len(X_.shape) == 2 and X_.shape == (m, n), (X_.shape, (m, n)) + iparams["batch_index"] = i + worker = LOBPCG(A_, B_, X_, iK, iparams, fparams, bparams, method, tracker) + worker.run() + bE[i] = worker.E[:k] + bXret[i] = worker.X[:, :k] + + if not torch.jit.is_scripting(): + LOBPCG.call_tracker = LOBPCG_call_tracker_orig # type: ignore[method-assign] + + return bE.reshape(A.shape[:-2] + (k,)), bXret.reshape(A.shape[:-2] + (m, k)) + + X = torch.randn((m, n), dtype=dtype, device=device) if X is None else X + assert len(X.shape) == 2 and X.shape == (m, n), (X.shape, (m, n)) + + worker = LOBPCG(A, B, X, iK, iparams, fparams, bparams, method, tracker) + + worker.run() + + if not torch.jit.is_scripting(): + LOBPCG.call_tracker = LOBPCG_call_tracker_orig # type: ignore[method-assign] + + return worker.E[:k], worker.X[:, :k] + + +class LOBPCG: + """Worker class of LOBPCG methods.""" + + def __init__( + self, + A: Optional[Tensor], + B: Optional[Tensor], + X: Tensor, + iK: Optional[Tensor], + iparams: dict[str, int], + fparams: dict[str, float], + bparams: dict[str, bool], + method: str, + tracker: None, + ) -> None: + # constant parameters + self.A = A + self.B = B + self.iK = iK + self.iparams = iparams + self.fparams = fparams + self.bparams = bparams + self.method = method + self.tracker = tracker + m = iparams["m"] + n = iparams["n"] + + # variable parameters + self.X = X + self.E = torch.zeros((n,), dtype=X.dtype, device=X.device) + self.R = torch.zeros((m, n), dtype=X.dtype, device=X.device) + self.S = torch.zeros((m, 3 * n), dtype=X.dtype, device=X.device) + self.tvars: dict[str, Tensor] = {} + self.ivars: dict[str, int] = {"istep": 0} + self.fvars: dict[str, float] = {"_": 0.0} + self.bvars: dict[str, bool] = {"_": False} + + def __str__(self): + lines = ["LOPBCG:"] + lines += [f" iparams={self.iparams}"] + lines += [f" fparams={self.fparams}"] + lines += [f" bparams={self.bparams}"] + lines += [f" ivars={self.ivars}"] + lines += [f" fvars={self.fvars}"] + lines += [f" bvars={self.bvars}"] + lines += [f" tvars={self.tvars}"] + lines += [f" A={self.A}"] + lines += [f" B={self.B}"] + lines += [f" iK={self.iK}"] + lines += [f" X={self.X}"] + lines += [f" E={self.E}"] + r = "" + for line in lines: + r += line + "\n" + return r + + def update(self): + """Set and update iteration variables.""" + if self.ivars["istep"] == 0: + X_norm = float(torch.norm(self.X)) + iX_norm = X_norm**-1 + A_norm = float(torch.norm(_utils.matmul(self.A, self.X))) * iX_norm + B_norm = float(torch.norm(_utils.matmul(self.B, self.X))) * iX_norm + self.fvars["X_norm"] = X_norm + self.fvars["A_norm"] = A_norm + self.fvars["B_norm"] = B_norm + self.ivars["iterations_left"] = self.iparams["niter"] + self.ivars["converged_count"] = 0 + self.ivars["converged_end"] = 0 + + if self.method == "ortho": + self._update_ortho() + else: + self._update_basic() + + self.ivars["iterations_left"] = self.ivars["iterations_left"] - 1 + self.ivars["istep"] = self.ivars["istep"] + 1 + + def update_residual(self): + """Update residual R from A, B, X, E.""" + mm = _utils.matmul + self.R = mm(self.A, self.X) - mm(self.B, self.X) * self.E + + def update_converged_count(self): + """Determine the number of converged eigenpairs using backward stable + convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018]. + + Users may redefine this method for custom convergence criteria. + """ + # (...) -> int + prev_count = self.ivars["converged_count"] + tol = self.fparams["tol"] + A_norm = self.fvars["A_norm"] + B_norm = self.fvars["B_norm"] + E, X, R = self.E, self.X, self.R + rerr = torch.norm(R, 2, (0,)) / ( + torch.norm(X, 2, (0,)) * (A_norm + torch.abs(E[: X.shape[-1]]) * B_norm) + ) + converged = rerr < tol + count = 0 + for b in converged: + if not b: + # ignore convergence of following pairs to ensure + # strict ordering of eigenpairs + break + count += 1 + assert count >= prev_count, ( + f"the number of converged eigenpairs (was {prev_count}, got {count}) cannot decrease" + ) + self.ivars["converged_count"] = count + self.tvars["rerr"] = rerr + return count + + def stop_iteration(self): + """Return True to stop iterations. + + Note that tracker (if defined) can force-stop iterations by + setting ``worker.bvars['force_stop'] = True``. + """ + return ( + self.bvars.get("force_stop", False) + or self.ivars["iterations_left"] == 0 + or self.ivars["converged_count"] >= self.iparams["k"] + ) + + def run(self): + """Run LOBPCG iterations. + + Use this method as a template for implementing LOBPCG + iteration scheme with custom tracker that is compatible with + TorchScript. + """ + self.update() + + if not torch.jit.is_scripting() and self.tracker is not None: + self.call_tracker() + + while not self.stop_iteration(): + self.update() + + if not torch.jit.is_scripting() and self.tracker is not None: + self.call_tracker() + + @torch.jit.unused + def call_tracker(self): + """Interface for tracking iteration process in Python mode. + + Tracking the iteration process is disabled in TorchScript + mode. In fact, one should specify tracker=None when JIT + compiling functions using lobpcg. + """ + # do nothing when in TorchScript mode + + # Internal methods + + def _update_basic(self): + """ + Update or initialize iteration variables when `method == "basic"`. + """ + mm = torch.matmul + ns = self.ivars["converged_end"] + nc = self.ivars["converged_count"] + n = self.iparams["n"] + largest = self.bparams["largest"] + + if self.ivars["istep"] == 0: + Ri = self._get_rayleigh_ritz_transform(self.X) + M = _utils.qform(_utils.qform(self.A, self.X), Ri) + E, Z = _utils.symeig(M, largest) + self.X[:] = mm(self.X, mm(Ri, Z)) + self.E[:] = E + np = 0 + self.update_residual() + nc = self.update_converged_count() + self.S[..., :n] = self.X + + W = _utils.matmul(self.iK, self.R) + self.ivars["converged_end"] = ns = n + np + W.shape[-1] + self.S[:, n + np : ns] = W + else: + S_ = self.S[:, nc:ns] + Ri = self._get_rayleigh_ritz_transform(S_) + M = _utils.qform(_utils.qform(self.A, S_), Ri) + E_, Z = _utils.symeig(M, largest) + self.X[:, nc:] = mm(S_, mm(Ri, Z[:, : n - nc])) + self.E[nc:] = E_[: n - nc] + P = mm(S_, mm(Ri, Z[:, n : 2 * n - nc])) + np = P.shape[-1] + + self.update_residual() + nc = self.update_converged_count() + self.S[..., :n] = self.X + self.S[:, n : n + np] = P + W = _utils.matmul(self.iK, self.R[:, nc:]) + + self.ivars["converged_end"] = ns = n + np + W.shape[-1] + self.S[:, n + np : ns] = W + + def _update_ortho(self): + """ + Update or initialize iteration variables when `method == "ortho"`. + """ + mm = torch.matmul + ns = self.ivars["converged_end"] + nc = self.ivars["converged_count"] + n = self.iparams["n"] + largest = self.bparams["largest"] + + if self.ivars["istep"] == 0: + Ri = self._get_rayleigh_ritz_transform(self.X) + M = _utils.qform(_utils.qform(self.A, self.X), Ri) + _E, Z = _utils.symeig(M, largest) + self.X = mm(self.X, mm(Ri, Z)) + self.update_residual() + np = 0 + nc = self.update_converged_count() + self.S[:, :n] = self.X + W = self._get_ortho(self.R, self.X) + ns = self.ivars["converged_end"] = n + np + W.shape[-1] + self.S[:, n + np : ns] = W + + else: + S_ = self.S[:, nc:ns] + # Rayleigh-Ritz procedure + E_, Z = _utils.symeig(_utils.qform(self.A, S_), largest) + + # Update E, X, P + self.X[:, nc:] = mm(S_, Z[:, : n - nc]) + self.E[nc:] = E_[: n - nc] + P = mm(S_, mm(Z[:, n - nc :], _utils.basis(Z[: n - nc, n - nc :].mT))) + np = P.shape[-1] + + # check convergence + self.update_residual() + nc = self.update_converged_count() + + # update S + self.S[:, :n] = self.X + self.S[:, n : n + np] = P + W = self._get_ortho(self.R[:, nc:], self.S[:, : n + np]) + ns = self.ivars["converged_end"] = n + np + W.shape[-1] + self.S[:, n + np : ns] = W + + def _get_rayleigh_ritz_transform(self, S): + """Return a transformation matrix that is used in Rayleigh-Ritz + procedure for reducing a general eigenvalue problem :math:`(S^TAS) + C = (S^TBS) C E` to a standard eigenvalue problem :math: `(Ri^T + S^TAS Ri) Z = Z E` where `C = Ri Z`. + + .. note:: In the original Rayleight-Ritz procedure in + [DuerschEtal2018], the problem is formulated as follows:: + + SAS = S^T A S + SBS = S^T B S + D = () ** -1/2 + R^T R = Cholesky(D SBS D) + Ri = D R^-1 + solve symeig problem Ri^T SAS Ri Z = Theta Z + C = Ri Z + + To reduce the number of matrix products (denoted by empty + space between matrices), here we introduce element-wise + products (denoted by symbol `*`) so that the Rayleight-Ritz + procedure becomes:: + + SAS = S^T A S + SBS = S^T B S + d = () ** -1/2 # this is 1-d column vector + dd = d d^T # this is 2-d matrix + R^T R = Cholesky(dd * SBS) + Ri = R^-1 * d # broadcasting + solve symeig problem Ri^T SAS Ri Z = Theta Z + C = Ri Z + + where `dd` is 2-d matrix that replaces matrix products `D M + D` with one element-wise product `M * dd`; and `d` replaces + matrix product `D M` with element-wise product `M * + d`. Also, creating the diagonal matrix `D` is avoided. + + Args: + S (Tensor): the matrix basis for the search subspace, size is + :math:`(m, n)`. + + Returns: + Ri (tensor): upper-triangular transformation matrix of size + :math:`(n, n)`. + + """ + B = self.B + SBS = _utils.qform(B, S) + d_row = SBS.diagonal(0, -2, -1) ** -0.5 + d_col = d_row.reshape(d_row.shape[0], 1) + # TODO use torch.linalg.cholesky_solve once it is implemented + R = torch.linalg.cholesky((SBS * d_row) * d_col, upper=True) + return torch.linalg.solve_triangular( + R, d_row.diag_embed(), upper=True, left=False + ) + + def _get_svqb(self, U: Tensor, drop: bool, tau: float) -> Tensor: + """Return B-orthonormal U. + + .. note:: When `drop` is `False` then `svqb` is based on the + Algorithm 4 from [DuerschPhD2015] that is a slight + modification of the corresponding algorithm + introduced in [StathopolousWu2002]. + + Args: + + U (Tensor) : initial approximation, size is (m, n) + drop (bool) : when True, drop columns that + contribution to the `span([U])` is small. + tau (float) : positive tolerance + + Returns: + + U (Tensor) : B-orthonormal columns (:math:`U^T B U = I`), size + is (m, n1), where `n1 = n` if `drop` is `False, + otherwise `n1 <= n`. + + """ + if torch.numel(U) == 0: + return U + UBU = _utils.qform(self.B, U) + d = UBU.diagonal(0, -2, -1) + + # Detect and drop exact zero columns from U. While the test + # `abs(d) == 0` is unlikely to be True for random data, it is + # possible to construct input data to lobpcg where it will be + # True leading to a failure (notice the `d ** -0.5` operation + # in the original algorithm). To prevent the failure, we drop + # the exact zero columns here and then continue with the + # original algorithm below. + nz = torch.where(abs(d) != 0.0) + assert len(nz) == 1, nz + if len(nz[0]) < len(d): + U = U[:, nz[0]] + if torch.numel(U) == 0: + return U + UBU = _utils.qform(self.B, U) + d = UBU.diagonal(0, -2, -1) + nz = torch.where(abs(d) != 0.0) + assert len(nz[0]) == len(d) + + # The original algorithm 4 from [DuerschPhD2015]. + d_col = (d**-0.5).reshape(d.shape[0], 1) + DUBUD = (UBU * d_col) * d_col.mT + E, Z = _utils.symeig(DUBUD) + t = tau * abs(E).max() + if drop: + keep = torch.where(E > t) + assert len(keep) == 1, keep + E = E[keep[0]] + Z = Z[:, keep[0]] + d_col = d_col[keep[0]] + else: + E[(torch.where(E < t))[0]] = t + + return torch.matmul(U * d_col.mT, Z * E**-0.5) + + def _get_ortho(self, U, V): + """Return B-orthonormal U with columns are B-orthogonal to V. + + .. note:: When `bparams["ortho_use_drop"] == False` then + `_get_ortho` is based on the Algorithm 3 from + [DuerschPhD2015] that is a slight modification of + the corresponding algorithm introduced in + [StathopolousWu2002]. Otherwise, the method + implements Algorithm 6 from [DuerschPhD2015] + + .. note:: If all U columns are B-collinear to V then the + returned tensor U will be empty. + + Args: + + U (Tensor) : initial approximation, size is (m, n) + V (Tensor) : B-orthogonal external basis, size is (m, k) + + Returns: + + U (Tensor) : B-orthonormal columns (:math:`U^T B U = I`) + such that :math:`V^T B U=0`, size is (m, n1), + where `n1 = n` if `drop` is `False, otherwise + `n1 <= n`. + """ + mm = torch.matmul + mm_B = _utils.matmul + m = self.iparams["m"] + tau_ortho = self.fparams["ortho_tol"] + tau_drop = self.fparams["ortho_tol_drop"] + tau_replace = self.fparams["ortho_tol_replace"] + i_max = self.iparams["ortho_i_max"] + j_max = self.iparams["ortho_j_max"] + # when use_drop==True, enable dropping U columns that have + # small contribution to the `span([U, V])`. + use_drop = self.bparams["ortho_use_drop"] + + # clean up variables from the previous call + for vkey in list(self.fvars.keys()): + if vkey.startswith("ortho_") and vkey.endswith("_rerr"): + self.fvars.pop(vkey) + self.ivars.pop("ortho_i", 0) + self.ivars.pop("ortho_j", 0) + + BV_norm = torch.norm(mm_B(self.B, V)) + BU = mm_B(self.B, U) + VBU = mm(V.mT, BU) + i = j = 0 + for i in range(i_max): + U = U - mm(V, VBU) + drop = False + tau_svqb = tau_drop + for j in range(j_max): + if use_drop: + U = self._get_svqb(U, drop, tau_svqb) + drop = True + tau_svqb = tau_replace + else: + U = self._get_svqb(U, False, tau_replace) + if torch.numel(U) == 0: + # all initial U columns are B-collinear to V + self.ivars["ortho_i"] = i + self.ivars["ortho_j"] = j + return U + BU = mm_B(self.B, U) + UBU = mm(U.mT, BU) + U_norm = torch.norm(U) + BU_norm = torch.norm(BU) + R = UBU - torch.eye(UBU.shape[-1], device=UBU.device, dtype=UBU.dtype) + R_norm = torch.norm(R) + # https://github.com/pytorch/pytorch/issues/33810 workaround: + rerr = float(R_norm) * float(BU_norm * U_norm) ** -1 + vkey = f"ortho_UBUmI_rerr[{i}, {j}]" + self.fvars[vkey] = rerr + if rerr < tau_ortho: + break + VBU = mm(V.mT, BU) + VBU_norm = torch.norm(VBU) + U_norm = torch.norm(U) + rerr = float(VBU_norm) * float(BV_norm * U_norm) ** -1 + vkey = f"ortho_VBU_rerr[{i}]" + self.fvars[vkey] = rerr + if rerr < tau_ortho: + break + if m < U.shape[-1] + V.shape[-1]: + # TorchScript needs the class var to be assigned to a local to + # do optional type refinement + B = self.B + assert B is not None + raise ValueError( + "Overdetermined shape of U:" + f" #B-cols(={B.shape[-1]}) >= #U-cols(={U.shape[-1]}) + #V-cols(={V.shape[-1]}) must hold" + ) + self.ivars["ortho_i"] = i + self.ivars["ortho_j"] = j + return U + + +# Calling tracker is separated from LOBPCG definitions because +# TorchScript does not support user-defined callback arguments: +LOBPCG_call_tracker_orig = LOBPCG.call_tracker + + +def LOBPCG_call_tracker(self): + self.tracker(self) diff --git a/venv/lib/python3.10/site-packages/torch/_lowrank.py b/venv/lib/python3.10/site-packages/torch/_lowrank.py new file mode 100644 index 0000000000000000000000000000000000000000..182883cfc5e593fbfcc80c021a2803dfbd197d94 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_lowrank.py @@ -0,0 +1,294 @@ +"""Implement various linear algebra algorithms for low rank matrices.""" + +__all__ = ["svd_lowrank", "pca_lowrank"] + +from typing import Optional + +import torch +from torch import _linalg_utils as _utils, Tensor +from torch.overrides import handle_torch_function, has_torch_function + + +def get_approximate_basis( + A: Tensor, + q: int, + niter: Optional[int] = 2, + M: Optional[Tensor] = None, +) -> Tensor: + """Return tensor :math:`Q` with :math:`q` orthonormal columns such + that :math:`Q Q^H A` approximates :math:`A`. If :math:`M` is + specified, then :math:`Q` is such that :math:`Q Q^H (A - M)` + approximates :math:`A - M`. without instantiating any tensors + of the size of :math:`A` or :math:`M`. + + .. note:: The implementation is based on the Algorithm 4.4 from + Halko et al., 2009. + + .. note:: For an adequate approximation of a k-rank matrix + :math:`A`, where k is not known in advance but could be + estimated, the number of :math:`Q` columns, q, can be + chosen according to the following criteria: in general, + :math:`k <= q <= min(2*k, m, n)`. For large low-rank + matrices, take :math:`q = k + 5..10`. If k is + relatively small compared to :math:`min(m, n)`, choosing + :math:`q = k + 0..2` may be sufficient. + + .. note:: To obtain repeatable results, reset the seed for the + pseudorandom number generator + + Args:: + A (Tensor): the input tensor of size :math:`(*, m, n)` + + q (int): the dimension of subspace spanned by :math:`Q` + columns. + + niter (int, optional): the number of subspace iterations to + conduct; ``niter`` must be a + nonnegative integer. In most cases, the + default value 2 is more than enough. + + M (Tensor, optional): the input tensor's mean of size + :math:`(*, m, n)`. + + References:: + - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding + structure with randomness: probabilistic algorithms for + constructing approximate matrix decompositions, + arXiv:0909.4061 [math.NA; math.PR], 2009 (available at + `arXiv `_). + """ + + niter = 2 if niter is None else niter + dtype = _utils.get_floating_dtype(A) if not A.is_complex() else A.dtype + matmul = _utils.matmul + + R = torch.randn(A.shape[-1], q, dtype=dtype, device=A.device) + + # The following code could be made faster using torch.geqrf + torch.ormqr + # but geqrf is not differentiable + + X = matmul(A, R) + if M is not None: + X = X - matmul(M, R) + Q = torch.linalg.qr(X).Q + for _ in range(niter): + X = matmul(A.mH, Q) + if M is not None: + X = X - matmul(M.mH, Q) + Q = torch.linalg.qr(X).Q + X = matmul(A, Q) + if M is not None: + X = X - matmul(M, Q) + Q = torch.linalg.qr(X).Q + return Q + + +def svd_lowrank( + A: Tensor, + q: Optional[int] = 6, + niter: Optional[int] = 2, + M: Optional[Tensor] = None, +) -> tuple[Tensor, Tensor, Tensor]: + r"""Return the singular value decomposition ``(U, S, V)`` of a matrix, + batches of matrices, or a sparse matrix :math:`A` such that + :math:`A \approx U \operatorname{diag}(S) V^{\text{H}}`. In case :math:`M` is given, then + SVD is computed for the matrix :math:`A - M`. + + .. note:: The implementation is based on the Algorithm 5.1 from + Halko et al., 2009. + + .. note:: For an adequate approximation of a k-rank matrix + :math:`A`, where k is not known in advance but could be + estimated, the number of :math:`Q` columns, q, can be + chosen according to the following criteria: in general, + :math:`k <= q <= min(2*k, m, n)`. For large low-rank + matrices, take :math:`q = k + 5..10`. If k is + relatively small compared to :math:`min(m, n)`, choosing + :math:`q = k + 0..2` may be sufficient. + + .. note:: This is a randomized method. To obtain repeatable results, + set the seed for the pseudorandom number generator + + .. note:: In general, use the full-rank SVD implementation + :func:`torch.linalg.svd` for dense matrices due to its 10x + higher performance characteristics. The low-rank SVD + will be useful for huge sparse matrices that + :func:`torch.linalg.svd` cannot handle. + + Args:: + A (Tensor): the input tensor of size :math:`(*, m, n)` + + q (int, optional): a slightly overestimated rank of A. + + niter (int, optional): the number of subspace iterations to + conduct; niter must be a nonnegative + integer, and defaults to 2 + + M (Tensor, optional): the input tensor's mean of size + :math:`(*, m, n)`, which will be broadcasted + to the size of A in this function. + + References:: + - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding + structure with randomness: probabilistic algorithms for + constructing approximate matrix decompositions, + arXiv:0909.4061 [math.NA; math.PR], 2009 (available at + `arXiv `_). + + """ + if not torch.jit.is_scripting(): + tensor_ops = (A, M) + if not set(map(type, tensor_ops)).issubset( + (torch.Tensor, type(None)) + ) and has_torch_function(tensor_ops): + return handle_torch_function( + svd_lowrank, tensor_ops, A, q=q, niter=niter, M=M + ) + return _svd_lowrank(A, q=q, niter=niter, M=M) + + +def _svd_lowrank( + A: Tensor, + q: Optional[int] = 6, + niter: Optional[int] = 2, + M: Optional[Tensor] = None, +) -> tuple[Tensor, Tensor, Tensor]: + # Algorithm 5.1 in Halko et al., 2009 + + q = 6 if q is None else q + m, n = A.shape[-2:] + matmul = _utils.matmul + if M is not None: + M = M.broadcast_to(A.size()) + + # Assume that A is tall + if m < n: + A = A.mH + if M is not None: + M = M.mH + + Q = get_approximate_basis(A, q, niter=niter, M=M) + B = matmul(Q.mH, A) + if M is not None: + B = B - matmul(Q.mH, M) + U, S, Vh = torch.linalg.svd(B, full_matrices=False) + V = Vh.mH + U = Q.matmul(U) + + if m < n: + U, V = V, U + + return U, S, V + + +def pca_lowrank( + A: Tensor, + q: Optional[int] = None, + center: bool = True, + niter: int = 2, +) -> tuple[Tensor, Tensor, Tensor]: + r"""Performs linear Principal Component Analysis (PCA) on a low-rank + matrix, batches of such matrices, or sparse matrix. + + This function returns a namedtuple ``(U, S, V)`` which is the + nearly optimal approximation of a singular value decomposition of + a centered matrix :math:`A` such that :math:`A \approx U \operatorname{diag}(S) V^{\text{H}}` + + .. note:: The relation of ``(U, S, V)`` to PCA is as follows: + + - :math:`A` is a data matrix with ``m`` samples and + ``n`` features + + - the :math:`V` columns represent the principal directions + + - :math:`S ** 2 / (m - 1)` contains the eigenvalues of + :math:`A^T A / (m - 1)` which is the covariance of + ``A`` when ``center=True`` is provided. + + - ``matmul(A, V[:, :k])`` projects data to the first k + principal components + + .. note:: Different from the standard SVD, the size of returned + matrices depend on the specified rank and q + values as follows: + + - :math:`U` is m x q matrix + + - :math:`S` is q-vector + + - :math:`V` is n x q matrix + + .. note:: To obtain repeatable results, reset the seed for the + pseudorandom number generator + + Args: + + A (Tensor): the input tensor of size :math:`(*, m, n)` + + q (int, optional): a slightly overestimated rank of + :math:`A`. By default, ``q = min(6, m, + n)``. + + center (bool, optional): if True, center the input tensor, + otherwise, assume that the input is + centered. + + niter (int, optional): the number of subspace iterations to + conduct; niter must be a nonnegative + integer, and defaults to 2. + + References:: + + - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding + structure with randomness: probabilistic algorithms for + constructing approximate matrix decompositions, + arXiv:0909.4061 [math.NA; math.PR], 2009 (available at + `arXiv `_). + + """ + + if not torch.jit.is_scripting(): + if type(A) is not torch.Tensor and has_torch_function((A,)): + return handle_torch_function( + pca_lowrank, (A,), A, q=q, center=center, niter=niter + ) + + (m, n) = A.shape[-2:] + + if q is None: + q = min(6, m, n) + elif not (q >= 0 and q <= min(m, n)): + raise ValueError( + f"q(={q}) must be non-negative integer and not greater than min(m, n)={min(m, n)}" + ) + if not (niter >= 0): + raise ValueError(f"niter(={niter}) must be non-negative integer") + + dtype = _utils.get_floating_dtype(A) + + if not center: + return _svd_lowrank(A, q, niter=niter, M=None) + + if _utils.is_sparse(A): + if len(A.shape) != 2: + raise ValueError("pca_lowrank input is expected to be 2-dimensional tensor") + c = torch.sparse.sum(A, dim=(-2,)) / m + # reshape c + column_indices = c.indices()[0] + indices = torch.zeros( + 2, + len(column_indices), + dtype=column_indices.dtype, + device=column_indices.device, + ) + indices[0] = column_indices + C_t = torch.sparse_coo_tensor( + indices, c.values(), (n, 1), dtype=dtype, device=A.device + ) + + ones_m1_t = torch.ones(A.shape[:-2] + (1, m), dtype=dtype, device=A.device) + M = torch.sparse.mm(C_t, ones_m1_t).mT + return _svd_lowrank(A, q, niter=niter, M=M) + else: + C = A.mean(dim=(-2,), keepdim=True) + return _svd_lowrank(A - C, q, niter=niter, M=None) diff --git a/venv/lib/python3.10/site-packages/torch/_meta_registrations.py b/venv/lib/python3.10/site-packages/torch/_meta_registrations.py new file mode 100644 index 0000000000000000000000000000000000000000..9202b4da41d28672b3e60348934d2aa247388cd8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_meta_registrations.py @@ -0,0 +1,7968 @@ +# mypy: allow-untyped-defs +import math +from collections.abc import Sequence +from enum import Enum +from functools import wraps +from typing import Callable, Optional, TypeVar, Union +from typing_extensions import ParamSpec + +import torch +import torch._prims_common as utils +from torch import SymBool, SymFloat, Tensor +from torch._decomp import ( + _add_op_to_registry, + _convert_out_params, + global_decomposition_table, + meta_table, +) +from torch._ops import OpOverload +from torch._prims import _prim_elementwise_meta, ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND +from torch._prims_common import ( + BoolLike, + corresponding_complex_dtype, + corresponding_real_dtype, + elementwise_dtypes, + ELEMENTWISE_TYPE_PROMOTION_KIND, + FloatLike, + IntLike, + make_contiguous_strides_for, + Number, + suggest_memory_format, + TensorLike, +) +from torch._prims_common.wrappers import ( + _maybe_convert_to_dtype, + _maybe_resize_out, + _resize_output_check, + _safe_copy_out, + out_wrapper, +) +from torch._refs import _broadcast_shapes, _maybe_broadcast +from torch.fx.experimental import _config as exp_config +from torch.utils import _pytree as pytree + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +aten = torch.ops.aten + +_meta_lib_dont_use_me_use_register_meta = torch.library.Library("aten", "IMPL", "Meta") +MODE_SUM, MODE_MEAN, MODE_MAX = range(3) + + +def register_meta(op) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + def wrapper(fn): + fn = _convert_out_params(fn) + + def register(op): + _add_op_to_registry(meta_table, op, fn) + + pytree.tree_map_(register, op) + return fn + + return wrapper + + +def elementwise_meta( + *args, + type_promotion: ELEMENTWISE_TYPE_PROMOTION_KIND, +): + # Perform type promotion, as this is expected from prim_metafunction + _, result_dtype = utils.elementwise_dtypes( + *args, + type_promotion_kind=type_promotion, + ) + args = [_maybe_convert_to_dtype(x, result_dtype) for x in args] + + # Broadcast + args = _maybe_broadcast(*args) + + # Perform prim checks + return _prim_elementwise_meta( + *args, type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT + ) + + +def toRealValueType(dtype): + from_complex = { + torch.complex32: torch.half, + torch.cfloat: torch.float, + torch.cdouble: torch.double, + } + return from_complex.get(dtype, dtype) + + +def check_inplace_broadcast(self_shape, *args_shape): + broadcasted_shape = tuple(_broadcast_shapes(self_shape, *args_shape)) + torch._check( + broadcasted_shape == self_shape, + lambda: f"output with shape {self_shape} doesn't match the broadcast shape {broadcasted_shape}", + ) + + +@register_meta([aten.linspace, aten.logspace]) +@out_wrapper() +def meta_linspace_logspace( + start, + end, + steps, + base=None, + dtype=None, + device=None, + layout=torch.strided, + pin_memory=False, + requires_grad=False, +): + if isinstance(start, torch.Tensor): + torch._check( + start.dim() == 0, + lambda: "linspace only supports 0-dimensional start and end tensors", + ) + if isinstance(end, torch.Tensor): + torch._check( + end.dim() == 0, + lambda: "linspace only supports 0-dimensional start and end tensors", + ) + + if any(isinstance(arg, complex) for arg in (start, end, steps)): + default_complex_dtype = utils.corresponding_complex_dtype( + torch.get_default_dtype() + ) + if dtype is None: + dtype = default_complex_dtype + else: + torch._check( + utils.is_complex_dtype(dtype), + lambda: f"linspace(): inferred dtype {default_complex_dtype} can't be safely cast to passed dtype {dtype}", + ) + else: + dtype = dtype or torch.get_default_dtype() + assert isinstance(dtype, torch.dtype) + + # steps does not participate in the computation of the dtype + torch._check_type( + isinstance(steps, IntLike), + lambda: f"received an invalid combination of arguments - got \ +({type(start).__name__}, {type(end).__name__}, {type(steps).__name__})", + ) + assert isinstance(steps, IntLike) # for mypy + torch._check(steps >= 0, lambda: "number of steps must be non-negative") + + return torch.empty( + (steps,), # type: ignore[arg-type] + dtype=dtype, + layout=layout, + device="meta", + pin_memory=pin_memory, + requires_grad=requires_grad, + ) + + +@register_meta([aten.take.default, aten.take.out]) +@out_wrapper() +def meta_take(self, index): + # Type and device checks + torch._check( + index.dtype == torch.long, + lambda: f"take(): Expected a long tensor for index, but got {index.dtype}", + ) + # Index checks + torch._check_index( + not (self.numel() == 0 and index.numel() != 0), + lambda: "take(): tried to take from an empty tensor", + ) + return self.new_empty(index.shape) + + +@register_meta([aten.linalg_cross.default, aten.linalg_cross.out]) +@out_wrapper() +def linalg_cross(self, other, *, dim=-1): + x_d = self.ndim + y_d = other.ndim + torch._check( + x_d == y_d, + lambda: "linalg.cross: inputs must have the same number of dimensions.", + ) + torch._check( + self.size(dim) == 3 and other.size(dim) == 3, + lambda: ( + f"linalg.cross: inputs dimension {dim} must have length 3. " + f"Got {self.size(dim)} and {other.size(dim)}" + ), + ) + out_shape = _broadcast_shapes(self.shape, other.shape) + return self.new_empty(out_shape) + + +@register_meta(aten.linalg_matrix_exp) +@out_wrapper() +def linalg_matrix_exp(self): + squareCheckInputs(self, "linalg.matrix_exp") + checkFloatingOrComplex(self, "linalg.matrix_exp") + return torch.empty_like(self, memory_format=torch.contiguous_format) + + +@register_meta( + [aten.cummax.default, aten.cummax.out, aten.cummin.default, aten.cummin.out] +) +@out_wrapper("values", "indices") +def cummaxmin(self, dim): + values = torch.empty(self.shape, device=self.device, dtype=self.dtype) + indices = torch.empty(self.shape, device=self.device, dtype=torch.int64) + if self.numel() != 0 and self.ndim != 0: + # Checks that dim is within bounds + maybe_wrap_dim(dim, self.ndim) + return values, indices + + +@register_meta([aten.logcumsumexp.default, aten.logcumsumexp.out]) +@out_wrapper() +def logcumsumexp(self, dim): + # Checks that dim is within bounds + maybe_wrap_dim(dim, self.ndim) + return torch.empty_like(self, memory_format=torch.contiguous_format) + + +# Stride-related code from _exec_fft in aten/src/ATen/native/mkl/SpectralOps.cpp +# and aten/src/ATen/cuda/SpectralOps.cpp +# +# Although the actual FFT launch is different, all the permuting code appears +# to be the same +def _exec_fft(out, self, out_sizes, dim, *, forward): + ndim = self.ndim + signal_ndim = len(dim) + batch_dims = ndim - signal_ndim + + # Permute dimensions so batch dimensions come first, and in stride order + dim_permute = list(range(ndim)) + + is_transformed_dim = [False for _ in range(ndim)] + for d in dim: + is_transformed_dim[d] = True + + # std::partition + left, right = [], [] + for d in dim_permute: + if not is_transformed_dim[d]: + left.append(d) + else: + right.append(d) + dim_permute = left + right + batch_end = len(left) + + self_strides = self.stride() + tmp = dim_permute[:batch_end] + tmp.sort(key=lambda x: self_strides[x], reverse=True) + dim_permute = tmp + dim_permute[batch_end:] + input = self.permute(dim_permute) + + # Collapse batch dimensions into a single dimension + batched_sizes = [-1] + list(input.shape[batch_dims:]) + input = input.reshape(batched_sizes) + + batch_size = input.size(0) + batched_sizes[0] = batch_size + batched_out_sizes = list(batched_sizes) + for i in range(len(dim)): + batched_out_sizes[i + 1] = out_sizes[dim[i]] + out.resize_(batched_out_sizes, memory_format=torch.contiguous_format) + + # Inplace reshaping to original batch shape and inverting the dimension permutation + out_strides = [0 for _ in range(ndim)] + batch_numel = 1 + i = batch_dims - 1 + while i >= 0: + out_strides[dim_permute[i]] = batch_numel * out.stride(0) + batch_numel *= out_sizes[dim_permute[i]] + i -= 1 + for i in range(batch_dims, ndim): + out_strides[dim_permute[i]] = out.stride(1 + (i - batch_dims)) + out.as_strided_(out_sizes, out_strides, out.storage_offset()) + + return out + + +def _sort_dims(self: Tensor, dim: list[int], exclude_last: bool = False): + sorted_dims = list(dim) + self_strides = self.stride() + sorted_dims[: len(sorted_dims) - int(exclude_last)].sort( + key=lambda i: self_strides[i] + ) + return sorted_dims + + +# See _fft_c2c_cufft in aten/src/ATen/native/cuda/SpectralOps.cpp +# and _fft_c2c_mkl in aten/src/ATen/native/mkl/SpectralOps.cpp +@register_meta([aten._fft_c2c.default, aten._fft_c2c.out]) +@out_wrapper() +def meta_fft_c2c(self, dim, normalization, forward): + torch._check(self.dtype.is_complex) + if not dim: + return self.clone() + + sorted_dims = _sort_dims(self, dim) + out = self.new_empty(self.size()) + return _exec_fft(out, self, self.size(), sorted_dims, forward=forward) + + +cufft_max_ndim = 3 + + +def use_optimized_cufft_path(dim: list[int]): + if len(dim) > cufft_max_ndim or (len(dim) >= 2 and dim[0] == 0 and dim[1] == 1): + return False + else: + return True + + +@register_meta([aten._fft_r2c.default, aten._fft_r2c.out]) +@out_wrapper() +def meta_fft_r2c(self, dim, normalization, onesided): + torch._check(self.dtype.is_floating_point) + input_sizes = list(self.size()) + out_sizes = list(input_sizes) + last_dim = dim[-1] + last_dim_halfsize = input_sizes[last_dim] // 2 + 1 + onesided_sizes = list(input_sizes) + onesided_sizes[last_dim] = last_dim_halfsize + + if onesided: + out_sizes[last_dim] = last_dim_halfsize + + if device_hint(self) == "cuda" or device_hint(self) == "xpu": + # _fft_r2c_cufft in aten/src/ATen/native/cuda/SpectralOps.cpp + # _fft_r2c_xpu in torch-xpu-ops/src/ATen/native/xpu/SpectralOps.cpp + output = self.new_empty( + out_sizes, dtype=utils.corresponding_complex_dtype(self.dtype) + ) + + working_tensor = self + if device_hint(self) == "cuda" and use_optimized_cufft_path(dim): + _exec_fft(output, working_tensor, out_sizes, dim, forward=True) + else: + # First do the R2C transform on the last dimension + target_sizes = out_sizes if len(dim) == 1 else onesided_sizes + _exec_fft(output, working_tensor, target_sizes, [last_dim], forward=True) + if len(dim) > 1: + working_tensor = self.new_empty( + out_sizes, dtype=utils.corresponding_complex_dtype(self.dtype) + ) + + # Then any remaining C2C transforms + sorted_dims = dim[:-1] + while sorted_dims: + output, working_tensor = working_tensor, output + strides = working_tensor.stride() + sorted_dims.sort( + key=lambda i: strides[i], reverse=True + ) # NB reverse! Not sure if this is og bug + max_dims = min(cufft_max_ndim, len(sorted_dims)) + last_dims = sorted_dims[len(sorted_dims) - max_dims :] + _exec_fft( + output, working_tensor, onesided_sizes, last_dims, forward=True + ) + sorted_dims = sorted_dims[: len(sorted_dims) - max_dims] + + if not onesided: + if output.size(last_dim) != out_sizes[last_dim]: + working_tensor.resize_(out_sizes, memory_format=torch.contiguous_format) + output = working_tensor + + return output + + else: + return self.new_empty( + out_sizes, dtype=utils.corresponding_complex_dtype(self.dtype) + ) + + +@register_meta(aten.randperm.generator_out) +def meta_randperm(n, *, generator=None, out): + return _maybe_resize_out(out, torch.Size([n])) + + +@register_meta(aten.randperm.default) +def meta_randperm_default( + n, + *, + dtype=torch.long, + layout=None, + device=None, + pin_memory=None, +): + return torch.empty( + n, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta([aten.randint.default, aten.randint.out]) +@out_wrapper() +def meta_randint( + high, + size, + *, + dtype=torch.long, + layout=None, + device=None, + pin_memory=None, +): + low = 0 + torch._check( + high > low, + lambda: f"random_ expects 'from' to be less than 'to', but got from={low} >= to={high}", + ) + return torch.empty( + size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta([aten.randint.low, aten.randint.low_out]) +@out_wrapper() +def meta_randint_low( + low, + high, + size, + *, + dtype=torch.long, + layout=None, + device=None, + pin_memory=None, +): + torch._check( + high > low, + lambda: f"random_ expects 'from' to be less than 'to', but got from={low} >= to={high}", + ) + return torch.empty( + size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta([aten.rand.default, aten.rand.out]) +@out_wrapper() +def meta_rand_default(size, *, dtype=None, layout=None, device=None, pin_memory=None): + return torch.empty( + size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta([aten._fft_c2r.default, aten._fft_c2r.out]) +@out_wrapper() +def meta_fft_c2r(self: Tensor, dim: list[int], normalization: int, lastdim: int): + # _fft_c2r_mkl + torch._check(self.dtype.is_complex) + + if device_hint(self) == "cuda": + out_sizes = list(self.size()) + out_sizes[dim[-1]] = lastdim + + output = self.new_empty(out_sizes, dtype=toRealValueType(self.dtype)) + + if use_optimized_cufft_path(dim): + return _exec_fft( + output, + self.clone(memory_format=torch.contiguous_format), + out_sizes, + dim, + forward=False, + ) + else: + # First complete any C2C transforms + if len(dim) > 1: + temp = meta_fft_c2c(self, dim[:-1], 0, lastdim) # fft_norm_mode::none + else: + temp = self.clone(memory_format=torch.contiguous_format) + return _exec_fft(output, temp, out_sizes, [dim[-1]], forward=False) + + else: + input = self + if len(dim) > 1: + c2c_dims = dim[:-1] + input = meta_fft_c2c(self, c2c_dims, normalization, forward=False) + dim = dim[-1:] + + out_sizes = list(input.size()) + out_sizes[dim[-1]] = lastdim + out = self.new_empty(out_sizes, dtype=toRealValueType(self.dtype)) + return _exec_fft(out, input, out_sizes, dim, forward=False) + + +@register_meta(aten.copy_.default) +def meta_copy_(self, src, non_blocking=False): + # This code simulates the original decomp from inductor, + # which runs most of the meta checks that we care about. + # In theory, we should make this more robust by carefully + # auditing our C++ copy_() kernel and copying the checks here. + from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols + + # TODO: Ideally, we'd insert a deferred runtime assert here, but if we are + # calling an actual copy_, you'll get that automatically + # https://github.com/pytorch/pytorch/issues/122477 + if ( + not free_unbacked_symbols(self) and torch._debug_has_internal_overlap(self) == 1 + ): # 1 == MemOverlap::Yes + raise RuntimeError( + "more than one element of the written-to tensor refers to a single memory location" + ) + + if isinstance(src, Tensor): + intermediate = src.to(self, non_blocking) + if self.size() != intermediate.size(): + aten.expand_copy.default(intermediate, self.size()) + return self + + +def inferUnsqueezeGeometry(tensor, dim): + result_sizes = list(tensor.size()) + result_strides = list(tensor.stride()) + new_stride = 1 if dim >= tensor.dim() else result_sizes[dim] * result_strides[dim] + result_sizes.insert(dim, 1) + result_strides.insert(dim, new_stride) + return result_sizes, result_strides + + +@register_meta(aten.unsqueeze_.default) +def meta_unsqueeze_(self, dim): + dim = maybe_wrap_dim(dim, self.dim() + 1) + g_sizes, g_strides = inferUnsqueezeGeometry(self, dim) + self.as_strided_(g_sizes, g_strides) + return self + + +@register_meta(aten._sparse_semi_structured_linear) +def meta_sparse_structured_linear( + input: Tensor, + weight: Tensor, + _meta: Tensor, + bias: Optional[Tensor] = None, + _activation_opt: Optional[str] = None, + out_dtype: Optional[torch.dtype] = None, +): + output_sizes = list(input.shape) + if bias is not None: + assert weight.size(0) == bias.size(0), "output size mismatch" + assert weight.size(1) == input.size(-1) / 2 + output_sizes[-1] = weight.size(0) + + # see: https://github.com/pytorch/pytorch/pull/114477#issuecomment-1830121375 + # We assume that we have already squashed the inputs into a 2-D tensor + # Then, as the output is transposed, we need to propagate the transposed + # stride information to the output tensor + assert len(input.shape) == 2, "we can only handle the squashed input case" + transposed_strides = (1, input.size(0)) + + if out_dtype is not None: + assert input.dtype == torch.int8 and out_dtype == torch.int32, ( + "out_dtype is only supported for i8i8->i32 linear operator" + ) + output = input.new_empty( + output_sizes, + dtype=input.dtype if out_dtype is None else out_dtype, + ).as_strided(output_sizes, transposed_strides) + + return output + + +@register_meta(aten._sparse_semi_structured_mm) +def meta_sparse_structured_mm( + mat1: Tensor, + mat1_meta: Tensor, + mat2: Tensor, + out_dtype: Optional[torch.dtype] = None, +): + assert len(mat1.shape) == 2 + assert len(mat1_meta.shape) == 2 + assert len(mat2.shape) == 2 + assert mat1.size(1) == mat2.size(0) / 2 + output_sizes = [mat1.size(0), mat2.size(1)] + + if out_dtype is not None: + assert mat2.dtype == torch.int8 and out_dtype == torch.int32, ( + "out_dtype is only supported for i8i8->i32 linear operator" + ) + output = mat2.new_empty( + output_sizes, + dtype=mat2.dtype if out_dtype is None else out_dtype, + ) + + return output + + +@register_meta(aten._sparse_semi_structured_addmm) +def meta_sparse_structured_addmm( + input: Tensor, + mat1: Tensor, + mat1_meta: Tensor, + mat2: Tensor, + *, + alpha=1, + beta=1, + out_dtype: Optional[torch.dtype] = None, +): + assert len(input.shape) == 1, ( + "only input broadcasted to columns of mat1 * mat2 product is supported" + ) + assert len(mat1.shape) == 2 + assert len(mat1_meta.shape) == 2 + assert len(mat2.shape) == 2 + assert input.size(0) == mat1.size(0), ( + "only input broadcasted to columns of mat1 * mat2 product is supported" + ) + assert mat1.size(1) == mat2.size(0) / 2 + output_sizes = [mat1.size(0), mat2.size(1)] + + if out_dtype is not None: + assert mat2.dtype == torch.int8 and out_dtype == torch.int32, ( + "out_dtype is only supported for i8i8->i32 linear operator" + ) + output = mat2.new_empty( + output_sizes, + dtype=mat2.dtype if out_dtype is None else out_dtype, + ) + + return output + + +@register_meta(aten._cslt_sparse_mm) +def meta__cslt_sparse_mm( + compressed_A: torch.Tensor, + dense_B: torch.Tensor, + bias: Optional[Tensor] = None, + alpha: Optional[Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + transpose_result: bool = False, + alg_id: int = 0, + split_k: int = 1, + split_k_mode: int = -1, +): + assert dense_B.dtype in { + torch.float32, + torch.float16, + torch.bfloat16, + torch.int8, + torch.float8_e4m3fn, + }, "_cslt_sparse_mm only supports fp16, bf16, int8, and fp8e4m3" + assert compressed_A.dtype == dense_B.dtype, "inputs must have the same dtype" + assert len(dense_B.shape) == 2, "_cslt_sparse_mm only supports 2d inputs" + + is_8bit_input_type = compressed_A.dtype in [torch.int8, torch.float8_e4m3fn] + compression_factor = 10 if is_8bit_input_type else 9 + + if is_8bit_input_type: + assert not dense_B.is_contiguous(), ( + "dense input must be transposed for 8bit dtypes" + ) + + k = dense_B.size(0) + n = dense_B.size(1) + m = (compressed_A.numel() * 16) // (compression_factor * k) + if bias is not None: + assert m == bias.size(0) + + if out_dtype is not None: + assert is_8bit_input_type and out_dtype in { + torch.float16, + torch.bfloat16, + torch.int32, + torch.float8_e4m3fn, + }, ( + "out_dtype is not supported for {compressed_A.dtype} x {dense_B.dtype} -> {out_dtype} matmul!" + ) + output_shape = (n, m) if transpose_result else (m, n) + return dense_B.new_empty(output_shape, dtype=out_dtype) + + +@register_meta(aten.index_reduce.default) +def meta_index_reduce( + self: Tensor, + dim: int, + index: Tensor, + source: torch.Tensor, + reduce: str, + *, + include_self: bool = True, +) -> Tensor: + return torch.empty_like(self, memory_format=torch.contiguous_format) + + +@register_meta(aten.index_reduce_.default) +def meta_index_reduce_( + self: Tensor, + dim: int, + index: Tensor, + source: torch.Tensor, + reduce: str, + *, + include_self: bool = True, +) -> Tensor: + return self + + +# Implementations below are taken from https://github.com/albanD/subclass_zoo/blob/main/python_meta_tensor.py +@out_wrapper() +@register_meta(aten.index_select.default) +def meta_index_select(self, dim, index): + result_size = list(self.size()) + if self.dim() > 0: + result_size[dim] = index.numel() + return self.new_empty(result_size) + + +@register_meta(aten.segment_reduce.default) +def meta_segment_reduce( + data: Tensor, + reduce: str, + *, + lengths: Optional[Tensor] = None, + indices: Optional[Tensor] = None, + offsets: Optional[Tensor] = None, + axis: int = 0, + unsafe: bool = False, + initial=None, +) -> Tensor: + if indices is not None: + raise NotImplementedError( + "segment_reduce(): indices based reduction is not supported yet." + ) + + def segment_reduce_lengths_tensor(lengths_shape): + return torch.empty( + lengths_shape + data.shape[axis + 1 :], + dtype=data.dtype, + device="meta", + memory_format=torch.contiguous_format, + ) + + if lengths is not None: + return segment_reduce_lengths_tensor(lengths.shape) + # FIXME should probably check that lengths and offset aren't both set, but + # the ATen implementation neglects this too + if offsets is not None: + # lengths == torch.diff(offsets) + lengths_shape = offsets.shape[:-1] + (offsets.shape[-1] - 1,) + return segment_reduce_lengths_tensor(lengths_shape) + raise RuntimeError("segment_reduce(): Either lengths or offsets must be defined.") + + +@register_meta([aten.max.default, aten.max.unary_out]) +@out_wrapper() +def meta_max(self): + return self.new_empty(()) + + +@register_meta(aten.max.dim) +def meta_max_dim(self, dim, keepdim=False): + dim = utils.reduction_dims(self.shape, (dim,)) + output_shape = _compute_reduction_shape(self, dim, keepdim) + return ( + self.new_empty(output_shape), + self.new_empty(output_shape, dtype=torch.long), + ) + + +@register_meta([aten.min.default, aten.min.unary_out]) +@out_wrapper() +def meta_min(self): + return self.new_empty(()) + + +@register_meta(aten.min.dim) +def meta_min_dim(self, dim, keepdim=False): + dim = utils.reduction_dims(self.shape, (dim,)) + output_shape = _compute_reduction_shape(self, dim, keepdim) + return ( + self.new_empty(output_shape), + self.new_empty(output_shape, dtype=torch.long), + ) + + +@register_meta(aten.angle.default) +def meta_angle(self): + if self.is_complex(): + result_dtype = corresponding_real_dtype(self.dtype) + else: + _, result_dtype = elementwise_dtypes( + self, + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, + ) + return torch.empty_like(self, dtype=result_dtype) + + +@register_meta(aten.angle.out) +def meta_angle_out(self, out): + torch._resize_output_(out, self.size(), self.device) + return out.copy_(torch.angle(self)) + + +@register_meta(aten._assert_async.default) +def assert_async(val): + return + + +@register_meta(aten._assert_async.msg) +def assert_async_meta(val, assert_msg): + return + + +@register_meta(aten._print.default) +def print_meta(s): + return + + +@register_meta(aten._make_dep_token.default) +def make_dep_token( + *, + dtype=None, + layout=None, + device=None, + pin_memory=None, + memory_format=None, +): + return torch.empty(0, device="meta") + + +@register_meta(aten.sym_constrain_range.default) +def sym_constrain_range(size, min=None, max=None): + # Avoid importing sympy at a module level + from torch.fx.experimental.symbolic_shapes import constrain_range + + if isinstance(size, (SymFloat, SymBool)): + raise ValueError("Constraining SymFloat or Symbool is nyi") + constrain_range(size, min=min, max=max) + + +@register_meta(aten._functional_sym_constrain_range.default) +def functional_sym_constrain_range(size, min=None, max=None, dep_token=None): + aten.sym_constrain_range(size, min=min, max=max) + return dep_token + + +@register_meta(aten.sym_constrain_range_for_size.default) +def sym_constrain_range_for_size(size, min=None, max=None): + # Avoid importing sympy at a module level + from torch.fx.experimental.symbolic_shapes import _constrain_range_for_size + + if min is None and max is None: + torch._check_is_size(size) + return + + if isinstance(size, (SymFloat, SymBool)): + raise ValueError("Constraining SymFloat or Symbool is nyi") + if type(size) is int: + if min is not None: + torch._check(size >= min) + if max is not None: + torch._check(size <= max) + return + _constrain_range_for_size(size, min=min, max=max) + + +@register_meta(aten._functional_sym_constrain_range_for_size.default) +def functional_sym_constrain_range_for_size(size, min, max, dep_token): + aten.sym_constrain_range_for_size(size, min=min, max=max) + return dep_token + + +@register_meta(aten._functional_assert_async.msg) +def functional_assert_async_meta(val, assert_msg, dep_token): + return dep_token + + +# From aten/src/ATen/native/LinearAlgebraUtils.h +def squareCheckInputs(self: Tensor, f_name: str): + assert self.dim() >= 2, ( + f"{f_name}: The input tensor must have at least 2 dimensions." + ) + assert self.size(-1) == self.size(-2), ( + f"{f_name}: A must be batches of square matrices, but they are {self.size(-2)} by {self.size(-1)} matrices" + ) + + +# Validates input shapes and devices +# for linear solve methods (solve, cholesky_solve, lu_solve, triangular_solve) +# From aten/src/ATen/native/LinearAlgebraUtils.h +def linearSolveCheckInputs(self: Tensor, A: Tensor, name: str): + torch._check( + self.device == A.device, + lambda: ( + f"Expected b and A to be on the same device, but found b on " + f"{self.device} and A on {A.device} instead." + ), + ) + + torch._check( + self.dtype == A.dtype, + lambda: ( + f"Expected b and A to have the same dtype, but found b of type " + f"{self.dtype} and A of type {A.dtype} instead." + ), + ) + + torch._check( + A.size(-1) == A.size(-2), + lambda: ( + f"A must be batches of square matrices, " + f"but they are {A.size(-2)} by {A.size(-1)} matrices" + ), + ) + + torch._check( + A.size(-1) == self.size(-2), + lambda: ( + f"Incompatible matrix sizes for {name}: each A " + f"matrix is {A.size(-1)} by {A.size(-1)}" + f" but each b matrix is {self.size(-2)} by {self.size(-1)}" + ), + ) + + +# From aten/src/ATen/native/LinearAlgebraUtils.h +def checkFloatingOrComplex( + t: Tensor, + f_name: str, + allow_low_precision_dtypes: bool = True, +): + dtype = t.dtype + torch._check( + t.is_floating_point() or t.is_complex(), + lambda: f"{f_name}: Expected a floating point or complex tensor as input. Got {dtype}", + ) + if not allow_low_precision_dtypes: + torch._check( + dtype in (torch.float, torch.double, torch.cfloat, torch.cdouble), + lambda: f"{f_name}: Low precision dtypes not supported. Got {dtype}", + ) + + +# From aten/src/ATen/native/LinearAlgebraUtils.h +def checkIsMatrix(A: Tensor, f_name: str, arg_name: str = "A"): + torch._check( + A.dim() >= 2, + lambda: f"{f_name}: The input tensor {arg_name} must have at least 2 dimensions.", + ) + + +def checkInputsSolver(A: Tensor, B: Tensor, left: bool, f_name: str): + squareCheckInputs(A, f_name) + checkIsMatrix(B, f_name) + torch._check( + A.size(-2) == B.size(-2) if left else A.size(-1) == B.size(-1), + lambda: ( + f"{f_name}: Incompatible shapes of A and B for the equation " + f"{'AX = B' if left else 'XA = B'}" + f" ({A.size(-2)}x{A.size(-1)} and {B.size(-2)}x{B.size(-1)})" + ), + ) + + +def checkSameDevice( + fn_name: str, + result: Tensor, + input: Tensor, + result_name: str = "result", +): + torch._check( + result.device == input.device, + lambda: ( + f"{fn_name}: Expected {result_name} and input tensors to be on the same device, but got " + f"{result_name} on {result.device} and input on {input.device}" + ), + ) + + +def checkUplo(UPLO: str): + UPLO_uppercase = UPLO.upper() + torch._check( + len(UPLO) == 1 and (UPLO_uppercase == "U" or UPLO_uppercase == "L"), + lambda: f"Expected UPLO argument to be 'L' or 'U', but got {UPLO}", + ) + + +@register_meta([aten._linalg_eigh.default, aten._linalg_eigh.eigenvalues]) +@out_wrapper("eigenvalues", "eigenvectors") +def meta__linalg_eigh(A: Tensor, UPLO: str = "L", compute_v: bool = True): + squareCheckInputs(A, "linalg.eigh") + checkUplo(UPLO) + + shape = list(A.shape) + if compute_v: + vecs = A.new_empty(shape) + vecs.as_strided_(shape, make_contiguous_strides_for(shape, row_major=False)) + else: + vecs = A.new_empty([0]) + + shape.pop() + vals = A.new_empty(shape, dtype=toRealValueType(A.dtype)) + + return vals, vecs + + +@register_meta([aten._linalg_eigvals.default, aten.linalg_eigvals.out]) +@out_wrapper() +def meta__linalg_eigvals(input: Tensor) -> Tensor: + squareCheckInputs(input, "linalg.eigvals") + complex_dtype = ( + input.dtype + if utils.is_complex_dtype(input.dtype) + else utils.corresponding_complex_dtype(input.dtype) + ) + return input.new_empty(input.shape[:-1], dtype=complex_dtype) + + +@register_meta([aten.linalg_eig]) +@out_wrapper("eigenvalues", "eigenvectors") +def meta_linalg_eig(input: Tensor): + squareCheckInputs(input, "linalg.eig") + complex_dtype = ( + input.dtype + if utils.is_complex_dtype(input.dtype) + else utils.corresponding_complex_dtype(input.dtype) + ) + values = input.new_empty(input.shape[:-1], dtype=complex_dtype) + vectors = input.new_empty(input.shape, dtype=complex_dtype) + return values, vectors + + +def cloneBatchedColumnMajor(src: Tensor) -> Tensor: + return src.mT.clone(memory_format=torch.contiguous_format).transpose(-2, -1) + + +@register_meta(aten._cholesky_solve_helper) +@out_wrapper() +def _cholesky_solve_helper(self: Tensor, A: Tensor, upper: bool) -> Tensor: + return cloneBatchedColumnMajor(self) + + +@register_meta(aten.cholesky_solve) +@out_wrapper() +def cholesky_solve(self: Tensor, A: Tensor, upper: bool = False) -> Tensor: + torch._check( + self.ndim >= 2, + lambda: f"b should have at least 2 dimensions, but has {self.ndim} dimensions instead", + ) + torch._check( + A.ndim >= 2, + lambda: f"u should have at least 2 dimensions, but has {A.ndim} dimensions instead", + ) + self_broadcasted, A_broadcasted = _linalg_broadcast_batch_dims_name( + self, A, "cholesky_solve" + ) + return _cholesky_solve_helper(self_broadcasted, A_broadcasted, upper) + + +@register_meta(aten.cholesky) +@out_wrapper() +def cholesky(self: Tensor, upper: bool = False) -> Tensor: + if self.numel() == 0: + return torch.empty_like(self, memory_format=torch.legacy_contiguous_format) + squareCheckInputs(self, "cholesky") + return cloneBatchedColumnMajor(self) + + +@register_meta(aten.cholesky_inverse) +@out_wrapper() +def cholesky_inverse(self: Tensor, upper: bool = False) -> Tensor: + squareCheckInputs(self, "cholesky_inverse") + return cloneBatchedColumnMajor(self) + + +# From aten/src/ATen/native/BatchLinearAlgebra.cpp +@register_meta(aten.linalg_cholesky_ex.default) +def linalg_cholesky_ex(A: Tensor, upper: bool = False, check_errors: bool = False): + squareCheckInputs(A, "linalg.cholesky") + checkFloatingOrComplex(A, "linalg.cholesky") + + A_shape = A.shape + ndim = len(A_shape) + + # L + L_strides = make_contiguous_strides_for(A_shape, False) + L = A.new_empty(A_shape) + L.as_strided_(A_shape, L_strides) + + # infos + infos = A.new_empty(A_shape[0 : ndim - 2], dtype=torch.int32) + return L, infos + + +@register_meta( + [aten.linalg_householder_product.default, aten.linalg_householder_product.out] +) +@out_wrapper() +def linalg_householder_product(input: Tensor, tau: Tensor) -> Tensor: + torch._check( + input.ndim >= 2, + lambda: "torch.linalg.householder_product: input must have at least 2 dimensions.", + ) + torch._check( + input.size(-2) >= input.size(-1), + lambda: "torch.linalg.householder_product: input.shape[-2] must be greater than or equal to input.shape[-1]", + ) + torch._check( + input.size(-1) >= tau.size(-1), + lambda: "torch.linalg.householder_product: input.shape[-1] must be greater than or equal to tau.shape[-1]", + ) + + torch._check( + input.ndim - tau.ndim == 1, + lambda: ( + f"torch.linalg.householder_product: Expected tau to have one dimension less than input, " + f"but got tau.ndim equal to {tau.ndim} and input.ndim is equal to {input.ndim}" + ), + ) + if input.ndim > 2: + expected_batch_tau_shape = input.shape[:-2] + actual_batch_tau_shape = tau.shape[:-1] + torch._check( + actual_batch_tau_shape == expected_batch_tau_shape, + lambda: ( + f"torch.linalg.householder_product: Expected batch dimensions of tau to be " + f"equal to input.shape[:-2], but got {actual_batch_tau_shape}" + ), + ) + + torch._check( + tau.dtype == input.dtype, + lambda: ( + f"torch.linalg.householder_product: tau dtype {tau.dtype}" + f" does not match input dtype {input.dtype}" + ), + ) + checkSameDevice("torch.linalg.householder_product", tau, input, "tau") + + return torch.empty_strided( + size=input.shape, + stride=make_contiguous_strides_for(input.shape, row_major=False), + dtype=input.dtype, + device=input.device, + ) + + +# From aten/src/ATen/native/BatchLinearAlgebra.cpp +@register_meta(aten.linalg_inv_ex.default) +def linalg_inv_ex_meta(A: Tensor, check_errors: bool = False): + squareCheckInputs(A, "linalg.inv_ex") + checkFloatingOrComplex(A, "linalg.inv_ex", allow_low_precision_dtypes=False) + + L = A.new_empty(A.shape) + L.as_strided_(A.shape, make_contiguous_strides_for(A.shape, row_major=False)) + + infos = A.new_empty(A.shape[:-2], dtype=torch.int32) + return L, infos + + +@register_meta([aten.linalg_ldl_factor_ex.default, aten.linalg_ldl_factor_ex.out]) +@out_wrapper("LD", "pivots", "info") +def linalg_ldl_factor_ex_meta( + self: Tensor, + *, + hermitian: bool = False, + check_errors: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + squareCheckInputs(self, "torch.linalg.ldl_factor_ex") + checkFloatingOrComplex(self, "torch.linalg.ldl_factor_ex") + LD = torch.empty_strided( + size=self.shape, + stride=make_contiguous_strides_for(self.shape, row_major=False), + dtype=self.dtype, + device=self.device, + ) + pivots = self.new_empty(self.shape[:-1], dtype=torch.int) + info = self.new_empty(self.shape[:-2], dtype=torch.int) + return LD, pivots, info + + +@register_meta([aten.linalg_ldl_solve.default, aten.linalg_ldl_solve.out]) +@out_wrapper() +def linalg_ldl_solve_meta( + LD: Tensor, + pivots: Tensor, + B: Tensor, + *, + hermitian: bool = False, +) -> Tensor: + squareCheckInputs(LD, "torch.linalg.ldl_solve") + checkFloatingOrComplex(LD, "torch.linalg.ldl_solve") + linearSolveCheckInputs(B, LD, "torch.linalg.ldl_solve") + torch._check( + B.ndim >= 2, + lambda: ( + f"torch.linalg.ldl_solve: Expected B to have at least 2 dimensions, " + f"but it has {B.ndim} dimensions instead" + ), + ) + expected_pivots_shape = LD.shape[:-1] + torch._check( + expected_pivots_shape == pivots.shape, + lambda: ( + f"torch.linalg.ldl_solve: Expected LD.shape[:-1] and pivots.shape to be the same, " + f"but got pivots with shape {pivots.shape} instead" + ), + ) + torch._check( + utils.is_integer_dtype(pivots.dtype), + lambda: f"torch.linalg.ldl_solve: Expected pivots to be integers. Got {pivots.dtype}", + ) + torch._check( + LD.dtype == B.dtype, + lambda: f"torch.linalg.ldl_solve: LD dtype {LD.dtype} does not match b dtype {B.dtype}", + ) + B_broadcast_size, _ = _linalg_broadcast_batch_dims(B, LD) + return torch.empty_strided( + size=B_broadcast_size, + stride=make_contiguous_strides_for(B_broadcast_size, row_major=False), + dtype=B.dtype, + device=B.device, + ) + + +@register_meta([aten.linalg_lu.default, aten.linalg_lu.out]) +@out_wrapper("P", "L", "U") +def linalg_lu_meta(A: Tensor, *, pivot: bool = True) -> tuple[Tensor, Tensor, Tensor]: + torch._check( + A.ndim >= 2, + lambda: f"linalg.lu: Expected tensor with 2 or more dimensions. Got size: {A.shape} instead", + ) + + sizes = list(A.shape) + m = sizes[-2] + n = sizes[-1] + k = min(m, n) + + sizes[-1] = m + if pivot: + P = A.new_empty(sizes) + else: + P = A.new_empty([0]) + + sizes[-1] = k + L = A.new_empty(sizes) + + sizes[-2] = k + sizes[-1] = n + U = A.new_empty(sizes) + return P, L, U + + +@register_meta([aten.linalg_lu_factor_ex.default, aten.linalg_lu_factor_ex.out]) +@out_wrapper("LU", "pivots", "info") +def linalg_lu_factor_ex_meta( + A: Tensor, + *, + pivot: bool = True, + check_errors: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + torch._check( + A.ndim >= 2, + lambda: f"torch.lu_factor: Expected tensor with 2 or more dimensions. Got size: {A.shape} instead", + ) + + sizes = list(A.shape) + m = sizes[-2] + n = sizes[-1] + + LU = torch.empty_strided( + size=sizes, + stride=make_contiguous_strides_for(sizes, row_major=False), + dtype=A.dtype, + device=A.device, + ) + + # Sets sizes to the size of pivots + sizes.pop() + sizes[-1] = min(m, n) + pivots = A.new_empty(sizes, dtype=torch.int) + + # Sets sizes to the size of info + sizes.pop() + info = A.new_empty(sizes, dtype=torch.int) + + return LU, pivots, info + + +@register_meta([aten.linalg_lu_solve.default, aten.linalg_lu_solve.out]) +@out_wrapper() +def linalg_lu_solve_meta( + LU: Tensor, + pivots: Tensor, + B: Tensor, + *, + left: bool = True, + adjoint: bool = False, +) -> Tensor: + # dtype + checkFloatingOrComplex(LU, "torch.linalg.lu_solve") + torch._check( + LU.dtype == B.dtype, + lambda: ( + f"linalg.lu_solve: Expected LU and B to have the same dtype, " + f"but found LU of type {LU.dtype} and B of type {B.dtype} instead" + ), + ) + torch._check( + pivots.dtype == torch.int, + lambda: "linalg.lu_solve: pivots should be a Tensor of scalar type torch.int32", + ) + + # matrix shapes + squareCheckInputs(LU, "torch.linalg.lu_solve") + checkInputsSolver(LU, B, left, "linalg.lu_solve") + torch._check( + LU.size(-1) == pivots.size(-1), + lambda: "linalg.lu_solve: Number of pivots per batch should be same as the dimension of the matrix", + ) + + # batches + torch._check( + LU.shape[:-1] == pivots.shape, + lambda: ( + f"linalg.lu_solve: Expected LU.shape[:-1] and pivots.shape to be the same, " + f"but got pivots with shape {pivots.shape} instead" + ), + ) + + B_broadcast_size, _ = _linalg_broadcast_batch_dims(B, LU) + + result = torch.empty_strided( + size=B_broadcast_size, + stride=make_contiguous_strides_for(B_broadcast_size, row_major=not left), + dtype=B.dtype, + device=B.device, + ) + + if result.numel() != 0 and not left: + if result.is_complex(): + result = result.conj() + + return result + + +@register_meta(aten.lu_unpack) +@out_wrapper("P", "L", "U") +def lu_unpack_meta( + LU: Tensor, + pivots: Tensor, + unpack_data: bool = True, + unpack_pivots: bool = True, +) -> tuple[Tensor, Tensor, Tensor]: + torch._check( + LU.ndim >= 2, + lambda: f"torch.lu_unpack: Expected tensor with 2 or more dimensions. Got size: {LU.shape} instead", + ) + if unpack_pivots: + torch._check( + pivots.dtype == torch.int32, + lambda: ( + "torch.lu_unpack: LU_pivots is expected to be a contiguous tensor of torch.int32 dtype.\n" + "Note: this function is intended to be used with the output produced by torch.linalg.lu_factor" + ), + ) + sizes = list(LU.shape) + m = sizes[-2] + n = sizes[-1] + k = min(m, n) + sizes[-1] = m + if unpack_pivots: + P = LU.new_empty(sizes) + else: + P = LU.new_empty([0]) + if unpack_data: + sizes[-1] = k + L = LU.new_empty(sizes) + sizes[-2] = k + sizes[-1] = n + U = LU.new_empty(sizes) + else: + L = LU.new_empty([0]) + U = LU.new_empty([0]) + return P, L, U + + +# parse the "mode" param in linalg_qr: return a tuple of bools (compute_q, reduced) +def _parse_qr_mode(mode: str) -> tuple[bool, bool]: + if mode == "reduced": + compute_q = True + reduced = True + elif mode == "complete": + compute_q = True + reduced = False + elif mode == "r": + compute_q = False + reduced = True # this is actually irrelevant in this mode + else: + torch._check( + False, + lambda: ( + f"qr received unrecognized mode '{mode}' " + f"but expected one of 'reduced' (default), 'r', or 'complete'" + ), + ) + return compute_q, reduced # type: ignore[possibly-undefined] + + +@register_meta([aten.linalg_qr.default, aten.linalg_qr.out]) +@out_wrapper("Q", "R") +def linalg_qr_meta(A: Tensor, mode: str = "reduced") -> tuple[Tensor, Tensor]: + checkIsMatrix(A, "linalg.qr") + checkFloatingOrComplex(A, "linalg.qr") + + compute_q, reduced_mode = _parse_qr_mode(mode) + + m = A.shape[-2] + n = A.shape[-1] + k = min(m, n) + + if compute_q: + Q_shape = list(A.shape) + Q_shape[-1] = k if reduced_mode else m + Q = A.new_empty(Q_shape) + Q.as_strided_(Q_shape, make_contiguous_strides_for(Q_shape, row_major=False)) + else: + Q = A.new_empty([0]) + + # For readability + R_shape = list(A.shape) + R_shape[-2] = k if reduced_mode or not compute_q else m + R = A.new_empty(R_shape) + R.as_strided_(R_shape, make_contiguous_strides_for(R_shape, row_major=False)) + return Q, R + + +@register_meta([aten._linalg_slogdet.default, aten._linalg_slogdet.sign]) +@out_wrapper("sign", "logabsdet", "LU", "pivots") +def _linalg_slogdet(A: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: + squareCheckInputs(A, "linalg.slogdet") + checkFloatingOrComplex(A, "linalg.slogdet", False) + shape = A.shape + sign = A.new_empty(shape[:-2]) + logabsdet = A.new_empty(shape[:-2], dtype=toRealValueType(A.dtype)) + LU = torch.empty_strided( + size=shape, + stride=make_contiguous_strides_for(shape, False), + dtype=A.dtype, + device=A.device, + ) + pivots = A.new_empty(shape[:-1], dtype=torch.int32) + return sign, logabsdet, LU, pivots + + +# From aten/src/ATen/native/BatchLinearAlgebra.cpp +# NOTE: matching defaults in aten/src/ATen/native/native_functions.yaml +@register_meta(aten._linalg_svd.default) +def _linalg_svd_meta( + A: Tensor, + full_matrices: bool = False, + compute_uv: bool = True, + driver: Optional[str] = None, +): + checkIsMatrix(A, "linalg.svd") + checkFloatingOrComplex(A, "linalg.svd") + + batch_dims = list(A.shape[:-2]) + m = A.shape[-2] + n = A.shape[-1] + k = min(m, n) + + if compute_uv: + U_shape = batch_dims + [m, m if full_matrices else k] + U = A.new_empty(U_shape) + U.as_strided_(U_shape, make_contiguous_strides_for(U_shape, row_major=False)) + + V_shape = batch_dims + [n if full_matrices else k, n] + V = A.new_empty(V_shape) + # NB: This checks for CUDA since there is no way to check for cuSolver. + # Also, this might not work correctly on CPU when fake_device is not + # available as device_hint just defaults to CUDA in that case. See + # _linalg_svd meta in core. + is_cuda = device_hint(A) == "cuda" + V.as_strided_(V_shape, make_contiguous_strides_for(V_shape, row_major=is_cuda)) + else: + # doesn't matter + U = A.new_empty([0]) + V = A.new_empty([0]) + + # S is always real, even when A is complex. + S = A.new_empty(batch_dims + [k], dtype=toRealValueType(A.dtype)) + return U, S, V + + +def _linalg_broadcast_batch_dims( + arg1: Tensor, + arg2: Tensor, +) -> tuple[list[int], list[int]]: + # broadcast the batch dimensions of arg1 and arg2. + arg1_batch_sizes = arg1.shape[:-2] + arg2_batch_sizes = arg2.shape[:-2] + expand_batch_portion = _broadcast_shapes(arg1_batch_sizes, arg2_batch_sizes) + + arg1_expand_size = list(expand_batch_portion) + arg1_expand_size += [arg1.size(-2), arg1.size(-1)] + + arg2_expand_size = list(expand_batch_portion) + arg2_expand_size += [arg2.size(-2), arg2.size(-1)] + return arg1_expand_size, arg2_expand_size + + +def _linalg_broadcast_batch_dims_name( + arg1: Tensor, + arg2: Tensor, + name: Optional[str], +) -> tuple[Tensor, Tensor]: + # If there's no name we assume we don't want to check the errors + if name: + linearSolveCheckInputs(arg1, arg2, name) + + arg1_expand_size, arg2_expand_size = _linalg_broadcast_batch_dims(arg1, arg2) + + arg1_broadcasted = ( + arg1 if arg1_expand_size == arg1.shape else arg1.expand(arg1_expand_size) + ) + arg2_broadcasted = ( + arg2 if arg2_expand_size == arg2.shape else arg2.expand(arg2_expand_size) + ) + return arg1_broadcasted, arg2_broadcasted + + +def linalg_solve_is_vector_rhs(input: Tensor, other: Tensor) -> bool: + expected_batched_rhs_shape = input.shape[:-1] + vector_case = other.ndim == 1 or ( + input.ndim - 1 == other.ndim and other.shape == expected_batched_rhs_shape + ) + return vector_case + + +@register_meta(aten._linalg_solve_ex) +def _linalg_solve_ex( + A: Tensor, + B: Tensor, + *, + left: bool = True, + check_errors: bool = False, + result: Optional[Tensor] = None, + LU: Optional[Tensor] = None, + pivots: Optional[Tensor] = None, + info: Optional[Tensor] = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + checkFloatingOrComplex(A, "linalg.solve") + torch._check( + A.dtype == B.dtype, + lambda: ( + f"linalg.solve: Expected A and B to have the same dtype, but found A of type " + f"{A.dtype} and B of type {B.dtype} instead" + ), + ) + vector_case = linalg_solve_is_vector_rhs(A, B) + B_ = B.unsqueeze(-1) if vector_case else B + checkInputsSolver(A, B_, left, "linalg.solve") + B_broad_shape, _ = _linalg_broadcast_batch_dims(B_, A) + torch._check( + left or not vector_case, + lambda: ( + "linalg.solve: Vector broadcasting of the left hand side is not supported for left=False. " + "In this case linalg.solve is equivalent to B / A.squeeze(-1)" + ), + ) + result_shape = B_broad_shape[:-1] if vector_case else B_broad_shape + result_ = torch.empty_strided( + size=result_shape, + stride=make_contiguous_strides_for(result_shape, not left), + dtype=B.dtype, + device=B.device, + ) + shape = A.shape + LU_ = torch.empty_strided( + size=shape, + stride=make_contiguous_strides_for(shape, False), + dtype=A.dtype, + device=A.device, + ) + pivots_ = A.new_empty(shape[:-1], dtype=torch.int32) + info_ = A.new_empty(shape[:-2], dtype=torch.int32) + out = (result, LU, pivots, info) + res = (result_, LU_, pivots_, info_) + if all(x is not None for x in out): + for r, o in zip(res, out): + # resize and copy operations are done in-place + _maybe_resize_out(o, r.shape) # type: ignore[arg-type] + # strides are not copied in out_wrapper + o.as_strided_(r.shape, r.stride()) # type: ignore[union-attr] + _safe_copy_out(copy_from=r, copy_to=o, exact_dtype=False) # type: ignore[arg-type] + return res + + +@register_meta([aten.linalg_solve_triangular.default, aten.linalg_solve_triangular.out]) +def linalg_solve_triangular_meta( + A: Tensor, + B: Tensor, + *, + upper: bool, + left: bool = True, + unitriangular: bool = False, + out: Optional[Tensor] = None, +) -> Tensor: + if out is None: + out = A.new_empty([0]) + assert isinstance(out, TensorLike) + checkInputsSolver(A, B, left, "linalg.solve_triangular") + B_, A_ = _linalg_broadcast_batch_dims_name(B, A, None) + avoid_copy_A = A_.transpose(-2, -1).is_contiguous() and A_.is_conj() + if avoid_copy_A: + out = _maybe_resize_out(out, B_.shape) + else: + # reimplementation of resize_output with result F-contig + if _resize_output_check(out, B_.shape): + out.resize_(B_.transpose(-2, -1).shape) + out.transpose_(-2, -1) + return out # type: ignore[return-value] + + +@register_meta(aten.triangular_solve) +@out_wrapper("X", "M", exact_dtype=True) +def triangular_solve_meta( + self: Tensor, + A: Tensor, + upper: bool = True, + transpose: bool = False, + unitriangular: bool = False, +) -> tuple[Tensor, Tensor]: + torch._check( + self.ndim >= 2, + lambda: ( + f"torch.triangular_solve: Expected b to have at least 2 dimensions, " + f"but it has {self.ndim} dimensions instead" + ), + ) + torch._check( + A.ndim >= 2, + lambda: ( + f"torch.triangular_solve: Expected A to have at least 2 dimensions, " + f"but it has {A.ndim} dimensions instead" + ), + ) + + linearSolveCheckInputs(self, A, "triangular_solve") + + if A.layout == torch.strided: + self_broadcast_size, A_broadcast_size = _linalg_broadcast_batch_dims(self, A) + solution = torch.empty_strided( + size=self_broadcast_size, + stride=make_contiguous_strides_for(self_broadcast_size, row_major=False), + dtype=self.dtype, + device=self.device, + ) + cloned_coefficient = torch.empty_strided( + size=A_broadcast_size, + stride=make_contiguous_strides_for(A_broadcast_size, row_major=False), + dtype=A.dtype, + device=A.device, + ) + elif A.layout == torch.sparse_csr or A.layout == torch.sparse_bsr: + solution = torch.empty_like(self) + cloned_coefficient = self.new_empty([0]) + else: + torch._check(False, lambda: "triangular_solve: Got an unexpected layout.") + return solution, cloned_coefficient # type: ignore[possibly-undefined] + + +# From aten/src/ATen/native/LinearAlgebra.cpp +@register_meta(aten._linalg_det.default) +def _linalg_det_meta(A): + squareCheckInputs(A, "linalg.det") + checkFloatingOrComplex(A, "linalg.det") + + det = A.new_empty(A.shape[:-2]) + + LU = A.new_empty(A.shape) + LU.as_strided_(A.shape, make_contiguous_strides_for(A.shape, row_major=False)) + + pivots = A.new_empty(A.shape[:-1], dtype=torch.int32) + return det, LU, pivots + + +@register_meta(aten.ormqr) +@out_wrapper() +def ormqr( + input: Tensor, + tau: Tensor, + other: Tensor, + left: bool = True, + transpose: bool = False, +) -> Tensor: + torch._check( + input.ndim >= 2, lambda: "torch.ormqr: input must have at least 2 dimensions." + ) + torch._check( + other.ndim >= 2, lambda: "torch.ormqr: other must have at least 2 dimensions." + ) + + left_size_condition = -2 if left else -1 + torch._check( + other.shape[left_size_condition] >= tau.shape[-1], + lambda: f"torch.ormqr: other.shape[{left_size_condition}] must be greater than or equal to tau.shape[-1]", + ) + torch._check( + other.shape[left_size_condition] == input.shape[-2], + lambda: f"torch.ormqr: other.shape[{left_size_condition}] must be equal to input.shape[-2]", + ) + + torch._check( + tau.shape[-1] <= input.shape[-1], + lambda: "torch.ormqr: tau.shape[-1] must be less than or equal to input.shape[-1]", + ) + + torch._check( + input.ndim - tau.ndim == 1, + lambda: ( + f"torch.ormqr: Expected tau to have one dimension less than input, " + f"but got tau.ndim equal to {tau.ndim} and input.ndim is equal to {input.ndim}" + ), + ) + torch._check( + input.ndim == other.ndim, + lambda: ( + f"torch.ormqr: Expected other to have the same number of dimensions as input, " + f"but got other.ndim equal to {other.ndim} and input.ndim is equal to {input.ndim}" + ), + ) + + if input.ndim > 2: + expected_batch_shape = input.shape[:-2] + actual_batch_tau_shape = tau.shape[:-1] + torch._check( + actual_batch_tau_shape == expected_batch_shape, + lambda: ( + f"torch.ormqr: Expected batch dimensions of tau to be " + f"equal to input.shape[:-2], but got {actual_batch_tau_shape}" + ), + ) + + actual_batch_other_shape = other.shape[:-2] + torch._check( + actual_batch_other_shape == expected_batch_shape, + lambda: ( + f"torch.ormqr: Expected batch dimensions of other to be " + f"equal to input.shape[:-2], but got {actual_batch_other_shape}" + ), + ) + + torch._check( + tau.dtype == input.dtype, + lambda: ( + f"torch.ormqr: Expected input and tau to have the same dtype, " + f"but input has dtype {input.dtype} and tau has dtype {tau.dtype}" + ), + ) + torch._check( + other.dtype == input.dtype, + lambda: ( + f"torch.ormqr: Expected input and other to have the same dtype, " + f"but input has dtype {input.dtype} and other has dtype {other.dtype}" + ), + ) + + checkSameDevice("torch.ormqr", tau, input, "tau") + checkSameDevice("torch.ormqr", other, input, "other") + + return torch.empty_strided( + size=other.shape, + stride=make_contiguous_strides_for(other.shape, row_major=False), + dtype=other.dtype, + device=other.device, + ) + + +def _padding_check_valid_input(input, padding, *, dim): + torch._check( + len(padding) == 2 * dim, + lambda: f"padding size is expected to be {2 * dim}, but got: {len(padding)}", + ) + + input_dim = input.ndim + + is_batch_mode = input_dim == (dim + 2) + + valid_batch_mode = is_batch_mode + valid_non_batch_mode = not is_batch_mode + + if is_batch_mode: + # allow batch size of 0-dim. + for d in range(1, input_dim): + valid_batch_mode = valid_batch_mode and input.size(d) != 0 + else: + for d in range(0, input_dim): + valid_non_batch_mode = valid_non_batch_mode and input.size(d) != 0 + + # allow empty batch size but not other dimensions. + torch._check( + valid_batch_mode or valid_non_batch_mode, + lambda: ( + f"Expected {dim + 1}D or {dim + 2}D (batch mode) tensor with possibly 0 batch size " + f"and other non-zero dimensions for input, but got: {input.shape}" + ), + ) + + +def _pad1d_common(input, padding, *, is_reflection): + dim_plane = 0 + dim_w = 1 + nbatch = 1 + + if input.ndim == 3: + nbatch = input.size(0) + dim_w += 1 + dim_plane += 1 + + _padding_check_valid_input(input, padding, dim=1) + + pad_l, pad_r = padding + + nplane = input.size(dim_plane) + input_w = input.size(dim_w) + output_w = input_w + pad_l + pad_r + + if is_reflection: + torch._check( + pad_l < input_w and pad_r < input_w, + lambda: ( + f"Argument #4: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_l}, {pad_r}) at dimension {dim_w} of input {input.shape}" + ), + ) + + torch._check( + output_w >= 1, + lambda: f"input (W: {input_w}) is too small. Calculated output W: {output_w}", + ) + + if input.ndim == 2: + return input.new_empty((nplane, output_w)) + else: + return input.new_empty((nbatch, nplane, output_w)) + + +@register_meta(aten.reflection_pad1d) +@out_wrapper() +def meta_reflection_pad1d(input, padding): + return _pad1d_common(input, padding, is_reflection=True) + + +@register_meta(aten.replication_pad1d) +@out_wrapper() +def meta_replication_pad1d(input, padding): + torch._check( + input.dtype != torch.bool, + lambda: f""""replication_pad1d" not implemented for '{input.dtype.__str__()}'""", + ) + return _pad1d_common(input, padding, is_reflection=False) + + +def _pad1d_backward_common(grad_output, input, padding, *, is_reflection): + dim_w = 1 + if not is_reflection: + torch._check(len(padding) == 2, lambda: "padding size is expected to be 2") + + if input.ndim == 3: + dim_w += 1 + + pad_l, pad_r = padding + + input_w = input.size(dim_w) + output_w = input_w + pad_l + pad_r + + if is_reflection: + torch._check( + pad_l < input_w and pad_r < input_w, + lambda: ( + f"Argument #4: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_l}, {pad_r}) at dimension {dim_w} of input {input.shape}" + ), + ) + + torch._check( + output_w == grad_output.size(dim_w), + lambda: f"grad_output width unexpected. Expected: {output_w}, Got: {grad_output.size(dim_w)}", + ) + + return input.new_empty(input.shape) + + +@register_meta(aten.reflection_pad1d_backward) +@out_wrapper("grad_input") +def meta_reflection_pad1d_backward(grad_output, input, padding): + return _pad1d_backward_common(grad_output, input, padding, is_reflection=True) + + +@register_meta(aten.replication_pad1d_backward) +@out_wrapper("grad_input") +def meta_replication_pad1d_backward(grad_output, input, padding): + return _pad1d_backward_common(grad_output, input, padding, is_reflection=False) + + +def _pad2d_common(input, padding, *, is_reflection): + dim_w = 2 + dim_h = 1 + dim_slices = 0 + nbatch = 1 + + _padding_check_valid_input(input, padding, dim=2) + + ndim = input.ndim + if ndim == 4: + nbatch = input.size(0) + dim_w += 1 + dim_h += 1 + dim_slices += 1 + + pad_l, pad_r, pad_t, pad_b = padding + + nplane = input.size(dim_slices) + input_h = input.size(dim_h) + input_w = input.size(dim_w) + output_h = input_h + pad_t + pad_b + output_w = input_w + pad_l + pad_r + + if is_reflection: + torch._check( + pad_l < input_w and pad_r < input_w, + lambda: ( + f"Argument #4: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_l}, {pad_r}) at dimension {dim_w} of input {input.shape}" + ), + ) + torch._check( + pad_t < input_h and pad_b < input_h, + lambda: ( + f"Argument #6: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_t}, {pad_b}) at dimension {dim_h} of input {input.shape}" + ), + ) + + torch._check( + output_w >= 1 or output_h >= 1, + lambda: ( + f"input (H: {input_h} W: {input_w}) is too small. " + f"Calculated output H: {output_h} W: {output_w}" + ), + ) + + if input.ndim == 3: + return input.new_empty((nplane, output_h, output_w)) + else: + return input.new_empty((nbatch, nplane, output_h, output_w)) + + +@register_meta(aten.reflection_pad2d) +@out_wrapper() +def meta_reflection_pad2d(input, padding): + return _pad2d_common(input, padding, is_reflection=True) + + +@register_meta(aten.replication_pad2d) +@out_wrapper() +def meta_replication_pad2d(input, padding): + torch._check( + input.dtype != torch.bool, + lambda: f""""replication_pad2d" not implemented for '{input.dtype.__str__()}'""", + ) + return _pad2d_common(input, padding, is_reflection=False) + + +@register_meta( + [ + aten.reflection_pad2d_backward.default, + aten.reflection_pad2d_backward.grad_input, + aten.replication_pad2d_backward.default, + aten.replication_pad2d_backward.grad_input, + ] +) +@out_wrapper("grad_input") +def meta_pad2d_backward(grad_output, self, padding): + dim_w = 2 + dim_h = 1 + dim_plane = 0 + + self_shape = self.shape + if self.dim() == 4: + dim_w += 1 + dim_h += 1 + dim_plane += 1 + + pad_l, pad_r, pad_t, pad_b = padding + + input_h = self_shape[dim_h] + input_w = self_shape[dim_w] + output_h = input_h + pad_t + pad_b + output_w = input_w + pad_l + pad_r + + torch._check( + output_w == grad_output.size(dim_w), + lambda: f"grad_output width unexpected. Expected: {output_w}, Got: {grad_output.size(dim_w)}", + ) + torch._check( + output_h == grad_output.size(dim_h), + lambda: f"grad_output height unexpected. Expected: {output_h}, Got: {grad_output.size(dim_h)}", + ) + return self.new_empty(self.shape) + + +def _pad3d_common(input, padding, *, is_reflection): + dim_w = 3 + dim_h = 2 + dim_d = 1 + dim_plane = 0 + + _padding_check_valid_input(input, padding, dim=3) + + batch_mode = input.ndim == 5 + if batch_mode: + nbatch = input.size(0) + dim_w += 1 + dim_h += 1 + dim_d += 1 + dim_plane += 1 + + pad_l, pad_r, pad_t, pad_b, pad_f, pad_bk = padding + + nplane = input.size(dim_plane) + input_d = input.size(dim_d) + input_h = input.size(dim_h) + input_w = input.size(dim_w) + output_d = input_d + pad_f + pad_bk + output_h = input_h + pad_t + pad_b + output_w = input_w + pad_l + pad_r + + if is_reflection: + torch._check( + pad_l < input_w and pad_r < input_w, + lambda: ( + f"Argument #4: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_l}, {pad_r}) at dimension {dim_w} of input {input.shape}" + ), + ) + torch._check( + pad_t < input_h and pad_b < input_h, + lambda: ( + f"Argument #6: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_t}, {pad_b}) at dimension {dim_h} of input {input.shape}" + ), + ) + torch._check( + pad_f < input_d and pad_bk < input_d, + lambda: ( + f"Argument #8: Padding size should be less than the corresponding input dimension, " + f"but got: padding ({pad_f}, {pad_bk}) at dimension {dim_d} of input {input.shape}" + ), + ) + + torch._check( + output_w >= 1 or output_h >= 1 or output_d >= 1, + lambda: ( + f"input (D: {input_d} H: {input_h} W: {input_w}) is too small. " + f"Calculated output D: {output_d} H: {output_h} W: {output_w}" + ), + ) + + if batch_mode: + return input.new_empty((nbatch, nplane, output_d, output_h, output_w)) # type: ignore[possibly-undefined] + else: + return input.new_empty((nplane, output_d, output_h, output_w)) + + +@register_meta(aten.reflection_pad3d) +@out_wrapper() +def meta_reflection_pad3d(input, padding): + return _pad3d_common(input, padding, is_reflection=True) + + +@register_meta(aten.replication_pad3d) +@out_wrapper() +def meta_replication_pad3d(input, padding): + torch._check( + input.dtype != torch.bool, + lambda: f""""replication_pad3d" not implemented for '{input.dtype.__str__()}'""", + ) + return _pad3d_common(input, padding, is_reflection=False) + + +@register_meta( + [ + aten.reflection_pad3d_backward.default, + aten.reflection_pad3d_backward.grad_input, + aten.replication_pad3d_backward.default, + aten.replication_pad3d_backward.grad_input, + ] +) +@out_wrapper("grad_input") +def meta_pad3d_backward(grad_output, input, padding): + torch._check(len(padding) == 6, lambda: "padding size is expected to be 6") + assert input.ndim > 3 + assert grad_output.ndim == input.ndim + + dim_w = 3 + dim_h = 2 + dim_d = 1 + + if input.ndim == 5: + dim_w += 1 + dim_h += 1 + dim_d += 1 + + pad_l, pad_r, pad_t, pad_b, pad_f, pad_bk = padding + + input_d = input.size(dim_d) + input_h = input.size(dim_h) + input_w = input.size(dim_w) + output_d = input_d + pad_f + pad_bk + output_h = input_h + pad_t + pad_b + output_w = input_w + pad_l + pad_r + + torch._check( + output_w == grad_output.size(dim_w), + lambda: f"grad_output width unexpected. Expected: {output_w}, Got: {grad_output.size(dim_w)}", + ) + torch._check( + output_h == grad_output.size(dim_h), + lambda: f"grad_output height unexpected. Expected: {output_h}, Got: {grad_output.size(dim_h)}", + ) + torch._check( + output_d == grad_output.size(dim_d), + lambda: f"grad_output depth unexpected. Expected: {output_d}, Got: {grad_output.size(dim_d)}", + ) + + return input.new_empty(input.shape) + + +@register_meta(aten._pdist_forward) +@out_wrapper() +def meta__pdist_forward(self: Tensor, p: float = 2) -> Tensor: + torch._check( + self.is_contiguous(), lambda: "_pdist_forward requires contiguous input" + ) + n = self.size(0) + if n <= 1: + return self.new_empty([0]).to(memory_format=torch.legacy_contiguous_format) # type: ignore[call-overload] + else: + return self.new_empty((n * (n - 1) // 2,)).to( + memory_format=torch.legacy_contiguous_format + ) # type: ignore[call-overload] + + +@register_meta(aten._pdist_backward) +@out_wrapper() +def meta__pdist_backward(grad: Tensor, self: Tensor, p: float, pdist: Tensor) -> Tensor: + torch._check( + self.is_contiguous(), lambda: "_pdist_backward requires self to be contiguous" + ) + torch._check( + pdist.is_contiguous(), lambda: "_pdist_backward requires pdist to be contiguous" + ) + return torch.empty_like(self, memory_format=torch.legacy_contiguous_format) + + +@register_meta([aten.baddbmm.default, aten.baddbmm.out]) +@out_wrapper(exact_dtype=True) +def meta_baddbmm(self, batch1, batch2, *, beta=1, alpha=1): + from torch.fx.experimental.symbolic_shapes import guard_or_true, sym_eq + + dim1 = batch1.size(0) + dim2 = batch1.size(1) + dim3 = batch2.size(2) + if guard_or_true(torch.sym_not(sym_eq(self.shape, (dim1, dim2, dim3)))): + self = self.expand((dim1, dim2, dim3)) + torch._check(batch1.dim() == 3, lambda: "batch1 must be a 3D tensor") + torch._check(batch2.dim() == 3, lambda: "batch2 must be a 3D tensor") + if not exp_config.skip_dtype_check_in_meta_registrations: + torch._check( + self.dtype == batch1.dtype == batch2.dtype, + lambda: f"Input dtypes must be the same, got: input: {self.dtype}, batch1: {batch1.dtype}, batch2: {batch2.dtype}", + ) + batch1_sizes = batch1.shape + batch2_sizes = batch2.shape + bs = batch1_sizes[0] + contraction_size = batch1_sizes[2] + torch._check( + batch2_sizes[0] == bs and batch2_sizes[1] == contraction_size, + lambda: ( + f"Expected size for first two dimensions of batch2 tensor to be: " + f"[{bs}, {contraction_size}] but got: [{batch2_sizes[0]}, {batch2_sizes[1]}]." + ), + ) + return self.new_empty(self.size()) + + +@register_meta([aten.bernoulli.default, aten.bernoulli.out]) +@out_wrapper() +def meta_bernoulli(self, *, generator=None): + # https://github.com/pytorch/pytorch/issues/88612 + return torch.empty_like(self, memory_format=torch.contiguous_format) + + +@register_meta(aten.bernoulli_.float) +def meta_bernoulli_(self, p=0.5, generator=None): + return self + + +@register_meta(aten.bernoulli.p) +def meta_bernoulli_p(self, p=0.5, generator=None): + # https://github.com/pytorch/pytorch/issues/88612 + return torch.empty_like(self, memory_format=torch.contiguous_format) + + +@register_meta([aten.poisson.default, aten.poisson.out]) +@out_wrapper() +def meta_poisson(self, generator=None): + return torch.empty_like(self) + + +@register_meta(aten._fused_moving_avg_obs_fq_helper.default) +def meta__fused_moving_avg_obs_fq_helper( + self, + observer_on, + fake_quant_on, + running_min, + running_max, + scale, + zero_point, + averaging_const, + quant_min, + quant_max, + ch_axis, + per_row_fake_quant=False, + symmetric_quant=False, +): + torch._check( + ch_axis < self.dim(), + lambda: "Error in fused_moving_avg_obs_fake_quant_cpu: ch_axis must be < self.dim()", + ) + mask = torch.empty_like(self, dtype=torch.bool) + return (torch.empty_like(self), mask) + + +@register_meta(aten.mm) +@out_wrapper(exact_dtype=True) +def meta_mm(a, b): + torch._check(a.dim() == 2, lambda: "a must be 2D") + torch._check(b.dim() == 2, lambda: "b must be 2D") + N, M1 = a.shape + M2, P = b.shape + torch._check( + M1 == M2, + lambda: f"a and b must have same reduction dim, but got [{N}, {M1}] X [{M2}, {P}].", + ) + return a.new_empty(N, P) + + +def _compute_reduction_shape(self, dims, keepdim): + if keepdim: + return tuple(self.shape[i] if i not in dims else 1 for i in range(self.ndim)) + + return utils.compute_reduction_output_shape(self.shape, dims) + + +# FakeTensors (meta tensors with a device) will report device as meta +# when running meta kernels. Here, access the "fake device" of FakeTensor if it +# exists so meta kernels which have diverge per device will be more +# accurate when run with FakeTensors +def device_hint(tensor) -> "str": + if isinstance(tensor, torch._subclasses.FakeTensor): + return tensor.fake_device.type + elif ( + hasattr(tensor, "device") + and hasattr(tensor.device, "type") + and tensor.device.type != "meta" + ): + return tensor.device.type + else: + return "cuda" # default to cuda + + +def calc_conv_nd_return_shape( + input_tensor: torch.Tensor, + weight: torch.Tensor, + stride: Union[list[int], int], + padding: Union[list[int], int], + dilation: Union[list[int], int], + is_transposed: bool, + groups: int, + output_padding: Optional[Union[list[int], int]] = None, +): + def _formula(ln: int, p: int, d: int, k: int, s: int) -> int: + """ + Formula to apply to calculate the length of some dimension of the output + + See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html + + Args: + ln: length of the dimension + p: padding in that dim + d: dilation in that dim + k: kernel size in that dim + s: stride in that dim + Returns: + The output length + """ + return (ln + 2 * p - d * (k - 1) - 1) // s + 1 + + def _formula_transposed(ln: int, p: int, d: int, k: int, s: int, op: int) -> int: + """ + Formula to apply to calculate the length of some dimension of the output + if transposed convolution is used. + See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html + + Args: + ln: length of the dimension + p: padding in that dim + d: dilation in that dim + k: kernel size in that dim + s: stride in that dim + op: output padding in that dim + + Returns: + The output length + """ + return (ln - 1) * s - 2 * p + d * (k - 1) + op + 1 + + kernel_size = weight.shape[2:] + dims = input_tensor.shape[2:] + if is_transposed: + out_channels = groups * weight.shape[1] + else: + out_channels = weight.shape[0] + if weight.shape[1] * groups != input_tensor.shape[1]: + raise RuntimeError("Invalid channel dimensions") + + ret_shape = [input_tensor.shape[0], out_channels] + if isinstance(stride, IntLike): + stride = [stride] * len(dims) + elif len(stride) == 1: + stride = [stride[0]] * len(dims) + + if isinstance(padding, IntLike): + padding = [padding] * len(dims) + elif len(padding) == 1: + padding = [padding[0]] * len(dims) + + if isinstance(dilation, IntLike): + dilation = [dilation] * len(dims) + elif len(dilation) == 1: + dilation = [dilation[0]] * len(dims) + + output_padding_list: Optional[list[int]] = None + if output_padding: + if isinstance(output_padding, IntLike): + output_padding_list = [output_padding] * len(dims) + elif len(output_padding) == 1: + output_padding_list = [output_padding[0]] * len(dims) + else: + output_padding_list = output_padding + + for i in range(len(dims)): + # If output_padding is present, we are dealing with a transposed convolution + if output_padding_list: + ret_shape.append( + _formula_transposed( + dims[i], + padding[i], + dilation[i], + kernel_size[i], + stride[i], + output_padding_list[i], + ) + ) + else: + ret_shape.append( + _formula(dims[i], padding[i], dilation[i], kernel_size[i], stride[i]) + ) + from torch.fx.experimental.symbolic_shapes import sym_or + + torch._check( + sym_or(*[x > 0 for x in ret_shape[2:]]), + lambda: f"Given input size per channel: {list(dims)}. " + f"Calculated output size per channel: {ret_shape[2:]}. " + f"Output size is too small", + ) + + return ret_shape + + +def is_channels_last(ten): + return torch._prims_common.suggest_memory_format(ten) == torch.channels_last + + +@register_meta(aten.miopen_batch_norm.default) +def meta_miopen_batch_norm( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + running_mean: Optional[torch.Tensor], + running_var: Optional[torch.Tensor], + training: bool, + exponential_average_factor: float, + epsilon: float, +): + # In batch norm the output is of the same shape as the input + out_shape = input_tensor.shape + + # If tensor is provided for running_mean and running_var then use this. If these are not + # provided then we return the shape of weight tensor. Similar to how this is handled in the decomposition + save_mean_shape = running_mean.shape if running_mean is not None else weight.shape + save_var_shape = running_var.shape if running_var is not None else weight.shape + + def pick_memory_format(): + if is_channels_last(input_tensor): + return torch.channels_last + if input_tensor.is_contiguous(memory_format=torch.contiguous_format): + return torch.contiguous_format + return torch.contiguous_format + + out = input_tensor.new_empty(out_shape).to(memory_format=pick_memory_format()) + + if training: + save_mean = input_tensor.new_empty(save_mean_shape) + save_var = input_tensor.new_empty(save_var_shape) + else: + save_mean = input_tensor.new_empty((0,)) + save_var = input_tensor.new_empty((0,)) + + return out, save_mean, save_var + + +@register_meta(aten.convolution.default) +def meta_conv( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: list[int], + padding: list[int], + dilation: list[int], + is_transposed: bool, + output_padding: list[int], + groups: int, +): + def pick_memory_format(): + if device_hint(input_tensor) == "cuda": + if is_channels_last(input_tensor) or is_channels_last(weight): + return torch.channels_last + else: + if is_channels_last(input_tensor): + return torch.channels_last + if input_tensor.is_contiguous(memory_format=torch.contiguous_format): + return torch.contiguous_format + elif input_tensor.is_contiguous(memory_format=torch.preserve_format): + return torch.preserve_format + + shape_out = calc_conv_nd_return_shape( + input_tensor, + weight, + stride, + padding, + dilation, + is_transposed, + groups, + output_padding if is_transposed else None, + ) + + input_channels_dim = 1 + output_channels_dim = 1 + if input_tensor.size(input_channels_dim) == 0: + shape_out[output_channels_dim] = 0 + + out = input_tensor.new_empty(shape_out) + out = out.to(memory_format=pick_memory_format()) # type: ignore[call-overload] + return out + + +if torch._C._has_mkldnn: + _meta_lib_dont_use_me_use_register_meta_for_mkldnn = torch.library.Library( + "mkldnn", "IMPL", "Meta" + ) + + @register_meta(torch.ops.mkldnn._convolution_pointwise.default) + def meta_mkldnn_convolution_default( + input_tensor, + weight, + bias, + padding, + stride, + dilation, + groups, + attr, + scalars, + algorithm, + ): + shape_out = calc_conv_nd_return_shape( + input_tensor, weight, stride, padding, dilation, False, groups, [] + ) + out = input_tensor.new_empty(shape_out) + out_memory_format = torch.channels_last + if input_tensor.dim() == 5: + out_memory_format = torch.channels_last_3d + out = out.to(memory_format=out_memory_format) # type: ignore[call-overload] + return out + + @register_meta(torch.ops.mkldnn._linear_pointwise.default) + def meta_linear_pointwise_default( + input_tensor, weight, bias, attr, scalars, algorithm + ): + return input_tensor.new_empty((*input_tensor.shape[:-1], weight.shape[0])) + + if torch._C.has_mkl: + _meta_lib_dont_use_me_use_register_meta_for_mkl = torch.library.Library( + "mkl", "IMPL", "Meta" + ) + + @register_meta(torch.ops.mkl._mkl_linear) + def meta_mkl_linear(input_tensor, packed_weight, orig_weight, bias, batch_size): + return input_tensor.new_empty( + (*input_tensor.shape[:-1], orig_weight.shape[0]) + ) + + _meta_lib_dont_use_me_use_register_meta_for_onednn = torch.library.Library( + "onednn", "IMPL", "Meta" + ) + + @register_meta(torch.ops.onednn.qconv2d_pointwise.default) + @register_meta(torch.ops.onednn.qconv_pointwise.default) + def meta_qconv_pointwise( + x, + x_scale, + x_zp, + w, # prepacked_weight + w_scale, + w_zp, + bias, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point, + output_dtype, + attr, + scalars, + algorithm, + ): + shape_out = calc_conv_nd_return_shape( + x, + w, + stride, + padding, + dilation, + False, + groups, + None, + ) + if output_dtype is None: + output_dtype = x.dtype + assert output_dtype in [ + torch.float32, + torch.bfloat16, + torch.uint8, + torch.int8, + torch.float8_e4m3fn, + ] + out = x.new_empty(shape_out, dtype=output_dtype) + assert len(shape_out) in [3, 4, 5], ( + "Expect output to be 3d/4d/5d for conv1d/2d/3d" + ) + format = { + 3: torch.contiguous_format, + 4: torch.channels_last, + 5: torch.channels_last_3d, + }[len(shape_out)] + out = out.to(memory_format=format) + return out + + @register_meta(torch.ops.onednn.qconv2d_pointwise.binary) + def meta_qconv2d_pointwise_binary( + x, + x_scale, + x_zp, + w, + w_scale, + w_zp, + accum, + bias, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point, + output_dtype, + accum_scale, + accum_zero_point, + binary_op_name, + alpha, + unary_op_name, + unary_op_args, + unary_op_algorithm, + ): + assert binary_op_name == "sum" + return accum + + @register_meta(torch.ops.onednn.qlinear_pointwise.default) + @register_meta(torch.ops.onednn.qlinear_pointwise.tensor) + def meta_qlinear_pointwise( + x, + x_scale, + x_zp, + w, + w_scale, + w_zp, + bias, + output_scale, + output_zero_point, + output_dtype, + post_op_name, + post_op_args, + post_op_algorithm, + ): + output_shape = list(x.shape) + # The weight has been transposed during the qlinear weight prepack process. + output_shape[-1] = w.shape[1] + assert output_dtype in [ + torch.float32, + torch.bfloat16, + torch.int8, + torch.uint8, + torch.float8_e4m3fn, + ] + out = x.new_empty(output_shape, dtype=output_dtype) + return out + + @register_meta(torch.ops.onednn.qlinear_pointwise.binary) + @register_meta(torch.ops.onednn.qlinear_pointwise.binary_tensor) + def meta_qlinear_pointwise_binary( + x, + x_scale, + x_zp, + w, + w_scale, + w_zp, + x_2, + bias, + output_scale, + output_zero_point, + output_dtype, + x2_scale, + x2_zp, + binary_op_name, + alpha, + unary_op_name, + unary_op_args, + unary_op_algorithm, + ): + if binary_op_name == "sum": + return x_2 + output_shape = list(x.shape) + # The weight has been transposed during the qlinear weight prepack process. + output_shape[-1] = w.shape[1] + assert output_dtype in [ + torch.float32, + torch.bfloat16, + torch.uint8, + torch.int8, + torch.float8_e4m3fn, + ] + out = x.new_empty(output_shape, dtype=output_dtype) + return out + + @register_meta(torch.ops.onednn.linear_dynamic_fp16.default) + @register_meta(torch.ops.onednn.linear_relu_dynamic_fp16.default) + def meta_linear_dynamic_fp16( + x, + w, + bias, + ): + output_shape = list(x.shape) + # The weight has been transposed during the qlinear weight prepack process. + output_shape[-1] = w.shape[1] + out = x.new_empty(output_shape) + return out + + _meta_lib_dont_use_me_use_register_meta_for_quantized = torch.library.Library( + "quantized", "IMPL", "Meta" + ) + + @register_meta(torch.ops.quantized.max_pool2d) + def meta_quantized_max_pool2d( + input, + kernel_size, + stride=(), + padding=(0,), + dilation=(1,), + ceil_mode=False, + ): + ( + nInputPlane, + outputHeight, + outputWidth, + ) = max_pool2d_checks_and_compute_shape( + input, kernel_size, stride, padding, dilation, ceil_mode + ) + nbatch = input.size(-4) if input.dim() == 4 else 1 + memory_format = torch.channels_last + if input.dim() == 3: + size = [nInputPlane, outputHeight, outputWidth] + else: + size = [nbatch, nInputPlane, outputHeight, outputWidth] + return torch.empty( + size, + dtype=input.dtype, + device=input.device, + memory_format=memory_format, + ) + + @register_meta(torch.ops.quantized.int4mm_packed_weight_cpu) + def meta_int4mm_packed_weight_cpu(x, w, q_group_size, q_scale_and_zeros): + torch._check(x.dim() == 2, f"x must be a 2D tensor, got {x.dim()}D") + torch._check(w.dim() == 2, f"w must be a 2D tensor, got {w.dim()}D") + torch._check( + x.dtype in [torch.float32, torch.float16, torch.bfloat16], + f"expected x to be f32/f16/bf16, got {x.dtype}", + ) + torch._check(w.dtype == torch.uint8, f"expected w to be uint8, got {w.dtype}") + torch._check( + q_group_size.dtype == torch.int64, + f"q_group_size must be int64, got {q_group_size.dtype}", + ) + torch._check( + q_scale_and_zeros.dtype == x.dtype, + f"q_scale_and_zeros must have the same dtype as x, got {q_scale_and_zeros.dtype}", + ) + return x.new_empty(x.size(0), w.size(0), dtype=x.dtype) + + +# from check_dim_size() in aten/src/ATen/TensorUtils.cpp. +def check_dim_size(tensor, dim, dim_size, size): + torch._check( + tensor.dim() == dim and tensor.shape[dim_size] == size, + lambda: f"Expected a tensor of dimension {dim} and tensor.size[{dim_size}] == {size}, " + + f"but got : dimension {tensor.dim()} and tensor.size[{dim_size}] = {tensor.shape[dim_size]}", + ) + + +@register_meta(aten.avg_pool2d.default) +def meta_avg_pool2d( + input, + kernel_size, + stride=(), + padding=(0,), + ceil_mode=False, + count_include_pad=True, + divisor_override=None, +): + def unpack(name, val): + torch._check( + len(val) in [1, 2], + lambda: f"avg_pool2d: {name} must either be a single int, or a tuple of two ints", + ) + H = val[0] + W = H if len(val) == 1 else val[1] + return H, W + + kH, kW = unpack("kernel_size", kernel_size) + torch._check( + len(stride) in [0, 1, 2], + lambda: "avg_pool2d: stride must either be omitted, a single int, or a tuple of two ints", + ) + torch._check( + input.dtype not in [torch.uint8, torch.uint16, torch.uint32, torch.uint64], + lambda: f""""avg_pool2d" not implemented for '{input.dtype.__str__()}'""", + ) + if len(stride) == 0: + dH, dW = kH, kW + elif len(stride) == 1: + dH, dW = stride[0], stride[0] + else: + dH, dW = unpack("stride", stride) + + padH, padW = unpack("padding", padding) + + torch._check( + divisor_override is None or divisor_override != 0, + lambda: "divisor must be not zero", + ) + + nbatch = input.size(-4) if input.dim() == 4 else 1 + nInputPlane = input.size(-3) + inputHeight = input.size(-2) + inputWidth = input.size(-1) + + outputHeight = pooling_output_shape(inputHeight, kH, padH, dH, 1, ceil_mode) + outputWidth = pooling_output_shape(inputWidth, kW, padW, dW, 1, ceil_mode) + + memory_format = utils.suggest_memory_format(input) + pool2d_shape_check( + input, + kH, + kW, + dH, + dW, + padH, + padW, + 1, + 1, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + memory_format, + ) + + if input.dim() == 3: + size = [nInputPlane, outputHeight, outputWidth] + else: + size = [nbatch, nInputPlane, outputHeight, outputWidth] + return torch.empty( + size, + dtype=input.dtype, + device=input.device, + memory_format=memory_format, + ) + + +# from avg_pool2d_backward_shape_check() in aten/src/ATen/native/Pool.h. +def avg_pool2d_backward_shape_check( + input, + gradOutput, + nbatch, + kH, + kW, + dH, + dW, + padH, + padW, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + mem_format, +): + pool2d_shape_check( + input, + kH, + kW, + dH, + dW, + padH, + padW, + 1, + 1, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + mem_format, + ) + + ndim = input.dim() + nOutputPlane = nInputPlane + + check_dim_size(gradOutput, ndim, ndim - 3, nOutputPlane) + check_dim_size(gradOutput, ndim, ndim - 2, outputHeight) + check_dim_size(gradOutput, ndim, ndim - 1, outputWidth) + + +# Don't override the C++ registration. +@register_meta(aten.avg_pool2d_backward.default) +def meta_avg_pool2d_backward( + gradOutput_, + input, + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, +): + # From aten/src/ATen/native/AveragePool2d.cpp structured kernel meta func. + torch._check( + len(kernel_size) == 1 or len(kernel_size) == 2, + lambda: "avg_pool2d: kernel_size must either be a single int, or a tuple of two ints", + ) + kH = kernel_size[0] + kW = kH if len(kernel_size) == 1 else kernel_size[1] + torch._check( + len(stride) == 0 or len(stride) == 1 or len(stride) == 2, + lambda: "avg_pool2d: stride must either be omitted, a single int, or a tuple of two ints", + ) + dH = kH if len(stride) == 0 else stride[0] + dW = kW if len(stride) == 0 else dH if len(stride) == 1 else stride[1] + torch._check( + len(padding) == 1 or len(padding) == 2, + lambda: "avg_pool2d: padding must either be a single int, or a tuple of two ints", + ) + padH = padding[0] + padW = padH if len(padding) == 1 else padding[1] + + torch._check( + divisor_override is None or divisor_override != 0, + lambda: "divisor must be not zero", + ) + + input_size = input.shape + nbatch = input_size[-4] if input.dim() == 4 else 1 + nInputPlane = input_size[-3] + inputHeight = input_size[-2] + inputWidth = input_size[-1] + + outputHeight = pooling_output_shape(inputHeight, kH, padH, dH, 1, ceil_mode) + outputWidth = pooling_output_shape(inputWidth, kW, padW, dW, 1, ceil_mode) + + mem_format = utils.suggest_memory_format(input) + + avg_pool2d_backward_shape_check( + input, + gradOutput_, + nbatch, + kH, + kW, + dH, + dW, + padH, + padW, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + mem_format, + ) + + return torch.empty( + input_size, + dtype=input.dtype, + device=input.device, + memory_format=mem_format, + ) + + +@register_meta(aten.avg_pool3d) +@out_wrapper() +def meta_avg_pool3d( + input, + kernel_size, + stride=(), + padding=(0,), + ceil_mode=False, + count_include_pad=True, + divisor_override=None, +): + torch._check( + len(kernel_size) in (1, 3), + lambda: "avg_pool3d: kernel_size must be a single int, or a tuple of three ints", + ) + kT = kernel_size[0] + kH = kT if len(kernel_size) == 1 else kernel_size[1] + kW = kT if len(kernel_size) == 1 else kernel_size[2] + + torch._check( + not stride or len(stride) in (1, 3), + lambda: "avg_pool3d: stride must be omitted, a single int, or a tuple of three ints", + ) + torch._check( + input.dtype not in [torch.uint8, torch.uint16, torch.uint32, torch.uint64], + lambda: f""""avg_pool3d" not implemented for '{input.dtype.__str__()}'""", + ) + dT = kT if not stride else stride[0] + dH = kH if not stride else (dT if len(stride) == 1 else stride[1]) + dW = kW if not stride else (dT if len(stride) == 1 else stride[2]) + + torch._check( + len(padding) in (1, 3), + lambda: "avg_pool3d: padding must be a single int, or a tuple of three ints", + ) + padT = padding[0] + padH = padT if len(padding) == 1 else padding[1] + padW = padT if len(padding) == 1 else padding[2] + + torch._check( + input.ndim in (4, 5), + lambda: "non-empty 4D or 5D (batch mode) tensor expected for input", + ) + + torch._check( + not divisor_override or divisor_override != 0, + lambda: "divisor must be not zero", + ) + + nbatch = input.size(0) + nslices = input.size(-4) + itime = input.size(-3) + iheight = input.size(-2) + iwidth = input.size(-1) + + otime = pooling_output_shape(itime, kT, padT, dT, 1, ceil_mode) + oheight = pooling_output_shape(iheight, kH, padH, dH, 1, ceil_mode) + owidth = pooling_output_shape(iwidth, kW, padW, dW, 1, ceil_mode) + + pool3d_shape_check( + input, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + padT, + padH, + padW, + 1, + 1, + 1, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + "avg_pool3d()", + check_input_size=True, + ) + + if input.ndim == 4: + return input.new_empty((nslices, otime, oheight, owidth)) + else: + return input.new_empty((nbatch, nslices, otime, oheight, owidth)) + + +@register_meta(aten.avg_pool3d_backward) +@out_wrapper("grad_input") +def meta_avg_pool3d_backward( + grad_output, + input, + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + divisor_override, +): + torch._check( + len(kernel_size) in (1, 3), + lambda: "avg_pool3d: kernel_size must be a single int, or a tuple of three ints", + ) + kT = kernel_size[0] + kH = kT if len(kernel_size) == 1 else kernel_size[1] + kW = kT if len(kernel_size) == 1 else kernel_size[2] + + torch._check( + not stride or len(stride) in (1, 3), + lambda: "avg_pool3d: stride must be omitted, a single int, or a tuple of three ints", + ) + dT = kT if not stride else stride[0] + dH = kH if not stride else (dT if len(stride) == 1 else stride[1]) + dW = kW if not stride else (dT if len(stride) == 1 else stride[2]) + + torch._check( + len(padding) in (1, 3), + lambda: "avg_pool3d: padding must be a single int, or a tuple of three ints", + ) + padT = padding[0] + padH = padT if len(padding) == 1 else padding[1] + padW = padT if len(padding) == 1 else padding[2] + + torch._check( + input.ndim in (4, 5), + lambda: "non-empty 4D or 5D (batch mode) tensor expected for input", + ) + + torch._check( + not divisor_override or divisor_override != 0, + lambda: "divisor must be not zero", + ) + + nslices = input.size(-4) + itime = input.size(-3) + iheight = input.size(-2) + iwidth = input.size(-1) + + otime_for_shape_check = pooling_output_shape(itime, kT, padT, dT, 1, ceil_mode) + oheight_for_shape_check = pooling_output_shape(iheight, kH, padH, dH, 1, ceil_mode) + owidth_for_shape_check = pooling_output_shape(iwidth, kW, padW, dW, 1, ceil_mode) + + avg_pool3d_backward_shape_check( + input, + grad_output, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + padT, + padH, + padW, + itime, + iheight, + iwidth, + otime_for_shape_check, + oheight_for_shape_check, + owidth_for_shape_check, + "avg_pool3d_backward()", + ) + + return input.new_empty(input.shape) + + +@register_meta(aten._adaptive_avg_pool2d.default) +def meta_adaptive_avg_pool2d(self, output_size): + torch._check( + self.ndim == 3 or self.ndim == 4, + lambda: f"Expected 3D or 4D tensor, but got {self.shape}", + ) + output_shape = self.shape[:-2] + tuple(output_size) + memory_format = utils.suggest_memory_format(self) + # need to set memory_format to preserve the memory format of the input + # channel last input should have channel last output + return torch.empty( + output_shape, + dtype=self.dtype, + device=self.device, + memory_format=memory_format, + ) + + +@register_meta(aten._adaptive_avg_pool3d.default) +def meta_adaptive_avg_pool3d(self, output_size): + torch._check( + self.ndim == 4 or self.ndim == 5, + lambda: f"Expected 4D or 5D tensor, but got {self.shape}", + ) + return self.new_empty(self.shape[:-3] + tuple(output_size)) + + +@register_meta(aten._adaptive_avg_pool2d_backward.default) +def meta__adaptive_avg_pool2d_backward(grad_out, self): + ndim = grad_out.ndim + for i in range(1, ndim): + torch._check( + grad_out.size(i) > 0, + lambda: f"adaptive_avg_pool2d_backward(): Expected grad_output to have non-zero \ + size for non-batch dimensions, {grad_out.shape} with dimension {i} being empty", + ) + torch._check( + ndim == 3 or ndim == 4, + lambda: f"adaptive_avg_pool2d_backward(): Expected 3D or 4D tensor, but got {self.shape}", + ) + torch._check( + self.dtype == grad_out.dtype, + lambda: f"expected dtype {self.dtype} for `grad_output` but got dtype {grad_out.dtype}", + ) + memory_format = torch.contiguous_format + if is_channels_last(self): + memory_format = torch.channels_last + return self.new_empty(self.shape).to(memory_format=memory_format) + + +@register_meta(aten._adaptive_avg_pool3d_backward) +@out_wrapper("grad_input") +def meta__adaptive_avg_pool3d_backward(grad_output, self): + _adaptive_pool_empty_output_check(grad_output, "adaptive_avg_pool3d_backward") + return torch.empty_like(self, memory_format=torch.legacy_contiguous_format) + + +def _adaptive_pool_empty_output_check(grad_output: Tensor, arg_name: str): + ndim = grad_output.ndim + for i in range(1, ndim): + torch._check( + grad_output.size(i) > 0, + lambda: ( + f"{arg_name}(): Expected grad_output to have non-zero size for non-batch dimensions, " + f"but grad_output has sizes {grad_output.shape} with dimension {i} being empty" + ), + ) + + +@register_meta(aten.adaptive_max_pool2d) +@out_wrapper("out", "indices") +def meta_adaptive_max_pool2d(input, output_size): + ndim = input.ndim + torch._check( + ndim in (3, 4), + lambda: f"adaptive_max_pool2d(): Expected 3D or 4D tensor, but got: {input.shape}", + ) + for i in range(1, ndim): + torch._check( + input.size(i) > 0, + lambda: ( + f"adaptive_max_pool2d(): Expected input to have non-zero size for non-batch dimensions, " + f"but input has sizes {input.shape} with dimension {i} being empty" + ), + ) + + torch._check( + len(output_size) == 2, + lambda: "adaptive_max_pool2d(): internal error: output_size.size() must be 2", + ) + + dimH = 1 + sizeB = 1 + sizeD = 0 + + if input.ndim == 4: + sizeB = input.size(0) + dimH += 1 + + sizeD = input.size(dimH - 1) + osizeH, osizeW = output_size + + if input.ndim == 3: + out_shape = (sizeD, osizeH, osizeW) + out = input.new_empty(out_shape) + indices = input.new_empty(out_shape, dtype=torch.int64) + return out, indices + else: + out_shape = (sizeB, sizeD, osizeH, osizeW) # type: ignore[assignment] + memory_format = utils.suggest_memory_format(input) + out = input.new_empty(out_shape).to(memory_format=memory_format) + indices = input.new_empty(out_shape, dtype=torch.int64).to( + memory_format=memory_format + ) + return out, indices + + +@register_meta(aten.adaptive_max_pool2d_backward) +@out_wrapper("grad_input") +def meta_adaptive_max_pool2d_backward(grad_output, input, indices): + ndim = grad_output.ndim + torch._check( + ndim in (3, 4), + lambda: f"adaptive_max_pooling2d_backward(): Expected 3D or 4D grad_output, but got: {grad_output.shape}", + ) + + _adaptive_pool_empty_output_check(grad_output, "adaptive_max_pool2d_backward") + + torch._check( + input.dtype == grad_output.dtype, + lambda: f"expected dtype {input.dtype} for `grad_output` but got dtype {grad_output.dtype}", + ) + + memory_format = utils.suggest_memory_format(input) + return input.new_empty(input.shape).to(memory_format=memory_format) + + +@register_meta(aten.adaptive_max_pool3d) +@out_wrapper("out", "indices") +def meta_adaptive_max_pool3d(input, output_size): + ndim = input.ndim + torch._check( + ndim in (4, 5), + lambda: f"adaptive_max_pool3d(): Expected 4D or 5D tensor, but got: {input.shape}", + ) + for i in range(1, ndim): + torch._check( + input.size(i) > 0, + lambda: ( + f"adaptive_max_pool3d(): Expected input to have non-zero size for non-batch dimensions, " + f"but input has sizes {input.shape} with dimension {i} being empty" + ), + ) + + torch._check( + len(output_size) == 3, + lambda: "adaptive_max_pool3d(): internal error: output_size.size() must be 3", + ) + + dimD = 0 + sizeB = 1 + sizeD = 0 + + if ndim == 5: + sizeB = input.size(0) + dimD += 1 + + sizeD = input.size(dimD) + osizeT, osizeH, osizeW = output_size + + if ndim == 4: + out_shape = (sizeD, osizeT, osizeH, osizeW) + else: + out_shape = (sizeB, sizeD, osizeT, osizeH, osizeW) # type: ignore[assignment] + + out = input.new_empty(out_shape) + indices = input.new_empty(out_shape, dtype=torch.int64) + + return out, indices + + +@register_meta(aten.adaptive_max_pool3d_backward) +@out_wrapper("grad_input") +def meta_adaptive_max_pool3d_backward(grad_output, input, indices): + _adaptive_pool_empty_output_check(grad_output, "adaptive_max_pool3d_backward") + return input.new_empty(input.shape) + + +@register_meta(aten.repeat_interleave.Tensor) +def meta_repeat_interleave_Tensor(repeats, output_size=None): + if output_size is None: + raise RuntimeError("cannot repeat_interleave a meta tensor without output_size") + return repeats.new_empty(output_size) + + +@register_meta([aten.complex.default, aten.complex.out]) +@out_wrapper() +def meta_complex(real, imag): + assert real.dtype.is_floating_point + assert imag.dtype.is_floating_point + result = elementwise_meta( + real.to(corresponding_complex_dtype(real.dtype)), + imag.to(corresponding_complex_dtype(imag.dtype)), + type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ) + return result + + +@register_meta([aten.nonzero_static.default, aten.nonzero_static.out]) +@out_wrapper() +def nonzero_static(self, *, size, fill_value: int = -1): + return self.new_empty((size, self.dim()), dtype=torch.long) + + +@register_meta([torch.ops.aten.nonzero.default, torch.ops.aten.nonzero.out]) +@out_wrapper() +def nonzero(self): + torch._check_not_implemented( + exp_config.meta_nonzero_assume_all_nonzero, + lambda: "The register_meta function for torch.nonzero() raises unimplemented by default, " + "as a correct data-independent implementation does not exist. This implementation " + "returns a fake value, assuming all elements of the tensor are non-zero. " + "To enable this registration, please set " + "'torch.fx.experimental._config.meta_nonzero_assume_all_nonzero' to True.", + ) + return torch.empty_strided( + (self.numel(), self.dim()), + (1, self.numel()), + dtype=torch.long, + device=self.device, + ) + + +@register_meta([aten.index.Tensor, aten._unsafe_index.Tensor]) +def meta_index_Tensor(self, indices): + torch._check(bool(indices), lambda: "at least one index must be provided") + # aten::index is the internal advanced indexing implementation + # checkIndexTensorTypes and expandTensors + result: list[Optional[Tensor]] = [] + for i, index in enumerate(indices): + if index is not None: + torch._check( + index.dtype in [torch.long, torch.int, torch.int8, torch.bool], + lambda: "tensors used as indices must be long, int, byte or bool tensors", + ) + if index.dtype in [torch.int8, torch.bool]: + nonzero = index.nonzero() + k = len(result) + torch._check_index( + k + index.ndim <= self.ndim, + lambda: f"too many indices for tensor of dimension {self.ndim}", + ) + for j in range(index.ndim): + torch._check_index( + index.shape[j] == self.shape[k + j], + lambda: f"The shape of the mask {index.shape} at index {i} " + f"does not match the shape of the indexed tensor {self.shape} at index {k + j}", + ) + result.append(nonzero.select(1, j)) + else: + result.append(index) + else: + result.append(index) + indices = result + torch._check( + len(indices) <= self.ndim, + lambda: f"too many indices for tensor of dimension {self.ndim} (got {len(indices)})", + ) + # expand_outplace + import torch._refs as refs # avoid import cycle in mypy + + indices = list(refs._maybe_broadcast(*indices)) + # add missing null tensors + while len(indices) < self.ndim: + indices.append(None) + + # hasContiguousSubspace + # true if all non-null tensors are adjacent + # See: + # https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing + # https://stackoverflow.com/questions/53841497/why-does-numpy-mixed-basic-advanced-indexing-depend-on-slice-adjacency + state = 0 + has_contiguous_subspace = False + for index in indices: + if state == 0: + if index is not None: + state = 1 + elif state == 1: + if index is None: + state = 2 + else: + if index is not None: + break + else: + has_contiguous_subspace = True + + # transposeToFront + # This is the logic that causes the newly inserted dimensions to show up + # at the beginning of the tensor, if they're not contiguous + if not has_contiguous_subspace: + dims = [] + transposed_indices = [] + for i, index in enumerate(indices): + if index is not None: + dims.append(i) + transposed_indices.append(index) + for i, index in enumerate(indices): + if index is None: + dims.append(i) + transposed_indices.append(index) + self = self.permute(dims) + indices = transposed_indices + + # AdvancedIndex::AdvancedIndex + # Now we can assume the indices have contiguous subspace + # This is simplified from AdvancedIndex which goes to more effort + # to put the input and indices in a form so that TensorIterator can + # take them. If we write a ref for this, probably that logic should + # get implemented + before_shape: list[int] = [] + after_shape: list[int] = [] + replacement_shape: list[int] = [] + for dim, index in enumerate(indices): + if index is None: + if replacement_shape: + after_shape.append(self.shape[dim]) + else: + before_shape.append(self.shape[dim]) + else: + replacement_shape = list(index.shape) + + def _restride_src(self): + """ + This follows restride_src in TensorAdvancedIndexing.cpp + """ + shape = before_shape + replacement_shape + after_shape + strides = list(self.stride()) + strides[len(before_shape) : len(self.shape) - len(after_shape)] = [0] * len( + replacement_shape + ) + return self.as_strided(shape, strides) + + out = self.new_empty(before_shape + replacement_shape + after_shape) + from torch.fx.experimental.symbolic_shapes import guard_or_false + + if guard_or_false(self.numel() == 0): + # No need to worry about the output strides if self is empty. + return out + + # Try to follow eager to decide the output stride based on self. + # Note that perm here is the reverse of the 'perm_' decided by + # TensorIteratorBase::reorder_dimensions + restrided_self = _restride_src(self) + perm = utils.compute_elementwise_output_logical_to_physical_perm(restrided_self) + + # Follow TensorIteratorBase::allocate_or_resize_outputs + if list(perm) != list(range(len(perm))): + perm_shape = utils.apply_perm(out.shape, perm) + new_stride = utils.make_contiguous_strides_for(perm_shape) + new_stride = utils.apply_perm(new_stride, utils.invert_perm(perm)) + out = out.as_strided(out.size(), new_stride) + return out + + +@register_meta([aten.convolution_backward.default]) +def meta_convolution_backward( + grad_output_, + input_, + weight_, + bias_sizes_opt, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output_mask, +): + # High level logic taken from slow_conv3d_backward_cpu which should + # be representative of all convolution_backward impls + backend_grad_input = None + backend_grad_weight = None + backend_grad_bias = None + + if output_mask[0]: + backend_grad_input = grad_output_.new_empty(input_.size()) + if output_mask[1]: + backend_grad_weight = grad_output_.new_empty(weight_.size()) + if output_mask[2]: + backend_grad_bias = grad_output_.new_empty(bias_sizes_opt) + + return (backend_grad_input, backend_grad_weight, backend_grad_bias) + + +@register_meta([aten.addbmm.default, aten.addbmm.out]) +@out_wrapper(exact_dtype=True) +def meta_addbmm(self, batch1, batch2, *, beta=1, alpha=1): + dim1 = batch1.size(1) + dim2 = batch2.size(2) + self = self.expand((dim1, dim2)) + torch._check(batch1.dim() == 3, lambda: "batch1 must be a 3D tensor") + torch._check(batch2.dim() == 3, lambda: "batch2 must be a 3D tensor") + torch._check( + batch1.size(0) == batch2.size(0), + lambda: f"batch1 and batch2 must have same number of batches, got {batch1.size(0)} and {batch2.size(0)}", + ) + torch._check( + batch1.size(2) == batch2.size(1), + lambda: ( + f"Incompatible matrix sizes for bmm ({batch1.size(1)}x{batch1.size(2)} " + f"and {batch2.size(1)}x{batch2.size(2)})" + ), + ) + torch._check( + self.size(0) == dim1 and self.size(1) == dim2, + lambda: "self tensor does not match matmul output shape", + ) + return self.new_empty(self.size()) + + +@register_meta([aten.randint_like.Tensor]) +def meta_randint_like(self, high, **kwargs): + return self.new_empty(self.size()) + + +@register_meta([aten._fused_adam_.default, aten._fused_adamw_.default]) +def meta__fused_adam_( + self, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + *, + lr, + beta1, + beta2, + weight_decay, + eps, + amsgrad, + maximize, + grad_scale=None, + found_inf=None, +): + for l in [self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps]: + torch._check( + isinstance(l, list), + lambda: f"exponent must be a tensor list but got {type(l)}", + ) + + +@register_meta([aten._fused_adam.default]) +def meta__fused_adam( + self, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + *, + lr, + beta1, + beta2, + weight_decay, + eps, + amsgrad, + maximize, + grad_scale=None, + found_inf=None, +): + for l in [self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps]: + torch._check( + isinstance(l, list), + lambda: f"exponent must be a tensor list but got {type(l)}", + ) + + def empty_like_list(tensor_list): + return [torch.empty_like(t) for t in tensor_list] + + return ( + empty_like_list(self), + empty_like_list(grads), + empty_like_list(exp_avgs), + empty_like_list(exp_avg_sqs), + empty_like_list(max_exp_avg_sqs), + ) + + +@register_meta([aten._int_mm]) +@out_wrapper() +def meta__int_mm(a, b): + torch._check(a.dim() == 2, lambda: "a must be a 2D tensor") + torch._check(b.dim() == 2, lambda: "b must be a 2D tensor") + torch._check( + a.dtype is torch.int8, + lambda: f"expected self to be int8, got {a.dtype}", + ) + torch._check( + b.dtype is torch.int8, + lambda: f"expected mat2 to be int8, got {b.dtype}", + ) + torch._check( + a.size(1) == b.size(0), + lambda: ( + f"Incompatible matrix sizes for _int_mm ({a.size(0)}x{a.size(1)} " + f"and {b.size(0)}x{b.size(1)})" + ), + ) + return a.new_empty((a.size(0), b.size(1)), dtype=torch.int32) + + +@register_meta([aten._convert_weight_to_int4pack]) +def meta__convert_weight_to_int4pack(w, inner_k_tiles): + torch._check(w.dim() == 2, lambda: "w must be a 2D tensor") + torch._check( + w.dtype is torch.uint8, + lambda: f"expected w to be uint8, got {w.dtype}", + ) + n = w.size(0) + k = w.size(1) * 2 # w is [n][k / 2] uint8 + return w.new_empty( + ( + n // 8, + k // (inner_k_tiles * 16), + 32, + inner_k_tiles // 2, + ), + dtype=torch.int32, + ) + + +@register_meta([aten._convert_weight_to_int4pack_for_cpu]) +def meta__convert_weight_to_int4pack_for_cpu(w, inner_k_tiles): + torch._check(w.dim() == 2, lambda: "w must be a 2D tensor") + torch._check( + w.dtype is torch.int32, + lambda: f"expected w to be int32, got {w.dtype}", + ) + n = w.size(0) + k = w.size(1) # w is [n][k] int32 + return w.new_empty( + (n, k // 2), + dtype=torch.uint8, + ) + + +@register_meta([aten._weight_int4pack_mm]) +def meta__weight_int4pack_mm(x, w, q_group_size, q_scale_and_zeros): + torch._check(x.dim() == 2, lambda: "x must be a 2D tensor") + torch._check(w.dim() == 4, lambda: "w must be a 4D tensor") + torch._check( + x.dtype in [torch.float32, torch.float16, torch.bfloat16], + lambda: f"expected x to be f32/f16/bf16, got {x.dtype}", + ) + torch._check( + w.dtype is torch.int32, + lambda: f"expected w to be int32, got {w.dtype}", + ) + return x.new_empty(x.size(0), w.size(0) * 8, dtype=x.dtype) + + +@register_meta([aten._weight_int4pack_mm_for_cpu]) +def meta__weight_int4pack_mm_for_cpu(x, w, q_group_size, q_scale_and_zeros): + torch._check(x.dim() == 2, lambda: "x must be a 2D tensor") + torch._check(w.dim() == 2, lambda: "w must be a 2D tensor") + torch._check( + x.dtype in [torch.float32, torch.float16, torch.bfloat16], + lambda: f"expected x to be f32/f16/bf16, got {x.dtype}", + ) + torch._check( + w.dtype is torch.uint8, + lambda: f"expected w to be uint8, got {w.dtype}", + ) + return x.new_empty(x.size(0), w.size(0), dtype=x.dtype) + + +@register_meta([aten._weight_int4pack_mm_with_scales_and_zeros]) +def _weight_int4pack_mm_with_scales_and_zeros(x, w, q_group_size, qScale, qZeros): + torch._check(x.dim() == 2, lambda: "x must be a 2D tensor") + torch._check(w.dim() == 2, lambda: "w must be a 2D tensor") + torch._check( + x.dtype in [torch.float32, torch.float16, torch.bfloat16], + lambda: f"expected x to be f32/f16/bf16, got {x.dtype}", + ) + torch._check( + w.dtype is torch.int32, + lambda: f"expected w to be int32, got {w.dtype}", + ) + return x.new_empty(x.size(0), w.size(0), dtype=x.dtype) + + +def kai_roundup(a: int, b: int) -> int: + return ((a + b - 1) // b) * b + + +def get_kai_packed_weight_size(n_bits, N, K, groupsize): + if n_bits == 4: + if groupsize == K: # channelwise + # dotprod params only [1x8x32_neon_dotprod] + kai_nr = 8 + kai_kr = 16 + kai_sr = 2 + kai_num_bytes_sum_rhs = 4 # sizeof(int32_t) + kai_num_bytes_multiplier_rhs = 4 # sizeof(float) + kai_num_bytes_bias = 4 # sizeof(float) + + def kai_k_roundedup(k, kr, sr): + # Since we pack a float and int32 value at the end of the row, + # we must make sure that k is a multiple of 4 for alignment + kr_sr_roundedup4 = kai_roundup(kr * sr, 4) + return kai_roundup(k, kr_sr_roundedup4) + + def kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4cxp_qsu4cxs1s0( + k, nr, kr, sr + ): + k_internal = kai_k_roundedup(k, kr, sr) + + assert (k_internal % 2) == 0, "k_internal must be even" + + return nr * ( + (k_internal // 2) + + kai_num_bytes_multiplier_rhs + + kai_num_bytes_sum_rhs + + kai_num_bytes_bias + ) + + def kai_get_rhs_packed_size_rhs_pack_nxk_qsi4cxp_qsu4cxs1s0( + n, k, nr, kr, sr + ): + num_rows = kai_roundup(n, nr) // nr + + return ( + num_rows + * kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4cxp_qsu4cxs1s0( + k, nr, kr, sr + ) + ) + + return kai_get_rhs_packed_size_rhs_pack_nxk_qsi4cxp_qsu4cxs1s0( + N, K, kai_nr, kai_kr, kai_sr + ) + elif groupsize % 32 == 0 and K % groupsize == 0: # groupwise + kai_nr = 8 + kai_kr = 16 + kai_sr = 2 + kai_num_bytes_sum_rhs = 4 + kai_num_bytes_bias = 4 + kai_nr_multiple_of = 4 + kai_bl_multiple_of = 32 + + def kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0( + n, k, nr, kr, sr, bl + ): + assert (bl % kr) == 0 + assert (nr % kai_nr_multiple_of) == 0 + assert (bl % kai_bl_multiple_of) == 0 + + num_rows = kai_roundup(n, nr) // nr + + return ( + num_rows + * kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0( + k, nr, kr, sr, bl + ) + ) + + def kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0( + k, nr, kr, sr, bl + ): + assert (bl % kr) == 0 + assert (nr % kai_nr_multiple_of) == 0 + assert (bl % kai_bl_multiple_of) == 0 + + # kr and sr are unused in the calculation + num_bytes_multiplier_rhs = kai_get_bf16_datatype_size_in_bytes() + num_blocks_per_row = kai_num_blocks_per_row(k, bl) + num_bytes_per_block = kai_num_bytes_per_block( + bl, num_bytes_multiplier_rhs + ) + + return nr * ( + (num_bytes_per_block * num_blocks_per_row) + + kai_num_bytes_sum_rhs + + kai_num_bytes_bias + ) + + # This function returns size of these datatypes stored as enum. We modify it to just return bf16 datatype + # https://gitlab.arm.com/kleidi/kleidiai/-/blob/main/kai/kai_common.h?ref_type=heads#L55 + def kai_get_bf16_datatype_size_in_bytes(): + return 2 # 2 bytes + + def kai_num_blocks_per_row(k, bl): + assert (bl % kai_bl_multiple_of) == 0 + return kai_roundup(k, bl) // bl + + def kai_num_bytes_per_block(bl, num_bytes_multiplier_rhs): + assert (bl % kai_bl_multiple_of) == 0 + return (bl // 2) + num_bytes_multiplier_rhs + + return kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0( + N, K, kai_nr, kai_kr, kai_sr, groupsize + ) + + +@register_meta([aten._dyn_quant_pack_4bit_weight]) +def meta__dyn_quant_pack_4bit_weight( + weights, scales_zeros, bias: Optional[Tensor], block_size, in_features, out_features +): + torch._check( + weights.dtype is torch.uint8, + lambda: f"expected w to be uint8, got {weights.dtype}", + ) + if torch.backends.kleidiai.is_available() and ( + (block_size == in_features and scales_zeros.dtype == torch.float) + or ( + block_size < in_features + and block_size % 32 == 0 + and in_features % block_size == 0 + and scales_zeros.dtype == torch.bfloat16 + ) + ): + packed_weight_size = get_kai_packed_weight_size( + 4, out_features, in_features, block_size + ) + return weights.new_empty(int(packed_weight_size), dtype=torch.uint8) + packed_weight_size = weights.numel() + scales_zeros.numel() + return weights.new_empty(packed_weight_size, dtype=torch.float) + + +@register_meta([aten._dyn_quant_matmul_4bit]) +def meta__dyn_quant_matmul_4bit( + inp, + packed_weights, + block_size, + in_features, + out_features, +): + torch._check(inp.dim() == 2, lambda: "input must be a 2D tensor") + torch._check( + inp.dtype in [torch.float32], + lambda: f"expected input to be f32, got {inp.dtype}", + ) + M = inp.size(0) + return inp.new_empty(M, out_features, dtype=inp.dtype) + + +@register_meta([aten._weight_int8pack_mm]) +def meta__weight_int8pack_mm(x, w, q_scales): + torch._check(x.dim() == 2, lambda: "x must be a 2D tensor") + torch._check( + x.dtype in [torch.float32, torch.float16, torch.bfloat16], + lambda: f"expected x to be f32/f16/bf16, got {x.dtype}", + ) + torch._check(w.dim() == 2, lambda: "w must be a 2D tensor") + torch._check( + w.dtype is torch.int8, + lambda: f"expected w to be int8, got {w.dtype}", + ) + return x.new_empty(x.size(0), w.size(0), dtype=x.dtype) + + +@register_meta(aten._cdist_forward.default) +def meta_cdist_forward(x1, x2, p, compute_mode): + torch._check( + x1.dim() >= 2, + lambda: f"cdist only supports at least 2D tensors, X1 got: {x1.dim()}D", + ) + torch._check( + x2.dim() >= 2, + lambda: f"cdist only supports at least 2D tensors, X2 got: {x2.dim()}D", + ) + torch._check( + x1.size(-1) == x2.size(-1), + lambda: f"X1 and X2 must have the same number of columns. X1: {x1.size(-1)} X2: {x2.size(-1)}", + ) + torch._check( + utils.is_float_dtype(x1.dtype), + lambda: "cdist only supports floating-point dtypes, X1 got: {x1.dtype}", + ) + torch._check( + utils.is_float_dtype(x2.dtype), + lambda: "cdist only supports floating-point dtypes, X2 got: {x2.dtype}", + ) + torch._check(p >= 0, lambda: "cdist only supports non-negative p values") + torch._check( + compute_mode in (None, 1, 2), + lambda: f"possible modes: None, 1, 2, but was: {compute_mode}", + ) + r1 = x1.size(-2) + r2 = x2.size(-2) + batch_tensor1 = x1.shape[:-2] + batch_tensor2 = x2.shape[:-2] + output_shape = list(torch.broadcast_shapes(batch_tensor1, batch_tensor2)) + output_shape.extend([r1, r2]) + return x1.new_empty(output_shape) + + +@register_meta(aten._cdist_backward) +@out_wrapper() +def meta_cdist_backward(grad, x1, x2, p, cdist): + c1 = x1.shape[-1] + r1 = x1.shape[-2] + r2 = x2.shape[-2] + batch_tensor1 = x1.shape[:-2] + batch_tensor2 = x2.shape[:-2] + expand_batch_portion = list(torch.broadcast_shapes(batch_tensor1, batch_tensor2)) + tensor1_expand_size = expand_batch_portion.copy() + tensor1_expand_size.extend([r1, c1]) + batch_product = math.prod(expand_batch_portion) + if r1 == 0 or r2 == 0 or c1 == 0 or batch_product == 0: + return torch.zeros_like(x1) + if tensor1_expand_size != list(x1.shape): + x1 = x1.expand(tensor1_expand_size) + return torch.empty_like(x1, memory_format=torch.contiguous_format) + + +# NB: This meta function accepts non-meta arguments! When this behavior +# was originally introduced this was accidental, but it is now load bearing +# as people are using this so that they can conveniently test code involving +# embeddings (feeding CPU tensor inputs with meta device EmbeddingBag module) +@register_meta(aten._embedding_bag.default) +def meta_embedding_bag( + weight, + indices, + offsets, + scale_grad_by_freq=False, + mode=0, + sparse=False, + per_sample_weights=None, + include_last_offset=False, + padding_idx=-1, +): + torch._check( + indices.dtype in (torch.long, torch.int), + lambda: f"expected indices to be long or int, got {indices.dtype}", + ) + torch._check( + offsets.dtype in (torch.long, torch.int), + lambda: f"expected offsets to be long or int, got {offsets.dtype}", + ) + torch._check( + utils.is_float_dtype(weight.dtype), + lambda: f"expected weight to be floating point type, got {weight.dtype}", + ) + + num_bags = offsets.size(0) + if include_last_offset: + torch._check( + num_bags >= 1, + lambda: "include_last_offset: numBags should be at least 1", + ) + num_bags -= 1 + + output = weight.new_empty(num_bags, weight.size(1)) + + if per_sample_weights is not None: + torch._check( + mode == MODE_SUM, + lambda: "embedding_bag: per_sample_weights only supported with mode='sum'", + ) + torch._check( + per_sample_weights.ndim == 1, + lambda: f"expected per_sample_weights to be 1D tensor, got {per_sample_weights.ndim}D", + ) + torch._check( + per_sample_weights.numel() == indices.numel(), + lambda: ( + f"expected per_sample_weights.numel() ({per_sample_weights.numel()} " + f"to be the same as indices.numel() ({indices.numel()})" + ), + ) + + def is_fast_path_index_select_scale(src, scale, output, padding_idx): + return ( + is_fast_path_index_select(src, output, padding_idx) and scale.stride(0) == 1 + ) + + def is_fast_path_index_select(src, output, padding_idx): + return ( + (src.dtype == torch.float or src.dtype == torch.half) + and src.stride(1) == 1 + and output.stride(1) == 1 + and padding_idx < 0 + ) + + def is_fast_path(src, scale, output, padding_idx): + if scale is not None: + return is_fast_path_index_select_scale(src, scale, output, padding_idx) + else: + return is_fast_path_index_select(src, output, padding_idx) + + if device_hint(offsets) != "cpu": + offset2bag = indices.new_empty(indices.size(0)) + bag_size = indices.new_empty(offsets.size()) + if mode == MODE_MAX: + max_indices = indices.new_empty(num_bags, weight.size(1)) + else: + max_indices = indices.new_empty(0) + else: + fast_path_sum = is_fast_path(weight, per_sample_weights, output, padding_idx) + if mode in (MODE_MEAN, MODE_MAX) or not fast_path_sum: + offset2bag = offsets.new_empty(indices.size(0)) + else: + offset2bag = offsets.new_empty(0) + bag_size = offsets.new_empty(num_bags) + # This part of the logic comes from make_max_indices_out in EmbeddingBag.cpp + numBags = offsets.shape[0] + if mode == MODE_MAX: + if include_last_offset: + torch._check( + numBags >= 1, + lambda: "include_last_offset: numBags should be at least 1", + ) + numBags -= 1 + max_indices = offsets.new_empty(numBags, weight.shape[1]) + else: + max_indices = offsets.new_empty(bag_size.size()) + return output, offset2bag, bag_size, max_indices + + +@register_meta(aten._embedding_bag_forward_only.default) +def meta_embedding_bag_forward_only(weight, indices, offsets, *args): + output, offset2bag, bag_size, max_indices = meta_embedding_bag( + weight, indices, offsets, *args + ) + if device_hint(offsets) == "cpu": + bag_size = offsets.new_empty(offsets.size()) + return output, offset2bag, bag_size, max_indices + + +def _get_reduction_dtype(input, dtype, promote_int_to_long=True): + # if specified, dtype takes precedence + if dtype: + return dtype + + if input.dtype.is_floating_point or input.dtype.is_complex: + return input.dtype + elif promote_int_to_long: + return torch.long + + return input.dtype + + +@register_meta([aten.nansum.default, aten.nansum.out]) +@out_wrapper() +def meta_nansum(input, dims=None, keepdim=False, *, dtype=None): + output_dtype = _get_reduction_dtype(input, dtype, promote_int_to_long=True) + dims = utils.reduction_dims(input.shape, dims) + output_shape = _compute_reduction_shape(input, dims, keepdim) + return input.new_empty(output_shape, dtype=output_dtype) + + +@register_meta([aten.median.default, aten.nanmedian.default]) +def meta_median(input): + output_shape = utils.compute_reduction_output_shape( + input.shape, tuple(range(input.dim())) + ) + return input.new_empty(output_shape) + + +@register_meta( + [ + aten.median.dim, + aten.median.dim_values, + aten.nanmedian.dim, + aten.nanmedian.dim_values, + aten.mode.default, + aten.mode.values, + ] +) +@out_wrapper("values", "indices") +def meta_median_mode_dim(input, dim=-1, keepdim=False): + if device_hint(input) == "cuda": + utils.alert_not_deterministic("median CUDA with indices output") + dim = utils.reduction_dims(input.shape, (dim,)) + output_shape = _compute_reduction_shape(input, dim, keepdim) + return ( + input.new_empty(output_shape), + input.new_empty(output_shape, dtype=torch.long), + ) + + +@register_meta(aten.logical_not_.default) +def meta_logical_not_(self): + return self + + +@register_meta(aten.repeat.default) +def meta_repeat(self, repeats): + torch._check( + len(repeats) >= self.dim(), + lambda: "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor", + ) + for i, rep in enumerate(repeats): + torch._check( + rep >= 0, + lambda: f"Repeats cannot be negative, found {rep} at index {i}", + ) + # Add new leading dimensions to the tensor if the + # number of target dimensions is larger than the + # number of source dimensions. + num_new_dimensions = len(repeats) - self.dim() + padded_size = (1,) * num_new_dimensions + tuple(self.shape) + target_size = [padded_size[i] * repeats[i] for i in range(len(repeats))] + return self.new_empty(target_size) + + +@register_meta(aten.zero_.default) +def meta_zero_(self): + return self + + +@register_meta( + [ + aten.mul_.Scalar, + aten.div_.Scalar, + aten.mul_.Tensor, + aten.div_.Tensor, + aten.logical_and_.default, + aten.logical_or_.default, + aten.logical_xor_.default, + ], +) +def meta_binop_inplace(self, other): + if isinstance(other, torch.Tensor): + check_inplace_broadcast(self.shape, other.shape) + return self + + +@register_meta( + [ + aten.add_.Scalar, + aten.sub_.Scalar, + aten.add_.Tensor, + aten.sub_.Tensor, + ], +) +def meta_binop_inplace_alpha(self, other, alpha=1): + """ + Some checks for inplace ops. + Checks for promotion rules for some dtypes. + int.add/sub_(float) and bool.add/sub_(others) are rejected. + Promoting in these in-place operations would require reallocating + and copying over elements, hence not allowed. + Checks for alpha param. + """ + + def is_integeric(arg): + if isinstance(arg, TensorLike): + return utils.is_integer_dtype(arg.dtype) + else: + return isinstance(arg, IntLike) + + def is_floatic(arg): + if isinstance(arg, TensorLike): + return utils.is_float_dtype(arg.dtype) + else: + return isinstance(arg, FloatLike) + + def is_booleanic(arg): + if isinstance(arg, TensorLike): + return utils.is_boolean_dtype(arg.dtype) + else: + return isinstance(arg, BoolLike) + + # Do not allow int+float->int in-place + if is_integeric(self) and is_floatic(other): + raise RuntimeError( + "Promotion of int.add/sub_(float) in in-place ops are not possible due to element size change." + ) + + # Do not allow bool+other->bool in-place + if is_booleanic(self) and not is_booleanic(other): + raise RuntimeError( + "Promotion of book.add/sub_(others) in in-place ops are not possible due to element size change." + ) + + if isinstance(other, torch.Tensor): + check_inplace_broadcast(self.shape, other.shape) + return self + + +@register_meta( + [ + aten.add.Scalar, + aten.sub.Scalar, + ], +) +def meta_binop_alpha(self, other, alpha=1): + return elementwise_meta( + self, other, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta([aten.round.default, aten.round.decimals]) +def meta_round(self, **kwargs): + return elementwise_meta( + self, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +def shift_dtype_check(fn_name, self, val): + torch._check( + utils.is_integer_dtype(self.dtype), + lambda: f"{fn_name}: Expected input tensor to have an integral dtype. Got {self.dtype}", + ) + if isinstance(val, torch.Tensor): + torch._check( + utils.is_integer_dtype(val.dtype), + lambda: f"{fn_name}: Expected shift value to have an integral dtype. Got {val.dtype}", + ) + else: + torch._check( + isinstance(val, IntLike), + lambda: f"{fn_name}: Expected shift value to be an int. Got {val}", + ) + + +@register_meta([aten.__rshift__.Tensor, aten.__rshift__.Scalar]) +def meta_rshifts(self, other): + shift_dtype_check("rshift", self, other) + return elementwise_meta( + self, other, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta([aten.__lshift__.Tensor, aten.__lshift__.Scalar]) +def meta_lshifts(self, other): + shift_dtype_check("lshift", self, other) + return elementwise_meta( + self, other, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta(aten.zero.default) +def meta_zero(self): + return self.new_empty(self.shape) + + +@register_meta([aten.fill_.Tensor, aten.fill_.Scalar]) +def meta_fill_(self, val): + return self + + +@register_meta([aten.fill.Tensor, aten.fill.Scalar]) +def meta_fill(self, val): + return torch.empty_like(self) + + +@register_meta(aten.relu_.default) +def meta_relu_(self): + return self + + +@register_meta(aten._add_relu.Tensor) +@out_wrapper() +def meta__add_relu(self, other, alpha=1) -> Tensor: + return elementwise_meta( + self, other, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta([aten.rrelu_with_noise]) +@out_wrapper() +def meta_rrelu_with_noise( + self, noise, lower=0.125, upper=0.3333333333333333, training=False, generator=None +): + return torch.empty_like(self) + + +@register_meta([aten.rrelu_with_noise_functional]) +def meta_rrelu_with_noise_functional( + self, noise, lower=0.125, upper=0.3333333333333333, training=False, generator=None +): + return torch.empty_like(self), torch.empty_like(noise) + + +@register_meta([aten.rrelu_with_noise_]) +def meta_rrelu_with_noise_( + self, lower=0.125, upper=0.3333333333333333, training=False, generator=None +): + return self + + +@register_meta([aten.index_put.default, aten._unsafe_index_put.default]) +def meta_index_put(self, indices, values, accumulate=False): + return torch.empty_like(self) + + +@register_meta(aten.masked_fill_.Scalar) +def meta_masked_fill_(self, mask, value): + check_inplace_broadcast(self.shape, mask.shape) + return self + + +@register_meta(aten._masked_scale.default) +def meta__masked_scale(self, mask, scale): + masked_scale = self.new_empty(self.size()).to( + memory_format=utils.suggest_memory_format(self) + ) + return masked_scale + + +@register_meta(aten.masked_scatter_) +def meta_masked_scatter_(self, mask, source): + torch._check( + mask.dtype in (torch.bool, torch.uint8), lambda: "Mask must be bool or uint8" + ) + torch._check( + self.dtype == source.dtype, + lambda: "masked_scatter: expected self and source to have same " + f"dtypes but got {self.dtype} and {source.dtype}", + ) + return self + + +@register_meta(aten.masked_scatter) +@out_wrapper() +def meta_masked_scatter(self, mask, source): + self, mask = _maybe_broadcast(self, mask) + output = torch.empty_like(self, memory_format=torch.contiguous_format) + return meta_masked_scatter_(output, mask, source) + + +@register_meta(aten.masked_scatter_backward) +def meta_masked_scatter_backward(self, mask, sizes): + return self.new_empty(sizes) + + +@register_meta(aten.index_put_.default) +def meta_index_put_(self, indices, values, accumulate=False): + return self + + +def common_meta_baddbmm_bmm(batch1, batch2, is_bmm, self_baddbmm=None, out_dtype=None): + torch._check(batch1.dim() == 3, lambda: "batch1 must be a 3D tensor") + torch._check(batch2.dim() == 3, lambda: "batch2 must be a 3D tensor") + + batch1_sizes = batch1.size() + batch2_sizes = batch2.size() + + bs = batch1_sizes[0] + contraction_size = batch1_sizes[2] + res_rows = batch1_sizes[1] + res_cols = batch2_sizes[2] + output_size = (bs, res_rows, res_cols) + + torch._check( + batch2_sizes[0] == bs and batch2_sizes[1] == contraction_size, + lambda: f"Expected size for first two dimensions of batch2 tensor to be: [{bs}" + f", {contraction_size}] but got: [{batch2_sizes[0]}, {batch2_sizes[1]}].", + ) + if out_dtype: + supported_out_dtype = ( + batch1.dtype == torch.float16 or batch1.dtype == torch.bfloat16 + ) and out_dtype == torch.float32 + torch._check( + out_dtype == batch1.dtype or supported_out_dtype, + lambda: "out_dtype only supported for torch.float32 output with float16/bfloat16 inputs or same as input dtypes", + ) + output = batch2.new_empty(output_size).to(out_dtype) + else: + # TODO: handle out + output = batch2.new_empty(output_size) + + if not is_bmm and self_baddbmm is not None: + torch._check(self_baddbmm.dim() == 3, lambda: "self must be a 3D tensor") + torch._check( + self_baddbmm.size() == output_size, + lambda: f"Expected an input tensor shape with shape {output_size} but got shape: {self_baddbmm.size()}", + ) + + return output + + +@register_meta(aten.bmm.default) +def meta_bmm(self, mat2): + return common_meta_baddbmm_bmm(self, mat2, True) + + +@register_meta(aten.bmm.dtype) +def meta_bmm_dtype(self, mat2, out_dtype): + return common_meta_baddbmm_bmm(self, mat2, True, out_dtype=out_dtype) + + +def div_rtn(x, y): + q = x // y + r = x % y + # WARNING: explicit bool conversion here is necessary; + # would be fixed by SymBool + if r != 0 and (bool(r < 0) != bool(y < 0)): + q -= 1 + return q + + +def pooling_output_shape_pad_lr( + inputSize, + kernelSize, + pad_l, + pad_r, + stride, + dilation, + ceil_mode, +): + outputSize = ( + div_rtn( + inputSize + + pad_l + + pad_r + - dilation * (kernelSize - 1) + - 1 + + (stride - 1 if ceil_mode else 0), + stride, + ) + + 1 + ) + if ceil_mode: + if (outputSize - 1) * stride >= inputSize + pad_l: + outputSize -= 1 + return outputSize + + +def pooling_output_shape(inputSize, kernelSize, pad, stride, dilation, ceil_mode): + torch._check(stride != 0, lambda: "stride should not be zero") + torch._check(pad >= 0, lambda: f"pad must be non-negative, but got pad: {pad}") + torch._check( + pad <= ((kernelSize - 1) * dilation + 1) // 2, + lambda: ( + f"pad should be at most half of effective kernel size, but got pad={pad}, " + f"kernel_size={kernelSize} and dilation={dilation}" + ), + ) + return pooling_output_shape_pad_lr( + inputSize, kernelSize, pad, pad, stride, dilation, ceil_mode + ) + + +def pool2d_shape_check( + input, + kH, + kW, + dH, + dW, + padH, + padW, + dilationH, + dilationW, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + memory_format, +): + ndim = input.dim() + nOutputPlane = nInputPlane + + torch._check( + kW > 0 and kH > 0, + lambda: "kernel size should be greater than zero, but got kH: {kH}, kW: {kW}", + ) + torch._check( + dW > 0 and dH > 0, + lambda: "stride should be greater than zero, but got dH: {dH}, dW: {dW}", + ) + torch._check( + dilationH > 0 and dilationW > 0, + lambda: "dilation should be greater than zero, but got dilationH: {dilationH}, dilationW: {dilationW}", + ) + + valid_dims = input.size(1) != 0 and input.size(2) != 0 + + if memory_format == torch.channels_last: + torch._check( + ndim == 4 and valid_dims and input.size(3) != 0, + lambda: "Expected 4D (batch mode) tensor expected for input with channels_last layout" + " with optional 0 dim batch size for input, but got: {input.size()}", + ) + else: + torch._check( + (ndim == 3 and input.size(0) != 0 and valid_dims) + or (ndim == 4 and valid_dims and input.size(3) != 0), + lambda: f"Expected 3D or 4D (batch mode) tensor with optional 0 dim batch size for input, but got: {input.size()}", + ) + + torch._check( + kW // 2 >= padW and kH // 2 >= padH, + lambda: "pad should be smaller than or equal to half of kernel size, but got " + f"padW = {padW}, padH = {padH}, kW = {kW}, kH = {kH}", + ) + + torch._check( + outputWidth >= 1 and outputHeight >= 1, + lambda: f"Given input size: ({nInputPlane}x{inputHeight}x{inputWidth}). " + f"Calculated output size: ({nOutputPlane}x{outputHeight}x{outputWidth}). " + "Output size is too small", + ) + + +def pool3d_shape_check( + input: Tensor, + nslices: int, + kT: int, + kH: int, + kW: int, + dT: int, + dH: int, + dW: int, + pT: int, + pH: int, + pW: int, + dilationT: int, + dilationH: int, + dilationW: int, + itime: int, + iheight: int, + iwidth: int, + otime: int, + oheight: int, + owidth: int, + fn_name: str, + check_input_size: bool = False, +): + ndim = input.ndim + + torch._check( + kT > 0 and kW > 0 and kH > 0, + lambda: ( + f"kernel size should be greater than zero, but got " + f"kT: {kT}, kH: {kH}, kW: {kW}" + ), + ) + torch._check( + dT > 0 and dW > 0 and dH > 0, + lambda: ( + f"stride should be greater than zero, but got dT: {dT}, dH: {dH}, dW: {dW}" + ), + ) + torch._check( + dilationT > 0 and dilationW > 0 and dilationH > 0, + lambda: ( + f"dilation should be greater than zero, but got " + f"dilationT: {dilationT}, dilationH: {dilationH}, dilationW: {dilationW}" + ), + ) + + torch._check( + ndim in (4, 5), + lambda: f"{fn_name}: Expected 4D or 5D tensor for input, but got: {input.shape}", + ) + + for i in range(ndim): + if ndim == 5 and i == 0: + # size of batch-dim can be 0. + continue + torch._check( + input.size(i) > 0, + lambda: ( + f"{fn_name}: Expected input's non-batch dimensions to have positive length," + f" but input has a shape of {input.shape}" + f" and non-batch dimension {input.size(i)} has length zero!" + ), + ) + + if check_input_size: # AveragePool3d + torch._check( + itime >= kT and iheight >= kH and iwidth >= kW, + lambda: ( + f"input image (T: {itime} H: {iheight} W: {iwidth}) smaller than " + f"kernel size (kT: {kT} kH: {kH} kW: {kW})" + ), + ) + + torch._check( + kT / 2 >= pT and kW / 2 >= pW and kH / 2 >= pH, + lambda: ( + f"pad should be smaller than or equal to half of kernel size, but got " + f"kT: {kT} kW: {kW} kH: {kH} padT: {pT} padW: {pW} padH: {pH}" + ), + ) + + torch._check( + otime >= 1 and owidth >= 1 and oheight >= 1, + lambda: ( + f"Given input size: ({nslices}x{itime}x{iheight}x{iwidth}). " + f"Calculated output size: ({nslices}x{otime}x{oheight}x{owidth}). " + f"Output size is too small" + ), + ) + + +def max_pool3d_backward_shape_check( + input, + grad_output, + indices, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + pT, + pH, + pW, + dilationT, + dilationH, + dilationW, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + fn_name, +): + ndim = input.ndim + + pool3d_shape_check( + input, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + pT, + pH, + pW, + dilationT, + dilationH, + dilationW, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + fn_name, + ) + + check_dim_size(grad_output, ndim, ndim - 4, nslices) + check_dim_size(grad_output, ndim, ndim - 3, otime) + check_dim_size(grad_output, ndim, ndim - 2, oheight) + check_dim_size(grad_output, ndim, ndim - 1, owidth) + + check_dim_size(indices, ndim, ndim - 4, nslices) + check_dim_size(indices, ndim, ndim - 3, otime) + check_dim_size(indices, ndim, ndim - 2, oheight) + check_dim_size(indices, ndim, ndim - 1, owidth) + + +def avg_pool3d_backward_shape_check( + input: Tensor, + grad_output: Tensor, + nslices: int, + kT: int, + kH: int, + kW: int, + dT: int, + dH: int, + dW: int, + pT: int, + pH: int, + pW: int, + itime: int, + iheight: int, + iwidth: int, + otime: int, + oheight: int, + owidth: int, + fn_name: str, +): + ndim = input.ndim + + pool3d_shape_check( + input, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + pT, + pH, + pW, + 1, + 1, + 1, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + fn_name, + True, + ) + + check_dim_size(grad_output, ndim, ndim - 4, nslices) + check_dim_size(grad_output, ndim, ndim - 3, otime) + check_dim_size(grad_output, ndim, ndim - 2, oheight) + check_dim_size(grad_output, ndim, ndim - 1, owidth) + + +def max_pool2d_checks_and_compute_shape( + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, +): + # Reference: aten/src/ATen/native/DilatedMaxPool2d.cpp + def unpack(name, val): + torch._check( + len(val) in [1, 2], + lambda: f"max_pool2d: {name} must either be a single int, or a tuple of two ints", + ) + H = val[0] + W = H if len(val) == 1 else val[1] + return H, W + + kH, kW = unpack("kernel_size", kernel_size) + + torch._check( + len(stride) in [0, 1, 2], + lambda: "max_pool2d: stride must either be omitted, a single int, or a tuple of two ints", + ) + if len(stride) == 0: + dH, dW = kH, kW + else: + dH, dW = unpack("stride", stride) + + padH, padW = unpack("padding", padding) + dilationH, dilationW = unpack("dilation", dilation) + nInputPlane = input.size(-3) + inputHeight = input.size(-2) + inputWidth = input.size(-1) + + memory_format = utils.suggest_memory_format(input) + if memory_format == torch.channels_last: + torch._check( + input.dim() == 4, + lambda: "non-empty 4D (batch mode) tensor expected for input with channels_last layout", + ) + elif memory_format == torch.contiguous_format: + torch._check( + input.dim() in [3, 4], + lambda: "non-empty 3D or 4D (batch mode) tensor expected for input", + ) + else: + torch._check( + False, + lambda: "Unsupported memory format. Supports only ChannelsLast, Contiguous", + ) + + outputHeight = pooling_output_shape(inputHeight, kH, padH, dH, dilationH, ceil_mode) + outputWidth = pooling_output_shape(inputWidth, kW, padW, dW, dilationW, ceil_mode) + + pool2d_shape_check( + input, + kH, + kW, + dH, + dW, + padH, + padW, + dilationH, + dilationW, + nInputPlane, + inputHeight, + inputWidth, + outputHeight, + outputWidth, + memory_format, + ) + + return nInputPlane, outputHeight, outputWidth + + +@register_meta(aten.max_pool2d_with_indices_backward.default) +def meta_max_pool2d_with_indices_backward( + grad_output, + self, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices, +): + ( + nInputPlane, + outputHeight, + outputWidth, + ) = max_pool2d_checks_and_compute_shape( + self, kernel_size, stride, padding, dilation, ceil_mode + ) + + torch._check( + self.dtype == grad_output.dtype, + lambda: f"Expected dtype {self.dtype} for `gradOutput` but got dtype {grad_output.dtype}", + ) + + nOutputPlane = nInputPlane + ndim = self.ndim + + def _check_dim_size(t): + check_dim_size(t, ndim, ndim - 3, nOutputPlane) + check_dim_size(t, ndim, ndim - 2, outputHeight) + check_dim_size(t, ndim, ndim - 1, outputWidth) + + _check_dim_size(grad_output) + _check_dim_size(indices) + + memory_format = utils.suggest_memory_format(self) + return torch.empty( + self.shape, + dtype=self.dtype, + device=self.device, + memory_format=memory_format, + ) + + +@register_meta(aten.max_pool2d_with_indices.default) +def meta_max_pool2d_with_indices( + input, + kernel_size, + stride=(), + padding=(0,), + dilation=(1,), + ceil_mode=False, +): + ( + nInputPlane, + outputHeight, + outputWidth, + ) = max_pool2d_checks_and_compute_shape( + input, kernel_size, stride, padding, dilation, ceil_mode + ) + + nbatch = input.size(-4) if input.dim() == 4 else 1 + memory_format = utils.suggest_memory_format(input) + if input.dim() == 3: + size = [nInputPlane, outputHeight, outputWidth] + else: + size = [nbatch, nInputPlane, outputHeight, outputWidth] + return ( + torch.empty( + size, + dtype=input.dtype, + device=input.device, + memory_format=memory_format, + ), + torch.empty( + size, + dtype=torch.int64, + device=input.device, + memory_format=memory_format, + ), + ) + + +@register_meta(aten.fractional_max_pool2d.default) +def meta_fractional_max_pool2d(self, kernel_size, output_size, random_samples): + torch._check( + self.ndim in (3, 4), + lambda: f"fractional_max_pool2d: Expected 3D or 4D tensor, but got: {self.ndim}", + ) + ndim = self.ndim + + for d in range(ndim - 3, ndim): + torch._check( + self.size(d) > 0, + f"fractional_max_pool2d: Expected input to have non-zero " + f" size for non-batch dimensions, but got {self.size()} with dimension {d} empty", + ) + + # the check and message are out of sync, but this matches the structured meta + torch._check( + len(kernel_size) == 2, + lambda: "fractional_max_pool2d: kernel_size must" + "either be a single int or tuple of Ints", + ) + torch._check( + len(output_size) == 2, + lambda: "fractional_max_pool2d: output_size must " + "either be a single int or tuple of Ints", + ) + + input_channels = self.size(-3) + input_height = self.size(-2) + input_width = self.size(-1) + if ndim == 4: + input_batch = self.size(0) + else: + input_batch = 1 + + torch._check( + self.dtype == random_samples.dtype, + lambda: "Expect _random_samples to have the same dtype as input", + ) + torch._check( + random_samples.ndim == 3, + lambda: f"Expect _random samples to have 3 dimensions got, {random_samples.ndim}", + ) + + n = random_samples.size(0) + c = random_samples.size(1) + d = random_samples.size(2) + torch._check( + n >= input_batch, + "Expect _random_samples.size(0) no less then input batch size.", + ) + torch._check( + c == input_channels, + lambda: "Expect _random_samples.size(1) equals to input channel size.", + ) + torch._check(d == 2, lambda: f"Expect _random_samples.size(2) equals to 2 got {d}.") + + torch._check( + output_size[0] + kernel_size[0] - 1 <= input_height, + lambda: f"fractional_max_pool2d: kernel height {kernel_size[0]} is too large relative to input height {input_height}", + ) + torch._check( + output_size[1] + kernel_size[1] - 1 <= input_width, + lambda: f"fractional_max_pool2d: kernel width {kernel_size[1]} is too large relative to input width {input_width}", + ) + + if self.dim() == 4: + size = [input_batch, input_channels, output_size[0], output_size[1]] + else: + size = [input_channels, output_size[0], output_size[1]] + + return ( + torch.empty( + size, + dtype=self.dtype, + device=self.device, + ), + torch.empty( + size, + dtype=torch.int64, + device=self.device, + ), + ) + + +@register_meta(aten.max_pool3d_with_indices) +@out_wrapper("out", "indices") +def meta_max_pool3d_with_indices( + input, + kernel_size, + stride=(), + padding=(0,), + dilation=(1,), + ceil_mode=False, +): + torch._check( + len(kernel_size) in (1, 3), + lambda: "max_pool3d: kernel_size must either be a single int, or a tuple of three ints", + ) + kT = kernel_size[0] + kH = kT if len(kernel_size) == 1 else kernel_size[1] + kW = kT if len(kernel_size) == 1 else kernel_size[2] + + torch._check( + not stride or len(stride) in (1, 3), + lambda: "max_pool3d: stride must either be omitted, a single int, or a tuple of three ints", + ) + dT = kT if not stride else stride[0] + dH = kH if not stride else (dT if len(stride) == 1 else stride[1]) + dW = kW if not stride else (dT if len(stride) == 1 else stride[2]) + + torch._check( + len(padding) in (1, 3), + lambda: "max_pool3d: padding must either be a single int, or a tuple of three ints", + ) + pT = padding[0] + pH = pT if len(padding) == 1 else padding[1] + pW = pT if len(padding) == 1 else padding[2] + + torch._check( + len(dilation) in (1, 3), + lambda: "max_pool3d: dilation must be either a single int, or a tuple of three ints", + ) + dilationT = dilation[0] + dilationH = dilationT if len(dilation) == 1 else dilation[1] + dilationW = dilationT if len(dilation) == 1 else dilation[2] + + torch._check( + input.ndim in (4, 5), + lambda: "non-empty 4D or 5D (batch mode) tensor expected for input", + ) + + nbatch = input.size(-5) if input.ndim == 5 else 1 + nslices = input.size(-4) + itime = input.size(-3) + iheight = input.size(-2) + iwidth = input.size(-1) + + otime = pooling_output_shape(itime, kT, pT, dT, dilationT, ceil_mode) + oheight = pooling_output_shape(iheight, kH, pH, dH, dilationH, ceil_mode) + owidth = pooling_output_shape(iwidth, kW, pW, dW, dilationW, ceil_mode) + + pool3d_shape_check( + input, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + pT, + pH, + pW, + dilationT, + dilationH, + dilationW, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + "max_pool3d_with_indices()", + ) + + channels_last = ( + input.ndim == 5 and utils.suggest_memory_format(input) == torch.channels_last_3d + ) + if input.ndim == 4: + input_channels_last_check = input.unsqueeze(0) + channels_last = ( + not input_channels_last_check.is_contiguous() + ) and input_channels_last_check.is_contiguous( + memory_format=torch.channels_last_3d + ) + out_shape = (nslices, otime, oheight, owidth) + else: + out_shape = (nbatch, nslices, otime, oheight, owidth) # type: ignore[assignment] + + out = input.new_empty(out_shape) + indices = input.new_empty(out_shape, dtype=torch.int64) + + if channels_last: + out = out.to(memory_format=torch.channels_last_3d) + indices = indices.to(memory_format=torch.channels_last_3d) + + return out, indices + + +@register_meta(aten.max_pool3d_with_indices_backward) +@out_wrapper("grad_input") +def meta_max_pool3d_with_indices_backward( + grad_output, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices, +): + torch._check( + len(kernel_size) in (1, 3), + lambda: "max_pool3d: kernel_size must either be a single int, or a tuple of three ints", + ) + kT = kernel_size[0] + kH = kT if len(kernel_size) == 1 else kernel_size[1] + kW = kT if len(kernel_size) == 1 else kernel_size[2] + + torch._check( + not stride or len(stride) in (1, 3), + lambda: "max_pool3d: stride must either be omitted, a single int, or a tuple of three ints", + ) + dT = kT if not stride else stride[0] + dH = kH if not stride else (dT if len(stride) == 1 else stride[1]) + dW = kW if not stride else (dT if len(stride) == 1 else stride[2]) + + torch._check( + len(padding) in (1, 3), + lambda: "max_pool3d: padding must either be a single int, or a tuple of three ints", + ) + pT = padding[0] + pH = pT if len(padding) == 1 else padding[1] + pW = pT if len(padding) == 1 else padding[2] + + torch._check( + len(dilation) in (1, 3), + lambda: "max_pool3d: dilation must be either a single int, or a tuple of three ints", + ) + dilationT = dilation[0] + dilationH = dilationT if len(dilation) == 1 else dilation[1] + dilationW = dilationT if len(dilation) == 1 else dilation[2] + + torch._check( + input.ndim in (4, 5), + lambda: "non-empty 4D or 5D (batch mode) tensor expected for input", + ) + + nslices = input.size(-4) + itime = input.size(-3) + iheight = input.size(-2) + iwidth = input.size(-1) + + otime = grad_output.size(-3) + oheight = grad_output.size(-2) + owidth = grad_output.size(-1) + + max_pool3d_backward_shape_check( + input, + grad_output, + indices, + nslices, + kT, + kH, + kW, + dT, + dH, + dW, + pT, + pH, + pW, + dilationT, + dilationH, + dilationW, + itime, + iheight, + iwidth, + otime, + oheight, + owidth, + "max_pool3d_with_indices_backward()", + ) + + channels_last = ( + input.ndim == 5 and utils.suggest_memory_format(input) == torch.channels_last_3d + ) + if input.ndim == 4: + input_channels_last_check = input.unsqueeze(0) + channels_last = ( + not input_channels_last_check.is_contiguous() + ) and input_channels_last_check.is_contiguous( + memory_format=torch.channels_last_3d + ) + + grad_input = input.new_empty(input.shape) + + if channels_last: + grad_input = grad_input.to(memory_format=torch.channels_last_3d) + + return grad_input + + +def check_grid_sampler_common(input: Tensor, grid: Tensor): + torch._check( + input.device == grid.device, + lambda: ( + f"grid_sampler(): expected input and grid to be on same device, but input " + f"is on {input.device} and grid is on {grid.device}" + ), + ) + torch._check( + input.layout == torch.strided and grid.layout == torch.strided, + lambda: ( + f"grid_sampler(): expected input and grid to have torch.strided layout, but " + f"input has {input.layout} and grid has {grid.layout}" + ), + ) + torch._check( + input.shape[0] == grid.shape[0], + lambda: ( + f"grid_sampler(): expected grid and input to have same batch size, but got " + f"input with sizes {input.shape} and grid with sizes {grid.shape}" + ), + ) + torch._check( + grid.shape[-1] == input.ndim - 2, + lambda: ( + f"grid_sampler(): expected grid to have size {input.ndim - 2} in last " + f"dimension, but got grid with sizes {grid.shape}" + ), + ) + + for i in range(2, input.ndim): + torch._check( + input.shape[i] > 0, + lambda: ( + f"grid_sampler(): expected input to have non-empty spatial dimensions, " + f"but input has sizes {input.shape} with dimension {i} being empty" + ), + ) + + +class GridSamplerInterpolation(Enum): + BILINEAR = 0 + NEAREST = 1 + BICUBIC = 2 + + +def check_grid_sampler_3d(input: Tensor, grid: Tensor, interpolation_mode: int): + torch._check( + input.ndim == 5 and input.ndim == grid.ndim, + lambda: ( + f"grid_sampler(): expected 5D input and grid with same number of " + f"dimensions, but got input with sizes {input.shape}" + f" and grid with sizes {grid.shape}" + ), + ) + torch._check( + not ( + input.ndim == 5 + and interpolation_mode == GridSamplerInterpolation.BICUBIC.value + ), + lambda: "grid_sampler(): bicubic interpolation only supports 4D input", + ) + + +@register_meta(aten.grid_sampler_2d_backward.default) +def grid_sampler_2d_backward_meta( + grad_output, + input, + grid, + interpolation_mode, + padding_mode, + align_corners, + output_mask, +): + input_requires_grad = output_mask[0] + if input_requires_grad: + grad_input = torch.zeros_like(input, memory_format=torch.contiguous_format) + else: + grad_input = None + grad_grid = torch.empty_like(grid, memory_format=torch.contiguous_format) + return (grad_input, grad_grid) + + +@register_meta(aten.grid_sampler_3d) +@out_wrapper() +def grid_sampler_3d( + input, + grid, + interpolation_mode, + padding_mode, + align_corners, +): + check_grid_sampler_common(input, grid) + check_grid_sampler_3d(input, grid, interpolation_mode) + N = input.shape[0] + C = input.shape[1] + out_D = grid.shape[1] + out_H = grid.shape[2] + out_W = grid.shape[3] + return input.new_empty((N, C, out_D, out_H, out_W)) + + +@register_meta(aten.grid_sampler_3d_backward) +@out_wrapper("grad_input", "grad_grid") +def grid_sampler_3d_backward( + grad_output, + input, + grid, + interpolation_mode, + padding_mode, + align_corners, + output_mask, +): + check_grid_sampler_common(input, grid) + check_grid_sampler_3d(input, grid, interpolation_mode) + input_requires_grad = output_mask[0] + if input_requires_grad: + grad_input = torch.zeros_like( + input, memory_format=torch.legacy_contiguous_format + ) + else: + grad_input = None + grad_grid = torch.empty_like(grid, memory_format=torch.legacy_contiguous_format) + return grad_input, grad_grid + + +@register_meta([aten.full.default]) +def full(size, fill_value, *args, **kwargs): + dtype = kwargs.get("dtype", None) + if not dtype: + dtype = utils.get_dtype(fill_value) + kwargs["dtype"] = dtype + return torch.empty(size, *args, **kwargs) + + +# zeros_like is special cased to work for sparse +@register_meta(aten.zeros_like.default) +def zeros_like( + self, + dtype=None, + layout=None, + device=None, + pin_memory=None, + memory_format=None, +): + if layout == torch.sparse_coo: + torch._check( + memory_format is None, + lambda: "memory format option is only supported by strided tensors", + ) + + res = torch.empty( + 0, + dtype=self.dtype if dtype is None else dtype, + layout=layout, + device=self.device if device is None else device, + pin_memory=pin_memory, + ) + + if self.is_sparse: + res.sparse_resize_and_clear_( + self.size(), self.sparse_dim(), self.dense_dim() + ) + else: + res.sparse_resize_and_clear_(self.size(), self.dim(), 0) + + res._coalesced_(True) + return res + res = aten.empty_like.default( + self, + dtype=dtype, + layout=layout, + device=device, + pin_memory=pin_memory, + memory_format=memory_format, + ) + # device can be not "meta" + res.fill_(0) + return res + + +@register_meta([aten.ones.default, aten.ones.out]) +@out_wrapper() +def meta_ones( + size, + *, + dtype=None, + layout=None, + device=None, + pin_memory=None, + requires_grad=False, +): + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.get_default_device() + if layout is None: + layout = torch.strided + return torch.empty( + size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta([aten.zeros.default, aten.zeros.out]) +@out_wrapper() +def meta_zeros( + size, + *, + dtype=None, + layout=None, + device=None, + pin_memory=None, + requires_grad=False, +): + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.get_default_device() + if layout is None: + layout = torch.strided + return torch.empty( + size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta(aten.select_scatter.default) +def meta_select_scatter(self, src, dim, index): + return utils.clone_preserve_strides(self) + + +@register_meta(aten.slice_scatter.default) +def meta_slice_scatter(self, src, dim=0, start=None, end=None, step=1): + return utils.clone_preserve_strides(self) + + +# TODO: Deduplicate this with canonicalize_dim +def maybe_wrap_dim(dim: int, dim_post_expr: int, wrap_scalar: bool = True): + if dim_post_expr <= 0: + assert wrap_scalar + dim_post_expr = 1 + min = -dim_post_expr + max = dim_post_expr - 1 + assert not (dim < min or dim > max), f"dim {dim} out of bounds ({min}, {max})" + if dim < 0: + dim += dim_post_expr + return dim + + +def ensure_nonempty_size(t, dim): + return 1 if t.dim() == 0 else t.shape[dim] + + +# From aten/src/ATen/native/ScatterGatherChecks.h +def gather_shape_check(self, dim, index): + self_dims = max(self.dim(), 1) + index_dims = max(index.dim(), 1) + torch._check( + self_dims == index_dims, + lambda: "Index tensor must have the same number of dimensions as input tensor", + ) + for i in range(self_dims): + if i != dim: + torch._check( + ensure_nonempty_size(index, i) <= ensure_nonempty_size(self, i), + lambda: f"Size does not match at dimension {i} expected index {index.shape}" + + f" to be no larger than self {self.shape} apart from dimension {dim}", + ) + + +@register_meta(aten.gather.default) +def meta_gather(self, dim, index, sparse_grad=False): + from torch.fx.experimental.symbolic_shapes import guard_or_false + + wrapped_dim = maybe_wrap_dim(dim, self.dim()) + is_index_empty = guard_or_false(index.numel() == 0) + if not is_index_empty: + torch._check( + index.dtype == torch.long or index.dtype == torch.int, + lambda: f"gather(): Expected dtype int32/int64 for index, but got {index.dtype}", + ) + gather_shape_check(self, wrapped_dim, index) + return self.new_empty(index.shape) + + +# From aten/src/ATen/native/TensorAdvancedIndexing.cpp +def get_operator_enum(reduce_, use_new_options=False): + if use_new_options: + if reduce_ == "sum": + return "REDUCE_ADD" + elif reduce_ == "prod": + return "REDUCE_MULTIPLY" + elif reduce_ == "mean": + return "REDUCE_MEAN" + elif reduce_ == "amax": + return "REDUCE_MAXIMUM" + elif reduce_ == "amin": + return "REDUCE_MINIMUM" + torch._check( + False, + lambda: "reduce argument must be either sum, prod, mean, amax or amin.", + ) + return + else: + if reduce_ == "add": + return "REDUCE_ADD" + elif reduce_ == "multiply": + return "REDUCE_MULTIPLY" + torch._check(False, lambda: "reduce argument must be either add or multiply.") + return + + +# From aten/src/ATen/native/ScatterGatherChecks.h +def scatter_gather_dtype_check(method_name, self, index, src_opt=None): + from torch.fx.experimental.symbolic_shapes import guard_or_true + + if guard_or_true(index.numel() != 0): + torch._check( + index.dtype == torch.long or index.dtype == torch.int, + lambda: f"{method_name}(): Expected dtype int32/int64 for index", + ) + + if src_opt is not None: + torch._check( + self.dtype == src_opt.dtype, + lambda: f"{method_name}(): Expected self.dtype to be equal to src.dtype", + ) + + +def ensure_nonempty_dim(dim): + return max(dim, 1) + + +# From aten/src/ATen/native/ScatterGatherChecks.h +def scatter_shape_check(self, dim, index, src_opt=None): + from torch.fx.experimental.symbolic_shapes import guard_or_false + + if guard_or_false(index.numel() == 0): + return + torch._check( + ensure_nonempty_dim(self.dim()) == ensure_nonempty_dim(index.dim()), + lambda: "Index tensor must have the same number of dimensions as self tensor", + ) + + is_wrong_shape = False + self_dims = ensure_nonempty_dim(self.dim()) + + # Check: index.size(d) <= self.size(d) for all d != dim + for d in range(self_dims): + index_d_size = ensure_nonempty_size(index, d) + if d == dim: + continue + if index_d_size > ensure_nonempty_size(self, d): + is_wrong_shape = True + break + + # Check: index.size(d) <= src.size(d) for all d if src is Tensor + if not is_wrong_shape and src_opt is not None: + for d in range(self_dims): + index_d_size = ensure_nonempty_size(index, d) + if index_d_size > ensure_nonempty_size(src_opt, d): + is_wrong_shape = True + break + + if src_opt is not None: + torch._check( + ensure_nonempty_dim(self.dim()) == ensure_nonempty_dim(index.dim()), + lambda: "Index tensor must have the same number of dimensions as self tensor", + ) + torch._check( + not is_wrong_shape, + lambda: f"Expected index {index.shape} to be no larger than self {self.shape}" + + f" apart from dimension {dim} and to be no larger than src {src_opt.shape}", + ) + else: + torch._check( + not is_wrong_shape, + lambda: f"Expected index {index.shape} to be no larger than self {self.shape}" + + f" apart from dimension {dim}", + ) + + +# From aten/src/ATen/native/TensorAdvancedIndexing.cpp +def scatter_meta_impl(self, dim, index, src=None, reduce_=None, use_new_options=False): + wrapped_dim = maybe_wrap_dim(dim, self.dim()) + scatter_gather_dtype_check("scatter", self, index, src) + scatter_shape_check(self, wrapped_dim, index, src) + if reduce_ is not None: + # Check if we have a valid reduce operator. + get_operator_enum(reduce_, use_new_options) + + +@register_meta(aten.scatter_add.default) +def meta_scatter_add(self, dim, index, src): + scatter_meta_impl(self, dim, index, src, "add") + return self.new_empty(self.shape) + + +@register_meta(aten.scatter_add_) +def meta_scatter_add_(self, dim, index, src): + scatter_meta_impl(self, dim, index, src, "add") + return self + + +@register_meta( + [ + aten.scatter.src, + aten.scatter.value, + aten.scatter.reduce, + aten.scatter.value_reduce, + ] +) +@out_wrapper() +def meta_scatter(self, dim, index, src_or_value, reduce=None): + src = src_or_value if isinstance(src_or_value, torch.Tensor) else None + scatter_meta_impl(self, dim, index, src, reduce) + return self.new_empty(self.shape) + + +@register_meta( + [ + aten.scatter_.src, + aten.scatter_.value, + aten.scatter_.reduce, + aten.scatter_.value_reduce, + ] +) +def meta_scatter_(self, dim, index, src_or_value, reduce=None): + src = src_or_value if isinstance(src_or_value, torch.Tensor) else None + scatter_meta_impl(self, dim, index, src, reduce) + return self + + +@register_meta([aten._scaled_dot_product_flash_attention]) +def meta__scaled_dot_product_flash_attention( + query: Tensor, + key: Tensor, + value: Tensor, + dropout_p: float = 0.0, + is_causal: bool = False, + return_debug_mask: bool = False, + scale: Optional[float] = None, +): + batch_size = query.size(0) + num_heads = query.size(1) + max_seqlen_batch_q = query.size(2) + head_dim = query.size(3) + max_seqlen_batch_k = key.size(2) + + query_t = query.transpose(1, 2) + attention = torch.empty_like(query_t).transpose(1, 2) + logsumexp = torch.empty( + (batch_size, num_heads, max_seqlen_batch_q), + dtype=torch.float, + device=query.device, + ) + + if return_debug_mask: + blocksize_c = 128 if head_dim > 64 else 256 + max_seqlen_k = math.ceil(max_seqlen_batch_q / blocksize_c) + if max_seqlen_batch_k <= 128: + max_seqlen_k = 128 + elif max_seqlen_batch_k <= 256: + max_seqlen_k = 256 + debug_mask = torch.empty( + (batch_size, num_heads, max_seqlen_batch_q, max_seqlen_k), + dtype=query.dtype, + device=query.device, + ) + else: + debug_mask = torch.empty(0, dtype=query.dtype, device=query.device) + + # Note [Seed and Offset]: device for seed and offset below depends on whether we are + # capturing or not, but at the time of tracing we don't know if we + # are going to use cudagraphs or not, so we return meta tensors here + # it's possible we'll need to have some special handling in inductor for sdpa + # See [Note] BC breaking change to flash seed/offset + if torch.version.hip and torch.cuda.is_available(): + # Maintain old path on AMD + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + else: + seed = torch.empty((2), dtype=torch.uint64, device="meta") + offset = torch.empty((), dtype=torch.uint64, device="meta") + + return ( + attention, + logsumexp, + None, + None, + max_seqlen_batch_q, + max_seqlen_batch_k, + seed, + offset, + debug_mask, + ) + + +def alloc_with_matching_layout( + query: Tensor, + res_shape: tuple[int, ...], +): + if tuple(query.shape) == res_shape: + query_t = query.transpose(1, 2) + res = torch.empty_like(query_t).transpose(1, 2) + else: + dim_order = sorted( + [0, 1, 2, 3], key=lambda idx: query.stride()[idx], reverse=True + ) + permuted_shape = [res_shape[idx] for idx in dim_order] + final_permute = [dim_order.index(i) for i in range(len(dim_order))] + res = torch.empty( + permuted_shape, dtype=query.dtype, device=query.device + ).permute(final_permute) + + return res + + +@register_meta([aten._scaled_dot_product_cudnn_attention]) +def meta__scaled_dot_product_cudnn_attention( + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Optional[Tensor], + compute_log_sumexp: bool, + dropout_p: float = 0.0, + is_causal: bool = False, + return_debug_mask: bool = False, + scale: Optional[float] = None, +): + B = query.size(0) + H = query.size(1) + S_Q = query.size(2) + S_KV = key.size(2) + D_V = value.size(-1) + + res_shape = (B, H, S_Q, D_V) + res = alloc_with_matching_layout(query, res_shape) + + logsum_exp = torch.empty( + (B, H, S_Q, 1), + dtype=torch.float, + device=query.device, + ) + + # See Note [Seed and Offset] + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + + return ( + res, + logsum_exp, + None, + None, + S_Q, + S_KV, + seed, + offset, + None, + ) + + +@register_meta([aten._scaled_dot_product_fused_attention_overrideable]) +def meta__scaled_dot_product_fused_attention_overrideable( + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Optional[Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + return_debug_mask: bool = False, + scale: Optional[float] = None, +): + B = query.size(0) + H_Q = query.size(1) + S_Q = query.size(2) + S_KV = key.size(2) + D_V = value.size(-1) + + res_shape = (B, H_Q, S_Q, D_V) + res = alloc_with_matching_layout(query, res_shape) + + logsum_exp = torch.empty( + (B, H_Q, S_Q), + dtype=torch.float, + device=query.device, + ) + + # See Note [Seed and Offset] + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + + return ( + res, + logsum_exp, + None, + None, + S_Q, + S_KV, + seed, + offset, + None, + ) + + +@register_meta( + [ + aten._scaled_dot_product_flash_attention_backward, + ] +) +def meta__scaled_dot_product_flash_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + out: Tensor, + logsumexp: Tensor, + cum_seq_q: Tensor, + cum_seq_k: Tensor, + max_q: int, + max_k: int, + dropout_p: float, + is_causal: bool, + philox_seed: Tensor, + philox_offset: Tensor, + scale: Optional[float] = None, +): + grad_q = torch.empty_like(query.transpose(1, 2)).transpose(1, 2) + grad_k = torch.empty_like(key.transpose(1, 2)).transpose(1, 2) + grad_v = torch.empty_like(value.transpose(1, 2)).transpose(1, 2) + return grad_q, grad_k, grad_v + + +@register_meta( + [ + aten._scaled_dot_product_flash_attention_for_cpu, + ] +) +def meta__scaled_dot_product_flash_attention_for_cpu( + query: Tensor, + key: Tensor, + value: Tensor, + dropout_p: float = 0.0, + is_causal: bool = False, + attn_mask: Optional[Tensor] = None, + scale: Optional[float] = None, +): + batch_size = query.size(0) + num_heads = query.size(1) + max_seqlen_batch_q = query.size(2) + + attention = torch.empty_like(query) + logsumexp = torch.empty( + ( + batch_size, + max_seqlen_batch_q, + num_heads, + ), + dtype=torch.float, + device=query.device, + ).transpose(1, 2) + return ( + attention, + logsumexp, + ) + + +@register_meta( + [ + aten._scaled_dot_product_flash_attention_for_cpu_backward, + ] +) +def meta__scaled_dot_product_flash_attention_for_cpu_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + out: Tensor, + logsumexp: Tensor, + dropout_p: float, + is_causal: bool, + attn_mask: Optional[Tensor] = None, + scale: Optional[float] = None, +): + # cpus's grad layout is different from cuda's, + # i.e. (batch_size, seq_len, num_heads, head_dim) + grad_q = torch.empty_permuted( + query.size(), + (0, 2, 1, 3), + dtype=query.dtype, + device=query.device, + ) + grad_k = torch.empty_permuted( + key.size(), + (0, 2, 1, 3), + dtype=key.dtype, + device=key.device, + ) + grad_v = torch.empty_permuted( + value.size(), + (0, 2, 1, 3), + dtype=value.dtype, + device=value.device, + ) + + return grad_q, grad_k, grad_v + + +@register_meta([aten._scaled_dot_product_attention_math_for_mps]) +def meta__scaled_dot_product_attention_math_for_mps( + query: Tensor, + key: Tensor, + value: Tensor, + attn_mask: Optional[Tensor] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + dropout_mask: Optional[Tensor] = None, + scale: Optional[float] = None, +) -> tuple[Tensor, Tensor]: + def ensure_4d(x): + if x.dim() == 3: + return x.unsqueeze(0), True + elif x.dim() > 4: + batch_size = 1 + for i in range(x.dim() - 3): + batch_size *= x.shape[i] + return x.view(batch_size, x.size(-3), x.size(-2), x.size(-1)), True + else: + return x, False + + q_, unsqueezed = ensure_4d(query) + k_, _ = ensure_4d(key) + v_, _ = ensure_4d(value) + + batch_size, num_head, q_size, head_size = q_.shape + _, k_size, max_seq_length, _ = k_.shape + + def sdpa_vector_fast_mps(): + out = q_.new_empty(q_.shape) + if unsqueezed: + out = out.view_as(query) + + attn = q_.new_empty((batch_size, num_head, q_size, max_seq_length)) + if unsqueezed: + if query.dim() == 3: + attn = attn.squeeze(0) + else: + shape = list(query.shape[:-3]) + attn.shape[1:4] + attn = attn.view(shape) + return out, attn + + def sdpa_vector_2pass_mps(): + blocks = 32 + out = q_.new_empty(q_.shape) + intermediate = q_.new_empty((batch_size, num_head, q_size, blocks, head_size)) + return out, intermediate + + if (max_seq_length >= 1024) or (k_size < q_size and max_seq_length >= 4096): + return sdpa_vector_2pass_mps() + else: + return sdpa_vector_fast_mps() + + +@register_meta([aten._scaled_dot_product_efficient_attention]) +def meta__scaled_dot_product_efficient_attention( + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Optional[Tensor], + compute_log_sumexp: bool, + dropout_p=0.0, + is_causal: bool = False, + scale: Optional[float] = None, +): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + B = query.size(0) + M = query.size(1) + num_heads = query.size(-2) + Kv = value.size(-1) + + res = torch.empty(B, M, num_heads, Kv, dtype=query.dtype, device=query.device) + + if torch.version.hip and torch.cuda.is_available(): + """Please see: https://github.com/pytorch/pytorch/issues/146848 + longsumexp last dim should be seq length + """ + logsumexp_dim = M if compute_log_sumexp else 0 + else: + logsumexp_dim = math.ceil(M / 32) * 32 if compute_log_sumexp else 0 + + logsum_exp = torch.empty( + (B, num_heads, logsumexp_dim), + dtype=torch.float, + device=query.device, + ) + + res = res.transpose(1, 2) + + # See Note [Seed and Offset]: + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + + return res, logsum_exp, seed, offset + + +@register_meta( + [ + aten._scaled_dot_product_efficient_attention_backward, + ] +) +def meta__scaled_dot_product_efficient_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + attn_bias: Optional[Tensor], + out: Tensor, + logsumexp: Tensor, + philox_seed: Tensor, + philox_offset: Tensor, + dropout_p: float, + grad_input_mask: list[bool], + is_causal: bool = False, + scale: Optional[float] = None, +): + batch_size = query.size(0) + num_heads = query.size(1) + max_q = query.size(2) + head_dim = query.size(3) + head_dim_v = value.size(3) + + max_k = key.size(2) + + grad_q = torch.empty_permuted( + (batch_size, num_heads, max_q, head_dim), + (0, 2, 1, 3), + dtype=query.dtype, + device=query.device, + ) + grad_k = torch.empty_permuted( + (batch_size, num_heads, max_k, head_dim), + (0, 2, 1, 3), + dtype=key.dtype, + device=key.device, + ) + grad_v = torch.empty_permuted( + (batch_size, num_heads, max_k, head_dim_v), + (0, 2, 1, 3), + dtype=value.dtype, + device=value.device, + ) + grad_bias = None + if attn_bias is not None and grad_input_mask[3]: + lastDim = attn_bias.size(-1) + lastDimAligned = lastDim if lastDim % 16 == 0 else lastDim + 16 - lastDim % 16 + new_sizes = list(attn_bias.size()) + new_sizes[-1] = lastDimAligned + grad_bias = torch.empty( + new_sizes, dtype=attn_bias.dtype, device=attn_bias.device + ) + grad_bias = grad_bias[..., :lastDim] + + return grad_q, grad_k, grad_v, grad_bias + + +@register_meta( + [ + aten._scaled_dot_product_cudnn_attention_backward, + ] +) +def meta__scaled_dot_product_cudnn_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + out: Tensor, + logsumexp: Tensor, + philox_seed: Tensor, + philox_offset: Tensor, + attn_bias: Tensor, + cum_seq_q: Tensor, + cum_seq_k: Tensor, + max_q: int, + max_k: int, + dropout_p: float, + is_causal: bool, + scale: Optional[float] = None, +): + grad_q = torch.empty_like(query) + grad_k = torch.empty_like(key) + grad_v = torch.empty_like(value) + return grad_q, grad_k, grad_v + + +@register_meta( + [ + aten._flash_attention_forward, + ] +) +def meta__flash_attention_forward( + query: Tensor, + key: Tensor, + value: Tensor, + cum_seq_q: Optional[Tensor], + cum_seq_k: Optional[Tensor], + max_q: int, + max_k: int, + dropout_p: float, + is_causal: bool, + return_debug_mask: bool, + scale: Optional[float] = None, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + seqused_k: Optional[Tensor] = None, + alibi_slopes: Optional[Tensor] = None, +): + # NB: there are two underlying paths: + # 1. normal dense path; expect 4D inputs of shape (batch_size, seqlen, num_heads, head_dim) + # 2. varseqlen path; expect 3D inputs of shape (total, num_heads, head_dim) where total + # includes all batch item sequences. cum_seq_q / cum_seq_k contain offsets into total + batch_size = query.size(0) if cum_seq_q is None else cum_seq_q.numel() - 1 + max_seqlen_batch_q = query.size(1) if cum_seq_q is None else max_q + max_seqlen_batch_k = key.size(1) if cum_seq_k is None else max_k + num_heads = query.size(-2) + head_dim = query.size(-1) + + # Cuda Path + attention = torch.empty_like(query) + if cum_seq_q is None: + logsumexp = torch.empty( + (batch_size, num_heads, max_seqlen_batch_q), + dtype=torch.float, + device=query.device, + ) + else: + total_q = query.size(0) + logsumexp = torch.empty( + (num_heads, total_q), dtype=torch.float, device=query.device + ) + + if return_debug_mask: + blocksize_c = 128 if head_dim > 64 else 256 + max_seqlen_k = math.ceil(max_seqlen_batch_q / blocksize_c) + if max_seqlen_batch_k <= 128: + max_seqlen_k = 128 + elif max_seqlen_batch_k <= 256: + max_seqlen_k = 256 + debug_mask = torch.empty( + (batch_size, num_heads, max_seqlen_batch_q, max_seqlen_k), + dtype=query.dtype, + device=query.device, + ) + else: + debug_mask = torch.empty(0, dtype=query.dtype, device=query.device) + + # See Note [Seed and Offset] + # See [Note] BC breaking change to flash seed/offset + seed, offset = None, None + if torch.version.hip and torch.cuda.is_available(): + # Maintain old path on AMD + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + else: + seed = torch.empty((2), dtype=torch.uint64, device="meta") + offset = torch.empty((), dtype=torch.uint64, device="meta") + return ( + attention, + logsumexp, + seed, + offset, + debug_mask, + ) + + +@register_meta( + [ + aten._flash_attention_backward, + ] +) +def meta__flash_attention_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + out: Tensor, + logsumexp: Tensor, + cum_seq_q: Tensor, + cum_seq_k: Tensor, + max_q: int, + max_k: int, + dropout_p: float, + is_causal: bool, + philox_seed: Tensor, + philox_offset: Tensor, + scale: Optional[float] = None, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, +): + grad_query = torch.empty_like(query) + grad_key = torch.empty_like(key) + grad_value = torch.empty_like(value) + + return grad_query, grad_key, grad_value + + +@register_meta( + [ + aten._efficient_attention_forward, + ] +) +def meta__efficient_attention_forward( + query: Tensor, + key: Tensor, + value: Tensor, + bias: Optional[Tensor], + cu_seqlens_q: Optional[Tensor], + cu_seqlens_k: Optional[Tensor], + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + dropout_p: float, + custom_mask_type: int, + compute_log_sumexp: bool = False, + scale: Optional[float] = None, + causal_diagonal: Optional[Tensor] = None, + seqlen_k: Optional[Tensor] = None, + window_size: Optional[int] = None, +): + B = query.size(0) + M = query.size(1) + N = key.size(1) + num_heads = query.size(-2) + Kv = value.size(-1) + + res = torch.empty(B, M, num_heads, Kv, dtype=query.dtype, device=query.device) + + logsumexp_batch_dim = cu_seqlens_q.size(0) - 1 if (cu_seqlens_q is not None) else B + actual_max_seqlen_q = M + if cu_seqlens_q is not None: + assert max_seqlen_q is not None + actual_max_seqlen_q = max_seqlen_q + actual_max_seqlen_k = max_seqlen_k if max_seqlen_k is not None else N + logsumexp_dim = ( + math.ceil(actual_max_seqlen_q / 32) * 32 if compute_log_sumexp else 0 + ) + logsum_exp = torch.empty( + (logsumexp_batch_dim, num_heads, logsumexp_dim), + dtype=torch.float, + device=query.device, + ) + + # See Note [Seed and Offset]: + seed = torch.empty((), dtype=torch.long, device="meta") + offset = torch.empty((), dtype=torch.long, device="meta") + + return res, logsum_exp, seed, offset, actual_max_seqlen_q, actual_max_seqlen_k + + +@register_meta( + [ + aten._efficient_attention_backward, + ] +) +def meta__efficient_attention_backward( + grad_out: Tensor, + query: Tensor, + key: Tensor, + value: Tensor, + bias: Optional[Tensor], + cu_seqlens_q: Optional[Tensor], + cu_seqlens_k: Optional[Tensor], + max_seqlen_q: torch.SymInt, + max_seqlen_k: torch.SymInt, + logsumexp: Tensor, + dropout_p: float, + philox_seed: Tensor, + philox_offset: Tensor, + custom_mask_type: int, + bias_requires_grad: bool, + scale: Optional[float] = None, + num_splits_key: Optional[int] = None, + shared_storage_dqdkdv: bool = False, +): + if shared_storage_dqdkdv: + torch._check( + query.shape[1] == key.shape[1], + lambda: "seqlen must match for `shared_storage_dqdkdv", + ) + torch._check( + query.shape[3] == key.shape[3], + lambda: "embedding dim must match for `shared_storage_dqdkdv", + ) + chunk = torch.empty( + (*query.shape[0:-2], 3, query.shape[-2], query.shape[-1]), + dtype=query.dtype, + device=query.device, + ) + grad_query = chunk.select(-3, 0) + grad_key = chunk.select(-3, 1) + grad_value = chunk.select(-3, 2) + else: + grad_query = torch.empty_like(query) + grad_key = torch.empty_like(key) + grad_value = torch.empty_like(value) + + if bias is not None: + lastDim = bias.size(-1) + lastDimAligned = lastDim if lastDim % 16 == 0 else lastDim + 16 - lastDim % 16 + new_sizes = list(bias.size()) + new_sizes[-1] = lastDimAligned + grad_bias = torch.empty(new_sizes, dtype=bias.dtype, device=bias.device) + grad_bias = grad_bias[..., :lastDim] + else: + grad_bias = torch.empty((), device=query.device) + + return grad_query, grad_key, grad_value, grad_bias + + +@register_meta([aten._scaled_mm.default]) +def meta_scaled_mm( + self: torch.Tensor, + mat2: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + bias: Optional[torch.Tensor] = None, + scale_result: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + use_fast_accum: bool = False, +): + def is_fp8_or_fp4_type(dtype): + return dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, + torch.float8_e5m2fnuz, + torch.float4_e2m1fn_x2, + ) + + torch._check( + self.dim() == 2 and mat2.dim() == 2, + lambda: f"Inputs must be 2D but got self.dim()={self.dim()} and mat2.dim()={mat2.dim()}", + ) + torch._check( + is_fp8_or_fp4_type(self.dtype) and is_fp8_or_fp4_type(mat2.dtype), + lambda: f"Expected both inputs to be fp8 or fp4 types but got self.dtype={self.dtype} and mat2.dtype={mat2.dtype}", + ) + + if device_hint(self) == "cuda": + + def is_row_major(stride): + return stride[0] > stride[1] and stride[1] == 1 + + def is_col_major(stride): + return stride[0] == 1 and stride[1] > 1 + + def has_zero_dim(tensor_2d): + return tensor_2d.size(0) == 0 or tensor_2d.size(1) == 0 + + torch._check( + is_row_major(self.stride()) or has_zero_dim(self), + lambda: f"self must be row_major, got stride {self.stride()}", + ) + torch._check( + is_col_major(mat2.stride()) or has_zero_dim(mat2), + lambda: f"mat2 must be col_major, got stride {mat2.stride()}", + ) + torch._check( + self.size(1) % 16 == 0, + lambda: f"Expected self.size(1) to be divisible by 16, but got self.size(1)={self.size(1)}", + ) + torch._check( + mat2.size(0) % 16 == 0 and mat2.size(1) % 16 == 0, + lambda: f"Expected both dimensions of mat2 to be divisible by 16 but got {mat2.shape}", + ) + + # determine scaling type and check input dimensions (refer to Blas.cpp op) + + m, _k = self.shape + n = mat2.size(1) + + is_blockwise_scaling = ( + scale_a.dtype == torch.float8_e8m0fnu + and scale_b.dtype == torch.float8_e8m0fnu + ) or ( + scale_a.dtype == torch.float8_e4m3fn + and scale_b.dtype == torch.float8_e4m3fn + ) + + if scale_a.numel() == 1 and scale_b.numel() == 1: + # tensorwise scaling + torch._check( + scale_a.dtype == torch.float32 and scale_b.dtype == torch.float32, + lambda: "For tensorwise scaling, both scale_a and scale_b must be float (fp32) tensors.", + ) + elif is_blockwise_scaling: + # blockwise scaling + + if scale_a.dtype == torch.float8_e4m3fn: + # NVIDIA's nvfp4 recipe: + # * block size is 16 elements packed (32 unpacked) + # * _k needs to be translated to the unpacked version + block_size_k = 16 + _k = _k * 2 + else: + block_size_k = 32 + + block_size_mn = 128 + + def ceil_div(a, b): + return (a + b - 1) // b + + num_k_blocks = ceil_div(_k, block_size_k) + padded_num_k_blocks = ceil_div(num_k_blocks, 4) * 4 + + expected_a_size = ( + block_size_mn * ceil_div(m, block_size_mn) * padded_num_k_blocks + ) + expected_b_size = ( + block_size_mn * ceil_div(n, block_size_mn) * padded_num_k_blocks + ) + + if ( + scale_a.numel() == expected_a_size + and scale_b.numel() == expected_b_size + ): + torch._check( + scale_a.is_contiguous(), + lambda: "scale_a must be contiguous", + ) + torch._check( + scale_b.is_contiguous(), + lambda: "scale_b must be contiguous", + ) + else: + torch._check( + False, + lambda: ( + "Invalid blockwise scaling configuration. " + f"For blockwise scaling, scale_a should have {expected_a_size} elements, got {scale_a.numel()}, " + f"scale_b should have {expected_b_size} elements, got {scale_b.numel()}." + ), + ) + else: + torch._check( + scale_a.dtype == torch.float32 and scale_b.dtype == torch.float32, + lambda: "For rowwise scaling, both scale_a and scale_b must be float (fp32) tensors.", + ) + # for rowwise scaling, enforce 2D input tensors + torch._check( + scale_a.dim() == 2 and scale_b.dim() == 2, + lambda: f"For non-tensorwise scaling, scale tensors must be 2D, but got {scale_a.dim()=} and {scale_b.dim()=}", + ) + + if ( + scale_a.size(0) == m + and scale_a.size(1) == 1 + and scale_b.size(0) == 1 + and scale_b.size(1) == n + ): + # rowwise scaling + torch._check( + scale_a.is_contiguous() and scale_b.is_contiguous(), + lambda: "Both scale_a and scale_b must be contiguous for rowwise scaling.", + ) + else: + # does not match any valid scaling type + torch._check( + False, + lambda: ( + "Invalid scaling configuration. " + "For tensorwise scaling, both scales should be scalar. " + f"For rowwise scaling, scale_a should be ({m}, 1), scale_b should be (1, {n}). " + f"Got scale_a.size()=({scale_a.size(0)}, {scale_a.size(1)}) " + f"and scale_b.size()=({scale_b.size(0)}, {scale_b.size(1)})" + ), + ) + + _out_dtype = out_dtype if out_dtype is not None else self.dtype + return torch.empty(self.size(0), mat2.size(1), dtype=_out_dtype, device=self.device) + + +@register_meta([aten.scatter_reduce.two, aten.scatter_reduce.two_out]) +@out_wrapper() +def meta_scatter_reduce_two(self, dim, index, src, reduce, include_self=True): + scatter_meta_impl(self, dim, index, src, reduce, use_new_options=True) + return self.new_empty(self.shape) + + +@register_meta(aten.scatter_reduce_.two) +def meta_scatter_reduce__two(self, dim, index, src, reduce, include_self=True): + scatter_meta_impl(self, dim, index, src, reduce, use_new_options=True) + return self + + +@register_meta([aten.multinomial.default, aten.multinomial.out]) +@out_wrapper() +def meta_multinomial(input, num_samples, replacement=False, *, generator=None): + torch._check( + 0 < input.dim() <= 2, + lambda: f"The probability distributions dimensions must be 1 or 2, but got {input.dim()}", + ) + if input.dim() == 1: + return torch.empty(num_samples, dtype=torch.long, device=input.device) + return torch.empty( + input.size(0), num_samples, dtype=torch.long, device=input.device + ) + + +def multiply_integers(vs): + r = 1 + for v in vs: + r *= v + return r + + +def upsample_common_check(input_size, output_size, num_spatial_dims): + torch._check( + len(output_size) == num_spatial_dims, + lambda: f"It is expected output_size equals to {num_spatial_dims}, but got size {len(output_size)}", + ) + expected_input_dims = num_spatial_dims + 2 # N, C, ... + torch._check( + len(input_size) == expected_input_dims, + lambda: f"It is expected input_size equals to {expected_input_dims}, but got size {len(input_size)}", + ) + + torch._check( + all(s > 0 for s in input_size[2:]) and all(s > 0 for s in output_size), + lambda: f"Input and output sizes should be greater than 0, but got " + f"input size {input_size} and output size {output_size}", + ) + + nbatch, channels = input_size[:2] + return (nbatch, channels, *output_size) + + +@register_meta( + [aten.upsample_nearest1d.default, aten._upsample_nearest_exact1d.default] +) +def upsample_nearest1d(input, output_size, scales=None): + torch._check( + input.numel() != 0 or multiply_integers(input.size()[1:]), + lambda: f"Non-empty 3D data tensor expected but got a tensor with sizes {input.size()}", + ) + full_output_size = upsample_common_check( + input.size(), output_size, num_spatial_dims=1 + ) + return input.new_empty(full_output_size).to( + memory_format=utils.suggest_memory_format(input) + ) + + +@register_meta( + [aten.upsample_nearest2d.default, aten._upsample_nearest_exact2d.default] +) +def upsample_nearest2d(input, output_size, scales_h=None, scales_w=None): + torch._check( + input.numel() != 0 or multiply_integers(input.size()[1:]), + lambda: f"Non-empty 4D data tensor expected but got a tensor with sizes {input.size()}", + ) + full_output_size = upsample_common_check( + input.size(), output_size, num_spatial_dims=2 + ) + output = input.new_empty(full_output_size) + + # convert output to correct memory format, if necessary + memory_format = utils.suggest_memory_format(input) + + # following "heuristic: only use channels_last path when it's faster than the contiguous path" + _, n_channels, _, _ = input.shape + if input.device.type == "cuda" and n_channels < 4: + memory_format = torch.contiguous_format + + output = output.contiguous(memory_format=memory_format) + + return output + + +@register_meta( + [ + aten.upsample_nearest2d_backward.default, + aten._upsample_nearest_exact2d_backward.default, + ] +) +def upsample_nearest2d_backward( + grad_output: Tensor, + output_size: Sequence[Union[int, torch.SymInt]], + input_size: Sequence[Union[int, torch.SymInt]], + scales_h: Optional[float] = None, + scales_w: Optional[float] = None, +): + full_output_size = upsample_common_check( + input_size, output_size, num_spatial_dims=2 + ) + torch._check( + grad_output.ndim == 4, + lambda: f"Expected grad_output to be a tensor of dimension 4 but got: dimension {grad_output.ndim}", + ) + for i in range(4): + torch._check( + grad_output.size(i) == full_output_size[i], + lambda: ( + f"Expected grad_output to have the same shape as output;" + f" output.size({i}) = {full_output_size[i]}" + f" but got grad_output.size({i}) = {grad_output.size(i)}" + ), + ) + + return grad_output.new_empty(input_size).to( + memory_format=utils.suggest_memory_format(grad_output) + ) # type: ignore[call-overload] + + +@register_meta( + [aten.upsample_nearest3d.default, aten._upsample_nearest_exact3d.default] +) +def upsample_nearest3d(input, output_size, scales_d=None, scales_h=None, scales_w=None): + torch._check( + input.numel() != 0 or multiply_integers(input.size()[1:]), + lambda: f"Non-empty 5D data tensor expected but got a tensor with sizes {input.size()}", + ) + full_output_size = upsample_common_check( + input.size(), output_size, num_spatial_dims=3 + ) + return input.new_empty(full_output_size).to( + memory_format=utils.suggest_memory_format(input) + ) + + +@register_meta( + [ + aten.sort.default, + aten.sort.stable, + aten.sort.values, + aten.sort.values_stable, + ] +) +def meta_sort(self, stable=None, dim=-1, descending=False, values=None, indices=None): + v, i = torch.empty_like(self), torch.empty_like(self, dtype=torch.int64) + if values is not None and indices is not None: + assert isinstance(values, TensorLike) + assert isinstance(indices, TensorLike) + # Makes sure values and indices have the same strides. For cases where + # these have different shapes, like (5, 10, 5) and (0) in msort. + out_shape = v.shape + out_stride = v.stride() + values = _maybe_resize_out(values, out_shape) + indices = _maybe_resize_out(indices, out_shape) + values.as_strided_(out_shape, out_stride) + indices.as_strided_(out_shape, out_stride) + _safe_copy_out(copy_from=v, copy_to=values) # type: ignore[arg-type] + _safe_copy_out(copy_from=i, copy_to=indices) # type: ignore[arg-type] + return values, indices + return v, i + + +def rnn_cell_checkSizes( + input_gates, + hidden_gates, + input_bias, + hidden_bias, + factor, + prev_hidden, +): + torch._check(input_gates.ndim == 2, lambda: f"{input_gates.ndim} != 2") + torch._check( + input_gates.shape == hidden_gates.shape, + lambda: f"{input_gates.shape} != {hidden_gates.shape}", + ) + gates_size = input_gates.size(1) + if input_bias is not None: + torch._check(input_bias.ndim == 1, lambda: f"{input_bias.ndim} != 1") + torch._check( + input_bias.numel() == gates_size, + lambda: f"{input_bias.numel()} != {gates_size}", + ) + torch._check( + input_bias.shape == hidden_bias.shape, + lambda: f"{input_bias.shape} != {hidden_bias.shape}", + ) + torch._check(prev_hidden.ndim == 2, lambda: f"{prev_hidden.ndim} != 2") + expected_prev_hidden_numel = input_gates.size(0) * gates_size // factor + torch._check( + prev_hidden.numel() == expected_prev_hidden_numel, + lambda: f"{prev_hidden.numel()} != {input_gates.size(0)} * {gates_size} // {factor} (aka {expected_prev_hidden_numel})", + ) + torch._check( + all( + x.device == input_gates.device + for x in [hidden_gates, input_bias, hidden_bias, prev_hidden] + ), + lambda: "expected all inputs to be same device", + ) + + +@register_meta(aten._thnn_fused_lstm_cell.default) +def _thnn_fused_lstm_cell_meta( + input_gates, + hidden_gates, + cx, + input_bias=None, + hidden_bias=None, +): + rnn_cell_checkSizes(input_gates, hidden_gates, input_bias, hidden_bias, 4, cx) + workspace = torch.empty_like(input_gates, memory_format=torch.contiguous_format) + hy = torch.empty_like(cx, memory_format=torch.contiguous_format) + cy = torch.empty_like(cx, memory_format=torch.contiguous_format) + return (hy, cy, workspace) + + +@register_meta(aten._cudnn_rnn.default) +def _cudnn_rnn( + input, + weight, + weight_stride0, + weight_buf, + hx, + cx, + mode, + hidden_size, + proj_size, + num_layers, + batch_first, + dropout, + train, + bidirectional, + batch_sizes, + dropout_state, +): + is_input_packed = len(batch_sizes) != 0 + if is_input_packed: + seq_length = len(batch_sizes) + mini_batch = batch_sizes[0] + batch_sizes_sum = input.shape[0] + else: + seq_length = input.shape[1] if batch_first else input.shape[0] + mini_batch = input.shape[0] if batch_first else input.shape[1] + batch_sizes_sum = -1 + + num_directions = 2 if bidirectional else 1 + out_size = proj_size if proj_size != 0 else hidden_size + if is_input_packed: + out_shape = [batch_sizes_sum, out_size * num_directions] + else: + out_shape = ( + [mini_batch, seq_length, out_size * num_directions] + if batch_first + else [seq_length, mini_batch, out_size * num_directions] + ) + output = input.new_empty(out_shape) + + cell_shape = [num_layers * num_directions, mini_batch, hidden_size] + if cx is None: + cy = torch.empty(0, device=input.device) + else: + cy = cx.new_empty(cell_shape) + + hy = hx.new_empty([num_layers * num_directions, mini_batch, out_size]) + + # TODO: Query cudnnGetRNNTrainingReserveSize (expose to python) + reserve_shape = 0 if train else 0 + reserve = input.new_empty(reserve_shape, dtype=torch.uint8) + + return output, hy, cy, reserve, weight_buf + + +@register_meta(aten.mkldnn_rnn_layer.default) +def mkldnn_rnn_layer( + input, + w0, + w1, + w2, + w3, + hx_, + cx_, + reverse, + batch_sizes, + mode, + hidden_size, + num_layers, + has_biases, + bidirectional, + batch_first, + train, +): + seq_length = input.shape[1] if batch_first else input.shape[0] + mini_batch = input.shape[0] if batch_first else input.shape[1] + output_chanels = hidden_size + out_shape = ( + [mini_batch, seq_length, output_chanels] + if batch_first + else [seq_length, mini_batch, output_chanels] + ) + output = input.new_empty(out_shape) + if hx_ is None: + hy = torch.empty(0, device=input.device) + else: + hy = hx_.new_empty(hx_.shape) + if cx_ is None: + cy = torch.empty(0, device=input.device) + else: + cy = cx_.new_empty(cx_.shape) + workspace = torch.empty(0, device=input.device, dtype=torch.uint8) + return output, hy, cy, workspace + + +def zero_numel_check_dims(self, dim, fn_name): + if self.ndim == 0: + torch._check_index( + dim == 0 or dim == -1, + lambda: f"{fn_name}: Expected reduction dim -1 or 0 for scalar but got {dim}", + ) + else: + torch._check_index( + self.size(dim) != 0, + lambda: f"{fn_name}: Expected reduction dim {dim} to have non-zero size.", + ) + + +# From aten/src/ATen/native/ReduceOps.cpp +def check_argmax_argmin(name, self, dim): + if dim is not None: + dim = maybe_wrap_dim(dim, self.dim()) + zero_numel_check_dims(self, dim, name) + else: + torch._check( + self.numel() != 0, + lambda: f"{name}: Expected reduction dim to be specified for input.numel() == 0.", + ) + + +@register_meta([aten.argmax.default, aten.argmin.default]) +def argmax_argmin_meta(self, dim=None, keepdim=False): + check_argmax_argmin("argmax", self, dim) + dims = utils.reduction_dims(self.shape, (dim,) if dim is not None else None) + shape = _compute_reduction_shape(self, dims, keepdim) + return self.new_empty(shape, dtype=torch.int64) + + +@register_meta(aten.scalar_tensor.default) +def scalar_tensor(s, dtype=None, layout=None, device=None, pin_memory=None): + # NB: It's always wrong to try to create a scalar tensor with the jagged layout. + # Rather than fix this everywhere, just use the strided layout and let NJT handle + # scalar tensor broadcasting. + if layout == torch.jagged: + layout = torch.strided + return torch.empty( + (), dtype=dtype, layout=layout, device=device, pin_memory=pin_memory + ) + + +@register_meta(aten.topk.default) +def topk_meta(self, k, dim=-1, largest=True, sorted=True): + # From aten/src/ATen/native/Sorting.cpp + dim = maybe_wrap_dim(dim, self.dim(), wrap_scalar=True) + sliceSize = 1 if self.dim() == 0 else self.size(dim) + torch._check_is_size(k) + torch._check(k <= sliceSize, lambda: "k not in range for dimension") + + topKSize = list(self.shape) + if len(topKSize) > 0: + topKSize[dim] = k + return self.new_empty(topKSize), self.new_empty(topKSize, dtype=torch.int64) + + +@register_meta(aten._segment_reduce_backward) +@out_wrapper() +def meta__segment_reduce_backward( + grad, output, data, reduce, lengths=None, offsets=None, axis=0, initial=None +): + assert lengths is not None or offsets is not None, ( + "segment_reduce(): Either lengths or offsets must be defined" + ) + data_contig = data.contiguous() + grad_contig = grad.contiguous() + return torch.empty_like( + data_contig, + dtype=grad_contig.dtype, + device=grad_contig.device, + layout=grad_contig.layout, + ) + + +@register_meta([aten.kthvalue.default, aten.kthvalue.values]) +@out_wrapper("values", "indices") +def kthvalue_meta(self, k, dim=-1, keepdim=False): + from torch.fx.experimental.symbolic_shapes import sym_and + + dim = maybe_wrap_dim(dim, self.dim(), wrap_scalar=True) + dimSize = self.size(dim) if self.dim() > 0 else 1 + torch._check( + sym_and(k >= 1, k <= dimSize), + lambda: f"kthvalue(): selected number k out of range for dimension {dim}", + ) + + shape = list(self.shape[:dim] + self.shape[dim + 1 :]) + if keepdim and self.dim() > 0: + shape.insert(dim, 1) + return self.new_empty(shape), self.new_empty(shape, dtype=torch.int64) + + +legacy_contiguous_memory_format = torch.contiguous_format + + +# From aten/src/ATen/native/cuda/RNN.cu +def checkLSTMBackwardSizes(grad_hy, grad_cy, cx, cy, workspace): + defined_grad = grad_hy if grad_hy is not None else grad_cy + torch._check(defined_grad.dim() == 2, lambda: "") + exp_size = defined_grad.size() + if grad_hy is not None: + torch._check(grad_hy.size() == exp_size, lambda: "") + if grad_cy is not None: + torch._check(grad_cy.size() == exp_size, lambda: "") + torch._check(cx.size() == exp_size, lambda: "") + torch._check(cy.size() == exp_size, lambda: "") + torch._check(workspace.dim() == 2, lambda: "") + torch._check(workspace.numel() == exp_size[0] * exp_size[1] * 4, lambda: "") + + +# From aten/src/ATen/native/cuda/RNN.cu +@register_meta(aten._thnn_fused_lstm_cell_backward_impl.default) +def _thnn_fused_lstm_cell_backward_impl(grad_hy, grad_cy, cx, cy, workspace, has_bias): + if grad_hy is None and grad_cy is None: + return None, None, None + checkLSTMBackwardSizes(grad_hy, grad_cy, cx, cy, workspace) + grad_gates = torch.empty_like( + workspace, memory_format=legacy_contiguous_memory_format + ) + grad_cx = torch.empty_like(cx, memory_format=legacy_contiguous_memory_format) + grad_bias = grad_gates.sum(0, keepdim=False) if has_bias else None + return grad_gates, grad_cx, grad_bias + + +# From aten/src/ATen/native/mps/operations/Linear.mm +@register_meta(aten.linear_backward.default) +def linear_backward(input_, grad_output_, weight_, output_mask): + grad_input = None + grad_weight = None + grad_bias = None + if output_mask[0]: + grad_input = grad_output_.new_empty(input_.size()) + if output_mask[1] or output_mask[2]: + grad_weight = grad_output_.new_empty((grad_output_.size(-1), input_.size(-1))) + grad_bias = grad_output_.new_empty(grad_output_.size(-1)) + return (grad_input, grad_weight, grad_bias) + + +@register_meta(aten.pixel_shuffle.default) +def meta_pixel_shuffle(self, upscale_factor): + assert ( + len(self.shape) > 2 and self.shape[-3] % (upscale_factor * upscale_factor) == 0 + ), ( + f"Invalid input shape for pixel_shuffle: {self.shape} with upscale_factor = {upscale_factor}" + ) + + def is_channels_last(ten): + return torch._prims_common.suggest_memory_format(ten) == torch.channels_last + + def pick_memory_format(): + if is_channels_last(self): + if device_hint(self) == "cuda": + return torch.contiguous_format + else: + return torch.channels_last + elif self.is_contiguous(memory_format=torch.contiguous_format): + return torch.contiguous_format + elif self.is_contiguous(memory_format=torch.preserve_format): + return torch.preserve_format + + C = self.shape[-3] // (upscale_factor * upscale_factor) + Hr = self.shape[-2] * upscale_factor + Wr = self.shape[-1] * upscale_factor + out_shape = (*self.shape[:-3], C, Hr, Wr) + + out = self.new_empty(out_shape) + out = out.to(memory_format=pick_memory_format()) # type: ignore[call-overload] + return out + + +@register_meta(aten.mkldnn_rnn_layer_backward.default) +def mkldnn_rnn_layer_backward( + input, + weight0, + weight1, + weight2, + weight3, + hx_, + cx_tmp, + output, + hy_, + cy_, + grad_output_r_opt, + grad_hy_r_opt, + grad_cy_r_opt, + reverse, + mode, + hidden_size, + num_layers, + has_biases, + train, + bidirectional, + batch_sizes, + batch_first, + workspace, +): + diff_x = input.new_empty(input.shape) + diff_hx = hx_.new_empty(hx_.shape) + diff_cx = cx_tmp.new_empty(cx_tmp.shape) + diff_w1 = weight0.new_empty(weight0.shape) + diff_w2 = weight1.new_empty(weight1.shape) + diff_b = weight2.new_empty(weight2.shape) + return diff_x, diff_w1, diff_w2, diff_b, diff_b, diff_hx, diff_cx + + +@register_meta([aten.bucketize.Tensor, aten.bucketize.Tensor_out]) +@out_wrapper() +def meta_bucketize(self, boundaries, *, out_int32=False, right=False): + return torch.empty_like( + self, + dtype=torch.int32 if out_int32 else torch.int64, + memory_format=torch.contiguous_format, + ) + + +@register_meta([aten.histc]) +@out_wrapper() +def meta_histc(input, bins=100, min=0, max=0): + fn_name = "histc()" + if device_hint(input) == "cpu": + torch._check( + input.is_floating_point(), + lambda: f"\"histogram_cpu\" not implemented for '{input.dtype}'", + ) + if device_hint(input) == "cuda" and input.is_floating_point(): + utils.alert_not_deterministic("_histc_cuda with floating point input") + torch._check( + isinstance(bins, IntLike), + lambda: f"{fn_name}: argument 'bins' must be int, not {type(bins)}", + ) + torch._check(bins > 0, lambda: f"{fn_name}: bins must be > 0, but got {bins}") + torch._check( + isinstance(min, Number), + lambda: f"{fn_name}: argument 'min' must be Number, not {type(min)}", + ) + torch._check( + isinstance(max, Number), + lambda: f"{fn_name}: argument 'max' must be Number, not {type(max)}", + ) + torch._check(max >= min, lambda: "{fn_name}: max must be larger than min") + return torch.empty(bins, device=input.device, dtype=input.dtype) + + +@register_meta( + [aten._upsample_bilinear2d_aa.default, aten._upsample_bicubic2d_aa.default] +) +def meta_upsample_bimode2d_aa( + input, + output_size, + align_corners, + scales_h=None, + scales_w=None, +): + full_output_size = upsample_common_check( + input.size(), output_size, num_spatial_dims=2 + ) + torch._check( + input.numel() != 0 or all(size > 0 for size in input.size()[1:]), + lambda: f"Non-empty 4D data tensor expected but got a tensor with sizes {input.size()}", + ) + return input.new_empty(full_output_size).to( + memory_format=utils.suggest_memory_format(input) + ) + + +@register_meta([aten._upsample_bilinear2d_aa_backward.default]) +def meta_upsample_bimode2d_aa_backward( + grad_output, + output_size, + input_size, + align_corners, + scales_h=None, + scales_w=None, +): + full_output_size = upsample_common_check( + input_size, output_size, num_spatial_dims=2 + ) + torch._check( + grad_output.ndim == 4, + lambda: f"Expected grad_output to be a tensor of dimension 4 but got: dimension {grad_output.ndim}", + ) + for i in range(4): + torch._check( + grad_output.shape[i] == full_output_size[i], + lambda: f""" +Expected grad_output to have the same shape as output; output.size({i}) = {full_output_size[i]} +but got grad_output_size({i}) = {grad_output.size(i)}""", + ) + return grad_output.new_empty(input_size).to( + memory_format=utils.suggest_memory_format(grad_output) + ) + + +# From aten/src/ATen/native/cuda/AmpKernels.cu +@register_meta(aten._amp_foreach_non_finite_check_and_unscale_.default) +def _amp_foreach_non_finite_check_and_unscale_(self, found_inf, inv_scale): + torch._check( + found_inf.numel() == 1, lambda: "found_inf must be a 1-element tensor." + ) + torch._check( + inv_scale.numel() == 1, lambda: "inv_scale must be a 1-element tensor." + ) + torch._check( + found_inf.dtype.is_floating_point, + lambda: "found_inf must be a float tensor.", + ) + torch._check( + inv_scale.dtype.is_floating_point, + lambda: "inv_scale must be a float tensor.", + ) + + +# From aten/src/ATen/native/UnaryOps.cpp +@register_meta([aten.nan_to_num.default, aten.nan_to_num.out]) +@out_wrapper() +def nan_to_num(self, nan=None, posinf=None, neginf=None): + return torch.empty_like(self) + + +@register_meta(torch.ops.aten.transpose_) +def transpose_(self, dim0, dim1): + assert self.layout not in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }, ( + f"torch.transpose_: in-place transposition is not supported for {self.layout} layout" + ) + + ndims = self.ndim + + dim0 = maybe_wrap_dim(dim0, ndims) + dim1 = maybe_wrap_dim(dim1, ndims) + + if dim0 == dim1: + return self + + size = list(self.size()) + stride = list(self.stride()) + + stride[dim0], stride[dim1] = stride[dim1], stride[dim0] + size[dim0], size[dim1] = size[dim1], size[dim0] + + self.as_strided_(size, stride) + return self + + +@register_meta(torch.ops.aten.t_) +def t_(self): + ndims = self.ndim + + if self.is_sparse: + sparse_dim = self.sparse_dim() + dense_dim = self.dense_dim() + assert sparse_dim <= 2 and dense_dim == 0, ( + f"t_ expects a tensor with <= 2 sparse and 0 dense dimensions, " + f"but got {sparse_dim} sparse and {dense_dim} dense dimensions" + ) + else: + assert self.dim() <= 2, ( + f"t_ expects a tensor with <= 2 dimensions, but self is {ndims}D" + ) + + return transpose_(self, 0, 0 if ndims < 2 else 1) + + +@register_meta(aten.searchsorted) +@out_wrapper() +def meta_searchsorted( + sorted_sequence, + self, + *, + out_int32=False, + right=False, + side=None, + sorter=None, +): + # If the sorted_sequence is not one-dimensional, its shape must match that of values + # in all but the last dimension. + torch._check( + len(sorted_sequence.shape) <= 1 + or sorted_sequence.shape[:-1] == self.shape[:-1], + lambda: ( + "torch.searchsorted(): boundaries tensor should be 1 dimension or the " + "first N-1 dimensions of boundaries tensor and input value tensor must " + f"match, but we got boundaries tensor {list(sorted_sequence.shape)} and " + f"input value tensor {list(self.shape)}" + ), + ) + + # If a sorter array is provided, its dimensions must exactly match sorted_sequence. + torch._check( + sorter is None or sorted_sequence.shape == sorter.shape, + lambda: ( + "torch.searchsorted(): boundary and sorter must have the same size, but " + f"got boundary tensor {list(sorted_sequence.shape)} and got sorter tensor " + f"{list(sorter.shape) if sorter is not None else []}" + ), + ) + + # Per the docs, if side == "left" and right is True, we error. + torch._check( + side != "left" or not right, + "torch.searchsorted(): side and right can't be set to opposites, got side of " + "left while right was True", + ) + + dtype = torch.int32 if out_int32 else torch.int64 + if isinstance(self, torch.Tensor): + return torch.empty_like( + self, dtype=dtype, memory_format=torch.contiguous_format + ) + else: # Scalar + return torch.empty((), dtype=dtype, device=sorted_sequence.device) + + +def _check_for_unsupported_isin_dtype(dtype): + torch._check( + dtype not in (torch.bool, torch.complex128, torch.complex64), + lambda: f"Unsupported input type encountered for isin(): {dtype}", + ) + + +@register_meta(aten.embedding_dense_backward) +def meta_embedding_dense_backward( + grad_output, + indices, + num_weights, + padding_idx, + scale_grad_by_freq, +): + grad_weight = grad_output.new_empty((num_weights, grad_output.size(-1))) + return grad_weight + + +@register_meta(aten._embedding_bag_backward) +def meta_embedding_bag_backward( + grad, + indices, + offsets, + offset2bag, + bag_size, + maximum_indices, + num_weights, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + padding_idx=-1, +): + if sparse: + return aten._embedding_bag_sparse_backward( + grad, + indices, + offsets, + offset2bag, + bag_size, + num_weights, + scale_grad_by_freq, + mode, + per_sample_weights, + padding_idx, + ) + else: + return meta_embedding_bag_dense_backward( + grad, + indices, + offset2bag, + bag_size, + maximum_indices, + num_weights, + scale_grad_by_freq, + mode, + per_sample_weights, + padding_idx, + ) + + +@register_meta(aten._embedding_bag_dense_backward) +def meta_embedding_bag_dense_backward( + grad, + indices, + offset2bag, + bag_size, + maximum_indices, + num_weights, + scale_grad_by_freq, + mode, + per_sample_weights, + padding_idx=-1, +): + torch._check( + grad.dtype in [torch.float16, torch.bfloat16, torch.float32, torch.float64], + lambda: f"Unsupported input type encountered: {grad.dtype}", + ) + if mode == MODE_MAX: + torch._check(maximum_indices is not None) + index_grad_weight = grad.new_empty((num_weights, grad.size(1))) + return index_grad_weight + + +@register_meta(aten._embedding_bag_per_sample_weights_backward) +def meta_embedding_bag_per_sample_weights_backward( + grad, + weight, + indices, + offsets, + offset2bag, + mode, + padding_idx=-1, +): + embedding_features = grad.size(1) + torch._check( + mode == MODE_SUM, + "embedding_bag_backward: per_sample_weights only supported for mode='sum'", + ) + torch._check(grad.dim() == 2) + torch._check(indices.dim() == 1) + num_samples = indices.size(0) + torch._check(weight.dim() == 2) + torch._check(weight.size(1) == embedding_features) + output = grad.new_empty((num_samples,)) + return output + + +@register_meta(aten.isin) +@out_wrapper() +def meta_isin(elements, test_elements, *, assume_unique=False, invert=False): + torch._check( + isinstance(elements, Tensor) or isinstance(test_elements, Tensor), + lambda: "At least one of elements and test_elements must be a Tensor.", + ) + if not isinstance(elements, Tensor): + elements = torch.tensor(elements, device=test_elements.device) + + if not isinstance(test_elements, Tensor): + test_elements = torch.tensor(test_elements, device=elements.device) + + _check_for_unsupported_isin_dtype(elements.dtype) + _check_for_unsupported_isin_dtype(test_elements.dtype) + return torch.empty_like(elements, dtype=torch.bool) + + +@register_meta(aten.polygamma) +@out_wrapper() +def meta_polygamma(n: int, self: Tensor) -> Tensor: + torch._check(n >= 0, lambda: "polygamma(n, x) does not support negative n.") + _, result_dtype = elementwise_dtypes( + self, + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, + ) + return torch.empty_like(self, dtype=result_dtype) + + +@register_meta(aten._local_scalar_dense) +def meta_local_scalar_dense(self: Tensor): + raise RuntimeError("Tensor.item() cannot be called on meta tensors") + + +@register_meta(aten.silu) +@out_wrapper(exact_dtype=True) +def silu(self: Tensor) -> Tensor: + return torch.empty_like(self) + + +@register_meta(aten.sigmoid) +@out_wrapper() +def sigmoid(self: Tensor) -> Tensor: + _, result_dtype = elementwise_dtypes( + self, + type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, + ) + return torch.empty_like(self, dtype=result_dtype) + + +def _create_grouped_mm_output_tensor(mat1, mat2, offs, out_dtype): + mat1_is_2d = mat1.dim() == 2 + mat2_is_2d = mat2.dim() == 2 + + if mat1_is_2d: + if mat2_is_2d: + out_size = [offs.size(0), mat1.size(0), mat2.size(1)] + else: + torch._check( + offs.size(0) == mat2.size(0), lambda: "matrix batch sizes have to match" + ) + out_size = [mat1.size(0), mat2.size(-1)] + else: + if mat2_is_2d: + torch._check( + offs.size(0) == mat1.size(0), lambda: "matrix batch sizes have to match" + ) + out_size = [mat1.size(1), mat2.size(1)] + else: + # regular bmm + torch._check( + mat1.size(0) == mat2.size(0), lambda: "batched dimension has to match" + ) + out_size = [mat1.size(0), mat1.size(1), mat2.size(-1)] + + out_dtype = out_dtype or mat1.dtype + + if torch.version.cuda: + alignment = 16 // out_dtype.itemsize + size_padded = (out_size[-1] + alignment - 1) // alignment * alignment + if mat1_is_2d == mat2_is_2d: + out_stride = [out_size[1] * size_padded, size_padded, 1] + else: + out_stride = [size_padded, 1] + out = torch.empty_strided( + out_size, out_stride, dtype=out_dtype, device=mat1.device + ) + else: + out = torch.empty(out_size, dtype=out_dtype, device=mat1.device) + return out + + +def _meta_grouped_mm_common( + mat_a: Tensor, + mat_b: Tensor, + scale_a: Optional[torch.Tensor], + scale_b: Optional[torch.Tensor], + offs: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + scale_result: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + use_fast_accum: bool = False, +): + torch._check( + (scale_a is None) == (scale_b is None), + lambda: "Either both scale factors are given, or none", + ) + scaled = scale_a is not None and scale_b is not None + + # Implementing all the checks from + # _grouped_mm_cuda()/_scaled_grouped_mm_cuda() code in + # aten/src/ATen/native/cuda/Blas.cpp. + + if scaled: + fp8_dtype = torch.float8_e4m3fnuz if torch.version.hip else torch.float8_e4m3fn + torch._check( + mat_a.dtype == fp8_dtype and mat_b.dtype == fp8_dtype, + lambda: f"Expected inputs of E4M3 FP8 type but got mat_a.dtype={mat_a.dtype} and mat_b.dtype={mat_b.dtype}.", # noqa: B950 + ) + else: + torch._check( + mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, + lambda: f"Expected inputs of BF16 type but got mat_a.dtype={mat_a.dtype} and mat_b.dtype={mat_b.dtype}.", # noqa: B950 + ) + + torch._check( + mat_a.dim() in [2, 3] and mat_b.dim() in [2, 3], + lambda: f"Multiplicands must be 2D or 3D but got mat_a.dim()={mat_a.dim()} and mat_b.dim()={mat_b.dim()}", # noqa: B950 + ) + + mat_a_is_2d = mat_a.dim() == 2 + mat_b_is_2d = mat_b.dim() == 2 + + if not mat_a_is_2d or not mat_b_is_2d: + torch._check( + mat_a.size(-1) == mat_b.size(-2), + "contraction dimension of mat_a and mat_b must match", + ) + + if scaled: + + def is_row_major(mat): + mat_stride = mat.stride() + return mat_stride[-2] > 1 and mat_stride[-1] == 1 + + def is_col_major(mat): + mat_stride = mat.stride() + return mat_stride[-2] == 1 and mat_stride[-1] > 1 + + torch._check( + is_row_major(mat_a), + lambda: f"Expected mat_a tensor to be row major in the last two dimensions, got strides {mat_a.stride()[-2:]}", # noqa: B950 + ) + torch._check( + is_col_major(mat_b), + lambda: f"Expected mat_b tensor to be column major in the last two dimensions, got strides {mat_b.stride()[-2:]}", # noqa: B950 + ) + + def check_valid_strides(mat_name, mat): + end_dim = mat.dim() - 1 + alignment = 16 // mat.element_size() + mat_stride = mat.stride() + if mat_stride[end_dim - 1] == 1 and mat_stride[end_dim] >= max( + 1, mat.shape[end_dim - 1] + ): + torch._check( + mat_stride[end_dim] % alignment == 0, + lambda: f"Expected {mat_name} stride along {end_dim} dim to be multiple of 16 bytes, got {mat_stride[end_dim]}.", # noqa: B950 + ) + elif mat_stride[end_dim] == 1 and mat_stride[end_dim - 1] >= max( + 1, mat.shape[end_dim] + ): + torch._check( + mat_stride[end_dim - 1] % alignment == 0, + lambda: f"Expected {mat_name} stride along {end_dim - 1} dim to be multiple of 16 bytes, got {mat_stride[end_dim - 1]}.", # noqa: B950 + ) + else: + torch._check( + False, + lambda: f"Invalid strides/sizes, got {mat_stride} for strides and {mat.shape} for sizes.", # noqa: B950 + ) + + check_valid_strides("mat_a", mat_a) + check_valid_strides("mat_b", mat_b) + + if scale_a is not None and scale_b is not None: + torch._check( + (scale_a.dtype == torch.float32 and scale_b.dtype == torch.float32) + or ( + scale_a.dtype == torch.float8_e8m0fnu + and scale_b.dtype == torch.float8_e8m0fnu + ), + lambda: f"For FP8 scales must both be float32, or for MXFP8 both scales must be float8_e8m0fnu. Got scale_a.dtype={scale_a.dtype} and scale_b.dtype={scale_b.dtype}.", # noqa: B950 + ) + is_mxfp8 = ( + scale_a.dtype == torch.float8_e8m0fnu + and scale_b.dtype == torch.float8_e8m0fnu + ) + + def round_up(x, y): + """Rounds up x to nearest multiple of y""" + return ((x + y - 1) // y) * y + + def check_scale(scale_name, scale, mat, scaled_dim, scale_multiplier=1): + if mat.dim() == 2: + torch._check( + scale.is_contiguous(), + lambda: f"Expected {scale_name} to be contiguous.", + ) + # For MXFP8, 2d tensors have variable size groups represented as subtensors, + # that are converted to blocked padded format individually. At compile time we don't know + # the group sizes yet, so we don't know the expect size of the blocked format scale. + # This limits what we can check here. + if is_mxfp8: + torch._check( + scale.dim() == mat.dim(), + lambda: f"For MXFP8, scale must have same number of dimensions as target tensor, but {scale_name} has mat.ndim={mat.ndim} and scale.ndim={scale.ndim}", # noqa: B950 + ) + else: + torch._check( + scale.dim() == 1, + lambda: f"Expected {scale_name} to be 1D tensor, but got {scale.dim()}D tensor.", + ) + torch._check( + scale.shape[0] == mat.shape[scaled_dim] * scale_multiplier, + lambda: f"Expected {scale_name} to have {mat.shape[scaled_dim] * scale_multiplier} elements, got {scale.shape[0]} elements.", # noqa: B950 + ) + else: + torch._check( + scale.stride(-1) == 1, + lambda: f"Expected {scale_name} to be contiguous in the last dimension.", + ) + torch._check( + scale.shape[0] == mat.shape[0], + lambda: f"Expected {scale_name} batch dimension to be {mat.shape[0]}, got {scale.shape[0]}.", + ) + # For MXFP8, 3d tensors have static 'groups' (stack of 2d tensors) so we can know the expected blocked + # scale sizes at compile time. + if is_mxfp8: + torch._check( + mat.ndim == scale.ndim, + lambda: f"For MXFP8, scale should have same number of dimensions as target tensor, but {scale_name} has mat.ndim={mat.ndim} and scale.ndim={scale.ndim}", # noqa: B950 + ) + # TODO: This logic only holds for RHS tensor in 2d-3d case. + # We'll need to update it to handle LHS 3d tensor in 3d-2d and 3d-3d cases. + G, K, N = scale.shape + block_size = 32 + blocked_K = round_up(K / block_size, 4) + blocked_N = round_up(N, 128) + torch._check( + mat.shape[-2] == blocked_K and mat.shape[-1] == blocked_N, + lambda: f"For MXFP8, expected mat.shape={mat.shape} to have scale shape of ({G},{blocked_K},{blocked_N}), but got {scale.shape}", # noqa: B950 + ) + else: + torch._check( + scale.dim() == 2, + lambda: f"Expected {scale_name} to be 2D tensor, but got {scale.dim()}D tensor.", + ) + torch._check( + scale.shape[1] == mat.shape[1 + scaled_dim], + lambda: f"Expected {scale_name} non-batch dimension to be {mat.shape[1 + scaled_dim]}, got {scale.shape[1]}.", # noqa: B950 + ) + + scale_multiplier = ( + offs.shape[0] if offs is not None and mat_a_is_2d and mat_b_is_2d else 1 + ) + check_scale("scale_a", scale_a, mat_a, 0, scale_multiplier) + check_scale("scale_b", scale_b, mat_b, 1, scale_multiplier) + + torch._check( + scale_result is None, + lambda: "Scale result tensor provided, but it is not supported yet.", + ) + + if mat_a_is_2d or mat_b_is_2d: + torch._check( + offs is not None, + lambda: f"Offsets tensor not provided, but is needed for {mat_a.dim()}D/{mat_b.dim()}D multiplicand layouts.", + ) + if offs is not None: # to silence Mypy + torch._check( + offs.dim() == 1, + lambda: f"Offsets tensor must be 1D, but got offs.dim()={offs.dim()}.", + ) + torch._check( + offs.dtype == torch.int32, + lambda: f"Offsets tensor must be integer (int32) tensor, but got {offs.dtype}.", + ) + else: + torch._check( + offs is None, + lambda: "Offsets tensor provided, but is not needed for 3D/3D multiplicand layouts.", + ) + + torch._check( + bias is None, + lambda: "Bias tensor provided, but it is not supported yet.", + ) + + torch._check( + out_dtype is None or out_dtype == torch.bfloat16, + lambda: "If output dtype provided, it must be torch.bfloat16.", + ) + + return _create_grouped_mm_output_tensor(mat_a, mat_b, offs, out_dtype) + + +@register_meta(aten._grouped_mm) +@out_wrapper() +def meta_grouped_mm( + mat_a: Tensor, + mat_b: Tensor, + offs: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + out_dtype: Optional[torch.dtype] = None, +) -> Tensor: + return _meta_grouped_mm_common( + mat_a, + mat_b, + scale_a=None, + scale_b=None, + offs=offs, + bias=bias, + scale_result=None, + out_dtype=out_dtype, + ) + + +@register_meta([aten._scaled_grouped_mm]) +def meta_scaled_grouped_mm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + offs: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + scale_result: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + use_fast_accum: bool = False, +): + return _meta_grouped_mm_common( + mat_a, + mat_b, + scale_a=scale_a, + scale_b=scale_b, + offs=offs, + bias=bias, + scale_result=scale_result, + out_dtype=out_dtype, + use_fast_accum=use_fast_accum, + ) + + +@register_meta(aten._softmax) +@out_wrapper() +def softmax(x: Tensor, dim: int, half_to_float: bool) -> Tensor: + if half_to_float: + assert x.dtype == torch.half + computation_dtype, result_dtype = utils.elementwise_dtypes( + x, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + result_dtype = result_dtype if not half_to_float else computation_dtype + res = torch.empty_like(x, dtype=result_dtype, memory_format=torch.contiguous_format) + return res + + +@register_meta(aten.constant_pad_nd) +@out_wrapper() +def _constant_pad_nd_meta(input, pad, value=0): + # same checks as decomposition in torch/_refs/__init__.py:constant_pad_nd() + torch._check( + len(pad) % 2 == 0, + lambda: f"Length of pad must be even but instead it equals {len(pad)}", + ) + + input_sizes = input.shape + l_inp = len(input_sizes) + l_pad = len(pad) // 2 + l_diff = l_inp - l_pad + + torch._check( + l_inp >= l_pad, + lambda: "Length of pad should be no more than twice the number of " + f"dimensions of the input. Pad length is {len(pad)} while the input has " + f"{l_inp} dimensions.", + ) + + if all(isinstance(p, utils.IntWithoutSymInt) and p <= 0 for p in pad): + c_input = input + for i in range(l_diff, l_inp): + pad_idx = 2 * (l_inp - i - 1) + if pad[pad_idx] < 0: + c_input = c_input.narrow( + i, -pad[pad_idx], c_input.shape[i] + pad[pad_idx] + ) + + if pad[pad_idx + 1] < 0: + c_input = c_input.narrow(i, 0, c_input.shape[i] + pad[pad_idx + 1]) + + return c_input.clone() + + new_shape = list(input_sizes[:l_diff]) + for i in range(l_pad): + pad_idx = len(pad) - ((i + 1) * 2) + new_dim = input_sizes[l_diff + i] + pad[pad_idx] + pad[pad_idx + 1] + torch._check( + new_dim >= 0, + lambda: f"The input size {input_sizes[l_diff + i]}, plus negative padding " + f"{pad[pad_idx]} and {pad[pad_idx + 1]} resulted in a negative output size, " + f"which is invalid. Check dimension {l_diff + i} of your input.", + ) + new_shape.append(new_dim) + + return torch.empty( + new_shape, + dtype=input.dtype, + device=input.device, + requires_grad=input.requires_grad, + memory_format=suggest_memory_format(input), + ) + + +@register_meta(aten.embedding) +@out_wrapper() +def embedding( + weight: Tensor, + indices: Tensor, + padding_idx: int = -1, + scale_grad_by_freq: bool = False, + sparse: bool = False, +) -> Tensor: + assert weight.dim() == 2, "'weight' must be 2-D" + weight_shape = weight.shape + indices_shape = indices.shape + + if indices.ndim == 0: + out_shape: tuple[int, ...] = (weight_shape[1],) + elif indices.ndim == 1: + out_shape = (indices_shape[0], weight_shape[1]) + else: + out_shape = (*indices_shape, weight_shape[1]) + + out_dtype = weight.dtype + return weight.new_empty(out_shape, dtype=out_dtype) + + +@register_meta(aten._jagged_to_padded_dense_forward.default) +def meta__jagged_to_padded_dense_forward( + values: Tensor, + offsets: list[Tensor], + max_lengths: list[int], + padding_value: float = 0.0, +): + # only one jagged dim is supported for now + assert len(offsets) == 1 + assert len(max_lengths) == 1 + + B = offsets[0].shape[0] - 1 + S = max_lengths[0] + output_shape = (B, S, *values.shape[1:]) + return values.new_empty(output_shape) + + +def _create_unary_float_meta_func(func): + @register_meta(func) + @out_wrapper() + def _f(x): + return elementwise_meta( + x, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ) + + return _f + + +def _create_binary_float_meta_func(func): + @register_meta(func) + @out_wrapper() + def _f(x, y): + return elementwise_meta( + x, y, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ) + + return _f + + +_create_unary_float_meta_func(aten.special_airy_ai) +_create_unary_float_meta_func(aten.special_bessel_y0) +_create_unary_float_meta_func(aten.special_bessel_y1) +_create_unary_float_meta_func(aten.special_modified_bessel_i0) +_create_unary_float_meta_func(aten.special_modified_bessel_i1) +_create_unary_float_meta_func(aten.special_modified_bessel_k0) +_create_unary_float_meta_func(aten.special_modified_bessel_k1) +_create_unary_float_meta_func(aten.special_scaled_modified_bessel_k0) +_create_unary_float_meta_func(aten.special_scaled_modified_bessel_k1) + + +_create_binary_float_meta_func(aten.special_chebyshev_polynomial_t) +_create_binary_float_meta_func(aten.special_chebyshev_polynomial_u) +_create_binary_float_meta_func(aten.special_chebyshev_polynomial_v) +_create_binary_float_meta_func(aten.special_chebyshev_polynomial_w) +_create_binary_float_meta_func(aten.special_shifted_chebyshev_polynomial_t) +_create_binary_float_meta_func(aten.special_shifted_chebyshev_polynomial_u) +_create_binary_float_meta_func(aten.special_shifted_chebyshev_polynomial_v) +_create_binary_float_meta_func(aten.special_shifted_chebyshev_polynomial_w) +_create_binary_float_meta_func(aten.special_hermite_polynomial_h) +_create_binary_float_meta_func(aten.special_hermite_polynomial_he) +_create_binary_float_meta_func(aten.special_laguerre_polynomial_l) +_create_binary_float_meta_func(aten.special_legendre_polynomial_p) + + +def _register_inplace_meta(fn): + @wraps(fn) + def _fn(self, *args, **kwargs): + out = fn(self, *args, **kwargs) + check_inplace_broadcast(self.shape, out.shape) + return self + + inplace_name = f"{fn.__name__}_" + _fn.__name__ = inplace_name + _fn = register_meta(getattr(aten, inplace_name))(_fn) # type: ignore[assignment] + + return _fn + + +@register_meta(aten.lerp) +@out_wrapper() +def lerp(start, end, weight): + torch._check( + start.dtype == end.dtype, + lambda: f"expected dtype {start.dtype} for `end`, but got dtype {end.dtype}", + ) + args = [start, end] + if isinstance(weight, TensorLike): + if weight.ndim != 0: + torch._check( + start.dtype == weight.dtype, + lambda: f"expected dtype {start.dtype} for `weight`, but got dtype {weight.dtype}", + ) + args.append(weight) + return elementwise_meta( + *args, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta(aten.addcmul) +@out_wrapper() +def addcmul(input, tensor1, tensor2, *, value=1): + return elementwise_meta( + input, tensor1, tensor2, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +@register_meta(aten.addcdiv) +@out_wrapper() +def addcdiv(input, tensor1, tensor2, *, value=1): + torch._check( + not ( + utils.is_integer_dtype(tensor1.dtype) + and utils.is_integer_dtype(tensor2.dtype) + ), + lambda: ( + "Integer division with addcdiv is no longer supported, and in a future ", + "release addcdiv will perform a true division of tensor1 and tensor2. ", + "The historic addcdiv behavior can be implemented as ", + "(input + value * torch.trunc(tensor1 / tensor2)).to(input.dtype) ", + "for integer inputs and as ", + "(input + value * tensor1 / tensor2) for float inputs. ", + "The future addcdiv behavior is just the latter implementation: ", + "(input + value * tensor1 / tensor2), for all dtypes.", + ), + ) + return elementwise_meta( + input, tensor1, tensor2, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ) + + +lerp_ = _register_inplace_meta(aten.lerp) +addcmul_ = _register_inplace_meta(aten.addcmul) +addcdiv_ = _register_inplace_meta(aten.addcdiv) + + +# We must also trigger meta registrations from PrimTorch ref +# decompositions +import torch._refs +import torch._refs.nn.functional +import torch._refs.special + + +def activate_meta(): + activate_meta_table = {} + + # For a given op, we pick the most specific decomp function from + # global_decomp_table in the precedence order of meta > post_autograd > pre_autograd + for type in ["meta", "post_autograd", "pre_autograd"]: + registry = global_decomposition_table[type] + + for opo in registry: + if opo not in activate_meta_table: + activate_meta_table[opo] = registry[opo] + + for op_overload, fn in activate_meta_table.items(): + # Don't register meta for HigherOrderOp's decomp. + # We can reconsider this in the future, but in general, + # the way you do a meta for a HigherOrderOp is different from + # OpOverload. + if isinstance(op_overload, torch._ops.HigherOrderOperator): + continue + assert isinstance(op_overload, OpOverload) + + op_overload.py_impl(torch._C.DispatchKey.Meta)(fn) + + if torch._C._dispatch_has_kernel_for_dispatch_key( + op_overload.name(), "CompositeImplicitAutograd" + ): + # Internally, we shouldn't be registering meta kernels for any operators that + # have CompositeImplicitAutograd kernels. + # Instead, we should be letting those decompositions run, and writing meta kernels + # only for the base operators. + if op_overload in global_decomposition_table["meta"]: + raise RuntimeError( + f"{op_overload} is a CompositeImplicitAutograd op, we shouldn't " + "register meta function for it. Instead, we should let the decomposition run and write " + "meta kernels for the base operators." + ) + elif op_overload.is_view: + # Attempting to register a python meta kernel for a view operator. + # We shouldn't do this, because the output will report as not having aliased storages. + # All view ops have meta kernels in C++ today, so we should use those instead. + pass + elif ( + op_overload.name() + in { + "aten::empty_strided", # causing infinite recursion, test_meta.py + "aten::clone", # causing infinite recursion + "aten::_to_copy", # causing infinite recursion, test_serialization.py -k test_tensor_subclass_getstate_overwrite # noqa: B950 + "aten::copy_", # Exception not raised, test_torch.py -k test_storage_meta_errors_cpu_int64 # noqa: B950 + "aten::constant_pad_nd", # requires_grad mismatch, test_ops.py -k test_fake_crossref_backward_amp_istft_cuda_float32 # noqa: B950 + "aten::rot90", # requires_grad mismatch! test_ops.py -k test_fake_crossref_backward_amp_rot90_cuda_float32 # noqa: B950 + "aten::as_strided_scatter", # requires_grad mismatch, test_ops.py -k test_fake_crossref_backward_no_amp_as_strided_scatter_cuda_float32 # noqa: B950 + } + ): + pass + else: + if "mkldnn::" in op_overload.name(): + _meta_lib_dont_use_me_use_register_meta_for_mkldnn.impl(op_overload, fn) + elif "mkl::" in op_overload.name(): + _meta_lib_dont_use_me_use_register_meta_for_mkl.impl(op_overload, fn) + elif "onednn::" in op_overload.name(): + _meta_lib_dont_use_me_use_register_meta_for_onednn.impl(op_overload, fn) + elif "quantized::" in op_overload.name(): + _meta_lib_dont_use_me_use_register_meta_for_quantized.impl( + op_overload, fn + ) + else: + _meta_lib_dont_use_me_use_register_meta.impl(op_overload, fn) + + +activate_meta() diff --git a/venv/lib/python3.10/site-packages/torch/_namedtensor_internals.py b/venv/lib/python3.10/site-packages/torch/_namedtensor_internals.py new file mode 100644 index 0000000000000000000000000000000000000000..16d04f181525d45259bcb77d1e215c1cd1118632 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_namedtensor_internals.py @@ -0,0 +1,159 @@ +# mypy: allow-untyped-defs +from collections import OrderedDict + + +""" +This file contains helper functions that implement experimental functionality +for named tensors in python. All of these are experimental, unstable, and +subject to change or deletion. +""" + + +def check_serializing_named_tensor(tensor): + if tensor.has_names(): + raise RuntimeError( + "NYI: Named tensors don't support serialization. Please drop " + "names via `tensor = tensor.rename(None)` before serialization." + ) + + +def build_dim_map(tensor): + """Returns a map of { dim: dim_name } where dim is a name if the dim is named + and the dim index otherwise.""" + return OrderedDict( + [(idx if name is None else name, name) for idx, name in enumerate(tensor.names)] + ) + + +def unzip_namedshape(namedshape): + if isinstance(namedshape, OrderedDict): + namedshape = namedshape.items() + if not hasattr(namedshape, "__iter__") and not isinstance(namedshape, tuple): + raise RuntimeError( + f"Expected namedshape to be OrderedDict or iterable of tuples, got: {type(namedshape)}" + ) + if len(namedshape) == 0: + raise RuntimeError("Expected namedshape to non-empty.") + return zip(*namedshape) + + +def namer_api_name(inplace): + if inplace: + return "rename_" + else: + return "rename" + + +def is_ellipsis(item): + return item == Ellipsis or item == "..." + + +def single_ellipsis_index(names, fn_name): + ellipsis_indices = [i for i, name in enumerate(names) if is_ellipsis(name)] + if len(ellipsis_indices) >= 2: + raise RuntimeError( + f"{fn_name}: More than one Ellipsis ('...') found in names (" + f"{names}). This function supports up to one Ellipsis." + ) + if len(ellipsis_indices) == 1: + return ellipsis_indices[0] + return None + + +def expand_single_ellipsis(numel_pre_glob, numel_post_glob, names): + return names[numel_pre_glob : len(names) - numel_post_glob] + + +def replace_ellipsis_by_position(ellipsis_idx, names, tensor_names): + globbed_names = expand_single_ellipsis( + ellipsis_idx, len(names) - ellipsis_idx - 1, tensor_names + ) + return names[:ellipsis_idx] + globbed_names + names[ellipsis_idx + 1 :] + + +def resolve_ellipsis(names, tensor_names, fn_name): + """ + Expands ... inside `names` to be equal to a list of names from `tensor_names`. + """ + ellipsis_idx = single_ellipsis_index(names, fn_name) + if ellipsis_idx is None: + return names + return replace_ellipsis_by_position(ellipsis_idx, names, tensor_names) + + +def update_names_with_list(tensor, names, inplace): + # Special case for tensor.rename(None) + if len(names) == 1 and names[0] is None: + return tensor._update_names(None, inplace) + + return tensor._update_names( + resolve_ellipsis(names, tensor.names, namer_api_name(inplace)), inplace + ) + + +def update_names_with_mapping(tensor, rename_map, inplace): + dim_map = build_dim_map(tensor) + for old_dim in rename_map.keys(): + new_dim = rename_map[old_dim] + if old_dim in dim_map.keys(): + dim_map[old_dim] = new_dim + else: + raise RuntimeError( + f"{namer_api_name(inplace)}: Tried to rename dim '{old_dim}' to dim " + f"{new_dim} in Tensor[{tensor.names}] but dim '{old_dim}' does not exist" + ) + return tensor._update_names(tuple(dim_map.values()), inplace) + + +def update_names(tensor, names, rename_map, inplace): + """There are two usages: + + tensor.rename(*names) returns a view on tensor with named dims `names`. + `names` must be of length `tensor.dim()`; otherwise, if '...' is in `names`, + then it is expanded greedily to be equal to the corresponding names from + `tensor.names`. + + For example, + ``` + >>> # xdoctest: +SKIP + >>> x = torch.empty(2, 3, 5, 7, names=('N', 'C', 'H', 'W')) + >>> x.rename('...', 'height', 'width').names + ('N', 'C', 'height', 'width') + + >>> # xdoctest: +SKIP + >>> x.rename('batch', '...', 'width').names + ('batch', 'C', 'H', 'width') + + ``` + + tensor.rename(**rename_map) returns a view on tensor that has rename dims + as specified in the mapping `rename_map`. + + For example, + ``` + >>> # xdoctest: +SKIP + >>> x = torch.empty(2, 3, 5, 7, names=('N', 'C', 'H', 'W')) + >>> x.rename(W='width', H='height').names + ('N', 'C', 'height', 'width') + + ``` + + Finally, tensor.rename has an in-place version called tensor.rename_. + """ + has_names = len(names) > 0 + has_rename_pairs = bool(rename_map) + if has_names and has_rename_pairs: + raise RuntimeError( + f"{namer_api_name(inplace)}: This function takes either positional " + f"args or keyword args, but not both. Use tensor.{namer_api_name(inplace)}(*names) " + f"to name dims and tensor.{namer_api_name(inplace)}(**rename_map) to rename " + "dims." + ) + + # Special case for tensor.rename(*[]), which is valid for a 0 dim tensor. + if not has_names and not has_rename_pairs: + return update_names_with_list(tensor, names, inplace) + + if has_names: + return update_names_with_list(tensor, names, inplace) + return update_names_with_mapping(tensor, rename_map, inplace) diff --git a/venv/lib/python3.10/site-packages/torch/_ops.py b/venv/lib/python3.10/site-packages/torch/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b351aa17dfa7624f7f0187f98823efd5983d3ea3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_ops.py @@ -0,0 +1,1495 @@ +# mypy: allow-untyped-defs +import abc +import contextlib +import ctypes +import importlib +import inspect +import sys +import types +from collections.abc import Iterator +from functools import cached_property +from typing import ( + Any, + Callable, + ClassVar, + final, + Generic, + Optional, + TYPE_CHECKING, + Union, +) +from typing_extensions import Concatenate, ParamSpec, TypeVar + +import torch +import torch.utils._pytree as pytree +from torch import _utils_internal +from torch._C import _dispatch_is_included_in_alias as is_included_in_alias, DispatchKey +from torch._functorch.pyfunctorch import dispatch_functorch, TransformType +from torch.utils._python_dispatch import TorchDispatchMode + + +if TYPE_CHECKING: + from torch._subclasses.functional_tensor import BaseFunctionalizeAPI + + +_T = TypeVar("_T", default=Any) +_P = ParamSpec("_P", default=...) + + +# Query `hasattr` only once. +_SET_GLOBAL_FLAGS = hasattr(sys, "getdlopenflags") and hasattr(sys, "setdlopenflags") + + +@contextlib.contextmanager +def dl_open_guard(): + """ + Context manager to set the RTLD_GLOBAL dynamic linker flag while we open a + shared library to load custom operators. + """ + if not _SET_GLOBAL_FLAGS: + yield + return + old_flags = sys.getdlopenflags() + sys.setdlopenflags(old_flags | ctypes.RTLD_GLOBAL) + try: + yield + finally: + sys.setdlopenflags(old_flags) + + +class OperatorBase: + """ + Base class for OpOverload (which represents C++ ATen operators) and HigherOrderOperator + (which represents Python-only operators that are unrepresentable in TorchScript). + """ + + def __init__(self): + # The dispatch cache precomputes a mapping of dispatch key that the + # dispatcher wants to dispatch to, to an actual implementation of the + # dispatch key. Confusingly, the actual implementation could *also* be a + # dispatch key, but in this case, this refers to the C++ kernel that + # was registered to some dispatch key. Aliases are permitted in the + # latter but not the former; for example, you might lookup the + # entry for AutogradCPU, and this maps you to the Autograd key for + # the generic autograd kernel that works for all devices. Since this + # is the Python dispatcher, you can also put an arbitrary Python + # callable to call instead. This handler gets precisely the + # args/kwargs that the operator was __call__'ed with. + # NB: This name is hard-coded in torch/csrc/autograd/python_variable.cpp + # for use with OpOverload; cache lookup is done entirely from C++ + # for speed. + # TODO: The cache is NOT currently used by HigherOrderOperator, but it should! + self._dispatch_cache: dict[ + DispatchKey, Union[DispatchKey, Callable[..., Any]] + ] = {} + + # This table allows you to override the behavior of a particular + # dispatch key to call a custom Python function, rather than the + # ordinary C++ configured behavior. This is the raison d'etre of # codespell:ignore + # Python dispatcher: to let you program the dispatcher from Python + # in case you need something unusual, and don't want to clobber + # the existing registrations using the Python operator registration + # API. + self.py_kernels: dict[DispatchKey, Callable[..., Any]] = {} + + # This table allows you to override the behavior of a particular + # operator for a particular TorchDispatchMode. In practice, + # we are using this mostly for ProxyTensorMode. Modes can be + # thought of as an open world extension of dispatch keys, so it + # makes sense that you should be able to register them, the same + # way you can register dispatch keys. + self.python_key_table: dict[ + type[Union[TorchDispatchMode, torch.Tensor]], Callable[..., Any] + ] = {} + + # This table allows you to override the behavior of functorch + # transformations. NB: this currently only does something for + # HigherOrderOperator + self.functorch_table = {} + + def __call__(self, *args, **kwargs): + raise NotImplementedError + + def has_kernel_for_dispatch_key(self, k): + return k in self.py_kernels + + def has_kernel_for_any_dispatch_key(self, ks): + for k in self.py_kernels: + if not torch._C._dispatch_is_alias_key(k) and ks.has(k): + return True + return False + + def py_impl( + self, + k: Union[ + type[TorchDispatchMode], + type[torch.Tensor], + TransformType, + DispatchKey, + ], + ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + def inner(fn: Callable[_P, _T]) -> Callable[_P, _T]: + if inspect.isclass(k) and ( + issubclass(k, TorchDispatchMode) or issubclass(k, torch.Tensor) + ): + assert k not in self.python_key_table + # TODO(voz): Should we replace setting DispatchKey.Python entirely with setting mode keys? + self.python_key_table[k] = fn + self._dispatch_cache.clear() + return fn + + if isinstance(k, TransformType): + assert k not in self.functorch_table + self.functorch_table[k] = fn + return fn + + assert isinstance(k, DispatchKey) + assert k != DispatchKey.Python, ( + "Please register a mode for the DispatchKey.Python key instead." + ) + + if k in self.py_kernels: + raise RuntimeError( + f"Trying to override a python impl for {k} on operator {self.name()}" + ) + self.py_kernels[k] = fn + self._dispatch_cache.clear() + return fn + + return inner + + # Registers an implementation to all **3** variants of functionalization that we have: + # - DispatchKey.Functionalize + # - functorch.TransformType.Functionalize + # - FunctionalTensorMode + # Example: + # @py_functionalize_impl + # def functionalize_rule(ctx, inner_f, *args): + # args_unwrapped = ctx.unwrap_tensors(args) + # with ctx.redispatch_to_next(): + # out = ctx.functionalize(inner_f)(*args_unwrapped) + # return ctx.wrap_tensors(out) + def py_functionalize_impl( + self, fn: Callable[Concatenate["BaseFunctionalizeAPI", _P], _T] + ) -> Callable[Concatenate["BaseFunctionalizeAPI", _P], _T]: + from torch._subclasses.functional_tensor import ( + CppFunctionalizeAPI, + FunctionalTensorMode, + FunctorchFunctionalizeAPI, + PythonFunctionalizeAPI, + ) + + # Construct our three flavors of functionalization, + # each of which have slightly different wrap/unwrap/redispatch policies + def functionalize_dk_fn(*args: _P.args, **kwargs: _P.kwargs) -> _T: + return fn(CppFunctionalizeAPI(), *args, **kwargs) + + def functionalize_dispatch_mode_fn( + mode: Optional[FunctionalTensorMode], *args: _P.args, **kwargs: _P.kwargs + ) -> _T: + return fn(PythonFunctionalizeAPI(mode), *args, **kwargs) + + def functionalize_functorch_fn( + interpreter, *args: _P.args, **kwargs: _P.kwargs + ) -> _T: + return fn(FunctorchFunctionalizeAPI(interpreter), *args, **kwargs) + + self.py_impl(DispatchKey.Functionalize)(functionalize_dk_fn) + self.py_impl(FunctionalTensorMode)(functionalize_dispatch_mode_fn) + self.py_impl(TransformType.Functionalize)(functionalize_functorch_fn) + + return fn + + def name(self): + raise NotImplementedError + + +# Equivalent to computeDispatchTableEntryWithDebug +def resolve_key(op: OperatorBase, k: DispatchKey): # type: ignore[valid-type] + # 1. (Direct) operator registration + if op.has_kernel_for_dispatch_key(k): + return k + # 2.1 Use CompositeExplicitAutogradNonFunctional kernel if available + cand = DispatchKey.CompositeExplicitAutogradNonFunctional + if ( + k == DispatchKey.Undefined or is_included_in_alias(k, cand) + ) and op.has_kernel_for_dispatch_key(cand): + return cand + # 2.2 Use CompositeExplicitAutograd kernel if available + cand = DispatchKey.CompositeExplicitAutograd + if ( + k == DispatchKey.Undefined or is_included_in_alias(k, cand) + ) and op.has_kernel_for_dispatch_key(cand): + return cand + has_backend_kernel = op.has_kernel_for_any_dispatch_key( + torch._C._dispatch_get_backend_keyset_from_autograd(k) + ) or op.has_kernel_for_dispatch_key(DispatchKey.CompositeExplicitAutograd) + # 2.3. Use CompositeImplicitAutograd kernel if available + cand = DispatchKey.CompositeImplicitAutogradNestedTensor + if ( + (k != DispatchKey.Undefined and is_included_in_alias(k, cand)) + and op.has_kernel_for_dispatch_key(cand) + and not has_backend_kernel + ): + return cand + cand = DispatchKey.CompositeImplicitAutograd + if ( + k == DispatchKey.Undefined or is_included_in_alias(k, cand) + ) and op.has_kernel_for_dispatch_key(cand): + if k == DispatchKey.AutogradOther and op.has_kernel_for_any_dispatch_key( + torch._C._dispatch_autogradother_backends + ): + raise RuntimeError("ambiguous autogradother kernel") + elif not has_backend_kernel: + return cand + # 2.4. For autograd backend keys, use kernel from DispatchKey::Autograd if available + cand = DispatchKey.Autograd + if is_included_in_alias(k, cand) and op.has_kernel_for_dispatch_key(cand): + return cand + # 2.5 Use kernel from DispatchKey::FuncTorchBatchedDecomposition if available + cand = DispatchKey.FuncTorchBatchedDecomposition + if is_included_in_alias(k, cand) and op.has_kernel_for_dispatch_key(cand): + return cand + # Backend fallback + if torch._C._dispatch_has_backend_fallback(k): + # The dispatch key itself will implicitly route to backend fallback. + # This is probably not great for the pure Python implementation. + return k + raise NotImplementedError(f"could not find kernel for {op} at dispatch key {k}") + + +_higher_order_ops: dict[str, "HigherOrderOperator"] = {} + +_HIGHER_ORDER_OP_DEFAULT_FALLTHROUGH_DISPATCH_KEYS = [ + DispatchKey.PythonDispatcher, # type: ignore[attr-defined] + DispatchKey.PythonTLSSnapshot, # type: ignore[attr-defined] + DispatchKey.ADInplaceOrView, + DispatchKey.BackendSelect, + DispatchKey.AutocastCPU, # type: ignore[attr-defined] + DispatchKey.AutocastCUDA, # type: ignore[attr-defined] + DispatchKey.AutocastXPU, # type: ignore[attr-defined] +] + + +class HigherOrderOperator(OperatorBase, abc.ABC): + # The HigherOrderOperator will appear as torch.ops.higher_order.{name} + # + # If you're creating a new HigherOrderOperator, please do not change the + # default. Adding operators to the global torch.ops namespace is a bad + # practice due to name collisions. + def __init__(self, name, *, cacheable=False): + super().__init__() + if type(self) is HigherOrderOperator: + raise RuntimeError( + "Direct instantiation of HigherOrderOperator is not allowed. Please subclass it." + ) + self._name = name + + # Make _OPNamespace not scream, this whole name based association needs a good hard look + self.__name__ = name + _higher_order_ops[name] = self + self._ns = "higher_order" + self.__module__ = "torch.ops.higher_order" + self._cacheable = cacheable + + self.non_fallthrough_keys = torch._C._dispatch_keyset_full() + + for dispatch_key in _HIGHER_ORDER_OP_DEFAULT_FALLTHROUGH_DISPATCH_KEYS: + self.fallthrough(dispatch_key) + + # [NOTE] We have to register pre-dispatch key implementation + # because sometimes HOP use aot-dispatch tracing to detect certain + # mutations. This is problematic when we are functionalizing HOP + # during pre-dispatch because when the inner tracer starts, it will see + # that PreDispatch key is still active. In that case, we just redispatch + # it to next key. This is only safe to do when PreDispatch key stack has no + # active modes. + + def py_impl( + self, + k: Union[ + type[TorchDispatchMode], + type[torch.Tensor], + TransformType, + DispatchKey, + ], + ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + if isinstance(k, DispatchKey) and not self.non_fallthrough_keys.has(k): + self.non_fallthrough_keys = self.non_fallthrough_keys.add(k) + return super().py_impl(k) + + def py_autograd_impl( + self, + fn: Callable[_P, _T], + ) -> Callable[_P, _T]: + def maybe_run_autograd(*args: _P.args, **kwargs: _P.kwargs) -> _T: + if not torch.is_grad_enabled() or pytree.tree_all_only( + torch.Tensor, + lambda t: not t.requires_grad, # type: ignore[union-attr] + (*args, kwargs), + ): + with torch._C._AutoDispatchBelowAutograd(): + return self(*args, **kwargs) + + from torch._higher_order_ops.utils import _has_gen_schema + + if _has_gen_schema(self): + schema = self.gen_schema(*args, **kwargs) + if any(arg.is_write for arg in schema.arguments): + raise RuntimeError( + f"The {self.name()} HigherOrderOperator does not currently support training " + "with in-place input or buffer mutations " + "If you require this feature, please submit an issue to PyTorch. " + "Alternatively, consider creating your own custom autograd.Function. " + ) + + return fn(*args, **kwargs) + + self.py_impl(DispatchKey.Autograd)(maybe_run_autograd) + + return fn + + @property + def namespace(self): + return self._ns + + @final + def cacheable(self) -> bool: + from torch._functorch.autograd_function import AutogradFunctionApply + + return ( + self._cacheable + or f"{self.__module__}.{self.__name__}" + in torch._inductor.config.unsafe_marked_cacheable_functions + or ( + isinstance(self, AutogradFunctionApply) + and torch._functorch.config.autograd_cache_allow_custom_autograd_functions + ) + ) + + def fallthrough(self, dispatch_key): + self.non_fallthrough_keys = self.non_fallthrough_keys.remove(dispatch_key) + + # Use positional-only argument to avoid naming collide with custom ops arguments + # that are named "self". + def dispatch(self, /, dispatch_key, *args, **kwargs): + from torch.utils._python_dispatch import _get_current_dispatch_mode + + if dispatch_key in self._dispatch_cache: + kernel = self._dispatch_cache[dispatch_key] + assert not isinstance(kernel, DispatchKey) + return kernel(*args, **kwargs) + + if dispatch_key == DispatchKey.FuncTorchDynamicLayerFrontMode: + return dispatch_functorch(self, args, kwargs) + + if dispatch_key == DispatchKey.Python: + # Keep the following 1:1 with handle_torch_function_no_python_arg_parser + # in torch/csrc/utils/python_arg_parser.cpp + + overloaded_args_list = [] + + def has_python_key(tensor): + return torch._C._dispatch_keys(tensor).has("Python") + + def check_overloaded(arg): + if isinstance(arg, torch.Tensor) and has_python_key(arg): + overloaded_args_list.append(arg) + + for arg in (*args, *kwargs.values()): + check_overloaded(arg) + if isinstance(arg, (list, tuple)): + for a in arg: + check_overloaded(a) + + overloaded_args = tuple(overloaded_args_list) + + # Step 1: dispatch on any user TorchDispatchModes + from torch.utils._python_dispatch import _pop_mode_temporarily + + curr_mode = _get_current_dispatch_mode() + if curr_mode is not None: + if type(curr_mode) in self.python_key_table: + handler = self.python_key_table[type(curr_mode)] + with _pop_mode_temporarily() as mode: + # "natural" calling convention: (mode, *args, **kwargs) + # TODO(rzou): we should support torch_dispatch calling convention too. + result = handler(mode, *args, **kwargs) + else: + if curr_mode.supports_higher_order_operators: + with _pop_mode_temporarily() as mode: + return curr_mode.__torch_dispatch__(self, [], args, kwargs) + else: + raise NotImplementedError( + f"There was no rule registered for HigherOrderOperator {self._name} and mode {curr_mode}." + f"Hint: set {curr_mode}'s supports_higher_order_operators to True." + f" This causes all higher order operators to pass through {curr_mode}'s __torch_dispatch__," + f" so handle them accordingly by" + f" adding support for HigerOrderOperators (in this case, {self._name}) in" + f" {curr_mode}.__torch_dispatch__ or" + f" returning NotImplemented when not supported." + ) + if result is not NotImplemented: + return result + + # Step 2: dispatch on any subclasses + for arg in overloaded_args: + subclass_type = type(arg) + if ( + subclass_type.__torch_dispatch__ + == torch._C._disabled_torch_dispatch_impl + ): + continue + + # In some case, people are using FakeTensor without a FakeTensorMode. + # For example, some sparse arch model has a mix of FakeTensor and real + # tensor for weights during lowering, and ppl tends to run eager evaluation + # on the model without setting up the FakeTensorMode. + # In this case, we pull FakeTensorMode impl. + if subclass_type is torch._subclasses.fake_tensor.FakeTensor: + subclass_type = torch._subclasses.fake_tensor.FakeTensorMode # type: ignore[assignment] + handler = self.python_key_table[subclass_type] + result = handler(arg.fake_mode, *args, **kwargs) # type: ignore[attr-defined] + return result + + if subclass_type in self.python_key_table: + handler = self.python_key_table[subclass_type] + # "natural" calling convention: (*args, **kwargs) + # TODO(rzou): we should support torch_dispatch calling convention too. + result = handler(*args, **kwargs) + else: + raise NotImplementedError( + f"There was no rule registered for HOP {self._name} and subclass {subclass_type}. " + f"We recommend filing an issue." + ) + if result is not NotImplemented: + return result + + # All handlers returned NotImplemented + raise TypeError( + f"HigherOrderOperator '{self._name}' is not supported for the given input types. " + f"This typically happens when using custom tensor types or dispatch modes that don't " + f"have implementations for this operation.\n\n" + f"Current mode: {curr_mode}\n" + f"Input types: {[type(a).__name__ for a in overloaded_args]}\n\n" + f"To fix this, can add support for '{self._name}' in {curr_mode}'s __torch_dispatch__\n" + ) + + functionality_key = torch._C._to_functionality_key(dispatch_key) # type: ignore[attr-defined] + if functionality_key == DispatchKey.PreDispatch: + from torch.utils._python_dispatch import _pop_mode_temporarily + + # The check for Python in the exclude set is so we properly respect `with no_dispatch()` + # calls inside of a mode. + if ( + _len_torch_dispatch_stack_pre_dispatch() > 0 + ) and not torch._C._dispatch_tls_is_dispatch_key_excluded( + DispatchKey.Python + ): + curr_mode = _get_current_dispatch_mode_pre_dispatch() + assert curr_mode is not None, ( + "Illegal invocation of dispatch on DispatchKey.PreDispatch without a mode." + ) + assert type(curr_mode) in self.python_key_table, ( + f"Current active mode {curr_mode} not registered" + ) + handler = self.python_key_table[type(curr_mode)] + with _pop_mode_temporarily(functionality_key) as mode: + return handler(mode, *args, **kwargs) + + final_key = resolve_key(self, dispatch_key) + + # This can current fail due to backend fallbacks. You just have to + # register them by hand for HigherOrderOperator. + if final_key not in self.py_kernels: + raise NotImplementedError( + f"could not find kernel for HigherOrderOperator {self._name} " + f"at dispatch key {final_key} (resolved from {dispatch_key})" + ) + + # [NOTE] We shouldn't cache PreDispatch kernel here because depending + # on what modes are active, predispatch behaviour is different. + # Also we do same thing for normal ops: + # See Note [Not Caching Per-Dispatch-Key Mode Handlers] + if dispatch_key != DispatchKey.PreDispatch: + self._dispatch_cache[dispatch_key] = self.py_kernels[final_key] + kernel = self.py_kernels[final_key] + # It's illegal to register DispatchKey to py_kernels, since there's no + # C++ kernel to call into + assert not isinstance(kernel, DispatchKey) + return kernel(*args, **kwargs) + + @abc.abstractmethod + def __call__(self, /, *args, **kwargs): + def wrapper(): + flat_args = _to_flat_tuple(args, kwargs) + if torch.overrides.has_torch_function(flat_args): + return torch.overrides.handle_torch_function( + self, flat_args, *args, **kwargs + ) + + dispatch_key_set = _compute_keyset(args, kwargs, self.non_fallthrough_keys) + return self.dispatch( + dispatch_key_set.highestPriorityTypeId(), *args, **kwargs + ) + + return wrapper() + + # NOTE [HigherOrderOprator Schema] + # Each invocation of a HigherOrderOperator (hop) should have its own schema because + # the subgraphs and the arguments can be different even for the same hop. + # + # Each hop should implement its own gen_schema method, which should + # take the same input as the __call__ method and returns a FunctionSchema. + # The schema provides a unified way to check if the hop mutates its inputs, + # which can be useful in implementing optimizations. + # + # If the hop doesn't implement the gen_schema method, + # we expect it to be functional. It should not mutate its inputs and there + # are no input, output aliasing via views or direct referencing. + def gen_schema(self, *args, **kwargs): + raise NotImplementedError( + f"HigherOrderOperator {self._name} does not implement a gen_schema. " + f"This is OK as long as the hop is functional. " + f"e.g. it should not mutate its inputs and there are no input, output aliasing " + f"via views or direct referencing." + ) + + def __str__(self): + return f"{self.name()}" + + def name(self): + return self._name + + +def _to_flat_tuple(args, kwargs): + return pytree.arg_tree_leaves(*args, **kwargs) + + +def _compute_keyset(args, kwargs, non_fallthrough_keys): + tensors = _get_tensors(args, kwargs) + return key_extractor(tensors, non_fallthrough_keys) + + +def _get_tensors(args, kwargs): + flat_all = _to_flat_tuple(args, kwargs) + tensor_args = [t for t in flat_all if isinstance(t, torch.Tensor)] + return tuple(tensor_args) + + +# Note - this should maintain identical impl to the C++ dispatcher key extraction logic +# at ATen/core/dispatch/DispatchKeyExtractor.h +def key_extractor(tensors, key_mask): + key_set = torch._C._dispatch_tls_local_include_set() + for tensor in tensors: + key_set = key_set | torch._C._dispatch_keys(tensor) + key_set = key_set - torch._C._dispatch_tls_local_exclude_set() + key_set = key_set & key_mask + return key_set + + +# Mode stack for PreDispatchKey +# it should always have three keys with +# priority given to FunctionalTensorMode and +# then ProxyTorchDispatchMode. It means that +# slot 0 belongs to ProxyTorchDispatchMode and +# slot 1 belongs to FunctionalTensorMode. +# +# SchemaCheckMode is separate from the other 2, +# and is only valid when the stack is empty. +# SchemaCheckMode is for testing purposes, and +# is meant to run in eager mode on concrete inputs, +# checking for incorrect schemas in regards to +# aliasing or mutating ops. +class _ModeStackStateForPreDispatch: + def __init__(self): + self.__infra_modes = [None, None] + self._schema_check_mode = None + + def set(self, index, mode): + assert index < len(self.__infra_modes) + self.__infra_modes[index] = mode + + def get(self, index): + assert index < len(self.__infra_modes) + return self.__infra_modes[index] + + def count(self): + return len([i for i in self.__infra_modes if i is not None]) + int( + self._schema_check_mode is not None + ) + + +_mode_stack_state_for_pre_dispatch = _ModeStackStateForPreDispatch() + + +def unset_mode_pre_dispatch(mode_key, schema_check=False): + current_mode_stack_pre_dispatch = mode_stack_state_for_pre_dispatch() + assert mode_key is None or mode_key in ( + torch._C._TorchDispatchModeKey.PROXY, + torch._C._TorchDispatchModeKey.FUNCTIONAL, + ) + if schema_check: + assert mode_key is None + + def _unset_mode(): + if mode_key == torch._C._TorchDispatchModeKey.PROXY: + current_mode = current_mode_stack_pre_dispatch.get(0) + mode_stack_state_for_pre_dispatch().set(0, None) + return current_mode + elif mode_key == torch._C._TorchDispatchModeKey.FUNCTIONAL: + current_mode = current_mode_stack_pre_dispatch.get(1) + mode_stack_state_for_pre_dispatch().set(1, None) + return current_mode + else: + current_mode = mode_stack_state_for_pre_dispatch()._schema_check_mode + mode_stack_state_for_pre_dispatch()._schema_check_mode = None + return current_mode + + current_mode = _unset_mode() + + new_pre_dispatch_len = _len_torch_dispatch_stack_pre_dispatch() + # When we are unsetting a mode, we need to check if there is + # active mode left on the PreDispatch key. If there is nothing + # active, we need to remove PreDispatch key from local dispatch include + # set. + if new_pre_dispatch_len == 0: + torch._C._dispatch_tls_set_dispatch_key_included(DispatchKey.PreDispatch, False) + + return current_mode + + +def _set_mode_pre_dispatch(mode): + from torch._subclasses.functional_tensor import FunctionalTensorMode + from torch._subclasses.schema_check_mode import SchemaCheckMode + from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode + + assert isinstance( + mode, + ( + FunctionalTensorMode, + ProxyTorchDispatchMode, + SchemaCheckMode, + ), + ) + + previous_mode_stack_len = _len_torch_dispatch_stack_pre_dispatch() + if isinstance(mode, SchemaCheckMode): + current_mode = mode_stack_state_for_pre_dispatch()._schema_check_mode + if previous_mode_stack_len > 0: + raise AssertionError( + "SchemaCheckMode for pre-dispatch must be used exclusively, found other modes on the stack" + ) + mode_stack_state_for_pre_dispatch()._schema_check_mode = mode + elif isinstance(mode, FunctionalTensorMode): + current_mode = mode_stack_state_for_pre_dispatch().get(1) + assert current_mode is None + mode_stack_state_for_pre_dispatch().set(1, mode) + else: + current_mode = mode_stack_state_for_pre_dispatch().get(0) + assert current_mode is None + mode_stack_state_for_pre_dispatch().set(0, mode) + + # When we are setting a mode, we need to check if there is + # active mode left on the PreDispatch key. If there was nothing + # active before setting this mode, it means that PreDispatch key + # was turned off. So we need to turn it on again. + if previous_mode_stack_len == 0: + torch._C._dispatch_tls_set_dispatch_key_included(DispatchKey.PreDispatch, True) + + +def _pop_mode_from_pre_dispatch(): + mode_stack = mode_stack_state_for_pre_dispatch() + pre_dispatch_len = _len_torch_dispatch_stack_pre_dispatch() + + if pre_dispatch_len == 0: + raise AssertionError("Trying to pop empty mode stack") + + if mode_stack._schema_check_mode is not None: + return unset_mode_pre_dispatch(None, schema_check=True) + if mode_stack.get(1) is not None: + return unset_mode_pre_dispatch(torch._C._TorchDispatchModeKey.FUNCTIONAL) + if mode_stack.get(0) is not None: + return unset_mode_pre_dispatch(torch._C._TorchDispatchModeKey.PROXY) + + +def _len_torch_dispatch_stack_pre_dispatch(): + return mode_stack_state_for_pre_dispatch().count() + + +def _get_dispatch_mode_pre_dispatch(mode_key): + assert mode_key in ( + torch._C._TorchDispatchModeKey.PROXY, + torch._C._TorchDispatchModeKey.FUNCTIONAL, + ) + if mode_key == torch._C._TorchDispatchModeKey.PROXY: + return mode_stack_state_for_pre_dispatch().get(0) + else: + return mode_stack_state_for_pre_dispatch().get(1) + + +def _get_current_dispatch_mode_pre_dispatch(): + if mode_stack_state_for_pre_dispatch()._schema_check_mode is not None: + return mode_stack_state_for_pre_dispatch()._schema_check_mode + else: + stack_len = mode_stack_state_for_pre_dispatch().count() + if stack_len == 2: + return mode_stack_state_for_pre_dispatch().get(1) + if stack_len == 1: + return ( + mode_stack_state_for_pre_dispatch().get(1) + if mode_stack_state_for_pre_dispatch().get(1) is not None + else mode_stack_state_for_pre_dispatch().get(0) + ) + return None + + +def mode_stack_state_for_pre_dispatch(): + global _mode_stack_state_for_pre_dispatch + return _mode_stack_state_for_pre_dispatch + + +cached_ops: set["OpOverload"] = set() + + +def add_cached_op(op_overload): + global cached_ops + cached_ops.add(op_overload) + + +def reset_cached_ops(): + global cached_ops + cached_ops.clear() + + +def get_cached_ops(): + global cached_ops + return cached_ops + + +# Each OpOverload object contains pointer to a specific operator overload, a pointer to the parent `OpOverloadPacket` object. +# You can obtain an OpOverload object through attribute query on OpOverloadPacket. +class OpOverload(OperatorBase, Generic[_P, _T]): + def __init__( + self, + overloadpacket: "OpOverloadPacket", + op: Callable[_P, _T], + op_dk: Callable[Concatenate[DispatchKey, _P], _T], + schema: torch._C.FunctionSchema, + tags: list[Any], + ) -> None: + super().__init__() + self._op = op + self._op_dk = op_dk + self._schema = schema + self._overloadpacket = overloadpacket + self._tags = tags + self._overloadname = ( + "default" if schema.overload_name == "" else schema.overload_name + ) + if tags: + self._nondeterministic_seeded = torch.Tag.nondeterministic_seeded in tags + self._name = self._schema.name + if schema.overload_name: + self._name += "." + schema.overload_name + self.__name__ = f"{self._schema.name.split('::')[1]}.{self._overloadname}" + self.__module__ = overloadpacket.__module__ + op.__module__ = overloadpacket.__module__ + self.__qualname__ = self._name + self.__annotations__ = {} + + # If the OpOverload was constructed from a Library.def in Python. + self._defined_in_python = self.__qualname__ in torch.library._defs + + # Logic replicated from aten/src/ATen/native/MathBitsFallback.h + is_write = None + for a in self._schema.arguments: + if a.alias_info is None: + continue + if is_write is None: + is_write = a.alias_info.is_write + else: + # We will conservatively call mixed mutable/non-mutable + # aliased inputs as NOT a view + is_write = a.alias_info.is_write or is_write + self.is_view = is_write is not None and not is_write + + @cached_property + def _namespace(self) -> str: + return self._schema.name.split("::", maxsplit=1)[0] + + @cached_property + def _opname(self) -> str: + return self._schema.name.split("::", maxsplit=1)[1] + + @cached_property + def _handle(self) -> torch._C._DispatchOperatorHandle: + return torch._C._dispatch_find_schema_or_throw( + self._schema.name, self._schema.overload_name + ) + + # it's a no-op since OpOverload object is immutable and must be unique for a given op overload. + def __deepcopy__(self, memo=None): + return self + + def __repr__(self): + return f"" + + # Use positional-only argument to avoid naming collision with aten ops arguments + # that are named "self". This way, all the aten ops can be called by kwargs. + def __call__(self, /, *args: _P.args, **kwargs: _P.kwargs) -> _T: + return self._op(*args, **kwargs) + + # Use positional-only argument to avoid naming collision with aten ops arguments + # that are named "self". This way, all the aten ops can be called by kwargs. + def redispatch( + self, /, keyset: torch._C.DispatchKeySet, *args: _P.args, **kwargs: _P.kwargs + ) -> _T: + return self._handle.redispatch_boxed(keyset, *args, **kwargs) # type: ignore[return-value] + + def __hash__(self): + return hash(self._op) + + # `my_namespace.my_op_name.overload_name` + def __str__(self): + return "{}.{}.{}".format(*self._schema.name.split("::"), self._overloadname) + + def has_kernel_for_dispatch_key(self, k: DispatchKey) -> bool: + return super().has_kernel_for_dispatch_key( + k + ) or torch._C._dispatch_has_kernel_for_dispatch_key(self.name(), k) + + def has_kernel_for_any_dispatch_key(self, ks: torch._C.DispatchKeySet) -> bool: + return torch._C._dispatch_has_kernel_for_any_dispatch_key( + self.name(), ks + ) or super().has_kernel_for_any_dispatch_key(ks) + + @property + def namespace(self) -> str: + return self._namespace + + def _can_decompose(self) -> bool: + dk = DispatchKey.CompositeImplicitAutograd + return dk in self.py_kernels or torch._C._dispatch_has_kernel_for_dispatch_key( + self.name(), dk + ) + + def decompose(self, *args: _P.args, **kwargs: _P.kwargs) -> _T: + dk = DispatchKey.CompositeImplicitAutograd + if dk in self.py_kernels: + # NB: This branch is not too necessary anymore, because we can + # apply Python CompositeImplicitAutograd *before* tracing + # using Python dispatcher (also taking advantage of the autograd + # formula). But it's included for completeness + return self.py_kernels[dk](*args, **kwargs) + elif torch._C._dispatch_has_kernel_for_dispatch_key(self.name(), dk): + return self._op_dk(dk, *args, **kwargs) + else: + return NotImplemented + + # Remove a dispatch key from the dispatch cache. This will force it to get + # recomputed the next time. Does nothing + # WARNING: if you register a dispatch key to py_kernels of an OpOverload, + # calling _del_dispatch on that key is NOT sufficient to apply your change, + # because a single registration may affect MULTIPLE dispatch keys (e.g., + # registering Autograd affects AutogradCPU). del_dispatch is to be used + # only if you are specifically modifying how get_dispatch handles a + # particular input 'key'. + def _uncache_dispatch(self, key: DispatchKey) -> None: + self._dispatch_cache.pop(key, None) + + # This implements the pre-computation logic for the Python dispatcher. + def _get_dispatch(self, key: DispatchKey) -> Union[DispatchKey, Callable[_P, _T]]: + # This is only called upon a cache miss + assert key not in self._dispatch_cache, f"{self} {key}" + + if key == DispatchKey.Python: + if not isinstance(self, TorchBindOpOverload) and not self.python_key_table: + self._dispatch_cache[key] = key + add_cached_op(self) + return key + + def handler(*args: _P.args, **kwargs: _P.kwargs) -> _T: + from torch.utils._python_dispatch import _get_current_dispatch_mode + + # TODO: We also need to handle tensor subclasses here + # TODO(voz): We should walk all the nodes here / turn it into a list, topmode is ok for now. + curr_mode = type(_get_current_dispatch_mode()) + assert curr_mode is not None, ( + "Illegal invocation of dispatch on DispatchKey.Python without a mode." + ) + + if curr_mode not in self.python_key_table: + if isinstance(self, TorchBindOpOverload): + with ( + torch.utils._python_dispatch._pop_mode_temporarily() as mode + ): + return torch._library.utils.handle_dispatch_mode( + mode, self, *args, **kwargs + ) + else: + return self._op_dk(key, *args, **kwargs) + + with torch.utils._python_dispatch._pop_mode_temporarily() as mode: + return self.python_key_table[curr_mode](mode, *args, **kwargs) + + self._dispatch_cache[key] = handler + add_cached_op(self) + return handler + + functionality_key = torch._C._to_functionality_key(key) # type: ignore[attr-defined] + if functionality_key == DispatchKey.PreDispatch: + curr_stack_len = _len_torch_dispatch_stack_pre_dispatch() + # The check for Python in the exclude set is so we properly respect `with no_dispatch()` + # calls inside of a mode. + if ( + curr_stack_len > 0 + and not torch._C._dispatch_tls_is_dispatch_key_excluded( + DispatchKey.Python + ) + ): + + def handler(*args: _P.args, **kwargs: _P.kwargs) -> _T: + @contextlib.contextmanager + def _temporarily_pop_modes_from_pre_dispatch(): + top_mode = _pop_mode_from_pre_dispatch() + try: + yield top_mode + finally: + _set_mode_pre_dispatch(top_mode) + + with _temporarily_pop_modes_from_pre_dispatch() as curr_mode: + return torch._library.utils.handle_dispatch_mode( + curr_mode, self, *args, **kwargs + ) + + # Note [Not Caching Per-Dispatch-Key Mode Handlers] + # Note that we're not caching this handler. There isn't really a point, since the slow bit + # is the handler itself (in python). + # Also, not caching means that we don't have to reset the cache when any existing + # modes go out of scope (which in of itself takes time to loop through all operators). + return handler + + final_key = resolve_key(self, key) + + # See Note [Not Caching Per-Dispatch-Key Mode Handlers] + cache_result = key != DispatchKey.PreDispatch + + # TODO: We could potentially have lots of debugging wrappers against + # dispatch keys; design some general registration mechanism instead of + # having if statement for each of them + if key == DispatchKey.Functionalize: + import torch._dispatch.python as pydispatch + + if pydispatch.CROSSREF_FUNCTIONALIZE: + handler = pydispatch.make_crossref_functionalize(self, final_key) # type: ignore[assignment] + if cache_result: + self._dispatch_cache[key] = handler + add_cached_op(self) + return handler + + r = self.py_kernels.get(final_key, final_key) + if cache_result: + self._dispatch_cache[key] = r + add_cached_op(self) + return r + + def name(self): + return self._name + + @property + def overloadpacket(self): + return self._overloadpacket + + @property + def op(self): + return self._op + + @property + def tags(self): + return self._tags + + # TODO: add more methods to expose information about input and output arguments + + +# TorchBindOpOverload are those custom ops which have at least one overload's +# schema consists of torch.ScriptObject (i.e. custom class) input. +# TorchBindOpOverload will skip C++ dispatcher and purely dispatched in python +# when its inputs contain FakeScriptObject in a similar way as higher order ops. +class TorchBindOpOverload(OpOverload[_P, _T]): + def _fallthrough_keys(self) -> list[DispatchKey]: + # TODO: we should be calling the fallback for these, but a fallthrough is almost close + # enough to the fallback in most cases that we care about. + _DEFAULT_FALLTHROUGH_KEYS = [ + DispatchKey.Autograd, + DispatchKey.AutogradCPU, + DispatchKey.AutogradCUDA, + DispatchKey.ADInplaceOrView, + DispatchKey.BackendSelect, + DispatchKey.PythonTLSSnapshot, + DispatchKey.PythonDispatcher, + ] + + def _may_use_fallthrough_instead_of_fallback(key: DispatchKey): + if torch._C._dispatch_has_kernel_for_dispatch_key(self.name(), key): + return torch._C._dispatch_kernel_for_dispatch_key_is_fallthrough( + self.name(), key + ) + + return ( + key not in self.py_kernels + or self.py_kernels[key] is torch.library.fallthrough_kernel + ) + + return [ + key + for key in _DEFAULT_FALLTHROUGH_KEYS + if _may_use_fallthrough_instead_of_fallback(key) + ] + + @contextlib.contextmanager + def _register_as_effectful_op_temporarily(self): + from torch._higher_order_ops.effects import ( + _EffectType, + _register_effectful_op, + SIDE_EFFECTS, + ) + + try: + if self not in SIDE_EFFECTS: + _register_effectful_op(self, _EffectType.ORDERED) + yield + finally: + if self in SIDE_EFFECTS: + del SIDE_EFFECTS[self] + + # Use positional-only argument to avoid naming collision with aten ops arguments + # that are named "self". This way, all the aten ops can be called by kwargs. + def __call__(self, /, *args: _P.args, **kwargs: _P.kwargs) -> _T: + if _must_dispatch_in_python(args, kwargs): + # When any inputs are FakeScriptObject, we need to + # skip c++ dispatcher and dispatch in python through _get_dispatch of python_dispatcher + # because C++ dispatcher will check the schema and cannot recognize FakeScriptObject. + # + # Note: + # 1. We only register the torchbind op temporarily as effectful op because we only want + # the effect token functionalization logic to be applied during tracing. Otherwise, the behavior + # of the eagerly executing the op might change after tracing. + # 2. We don't want to register the op as effectful for all torchbind ops in ctor because this might + # cause unexpected behavior for some autograd.profiler ops e.g. profiler._record_function_exit._RecordFunction. + with self._register_as_effectful_op_temporarily(): + return self._dispatch_in_python( + self._fallthrough_keys(), *args, **kwargs + ) + return self._op(*args, **kwargs) + + def _dispatch_in_python( + self, fallthrough_keys: list[DispatchKey], *args: _P.args, **kwargs: _P.kwargs + ) -> _T: + non_fallthrough_keys = torch._C._dispatch_keyset_full() + for key in fallthrough_keys: + non_fallthrough_keys = non_fallthrough_keys.remove(key) + + dispatch_key_set = _compute_keyset(args, kwargs, non_fallthrough_keys) + dispatch_key = dispatch_key_set.highestPriorityTypeId() + + handler = ( + self._get_dispatch(dispatch_key) + if dispatch_key not in self._dispatch_cache + else self._dispatch_cache[dispatch_key] + ) + + if isinstance(handler, DispatchKey): + # fallthrough keys can be registered at runtime via torch.library.impl + # so need to add it to fallthrough_keys and re-dispatch. + if torch._C._dispatch_kernel_for_dispatch_key_is_fallthrough( + self.name(), dispatch_key + ): + return self._dispatch_in_python( + fallthrough_keys + [dispatch_key], + *args, + **kwargs, + ) + + raise RuntimeError( + f"Torchbind op {self} received a FakeScriptObject input when dispatching {handler}." + f" but no python implementation is found." + f" Please file an issue on this when you encounter this error." + f" This error can happen when you export or compile the model." + f" It can still happen even if a C++ implementation for {dispatch_key}. " + f" has been registered. That's because FakeScriptObject purely lives in python and cannot work " + f" with a C++ implementation." + ) + + assert isinstance(handler, Callable) # type: ignore[arg-type] + return handler(*args, **kwargs) + + +def _must_dispatch_in_python(args, kwargs): + return pytree.tree_any( + lambda obj: isinstance( + obj, torch._library.fake_class_registry.FakeScriptObject + ), + (args, kwargs), + ) + + +def _has_script_object_arg(schema: torch.FunctionSchema) -> bool: + return any(isinstance(arg.type, torch.ClassType) for arg in schema.arguments) + + +# OpOverloadPacket class contains pointer to a base unresolved operator that doesn't correspond to a specific operator +# You can obtain an OpOverload object through attribute query. +class OpOverloadPacket(Generic[_P, _T]): + __file__: ClassVar[str] = "torch.ops" + + def __init__( + self, + qualified_op_name: str, + op_name: str, + op: Callable[_P, _T], + overload_names: list[str], + ) -> None: + # These attributes are accessible on the object through the properties + # defined below but are immutable + self._qualified_op_name = qualified_op_name + self.__name__ = op_name + self._op = op + self._overload_names = overload_names + self._dir: list[str] = [] + self._has_torchbind_op_overload = any( + _has_script_object_arg(schema) for schema in self._schemas.values() + ) + + # it's a no-op since OpOverloadPacket object is immutable and must be unique for a given op. + def __deepcopy__(self, memo=None): + return self + + def __repr__(self): + return "".format( + *self._qualified_op_name.split("::") + ) + + def __hash__(self): + return hash(self._op) + + def __str__(self): + return "{}.{}".format(*self._qualified_op_name.split("::")) + + @property + def op(self): + return self._op + + @property + def _schemas(self): + return { + overload_name: torch._C._get_schema(self._qualified_op_name, overload_name) + for overload_name in self._overload_names + } + + def __getattr__(self, key: str) -> OpOverload[_P, _T]: + # ensure that query for dunder attributes that does not exist on + # opoverloadpacket but instead exists on the self._op object does not unnecessarily call + # `_get_operation_overload` (which is an expensive operation). + # This is done to prevent any potential slowdown. This list can be extended + # if there exists other attributes like `__name__` that only exist on self._op and not on the + # opoverloadpacket. + # This is ok since we are guaranteed that an overload name for an aten op can't start with '__' + try: + if key.startswith("__"): + return getattr(self._op, key) + except AttributeError: + # for consistency because it seems weird to + # throw an attribute error with a message containing + # an object name different from the one the attribute + # query was performed on. + raise AttributeError( + f"'{str(self)}' can't have an overload name beginning with '__' and the " + f"underlying op {str(self._op)} has no attribute {key} either." + ) from None + + try: + # This is ok since we are guaranteed that an overload name for an aten op can't be 'default' + use_key = "" if key == "default" else key + # TODO: disallow access to overloads registered by JIT + op_dk_tags = torch._C._get_operation_overload( + self._qualified_op_name, use_key + ) + if op_dk_tags is None: + raise AttributeError( + f"The underlying op of '{str(self)}' has no overload name '{key}'" + ) + + op_, op_dk_, tags = op_dk_tags + schema = torch._C._get_schema(self._qualified_op_name, use_key) + overload: OpOverload[_P, _T] = ( + OpOverload(self, op_, op_dk_, schema, tags) + if not _has_script_object_arg(schema) + else TorchBindOpOverload(self, op_, op_dk_, schema, tags) + ) + # cache the overload object + setattr(self, key, overload) + self._dir.append(key) + return overload + except RuntimeError: + raise AttributeError( + f"The underlying op of '{str(self)}' has no overload name '{key}'" + ) from None + + def __iter__(self) -> Iterator[str]: + return iter(self._dir) + + # Use positional-only argument to avoid naming collision with aten ops arguments + # that are named "self". This way, all the aten ops can be called by kwargs. + def __call__(self, /, *args: _P.args, **kwargs: _P.kwargs) -> _T: + # overloading __call__ to ensure torch.ops.foo.bar() + # is still callable from JIT + # We save the function ptr as the `op` attribute on + # OpOverloadPacket to access it here. + + # Directly calling OverloadPacket goes into C++, which will check + # the schema and cause an error for torchbind op when inputs consist of FakeScriptObject so we + # intercept it here and call TorchBindOpverload instead. + if self._has_torchbind_op_overload and _must_dispatch_in_python(args, kwargs): + return _call_overload_packet_from_python(self, *args, **kwargs) + return self._op(*args, **kwargs) + + # TODO: use this to make a __dir__ + def overloads(self): + return [n if n else "default" for n in self._overload_names] + + +# Note - this mirrors the logic of the cpp_function defined in jit/python/init.cpp +# _jit_get_operations, which calls _get_operation_for_overload_or_packet. +def _call_overload_packet_from_python( + op: OpOverloadPacket[_P, _T], *args: _P.args, **kwargs: _P.kwargs +) -> _T: + # Reuse the torch function handling logic in cpp + torch_function_called, ret = torch._C._maybe_call_torch_function_for_op_packet( + op, *args, **kwargs + ) + + if torch_function_called: + return ret + + # The following mirrors getOpWithStack. + # In cpp, we do a schema matching for the arguments, and call ToIValue to + # to check whether the arguments are valid. But need to do similar things here + # and check the schema whether the FakeScriptObject is the corresponding fake class + # of the actual class used in schema. + exceptions = {} + found_op = None + for overload_name in op.overloads(): + op_overload = getattr(op, overload_name) + try: + _ = torch._C._check_schema_allow_fake_script_object( + op_overload._schema, *args, **kwargs + ) + found_op = op_overload + break + except RuntimeError as e: + exceptions[overload_name] = e + + if found_op: + return found_op(*args, **kwargs) + + err_msg = ( + f"Fail to match any TorchBindOverload of {op} with following exceptions:\n" + ) + for key, msg in exceptions.items(): + err_msg += f"Overload name {key}:\n {msg}\n" + raise RuntimeError(err_msg) + + +# Resolution of torch.fn is different from torch.ops.aten.fn +# torch.fn uses the Python argparser, matches with the +# appropriate schema, and calls into the unboxed version of the method +# torch.ops.aten.fn resolution is done via the mechanism defined in JIT. +# JIT creates a stack of all the overloads and then tries to match the +# correct one at runtime and always calls into the boxed version of the method +# Autograd codegen creates VariableType, TracerType, +# inplace or view type and python bindings. +# Aten codegen generates tensor methods for the tensor class. + +# _OpNamespace is a subclass of ModuleType because the torch script +# allows attribute lookups on modules only. Since we want torch.ops.foo.bar() +# to work from script, we need to ensure ops and foo are modules + + +class _OpNamespace(types.ModuleType): + """ + An op namespace to dynamically bind Operators into Python. + + Say a user has created a custom Operator called "my_namespace::my_op". To + call this op, the user will write torch.ops.my_namespace.my_op(...). + At startup, this operation will not yet be bound into Python. Instead, the + following sequence of magic tricks will occur: + 1. `torch.ops.my_namespace` will invoke the `__getattr__` magic method + on the `torch.ops` object, which will create a new `_OpNamespace` + object called `my_namespace` and set it as an attribute on the `ops` + object. + 2. `torch.ops.my_namespace.my_op` will then invoke `__getattr__` on + the `my_namespace` object, which will retrieve the operation via + `torch.get_operation`, a function bound from C++, and then in a similar + fashion bind this new object onto the `my_namespace` object. + 3. `torch.ops.my_namespace.my_op(...)` then calls this new operation + and subsequent accesses will incur no further lookup (the namespace and + operation will already exist). + """ + + __file__ = "torch.ops" + + def __init__(self, name: str) -> None: + super().__init__("torch.ops." + name) + self.name = name + self._dir: list[str] = [] + + def __iter__(self) -> Iterator[str]: + return iter(self._dir) + + def __getattr__(self, op_name: str) -> OpOverloadPacket: + if op_name in ("__origin__", "__self__"): + raise AttributeError( + f"Invalid attribute '{op_name}' for '_OpNamespace' '{self.name}'" + ) + + # Get the op `my_namespace::my_op` if available. This will also check + # for overloads and raise an exception if there are more than one. + namespace_name = self.name + qualified_op_name = f"{namespace_name}::{op_name}" + module_name = self.__module__ + "." + namespace_name + + try: + op, overload_names = _get_packet(qualified_op_name, module_name) + if op is None: + raise AttributeError( + f"'_OpNamespace' '{self.name}' object has no attribute '{op_name}'" + ) + except RuntimeError as e: + # Turn this into AttributeError so getattr(obj, key, default) + # works (this is called by TorchScript with __origin__) + raise AttributeError( + f"'_OpNamespace' '{self.name}' object has no attribute '{op_name}'" + ) from e + + op.__module__ = module_name + opoverloadpacket = OpOverloadPacket( + qualified_op_name, op_name, op, overload_names + ) + opoverloadpacket.__module__ = self.__module__ + "." + namespace_name + # cache the opoverloadpacket to ensure that each op corresponds to + # a unique OpOverloadPacket object + setattr(self, op_name, opoverloadpacket) + self._dir.append(op_name) + return opoverloadpacket + + +def _get_packet(qualname, op_module): + op, overload_names = torch._C._jit_get_operation(qualname) + if op is not None: + # let the script frontend know that op is identical to the builtin op + # with qualified_op_name + torch.jit._builtins._register_builtin(op, qualname) + op.__module__ = op_module + return op, overload_names + + +def _refresh_packet(packet): + op, overload_names = _get_packet(packet._qualified_op_name, packet._op.__module__) + assert op is not None + packet._op = op + packet._overload_names = overload_names + + +class _HigherOrderNamespace(types.ModuleType): + __file__ = "torch.ops" + + def __init__(self) -> None: + super().__init__("torch.ops.higher_order") + self._dir: list[str] = [] + + def __iter__(self) -> Iterator[str]: + return iter(self._dir) + + def __getattr__(self, name: str) -> HigherOrderOperator: + # Following _OpNamespace.__getattr__, we cache the op on this object. + op = _higher_order_ops.get(name, None) + if op is None: + raise AttributeError( + f"'_HigherOrderNamespace' 'torch.ops.higher_order' object has no attribute '{name}'" + ) + setattr(self, name, op) + self._dir.append(name) + return op + + +class _Ops(types.ModuleType): + __file__ = "_ops.py" + + def __init__(self): + super().__init__("torch.ops") + self.loaded_libraries = set() + self.higher_order = _HigherOrderNamespace() + self._dir = [] + + def __getattr__(self, name: str) -> _OpNamespace: + # Here we are creating `torch.ops.my_namespace` + namespace = _OpNamespace(name) + setattr(self, name, namespace) + self._dir.append(name) + return namespace + + def __iter__(self) -> Iterator[str]: + return iter(self._dir) + + def import_module(self, module): + """ + Imports a Python module that has torch.library registrations. + + Generally, to extend PyTorch with custom operators, a user will + create a Python module whose import triggers registration of + the custom operators via a torch.ops.load_library call or a call + to one or more torch.library.* APIs. + + It is unexpected for Python modules to have side effects, so some + linters and formatters will complain. Use this API to import Python + modules that contain these torch.library side effects. + + Args: + module (str): The name of the Python module to import + + """ + importlib.import_module(module) + + def load_library(self, path): + """ + Loads a shared library from the given path into the current process. + + The library being loaded may run global initialization code to register + custom operators with the PyTorch JIT runtime. This allows dynamically + loading custom operators. For this, you should compile your operator + and the static registration code into a shared library object, and then + call ``torch.ops.load_library('path/to/libcustom.so')`` to load the + shared object. + + After the library is loaded, it is added to the + ``torch.ops.loaded_libraries`` attribute, a set that may be inspected + for the paths of all libraries loaded using this function. + + Args: + path (str): A path to a shared library to load. + """ + path = _utils_internal.resolve_library_path(path) + with dl_open_guard(): + # Import the shared library into the process, thus running its + # static (global) initialization code in order to register custom + # operators with the JIT. + try: + ctypes.CDLL(path) + except Exception as e: + raise OSError(f"Could not load this library: {path}") from e + self.loaded_libraries.add(path) + + +# The ops "namespace" +ops: _Ops = _Ops() diff --git a/venv/lib/python3.10/site-packages/torch/_python_dispatcher.py b/venv/lib/python3.10/site-packages/torch/_python_dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..d2d4fbbf621e560ce0e8ee5afddfb0420b3d949c --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_python_dispatcher.py @@ -0,0 +1,182 @@ +# mypy: allow-untyped-defs +import re + +import torch._C as C + + +""" +PythonDispatcher class is a thin python-binding to C++ dispatcher and it +is designed to show how dispatcher precompute works. In particular, +it shows for a certain op `foo`, what the computed dispatch table looks +like after user register their kernels to certains dispatch keys. + +In the real C++ dispatcher we support many dispatch keys for different +functionalities. For simplicity PythonDispatcher only supports dispatch +keys for a single example of each use case. These use cases are listed below: + +- CPU/AutogradCPU: represents in-tree backends which we usually have dedicated inference & + autograd kernel in pytorch core library. + E.g. CPU, CUDA +- FPGA/AutogradOther: represents in-tree backends which we usually have backend specific + inference kernels, but they share the same autograd kernel specified in AutogradOther. + E.g. FPGA, SparseCsrCPU +- XLA/AutogradXLA: represents out-of-tree backends which we don't have either inference or autograd + kernel defined in pytorch core library. Backend owner is responsible for registering both + inference & autograd kernels in their extensions(e.g. torch-xla) for the operators they support. + E.g. XLA, XPU, MPS +- CompositeExplicitAutograd: alias key mapped to inference kernels of all backends like CPU, CUDA, XLA etc. + Kernels registered to this key MUST work for inference for all backends. +- Autograd: alias key mapped to autograd of all backends like AutogradCPU, AutogradXLA, AutogradOther. + Kernels registered to this key MUST work for autograd for all backends. +- CompositeImplicitAutograd: alias key CompositeImplicitAutograd = CompositeExplicitAutograd + Autograd + Kernels registered to this key MUST work for both inference + autograd for all backends. + +Note we only allow registrations to alias keys inside pytorch core library. E.g +you shouldn't register a CompositeImplicitAutograd or CompositeExplicitAutograd +kernel from torch-xla extension, instead you should upstream the kernel into +pytorch/pytorch repo so that it's available for all backends and continuously +tested even without the extension. + +Usage: + dispatcher = PythonDispatcher() + dispatcher.register(["CPU", "XLA", "CompositeImplicitAutograd"]) + print(dispatcher.dispatchTable()) # This tells you exactly which kernel is used for certain backend. + # For more debugging information + # print(dispatcher.keys()) + # print(dispatcher.registrations()) + # print(dispatcher.rawRegistrations()) + # print(dispatcher.rawDispatchTable()) +PythonDispatcher calls C++ dispatcher under the hood for to precompute dispatch table. +This file only provides the simplified API for developers, relevant test code is located in +test/test_dispatch.py +""" + + +class PythonDispatcher: + namespace = "__test__" + name = "foo" + # fmt: off + runtime_keys = [ + "CPU", "AutogradCPU", + "FPGA", "AutogradOther", + "XLA", "AutogradXLA", + "Lazy", "AutogradLazy", + ] + # fmt: on + alias_keys = [ + "CompositeExplicitAutograd", + "Autograd", + "CompositeImplicitAutograd", + ] + supported_keys = runtime_keys + alias_keys + + def __init__(self) -> None: + C._dispatch_check_invariants(self.name) # type: ignore[attr-defined] + self.ref = C._dispatch_library("FRAGMENT", self.namespace, "") + self.ref.def_("foo(Tensor x) -> Tensor") + + """ + Returns a list of dispatch keys supported by PythonDispatcher. + You can register kernels to these keys. + """ + + def keys(self): + return self.supported_keys + + """ + Register kernels to the target dispatchKeys. + dispatchKeys(list[str]): a list of dispatch keys that you want to register + your own kernel. Note that you don't need to write the kernel yourself in + this PythonDispatcher.E.g. for CPU key, a kernel(e.g fn_CPU for CPU) is + automatically generated and registered. + """ + + def register(self, dispatchKeys): + # Overridden is not supported and triggers a warning in C++ dispatcher. + if len(set(dispatchKeys)) != len(dispatchKeys): + raise RuntimeError( + f"Overridden is not allowed but found duplicates in {dispatchKeys}." + ) + # We currently forbid this in codegen instead of C++ dispatcher. + if ( + "CompositeImplicitAutograd" in dispatchKeys + and "CompositeExplicitAutograd" in dispatchKeys + ): + raise RuntimeError( + "Registration to both CompositeImplicitAutograd and CompositeExplicitAutograd is not allowed." + ) + for key in dispatchKeys: + if key not in self.supported_keys: + raise RuntimeError( + f"{key} is not supported, please select a dispatch key in {self.supported_keys}." + ) + self.ref.impl_t_t("foo", dispatch=key, debug="fn_" + key) + + """ + Helper function to format (key, kernel). + """ + + def _format_line(self, key, kernel): + return f"{key:<15} {kernel}\n" + + """ + Helper function to print a table header. + """ + + def _format_header(self, header): + s = f""" +{header} +""" + s += self._format_line("key", "kernel") + s += "---------------------------\n" + return s + + """ + Returns raw output of all registration info for debugging only. + Use registrations() for a simplified version. + """ + + def rawRegistrations(self): + return C._dispatch_dump(f"{self.namespace}::{self.name}") # type: ignore[attr-defined] + + """ + Returns raw output of computed dispatch table for debugging only. + Use dispatchTable() for a simplified version. + """ + + def rawDispatchTable(self): + return C._dispatch_dump_table(f"{self.namespace}::{self.name}") # type: ignore[attr-defined] + + """ + Returns a table(str) including all the registrations from users. + Note this includes registrations to both runtime keys and alias keys. + """ + + def registrations(self): + output = self._format_header("Registered Kernels") + state = self.rawRegistrations() + state_entries = state.split("\n") + for line in state_entries: + first = line.split(":")[0] + if any(first.startswith(k) for k in self.supported_keys): + kernel = line.split("::")[0].split(" ")[1] + output += self._format_line(first, kernel) + return output + + """ + Returns the computed dispatch table(str). Note this only include + runtime keys, registrations to alias keys have been decoded to their + mapped runtime keys. + """ + + def dispatchTable(self): + output = self._format_header("Computed Dispatch Table") + table = self.rawDispatchTable() + table_entries = table.split("\n") + regex = re.compile(r"registered at .*FallbackKernel\.cpp.*(\[)") + for line in table_entries: + k = line.split(":")[0] + if k in self.runtime_keys: + entry = regex.sub("[", line) + output += self._format_line(k, entry.split(": ")[1]) + return output diff --git a/venv/lib/python3.10/site-packages/torch/_size_docs.py b/venv/lib/python3.10/site-packages/torch/_size_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..e30240a1e6f6748d34c673489f225a51c6fe8b9d --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_size_docs.py @@ -0,0 +1,39 @@ +"""Adds docstrings to torch.Size functions""" + +import torch._C +from torch._C import _add_docstr as add_docstr + + +def add_docstr_all(method: str, docstr: str) -> None: + add_docstr(getattr(torch._C.Size, method), docstr) + + +add_docstr_all( + "numel", + """ +numel() -> int + +Returns the number of elements a :class:`torch.Tensor` with the given size would contain. + +More formally, for a tensor ``x = tensor.ones(10, 10)`` with size ``s = torch.Size([10, 10])``, +``x.numel() == x.size().numel() == s.numel() == 100`` holds true. + +Example:: + + >>> x=torch.ones(10, 10) + >>> s=x.size() + >>> s + torch.Size([10, 10]) + >>> s.numel() + 100 + >>> x.numel() == s.numel() + True + + +.. warning:: + + This function does not return the number of dimensions described by :class:`torch.Size`, but instead the number + of elements a :class:`torch.Tensor` with that size would contain. + +""", +) diff --git a/venv/lib/python3.10/site-packages/torch/_sources.py b/venv/lib/python3.10/site-packages/torch/_sources.py new file mode 100644 index 0000000000000000000000000000000000000000..1327729a717b106e1a3612ba1c69c1db14e9c9fa --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_sources.py @@ -0,0 +1,138 @@ +# mypy: allow-untyped-defs +import ast +import functools +import inspect +from textwrap import dedent +from typing import Any, NamedTuple, Optional + +from torch._C import ErrorReport +from torch._C._jit_tree_views import SourceRangeFactory + + +def get_source_lines_and_file( + obj: Any, + error_msg: Optional[str] = None, +) -> tuple[list[str], int, Optional[str]]: + """ + Wrapper around inspect.getsourcelines and inspect.getsourcefile. + + Returns: (sourcelines, file_lino, filename) + """ + filename = None # in case getsourcefile throws + try: + filename = inspect.getsourcefile(obj) + sourcelines, file_lineno = inspect.getsourcelines(obj) + except OSError as e: + msg = ( + f"Can't get source for {obj}. TorchScript requires source access in " + "order to carry out compilation, make sure original .py files are " + "available." + ) + if error_msg: + msg += "\n" + error_msg + raise OSError(msg) from e + + return sourcelines, file_lineno, filename + + +def normalize_source_lines(sourcelines: list[str]) -> list[str]: + """ + This helper function accepts a list of source lines. It finds the + indentation level of the function definition (`def`), then it indents + all lines in the function body to a point at or greater than that + level. This allows for comments and continued string literals that + are at a lower indentation than the rest of the code. + Args: + sourcelines: function source code, separated into lines by + the '\n' character + Returns: + A list of source lines that have been correctly aligned + """ + + def remove_prefix(text, prefix): + return text[text.startswith(prefix) and len(prefix) :] + + # Find the line and line number containing the function definition + idx = None + for i, l in enumerate(sourcelines): + if l.lstrip().startswith("def"): + idx = i + break + + # This will happen when the function is a lambda- we won't find "def" anywhere in the source + # lines in that case. Currently trying to JIT compile a lambda will throw an error up in + # `parse_def()`, but we might want to handle this case in the future. + if idx is None: + return sourcelines + + # Get a string representing the amount of leading whitespace + fn_def = sourcelines[idx] + whitespace = fn_def.split("def")[0] + + # Add this leading whitespace to all lines before and after the `def` + aligned_prefix = [ + whitespace + remove_prefix(s, whitespace) for s in sourcelines[:idx] + ] + aligned_suffix = [ + whitespace + remove_prefix(s, whitespace) for s in sourcelines[idx + 1 :] + ] + + # Put it together again + aligned_prefix.append(fn_def) + return aligned_prefix + aligned_suffix + + +# Thin wrapper around SourceRangeFactory to store extra metadata +# about the function-to-be-compiled. +class SourceContext(SourceRangeFactory): + def __init__( + self, + source, + filename, + file_lineno, + leading_whitespace_len, + uses_true_division=True, + funcname=None, + ): + super().__init__(source, filename, file_lineno, leading_whitespace_len) + self.uses_true_division = uses_true_division + self.filename = filename + self.funcname = funcname + + +@functools.cache +def make_source_context(*args): + return SourceContext(*args) + + +def fake_range(): + return SourceContext("", None, 0, 0).make_raw_range(0, 1) + + +class ParsedDef(NamedTuple): + ast: ast.Module + ctx: SourceContext + source: str + filename: Optional[str] + file_lineno: int + + +def parse_def(fn): + sourcelines, file_lineno, filename = get_source_lines_and_file( + fn, ErrorReport.call_stack() + ) + sourcelines = normalize_source_lines(sourcelines) + source = "".join(sourcelines) + dedent_src = dedent(source) + py_ast = ast.parse(dedent_src) + if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef): + raise RuntimeError( + f"Expected a single top-level function: {filename}:{file_lineno}" + ) + leading_whitespace_len = len(source.split("\n", 1)[0]) - len( + dedent_src.split("\n", 1)[0] + ) + ctx = make_source_context( + source, filename, file_lineno, leading_whitespace_len, True, fn.__name__ + ) + return ParsedDef(py_ast, ctx, source, filename, file_lineno) diff --git a/venv/lib/python3.10/site-packages/torch/_storage_docs.py b/venv/lib/python3.10/site-packages/torch/_storage_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..f0d16bc4250ffb2a383af727c974e9f910a5b2a5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_storage_docs.py @@ -0,0 +1,42 @@ +# mypy: allow-untyped-defs +"""Adds docstrings to Storage functions""" + +import torch._C +from torch._C import _add_docstr as add_docstr + + +storage_classes = ["StorageBase"] + + +def add_docstr_all(method, docstr): + for cls_name in storage_classes: + cls = getattr(torch._C, cls_name) + try: + add_docstr(getattr(cls, method), docstr) + except AttributeError: + pass + + +add_docstr_all( + "from_file", + """ +from_file(filename, shared=False, nbytes=0) -> Storage + +Creates a CPU storage backed by a memory-mapped file. + +If ``shared`` is ``True``, then memory is shared between all processes. +All changes are written to the file. If ``shared`` is ``False``, then the changes on +the storage do not affect the file. + +``nbytes`` is the number of bytes of storage. If ``shared`` is ``False``, +then the file must contain at least ``nbytes`` bytes. If ``shared`` is +``True`` the file will be created if needed. (Note that for ``UntypedStorage`` +this argument differs from that of ``TypedStorage.from_file``) + +Args: + filename (str): file name to map + shared (bool): whether to share memory (whether ``MAP_SHARED`` or ``MAP_PRIVATE`` is passed to the + underlying `mmap(2) call `_) + nbytes (int): number of bytes of storage +""", +) diff --git a/venv/lib/python3.10/site-packages/torch/_streambase.py b/venv/lib/python3.10/site-packages/torch/_streambase.py new file mode 100644 index 0000000000000000000000000000000000000000..9d71120c959b14b09309738c84d788d60e7db326 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_streambase.py @@ -0,0 +1,20 @@ +from typing_extensions import deprecated + +import torch + + +# Preserved only for BC reasons +@deprecated( + "`torch._streambase._StreamBase` is deprecated. Please use `torch.Stream` instead.", + category=FutureWarning, +) +class _StreamBase(torch.Stream): + pass + + +@deprecated( + "`torch._streambase._EventBase` is deprecated. Please use `torch.Event` instead.", + category=FutureWarning, +) +class _EventBase(torch.Event): + pass diff --git a/venv/lib/python3.10/site-packages/torch/_tensor.py b/venv/lib/python3.10/site-packages/torch/_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..6cebed28b8b0da506774058a0356da5e903c7e2f --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_tensor.py @@ -0,0 +1,1837 @@ +# mypy: allow-untyped-defs +import copyreg +import enum +import functools +import warnings +from collections import OrderedDict +from copy import deepcopy +from numbers import Number +from typing import Any, Callable, cast, Optional, TypeVar, Union +from typing_extensions import Concatenate, ParamSpec + +import torch +import torch._C as _C +from torch._namedtensor_internals import ( + check_serializing_named_tensor, + is_ellipsis, + resolve_ellipsis, + single_ellipsis_index, + unzip_namedshape, + update_names, +) +from torch.overrides import ( + get_default_nowrap_functions, + handle_torch_function, + has_torch_function, + has_torch_function_unary, + has_torch_function_variadic, +) + + +_P = ParamSpec("_P") +_TensorLike = TypeVar("_TensorLike", bound=_C.TensorBase) + + +def _handle_torch_function_and_wrap_type_error_to_not_implemented( + f: Callable[Concatenate[_TensorLike, _P], "Tensor"], +) -> Callable[Concatenate[_TensorLike, _P], "Tensor"]: + @functools.wraps(f) + def wrapped(self: _TensorLike, *args: _P.args, **kwargs: _P.kwargs) -> "Tensor": + try: + # See https://github.com/pytorch/pytorch/issues/75462 + sargs = self, *args + if has_torch_function(sargs): + return handle_torch_function(wrapped, sargs, *sargs, **kwargs) + return f(self, *args, **kwargs) + except TypeError: + return NotImplemented + + return wrapped + + +# Should not be used, this is kept only for BC of loading old serialized Tensor subclasses +def _rebuild_from_type(func, type, args, dict): + if type is Tensor: + return func(*args) + + ret = func(*args).as_subclass(type) + ret.__dict__ = dict + return ret + + +def _rebuild_from_type_v2(func, new_type, args, state): + ret = func(*args) + if type(ret) is not new_type: + ret = ret.as_subclass(new_type) + # Tensor does define __setstate__ even though it doesn't define + # __getstate__. So only use __setstate__ if it is NOT the one defined + # on Tensor + if ( + getattr(ret.__class__, "__setstate__", Tensor.__setstate__) + is not Tensor.__setstate__ + ): + ret.__setstate__(state) + else: + ret = torch._utils._set_obj_state(ret, state) + return ret + + +def _dtype_to_typestr(dtype): + # CUDA devices are little-endian and tensors are stored in native byte + # order. 1-byte entries are endian-agnostic. + return { + torch.complex64: " torch.TypedStorage + + Returns the underlying :class:`TypedStorage`. + + .. warning:: + + :class:`TypedStorage` is deprecated. It will be removed in the future, and + :class:`UntypedStorage` will be the only storage class. To access the + :class:`UntypedStorage` directly, use :attr:`Tensor.untyped_storage()`. + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.storage, (self,), self) + + torch.storage._warn_typed_storage_removal(stacklevel=2) + return self._typed_storage() + + # For internal use only, to avoid raising deprecation warning + def _typed_storage(self): + untyped_storage = self.untyped_storage() + return torch.TypedStorage( + wrap_storage=untyped_storage, dtype=self.dtype, _internal=True + ) + + def _reduce_ex_internal(self, proto): + check_serializing_named_tensor(self) + + from torch.utils.hooks import warn_if_has_hooks + + # See Note [Don't serialize hooks] + warn_if_has_hooks(self) + backward_hooks: dict[Any, Any] = OrderedDict() + + skip_data = torch.serialization._serialization_tls.skip_data + materialize_fake_tensors = ( + torch.serialization._serialization_tls.materialize_fake_tensors + ) + + if self.device.type in ["xla", "maia", "mtia"] or ( + not torch._C._has_storage(self) + and self.device.type == torch._C._get_privateuse1_backend_name() + ): + if skip_data: + raise RuntimeError( + "Cannot serialize tensors on backends with no storage under skip_data context manager" + ) + cpu_tensor = self.cpu() + return ( + torch._utils._rebuild_device_tensor_from_cpu_tensor, + (cpu_tensor, self.dtype, str(self.device), self.requires_grad), + ) + if self.device.type == "meta": + # NB: This implementation BREAKS storage sharing. Current + # hypothesis is that no one cares for meta tensors. + if skip_data: + warnings.warn( + "Serializing tensors on the meta device under skip_data context manager is a no-op" + ) + arg_meta = ( + self.dtype, + tuple(self.size()), + self.stride(), + self.requires_grad, + ) + return (torch._utils._rebuild_meta_tensor_no_storage, arg_meta) + if self.is_quantized: + if skip_data: + raise RuntimeError( + "Cannot serialize qtensor under skip_data context manager, file an issue if you need this feature" + ) + # quantizer_params can be different type based on torch attribute + quantizer_params: Union[ + tuple[torch.qscheme, float, int], tuple[Any, Tensor, Tensor, int] + ] + if self.qscheme() == torch.per_tensor_affine: + quantizer_params = ( + torch.per_tensor_affine, + self.q_scale(), + self.q_zero_point(), + ) + elif self.qscheme() in ( + torch.per_channel_affine, + torch.per_channel_affine_float_qparams, + ): + # convert scales and zero points to tuple to avoid recursive calls + # when/if we get multi-axis quantized tensors in the future, the shape + # is recoverable from the main tensor shape + quantizer_params = ( + torch.per_channel_affine, + self.q_per_channel_scales(), + self.q_per_channel_zero_points(), + self.q_per_channel_axis(), + ) + else: + raise RuntimeError( + f"Serialization is not supported for tensors of type {self.qscheme()}" + ) + # TODO: Once we decide to break serialization FC, no longer + # need to wrap with TypedStorage + args_qtensor = ( + torch.storage.TypedStorage( + wrap_storage=self._typed_storage()._untyped_storage, + dtype=self.dtype, + _internal=True, + ), + self.storage_offset(), + tuple(self.size()), + self.stride(), + quantizer_params, + self.requires_grad, + backward_hooks, + ) + return (torch._utils._rebuild_qtensor, args_qtensor) + elif self.is_sparse: + if self.layout == torch.sparse_coo: + args_sparse = ( + self.layout, + (self._indices(), self._values(), self.size(), self.is_coalesced()), + ) + else: + raise NotImplementedError( + f"sparse tensor __reduce_ex__ for layout `{self.layout}`" + ) + return (torch._utils._rebuild_sparse_tensor, args_sparse) + elif self.layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + if self.layout in {torch.sparse_csr, torch.sparse_bsr}: + compressed_indices, plain_indices = ( + self.crow_indices(), + self.col_indices(), + ) + else: + compressed_indices, plain_indices = ( + self.ccol_indices(), + self.row_indices(), + ) + args_sparse_compressed = ( + self.layout, + ( + compressed_indices, + plain_indices, + self.values(), + self.size(), + ), + ) + return (torch._utils._rebuild_sparse_tensor, args_sparse_compressed) + elif self.is_nested: + if skip_data: + raise RuntimeError( + "Cannot serialize nested tensor under skip_data context manager, file an issue if you need this feature" + ) + args_nested = ( + # NB: values() currently returns the storage as a buffer in an unsafe way. + # Ideally, we'd use a private API for this instead. TODO: Switch to this if + # we ever get around to adding it. + self.values(), + self._nested_tensor_size(), + self._nested_tensor_strides(), + self._nested_tensor_storage_offsets(), + ) + return (torch._utils._rebuild_nested_tensor, args_nested) + elif ( + type(self) is not torch.Tensor + and type(self).__torch_dispatch__ is not torch.Tensor.__torch_dispatch__ + and ( + isinstance(self, torch._subclasses.functional_tensor.FunctionalTensor) + or ( + not isinstance(self, torch._subclasses.fake_tensor.FakeTensor) + and self.data_ptr() == 0 + ) + ) + ): + arg_wrapper_subclass = ( + type(self), + self.dtype, + tuple(self.size()), + self.stride(), + self.storage_offset(), + self.layout, + self.device, + self.requires_grad, + ) + return (torch._utils._rebuild_wrapper_subclass, arg_wrapper_subclass) + elif ( + type(self) is not torch.Tensor + and type(self).__torch_dispatch__ is not torch.Tensor.__torch_dispatch__ + and ( + isinstance(self, torch._subclasses.fake_tensor.FakeTensor) + and not (skip_data and materialize_fake_tensors) + ) + ): + arg_wrapper_subclass = ( + type(self), + self.dtype, + tuple(self.size()), + self.stride(), + self.storage_offset(), + self.layout, + self.device, + self.requires_grad, + ) + return (torch._utils._rebuild_wrapper_subclass, arg_wrapper_subclass) + else: + v3_dtypes = torch.storage._new_dtypes() + if self.dtype in v3_dtypes: + rebuild_func = torch._utils._rebuild_tensor_v3 + storage = self.untyped_storage() + else: + # TODO: Once we decide to break serialization FC, no longer + # need to wrap with TypedStorage + rebuild_func = torch._utils._rebuild_tensor_v2 # type: ignore[assignment] + storage = torch.storage.TypedStorage( + wrap_storage=self._typed_storage()._untyped_storage, + dtype=self.dtype, + _internal=True, + ) # type: ignore[assignment] + + # TODO: remove hasattr, it's a hack to support versions of torch that + # don't have _subclasses + if ( + hasattr(torch, "_subclasses") + and isinstance(self, torch._subclasses.fake_tensor.FakeTensor) + and skip_data + ): + storage._fake_device = self.device + + args = ( + storage, + self.storage_offset(), + tuple(self.size()), + self.stride(), + self.requires_grad, + backward_hooks, + ) # previously was self._backward_hooks + + if isinstance(storage, torch.storage.UntypedStorage): + args = args + (self.dtype,) # type: ignore[assignment] + + metadata = torch._utils.get_tensor_metadata(self) + if metadata: + args = args + (metadata,) # type: ignore[assignment] + + return (rebuild_func, args) + + def __setstate__(self, state): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__setstate__, (self,), self, state) + # Warning: this method is NOT called when you torch.load() a tensor; + # that is managed by _rebuild_tensor_v2 + if not self.is_leaf: + raise RuntimeError("__setstate__ can be only called on leaf Tensors") + if len(state) == 4: + # legacy serialization of Tensor + self.set_(*state) + return + elif len(state) == 5: + # legacy serialization of Variable + self.data = state[0] + state = (state[3], state[4], state[2]) + # The setting of _backward_hooks is expected to be a no-op. + # See Note [Don't serialize hooks] + self.requires_grad, _, self._backward_hooks = state + + def __repr__(self, *, tensor_contents=None): + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.__repr__, (self,), self, tensor_contents=tensor_contents + ) + # All strings are unicode in Python 3. + return torch._tensor_str._str(self, tensor_contents=tensor_contents) + + def backward( + self, gradient=None, retain_graph=None, create_graph=False, inputs=None + ): + r"""Computes the gradient of current tensor wrt graph leaves. + + The graph is differentiated using the chain rule. If the tensor is + non-scalar (i.e. its data has more than one element) and requires + gradient, the function additionally requires specifying a ``gradient``. + It should be a tensor of matching type and shape, that represents + the gradient of the differentiated function w.r.t. ``self``. + + This function accumulates gradients in the leaves - you might need to zero + ``.grad`` attributes or set them to ``None`` before calling it. + See :ref:`Default gradient layouts` + for details on the memory layout of accumulated gradients. + + .. note:: + + If you run any forward ops, create ``gradient``, and/or call ``backward`` + in a user-specified CUDA stream context, see + :ref:`Stream semantics of backward passes`. + + .. note:: + + When ``inputs`` are provided and a given input is not a leaf, + the current implementation will call its grad_fn (though it is not strictly needed to get this gradients). + It is an implementation detail on which the user should not rely. + See https://github.com/pytorch/pytorch/pull/60521#issuecomment-867061780 for more details. + + Args: + gradient (Tensor, optional): The gradient of the function + being differentiated w.r.t. ``self``. + This argument can be omitted if ``self`` is a scalar. Defaults to ``None``. + retain_graph (bool, optional): If ``False``, the graph used to compute the grads will be freed; + If ``True``, it will be retained. The default is ``None``, in which case the value is inferred from ``create_graph`` + (i.e., the graph is retained only when higher-order derivative tracking is requested). Note that in nearly all cases + setting this option to True is not needed and often can be worked around in a much more efficient way. + create_graph (bool, optional): If ``True``, graph of the derivative will + be constructed, allowing to compute higher order derivative + products. Defaults to ``False``. + inputs (Sequence[Tensor], optional): Inputs w.r.t. which the gradient will be + accumulated into ``.grad``. All other tensors will be ignored. If not + provided, the gradient is accumulated into all the leaf Tensors that were + used to compute the :attr:`tensors`. Defaults to ``None``. + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.backward, + (self,), + self, + gradient=gradient, + retain_graph=retain_graph, + create_graph=create_graph, + inputs=inputs, + ) + torch.autograd.backward( + self, gradient, retain_graph, create_graph, inputs=inputs + ) + + def register_hook(self, hook): + r"""Registers a backward hook. + + The hook will be called every time a gradient with respect to the + Tensor is computed. The hook should have the following signature:: + + hook(grad) -> Tensor or None + + + The hook should not modify its argument, but it can optionally return + a new gradient which will be used in place of :attr:`grad`. + + This function returns a handle with a method ``handle.remove()`` + that removes the hook from the module. + + .. note:: + See :ref:`backward-hooks-execution` for more information on how when this hook + is executed, and how its execution is ordered relative to other hooks. + + Example:: + + >>> v = torch.tensor([0., 0., 0.], requires_grad=True) + >>> h = v.register_hook(lambda grad: grad * 2) # double the gradient + >>> v.backward(torch.tensor([1., 2., 3.])) + >>> v.grad + + 2 + 4 + 6 + [torch.FloatTensor of size (3,)] + + >>> h.remove() # removes the hook + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.register_hook, (self,), self, hook) + if not self.requires_grad: + raise RuntimeError( + "cannot register a hook on a tensor that doesn't require gradient" + ) + if self._backward_hooks is None: + self._backward_hooks = OrderedDict() + if self.grad_fn is not None: + self.grad_fn._register_hook_dict(self) + + from torch.utils.hooks import RemovableHandle + + handle = RemovableHandle(self._backward_hooks) + self._backward_hooks[handle.id] = hook + return handle + + def register_post_accumulate_grad_hook(self, hook): + r"""Registers a backward hook that runs after grad accumulation. + + The hook will be called after all gradients for a tensor have been accumulated, + meaning that the .grad field has been updated on that tensor. The post + accumulate grad hook is ONLY applicable for leaf tensors (tensors without a + .grad_fn field). Registering this hook on a non-leaf tensor will error! + + The hook should have the following signature:: + + hook(param: Tensor) -> None + + Note that, unlike other autograd hooks, this hook operates on the tensor + that requires grad and not the grad itself. The hook can in-place modify + and access its Tensor argument, including its .grad field. + + This function returns a handle with a method ``handle.remove()`` + that removes the hook from the module. + + .. note:: + See :ref:`backward-hooks-execution` for more information on how when this hook + is executed, and how its execution is ordered relative to other hooks. Since + this hook runs during the backward pass, it will run in no_grad mode (unless + create_graph is True). You can use torch.enable_grad() to re-enable autograd + within the hook if you need it. + + Example:: + + >>> v = torch.tensor([0., 0., 0.], requires_grad=True) + >>> lr = 0.01 + >>> # simulate a simple SGD update + >>> h = v.register_post_accumulate_grad_hook(lambda p: p.add_(p.grad, alpha=-lr)) + >>> v.backward(torch.tensor([1., 2., 3.])) + >>> v + tensor([-0.0100, -0.0200, -0.0300], requires_grad=True) + + >>> h.remove() # removes the hook + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.register_post_accumulate_grad_hook, (self,), self, hook + ) + if not self.requires_grad: + raise RuntimeError( + "cannot register a hook on a tensor that doesn't require gradient" + ) + if self.grad_fn is not None: + raise RuntimeError( + "post accumulate grad hooks cannot be registered on non-leaf tensors" + ) + if self._post_accumulate_grad_hooks is None: + self._post_accumulate_grad_hooks: dict[Any, Any] = OrderedDict() + + from torch.utils.hooks import RemovableHandle + + handle = RemovableHandle(self._post_accumulate_grad_hooks) + self._post_accumulate_grad_hooks[handle.id] = hook + return handle + + def reinforce(self, reward): + def trim(str): + return "\n".join([line.strip() for line in str.split("\n")]) + + raise RuntimeError( + trim( + r"""reinforce() was removed. + Use torch.distributions instead. + See https://pytorch.org/docs/main/distributions.html + + Instead of: + + probs = policy_network(state) + action = probs.multinomial() + next_state, reward = env.step(action) + action.reinforce(reward) + action.backward() + + Use: + + probs = policy_network(state) + # NOTE: categorical is equivalent to what used to be called multinomial + m = torch.distributions.Categorical(probs) + action = m.sample() + next_state, reward = env.step(action) + loss = -m.log_prob(action) * reward + loss.backward() + """ + ) + ) + + detach = _C._add_docstr( + _C.TensorBase.detach, + r""" + Returns a new Tensor, detached from the current graph. + + The result will never require gradient. + + This method also affects forward mode AD gradients and the result will never + have forward mode AD gradients. + + .. note:: + + Returned Tensor shares the same storage with the original one. + In-place modifications on either of them will be seen, and may trigger + errors in correctness checks. + """, + ) + + detach_ = _C._add_docstr( + _C.TensorBase.detach_, + r""" + Detaches the Tensor from the graph that created it, making it a leaf. + Views cannot be detached in-place. + + This method also affects forward mode AD gradients and the result will never + have forward mode AD gradients. + """, + ) + + def is_shared(self): + r"""Checks if tensor is in shared memory. + + This is always ``True`` for CUDA tensors. + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.is_shared, (self,), self) + return self._typed_storage()._is_shared() + + def share_memory_(self): + r"""Moves the underlying storage to shared memory. + + This is a no-op if the underlying storage is already in shared memory + and for CUDA tensors. Tensors in shared memory cannot be resized. + + See :meth:`torch.UntypedStorage.share_memory_` for more details. + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.share_memory_, (self,), self) + self._typed_storage()._share_memory_() + return self + + def module_load(self, other, assign=False): + r"""Defines how to transform ``other`` when loading it into ``self`` in :meth:`~nn.Module.load_state_dict`. + + Used when :func:`~torch.__future__.get_swap_module_params_on_conversion` is ``True``. + + It is expected that ``self`` is a parameter or buffer in an ``nn.Module`` and ``other`` is the + value in the state dictionary with the corresponding key, this method defines + how ``other`` is remapped before being swapped with ``self`` via + :func:`~torch.utils.swap_tensors` in :meth:`~nn.Module.load_state_dict`. + + .. note:: + This method should always return a new object that is not ``self`` or ``other``. + For example, the default implementation returns ``self.copy_(other).detach()`` + if ``assign`` is ``False`` or ``other.detach()`` if ``assign`` is ``True``. + + Args: + other (Tensor): value in state dict with key corresponding to ``self`` + assign (bool): the assign argument passed to :meth:`nn.Module.load_state_dict` + + """ + if has_torch_function_variadic(self, other): + return handle_torch_function( + Tensor.module_load, (self, other), self, other, assign=assign + ) + + if assign: + return other.detach() + else: + return self.copy_(other).detach() + + def __reversed__(self): + r"""Reverses the tensor along dimension 0.""" + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__reversed__, (self,), self) + if self.dim() == 0: + return self + else: + return self.flip(0) + + def norm( + self, + p: Optional[Union[float, str]] = "fro", + dim=None, + keepdim=False, + dtype=None, + ): + r"""See :func:`torch.norm`""" + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.norm, (self,), self, p=p, dim=dim, keepdim=keepdim, dtype=dtype + ) + return torch.norm(self, p, dim, keepdim, dtype=dtype) + + def solve(self, other): + from torch._linalg_utils import solve + + return solve(self, other) + + def lstsq(self, other): + from torch._linalg_utils import lstsq + + return lstsq(self, other) + + def eig(self, eigenvectors=False): + from torch._linalg_utils import eig + + return eig(self, eigenvectors=eigenvectors) + + def symeig(self, eigenvectors=False): + from torch._linalg_utils import _symeig + + return _symeig(self, eigenvectors=eigenvectors) + + def lu(self, pivot=True, get_infos=False): + r"""See :func:`torch.lu`""" + # If get_infos is True, then we don't need to check for errors and vice versa + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.lu, (self,), self, pivot=pivot, get_infos=get_infos + ) + + LU, pivots, infos = torch._lu_with_info( + self, pivot=pivot, check_errors=(not get_infos) + ) + if get_infos: + return LU, pivots, infos + else: + return LU, pivots + + def stft( + self, + n_fft: int, + hop_length: Optional[int] = None, + win_length: Optional[int] = None, + window: "Optional[Tensor]" = None, + center: bool = True, + pad_mode: str = "reflect", + normalized: bool = False, + onesided: Optional[bool] = None, + return_complex: Optional[bool] = None, + align_to_window: Optional[bool] = None, + ): + r"""See :func:`torch.stft` + + .. warning:: + This function changed signature at version 0.4.1. Calling with + the previous signature may cause error or return incorrect result. + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.stft, + (self,), + self, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window, + center=center, + pad_mode=pad_mode, + normalized=normalized, + onesided=onesided, + return_complex=return_complex, + align_to_window=align_to_window, + ) + return torch.stft( + self, + n_fft, + hop_length, + win_length, + window, + center, + pad_mode, + normalized, + onesided, + return_complex=return_complex, + align_to_window=align_to_window, + ) + + def istft( + self, + n_fft: int, + hop_length: Optional[int] = None, + win_length: Optional[int] = None, + window: "Optional[Tensor]" = None, + center: bool = True, + normalized: bool = False, + onesided: Optional[bool] = None, + length: Optional[int] = None, + return_complex: bool = False, + ): + r"""See :func:`torch.istft`""" + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.istft, + (self,), + self, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window, + center=center, + normalized=normalized, + onesided=onesided, + length=length, + return_complex=return_complex, + ) + return torch.istft( + self, + n_fft, + hop_length, + win_length, + window, + center, + normalized, + onesided, + length, + return_complex=return_complex, + ) + + def resize(self, *sizes): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.resize, (self,), self, *sizes) + warnings.warn("non-inplace resize is deprecated") + from torch.autograd._functions import Resize + + return Resize.apply(self, sizes) + + def resize_as(self, tensor): + if has_torch_function_variadic(self, tensor): + return handle_torch_function(Tensor.resize_as, (self, tensor), self, tensor) + warnings.warn("non-inplace resize_as is deprecated") + from torch.autograd._functions import Resize + + return Resize.apply(self, tensor.size()) + + def split(self, split_size, dim=0): + r"""See :func:`torch.split`""" + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.split, (self,), self, split_size, dim=dim + ) + if isinstance(split_size, Tensor): + try: + split_size = int(split_size) + except ValueError: + pass + + if isinstance(split_size, (int, torch.SymInt)): + return torch._VF.split(self, split_size, dim) # type: ignore[attr-defined] + else: + return torch._VF.split_with_sizes(self, split_size, dim) + + def unique(self, sorted=True, return_inverse=False, return_counts=False, dim=None): + r"""Returns the unique elements of the input tensor. + + See :func:`torch.unique` + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.unique, + (self,), + self, + sorted=sorted, + return_inverse=return_inverse, + return_counts=return_counts, + dim=dim, + ) + return torch.unique( + self, + sorted=sorted, + return_inverse=return_inverse, + return_counts=return_counts, + dim=dim, + ) + + def unique_consecutive(self, return_inverse=False, return_counts=False, dim=None): + r"""Eliminates all but the first element from every consecutive group of equivalent elements. + + See :func:`torch.unique_consecutive` + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.unique_consecutive, + (self,), + self, + return_inverse=return_inverse, + return_counts=return_counts, + dim=dim, + ) + return torch.unique_consecutive( + self, return_inverse=return_inverse, return_counts=return_counts, dim=dim + ) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rsub__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor": + return _C._VariableFunctions.rsub(self, other) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rdiv__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor": + return self.reciprocal() * other + + __rtruediv__ = __rdiv__ + __itruediv__ = _C.TensorBase.__idiv__ + + __pow__ = cast( + Callable[ + ["torch._C.TensorBase", Union["Tensor", int, float, bool, complex]], + "Tensor", + ], + _handle_torch_function_and_wrap_type_error_to_not_implemented( + _C.TensorBase.pow + ), + ) + + __ipow__ = _handle_torch_function_and_wrap_type_error_to_not_implemented( + _C.TensorBase.pow_ + ) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rmod__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor": + return torch.remainder(other, self) + + def __format__(self, format_spec): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__format__, (self,), self, format_spec) + if self.dim() == 0 and not self.is_meta and type(self) is Tensor: + # Use detach() here to avoid the warning when converting a scalar Tensor that + # requires gradients to a python number. It is ok for formatting. + return self.detach().item().__format__(format_spec) + return object.__format__(self, format_spec) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rpow__(self, other: Union["Tensor", int, float, bool, complex]) -> "Tensor": + return torch.pow(other, self) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __floordiv__(self, other: Union["Tensor", int, float, bool]) -> "Tensor": # type: ignore[override] + # TODO(rec): the superclass says it accepts complex here, + # but torch.floor_divide says it doesn't. + return torch.floor_divide(self, other) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rfloordiv__(self, other: Union["Tensor", int, float, bool]) -> "Tensor": # type: ignore[override] + return torch.floor_divide(other, self) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rlshift__( + self, other: Union["Tensor", int, float, bool, complex] + ) -> "Tensor": + return torch.bitwise_left_shift(other, self) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rrshift__( + self, other: Union["Tensor", int, float, bool, complex] + ) -> "Tensor": + return torch.bitwise_right_shift(other, self) + + @_handle_torch_function_and_wrap_type_error_to_not_implemented + def __rmatmul__(self, other: "Tensor") -> "Tensor": + return torch.matmul(other, self) + + __pos__ = _C.TensorBase.positive + __neg__ = _C.TensorBase.neg + __abs__ = _C.TensorBase.abs + + def __len__(self): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__len__, (self,), self) + if self.dim() == 0: + raise TypeError("len() of a 0-d tensor") + if torch._C._get_tracing_state(): + warnings.warn( + "Using len to get tensor shape might cause the trace to be incorrect. " + "Recommended usage would be tensor.shape[0]. " + "Passing a tensor of different shape might lead to errors or silently give " + "incorrect results.", + category=torch.jit.TracerWarning, + stacklevel=2, + ) + return self.shape[0] + + def __iter__(self): + # NB: we use 'imap' and not 'map' here, so that in Python 2 we get a + # generator and don't eagerly perform all the indexes. This could + # save us work, and also helps keep trace ordering deterministic + # (e.g., if you zip(*hiddens), the eager map will force all the + # indexes of hiddens[0] before hiddens[1], while the generator + # map will interleave them.) + # NB: We have intentionally skipped __torch_function__ dispatch here. + # See gh-54457 + if self.dim() == 0: + raise TypeError("iteration over a 0-d tensor") + if torch._C._get_tracing_state(): + warnings.warn( + "Iterating over a tensor might cause the trace to be incorrect. " + "Passing a tensor of different shape won't change the number of " + "iterations executed (and might lead to errors or silently give " + "incorrect results).", + category=torch.jit.TracerWarning, + stacklevel=2, + ) + return iter(self.unbind(0)) + + def __hash__(self): + # Do NOT handle __torch_function__ here as user's default + # implementation that handle most functions will most likely do it wrong. + # It can be easily overridden by defining this method on the user + # subclass if needed. + return id(self) + + def __dir__(self): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__dir__, (self,), self) + tensor_methods = dir(self.__class__) + tensor_methods.remove("volatile") # deprecated + attrs = list(self.__dict__.keys()) + keys = tensor_methods + attrs + + # property only available dense, cuda tensors + if (not self.is_cuda) or self.is_sparse: + keys.remove("__cuda_array_interface__") + + return sorted(keys) + + # Numpy array interface, to support `numpy.asarray(tensor) -> ndarray` + __array_priority__ = 1000 # prefer Tensor ops over numpy ones + + def __array__(self, dtype=None): + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__array__, (self,), self, dtype=dtype) + if dtype is None: + return self.numpy() + else: + return self.numpy().astype(dtype, copy=False) + + # Wrap Numpy array again in a suitable tensor when done, to support e.g. + # `numpy.sin(tensor) -> tensor` or `numpy.greater(tensor, 0) -> ByteTensor` + def __array_wrap__(self, array): + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.__array_wrap__, (self,), self, array=array + ) + if array.dtype == bool: + # Workaround, torch has no built-in bool tensor + array = array.astype("uint8") + return torch.from_numpy(array) + + def __contains__(self, element: Any, /) -> bool: + r"""Check if `element` is present in tensor + + Args: + element (Tensor or scalar): element to be checked + for presence in current tensor" + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__contains__, (self,), self, element) + if isinstance( + element, (torch.Tensor, Number, torch.SymInt, torch.SymFloat, torch.SymBool) + ): + # type hint doesn't understand the __contains__ result array + return bool((element == self).any().item()) # type: ignore[union-attr] + + raise RuntimeError( + f"Tensor.__contains__ only supports Tensor or scalar, but you passed in a {type(element)}." + ) + + @property + def __cuda_array_interface__(self): + """Array view description for cuda tensors. + + See: + https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html + """ + if has_torch_function_unary(self): + # TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185 + return handle_torch_function( + Tensor.__cuda_array_interface__.__get__, # type: ignore[attr-defined] + (self,), + self, + ) + + # raise AttributeError for unsupported tensors, so that + # hasattr(cpu_tensor, "__cuda_array_interface__") is False. + if not self.is_cuda: + raise AttributeError( + f"Can't get __cuda_array_interface__ on non-CUDA tensor type: {self.type()} " + "If CUDA data is required use tensor.cuda() to copy tensor to device memory." + ) + + if self.is_sparse: + raise AttributeError( + f"Can't get __cuda_array_interface__ on sparse type: {self.type()} " + "Use Tensor.to_dense() to convert to a dense tensor first." + ) + + # RuntimeError, matching tensor.__array__() behavior. + if self.requires_grad: + raise RuntimeError( + "Can't get __cuda_array_interface__ on Variable that requires grad. " + "If gradients aren't required, use var.detach() to get Variable that doesn't require grad." + ) + + typestr = _dtype_to_typestr(self.dtype) + itemsize = self.element_size() + shape = tuple(self.shape) + if self.is_contiguous(): + # __cuda_array_interface__ v2 requires the strides to be omitted + # (either not set or set to None) for C-contiguous arrays. + strides = None + else: + strides = tuple(s * itemsize for s in self.stride()) + data_ptr = self.data_ptr() if self.numel() > 0 else 0 + data = (data_ptr, False) # read-only is false + + return dict(typestr=typestr, shape=shape, strides=strides, data=data, version=2) + + def storage_type(self): + r"""storage_type() -> type + + Returns the type of the underlying storage. + + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.storage_type, (self,), self) + + torch.storage._warn_typed_storage_removal() + + return self._typed_storage()._get_legacy_storage_class() + + def refine_names(self, *names): + r"""Refines the dimension names of :attr:`self` according to :attr:`names`. + + Refining is a special case of renaming that "lifts" unnamed dimensions. + A ``None`` dim can be refined to have any name; a named dim can only be + refined to have the same name. + + Because named tensors can coexist with unnamed tensors, refining names + gives a nice way to write named-tensor-aware code that works with both + named and unnamed tensors. + + :attr:`names` may contain up to one Ellipsis (``...``). + The Ellipsis is expanded greedily; it is expanded in-place to fill + :attr:`names` to the same length as ``self.dim()`` using names from the + corresponding indices of ``self.names``. + + Python 2 does not support Ellipsis but one may use a string literal + instead (``'...'``). + + Args: + names (iterable of str): The desired names of the output tensor. May + contain up to one Ellipsis. + + Examples:: + + >>> imgs = torch.randn(32, 3, 128, 128) + >>> named_imgs = imgs.refine_names('N', 'C', 'H', 'W') + >>> named_imgs.names + ('N', 'C', 'H', 'W') + + >>> tensor = torch.randn(2, 3, 5, 7, 11) + >>> tensor = tensor.refine_names('A', ..., 'B', 'C') + >>> tensor.names + ('A', None, None, 'B', 'C') + + .. warning:: + The named tensor API is experimental and subject to change. + + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.refine_names, (self,), self, *names) + names = resolve_ellipsis(names, self.names, "refine_names") + return super().refine_names(names) + + def align_to(self, *names): + r"""Permutes the dimensions of the :attr:`self` tensor to match the order + specified in :attr:`names`, adding size-one dims for any new names. + + All of the dims of :attr:`self` must be named in order to use this method. + The resulting tensor is a view on the original tensor. + + All dimension names of :attr:`self` must be present in :attr:`names`. + :attr:`names` may contain additional names that are not in ``self.names``; + the output tensor has a size-one dimension for each of those new names. + + :attr:`names` may contain up to one Ellipsis (``...``). + The Ellipsis is expanded to be equal to all dimension names of :attr:`self` + that are not mentioned in :attr:`names`, in the order that they appear + in :attr:`self`. + + Python 2 does not support Ellipsis but one may use a string literal + instead (``'...'``). + + Args: + names (iterable of str): The desired dimension ordering of the + output tensor. May contain up to one Ellipsis that is expanded + to all unmentioned dim names of :attr:`self`. + + Examples:: + + >>> tensor = torch.randn(2, 2, 2, 2, 2, 2) + >>> named_tensor = tensor.refine_names('A', 'B', 'C', 'D', 'E', 'F') + + # Move the F and E dims to the front while keeping the rest in order + >>> named_tensor.align_to('F', 'E', ...) + + .. warning:: + The named tensor API is experimental and subject to change. + + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.align_to, (self,), self, *names) + ellipsis_idx = single_ellipsis_index(names, "align_to") + if ellipsis_idx is None: + return super().align_to(names) + return super().align_to( + [name for name in names if not is_ellipsis(name)], ellipsis_idx + ) + + def unflatten(self, dim, sizes): # type: ignore[override] + r""" + unflatten(dim, sizes) -> Tensor + + See :func:`torch.unflatten`. + + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.unflatten, (self,), self, dim, sizes) + + if not sizes: + raise RuntimeError("unflatten: sizes must be non-empty") + + names = None + if isinstance(sizes, OrderedDict) or ( + isinstance(sizes, (tuple, list)) and isinstance(sizes[0], (tuple, list)) + ): + names, sizes = unzip_namedshape(sizes) + return super().unflatten(dim, sizes, names) + else: + return super().unflatten(dim, sizes) + + def rename_(self, *names, **rename_map): + """In-place version of :meth:`~Tensor.rename`.""" + + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.rename_, (self,), self, *names, **rename_map + ) + + # Note [rename_ / rename API] + # The Python API for these is different from the C++ API. In Python: + # 1) tensor.rename(*names) takes a vararglist of names + # 2) tensor.rename(**rename_map) takes a map of names to rename. + # C++ is static, making it difficult to implement similar behavior. + return update_names(self, names, rename_map, inplace=True) + + def rename(self, *names, **rename_map): + """Renames dimension names of :attr:`self`. + + There are two main usages: + + ``self.rename(**rename_map)`` returns a view on tensor that has dims + renamed as specified in the mapping :attr:`rename_map`. + + ``self.rename(*names)`` returns a view on tensor, renaming all + dimensions positionally using :attr:`names`. + Use ``self.rename(None)`` to drop names on a tensor. + + One cannot specify both positional args :attr:`names` and keyword args + :attr:`rename_map`. + + Examples:: + + >>> imgs = torch.rand(2, 3, 5, 7, names=('N', 'C', 'H', 'W')) + >>> renamed_imgs = imgs.rename(N='batch', C='channels') + >>> renamed_imgs.names + ('batch', 'channels', 'H', 'W') + + >>> renamed_imgs = imgs.rename(None) + >>> renamed_imgs.names + (None, None, None, None) + + >>> renamed_imgs = imgs.rename('batch', 'channel', 'height', 'width') + >>> renamed_imgs.names + ('batch', 'channel', 'height', 'width') + + .. warning:: + The named tensor API is experimental and subject to change. + + """ + if has_torch_function_unary(self): + return handle_torch_function( + Tensor.rename, (self,), self, *names, **rename_map + ) + + # See Note [rename_ / rename API] + return update_names(self, names, rename_map, inplace=False) + + def to_sparse_coo(self): + """Convert a tensor to :ref:`coordinate format `. + + Examples:: + + >>> dense = torch.randn(5, 5) + >>> sparse = dense.to_sparse_coo() + >>> sparse._nnz() + 25 + + """ + return self.to_sparse() + + def dim_order( + self, *, ambiguity_check: Union[bool, list[torch.memory_format]] = False + ): + """ + dim_order(ambiguity_check=False) -> tuple + + Returns the uniquely determined tuple of int describing the dim order or + physical layout of :attr:`self`. + + The dim order represents how dimensions are laid out in memory of dense tensors, + starting from the outermost to the innermost dimension. + + Note that the dim order may not always be uniquely determined. + If `ambiguity_check` is True, this function raises a RuntimeError when the dim order cannot be uniquely determined; + If `ambiguity_check` is a list of memory formats, this function raises a RuntimeError when tensor can not be interpreted + into exactly one of the given memory formats, or it cannot be uniquely determined. + If `ambiguity_check` is False, it will return one of legal dim order(s) without checking its uniqueness. + Otherwise, it will raise TypeError. + + Args: + ambiguity_check (bool or List[torch.memory_format]): The check method for ambiguity of dim order. + + Examples:: + + >>> torch.empty((2, 3, 5, 7)).dim_order() + (0, 1, 2, 3) + >>> torch.empty((2, 3, 5, 7)).transpose(1, 2).dim_order() + (0, 2, 1, 3) + >>> torch.empty((2, 3, 5, 7), memory_format=torch.channels_last).dim_order() + (0, 2, 3, 1) + >>> torch.empty((1, 2, 3, 4)).dim_order() + (0, 1, 2, 3) + >>> try: + ... torch.empty((1, 2, 3, 4)).dim_order(ambiguity_check=True) + ... except RuntimeError as e: + ... print(e) + The tensor does not have unique dim order, or cannot map to exact one of the given memory formats. + >>> torch.empty((1, 2, 3, 4)).dim_order( + ... ambiguity_check=[torch.contiguous_format, torch.channels_last] + ... ) # It can be mapped to contiguous format + (0, 1, 2, 3) + >>> try: + ... torch.empty((1, 2, 3, 4)).dim_order(ambiguity_check="ILLEGAL") + ... except TypeError as e: + ... print(e) + The ambiguity_check argument must be a bool or a list of memory formats. + + .. warning:: + The dim_order tensor API is experimental and subject to change. + """ + if has_torch_function_unary(self): + return handle_torch_function(Tensor.dim_order, (self,), self) + + if self.is_sparse: + raise AttributeError( + f"Can't get dim order on sparse type: {self.type()} " + "Use Tensor.to_dense() to convert to a dense tensor first." + ) + + # Sanity check ambiguity_check data types + if not isinstance(ambiguity_check, bool): + if not isinstance(ambiguity_check, list): + raise TypeError( + "The ambiguity_check argument must be a bool or a list of memory formats." + ) + for memory_format in ambiguity_check: + if not isinstance(memory_format, torch.memory_format): + raise TypeError( + "The ambiguity_check argument must be a bool or a list of memory formats." + ) + + def invalid_unique_memory_format(tensor, valid_memory_formats): + """ + Returns True if the tensor cannot be uniquely mapped to any of the given memory formats, False otherwise. + """ + + n_legality = 0 + + for memory_format in valid_memory_formats: + if tensor.is_contiguous(memory_format=memory_format): + n_legality += 1 + + return n_legality != 1 + + def has_multiple_dim_order(tensor): + """ + Returns True if there're multiple legal dim orders for given tensor, False otherwise. + + The tensor is considered to have multiple legal dim orders if either of the following conditions is met: + + * Singleton Dimensions: There's at least one singleteon dimension in the tensor. + Since their size is 1, they don't affect the memory offset (stride * index + is zero because index is always zero). Therefore, they can be placed anywhere + in the dimension order without changing how data is accessed. + * Same strides: Strides reflect how the tensor is stored in memory. + If any two dimensions have the same stride, swapping these dimensions won't + change how data is accessed, leading to multiple correct dimension orders. + """ + + sizes = tensor.size() + strides = tensor.stride() + + # Check if there are any duplicate strides + has_duplicate_strides = any( + earlier == later for earlier, later in zip(strides, strides[1:]) + ) + + # Check if there are any singleton dimensions + has_singleton_dims = any(size == 1 for size in sizes) + + return has_duplicate_strides or has_singleton_dims + + valid_memory_formats = ( + ambiguity_check if isinstance(ambiguity_check, list) else [] + ) + check_multiple_dim_order = ( + ambiguity_check if isinstance(ambiguity_check, bool) else True + ) + + if ( + check_multiple_dim_order and has_multiple_dim_order(self) + ) and invalid_unique_memory_format(self, valid_memory_formats): + raise RuntimeError( + "The tensor does not have unique dim order, or cannot map to exact one of the given memory formats." + ) + + import torch._prims_common as utils + + return tuple(utils.compute_elementwise_output_logical_to_physical_perm(self)) + + def _update_names(self, names, inplace): + if has_torch_function_unary(self): + return handle_torch_function( + Tensor._update_names, (self,), self, names, inplace + ) + + # See Note [rename_ / rename API] + if inplace: + return super().rename_(names) + else: + return super().rename(names) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """ + This __torch_function__ implementation wraps subclasses such that + methods called on subclasses return a subclass instance instead of + a ``torch.Tensor`` instance. + + One corollary to this is that you need coverage for torch.Tensor + methods if implementing __torch_function__ for subclasses. + + We recommend always calling ``super().__torch_function__`` as the base + case when doing the above. + + While not mandatory, we recommend making `__torch_function__` a classmethod. + """ + if kwargs is None: + kwargs = {} + + if not all(issubclass(cls, t) for t in types): + return NotImplemented + + with _C.DisableTorchFunctionSubclass(): + ret = func(*args, **kwargs) + if func in get_default_nowrap_functions(): + return ret + else: + return _convert(ret, cls) + + __torch_dispatch__ = _C._disabled_torch_dispatch_impl + + def __dlpack__( + self, + *, + stream: Optional[Any] = None, + max_version: Optional[tuple[int, int]] = None, + dl_device: Optional[tuple[enum.IntEnum, int]] = None, + copy: Optional[bool] = None, + ): + """ + Creates a DLpack `capsule https://data-apis.org/array-api/latest/design_topics/data_interchange.html#data-interchange`_ + of the current tensor to be exported to other libraries. + + This function will be called from the `from_dlpack` method + of the library that will consume the capsule. `from_dlpack` passes the current + stream to this method as part of the specification. + + Args: + stream (integer or None): An optional Python integer representing a + pointer to a CUDA stream. The current stream is synchronized with + this stream before the capsule is created, and since the capsule + shares its storage with the tensor this make it safe to access from + both streams. If None or -1 is passed then no synchronization is performed. + If 1 (on CUDA) or 0 (on ROCM) then the default stream is used for + synchronization. + + max_version (tuple[int, int] or None): An optional Python tuple with + 2 integers, representing the maximum version the caller supports. If + None (default), PyTorch will fallback to DLPack 0.8. + + dl_device (tuple[DLDeviceType, int] or None): An optional tuple specifying + in which device the exported DLPack capsule should be on. If None (default), + the exported DLPack capsule will be on the same device as ``self``. + + copy (bool or None): An optional boolean indicating whether or not to copy + ``self``. If None, PyTorch will copy only if necessary. + """ + if has_torch_function_unary(self): + args = (self,) + kwargs = { + "stream": stream, + "max_version": max_version, + "dl_device": dl_device, + "copy": copy, + } + return handle_torch_function(Tensor.__dlpack__, (self,), *args, **kwargs) + + # DLPack capsules can't capture all of PyTorch's semantics, + # so we prohibit exporting tensors that would lose their properties like + # requires_grad and having the conjugate bit set. + if self.requires_grad: + raise BufferError( + "Can't export tensors that require gradient, use tensor.detach()" + ) + if self.is_conj(): + raise BufferError("Can't export tensors with the conjugate bit set") + if self.layout != torch.strided: + raise BufferError( + "Can't export tensors with layout other than torch.strided" + ) + + if ( + self.device.type == "cuda" + and self.device.index != torch.cuda.current_device() + ): + raise BufferError( + "Can't export tensors on a different CUDA device index. " + f"Expected: {self.device.index}. " + f"Current device: {torch.cuda.current_device()}." + ) + + if stream is not None and type(stream) is not int: + # Stream pointers in CUDA/ROCm are uniquely numbered and can + # be retrieved from their integer value. + raise TypeError("stream must be ``int`` or ``none``") + elif self.device.type == "cuda" and stream != -1: + # NB: This logic handles the special case values for default + # streams and must be kept in sync with from_dlpack in + # torch/utils/dlpack.py + is_rocm = torch.version.hip is not None + is_cuda = not is_rocm + + if stream is None or (is_rocm and stream == 0) or (is_cuda and stream == 1): + stream = torch.cuda.default_stream() + else: + if is_cuda and stream == 2: + raise BufferError("per-thread default stream is not supported.") + + device_str = "CUDA" if is_cuda else "ROCm" + assert (is_cuda and stream != 0) or ( + is_rocm and stream not in (1, 2) + ), f"unsupported stream on {device_str}: {stream}." + + stream = torch.cuda.ExternalStream(stream) + + # Only synchronize on different streams + current_stream = torch.cuda.current_stream() + if stream != current_stream: + event = torch.cuda.Event() + event.record(current_stream) + stream.wait_event(event) + elif self.device.type == "cpu": + assert stream is None, "stream should be None on cpu." + + if self.device.type == "xla": + import torch_xla + import torch_xla.utils.dlpack as xla_dlpack + + if ( + len(torch_xla.real_devices()) <= 0 + or "cuda" not in torch_xla.real_devices()[0].lower() + ): + raise RuntimeError( + "Can't export to dlpack an XLA tensor that is not on CUDA." + ) + + # Does not support DLPack 1.0, yet. + return xla_dlpack.to_dlpack(self) + + if max_version is None or max_version[0] < 1: + # Fallback to the old, unversioned variant. + return _C._to_dlpack(self, dl_device=dl_device, copy=copy) + + return _C._to_dlpack_versioned(self, dl_device=dl_device, copy=copy) + + def __dlpack_device__(self) -> tuple[enum.IntEnum, int]: + if has_torch_function_unary(self): + return handle_torch_function(Tensor.__dlpack_device__, (self,), self) + + from torch.utils.dlpack import DLDeviceType + + device = self.device + idx = device.index if device.index is not None else 0 + torch_device_type = device.type + if torch_device_type == "cuda" and torch.version.hip is not None: + device_type = DLDeviceType.kDLROCM + elif torch_device_type == "cpu" and self.is_pinned(): + device_type = DLDeviceType.kDLCUDAHost + elif torch_device_type == "cuda": + device_type = DLDeviceType.kDLCUDA + elif torch_device_type == "cpu": + device_type = DLDeviceType.kDLCPU + elif torch_device_type == "xpu": + device_type = DLDeviceType.kDLOneAPI + elif self.device.type == "privateuse1": + device_type = DLDeviceType.kDLExtDev + elif torch_device_type == "xla": + import torch_xla + + if ( + len(torch_xla.real_devices()) <= 0 + or "cuda" not in torch_xla.real_devices()[0].lower() + ): + raise ValueError(f"Unknown device type {torch_device_type} for Dlpack") + + device_type = DLDeviceType.kDLCUDA + elif torch_device_type == "mps": + device_type = DLDeviceType.kDLMetal + else: + raise ValueError(f"Unknown device type {torch_device_type} for Dlpack") + return (device_type, idx) + + __module__ = "torch" + + +def _convert(ret, cls): + if cls is Tensor: + return ret + + if isinstance(ret, Tensor) and not isinstance(ret, cls): + ret = ret.as_subclass(cls) + + if isinstance(ret, (tuple, list)): + # Also handles things like namedtuples + ret = type(ret)(_convert(r, cls) for r in ret) + + return ret diff --git a/venv/lib/python3.10/site-packages/torch/_tensor_docs.py b/venv/lib/python3.10/site-packages/torch/_tensor_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..33a2184b71f5e0c249643dccf11145cc93ca0af4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_tensor_docs.py @@ -0,0 +1,6974 @@ +# mypy: allow-untyped-defs +"""Adds docstrings to Tensor functions""" + +import torch._C +from torch._C import _add_docstr as add_docstr +from torch._torch_docs import parse_kwargs, reproducibility_notes + + +def add_docstr_all(method: str, docstr: str) -> None: + add_docstr(getattr(torch._C.TensorBase, method), docstr) + + +common_args = parse_kwargs( + """ + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. +""" +) + +new_common_args = parse_kwargs( + """ + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + dtype (:class:`torch.dtype`, optional): the desired type of returned tensor. + Default: if None, same :class:`torch.dtype` as this tensor. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if None, same :class:`torch.device` as this tensor. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. +""" +) + +add_docstr_all( + "new_tensor", + """ +new_tensor(data, *, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a new Tensor with :attr:`data` as the tensor data. +By default, the returned Tensor has the same :class:`torch.dtype` and +:class:`torch.device` as this tensor. + +.. warning:: + + :func:`new_tensor` always copies :attr:`data`. If you have a Tensor + ``data`` and want to avoid a copy, use :func:`torch.Tensor.requires_grad_` + or :func:`torch.Tensor.detach`. + If you have a numpy array and want to avoid a copy, use + :func:`torch.from_numpy`. + +.. warning:: + + When data is a tensor `x`, :func:`new_tensor()` reads out 'the data' from whatever it is passed, + and constructs a leaf variable. Therefore ``tensor.new_tensor(x)`` is equivalent to ``x.detach().clone()`` + and ``tensor.new_tensor(x, requires_grad=True)`` is equivalent to ``x.detach().clone().requires_grad_(True)``. + The equivalents using ``detach()`` and ``clone()`` are recommended. + +Args: + data (array_like): The returned Tensor copies :attr:`data`. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.ones((2,), dtype=torch.int8) + >>> data = [[0, 1], [2, 3]] + >>> tensor.new_tensor(data) + tensor([[ 0, 1], + [ 2, 3]], dtype=torch.int8) + +""".format(**new_common_args), +) + +add_docstr_all( + "new_full", + """ +new_full(size, fill_value, *, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a Tensor of size :attr:`size` filled with :attr:`fill_value`. +By default, the returned Tensor has the same :class:`torch.dtype` and +:class:`torch.device` as this tensor. + +Args: + fill_value (scalar): the number to fill the output tensor with. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.ones((2,), dtype=torch.float64) + >>> tensor.new_full((3, 4), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416, 3.1416]], dtype=torch.float64) + +""".format(**new_common_args), +) + +add_docstr_all( + "new_empty", + """ +new_empty(size, *, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a Tensor of size :attr:`size` filled with uninitialized data. +By default, the returned Tensor has the same :class:`torch.dtype` and +:class:`torch.device` as this tensor. + +Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.ones(()) + >>> tensor.new_empty((2, 3)) + tensor([[ 5.8182e-18, 4.5765e-41, -1.0545e+30], + [ 3.0949e-41, 4.4842e-44, 0.0000e+00]]) + +""".format(**new_common_args), +) + +add_docstr_all( + "new_empty_strided", + """ +new_empty_strided(size, stride, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a Tensor of size :attr:`size` and strides :attr:`stride` filled with +uninitialized data. By default, the returned Tensor has the same +:class:`torch.dtype` and :class:`torch.device` as this tensor. + +Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.ones(()) + >>> tensor.new_empty_strided((2, 3), (3, 1)) + tensor([[ 5.8182e-18, 4.5765e-41, -1.0545e+30], + [ 3.0949e-41, 4.4842e-44, 0.0000e+00]]) + +""".format(**new_common_args), +) + +add_docstr_all( + "new_ones", + """ +new_ones(size, *, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a Tensor of size :attr:`size` filled with ``1``. +By default, the returned Tensor has the same :class:`torch.dtype` and +:class:`torch.device` as this tensor. + +Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.tensor((), dtype=torch.int32) + >>> tensor.new_ones((2, 3)) + tensor([[ 1, 1, 1], + [ 1, 1, 1]], dtype=torch.int32) + +""".format(**new_common_args), +) + +add_docstr_all( + "new_zeros", + """ +new_zeros(size, *, dtype=None, device=None, requires_grad=False, layout=torch.strided, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a Tensor of size :attr:`size` filled with ``0``. +By default, the returned Tensor has the same :class:`torch.dtype` and +:class:`torch.device` as this tensor. + +Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + +Keyword args: + {dtype} + {device} + {requires_grad} + {layout} + {pin_memory} + +Example:: + + >>> tensor = torch.tensor((), dtype=torch.float64) + >>> tensor.new_zeros((2, 3)) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]], dtype=torch.float64) + +""".format(**new_common_args), +) + +add_docstr_all( + "abs", + r""" +abs() -> Tensor + +See :func:`torch.abs` +""", +) + +add_docstr_all( + "abs_", + r""" +abs_() -> Tensor + +In-place version of :meth:`~Tensor.abs` +""", +) + +add_docstr_all( + "absolute", + r""" +absolute() -> Tensor + +Alias for :func:`abs` +""", +) + +add_docstr_all( + "absolute_", + r""" +absolute_() -> Tensor + +In-place version of :meth:`~Tensor.absolute` +Alias for :func:`abs_` +""", +) + +add_docstr_all( + "acos", + r""" +acos() -> Tensor + +See :func:`torch.acos` +""", +) + +add_docstr_all( + "acos_", + r""" +acos_() -> Tensor + +In-place version of :meth:`~Tensor.acos` +""", +) + +add_docstr_all( + "arccos", + r""" +arccos() -> Tensor + +See :func:`torch.arccos` +""", +) + +add_docstr_all( + "arccos_", + r""" +arccos_() -> Tensor + +In-place version of :meth:`~Tensor.arccos` +""", +) + +add_docstr_all( + "acosh", + r""" +acosh() -> Tensor + +See :func:`torch.acosh` +""", +) + +add_docstr_all( + "acosh_", + r""" +acosh_() -> Tensor + +In-place version of :meth:`~Tensor.acosh` +""", +) + +add_docstr_all( + "arccosh", + r""" +acosh() -> Tensor + +See :func:`torch.arccosh` +""", +) + +add_docstr_all( + "arccosh_", + r""" +acosh_() -> Tensor + +In-place version of :meth:`~Tensor.arccosh` +""", +) + +add_docstr_all( + "add", + r""" +add(other, *, alpha=1) -> Tensor + +Add a scalar or tensor to :attr:`self` tensor. If both :attr:`alpha` +and :attr:`other` are specified, each element of :attr:`other` is scaled by +:attr:`alpha` before being used. + +When :attr:`other` is a tensor, the shape of :attr:`other` must be +:ref:`broadcastable ` with the shape of the underlying +tensor + +See :func:`torch.add` +""", +) + +add_docstr_all( + "add_", + r""" +add_(other, *, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.add` +""", +) + +add_docstr_all( + "addbmm", + r""" +addbmm(batch1, batch2, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.addbmm` +""", +) + +add_docstr_all( + "addbmm_", + r""" +addbmm_(batch1, batch2, *, beta=1, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.addbmm` +""", +) + +add_docstr_all( + "addcdiv", + r""" +addcdiv(tensor1, tensor2, *, value=1) -> Tensor + +See :func:`torch.addcdiv` +""", +) + +add_docstr_all( + "addcdiv_", + r""" +addcdiv_(tensor1, tensor2, *, value=1) -> Tensor + +In-place version of :meth:`~Tensor.addcdiv` +""", +) + +add_docstr_all( + "addcmul", + r""" +addcmul(tensor1, tensor2, *, value=1) -> Tensor + +See :func:`torch.addcmul` +""", +) + +add_docstr_all( + "addcmul_", + r""" +addcmul_(tensor1, tensor2, *, value=1) -> Tensor + +In-place version of :meth:`~Tensor.addcmul` +""", +) + +add_docstr_all( + "addmm", + r""" +addmm(mat1, mat2, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.addmm` +""", +) + +add_docstr_all( + "addmm_", + r""" +addmm_(mat1, mat2, *, beta=1, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.addmm` +""", +) + +add_docstr_all( + "addmv", + r""" +addmv(mat, vec, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.addmv` +""", +) + +add_docstr_all( + "addmv_", + r""" +addmv_(mat, vec, *, beta=1, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.addmv` +""", +) + +add_docstr_all( + "sspaddmm", + r""" +sspaddmm(mat1, mat2, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.sspaddmm` +""", +) + +add_docstr_all( + "smm", + r""" +smm(mat) -> Tensor + +See :func:`torch.smm` +""", +) + +add_docstr_all( + "addr", + r""" +addr(vec1, vec2, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.addr` +""", +) + +add_docstr_all( + "addr_", + r""" +addr_(vec1, vec2, *, beta=1, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.addr` +""", +) + +add_docstr_all( + "align_as", + r""" +align_as(other) -> Tensor + +Permutes the dimensions of the :attr:`self` tensor to match the dimension order +in the :attr:`other` tensor, adding size-one dims for any new names. + +This operation is useful for explicit broadcasting by names (see examples). + +All of the dims of :attr:`self` must be named in order to use this method. +The resulting tensor is a view on the original tensor. + +All dimension names of :attr:`self` must be present in ``other.names``. +:attr:`other` may contain named dimensions that are not in ``self.names``; +the output tensor has a size-one dimension for each of those new names. + +To align a tensor to a specific order, use :meth:`~Tensor.align_to`. + +Examples:: + + # Example 1: Applying a mask + >>> mask = torch.randint(2, [127, 128], dtype=torch.bool).refine_names('W', 'H') + >>> imgs = torch.randn(32, 128, 127, 3, names=('N', 'H', 'W', 'C')) + >>> imgs.masked_fill_(mask.align_as(imgs), 0) + + + # Example 2: Applying a per-channel-scale + >>> def scale_channels(input, scale): + >>> scale = scale.refine_names('C') + >>> return input * scale.align_as(input) + + >>> num_channels = 3 + >>> scale = torch.randn(num_channels, names=('C',)) + >>> imgs = torch.rand(32, 128, 128, num_channels, names=('N', 'H', 'W', 'C')) + >>> more_imgs = torch.rand(32, num_channels, 128, 128, names=('N', 'C', 'H', 'W')) + >>> videos = torch.randn(3, num_channels, 128, 128, 128, names=('N', 'C', 'H', 'W', 'D')) + + # scale_channels is agnostic to the dimension order of the input + >>> scale_channels(imgs, scale) + >>> scale_channels(more_imgs, scale) + >>> scale_channels(videos, scale) + +.. warning:: + The named tensor API is experimental and subject to change. + +""", +) + +add_docstr_all( + "all", + r""" +all(dim=None, keepdim=False) -> Tensor + +See :func:`torch.all` +""", +) + +add_docstr_all( + "allclose", + r""" +allclose(other, rtol=1e-05, atol=1e-08, equal_nan=False) -> Tensor + +See :func:`torch.allclose` +""", +) + +add_docstr_all( + "angle", + r""" +angle() -> Tensor + +See :func:`torch.angle` +""", +) + +add_docstr_all( + "any", + r""" +any(dim=None, keepdim=False) -> Tensor + +See :func:`torch.any` +""", +) + +add_docstr_all( + "apply_", + r""" +apply_(callable) -> Tensor + +Applies the function :attr:`callable` to each element in the tensor, replacing +each element with the value returned by :attr:`callable`. + +.. note:: + + This function only works with CPU tensors and should not be used in code + sections that require high performance. +""", +) + +add_docstr_all( + "asin", + r""" +asin() -> Tensor + +See :func:`torch.asin` +""", +) + +add_docstr_all( + "asin_", + r""" +asin_() -> Tensor + +In-place version of :meth:`~Tensor.asin` +""", +) + +add_docstr_all( + "arcsin", + r""" +arcsin() -> Tensor + +See :func:`torch.arcsin` +""", +) + +add_docstr_all( + "arcsin_", + r""" +arcsin_() -> Tensor + +In-place version of :meth:`~Tensor.arcsin` +""", +) + +add_docstr_all( + "asinh", + r""" +asinh() -> Tensor + +See :func:`torch.asinh` +""", +) + +add_docstr_all( + "asinh_", + r""" +asinh_() -> Tensor + +In-place version of :meth:`~Tensor.asinh` +""", +) + +add_docstr_all( + "arcsinh", + r""" +arcsinh() -> Tensor + +See :func:`torch.arcsinh` +""", +) + +add_docstr_all( + "arcsinh_", + r""" +arcsinh_() -> Tensor + +In-place version of :meth:`~Tensor.arcsinh` +""", +) + +add_docstr_all( + "as_strided", + r""" +as_strided(size, stride, storage_offset=None) -> Tensor + +See :func:`torch.as_strided` +""", +) + +add_docstr_all( + "as_strided_", + r""" +as_strided_(size, stride, storage_offset=None) -> Tensor + +In-place version of :meth:`~Tensor.as_strided` +""", +) + +add_docstr_all( + "atan", + r""" +atan() -> Tensor + +See :func:`torch.atan` +""", +) + +add_docstr_all( + "atan_", + r""" +atan_() -> Tensor + +In-place version of :meth:`~Tensor.atan` +""", +) + +add_docstr_all( + "arctan", + r""" +arctan() -> Tensor + +See :func:`torch.arctan` +""", +) + +add_docstr_all( + "arctan_", + r""" +arctan_() -> Tensor + +In-place version of :meth:`~Tensor.arctan` +""", +) + +add_docstr_all( + "atan2", + r""" +atan2(other) -> Tensor + +See :func:`torch.atan2` +""", +) + +add_docstr_all( + "atan2_", + r""" +atan2_(other) -> Tensor + +In-place version of :meth:`~Tensor.atan2` +""", +) + +add_docstr_all( + "arctan2", + r""" +arctan2(other) -> Tensor + +See :func:`torch.arctan2` +""", +) + +add_docstr_all( + "arctan2_", + r""" +atan2_(other) -> Tensor + +In-place version of :meth:`~Tensor.arctan2` +""", +) + +add_docstr_all( + "atanh", + r""" +atanh() -> Tensor + +See :func:`torch.atanh` +""", +) + +add_docstr_all( + "atanh_", + r""" +atanh_(other) -> Tensor + +In-place version of :meth:`~Tensor.atanh` +""", +) + +add_docstr_all( + "arctanh", + r""" +arctanh() -> Tensor + +See :func:`torch.arctanh` +""", +) + +add_docstr_all( + "arctanh_", + r""" +arctanh_(other) -> Tensor + +In-place version of :meth:`~Tensor.arctanh` +""", +) + +add_docstr_all( + "baddbmm", + r""" +baddbmm(batch1, batch2, *, beta=1, alpha=1) -> Tensor + +See :func:`torch.baddbmm` +""", +) + +add_docstr_all( + "baddbmm_", + r""" +baddbmm_(batch1, batch2, *, beta=1, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.baddbmm` +""", +) + +add_docstr_all( + "bernoulli", + r""" +bernoulli(*, generator=None) -> Tensor + +Returns a result tensor where each :math:`\texttt{result[i]}` is independently +sampled from :math:`\text{Bernoulli}(\texttt{self[i]})`. :attr:`self` must have +floating point ``dtype``, and the result will have the same ``dtype``. + +See :func:`torch.bernoulli` +""", +) + +add_docstr_all( + "bernoulli_", + r""" +bernoulli_(p=0.5, *, generator=None) -> Tensor + +Fills each location of :attr:`self` with an independent sample from +:math:`\text{Bernoulli}(\texttt{p})`. :attr:`self` can have integral +``dtype``. + +:attr:`p` should either be a scalar or tensor containing probabilities to be +used for drawing the binary random number. + +If it is a tensor, the :math:`\text{i}^{th}` element of :attr:`self` tensor +will be set to a value sampled from +:math:`\text{Bernoulli}(\texttt{p\_tensor[i]})`. In this case `p` must have +floating point ``dtype``. + +See also :meth:`~Tensor.bernoulli` and :func:`torch.bernoulli` +""", +) + +add_docstr_all( + "bincount", + r""" +bincount(weights=None, minlength=0) -> Tensor + +See :func:`torch.bincount` +""", +) + +add_docstr_all( + "bitwise_not", + r""" +bitwise_not() -> Tensor + +See :func:`torch.bitwise_not` +""", +) + +add_docstr_all( + "bitwise_not_", + r""" +bitwise_not_() -> Tensor + +In-place version of :meth:`~Tensor.bitwise_not` +""", +) + +add_docstr_all( + "bitwise_and", + r""" +bitwise_and() -> Tensor + +See :func:`torch.bitwise_and` +""", +) + +add_docstr_all( + "bitwise_and_", + r""" +bitwise_and_() -> Tensor + +In-place version of :meth:`~Tensor.bitwise_and` +""", +) + +add_docstr_all( + "bitwise_or", + r""" +bitwise_or() -> Tensor + +See :func:`torch.bitwise_or` +""", +) + +add_docstr_all( + "bitwise_or_", + r""" +bitwise_or_() -> Tensor + +In-place version of :meth:`~Tensor.bitwise_or` +""", +) + +add_docstr_all( + "bitwise_xor", + r""" +bitwise_xor() -> Tensor + +See :func:`torch.bitwise_xor` +""", +) + +add_docstr_all( + "bitwise_xor_", + r""" +bitwise_xor_() -> Tensor + +In-place version of :meth:`~Tensor.bitwise_xor` +""", +) + +add_docstr_all( + "bitwise_left_shift", + r""" +bitwise_left_shift(other) -> Tensor + +See :func:`torch.bitwise_left_shift` +""", +) + +add_docstr_all( + "bitwise_left_shift_", + r""" +bitwise_left_shift_(other) -> Tensor + +In-place version of :meth:`~Tensor.bitwise_left_shift` +""", +) + +add_docstr_all( + "bitwise_right_shift", + r""" +bitwise_right_shift(other) -> Tensor + +See :func:`torch.bitwise_right_shift` +""", +) + +add_docstr_all( + "bitwise_right_shift_", + r""" +bitwise_right_shift_(other) -> Tensor + +In-place version of :meth:`~Tensor.bitwise_right_shift` +""", +) + +add_docstr_all( + "broadcast_to", + r""" +broadcast_to(shape) -> Tensor + +See :func:`torch.broadcast_to`. +""", +) + +add_docstr_all( + "logical_and", + r""" +logical_and() -> Tensor + +See :func:`torch.logical_and` +""", +) + +add_docstr_all( + "logical_and_", + r""" +logical_and_() -> Tensor + +In-place version of :meth:`~Tensor.logical_and` +""", +) + +add_docstr_all( + "logical_not", + r""" +logical_not() -> Tensor + +See :func:`torch.logical_not` +""", +) + +add_docstr_all( + "logical_not_", + r""" +logical_not_() -> Tensor + +In-place version of :meth:`~Tensor.logical_not` +""", +) + +add_docstr_all( + "logical_or", + r""" +logical_or() -> Tensor + +See :func:`torch.logical_or` +""", +) + +add_docstr_all( + "logical_or_", + r""" +logical_or_() -> Tensor + +In-place version of :meth:`~Tensor.logical_or` +""", +) + +add_docstr_all( + "logical_xor", + r""" +logical_xor() -> Tensor + +See :func:`torch.logical_xor` +""", +) + +add_docstr_all( + "logical_xor_", + r""" +logical_xor_() -> Tensor + +In-place version of :meth:`~Tensor.logical_xor` +""", +) + +add_docstr_all( + "bmm", + r""" +bmm(batch2) -> Tensor + +See :func:`torch.bmm` +""", +) + +add_docstr_all( + "cauchy_", + r""" +cauchy_(median=0, sigma=1, *, generator=None) -> Tensor + +Fills the tensor with numbers drawn from the Cauchy distribution: + +.. math:: + + f(x) = \dfrac{1}{\pi} \dfrac{\sigma}{(x - \text{median})^2 + \sigma^2} + +.. note:: + Sigma (:math:`\sigma`) is used to denote the scale parameter in Cauchy distribution. +""", +) + +add_docstr_all( + "ceil", + r""" +ceil() -> Tensor + +See :func:`torch.ceil` +""", +) + +add_docstr_all( + "ceil_", + r""" +ceil_() -> Tensor + +In-place version of :meth:`~Tensor.ceil` +""", +) + +add_docstr_all( + "cholesky", + r""" +cholesky(upper=False) -> Tensor + +See :func:`torch.cholesky` +""", +) + +add_docstr_all( + "cholesky_solve", + r""" +cholesky_solve(input2, upper=False) -> Tensor + +See :func:`torch.cholesky_solve` +""", +) + +add_docstr_all( + "cholesky_inverse", + r""" +cholesky_inverse(upper=False) -> Tensor + +See :func:`torch.cholesky_inverse` +""", +) + +add_docstr_all( + "clamp", + r""" +clamp(min=None, max=None) -> Tensor + +See :func:`torch.clamp` +""", +) + +add_docstr_all( + "clamp_", + r""" +clamp_(min=None, max=None) -> Tensor + +In-place version of :meth:`~Tensor.clamp` +""", +) + +add_docstr_all( + "clip", + r""" +clip(min=None, max=None) -> Tensor + +Alias for :meth:`~Tensor.clamp`. +""", +) + +add_docstr_all( + "clip_", + r""" +clip_(min=None, max=None) -> Tensor + +Alias for :meth:`~Tensor.clamp_`. +""", +) + +add_docstr_all( + "clone", + r""" +clone(*, memory_format=torch.preserve_format) -> Tensor + +See :func:`torch.clone` +""".format(**common_args), +) + +add_docstr_all( + "coalesce", + r""" +coalesce() -> Tensor + +Returns a coalesced copy of :attr:`self` if :attr:`self` is an +:ref:`uncoalesced tensor `. + +Returns :attr:`self` if :attr:`self` is a coalesced tensor. + +.. warning:: + Throws an error if :attr:`self` is not a sparse COO tensor. +""", +) + +add_docstr_all( + "contiguous", + r""" +contiguous(memory_format=torch.contiguous_format) -> Tensor + +Returns a contiguous in memory tensor containing the same data as :attr:`self` tensor. If +:attr:`self` tensor is already in the specified memory format, this function returns the +:attr:`self` tensor. + +Args: + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. +""", +) + +add_docstr_all( + "copy_", + r""" +copy_(src, non_blocking=False) -> Tensor + +Copies the elements from :attr:`src` into :attr:`self` tensor and returns +:attr:`self`. + +The :attr:`src` tensor must be :ref:`broadcastable ` +with the :attr:`self` tensor. It may be of a different data type or reside on a +different device. + +Args: + src (Tensor): the source tensor to copy from + non_blocking (bool, optional): if ``True`` and this copy is between CPU and GPU, + the copy may occur asynchronously with respect to the host. For other + cases, this argument has no effect. Default: ``False`` +""", +) + +add_docstr_all( + "conj", + r""" +conj() -> Tensor + +See :func:`torch.conj` +""", +) + +add_docstr_all( + "conj_physical", + r""" +conj_physical() -> Tensor + +See :func:`torch.conj_physical` +""", +) + +add_docstr_all( + "conj_physical_", + r""" +conj_physical_() -> Tensor + +In-place version of :meth:`~Tensor.conj_physical` +""", +) + +add_docstr_all( + "resolve_conj", + r""" +resolve_conj() -> Tensor + +See :func:`torch.resolve_conj` +""", +) + +add_docstr_all( + "resolve_neg", + r""" +resolve_neg() -> Tensor + +See :func:`torch.resolve_neg` +""", +) + +add_docstr_all( + "copysign", + r""" +copysign(other) -> Tensor + +See :func:`torch.copysign` +""", +) + +add_docstr_all( + "copysign_", + r""" +copysign_(other) -> Tensor + +In-place version of :meth:`~Tensor.copysign` +""", +) + +add_docstr_all( + "cos", + r""" +cos() -> Tensor + +See :func:`torch.cos` +""", +) + +add_docstr_all( + "cos_", + r""" +cos_() -> Tensor + +In-place version of :meth:`~Tensor.cos` +""", +) + +add_docstr_all( + "cosh", + r""" +cosh() -> Tensor + +See :func:`torch.cosh` +""", +) + +add_docstr_all( + "cosh_", + r""" +cosh_() -> Tensor + +In-place version of :meth:`~Tensor.cosh` +""", +) + +add_docstr_all( + "cpu", + r""" +cpu(memory_format=torch.preserve_format) -> Tensor + +Returns a copy of this object in CPU memory. + +If this object is already in CPU memory, +then no copy is performed and the original object is returned. + +Args: + {memory_format} + +""".format(**common_args), +) + +add_docstr_all( + "count_nonzero", + r""" +count_nonzero(dim=None) -> Tensor + +See :func:`torch.count_nonzero` +""", +) + +add_docstr_all( + "cov", + r""" +cov(*, correction=1, fweights=None, aweights=None) -> Tensor + +See :func:`torch.cov` +""", +) + +add_docstr_all( + "corrcoef", + r""" +corrcoef() -> Tensor + +See :func:`torch.corrcoef` +""", +) + +add_docstr_all( + "cross", + r""" +cross(other, dim=None) -> Tensor + +See :func:`torch.cross` +""", +) + +add_docstr_all( + "cuda", + r""" +cuda(device=None, non_blocking=False, memory_format=torch.preserve_format) -> Tensor + +Returns a copy of this object in CUDA memory. + +If this object is already in CUDA memory and on the correct device, +then no copy is performed and the original object is returned. + +Args: + device (:class:`torch.device`, optional): The destination GPU device. + Defaults to the current CUDA device. + non_blocking (bool, optional): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. + Otherwise, the argument has no effect. Default: ``False``. + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "mtia", + r""" +mtia(device=None, non_blocking=False, memory_format=torch.preserve_format) -> Tensor + +Returns a copy of this object in MTIA memory. + +If this object is already in MTIA memory and on the correct device, +then no copy is performed and the original object is returned. + +Args: + device (:class:`torch.device`, optional): The destination MTIA device. + Defaults to the current MTIA device. + non_blocking (bool, optional): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. + Otherwise, the argument has no effect. Default: ``False``. + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "ipu", + r""" +ipu(device=None, non_blocking=False, memory_format=torch.preserve_format) -> Tensor + +Returns a copy of this object in IPU memory. + +If this object is already in IPU memory and on the correct device, +then no copy is performed and the original object is returned. + +Args: + device (:class:`torch.device`, optional): The destination IPU device. + Defaults to the current IPU device. + non_blocking (bool, optional): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. + Otherwise, the argument has no effect. Default: ``False``. + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "xpu", + r""" +xpu(device=None, non_blocking=False, memory_format=torch.preserve_format) -> Tensor + +Returns a copy of this object in XPU memory. + +If this object is already in XPU memory and on the correct device, +then no copy is performed and the original object is returned. + +Args: + device (:class:`torch.device`, optional): The destination XPU device. + Defaults to the current XPU device. + non_blocking (bool, optional): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. + Otherwise, the argument has no effect. Default: ``False``. + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "logcumsumexp", + r""" +logcumsumexp(dim) -> Tensor + +See :func:`torch.logcumsumexp` +""", +) + +add_docstr_all( + "cummax", + r""" +cummax(dim) -> (Tensor, Tensor) + +See :func:`torch.cummax` +""", +) + +add_docstr_all( + "cummin", + r""" +cummin(dim) -> (Tensor, Tensor) + +See :func:`torch.cummin` +""", +) + +add_docstr_all( + "cumprod", + r""" +cumprod(dim, dtype=None) -> Tensor + +See :func:`torch.cumprod` +""", +) + +add_docstr_all( + "cumprod_", + r""" +cumprod_(dim, dtype=None) -> Tensor + +In-place version of :meth:`~Tensor.cumprod` +""", +) + +add_docstr_all( + "cumsum", + r""" +cumsum(dim, dtype=None) -> Tensor + +See :func:`torch.cumsum` +""", +) + +add_docstr_all( + "cumsum_", + r""" +cumsum_(dim, dtype=None) -> Tensor + +In-place version of :meth:`~Tensor.cumsum` +""", +) + +add_docstr_all( + "data_ptr", + r""" +data_ptr() -> int + +Returns the address of the first element of :attr:`self` tensor. +""", +) + +add_docstr_all( + "dequantize", + r""" +dequantize() -> Tensor + +Given a quantized Tensor, dequantize it and return the dequantized float Tensor. +""", +) + +add_docstr_all( + "dense_dim", + r""" +dense_dim() -> int + +Return the number of dense dimensions in a :ref:`sparse tensor ` :attr:`self`. + +.. note:: + Returns ``len(self.shape)`` if :attr:`self` is not a sparse tensor. + +See also :meth:`Tensor.sparse_dim` and :ref:`hybrid tensors `. +""", +) + +add_docstr_all( + "diag", + r""" +diag(diagonal=0) -> Tensor + +See :func:`torch.diag` +""", +) + +add_docstr_all( + "diag_embed", + r""" +diag_embed(offset=0, dim1=-2, dim2=-1) -> Tensor + +See :func:`torch.diag_embed` +""", +) + +add_docstr_all( + "diagflat", + r""" +diagflat(offset=0) -> Tensor + +See :func:`torch.diagflat` +""", +) + +add_docstr_all( + "diagonal", + r""" +diagonal(offset=0, dim1=0, dim2=1) -> Tensor + +See :func:`torch.diagonal` +""", +) + +add_docstr_all( + "diagonal_scatter", + r""" +diagonal_scatter(src, offset=0, dim1=0, dim2=1) -> Tensor + +See :func:`torch.diagonal_scatter` +""", +) + +add_docstr_all( + "as_strided_scatter", + r""" +as_strided_scatter(src, size, stride, storage_offset=None) -> Tensor + +See :func:`torch.as_strided_scatter` +""", +) + +add_docstr_all( + "fill_diagonal_", + r""" +fill_diagonal_(fill_value, wrap=False) -> Tensor + +Fill the main diagonal of a tensor that has at least 2-dimensions. +When dims>2, all dimensions of input must be of equal length. +This function modifies the input tensor in-place, and returns the input tensor. + +Arguments: + fill_value (Scalar): the fill value + wrap (bool, optional): the diagonal 'wrapped' after N columns for tall matrices. Default: ``False`` + +Example:: + + >>> a = torch.zeros(3, 3) + >>> a.fill_diagonal_(5) + tensor([[5., 0., 0.], + [0., 5., 0.], + [0., 0., 5.]]) + >>> b = torch.zeros(7, 3) + >>> b.fill_diagonal_(5) + tensor([[5., 0., 0.], + [0., 5., 0.], + [0., 0., 5.], + [0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]]) + >>> c = torch.zeros(7, 3) + >>> c.fill_diagonal_(5, wrap=True) + tensor([[5., 0., 0.], + [0., 5., 0.], + [0., 0., 5.], + [0., 0., 0.], + [5., 0., 0.], + [0., 5., 0.], + [0., 0., 5.]]) + +""", +) + +add_docstr_all( + "floor_divide", + r""" +floor_divide(value) -> Tensor + +See :func:`torch.floor_divide` +""", +) + +add_docstr_all( + "floor_divide_", + r""" +floor_divide_(value) -> Tensor + +In-place version of :meth:`~Tensor.floor_divide` +""", +) + +add_docstr_all( + "diff", + r""" +diff(n=1, dim=-1, prepend=None, append=None) -> Tensor + +See :func:`torch.diff` +""", +) + +add_docstr_all( + "digamma", + r""" +digamma() -> Tensor + +See :func:`torch.digamma` +""", +) + +add_docstr_all( + "digamma_", + r""" +digamma_() -> Tensor + +In-place version of :meth:`~Tensor.digamma` +""", +) + +add_docstr_all( + "dim", + r""" +dim() -> int + +Returns the number of dimensions of :attr:`self` tensor. +""", +) + +add_docstr_all( + "dist", + r""" +dist(other, p=2) -> Tensor + +See :func:`torch.dist` +""", +) + +add_docstr_all( + "div", + r""" +div(value, *, rounding_mode=None) -> Tensor + +See :func:`torch.div` +""", +) + +add_docstr_all( + "div_", + r""" +div_(value, *, rounding_mode=None) -> Tensor + +In-place version of :meth:`~Tensor.div` +""", +) + +add_docstr_all( + "divide", + r""" +divide(value, *, rounding_mode=None) -> Tensor + +See :func:`torch.divide` +""", +) + +add_docstr_all( + "divide_", + r""" +divide_(value, *, rounding_mode=None) -> Tensor + +In-place version of :meth:`~Tensor.divide` +""", +) + +add_docstr_all( + "dot", + r""" +dot(other) -> Tensor + +See :func:`torch.dot` +""", +) + +add_docstr_all( + "element_size", + r""" +element_size() -> int + +Returns the size in bytes of an individual element. + +Example:: + + >>> torch.tensor([]).element_size() + 4 + >>> torch.tensor([], dtype=torch.uint8).element_size() + 1 + +""", +) + +add_docstr_all( + "eq", + r""" +eq(other) -> Tensor + +See :func:`torch.eq` +""", +) + +add_docstr_all( + "eq_", + r""" +eq_(other) -> Tensor + +In-place version of :meth:`~Tensor.eq` +""", +) + +add_docstr_all( + "equal", + r""" +equal(other) -> bool + +See :func:`torch.equal` +""", +) + +add_docstr_all( + "erf", + r""" +erf() -> Tensor + +See :func:`torch.erf` +""", +) + +add_docstr_all( + "erf_", + r""" +erf_() -> Tensor + +In-place version of :meth:`~Tensor.erf` +""", +) + +add_docstr_all( + "erfc", + r""" +erfc() -> Tensor + +See :func:`torch.erfc` +""", +) + +add_docstr_all( + "erfc_", + r""" +erfc_() -> Tensor + +In-place version of :meth:`~Tensor.erfc` +""", +) + +add_docstr_all( + "erfinv", + r""" +erfinv() -> Tensor + +See :func:`torch.erfinv` +""", +) + +add_docstr_all( + "erfinv_", + r""" +erfinv_() -> Tensor + +In-place version of :meth:`~Tensor.erfinv` +""", +) + +add_docstr_all( + "exp", + r""" +exp() -> Tensor + +See :func:`torch.exp` +""", +) + +add_docstr_all( + "exp_", + r""" +exp_() -> Tensor + +In-place version of :meth:`~Tensor.exp` +""", +) + +add_docstr_all( + "exp2", + r""" +exp2() -> Tensor + +See :func:`torch.exp2` +""", +) + +add_docstr_all( + "exp2_", + r""" +exp2_() -> Tensor + +In-place version of :meth:`~Tensor.exp2` +""", +) + +add_docstr_all( + "expm1", + r""" +expm1() -> Tensor + +See :func:`torch.expm1` +""", +) + +add_docstr_all( + "expm1_", + r""" +expm1_() -> Tensor + +In-place version of :meth:`~Tensor.expm1` +""", +) + +add_docstr_all( + "exponential_", + r""" +exponential_(lambd=1, *, generator=None) -> Tensor + +Fills :attr:`self` tensor with elements drawn from the PDF (probability density function): + +.. math:: + + f(x) = \lambda e^{-\lambda x}, x > 0 + +.. note:: + In probability theory, exponential distribution is supported on interval [0, :math:`\inf`) (i.e., :math:`x >= 0`) + implying that zero can be sampled from the exponential distribution. + However, :func:`torch.Tensor.exponential_` does not sample zero, + which means that its actual support is the interval (0, :math:`\inf`). + + Note that :func:`torch.distributions.exponential.Exponential` is supported on the interval [0, :math:`\inf`) and can sample zero. +""", +) + +add_docstr_all( + "fill_", + r""" +fill_(value) -> Tensor + +Fills :attr:`self` tensor with the specified value. +""", +) + +add_docstr_all( + "floor", + r""" +floor() -> Tensor + +See :func:`torch.floor` +""", +) + +add_docstr_all( + "flip", + r""" +flip(dims) -> Tensor + +See :func:`torch.flip` +""", +) + +add_docstr_all( + "fliplr", + r""" +fliplr() -> Tensor + +See :func:`torch.fliplr` +""", +) + +add_docstr_all( + "flipud", + r""" +flipud() -> Tensor + +See :func:`torch.flipud` +""", +) + +add_docstr_all( + "roll", + r""" +roll(shifts, dims) -> Tensor + +See :func:`torch.roll` +""", +) + +add_docstr_all( + "floor_", + r""" +floor_() -> Tensor + +In-place version of :meth:`~Tensor.floor` +""", +) + +add_docstr_all( + "fmod", + r""" +fmod(divisor) -> Tensor + +See :func:`torch.fmod` +""", +) + +add_docstr_all( + "fmod_", + r""" +fmod_(divisor) -> Tensor + +In-place version of :meth:`~Tensor.fmod` +""", +) + +add_docstr_all( + "frac", + r""" +frac() -> Tensor + +See :func:`torch.frac` +""", +) + +add_docstr_all( + "frac_", + r""" +frac_() -> Tensor + +In-place version of :meth:`~Tensor.frac` +""", +) + +add_docstr_all( + "frexp", + r""" +frexp(input) -> (Tensor mantissa, Tensor exponent) + +See :func:`torch.frexp` +""", +) + +add_docstr_all( + "flatten", + r""" +flatten(start_dim=0, end_dim=-1) -> Tensor + +See :func:`torch.flatten` +""", +) + +add_docstr_all( + "gather", + r""" +gather(dim, index) -> Tensor + +See :func:`torch.gather` +""", +) + +add_docstr_all( + "gcd", + r""" +gcd(other) -> Tensor + +See :func:`torch.gcd` +""", +) + +add_docstr_all( + "gcd_", + r""" +gcd_(other) -> Tensor + +In-place version of :meth:`~Tensor.gcd` +""", +) + +add_docstr_all( + "ge", + r""" +ge(other) -> Tensor + +See :func:`torch.ge`. +""", +) + +add_docstr_all( + "ge_", + r""" +ge_(other) -> Tensor + +In-place version of :meth:`~Tensor.ge`. +""", +) + +add_docstr_all( + "greater_equal", + r""" +greater_equal(other) -> Tensor + +See :func:`torch.greater_equal`. +""", +) + +add_docstr_all( + "greater_equal_", + r""" +greater_equal_(other) -> Tensor + +In-place version of :meth:`~Tensor.greater_equal`. +""", +) + +add_docstr_all( + "geometric_", + r""" +geometric_(p, *, generator=None) -> Tensor + +Fills :attr:`self` tensor with elements drawn from the geometric distribution: + +.. math:: + + P(X=k) = (1 - p)^{k - 1} p, k = 1, 2, ... + +.. note:: + :func:`torch.Tensor.geometric_` `k`-th trial is the first success hence draws samples in :math:`\{1, 2, \ldots\}`, whereas + :func:`torch.distributions.geometric.Geometric` :math:`(k+1)`-th trial is the first success + hence draws samples in :math:`\{0, 1, \ldots\}`. +""", +) + +add_docstr_all( + "geqrf", + r""" +geqrf() -> (Tensor, Tensor) + +See :func:`torch.geqrf` +""", +) + +add_docstr_all( + "ger", + r""" +ger(vec2) -> Tensor + +See :func:`torch.ger` +""", +) + +add_docstr_all( + "inner", + r""" +inner(other) -> Tensor + +See :func:`torch.inner`. +""", +) + +add_docstr_all( + "outer", + r""" +outer(vec2) -> Tensor + +See :func:`torch.outer`. +""", +) + +add_docstr_all( + "hypot", + r""" +hypot(other) -> Tensor + +See :func:`torch.hypot` +""", +) + +add_docstr_all( + "hypot_", + r""" +hypot_(other) -> Tensor + +In-place version of :meth:`~Tensor.hypot` +""", +) + +add_docstr_all( + "i0", + r""" +i0() -> Tensor + +See :func:`torch.i0` +""", +) + +add_docstr_all( + "i0_", + r""" +i0_() -> Tensor + +In-place version of :meth:`~Tensor.i0` +""", +) + +add_docstr_all( + "igamma", + r""" +igamma(other) -> Tensor + +See :func:`torch.igamma` +""", +) + +add_docstr_all( + "igamma_", + r""" +igamma_(other) -> Tensor + +In-place version of :meth:`~Tensor.igamma` +""", +) + +add_docstr_all( + "igammac", + r""" +igammac(other) -> Tensor +See :func:`torch.igammac` +""", +) + +add_docstr_all( + "igammac_", + r""" +igammac_(other) -> Tensor +In-place version of :meth:`~Tensor.igammac` +""", +) + +add_docstr_all( + "indices", + r""" +indices() -> Tensor + +Return the indices tensor of a :ref:`sparse COO tensor `. + +.. warning:: + Throws an error if :attr:`self` is not a sparse COO tensor. + +See also :meth:`Tensor.values`. + +.. note:: + This method can only be called on a coalesced sparse tensor. See + :meth:`Tensor.coalesce` for details. +""", +) + +add_docstr_all( + "get_device", + r""" +get_device() -> Device ordinal (Integer) + +For CUDA tensors, this function returns the device ordinal of the GPU on which the tensor resides. +For CPU tensors, this function returns `-1`. + +Example:: + + >>> x = torch.randn(3, 4, 5, device='cuda:0') + >>> x.get_device() + 0 + >>> x.cpu().get_device() + -1 +""", +) + +add_docstr_all( + "values", + r""" +values() -> Tensor + +Return the values tensor of a :ref:`sparse COO tensor `. + +.. warning:: + Throws an error if :attr:`self` is not a sparse COO tensor. + +See also :meth:`Tensor.indices`. + +.. note:: + This method can only be called on a coalesced sparse tensor. See + :meth:`Tensor.coalesce` for details. +""", +) + +add_docstr_all( + "gt", + r""" +gt(other) -> Tensor + +See :func:`torch.gt`. +""", +) + +add_docstr_all( + "gt_", + r""" +gt_(other) -> Tensor + +In-place version of :meth:`~Tensor.gt`. +""", +) + +add_docstr_all( + "greater", + r""" +greater(other) -> Tensor + +See :func:`torch.greater`. +""", +) + +add_docstr_all( + "greater_", + r""" +greater_(other) -> Tensor + +In-place version of :meth:`~Tensor.greater`. +""", +) + +add_docstr_all( + "has_names", + r""" +Is ``True`` if any of this tensor's dimensions are named. Otherwise, is ``False``. +""", +) + +add_docstr_all( + "hardshrink", + r""" +hardshrink(lambd=0.5) -> Tensor + +See :func:`torch.nn.functional.hardshrink` +""", +) + +add_docstr_all( + "heaviside", + r""" +heaviside(values) -> Tensor + +See :func:`torch.heaviside` +""", +) + +add_docstr_all( + "heaviside_", + r""" +heaviside_(values) -> Tensor + +In-place version of :meth:`~Tensor.heaviside` +""", +) + +add_docstr_all( + "histc", + r""" +histc(bins=100, min=0, max=0) -> Tensor + +See :func:`torch.histc` +""", +) + +add_docstr_all( + "histogram", + r""" +histogram(input, bins, *, range=None, weight=None, density=False) -> (Tensor, Tensor) + +See :func:`torch.histogram` +""", +) + +add_docstr_all( + "index_add_", + r""" +index_add_(dim, index, source, *, alpha=1) -> Tensor + +Accumulate the elements of :attr:`alpha` times ``source`` into the :attr:`self` +tensor by adding to the indices in the order given in :attr:`index`. For example, +if ``dim == 0``, ``index[i] == j``, and ``alpha=-1``, then the ``i``\ th row of +``source`` is subtracted from the ``j``\ th row of :attr:`self`. + +The :attr:`dim`\ th dimension of ``source`` must have the same size as the +length of :attr:`index` (which must be a vector), and all other dimensions must +match :attr:`self`, or an error will be raised. + +For a 3-D tensor the output is given as:: + + self[index[i], :, :] += alpha * src[i, :, :] # if dim == 0 + self[:, index[i], :] += alpha * src[:, i, :] # if dim == 1 + self[:, :, index[i]] += alpha * src[:, :, i] # if dim == 2 + +Note: + {forward_reproducibility_note} + +Args: + dim (int): dimension along which to index + index (Tensor): indices of ``source`` to select from, + should have dtype either `torch.int64` or `torch.int32` + source (Tensor): the tensor containing values to add + +Keyword args: + alpha (Number): the scalar multiplier for ``source`` + +Example:: + + >>> x = torch.ones(5, 3) + >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) + >>> index = torch.tensor([0, 4, 2]) + >>> x.index_add_(0, index, t) + tensor([[ 2., 3., 4.], + [ 1., 1., 1.], + [ 8., 9., 10.], + [ 1., 1., 1.], + [ 5., 6., 7.]]) + >>> x.index_add_(0, index, t, alpha=-1) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.]]) +""".format(**reproducibility_notes), +) + +add_docstr_all( + "index_copy_", + r""" +index_copy_(dim, index, tensor) -> Tensor + +Copies the elements of :attr:`tensor` into the :attr:`self` tensor by selecting +the indices in the order given in :attr:`index`. For example, if ``dim == 0`` +and ``index[i] == j``, then the ``i``\ th row of :attr:`tensor` is copied to the +``j``\ th row of :attr:`self`. + +The :attr:`dim`\ th dimension of :attr:`tensor` must have the same size as the +length of :attr:`index` (which must be a vector), and all other dimensions must +match :attr:`self`, or an error will be raised. + +.. note:: + If :attr:`index` contains duplicate entries, multiple elements from + :attr:`tensor` will be copied to the same index of :attr:`self`. The result + is nondeterministic since it depends on which copy occurs last. + +Args: + dim (int): dimension along which to index + index (LongTensor): indices of :attr:`tensor` to select from + tensor (Tensor): the tensor containing values to copy + +Example:: + + >>> x = torch.zeros(5, 3) + >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) + >>> index = torch.tensor([0, 4, 2]) + >>> x.index_copy_(0, index, t) + tensor([[ 1., 2., 3.], + [ 0., 0., 0.], + [ 7., 8., 9.], + [ 0., 0., 0.], + [ 4., 5., 6.]]) +""", +) + +add_docstr_all( + "index_fill_", + r""" +index_fill_(dim, index, value) -> Tensor + +Fills the elements of the :attr:`self` tensor with value :attr:`value` by +selecting the indices in the order given in :attr:`index`. + +Args: + dim (int): dimension along which to index + index (LongTensor): indices of :attr:`self` tensor to fill in + value (float): the value to fill with + +Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float) + >>> index = torch.tensor([0, 2]) + >>> x.index_fill_(1, index, -1) + tensor([[-1., 2., -1.], + [-1., 5., -1.], + [-1., 8., -1.]]) +""", +) + +add_docstr_all( + "index_put_", + r""" +index_put_(indices, values, accumulate=False) -> Tensor + +Puts values from the tensor :attr:`values` into the tensor :attr:`self` using +the indices specified in :attr:`indices` (which is a tuple of Tensors). The +expression ``tensor.index_put_(indices, values)`` is equivalent to +``tensor[indices] = values``. Returns :attr:`self`. + +If :attr:`accumulate` is ``True``, the elements in :attr:`values` are added to +:attr:`self`. If accumulate is ``False``, the behavior is undefined if indices +contain duplicate elements. + +Args: + indices (tuple of LongTensor): tensors used to index into `self`. + values (Tensor): tensor of same dtype as `self`. + accumulate (bool): whether to accumulate into self +""", +) + +add_docstr_all( + "index_put", + r""" +index_put(indices, values, accumulate=False) -> Tensor + +Out-place version of :meth:`~Tensor.index_put_`. +""", +) + +add_docstr_all( + "index_reduce_", + r""" +index_reduce_(dim, index, source, reduce, *, include_self=True) -> Tensor + +Accumulate the elements of ``source`` into the :attr:`self` +tensor by accumulating to the indices in the order given in :attr:`index` +using the reduction given by the ``reduce`` argument. For example, if ``dim == 0``, +``index[i] == j``, ``reduce == prod`` and ``include_self == True`` then the ``i``\ th +row of ``source`` is multiplied by the ``j``\ th row of :attr:`self`. If +:obj:`include_self="True"`, the values in the :attr:`self` tensor are included +in the reduction, otherwise, rows in the :attr:`self` tensor that are accumulated +to are treated as if they were filled with the reduction identities. + +The :attr:`dim`\ th dimension of ``source`` must have the same size as the +length of :attr:`index` (which must be a vector), and all other dimensions must +match :attr:`self`, or an error will be raised. + +For a 3-D tensor with :obj:`reduce="prod"` and :obj:`include_self=True` the +output is given as:: + + self[index[i], :, :] *= src[i, :, :] # if dim == 0 + self[:, index[i], :] *= src[:, i, :] # if dim == 1 + self[:, :, index[i]] *= src[:, :, i] # if dim == 2 + +Note: + {forward_reproducibility_note} + +.. note:: + + This function only supports floating point tensors. + +.. warning:: + + This function is in beta and may change in the near future. + +Args: + dim (int): dimension along which to index + index (Tensor): indices of ``source`` to select from, + should have dtype either `torch.int64` or `torch.int32` + source (FloatTensor): the tensor containing values to accumulate + reduce (str): the reduction operation to apply + (:obj:`"prod"`, :obj:`"mean"`, :obj:`"amax"`, :obj:`"amin"`) + +Keyword args: + include_self (bool): whether the elements from the ``self`` tensor are + included in the reduction + +Example:: + + >>> x = torch.empty(5, 3).fill_(2) + >>> t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=torch.float) + >>> index = torch.tensor([0, 4, 2, 0]) + >>> x.index_reduce_(0, index, t, 'prod') + tensor([[20., 44., 72.], + [ 2., 2., 2.], + [14., 16., 18.], + [ 2., 2., 2.], + [ 8., 10., 12.]]) + >>> x = torch.empty(5, 3).fill_(2) + >>> x.index_reduce_(0, index, t, 'prod', include_self=False) + tensor([[10., 22., 36.], + [ 2., 2., 2.], + [ 7., 8., 9.], + [ 2., 2., 2.], + [ 4., 5., 6.]]) +""".format(**reproducibility_notes), +) + +add_docstr_all( + "index_select", + r""" +index_select(dim, index) -> Tensor + +See :func:`torch.index_select` +""", +) + +add_docstr_all( + "sparse_mask", + r""" +sparse_mask(mask) -> Tensor + +Returns a new :ref:`sparse tensor ` with values from a +strided tensor :attr:`self` filtered by the indices of the sparse +tensor :attr:`mask`. The values of :attr:`mask` sparse tensor are +ignored. :attr:`self` and :attr:`mask` tensors must have the same +shape. + +.. note:: + + The returned sparse tensor might contain duplicate values if :attr:`mask` + is not coalesced. It is therefore advisable to pass ``mask.coalesce()`` + if such behavior is not desired. + +.. note:: + + The returned sparse tensor has the same indices as the sparse tensor + :attr:`mask`, even when the corresponding values in :attr:`self` are + zeros. + +Args: + mask (Tensor): a sparse tensor whose indices are used as a filter + +Example:: + + >>> nse = 5 + >>> dims = (5, 5, 2, 2) + >>> I = torch.cat([torch.randint(0, dims[0], size=(nse,)), + ... torch.randint(0, dims[1], size=(nse,))], 0).reshape(2, nse) + >>> V = torch.randn(nse, dims[2], dims[3]) + >>> S = torch.sparse_coo_tensor(I, V, dims).coalesce() + >>> D = torch.randn(dims) + >>> D.sparse_mask(S) + tensor(indices=tensor([[0, 0, 0, 2], + [0, 1, 4, 3]]), + values=tensor([[[ 1.6550, 0.2397], + [-0.1611, -0.0779]], + + [[ 0.2326, -1.0558], + [ 1.4711, 1.9678]], + + [[-0.5138, -0.0411], + [ 1.9417, 0.5158]], + + [[ 0.0793, 0.0036], + [-0.2569, -0.1055]]]), + size=(5, 5, 2, 2), nnz=4, layout=torch.sparse_coo) +""", +) + +add_docstr_all( + "inverse", + r""" +inverse() -> Tensor + +See :func:`torch.inverse` +""", +) + +add_docstr_all( + "isnan", + r""" +isnan() -> Tensor + +See :func:`torch.isnan` +""", +) + +add_docstr_all( + "isinf", + r""" +isinf() -> Tensor + +See :func:`torch.isinf` +""", +) + +add_docstr_all( + "isposinf", + r""" +isposinf() -> Tensor + +See :func:`torch.isposinf` +""", +) + +add_docstr_all( + "isneginf", + r""" +isneginf() -> Tensor + +See :func:`torch.isneginf` +""", +) + +add_docstr_all( + "isfinite", + r""" +isfinite() -> Tensor + +See :func:`torch.isfinite` +""", +) + +add_docstr_all( + "isclose", + r""" +isclose(other, rtol=1e-05, atol=1e-08, equal_nan=False) -> Tensor + +See :func:`torch.isclose` +""", +) + +add_docstr_all( + "isreal", + r""" +isreal() -> Tensor + +See :func:`torch.isreal` +""", +) + +add_docstr_all( + "is_coalesced", + r""" +is_coalesced() -> bool + +Returns ``True`` if :attr:`self` is a :ref:`sparse COO tensor +` that is coalesced, ``False`` otherwise. + +.. warning:: + Throws an error if :attr:`self` is not a sparse COO tensor. + +See :meth:`coalesce` and :ref:`uncoalesced tensors `. +""", +) + +add_docstr_all( + "is_contiguous", + r""" +is_contiguous(memory_format=torch.contiguous_format) -> bool + +Returns True if :attr:`self` tensor is contiguous in memory in the order specified +by memory format. + +Args: + memory_format (:class:`torch.memory_format`, optional): Specifies memory allocation + order. Default: ``torch.contiguous_format``. +""", +) + +add_docstr_all( + "is_pinned", + r""" +Returns true if this tensor resides in pinned memory. +By default, the device pinned memory on will be the current :ref:`accelerator`. +""", +) + +add_docstr_all( + "is_floating_point", + r""" +is_floating_point() -> bool + +Returns True if the data type of :attr:`self` is a floating point data type. +""", +) + +add_docstr_all( + "is_complex", + r""" +is_complex() -> bool + +Returns True if the data type of :attr:`self` is a complex data type. +""", +) + +add_docstr_all( + "is_inference", + r""" +is_inference() -> bool + +See :func:`torch.is_inference` +""", +) + +add_docstr_all( + "is_conj", + r""" +is_conj() -> bool + +Returns True if the conjugate bit of :attr:`self` is set to true. +""", +) + +add_docstr_all( + "is_neg", + r""" +is_neg() -> bool + +Returns True if the negative bit of :attr:`self` is set to true. +""", +) + +add_docstr_all( + "is_signed", + r""" +is_signed() -> bool + +Returns True if the data type of :attr:`self` is a signed data type. +""", +) + +add_docstr_all( + "is_set_to", + r""" +is_set_to(tensor) -> bool + +Returns True if both tensors are pointing to the exact same memory (same +storage, offset, size and stride). +""", +) + +add_docstr_all( + "item", + r""" +item() -> number + +Returns the value of this tensor as a standard Python number. This only works +for tensors with one element. For other cases, see :meth:`~Tensor.tolist`. + +This operation is not differentiable. + +Example:: + + >>> x = torch.tensor([1.0]) + >>> x.item() + 1.0 + +""", +) + +add_docstr_all( + "kron", + r""" +kron(other) -> Tensor + +See :func:`torch.kron` +""", +) + +add_docstr_all( + "kthvalue", + r""" +kthvalue(k, dim=None, keepdim=False) -> (Tensor, LongTensor) + +See :func:`torch.kthvalue` +""", +) + +add_docstr_all( + "ldexp", + r""" +ldexp(other) -> Tensor + +See :func:`torch.ldexp` +""", +) + +add_docstr_all( + "ldexp_", + r""" +ldexp_(other) -> Tensor + +In-place version of :meth:`~Tensor.ldexp` +""", +) + +add_docstr_all( + "lcm", + r""" +lcm(other) -> Tensor + +See :func:`torch.lcm` +""", +) + +add_docstr_all( + "lcm_", + r""" +lcm_(other) -> Tensor + +In-place version of :meth:`~Tensor.lcm` +""", +) + +add_docstr_all( + "le", + r""" +le(other) -> Tensor + +See :func:`torch.le`. +""", +) + +add_docstr_all( + "le_", + r""" +le_(other) -> Tensor + +In-place version of :meth:`~Tensor.le`. +""", +) + +add_docstr_all( + "less_equal", + r""" +less_equal(other) -> Tensor + +See :func:`torch.less_equal`. +""", +) + +add_docstr_all( + "less_equal_", + r""" +less_equal_(other) -> Tensor + +In-place version of :meth:`~Tensor.less_equal`. +""", +) + +add_docstr_all( + "lerp", + r""" +lerp(end, weight) -> Tensor + +See :func:`torch.lerp` +""", +) + +add_docstr_all( + "lerp_", + r""" +lerp_(end, weight) -> Tensor + +In-place version of :meth:`~Tensor.lerp` +""", +) + +add_docstr_all( + "lgamma", + r""" +lgamma() -> Tensor + +See :func:`torch.lgamma` +""", +) + +add_docstr_all( + "lgamma_", + r""" +lgamma_() -> Tensor + +In-place version of :meth:`~Tensor.lgamma` +""", +) + +add_docstr_all( + "log", + r""" +log() -> Tensor + +See :func:`torch.log` +""", +) + +add_docstr_all( + "log_", + r""" +log_() -> Tensor + +In-place version of :meth:`~Tensor.log` +""", +) + +add_docstr_all( + "log10", + r""" +log10() -> Tensor + +See :func:`torch.log10` +""", +) + +add_docstr_all( + "log10_", + r""" +log10_() -> Tensor + +In-place version of :meth:`~Tensor.log10` +""", +) + +add_docstr_all( + "log1p", + r""" +log1p() -> Tensor + +See :func:`torch.log1p` +""", +) + +add_docstr_all( + "log1p_", + r""" +log1p_() -> Tensor + +In-place version of :meth:`~Tensor.log1p` +""", +) + +add_docstr_all( + "log2", + r""" +log2() -> Tensor + +See :func:`torch.log2` +""", +) + +add_docstr_all( + "log2_", + r""" +log2_() -> Tensor + +In-place version of :meth:`~Tensor.log2` +""", +) + +add_docstr_all( + "logaddexp", + r""" +logaddexp(other) -> Tensor + +See :func:`torch.logaddexp` +""", +) + +add_docstr_all( + "logaddexp2", + r""" +logaddexp2(other) -> Tensor + +See :func:`torch.logaddexp2` +""", +) + +add_docstr_all( + "log_normal_", + r""" +log_normal_(mean=1, std=2, *, generator=None) + +Fills :attr:`self` tensor with numbers samples from the log-normal distribution +parameterized by the given mean :math:`\mu` and standard deviation +:math:`\sigma`. Note that :attr:`mean` and :attr:`std` are the mean and +standard deviation of the underlying normal distribution, and not of the +returned distribution: + +.. math:: + + f(x) = \dfrac{1}{x \sigma \sqrt{2\pi}}\ e^{-\frac{(\ln x - \mu)^2}{2\sigma^2}} +""", +) + +add_docstr_all( + "logsumexp", + r""" +logsumexp(dim, keepdim=False) -> Tensor + +See :func:`torch.logsumexp` +""", +) + +add_docstr_all( + "lt", + r""" +lt(other) -> Tensor + +See :func:`torch.lt`. +""", +) + +add_docstr_all( + "lt_", + r""" +lt_(other) -> Tensor + +In-place version of :meth:`~Tensor.lt`. +""", +) + +add_docstr_all( + "less", + r""" +lt(other) -> Tensor + +See :func:`torch.less`. +""", +) + +add_docstr_all( + "less_", + r""" +less_(other) -> Tensor + +In-place version of :meth:`~Tensor.less`. +""", +) + +add_docstr_all( + "lu_solve", + r""" +lu_solve(LU_data, LU_pivots) -> Tensor + +See :func:`torch.lu_solve` +""", +) + +add_docstr_all( + "map_", + r""" +map_(tensor, callable) + +Applies :attr:`callable` for each element in :attr:`self` tensor and the given +:attr:`tensor` and stores the results in :attr:`self` tensor. :attr:`self` tensor and +the given :attr:`tensor` must be :ref:`broadcastable `. + +The :attr:`callable` should have the signature:: + + def callable(a, b) -> number +""", +) + +add_docstr_all( + "masked_scatter_", + r""" +masked_scatter_(mask, source) + +Copies elements from :attr:`source` into :attr:`self` tensor at positions where +the :attr:`mask` is True. Elements from :attr:`source` are copied into :attr:`self` +starting at position 0 of :attr:`source` and continuing in order one-by-one for each +occurrence of :attr:`mask` being True. +The shape of :attr:`mask` must be :ref:`broadcastable ` +with the shape of the underlying tensor. The :attr:`source` should have at least +as many elements as the number of ones in :attr:`mask`. + +Args: + mask (BoolTensor): the boolean mask + source (Tensor): the tensor to copy from + +.. note:: + + The :attr:`mask` operates on the :attr:`self` tensor, not on the given + :attr:`source` tensor. + +Example: + + >>> self = torch.tensor([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]) + >>> mask = torch.tensor( + ... [[0, 0, 0, 1, 1], [1, 1, 0, 1, 1]], + ... dtype=torch.bool, + ... ) + >>> source = torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]) + >>> self.masked_scatter_(mask, source) + tensor([[0, 0, 0, 0, 1], + [2, 3, 0, 4, 5]]) + +""", +) + +add_docstr_all( + "masked_fill_", + r""" +masked_fill_(mask, value) + +Fills elements of :attr:`self` tensor with :attr:`value` where :attr:`mask` is +True. The shape of :attr:`mask` must be +:ref:`broadcastable ` with the shape of the underlying +tensor. + +Args: + mask (BoolTensor): the boolean mask + value (float): the value to fill in with +""", +) + +add_docstr_all( + "masked_select", + r""" +masked_select(mask) -> Tensor + +See :func:`torch.masked_select` +""", +) + +add_docstr_all( + "matrix_power", + r""" +matrix_power(n) -> Tensor + +.. note:: :meth:`~Tensor.matrix_power` is deprecated, use :func:`torch.linalg.matrix_power` instead. + +Alias for :func:`torch.linalg.matrix_power` +""", +) + +add_docstr_all( + "matrix_exp", + r""" +matrix_exp() -> Tensor + +See :func:`torch.matrix_exp` +""", +) + +add_docstr_all( + "max", + r""" +max(dim=None, keepdim=False) -> Tensor or (Tensor, Tensor) + +See :func:`torch.max` +""", +) + +add_docstr_all( + "amax", + r""" +amax(dim=None, keepdim=False) -> Tensor + +See :func:`torch.amax` +""", +) + +add_docstr_all( + "maximum", + r""" +maximum(other) -> Tensor + +See :func:`torch.maximum` +""", +) + +add_docstr_all( + "fmax", + r""" +fmax(other) -> Tensor + +See :func:`torch.fmax` +""", +) + +add_docstr_all( + "argmax", + r""" +argmax(dim=None, keepdim=False) -> LongTensor + +See :func:`torch.argmax` +""", +) + +add_docstr_all( + "argwhere", + r""" +argwhere() -> Tensor + +See :func:`torch.argwhere` +""", +) + +add_docstr_all( + "mean", + r""" +mean(dim=None, keepdim=False, *, dtype=None) -> Tensor + +See :func:`torch.mean` +""", +) + +add_docstr_all( + "nanmean", + r""" +nanmean(dim=None, keepdim=False, *, dtype=None) -> Tensor + +See :func:`torch.nanmean` +""", +) + +add_docstr_all( + "median", + r""" +median(dim=None, keepdim=False) -> (Tensor, LongTensor) + +See :func:`torch.median` +""", +) + +add_docstr_all( + "nanmedian", + r""" +nanmedian(dim=None, keepdim=False) -> (Tensor, LongTensor) + +See :func:`torch.nanmedian` +""", +) + +add_docstr_all( + "min", + r""" +min(dim=None, keepdim=False) -> Tensor or (Tensor, Tensor) + +See :func:`torch.min` +""", +) + +add_docstr_all( + "amin", + r""" +amin(dim=None, keepdim=False) -> Tensor + +See :func:`torch.amin` +""", +) + +add_docstr_all( + "minimum", + r""" +minimum(other) -> Tensor + +See :func:`torch.minimum` +""", +) + +add_docstr_all( + "aminmax", + r""" +aminmax(*, dim=None, keepdim=False) -> (Tensor min, Tensor max) + +See :func:`torch.aminmax` +""", +) + +add_docstr_all( + "fmin", + r""" +fmin(other) -> Tensor + +See :func:`torch.fmin` +""", +) + +add_docstr_all( + "argmin", + r""" +argmin(dim=None, keepdim=False) -> LongTensor + +See :func:`torch.argmin` +""", +) + +add_docstr_all( + "mm", + r""" +mm(mat2) -> Tensor + +See :func:`torch.mm` +""", +) + +add_docstr_all( + "mode", + r""" +mode(dim=None, keepdim=False) -> (Tensor, LongTensor) + +See :func:`torch.mode` +""", +) + +add_docstr_all( + "movedim", + r""" +movedim(source, destination) -> Tensor + +See :func:`torch.movedim` +""", +) + +add_docstr_all( + "moveaxis", + r""" +moveaxis(source, destination) -> Tensor + +See :func:`torch.moveaxis` +""", +) + +add_docstr_all( + "mul", + r""" +mul(value) -> Tensor + +See :func:`torch.mul`. +""", +) + +add_docstr_all( + "mul_", + r""" +mul_(value) -> Tensor + +In-place version of :meth:`~Tensor.mul`. +""", +) + +add_docstr_all( + "multiply", + r""" +multiply(value) -> Tensor + +See :func:`torch.multiply`. +""", +) + +add_docstr_all( + "multiply_", + r""" +multiply_(value) -> Tensor + +In-place version of :meth:`~Tensor.multiply`. +""", +) + +add_docstr_all( + "multinomial", + r""" +multinomial(num_samples, replacement=False, *, generator=None) -> Tensor + +See :func:`torch.multinomial` +""", +) + +add_docstr_all( + "mv", + r""" +mv(vec) -> Tensor + +See :func:`torch.mv` +""", +) + +add_docstr_all( + "mvlgamma", + r""" +mvlgamma(p) -> Tensor + +See :func:`torch.mvlgamma` +""", +) + +add_docstr_all( + "mvlgamma_", + r""" +mvlgamma_(p) -> Tensor + +In-place version of :meth:`~Tensor.mvlgamma` +""", +) + +add_docstr_all( + "narrow", + r""" +narrow(dimension, start, length) -> Tensor + +See :func:`torch.narrow`. +""", +) + +add_docstr_all( + "narrow_copy", + r""" +narrow_copy(dimension, start, length) -> Tensor + +See :func:`torch.narrow_copy`. +""", +) + +add_docstr_all( + "ndimension", + r""" +ndimension() -> int + +Alias for :meth:`~Tensor.dim()` +""", +) + +add_docstr_all( + "nan_to_num", + r""" +nan_to_num(nan=0.0, posinf=None, neginf=None) -> Tensor + +See :func:`torch.nan_to_num`. +""", +) + +add_docstr_all( + "nan_to_num_", + r""" +nan_to_num_(nan=0.0, posinf=None, neginf=None) -> Tensor + +In-place version of :meth:`~Tensor.nan_to_num`. +""", +) + +add_docstr_all( + "ne", + r""" +ne(other) -> Tensor + +See :func:`torch.ne`. +""", +) + +add_docstr_all( + "ne_", + r""" +ne_(other) -> Tensor + +In-place version of :meth:`~Tensor.ne`. +""", +) + +add_docstr_all( + "not_equal", + r""" +not_equal(other) -> Tensor + +See :func:`torch.not_equal`. +""", +) + +add_docstr_all( + "not_equal_", + r""" +not_equal_(other) -> Tensor + +In-place version of :meth:`~Tensor.not_equal`. +""", +) + +add_docstr_all( + "neg", + r""" +neg() -> Tensor + +See :func:`torch.neg` +""", +) + +add_docstr_all( + "negative", + r""" +negative() -> Tensor + +See :func:`torch.negative` +""", +) + +add_docstr_all( + "neg_", + r""" +neg_() -> Tensor + +In-place version of :meth:`~Tensor.neg` +""", +) + +add_docstr_all( + "negative_", + r""" +negative_() -> Tensor + +In-place version of :meth:`~Tensor.negative` +""", +) + +add_docstr_all( + "nelement", + r""" +nelement() -> int + +Alias for :meth:`~Tensor.numel` +""", +) + +add_docstr_all( + "nextafter", + r""" +nextafter(other) -> Tensor +See :func:`torch.nextafter` +""", +) + +add_docstr_all( + "nextafter_", + r""" +nextafter_(other) -> Tensor +In-place version of :meth:`~Tensor.nextafter` +""", +) + +add_docstr_all( + "nonzero", + r""" +nonzero() -> LongTensor + +See :func:`torch.nonzero` +""", +) + +add_docstr_all( + "nonzero_static", + r""" +nonzero_static(input, *, size, fill_value=-1) -> Tensor + +Returns a 2-D tensor where each row is the index for a non-zero value. +The returned Tensor has the same `torch.dtype` as `torch.nonzero()`. + +Args: + input (Tensor): the input tensor to count non-zero elements. + +Keyword args: + size (int): the size of non-zero elements expected to be included in the out + tensor. Pad the out tensor with `fill_value` if the `size` is larger + than total number of non-zero elements, truncate out tensor if `size` + is smaller. The size must be a non-negative integer. + fill_value (int, optional): the value to fill the output tensor with when `size` is larger + than the total number of non-zero elements. Default is `-1` to represent + invalid index. + +Example: + + # Example 1: Padding + >>> input_tensor = torch.tensor([[1, 0], [3, 2]]) + >>> static_size = 4 + >>> t = torch.nonzero_static(input_tensor, size=static_size) + tensor([[ 0, 0], + [ 1, 0], + [ 1, 1], + [ -1, -1]], dtype=torch.int64) + + # Example 2: Truncating + >>> input_tensor = torch.tensor([[1, 0], [3, 2]]) + >>> static_size = 2 + >>> t = torch.nonzero_static(input_tensor, size=static_size) + tensor([[ 0, 0], + [ 1, 0]], dtype=torch.int64) + + # Example 3: 0 size + >>> input_tensor = torch.tensor([10]) + >>> static_size = 0 + >>> t = torch.nonzero_static(input_tensor, size=static_size) + tensor([], size=(0, 1), dtype=torch.int64) + + # Example 4: 0 rank input + >>> input_tensor = torch.tensor(10) + >>> static_size = 2 + >>> t = torch.nonzero_static(input_tensor, size=static_size) + tensor([], size=(2, 0), dtype=torch.int64) +""", +) + +add_docstr_all( + "norm", + r""" +norm(p=2, dim=None, keepdim=False) -> Tensor + +See :func:`torch.norm` +""", +) + +add_docstr_all( + "normal_", + r""" +normal_(mean=0, std=1, *, generator=None) -> Tensor + +Fills :attr:`self` tensor with elements samples from the normal distribution +parameterized by :attr:`mean` and :attr:`std`. +""", +) + +add_docstr_all( + "numel", + r""" +numel() -> int + +See :func:`torch.numel` +""", +) + +add_docstr_all( + "numpy", + r""" +numpy(*, force=False) -> numpy.ndarray + +Returns the tensor as a NumPy :class:`ndarray`. + +If :attr:`force` is ``False`` (the default), the conversion +is performed only if the tensor is on the CPU, does not require grad, +does not have its conjugate bit set, and is a dtype and layout that +NumPy supports. The returned ndarray and the tensor will share their +storage, so changes to the tensor will be reflected in the ndarray +and vice versa. + +If :attr:`force` is ``True`` this is equivalent to +calling ``t.detach().cpu().resolve_conj().resolve_neg().numpy()``. +If the tensor isn't on the CPU or the conjugate or negative bit is set, +the tensor won't share its storage with the returned ndarray. +Setting :attr:`force` to ``True`` can be a useful shorthand. + +Args: + force (bool): if ``True``, the ndarray may be a copy of the tensor + instead of always sharing memory, defaults to ``False``. +""", +) + +add_docstr_all( + "orgqr", + r""" +orgqr(input2) -> Tensor + +See :func:`torch.orgqr` +""", +) + +add_docstr_all( + "ormqr", + r""" +ormqr(input2, input3, left=True, transpose=False) -> Tensor + +See :func:`torch.ormqr` +""", +) + +add_docstr_all( + "permute", + r""" +permute(*dims) -> Tensor + +See :func:`torch.permute` +""", +) + +add_docstr_all( + "polygamma", + r""" +polygamma(n) -> Tensor + +See :func:`torch.polygamma` +""", +) + +add_docstr_all( + "polygamma_", + r""" +polygamma_(n) -> Tensor + +In-place version of :meth:`~Tensor.polygamma` +""", +) + +add_docstr_all( + "positive", + r""" +positive() -> Tensor + +See :func:`torch.positive` +""", +) + +add_docstr_all( + "pow", + r""" +pow(exponent) -> Tensor + +See :func:`torch.pow` +""", +) + +add_docstr_all( + "pow_", + r""" +pow_(exponent) -> Tensor + +In-place version of :meth:`~Tensor.pow` +""", +) + +add_docstr_all( + "float_power", + r""" +float_power(exponent) -> Tensor + +See :func:`torch.float_power` +""", +) + +add_docstr_all( + "float_power_", + r""" +float_power_(exponent) -> Tensor + +In-place version of :meth:`~Tensor.float_power` +""", +) + +add_docstr_all( + "prod", + r""" +prod(dim=None, keepdim=False, dtype=None) -> Tensor + +See :func:`torch.prod` +""", +) + +add_docstr_all( + "put_", + r""" +put_(index, source, accumulate=False) -> Tensor + +Copies the elements from :attr:`source` into the positions specified by +:attr:`index`. For the purpose of indexing, the :attr:`self` tensor is treated as if +it were a 1-D tensor. + +:attr:`index` and :attr:`source` need to have the same number of elements, but not necessarily +the same shape. + +If :attr:`accumulate` is ``True``, the elements in :attr:`source` are added to +:attr:`self`. If accumulate is ``False``, the behavior is undefined if :attr:`index` +contain duplicate elements. + +Args: + index (LongTensor): the indices into self + source (Tensor): the tensor containing values to copy from + accumulate (bool, optional): whether to accumulate into self. Default: ``False`` + +Example:: + + >>> src = torch.tensor([[4, 3, 5], + ... [6, 7, 8]]) + >>> src.put_(torch.tensor([1, 3]), torch.tensor([9, 10])) + tensor([[ 4, 9, 5], + [ 10, 7, 8]]) +""", +) + +add_docstr_all( + "put", + r""" +put(input, index, source, accumulate=False) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.put_`. +`input` corresponds to `self` in :meth:`torch.Tensor.put_`. +""", +) + +add_docstr_all( + "qr", + r""" +qr(some=True) -> (Tensor, Tensor) + +See :func:`torch.qr` +""", +) + +add_docstr_all( + "qscheme", + r""" +qscheme() -> torch.qscheme + +Returns the quantization scheme of a given QTensor. +""", +) + +add_docstr_all( + "quantile", + r""" +quantile(q, dim=None, keepdim=False, *, interpolation='linear') -> Tensor + +See :func:`torch.quantile` +""", +) + +add_docstr_all( + "nanquantile", + r""" +nanquantile(q, dim=None, keepdim=False, *, interpolation='linear') -> Tensor + +See :func:`torch.nanquantile` +""", +) + +add_docstr_all( + "q_scale", + r""" +q_scale() -> float + +Given a Tensor quantized by linear(affine) quantization, +returns the scale of the underlying quantizer(). +""", +) + +add_docstr_all( + "q_zero_point", + r""" +q_zero_point() -> int + +Given a Tensor quantized by linear(affine) quantization, +returns the zero_point of the underlying quantizer(). +""", +) + +add_docstr_all( + "q_per_channel_scales", + r""" +q_per_channel_scales() -> Tensor + +Given a Tensor quantized by linear (affine) per-channel quantization, +returns a Tensor of scales of the underlying quantizer. It has the number of +elements that matches the corresponding dimensions (from q_per_channel_axis) of +the tensor. +""", +) + +add_docstr_all( + "q_per_channel_zero_points", + r""" +q_per_channel_zero_points() -> Tensor + +Given a Tensor quantized by linear (affine) per-channel quantization, +returns a tensor of zero_points of the underlying quantizer. It has the number of +elements that matches the corresponding dimensions (from q_per_channel_axis) of +the tensor. +""", +) + +add_docstr_all( + "q_per_channel_axis", + r""" +q_per_channel_axis() -> int + +Given a Tensor quantized by linear (affine) per-channel quantization, +returns the index of dimension on which per-channel quantization is applied. +""", +) + +add_docstr_all( + "random_", + r""" +random_(from=0, to=None, *, generator=None) -> Tensor + +Fills :attr:`self` tensor with numbers sampled from the discrete uniform +distribution over ``[from, to - 1]``. If not specified, the values are usually +only bounded by :attr:`self` tensor's data type. However, for floating point +types, if unspecified, range will be ``[0, 2^mantissa]`` to ensure that every +value is representable. For example, `torch.tensor(1, dtype=torch.double).random_()` +will be uniform in ``[0, 2^53]``. +""", +) + +add_docstr_all( + "rad2deg", + r""" +rad2deg() -> Tensor + +See :func:`torch.rad2deg` +""", +) + +add_docstr_all( + "rad2deg_", + r""" +rad2deg_() -> Tensor + +In-place version of :meth:`~Tensor.rad2deg` +""", +) + +add_docstr_all( + "deg2rad", + r""" +deg2rad() -> Tensor + +See :func:`torch.deg2rad` +""", +) + +add_docstr_all( + "deg2rad_", + r""" +deg2rad_() -> Tensor + +In-place version of :meth:`~Tensor.deg2rad` +""", +) + +add_docstr_all( + "ravel", + r""" +ravel() -> Tensor + +see :func:`torch.ravel` +""", +) + +add_docstr_all( + "reciprocal", + r""" +reciprocal() -> Tensor + +See :func:`torch.reciprocal` +""", +) + +add_docstr_all( + "reciprocal_", + r""" +reciprocal_() -> Tensor + +In-place version of :meth:`~Tensor.reciprocal` +""", +) + +add_docstr_all( + "record_stream", + r""" +record_stream(stream) + +Marks the tensor as having been used by this stream. When the tensor +is deallocated, ensure the tensor memory is not reused for another tensor +until all work queued on :attr:`stream` at the time of deallocation is +complete. + +.. note:: + + The caching allocator is aware of only the stream where a tensor was + allocated. Due to the awareness, it already correctly manages the life + cycle of tensors on only one stream. But if a tensor is used on a stream + different from the stream of origin, the allocator might reuse the memory + unexpectedly. Calling this method lets the allocator know which streams + have used the tensor. + +.. warning:: + + This method is most suitable for use cases where you are providing a + function that created a tensor on a side stream, and want users to be able + to make use of the tensor without having to think carefully about stream + safety when making use of them. These safety guarantees come at some + performance and predictability cost (analogous to the tradeoff between GC + and manual memory management), so if you are in a situation where + you manage the full lifetime of your tensors, you may consider instead + manually managing CUDA events so that calling this method is not necessary. + In particular, when you call this method, on later allocations the + allocator will poll the recorded stream to see if all operations have + completed yet; you can potentially race with side stream computation and + non-deterministically reuse or fail to reuse memory for an allocation. + + You can safely use tensors allocated on side streams without + :meth:`~Tensor.record_stream`; you must manually ensure that + any non-creation stream uses of a tensor are synced back to the creation + stream before you deallocate the tensor. As the CUDA caching allocator + guarantees that the memory will only be reused with the same creation stream, + this is sufficient to ensure that writes to future reallocations of the + memory will be delayed until non-creation stream uses are done. + (Counterintuitively, you may observe that on the CPU side we have already + reallocated the tensor, even though CUDA kernels on the old tensor are + still in progress. This is fine, because CUDA operations on the new + tensor will appropriately wait for the old operations to complete, as they + are all on the same stream.) + + Concretely, this looks like this:: + + with torch.cuda.stream(s0): + x = torch.zeros(N) + + s1.wait_stream(s0) + with torch.cuda.stream(s1): + y = some_comm_op(x) + + ... some compute on s0 ... + + # synchronize creation stream s0 to side stream s1 + # before deallocating x + s0.wait_stream(s1) + del x + + Note that some discretion is required when deciding when to perform + ``s0.wait_stream(s1)``. In particular, if we were to wait immediately + after ``some_comm_op``, there wouldn't be any point in having the side + stream; it would be equivalent to have run ``some_comm_op`` on ``s0``. + Instead, the synchronization must be placed at some appropriate, later + point in time where you expect the side stream ``s1`` to have finished + work. This location is typically identified via profiling, e.g., using + Chrome traces produced + :meth:`torch.autograd.profiler.profile.export_chrome_trace`. If you + place the wait too early, work on s0 will block until ``s1`` has finished, + preventing further overlapping of communication and computation. If you + place the wait too late, you will use more memory than is strictly + necessary (as you are keeping ``x`` live for longer.) For a concrete + example of how this guidance can be applied in practice, see this post: + `FSDP and CUDACachingAllocator + `_. +""", +) + +add_docstr_all( + "remainder", + r""" +remainder(divisor) -> Tensor + +See :func:`torch.remainder` +""", +) + +add_docstr_all( + "remainder_", + r""" +remainder_(divisor) -> Tensor + +In-place version of :meth:`~Tensor.remainder` +""", +) + +add_docstr_all( + "renorm", + r""" +renorm(p, dim, maxnorm) -> Tensor + +See :func:`torch.renorm` +""", +) + +add_docstr_all( + "renorm_", + r""" +renorm_(p, dim, maxnorm) -> Tensor + +In-place version of :meth:`~Tensor.renorm` +""", +) + +add_docstr_all( + "repeat", + r""" +repeat(*repeats) -> Tensor + +Repeats this tensor along the specified dimensions. + +Unlike :meth:`~Tensor.expand`, this function copies the tensor's data. + +.. warning:: + + :meth:`~Tensor.repeat` behaves differently from + `numpy.repeat `_, + but is more similar to + `numpy.tile `_. + For the operator similar to `numpy.repeat`, see :func:`torch.repeat_interleave`. + +Args: + repeat (torch.Size, int..., tuple of int or list of int): The number of times to repeat this tensor along each dimension + +Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.repeat(4, 2) + tensor([[ 1, 2, 3, 1, 2, 3], + [ 1, 2, 3, 1, 2, 3], + [ 1, 2, 3, 1, 2, 3], + [ 1, 2, 3, 1, 2, 3]]) + >>> x.repeat(4, 2, 1).size() + torch.Size([4, 2, 3]) +""", +) + +add_docstr_all( + "repeat_interleave", + r""" +repeat_interleave(repeats, dim=None, *, output_size=None) -> Tensor + +See :func:`torch.repeat_interleave`. +""", +) + +add_docstr_all( + "requires_grad_", + r""" +requires_grad_(requires_grad=True) -> Tensor + +Change if autograd should record operations on this tensor: sets this tensor's +:attr:`requires_grad` attribute in-place. Returns this tensor. + +:func:`requires_grad_`'s main use case is to tell autograd to begin recording +operations on a Tensor ``tensor``. If ``tensor`` has ``requires_grad=False`` +(because it was obtained through a DataLoader, or required preprocessing or +initialization), ``tensor.requires_grad_()`` makes it so that autograd will +begin to record operations on ``tensor``. + +Args: + requires_grad (bool): If autograd should record operations on this tensor. + Default: ``True``. + +Example:: + + >>> # Let's say we want to preprocess some saved weights and use + >>> # the result as new weights. + >>> saved_weights = [0.1, 0.2, 0.3, 0.25] + >>> loaded_weights = torch.tensor(saved_weights) + >>> weights = preprocess(loaded_weights) # some function + >>> weights + tensor([-0.5503, 0.4926, -2.1158, -0.8303]) + + >>> # Now, start to record operations done to weights + >>> weights.requires_grad_() + >>> out = weights.pow(2).sum() + >>> out.backward() + >>> weights.grad + tensor([-1.1007, 0.9853, -4.2316, -1.6606]) + +""", +) + +add_docstr_all( + "reshape", + r""" +reshape(*shape) -> Tensor + +Returns a tensor with the same data and number of elements as :attr:`self` +but with the specified shape. This method returns a view if :attr:`shape` is +compatible with the current shape. See :meth:`torch.Tensor.view` on when it is +possible to return a view. + +See :func:`torch.reshape` + +Args: + shape (tuple of ints or int...): the desired shape + +""", +) + +add_docstr_all( + "reshape_as", + r""" +reshape_as(other) -> Tensor + +Returns this tensor as the same shape as :attr:`other`. +``self.reshape_as(other)`` is equivalent to ``self.reshape(other.sizes())``. +This method returns a view if ``other.sizes()`` is compatible with the current +shape. See :meth:`torch.Tensor.view` on when it is possible to return a view. + +Please see :meth:`reshape` for more information about ``reshape``. + +Args: + other (:class:`torch.Tensor`): The result tensor has the same shape + as :attr:`other`. +""", +) + +add_docstr_all( + "resize_", + r""" +resize_(*sizes, memory_format=torch.contiguous_format) -> Tensor + +Resizes :attr:`self` tensor to the specified size. If the number of elements is +larger than the current storage size, then the underlying storage is resized +to fit the new number of elements. If the number of elements is smaller, the +underlying storage is not changed. Existing elements are preserved but any new +memory is uninitialized. + +.. warning:: + + This is a low-level method. The storage is reinterpreted as C-contiguous, + ignoring the current strides (unless the target size equals the current + size, in which case the tensor is left unchanged). For most purposes, you + will instead want to use :meth:`~Tensor.view()`, which checks for + contiguity, or :meth:`~Tensor.reshape()`, which copies data if needed. To + change the size in-place with custom strides, see :meth:`~Tensor.set_()`. + +.. note:: + + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, new elements are initialized to prevent nondeterministic behavior + from using the result as an input to an operation. Floating point and + complex values are set to NaN, and integer values are set to the maximum + value. + +Args: + sizes (torch.Size or int...): the desired size + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + Tensor. Default: ``torch.contiguous_format``. Note that memory format of + :attr:`self` is going to be unaffected if ``self.size()`` matches ``sizes``. + +Example:: + + >>> x = torch.tensor([[1, 2], [3, 4], [5, 6]]) + >>> x.resize_(2, 2) + tensor([[ 1, 2], + [ 3, 4]]) +""", +) + +add_docstr_all( + "resize_as_", + r""" +resize_as_(tensor, memory_format=torch.contiguous_format) -> Tensor + +Resizes the :attr:`self` tensor to be the same size as the specified +:attr:`tensor`. This is equivalent to ``self.resize_(tensor.size())``. + +Args: + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + Tensor. Default: ``torch.contiguous_format``. Note that memory format of + :attr:`self` is going to be unaffected if ``self.size()`` matches ``tensor.size()``. + +""", +) + +add_docstr_all( + "rot90", + r""" +rot90(k, dims) -> Tensor + +See :func:`torch.rot90` +""", +) + +add_docstr_all( + "round", + r""" +round(decimals=0) -> Tensor + +See :func:`torch.round` +""", +) + +add_docstr_all( + "round_", + r""" +round_(decimals=0) -> Tensor + +In-place version of :meth:`~Tensor.round` +""", +) + +add_docstr_all( + "rsqrt", + r""" +rsqrt() -> Tensor + +See :func:`torch.rsqrt` +""", +) + +add_docstr_all( + "rsqrt_", + r""" +rsqrt_() -> Tensor + +In-place version of :meth:`~Tensor.rsqrt` +""", +) + +add_docstr_all( + "scatter_", + r""" +scatter_(dim, index, src, *, reduce=None) -> Tensor + +Writes all values from the tensor :attr:`src` into :attr:`self` at the indices +specified in the :attr:`index` tensor. For each value in :attr:`src`, its output +index is specified by its index in :attr:`src` for ``dimension != dim`` and by +the corresponding value in :attr:`index` for ``dimension = dim``. + +For a 3-D tensor, :attr:`self` is updated as:: + + self[index[i][j][k]][j][k] = src[i][j][k] # if dim == 0 + self[i][index[i][j][k]][k] = src[i][j][k] # if dim == 1 + self[i][j][index[i][j][k]] = src[i][j][k] # if dim == 2 + +This is the reverse operation of the manner described in :meth:`~Tensor.gather`. + +It is also required that +``index.size(d) <= src.size(d)`` for all dimensions ``d``, and that +``index.size(d) <= self.size(d)`` for all dimensions ``d != dim``. +Note that ``input`` and ``index`` do not broadcast against each other for NPUs, +so when running on NPUs, :attr:`input` and :attr:`index` must have the same number of dimensions. +Standard broadcasting occurs in all other cases. + +Moreover, as for :meth:`~Tensor.gather`, the values of :attr:`index` must be +between ``0`` and ``self.size(dim) - 1`` inclusive. + +.. warning:: + + When indices are not unique, the behavior is non-deterministic (one of the + values from ``src`` will be picked arbitrarily) and the gradient will be + incorrect (it will be propagated to all locations in the source that + correspond to the same index)! + +.. note:: + + The backward pass is implemented only for ``src.shape == index.shape``. + +Additionally accepts an optional :attr:`reduce` argument that allows +specification of an optional reduction operation, which is applied to all +values in the tensor :attr:`src` into :attr:`self` at the indices +specified in the :attr:`index`. For each value in :attr:`src`, the reduction +operation is applied to an index in :attr:`self` which is specified by +its index in :attr:`src` for ``dimension != dim`` and by the corresponding +value in :attr:`index` for ``dimension = dim``. + +Given a 3-D tensor and reduction using the multiplication operation, :attr:`self` +is updated as:: + + self[index[i][j][k]][j][k] *= src[i][j][k] # if dim == 0 + self[i][index[i][j][k]][k] *= src[i][j][k] # if dim == 1 + self[i][j][index[i][j][k]] *= src[i][j][k] # if dim == 2 + +Reducing with the addition operation is the same as using +:meth:`~torch.Tensor.scatter_add_`. + +.. warning:: + The reduce argument with Tensor ``src`` is deprecated and will be removed in + a future PyTorch release. Please use :meth:`~torch.Tensor.scatter_reduce_` + instead for more reduction options. + +Args: + dim (int): the axis along which to index + index (LongTensor): the indices of elements to scatter, can be either empty + or of the same dimensionality as ``src``. When empty, the operation + returns ``self`` unchanged. + src (Tensor): the source element(s) to scatter. + +Keyword args: + reduce (str, optional): reduction operation to apply, can be either + ``'add'`` or ``'multiply'``. + +Example:: + + >>> src = torch.arange(1, 11).reshape((2, 5)) + >>> src + tensor([[ 1, 2, 3, 4, 5], + [ 6, 7, 8, 9, 10]]) + >>> index = torch.tensor([[0, 1, 2, 0]]) + >>> torch.zeros(3, 5, dtype=src.dtype).scatter_(0, index, src) + tensor([[1, 0, 0, 4, 0], + [0, 2, 0, 0, 0], + [0, 0, 3, 0, 0]]) + >>> index = torch.tensor([[0, 1, 2], [0, 1, 4]]) + >>> torch.zeros(3, 5, dtype=src.dtype).scatter_(1, index, src) + tensor([[1, 2, 3, 0, 0], + [6, 7, 0, 0, 8], + [0, 0, 0, 0, 0]]) + + >>> torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]), + ... 1.23, reduce='multiply') + tensor([[2.0000, 2.0000, 2.4600, 2.0000], + [2.0000, 2.0000, 2.0000, 2.4600]]) + >>> torch.full((2, 4), 2.).scatter_(1, torch.tensor([[2], [3]]), + ... 1.23, reduce='add') + tensor([[2.0000, 2.0000, 3.2300, 2.0000], + [2.0000, 2.0000, 2.0000, 3.2300]]) + +.. function:: scatter_(dim, index, value, *, reduce=None) -> Tensor: + :noindex: + +Writes the value from :attr:`value` into :attr:`self` at the indices +specified in the :attr:`index` tensor. This operation is equivalent to the previous version, +with the :attr:`src` tensor filled entirely with :attr:`value`. + +Args: + dim (int): the axis along which to index + index (LongTensor): the indices of elements to scatter, can be either empty + or of the same dimensionality as ``src``. When empty, the operation + returns ``self`` unchanged. + value (Scalar): the value to scatter. + +Keyword args: + reduce (str, optional): reduction operation to apply, can be either + ``'add'`` or ``'multiply'``. + +Example:: + + >>> index = torch.tensor([[0, 1]]) + >>> value = 2 + >>> torch.zeros(3, 5).scatter_(0, index, value) + tensor([[2., 0., 0., 0., 0.], + [0., 2., 0., 0., 0.], + [0., 0., 0., 0., 0.]]) +""", +) + +add_docstr_all( + "scatter_add_", + r""" +scatter_add_(dim, index, src) -> Tensor + +Adds all values from the tensor :attr:`src` into :attr:`self` at the indices +specified in the :attr:`index` tensor in a similar fashion as +:meth:`~torch.Tensor.scatter_`. For each value in :attr:`src`, it is added to +an index in :attr:`self` which is specified by its index in :attr:`src` +for ``dimension != dim`` and by the corresponding value in :attr:`index` for +``dimension = dim``. + +For a 3-D tensor, :attr:`self` is updated as:: + + self[index[i][j][k]][j][k] += src[i][j][k] # if dim == 0 + self[i][index[i][j][k]][k] += src[i][j][k] # if dim == 1 + self[i][j][index[i][j][k]] += src[i][j][k] # if dim == 2 + +:attr:`self`, :attr:`index` and :attr:`src` should have same number of +dimensions. It is also required that ``index.size(d) <= src.size(d)`` for all +dimensions ``d``, and that ``index.size(d) <= self.size(d)`` for all dimensions +``d != dim``. Note that ``index`` and ``src`` do not broadcast. +When :attr:`index` is empty, we always return the original tensor +without further error checking. + +Note: + {forward_reproducibility_note} + +.. note:: + + The backward pass is implemented only for ``src.shape == index.shape``. + +Args: + dim (int): the axis along which to index + index (LongTensor): the indices of elements to scatter and add, can be + either empty or of the same dimensionality as ``src``. When empty, the + operation returns ``self`` unchanged. + src (Tensor): the source elements to scatter and add + +Example:: + + >>> src = torch.ones((2, 5)) + >>> index = torch.tensor([[0, 1, 2, 0, 0]]) + >>> torch.zeros(3, 5, dtype=src.dtype).scatter_add_(0, index, src) + tensor([[1., 0., 0., 1., 1.], + [0., 1., 0., 0., 0.], + [0., 0., 1., 0., 0.]]) + >>> index = torch.tensor([[0, 1, 2, 0, 0], [0, 1, 2, 2, 2]]) + >>> torch.zeros(3, 5, dtype=src.dtype).scatter_add_(0, index, src) + tensor([[2., 0., 0., 1., 1.], + [0., 2., 0., 0., 0.], + [0., 0., 2., 1., 1.]]) + +""".format(**reproducibility_notes), +) + +add_docstr_all( + "scatter_reduce_", + r""" +scatter_reduce_(dim, index, src, reduce, *, include_self=True) -> Tensor + +Reduces all values from the :attr:`src` tensor to the indices specified in +the :attr:`index` tensor in the :attr:`self` tensor using the applied reduction +defined via the :attr:`reduce` argument (:obj:`"sum"`, :obj:`"prod"`, :obj:`"mean"`, +:obj:`"amax"`, :obj:`"amin"`). For each value in :attr:`src`, it is reduced to an +index in :attr:`self` which is specified by its index in :attr:`src` for +``dimension != dim`` and by the corresponding value in :attr:`index` for +``dimension = dim``. If :obj:`include_self="True"`, the values in the :attr:`self` +tensor are included in the reduction. + +:attr:`self`, :attr:`index` and :attr:`src` should all have +the same number of dimensions. It is also required that +``index.size(d) <= src.size(d)`` for all dimensions ``d``, and that +``index.size(d) <= self.size(d)`` for all dimensions ``d != dim``. +Note that ``index`` and ``src`` do not broadcast. + +For a 3-D tensor with :obj:`reduce="sum"` and :obj:`include_self=True` the +output is given as:: + + self[index[i][j][k]][j][k] += src[i][j][k] # if dim == 0 + self[i][index[i][j][k]][k] += src[i][j][k] # if dim == 1 + self[i][j][index[i][j][k]] += src[i][j][k] # if dim == 2 + +Note: + {forward_reproducibility_note} + +.. note:: + + The backward pass is implemented only for ``src.shape == index.shape``. + +.. warning:: + + This function is in beta and may change in the near future. + +Args: + dim (int): the axis along which to index + index (LongTensor): the indices of elements to scatter and reduce. + src (Tensor): the source elements to scatter and reduce + reduce (str): the reduction operation to apply for non-unique indices + (:obj:`"sum"`, :obj:`"prod"`, :obj:`"mean"`, :obj:`"amax"`, :obj:`"amin"`) + include_self (bool): whether elements from the :attr:`self` tensor are + included in the reduction + +Example:: + + >>> src = torch.tensor([1., 2., 3., 4., 5., 6.]) + >>> index = torch.tensor([0, 1, 0, 1, 2, 1]) + >>> input = torch.tensor([1., 2., 3., 4.]) + >>> input.scatter_reduce(0, index, src, reduce="sum") + tensor([5., 14., 8., 4.]) + >>> input.scatter_reduce(0, index, src, reduce="sum", include_self=False) + tensor([4., 12., 5., 4.]) + >>> input2 = torch.tensor([5., 4., 3., 2.]) + >>> input2.scatter_reduce(0, index, src, reduce="amax") + tensor([5., 6., 5., 2.]) + >>> input2.scatter_reduce(0, index, src, reduce="amax", include_self=False) + tensor([3., 6., 5., 2.]) + + +""".format(**reproducibility_notes), +) + +add_docstr_all( + "select", + r""" +select(dim, index) -> Tensor + +See :func:`torch.select` +""", +) + +add_docstr_all( + "select_scatter", + r""" +select_scatter(src, dim, index) -> Tensor + +See :func:`torch.select_scatter` +""", +) + +add_docstr_all( + "slice_scatter", + r""" +slice_scatter(src, dim=0, start=None, end=None, step=1) -> Tensor + +See :func:`torch.slice_scatter` +""", +) + +add_docstr_all( + "set_", + r""" +set_(source=None, storage_offset=0, size=None, stride=None) -> Tensor + +Sets the underlying storage, size, and strides. If :attr:`source` is a tensor, +:attr:`self` tensor will share the same storage and have the same size and +strides as :attr:`source`. Changes to elements in one tensor will be reflected +in the other. + +If :attr:`source` is a :class:`~torch.Storage`, the method sets the underlying +storage, offset, size, and stride. + +Args: + source (Tensor or Storage): the tensor or storage to use + storage_offset (int, optional): the offset in the storage + size (torch.Size, optional): the desired size. Defaults to the size of the source. + stride (tuple, optional): the desired stride. Defaults to C-contiguous strides. +""", +) + +add_docstr_all( + "sigmoid", + r""" +sigmoid() -> Tensor + +See :func:`torch.sigmoid` +""", +) + +add_docstr_all( + "sigmoid_", + r""" +sigmoid_() -> Tensor + +In-place version of :meth:`~Tensor.sigmoid` +""", +) + +add_docstr_all( + "logit", + r""" +logit() -> Tensor + +See :func:`torch.logit` +""", +) + +add_docstr_all( + "logit_", + r""" +logit_() -> Tensor + +In-place version of :meth:`~Tensor.logit` +""", +) + +add_docstr_all( + "sign", + r""" +sign() -> Tensor + +See :func:`torch.sign` +""", +) + +add_docstr_all( + "sign_", + r""" +sign_() -> Tensor + +In-place version of :meth:`~Tensor.sign` +""", +) + +add_docstr_all( + "signbit", + r""" +signbit() -> Tensor + +See :func:`torch.signbit` +""", +) + +add_docstr_all( + "sgn", + r""" +sgn() -> Tensor + +See :func:`torch.sgn` +""", +) + +add_docstr_all( + "sgn_", + r""" +sgn_() -> Tensor + +In-place version of :meth:`~Tensor.sgn` +""", +) + +add_docstr_all( + "sin", + r""" +sin() -> Tensor + +See :func:`torch.sin` +""", +) + +add_docstr_all( + "sin_", + r""" +sin_() -> Tensor + +In-place version of :meth:`~Tensor.sin` +""", +) + +add_docstr_all( + "sinc", + r""" +sinc() -> Tensor + +See :func:`torch.sinc` +""", +) + +add_docstr_all( + "sinc_", + r""" +sinc_() -> Tensor + +In-place version of :meth:`~Tensor.sinc` +""", +) + +add_docstr_all( + "sinh", + r""" +sinh() -> Tensor + +See :func:`torch.sinh` +""", +) + +add_docstr_all( + "sinh_", + r""" +sinh_() -> Tensor + +In-place version of :meth:`~Tensor.sinh` +""", +) + +add_docstr_all( + "size", + r""" +size(dim=None) -> torch.Size or int + +Returns the size of the :attr:`self` tensor. If ``dim`` is not specified, +the returned value is a :class:`torch.Size`, a subclass of :class:`tuple`. +If ``dim`` is specified, returns an int holding the size of that dimension. + +Args: + dim (int, optional): The dimension for which to retrieve the size. + +Example:: + + >>> t = torch.empty(3, 4, 5) + >>> t.size() + torch.Size([3, 4, 5]) + >>> t.size(dim=1) + 4 + +""", +) + +add_docstr_all( + "shape", + r""" +shape() -> torch.Size + +Returns the size of the :attr:`self` tensor. Alias for :attr:`size`. + +See also :meth:`Tensor.size`. + +Example:: + + >>> t = torch.empty(3, 4, 5) + >>> t.size() + torch.Size([3, 4, 5]) + >>> t.shape + torch.Size([3, 4, 5]) + +""", +) + +add_docstr_all( + "sort", + r""" +sort(dim=-1, descending=False) -> (Tensor, LongTensor) + +See :func:`torch.sort` +""", +) + +add_docstr_all( + "msort", + r""" +msort() -> Tensor + +See :func:`torch.msort` +""", +) + +add_docstr_all( + "argsort", + r""" +argsort(dim=-1, descending=False) -> LongTensor + +See :func:`torch.argsort` +""", +) + +add_docstr_all( + "sparse_dim", + r""" +sparse_dim() -> int + +Return the number of sparse dimensions in a :ref:`sparse tensor ` :attr:`self`. + +.. note:: + Returns ``0`` if :attr:`self` is not a sparse tensor. + +See also :meth:`Tensor.dense_dim` and :ref:`hybrid tensors `. +""", +) + +add_docstr_all( + "sparse_resize_", + r""" +sparse_resize_(size, sparse_dim, dense_dim) -> Tensor + +Resizes :attr:`self` :ref:`sparse tensor ` to the desired +size and the number of sparse and dense dimensions. + +.. note:: + If the number of specified elements in :attr:`self` is zero, then + :attr:`size`, :attr:`sparse_dim`, and :attr:`dense_dim` can be any + size and positive integers such that ``len(size) == sparse_dim + + dense_dim``. + + If :attr:`self` specifies one or more elements, however, then each + dimension in :attr:`size` must not be smaller than the corresponding + dimension of :attr:`self`, :attr:`sparse_dim` must equal the number + of sparse dimensions in :attr:`self`, and :attr:`dense_dim` must + equal the number of dense dimensions in :attr:`self`. + +.. warning:: + Throws an error if :attr:`self` is not a sparse tensor. + +Args: + size (torch.Size): the desired size. If :attr:`self` is non-empty + sparse tensor, the desired size cannot be smaller than the + original size. + sparse_dim (int): the number of sparse dimensions + dense_dim (int): the number of dense dimensions +""", +) + +add_docstr_all( + "sparse_resize_and_clear_", + r""" +sparse_resize_and_clear_(size, sparse_dim, dense_dim) -> Tensor + +Removes all specified elements from a :ref:`sparse tensor +` :attr:`self` and resizes :attr:`self` to the desired +size and the number of sparse and dense dimensions. + +.. warning: + Throws an error if :attr:`self` is not a sparse tensor. + +Args: + size (torch.Size): the desired size. + sparse_dim (int): the number of sparse dimensions + dense_dim (int): the number of dense dimensions +""", +) + +add_docstr_all( + "sqrt", + r""" +sqrt() -> Tensor + +See :func:`torch.sqrt` +""", +) + +add_docstr_all( + "sqrt_", + r""" +sqrt_() -> Tensor + +In-place version of :meth:`~Tensor.sqrt` +""", +) + +add_docstr_all( + "square", + r""" +square() -> Tensor + +See :func:`torch.square` +""", +) + +add_docstr_all( + "square_", + r""" +square_() -> Tensor + +In-place version of :meth:`~Tensor.square` +""", +) + +add_docstr_all( + "squeeze", + r""" +squeeze(dim=None) -> Tensor + +See :func:`torch.squeeze` +""", +) + +add_docstr_all( + "squeeze_", + r""" +squeeze_(dim=None) -> Tensor + +In-place version of :meth:`~Tensor.squeeze` +""", +) + +add_docstr_all( + "std", + r""" +std(dim=None, *, correction=1, keepdim=False) -> Tensor + +See :func:`torch.std` +""", +) + +add_docstr_all( + "storage_offset", + r""" +storage_offset() -> int + +Returns :attr:`self` tensor's offset in the underlying storage in terms of +number of storage elements (not bytes). + +Example:: + + >>> x = torch.tensor([1, 2, 3, 4, 5]) + >>> x.storage_offset() + 0 + >>> x[3:].storage_offset() + 3 + +""", +) + +add_docstr_all( + "untyped_storage", + r""" +untyped_storage() -> torch.UntypedStorage + +Returns the underlying :class:`UntypedStorage`. +""", +) + +add_docstr_all( + "stride", + r""" +stride(dim) -> tuple or int + +Returns the stride of :attr:`self` tensor. + +Stride is the jump necessary to go from one element to the next one in the +specified dimension :attr:`dim`. A tuple of all strides is returned when no +argument is passed in. Otherwise, an integer value is returned as the stride in +the particular dimension :attr:`dim`. + +Args: + dim (int, optional): the desired dimension in which stride is required + +Example:: + + >>> x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]) + >>> x.stride() + (5, 1) + >>> x.stride(0) + 5 + >>> x.stride(-1) + 1 + +""", +) + +add_docstr_all( + "sub", + r""" +sub(other, *, alpha=1) -> Tensor + +See :func:`torch.sub`. +""", +) + +add_docstr_all( + "sub_", + r""" +sub_(other, *, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.sub` +""", +) + +add_docstr_all( + "subtract", + r""" +subtract(other, *, alpha=1) -> Tensor + +See :func:`torch.subtract`. +""", +) + +add_docstr_all( + "subtract_", + r""" +subtract_(other, *, alpha=1) -> Tensor + +In-place version of :meth:`~Tensor.subtract`. +""", +) + +add_docstr_all( + "sum", + r""" +sum(dim=None, keepdim=False, dtype=None) -> Tensor + +See :func:`torch.sum` +""", +) + +add_docstr_all( + "nansum", + r""" +nansum(dim=None, keepdim=False, dtype=None) -> Tensor + +See :func:`torch.nansum` +""", +) + +add_docstr_all( + "svd", + r""" +svd(some=True, compute_uv=True) -> (Tensor, Tensor, Tensor) + +See :func:`torch.svd` +""", +) + +add_docstr_all( + "swapdims", + r""" +swapdims(dim0, dim1) -> Tensor + +See :func:`torch.swapdims` +""", +) + +add_docstr_all( + "swapdims_", + r""" +swapdims_(dim0, dim1) -> Tensor + +In-place version of :meth:`~Tensor.swapdims` +""", +) + +add_docstr_all( + "swapaxes", + r""" +swapaxes(axis0, axis1) -> Tensor + +See :func:`torch.swapaxes` +""", +) + +add_docstr_all( + "swapaxes_", + r""" +swapaxes_(axis0, axis1) -> Tensor + +In-place version of :meth:`~Tensor.swapaxes` +""", +) + +add_docstr_all( + "t", + r""" +t() -> Tensor + +See :func:`torch.t` +""", +) + +add_docstr_all( + "t_", + r""" +t_() -> Tensor + +In-place version of :meth:`~Tensor.t` +""", +) + +add_docstr_all( + "tile", + r""" +tile(dims) -> Tensor + +See :func:`torch.tile` +""", +) + +add_docstr_all( + "to", + r""" +to(*args, **kwargs) -> Tensor + +Performs Tensor dtype and/or device conversion. A :class:`torch.dtype` and :class:`torch.device` are +inferred from the arguments of ``self.to(*args, **kwargs)``. + +.. note:: + + If the ``self`` Tensor already + has the correct :class:`torch.dtype` and :class:`torch.device`, then ``self`` is returned. + Otherwise, the returned tensor is a copy of ``self`` with the desired + :class:`torch.dtype` and :class:`torch.device`. + +.. note:: + + If ``self`` requires gradients (``requires_grad=True``) but the target + ``dtype`` specified is an integer type, the returned tensor will implicitly + set ``requires_grad=False``. This is because only tensors with + floating-point or complex dtypes can require gradients. + +Here are the ways to call ``to``: + +.. method:: to(dtype, non_blocking=False, copy=False, memory_format=torch.preserve_format) -> Tensor + :noindex: + + Returns a Tensor with the specified :attr:`dtype` + + Args: + {memory_format} + +.. note:: + + According to `C++ type conversion rules `_, + converting floating point value to integer type will truncate the fractional part. + If the truncated value cannot fit into the target type (e.g., casting ``torch.inf`` to ``torch.long``), + the behavior is undefined and the result may vary across platforms. + +.. method:: to(device=None, dtype=None, non_blocking=False, copy=False, memory_format=torch.preserve_format) -> Tensor + :noindex: + + Returns a Tensor with the specified :attr:`device` and (optional) + :attr:`dtype`. If :attr:`dtype` is ``None`` it is inferred to be ``self.dtype``. + When :attr:`non_blocking` is set to ``True``, the function attempts to perform + the conversion asynchronously with respect to the host, if possible. This + asynchronous behavior applies to both pinned and pageable memory. However, + caution is advised when using this feature. For more information, refer to the + `tutorial on good usage of non_blocking and pin_memory `__. + When :attr:`copy` is set, a new Tensor is created even when the Tensor + already matches the desired conversion. + + Args: + {memory_format} + +.. method:: to(other, non_blocking=False, copy=False) -> Tensor + :noindex: + + Returns a Tensor with same :class:`torch.dtype` and :class:`torch.device` as + the Tensor :attr:`other`. + When :attr:`non_blocking` is set to ``True``, the function attempts to perform + the conversion asynchronously with respect to the host, if possible. This + asynchronous behavior applies to both pinned and pageable memory. However, + caution is advised when using this feature. For more information, refer to the + `tutorial on good usage of non_blocking and pin_memory `__. + When :attr:`copy` is set, a new Tensor is created even when the Tensor + already matches the desired conversion. + +Example:: + + >>> tensor = torch.randn(2, 2) # Initially dtype=float32, device=cpu + >>> tensor.to(torch.float64) + tensor([[-0.5044, 0.0005], + [ 0.3310, -0.0584]], dtype=torch.float64) + + >>> cuda0 = torch.device('cuda:0') + >>> tensor.to(cuda0) + tensor([[-0.5044, 0.0005], + [ 0.3310, -0.0584]], device='cuda:0') + + >>> tensor.to(cuda0, dtype=torch.float64) + tensor([[-0.5044, 0.0005], + [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0') + + >>> other = torch.randn((), dtype=torch.float64, device=cuda0) + >>> tensor.to(other, non_blocking=True) + tensor([[-0.5044, 0.0005], + [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0') +""".format(**common_args), +) + +add_docstr_all( + "byte", + r""" +byte(memory_format=torch.preserve_format) -> Tensor + +``self.byte()`` is equivalent to ``self.to(torch.uint8)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "bool", + r""" +bool(memory_format=torch.preserve_format) -> Tensor + +``self.bool()`` is equivalent to ``self.to(torch.bool)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "char", + r""" +char(memory_format=torch.preserve_format) -> Tensor + +``self.char()`` is equivalent to ``self.to(torch.int8)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "bfloat16", + r""" +bfloat16(memory_format=torch.preserve_format) -> Tensor +``self.bfloat16()`` is equivalent to ``self.to(torch.bfloat16)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "double", + r""" +double(memory_format=torch.preserve_format) -> Tensor + +``self.double()`` is equivalent to ``self.to(torch.float64)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "float", + r""" +float(memory_format=torch.preserve_format) -> Tensor + +``self.float()`` is equivalent to ``self.to(torch.float32)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "cdouble", + r""" +cdouble(memory_format=torch.preserve_format) -> Tensor + +``self.cdouble()`` is equivalent to ``self.to(torch.complex128)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "cfloat", + r""" +cfloat(memory_format=torch.preserve_format) -> Tensor + +``self.cfloat()`` is equivalent to ``self.to(torch.complex64)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "chalf", + r""" +chalf(memory_format=torch.preserve_format) -> Tensor + +``self.chalf()`` is equivalent to ``self.to(torch.complex32)``. See :func:`to`. + +Args: + {memory_format} + """.format(**common_args), +) + +add_docstr_all( + "half", + r""" +half(memory_format=torch.preserve_format) -> Tensor + +``self.half()`` is equivalent to ``self.to(torch.float16)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "int", + r""" +int(memory_format=torch.preserve_format) -> Tensor + +``self.int()`` is equivalent to ``self.to(torch.int32)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "int_repr", + r""" +int_repr() -> Tensor + +Given a quantized Tensor, +``self.int_repr()`` returns a CPU Tensor with uint8_t as data type that stores the +underlying uint8_t values of the given Tensor. +""", +) + + +add_docstr_all( + "long", + r""" +long(memory_format=torch.preserve_format) -> Tensor + +``self.long()`` is equivalent to ``self.to(torch.int64)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "short", + r""" +short(memory_format=torch.preserve_format) -> Tensor + +``self.short()`` is equivalent to ``self.to(torch.int16)``. See :func:`to`. + +Args: + {memory_format} +""".format(**common_args), +) + +add_docstr_all( + "take", + r""" +take(indices) -> Tensor + +See :func:`torch.take` +""", +) + +add_docstr_all( + "take_along_dim", + r""" +take_along_dim(indices, dim) -> Tensor + +See :func:`torch.take_along_dim` +""", +) + +add_docstr_all( + "tan", + r""" +tan() -> Tensor + +See :func:`torch.tan` +""", +) + +add_docstr_all( + "tan_", + r""" +tan_() -> Tensor + +In-place version of :meth:`~Tensor.tan` +""", +) + +add_docstr_all( + "tanh", + r""" +tanh() -> Tensor + +See :func:`torch.tanh` +""", +) + +add_docstr_all( + "softmax", + r""" +softmax(dim) -> Tensor + +Alias for :func:`torch.nn.functional.softmax`. +""", +) + +add_docstr_all( + "tanh_", + r""" +tanh_() -> Tensor + +In-place version of :meth:`~Tensor.tanh` +""", +) + +add_docstr_all( + "tolist", + r""" +tolist() -> list or number + +Returns the tensor as a (nested) list. For scalars, a standard +Python number is returned, just like with :meth:`~Tensor.item`. +Tensors are automatically moved to the CPU first if necessary. + +This operation is not differentiable. + +Examples:: + + >>> a = torch.randn(2, 2) + >>> a.tolist() + [[0.012766935862600803, 0.5415473580360413], + [-0.08909505605697632, 0.7729271650314331]] + >>> a[0,0].tolist() + 0.012766935862600803 +""", +) + +add_docstr_all( + "topk", + r""" +topk(k, dim=None, largest=True, sorted=True) -> (Tensor, LongTensor) + +See :func:`torch.topk` +""", +) + +add_docstr_all( + "to_dense", + r""" +to_dense(dtype=None, *, masked_grad=True) -> Tensor + +Creates a strided copy of :attr:`self` if :attr:`self` is not a strided tensor, otherwise returns :attr:`self`. + +Keyword args: + {dtype} + masked_grad (bool, optional): If set to ``True`` (default) and + :attr:`self` has a sparse layout then the backward of + :meth:`to_dense` returns ``grad.sparse_mask(self)``. + +Example:: + + >>> s = torch.sparse_coo_tensor( + ... torch.tensor([[1, 1], + ... [0, 2]]), + ... torch.tensor([9, 10]), + ... size=(3, 3)) + >>> s.to_dense() + tensor([[ 0, 0, 0], + [ 9, 0, 10], + [ 0, 0, 0]]) +""", +) + +add_docstr_all( + "to_sparse", + r""" +to_sparse(sparseDims) -> Tensor + +Returns a sparse copy of the tensor. PyTorch supports sparse tensors in +:ref:`coordinate format `. + +Args: + sparseDims (int, optional): the number of sparse dimensions to include in the new sparse tensor + +Example:: + + >>> d = torch.tensor([[0, 0, 0], [9, 0, 10], [0, 0, 0]]) + >>> d + tensor([[ 0, 0, 0], + [ 9, 0, 10], + [ 0, 0, 0]]) + >>> d.to_sparse() + tensor(indices=tensor([[1, 1], + [0, 2]]), + values=tensor([ 9, 10]), + size=(3, 3), nnz=2, layout=torch.sparse_coo) + >>> d.to_sparse(1) + tensor(indices=tensor([[1]]), + values=tensor([[ 9, 0, 10]]), + size=(3, 3), nnz=1, layout=torch.sparse_coo) + +.. method:: to_sparse(*, layout=None, blocksize=None, dense_dim=None) -> Tensor + :noindex: + +Returns a sparse tensor with the specified layout and blocksize. If +the :attr:`self` is strided, the number of dense dimensions could be +specified, and a hybrid sparse tensor will be created, with +`dense_dim` dense dimensions and `self.dim() - 2 - dense_dim` batch +dimension. + +.. note:: If the :attr:`self` layout and blocksize parameters match + with the specified layout and blocksize, return + :attr:`self`. Otherwise, return a sparse tensor copy of + :attr:`self`. + +Args: + + layout (:class:`torch.layout`, optional): The desired sparse + layout. One of ``torch.sparse_coo``, ``torch.sparse_csr``, + ``torch.sparse_csc``, ``torch.sparse_bsr``, or + ``torch.sparse_bsc``. Default: if ``None``, + ``torch.sparse_coo``. + + blocksize (list, tuple, :class:`torch.Size`, optional): Block size + of the resulting BSR or BSC tensor. For other layouts, + specifying the block size that is not ``None`` will result in a + RuntimeError exception. A block size must be a tuple of length + two such that its items evenly divide the two sparse dimensions. + + dense_dim (int, optional): Number of dense dimensions of the + resulting CSR, CSC, BSR or BSC tensor. This argument should be + used only if :attr:`self` is a strided tensor, and must be a + value between 0 and dimension of :attr:`self` tensor minus two. + +Example:: + + >>> x = torch.tensor([[1, 0], [0, 0], [2, 3]]) + >>> x.to_sparse(layout=torch.sparse_coo) + tensor(indices=tensor([[0, 2, 2], + [0, 0, 1]]), + values=tensor([1, 2, 3]), + size=(3, 2), nnz=3, layout=torch.sparse_coo) + >>> x.to_sparse(layout=torch.sparse_bsr, blocksize=(1, 2)) + tensor(crow_indices=tensor([0, 1, 1, 2]), + col_indices=tensor([0, 0]), + values=tensor([[[1, 0]], + [[2, 3]]]), size=(3, 2), nnz=2, layout=torch.sparse_bsr) + >>> x.to_sparse(layout=torch.sparse_bsr, blocksize=(2, 1)) + RuntimeError: Tensor size(-2) 3 needs to be divisible by blocksize[0] 2 + >>> x.to_sparse(layout=torch.sparse_csr, blocksize=(3, 1)) + RuntimeError: to_sparse for Strided to SparseCsr conversion does not use specified blocksize + + >>> x = torch.tensor([[[1], [0]], [[0], [0]], [[2], [3]]]) + >>> x.to_sparse(layout=torch.sparse_csr, dense_dim=1) + tensor(crow_indices=tensor([0, 1, 1, 3]), + col_indices=tensor([0, 0, 1]), + values=tensor([[1], + [2], + [3]]), size=(3, 2, 1), nnz=3, layout=torch.sparse_csr) + +""", +) + +add_docstr_all( + "to_sparse_csr", + r""" +to_sparse_csr(dense_dim=None) -> Tensor + +Convert a tensor to compressed row storage format (CSR). Except for +strided tensors, only works with 2D tensors. If the :attr:`self` is +strided, then the number of dense dimensions could be specified, and a +hybrid CSR tensor will be created, with `dense_dim` dense dimensions +and `self.dim() - 2 - dense_dim` batch dimension. + +Args: + + dense_dim (int, optional): Number of dense dimensions of the + resulting CSR tensor. This argument should be used only if + :attr:`self` is a strided tensor, and must be a value between 0 + and dimension of :attr:`self` tensor minus two. + +Example:: + + >>> dense = torch.randn(5, 5) + >>> sparse = dense.to_sparse_csr() + >>> sparse._nnz() + 25 + + >>> dense = torch.zeros(3, 3, 1, 1) + >>> dense[0, 0] = dense[1, 2] = dense[2, 1] = 1 + >>> dense.to_sparse_csr(dense_dim=2) + tensor(crow_indices=tensor([0, 1, 2, 3]), + col_indices=tensor([0, 2, 1]), + values=tensor([[[1.]], + + [[1.]], + + [[1.]]]), size=(3, 3, 1, 1), nnz=3, + layout=torch.sparse_csr) + +""", +) + +add_docstr_all( + "to_sparse_csc", + r""" +to_sparse_csc() -> Tensor + +Convert a tensor to compressed column storage (CSC) format. Except +for strided tensors, only works with 2D tensors. If the :attr:`self` +is strided, then the number of dense dimensions could be specified, +and a hybrid CSC tensor will be created, with `dense_dim` dense +dimensions and `self.dim() - 2 - dense_dim` batch dimension. + +Args: + + dense_dim (int, optional): Number of dense dimensions of the + resulting CSC tensor. This argument should be used only if + :attr:`self` is a strided tensor, and must be a value between 0 + and dimension of :attr:`self` tensor minus two. + +Example:: + + >>> dense = torch.randn(5, 5) + >>> sparse = dense.to_sparse_csc() + >>> sparse._nnz() + 25 + + >>> dense = torch.zeros(3, 3, 1, 1) + >>> dense[0, 0] = dense[1, 2] = dense[2, 1] = 1 + >>> dense.to_sparse_csc(dense_dim=2) + tensor(ccol_indices=tensor([0, 1, 2, 3]), + row_indices=tensor([0, 2, 1]), + values=tensor([[[1.]], + + [[1.]], + + [[1.]]]), size=(3, 3, 1, 1), nnz=3, + layout=torch.sparse_csc) + +""", +) + +add_docstr_all( + "to_sparse_bsr", + r""" +to_sparse_bsr(blocksize, dense_dim) -> Tensor + +Convert a tensor to a block sparse row (BSR) storage format of given +blocksize. If the :attr:`self` is strided, then the number of dense +dimensions could be specified, and a hybrid BSR tensor will be +created, with `dense_dim` dense dimensions and `self.dim() - 2 - +dense_dim` batch dimension. + +Args: + + blocksize (list, tuple, :class:`torch.Size`, optional): Block size + of the resulting BSR tensor. A block size must be a tuple of + length two such that its items evenly divide the two sparse + dimensions. + + dense_dim (int, optional): Number of dense dimensions of the + resulting BSR tensor. This argument should be used only if + :attr:`self` is a strided tensor, and must be a value between 0 + and dimension of :attr:`self` tensor minus two. + +Example:: + + >>> dense = torch.randn(10, 10) + >>> sparse = dense.to_sparse_csr() + >>> sparse_bsr = sparse.to_sparse_bsr((5, 5)) + >>> sparse_bsr.col_indices() + tensor([0, 1, 0, 1]) + + >>> dense = torch.zeros(4, 3, 1) + >>> dense[0:2, 0] = dense[0:2, 2] = dense[2:4, 1] = 1 + >>> dense.to_sparse_bsr((2, 1), 1) + tensor(crow_indices=tensor([0, 2, 3]), + col_indices=tensor([0, 2, 1]), + values=tensor([[[[1.]], + + [[1.]]], + + + [[[1.]], + + [[1.]]], + + + [[[1.]], + + [[1.]]]]), size=(4, 3, 1), nnz=3, + layout=torch.sparse_bsr) + +""", +) + +add_docstr_all( + "to_sparse_bsc", + r""" +to_sparse_bsc(blocksize, dense_dim) -> Tensor + +Convert a tensor to a block sparse column (BSC) storage format of +given blocksize. If the :attr:`self` is strided, then the number of +dense dimensions could be specified, and a hybrid BSC tensor will be +created, with `dense_dim` dense dimensions and `self.dim() - 2 - +dense_dim` batch dimension. + +Args: + + blocksize (list, tuple, :class:`torch.Size`, optional): Block size + of the resulting BSC tensor. A block size must be a tuple of + length two such that its items evenly divide the two sparse + dimensions. + + dense_dim (int, optional): Number of dense dimensions of the + resulting BSC tensor. This argument should be used only if + :attr:`self` is a strided tensor, and must be a value between 0 + and dimension of :attr:`self` tensor minus two. + +Example:: + + >>> dense = torch.randn(10, 10) + >>> sparse = dense.to_sparse_csr() + >>> sparse_bsc = sparse.to_sparse_bsc((5, 5)) + >>> sparse_bsc.row_indices() + tensor([0, 1, 0, 1]) + + >>> dense = torch.zeros(4, 3, 1) + >>> dense[0:2, 0] = dense[0:2, 2] = dense[2:4, 1] = 1 + >>> dense.to_sparse_bsc((2, 1), 1) + tensor(ccol_indices=tensor([0, 1, 2, 3]), + row_indices=tensor([0, 1, 0]), + values=tensor([[[[1.]], + + [[1.]]], + + + [[[1.]], + + [[1.]]], + + + [[[1.]], + + [[1.]]]]), size=(4, 3, 1), nnz=3, + layout=torch.sparse_bsc) + +""", +) + +add_docstr_all( + "to_mkldnn", + r""" +to_mkldnn() -> Tensor +Returns a copy of the tensor in ``torch.mkldnn`` layout. + +""", +) + +add_docstr_all( + "trace", + r""" +trace() -> Tensor + +See :func:`torch.trace` +""", +) + +add_docstr_all( + "transpose", + r""" +transpose(dim0, dim1) -> Tensor + +See :func:`torch.transpose` +""", +) + +add_docstr_all( + "transpose_", + r""" +transpose_(dim0, dim1) -> Tensor + +In-place version of :meth:`~Tensor.transpose` +""", +) + +add_docstr_all( + "triangular_solve", + r""" +triangular_solve(A, upper=True, transpose=False, unitriangular=False) -> (Tensor, Tensor) + +See :func:`torch.triangular_solve` +""", +) + +add_docstr_all( + "tril", + r""" +tril(diagonal=0) -> Tensor + +See :func:`torch.tril` +""", +) + +add_docstr_all( + "tril_", + r""" +tril_(diagonal=0) -> Tensor + +In-place version of :meth:`~Tensor.tril` +""", +) + +add_docstr_all( + "triu", + r""" +triu(diagonal=0) -> Tensor + +See :func:`torch.triu` +""", +) + +add_docstr_all( + "triu_", + r""" +triu_(diagonal=0) -> Tensor + +In-place version of :meth:`~Tensor.triu` +""", +) + +add_docstr_all( + "true_divide", + r""" +true_divide(value) -> Tensor + +See :func:`torch.true_divide` +""", +) + +add_docstr_all( + "true_divide_", + r""" +true_divide_(value) -> Tensor + +In-place version of :meth:`~Tensor.true_divide_` +""", +) + +add_docstr_all( + "trunc", + r""" +trunc() -> Tensor + +See :func:`torch.trunc` +""", +) + +add_docstr_all( + "fix", + r""" +fix() -> Tensor + +See :func:`torch.fix`. +""", +) + +add_docstr_all( + "trunc_", + r""" +trunc_() -> Tensor + +In-place version of :meth:`~Tensor.trunc` +""", +) + +add_docstr_all( + "fix_", + r""" +fix_() -> Tensor + +In-place version of :meth:`~Tensor.fix` +""", +) + +add_docstr_all( + "type", + r""" +type(dtype=None, non_blocking=False, **kwargs) -> str or Tensor +Returns the type if `dtype` is not provided, else casts this object to +the specified type. + +If this is already of the correct type, no copy is performed and the +original object is returned. + +Args: + dtype (dtype or string): The desired type + non_blocking (bool): If ``True``, and the source is in pinned memory + and destination is on the GPU or vice versa, the copy is performed + asynchronously with respect to the host. Otherwise, the argument + has no effect. + **kwargs: For compatibility, may contain the key ``async`` in place of + the ``non_blocking`` argument. The ``async`` arg is deprecated. +""", +) + +add_docstr_all( + "type_as", + r""" +type_as(tensor) -> Tensor + +Returns this tensor cast to the type of the given tensor. + +This is a no-op if the tensor is already of the correct type. This is +equivalent to ``self.type(tensor.type())`` + +Args: + tensor (Tensor): the tensor which has the desired type +""", +) + +add_docstr_all( + "unfold", + r""" +unfold(dimension, size, step) -> Tensor + +Returns a view of the original tensor which contains all slices of size :attr:`size` from +:attr:`self` tensor in the dimension :attr:`dimension`. + +Step between two slices is given by :attr:`step`. + +If `sizedim` is the size of dimension :attr:`dimension` for :attr:`self`, the size of +dimension :attr:`dimension` in the returned tensor will be +`(sizedim - size) / step + 1`. + +An additional dimension of size :attr:`size` is appended in the returned tensor. + +Args: + dimension (int): dimension in which unfolding happens + size (int): the size of each slice that is unfolded + step (int): the step between each slice + +Example:: + + >>> x = torch.arange(1., 8) + >>> x + tensor([ 1., 2., 3., 4., 5., 6., 7.]) + >>> x.unfold(0, 2, 1) + tensor([[ 1., 2.], + [ 2., 3.], + [ 3., 4.], + [ 4., 5.], + [ 5., 6.], + [ 6., 7.]]) + >>> x.unfold(0, 2, 2) + tensor([[ 1., 2.], + [ 3., 4.], + [ 5., 6.]]) +""", +) + +add_docstr_all( + "uniform_", + r""" +uniform_(from=0, to=1, *, generator=None) -> Tensor + +Fills :attr:`self` tensor with numbers sampled from the continuous uniform +distribution: + +.. math:: + f(x) = \dfrac{1}{\text{to} - \text{from}} +""", +) + +add_docstr_all( + "unsqueeze", + r""" +unsqueeze(dim) -> Tensor + +See :func:`torch.unsqueeze` +""", +) + +add_docstr_all( + "unsqueeze_", + r""" +unsqueeze_(dim) -> Tensor + +In-place version of :meth:`~Tensor.unsqueeze` +""", +) + +add_docstr_all( + "var", + r""" +var(dim=None, *, correction=1, keepdim=False) -> Tensor + +See :func:`torch.var` +""", +) + +add_docstr_all( + "vdot", + r""" +vdot(other) -> Tensor + +See :func:`torch.vdot` +""", +) + +add_docstr_all( + "view", + r""" +view(*shape) -> Tensor + +Returns a new tensor with the same data as the :attr:`self` tensor but of a +different :attr:`shape`. + +The returned tensor shares the same data and must have the same number +of elements, but may have a different size. For a tensor to be viewed, the new +view size must be compatible with its original size and stride, i.e., each new +view dimension must either be a subspace of an original dimension, or only span +across original dimensions :math:`d, d+1, \dots, d+k` that satisfy the following +contiguity-like condition that :math:`\forall i = d, \dots, d+k-1`, + +.. math:: + + \text{stride}[i] = \text{stride}[i+1] \times \text{size}[i+1] + +Otherwise, it will not be possible to view :attr:`self` tensor as :attr:`shape` +without copying it (e.g., via :meth:`contiguous`). When it is unclear whether a +:meth:`view` can be performed, it is advisable to use :meth:`reshape`, which +returns a view if the shapes are compatible, and copies (equivalent to calling +:meth:`contiguous`) otherwise. + +Args: + shape (torch.Size or int...): the desired size + +Example:: + + >>> x = torch.randn(4, 4) + >>> x.size() + torch.Size([4, 4]) + >>> y = x.view(16) + >>> y.size() + torch.Size([16]) + >>> z = x.view(-1, 8) # the size -1 is inferred from other dimensions + >>> z.size() + torch.Size([2, 8]) + + >>> a = torch.randn(1, 2, 3, 4) + >>> a.size() + torch.Size([1, 2, 3, 4]) + >>> b = a.transpose(1, 2) # Swaps 2nd and 3rd dimension + >>> b.size() + torch.Size([1, 3, 2, 4]) + >>> c = a.view(1, 3, 2, 4) # Does not change tensor layout in memory + >>> c.size() + torch.Size([1, 3, 2, 4]) + >>> torch.equal(b, c) + False + + +.. method:: view(dtype) -> Tensor + :noindex: + +Returns a new tensor with the same data as the :attr:`self` tensor but of a +different :attr:`dtype`. + +If the element size of :attr:`dtype` is different than that of ``self.dtype``, +then the size of the last dimension of the output will be scaled +proportionally. For instance, if :attr:`dtype` element size is twice that of +``self.dtype``, then each pair of elements in the last dimension of +:attr:`self` will be combined, and the size of the last dimension of the output +will be half that of :attr:`self`. If :attr:`dtype` element size is half that +of ``self.dtype``, then each element in the last dimension of :attr:`self` will +be split in two, and the size of the last dimension of the output will be +double that of :attr:`self`. For this to be possible, the following conditions +must be true: + + * ``self.dim()`` must be greater than 0. + * ``self.stride(-1)`` must be 1. + +Additionally, if the element size of :attr:`dtype` is greater than that of +``self.dtype``, the following conditions must be true as well: + + * ``self.size(-1)`` must be divisible by the ratio between the element + sizes of the dtypes. + * ``self.storage_offset()`` must be divisible by the ratio between the + element sizes of the dtypes. + * The strides of all dimensions, except the last dimension, must be + divisible by the ratio between the element sizes of the dtypes. + +If any of the above conditions are not met, an error is thrown. + +.. warning:: + + This overload is not supported by TorchScript, and using it in a Torchscript + program will cause undefined behavior. + + +Args: + dtype (:class:`torch.dtype`): the desired dtype + +Example:: + + >>> x = torch.randn(4, 4) + >>> x + tensor([[ 0.9482, -0.0310, 1.4999, -0.5316], + [-0.1520, 0.7472, 0.5617, -0.8649], + [-2.4724, -0.0334, -0.2976, -0.8499], + [-0.2109, 1.9913, -0.9607, -0.6123]]) + >>> x.dtype + torch.float32 + + >>> y = x.view(torch.int32) + >>> y + tensor([[ 1064483442, -1124191867, 1069546515, -1089989247], + [-1105482831, 1061112040, 1057999968, -1084397505], + [-1071760287, -1123489973, -1097310419, -1084649136], + [-1101533110, 1073668768, -1082790149, -1088634448]], + dtype=torch.int32) + >>> y[0, 0] = 1000000000 + >>> x + tensor([[ 0.0047, -0.0310, 1.4999, -0.5316], + [-0.1520, 0.7472, 0.5617, -0.8649], + [-2.4724, -0.0334, -0.2976, -0.8499], + [-0.2109, 1.9913, -0.9607, -0.6123]]) + + >>> x.view(torch.cfloat) + tensor([[ 0.0047-0.0310j, 1.4999-0.5316j], + [-0.1520+0.7472j, 0.5617-0.8649j], + [-2.4724-0.0334j, -0.2976-0.8499j], + [-0.2109+1.9913j, -0.9607-0.6123j]]) + >>> x.view(torch.cfloat).size() + torch.Size([4, 2]) + + >>> x.view(torch.uint8) + tensor([[ 0, 202, 154, 59, 182, 243, 253, 188, 185, 252, 191, 63, 240, 22, + 8, 191], + [227, 165, 27, 190, 128, 72, 63, 63, 146, 203, 15, 63, 22, 106, + 93, 191], + [205, 59, 30, 192, 112, 206, 8, 189, 7, 95, 152, 190, 12, 147, + 89, 191], + [ 43, 246, 87, 190, 235, 226, 254, 63, 111, 240, 117, 191, 177, 191, + 28, 191]], dtype=torch.uint8) + >>> x.view(torch.uint8).size() + torch.Size([4, 16]) +""", +) + +add_docstr_all( + "view_as", + r""" +view_as(other) -> Tensor + +View this tensor as the same size as :attr:`other`. +``self.view_as(other)`` is equivalent to ``self.view(other.size())``. + +Please see :meth:`~Tensor.view` for more information about ``view``. + +Args: + other (:class:`torch.Tensor`): The result tensor has the same size + as :attr:`other`. +""", +) + +add_docstr_all( + "expand", + r""" +expand(*sizes) -> Tensor + +Returns a new view of the :attr:`self` tensor with singleton dimensions expanded +to a larger size. + +Passing -1 as the size for a dimension means not changing the size of +that dimension. + +Tensor can be also expanded to a larger number of dimensions, and the +new ones will be appended at the front. For the new dimensions, the +size cannot be set to -1. + +Expanding a tensor does not allocate new memory, but only creates a +new view on the existing tensor where a dimension of size one is +expanded to a larger size by setting the ``stride`` to 0. Any dimension +of size 1 can be expanded to an arbitrary value without allocating new +memory. + +Args: + *sizes (torch.Size or int...): the desired expanded size + +.. warning:: + + More than one element of an expanded tensor may refer to a single + memory location. As a result, in-place operations (especially ones that + are vectorized) may result in incorrect behavior. If you need to write + to the tensors, please clone them first. + +Example:: + + >>> x = torch.tensor([[1], [2], [3]]) + >>> x.size() + torch.Size([3, 1]) + >>> x.expand(3, 4) + tensor([[ 1, 1, 1, 1], + [ 2, 2, 2, 2], + [ 3, 3, 3, 3]]) + >>> x.expand(-1, 4) # -1 means not changing the size of that dimension + tensor([[ 1, 1, 1, 1], + [ 2, 2, 2, 2], + [ 3, 3, 3, 3]]) +""", +) + +add_docstr_all( + "expand_as", + r""" +expand_as(other) -> Tensor + +Expand this tensor to the same size as :attr:`other`. +``self.expand_as(other)`` is equivalent to ``self.expand(other.size())``. + +Please see :meth:`~Tensor.expand` for more information about ``expand``. + +Args: + other (:class:`torch.Tensor`): The result tensor has the same size + as :attr:`other`. +""", +) + +add_docstr_all( + "sum_to_size", + r""" +sum_to_size(*size) -> Tensor + +Sum ``this`` tensor to :attr:`size`. +:attr:`size` must be broadcastable to ``this`` tensor size. + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. +""", +) + + +add_docstr_all( + "zero_", + r""" +zero_() -> Tensor + +Fills :attr:`self` tensor with zeros. +""", +) + +add_docstr_all( + "matmul", + r""" +matmul(tensor2) -> Tensor + +See :func:`torch.matmul` +""", +) + +add_docstr_all( + "chunk", + r""" +chunk(chunks, dim=0) -> List of Tensors + +See :func:`torch.chunk` +""", +) + +add_docstr_all( + "unsafe_chunk", + r""" +unsafe_chunk(chunks, dim=0) -> List of Tensors + +See :func:`torch.unsafe_chunk` +""", +) + +add_docstr_all( + "unsafe_split", + r""" +unsafe_split(split_size, dim=0) -> List of Tensors + +See :func:`torch.unsafe_split` +""", +) + +add_docstr_all( + "tensor_split", + r""" +tensor_split(indices_or_sections, dim=0) -> List of Tensors + +See :func:`torch.tensor_split` +""", +) + +add_docstr_all( + "hsplit", + r""" +hsplit(split_size_or_sections) -> List of Tensors + +See :func:`torch.hsplit` +""", +) + +add_docstr_all( + "vsplit", + r""" +vsplit(split_size_or_sections) -> List of Tensors + +See :func:`torch.vsplit` +""", +) + +add_docstr_all( + "dsplit", + r""" +dsplit(split_size_or_sections) -> List of Tensors + +See :func:`torch.dsplit` +""", +) + +add_docstr_all( + "stft", + r""" +stft(frame_length, hop, fft_size=None, return_onesided=True, window=None, + pad_end=0, align_to_window=None) -> Tensor + +See :func:`torch.stft` +""", +) + +add_docstr_all( + "istft", + r""" +istft(n_fft, hop_length=None, win_length=None, window=None, + center=True, normalized=False, onesided=True, length=None) -> Tensor + +See :func:`torch.istft` +""", +) + +add_docstr_all( + "det", + r""" +det() -> Tensor + +See :func:`torch.det` +""", +) + +add_docstr_all( + "where", + r""" +where(condition, y) -> Tensor + +``self.where(condition, y)`` is equivalent to ``torch.where(condition, self, y)``. +See :func:`torch.where` +""", +) + +add_docstr_all( + "logdet", + r""" +logdet() -> Tensor + +See :func:`torch.logdet` +""", +) + +add_docstr_all( + "slogdet", + r""" +slogdet() -> (Tensor, Tensor) + +See :func:`torch.slogdet` +""", +) + +add_docstr_all( + "unbind", + r""" +unbind(dim=0) -> seq + +See :func:`torch.unbind` +""", +) + +add_docstr_all( + "pin_memory", + r""" +pin_memory() -> Tensor + +Copies the tensor to pinned memory, if it's not already pinned. +By default, the device pinned memory on will be the current :ref:`accelerator`. +""", +) + +add_docstr_all( + "pinverse", + r""" +pinverse() -> Tensor + +See :func:`torch.pinverse` +""", +) + +add_docstr_all( + "index_add", + r""" +index_add(dim, index, source, *, alpha=1) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.index_add_`. +""", +) + +add_docstr_all( + "index_copy", + r""" +index_copy(dim, index, tensor2) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.index_copy_`. +""", +) + +add_docstr_all( + "index_fill", + r""" +index_fill(dim, index, value) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.index_fill_`. +""", +) + +add_docstr_all( + "scatter", + r""" +scatter(dim, index, src) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_` +""", +) + +add_docstr_all( + "scatter_add", + r""" +scatter_add(dim, index, src) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_add_` +""", +) + +add_docstr_all( + "scatter_reduce", + r""" +scatter_reduce(dim, index, src, reduce, *, include_self=True) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_reduce_` +""", +) + +add_docstr_all( + "masked_scatter", + r""" +masked_scatter(mask, tensor) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.masked_scatter_` + +.. note:: + + The inputs :attr:`self` and :attr:`mask` + :ref:`broadcast `. + +Example: + + >>> self = torch.tensor([0, 0, 0, 0, 0]) + >>> mask = torch.tensor( + ... [[0, 0, 0, 1, 1], [1, 1, 0, 1, 1]], + ... dtype=torch.bool, + ... ) + >>> source = torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]) + >>> self.masked_scatter(mask, source) + tensor([[0, 0, 0, 0, 1], + [2, 3, 0, 4, 5]]) + +""", +) + +add_docstr_all( + "xlogy", + r""" +xlogy(other) -> Tensor + +See :func:`torch.xlogy` +""", +) + +add_docstr_all( + "xlogy_", + r""" +xlogy_(other) -> Tensor + +In-place version of :meth:`~Tensor.xlogy` +""", +) + +add_docstr_all( + "masked_fill", + r""" +masked_fill(mask, value) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.masked_fill_` +""", +) + +add_docstr_all( + "grad", + r""" +This attribute is ``None`` by default and becomes a Tensor the first time a call to +:func:`backward` computes gradients for ``self``. +The attribute will then contain the gradients computed and future calls to +:func:`backward` will accumulate (add) gradients into it. +""", +) + +add_docstr_all( + "retain_grad", + r""" +retain_grad() -> None + +Enables this Tensor to have their :attr:`grad` populated during +:func:`backward`. This is a no-op for leaf tensors. +""", +) + +add_docstr_all( + "retains_grad", + r""" +Is ``True`` if this Tensor is non-leaf and its :attr:`grad` is enabled to be +populated during :func:`backward`, ``False`` otherwise. +""", +) + +add_docstr_all( + "requires_grad", + r""" +Is ``True`` if gradients need to be computed for this Tensor, ``False`` otherwise. + +.. note:: + + The fact that gradients need to be computed for a Tensor do not mean that the :attr:`grad` + attribute will be populated, see :attr:`is_leaf` for more details. + +""", +) + +add_docstr_all( + "is_leaf", + r""" +All Tensors that have :attr:`requires_grad` which is ``False`` will be leaf Tensors by convention. + +For Tensors that have :attr:`requires_grad` which is ``True``, they will be leaf Tensors if they were +created by the user. This means that they are not the result of an operation and so +:attr:`grad_fn` is None. + +Only leaf Tensors will have their :attr:`grad` populated during a call to :func:`backward`. +To get :attr:`grad` populated for non-leaf Tensors, you can use :func:`retain_grad`. + +Example:: + + >>> a = torch.rand(10, requires_grad=True) + >>> a.is_leaf + True + >>> b = torch.rand(10, requires_grad=True).cuda() + >>> b.is_leaf + False + # b was created by the operation that cast a cpu Tensor into a cuda Tensor + >>> c = torch.rand(10, requires_grad=True) + 2 + >>> c.is_leaf + False + # c was created by the addition operation + >>> d = torch.rand(10).cuda() + >>> d.is_leaf + True + # d does not require gradients and so has no operation creating it (that is tracked by the autograd engine) + >>> e = torch.rand(10).cuda().requires_grad_() + >>> e.is_leaf + True + # e requires gradients and has no operations creating it + >>> f = torch.rand(10, requires_grad=True, device="cuda") + >>> f.is_leaf + True + # f requires grad, has no operation creating it + + +""", +) + +add_docstr_all( + "names", + r""" +Stores names for each of this tensor's dimensions. + +``names[idx]`` corresponds to the name of tensor dimension ``idx``. +Names are either a string if the dimension is named or ``None`` if the +dimension is unnamed. + +Dimension names may contain characters or underscore. Furthermore, a dimension +name must be a valid Python variable name (i.e., does not start with underscore). + +Tensors may not have two named dimensions with the same name. + +.. warning:: + The named tensor API is experimental and subject to change. + +""", +) + +add_docstr_all( + "is_cuda", + r""" +Is ``True`` if the Tensor is stored on the GPU, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_cpu", + r""" +Is ``True`` if the Tensor is stored on the CPU, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_xla", + r""" +Is ``True`` if the Tensor is stored on an XLA device, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_ipu", + r""" +Is ``True`` if the Tensor is stored on the IPU, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_xpu", + r""" +Is ``True`` if the Tensor is stored on the XPU, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_quantized", + r""" +Is ``True`` if the Tensor is quantized, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_meta", + r""" +Is ``True`` if the Tensor is a meta tensor, ``False`` otherwise. Meta tensors +are like normal tensors, but they carry no data. +""", +) + +add_docstr_all( + "is_mps", + r""" +Is ``True`` if the Tensor is stored on the MPS device, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_sparse", + r""" +Is ``True`` if the Tensor uses sparse COO storage layout, ``False`` otherwise. +""", +) + +add_docstr_all( + "is_sparse_csr", + r""" +Is ``True`` if the Tensor uses sparse CSR storage layout, ``False`` otherwise. +""", +) + +add_docstr_all( + "device", + r""" +Is the :class:`torch.device` where this Tensor is. +""", +) + +add_docstr_all( + "ndim", + r""" +Alias for :meth:`~Tensor.dim()` +""", +) + +add_docstr_all( + "itemsize", + r""" +Alias for :meth:`~Tensor.element_size()` +""", +) + +add_docstr_all( + "nbytes", + r""" +Returns the number of bytes consumed by the "view" of elements of the Tensor +if the Tensor does not use sparse storage layout. +Defined to be :meth:`~Tensor.numel()` * :meth:`~Tensor.element_size()` +""", +) + +add_docstr_all( + "T", + r""" +Returns a view of this tensor with its dimensions reversed. + +If ``n`` is the number of dimensions in ``x``, +``x.T`` is equivalent to ``x.permute(n-1, n-2, ..., 0)``. + +.. warning:: + The use of :func:`Tensor.T` on tensors of dimension other than 2 to reverse their shape + is deprecated and it will throw an error in a future release. Consider :attr:`~.Tensor.mT` + to transpose batches of matrices or `x.permute(*torch.arange(x.ndim - 1, -1, -1))` to reverse + the dimensions of a tensor. +""", +) + +add_docstr_all( + "H", + r""" +Returns a view of a matrix (2-D tensor) conjugated and transposed. + +``x.H`` is equivalent to ``x.transpose(0, 1).conj()`` for complex matrices and +``x.transpose(0, 1)`` for real matrices. + +.. seealso:: + + :attr:`~.Tensor.mH`: An attribute that also works on batches of matrices. +""", +) + +add_docstr_all( + "mT", + r""" +Returns a view of this tensor with the last two dimensions transposed. + +``x.mT`` is equivalent to ``x.transpose(-2, -1)``. +""", +) + +add_docstr_all( + "mH", + r""" +Accessing this property is equivalent to calling :func:`adjoint`. +""", +) + +add_docstr_all( + "adjoint", + r""" +adjoint() -> Tensor + +Alias for :func:`adjoint` +""", +) + +add_docstr_all( + "real", + r""" +Returns a new tensor containing real values of the :attr:`self` tensor for a complex-valued input tensor. +The returned tensor and :attr:`self` share the same underlying storage. + +Returns :attr:`self` if :attr:`self` is a real-valued tensor tensor. + +Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.real + tensor([ 0.3100, -0.5445, -1.6492, -0.0638]) + +""", +) + +add_docstr_all( + "imag", + r""" +Returns a new tensor containing imaginary values of the :attr:`self` tensor. +The returned tensor and :attr:`self` share the same underlying storage. + +.. warning:: + :func:`imag` is only supported for tensors with complex dtypes. + +Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.imag + tensor([ 0.3553, -0.7896, -0.0633, -0.8119]) + +""", +) + +add_docstr_all( + "as_subclass", + r""" +as_subclass(cls) -> Tensor + +Makes a ``cls`` instance with the same data pointer as ``self``. Changes +in the output mirror changes in ``self``, and the output stays attached +to the autograd graph. ``cls`` must be a subclass of ``Tensor``. +""", +) + +add_docstr_all( + "crow_indices", + r""" +crow_indices() -> IntTensor + +Returns the tensor containing the compressed row indices of the :attr:`self` +tensor when :attr:`self` is a sparse CSR tensor of layout ``sparse_csr``. +The ``crow_indices`` tensor is strictly of shape (:attr:`self`.size(0) + 1) +and of type ``int32`` or ``int64``. When using MKL routines such as sparse +matrix multiplication, it is necessary to use ``int32`` indexing in order +to avoid downcasting and potentially losing information. + +Example:: + + >>> csr = torch.eye(5,5).to_sparse_csr() + >>> csr.crow_indices() + tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32) + +""", +) + +add_docstr_all( + "col_indices", + r""" +col_indices() -> IntTensor + +Returns the tensor containing the column indices of the :attr:`self` +tensor when :attr:`self` is a sparse CSR tensor of layout ``sparse_csr``. +The ``col_indices`` tensor is strictly of shape (:attr:`self`.nnz()) +and of type ``int32`` or ``int64``. When using MKL routines such as sparse +matrix multiplication, it is necessary to use ``int32`` indexing in order +to avoid downcasting and potentially losing information. + +Example:: + + >>> csr = torch.eye(5,5).to_sparse_csr() + >>> csr.col_indices() + tensor([0, 1, 2, 3, 4], dtype=torch.int32) + +""", +) + +add_docstr_all( + "to_padded_tensor", + r""" +to_padded_tensor(padding, output_size=None) -> Tensor +See :func:`to_padded_tensor` +""", +) diff --git a/venv/lib/python3.10/site-packages/torch/_tensor_str.py b/venv/lib/python3.10/site-packages/torch/_tensor_str.py new file mode 100644 index 0000000000000000000000000000000000000000..c9262e1b2ee0a53df695af3a9dfbe8301e8bbae2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_tensor_str.py @@ -0,0 +1,722 @@ +# mypy: allow-untyped-defs +import contextlib +import dataclasses +import math +import textwrap +from typing import Any, Optional + +import torch +from torch import inf + + +@dataclasses.dataclass +class __PrinterOptions: + precision: int = 4 + threshold: float = 1000 + edgeitems: int = 3 + linewidth: int = 80 + sci_mode: Optional[bool] = None + + +PRINT_OPTS = __PrinterOptions() + + +# We could use **kwargs, but this will give better docs +def set_printoptions( + precision=None, + threshold=None, + edgeitems=None, + linewidth=None, + profile=None, + sci_mode=None, +): + r"""Set options for printing. Items shamelessly taken from NumPy + + Args: + precision: Number of digits of precision for floating point output + (default = 4). + threshold: Total number of array elements which trigger summarization + rather than full `repr` (default = 1000). + edgeitems: Number of array items in summary at beginning and end of + each dimension (default = 3). + linewidth: The number of characters per line for the purpose of + inserting line breaks (default = 80). Thresholded matrices will + ignore this parameter. + profile: Sane defaults for pretty printing. Can override with any of + the above options. (any one of `default`, `short`, `full`) + sci_mode: Enable (True) or disable (False) scientific notation. If + None (default) is specified, the value is defined by + `torch._tensor_str._Formatter`. This value is automatically chosen + by the framework. + + Example:: + + >>> # Limit the precision of elements + >>> torch.set_printoptions(precision=2) + >>> torch.tensor([1.12345]) + tensor([1.12]) + >>> # Limit the number of elements shown + >>> torch.set_printoptions(threshold=5) + >>> torch.arange(10) + tensor([0, 1, 2, ..., 7, 8, 9]) + >>> # Restore defaults + >>> torch.set_printoptions(profile='default') + >>> torch.tensor([1.12345]) + tensor([1.1235]) + >>> torch.arange(10) + tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + """ + if profile is not None: + if profile == "default": + PRINT_OPTS.precision = 4 + PRINT_OPTS.threshold = 1000 + PRINT_OPTS.edgeitems = 3 + PRINT_OPTS.linewidth = 80 + elif profile == "short": + PRINT_OPTS.precision = 2 + PRINT_OPTS.threshold = 1000 + PRINT_OPTS.edgeitems = 2 + PRINT_OPTS.linewidth = 80 + elif profile == "full": + PRINT_OPTS.precision = 4 + PRINT_OPTS.threshold = inf + PRINT_OPTS.edgeitems = 3 + PRINT_OPTS.linewidth = 80 + + if precision is not None: + PRINT_OPTS.precision = precision + if threshold is not None: + PRINT_OPTS.threshold = threshold + if edgeitems is not None: + PRINT_OPTS.edgeitems = edgeitems + if linewidth is not None: + PRINT_OPTS.linewidth = linewidth + PRINT_OPTS.sci_mode = sci_mode + + +def get_printoptions() -> dict[str, Any]: + r"""Gets the current options for printing, as a dictionary that + can be passed as ``**kwargs`` to set_printoptions(). + """ + return dataclasses.asdict(PRINT_OPTS) + + +@contextlib.contextmanager +def printoptions(**kwargs): + r"""Context manager that temporarily changes the print options. Accepted + arguments are same as :func:`set_printoptions`.""" + old_kwargs = get_printoptions() + set_printoptions(**kwargs) + try: + yield + finally: + set_printoptions(**old_kwargs) + + +def tensor_totype(t): + dtype = ( + torch.float + if ( + t.is_mps + or (t.is_xpu and not torch.xpu.get_device_properties(t.device).has_fp64) + or t.is_maia + ) + else torch.double + ) + return t.to(dtype=dtype) + + +class _Formatter: + def __init__(self, tensor): + self.floating_dtype = tensor.dtype.is_floating_point + self.int_mode = True + self.sci_mode = False + self.max_width = 1 + + with torch.no_grad(): + tensor_view = tensor.reshape(-1) + + if not self.floating_dtype: + for value in tensor_view: + value_str = f"{value}" + self.max_width = max(self.max_width, len(value_str)) + + else: + if tensor.dtype == torch.float4_e2m1fn_x2: # type: ignore[attr-defined] + # torch.float4_e2m1fn_x2 is special and does not support the casts necessary + # to print it, we choose to display the uint8 representation here for + # convenience of being able to print a tensor. + # TODO(#146647): extend this to other dtypes without casts defined, such + # as the bits, uint1..7 and int1..7 dtypes. + tensor_view = tensor_view.view(torch.uint8) + + nonzero_finite_vals = torch.masked_select( + tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0) + ) + + if nonzero_finite_vals.numel() == 0: + # no valid number, do nothing + return + + if tensor.dtype == torch.float8_e8m0fnu: # type: ignore[attr-defined] + # float8_e8m0fnu is special and does not define arithmetic ops, + # and printing code further in this file assumes the existence + # of various arithmetic ops to figure out what to print. We hack + # and convert to float here to make printing work correctly. + # TODO(#113663): also add the other float8 dtypes here after arithmetic + # support for them is removed + nonzero_finite_vals = nonzero_finite_vals.float() + + # Convert to double (or float) for easy calculation. HalfTensor overflows with 1e8, and there's no div() on CPU. + nonzero_finite_abs = tensor_totype(nonzero_finite_vals.abs()) + nonzero_finite_min = tensor_totype(nonzero_finite_abs.min()) + nonzero_finite_max = tensor_totype(nonzero_finite_abs.max()) + + for value in nonzero_finite_vals: + if value != torch.ceil(value): + self.int_mode = False + break + + self.sci_mode = ( + nonzero_finite_max / nonzero_finite_min > 1000.0 + or nonzero_finite_max > 1.0e8 + or nonzero_finite_min < 1.0e-4 + if PRINT_OPTS.sci_mode is None + else PRINT_OPTS.sci_mode + ) + + if self.int_mode: + # in int_mode for floats, all numbers are integers, and we append a decimal to nonfinites + # to indicate that the tensor is of floating type. add 1 to the len to account for this. + if self.sci_mode: + for value in nonzero_finite_vals: + value_str = f"{{:.{PRINT_OPTS.precision}e}}".format(value) + self.max_width = max(self.max_width, len(value_str)) + else: + for value in nonzero_finite_vals: + value_str = f"{value:.0f}" + self.max_width = max(self.max_width, len(value_str) + 1) + else: + # Check if scientific representation should be used. + if self.sci_mode: + for value in nonzero_finite_vals: + value_str = f"{{:.{PRINT_OPTS.precision}e}}".format(value) + self.max_width = max(self.max_width, len(value_str)) + else: + for value in nonzero_finite_vals: + value_str = f"{{:.{PRINT_OPTS.precision}f}}".format(value) + self.max_width = max(self.max_width, len(value_str)) + + def width(self): + return self.max_width + + def format(self, value): + if self.floating_dtype: + if self.sci_mode: + ret = f"{{:{self.max_width}.{PRINT_OPTS.precision}e}}".format(value) + elif self.int_mode: + ret = f"{value:.0f}" + if not (math.isinf(value) or math.isnan(value)): + ret += "." + else: + ret = f"{{:.{PRINT_OPTS.precision}f}}".format(value) + else: + ret = f"{value}" + return (self.max_width - len(ret)) * " " + ret + + +def _scalar_str(self, formatter1, formatter2=None): + if formatter2 is not None: + real_str = _scalar_str(self.real, formatter1) + imag_str = (_scalar_str(self.imag, formatter2) + "j").lstrip() + # handles negative numbers, +0.0, -0.0 + if imag_str[0] == "+" or imag_str[0] == "-": + return real_str + imag_str + else: + return real_str + "+" + imag_str + else: + return formatter1.format(self.item()) + + +def _vector_str(self, indent, summarize, formatter1, formatter2=None): + # length includes spaces and comma between elements + element_length = formatter1.width() + 2 + if formatter2 is not None: + # width for imag_formatter + an extra j for complex + element_length += formatter2.width() + 1 + + elements_per_line = max( + 1, int(math.floor((PRINT_OPTS.linewidth - indent) / (element_length))) + ) + + def _val_formatter(val, formatter1=formatter1, formatter2=formatter2): + if formatter2 is not None: + real_str = formatter1.format(val.real) + imag_str = (formatter2.format(val.imag) + "j").lstrip() + # handles negative numbers, +0.0, -0.0 + if imag_str[0] == "+" or imag_str[0] == "-": + return real_str + imag_str + else: + return real_str + "+" + imag_str + else: + return formatter1.format(val) + + if self.dtype == torch.float4_e2m1fn_x2: # type: ignore[attr-defined] + # torch.float4_e2m1fn_x2 is special and does not support the casts necessary + # to print it, we choose to display the uint8 representation here for + # convenience of being able to print a tensor. + # TODO(#146647): extend this to other dtypes without casts defined, such + # as the bits, uint1..7 and int1..7 dtypes. + self = self.view(torch.uint8) + + if summarize and not PRINT_OPTS.edgeitems: + # Deal with edge case that negative zero is zero + data = ["..."] + elif summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems: + data = ( + [_val_formatter(val) for val in self[: PRINT_OPTS.edgeitems].tolist()] + + [" ..."] + + [_val_formatter(val) for val in self[-PRINT_OPTS.edgeitems :].tolist()] + ) + else: + data = [_val_formatter(val) for val in self.tolist()] + + data_lines = [ + data[i : i + elements_per_line] for i in range(0, len(data), elements_per_line) + ] + lines = [", ".join(line) for line in data_lines] + return "[" + ("," + "\n" + " " * (indent + 1)).join(lines) + "]" + + +# formatter2 is only used for printing complex tensors. +# For complex tensors, formatter1 and formatter2 are the formatters for tensor.real +# and tensor.imag respesectively +def _tensor_str_with_formatter(self, indent, summarize, formatter1, formatter2=None): + dim = self.dim() + + if dim == 0: + return _scalar_str(self, formatter1, formatter2) + + if dim == 1: + return _vector_str(self, indent, summarize, formatter1, formatter2) + + if summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems: + slices = ( + [ + _tensor_str_with_formatter( + self[i], indent + 1, summarize, formatter1, formatter2 + ) + for i in range(0, PRINT_OPTS.edgeitems) + ] + + ["..."] + + [ + _tensor_str_with_formatter( + self[i], indent + 1, summarize, formatter1, formatter2 + ) + for i in range(len(self) - PRINT_OPTS.edgeitems, len(self)) + ] + ) + else: + slices = [ + _tensor_str_with_formatter( + self[i], indent + 1, summarize, formatter1, formatter2 + ) + for i in range(0, self.size(0)) + ] + + tensor_str = ("," + "\n" * (dim - 1) + " " * (indent + 1)).join(slices) + return "[" + tensor_str + "]" + + +def _tensor_str(self, indent): + if self.numel() == 0: + return "[]" + + if self.has_names(): + # There are two main codepaths (possibly more) that tensor printing goes through: + # - tensor data can fit comfortably on screen + # - tensor data needs to be summarized + # Some of the codepaths don't fully support named tensors, so we send in + # an unnamed tensor to the formatting code as a workaround. + self = self.rename(None) + + summarize = self.numel() > PRINT_OPTS.threshold + + if self._is_zerotensor(): + self = self.clone() + + # handle the negative bit + if self.is_neg(): + self = self.resolve_neg() + + # TODO: Remove me when `masked_select` is implemented for FP8 + if self.dtype in [ + torch.float8_e5m2, + torch.float8_e5m2fnuz, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ]: + self = self.half() + + if self.dtype.is_complex: + # handle the conjugate bit + self = self.resolve_conj() + real_formatter = _Formatter( + get_summarized_data(self.real) if summarize else self.real + ) + imag_formatter = _Formatter( + get_summarized_data(self.imag) if summarize else self.imag + ) + return _tensor_str_with_formatter( + self, indent, summarize, real_formatter, imag_formatter + ) + else: + formatter = _Formatter(get_summarized_data(self) if summarize else self) + return _tensor_str_with_formatter(self, indent, summarize, formatter) + + +def _add_suffixes(tensor_str, suffixes, indent, force_newline): + tensor_strs = [tensor_str] + last_line_len = len(tensor_str) - tensor_str.rfind("\n") + 1 + for suffix in suffixes: + suffix_len = len(suffix) + if force_newline or last_line_len + suffix_len + 2 > PRINT_OPTS.linewidth: + tensor_strs.append(",\n" + " " * indent + suffix) + last_line_len = indent + suffix_len + force_newline = False + else: + tensor_strs.append(", " + suffix) + last_line_len += suffix_len + 2 + tensor_strs.append(")") + return "".join(tensor_strs) + + +def get_summarized_data(self): + dim = self.dim() + if dim == 0: + return self + if dim == 1: + if self.size(0) > 2 * PRINT_OPTS.edgeitems: + return torch.cat( + (self[: PRINT_OPTS.edgeitems], self[-PRINT_OPTS.edgeitems :]) + ) + else: + return self + if not PRINT_OPTS.edgeitems: + return self.new_empty([0] * self.dim()) + elif self.size(0) > 2 * PRINT_OPTS.edgeitems: + start = [self[i] for i in range(0, PRINT_OPTS.edgeitems)] + end = [self[i] for i in range(len(self) - PRINT_OPTS.edgeitems, len(self))] + return torch.stack([get_summarized_data(x) for x in (start + end)]) + else: + return torch.stack([get_summarized_data(x) for x in self]) + + +def _str_intern(inp, *, tensor_contents=None): + if torch._C._functorch.is_functorch_wrapped_tensor(inp): + return _functorch_wrapper_str_intern(inp, tensor_contents=tensor_contents) + is_plain_tensor = type(inp) is torch.Tensor or type(inp) is torch.nn.Parameter + if inp.is_nested: + prefix = "nested_tensor(" + elif is_plain_tensor: + prefix = "tensor(" + else: + prefix = f"{type(inp).__name__}(" + indent = len(prefix) + suffixes = [] + custom_contents_provided = tensor_contents is not None + if custom_contents_provided: + tensor_str = tensor_contents + + # This is used to extract the primal value and thus disable the forward AD + # within this function. + # TODO(albanD) This needs to be updated when more than one level is supported + self, tangent = torch.autograd.forward_ad.unpack_dual(inp) + + # Note [Print tensor device]: + # A general logic here is we only print device when it doesn't match + # the device specified in default tensor type. + # Currently torch.set_default_tensor_type() only supports CPU/CUDA, thus + # torch._C._get_default_device() only returns either cpu or cuda. + # In other cases, we don't have a way to set them as default yet, + # and we should always print out device for them. + if ( + self.device.type != torch._C._get_default_device() + or ( + self.device.type == "cuda" + and torch.cuda.current_device() != self.device.index + ) + or (self.device.type == "mps") + ): + suffixes.append("device='" + str(self.device) + "'") + + # Tensor printing performs tensor operations like slice, indexing, etc to make it in a + # representable format. These operations on ipu/xla/lazy/mtia tensor results in compilations. Hence, + # to avoid compilations, copying the tensor to cpu before printing. + if self.device.type in ["xla", "lazy", "ipu", "mtia"]: + self = self.to("cpu") + + # TODO: add an API to map real -> complex dtypes + _default_complex_dtype = ( + torch.cdouble if torch.get_default_dtype() == torch.double else torch.cfloat + ) + has_default_dtype = self.dtype in ( + torch.get_default_dtype(), + _default_complex_dtype, + torch.int64, + torch.bool, + ) + if self.is_sparse: + suffixes.append("size=" + str(tuple(self.shape))) + from torch._subclasses.fake_tensor import FakeTensor + + is_meta = self.is_meta or isinstance(self, FakeTensor) + if not is_meta: + suffixes.append("nnz=" + str(self._nnz())) + if not has_default_dtype: + suffixes.append("dtype=" + str(self.dtype)) + if not custom_contents_provided: + indices_prefix = "indices=tensor(" + indices = self._indices().detach() + if is_meta: + indices_str = "..." + else: + indices_str = _tensor_str(indices, indent + len(indices_prefix)) + if is_meta or indices.numel() == 0: + indices_str += ", size=" + str(tuple(indices.shape)) + values_prefix = "values=tensor(" + values = self._values().detach() + if is_meta: + values_str = "..." + else: + values_str = _tensor_str(values, indent + len(values_prefix)) + if is_meta or values.numel() == 0: + values_str += ", size=" + str(tuple(values.shape)) + tensor_str = ( + indices_prefix + + indices_str + + "),\n" + + " " * indent + + values_prefix + + values_str + + ")" + ) + elif self.layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + from torch._subclasses.fake_tensor import FakeTensor + + suffixes.append("size=" + str(tuple(self.shape))) + is_meta = self.is_meta or isinstance(self, FakeTensor) + if not is_meta: + suffixes.append("nnz=" + str(self._nnz())) + if not has_default_dtype: + suffixes.append("dtype=" + str(self.dtype)) + if not custom_contents_provided: + compressed_indices_method, plain_indices_method = { + torch.sparse_csr: (torch.Tensor.crow_indices, torch.Tensor.col_indices), + torch.sparse_csc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices), + torch.sparse_bsr: (torch.Tensor.crow_indices, torch.Tensor.col_indices), + torch.sparse_bsc: (torch.Tensor.ccol_indices, torch.Tensor.row_indices), + }[self.layout] + if self.layout in {torch.sparse_csr, torch.sparse_bsr}: + cdimname, pdimname = "row", "column" + else: + cdimname, pdimname = "column", "row" + compressed_indices_prefix = f"c{cdimname[:3]}_indices=tensor(" + compressed_indices = compressed_indices_method(self).detach() + if is_meta: + compressed_indices_str = "..." + else: + compressed_indices_str = _tensor_str( + compressed_indices, indent + len(compressed_indices_prefix) + ) + if compressed_indices.numel() == 0 or is_meta: + compressed_indices_str += ", size=" + str( + tuple(compressed_indices.shape) + ) + plain_indices_prefix = f"{pdimname[:3]}_indices=tensor(" + plain_indices = plain_indices_method(self).detach() + if is_meta: + plain_indices_str = "..." + else: + plain_indices_str = _tensor_str( + plain_indices, indent + len(plain_indices_prefix) + ) + if plain_indices.numel() == 0 or is_meta: + plain_indices_str += ", size=" + str(tuple(plain_indices.shape)) + values_prefix = "values=tensor(" + values = self.values().detach() + if is_meta: + values_str = "..." + else: + values_str = _tensor_str(values, indent + len(values_prefix)) + if values.numel() == 0 or is_meta: + values_str += ", size=" + str(tuple(values.shape)) + tensor_str = ( + compressed_indices_prefix + + compressed_indices_str + + "),\n" + + " " * indent + + plain_indices_prefix + + plain_indices_str + + "),\n" + + " " * indent + + values_prefix + + values_str + + ")" + ) + elif self.is_quantized: + suffixes.append("size=" + str(tuple(self.shape))) + if not has_default_dtype: + suffixes.append("dtype=" + str(self.dtype)) + suffixes.append("quantization_scheme=" + str(self.qscheme())) + if ( + self.qscheme() == torch.per_tensor_affine + or self.qscheme() == torch.per_tensor_symmetric + ): + suffixes.append("scale=" + str(self.q_scale())) + suffixes.append("zero_point=" + str(self.q_zero_point())) + elif ( + self.qscheme() == torch.per_channel_affine + or self.qscheme() == torch.per_channel_symmetric + or self.qscheme() == torch.per_channel_affine_float_qparams + ): + suffixes.append("scale=" + str(self.q_per_channel_scales())) + suffixes.append("zero_point=" + str(self.q_per_channel_zero_points())) + suffixes.append("axis=" + str(self.q_per_channel_axis())) + if not custom_contents_provided: + tensor_str = _tensor_str(self.dequantize(), indent) + elif self.is_nested: + if not custom_contents_provided: + + def indented_str(s, indent): + return "\n".join(f" {line}" for line in s.split("\n")) + + strs = ",\n".join( + indented_str(str(t), indent + 1) + for t in torch.ops.aten.unbind.int(self, 0) + ) + tensor_str = f"[\n{strs}\n]" + elif torch._is_functional_tensor(self): + prefix = "_to_functional_tensor(" + tensor_str = repr(torch._from_functional_tensor(self)) + else: + # Circular import problem, so we import it here + from torch._subclasses.fake_tensor import FakeTensor + + if self.is_meta or isinstance(self, FakeTensor): + suffixes.append("size=" + str(tuple(self.shape))) + if self.dtype != torch.get_default_dtype(): + suffixes.append("dtype=" + str(self.dtype)) + # TODO: This implies that ellipses is valid syntax for allocating + # a meta tensor or FakeTensor, which it could be, but it isn't right now + if not custom_contents_provided: + tensor_str = "..." + else: + if self.numel() == 0 and not self.is_sparse: + # Explicitly print the shape if it is not (0,), to match NumPy behavior + if self.dim() != 1: + suffixes.append("size=" + str(tuple(self.shape))) + + # In an empty tensor, there are no elements to infer if the dtype + # should be int64, so it must be shown explicitly. + if self.dtype != torch.get_default_dtype(): + suffixes.append("dtype=" + str(self.dtype)) + if not custom_contents_provided: + tensor_str = "[]" + else: + if not PRINT_OPTS.edgeitems: + suffixes.append("size=" + str(tuple(self.shape))) + + if not has_default_dtype: + suffixes.append("dtype=" + str(self.dtype)) + + if not custom_contents_provided: + if self.layout != torch.strided: + tensor_str = _tensor_str(self.to_dense(), indent) + else: + tensor_str = _tensor_str(self, indent) + + if self.layout != torch.strided: + suffixes.append("layout=" + str(self.layout)) + + # Use inp here to get the original grad_fn and not the one generated by the forward grad + # unpacking. + grad_fn_name = None + try: + grad_fn = inp.grad_fn + except RuntimeError: + # Accessing the grad_fn calls rebasing logic which would cause an error + # if that tensor is a view created in no-grad mode modified in-place in + # no-grad mode. See: https://github.com/pytorch/pytorch/issues/99968 + grad_fn_name = "Invalid" + + if grad_fn_name is None and grad_fn is not None: # type: ignore[possibly-undefined] + grad_fn_name = type(grad_fn).__name__ + if grad_fn_name == "CppFunction": + grad_fn_name = grad_fn.name().rsplit("::", 1)[-1] + + if grad_fn_name is not None: + suffixes.append(f"grad_fn=<{grad_fn_name}>") + elif inp.requires_grad: + suffixes.append("requires_grad=True") + + if self.has_names(): + suffixes.append(f"names={self.names}") + + if tangent is not None: + suffixes.append(f"tangent={tangent}") + + string_repr = _add_suffixes( + prefix + tensor_str, # type: ignore[possibly-undefined] + suffixes, + indent, + force_newline=self.is_sparse, + ) + + # Check if this instance is flagged as a parameter and change the repr accordingly. + # Unfortunately, this function has to be aware of this detail. + # NB: This is currently skipped for plain tensor parameters to maintain BC. In the future, + # this should be done for those as well to produce a valid repr. + if isinstance(self, torch.nn.Parameter) and not is_plain_tensor: + string_repr = f"Parameter({string_repr})" + + return string_repr + + +def _functorch_wrapper_str_intern(tensor, *, tensor_contents=None): + level = torch._C._functorch.maybe_get_level(tensor) + assert level != -1 + + if torch._C._functorch.is_functionaltensor(tensor): + # Since we're unwrapping the FunctionalTensorWrapper, we need to make sure + # that it's up to date first + torch._sync(tensor) + + value = torch._C._functorch.get_unwrapped(tensor) + value_repr = repr(value) + + indented_value_repr = textwrap.indent(value_repr, " " * 4) + if torch._C._functorch.is_batchedtensor(tensor): + bdim = torch._C._functorch.maybe_get_bdim(tensor) + assert bdim != -1 + return ( + f"BatchedTensor(lvl={level}, bdim={bdim}, value=\n{indented_value_repr}\n)" + ) + if torch._C._functorch.is_gradtrackingtensor(tensor): + return f"GradTrackingTensor(lvl={level}, value=\n{indented_value_repr}\n)" + if torch._C._functorch.is_functionaltensor(tensor): + return f"FunctionalTensor(lvl={level}, value=\\\n{value_repr})" + + raise ValueError("We don't know how to print this, please file us an issue") + + +def _str(self, *, tensor_contents=None): + with torch.no_grad(), torch.utils._python_dispatch._disable_current_modes(): + guard = torch._C._DisableFuncTorch() # noqa: F841 + return _str_intern(self, tensor_contents=tensor_contents) diff --git a/venv/lib/python3.10/site-packages/torch/_thread_safe_fork.py b/venv/lib/python3.10/site-packages/torch/_thread_safe_fork.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/torch/_torch_docs.py b/venv/lib/python3.10/site-packages/torch/_torch_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..68c3fe31c5bf037b46f9d771bc2fc359ba621843 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_torch_docs.py @@ -0,0 +1,14383 @@ +# mypy: allow-untyped-defs +"""Adds docstrings to functions defined in the torch._C module.""" + +import re + +import torch._C +from torch._C import _add_docstr as add_docstr + + +def parse_kwargs(desc): + r"""Map a description of args to a dictionary of {argname: description}. + + Input: + (' weight (Tensor): a weight tensor\n' + + ' Some optional description') + Output: { + 'weight': \ + 'weight (Tensor): a weight tensor\n Some optional description' + } + """ + # Split on exactly 4 spaces after a newline + regx = re.compile(r"\n\s{4}(?!\s)") + kwargs = [section.strip() for section in regx.split(desc)] + kwargs = [section for section in kwargs if len(section) > 0] + return {desc.split(" ")[0]: desc for desc in kwargs} + + +def merge_dicts(*dicts): + """Merge dictionaries into a single dictionary.""" + return {x: d[x] for d in dicts for x in d} + + +common_args = parse_kwargs( + """ + input (Tensor): the input tensor. + generator (:class:`torch.Generator`, optional): a pseudorandom number generator for sampling + out (Tensor, optional): the output tensor. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned tensor. Default: ``torch.preserve_format``. +""" +) + +reduceops_common_args = merge_dicts( + common_args, + parse_kwargs( + """ + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is casted to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + keepdim (bool): whether the output tensor has :attr:`dim` retained or not. +""" + ), + { + "opt_keepdim": """ + keepdim (bool, optional): whether the output tensor has :attr:`dim` retained or not. Default: ``False``. +""" + }, +) + +multi_dim_common = merge_dicts( + reduceops_common_args, + parse_kwargs( + """ + dim (int or tuple of ints): the dimension or dimensions to reduce. +""" + ), + { + "keepdim_details": """ +If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension(s) :attr:`dim` where it is of size 1. +Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in the +output tensor having 1 (or ``len(dim)``) fewer dimension(s). +""" + }, + { + "opt_dim": """ + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. +""" + }, + { + "opt_dim_all_reduce": """ + dim (int or tuple of ints, optional): the dimension or dimensions to reduce. + If ``None``, all dimensions are reduced. +""" + }, +) + +single_dim_common = merge_dicts( + reduceops_common_args, + parse_kwargs( + """ + dim (int): the dimension to reduce. +""" + ), + { + "opt_dim": """ + dim (int, optional): the dimension to reduce. +""" + }, + { + "opt_dim_all_reduce": """ + dim (int, optional): the dimension to reduce. + If ``None``, all dimensions are reduced. +""" + }, + { + "opt_dim_without_none": """ + dim (int, optional): the dimension to reduce. If omitted, all dimensions are reduced. Explicit ``None`` is not supported. +""" + }, + { + "keepdim_details": """If :attr:`keepdim` is ``True``, the output tensor is of the same size +as :attr:`input` except in the dimension :attr:`dim` where it is of size 1. +Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in +the output tensor having 1 fewer dimension than :attr:`input`.""" + }, +) + +factory_common_args = merge_dicts( + common_args, + parse_kwargs( + """ + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default (see :func:`torch.set_default_dtype`). + layout (:class:`torch.layout`, optional): the desired layout of returned Tensor. + Default: ``torch.strided``. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.contiguous_format``. + check_invariants (bool, optional): If sparse tensor invariants are checked. + Default: as returned by :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled`, + initially False. +""" + ), + { + "sparse_factory_device_note": """\ +.. note:: + + If the ``device`` argument is not specified the device of the given + :attr:`values` and indices tensor(s) must match. If, however, the + argument is specified the input Tensors will be converted to the + given device and in turn determine the device of the constructed + sparse tensor.""" + }, +) + +factory_like_common_args = parse_kwargs( + """ + input (Tensor): the size of :attr:`input` will determine size of the output tensor. + layout (:class:`torch.layout`, optional): the desired layout of returned tensor. + Default: if ``None``, defaults to the layout of :attr:`input`. + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: if ``None``, defaults to the dtype of :attr:`input`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, defaults to the device of :attr:`input`. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. + memory_format (:class:`torch.memory_format`, optional): the desired memory format of + returned Tensor. Default: ``torch.preserve_format``. +""" +) + +factory_data_common_args = parse_kwargs( + """ + data (array_like): Initial data for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if ``None``, infers data type from :attr:`data`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + requires_grad (bool, optional): If autograd should record operations on the + returned tensor. Default: ``False``. + pin_memory (bool, optional): If set, returned tensor would be allocated in + the pinned memory. Works only for CPU tensors. Default: ``False``. +""" +) + +tf32_notes = { + "tf32_note": """This operator supports :ref:`TensorFloat32`.""" +} + +rocm_fp16_notes = { + "rocm_fp16_note": """On certain ROCm devices, when using float16 inputs this module will use \ +:ref:`different precision` for backward.""" +} + +reproducibility_notes: dict[str, str] = { + "forward_reproducibility_note": """This operation may behave nondeterministically when given tensors on \ +a CUDA device. See :doc:`/notes/randomness` for more information.""", + "backward_reproducibility_note": """This operation may produce nondeterministic gradients when given tensors on \ +a CUDA device. See :doc:`/notes/randomness` for more information.""", + "cudnn_reproducibility_note": """In some circumstances when given tensors on a CUDA device \ +and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is \ +undesirable, you can try to make the operation deterministic (potentially at \ +a performance cost) by setting ``torch.backends.cudnn.deterministic = True``. \ +See :doc:`/notes/randomness` for more information.""", +} + +sparse_support_notes = { + "sparse_beta_warning": """ +.. warning:: + Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, + or may not have autograd support. If you notice missing functionality please + open a feature request.""", +} + +add_docstr( + torch.abs, + r""" +abs(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Computes the absolute value of each element in :attr:`input`. + +.. math:: + \text{out}_{i} = |\text{input}_{i}| +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.abs(torch.tensor([-1, -2, 3])) + tensor([ 1, 2, 3]) +""".format(**common_args), +) + +add_docstr( + torch.absolute, + r""" +absolute(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.abs` +""", +) + +add_docstr( + torch.acos, + r""" +acos(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Computes the inverse cosine of each element in :attr:`input`. + +.. math:: + \text{out}_{i} = \cos^{-1}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.3348, -0.5889, 0.2005, -0.1584]) + >>> torch.acos(a) + tensor([ 1.2294, 2.2004, 1.3690, 1.7298]) +""".format(**common_args), +) + +add_docstr( + torch.arccos, + r""" +arccos(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.acos`. +""", +) + +add_docstr( + torch.acosh, + r""" +acosh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the inverse hyperbolic cosine of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \cosh^{-1}(\text{input}_{i}) + +Note: + The domain of the inverse hyperbolic cosine is `[1, inf)` and values outside this range + will be mapped to ``NaN``, except for `+ INF` for which the output is mapped to `+ INF`. +""" + + r""" +Args: + {input} + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.randn(4).uniform_(1, 2) + >>> a + tensor([ 1.3192, 1.9915, 1.9674, 1.7151 ]) + >>> torch.acosh(a) + tensor([ 0.7791, 1.3120, 1.2979, 1.1341 ]) +""".format(**common_args), +) + +add_docstr( + torch.arccosh, + r""" +arccosh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.acosh`. +""", +) + +add_docstr( + torch.index_add, + r""" +index_add(input: Tensor, dim: int, index: Tensor, source: Tensor, *, alpha: Union[Number, _complex] = 1, out: Optional[Tensor]) -> Tensor # noqa: B950 + +See :meth:`~Tensor.index_add_` for function description. +""", +) + +add_docstr( + torch.index_copy, + r""" +index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor, *, out: Optional[Tensor]) -> Tensor + +See :meth:`~Tensor.index_add_` for function description. +""", +) + +add_docstr( + torch.index_reduce, + r""" +index_reduce(input: Tensor, dim: int, index: Tensor, source: Tensor, reduce: str, *, include_self: bool = True, out: Optional[Tensor]) -> Tensor # noqa: B950 + +See :meth:`~Tensor.index_reduce_` for function description. +""", +) + +add_docstr( + torch.add, + r""" +add(input, other, *, alpha=1, out=None) -> Tensor + +Adds :attr:`other`, scaled by :attr:`alpha`, to :attr:`input`. + +.. math:: + \text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i +""" + + r""" + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer, float, and complex inputs. + +Args: + {input} + other (Tensor or Number): the tensor or number to add to :attr:`input`. + +Keyword arguments: + alpha (Number): the multiplier for :attr:`other`. + {out} + +Examples:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.0202, 1.0985, 1.3506, -0.6056]) + >>> torch.add(a, 20) + tensor([ 20.0202, 21.0985, 21.3506, 19.3944]) + + >>> b = torch.randn(4) + >>> b + tensor([-0.9732, -0.3497, 0.6245, 0.4022]) + >>> c = torch.randn(4, 1) + >>> c + tensor([[ 0.3743], + [-1.7724], + [-0.5811], + [-0.8017]]) + >>> torch.add(b, c, alpha=10) + tensor([[ 2.7695, 3.3930, 4.3672, 4.1450], + [-18.6971, -18.0736, -17.0994, -17.3216], + [ -6.7845, -6.1610, -5.1868, -5.4090], + [ -8.9902, -8.3667, -7.3925, -7.6147]]) +""".format(**common_args), +) + +add_docstr( + torch.addbmm, + r""" +addbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None) -> Tensor + +Performs a batch matrix-matrix product of matrices stored +in :attr:`batch1` and :attr:`batch2`, +with a reduced add step (all matrix multiplications get accumulated +along the first dimension). +:attr:`input` is added to the final result. + +:attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the +same number of matrices. + +If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a +:math:`(b \times m \times p)` tensor, :attr:`input` must be +:ref:`broadcastable ` with a :math:`(n \times p)` tensor +and :attr:`out` will be a :math:`(n \times p)` tensor. + +.. math:: + out = \beta\ \text{input} + \alpha\ (\sum_{i=0}^{b-1} \text{batch1}_i \mathbin{@} \text{batch2}_i) + +If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in +it will not be propagated. +""" + + r""" +For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and :attr:`alpha` +must be real numbers, otherwise they should be integers. + +{tf32_note} + +{rocm_fp16_note} + +Args: + input (Tensor): matrix to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for `batch1 @ batch2` (:math:`\alpha`) + {out} + +Example:: + + >>> M = torch.randn(3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.addbmm(M, batch1, batch2) + tensor([[ 6.6311, 0.0503, 6.9768, -12.0362, -2.1653], + [ -4.8185, -1.4255, -6.6760, 8.9453, 2.5743], + [ -3.8202, 4.3691, 1.0943, -1.1109, 5.4730]]) +""".format(**common_args, **tf32_notes, **rocm_fp16_notes), +) + +add_docstr( + torch.addcdiv, + r""" +addcdiv(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + +Performs the element-wise division of :attr:`tensor1` by :attr:`tensor2`, +multiplies the result by the scalar :attr:`value` and adds it to :attr:`input`. + +.. warning:: + Integer division with addcdiv is no longer supported, and in a future + release addcdiv will perform a true division of tensor1 and tensor2. + The historic addcdiv behavior can be implemented as + (input + value * torch.trunc(tensor1 / tensor2)).to(input.dtype) + for integer inputs and as (input + value * tensor1 / tensor2) for float inputs. + The future addcdiv behavior is just the latter implementation: + (input + value * tensor1 / tensor2), for all dtypes. + +.. math:: + \text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i} +""" + + r""" + +The shapes of :attr:`input`, :attr:`tensor1`, and :attr:`tensor2` must be +:ref:`broadcastable `. + +For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be +a real number, otherwise an integer. + +Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the numerator tensor + tensor2 (Tensor): the denominator tensor + +Keyword args: + value (Number, optional): multiplier for :math:`\text{{tensor1}} / \text{{tensor2}}` + {out} + +Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcdiv(t, t1, t2, value=0.1) + tensor([[-0.2312, -3.6496, 0.1312], + [-1.0428, 3.4292, -0.1030], + [-0.5369, -0.9829, 0.0430]]) +""".format(**common_args), +) + +add_docstr( + torch.addcmul, + r""" +addcmul(input, tensor1, tensor2, *, value=1, out=None) -> Tensor + +Performs the element-wise multiplication of :attr:`tensor1` +by :attr:`tensor2`, multiplies the result by the scalar :attr:`value` +and adds it to :attr:`input`. + +.. math:: + \text{out}_i = \text{input}_i + \text{value} \times \text{tensor1}_i \times \text{tensor2}_i +""" + + r""" +The shapes of :attr:`tensor`, :attr:`tensor1`, and :attr:`tensor2` must be +:ref:`broadcastable `. + +For inputs of type `FloatTensor` or `DoubleTensor`, :attr:`value` must be +a real number, otherwise an integer. + +Args: + input (Tensor): the tensor to be added + tensor1 (Tensor): the tensor to be multiplied + tensor2 (Tensor): the tensor to be multiplied + +Keyword args: + value (Number, optional): multiplier for :math:`tensor1 .* tensor2` + {out} + +Example:: + + >>> t = torch.randn(1, 3) + >>> t1 = torch.randn(3, 1) + >>> t2 = torch.randn(1, 3) + >>> torch.addcmul(t, t1, t2, value=0.1) + tensor([[-0.8635, -0.6391, 1.6174], + [-0.7617, -0.5879, 1.7388], + [-0.8353, -0.6249, 1.6511]]) +""".format(**common_args), +) + +add_docstr( + torch.addmm, + r""" +addmm(input, mat1, mat2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + +Performs a matrix multiplication of the matrices :attr:`mat1` and :attr:`mat2`. +The matrix :attr:`input` is added to the final result. + +If :attr:`mat1` is a :math:`(n \times m)` tensor, :attr:`mat2` is a +:math:`(m \times p)` tensor, then :attr:`input` must be +:ref:`broadcastable ` with a :math:`(n \times p)` tensor +and :attr:`out` will be a :math:`(n \times p)` tensor. + +:attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between +:attr:`mat1` and :attr:`mat2` and the added matrix :attr:`input` respectively. + +.. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat1}_i \mathbin{@} \text{mat2}_i) + +If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in +it will not be propagated. +""" + + r""" +For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and +:attr:`alpha` must be real numbers, otherwise they should be integers. + +This operation has support for arguments with :ref:`sparse layouts`. If +:attr:`input` is sparse the result will have the same layout and if :attr:`out` +is provided it must have the same layout as :attr:`input`. + +{sparse_beta_warning} + +{tf32_note} + +{rocm_fp16_note} + +Args: + input (Tensor): matrix to be added + mat1 (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + {out} + +Example:: + + >>> M = torch.randn(2, 3) + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.addmm(M, mat1, mat2) + tensor([[-4.8716, 1.4671, -1.3746], + [ 0.7573, -3.9555, -2.8681]]) +""".format(**common_args, **tf32_notes, **rocm_fp16_notes, **sparse_support_notes), +) + +add_docstr( + torch.adjoint, + r""" +adjoint(input: Tensor) -> Tensor +Returns a view of the tensor conjugated and with the last two dimensions transposed. + +``x.adjoint()`` is equivalent to ``x.transpose(-2, -1).conj()`` for complex tensors and +to ``x.transpose(-2, -1)`` for real tensors. + +Args: + {input} + +Example:: + + >>> x = torch.arange(4, dtype=torch.float) + >>> A = torch.complex(x, x).reshape(2, 2) + >>> A + tensor([[0.+0.j, 1.+1.j], + [2.+2.j, 3.+3.j]]) + >>> A.adjoint() + tensor([[0.-0.j, 2.-2.j], + [1.-1.j, 3.-3.j]]) + >>> (A.adjoint() == A.mH).all() + tensor(True) +""", +) + +add_docstr( + torch.sspaddmm, + r""" +sspaddmm(input, mat1, mat2, *, beta=1, alpha=1, out=None) -> Tensor + +Matrix multiplies a sparse tensor :attr:`mat1` with a dense tensor +:attr:`mat2`, then adds the sparse tensor :attr:`input` to the result. + +Note: This function is equivalent to :func:`torch.addmm`, except +:attr:`input` and :attr:`mat1` are sparse. + +Args: + input (Tensor): a sparse matrix to be added + mat1 (Tensor): a sparse matrix to be matrix multiplied + mat2 (Tensor): a dense matrix to be matrix multiplied + +Keyword args: + beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + {out} +""".format(**common_args), +) + +add_docstr( + torch.smm, + r""" +smm(input, mat) -> Tensor + +Performs a matrix multiplication of the sparse matrix :attr:`input` +with the dense matrix :attr:`mat`. + +Args: + input (Tensor): a sparse matrix to be matrix multiplied + mat (Tensor): a dense matrix to be matrix multiplied +""", +) + +add_docstr( + torch.addmv, + r""" +addmv(input, mat, vec, *, beta=1, alpha=1, out=None) -> Tensor + +Performs a matrix-vector product of the matrix :attr:`mat` and +the vector :attr:`vec`. +The vector :attr:`input` is added to the final result. + +If :attr:`mat` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of +size `m`, then :attr:`input` must be +:ref:`broadcastable ` with a 1-D tensor of size `n` and +:attr:`out` will be 1-D tensor of size `n`. + +:attr:`alpha` and :attr:`beta` are scaling factors on matrix-vector product between +:attr:`mat` and :attr:`vec` and the added tensor :attr:`input` respectively. + +.. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{mat} \mathbin{@} \text{vec}) + +If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in +it will not be propagated. +""" + + r""" +For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and +:attr:`alpha` must be real numbers, otherwise they should be integers. + +Args: + input (Tensor): vector to be added + mat (Tensor): matrix to be matrix multiplied + vec (Tensor): vector to be matrix multiplied + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat @ vec` (:math:`\alpha`) + {out} + +Example:: + + >>> M = torch.randn(2) + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.addmv(M, mat, vec) + tensor([-0.3768, -5.5565]) +""".format(**common_args), +) + +add_docstr( + torch.addr, + r""" +addr(input, vec1, vec2, *, beta=1, alpha=1, out=None) -> Tensor + +Performs the outer-product of vectors :attr:`vec1` and :attr:`vec2` +and adds it to the matrix :attr:`input`. + +Optional values :attr:`beta` and :attr:`alpha` are scaling factors on the +outer product between :attr:`vec1` and :attr:`vec2` and the added matrix +:attr:`input` respectively. + +.. math:: + \text{out} = \beta\ \text{input} + \alpha\ (\text{vec1} \otimes \text{vec2}) + +If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in +it will not be propagated. +""" + + r""" +If :attr:`vec1` is a vector of size `n` and :attr:`vec2` is a vector +of size `m`, then :attr:`input` must be +:ref:`broadcastable ` with a matrix of size +:math:`(n \times m)` and :attr:`out` will be a matrix of size +:math:`(n \times m)`. + +Args: + input (Tensor): matrix to be added + vec1 (Tensor): the first vector of the outer product + vec2 (Tensor): the second vector of the outer product + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{{vec1}} \otimes \text{{vec2}}` (:math:`\alpha`) + {out} + +Example:: + + >>> vec1 = torch.arange(1., 4.) + >>> vec2 = torch.arange(1., 3.) + >>> M = torch.zeros(3, 2) + >>> torch.addr(M, vec1, vec2) + tensor([[ 1., 2.], + [ 2., 4.], + [ 3., 6.]]) +""".format(**common_args), +) + +add_docstr( + torch.allclose, + r""" +allclose(input: Tensor, other: Tensor, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) -> bool + +This function checks if :attr:`input` and :attr:`other` satisfy the condition: + +.. math:: + \lvert \text{input}_i - \text{other}_i \rvert \leq \texttt{atol} + \texttt{rtol} \times \lvert \text{other}_i \rvert +""" + + r""" +elementwise, for all elements of :attr:`input` and :attr:`other`. The behaviour of this function is analogous to +`numpy.allclose `_ + +Args: + input (Tensor): first tensor to compare + other (Tensor): second tensor to compare + atol (float, optional): absolute tolerance. Default: 1e-08 + rtol (float, optional): relative tolerance. Default: 1e-05 + equal_nan (bool, optional): if ``True``, then two ``NaN`` s will be considered equal. Default: ``False`` + +Example:: + + >>> torch.allclose(torch.tensor([10000., 1e-07]), torch.tensor([10000.1, 1e-08])) + False + >>> torch.allclose(torch.tensor([10000., 1e-08]), torch.tensor([10000.1, 1e-09])) + True + >>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')])) + False + >>> torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]), equal_nan=True) + True +""", +) + +add_docstr( + torch.all, + r""" +all(input: Tensor, *, out=None) -> Tensor + +Tests if all elements in :attr:`input` evaluate to `True`. + +.. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.all(a) + tensor(False, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.all(a) + tensor(False) + +.. function:: all(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + +For each row of :attr:`input` in the given dimension :attr:`dim`, +returns `True` if all elements in the row evaluate to `True` and `False` otherwise. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(4, 2).bool() + >>> a + tensor([[True, True], + [True, False], + [True, True], + [True, True]], dtype=torch.bool) + >>> torch.all(a, dim=1) + tensor([ True, False, True, True], dtype=torch.bool) + >>> torch.all(a, dim=0) + tensor([ True, False], dtype=torch.bool) +""".format(**multi_dim_common), +) + +add_docstr( + torch.any, + r""" +any(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Tests if any element in :attr:`input` evaluates to `True`. + +.. note:: This function matches the behaviour of NumPy in returning + output of dtype `bool` for all supported dtypes except `uint8`. + For `uint8` the dtype of output is `uint8` itself. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(1, 2).bool() + >>> a + tensor([[False, True]], dtype=torch.bool) + >>> torch.any(a) + tensor(True, dtype=torch.bool) + >>> a = torch.arange(0, 3) + >>> a + tensor([0, 1, 2]) + >>> torch.any(a) + tensor(True) + +.. function:: any(input, dim, keepdim=False, *, out=None) -> Tensor + :noindex: + +For each row of :attr:`input` in the given dimension :attr:`dim`, +returns `True` if any element in the row evaluate to `True` and `False` otherwise. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4, 2) < 0 + >>> a + tensor([[ True, True], + [False, True], + [ True, True], + [False, False]]) + >>> torch.any(a, 1) + tensor([ True, True, True, False]) + >>> torch.any(a, 0) + tensor([True, True]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.angle, + r""" +angle(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Computes the element-wise angle (in radians) of the given :attr:`input` tensor. + +.. math:: + \text{out}_{i} = angle(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +.. note:: Starting in PyTorch 1.8, angle returns pi for negative real numbers, + zero for non-negative real numbers, and propagates NaNs. Previously + the function would return zero for all real numbers and not propagate + floating-point NaNs. + +Example:: + + >>> torch.angle(torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]))*180/3.14159 + tensor([ 135., 135, -45]) +""".format(**common_args), +) + +add_docstr( + torch.as_strided, + r""" +as_strided(input, size, stride, storage_offset=None) -> Tensor + +Create a view of an existing `torch.Tensor` :attr:`input` with specified +:attr:`size`, :attr:`stride` and :attr:`storage_offset`. + +.. warning:: + Prefer using other view functions, like :meth:`torch.Tensor.view` or + :meth:`torch.Tensor.expand`, to setting a view's strides manually with + `as_strided`, as this function will throw an error on non-standard Pytorch + backends (that do not have a concept of stride) and the result will depend + on the current layout in memory. The constructed view must only refer to + elements within the Tensor's storage or a runtime error will be thrown. + If the generated view is "overlapped" (with multiple indices referring to + the same element in memory), the behavior of inplace operations on this view + is undefined (and might not throw runtime errors). + +Args: + {input} + size (tuple or ints): the shape of the output tensor + stride (tuple or ints): the stride of the output tensor + storage_offset (int, optional): the offset in the underlying storage of the output tensor. + If ``None``, the storage_offset of the output tensor will match the input tensor. + +Example:: + + >>> x = torch.randn(3, 3) + >>> x + tensor([[ 0.9039, 0.6291, 1.0795], + [ 0.1586, 2.1939, -0.4900], + [-0.1909, -0.7503, 1.9355]]) + >>> t = torch.as_strided(x, (2, 2), (1, 2)) + >>> t + tensor([[0.9039, 1.0795], + [0.6291, 0.1586]]) + >>> t = torch.as_strided(x, (2, 2), (1, 2), 1) + tensor([[0.6291, 0.1586], + [1.0795, 2.1939]]) +""".format(**common_args), +) + +add_docstr( + torch.as_tensor, + r""" +as_tensor(data: Any, dtype: Optional[dtype] = None, device: Optional[DeviceLikeType]) -> Tensor + +Converts :attr:`data` into a tensor, sharing data and preserving autograd +history if possible. + +If :attr:`data` is already a tensor with the requested dtype and device +then :attr:`data` itself is returned, but if :attr:`data` is a +tensor with a different dtype or device then it's copied as if using +`data.to(dtype=dtype, device=device)`. + +If :attr:`data` is a NumPy array (an ndarray) with the same dtype and device then a +tensor is constructed using :func:`torch.from_numpy`. + +If :attr:`data` is a CuPy array, the returned tensor will be located on the same device as the CuPy array unless +specifically overwritten by :attr:`device` or a default device. The device of the CuPy array is inferred from the +pointer of the array using `cudaPointerGetAttributes` unless :attr:`device` is provided with an explicit device index. + +.. seealso:: + + :func:`torch.tensor` never shares its data and creates a new "leaf tensor" (see :doc:`/notes/autograd`). + + +Args: + {data} + {dtype} + device (:class:`torch.device`, optional): the device of the constructed tensor. If None and data is a tensor + then the device of data is used. If None and data is not a tensor then + the result tensor is constructed on the current device. + + +Example:: + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.as_tensor(a) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.as_tensor(a, device=torch.device('cuda')) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([1, 2, 3]) +""".format(**factory_data_common_args), +) + +add_docstr( + torch.asin, + r""" +asin(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the arcsine of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \sin^{-1}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.5962, 1.4985, -0.4396, 1.4525]) + >>> torch.asin(a) + tensor([-0.6387, nan, -0.4552, nan]) +""".format(**common_args), +) + +add_docstr( + torch.arcsin, + r""" +arcsin(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.asin`. +""", +) + +add_docstr( + torch.asinh, + r""" +asinh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the inverse hyperbolic sine of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \sinh^{-1}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.1606, -1.4267, -1.0899, -1.0250 ]) + >>> torch.asinh(a) + tensor([ 0.1599, -1.1534, -0.9435, -0.8990 ]) +""".format(**common_args), +) + +add_docstr( + torch.arcsinh, + r""" +arcsinh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.asinh`. +""", +) + +add_docstr( + torch.atan, + r""" +atan(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the arctangent of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \tan^{-1}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.2341, 0.2539, -0.6256, -0.6448]) + >>> torch.atan(a) + tensor([ 0.2299, 0.2487, -0.5591, -0.5727]) +""".format(**common_args), +) + +add_docstr( + torch.arctan, + r""" +arctan(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.atan`. +""", +) + +add_docstr( + torch.atan2, + r""" +atan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + +Element-wise arctangent of :math:`\text{{input}}_{{i}} / \text{{other}}_{{i}}` +with consideration of the quadrant. Returns a new tensor with the signed angles +in radians between vector :math:`(\text{{other}}_{{i}}, \text{{input}}_{{i}})` +and vector :math:`(1, 0)`. (Note that :math:`\text{{other}}_{{i}}`, the second +parameter, is the x-coordinate, while :math:`\text{{input}}_{{i}}`, the first +parameter, is the y-coordinate.) + +The shapes of ``input`` and ``other`` must be +:ref:`broadcastable `. + +Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.9041, 0.0196, -0.3108, -2.4423]) + >>> torch.atan2(a, torch.randn(4)) + tensor([ 0.9833, 0.0811, -1.9743, -1.4151]) +""".format(**common_args), +) + +add_docstr( + torch.arctan2, + r""" +arctan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor +Alias for :func:`torch.atan2`. +""", +) + +add_docstr( + torch.atanh, + r""" +atanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the inverse hyperbolic tangent of the elements of :attr:`input`. + +Note: + The domain of the inverse hyperbolic tangent is `(-1, 1)` and values outside this range + will be mapped to ``NaN``, except for the values `1` and `-1` for which the output is + mapped to `+/-INF` respectively. + +.. math:: + \text{out}_{i} = \tanh^{-1}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.randn(4).uniform_(-1, 1) + >>> a + tensor([ -0.9385, 0.2968, -0.8591, -0.1871 ]) + >>> torch.atanh(a) + tensor([ -1.7253, 0.3060, -1.2899, -0.1893 ]) +""".format(**common_args), +) + +add_docstr( + torch.arctanh, + r""" +arctanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Alias for :func:`torch.atanh`. +""", +) + +add_docstr( + torch.asarray, + r""" +asarray(obj: Any, *, dtype: Optional[dtype], device: Optional[DeviceLikeType], copy: Optional[bool] = None, requires_grad: bool = False) -> Tensor # noqa: B950 + +Converts :attr:`obj` to a tensor. + +:attr:`obj` can be one of: + +1. a tensor +2. a NumPy array or a NumPy scalar +3. a DLPack capsule +4. an object that implements Python's buffer protocol +5. a scalar +6. a sequence of scalars + +When :attr:`obj` is a tensor, NumPy array, or DLPack capsule the returned tensor will, +by default, not require a gradient, have the same datatype as :attr:`obj`, be on the +same device, and share memory with it. These properties can be controlled with the +:attr:`dtype`, :attr:`device`, :attr:`copy`, and :attr:`requires_grad` keyword arguments. +If the returned tensor is of a different datatype, on a different device, or a copy is +requested then it will not share its memory with :attr:`obj`. If :attr:`requires_grad` +is ``True`` then the returned tensor will require a gradient, and if :attr:`obj` is +also a tensor with an autograd history then the returned tensor will have the same history. + +When :attr:`obj` is not a tensor, NumPy array, or DLPack capsule but implements Python's +buffer protocol then the buffer is interpreted as an array of bytes grouped according to +the size of the datatype passed to the :attr:`dtype` keyword argument. (If no datatype is +passed then the default floating point datatype is used, instead.) The returned tensor +will have the specified datatype (or default floating point datatype if none is specified) +and, by default, be on the CPU device and share memory with the buffer. + +When :attr:`obj` is a NumPy scalar, the returned tensor will be a 0-dimensional tensor on +the CPU and that doesn't share its memory (i.e. ``copy=True``). By default datatype will +be the PyTorch datatype corresponding to the NumPy's scalar's datatype. + +When :attr:`obj` is none of the above but a scalar, or a sequence of scalars then the +returned tensor will, by default, infer its datatype from the scalar values, be on the +current default device, and not share its memory. + +.. seealso:: + + :func:`torch.tensor` creates a tensor that always copies the data from the input object. + :func:`torch.from_numpy` creates a tensor that always shares memory from NumPy arrays. + :func:`torch.frombuffer` creates a tensor that always shares memory from objects that + implement the buffer protocol. + :func:`torch.from_dlpack` creates a tensor that always shares memory from + DLPack capsules. + +Args: + obj (object): a tensor, NumPy array, DLPack Capsule, object that implements Python's + buffer protocol, scalar, or sequence of scalars. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the datatype of the returned tensor. + Default: ``None``, which causes the datatype of the returned tensor to be + inferred from :attr:`obj`. + copy (bool, optional): controls whether the returned tensor shares memory with :attr:`obj`. + Default: ``None``, which causes the returned tensor to share memory with :attr:`obj` + whenever possible. If ``True`` then the returned tensor does not share its memory. + If ``False`` then the returned tensor shares its memory with :attr:`obj` and an + error is thrown if it cannot. + device (:class:`torch.device`, optional): the device of the returned tensor. + Default: ``None``, which causes the device of :attr:`obj` to be used. Or, if + :attr:`obj` is a Python sequence, the current default device will be used. + requires_grad (bool, optional): whether the returned tensor requires grad. + Default: ``False``, which causes the returned tensor not to require a gradient. + If ``True``, then the returned tensor will require a gradient, and if :attr:`obj` + is also a tensor with an autograd history then the returned tensor will have + the same history. + +Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> # Shares memory with tensor 'a' + >>> b = torch.asarray(a) + >>> a.data_ptr() == b.data_ptr() + True + >>> # Forces memory copy + >>> c = torch.asarray(a, copy=True) + >>> a.data_ptr() == c.data_ptr() + False + + >>> a = torch.tensor([1., 2., 3.], requires_grad=True) + >>> b = a + 2 + >>> b + tensor([3., 4., 5.], grad_fn=) + >>> # Shares memory with tensor 'b', with no grad + >>> c = torch.asarray(b) + >>> c + tensor([3., 4., 5.]) + >>> # Shares memory with tensor 'b', retaining autograd history + >>> d = torch.asarray(b, requires_grad=True) + >>> d + tensor([3., 4., 5.], grad_fn=) + + >>> array = numpy.array([1, 2, 3]) + >>> # Shares memory with array 'array' + >>> t1 = torch.asarray(array) + >>> array.__array_interface__['data'][0] == t1.data_ptr() + True + >>> # Copies memory due to dtype mismatch + >>> t2 = torch.asarray(array, dtype=torch.float32) + >>> array.__array_interface__['data'][0] == t2.data_ptr() + False + + >>> scalar = numpy.float64(0.5) + >>> torch.asarray(scalar) + tensor(0.5000, dtype=torch.float64) +""", +) + +add_docstr( + torch.baddbmm, + r""" +baddbmm(input, batch1, batch2, out_dtype=None, *, beta=1, alpha=1, out=None) -> Tensor + +Performs a batch matrix-matrix product of matrices in :attr:`batch1` +and :attr:`batch2`. +:attr:`input` is added to the final result. + +:attr:`batch1` and :attr:`batch2` must be 3-D tensors each containing the same +number of matrices. + +If :attr:`batch1` is a :math:`(b \times n \times m)` tensor, :attr:`batch2` is a +:math:`(b \times m \times p)` tensor, then :attr:`input` must be +:ref:`broadcastable ` with a +:math:`(b \times n \times p)` tensor and :attr:`out` will be a +:math:`(b \times n \times p)` tensor. Both :attr:`alpha` and :attr:`beta` mean the +same as the scaling factors used in :meth:`torch.addbmm`. + +.. math:: + \text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i) + +If :attr:`beta` is 0, then the content of :attr:`input` will be ignored, and `nan` and `inf` in +it will not be propagated. +""" + + r""" +For inputs of type `FloatTensor` or `DoubleTensor`, arguments :attr:`beta` and +:attr:`alpha` must be real numbers, otherwise they should be integers. + +{tf32_note} + +{rocm_fp16_note} + +Args: + input (Tensor): the tensor to be added + batch1 (Tensor): the first batch of matrices to be multiplied + batch2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`\text{{batch1}} \mathbin{{@}} \text{{batch2}}` (:math:`\alpha`) + {out} + +Example:: + + >>> M = torch.randn(10, 3, 5) + >>> batch1 = torch.randn(10, 3, 4) + >>> batch2 = torch.randn(10, 4, 5) + >>> torch.baddbmm(M, batch1, batch2).size() + torch.Size([10, 3, 5]) +""".format(**common_args, **tf32_notes, **rocm_fp16_notes), +) + +add_docstr( + torch.bernoulli, + r""" +bernoulli(input: Tensor, *, generator: Optional[Generator], out: Optional[Tensor]) -> Tensor + +Draws binary random numbers (0 or 1) from a Bernoulli distribution. + +The :attr:`input` tensor should be a tensor containing probabilities +to be used for drawing the binary random number. +Hence, all values in :attr:`input` have to be in the range: +:math:`0 \leq \text{input}_i \leq 1`. + +The :math:`\text{i}^{th}` element of the output tensor will draw a +value :math:`1` according to the :math:`\text{i}^{th}` probability value given +in :attr:`input`. + +.. math:: + \text{out}_{i} \sim \mathrm{Bernoulli}(p = \text{input}_{i}) +""" + + r""" +The returned :attr:`out` tensor only has values 0 or 1 and is of the same +shape as :attr:`input`. + +:attr:`out` can have integral ``dtype``, but :attr:`input` must have floating +point ``dtype``. + +Args: + input (Tensor): the input tensor of probability values for the Bernoulli distribution + +Keyword args: + {generator} + {out} + +Example:: + + >>> a = torch.empty(3, 3).uniform_(0, 1) # generate a uniform random matrix with range [0, 1] + >>> a + tensor([[ 0.1737, 0.0950, 0.3609], + [ 0.7148, 0.0289, 0.2676], + [ 0.9456, 0.8937, 0.7202]]) + >>> torch.bernoulli(a) + tensor([[ 1., 0., 0.], + [ 0., 0., 0.], + [ 1., 1., 1.]]) + + >>> a = torch.ones(3, 3) # probability of drawing "1" is 1 + >>> torch.bernoulli(a) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.], + [ 1., 1., 1.]]) + >>> a = torch.zeros(3, 3) # probability of drawing "1" is 0 + >>> torch.bernoulli(a) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.], + [ 0., 0., 0.]]) +""".format(**common_args), +) + +add_docstr( + torch.bincount, + r""" +bincount(input, weights=None, minlength=0) -> Tensor + +Count the frequency of each value in an array of non-negative ints. + +The number of bins (size 1) is one larger than the largest value in +:attr:`input` unless :attr:`input` is empty, in which case the result is a +tensor of size 0. If :attr:`minlength` is specified, the number of bins is at least +:attr:`minlength` and if :attr:`input` is empty, then the result is tensor of size +:attr:`minlength` filled with zeros. If ``n`` is the value at position ``i``, +``out[n] += weights[i]`` if :attr:`weights` is specified else +``out[n] += 1``. + +Note: + {backward_reproducibility_note} + +Arguments: + input (Tensor): 1-d int tensor + weights (Tensor): optional, weight for each value in the input tensor. + Should be of same size as input tensor. + minlength (int): optional, minimum number of bins. Should be non-negative. + +Returns: + output (Tensor): a tensor of shape ``Size([max(input) + 1])`` if + :attr:`input` is non-empty, else ``Size(0)`` + +Example:: + + >>> input = torch.randint(0, 8, (5,), dtype=torch.int64) + >>> weights = torch.linspace(0, 1, steps=5) + >>> input, weights + (tensor([4, 3, 6, 3, 4]), + tensor([ 0.0000, 0.2500, 0.5000, 0.7500, 1.0000]) + + >>> torch.bincount(input) + tensor([0, 0, 0, 2, 2, 0, 1]) + + >>> input.bincount(weights) + tensor([0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 0.0000, 0.5000]) +""".format(**reproducibility_notes), +) + +add_docstr( + torch.bitwise_not, + r""" +bitwise_not(input, *, out=None) -> Tensor + +Computes the bitwise NOT of the given input tensor. The input tensor must be of +integral or Boolean types. For bool tensors, it computes the logical NOT. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_not(torch.tensor([-1, -2, 3], dtype=torch.int8)) + tensor([ 0, 1, -4], dtype=torch.int8) +""".format(**common_args), +) + +add_docstr( + torch.bmm, + r""" +bmm(input, mat2, out_dtype=None, *, out=None) -> Tensor + +Performs a batch matrix-matrix product of matrices stored in :attr:`input` +and :attr:`mat2`. + +:attr:`input` and :attr:`mat2` must be 3-D tensors each containing +the same number of matrices. + +If :attr:`input` is a :math:`(b \times n \times m)` tensor, :attr:`mat2` is a +:math:`(b \times m \times p)` tensor, :attr:`out` will be a +:math:`(b \times n \times p)` tensor. + +.. math:: + \text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i +""" + + r""" +{tf32_note} + +{rocm_fp16_note} + +.. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + +Args: + input (Tensor): the first batch of matrices to be multiplied + mat2 (Tensor): the second batch of matrices to be multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + +Keyword Args: + {out} + +Example:: + + >>> input = torch.randn(10, 3, 4) + >>> mat2 = torch.randn(10, 4, 5) + >>> res = torch.bmm(input, mat2) + >>> res.size() + torch.Size([10, 3, 5]) +""".format(**common_args, **tf32_notes, **rocm_fp16_notes), +) + +add_docstr( + torch.bitwise_and, + r""" +bitwise_and(input, other, *, out=None) -> Tensor + +Computes the bitwise AND of :attr:`input` and :attr:`other`. The input tensor must be of +integral or Boolean types. For bool tensors, it computes the logical AND. + +Args: + input: the first input tensor + other: the second input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_and(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([1, 0, 3], dtype=torch.int8) + >>> torch.bitwise_and(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ False, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.bitwise_or, + r""" +bitwise_or(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + +Computes the bitwise OR of :attr:`input` and :attr:`other`. The input tensor must be of +integral or Boolean types. For bool tensors, it computes the logical OR. + +Args: + input: the first input tensor + other: the second input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_or(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -2, 3], dtype=torch.int8) + >>> torch.bitwise_or(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.bitwise_xor, + r""" +bitwise_xor(input, other, *, out=None) -> Tensor + +Computes the bitwise XOR of :attr:`input` and :attr:`other`. The input tensor must be of +integral or Boolean types. For bool tensors, it computes the logical XOR. + +Args: + input: the first input tensor + other: the second input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_xor(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 0], dtype=torch.int8) + >>> torch.bitwise_xor(torch.tensor([True, True, False]), torch.tensor([False, True, False])) + tensor([ True, False, False]) +""".format(**common_args), +) + +add_docstr( + torch.bitwise_left_shift, + r""" +bitwise_left_shift(input, other, *, out=None) -> Tensor + +Computes the left arithmetic shift of :attr:`input` by :attr:`other` bits. +The input tensor must be of integral type. This operator supports +:ref:`broadcasting to a common shape ` and +:ref:`type promotion `. + +The operation applied is: + +.. math:: + \text{{out}}_i = \text{{input}}_i << \text{{other}}_i + +Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_left_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-2, -2, 24], dtype=torch.int8) +""".format(**common_args), +) + +add_docstr( + torch.bitwise_right_shift, + r""" +bitwise_right_shift(input, other, *, out=None) -> Tensor + +Computes the right arithmetic shift of :attr:`input` by :attr:`other` bits. +The input tensor must be of integral type. This operator supports +:ref:`broadcasting to a common shape ` and +:ref:`type promotion `. +In any case, if the value of the right operand is negative or is greater +or equal to the number of bits in the promoted left operand, the behavior is undefined. + +The operation applied is: + +.. math:: + \text{{out}}_i = \text{{input}}_i >> \text{{other}}_i + +Args: + input (Tensor or Scalar): the first input tensor + other (Tensor or Scalar): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.bitwise_right_shift(torch.tensor([-2, -7, 31], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)) + tensor([-1, -7, 3], dtype=torch.int8) +""".format(**common_args), +) + +add_docstr( + torch.broadcast_to, + r""" +broadcast_to(input, shape) -> Tensor + +Broadcasts :attr:`input` to the shape :attr:`\shape`. +Equivalent to calling ``input.expand(shape)``. See :meth:`~Tensor.expand` for details. + +Args: + {input} + shape (list, tuple, or :class:`torch.Size`): the new shape. + +Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> torch.broadcast_to(x, (3, 3)) + tensor([[1, 2, 3], + [1, 2, 3], + [1, 2, 3]]) +""".format(**common_args), +) + +add_docstr( + torch.stack, + r""" +stack(tensors, dim=0, *, out=None) -> Tensor + +Concatenates a sequence of tensors along a new dimension. + +All tensors need to be of the same size. + +.. seealso:: + + :func:`torch.cat` concatenates the given sequence along an existing dimension. + +Arguments: + tensors (sequence of Tensors): sequence of tensors to concatenate + dim (int, optional): dimension to insert. Has to be between 0 and the number + of dimensions of concatenated tensors (inclusive). Default: 0 + +Keyword args: + {out} + +Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]]) + >>> torch.stack((x, x)) # same as torch.stack((x, x), dim=0) + tensor([[[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]], + + [[ 0.3367, 0.1288, 0.2345], + [ 0.2303, -1.1229, -0.1863]]]) + >>> torch.stack((x, x)).size() + torch.Size([2, 2, 3]) + >>> torch.stack((x, x), dim=1) + tensor([[[ 0.3367, 0.1288, 0.2345], + [ 0.3367, 0.1288, 0.2345]], + + [[ 0.2303, -1.1229, -0.1863], + [ 0.2303, -1.1229, -0.1863]]]) + >>> torch.stack((x, x), dim=2) + tensor([[[ 0.3367, 0.3367], + [ 0.1288, 0.1288], + [ 0.2345, 0.2345]], + + [[ 0.2303, 0.2303], + [-1.1229, -1.1229], + [-0.1863, -0.1863]]]) + >>> torch.stack((x, x), dim=-1) + tensor([[[ 0.3367, 0.3367], + [ 0.1288, 0.1288], + [ 0.2345, 0.2345]], + + [[ 0.2303, 0.2303], + [-1.1229, -1.1229], + [-0.1863, -0.1863]]]) +""".format(**common_args), +) + +add_docstr( + torch.hstack, + r""" +hstack(tensors, *, out=None) -> Tensor + +Stack tensors in sequence horizontally (column wise). + +This is equivalent to concatenation along the first axis for 1-D tensors, and along the second axis for all other tensors. + +Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.hstack((a,b)) + tensor([1, 2, 3, 4, 5, 6]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.hstack((a,b)) + tensor([[1, 4], + [2, 5], + [3, 6]]) + +""".format(**common_args), +) + +add_docstr( + torch.vstack, + r""" +vstack(tensors, *, out=None) -> Tensor + +Stack tensors in sequence vertically (row wise). + +This is equivalent to concatenation along the first axis after all 1-D tensors have been reshaped by :func:`torch.atleast_2d`. + +Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.vstack((a,b)) + tensor([[1, 2, 3], + [4, 5, 6]]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.vstack((a,b)) + tensor([[1], + [2], + [3], + [4], + [5], + [6]]) + + +""".format(**common_args), +) + +add_docstr( + torch.dstack, + r""" +dstack(tensors, *, out=None) -> Tensor + +Stack tensors in sequence depthwise (along third axis). + +This is equivalent to concatenation along the third axis after 1-D and 2-D tensors have been reshaped by :func:`torch.atleast_3d`. + +Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.dstack((a,b)) + tensor([[[1, 4], + [2, 5], + [3, 6]]]) + >>> a = torch.tensor([[1],[2],[3]]) + >>> b = torch.tensor([[4],[5],[6]]) + >>> torch.dstack((a,b)) + tensor([[[1, 4]], + [[2, 5]], + [[3, 6]]]) + + +""".format(**common_args), +) + +add_docstr( + torch.tensor_split, + r""" +tensor_split(input, indices_or_sections, dim=0) -> List of Tensors + +Splits a tensor into multiple sub-tensors, all of which are views of :attr:`input`, +along dimension :attr:`dim` according to the indices or number of sections specified +by :attr:`indices_or_sections`. This function is based on NumPy's +:func:`numpy.array_split`. + +Args: + input (Tensor): the tensor to split + indices_or_sections (Tensor, int or list or tuple of ints): + If :attr:`indices_or_sections` is an integer ``n`` or a zero dimensional long tensor + with value ``n``, :attr:`input` is split into ``n`` sections along dimension :attr:`dim`. + If :attr:`input` is divisible by ``n`` along dimension :attr:`dim`, each + section will be of equal size, :code:`input.size(dim) / n`. If :attr:`input` + is not divisible by ``n``, the sizes of the first :code:`int(input.size(dim) % n)` + sections will have size :code:`int(input.size(dim) / n) + 1`, and the rest will + have size :code:`int(input.size(dim) / n)`. + + If :attr:`indices_or_sections` is a list or tuple of ints, or a one-dimensional long + tensor, then :attr:`input` is split along dimension :attr:`dim` at each of the indices + in the list, tuple or tensor. For instance, :code:`indices_or_sections=[2, 3]` and :code:`dim=0` + would result in the tensors :code:`input[:2]`, :code:`input[2:3]`, and :code:`input[3:]`. + + If :attr:`indices_or_sections` is a tensor, it must be a zero-dimensional or one-dimensional + long tensor on the CPU. + + dim (int, optional): dimension along which to split the tensor. Default: ``0`` + +Example:: + + >>> x = torch.arange(8) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7])) + + >>> x = torch.arange(7) + >>> torch.tensor_split(x, 3) + (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6])) + >>> torch.tensor_split(x, (1, 6)) + (tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6])) + + >>> x = torch.arange(14).reshape(2, 7) + >>> x + tensor([[ 0, 1, 2, 3, 4, 5, 6], + [ 7, 8, 9, 10, 11, 12, 13]]) + >>> torch.tensor_split(x, 3, dim=1) + (tensor([[0, 1, 2], + [7, 8, 9]]), + tensor([[ 3, 4], + [10, 11]]), + tensor([[ 5, 6], + [12, 13]])) + >>> torch.tensor_split(x, (1, 6), dim=1) + (tensor([[0], + [7]]), + tensor([[ 1, 2, 3, 4, 5], + [ 8, 9, 10, 11, 12]]), + tensor([[ 6], + [13]])) +""", +) + +add_docstr( + torch.chunk, + r""" +chunk(input: Tensor, chunks: int, dim: int = 0) -> Tuple[Tensor, ...] + +Attempts to split a tensor into the specified number of chunks. Each chunk is a view of +the input tensor. + + +.. note:: + + This function may return fewer than the specified number of chunks! + +.. seealso:: + + :func:`torch.tensor_split` a function that always returns exactly the specified number of chunks + +If the tensor size along the given dimension :attr:`dim` is divisible by :attr:`chunks`, +all returned chunks will be the same size. +If the tensor size along the given dimension :attr:`dim` is not divisible by :attr:`chunks`, +all returned chunks will be the same size, except the last one. +If such division is not possible, this function may return fewer +than the specified number of chunks. + +Arguments: + input (Tensor): the tensor to split + chunks (int): number of chunks to return + dim (int): dimension along which to split the tensor + +Example: + >>> torch.arange(11).chunk(6) + (tensor([0, 1]), + tensor([2, 3]), + tensor([4, 5]), + tensor([6, 7]), + tensor([8, 9]), + tensor([10])) + >>> torch.arange(12).chunk(6) + (tensor([0, 1]), + tensor([2, 3]), + tensor([4, 5]), + tensor([6, 7]), + tensor([8, 9]), + tensor([10, 11])) + >>> torch.arange(13).chunk(6) + (tensor([0, 1, 2]), + tensor([3, 4, 5]), + tensor([6, 7, 8]), + tensor([ 9, 10, 11]), + tensor([12])) +""", +) + +add_docstr( + torch.unsafe_chunk, + r""" +unsafe_chunk(input, chunks, dim=0) -> List of Tensors + +Works like :func:`torch.chunk` but without enforcing the autograd restrictions +on inplace modification of the outputs. + +.. warning:: + This function is safe to use as long as only the input, or only the outputs + are modified inplace after calling this function. It is user's + responsibility to ensure that is the case. If both the input and one or more + of the outputs are modified inplace, gradients computed by autograd will be + silently incorrect. +""", +) + +add_docstr( + torch.unsafe_split, + r""" +unsafe_split(tensor, split_size_or_sections, dim=0) -> List of Tensors + +Works like :func:`torch.split` but without enforcing the autograd restrictions +on inplace modification of the outputs. + +.. warning:: + This function is safe to use as long as only the input, or only the outputs + are modified inplace after calling this function. It is user's + responsibility to ensure that is the case. If both the input and one or more + of the outputs are modified inplace, gradients computed by autograd will be + silently incorrect. +""", +) + +add_docstr( + torch.hsplit, + r""" +hsplit(input, indices_or_sections) -> List of Tensors + +Splits :attr:`input`, a tensor with one or more dimensions, into multiple tensors +horizontally according to :attr:`indices_or_sections`. Each split is a view of +:attr:`input`. + +If :attr:`input` is one dimensional this is equivalent to calling +torch.tensor_split(input, indices_or_sections, dim=0) (the split dimension is +zero), and if :attr:`input` has two or more dimensions it's equivalent to calling +torch.tensor_split(input, indices_or_sections, dim=1) (the split dimension is 1), +except that if :attr:`indices_or_sections` is an integer it must evenly divide +the split dimension or a runtime error will be thrown. + +This function is based on NumPy's :func:`numpy.hsplit`. + +Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + +Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.hsplit(t, 2) + (tensor([[ 0., 1.], + [ 4., 5.], + [ 8., 9.], + [12., 13.]]), + tensor([[ 2., 3.], + [ 6., 7.], + [10., 11.], + [14., 15.]])) + >>> torch.hsplit(t, [3, 6]) + (tensor([[ 0., 1., 2.], + [ 4., 5., 6.], + [ 8., 9., 10.], + [12., 13., 14.]]), + tensor([[ 3.], + [ 7.], + [11.], + [15.]]), + tensor([], size=(4, 0))) + +""", +) + +add_docstr( + torch.vsplit, + r""" +vsplit(input, indices_or_sections) -> List of Tensors + +Splits :attr:`input`, a tensor with two or more dimensions, into multiple tensors +vertically according to :attr:`indices_or_sections`. Each split is a view of +:attr:`input`. + +This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=0) +(the split dimension is 0), except that if :attr:`indices_or_sections` is an integer +it must evenly divide the split dimension or a runtime error will be thrown. + +This function is based on NumPy's :func:`numpy.vsplit`. + +Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + +Example:: + + >>> t = torch.arange(16.0).reshape(4,4) + >>> t + tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.], + [12., 13., 14., 15.]]) + >>> torch.vsplit(t, 2) + (tensor([[0., 1., 2., 3.], + [4., 5., 6., 7.]]), + tensor([[ 8., 9., 10., 11.], + [12., 13., 14., 15.]])) + >>> torch.vsplit(t, [3, 6]) + (tensor([[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.], + [ 8., 9., 10., 11.]]), + tensor([[12., 13., 14., 15.]]), + tensor([], size=(0, 4))) + +""", +) + +add_docstr( + torch.dsplit, + r""" +dsplit(input, indices_or_sections) -> List of Tensors + +Splits :attr:`input`, a tensor with three or more dimensions, into multiple tensors +depthwise according to :attr:`indices_or_sections`. Each split is a view of +:attr:`input`. + +This is equivalent to calling torch.tensor_split(input, indices_or_sections, dim=2) +(the split dimension is 2), except that if :attr:`indices_or_sections` is an integer +it must evenly divide the split dimension or a runtime error will be thrown. + +This function is based on NumPy's :func:`numpy.dsplit`. + +Args: + input (Tensor): tensor to split. + indices_or_sections (int or list or tuple of ints): See argument in :func:`torch.tensor_split`. + +Example:: + + >>> t = torch.arange(16.0).reshape(2, 2, 4) + >>> t + tensor([[[ 0., 1., 2., 3.], + [ 4., 5., 6., 7.]], + [[ 8., 9., 10., 11.], + [12., 13., 14., 15.]]]) + >>> torch.dsplit(t, 2) + (tensor([[[ 0., 1.], + [ 4., 5.]], + [[ 8., 9.], + [12., 13.]]]), + tensor([[[ 2., 3.], + [ 6., 7.]], + [[10., 11.], + [14., 15.]]])) + + >>> torch.dsplit(t, [3, 6]) + (tensor([[[ 0., 1., 2.], + [ 4., 5., 6.]], + [[ 8., 9., 10.], + [12., 13., 14.]]]), + tensor([[[ 3.], + [ 7.]], + [[11.], + [15.]]]), + tensor([], size=(2, 2, 0))) + +""", +) + +add_docstr( + torch.can_cast, + r""" +can_cast(from_, to) -> bool + +Determines if a type conversion is allowed under PyTorch casting rules +described in the type promotion :ref:`documentation `. + +Args: + from\_ (dtype): The original :class:`torch.dtype`. + to (dtype): The target :class:`torch.dtype`. + +Example:: + + >>> torch.can_cast(torch.double, torch.float) + True + >>> torch.can_cast(torch.float, torch.int) + False +""", +) + +add_docstr( + torch.corrcoef, + r""" +corrcoef(input) -> Tensor + +Estimates the Pearson product-moment correlation coefficient matrix of the variables given by the :attr:`input` matrix, +where rows are the variables and columns are the observations. + +.. note:: + + The correlation coefficient matrix R is computed using the covariance matrix C as given by + :math:`R_{ij} = \frac{ C_{ij} } { \sqrt{ C_{ii} * C_{jj} } }` + +.. note:: + + Due to floating point rounding, the resulting array may not be Hermitian and its diagonal elements may not be 1. + The real and imaginary values are clipped to the interval [-1, 1] in an attempt to improve this situation. + +Args: + input (Tensor): A 2D matrix containing multiple variables and observations, or a + Scalar or 1D vector representing a single variable. + +Returns: + (Tensor) The correlation coefficient matrix of the variables. + +.. seealso:: + + :func:`torch.cov` covariance matrix. + +Example:: + + >>> x = torch.tensor([[0, 1, 2], [2, 1, 0]]) + >>> torch.corrcoef(x) + tensor([[ 1., -1.], + [-1., 1.]]) + >>> x = torch.randn(2, 4) + >>> x + tensor([[-0.2678, -0.0908, -0.3766, 0.2780], + [-0.5812, 0.1535, 0.2387, 0.2350]]) + >>> torch.corrcoef(x) + tensor([[1.0000, 0.3582], + [0.3582, 1.0000]]) + >>> torch.corrcoef(x[0]) + tensor(1.) +""", +) + +add_docstr( + torch.cov, + r""" +cov(input, *, correction=1, fweights=None, aweights=None) -> Tensor + +Estimates the covariance matrix of the variables given by the :attr:`input` matrix, where rows are +the variables and columns are the observations. + +A covariance matrix is a square matrix giving the covariance of each pair of variables. The diagonal contains +the variance of each variable (covariance of a variable with itself). By definition, if :attr:`input` represents +a single variable (Scalar or 1D) then its variance is returned. + +The sample covariance of the variables :math:`x` and :math:`y` is given by: + +.. math:: + \text{cov}(x,y) = \frac{\sum^{N}_{i = 1}(x_{i} - \bar{x})(y_{i} - \bar{y})}{\max(0,~N~-~\delta N)} + +where :math:`\bar{x}` and :math:`\bar{y}` are the simple means of the :math:`x` and :math:`y` respectively, and +:math:`\delta N` is the :attr:`correction`. + +If :attr:`fweights` and/or :attr:`aweights` are provided, the weighted covariance +is calculated, which is given by: + +.. math:: + \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}w_i(x_{i} - \mu_x^*)(y_{i} - \mu_y^*)} + {\max(0,~\sum^{N}_{i = 1}w_i~-~\frac{\sum^{N}_{i = 1}w_ia_i}{\sum^{N}_{i = 1}w_i}~\delta N)} + +where :math:`w` denotes :attr:`fweights` or :attr:`aweights` (``f`` and ``a`` for brevity) based on whichever is +provided, or :math:`w = f \times a` if both are provided, and +:math:`\mu_x^* = \frac{\sum^{N}_{i = 1}w_ix_{i} }{\sum^{N}_{i = 1}w_i}` is the weighted mean of the variable. If not +provided, ``f`` and/or ``a`` can be seen as a :math:`\mathbb{1}` vector of appropriate size. + +Args: + input (Tensor): A 2D matrix containing multiple variables and observations, or a + Scalar or 1D vector representing a single variable. + +Keyword Args: + correction (int, optional): difference between the sample size and sample degrees of freedom. + Defaults to Bessel's correction, ``correction = 1`` which returns the unbiased estimate, + even if both :attr:`fweights` and :attr:`aweights` are specified. ``correction = 0`` + will return the simple average. Defaults to ``1``. + fweights (tensor, optional): A Scalar or 1D tensor of observation vector frequencies representing the number of + times each observation should be repeated. Its numel must equal the number of columns of :attr:`input`. + Must have integral dtype. Ignored if ``None``. Defaults to ``None``. + aweights (tensor, optional): A Scalar or 1D array of observation vector weights. + These relative weights are typically large for observations considered "important" and smaller for + observations considered less "important". Its numel must equal the number of columns of :attr:`input`. + Must have floating point dtype. Ignored if ``None``. Defaults to ``None``. + +Returns: + (Tensor) The covariance matrix of the variables. + +.. seealso:: + + :func:`torch.corrcoef` normalized covariance matrix. + +Example:: + + >>> x = torch.tensor([[0, 2], [1, 1], [2, 0]]).T + >>> x + tensor([[0, 1, 2], + [2, 1, 0]]) + >>> torch.cov(x) + tensor([[ 1., -1.], + [-1., 1.]]) + >>> torch.cov(x, correction=0) + tensor([[ 0.6667, -0.6667], + [-0.6667, 0.6667]]) + >>> fw = torch.randint(1, 10, (3,)) + >>> fw + tensor([1, 6, 9]) + >>> aw = torch.rand(3) + >>> aw + tensor([0.4282, 0.0255, 0.4144]) + >>> torch.cov(x, fweights=fw, aweights=aw) + tensor([[ 0.4169, -0.4169], + [-0.4169, 0.4169]]) +""", +) + +add_docstr( + torch.cat, + r""" +cat(tensors, dim=0, *, out=None) -> Tensor + +Concatenates the given sequence of tensors in :attr:`tensors` in the given dimension. +All tensors must either have the same shape (except in the concatenating +dimension) or be a 1-D empty tensor with size ``(0,)``. + +:func:`torch.cat` can be seen as an inverse operation for :func:`torch.split` +and :func:`torch.chunk`. + +:func:`torch.cat` can be best understood via examples. + +.. seealso:: + + :func:`torch.stack` concatenates the given sequence along a new dimension. + +Args: + tensors (sequence of Tensors): Non-empty tensors provided must have the same shape, + except in the cat dimension. + + dim (int, optional): the dimension over which the tensors are concatenated + +Keyword args: + {out} + +Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 0) + tensor([[ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497], + [ 0.6580, -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497]]) + >>> torch.cat((x, x, x), 1) + tensor([[ 0.6580, -1.0969, -0.4614, 0.6580, -1.0969, -0.4614, 0.6580, + -1.0969, -0.4614], + [-0.1034, -0.5790, 0.1497, -0.1034, -0.5790, 0.1497, -0.1034, + -0.5790, 0.1497]]) +""".format(**common_args), +) + +add_docstr( + torch.concat, + r""" +concat(tensors, dim=0, *, out=None) -> Tensor + +Alias of :func:`torch.cat`. +""", +) + +add_docstr( + torch.concatenate, + r""" +concatenate(tensors, axis=0, out=None) -> Tensor + +Alias of :func:`torch.cat`. +""", +) + +add_docstr( + torch.ceil, + r""" +ceil(input, *, out=None) -> Tensor + +Returns a new tensor with the ceil of the elements of :attr:`input`, +the smallest integer greater than or equal to each element. + +For integer inputs, follows the array-api convention of returning a +copy of the input tensor. + +.. math:: + \text{out}_{i} = \left\lceil \text{input}_{i} \right\rceil +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.6341, -1.4208, -1.0900, 0.5826]) + >>> torch.ceil(a) + tensor([-0., -1., -1., 1.]) +""".format(**common_args), +) + +add_docstr( + torch.real, + r""" +real(input) -> Tensor + +Returns a new tensor containing real values of the :attr:`self` tensor. +The returned tensor and :attr:`self` share the same underlying storage. + +Args: + {input} + +Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.real + tensor([ 0.3100, -0.5445, -1.6492, -0.0638]) + +""".format(**common_args), +) + +add_docstr( + torch.imag, + r""" +imag(input) -> Tensor + +Returns a new tensor containing imaginary values of the :attr:`self` tensor. +The returned tensor and :attr:`self` share the same underlying storage. + +.. warning:: + :func:`imag` is only supported for tensors with complex dtypes. + +Args: + {input} + +Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.3100+0.3553j), (-0.5445-0.7896j), (-1.6492-0.0633j), (-0.0638-0.8119j)]) + >>> x.imag + tensor([ 0.3553, -0.7896, -0.0633, -0.8119]) + +""".format(**common_args), +) + +add_docstr( + torch.view_as_real, + r""" +view_as_real(input) -> Tensor + +Returns a view of :attr:`input` as a real tensor. For an input complex tensor of +:attr:`size` :math:`m1, m2, \dots, mi`, this function returns a new +real tensor of size :math:`m1, m2, \dots, mi, 2`, where the last dimension of size 2 +represents the real and imaginary components of complex numbers. + +.. warning:: + :func:`view_as_real` is only supported for tensors with ``complex dtypes``. + +Args: + {input} + +Example:: + + >>> x=torch.randn(4, dtype=torch.cfloat) + >>> x + tensor([(0.4737-0.3839j), (-0.2098-0.6699j), (0.3470-0.9451j), (-0.5174-1.3136j)]) + >>> torch.view_as_real(x) + tensor([[ 0.4737, -0.3839], + [-0.2098, -0.6699], + [ 0.3470, -0.9451], + [-0.5174, -1.3136]]) +""".format(**common_args), +) + +add_docstr( + torch.view_as_complex, + r""" +view_as_complex(input) -> Tensor + +Returns a view of :attr:`input` as a complex tensor. For an input complex +tensor of :attr:`size` :math:`m1, m2, \dots, mi, 2`, this function returns a +new complex tensor of :attr:`size` :math:`m1, m2, \dots, mi` where the last +dimension of the input tensor is expected to represent the real and imaginary +components of complex numbers. + +.. warning:: + :func:`view_as_complex` is only supported for tensors with + :class:`torch.dtype` ``torch.float64`` and ``torch.float32``. The input is + expected to have the last dimension of :attr:`size` 2. In addition, the + tensor must have a `stride` of 1 for its last dimension. The strides of all + other dimensions must be even numbers. + +Args: + {input} + +Example:: + + >>> x=torch.randn(4, 2) + >>> x + tensor([[ 1.6116, -0.5772], + [-1.4606, -0.9120], + [ 0.0786, -1.7497], + [-0.6561, -1.6623]]) + >>> torch.view_as_complex(x) + tensor([(1.6116-0.5772j), (-1.4606-0.9120j), (0.0786-1.7497j), (-0.6561-1.6623j)]) +""".format(**common_args), +) + +add_docstr( + torch.reciprocal, + r""" +reciprocal(input, *, out=None) -> Tensor + +Returns a new tensor with the reciprocal of the elements of :attr:`input` + +.. math:: + \text{out}_{i} = \frac{1}{\text{input}_{i}} + +.. note:: + Unlike NumPy's reciprocal, torch.reciprocal supports integral inputs. Integral + inputs to reciprocal are automatically :ref:`promoted ` to + the default scalar type. +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.4595, -2.1219, -1.4314, 0.7298]) + >>> torch.reciprocal(a) + tensor([-2.1763, -0.4713, -0.6986, 1.3702]) +""".format(**common_args), +) + +add_docstr( + torch.cholesky, + r""" +cholesky(input, upper=False, *, out=None) -> Tensor + +Computes the Cholesky decomposition of a symmetric positive-definite +matrix :math:`A` or for batches of symmetric positive-definite matrices. + +If :attr:`upper` is ``True``, the returned matrix ``U`` is upper-triangular, and +the decomposition has the form: + +.. math:: + + A = U^TU + +If :attr:`upper` is ``False``, the returned matrix ``L`` is lower-triangular, and +the decomposition has the form: + +.. math:: + + A = LL^T + +If :attr:`upper` is ``True``, and :math:`A` is a batch of symmetric positive-definite +matrices, then the returned tensor will be composed of upper-triangular Cholesky factors +of each of the individual matrices. Similarly, when :attr:`upper` is ``False``, the returned +tensor will be composed of lower-triangular Cholesky factors of each of the individual +matrices. + +.. warning:: + + :func:`torch.cholesky` is deprecated in favor of :func:`torch.linalg.cholesky` + and will be removed in a future PyTorch release. + + ``L = torch.cholesky(A)`` should be replaced with + + .. code:: python + + L = torch.linalg.cholesky(A) + + ``U = torch.cholesky(A, upper=True)`` should be replaced with + + .. code:: python + + U = torch.linalg.cholesky(A).mH + + This transform will produce equivalent results for all valid (symmetric positive definite) inputs. + +Args: + input (Tensor): the input tensor :math:`A` of size :math:`(*, n, n)` where `*` is zero or more + batch dimensions consisting of symmetric positive-definite matrices. + upper (bool, optional): flag that indicates whether to return a + upper or lower triangular matrix. Default: ``False`` + +Keyword args: + out (Tensor, optional): the output matrix + +Example:: + + >>> a = torch.randn(3, 3) + >>> a = a @ a.mT + 1e-3 # make symmetric positive-definite + >>> l = torch.cholesky(a) + >>> a + tensor([[ 2.4112, -0.7486, 1.4551], + [-0.7486, 1.3544, 0.1294], + [ 1.4551, 0.1294, 1.6724]]) + >>> l + tensor([[ 1.5528, 0.0000, 0.0000], + [-0.4821, 1.0592, 0.0000], + [ 0.9371, 0.5487, 0.7023]]) + >>> l @ l.mT + tensor([[ 2.4112, -0.7486, 1.4551], + [-0.7486, 1.3544, 0.1294], + [ 1.4551, 0.1294, 1.6724]]) + >>> a = torch.randn(3, 2, 2) # Example for batched input + >>> a = a @ a.mT + 1e-03 # make symmetric positive-definite + >>> l = torch.cholesky(a) + >>> z = l @ l.mT + >>> torch.dist(z, a) + tensor(2.3842e-07) +""", +) + +add_docstr( + torch.cholesky_solve, + r""" +cholesky_solve(B, L, upper=False, *, out=None) -> Tensor + +Computes the solution of a system of linear equations with complex Hermitian +or real symmetric positive-definite lhs given its Cholesky decomposition. + +Let :math:`A` be a complex Hermitian or real symmetric positive-definite matrix, +and :math:`L` its Cholesky decomposition such that: + +.. math:: + + A = LL^{\text{H}} + +where :math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, +and the transpose when :math:`L` is real-valued. + +Returns the solution :math:`X` of the following linear system: + +.. math:: + + AX = B + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :math:`A` or :math:`B` is a batch of matrices +then the output has the same batch dimensions. + +Args: + B (Tensor): right-hand side tensor of shape `(*, n, k)` + where :math:`*` is zero or more batch dimensions + L (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of lower or upper triangular Cholesky decompositions of + symmetric or Hermitian positive-definite matrices. + upper (bool, optional): flag that indicates whether :math:`L` is lower triangular + or upper triangular. Default: ``False``. + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Example:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.T + torch.eye(3) * 1e-3 # Creates a symmetric positive-definite matrix + >>> L = torch.linalg.cholesky(A) # Extract Cholesky decomposition + >>> B = torch.randn(3, 2) + >>> torch.cholesky_solve(B, L) + tensor([[ -8.1625, 19.6097], + [ -5.8398, 14.2387], + [ -4.3771, 10.4173]]) + >>> A.inverse() @ B + tensor([[ -8.1626, 19.6097], + [ -5.8398, 14.2387], + [ -4.3771, 10.4173]]) + + >>> A = torch.randn(3, 2, 2, dtype=torch.complex64) + >>> A = A @ A.mH + torch.eye(2) * 1e-3 # Batch of Hermitian positive-definite matrices + >>> L = torch.linalg.cholesky(A) + >>> B = torch.randn(2, 1, dtype=torch.complex64) + >>> X = torch.cholesky_solve(B, L) + >>> torch.dist(X, A.inverse() @ B) + tensor(1.6881e-5) +""", +) + +add_docstr( + torch.cholesky_inverse, + r""" +cholesky_inverse(L, upper=False, *, out=None) -> Tensor + +Computes the inverse of a complex Hermitian or real symmetric +positive-definite matrix given its Cholesky decomposition. + +Let :math:`A` be a complex Hermitian or real symmetric positive-definite matrix, +and :math:`L` its Cholesky decomposition such that: + +.. math:: + + A = LL^{\text{H}} + +where :math:`L^{\text{H}}` is the conjugate transpose when :math:`L` is complex, +and the transpose when :math:`L` is real-valued. + +Computes the inverse matrix :math:`A^{-1}`. + +Supports input of float, double, cfloat and cdouble dtypes. +Also supports batches of matrices, and if :math:`A` is a batch of matrices +then the output has the same batch dimensions. + +Args: + L (Tensor): tensor of shape `(*, n, n)` where `*` is zero or more batch dimensions + consisting of lower or upper triangular Cholesky decompositions of + symmetric or Hermitian positive-definite matrices. + upper (bool, optional): flag that indicates whether :math:`L` is lower triangular + or upper triangular. Default: ``False`` + +Keyword args: + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Example:: + + >>> A = torch.randn(3, 3) + >>> A = A @ A.T + torch.eye(3) * 1e-3 # Creates a symmetric positive-definite matrix + >>> L = torch.linalg.cholesky(A) # Extract Cholesky decomposition + >>> torch.cholesky_inverse(L) + tensor([[ 1.9314, 1.2251, -0.0889], + [ 1.2251, 2.4439, 0.2122], + [-0.0889, 0.2122, 0.1412]]) + >>> A.inverse() + tensor([[ 1.9314, 1.2251, -0.0889], + [ 1.2251, 2.4439, 0.2122], + [-0.0889, 0.2122, 0.1412]]) + + >>> A = torch.randn(3, 2, 2, dtype=torch.complex64) + >>> A = A @ A.mH + torch.eye(2) * 1e-3 # Batch of Hermitian positive-definite matrices + >>> L = torch.linalg.cholesky(A) + >>> torch.dist(torch.inverse(A), torch.cholesky_inverse(L)) + tensor(5.6358e-7) +""", +) + +add_docstr( + torch.clone, + r""" +clone(input, *, memory_format=torch.preserve_format) -> Tensor + +Returns a copy of :attr:`input`. + +.. note:: + + This function is differentiable, so gradients will flow back from the + result of this operation to :attr:`input`. To create a tensor without an + autograd relationship to :attr:`input` see :meth:`~Tensor.detach`. + + In addition, when ``torch.preserve_format`` is used: + If the input tensor is dense (i.e., non-overlapping strided), + its memory format (including strides) is retained. + Otherwise (e.g., a non-dense view like a stepped slice), + the output is converted to the dense (contiguous) format. + +Args: + {input} + +Keyword args: + {memory_format} +""".format(**common_args), +) + +add_docstr( + torch.clamp, + r""" +clamp(input, min=None, max=None, *, out=None) -> Tensor + +Clamps all elements in :attr:`input` into the range `[` :attr:`min`, :attr:`max` `]`. +Letting min_value and max_value be :attr:`min` and :attr:`max`, respectively, this returns: + +.. math:: + y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i) + +If :attr:`min` is ``None``, there is no lower bound. +Or, if :attr:`max` is ``None`` there is no upper bound. +""" + + r""" + +.. note:: + If :attr:`min` is greater than :attr:`max` :func:`torch.clamp(..., min, max) ` + sets all elements in :attr:`input` to the value of :attr:`max`. + +Args: + {input} + min (Number or Tensor, optional): lower-bound of the range to be clamped to + max (Number or Tensor, optional): upper-bound of the range to be clamped to + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-1.7120, 0.1734, -0.0478, -0.0922]) + >>> torch.clamp(a, min=-0.5, max=0.5) + tensor([-0.5000, 0.1734, -0.0478, -0.0922]) + + >>> min = torch.linspace(-1, 1, steps=4) + >>> torch.clamp(a, min=min) + tensor([-1.0000, 0.1734, 0.3333, 1.0000]) + +""".format(**common_args), +) + +add_docstr( + torch.clip, + r""" +clip(input, min=None, max=None, *, out=None) -> Tensor + +Alias for :func:`torch.clamp`. +""", +) + +add_docstr( + torch.column_stack, + r""" +column_stack(tensors, *, out=None) -> Tensor + +Creates a new tensor by horizontally stacking the tensors in :attr:`tensors`. + +Equivalent to ``torch.hstack(tensors)``, except each zero or one dimensional tensor ``t`` +in :attr:`tensors` is first reshaped into a ``(t.numel(), 1)`` column before being stacked horizontally. + +Args: + tensors (sequence of Tensors): sequence of tensors to concatenate + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5, 6]) + >>> torch.column_stack((a, b)) + tensor([[1, 4], + [2, 5], + [3, 6]]) + >>> a = torch.arange(5) + >>> b = torch.arange(10).reshape(5, 2) + >>> torch.column_stack((a, b, b)) + tensor([[0, 0, 1, 0, 1], + [1, 2, 3, 2, 3], + [2, 4, 5, 4, 5], + [3, 6, 7, 6, 7], + [4, 8, 9, 8, 9]]) + +""".format(**common_args), +) + +add_docstr( + torch.complex, + r""" +complex(real, imag, *, out=None) -> Tensor + +Constructs a complex tensor with its real part equal to :attr:`real` and its +imaginary part equal to :attr:`imag`. + +Args: + real (Tensor): The real part of the complex tensor. Must be half, float or double. + imag (Tensor): The imaginary part of the complex tensor. Must be same dtype + as :attr:`real`. + +Keyword args: + out (Tensor): If the inputs are ``torch.float32``, must be + ``torch.complex64``. If the inputs are ``torch.float64``, must be + ``torch.complex128``. + +Example:: + + >>> real = torch.tensor([1, 2], dtype=torch.float32) + >>> imag = torch.tensor([3, 4], dtype=torch.float32) + >>> z = torch.complex(real, imag) + >>> z + tensor([(1.+3.j), (2.+4.j)]) + >>> z.dtype + torch.complex64 + +""", +) + +add_docstr( + torch.polar, + r""" +polar(abs, angle, *, out=None) -> Tensor + +Constructs a complex tensor whose elements are Cartesian coordinates +corresponding to the polar coordinates with absolute value :attr:`abs` and angle +:attr:`angle`. + +.. math:: + \text{out} = \text{abs} \cdot \cos(\text{angle}) + \text{abs} \cdot \sin(\text{angle}) \cdot j + +.. note:: + `torch.polar` is similar to + `std::polar `_ + and does not compute the polar decomposition + of a complex tensor like Python's `cmath.polar` and SciPy's `linalg.polar` do. + The behavior of this function is undefined if `abs` is negative or NaN, or if `angle` is + infinite. + +""" + + r""" +Args: + abs (Tensor): The absolute value the complex tensor. Must be float or double. + angle (Tensor): The angle of the complex tensor. Must be same dtype as + :attr:`abs`. + +Keyword args: + out (Tensor): If the inputs are ``torch.float32``, must be + ``torch.complex64``. If the inputs are ``torch.float64``, must be + ``torch.complex128``. + +Example:: + + >>> import numpy as np + >>> abs = torch.tensor([1, 2], dtype=torch.float64) + >>> angle = torch.tensor([np.pi / 2, 5 * np.pi / 4], dtype=torch.float64) + >>> z = torch.polar(abs, angle) + >>> z + tensor([(0.0000+1.0000j), (-1.4142-1.4142j)], dtype=torch.complex128) +""", +) + +add_docstr( + torch.conj_physical, + r""" +conj_physical(input, *, out=None) -> Tensor + +Computes the element-wise conjugate of the given :attr:`input` tensor. +If :attr:`input` has a non-complex dtype, this function just returns :attr:`input`. + +.. note:: + This performs the conjugate operation regardless of the fact conjugate bit is set or not. + +.. warning:: In the future, :func:`torch.conj_physical` may return a non-writeable view for an :attr:`input` of + non-complex dtype. It's recommended that programs not modify the tensor returned by :func:`torch.conj_physical` + when :attr:`input` is of non-complex dtype to be compatible with this change. + +.. math:: + \text{out}_{i} = conj(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.conj_physical(torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j])) + tensor([-1 - 1j, -2 - 2j, 3 + 3j]) +""".format(**common_args), +) + +add_docstr( + torch.conj, + r""" +conj(input) -> Tensor + +Returns a view of :attr:`input` with a flipped conjugate bit. If :attr:`input` has a non-complex dtype, +this function just returns :attr:`input`. + +.. note:: + :func:`torch.conj` performs a lazy conjugation, but the actual conjugated tensor can be materialized + at any time using :func:`torch.resolve_conj`. + +.. warning:: In the future, :func:`torch.conj` may return a non-writeable view for an :attr:`input` of + non-complex dtype. It's recommended that programs not modify the tensor returned by :func:`torch.conj_physical` + when :attr:`input` is of non-complex dtype to be compatible with this change. + +Args: + {input} + +Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> x.is_conj() + False + >>> y = torch.conj(x) + >>> y.is_conj() + True +""".format(**common_args), +) + +add_docstr( + torch.resolve_conj, + r""" +resolve_conj(input) -> Tensor + +Returns a new tensor with materialized conjugation if :attr:`input`'s conjugate bit is set to `True`, +else returns :attr:`input`. The output tensor will always have its conjugate bit set to `False`. + +Args: + {input} + +Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> y = x.conj() + >>> y.is_conj() + True + >>> z = y.resolve_conj() + >>> z + tensor([-1 - 1j, -2 - 2j, 3 + 3j]) + >>> z.is_conj() + False +""".format(**common_args), +) + +add_docstr( + torch.resolve_neg, + r""" +resolve_neg(input) -> Tensor + +Returns a new tensor with materialized negation if :attr:`input`'s negative bit is set to `True`, +else returns :attr:`input`. The output tensor will always have its negative bit set to `False`. + +Args: + {input} + +Example:: + + >>> x = torch.tensor([-1 + 1j, -2 + 2j, 3 - 3j]) + >>> y = x.conj() + >>> z = y.imag + >>> z.is_neg() + True + >>> out = z.resolve_neg() + >>> out + tensor([-1., -2., 3.]) + >>> out.is_neg() + False +""".format(**common_args), +) + +add_docstr( + torch.copysign, + r""" +copysign(input, other, *, out=None) -> Tensor + +Create a new floating-point tensor with the magnitude of :attr:`input` and the sign of :attr:`other`, elementwise. + +.. math:: + \text{out}_{i} = \begin{cases} + -|\text{input}_{i}| & \text{if } \text{other}_{i} \leq -0.0 \\ + |\text{input}_{i}| & \text{if } \text{other}_{i} \geq 0.0 \\ + \end{cases} +""" + + r""" + +Supports :ref:`broadcasting to a common shape `, +and integer and float inputs. + +Args: + input (Tensor): magnitudes. + other (Tensor or Number): contains value(s) whose signbit(s) are + applied to the magnitudes in :attr:`input`. + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(5) + >>> a + tensor([-1.2557, -0.0026, -0.5387, 0.4740, -0.9244]) + >>> torch.copysign(a, 1) + tensor([1.2557, 0.0026, 0.5387, 0.4740, 0.9244]) + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.7079, 0.2778, -1.0249, 0.5719], + [-0.0059, -0.2600, -0.4475, -1.3948], + [ 0.3667, -0.9567, -2.5757, -0.1751], + [ 0.2046, -0.0742, 0.2998, -0.1054]]) + >>> b = torch.randn(4) + tensor([ 0.2373, 0.3120, 0.3190, -1.1128]) + >>> torch.copysign(a, b) + tensor([[ 0.7079, 0.2778, 1.0249, -0.5719], + [ 0.0059, 0.2600, 0.4475, -1.3948], + [ 0.3667, 0.9567, 2.5757, -0.1751], + [ 0.2046, 0.0742, 0.2998, -0.1054]]) + >>> a = torch.tensor([1.]) + >>> b = torch.tensor([-0.]) + >>> torch.copysign(a, b) + tensor([-1.]) + +.. note:: + copysign handles signed zeros. If the other argument has a negative zero (-0), + the corresponding output value will be negative. + +""".format(**common_args), +) + +add_docstr( + torch.cos, + r""" +cos(input, *, out=None) -> Tensor + +Returns a new tensor with the cosine of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \cos(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 1.4309, 1.2706, -0.8562, 0.9796]) + >>> torch.cos(a) + tensor([ 0.1395, 0.2957, 0.6553, 0.5574]) +""".format(**common_args), +) + +add_docstr( + torch.cosh, + r""" +cosh(input, *, out=None) -> Tensor + +Returns a new tensor with the hyperbolic cosine of the elements of +:attr:`input`. + +.. math:: + \text{out}_{i} = \cosh(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.1632, 1.1835, -0.6979, -0.7325]) + >>> torch.cosh(a) + tensor([ 1.0133, 1.7860, 1.2536, 1.2805]) + +.. note:: + When :attr:`input` is on the CPU, the implementation of torch.cosh may use + the Sleef library, which rounds very large results to infinity or negative + infinity. See `here `_ for details. +""".format(**common_args), +) + +add_docstr( + torch.cross, + r""" +cross(input, other, dim=None, *, out=None) -> Tensor + + +Returns the cross product of vectors in dimension :attr:`dim` of :attr:`input` +and :attr:`other`. + +Supports input of float, double, cfloat and cdouble dtypes. Also supports batches +of vectors, for which it computes the product along the dimension :attr:`dim`. +In this case, the output has the same batch dimensions as the inputs. + +.. warning:: + If :attr:`dim` is not given, it defaults to the first dimension found + with the size 3. Note that this might be unexpected. + + This behavior is deprecated and will be changed to match that of :func:`torch.linalg.cross` + in a future release. + +.. seealso:: + :func:`torch.linalg.cross` which has dim=-1 as default. + + +Args: + {input} + other (Tensor): the second input tensor + dim (int, optional): the dimension to take the cross-product in. + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4, 3) + >>> a + tensor([[-0.3956, 1.1455, 1.6895], + [-0.5849, 1.3672, 0.3599], + [-1.1626, 0.7180, -0.0521], + [-0.1339, 0.9902, -2.0225]]) + >>> b = torch.randn(4, 3) + >>> b + tensor([[-0.0257, -1.4725, -1.2251], + [-1.1479, -0.7005, -1.9757], + [-1.3904, 0.3726, -1.1836], + [-0.9688, -0.7153, 0.2159]]) + >>> torch.cross(a, b, dim=1) + tensor([[ 1.0844, -0.5281, 0.6120], + [-2.4490, -1.5687, 1.9792], + [-0.8304, -1.3037, 0.5650], + [-1.2329, 1.9883, 1.0551]]) + >>> torch.cross(a, b) + tensor([[ 1.0844, -0.5281, 0.6120], + [-2.4490, -1.5687, 1.9792], + [-0.8304, -1.3037, 0.5650], + [-1.2329, 1.9883, 1.0551]]) +""".format(**common_args), +) + +add_docstr( + torch.logcumsumexp, + r""" +logcumsumexp(input, dim, *, out=None) -> Tensor +Returns the logarithm of the cumulative summation of the exponentiation of +elements of :attr:`input` in the dimension :attr:`dim`. + +For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{{logcumsumexp}}(x)_{{ij}} = \log \sum\limits_{{k=0}}^{{j}} \exp(x_{{ik}}) + +Args: + {input} + dim (int): the dimension to do the operation over + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(10) + >>> torch.logcumsumexp(a, dim=0) + tensor([-0.42296738, -0.04462666, 0.86278635, 0.94622083, 1.05277811, + 1.39202815, 1.83525007, 1.84492621, 2.06084887, 2.06844475])) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.cummax, + r""" +cummax(input, dim, *, out=None) -> (Tensor, LongTensor) +Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative maximum of +elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index +location of each maximum value found in the dimension :attr:`dim`. + +.. math:: + y_i = max(x_1, x_2, x_3, \dots, x_i) + +Args: + {input} + dim (int): the dimension to do the operation over + +Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + +Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.3449, -1.5447, 0.0685, -1.5104, -1.1706, 0.2259, 1.4696, -1.3284, + 1.9946, -0.8209]) + >>> torch.cummax(a, dim=0) + torch.return_types.cummax( + values=tensor([-0.3449, -0.3449, 0.0685, 0.0685, 0.0685, 0.2259, 1.4696, 1.4696, + 1.9946, 1.9946]), + indices=tensor([0, 0, 2, 2, 2, 5, 6, 6, 8, 8])) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.cummin, + r""" +cummin(input, dim, *, out=None) -> (Tensor, LongTensor) +Returns a namedtuple ``(values, indices)`` where ``values`` is the cumulative minimum of +elements of :attr:`input` in the dimension :attr:`dim`. And ``indices`` is the index +location of each maximum value found in the dimension :attr:`dim`. + +.. math:: + y_i = min(x_1, x_2, x_3, \dots, x_i) + +Args: + {input} + dim (int): the dimension to do the operation over + +Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + +Example:: + + >>> a = torch.randn(10) + >>> a + tensor([-0.2284, -0.6628, 0.0975, 0.2680, -1.3298, -0.4220, -0.3885, 1.1762, + 0.9165, 1.6684]) + >>> torch.cummin(a, dim=0) + torch.return_types.cummin( + values=tensor([-0.2284, -0.6628, -0.6628, -0.6628, -1.3298, -1.3298, -1.3298, -1.3298, + -1.3298, -1.3298]), + indices=tensor([0, 1, 1, 1, 4, 4, 4, 4, 4, 4])) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.cumprod, + r""" +cumprod(input, dim, *, dtype=None, out=None) -> Tensor + +Returns the cumulative product of elements of :attr:`input` in the dimension +:attr:`dim`. + +For example, if :attr:`input` is a vector of size N, the result will also be +a vector of size N, with elements. + +.. math:: + y_i = x_1 \times x_2\times x_3\times \dots \times x_i + +Args: + {input} + dim (int): the dimension to do the operation over + +Keyword args: + {dtype} + {out} + +Example:: + + >>> a = torch.randn(10) + >>> a + tensor([ 0.6001, 0.2069, -0.1919, 0.9792, 0.6727, 1.0062, 0.4126, + -0.2129, -0.4206, 0.1968]) + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0158, -0.0065, + 0.0014, -0.0006, -0.0001]) + + >>> a[5] = 0.0 + >>> torch.cumprod(a, dim=0) + tensor([ 0.6001, 0.1241, -0.0238, -0.0233, -0.0157, -0.0000, -0.0000, + 0.0000, -0.0000, -0.0000]) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.cumsum, + r""" +cumsum(input, dim, *, dtype=None, out=None) -> Tensor + +Returns the cumulative sum of elements of :attr:`input` in the dimension +:attr:`dim`. + +For example, if :attr:`input` is a vector of size N, the result will also be +a vector of size N, with elements. + +.. math:: + y_i = x_1 + x_2 + x_3 + \dots + x_i + +Args: + {input} + dim (int): the dimension to do the operation over + +Keyword args: + {dtype} + {out} + +Example:: + + >>> a = torch.randint(1, 20, (10,)) + >>> a + tensor([13, 7, 3, 10, 13, 3, 15, 10, 9, 10]) + >>> torch.cumsum(a, dim=0) + tensor([13, 20, 23, 33, 46, 49, 64, 74, 83, 93]) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.count_nonzero, + r""" +count_nonzero(input, dim=None) -> Tensor + +Counts the number of non-zero values in the tensor :attr:`input` along the given :attr:`dim`. +If no dim is specified then all non-zeros in the tensor are counted. + +Args: + {input} + dim (int or tuple of ints, optional): Dim or tuple of dims along which to count non-zeros. + +Example:: + + >>> x = torch.zeros(3,3) + >>> x[torch.randn(3,3) > 0.5] = 1 + >>> x + tensor([[0., 1., 1.], + [0., 0., 0.], + [0., 0., 1.]]) + >>> torch.count_nonzero(x) + tensor(3) + >>> torch.count_nonzero(x, dim=0) + tensor([0, 1, 2]) +""".format(**reduceops_common_args), +) + +add_docstr( + torch.dequantize, + r""" +dequantize(tensor) -> Tensor + +Returns an fp32 Tensor by dequantizing a quantized Tensor + +Args: + tensor (Tensor): A quantized Tensor + +.. function:: dequantize(tensors) -> sequence of Tensors + :noindex: + +Given a list of quantized Tensors, dequantize them and return a list of fp32 Tensors + +Args: + tensors (sequence of Tensors): A list of quantized Tensors +""", +) + +add_docstr( + torch.diag, + r""" +diag(input, diagonal=0, *, out=None) -> Tensor + +- If :attr:`input` is a vector (1-D tensor), then returns a 2-D square tensor + with the elements of :attr:`input` as the diagonal. +- If :attr:`input` is a matrix (2-D tensor), then returns a 1-D tensor with + the diagonal elements of :attr:`input`. + +The argument :attr:`diagonal` controls which diagonal to consider: + +- If :attr:`diagonal` = 0, it is the main diagonal. +- If :attr:`diagonal` > 0, it is above the main diagonal. +- If :attr:`diagonal` < 0, it is below the main diagonal. + +Args: + {input} + diagonal (int, optional): the diagonal to consider + +Keyword args: + {out} + +.. seealso:: + + :func:`torch.diagonal` always returns the diagonal of its input. + + :func:`torch.diagflat` always constructs a tensor with diagonal elements + specified by the input. + +Examples: + +Get the square matrix where the input vector is the diagonal:: + + >>> a = torch.randn(3) + >>> a + tensor([ 0.5950,-0.0872, 2.3298]) + >>> torch.diag(a) + tensor([[ 0.5950, 0.0000, 0.0000], + [ 0.0000,-0.0872, 0.0000], + [ 0.0000, 0.0000, 2.3298]]) + >>> torch.diag(a, 1) + tensor([[ 0.0000, 0.5950, 0.0000, 0.0000], + [ 0.0000, 0.0000,-0.0872, 0.0000], + [ 0.0000, 0.0000, 0.0000, 2.3298], + [ 0.0000, 0.0000, 0.0000, 0.0000]]) + +Get the k-th diagonal of a given matrix:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-0.4264, 0.0255,-0.1064], + [ 0.8795,-0.2429, 0.1374], + [ 0.1029,-0.6482,-1.6300]]) + >>> torch.diag(a, 0) + tensor([-0.4264,-0.2429,-1.6300]) + >>> torch.diag(a, 1) + tensor([ 0.0255, 0.1374]) +""".format(**common_args), +) + +add_docstr( + torch.diag_embed, + r""" +diag_embed(input, offset=0, dim1=-2, dim2=-1) -> Tensor + +Creates a tensor whose diagonals of certain 2D planes (specified by +:attr:`dim1` and :attr:`dim2`) are filled by :attr:`input`. +To facilitate creating batched diagonal matrices, the 2D planes formed by +the last two dimensions of the returned tensor are chosen by default. + +The argument :attr:`offset` controls which diagonal to consider: + +- If :attr:`offset` = 0, it is the main diagonal. +- If :attr:`offset` > 0, it is above the main diagonal. +- If :attr:`offset` < 0, it is below the main diagonal. + +The size of the new matrix will be calculated to make the specified diagonal +of the size of the last input dimension. +Note that for :attr:`offset` other than :math:`0`, the order of :attr:`dim1` +and :attr:`dim2` matters. Exchanging them is equivalent to changing the +sign of :attr:`offset`. + +Applying :meth:`torch.diagonal` to the output of this function with +the same arguments yields a matrix identical to input. However, +:meth:`torch.diagonal` has different default dimensions, so those +need to be explicitly specified. + +Args: + {input} Must be at least 1-dimensional. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: -2. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: -1. + +Example:: + + >>> a = torch.randn(2, 3) + >>> torch.diag_embed(a) + tensor([[[ 1.5410, 0.0000, 0.0000], + [ 0.0000, -0.2934, 0.0000], + [ 0.0000, 0.0000, -2.1788]], + + [[ 0.5684, 0.0000, 0.0000], + [ 0.0000, -1.0845, 0.0000], + [ 0.0000, 0.0000, -1.3986]]]) + + >>> torch.diag_embed(a, offset=1, dim1=0, dim2=2) + tensor([[[ 0.0000, 1.5410, 0.0000, 0.0000], + [ 0.0000, 0.5684, 0.0000, 0.0000]], + + [[ 0.0000, 0.0000, -0.2934, 0.0000], + [ 0.0000, 0.0000, -1.0845, 0.0000]], + + [[ 0.0000, 0.0000, 0.0000, -2.1788], + [ 0.0000, 0.0000, 0.0000, -1.3986]], + + [[ 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.0000, 0.0000, 0.0000, 0.0000]]]) +""".format(**common_args), +) + + +add_docstr( + torch.diagflat, + r""" +diagflat(input, offset=0) -> Tensor + +- If :attr:`input` is a vector (1-D tensor), then returns a 2-D square tensor + with the elements of :attr:`input` as the diagonal. +- If :attr:`input` is a tensor with more than one dimension, then returns a + 2-D tensor with diagonal elements equal to a flattened :attr:`input`. + +The argument :attr:`offset` controls which diagonal to consider: + +- If :attr:`offset` = 0, it is the main diagonal. +- If :attr:`offset` > 0, it is above the main diagonal. +- If :attr:`offset` < 0, it is below the main diagonal. + +Args: + {input} + offset (int, optional): the diagonal to consider. Default: 0 (main + diagonal). + +Examples:: + + >>> a = torch.randn(3) + >>> a + tensor([-0.2956, -0.9068, 0.1695]) + >>> torch.diagflat(a) + tensor([[-0.2956, 0.0000, 0.0000], + [ 0.0000, -0.9068, 0.0000], + [ 0.0000, 0.0000, 0.1695]]) + >>> torch.diagflat(a, 1) + tensor([[ 0.0000, -0.2956, 0.0000, 0.0000], + [ 0.0000, 0.0000, -0.9068, 0.0000], + [ 0.0000, 0.0000, 0.0000, 0.1695], + [ 0.0000, 0.0000, 0.0000, 0.0000]]) + + >>> a = torch.randn(2, 2) + >>> a + tensor([[ 0.2094, -0.3018], + [-0.1516, 1.9342]]) + >>> torch.diagflat(a) + tensor([[ 0.2094, 0.0000, 0.0000, 0.0000], + [ 0.0000, -0.3018, 0.0000, 0.0000], + [ 0.0000, 0.0000, -0.1516, 0.0000], + [ 0.0000, 0.0000, 0.0000, 1.9342]]) +""".format(**common_args), +) + +add_docstr( + torch.diagonal, + r""" +diagonal(input, offset=0, dim1=0, dim2=1) -> Tensor + +Returns a partial view of :attr:`input` with the its diagonal elements +with respect to :attr:`dim1` and :attr:`dim2` appended as a dimension +at the end of the shape. + +The argument :attr:`offset` controls which diagonal to consider: + +- If :attr:`offset` = 0, it is the main diagonal. +- If :attr:`offset` > 0, it is above the main diagonal. +- If :attr:`offset` < 0, it is below the main diagonal. + +Applying :meth:`torch.diag_embed` to the output of this function with +the same arguments yields a diagonal matrix with the diagonal entries +of the input. However, :meth:`torch.diag_embed` has different default +dimensions, so those need to be explicitly specified. + +Args: + {input} Must be at least 2-dimensional. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: 0. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: 1. + +.. note:: To take a batch diagonal, pass in dim1=-2, dim2=-1. + +Examples:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-1.0854, 1.1431, -0.1752], + [ 0.8536, -0.0905, 0.0360], + [ 0.6927, -0.3735, -0.4945]]) + + + >>> torch.diagonal(a) + tensor([-1.0854, -0.0905, -0.4945]) + + + >>> torch.diagonal(a, 1) + tensor([ 1.1431, 0.0360]) + + >>> b = torch.randn(2, 5) + >>> b + tensor([[-1.7948, -1.2731, -0.3181, 2.0200, -1.6745], + [ 1.8262, -1.5049, 0.4114, 1.0704, -1.2607]]) + + >>> torch.diagonal(b, 1, 1, 0) + tensor([1.8262]) + + >>> x = torch.randn(2, 5, 4, 2) + >>> torch.diagonal(x, offset=-1, dim1=1, dim2=2) + tensor([[[-1.2631, 0.3755, -1.5977, -1.8172], + [-1.1065, 1.0401, -0.2235, -0.7938]], + + [[-1.7325, -0.3081, 0.6166, 0.2335], + [ 1.0500, 0.7336, -0.3836, -1.1015]]]) +""".format(**common_args), +) + +add_docstr( + torch.diagonal_scatter, + r""" +diagonal_scatter(input, src, offset=0, dim1=0, dim2=1) -> Tensor + +Embeds the values of the :attr:`src` tensor into :attr:`input` along +the diagonal elements of :attr:`input`, with respect to :attr:`dim1` +and :attr:`dim2`. + +This function returns a tensor with fresh storage; it does not +return a view. + +The argument :attr:`offset` controls which diagonal to consider: + +- If :attr:`offset` = 0, it is the main diagonal. +- If :attr:`offset` > 0, it is above the main diagonal. +- If :attr:`offset` < 0, it is below the main diagonal. + +Args: + {input} Must be at least 2-dimensional. + src (Tensor): the tensor to embed into :attr:`input`. + offset (int, optional): which diagonal to consider. Default: 0 + (main diagonal). + dim1 (int, optional): first dimension with respect to which to + take diagonal. Default: 0. + dim2 (int, optional): second dimension with respect to which to + take diagonal. Default: 1. + +.. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + ``torch.diagonal(input, offset, dim1, dim2)`` + +Examples:: + + >>> a = torch.zeros(3, 3) + >>> a + tensor([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]]) + + >>> torch.diagonal_scatter(a, torch.ones(3), 0) + tensor([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + + >>> torch.diagonal_scatter(a, torch.ones(2), 1) + tensor([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) +""".format(**common_args), +) + +add_docstr( + torch.as_strided_scatter, + r""" +as_strided_scatter(input, src, size, stride, storage_offset=None) -> Tensor + +Embeds the values of the :attr:`src` tensor into :attr:`input` along +the elements corresponding to the result of calling +input.as_strided(size, stride, storage_offset). + +This function returns a tensor with fresh storage; it does not +return a view. + +Args: + {input} + size (tuple or ints): the shape of the output tensor + stride (tuple or ints): the stride of the output tensor + storage_offset (int, optional): the offset in the underlying storage of the output tensor + +.. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + `torch.as_strided(input, size, stride, storage_offset)` + +Example:: + + >>> a = torch.arange(4).reshape(2, 2) + 1 + >>> a + tensor([[1, 2], + [3, 4]]) + >>> b = torch.zeros(3, 3) + >>> b + tensor([[0., 0., 0.], + [0., 0., 0.], + [0., 0., 0.]]) + >>> torch.as_strided_scatter(b, a, (2, 2), (1, 2)) + tensor([[1., 3., 2.], + [4., 0., 0.], + [0., 0., 0.]]) + +""".format(**common_args), +) + +add_docstr( + torch.diff, + r""" +diff(input, n=1, dim=-1, prepend=None, append=None) -> Tensor + +Computes the n-th forward difference along the given dimension. + +The first-order differences are given by `out[i] = input[i + 1] - input[i]`. Higher-order +differences are calculated by using :func:`torch.diff` recursively. + +Args: + input (Tensor): the tensor to compute the differences on + n (int, optional): the number of times to recursively compute the difference + dim (int, optional): the dimension to compute the difference along. + Default is the last dimension. + prepend, append (Tensor, optional): values to prepend or append to + :attr:`input` along :attr:`dim` before computing the difference. + Their dimensions must be equivalent to that of input, and their shapes + must match input's shape except on :attr:`dim`. + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 3, 2]) + >>> torch.diff(a) + tensor([ 2, -1]) + >>> b = torch.tensor([4, 5]) + >>> torch.diff(a, append=b) + tensor([ 2, -1, 2, 1]) + >>> c = torch.tensor([[1, 2, 3], [3, 4, 5]]) + >>> torch.diff(c, dim=0) + tensor([[2, 2, 2]]) + >>> torch.diff(c, dim=1) + tensor([[1, 1], + [1, 1]]) +""".format(**common_args), +) + +add_docstr( + torch.digamma, + r""" +digamma(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.digamma`. +""", +) + +add_docstr( + torch.dist, + r""" +dist(input, other, p=2) -> Tensor + +Returns the p-norm of (:attr:`input` - :attr:`other`) + +The shapes of :attr:`input` and :attr:`other` must be +:ref:`broadcastable `. + +Args: + {input} + other (Tensor): the Right-hand-side input tensor + p (float, optional): the norm to be computed + +Example:: + + >>> x = torch.randn(4) + >>> x + tensor([-1.5393, -0.8675, 0.5916, 1.6321]) + >>> y = torch.randn(4) + >>> y + tensor([ 0.0967, -1.0511, 0.6295, 0.8360]) + >>> torch.dist(x, y, 3.5) + tensor(1.6727) + >>> torch.dist(x, y, 3) + tensor(1.6973) + >>> torch.dist(x, y, 0) + tensor(4.) + >>> torch.dist(x, y, 1) + tensor(2.6537) +""".format(**common_args), +) + +add_docstr( + torch.div, + r""" +div(input, other, *, rounding_mode=None, out=None) -> Tensor + +Divides each element of the input ``input`` by the corresponding element of +:attr:`other`. + +.. math:: + \text{{out}}_i = \frac{{\text{{input}}_i}}{{\text{{other}}_i}} + +.. note:: + By default, this performs a "true" division like Python 3. + See the :attr:`rounding_mode` argument for floor division. + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer, float, and complex inputs. +Always promotes integer types to the default scalar type. + +Args: + input (Tensor): the dividend + other (Tensor or Number): the divisor + +Keyword args: + rounding_mode (str, optional): Type of rounding applied to the result: + + * None - default behavior. Performs no rounding and, if both :attr:`input` and + :attr:`other` are integer types, promotes the inputs to the default scalar type. + Equivalent to true division in Python (the ``/`` operator) and NumPy's ``np.true_divide``. + * ``"trunc"`` - rounds the results of the division towards zero. + Equivalent to C-style integer division. + * ``"floor"`` - rounds the results of the division down. + Equivalent to floor division in Python (the ``//`` operator) and NumPy's ``np.floor_divide``. + + {out} + +Examples:: + + >>> x = torch.tensor([ 0.3810, 1.2774, -0.2972, -0.3719, 0.4637]) + >>> torch.div(x, 0.5) + tensor([ 0.7620, 2.5548, -0.5944, -0.7438, 0.9274]) + + >>> a = torch.tensor([[-0.3711, -1.9353, -0.4605, -0.2917], + ... [ 0.1815, -1.0111, 0.9805, -1.5923], + ... [ 0.1062, 1.4581, 0.7759, -1.2344], + ... [-0.1830, -0.0313, 1.1908, -1.4757]]) + >>> b = torch.tensor([ 0.8032, 0.2930, -0.8113, -0.2308]) + >>> torch.div(a, b) + tensor([[-0.4620, -6.6051, 0.5676, 1.2639], + [ 0.2260, -3.4509, -1.2086, 6.8990], + [ 0.1322, 4.9764, -0.9564, 5.3484], + [-0.2278, -0.1068, -1.4678, 6.3938]]) + + >>> torch.div(a, b, rounding_mode='trunc') + tensor([[-0., -6., 0., 1.], + [ 0., -3., -1., 6.], + [ 0., 4., -0., 5.], + [-0., -0., -1., 6.]]) + + >>> torch.div(a, b, rounding_mode='floor') + tensor([[-1., -7., 0., 1.], + [ 0., -4., -2., 6.], + [ 0., 4., -1., 5.], + [-1., -1., -2., 6.]]) + +""".format(**common_args), +) + +add_docstr( + torch.divide, + r""" +divide(input, other, *, rounding_mode=None, out=None) -> Tensor + +Alias for :func:`torch.div`. +""", +) + +add_docstr( + torch.dot, + r""" +dot(input, tensor, *, out=None) -> Tensor + +Computes the dot product of two 1D tensors. + +.. note:: + + Unlike NumPy's dot, torch.dot intentionally only supports computing the dot product + of two 1D tensors with the same number of elements. + +Args: + input (Tensor): first tensor in the dot product, must be 1D. + tensor (Tensor): second tensor in the dot product, must be 1D. + +Keyword args: + {out} + +Example:: + + >>> torch.dot(torch.tensor([2, 3]), torch.tensor([2, 1])) + tensor(7) + + >>> t1, t2 = torch.tensor([0, 1]), torch.tensor([2, 3]) + >>> torch.dot(t1, t2) + tensor(3) +""".format(**common_args), +) + +add_docstr( + torch.vdot, + r""" +vdot(input, other, *, out=None) -> Tensor + +Computes the dot product of two 1D vectors along a dimension. + +In symbols, this function computes + +.. math:: + + \sum_{i=1}^n \overline{x_i}y_i. + +where :math:`\overline{x_i}` denotes the conjugate for complex +vectors, and it is the identity for real vectors. + +.. note:: + + Unlike NumPy's vdot, torch.vdot intentionally only supports computing the dot product + of two 1D tensors with the same number of elements. + +.. seealso:: + + :func:`torch.linalg.vecdot` computes the dot product of two batches of vectors along a dimension. + +Args: + input (Tensor): first tensor in the dot product, must be 1D. Its conjugate is used if it's complex. + other (Tensor): second tensor in the dot product, must be 1D. + +Keyword args: +""" + + rf""" +.. note:: {common_args["out"]} +""" + + r""" + +Example:: + + >>> torch.vdot(torch.tensor([2, 3]), torch.tensor([2, 1])) + tensor(7) + >>> a = torch.tensor((1 +2j, 3 - 1j)) + >>> b = torch.tensor((2 +1j, 4 - 0j)) + >>> torch.vdot(a, b) + tensor([16.+1.j]) + >>> torch.vdot(b, a) + tensor([16.-1.j]) +""", +) + +add_docstr( + torch.eq, + r""" +eq(input, other, *, out=None) -> Tensor + +Computes element-wise equality + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is equal to :attr:`other` and False elsewhere + +Example:: + + >>> torch.eq(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[ True, False], + [False, True]]) +""".format(**common_args), +) + +add_docstr( + torch.equal, + r""" +equal(input, other) -> bool + +``True`` if two tensors have the same size and elements, ``False`` otherwise. + +.. note:: + + Tensors containing NaNs are never equal to each other. Additionally, this function does not + differentiate between the data types of the tensors during comparison. For more thorough tensor checks, + use :meth:`torch.testing.assert_close`. + +Example:: + + >>> torch.equal(torch.tensor([1, 2]), torch.tensor([1, 2])) + True + >>> torch.equal(torch.tensor([3, torch.nan]), torch.tensor([3, torch.nan])) + False + >>> torch.equal(torch.tensor([1, 2, 3], dtype=torch.int32), torch.tensor([1, 2, 3], dtype=torch.float32)) + True +""", +) + +add_docstr( + torch.erf, + r""" +erf(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.erf`. +""", +) + +add_docstr( + torch.erfc, + r""" +erfc(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.erfc`. +""", +) + +add_docstr( + torch.erfinv, + r""" +erfinv(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.erfinv`. +""", +) + +add_docstr( + torch.exp, + r""" +exp(input, *, out=None) -> Tensor + +Returns a new tensor with the exponential of the elements +of the input tensor :attr:`input`. + +.. math:: + y_{i} = e^{x_{i}} +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.exp(torch.tensor([0, math.log(2.)])) + tensor([ 1., 2.]) +""".format(**common_args), +) + +add_docstr( + torch.exp2, + r""" +exp2(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.exp2`. +""", +) + +add_docstr( + torch.expm1, + r""" +expm1(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.expm1`. +""", +) + +add_docstr( + torch.eye, + r""" +eye(n, m=None, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a 2-D tensor with ones on the diagonal and zeros elsewhere. + +Args: + n (int): the number of rows + m (int, optional): the number of columns with default being :attr:`n` + +Keyword arguments: + {out} + {dtype} + {layout} + {device} + {requires_grad} + +Returns: + Tensor: A 2-D tensor with ones on the diagonal and zeros elsewhere + +Example:: + + >>> torch.eye(3) + tensor([[ 1., 0., 0.], + [ 0., 1., 0.], + [ 0., 0., 1.]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.floor, + r""" +floor(input, *, out=None) -> Tensor + +Returns a new tensor with the floor of the elements of :attr:`input`, +the largest integer less than or equal to each element. + +For integer inputs, follows the array-api convention of returning a +copy of the input tensor. + +.. math:: + \text{out}_{i} = \left\lfloor \text{input}_{i} \right\rfloor +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.8166, 1.5308, -0.2530, -0.2091]) + >>> torch.floor(a) + tensor([-1., 1., -1., -1.]) +""".format(**common_args), +) + +add_docstr( + torch.floor_divide, + r""" +floor_divide(input, other, *, out=None) -> Tensor + +.. note:: + + Before PyTorch 1.13 :func:`torch.floor_divide` incorrectly performed + truncation division. To restore the previous behavior use + :func:`torch.div` with ``rounding_mode='trunc'``. + +Computes :attr:`input` divided by :attr:`other`, elementwise, and floors +the result. + +.. math:: + \text{{out}}_i = \text{floor} \left( \frac{{\text{{input}}_i}}{{\text{{other}}_i}} \right) + +""" + + r""" + +Supports broadcasting to a common shape, type promotion, and integer and float inputs. + +Args: + input (Tensor or Number): the dividend + other (Tensor or Number): the divisor + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([4.0, 3.0]) + >>> b = torch.tensor([2.0, 2.0]) + >>> torch.floor_divide(a, b) + tensor([2.0, 1.0]) + >>> torch.floor_divide(a, 1.4) + tensor([2.0, 2.0]) +""".format(**common_args), +) + +add_docstr( + torch.fmod, + r""" +fmod(input, other, *, out=None) -> Tensor + +Applies C++'s `std::fmod `_ entrywise. +The result has the same sign as the dividend :attr:`input` and its absolute value +is less than that of :attr:`other`. + +This function may be defined in terms of :func:`torch.div` as + +.. code:: python + + torch.fmod(a, b) == a - a.div(b, rounding_mode="trunc") * b + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer and float inputs. + +.. note:: + + When the divisor is zero, returns ``NaN`` for floating point dtypes + on both CPU and GPU; raises ``RuntimeError`` for integer division by + zero on CPU; Integer division by zero on GPU may return any value. + +.. note:: + + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + +.. seealso:: + + :func:`torch.remainder` which implements Python's modulus operator. + This one is defined using division rounding down the result. + +Args: + input (Tensor): the dividend + other (Tensor or Scalar): the divisor + +Keyword args: + {out} + +Example:: + + >>> torch.fmod(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([-1., -0., -1., 1., 0., 1.]) + >>> torch.fmod(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([1.0000, 0.5000, 0.0000, 1.0000, 0.5000]) + +""".format(**common_args), +) + +add_docstr( + torch.frac, + r""" +frac(input, *, out=None) -> Tensor + +Computes the fractional portion of each element in :attr:`input`. + +.. math:: + \text{out}_{i} = \text{input}_{i} - \left\lfloor |\text{input}_{i}| \right\rfloor * \operatorname{sgn}(\text{input}_{i}) + +Example:: + + >>> torch.frac(torch.tensor([1, 2.5, -3.2])) + tensor([ 0.0000, 0.5000, -0.2000]) +""", +) + +add_docstr( + torch.frexp, + r""" +frexp(input, *, out=None) -> (Tensor mantissa, Tensor exponent) + +Decomposes :attr:`input` into mantissa and exponent tensors +such that :math:`\text{input} = \text{mantissa} \times 2^{\text{exponent}}`. + +The range of mantissa is the open interval (-1, 1). + +Supports float inputs. + +Args: + input (Tensor): the input tensor + + +Keyword args: + out (tuple, optional): the output tensors + +Example:: + + >>> x = torch.arange(9.) + >>> mantissa, exponent = torch.frexp(x) + >>> mantissa + tensor([0.0000, 0.5000, 0.5000, 0.7500, 0.5000, 0.6250, 0.7500, 0.8750, 0.5000]) + >>> exponent + tensor([0, 1, 2, 2, 3, 3, 3, 3, 4], dtype=torch.int32) + >>> torch.ldexp(mantissa, exponent) + tensor([0., 1., 2., 3., 4., 5., 6., 7., 8.]) +""", +) + +add_docstr( + torch.from_numpy, + r""" +from_numpy(ndarray) -> Tensor + +Creates a :class:`Tensor` from a :class:`numpy.ndarray`. + +The returned tensor and :attr:`ndarray` share the same memory. Modifications to +the tensor will be reflected in the :attr:`ndarray` and vice versa. The returned +tensor is not resizable. + +It currently accepts :attr:`ndarray` with dtypes of ``numpy.float64``, +``numpy.float32``, ``numpy.float16``, ``numpy.complex64``, ``numpy.complex128``, +``numpy.int64``, ``numpy.int32``, ``numpy.int16``, ``numpy.int8``, ``numpy.uint8``, +and ``bool``. + +.. warning:: + Writing to a tensor created from a read-only NumPy array is not supported and will result in undefined behavior. + +Example:: + + >>> a = numpy.array([1, 2, 3]) + >>> t = torch.from_numpy(a) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) +""", +) + +add_docstr( + torch.frombuffer, + r""" +frombuffer(buffer, *, dtype, count=-1, offset=0, requires_grad=False) -> Tensor + +Creates a 1-dimensional :class:`Tensor` from an object that implements +the Python buffer protocol. + +Skips the first :attr:`offset` bytes in the buffer, and interprets the rest of +the raw bytes as a 1-dimensional tensor of type :attr:`dtype` with :attr:`count` +elements. + +Note that either of the following must be true: + +1. :attr:`count` is a positive non-zero number, and the total number of bytes +in the buffer is more than :attr:`offset` plus :attr:`count` times the size +(in bytes) of :attr:`dtype`. + +2. :attr:`count` is negative, and the length (number of bytes) of the buffer +subtracted by the :attr:`offset` is a multiple of the size (in bytes) of +:attr:`dtype`. + +The returned tensor and buffer share the same memory. Modifications to +the tensor will be reflected in the buffer and vice versa. The returned +tensor is not resizable. + +.. note:: + This function increments the reference count for the object that + owns the shared memory. Therefore, such memory will not be deallocated + before the returned tensor goes out of scope. + +.. warning:: + This function's behavior is undefined when passed an object implementing + the buffer protocol whose data is not on the CPU. Doing so is likely to + cause a segmentation fault. + +.. warning:: + This function does not try to infer the :attr:`dtype` (hence, it is not + optional). Passing a different :attr:`dtype` than its source may result + in unexpected behavior. + +Args: + buffer (object): a Python object that exposes the buffer interface. + +Keyword args: + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + count (int, optional): the number of desired elements to be read. + If negative, all the elements (until the end of the buffer) will be + read. Default: -1. + offset (int, optional): the number of bytes to skip at the start of + the buffer. Default: 0. + {requires_grad} + +Example:: + + >>> import array + >>> a = array.array('i', [1, 2, 3]) + >>> t = torch.frombuffer(a, dtype=torch.int32) + >>> t + tensor([ 1, 2, 3]) + >>> t[0] = -1 + >>> a + array([-1, 2, 3]) + + >>> # Interprets the signed char bytes as 32-bit integers. + >>> # Each 4 signed char elements will be interpreted as + >>> # 1 signed 32-bit integer. + >>> import array + >>> a = array.array('b', [-1, 0, 0, 0]) + >>> torch.frombuffer(a, dtype=torch.int32) + tensor([255], dtype=torch.int32) +""".format(**factory_common_args), +) + +add_docstr( + torch.from_file, + r""" +from_file(filename, shared=None, size=0, *, dtype=None, layout=None, device=None, pin_memory=False) + +Creates a CPU tensor with a storage backed by a memory-mapped file. + +If ``shared`` is True, then memory is shared between processes. All changes are written to the file. +If ``shared`` is False, then changes to the tensor do not affect the file. + +``size`` is the number of elements in the Tensor. If ``shared`` is ``False``, then the file must contain +at least ``size * sizeof(dtype)`` bytes. If ``shared`` is ``True`` the file will be created if needed. + +.. note:: + Only CPU tensors can be mapped to files. + +.. note:: + For now, tensors with storages backed by a memory-mapped file cannot be created in pinned memory. + + +Args: + filename (str): file name to map + shared (bool): whether to share memory (whether ``MAP_SHARED`` or ``MAP_PRIVATE`` is passed to the + underlying `mmap(2) call `_) + size (int): number of elements in the tensor + +Keyword args: + {dtype} + {layout} + {device} + {pin_memory} + +Example:: + + >>> t = torch.randn(2, 5, dtype=torch.float64) + >>> t.numpy().tofile('storage.pt') + >>> t_mapped = torch.from_file('storage.pt', shared=False, size=10, dtype=torch.float64) + """.format(**factory_common_args), +) + +add_docstr( + torch.flatten, + r""" +flatten(input, start_dim=0, end_dim=-1) -> Tensor + +Flattens :attr:`input` by reshaping it into a one-dimensional tensor. If :attr:`start_dim` or :attr:`end_dim` +are passed, only dimensions starting with :attr:`start_dim` and ending with :attr:`end_dim` are flattened. +The order of elements in :attr:`input` is unchanged. + +Unlike NumPy's flatten, which always copies input's data, this function may return the original object, a view, +or copy. If no dimensions are flattened, then the original object :attr:`input` is returned. Otherwise, if input can +be viewed as the flattened shape, then that view is returned. Finally, only if the input cannot be viewed as the +flattened shape is input's data copied. See :meth:`torch.Tensor.view` for details on when a view will be returned. + +.. note:: + Flattening a zero-dimensional tensor will return a one-dimensional view. + +Args: + {input} + start_dim (int): the first dim to flatten + end_dim (int): the last dim to flatten + +Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.flatten(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) + >>> torch.flatten(t, start_dim=1) + tensor([[1, 2, 3, 4], + [5, 6, 7, 8]]) +""".format(**common_args), +) + +add_docstr( + torch.unflatten, + r""" +unflatten(input, dim, sizes) -> Tensor + +Expands a dimension of the input tensor over multiple dimensions. + +.. seealso:: + + :func:`torch.flatten` the inverse of this function. It coalesces several dimensions into one. + +Args: + {input} + dim (int): Dimension to be unflattened, specified as an index into + ``input.shape``. + sizes (Tuple[int]): New shape of the unflattened dimension. + One of its elements can be `-1` in which case the corresponding output + dimension is inferred. Otherwise, the product of ``sizes`` *must* + equal ``input.shape[dim]``. + +Returns: + A View of input with the specified dimension unflattened. + +Examples:: + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (2, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(3, 4, 1), 1, (-1, 2)).shape + torch.Size([3, 2, 2, 1]) + >>> torch.unflatten(torch.randn(5, 12, 3), -2, (2, 2, 3, 1, 1)).shape + torch.Size([5, 2, 2, 3, 1, 1, 3]) +""".format(**common_args), +) + +add_docstr( + torch.gather, + r""" +gather(input, dim, index, *, sparse_grad=False, out=None) -> Tensor + +Gathers values along an axis specified by `dim`. + +For a 3-D tensor the output is specified by:: + + out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 + out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 + out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 + +:attr:`input` and :attr:`index` must have the same number of dimensions. +It is also required that ``index.size(d) <= input.size(d)`` for all +dimensions ``d != dim``. :attr:`out` will have the same shape as :attr:`index`. +Note that ``input`` and ``index`` do not broadcast against each other. +When :attr:`index` is empty, we always return an empty output with the same shape +without further error checking. + +Args: + input (Tensor): the source tensor + dim (int): the axis along which to index + index (LongTensor): the indices of elements to gather + +Keyword arguments: + sparse_grad (bool, optional): If ``True``, gradient w.r.t. :attr:`input` will be a sparse tensor. + out (Tensor, optional): the destination tensor + +Example:: + + >>> t = torch.tensor([[1, 2], [3, 4]]) + >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) + tensor([[ 1, 1], + [ 4, 3]]) +""", +) + + +add_docstr( + torch.gcd, + r""" +gcd(input, other, *, out=None) -> Tensor + +Computes the element-wise greatest common divisor (GCD) of :attr:`input` and :attr:`other`. + +Both :attr:`input` and :attr:`other` must have integer types. + +.. note:: + This defines :math:`gcd(0, 0) = 0`. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.tensor([5, 10, 15]) + >>> b = torch.tensor([3, 4, 5]) + >>> torch.gcd(a, b) + tensor([1, 2, 5]) + >>> c = torch.tensor([3]) + >>> torch.gcd(a, c) + tensor([1, 1, 3]) +""".format(**common_args), +) + +add_docstr( + torch.ge, + r""" +ge(input, other, *, out=None) -> Tensor + +Computes :math:`\text{input} \geq \text{other}` element-wise. +""" + + r""" + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is greater than or equal to :attr:`other` and False elsewhere + +Example:: + + >>> torch.ge(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, True], [False, True]]) +""".format(**common_args), +) + +add_docstr( + torch.greater_equal, + r""" +greater_equal(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.ge`. +""", +) + +add_docstr( + torch.gradient, + r""" +gradient(input, *, spacing=1, dim=None, edge_order=1) -> List of Tensors + +Estimates the gradient of a function :math:`g : \mathbb{R}^n \rightarrow \mathbb{R}` in +one or more dimensions using the `second-order accurate central differences method +`_ and +either first or second order estimates at the boundaries. + +The gradient of :math:`g` is estimated using samples. By default, when :attr:`spacing` is not +specified, the samples are entirely described by :attr:`input`, and the mapping of input coordinates +to an output is the same as the tensor's mapping of indices to values. For example, for a three-dimensional +:attr:`input` the function described is :math:`g : \mathbb{R}^3 \rightarrow \mathbb{R}`, and +:math:`g(1, 2, 3)\ == input[1, 2, 3]`. + +When :attr:`spacing` is specified, it modifies the relationship between :attr:`input` and input coordinates. +This is detailed in the "Keyword Arguments" section below. + +The gradient is estimated by estimating each partial derivative of :math:`g` independently. This estimation is +accurate if :math:`g` is in :math:`C^3` (it has at least 3 continuous derivatives), and the estimation can be +improved by providing closer samples. Mathematically, the value at each interior point of a partial derivative +is estimated using `Taylor's theorem with remainder `_. +Letting :math:`x` be an interior point with :math:`x-h_l` and :math:`x+h_r` be points neighboring +it to the left and right respectively, :math:`f(x+h_r)` and :math:`f(x-h_l)` can be estimated using: + +.. math:: + \begin{aligned} + f(x+h_r) = f(x) + h_r f'(x) + {h_r}^2 \frac{f''(x)}{2} + {h_r}^3 \frac{f'''(\xi_1)}{6}, \xi_1 \in (x, x+h_r) \\ + f(x-h_l) = f(x) - h_l f'(x) + {h_l}^2 \frac{f''(x)}{2} - {h_l}^3 \frac{f'''(\xi_2)}{6}, \xi_2 \in (x, x-h_l) \\ + \end{aligned} + +Using the fact that :math:`f \in C^3` and solving the linear system, we derive: + +.. math:: + f'(x) \approx \frac{ {h_l}^2 f(x+h_r) - {h_r}^2 f(x-h_l) + + ({h_r}^2-{h_l}^2 ) f(x) }{ {h_r} {h_l}^2 + {h_r}^2 {h_l} } + +.. note:: + We estimate the gradient of functions in complex domain + :math:`g : \mathbb{C}^n \rightarrow \mathbb{C}` in the same way. + +The value of each partial derivative at the boundary points is computed differently. See edge_order below. + +Args: + input (``Tensor``): the tensor that represents the values of the function + +Keyword args: + spacing (``scalar``, ``list of scalar``, ``list of Tensor``, optional): :attr:`spacing` can be used to modify + how the :attr:`input` tensor's indices relate to sample coordinates. If :attr:`spacing` is a scalar then + the indices are multiplied by the scalar to produce the coordinates. For example, if :attr:`spacing=2` the + indices (1, 2, 3) become coordinates (2, 4, 6). If :attr:`spacing` is a list of scalars then the corresponding + indices are multiplied. For example, if :attr:`spacing=(2, -1, 3)` the indices (1, 2, 3) become coordinates (2, -2, 9). + Finally, if :attr:`spacing` is a list of one-dimensional tensors then each tensor specifies the coordinates for + the corresponding dimension. For example, if the indices are (1, 2, 3) and the tensors are (t0, t1, t2), then + the coordinates are (t0[1], t1[2], t2[3]) + + dim (``int``, ``list of int``, optional): the dimension or dimensions to approximate the gradient over. By default + the partial gradient in every dimension is computed. Note that when :attr:`dim` is specified the elements of + the :attr:`spacing` argument must correspond with the specified dims." + + edge_order (``int``, optional): 1 or 2, for `first-order + `_ or + `second-order `_ + estimation of the boundary ("edge") values, respectively. Note that when :attr:`edge_order` is specified, each + dimension size of :attr:`input` should be at least edge_order+1 + +Examples:: + + >>> # Estimates the gradient of f(x)=x^2 at points [-2, -1, 2, 4] + >>> coordinates = (torch.tensor([-2., -1., 1., 4.]),) + >>> values = torch.tensor([4., 1., 1., 16.], ) + >>> torch.gradient(values, spacing = coordinates) + (tensor([-3., -2., 2., 5.]),) + + >>> # Estimates the gradient of the R^2 -> R function whose samples are + >>> # described by the tensor t. Implicit coordinates are [0, 1] for the outermost + >>> # dimension and [0, 1, 2, 3] for the innermost dimension, and function estimates + >>> # partial derivative for both dimensions. + >>> t = torch.tensor([[1, 2, 4, 8], [10, 20, 40, 80]]) + >>> torch.gradient(t) + (tensor([[ 9., 18., 36., 72.], + [ 9., 18., 36., 72.]]), + tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]])) + + >>> # A scalar value for spacing modifies the relationship between tensor indices + >>> # and input coordinates by multiplying the indices to find the + >>> # coordinates. For example, below the indices of the innermost + >>> # 0, 1, 2, 3 translate to coordinates of [0, 2, 4, 6], and the indices of + >>> # the outermost dimension 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = 2.0) # dim = None (implicitly [0, 1]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.5000, 0.7500, 1.5000, 2.0000], + [ 5.0000, 7.5000, 15.0000, 20.0000]])) + >>> # doubling the spacing between samples halves the estimated partial gradients. + + >>> + >>> # Estimates only the partial derivative for dimension 1 + >>> torch.gradient(t, dim = 1) # spacing = None (implicitly 1.) + (tensor([[ 1.0000, 1.5000, 3.0000, 4.0000], + [10.0000, 15.0000, 30.0000, 40.0000]]),) + + >>> # When spacing is a list of scalars, the relationship between the tensor + >>> # indices and input coordinates changes based on dimension. + >>> # For example, below, the indices of the innermost dimension 0, 1, 2, 3 translate + >>> # to coordinates of [0, 3, 6, 9], and the indices of the outermost dimension + >>> # 0, 1 translate to coordinates of [0, 2]. + >>> torch.gradient(t, spacing = [3., 2.]) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + + >>> # The following example is a replication of the previous one with explicit + >>> # coordinates. + >>> coords = (torch.tensor([0, 2]), torch.tensor([0, 3, 6, 9])) + >>> torch.gradient(t, spacing = coords) + (tensor([[ 4.5000, 9.0000, 18.0000, 36.0000], + [ 4.5000, 9.0000, 18.0000, 36.0000]]), + tensor([[ 0.3333, 0.5000, 1.0000, 1.3333], + [ 3.3333, 5.0000, 10.0000, 13.3333]])) + +""", +) + +add_docstr( + torch.geqrf, + r""" +geqrf(input, *, out=None) -> (Tensor, Tensor) + +This is a low-level function for calling LAPACK's geqrf directly. This function +returns a namedtuple (a, tau) as defined in `LAPACK documentation for geqrf`_ . + +Computes a QR decomposition of :attr:`input`. +Both `Q` and `R` matrices are stored in the same output tensor `a`. +The elements of `R` are stored on and above the diagonal. +Elementary reflectors (or Householder vectors) implicitly defining matrix `Q` +are stored below the diagonal. +The results of this function can be used together with :func:`torch.linalg.householder_product` +to obtain the `Q` matrix or +with :func:`torch.ormqr`, which uses an implicit representation of the `Q` matrix, +for an efficient matrix-matrix multiplication. + +See `LAPACK documentation for geqrf`_ for further details. + +.. note:: + See also :func:`torch.linalg.qr`, which computes Q and R matrices, and :func:`torch.linalg.lstsq` + with the ``driver="gels"`` option for a function that can solve matrix equations using a QR decomposition. + +Args: + input (Tensor): the input matrix + +Keyword args: + out (tuple, optional): the output tuple of (Tensor, Tensor). Ignored if `None`. Default: `None`. + +.. _LAPACK documentation for geqrf: + http://www.netlib.org/lapack/explore-html/df/dc5/group__variants_g_ecomputational_ga3766ea903391b5cf9008132f7440ec7b.html + +""", +) + +add_docstr( + torch.inner, + r""" +inner(input, other, *, out=None) -> Tensor + +Computes the dot product for 1D tensors. For higher dimensions, sums the product +of elements from :attr:`input` and :attr:`other` along their last dimension. + +.. note:: + + If either :attr:`input` or :attr:`other` is a scalar, the result is equivalent + to `torch.mul(input, other)`. + + If both :attr:`input` and :attr:`other` are non-scalars, the size of their last + dimension must match and the result is equivalent to `torch.tensordot(input, + other, dims=([-1], [-1]))` + +Args: + input (Tensor): First input tensor + other (Tensor): Second input tensor + +Keyword args: + out (Tensor, optional): Optional output tensor to write result into. The output + shape is `input.shape[:-1] + other.shape[:-1]`. + +Example:: + + # Dot product + >>> torch.inner(torch.tensor([1, 2, 3]), torch.tensor([0, 2, 1])) + tensor(7) + + # Multidimensional input tensors + >>> a = torch.randn(2, 3) + >>> a + tensor([[0.8173, 1.0874, 1.1784], + [0.3279, 0.1234, 2.7894]]) + >>> b = torch.randn(2, 4, 3) + >>> b + tensor([[[-0.4682, -0.7159, 0.1506], + [ 0.4034, -0.3657, 1.0387], + [ 0.9892, -0.6684, 0.1774], + [ 0.9482, 1.3261, 0.3917]], + + [[ 0.4537, 0.7493, 1.1724], + [ 0.2291, 0.5749, -0.2267], + [-0.7920, 0.3607, -0.3701], + [ 1.3666, -0.5850, -1.7242]]]) + >>> torch.inner(a, b) + tensor([[[-0.9837, 1.1560, 0.2907, 2.6785], + [ 2.5671, 0.5452, -0.6912, -1.5509]], + + [[ 0.1782, 2.9843, 0.7366, 1.5672], + [ 3.5115, -0.4864, -1.2476, -4.4337]]]) + + # Scalar input + >>> torch.inner(a, torch.tensor(2)) + tensor([[1.6347, 2.1748, 2.3567], + [0.6558, 0.2469, 5.5787]]) +""", +) + +add_docstr( + torch.outer, + r""" +outer(input, vec2, *, out=None) -> Tensor + +Outer product of :attr:`input` and :attr:`vec2`. +If :attr:`input` is a vector of size :math:`n` and :attr:`vec2` is a vector of +size :math:`m`, then :attr:`out` must be a matrix of size :math:`(n \times m)`. + +.. note:: This function does not :ref:`broadcast `. + +Args: + input (Tensor): 1-D input vector + vec2 (Tensor): 1-D input vector + +Keyword args: + out (Tensor, optional): optional output matrix + +Example:: + + >>> v1 = torch.arange(1., 5.) + >>> v2 = torch.arange(1., 4.) + >>> torch.outer(v1, v2) + tensor([[ 1., 2., 3.], + [ 2., 4., 6.], + [ 3., 6., 9.], + [ 4., 8., 12.]]) +""", +) + +add_docstr( + torch.ger, + r""" +ger(input, vec2, *, out=None) -> Tensor + +Alias of :func:`torch.outer`. + +.. warning:: + This function is deprecated and will be removed in a future PyTorch release. + Use :func:`torch.outer` instead. +""", +) + +add_docstr( + torch.get_default_dtype, + r""" +get_default_dtype() -> torch.dtype + +Get the current default floating point :class:`torch.dtype`. + +Example:: + + >>> torch.get_default_dtype() # initial default for floating point is torch.float32 + torch.float32 + >>> torch.set_default_dtype(torch.float64) + >>> torch.get_default_dtype() # default is now changed to torch.float64 + torch.float64 + +""", +) + +add_docstr( + torch.get_num_threads, + r""" +get_num_threads() -> int + +Returns the number of threads used for parallelizing CPU operations +""", +) + +add_docstr( + torch.get_num_interop_threads, + r""" +get_num_interop_threads() -> int + +Returns the number of threads used for inter-op parallelism on CPU +(e.g. in JIT interpreter) +""", +) + +add_docstr( + torch.gt, + r""" +gt(input, other, *, out=None) -> Tensor + +Computes :math:`\text{input} > \text{other}` element-wise. +""" + + r""" + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is greater than :attr:`other` and False elsewhere + +Example:: + + >>> torch.gt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [False, False]]) +""".format(**common_args), +) + +add_docstr( + torch.greater, + r""" +greater(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.gt`. +""", +) + +add_docstr( + torch.hash_tensor, + r""" +hash_tensor(input, *, mode=0) -> Tensor + +Returns a hash of all elements in the :attr:`input` tensor. + +Currently only mode=0 (reduction via xor) is supported. The output will always +be of type ``torch.uint64``. The elements of ``input`` are upcasted to their +64 bit float / integer equivalent and bitcasted to ``torch.uint64`` before +reduction via xor. + +Args: + {input} + +Keyword Args: + mode (int) : The hash to use. Default: 0 (xor_reduction) + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.1918, -1.1813, 0.3373]]) + >>> torch.hash_tensor(a) + tensor(13822780554648485888, dtype=torch.uint64) + +.. function:: hash_tensor(input, dim, *, keepdim=False, mode=0) -> Tensor + :noindex: + +Returns the hash of each row of the :attr:`input` tensor in the given +dimension :attr:`dim` given by mode. If :attr:`dim` is a list of dimensions, +reduce over all of them. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword Args: + mode (int) : The hash to use. Default: 0 (xor_reduction) + +Example:: + + >>> a = torch.randn(2, 4) + >>> a + tensor([[ 0.1317, -0.5554, -1.4724, -1.1391], + [ 0.0778, -0.6070, 0.6375, 0.1798]]) + >>> torch.hash_tensor(a, 1) + tensor([9233691267014066176, 9255993250844508160], dtype=torch.uint64) +""".format(**multi_dim_common), +) + +add_docstr( + torch.histc, + r""" +histc(input, bins=100, min=0, max=0, *, out=None) -> Tensor + +Computes the histogram of a tensor. + +The elements are sorted into equal width bins between :attr:`min` and +:attr:`max`. If :attr:`min` and :attr:`max` are both zero, the minimum and +maximum values of the data are used. + +Elements lower than min and higher than max and ``NaN`` elements are ignored. + +Args: + {input} + bins (int): number of histogram bins + min (Scalar): lower end of the range (inclusive) + max (Scalar): upper end of the range (inclusive) + +Keyword args: + {out} + +Returns: + Tensor: Histogram represented as a tensor + +Example:: + + >>> torch.histc(torch.tensor([1., 2, 1]), bins=4, min=0, max=3) + tensor([ 0., 2., 1., 0.]) +""".format(**common_args), +) + +add_docstr( + torch.histogram, + r""" +histogram(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor) + +Computes a histogram of the values in a tensor. + +:attr:`bins` can be an integer or a 1D tensor. + +If :attr:`bins` is an int, it specifies the number of equal-width bins. +By default, the lower and upper range of the bins is determined by the +minimum and maximum elements of the input tensor. The :attr:`range` +argument can be provided to specify a range for the bins. + +If :attr:`bins` is a 1D tensor, it specifies the sequence of bin edges +including the rightmost edge. It should contain at least 2 elements +and its elements should be increasing. + +Args: + {input} + bins: int or 1D Tensor. If int, defines the number of equal-width bins. If tensor, + defines the sequence of bin edges including the rightmost edge. + +Keyword args: + range (tuple of float): Defines the range of the bins. + weight (Tensor): If provided, weight should have the same shape as input. Each value in + input contributes its associated weight towards its bin's result. + density (bool): If False, the result will contain the count (or total weight) in each bin. + If True, the result is the value of the probability density function over the bins, + normalized such that the integral over the range of the bins is 1. + {out} (tuple, optional): The result tuple of two output tensors (hist, bin_edges). + +Returns: + hist (Tensor): 1D Tensor containing the values of the histogram. + bin_edges(Tensor): 1D Tensor containing the edges of the histogram bins. + +Example:: + + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.])) + (tensor([ 0., 5., 2., 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) + >>> torch.histogram(torch.tensor([1., 2, 1]), bins=4, range=(0., 3.), weight=torch.tensor([1., 2., 4.]), density=True) + (tensor([ 0., 0.9524, 0.3810, 0.]), tensor([0., 0.75, 1.5, 2.25, 3.])) +""".format(**common_args), +) + +add_docstr( + torch.histogramdd, + r""" +histogramdd(input, bins, *, range=None, weight=None, density=False, out=None) -> (Tensor, Tensor[]) + +Computes a multi-dimensional histogram of the values in a tensor. + +Interprets the elements of an input tensor whose innermost dimension has size N +as a collection of N-dimensional points. Maps each of the points into a set of +N-dimensional bins and returns the number of points (or total weight) in each bin. + +:attr:`input` must be a tensor with at least 2 dimensions. +If input has shape (M, N), each of its M rows defines a point in N-dimensional space. +If input has three or more dimensions, all but the last dimension are flattened. + +Each dimension is independently associated with its own strictly increasing sequence +of bin edges. Bin edges may be specified explicitly by passing a sequence of 1D +tensors. Alternatively, bin edges may be constructed automatically by passing a +sequence of integers specifying the number of equal-width bins in each dimension. + +For each N-dimensional point in input: + - Each of its coordinates is binned independently among the bin edges + corresponding to its dimension + - Binning results are combined to identify the N-dimensional bin (if any) + into which the point falls + - If the point falls into a bin, the bin's count (or total weight) is incremented + - Points which do not fall into any bin do not contribute to the output + +:attr:`bins` can be a sequence of N 1D tensors, a sequence of N ints, or a single int. + +If :attr:`bins` is a sequence of N 1D tensors, it explicitly specifies the N sequences +of bin edges. Each 1D tensor should contain a strictly increasing sequence with at +least one element. A sequence of K bin edges defines K-1 bins, explicitly specifying +the left and right edges of all bins. Every bin is inclusive of its left edge. Only +the rightmost bin is inclusive of its right edge. + +If :attr:`bins` is a sequence of N ints, it specifies the number of equal-width bins +in each dimension. By default, the leftmost and rightmost bin edges in each dimension +are determined by the minimum and maximum elements of the input tensor in the +corresponding dimension. The :attr:`range` argument can be provided to manually +specify the leftmost and rightmost bin edges in each dimension. + +If :attr:`bins` is an int, it specifies the number of equal-width bins for all dimensions. + +.. note:: + See also :func:`torch.histogram`, which specifically computes 1D histograms. + While :func:`torch.histogramdd` infers the dimensionality of its bins and + binned values from the shape of :attr:`input`, :func:`torch.histogram` + accepts and flattens :attr:`input` of any shape. + +Args: + {input} + bins: Tensor[], int[], or int. + If Tensor[], defines the sequences of bin edges. + If int[], defines the number of equal-width bins in each dimension. + If int, defines the number of equal-width bins for all dimensions. +Keyword args: + range (sequence of float): Defines the leftmost and rightmost bin edges + in each dimension. + weight (Tensor): By default, each value in the input has weight 1. If a weight + tensor is passed, each N-dimensional coordinate in input + contributes its associated weight towards its bin's result. + The weight tensor should have the same shape as the :attr:`input` + tensor excluding its innermost dimension N. + density (bool): If False (default), the result will contain the count (or total weight) + in each bin. If True, each count (weight) is divided by the total count + (total weight), then divided by the volume of its associated bin. +Returns: + hist (Tensor): N-dimensional Tensor containing the values of the histogram. + bin_edges(Tensor[]): sequence of N 1D Tensors containing the bin edges. + +Example:: + + >>> torch.histogramdd(torch.tensor([[0., 1.], [1., 0.], [2., 0.], [2., 2.]]), bins=[3, 3], + ... weight=torch.tensor([1., 2., 4., 8.])) + torch.return_types.histogramdd( + hist=tensor([[0., 1., 0.], + [2., 0., 0.], + [4., 0., 8.]]), + bin_edges=(tensor([0.0000, 0.6667, 1.3333, 2.0000]), + tensor([0.0000, 0.6667, 1.3333, 2.0000]))) + + >>> torch.histogramdd(torch.tensor([[0., 0.], [1., 1.], [2., 2.]]), bins=[2, 2], + ... range=[0., 1., 0., 1.], density=True) + torch.return_types.histogramdd( + hist=tensor([[2., 0.], + [0., 2.]]), + bin_edges=(tensor([0.0000, 0.5000, 1.0000]), + tensor([0.0000, 0.5000, 1.0000]))) + +""".format(**common_args), +) +# TODO: Fix via https://github.com/pytorch/pytorch/issues/75798 +torch.histogramdd.__module__ = "torch" + +add_docstr( + torch.hypot, + r""" +hypot(input, other, *, out=None) -> Tensor + +Given the legs of a right triangle, return its hypotenuse. + +.. math:: + \text{out}_{i} = \sqrt{\text{input}_{i}^{2} + \text{other}_{i}^{2}} + +The shapes of ``input`` and ``other`` must be +:ref:`broadcastable `. +""" + + r""" +Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.hypot(torch.tensor([4.0]), torch.tensor([3.0, 4.0, 5.0])) + tensor([5.0000, 5.6569, 6.4031]) + +""".format(**common_args), +) + +add_docstr( + torch.i0, + r""" +i0(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.i0`. +""", +) + +add_docstr( + torch.igamma, + r""" +igamma(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.special.gammainc`. +""", +) + +add_docstr( + torch.igammac, + r""" +igammac(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.special.gammaincc`. +""", +) + +add_docstr( + torch.index_select, + r""" +index_select(input, dim, index, *, out=None) -> Tensor + +Returns a new tensor which indexes the :attr:`input` tensor along dimension +:attr:`dim` using the entries in :attr:`index` which is a `LongTensor`. + +The returned tensor has the same number of dimensions as the original tensor +(:attr:`input`). The :attr:`dim`\ th dimension has the same size as the length +of :attr:`index`; other dimensions have the same size as in the original tensor. + +.. note:: The returned tensor does **not** use the same storage as the original + tensor. If :attr:`out` has a different shape than expected, we + silently change it to the correct shape, reallocating the underlying + storage if necessary. + +Args: + {input} + dim (int): the dimension in which we index + index (IntTensor or LongTensor): the 1-D tensor containing the indices to index + +Keyword args: + {out} + +Example:: + + >>> x = torch.randn(3, 4) + >>> x + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> indices = torch.tensor([0, 2]) + >>> torch.index_select(x, 0, indices) + tensor([[ 0.1427, 0.0231, -0.5414, -1.0009], + [-1.1734, -0.6571, 0.7230, -0.6004]]) + >>> torch.index_select(x, 1, indices) + tensor([[ 0.1427, -0.5414], + [-0.4664, -0.1228], + [-1.1734, 0.7230]]) +""".format(**common_args), +) + +add_docstr( + torch.inverse, + r""" +inverse(input, *, out=None) -> Tensor + +Alias for :func:`torch.linalg.inv` +""", +) + +add_docstr( + torch.isin, + r""" +isin(elements, test_elements, *, assume_unique=False, invert=False) -> Tensor + +Tests if each element of :attr:`elements` is in :attr:`test_elements`. Returns +a boolean tensor of the same shape as :attr:`elements` that is True for elements +in :attr:`test_elements` and False otherwise. + +.. note:: + One of :attr:`elements` or :attr:`test_elements` can be a scalar, but not both. + +Args: + elements (Tensor or Scalar): Input elements + test_elements (Tensor or Scalar): Values against which to test for each input element + assume_unique (bool, optional): If True, assumes both :attr:`elements` and + :attr:`test_elements` contain unique elements, which can speed up the + calculation. Default: False + invert (bool, optional): If True, inverts the boolean return tensor, resulting in True + values for elements *not* in :attr:`test_elements`. Default: False + +Returns: + A boolean tensor of the same shape as :attr:`elements` that is True for elements in + :attr:`test_elements` and False otherwise + +Example: + >>> torch.isin(torch.tensor([[1, 2], [3, 4]]), torch.tensor([2, 3])) + tensor([[False, True], + [ True, False]]) +""", +) + +add_docstr( + torch.isinf, + r""" +isinf(input) -> Tensor + +Tests if each element of :attr:`input` is infinite +(positive or negative infinity) or not. + +.. note:: + Complex values are infinite when their real or imaginary part is + infinite. + +Args: + {input} + +Returns: + A boolean tensor that is True where :attr:`input` is infinite and False elsewhere + +Example:: + + >>> torch.isinf(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')])) + tensor([False, True, False, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.isposinf, + r""" +isposinf(input, *, out=None) -> Tensor +Tests if each element of :attr:`input` is positive infinity or not. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([-float('inf'), float('inf'), 1.2]) + >>> torch.isposinf(a) + tensor([False, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.isneginf, + r""" +isneginf(input, *, out=None) -> Tensor +Tests if each element of :attr:`input` is negative infinity or not. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([-float('inf'), float('inf'), 1.2]) + >>> torch.isneginf(a) + tensor([ True, False, False]) +""".format(**common_args), +) + +add_docstr( + torch.isclose, + r""" +isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False) -> Tensor + +Returns a new tensor with boolean elements representing if each element of +:attr:`input` is "close" to the corresponding element of :attr:`other`. +Closeness is defined as: + +.. math:: + \lvert \text{input}_i - \text{other}_i \rvert \leq \texttt{rtol} \times \lvert \text{other}_i \rvert + \texttt{atol} +""" + + r""" + +where :attr:`input` and :attr:`other` are finite. Where :attr:`input` +and/or :attr:`other` are nonfinite they are close if and only if +they are equal, with NaNs being considered equal to each other when +:attr:`equal_nan` is True. + +Args: + input (Tensor): first tensor to compare + other (Tensor): second tensor to compare + rtol (float, optional): relative tolerance. Default: 1e-05 + atol (float, optional): absolute tolerance. Default: 1e-08 + equal_nan (bool, optional): if ``True``, then two ``NaN`` s will be considered equal. Default: ``False`` + +Examples:: + + >>> torch.isclose(torch.tensor((1., 2, 3)), torch.tensor((1 + 1e-10, 3, 4))) + tensor([ True, False, False]) + >>> torch.isclose(torch.tensor((float('inf'), 4)), torch.tensor((float('inf'), 6)), rtol=.5) + tensor([True, True]) +""", +) + +add_docstr( + torch.isfinite, + r""" +isfinite(input) -> Tensor + +Returns a new tensor with boolean elements representing if each element is `finite` or not. + +Real values are finite when they are not NaN, negative infinity, or infinity. +Complex values are finite when both their real and imaginary parts are finite. + +Args: + {input} + +Returns: + A boolean tensor that is True where :attr:`input` is finite and False elsewhere + +Example:: + + >>> torch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')])) + tensor([True, False, True, False, False]) +""".format(**common_args), +) + +add_docstr( + torch.isnan, + r""" +isnan(input) -> Tensor + +Returns a new tensor with boolean elements representing if each element of :attr:`input` +is NaN or not. Complex values are considered NaN when either their real +and/or imaginary part is NaN. + +Arguments: + {input} + +Returns: + A boolean tensor that is True where :attr:`input` is NaN and False elsewhere + +Example:: + + >>> torch.isnan(torch.tensor([1, float('nan'), 2])) + tensor([False, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.isreal, + r""" +isreal(input) -> Tensor + +Returns a new tensor with boolean elements representing if each element of :attr:`input` is real-valued or not. +All real-valued types are considered real. Complex values are considered real when their imaginary part is 0. + +Arguments: + {input} + +Returns: + A boolean tensor that is True where :attr:`input` is real and False elsewhere + +Example:: + + >>> torch.isreal(torch.tensor([1, 1+1j, 2+0j])) + tensor([True, False, True]) +""".format(**common_args), +) + +add_docstr( + torch.is_floating_point, + r""" +is_floating_point(input: Tensor) -> bool + +Returns True if the data type of :attr:`input` is a floating point data type i.e., +one of ``torch.float64``, ``torch.float32``, ``torch.float16``, and ``torch.bfloat16``. + +Args: + {input} + +Example:: + + >>> torch.is_floating_point(torch.tensor([1.0, 2.0, 3.0])) + True + >>> torch.is_floating_point(torch.tensor([1, 2, 3], dtype=torch.int32)) + False + >>> torch.is_floating_point(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)) + True + >>> torch.is_floating_point(torch.tensor([1, 2, 3], dtype=torch.complex64)) + False +""".format(**common_args), +) + +add_docstr( + torch.is_complex, + r""" +is_complex(input: Tensor) -> bool + +Returns True if the data type of :attr:`input` is a complex data type i.e., +one of ``torch.complex64``, and ``torch.complex128``. + +Args: + {input} + +Example:: + + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.complex64)) + True + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.complex128)) + True + >>> torch.is_complex(torch.tensor([1, 2, 3], dtype=torch.int32)) + False + >>> torch.is_complex(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)) + False +""".format(**common_args), +) + +add_docstr( + torch.is_grad_enabled, + r""" +is_grad_enabled() -> (bool) + +Returns True if grad mode is currently enabled. +""".format(**common_args), +) + +add_docstr( + torch.is_inference_mode_enabled, + r""" +is_inference_mode_enabled() -> (bool) + +Returns True if inference mode is currently enabled. +""".format(**common_args), +) + +add_docstr( + torch.is_inference, + r""" +is_inference(input) -> (bool) + +Returns True if :attr:`input` is an inference tensor. + +A non-view tensor is an inference tensor if and only if it was +allocated during inference mode. A view tensor is an inference +tensor if and only if the tensor it is a view of is an inference tensor. + +For details on inference mode please see +`Inference Mode `_. + +Args: + {input} +""".format(**common_args), +) + +add_docstr( + torch.is_conj, + r""" +is_conj(input) -> (bool) + +Returns True if the :attr:`input` is a conjugated tensor, i.e. its conjugate bit is set to `True`. + +Args: + {input} +""".format(**common_args), +) + +add_docstr( + torch.is_nonzero, + r""" +is_nonzero(input) -> (bool) + +Returns True if the :attr:`input` is a single element tensor which is not equal to zero +after type conversions. +i.e. not equal to ``torch.tensor([0.])`` or ``torch.tensor([0])`` or +``torch.tensor([False])``. +Throws a ``RuntimeError`` if ``torch.numel() != 1`` (even in case +of sparse tensors). + +Args: + {input} + +Examples:: + + >>> torch.is_nonzero(torch.tensor([0.])) + False + >>> torch.is_nonzero(torch.tensor([1.5])) + True + >>> torch.is_nonzero(torch.tensor([False])) + False + >>> torch.is_nonzero(torch.tensor([3])) + True + >>> torch.is_nonzero(torch.tensor([1, 3, 5])) + Traceback (most recent call last): + ... + RuntimeError: Boolean value of Tensor with more than one value is ambiguous + >>> torch.is_nonzero(torch.tensor([])) + Traceback (most recent call last): + ... + RuntimeError: Boolean value of Tensor with no values is ambiguous +""".format(**common_args), +) + +add_docstr( + torch.kron, + r""" +kron(input, other, *, out=None) -> Tensor + +Computes the Kronecker product, denoted by :math:`\otimes`, of :attr:`input` and :attr:`other`. + +If :attr:`input` is a :math:`(a_0 \times a_1 \times \dots \times a_n)` tensor and :attr:`other` is a +:math:`(b_0 \times b_1 \times \dots \times b_n)` tensor, the result will be a +:math:`(a_0*b_0 \times a_1*b_1 \times \dots \times a_n*b_n)` tensor with the following entries: + +.. math:: + (\text{input} \otimes \text{other})_{k_0, k_1, \dots, k_n} = + \text{input}_{i_0, i_1, \dots, i_n} * \text{other}_{j_0, j_1, \dots, j_n}, + +where :math:`k_t = i_t * b_t + j_t` for :math:`0 \leq t \leq n`. +If one tensor has fewer dimensions than the other it is unsqueezed until it has the same number of dimensions. + +Supports real-valued and complex-valued inputs. + +.. note:: + This function generalizes the typical definition of the Kronecker product for two matrices to two tensors, + as described above. When :attr:`input` is a :math:`(m \times n)` matrix and :attr:`other` is a + :math:`(p \times q)` matrix, the result will be a :math:`(p*m \times q*n)` block matrix: + + .. math:: + \mathbf{A} \otimes \mathbf{B}=\begin{bmatrix} + a_{11} \mathbf{B} & \cdots & a_{1 n} \mathbf{B} \\ + \vdots & \ddots & \vdots \\ + a_{m 1} \mathbf{B} & \cdots & a_{m n} \mathbf{B} \end{bmatrix} + + where :attr:`input` is :math:`\mathbf{A}` and :attr:`other` is :math:`\mathbf{B}`. + +Arguments: + input (Tensor) + other (Tensor) + +Keyword args: + out (Tensor, optional): The output tensor. Ignored if ``None``. Default: ``None`` + +Examples:: + + >>> mat1 = torch.eye(2) + >>> mat2 = torch.ones(2, 2) + >>> torch.kron(mat1, mat2) + tensor([[1., 1., 0., 0.], + [1., 1., 0., 0.], + [0., 0., 1., 1.], + [0., 0., 1., 1.]]) + + >>> mat1 = torch.eye(2) + >>> mat2 = torch.arange(1, 5).reshape(2, 2) + >>> torch.kron(mat1, mat2) + tensor([[1., 2., 0., 0.], + [3., 4., 0., 0.], + [0., 0., 1., 2.], + [0., 0., 3., 4.]]) +""", +) + +add_docstr( + torch.kthvalue, + r""" +kthvalue(input, k, dim=None, keepdim=False, *, out=None) -> (Tensor, LongTensor) + +Returns a namedtuple ``(values, indices)`` where ``values`` is the :attr:`k` th +smallest element of each row of the :attr:`input` tensor in the given dimension +:attr:`dim`. And ``indices`` is the index location of each element found. + +If :attr:`dim` is not given, the last dimension of the `input` is chosen. + +If :attr:`keepdim` is ``True``, both the :attr:`values` and :attr:`indices` tensors +are the same size as :attr:`input`, except in the dimension :attr:`dim` where +they are of size 1. Otherwise, :attr:`dim` is squeezed +(see :func:`torch.squeeze`), resulting in both the :attr:`values` and +:attr:`indices` tensors having 1 fewer dimension than the :attr:`input` tensor. + +.. note:: + When :attr:`input` is a CUDA tensor and there are multiple valid + :attr:`k` th values, this function may nondeterministically return + :attr:`indices` for any of them. + +Args: + {input} + k (int): k for the k-th smallest element + dim (int, optional): the dimension to find the kth value along + {opt_keepdim} + +Keyword args: + out (tuple, optional): the output tuple of (Tensor, LongTensor) + can be optionally given to be used as output buffers + +Example:: + + >>> x = torch.arange(1., 6.) + >>> x + tensor([ 1., 2., 3., 4., 5.]) + >>> torch.kthvalue(x, 4) + torch.return_types.kthvalue(values=tensor(4.), indices=tensor(3)) + + >>> x=torch.arange(1.,7.).resize_(2,3) + >>> x + tensor([[ 1., 2., 3.], + [ 4., 5., 6.]]) + >>> torch.kthvalue(x, 2, 0, True) + torch.return_types.kthvalue(values=tensor([[4., 5., 6.]]), indices=tensor([[1, 1, 1]])) +""".format(**single_dim_common), +) + +add_docstr( + torch.lcm, + r""" +lcm(input, other, *, out=None) -> Tensor + +Computes the element-wise least common multiple (LCM) of :attr:`input` and :attr:`other`. + +Both :attr:`input` and :attr:`other` must have integer types. + +.. note:: + This defines :math:`lcm(0, 0) = 0` and :math:`lcm(0, a) = 0`. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.tensor([5, 10, 15]) + >>> b = torch.tensor([3, 4, 5]) + >>> torch.lcm(a, b) + tensor([15, 20, 15]) + >>> c = torch.tensor([3]) + >>> torch.lcm(a, c) + tensor([15, 30, 15]) +""".format(**common_args), +) + +add_docstr( + torch.ldexp, + r""" +ldexp(input, other, *, out=None) -> Tensor + +Multiplies :attr:`input` by 2 ** :attr:`other`. + +.. math:: + \text{{out}}_i = \text{{input}}_i * 2^\text{{other}}_i +""" + + r""" + +Typically this function is used to construct floating point numbers by multiplying +mantissas in :attr:`input` with integral powers of two created from the exponents +in :attr:`other`. + +Args: + {input} + other (Tensor): a tensor of exponents, typically integers. + +Keyword args: + {out} + +Example:: + + >>> torch.ldexp(torch.tensor([1.]), torch.tensor([1])) + tensor([2.]) + >>> torch.ldexp(torch.tensor([1.0]), torch.tensor([1, 2, 3, 4])) + tensor([ 2., 4., 8., 16.]) + + +""".format(**common_args), +) + +add_docstr( + torch.le, + r""" +le(input, other, *, out=None) -> Tensor + +Computes :math:`\text{input} \leq \text{other}` element-wise. +""" + + r""" + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or Scalar): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is less than or equal to + :attr:`other` and False elsewhere + +Example:: + + >>> torch.le(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[True, False], [True, True]]) +""".format(**common_args), +) + +add_docstr( + torch.less_equal, + r""" +less_equal(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.le`. +""", +) + +add_docstr( + torch.lerp, + r""" +lerp(input, end, weight, *, out=None) + +Does a linear interpolation of two tensors :attr:`start` (given by :attr:`input`) and :attr:`end` based +on a scalar or tensor :attr:`weight` and returns the resulting :attr:`out` tensor. + +.. math:: + \text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i) +""" + + r""" +The shapes of :attr:`start` and :attr:`end` must be +:ref:`broadcastable `. If :attr:`weight` is a tensor, then +the shapes of :attr:`weight`, :attr:`start`, and :attr:`end` must be :ref:`broadcastable `. + +Args: + input (Tensor): the tensor with the starting points + end (Tensor): the tensor with the ending points + weight (float or tensor): the weight for the interpolation formula + +Keyword args: + {out} + +Example:: + + >>> start = torch.arange(1., 5.) + >>> end = torch.empty(4).fill_(10) + >>> start + tensor([ 1., 2., 3., 4.]) + >>> end + tensor([ 10., 10., 10., 10.]) + >>> torch.lerp(start, end, 0.5) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) + >>> torch.lerp(start, end, torch.full_like(start, 0.5)) + tensor([ 5.5000, 6.0000, 6.5000, 7.0000]) +""".format(**common_args), +) + +add_docstr( + torch.lgamma, + r""" +lgamma(input, *, out=None) -> Tensor + +Computes the natural logarithm of the absolute value of the gamma function on :attr:`input`. + +.. math:: + \text{out}_{i} = \ln |\Gamma(\text{input}_{i})| +""" + + """ +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.arange(0.5, 2, 0.5) + >>> torch.lgamma(a) + tensor([ 0.5724, 0.0000, -0.1208]) +""".format(**common_args), +) + +add_docstr( + torch.linspace, + r""" +linspace(start, end, steps, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly +spaced from :attr:`start` to :attr:`end`, inclusive. That is, the value are: + +.. math:: + (\text{start}, + \text{start} + \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \ldots, + \text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{\text{steps} - 1}, + \text{end}) +""" + + """ + +From PyTorch 1.11 linspace requires the steps argument. Use steps=100 to restore the previous behavior. + +Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + +Keyword arguments: + {out} + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + {layout} + {device} + {requires_grad} + + +Example:: + + >>> torch.linspace(3, 10, steps=5) + tensor([ 3.0000, 4.7500, 6.5000, 8.2500, 10.0000]) + >>> torch.linspace(-10, 10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=5) + tensor([-10., -5., 0., 5., 10.]) + >>> torch.linspace(start=-10, end=10, steps=1) + tensor([-10.]) +""".format(**factory_common_args), +) + +add_docstr( + torch.log, + r""" +log(input, *, out=None) -> Tensor + +Returns a new tensor with the natural logarithm of the elements +of :attr:`input`. + +.. math:: + y_{i} = \log_{e} (x_{i}) +""" + + r""" + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(5) * 5 + >>> a + tensor([4.7767, 4.3234, 1.2156, 0.2411, 4.5739]) + >>> torch.log(a) + tensor([ 1.5637, 1.4640, 0.1952, -1.4226, 1.5204]) +""".format(**common_args), +) + +add_docstr( + torch.log10, + r""" +log10(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the logarithm to the base 10 of the elements +of :attr:`input`. + +.. math:: + y_{i} = \log_{10} (x_{i}) +""" + + r""" + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(5) + >>> a + tensor([ 0.5224, 0.9354, 0.7257, 0.1301, 0.2251]) + + + >>> torch.log10(a) + tensor([-0.2820, -0.0290, -0.1392, -0.8857, -0.6476]) + +""".format(**common_args), +) + +add_docstr( + torch.log1p, + r""" +log1p(input, *, out=None) -> Tensor + +Returns a new tensor with the natural logarithm of (1 + :attr:`input`). + +.. math:: + y_i = \log_{e} (x_i + 1) +""" + + r""" +.. note:: This function is more accurate than :func:`torch.log` for small + values of :attr:`input` + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(5) + >>> a + tensor([-1.0090, -0.9923, 1.0249, -0.5372, 0.2492]) + >>> torch.log1p(a) + tensor([ nan, -4.8653, 0.7055, -0.7705, 0.2225]) +""".format(**common_args), +) + +add_docstr( + torch.log2, + r""" +log2(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the logarithm to the base 2 of the elements +of :attr:`input`. + +.. math:: + y_{i} = \log_{2} (x_{i}) +""" + + r""" + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(5) + >>> a + tensor([ 0.8419, 0.8003, 0.9971, 0.5287, 0.0490]) + + + >>> torch.log2(a) + tensor([-0.2483, -0.3213, -0.0042, -0.9196, -4.3504]) + +""".format(**common_args), +) + +add_docstr( + torch.logaddexp, + r""" +logaddexp(input, other, *, out=None) -> Tensor + +Logarithm of the sum of exponentiations of the inputs. + +Calculates pointwise :math:`\log\left(e^x + e^y\right)`. This function is useful +in statistics where the calculated probabilities of events may be so small as to +exceed the range of normal floating point numbers. In such cases the logarithm +of the calculated probability is stored. This function allows adding +probabilities stored in such a fashion. + +This op should be disambiguated with :func:`torch.logsumexp` which performs a +reduction on a single tensor. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword arguments: + {out} + +Example:: + + >>> torch.logaddexp(torch.tensor([-1.0]), torch.tensor([-1.0, -2, -3])) + tensor([-0.3069, -0.6867, -0.8731]) + >>> torch.logaddexp(torch.tensor([-100.0, -200, -300]), torch.tensor([-1.0, -2, -3])) + tensor([-1., -2., -3.]) + >>> torch.logaddexp(torch.tensor([1.0, 2000, 30000]), torch.tensor([-1.0, -2, -3])) + tensor([1.1269e+00, 2.0000e+03, 3.0000e+04]) +""".format(**common_args), +) + +add_docstr( + torch.logaddexp2, + r""" +logaddexp2(input, other, *, out=None) -> Tensor + +Logarithm of the sum of exponentiations of the inputs in base-2. + +Calculates pointwise :math:`\log_2\left(2^x + 2^y\right)`. See +:func:`torch.logaddexp` for more details. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword arguments: + {out} +""".format(**common_args), +) + +add_docstr( + torch.xlogy, + r""" +xlogy(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.special.xlogy`. +""", +) + +add_docstr( + torch.logical_and, + r""" +logical_and(input, other, *, out=None) -> Tensor + +Computes the element-wise logical AND of the given input tensors. Zeros are treated as ``False`` and nonzeros are +treated as ``True``. + +Args: + {input} + other (Tensor): the tensor to compute AND with + +Keyword args: + {out} + +Example:: + + >>> torch.logical_and(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([ True, False, False]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_and(a, b) + tensor([False, False, True, False]) + >>> torch.logical_and(a.double(), b.double()) + tensor([False, False, True, False]) + >>> torch.logical_and(a.double(), b) + tensor([False, False, True, False]) + >>> torch.logical_and(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([False, False, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.logical_not, + r""" +logical_not(input, *, out=None) -> Tensor + +Computes the element-wise logical NOT of the given input tensor. If not specified, the output tensor will have the bool +dtype. If the input tensor is not a bool tensor, zeros are treated as ``False`` and non-zeros are treated as ``True``. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.logical_not(torch.tensor([True, False])) + tensor([False, True]) + >>> torch.logical_not(torch.tensor([0, 1, -10], dtype=torch.int8)) + tensor([ True, False, False]) + >>> torch.logical_not(torch.tensor([0., 1.5, -10.], dtype=torch.double)) + tensor([ True, False, False]) + >>> torch.logical_not(torch.tensor([0., 1., -10.], dtype=torch.double), out=torch.empty(3, dtype=torch.int16)) + tensor([1, 0, 0], dtype=torch.int16) +""".format(**common_args), +) + +add_docstr( + torch.logical_or, + r""" +logical_or(input, other, *, out=None) -> Tensor + +Computes the element-wise logical OR of the given input tensors. Zeros are treated as ``False`` and nonzeros are +treated as ``True``. + +Args: + {input} + other (Tensor): the tensor to compute OR with + +Keyword args: + {out} + +Example:: + + >>> torch.logical_or(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([ True, False, True]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_or(a, b) + tensor([ True, True, True, False]) + >>> torch.logical_or(a.double(), b.double()) + tensor([ True, True, True, False]) + >>> torch.logical_or(a.double(), b) + tensor([ True, True, True, False]) + >>> torch.logical_or(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([ True, True, True, False]) +""".format(**common_args), +) + +add_docstr( + torch.logical_xor, + r""" +logical_xor(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor + +Computes the element-wise logical XOR of the given input tensors. Zeros are treated as ``False`` and nonzeros are +treated as ``True``. + +Args: + {input} + other (Tensor): the tensor to compute XOR with + +Keyword args: + {out} + +Example:: + + >>> torch.logical_xor(torch.tensor([True, False, True]), torch.tensor([True, False, False])) + tensor([False, False, True]) + >>> a = torch.tensor([0, 1, 10, 0], dtype=torch.int8) + >>> b = torch.tensor([4, 0, 1, 0], dtype=torch.int8) + >>> torch.logical_xor(a, b) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a.double(), b.double()) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a.double(), b) + tensor([ True, True, False, False]) + >>> torch.logical_xor(a, b, out=torch.empty(4, dtype=torch.bool)) + tensor([ True, True, False, False]) +""".format(**common_args), +) + +add_docstr( + torch.logspace, + """ +logspace(start, end, steps, base=10.0, *, \ + out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor +""" + + r""" + +Creates a one-dimensional tensor of size :attr:`steps` whose values are evenly +spaced from :math:`{{\text{{base}}}}^{{\text{{start}}}}` to +:math:`{{\text{{base}}}}^{{\text{{end}}}}`, inclusive, on a logarithmic scale +with base :attr:`base`. That is, the values are: + +.. math:: + (\text{base}^{\text{start}}, + \text{base}^{(\text{start} + \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \ldots, + \text{base}^{(\text{start} + (\text{steps} - 2) * \frac{\text{end} - \text{start}}{ \text{steps} - 1})}, + \text{base}^{\text{end}}) +""" + + """ + + +From PyTorch 1.11 logspace requires the steps argument. Use steps=100 to restore the previous behavior. + +Args: + start (float or Tensor): the starting value for the set of points. If `Tensor`, it must be 0-dimensional + end (float or Tensor): the ending value for the set of points. If `Tensor`, it must be 0-dimensional + steps (int): size of the constructed tensor + base (float, optional): base of the logarithm function. Default: ``10.0``. + +Keyword arguments: + {out} + dtype (torch.dtype, optional): the data type to perform the computation in. + Default: if None, uses the global default dtype (see torch.get_default_dtype()) + when both :attr:`start` and :attr:`end` are real, + and corresponding complex dtype when either is complex. + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.logspace(start=-10, end=10, steps=5) + tensor([ 1.0000e-10, 1.0000e-05, 1.0000e+00, 1.0000e+05, 1.0000e+10]) + >>> torch.logspace(start=0.1, end=1.0, steps=5) + tensor([ 1.2589, 2.1135, 3.5481, 5.9566, 10.0000]) + >>> torch.logspace(start=0.1, end=1.0, steps=1) + tensor([1.2589]) + >>> torch.logspace(start=2, end=2, steps=1, base=2) + tensor([4.0]) +""".format(**factory_common_args), +) + +add_docstr( + torch.logsumexp, + r""" +logsumexp(input, dim, keepdim=False, *, out=None) + +Returns the log of summed exponentials of each row of the :attr:`input` +tensor in the given dimension :attr:`dim`. The computation is numerically +stabilized. + +For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is + + .. math:: + \text{{logsumexp}}(x)_{{i}} = \log \sum_j \exp(x_{{ij}}) + +{keepdim_details} + +Args: + {input} + {dim} + {opt_keepdim} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(3, 3) + >>> torch.logsumexp(a, 1) + tensor([1.4907, 1.0593, 1.5696]) + >>> torch.dist(torch.logsumexp(a, 1), torch.log(torch.sum(torch.exp(a), 1))) + tensor(1.6859e-07) +""".format(**multi_dim_common), +) + +add_docstr( + torch.lt, + r""" +lt(input, other, *, out=None) -> Tensor + +Computes :math:`\text{input} < \text{other}` element-wise. +""" + + r""" + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is less than :attr:`other` and False elsewhere + +Example:: + + >>> torch.lt(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, False], [True, False]]) +""".format(**common_args), +) + +add_docstr( + torch.lu_unpack, + r""" +lu_unpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True, *, out=None) -> (Tensor, Tensor, Tensor) + +Unpacks the LU decomposition returned by :func:`~linalg.lu_factor` into the `P, L, U` matrices. + +.. seealso:: + + :func:`~linalg.lu` returns the matrices from the LU decomposition. Its gradient formula is more efficient + than that of doing :func:`~linalg.lu_factor` followed by :func:`~linalg.lu_unpack`. + +Args: + LU_data (Tensor): the packed LU factorization data + LU_pivots (Tensor): the packed LU factorization pivots + unpack_data (bool): flag indicating if the data should be unpacked. + If ``False``, then the returned ``L`` and ``U`` are empty tensors. + Default: ``True`` + unpack_pivots (bool): flag indicating if the pivots should be unpacked into a permutation matrix ``P``. + If ``False``, then the returned ``P`` is an empty tensor. + Default: ``True`` + +Keyword args: + out (tuple, optional): output tuple of three tensors. Ignored if `None`. + +Returns: + A namedtuple ``(P, L, U)`` + +Examples:: + + >>> A = torch.randn(2, 3, 3) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> P, L, U = torch.lu_unpack(LU, pivots) + >>> # We can recover A from the factorization + >>> A_ = P @ L @ U + >>> torch.allclose(A, A_) + True + + >>> # LU factorization of a rectangular matrix: + >>> A = torch.randn(2, 3, 2) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> P, L, U = torch.lu_unpack(LU, pivots) + >>> # P, L, U are the same as returned by linalg.lu + >>> P_, L_, U_ = torch.linalg.lu(A) + >>> torch.allclose(P, P_) and torch.allclose(L, L_) and torch.allclose(U, U_) + True + +""".format(**common_args), +) + +add_docstr( + torch.less, + r""" +less(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.lt`. +""", +) + +add_docstr( + torch.lu_solve, + r""" +lu_solve(b, LU_data, LU_pivots, *, out=None) -> Tensor + +Returns the LU solve of the linear system :math:`Ax = b` using the partially pivoted +LU factorization of A from :func:`~linalg.lu_factor`. + +This function supports ``float``, ``double``, ``cfloat`` and ``cdouble`` dtypes for :attr:`input`. + +.. warning:: + + :func:`torch.lu_solve` is deprecated in favor of :func:`torch.linalg.lu_solve`. + :func:`torch.lu_solve` will be removed in a future PyTorch release. + ``X = torch.lu_solve(B, LU, pivots)`` should be replaced with + + .. code:: python + + X = linalg.lu_solve(LU, pivots, B) + +Arguments: + b (Tensor): the RHS tensor of size :math:`(*, m, k)`, where :math:`*` + is zero or more batch dimensions. + LU_data (Tensor): the pivoted LU factorization of A from :meth:`~linalg.lu_factor` of size :math:`(*, m, m)`, + where :math:`*` is zero or more batch dimensions. + LU_pivots (IntTensor): the pivots of the LU factorization from :meth:`~linalg.lu_factor` of size :math:`(*, m)`, + where :math:`*` is zero or more batch dimensions. + The batch dimensions of :attr:`LU_pivots` must be equal to the batch dimensions of + :attr:`LU_data`. + +Keyword args: + {out} + +Example:: + + >>> A = torch.randn(2, 3, 3) + >>> b = torch.randn(2, 3, 1) + >>> LU, pivots = torch.linalg.lu_factor(A) + >>> x = torch.lu_solve(b, LU, pivots) + >>> torch.dist(A @ x, b) + tensor(1.00000e-07 * + 2.8312) +""".format(**common_args), +) + +add_docstr( + torch.masked_select, + r""" +masked_select(input, mask, *, out=None) -> Tensor + +Returns a new 1-D tensor which indexes the :attr:`input` tensor according to +the boolean mask :attr:`mask` which is a `BoolTensor`. + +The shapes of the :attr:`mask` tensor and the :attr:`input` tensor don't need +to match, but they must be :ref:`broadcastable `. + +.. note:: The returned tensor does **not** use the same storage + as the original tensor + +Args: + {input} + mask (BoolTensor): the tensor containing the binary mask to index with + +Keyword args: + {out} + +Example:: + + >>> x = torch.randn(3, 4) + >>> x + tensor([[ 0.3552, -2.3825, -0.8297, 0.3477], + [-1.2035, 1.2252, 0.5002, 0.6248], + [ 0.1307, -2.0608, 0.1244, 2.0139]]) + >>> mask = x.ge(0.5) + >>> mask + tensor([[False, False, False, False], + [False, True, True, True], + [False, False, False, True]]) + >>> torch.masked_select(x, mask) + tensor([ 1.2252, 0.5002, 0.6248, 2.0139]) +""".format(**common_args), +) + +add_docstr( + torch.matrix_power, + r""" +matrix_power(input, n, *, out=None) -> Tensor + +Alias for :func:`torch.linalg.matrix_power` +""", +) + +add_docstr( + torch.matrix_exp, + r""" +matrix_exp(A) -> Tensor + +Alias for :func:`torch.linalg.matrix_exp`. +""", +) + +add_docstr( + torch.max, + r""" +max(input, *, out=None) -> Tensor + +Returns the maximum value of all elements in the ``input`` tensor. + +.. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6763, 0.7445, -2.2369]]) + >>> torch.max(a) + tensor(0.7445) + +.. function:: max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + +Returns a namedtuple ``(values, indices)`` where ``values`` is the maximum +value of each row of the :attr:`input` tensor in the given dimension +:attr:`dim`. And ``indices`` is the index location of each maximum value found +(argmax). + +If ``keepdim`` is ``True``, the output tensors are of the same size +as ``input`` except in the dimension ``dim`` where they are of size 1. +Otherwise, ``dim`` is squeezed (see :func:`torch.squeeze`), resulting +in the output tensors having 1 fewer dimension than ``input``. + +.. note:: If there are multiple maximal values in a reduced row then + the indices of the first maximal value are returned. + +Args: + {input} + {opt_dim_without_none} + {opt_keepdim} + +Keyword args: + out (tuple, optional): the result tuple of two output tensors (max, max_indices) + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-1.2360, -0.2942, -0.1222, 0.8475], + [ 1.1949, -1.1127, -2.2379, -0.6702], + [ 1.5717, -0.9207, 0.1297, -1.8768], + [-0.6172, 1.0036, -0.6060, -0.2432]]) + >>> torch.max(a, 1) + torch.return_types.max(values=tensor([0.8475, 1.1949, 1.5717, 1.0036]), indices=tensor([3, 0, 0, 1])) + >>> a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + >>> a.max(dim=1, keepdim=True) + torch.return_types.max( + values=tensor([[2.], [4.]]), + indices=tensor([[1], [1]])) + >>> a.max(dim=1, keepdim=False) + torch.return_types.max( + values=tensor([2., 4.]), + indices=tensor([1, 1])) + +.. function:: max(input, other, *, out=None) -> Tensor + :noindex: + +See :func:`torch.maximum`. + +""".format(**single_dim_common), +) + +add_docstr( + torch.maximum, + r""" +maximum(input, other, *, out=None) -> Tensor + +Computes the element-wise maximum of :attr:`input` and :attr:`other`. + +.. note:: + If one of the elements being compared is a NaN, then that element is returned. + :func:`maximum` is not supported for tensors with complex dtypes. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor((1, 2, -1)) + >>> b = torch.tensor((3, 0, 4)) + >>> torch.maximum(a, b) + tensor([3, 2, 4]) +""".format(**common_args), +) + +add_docstr( + torch.fmax, + r""" +fmax(input, other, *, out=None) -> Tensor + +Computes the element-wise maximum of :attr:`input` and :attr:`other`. + +This is like :func:`torch.maximum` except it handles NaNs differently: +if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the maximum. +Only if both elements are NaN is NaN propagated. + +This function is a wrapper around C++'s ``std::fmax`` and is similar to NumPy's ``fmax`` function. + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer and floating-point inputs. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([9.7, float('nan'), 3.1, float('nan')]) + >>> b = torch.tensor([-2.2, 0.5, float('nan'), float('nan')]) + >>> torch.fmax(a, b) + tensor([9.7000, 0.5000, 3.1000, nan]) +""".format(**common_args), +) + +add_docstr( + torch.amax, + r""" +amax(input, dim, keepdim=False, *, out=None) -> Tensor + +Returns the maximum value of each slice of the :attr:`input` tensor in the given +dimension(s) :attr:`dim`. + +.. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.8177, 1.4878, -0.2491, 0.9130], + [-0.7158, 1.1775, 2.0992, 0.4817], + [-0.0053, 0.0164, -1.3738, -0.0507], + [ 1.9700, 1.1106, -1.0318, -1.0816]]) + >>> torch.amax(a, 1) + tensor([1.4878, 2.0992, 0.0164, 1.9700]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.argmax, + r""" +argmax(input) -> LongTensor + +Returns the indices of the maximum value of all elements in the :attr:`input` tensor. + +This is the second value returned by :meth:`torch.max`. See its +documentation for the exact semantics of this method. + +.. note:: If there are multiple maximal values then the indices of the first maximal value are returned. + +Args: + {input} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], + [-0.7401, -0.8805, -0.3402, -1.1936], + [ 0.4907, -1.3948, -1.0691, -0.3132], + [-1.6092, 0.5419, -0.2993, 0.3195]]) + >>> torch.argmax(a) + tensor(0) + +.. function:: argmax(input, dim, keepdim=False) -> LongTensor + :noindex: + +Returns the indices of the maximum values of a tensor across a dimension. + +This is the second value returned by :meth:`torch.max`. See its +documentation for the exact semantics of this method. + +Args: + {input} + {opt_dim} If ``None``, the argmax of the flattened input is returned. + {opt_keepdim} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 1.3398, 0.2663, -0.2686, 0.2450], + [-0.7401, -0.8805, -0.3402, -1.1936], + [ 0.4907, -1.3948, -1.0691, -0.3132], + [-1.6092, 0.5419, -0.2993, 0.3195]]) + >>> torch.argmax(a, dim=1) + tensor([ 0, 2, 0, 1]) +""".format(**single_dim_common), +) + +add_docstr( + torch.argwhere, + r""" +argwhere(input) -> Tensor + +Returns a tensor containing the indices of all non-zero elements of +:attr:`input`. Each row in the result contains the indices of a non-zero +element in :attr:`input`. The result is sorted lexicographically, with +the last index changing the fastest (C-style). + +If :attr:`input` has :math:`n` dimensions, then the resulting indices tensor +:attr:`out` is of size :math:`(z \times n)`, where :math:`z` is the total number of +non-zero elements in the :attr:`input` tensor. + +.. note:: + This function is similar to NumPy's `argwhere`. + + When :attr:`input` is on CUDA, this function causes host-device synchronization. + +Args: + {input} + +Example:: + + >>> t = torch.tensor([1, 0, 1]) + >>> torch.argwhere(t) + tensor([[0], + [2]]) + >>> t = torch.tensor([[1, 0, 1], [0, 1, 1]]) + >>> torch.argwhere(t) + tensor([[0, 0], + [0, 2], + [1, 1], + [1, 2]]) +""", +) + +add_docstr( + torch.mean, + r""" +mean(input, *, dtype=None) -> Tensor + +.. note:: + If the `input` tensor is empty, ``torch.mean()`` returns ``nan``. + This behavior is consistent with NumPy and follows the definition + that the mean over an empty set is undefined. + + +Returns the mean value of all elements in the :attr:`input` tensor. Input must be floating point or complex. + +Args: + input (Tensor): + the input tensor, either of floating point or complex dtype + +Keyword args: + {dtype} + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.2294, -0.5481, 1.3288]]) + >>> torch.mean(a) + tensor(0.3367) + +.. function:: mean(input, dim, keepdim=False, *, dtype=None, out=None) -> Tensor + :noindex: + +Returns the mean value of each row of the :attr:`input` tensor in the given +dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, +reduce over all of them. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {dtype} + {out} + +.. seealso:: + + :func:`torch.nanmean` computes the mean value of `non-NaN` elements. + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.3841, 0.6320, 0.4254, -0.7384], + [-0.9644, 1.0131, -0.6549, -1.4279], + [-0.2951, -1.3350, -0.7694, 0.5600], + [ 1.0842, -0.9580, 0.3623, 0.2343]]) + >>> torch.mean(a, 1) + tensor([-0.0163, -0.5085, -0.4599, 0.1807]) + >>> torch.mean(a, 1, True) + tensor([[-0.0163], + [-0.5085], + [-0.4599], + [ 0.1807]]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.nanmean, + r""" +nanmean(input, dim=None, keepdim=False, *, dtype=None, out=None) -> Tensor + +Computes the mean of all `non-NaN` elements along the specified dimensions. +Input must be floating point or complex. + +This function is identical to :func:`torch.mean` when there are no `NaN` values +in the :attr:`input` tensor. In the presence of `NaN`, :func:`torch.mean` will +propagate the `NaN` to the output whereas :func:`torch.nanmean` will ignore the +`NaN` values (`torch.nanmean(a)` is equivalent to `torch.mean(a[~a.isnan()])`). + +{keepdim_details} + +Args: + input (Tensor): the input tensor, either of floating point or complex dtype + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {dtype} + {out} + +.. seealso:: + + :func:`torch.mean` computes the mean value, propagating `NaN`. + +Example:: + + >>> x = torch.tensor([[torch.nan, 1, 2], [1, 2, 3]]) + >>> x.mean() + tensor(nan) + >>> x.nanmean() + tensor(1.8000) + >>> x.mean(dim=0) + tensor([ nan, 1.5000, 2.5000]) + >>> x.nanmean(dim=0) + tensor([1.0000, 1.5000, 2.5000]) + + # If all elements in the reduced dimensions are NaN then the result is NaN + >>> torch.tensor([torch.nan]).nanmean() + tensor(nan) +""".format(**multi_dim_common), +) + +add_docstr( + torch.median, + r""" +median(input) -> Tensor + +Returns the median of the values in :attr:`input`. + +.. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements. In this case the lower of the two medians is returned. To + compute the mean of both medians, use :func:`torch.quantile` with ``q=0.5`` instead. + +.. warning:: + This function produces deterministic (sub)gradients unlike ``median(dim=0)`` + +Args: + {input} + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 1.5219, -1.5212, 0.2202]]) + >>> torch.median(a) + tensor(0.2202) + +.. function:: median(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + +Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` +in the dimension :attr:`dim`, and ``indices`` contains the index of the median values found in the dimension :attr:`dim`. + +By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + +If :attr:`keepdim` is ``True``, the output tensors are of the same size +as :attr:`input` except in the dimension :attr:`dim` where they are of size 1. +Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in +the outputs tensor having 1 fewer dimension than :attr:`input`. + +.. note:: + The median is not unique for :attr:`input` tensors with an even number + of elements in the dimension :attr:`dim`. In this case the lower of the + two medians is returned. To compute the mean of both medians in + :attr:`input`, use :func:`torch.quantile` with ``q=0.5`` instead. + +.. warning:: + ``indices`` does not necessarily contain the first occurrence of each + median value found, unless it is unique. + The exact implementation details are device-specific. + Do not expect the same result when run on CPU and GPU in general. + For the same reason do not expect the gradients to be deterministic. + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + +Example:: + + >>> a = torch.randn(4, 5) + >>> a + tensor([[ 0.2505, -0.3982, -0.9948, 0.3518, -1.3131], + [ 0.3180, -0.6993, 1.0436, 0.0438, 0.2270], + [-0.2751, 0.7303, 0.2192, 0.3321, 0.2488], + [ 1.0778, -1.9510, 0.7048, 0.4742, -0.7125]]) + >>> torch.median(a, 1) + torch.return_types.median(values=tensor([-0.3982, 0.2270, 0.2488, 0.4742]), indices=tensor([1, 4, 4, 3])) +""".format(**single_dim_common), +) + +add_docstr( + torch.nanmedian, + r""" +nanmedian(input) -> Tensor + +Returns the median of the values in :attr:`input`, ignoring ``NaN`` values. + +This function is identical to :func:`torch.median` when there are no ``NaN`` values in :attr:`input`. +When :attr:`input` has one or more ``NaN`` values, :func:`torch.median` will always return ``NaN``, +while this function will return the median of the non-``NaN`` elements in :attr:`input`. +If all the elements in :attr:`input` are ``NaN`` it will also return ``NaN``. + +Args: + {input} + +Example:: + + >>> a = torch.tensor([1, float('nan'), 3, 2]) + >>> a.median() + tensor(nan) + >>> a.nanmedian() + tensor(2.) + +.. function:: nanmedian(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + +Returns a namedtuple ``(values, indices)`` where ``values`` contains the median of each row of :attr:`input` +in the dimension :attr:`dim`, ignoring ``NaN`` values, and ``indices`` contains the index of the median values +found in the dimension :attr:`dim`. + +This function is identical to :func:`torch.median` when there are no ``NaN`` values in a reduced row. When a reduced row has +one or more ``NaN`` values, :func:`torch.median` will always reduce it to ``NaN``, while this function will reduce it to the +median of the non-``NaN`` elements. If all the elements in a reduced row are ``NaN`` then it will be reduced to ``NaN``, too. + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + out ((Tensor, Tensor), optional): The first tensor will be populated with the median values and the second + tensor, which must have dtype long, with their indices in the dimension + :attr:`dim` of :attr:`input`. + +Example:: + + >>> a = torch.tensor([[2, 3, 1], [float('nan'), 1, float('nan')]]) + >>> a + tensor([[2., 3., 1.], + [nan, 1., nan]]) + >>> a.median(0) + torch.return_types.median(values=tensor([nan, 1., nan]), indices=tensor([1, 1, 1])) + >>> a.nanmedian(0) + torch.return_types.nanmedian(values=tensor([2., 1., 1.]), indices=tensor([0, 1, 0])) +""".format(**single_dim_common), +) + +add_docstr( + torch.quantile, + r""" +quantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + +Computes the q-th quantiles of each row of the :attr:`input` tensor along the dimension :attr:`dim`. + +To compute the quantile, we map q in [0, 1] to the range of indices [0, n] to find the location +of the quantile in the sorted input. If the quantile lies between two data points ``a < b`` with +indices ``i`` and ``j`` in the sorted order, result is computed according to the given +:attr:`interpolation` method as follows: + +- ``linear``: ``a + (b - a) * fraction``, where ``fraction`` is the fractional part of the computed quantile index. +- ``lower``: ``a``. +- ``higher``: ``b``. +- ``nearest``: ``a`` or ``b``, whichever's index is closer to the computed quantile index (rounding down for .5 fractions). +- ``midpoint``: ``(a + b) / 2``. + +If :attr:`q` is a 1D tensor, the first dimension of the output represents the quantiles and has size +equal to the size of :attr:`q`, the remaining dimensions are what remains from the reduction. + +.. note:: + By default :attr:`dim` is ``None`` resulting in the :attr:`input` tensor being flattened before computation. + +Args: + {input} + q (float or Tensor): a scalar or 1D tensor of values in the range [0, 1]. + {opt_dim} + {opt_keepdim} + +Keyword arguments: + interpolation (str, optional): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + {out} + +Example:: + + >>> a = torch.randn(2, 3) + >>> a + tensor([[ 0.0795, -1.2117, 0.9765], + [ 1.1707, 0.6706, 0.4884]]) + >>> q = torch.tensor([0.25, 0.5, 0.75]) + >>> torch.quantile(a, q, dim=1, keepdim=True) + tensor([[[-0.5661], + [ 0.5795]], + + [[ 0.0795], + [ 0.6706]], + + [[ 0.5280], + [ 0.9206]]]) + >>> torch.quantile(a, q, dim=1, keepdim=True).shape + torch.Size([3, 2, 1]) + >>> a = torch.arange(4.) + >>> a + tensor([0., 1., 2., 3.]) + >>> torch.quantile(a, 0.6, interpolation='linear') + tensor(1.8000) + >>> torch.quantile(a, 0.6, interpolation='lower') + tensor(1.) + >>> torch.quantile(a, 0.6, interpolation='higher') + tensor(2.) + >>> torch.quantile(a, 0.6, interpolation='midpoint') + tensor(1.5000) + >>> torch.quantile(a, 0.6, interpolation='nearest') + tensor(2.) + >>> torch.quantile(a, 0.4, interpolation='nearest') + tensor(1.) +""".format(**single_dim_common), +) + +add_docstr( + torch.nanquantile, + r""" +nanquantile(input, q, dim=None, keepdim=False, *, interpolation='linear', out=None) -> Tensor + +This is a variant of :func:`torch.quantile` that "ignores" ``NaN`` values, +computing the quantiles :attr:`q` as if ``NaN`` values in :attr:`input` did +not exist. If all values in a reduced row are ``NaN`` then the quantiles for +that reduction will be ``NaN``. See the documentation for :func:`torch.quantile`. + +Args: + {input} + q (float or Tensor): a scalar or 1D tensor of quantile values in the range [0, 1] + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword arguments: + interpolation (str): interpolation method to use when the desired quantile lies between two data points. + Can be ``linear``, ``lower``, ``higher``, ``midpoint`` and ``nearest``. + Default is ``linear``. + {out} + +Example:: + + >>> t = torch.tensor([float('nan'), 1, 2]) + >>> t.quantile(0.5) + tensor(nan) + >>> t.nanquantile(0.5) + tensor(1.5000) + >>> t = torch.tensor([[float('nan'), float('nan')], [1, 2]]) + >>> t + tensor([[nan, nan], + [1., 2.]]) + >>> t.nanquantile(0.5, dim=0) + tensor([1., 2.]) + >>> t.nanquantile(0.5, dim=1) + tensor([ nan, 1.5000]) +""".format(**single_dim_common), +) + +add_docstr( + torch.min, + r""" +min(input, *, out=None) -> Tensor + +Returns the minimum value of all elements in the :attr:`input` tensor. + +.. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.6750, 1.0857, 1.7197]]) + >>> torch.min(a) + tensor(0.6750) + +.. function:: min(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor) + :noindex: + +Returns a namedtuple ``(values, indices)`` where ``values`` is the minimum +value of each row of the :attr:`input` tensor in the given dimension +:attr:`dim`. And ``indices`` is the index location of each minimum value found +(argmin). + +If :attr:`keepdim` is ``True``, the output tensors are of the same size as +:attr:`input` except in the dimension :attr:`dim` where they are of size 1. +Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in +the output tensors having 1 fewer dimension than :attr:`input`. + +.. note:: If there are multiple minimal values in a reduced row then + the indices of the first minimal value are returned. + +Args: + {input} + {opt_dim_without_none} + {opt_keepdim} + +Keyword args: + out (tuple, optional): the tuple of two output tensors (min, min_indices) + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[-0.6248, 1.1334, -1.1899, -0.2803], + [-1.4644, -0.2635, -0.3651, 0.6134], + [ 0.2457, 0.0384, 1.0128, 0.7015], + [-0.1153, 2.9849, 2.1458, 0.5788]]) + >>> torch.min(a, 1) + torch.return_types.min(values=tensor([-1.1899, -1.4644, 0.0384, -0.1153]), indices=tensor([2, 0, 1, 0])) + +.. function:: min(input, other, *, out=None) -> Tensor + :noindex: + +See :func:`torch.minimum`. +""".format(**single_dim_common), +) + +add_docstr( + torch.minimum, + r""" +minimum(input, other, *, out=None) -> Tensor + +Computes the element-wise minimum of :attr:`input` and :attr:`other`. + +.. note:: + If one of the elements being compared is a NaN, then that element is returned. + :func:`minimum` is not supported for tensors with complex dtypes. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor((1, 2, -1)) + >>> b = torch.tensor((3, 0, 4)) + >>> torch.minimum(a, b) + tensor([1, 0, -1]) +""".format(**common_args), +) + +add_docstr( + torch.fmin, + r""" +fmin(input, other, *, out=None) -> Tensor + +Computes the element-wise minimum of :attr:`input` and :attr:`other`. + +This is like :func:`torch.minimum` except it handles NaNs differently: +if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the minimum. +Only if both elements are NaN is NaN propagated. + +This function is a wrapper around C++'s ``std::fmin`` and is similar to NumPy's ``fmin`` function. + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer and floating-point inputs. + +Args: + {input} + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([2.2, float('nan'), 2.1, float('nan')]) + >>> b = torch.tensor([-9.3, 0.1, float('nan'), float('nan')]) + >>> torch.fmin(a, b) + tensor([-9.3000, 0.1000, 2.1000, nan]) +""".format(**common_args), +) + +add_docstr( + torch.amin, + r""" +amin(input, dim, keepdim=False, *, out=None) -> Tensor + +Returns the minimum value of each slice of the :attr:`input` tensor in the given +dimension(s) :attr:`dim`. + +.. note:: + The difference between ``max``/``min`` and ``amax``/``amin`` is: + - ``amax``/``amin`` supports reducing on multiple dimensions, + - ``amax``/``amin`` does not return indices. + + Both ``amax``/``amin`` evenly distribute gradients between equal values + when there are multiple input elements with the same minimum or maximum value. + + For ``max``/``min``: + - If reduce over all dimensions(no dim specified), gradients evenly distribute between equally ``max``/``min`` values. + - If reduce over one specified axis, only propagate to the indexed element. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.6451, -0.4866, 0.2987, -1.3312], + [-0.5744, 1.2980, 1.8397, -0.2713], + [ 0.9128, 0.9214, -1.7268, -0.2995], + [ 0.9023, 0.4853, 0.9075, -1.6165]]) + >>> torch.amin(a, 1) + tensor([-1.3312, -0.5744, -1.7268, -1.6165]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.aminmax, + r""" +aminmax(input, *, dim=None, keepdim=False, out=None) -> (Tensor min, Tensor max) + +Computes the minimum and maximum values of the :attr:`input` tensor. + +Args: + input (Tensor): + The input tensor + +Keyword Args: + dim (Optional[int]): + The dimension along which to compute the values. If `None`, + computes the values over the entire :attr:`input` tensor. + Default is `None`. + keepdim (bool): + If `True`, the reduced dimensions will be kept in the output + tensor as dimensions with size 1 for broadcasting, otherwise + they will be removed, as if calling (:func:`torch.squeeze`). + Default is `False`. + out (Optional[Tuple[Tensor, Tensor]]): + Optional tensors on which to write the result. Must have the same + shape and dtype as the expected output. + Default is `None`. + +Returns: + A named tuple `(min, max)` containing the minimum and maximum values. + +Raises: + RuntimeError + If any of the dimensions to compute the values over has size 0. + +.. note:: + NaN values are propagated to the output if at least one value is NaN. + +.. seealso:: + :func:`torch.amin` computes just the minimum value + :func:`torch.amax` computes just the maximum value + +Example:: + + >>> torch.aminmax(torch.tensor([1, -3, 5])) + torch.return_types.aminmax( + min=tensor(-3), + max=tensor(5)) + + >>> # aminmax propagates NaNs + >>> torch.aminmax(torch.tensor([1, -3, 5, torch.nan])) + torch.return_types.aminmax( + min=tensor(nan), + max=tensor(nan)) + + >>> t = torch.arange(10).view(2, 5) + >>> t + tensor([[0, 1, 2, 3, 4], + [5, 6, 7, 8, 9]]) + >>> t.aminmax(dim=0, keepdim=True) + torch.return_types.aminmax( + min=tensor([[0, 1, 2, 3, 4]]), + max=tensor([[5, 6, 7, 8, 9]])) +""", +) + +add_docstr( + torch.argmin, + r""" +argmin(input, dim=None, keepdim=False) -> LongTensor + +Returns the indices of the minimum value(s) of the flattened tensor or along a dimension + +This is the second value returned by :meth:`torch.min`. See its +documentation for the exact semantics of this method. + +.. note:: If there are multiple minimal values then the indices of the first minimal value are returned. + +Args: + {input} + {opt_dim} If ``None``, the argmin of the flattened input is returned. + {opt_keepdim} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.1139, 0.2254, -0.1381, 0.3687], + [ 1.0100, -1.1975, -0.0102, -0.4732], + [-0.9240, 0.1207, -0.7506, -1.0213], + [ 1.7809, -1.2960, 0.9384, 0.1438]]) + >>> torch.argmin(a) + tensor(13) + >>> torch.argmin(a, dim=1) + tensor([ 2, 1, 3, 1]) + >>> torch.argmin(a, dim=1, keepdim=True) + tensor([[2], + [1], + [3], + [1]]) +""".format(**single_dim_common), +) + +add_docstr( + torch.mm, + r""" +mm(input, mat2, out_dtype=None, *, out=None) -> Tensor + +Performs a matrix multiplication of the matrices :attr:`input` and :attr:`mat2`. + +If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`mat2` is a +:math:`(m \times p)` tensor, :attr:`out` will be a :math:`(n \times p)` tensor. + +.. note:: This function does not :ref:`broadcast `. + For broadcasting matrix products, see :func:`torch.matmul`. + +Supports strided and sparse 2-D tensors as inputs, autograd with +respect to strided inputs. + +This operation has support for arguments with :ref:`sparse layouts`. +If :attr:`out` is provided its layout will be used. Otherwise, the result +layout will be deduced from that of :attr:`input`. + +{sparse_beta_warning} + +{tf32_note} + +{rocm_fp16_note} + +Args: + input (Tensor): the first matrix to be matrix multiplied + mat2 (Tensor): the second matrix to be matrix multiplied + out_dtype (dtype, optional): the dtype of the output tensor, + Supported only on CUDA and for torch.float32 given + torch.float16/torch.bfloat16 input dtypes + +Keyword args: + {out} + +Example:: + + >>> mat1 = torch.randn(2, 3) + >>> mat2 = torch.randn(3, 3) + >>> torch.mm(mat1, mat2) + tensor([[ 0.4851, 0.5037, -0.3633], + [-0.0760, -3.6705, 2.4784]]) +""".format(**common_args, **tf32_notes, **rocm_fp16_notes, **sparse_support_notes), +) + +add_docstr( + torch.hspmm, + r""" +hspmm(mat1, mat2, *, out=None) -> Tensor + +Performs a matrix multiplication of a :ref:`sparse COO matrix +` :attr:`mat1` and a strided matrix :attr:`mat2`. The +result is a (1 + 1)-dimensional :ref:`hybrid COO matrix +`. + +Args: + mat1 (Tensor): the first sparse matrix to be matrix multiplied + mat2 (Tensor): the second strided matrix to be matrix multiplied + +Keyword args: + {out} +""".format(**common_args), +) + +add_docstr( + torch.matmul, + r""" +matmul(input, other, *, out=None) -> Tensor + +Matrix product of two tensors. + +The behavior depends on the dimensionality of the tensors as follows: + +- If both tensors are 1-dimensional, the dot product (scalar) is returned. +- If both arguments are 2-dimensional, the matrix-matrix product is returned. +- If the first argument is 1-dimensional and the second argument is 2-dimensional, + a 1 is prepended to its dimension for the purpose of the matrix multiply. + After the matrix multiply, the prepended dimension is removed. +- If the first argument is 2-dimensional and the second argument is 1-dimensional, + the matrix-vector product is returned. +- If both arguments are at least 1-dimensional and at least one argument is + N-dimensional (where N > 2), then a batched matrix multiply is returned. If the first + argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the + batched matrix multiply and removed after. If the second argument is 1-dimensional, a + 1 is appended to its dimension for the purpose of the batched matrix multiply and removed after. + + The first N-2 dimensions of each argument, the batch dimensions, are + :ref:`broadcast ` (and thus must be broadcastable). + The last 2, the matrix dimensions, are handled as in the matrix-matrix product. + + For example, if :attr:`input` is a + :math:`(j \times 1 \times n \times m)` tensor and :attr:`other` is a :math:`(k \times m \times p)` + tensor, the batch dimensions are :math:`(j \times 1)` and :math:`(k)`, + and the matrix dimensions are :math:`(n \times m)` and :math:`(m \times p)`. + :attr:`out` will be a :math:`(j \times k \times n \times p)` tensor. + +This operation has support for arguments with :ref:`sparse layouts`. In particular the +matrix-matrix (both arguments 2-dimensional) supports sparse arguments with the same restrictions +as :func:`torch.mm` + +{sparse_beta_warning} + +{tf32_note} + +{rocm_fp16_note} + +.. note:: + + The 1-dimensional dot product version of this function does not support an :attr:`out` parameter. + +Arguments: + input (Tensor): the first tensor to be multiplied + other (Tensor): the second tensor to be multiplied + +Keyword args: + {out} + +Example:: + + >>> # vector x vector + >>> tensor1 = torch.randn(3) + >>> tensor2 = torch.randn(3) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([]) + >>> # matrix x vector + >>> tensor1 = torch.randn(3, 4) + >>> tensor2 = torch.randn(4) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([3]) + >>> # batched matrix x broadcasted vector + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(4) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3]) + >>> # batched matrix x batched matrix + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(10, 4, 5) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3, 5]) + >>> # batched matrix x broadcasted matrix + >>> tensor1 = torch.randn(10, 3, 4) + >>> tensor2 = torch.randn(4, 5) + >>> torch.matmul(tensor1, tensor2).size() + torch.Size([10, 3, 5]) + +""".format(**common_args, **tf32_notes, **rocm_fp16_notes, **sparse_support_notes), +) + +add_docstr( + torch.mode, + r""" +mode(input, dim=-1, keepdim=False, *, out=None) -> (Tensor, LongTensor) + +Returns a namedtuple ``(values, indices)`` where ``values`` is the mode +value of each row of the :attr:`input` tensor in the given dimension +:attr:`dim`, i.e. a value which appears most often +in that row, and ``indices`` is the index location of each mode value found. + +By default, :attr:`dim` is the last dimension of the :attr:`input` tensor. + +If :attr:`keepdim` is ``True``, the output tensors are of the same size as +:attr:`input` except in the dimension :attr:`dim` where they are of size 1. +Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting +in the output tensors having 1 fewer dimension than :attr:`input`. + +.. note:: This function is not defined for ``torch.cuda.Tensor`` yet. + +Args: + {input} + {opt_dim} + {opt_keepdim} + +Keyword args: + out (tuple, optional): the result tuple of two output tensors (values, indices) + +Example:: + + >>> b = torch.tensor([[0, 0, 0, 2, 0, 0, 2], + ... [0, 3, 0, 0, 2, 0, 1], + ... [2, 2, 2, 0, 0, 0, 3], + ... [2, 2, 3, 0, 1, 1, 0], + ... [1, 1, 0, 0, 2, 0, 2]]) + >>> torch.mode(b, 0) + torch.return_types.mode( + values=tensor([0, 2, 0, 0, 0, 0, 2]), + indices=tensor([1, 3, 4, 4, 2, 4, 4])) +""".format(**single_dim_common), +) + +add_docstr( + torch.mul, + r""" +mul(input, other, *, out=None) -> Tensor + +Multiplies :attr:`input` by :attr:`other`. + + +.. math:: + \text{out}_i = \text{input}_i \times \text{other}_i +""" + + r""" + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer, float, and complex inputs. + +Args: + {input} + other (Tensor or Number): the tensor or number to multiply input by. + +Keyword args: + {out} + +Examples:: + + >>> a = torch.randn(3) + >>> a + tensor([ 0.2015, -0.4255, 2.6087]) + >>> torch.mul(a, 100) + tensor([ 20.1494, -42.5491, 260.8663]) + + >>> b = torch.randn(4, 1) + >>> b + tensor([[ 1.1207], + [-0.3137], + [ 0.0700], + [ 0.8378]]) + >>> c = torch.randn(1, 4) + >>> c + tensor([[ 0.5146, 0.1216, -0.5244, 2.2382]]) + >>> torch.mul(b, c) + tensor([[ 0.5767, 0.1363, -0.5877, 2.5083], + [-0.1614, -0.0382, 0.1645, -0.7021], + [ 0.0360, 0.0085, -0.0367, 0.1567], + [ 0.4312, 0.1019, -0.4394, 1.8753]]) +""".format(**common_args), +) + +add_docstr( + torch.multiply, + r""" +multiply(input, other, *, out=None) + +Alias for :func:`torch.mul`. +""", +) + +add_docstr( + torch.multinomial, + r""" +multinomial(input, num_samples, replacement=False, *, generator=None, out=None) -> LongTensor + +Returns a tensor where each row contains :attr:`num_samples` indices sampled +from the multinomial (a stricter definition would be multivariate, +refer to :class:`torch.distributions.multinomial.Multinomial` for more details) +probability distribution located in the corresponding row +of tensor :attr:`input`. + +.. note:: + The rows of :attr:`input` do not need to sum to one (in which case we use + the values as weights), but must be non-negative, finite and have + a non-zero sum. + +Indices are ordered from left to right according to when each was sampled +(first samples are placed in first column). + +If :attr:`input` is a vector, :attr:`out` is a vector of size :attr:`num_samples`. + +If :attr:`input` is a matrix with `m` rows, :attr:`out` is an matrix of shape +:math:`(m \times \text{{num\_samples}})`. + +If replacement is ``True``, samples are drawn with replacement. + +If not, they are drawn without replacement, which means that when a +sample index is drawn for a row, it cannot be drawn again for that row. + +.. note:: + When drawn without replacement, :attr:`num_samples` must be lower than + number of non-zero elements in :attr:`input` (or the min number of non-zero + elements in each row of :attr:`input` if it is a matrix). + +Args: + input (Tensor): the input tensor containing probabilities + num_samples (int): number of samples to draw + replacement (bool, optional): whether to draw with replacement or not + +Keyword args: + {generator} + {out} + +Example:: + + >>> weights = torch.tensor([0, 10, 3, 0], dtype=torch.float) # create a tensor of weights + >>> torch.multinomial(weights, 2) + tensor([1, 2]) + >>> torch.multinomial(weights, 5) # ERROR! + RuntimeError: cannot sample n_sample > prob_dist.size(-1) samples without replacement + >>> torch.multinomial(weights, 4, replacement=True) + tensor([ 2, 1, 1, 1]) +""".format(**common_args), +) + +add_docstr( + torch.mv, + r""" +mv(input, vec, *, out=None) -> Tensor + +Performs a matrix-vector product of the matrix :attr:`input` and the vector +:attr:`vec`. + +If :attr:`input` is a :math:`(n \times m)` tensor, :attr:`vec` is a 1-D tensor of +size :math:`m`, :attr:`out` will be 1-D of size :math:`n`. + +.. note:: This function does not :ref:`broadcast `. + +Args: + input (Tensor): matrix to be multiplied + vec (Tensor): vector to be multiplied + +Keyword args: + {out} + +Example:: + + >>> mat = torch.randn(2, 3) + >>> vec = torch.randn(3) + >>> torch.mv(mat, vec) + tensor([ 1.0404, -0.6361]) +""".format(**common_args), +) + +add_docstr( + torch.mvlgamma, + r""" +mvlgamma(input, p, *, out=None) -> Tensor + +Alias for :func:`torch.special.multigammaln`. +""", +) + +add_docstr( + torch.movedim, + r""" +movedim(input, source, destination) -> Tensor + +Moves the dimension(s) of :attr:`input` at the position(s) in :attr:`source` +to the position(s) in :attr:`destination`. + +Other dimensions of :attr:`input` that are not explicitly moved remain in +their original order and appear at the positions not specified in :attr:`destination`. + +Args: + {input} + source (int or tuple of ints): Original positions of the dims to move. These must be unique. + destination (int or tuple of ints): Destination positions for each of the original dims. These must also be unique. + +Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.movedim(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.movedim(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.movedim(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.movedim(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) +""".format(**common_args), +) + +add_docstr( + torch.moveaxis, + r""" +moveaxis(input, source, destination) -> Tensor + +Alias for :func:`torch.movedim`. + +This function is equivalent to NumPy's moveaxis function. + +Examples:: + + >>> t = torch.randn(3,2,1) + >>> t + tensor([[[-0.3362], + [-0.8437]], + + [[-0.9627], + [ 0.1727]], + + [[ 0.5173], + [-0.1398]]]) + >>> torch.moveaxis(t, 1, 0).shape + torch.Size([2, 3, 1]) + >>> torch.moveaxis(t, 1, 0) + tensor([[[-0.3362], + [-0.9627], + [ 0.5173]], + + [[-0.8437], + [ 0.1727], + [-0.1398]]]) + >>> torch.moveaxis(t, (1, 2), (0, 1)).shape + torch.Size([2, 1, 3]) + >>> torch.moveaxis(t, (1, 2), (0, 1)) + tensor([[[-0.3362, -0.9627, 0.5173]], + + [[-0.8437, 0.1727, -0.1398]]]) +""".format(**common_args), +) + +add_docstr( + torch.swapdims, + r""" +swapdims(input, dim0, dim1) -> Tensor + +Alias for :func:`torch.transpose`. + +This function is equivalent to NumPy's swapaxes function. + +Examples:: + + >>> x = torch.tensor([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.swapdims(x, 0, 1) + tensor([[[0, 1], + [4, 5]], + + [[2, 3], + [6, 7]]]) + >>> torch.swapdims(x, 0, 2) + tensor([[[0, 4], + [2, 6]], + + [[1, 5], + [3, 7]]]) +""".format(**common_args), +) + +add_docstr( + torch.swapaxes, + r""" +swapaxes(input, axis0, axis1) -> Tensor + +Alias for :func:`torch.transpose`. + +This function is equivalent to NumPy's swapaxes function. + +Examples:: + + >>> x = torch.tensor([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.swapaxes(x, 0, 1) + tensor([[[0, 1], + [4, 5]], + + [[2, 3], + [6, 7]]]) + >>> torch.swapaxes(x, 0, 2) + tensor([[[0, 4], + [2, 6]], + + [[1, 5], + [3, 7]]]) +""".format(**common_args), +) + +add_docstr( + torch.narrow, + r""" +narrow(input, dim, start, length) -> Tensor + +Returns a new tensor that is a narrowed version of :attr:`input` tensor. The +dimension :attr:`dim` is input from :attr:`start` to ``start + length``. The +returned tensor and :attr:`input` tensor share the same underlying storage. + +Args: + input (Tensor): the tensor to narrow + dim (int): the dimension along which to narrow + start (int or Tensor): index of the element to start the narrowed dimension + from. Can be negative, which means indexing from the end of `dim`. If + `Tensor`, it must be an 0-dim integral `Tensor` (bools not allowed) + length (int): length of the narrowed dimension, must be weakly positive + +Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> torch.narrow(x, 0, 0, 2) + tensor([[ 1, 2, 3], + [ 4, 5, 6]]) + >>> torch.narrow(x, 1, 1, 2) + tensor([[ 2, 3], + [ 5, 6], + [ 8, 9]]) + >>> torch.narrow(x, -1, torch.tensor(-1), 1) + tensor([[3], + [6], + [9]]) +""", +) + +add_docstr( + torch.narrow_copy, + r""" +narrow_copy(input, dim, start, length, *, out=None) -> Tensor + +Same as :meth:`Tensor.narrow` except this returns a copy rather +than shared storage. This is primarily for sparse tensors, which +do not have a shared-storage narrow method. + +Args: + input (Tensor): the tensor to narrow + dim (int): the dimension along which to narrow + start (int): index of the element to start the narrowed dimension from. Can + be negative, which means indexing from the end of `dim` + length (int): length of the narrowed dimension, must be weakly positive + +Keyword args: + {out} + +Example:: + + >>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + >>> torch.narrow_copy(x, 0, 0, 2) + tensor([[ 1, 2, 3], + [ 4, 5, 6]]) + >>> torch.narrow_copy(x, 1, 1, 2) + tensor([[ 2, 3], + [ 5, 6], + [ 8, 9]]) + >>> s = torch.arange(16).reshape(2, 2, 2, 2).to_sparse(2) + >>> torch.narrow_copy(s, 0, 0, 1) + tensor(indices=tensor([[0, 0], + [0, 1]]), + values=tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]), + size=(1, 2, 2, 2), nnz=2, layout=torch.sparse_coo) + +.. seealso:: + + :func:`torch.narrow` for a non copy variant + +""".format(**common_args), +) + +add_docstr( + torch.nan_to_num, + r""" +nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None) -> Tensor + +Replaces :literal:`NaN`, positive infinity, and negative infinity values in :attr:`input` +with the values specified by :attr:`nan`, :attr:`posinf`, and :attr:`neginf`, respectively. +By default, :literal:`NaN`\ s are replaced with zero, positive infinity is replaced with the +greatest finite value representable by :attr:`input`'s dtype, and negative infinity +is replaced with the least finite value representable by :attr:`input`'s dtype. + +Args: + {input} + nan (Number, optional): the value to replace :literal:`NaN`\s with. Default is zero. + posinf (Number, optional): if a Number, the value to replace positive infinity values with. + If None, positive infinity values are replaced with the greatest finite value representable by :attr:`input`'s dtype. + Default is None. + neginf (Number, optional): if a Number, the value to replace negative infinity values with. + If None, negative infinity values are replaced with the lowest finite value representable by :attr:`input`'s dtype. + Default is None. + +Keyword args: + {out} + +Example:: + + >>> x = torch.tensor([float('nan'), float('inf'), -float('inf'), 3.14]) + >>> torch.nan_to_num(x) + tensor([ 0.0000e+00, 3.4028e+38, -3.4028e+38, 3.1400e+00]) + >>> torch.nan_to_num(x, nan=2.0) + tensor([ 2.0000e+00, 3.4028e+38, -3.4028e+38, 3.1400e+00]) + >>> torch.nan_to_num(x, nan=2.0, posinf=1.0) + tensor([ 2.0000e+00, 1.0000e+00, -3.4028e+38, 3.1400e+00]) + +""".format(**common_args), +) + +add_docstr( + torch.ne, + r""" +ne(input, other, *, out=None) -> Tensor + +Computes :math:`\text{input} \neq \text{other}` element-wise. +""" + + r""" + +The second argument can be a number or a tensor whose shape is +:ref:`broadcastable ` with the first argument. + +Args: + input (Tensor): the tensor to compare + other (Tensor or float): the tensor or value to compare + +Keyword args: + {out} + +Returns: + A boolean tensor that is True where :attr:`input` is not equal to :attr:`other` and False elsewhere + +Example:: + + >>> torch.ne(torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 1], [4, 4]])) + tensor([[False, True], [True, False]]) +""".format(**common_args), +) + +add_docstr( + torch.not_equal, + r""" +not_equal(input, other, *, out=None) -> Tensor + +Alias for :func:`torch.ne`. +""", +) + +add_docstr( + torch.neg, + r""" +neg(input, *, out=None) -> Tensor + +Returns a new tensor with the negative of the elements of :attr:`input`. + +.. math:: + \text{out} = -1 \times \text{input} +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(5) + >>> a + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) + >>> torch.neg(a) + tensor([-0.0090, 0.2262, 0.0682, 0.2866, -0.3940]) +""".format(**common_args), +) + +add_docstr( + torch.negative, + r""" +negative(input, *, out=None) -> Tensor + +Alias for :func:`torch.neg` +""", +) + +add_docstr( + torch.nextafter, + r""" +nextafter(input, other, *, out=None) -> Tensor + +Return the next floating-point value after :attr:`input` towards :attr:`other`, elementwise. + +The shapes of ``input`` and ``other`` must be +:ref:`broadcastable `. + +Args: + input (Tensor): the first input tensor + other (Tensor): the second input tensor + +Keyword args: + {out} + +Example:: + + >>> eps = torch.finfo(torch.float32).eps + >>> torch.nextafter(torch.tensor([1.0, 2.0]), torch.tensor([2.0, 1.0])) == torch.tensor([eps + 1, 2 - eps]) + tensor([True, True]) + +""".format(**common_args), +) + +add_docstr( + torch.nonzero, + r""" +nonzero(input, *, out=None, as_tuple=False) -> LongTensor or tuple of LongTensors + +.. note:: + :func:`torch.nonzero(..., as_tuple=False) ` (default) returns a + 2-D tensor where each row is the index for a nonzero value. + + :func:`torch.nonzero(..., as_tuple=True) ` returns a tuple of 1-D + index tensors, allowing for advanced indexing, so ``x[x.nonzero(as_tuple=True)]`` + gives all nonzero values of tensor ``x``. Of the returned tuple, each index tensor + contains nonzero indices for a certain dimension. + + See below for more details on the two behaviors. + + When :attr:`input` is on CUDA, :func:`torch.nonzero() ` causes + host-device synchronization. + +**When** :attr:`as_tuple` **is** ``False`` **(default)**: + +Returns a tensor containing the indices of all non-zero elements of +:attr:`input`. Each row in the result contains the indices of a non-zero +element in :attr:`input`. The result is sorted lexicographically, with +the last index changing the fastest (C-style). + +If :attr:`input` has :math:`n` dimensions, then the resulting indices tensor +:attr:`out` is of size :math:`(z \times n)`, where :math:`z` is the total number of +non-zero elements in the :attr:`input` tensor. + +**When** :attr:`as_tuple` **is** ``True``: + +Returns a tuple of 1-D tensors, one for each dimension in :attr:`input`, +each containing the indices (in that dimension) of all non-zero elements of +:attr:`input` . + +If :attr:`input` has :math:`n` dimensions, then the resulting tuple contains :math:`n` +tensors of size :math:`z`, where :math:`z` is the total number of +non-zero elements in the :attr:`input` tensor. + +As a special case, when :attr:`input` has zero dimensions and a nonzero scalar +value, it is treated as a one-dimensional tensor with one element. + +Args: + {input} + +Keyword args: + out (LongTensor, optional): the output tensor containing indices + +Returns: + LongTensor or tuple of LongTensor: If :attr:`as_tuple` is ``False``, the output + tensor containing indices. If :attr:`as_tuple` is ``True``, one 1-D tensor for + each dimension, containing the indices of each nonzero element along that + dimension. + +Example:: + + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1])) + tensor([[ 0], + [ 1], + [ 2], + [ 4]]) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]])) + tensor([[ 0, 0], + [ 1, 1], + [ 2, 2], + [ 3, 3]]) + >>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]), as_tuple=True) + (tensor([0, 1, 2, 4]),) + >>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0], + ... [0.0, 0.4, 0.0, 0.0], + ... [0.0, 0.0, 1.2, 0.0], + ... [0.0, 0.0, 0.0,-0.4]]), as_tuple=True) + (tensor([0, 1, 2, 3]), tensor([0, 1, 2, 3])) + >>> torch.nonzero(torch.tensor(5), as_tuple=True) + (tensor([0]),) +""".format(**common_args), +) + +add_docstr( + torch.normal, + r""" +normal(mean, std, *, generator=None, out=None) -> Tensor + +Returns a tensor of random numbers drawn from separate normal distributions +whose mean and standard deviation are given. + +The :attr:`mean` is a tensor with the mean of +each output element's normal distribution + +The :attr:`std` is a tensor with the standard deviation of +each output element's normal distribution + +The shapes of :attr:`mean` and :attr:`std` don't need to match, but the +total number of elements in each tensor need to be the same. + +.. note:: When the shapes do not match, the shape of :attr:`mean` + is used as the shape for the returned output tensor + +.. note:: When :attr:`std` is a CUDA tensor, this function synchronizes + its device with the CPU. + +Args: + mean (Tensor): the tensor of per-element means + std (Tensor): the tensor of per-element standard deviations + +Keyword args: + {generator} + {out} + +Example:: + + >>> torch.normal(mean=torch.arange(1., 11.), std=torch.arange(1, 0, -0.1)) + tensor([ 1.0425, 3.5672, 2.7969, 4.2925, 4.7229, 6.2134, + 8.0505, 8.1408, 9.0563, 10.0566]) + +.. function:: normal(mean=0.0, std, *, out=None) -> Tensor + :noindex: + +Similar to the function above, but the means are shared among all drawn +elements. + +Args: + mean (float, optional): the mean for all distributions + std (Tensor): the tensor of per-element standard deviations + +Keyword args: + {out} + +Example:: + + >>> torch.normal(mean=0.5, std=torch.arange(1., 6.)) + tensor([-1.2793, -1.0732, -2.0687, 5.1177, -1.2303]) + +.. function:: normal(mean, std=1.0, *, out=None) -> Tensor + :noindex: + +Similar to the function above, but the standard deviations are shared among +all drawn elements. + +Args: + mean (Tensor): the tensor of per-element means + std (float, optional): the standard deviation for all distributions + +Keyword args: + out (Tensor, optional): the output tensor + +Example:: + + >>> torch.normal(mean=torch.arange(1., 6.)) + tensor([ 1.1552, 2.6148, 2.6535, 5.8318, 4.2361]) + +.. function:: normal(mean, std, size, *, out=None) -> Tensor + :noindex: + +Similar to the function above, but the means and standard deviations are shared +among all drawn elements. The resulting tensor has size given by :attr:`size`. + +Args: + mean (float): the mean for all distributions + std (float): the standard deviation for all distributions + size (int...): a sequence of integers defining the shape of the output tensor. + +Keyword args: + {out} + +Example:: + + >>> torch.normal(2, 3, size=(1, 4)) + tensor([[-1.3987, -1.9544, 3.6048, 0.7909]]) +""".format(**common_args), +) + +add_docstr( + torch.numel, + r""" +numel(input: Tensor) -> int + +Returns the total number of elements in the :attr:`input` tensor. + +Args: + {input} + +Example:: + + >>> a = torch.randn(1, 2, 3, 4, 5) + >>> torch.numel(a) + 120 + >>> a = torch.zeros(4,4) + >>> torch.numel(a) + 16 + +""".format(**common_args), +) + +add_docstr( + torch.ones, + r""" +ones(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a tensor filled with the scalar value `1`, with the shape defined +by the variable argument :attr:`size`. + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + +Keyword arguments: + {out} + {dtype} + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.ones(2, 3) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) + + >>> torch.ones(5) + tensor([ 1., 1., 1., 1., 1.]) + +""".format(**factory_common_args), +) + +add_docstr( + torch.ones_like, + r""" +ones_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + +Returns a tensor filled with the scalar value `1`, with the same size as +:attr:`input`. ``torch.ones_like(input)`` is equivalent to +``torch.ones(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + +.. warning:: + As of 0.4, this function does not support an :attr:`out` keyword. As an alternative, + the old ``torch.ones_like(input, out=output)`` is equivalent to + ``torch.ones(input.size(), out=output)``. + +Args: + {input} + +Keyword arguments: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +Example:: + + >>> input = torch.empty(2, 3) + >>> torch.ones_like(input) + tensor([[ 1., 1., 1.], + [ 1., 1., 1.]]) +""".format(**factory_like_common_args), +) + +add_docstr( + torch.orgqr, + r""" +orgqr(input, tau) -> Tensor + +Alias for :func:`torch.linalg.householder_product`. +""", +) + +add_docstr( + torch.ormqr, + r""" +ormqr(input, tau, other, left=True, transpose=False, *, out=None) -> Tensor + +Computes the matrix-matrix multiplication of a product of Householder matrices with a general matrix. + +Multiplies a :math:`m \times n` matrix `C` (given by :attr:`other`) with a matrix `Q`, +where `Q` is represented using Householder reflectors `(input, tau)`. +See `Representation of Orthogonal or Unitary Matrices`_ for further details. + +If :attr:`left` is `True` then `op(Q)` times `C` is computed, otherwise the result is `C` times `op(Q)`. +When :attr:`left` is `True`, the implicit matrix `Q` has size :math:`m \times m`. +It has size :math:`n \times n` otherwise. +If :attr:`transpose` is `True` then `op` is the conjugate transpose operation, otherwise it's a no-op. + +Supports inputs of float, double, cfloat and cdouble dtypes. +Also supports batched inputs, and, if the input is batched, the output is batched with the same dimensions. + +.. seealso:: + :func:`torch.geqrf` can be used to form the Householder representation `(input, tau)` of matrix `Q` + from the QR decomposition. + +.. note:: + This function supports backward but it is only fast when ``(input, tau)`` do not require gradients + and/or ``tau.size(-1)`` is very small. + `` + +Args: + input (Tensor): tensor of shape `(*, mn, k)` where `*` is zero or more batch dimensions + and `mn` equals to `m` or `n` depending on the :attr:`left`. + tau (Tensor): tensor of shape `(*, min(mn, k))` where `*` is zero or more batch dimensions. + other (Tensor): tensor of shape `(*, m, n)` where `*` is zero or more batch dimensions. + left (bool): controls the order of multiplication. + transpose (bool): controls whether the matrix `Q` is conjugate transposed or not. + +Keyword args: + out (Tensor, optional): the output Tensor. Ignored if `None`. Default: `None`. + +.. _Representation of Orthogonal or Unitary Matrices: + https://www.netlib.org/lapack/lug/node128.html +""", +) + +add_docstr( + torch.permute, + r""" +permute(input, dims) -> Tensor + +Returns a view of the original tensor :attr:`input` with its dimensions permuted. + +Args: + {input} + dims (tuple of int): The desired ordering of dimensions + +Example: + >>> x = torch.randn(2, 3, 5) + >>> x.size() + torch.Size([2, 3, 5]) + >>> torch.permute(x, (2, 0, 1)).size() + torch.Size([5, 2, 3]) +""".format(**common_args), +) + +add_docstr( + torch.poisson, + r""" +poisson(input, generator=None) -> Tensor + +Returns a tensor of the same size as :attr:`input` with each element +sampled from a Poisson distribution with rate parameter given by the corresponding +element in :attr:`input` i.e., + +.. math:: + \text{{out}}_i \sim \text{{Poisson}}(\text{{input}}_i) + +:attr:`input` must be non-negative. + +Args: + input (Tensor): the input tensor containing the rates of the Poisson distribution + +Keyword args: + {generator} + +Example:: + + >>> rates = torch.rand(4, 4) * 5 # rate parameter between 0 and 5 + >>> torch.poisson(rates) + tensor([[9., 1., 3., 5.], + [8., 6., 6., 0.], + [0., 4., 5., 3.], + [2., 1., 4., 2.]]) +""".format(**common_args), +) + +add_docstr( + torch.polygamma, + r""" +polygamma(n, input, *, out=None) -> Tensor + +Alias for :func:`torch.special.polygamma`. +""", +) + +add_docstr( + torch.positive, + r""" +positive(input) -> Tensor + +Returns :attr:`input`. +Throws a runtime error if :attr:`input` is a bool tensor. +""" + + r""" +Args: + {input} + +Example:: + + >>> t = torch.randn(5) + >>> t + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) + >>> torch.positive(t) + tensor([ 0.0090, -0.2262, -0.0682, -0.2866, 0.3940]) +""".format(**common_args), +) + +add_docstr( + torch.pow, + r""" +pow(input, exponent, *, out=None) -> Tensor + +Takes the power of each element in :attr:`input` with :attr:`exponent` and +returns a tensor with the result. + +:attr:`exponent` can be either a single ``float`` number or a `Tensor` +with the same number of elements as :attr:`input`. + +When :attr:`exponent` is a scalar value, the operation applied is: + +.. math:: + \text{out}_i = x_i ^ \text{exponent} + +When :attr:`exponent` is a tensor, the operation applied is: + +.. math:: + \text{out}_i = x_i ^ {\text{exponent}_i} +""" + + r""" +When :attr:`exponent` is a tensor, the shapes of :attr:`input` +and :attr:`exponent` must be :ref:`broadcastable `. + +Args: + {input} + exponent (float or tensor): the exponent value + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.4331, 1.2475, 0.6834, -0.2791]) + >>> torch.pow(a, 2) + tensor([ 0.1875, 1.5561, 0.4670, 0.0779]) + >>> exp = torch.arange(1., 5.) + + >>> a = torch.arange(1., 5.) + >>> a + tensor([ 1., 2., 3., 4.]) + >>> exp + tensor([ 1., 2., 3., 4.]) + >>> torch.pow(a, exp) + tensor([ 1., 4., 27., 256.]) + +.. function:: pow(self, exponent, *, out=None) -> Tensor + :noindex: + +:attr:`self` is a scalar ``float`` value, and :attr:`exponent` is a tensor. +The returned tensor :attr:`out` is of the same shape as :attr:`exponent` + +The operation applied is: + +.. math:: + \text{{out}}_i = \text{{self}} ^ {{\text{{exponent}}_i}} + +Args: + self (float): the scalar base value for the power operation + exponent (Tensor): the exponent tensor + +Keyword args: + {out} + +Example:: + + >>> exp = torch.arange(1., 5.) + >>> base = 2 + >>> torch.pow(base, exp) + tensor([ 2., 4., 8., 16.]) +""".format(**common_args), +) + +add_docstr( + torch.float_power, + r""" +float_power(input, exponent, *, out=None) -> Tensor + +Raises :attr:`input` to the power of :attr:`exponent`, elementwise, in double precision. +If neither input is complex returns a ``torch.float64`` tensor, +and if one or more inputs is complex returns a ``torch.complex128`` tensor. + +.. note:: + This function always computes in double precision, unlike :func:`torch.pow`, + which implements more typical :ref:`type promotion `. + This is useful when the computation needs to be performed in a wider or more precise dtype, + or the results of the computation may contain fractional values not representable in the input dtypes, + like when an integer base is raised to a negative integer exponent. + +Args: + input (Tensor or Number): the base value(s) + exponent (Tensor or Number): the exponent value(s) + +Keyword args: + {out} + +Example:: + + >>> a = torch.randint(10, (4,)) + >>> a + tensor([6, 4, 7, 1]) + >>> torch.float_power(a, 2) + tensor([36., 16., 49., 1.], dtype=torch.float64) + + >>> a = torch.arange(1, 5) + >>> a + tensor([ 1, 2, 3, 4]) + >>> exp = torch.tensor([2, -3, 4, -5]) + >>> exp + tensor([ 2, -3, 4, -5]) + >>> torch.float_power(a, exp) + tensor([1.0000e+00, 1.2500e-01, 8.1000e+01, 9.7656e-04], dtype=torch.float64) +""".format(**common_args), +) + +add_docstr( + torch.prod, + r""" +prod(input: Tensor, *, dtype: Optional[_dtype]) -> Tensor + +Returns the product of all elements in the :attr:`input` tensor. + +Args: + {input} + +Keyword args: + {dtype} + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[-0.8020, 0.5428, -1.5854]]) + >>> torch.prod(a) + tensor(0.6902) + +.. function:: prod(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + +Returns the product of each row of the :attr:`input` tensor in the given +dimension :attr:`dim`. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {dtype} + +Example:: + + >>> a = torch.randn(4, 2) + >>> a + tensor([[ 0.5261, -0.3837], + [ 1.1857, -0.2498], + [-1.1646, 0.0705], + [ 1.1131, -1.0629]]) + >>> torch.prod(a, 1) + tensor([-0.2018, -0.2962, -0.0821, -1.1831]) +""".format(**single_dim_common), +) + +add_docstr( + torch.promote_types, + r""" +promote_types(type1, type2) -> dtype + +Returns the :class:`torch.dtype` with the smallest size and scalar kind that is +not smaller nor of lower kind than either `type1` or `type2`. See type promotion +:ref:`documentation ` for more information on the type +promotion logic. + +Args: + type1 (:class:`torch.dtype`) + type2 (:class:`torch.dtype`) + +Example:: + + >>> torch.promote_types(torch.int32, torch.float32) + torch.float32 + >>> torch.promote_types(torch.uint8, torch.long) + torch.long +""", +) + +add_docstr( + torch.qr, + r""" +qr(input: Tensor, some: bool = True, *, out: Union[Tensor, Tuple[Tensor, ...], List[Tensor], None]) -> (Tensor, Tensor) + +Computes the QR decomposition of a matrix or a batch of matrices :attr:`input`, +and returns a namedtuple (Q, R) of tensors such that :math:`\text{input} = Q R` +with :math:`Q` being an orthogonal matrix or batch of orthogonal matrices and +:math:`R` being an upper triangular matrix or batch of upper triangular matrices. + +If :attr:`some` is ``True``, then this function returns the thin (reduced) QR factorization. +Otherwise, if :attr:`some` is ``False``, this function returns the complete QR factorization. + +.. warning:: + + :func:`torch.qr` is deprecated in favor of :func:`torch.linalg.qr` + and will be removed in a future PyTorch release. The boolean parameter :attr:`some` has been + replaced with a string parameter :attr:`mode`. + + ``Q, R = torch.qr(A)`` should be replaced with + + .. code:: python + + Q, R = torch.linalg.qr(A) + + ``Q, R = torch.qr(A, some=False)`` should be replaced with + + .. code:: python + + Q, R = torch.linalg.qr(A, mode="complete") + +.. warning:: + If you plan to backpropagate through QR, note that the current backward implementation + is only well-defined when the first :math:`\min(input.size(-1), input.size(-2))` + columns of :attr:`input` are linearly independent. + This behavior will probably change once QR supports pivoting. + +.. note:: This function uses LAPACK for CPU inputs and MAGMA for CUDA inputs, + and may produce different (valid) decompositions on different device types + or different platforms. + +Args: + input (Tensor): the input tensor of size :math:`(*, m, n)` where `*` is zero or more + batch dimensions consisting of matrices of dimension :math:`m \times n`. + some (bool, optional): Set to ``True`` for reduced QR decomposition and ``False`` for + complete QR decomposition. If `k = min(m, n)` then: + + * ``some=True`` : returns `(Q, R)` with dimensions (m, k), (k, n) (default) + + * ``'some=False'``: returns `(Q, R)` with dimensions (m, m), (m, n) + +Keyword args: + out (tuple, optional): tuple of `Q` and `R` tensors. + The dimensions of `Q` and `R` are detailed in the description of :attr:`some` above. + +Example:: + + >>> a = torch.tensor([[12., -51, 4], [6, 167, -68], [-4, 24, -41]]) + >>> q, r = torch.qr(a) + >>> q + tensor([[-0.8571, 0.3943, 0.3314], + [-0.4286, -0.9029, -0.0343], + [ 0.2857, -0.1714, 0.9429]]) + >>> r + tensor([[ -14.0000, -21.0000, 14.0000], + [ 0.0000, -175.0000, 70.0000], + [ 0.0000, 0.0000, -35.0000]]) + >>> torch.mm(q, r).round() + tensor([[ 12., -51., 4.], + [ 6., 167., -68.], + [ -4., 24., -41.]]) + >>> torch.mm(q.t(), q).round() + tensor([[ 1., 0., 0.], + [ 0., 1., -0.], + [ 0., -0., 1.]]) + >>> a = torch.randn(3, 4, 5) + >>> q, r = torch.qr(a, some=False) + >>> torch.allclose(torch.matmul(q, r), a) + True + >>> torch.allclose(torch.matmul(q.mT, q), torch.eye(5)) + True +""", +) + +add_docstr( + torch.rad2deg, + r""" +rad2deg(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with each of the elements of :attr:`input` +converted from angles in radians to degrees. + +Args: + {input} + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.tensor([[3.142, -3.142], [6.283, -6.283], [1.570, -1.570]]) + >>> torch.rad2deg(a) + tensor([[ 180.0233, -180.0233], + [ 359.9894, -359.9894], + [ 89.9544, -89.9544]]) + +""".format(**common_args), +) + +add_docstr( + torch.deg2rad, + r""" +deg2rad(input, *, out=None) -> Tensor + +Returns a new tensor with each of the elements of :attr:`input` +converted from angles in degrees to radians. + +Args: + {input} + +Keyword arguments: + {out} + +Example:: + + >>> a = torch.tensor([[180.0, -180.0], [360.0, -360.0], [90.0, -90.0]]) + >>> torch.deg2rad(a) + tensor([[ 3.1416, -3.1416], + [ 6.2832, -6.2832], + [ 1.5708, -1.5708]]) + +""".format(**common_args), +) + +add_docstr( + torch.heaviside, + r""" +heaviside(input, values, *, out=None) -> Tensor + +Computes the Heaviside step function for each element in :attr:`input`. +The Heaviside step function is defined as: + +.. math:: + \text{{heaviside}}(input, values) = \begin{cases} + 0, & \text{if input < 0}\\ + values, & \text{if input == 0}\\ + 1, & \text{if input > 0} + \end{cases} +""" + + r""" + +Args: + {input} + values (Tensor): The values to use where :attr:`input` is zero. + +Keyword arguments: + {out} + +Example:: + + >>> input = torch.tensor([-1.5, 0, 2.0]) + >>> values = torch.tensor([0.5]) + >>> torch.heaviside(input, values) + tensor([0.0000, 0.5000, 1.0000]) + >>> values = torch.tensor([1.2, -2.0, 3.5]) + >>> torch.heaviside(input, values) + tensor([0., -2., 1.]) + +""".format(**common_args), +) + +add_docstr( + torch.rand, + """ +rand(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, \ +requires_grad=False, pin_memory=False) -> Tensor +""" + + r""" +Returns a tensor filled with random numbers from a uniform distribution +on the interval :math:`[0, 1)` + +The shape of the tensor is defined by the variable argument :attr:`size`. + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + +Keyword args: + {generator} + {out} + {dtype} + {layout} + {device} + {requires_grad} + {pin_memory} + +Example:: + + >>> torch.rand(4) + tensor([ 0.5204, 0.2503, 0.3525, 0.5673]) + >>> torch.rand(2, 3) + tensor([[ 0.8237, 0.5781, 0.6879], + [ 0.3816, 0.7249, 0.0998]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.rand_like, + r""" +rand_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + +Returns a tensor with the same size as :attr:`input` that is filled with +random numbers from a uniform distribution on the interval :math:`[0, 1)`. +``torch.rand_like(input)`` is equivalent to +``torch.rand(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + +Args: + {input} + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +""".format(**factory_like_common_args), +) + +add_docstr( + torch.randint, + """ +randint(low=0, high, size, \\*, generator=None, out=None, \ +dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a tensor filled with random integers generated uniformly +between :attr:`low` (inclusive) and :attr:`high` (exclusive). + +The shape of the tensor is defined by the variable argument :attr:`size`. + +.. note:: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + +Args: + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + size (tuple): a tuple defining the shape of the output tensor. + +Keyword args: + {generator} + {out} + dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if ``None``, + this function returns a tensor with dtype ``torch.int64``. + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.randint(3, 5, (3,)) + tensor([4, 3, 4]) + + + >>> torch.randint(10, (2, 2)) + tensor([[0, 2], + [5, 5]]) + + + >>> torch.randint(3, 10, (2, 2)) + tensor([[4, 5], + [6, 7]]) + + +""".format(**factory_common_args), +) + +add_docstr( + torch.randint_like, + """ +randint_like(input, low=0, high, \\*, dtype=None, layout=torch.strided, device=None, requires_grad=False, \ +memory_format=torch.preserve_format) -> Tensor + +Returns a tensor with the same shape as Tensor :attr:`input` filled with +random integers generated uniformly between :attr:`low` (inclusive) and +:attr:`high` (exclusive). + +.. note: + With the global dtype default (``torch.float32``), this function returns + a tensor with dtype ``torch.int64``. + +Args: + {input} + low (int, optional): Lowest integer to be drawn from the distribution. Default: 0. + high (int): One above the highest integer to be drawn from the distribution. + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +""".format(**factory_like_common_args), +) + +add_docstr( + torch.randn, + """ +randn(*size, *, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, \ +pin_memory=False) -> Tensor +""" + + r""" + +Returns a tensor filled with random numbers from a normal distribution +with mean `0` and variance `1` (also called the standard normal +distribution). + +.. math:: + \text{{out}}_{{i}} \sim \mathcal{{N}}(0, 1) + +For complex dtypes, the tensor is i.i.d. sampled from a `complex normal distribution`_ with zero mean and +unit variance as + +.. math:: + \text{{out}}_{{i}} \sim \mathcal{{CN}}(0, 1) + +This is equivalent to separately sampling the real :math:`(\operatorname{{Re}})` and imaginary +:math:`(\operatorname{{Im}})` part of :math:`\text{{out}}_i` as + +.. math:: + \operatorname{{Re}}(\text{{out}}_{{i}}) \sim \mathcal{{N}}(0, \frac{{1}}{{2}}),\quad + \operatorname{{Im}}(\text{{out}}_{{i}}) \sim \mathcal{{N}}(0, \frac{{1}}{{2}}) + +The shape of the tensor is defined by the variable argument :attr:`size`. + + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + +Keyword args: + {generator} + {out} + {dtype} + {layout} + {device} + {requires_grad} + {pin_memory} + +Example:: + + >>> torch.randn(4) + tensor([-2.1436, 0.9966, 2.3426, -0.6366]) + >>> torch.randn(2, 3) + tensor([[ 1.5954, 2.8929, -1.0923], + [ 1.1719, -0.4709, -0.1996]]) + +.. _complex normal distribution: https://en.wikipedia.org/wiki/Complex_normal_distribution +""".format(**factory_common_args), +) + +add_docstr( + torch.randn_like, + r""" +randn_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + +Returns a tensor with the same size as :attr:`input` that is filled with +random numbers from a normal distribution with mean 0 and variance 1. Please refer to :func:`torch.randn` for the +sampling process of complex dtypes. ``torch.randn_like(input)`` is equivalent to +``torch.randn(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + +Args: + {input} + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +""".format(**factory_like_common_args), +) + +add_docstr( + torch.randperm, + """ +randperm(n, *, generator=None, out=None, dtype=torch.int64,layout=torch.strided, \ +device=None, requires_grad=False, pin_memory=False) -> Tensor +""" + + r""" +Returns a random permutation of integers from ``0`` to ``n - 1``. + +Args: + n (int): the upper bound (exclusive) + +Keyword args: + {generator} + {out} + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: ``torch.int64``. + {layout} + {device} + {requires_grad} + {pin_memory} + +Example:: + + >>> torch.randperm(4) + tensor([2, 1, 0, 3]) +""".format(**factory_common_args), +) + +add_docstr( + torch.tensor, + r""" +tensor(data, *, dtype=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + +Constructs a tensor with no autograd history (also known as a "leaf tensor", see :doc:`/notes/autograd`) by copying :attr:`data`. + +.. warning:: + + When working with tensors prefer using :func:`torch.Tensor.clone`, + :func:`torch.Tensor.detach`, and :func:`torch.Tensor.requires_grad_` for + readability. Letting `t` be a tensor, ``torch.tensor(t)`` is equivalent to + ``t.detach().clone()``, and ``torch.tensor(t, requires_grad=True)`` + is equivalent to ``t.detach().clone().requires_grad_(True)``. + +.. seealso:: + + :func:`torch.as_tensor` preserves autograd history and avoids copies where possible. + :func:`torch.from_numpy` creates a tensor that shares storage with a NumPy array. + +Args: + {data} + +Keyword args: + {dtype} + device (:class:`torch.device`, optional): the device of the constructed tensor. If None and data is a tensor + then the device of data is used. If None and data is not a tensor then + the result tensor is constructed on the current device. + {requires_grad} + {pin_memory} + + +Example:: + + >>> torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]]) + tensor([[ 0.1000, 1.2000], + [ 2.2000, 3.1000], + [ 4.9000, 5.2000]]) + + >>> torch.tensor([0, 1]) # Type inference on data + tensor([ 0, 1]) + + >>> torch.tensor([[0.11111, 0.222222, 0.3333333]], + ... dtype=torch.float64, + ... device=torch.device('cuda:0')) # creates a double tensor on a CUDA device + tensor([[ 0.1111, 0.2222, 0.3333]], dtype=torch.float64, device='cuda:0') + + >>> torch.tensor(3.14159) # Create a zero-dimensional (scalar) tensor + tensor(3.1416) + + >>> torch.tensor([]) # Create an empty tensor (of size (0,)) + tensor([]) +""".format(**factory_data_common_args), +) + +add_docstr( + torch.range, + r""" +range(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a 1-D tensor of size :math:`\left\lfloor \frac{\text{end} - \text{start}}{\text{step}} \right\rfloor + 1` +with values from :attr:`start` to :attr:`end` with step :attr:`step`. Step is +the gap between two values in the tensor. + +.. math:: + \text{out}_{i+1} = \text{out}_i + \text{step}. +""" + + r""" +.. warning:: + This function is deprecated and will be removed in a future release because its behavior is inconsistent with + Python's range builtin. Instead, use :func:`torch.arange`, which produces values in [start, end). + +Args: + start (float, optional): the starting value for the set of points. Default: ``0``. + end (float): the ending value for the set of points + step (float, optional): the gap between each pair of adjacent points. Default: ``1``. + +Keyword args: + {out} + {dtype} If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `step` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.range(1, 4) + tensor([ 1., 2., 3., 4.]) + >>> torch.range(1, 4, 0.5) + tensor([ 1.0000, 1.5000, 2.0000, 2.5000, 3.0000, 3.5000, 4.0000]) +""".format(**factory_common_args), +) + +add_docstr( + torch.arange, + r""" +arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a 1-D tensor of size :math:`\left\lceil \frac{\text{end} - \text{start}}{\text{step}} \right\rceil` +with values from the interval ``[start, end)`` taken with common difference +:attr:`step` beginning from `start`. + +Note: When using floating-point dtypes (especially reduced precision types like ``bfloat16``), +the results may be affected by floating-point rounding behavior. Some values in the sequence +might not be exactly representable in certain floating-point formats, which can lead to +repeated values or unexpected rounding. For precise sequences, it is recommended to use +integer dtypes instead of floating-point dtypes. + +Note that non-integer :attr:`step` is subject to floating point rounding errors when +comparing against :attr:`end`; to avoid inconsistency, we advise subtracting a small epsilon from :attr:`end` +in such cases. + +.. math:: + \text{out}_{{i+1}} = \text{out}_{i} + \text{step} +""" + + r""" +Args: + start (Number, optional): the starting value for the set of points. Default: ``0``. + end (Number): the ending value for the set of points + step (Number, optional): the gap between each pair of adjacent points. Default: ``1``. + +Keyword args: + {out} + {dtype} If `dtype` is not given, infer the data type from the other input + arguments. If any of `start`, `end`, or `stop` are floating-point, the + `dtype` is inferred to be the default dtype, see + :meth:`~torch.get_default_dtype`. Otherwise, the `dtype` is inferred to + be `torch.int64`. + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.arange(5) + tensor([ 0, 1, 2, 3, 4]) + >>> torch.arange(1, 4) + tensor([ 1, 2, 3]) + >>> torch.arange(1, 2.5, 0.5) + tensor([ 1.0000, 1.5000, 2.0000]) +""".format(**factory_common_args), +) + +add_docstr( + torch.ravel, + r""" +ravel(input) -> Tensor + +Return a contiguous flattened tensor. A copy is made only if needed. + +Args: + {input} + +Example:: + + >>> t = torch.tensor([[[1, 2], + ... [3, 4]], + ... [[5, 6], + ... [7, 8]]]) + >>> torch.ravel(t) + tensor([1, 2, 3, 4, 5, 6, 7, 8]) +""".format(**common_args), +) + +add_docstr( + torch.remainder, + r""" +remainder(input, other, *, out=None) -> Tensor + +Computes +`Python's modulus operation `_ +entrywise. The result has the same sign as the divisor :attr:`other` and its absolute value +is less than that of :attr:`other`. + +It may also be defined in terms of :func:`torch.div` as + +.. code:: python + + torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer and float inputs. + +.. note:: + Complex inputs are not supported. In some cases, it is not mathematically + possible to satisfy the definition of a modulo operation with complex numbers. + See :func:`torch.fmod` for how division by zero is handled. + +.. seealso:: + + :func:`torch.fmod` which implements C++'s `std::fmod `_. + This one is defined in terms of division rounding towards zero. + +Args: + input (Tensor or Scalar): the dividend + other (Tensor or Scalar): the divisor + +Keyword args: + {out} + +Example:: + + >>> torch.remainder(torch.tensor([-3., -2, -1, 1, 2, 3]), 2) + tensor([ 1., 0., 1., 1., 0., 1.]) + >>> torch.remainder(torch.tensor([1, 2, 3, 4, 5]), -1.5) + tensor([ -0.5000, -1.0000, 0.0000, -0.5000, -1.0000 ]) +""".format(**common_args), +) + +add_docstr( + torch.renorm, + r""" +renorm(input, p, dim, maxnorm, *, out=None) -> Tensor + +Returns a tensor where each sub-tensor of :attr:`input` along dimension +:attr:`dim` is normalized such that the `p`-norm of the sub-tensor is lower +than the value :attr:`maxnorm` + +.. note:: If the norm of a row is lower than `maxnorm`, the row is unchanged + +Args: + {input} + p (float): the power for the norm computation + dim (int): the dimension to slice over to get the sub-tensors + maxnorm (float): the maximum norm to keep each sub-tensor under + +Keyword args: + {out} + +Example:: + + >>> x = torch.ones(3, 3) + >>> x[1].fill_(2) + tensor([ 2., 2., 2.]) + >>> x[2].fill_(3) + tensor([ 3., 3., 3.]) + >>> x + tensor([[ 1., 1., 1.], + [ 2., 2., 2.], + [ 3., 3., 3.]]) + >>> torch.renorm(x, 1, 0, 5) + tensor([[ 1.0000, 1.0000, 1.0000], + [ 1.6667, 1.6667, 1.6667], + [ 1.6667, 1.6667, 1.6667]]) +""".format(**common_args), +) + +add_docstr( + torch.reshape, + r""" +reshape(input, shape) -> Tensor + +Returns a tensor with the same data and number of elements as :attr:`input`, +but with the specified shape. When possible, the returned tensor will be a view +of :attr:`input`. Otherwise, it will be a copy. Contiguous inputs and inputs +with compatible strides can be reshaped without copying, but you should not +depend on the copying vs. viewing behavior. + +See :meth:`torch.Tensor.view` on when it is possible to return a view. + +A single dimension may be -1, in which case it's inferred from the remaining +dimensions and the number of elements in :attr:`input`. + +Args: + input (Tensor): the tensor to be reshaped + shape (tuple of int): the new shape + +Example:: + + >>> a = torch.arange(4.) + >>> torch.reshape(a, (2, 2)) + tensor([[ 0., 1.], + [ 2., 3.]]) + >>> b = torch.tensor([[0, 1], [2, 3]]) + >>> torch.reshape(b, (-1,)) + tensor([ 0, 1, 2, 3]) +""", +) + + +add_docstr( + torch.result_type, + r""" +result_type(tensor1, tensor2) -> dtype + +Returns the :class:`torch.dtype` that would result from performing an arithmetic +operation on the provided input tensors. See type promotion :ref:`documentation ` +for more information on the type promotion logic. + +Args: + tensor1 (Tensor or Number): an input tensor or number + tensor2 (Tensor or Number): an input tensor or number + +Example:: + + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.int), 1.0) + torch.float32 + >>> torch.result_type(torch.tensor([1, 2], dtype=torch.uint8), torch.tensor(1)) + torch.uint8 +""", +) + +add_docstr( + torch.row_stack, + r""" +row_stack(tensors, *, out=None) -> Tensor + +Alias of :func:`torch.vstack`. +""", +) + +add_docstr( + torch.round, + r""" +round(input, *, decimals=0, out=None) -> Tensor + +Rounds elements of :attr:`input` to the nearest integer. + +For integer inputs, follows the array-api convention of returning a +copy of the input tensor. +The return type of output is same as that of input's dtype. + +.. note:: + This function implements the "round half to even" to + break ties when a number is equidistant from two + integers (e.g. `round(2.5)` is 2). + + When the :attr:\`decimals\` argument is specified the + algorithm used is similar to NumPy's `around`. This + algorithm is fast but inexact and it can easily + overflow for low precision dtypes. + Eg. `round(tensor([10000], dtype=torch.float16), decimals=3)` is `inf`. + +.. seealso:: + :func:`torch.ceil`, which rounds up. + :func:`torch.floor`, which rounds down. + :func:`torch.trunc`, which rounds towards zero. + +Args: + {input} + decimals (int): Number of decimal places to round to (default: 0). + If decimals is negative, it specifies the number of positions + to the left of the decimal point. + +Keyword args: + {out} + +Example:: + + >>> torch.round(torch.tensor((4.7, -2.3, 9.1, -7.7))) + tensor([ 5., -2., 9., -8.]) + + >>> # Values equidistant from two integers are rounded towards the + >>> # the nearest even value (zero is treated as even) + >>> torch.round(torch.tensor([-0.5, 0.5, 1.5, 2.5])) + tensor([-0., 0., 2., 2.]) + + >>> # A positive decimals argument rounds to the to that decimal place + >>> torch.round(torch.tensor([0.1234567]), decimals=3) + tensor([0.1230]) + + >>> # A negative decimals argument rounds to the left of the decimal + >>> torch.round(torch.tensor([1200.1234567]), decimals=-3) + tensor([1000.]) +""".format(**common_args), +) + +add_docstr( + torch.rsqrt, + r""" +rsqrt(input, *, out=None) -> Tensor + +Returns a new tensor with the reciprocal of the square-root of each of +the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \frac{1}{\sqrt{\text{input}_{i}}} +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.0370, 0.2970, 1.5420, -0.9105]) + >>> torch.rsqrt(a) + tensor([ nan, 1.8351, 0.8053, nan]) +""".format(**common_args), +) + +add_docstr( + torch.scatter, + r""" +scatter(input, dim, index, src) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_` +""", +) + +add_docstr( + torch.scatter_add, + r""" +scatter_add(input, dim, index, src) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_add_` +""", +) + +add_docstr( + torch.scatter_reduce, + r""" +scatter_reduce(input, dim, index, src, reduce, *, include_self=True) -> Tensor + +Out-of-place version of :meth:`torch.Tensor.scatter_reduce_` +""", +) + +add_docstr( + torch.segment_reduce, + r""" +segment_reduce(data: Tensor, reduce: str, *, lengths: Tensor | None = None, indices: Tensor | None = None, offsets: Tensor | None = None, axis: _int = 0, unsafe: _bool = False, initial: Number | _complex | None = None) -> Tensor # noqa: B950 + +Perform a segment reduction operation on the input tensor along the specified axis. + +Args: + data (Tensor): The input tensor on which the segment reduction operation will be performed. + reduce (str): The type of reduction operation. Supported values are ``sum``, ``mean``, ``max``, ``min``, ``prod``. + +Keyword args: + lengths (Tensor, optional): Length of each segment. Default: ``None``. + offsets (Tensor, optional): Offset of each segment. Default: ``None``. + axis (int, optional): The axis perform reduction. Default: ``0``. + unsafe (bool, optional): Skip validation If `True`. Default: ``False``. + initial (Number, optional): The initial value for the reduction operation. Default: ``None``. + +Example:: + + >>> data = torch.tensor([[1, 2, 3, 4],[5, 6, 7, 8],[9, 10, 11, 12]], dtype=torch.float32, device='cuda') + >>> lengths = torch.tensor([2, 1], device='cuda') + >>> torch.segment_reduce(data, 'max', lengths=lengths) + tensor([[ 5., 6., 7., 8.], + [ 9., 10., 11., 12.]], device='cuda:0') +""", +) + +add_docstr( + torch.select, + r""" +select(input, dim, index) -> Tensor + +Slices the :attr:`input` tensor along the selected dimension at the given index. +This function returns a view of the original tensor with the given dimension removed. + +.. note:: If :attr:`input` is a sparse tensor and returning a view of + the tensor is not possible, a RuntimeError exception is + raised. In this is the case, consider using + :func:`torch.select_copy` function. + +Args: + {input} + dim (int): the dimension to slice + index (int): the index to select with + +.. note:: + + :meth:`select` is equivalent to slicing. For example, + ``tensor.select(0, index)`` is equivalent to ``tensor[index]`` and + ``tensor.select(2, index)`` is equivalent to ``tensor[:,:,index]``. +""".format(**common_args), +) + +add_docstr( + torch.select_scatter, + r""" +select_scatter(input, src, dim, index) -> Tensor + +Embeds the values of the :attr:`src` tensor into :attr:`input` at the given index. +This function returns a tensor with fresh storage; it does not create a view. + + +Args: + {input} + src (Tensor): The tensor to embed into :attr:`input` + dim (int): the dimension to insert the slice into. + index (int): the index to select with + +.. note:: + + :attr:`src` must be of the proper size in order to be embedded + into :attr:`input`. Specifically, it should have the same shape as + ``torch.select(input, dim, index)`` + +Example:: + + >>> a = torch.zeros(2, 2) + >>> b = torch.ones(2) + >>> a.select_scatter(b, 0, 0) + tensor([[1., 1.], + [0., 0.]]) +""".format(**common_args), +) + +add_docstr( + torch.slice_scatter, + r""" +slice_scatter(input, src, dim=0, start=None, end=None, step=1) -> Tensor + +Embeds the values of the :attr:`src` tensor into :attr:`input` at the given +dimension. +This function returns a tensor with fresh storage; it does not create a view. + + +Args: + {input} + src (Tensor): The tensor to embed into :attr:`input` + dim (int): the dimension to insert the slice into + start (Optional[int]): the start index of where to insert the slice + end (Optional[int]): the end index of where to insert the slice + step (int): the how many elements to skip in + +Example:: + + >>> a = torch.zeros(8, 8) + >>> b = torch.ones(2, 8) + >>> a.slice_scatter(b, start=6) + tensor([[0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1., 1., 1.]]) + + >>> b = torch.ones(8, 2) + >>> a.slice_scatter(b, dim=1, start=2, end=6, step=2) + tensor([[0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.], + [0., 0., 1., 0., 1., 0., 0., 0.]]) +""".format(**common_args), +) + +add_docstr( + torch.set_flush_denormal, + r""" +set_flush_denormal(mode) -> bool + +Disables denormal floating numbers on CPU. + +Returns ``True`` if your system supports flushing denormal numbers and it +successfully configures flush denormal mode. :meth:`~torch.set_flush_denormal` +is supported on x86 architectures supporting SSE3 and AArch64 architecture. + +Args: + mode (bool): Controls whether to enable flush denormal mode or not + +Example:: + + >>> torch.set_flush_denormal(True) + True + >>> torch.tensor([1e-323], dtype=torch.float64) + tensor([ 0.], dtype=torch.float64) + >>> torch.set_flush_denormal(False) + True + >>> torch.tensor([1e-323], dtype=torch.float64) + tensor(9.88131e-324 * + [ 1.0000], dtype=torch.float64) +""", +) + +add_docstr( + torch.set_num_threads, + r""" +set_num_threads(int) + +Sets the number of threads used for intraop parallelism on CPU. + +.. warning:: + To ensure that the correct number of threads is used, set_num_threads + must be called before running eager, JIT or autograd code. +""", +) + +add_docstr( + torch.set_num_interop_threads, + r""" +set_num_interop_threads(int) + +Sets the number of threads used for interop parallelism +(e.g. in JIT interpreter) on CPU. + +.. warning:: + Can only be called once and before any inter-op parallel work + is started (e.g. JIT execution). +""", +) + +add_docstr( + torch.sigmoid, + r""" +sigmoid(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.expit`. +""", +) + +add_docstr( + torch.logit, + r""" +logit(input, eps=None, *, out=None) -> Tensor + +Alias for :func:`torch.special.logit`. +""", +) + +add_docstr( + torch.sign, + r""" +sign(input, *, out=None) -> Tensor + +Returns a new tensor with the signs of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \operatorname{sgn}(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([0.7, -1.2, 0., 2.3]) + >>> a + tensor([ 0.7000, -1.2000, 0.0000, 2.3000]) + >>> torch.sign(a) + tensor([ 1., -1., 0., 1.]) +""".format(**common_args), +) + +add_docstr( + torch.signbit, + r""" +signbit(input, *, out=None) -> Tensor + +Tests if each element of :attr:`input` has its sign bit set or not. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([0.7, -1.2, 0., 2.3]) + >>> torch.signbit(a) + tensor([ False, True, False, False]) + >>> a = torch.tensor([-0.0, 0.0]) + >>> torch.signbit(a) + tensor([ True, False]) + +.. note:: + signbit handles signed zeros, so negative zero (-0) returns True. + +""".format(**common_args), +) + +add_docstr( + torch.sgn, + r""" +sgn(input, *, out=None) -> Tensor + +This function is an extension of torch.sign() to complex tensors. +It computes a new tensor whose elements have +the same angles as the corresponding elements of :attr:`input` and +absolute values (i.e. magnitudes) of one for complex tensors and +is equivalent to torch.sign() for non-complex tensors. + +.. math:: + \text{out}_{i} = \begin{cases} + 0 & |\text{{input}}_i| == 0 \\ + \frac{{\text{{input}}_i}}{|{\text{{input}}_i}|} & \text{otherwise} + \end{cases} + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> t = torch.tensor([3+4j, 7-24j, 0, 1+2j]) + >>> t.sgn() + tensor([0.6000+0.8000j, 0.2800-0.9600j, 0.0000+0.0000j, 0.4472+0.8944j]) +""".format(**common_args), +) + +add_docstr( + torch.sin, + r""" +sin(input, *, out=None) -> Tensor + +Returns a new tensor with the sine of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \sin(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-0.5461, 0.1347, -2.7266, -0.2746]) + >>> torch.sin(a) + tensor([-0.5194, 0.1343, -0.4032, -0.2711]) +""".format(**common_args), +) + +add_docstr( + torch.sinc, + r""" +sinc(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.sinc`. +""", +) + +add_docstr( + torch.sinh, + r""" +sinh(input, *, out=None) -> Tensor + +Returns a new tensor with the hyperbolic sine of the elements of +:attr:`input`. + +.. math:: + \text{out}_{i} = \sinh(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.5380, -0.8632, -0.1265, 0.9399]) + >>> torch.sinh(a) + tensor([ 0.5644, -0.9744, -0.1268, 1.0845]) + +.. note:: + When :attr:`input` is on the CPU, the implementation of torch.sinh may use + the Sleef library, which rounds very large results to infinity or negative + infinity. See `here `_ for details. +""".format(**common_args), +) + +add_docstr( + torch.sort, + r""" +sort(input, dim=-1, descending=False, *, stable=False, out=None) -> (Tensor, LongTensor) + +Sorts the elements of the :attr:`input` tensor along a given dimension +in ascending order by value. + +If :attr:`dim` is not given, the last dimension of the `input` is chosen. + +If :attr:`descending` is ``True`` then the elements are sorted in descending +order by value. + +If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving +the order of equivalent elements. + +A namedtuple of (values, indices) is returned, where the `values` are the +sorted values and `indices` are the indices of the elements in the original +`input` tensor. + +Args: + {input} + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + +Keyword args: + stable (bool, optional): makes the sorting routine stable, which guarantees that the order + of equivalent elements is preserved. + out (tuple, optional): the output tuple of (`Tensor`, `LongTensor`) that can + be optionally given to be used as output buffers + +Example:: + + >>> x = torch.randn(3, 4) + >>> sorted, indices = torch.sort(x) + >>> sorted + tensor([[-0.2162, 0.0608, 0.6719, 2.3332], + [-0.5793, 0.0061, 0.6058, 0.9497], + [-0.5071, 0.3343, 0.9553, 1.0960]]) + >>> indices + tensor([[ 1, 0, 2, 3], + [ 3, 1, 0, 2], + [ 0, 3, 1, 2]]) + + >>> sorted, indices = torch.sort(x, 0) + >>> sorted + tensor([[-0.5071, -0.2162, 0.6719, -0.5793], + [ 0.0608, 0.0061, 0.9497, 0.3343], + [ 0.6058, 0.9553, 1.0960, 2.3332]]) + >>> indices + tensor([[ 2, 0, 0, 1], + [ 0, 1, 1, 2], + [ 1, 2, 2, 0]]) + >>> x = torch.tensor([0, 1] * 9) + >>> x.sort() + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 2, 16, 4, 6, 14, 8, 0, 10, 12, 9, 17, 15, 13, 11, 7, 5, 3, 1])) + >>> x.sort(stable=True) + torch.return_types.sort( + values=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]), + indices=tensor([ 0, 2, 4, 6, 8, 10, 12, 14, 16, 1, 3, 5, 7, 9, 11, 13, 15, 17])) +""".format(**common_args), +) + +add_docstr( + torch.argsort, + r""" +argsort(input, dim=-1, descending=False, *, stable=False) -> Tensor + +Returns the indices that sort a tensor along a given dimension in ascending +order by value. + +This is the second value returned by :meth:`torch.sort`. See its documentation +for the exact semantics of this method. + +If :attr:`stable` is ``True`` then the sorting routine becomes stable, preserving +the order of equivalent elements. If ``False``, the relative order of values +which compare equal is not guaranteed. ``True`` is slower. + +Args: + {input} + dim (int, optional): the dimension to sort along + descending (bool, optional): controls the sorting order (ascending or descending) + +Keyword args: + stable (bool, optional): controls the relative order of equivalent elements + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0785, 1.5267, -0.8521, 0.4065], + [ 0.1598, 0.0788, -0.0745, -1.2700], + [ 1.2208, 1.0722, -0.7064, 1.2564], + [ 0.0669, -0.2318, -0.8229, -0.9280]]) + + + >>> torch.argsort(a, dim=1) + tensor([[2, 0, 3, 1], + [3, 2, 1, 0], + [2, 1, 0, 3], + [3, 2, 1, 0]]) +""".format(**common_args), +) + +add_docstr( + torch.msort, + r""" +msort(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Sorts the elements of the :attr:`input` tensor along its first dimension +in ascending order by value. + +.. note:: `torch.msort(t)` is equivalent to `torch.sort(t, dim=0)[0]`. + See also :func:`torch.sort`. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> t = torch.randn(3, 4) + >>> t + tensor([[-0.1321, 0.4370, -1.2631, -1.1289], + [-2.0527, -1.1250, 0.2275, 0.3077], + [-0.0881, -0.1259, -0.5495, 1.0284]]) + >>> torch.msort(t) + tensor([[-2.0527, -1.1250, -1.2631, -1.1289], + [-0.1321, -0.1259, -0.5495, 0.3077], + [-0.0881, 0.4370, 0.2275, 1.0284]]) +""".format(**common_args), +) + +add_docstr( + torch.sparse_compressed_tensor, + r"""sparse_compressed_tensor(compressed_indices, plain_indices, values, size=None, """ + r"""*, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + +Constructs a :ref:`sparse tensor in Compressed Sparse format - CSR, +CSC, BSR, or BSC - ` with specified values at +the given :attr:`compressed_indices` and :attr:`plain_indices`. Sparse +matrix multiplication operations in Compressed Sparse format are +typically faster than that for sparse tensors in COO format. Make you +have a look at :ref:`the note on the data type of the indices +`. + +{sparse_factory_device_note} + +Args: + compressed_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, compressed_dim_size + 1)``. The last element of + each batch is the number of non-zero elements or blocks. This + tensor encodes the index in ``values`` and ``plain_indices`` + depending on where the given compressed dimension (row or + column) starts. Each successive number in the tensor + subtracted by the number before it denotes the number of + elements or blocks in a given compressed dimension. + plain_indices (array_like): Plain dimension (column or row) + coordinates of each element or block in values. (B+1)-dimensional + tensor with the same length as values. + + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types. that + represents a (1+K)-dimensional (for CSR and CSC layouts) or + (1+2+K)-dimensional tensor (for BSR and BSC layouts) where + ``K`` is the number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` where ``blocksize[0] == + blocksize[1] == 1`` for CSR and CSC formats. If not provided, + the size will be inferred as the minimum size big enough to + hold all non-zero elements or blocks. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + layout (:class:`torch.layout`, required): the desired layout of + returned tensor: :attr:`torch.sparse_csr`, + :attr:`torch.sparse_csc`, :attr:`torch.sparse_bsr`, or + :attr:`torch.sparse_bsc`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + +Example:: + + >>> compressed_indices = [0, 2, 4] + >>> plain_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_compressed_tensor(torch.tensor(compressed_indices, dtype=torch.int64), + ... torch.tensor(plain_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double, layout=torch.sparse_csr) + tensor(crow_indices=tensor([0, 2, 4]), + col_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csr) +""".format(**factory_common_args), +) + +add_docstr( + torch.sparse_csr_tensor, + r"""sparse_csr_tensor(crow_indices, col_indices, values, size=None, """ + r"""*, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + +Constructs a :ref:`sparse tensor in CSR (Compressed Sparse Row) ` with specified +values at the given :attr:`crow_indices` and :attr:`col_indices`. Sparse matrix multiplication operations +in CSR format are typically faster than that for sparse tensors in COO format. Make you have a look +at :ref:`the note on the data type of the indices `. + +{sparse_factory_device_note} + +Args: + crow_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, nrows + 1)``. The last element of each batch + is the number of non-zeros. This tensor encodes the index in + values and col_indices depending on where the given row + starts. Each successive number in the tensor subtracted by the + number before it denotes the number of elements in a given + row. + col_indices (array_like): Column coordinates of each element in + values. (B+1)-dimensional tensor with the same length + as values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1+K)-dimensional tensor where ``K`` is the number + of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows, ncols, *densesize)``. If + not provided, the size will be inferred as the minimum size + big enough to hold all non-zero elements. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + +Example:: + + >>> crow_indices = [0, 2, 4] + >>> col_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_csr_tensor(torch.tensor(crow_indices, dtype=torch.int64), + ... torch.tensor(col_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(crow_indices=tensor([0, 2, 4]), + col_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csr) +""".format(**factory_common_args), +) + +add_docstr( + torch.sparse_csc_tensor, + r"""sparse_csc_tensor(ccol_indices, row_indices, values, size=None, """ + r"""*, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + +Constructs a :ref:`sparse tensor in CSC (Compressed Sparse Column) +` with specified values at the given +:attr:`ccol_indices` and :attr:`row_indices`. Sparse matrix +multiplication operations in CSC format are typically faster than that +for sparse tensors in COO format. Make you have a look at :ref:`the +note on the data type of the indices `. + +{sparse_factory_device_note} + +Args: + ccol_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, ncols + 1)``. The last element of each batch + is the number of non-zeros. This tensor encodes the index in + values and row_indices depending on where the given column + starts. Each successive number in the tensor subtracted by the + number before it denotes the number of elements in a given + column. + row_indices (array_like): Row coordinates of each element in + values. (B+1)-dimensional tensor with the same length as + values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1+K)-dimensional tensor where ``K`` is the number + of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows, ncols, *densesize)``. If + not provided, the size will be inferred as the minimum size + big enough to hold all non-zero elements. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + +Example:: + + >>> ccol_indices = [0, 2, 4] + >>> row_indices = [0, 1, 0, 1] + >>> values = [1, 2, 3, 4] + >>> torch.sparse_csc_tensor(torch.tensor(ccol_indices, dtype=torch.int64), + ... torch.tensor(row_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(ccol_indices=tensor([0, 2, 4]), + row_indices=tensor([0, 1, 0, 1]), + values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4, + dtype=torch.float64, layout=torch.sparse_csc) +""".format(**factory_common_args), +) + +add_docstr( + torch.sparse_bsr_tensor, + r"""sparse_bsr_tensor(crow_indices, col_indices, values, size=None, """ + r"""*, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + +Constructs a :ref:`sparse tensor in BSR (Block Compressed Sparse Row)) +` with specified 2-dimensional blocks at the given +:attr:`crow_indices` and :attr:`col_indices`. Sparse matrix +multiplication operations in BSR format are typically faster than that +for sparse tensors in COO format. Make you have a look at :ref:`the +note on the data type of the indices `. + +{sparse_factory_device_note} + +Args: + crow_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, nrowblocks + 1)``. The last element of each + batch is the number of non-zeros. This tensor encodes the + block index in values and col_indices depending on where the + given row block starts. Each successive number in the tensor + subtracted by the number before it denotes the number of + blocks in a given row. + col_indices (array_like): Column block coordinates of each block + in values. (B+1)-dimensional tensor with the same length as + values. + values (array_list): Initial values for the tensor. Can be a list, + tuple, NumPy ``ndarray``, scalar, and other types that + represents a (1 + 2 + K)-dimensional tensor where ``K`` is the + number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` where ``blocksize == + values.shape[1:3]``. If not provided, the size will be + inferred as the minimum size big enough to hold all non-zero + blocks. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + +Example:: + + >>> crow_indices = [0, 1, 2] + >>> col_indices = [0, 1] + >>> values = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> torch.sparse_bsr_tensor(torch.tensor(crow_indices, dtype=torch.int64), + ... torch.tensor(col_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(crow_indices=tensor([0, 1, 2]), + col_indices=tensor([0, 1]), + values=tensor([[[1., 2.], + [3., 4.]], + [[5., 6.], + [7., 8.]]]), size=(2, 2), nnz=2, dtype=torch.float64, + layout=torch.sparse_bsr) +""".format(**factory_common_args), +) + +add_docstr( + torch.sparse_bsc_tensor, + r"""sparse_bsc_tensor(ccol_indices, row_indices, values, size=None, """ + r"""*, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None) -> Tensor + +Constructs a :ref:`sparse tensor in BSC (Block Compressed Sparse +Column)) ` with specified 2-dimensional blocks at the +given :attr:`ccol_indices` and :attr:`row_indices`. Sparse matrix +multiplication operations in BSC format are typically faster than that +for sparse tensors in COO format. Make you have a look at :ref:`the +note on the data type of the indices `. + +{sparse_factory_device_note} + +Args: + ccol_indices (array_like): (B+1)-dimensional array of size + ``(*batchsize, ncolblocks + 1)``. The last element of each + batch is the number of non-zeros. This tensor encodes the + index in values and row_indices depending on where the given + column starts. Each successive number in the tensor subtracted + by the number before it denotes the number of elements in a + given column. + row_indices (array_like): Row block coordinates of each block in + values. (B+1)-dimensional tensor with the same length + as values. + values (array_list): Initial blocks for the tensor. Can be a list, + tuple, NumPy ``ndarray``, and other types that + represents a (1 + 2 + K)-dimensional tensor where ``K`` is the + number of dense dimensions. + size (list, tuple, :class:`torch.Size`, optional): Size of the + sparse tensor: ``(*batchsize, nrows * blocksize[0], ncols * + blocksize[1], *densesize)`` If not provided, the size will be + inferred as the minimum size big enough to hold all non-zero + blocks. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. Default: if None, infers data type from + :attr:`values`. + device (:class:`torch.device`, optional): the desired device of + returned tensor. Default: if None, uses the current device + for the default tensor type (see + :func:`torch.set_default_device`). :attr:`device` will be + the CPU for CPU tensor types and the current CUDA device for + CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + +Example:: + + >>> ccol_indices = [0, 1, 2] + >>> row_indices = [0, 1] + >>> values = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + >>> torch.sparse_bsc_tensor(torch.tensor(ccol_indices, dtype=torch.int64), + ... torch.tensor(row_indices, dtype=torch.int64), + ... torch.tensor(values), dtype=torch.double) + tensor(ccol_indices=tensor([0, 1, 2]), + row_indices=tensor([0, 1]), + values=tensor([[[1., 2.], + [3., 4.]], + [[5., 6.], + [7., 8.]]]), size=(2, 2), nnz=2, dtype=torch.float64, + layout=torch.sparse_bsc) +""".format(**factory_common_args), +) + +add_docstr( + torch.sparse_coo_tensor, + r"""sparse_coo_tensor(indices, values, size=None, """ + r"""*, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None, is_coalesced=None) -> Tensor + +Constructs a :ref:`sparse tensor in COO(rdinate) format +` with specified values at the given +:attr:`indices`. + +.. note:: + + This function returns an :ref:`uncoalesced tensor + ` when :attr:`is_coalesced` is + unspecified or ``None``. + +{sparse_factory_device_note} + +Args: + indices (array_like): Initial data for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. Will be cast to a :class:`torch.LongTensor` + internally. The indices are the coordinates of the non-zero values in the matrix, and thus + should be two-dimensional where the first dimension is the number of tensor dimensions and + the second dimension is the number of non-zero values. + values (array_like): Initial values for the tensor. Can be a list, tuple, + NumPy ``ndarray``, scalar, and other types. + size (list, tuple, or :class:`torch.Size`, optional): Size of the sparse tensor. If not + provided the size will be inferred as the minimum size big enough to hold all non-zero + elements. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + Default: if None, infers data type from :attr:`values`. + device (:class:`torch.device`, optional): the desired device of returned tensor. + Default: if None, uses the current device for the default tensor type + (see :func:`torch.set_default_device`). :attr:`device` will be the CPU + for CPU tensor types and the current CUDA device for CUDA tensor types. + {pin_memory} + {requires_grad} + {check_invariants} + is_coalesced (bool, optional): When``True``, the caller is + responsible for providing tensor indices that correspond to a + coalesced tensor. If the :attr:`check_invariants` flag is + False, no error will be raised if the prerequisites are not + met and this will lead to silently incorrect results. To force + coalescion please use :meth:`coalesce` on the resulting + Tensor. + Default: None: except for trivial cases (e.g. nnz < 2) the + resulting Tensor has is_coalesced set to ``False```. + +Example:: + + >>> i = torch.tensor([[0, 1, 1], + ... [2, 0, 2]]) + >>> v = torch.tensor([3, 4, 5], dtype=torch.float32) + >>> torch.sparse_coo_tensor(i, v, [2, 4]) + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + size=(2, 4), nnz=3, layout=torch.sparse_coo) + + >>> torch.sparse_coo_tensor(i, v) # Shape inference + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + size=(2, 3), nnz=3, layout=torch.sparse_coo) + + >>> torch.sparse_coo_tensor(i, v, [2, 4], + ... dtype=torch.float64, + ... device=torch.device('cuda:0')) + tensor(indices=tensor([[0, 1, 1], + [2, 0, 2]]), + values=tensor([3., 4., 5.]), + device='cuda:0', size=(2, 4), nnz=3, dtype=torch.float64, + layout=torch.sparse_coo) + + # Create an empty sparse tensor with the following invariants: + # 1. sparse_dim + dense_dim = len(SparseTensor.shape) + # 2. SparseTensor._indices().shape = (sparse_dim, nnz) + # 3. SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:]) + # + # For instance, to create an empty sparse tensor with nnz = 0, dense_dim = 0 and + # sparse_dim = 1 (hence indices is a 2D tensor of shape = (1, 0)) + >>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), [], [1]) + tensor(indices=tensor([], size=(1, 0)), + values=tensor([], size=(0,)), + size=(1,), nnz=0, layout=torch.sparse_coo) + + # and to create an empty sparse tensor with nnz = 0, dense_dim = 1 and + # sparse_dim = 1 + >>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), torch.empty([0, 2]), [1, 2]) + tensor(indices=tensor([], size=(1, 0)), + values=tensor([], size=(0, 2)), + size=(1, 2), nnz=0, layout=torch.sparse_coo) + +.. _torch.sparse: https://pytorch.org/docs/stable/sparse.html +""".format(**factory_common_args), +) + +add_docstr( + torch.sqrt, + r""" +sqrt(input, *, out=None) -> Tensor + +Returns a new tensor with the square-root of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \sqrt{\text{input}_{i}} +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-2.0755, 1.0226, 0.0831, 0.4806]) + >>> torch.sqrt(a) + tensor([ nan, 1.0112, 0.2883, 0.6933]) +""".format(**common_args), +) + +add_docstr( + torch.square, + r""" +square(input: Tensor, *, out: Optional[Tensor]) -> Tensor + +Returns a new tensor with the square of the elements of :attr:`input`. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-2.0755, 1.0226, 0.0831, 0.4806]) + >>> torch.square(a) + tensor([ 4.3077, 1.0457, 0.0069, 0.2310]) +""".format(**common_args), +) + +add_docstr( + torch.squeeze, + r""" +squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor + +Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed. + +For example, if `input` is of shape: +:math:`(A \times 1 \times B \times C \times 1 \times D)` then the `input.squeeze()` +will be of shape: :math:`(A \times B \times C \times D)`. + +When :attr:`dim` is given, a squeeze operation is done only in the given +dimension(s). If `input` is of shape: :math:`(A \times 1 \times B)`, +``squeeze(input, 0)`` leaves the tensor unchanged, but ``squeeze(input, 1)`` +will squeeze the tensor to the shape :math:`(A \times B)`. + +.. note:: The returned tensor shares the storage with the input tensor, + so changing the contents of one will change the contents of the other. + +.. warning:: If the tensor has a batch dimension of size 1, then `squeeze(input)` + will also remove the batch dimension, which can lead to unexpected + errors. Consider specifying only the dims you wish to be squeezed. + +Args: + {input} + dim (int or tuple of ints, optional): if given, the input will be squeezed + only in the specified dimensions. + + .. versionchanged:: 2.0 + :attr:`dim` now accepts tuples of dimensions. + +Example:: + + >>> x = torch.zeros(2, 1, 2, 1, 2) + >>> x.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x) + >>> y.size() + torch.Size([2, 2, 2]) + >>> y = torch.squeeze(x, 0) + >>> y.size() + torch.Size([2, 1, 2, 1, 2]) + >>> y = torch.squeeze(x, 1) + >>> y.size() + torch.Size([2, 2, 1, 2]) + >>> y = torch.squeeze(x, (1, 2, 3)) + torch.Size([2, 2, 2]) +""".format(**common_args), +) + +add_docstr( + torch.std, + r""" +std(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + +Calculates the standard deviation over the dimensions specified by :attr:`dim`. +:attr:`dim` can be a single dimension, list of dimensions, or ``None`` to +reduce over all dimensions. + +The standard deviation (:math:`\sigma`) is calculated as + +.. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + +where :math:`x` is the sample set of elements, :math:`\bar{x}` is the +sample mean, :math:`N` is the number of samples and :math:`\delta N` is +the :attr:`correction`. +""" + + r""" + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + +Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + {opt_keepdim} + {out} + +Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std(a, dim=1, keepdim=True) + tensor([[1.0311], + [0.7477], + [1.2204], + [0.9087]]) + +.. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + +""".format(**multi_dim_common), +) + +add_docstr( + torch.std_mean, + r""" +std_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + +Calculates the standard deviation and mean over the dimensions specified by +:attr:`dim`. :attr:`dim` can be a single dimension, list of dimensions, or +``None`` to reduce over all dimensions. + +The standard deviation (:math:`\sigma`) is calculated as + +.. math:: \sigma = \sqrt{\frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2} + +where :math:`x` is the sample set of elements, :math:`\bar{x}` is the +sample mean, :math:`N` is the number of samples and :math:`\delta N` is +the :attr:`correction`. + +""" + + r""" + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + +Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + {opt_keepdim} + {out} + +Returns: + A tuple (std, mean) containing the standard deviation and mean. + +Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.std_mean(a, dim=0, keepdim=True) + (tensor([[1.2620, 1.0028, 1.0957, 0.6038]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + +.. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + +""".format(**multi_dim_common), +) + +add_docstr( + torch.sub, + r""" +sub(input, other, *, alpha=1, out=None) -> Tensor + +Subtracts :attr:`other`, scaled by :attr:`alpha`, from :attr:`input`. + +.. math:: + \text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i +""" + + r""" + +Supports :ref:`broadcasting to a common shape `, +:ref:`type promotion `, and integer, float, and complex inputs. + +Args: + {input} + other (Tensor or Number): the tensor or number to subtract from :attr:`input`. + +Keyword args: + alpha (Number): the multiplier for :attr:`other`. + {out} + +Example:: + + >>> a = torch.tensor((1, 2)) + >>> b = torch.tensor((0, 1)) + >>> torch.sub(a, b, alpha=2) + tensor([1, 0]) +""".format(**common_args), +) + +add_docstr( + torch.subtract, + r""" +subtract(input, other, *, alpha=1, out=None) -> Tensor + +Alias for :func:`torch.sub`. +""", +) + +add_docstr( + torch.sum, + r""" +sum(input, *, dtype=None) -> Tensor + +Returns the sum of all elements in the :attr:`input` tensor. + +Args: + {input} + +Keyword args: + {dtype} + +.. note:: Use the `dtype` argument if you need the result in a specific tensor type. + Otherwise, the result type may be automatically promoted (e.g., from `torch.int32` to `torch.int64`). + +Example:: + + >>> a = torch.randn(1, 3) + >>> a + tensor([[ 0.1133, -0.9567, 0.2958]]) + >>> torch.sum(a) + tensor(-0.5475) + +.. function:: sum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + +Returns the sum of each row of the :attr:`input` tensor in the given +dimension :attr:`dim`. If :attr:`dim` is a list of dimensions, +reduce over all of them. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {dtype} + +Example:: + + >>> a = torch.randn(4, 4) + >>> a + tensor([[ 0.0569, -0.2475, 0.0737, -0.3429], + [-0.2993, 0.9138, 0.9337, -1.6864], + [ 0.1132, 0.7892, -0.1003, 0.5688], + [ 0.3637, -0.9906, -0.4752, -1.5197]]) + >>> torch.sum(a, 1) + tensor([-0.4598, -0.1381, 1.3708, -2.6217]) + >>> b = torch.arange(4 * 5 * 6).view(4, 5, 6) + >>> torch.sum(b, (2, 1)) + tensor([ 435., 1335., 2235., 3135.]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.nansum, + r""" +nansum(input, *, dtype=None) -> Tensor + +Returns the sum of all elements, treating Not a Numbers (NaNs) as zero. + +Args: + {input} + +Keyword args: + {dtype} + +Example:: + + >>> a = torch.tensor([1., 2., float('nan'), 4.]) + >>> torch.nansum(a) + tensor(7.) + +.. function:: nansum(input, dim, keepdim=False, *, dtype=None) -> Tensor + :noindex: + +Returns the sum of each row of the :attr:`input` tensor in the given +dimension :attr:`dim`, treating Not a Numbers (NaNs) as zero. +If :attr:`dim` is a list of dimensions, reduce over all of them. + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + {opt_keepdim} + +Keyword args: + {dtype} + +Example:: + + >>> torch.nansum(torch.tensor([1., float("nan")])) + tensor(1.) + >>> a = torch.tensor([[1, 2], [3., float("nan")]]) + >>> torch.nansum(a) + tensor(6.) + >>> torch.nansum(a, dim=0) + tensor([4., 2.]) + >>> torch.nansum(a, dim=1) + tensor([3., 3.]) +""".format(**multi_dim_common), +) + +add_docstr( + torch.svd, + r""" +svd(input, some=True, compute_uv=True, *, out=None) -> (Tensor, Tensor, Tensor) + +Computes the singular value decomposition of either a matrix or batch of +matrices :attr:`input`. The singular value decomposition is represented as a +namedtuple `(U, S, V)`, such that :attr:`input` :math:`= U \text{diag}(S) V^{\text{H}}`. +where :math:`V^{\text{H}}` is the transpose of `V` for real inputs, +and the conjugate transpose of `V` for complex inputs. +If :attr:`input` is a batch of matrices, then `U`, `S`, and `V` are also +batched with the same batch dimensions as :attr:`input`. + +If :attr:`some` is `True` (default), the method returns the reduced singular +value decomposition. In this case, if the last two dimensions of :attr:`input` are +`m` and `n`, then the returned `U` and `V` matrices will contain only +`min(n, m)` orthonormal columns. + +If :attr:`compute_uv` is `False`, the returned `U` and `V` will be +zero-filled matrices of shape `(m, m)` and `(n, n)` +respectively, and the same device as :attr:`input`. The argument :attr:`some` +has no effect when :attr:`compute_uv` is `False`. + +Supports :attr:`input` of float, double, cfloat and cdouble data types. +The dtypes of `U` and `V` are the same as :attr:`input`'s. `S` will +always be real-valued, even if :attr:`input` is complex. + +.. warning:: + + :func:`torch.svd` is deprecated in favor of :func:`torch.linalg.svd` + and will be removed in a future PyTorch release. + + ``U, S, V = torch.svd(A, some=some, compute_uv=True)`` (default) should be replaced with + + .. code:: python + + U, S, Vh = torch.linalg.svd(A, full_matrices=not some) + V = Vh.mH + + ``_, S, _ = torch.svd(A, some=some, compute_uv=False)`` should be replaced with + + .. code:: python + + S = torch.linalg.svdvals(A) + +.. note:: Differences with :func:`torch.linalg.svd`: + + * :attr:`some` is the opposite of + :func:`torch.linalg.svd`'s :attr:`full_matrices`. Note that + default value for both is `True`, so the default behavior is + effectively the opposite. + * :func:`torch.svd` returns `V`, whereas :func:`torch.linalg.svd` returns + `Vh`, that is, :math:`V^{\text{H}}`. + * If :attr:`compute_uv` is `False`, :func:`torch.svd` returns zero-filled + tensors for `U` and `Vh`, whereas :func:`torch.linalg.svd` returns + empty tensors. + +.. note:: The singular values are returned in descending order. If :attr:`input` is a batch of matrices, + then the singular values of each matrix in the batch are returned in descending order. + +.. note:: The `S` tensor can only be used to compute gradients if :attr:`compute_uv` is `True`. + +.. note:: When :attr:`some` is `False`, the gradients on `U[..., :, min(m, n):]` + and `V[..., :, min(m, n):]` will be ignored in the backward pass, as those vectors + can be arbitrary bases of the corresponding subspaces. + +.. note:: The implementation of :func:`torch.linalg.svd` on CPU uses LAPACK's routine `?gesdd` + (a divide-and-conquer algorithm) instead of `?gesvd` for speed. Analogously, + on GPU, it uses cuSOLVER's routines `gesvdj` and `gesvdjBatched` on CUDA 10.1.243 + and later, and MAGMA's routine `gesdd` on earlier versions of CUDA. + +.. note:: The returned `U` will not be contiguous. The matrix (or batch of matrices) will + be represented as a column-major matrix (i.e. Fortran-contiguous). + +.. warning:: The gradients with respect to `U` and `V` will only be finite when the input does not + have zero nor repeated singular values. + +.. warning:: If the distance between any two singular values is close to zero, the gradients with respect to + `U` and `V` will be numerically unstable, as they depends on + :math:`\frac{1}{\min_{i \neq j} \sigma_i^2 - \sigma_j^2}`. The same happens when the matrix + has small singular values, as these gradients also depend on `S^{-1}`. + +.. warning:: For complex-valued :attr:`input` the singular value decomposition is not unique, + as `U` and `V` may be multiplied by an arbitrary phase factor :math:`e^{i \phi}` on every column. + The same happens when :attr:`input` has repeated singular values, where one may multiply + the columns of the spanning subspace in `U` and `V` by a rotation matrix + and `the resulting vectors will span the same subspace`_. + Different platforms, like NumPy, or inputs on different device types, + may produce different `U` and `V` tensors. + +Args: + input (Tensor): the input tensor of size `(*, m, n)` where `*` is zero or more + batch dimensions consisting of `(m, n)` matrices. + some (bool, optional): controls whether to compute the reduced or full decomposition, and + consequently, the shape of returned `U` and `V`. Default: `True`. + compute_uv (bool, optional): controls whether to compute `U` and `V`. Default: `True`. + +Keyword args: + out (tuple, optional): the output tuple of tensors + +Example:: + + >>> a = torch.randn(5, 3) + >>> a + tensor([[ 0.2364, -0.7752, 0.6372], + [ 1.7201, 0.7394, -0.0504], + [-0.3371, -1.0584, 0.5296], + [ 0.3550, -0.4022, 1.5569], + [ 0.2445, -0.0158, 1.1414]]) + >>> u, s, v = torch.svd(a) + >>> u + tensor([[ 0.4027, 0.0287, 0.5434], + [-0.1946, 0.8833, 0.3679], + [ 0.4296, -0.2890, 0.5261], + [ 0.6604, 0.2717, -0.2618], + [ 0.4234, 0.2481, -0.4733]]) + >>> s + tensor([2.3289, 2.0315, 0.7806]) + >>> v + tensor([[-0.0199, 0.8766, 0.4809], + [-0.5080, 0.4054, -0.7600], + [ 0.8611, 0.2594, -0.4373]]) + >>> torch.dist(a, torch.mm(torch.mm(u, torch.diag(s)), v.t())) + tensor(8.6531e-07) + >>> a_big = torch.randn(7, 5, 3) + >>> u, s, v = torch.svd(a_big) + >>> torch.dist(a_big, torch.matmul(torch.matmul(u, torch.diag_embed(s)), v.mT)) + tensor(2.6503e-06) + +.. _the resulting vectors will span the same subspace: + (https://en.wikipedia.org/wiki/Singular_value_decomposition#Singular_values,_singular_vectors,_and_their_relation_to_the_SVD) +""", +) + + +add_docstr( + torch.t, + r""" +t(input) -> Tensor + +Expects :attr:`input` to be <= 2-D tensor and transposes dimensions 0 +and 1. + +0-D and 1-D tensors are returned as is. When input is a 2-D tensor this +is equivalent to ``transpose(input, 0, 1)``. + +Args: + {input} + +Example:: + + >>> x = torch.randn(()) + >>> x + tensor(0.1995) + >>> torch.t(x) + tensor(0.1995) + >>> x = torch.randn(3) + >>> x + tensor([ 2.4320, -0.4608, 0.7702]) + >>> torch.t(x) + tensor([ 2.4320, -0.4608, 0.7702]) + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 0.4875, 0.9158, -0.5872], + [ 0.3938, -0.6929, 0.6932]]) + >>> torch.t(x) + tensor([[ 0.4875, 0.3938], + [ 0.9158, -0.6929], + [-0.5872, 0.6932]]) + +See also :func:`torch.transpose`. +""".format(**common_args), +) + +add_docstr( + torch.flip, + r""" +flip(input, dims) -> Tensor + +Reverse the order of an n-D tensor along given axis in dims. + +.. note:: + `torch.flip` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.flip`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.flip` is expected to be slower than `np.flip`. + +Args: + {input} + dims (a list or tuple): axis to flip on + +Example:: + + >>> x = torch.arange(8).view(2, 2, 2) + >>> x + tensor([[[ 0, 1], + [ 2, 3]], + + [[ 4, 5], + [ 6, 7]]]) + >>> torch.flip(x, [0, 1]) + tensor([[[ 6, 7], + [ 4, 5]], + + [[ 2, 3], + [ 0, 1]]]) +""".format(**common_args), +) + +add_docstr( + torch.fliplr, + r""" +fliplr(input) -> Tensor + +Flip tensor in the left/right direction, returning a new tensor. + +Flip the entries in each row in the left/right direction. +Columns are preserved, but appear in a different order than before. + +Note: + Requires the tensor to be at least 2-D. + +.. note:: + `torch.fliplr` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.fliplr`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.fliplr` is expected to be slower than `np.fliplr`. + +Args: + input (Tensor): Must be at least 2-dimensional. + +Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.fliplr(x) + tensor([[1, 0], + [3, 2]]) +""".format(**common_args), +) + +add_docstr( + torch.flipud, + r""" +flipud(input) -> Tensor + +Flip tensor in the up/down direction, returning a new tensor. + +Flip the entries in each column in the up/down direction. +Rows are preserved, but appear in a different order than before. + +Note: + Requires the tensor to be at least 1-D. + +.. note:: + `torch.flipud` makes a copy of :attr:`input`'s data. This is different from NumPy's `np.flipud`, + which returns a view in constant time. Since copying a tensor's data is more work than viewing that data, + `torch.flipud` is expected to be slower than `np.flipud`. + +Args: + input (Tensor): Must be at least 1-dimensional. + +Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.flipud(x) + tensor([[2, 3], + [0, 1]]) +""".format(**common_args), +) + +add_docstr( + torch.roll, + r""" +roll(input, shifts, dims=None) -> Tensor + +Roll the tensor :attr:`input` along the given dimension(s). Elements that are +shifted beyond the last position are re-introduced at the first position. If +:attr:`dims` is `None`, the tensor will be flattened before rolling and then +restored to the original shape. + +Args: + {input} + shifts (int or tuple of ints): The number of places by which the elements + of the tensor are shifted. If shifts is a tuple, dims must be a tuple of + the same size, and each dimension will be rolled by the corresponding + value + dims (int or tuple of ints): Axis along which to roll + +Example:: + + >>> x = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]).view(4, 2) + >>> x + tensor([[1, 2], + [3, 4], + [5, 6], + [7, 8]]) + >>> torch.roll(x, 1) + tensor([[8, 1], + [2, 3], + [4, 5], + [6, 7]]) + >>> torch.roll(x, 1, 0) + tensor([[7, 8], + [1, 2], + [3, 4], + [5, 6]]) + >>> torch.roll(x, -1, 0) + tensor([[3, 4], + [5, 6], + [7, 8], + [1, 2]]) + >>> torch.roll(x, shifts=(2, 1), dims=(0, 1)) + tensor([[6, 5], + [8, 7], + [2, 1], + [4, 3]]) +""".format(**common_args), +) + +add_docstr( + torch.rot90, + r""" +rot90(input, k=1, dims=(0, 1)) -> Tensor + +Rotate an n-D tensor by 90 degrees in the plane specified by dims axis. +Rotation direction is from the first towards the second axis if k > 0, and from the second towards the first for k < 0. + +Args: + {input} + k (int): number of times to rotate. Default value is 1 + dims (a list or tuple): axis to rotate. Default value is [0, 1] + +Example:: + + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.rot90(x, 1, [0, 1]) + tensor([[1, 3], + [0, 2]]) + + >>> x = torch.arange(8).view(2, 2, 2) + >>> x + tensor([[[0, 1], + [2, 3]], + + [[4, 5], + [6, 7]]]) + >>> torch.rot90(x, 1, [1, 2]) + tensor([[[1, 3], + [0, 2]], + + [[5, 7], + [4, 6]]]) +""".format(**common_args), +) + +add_docstr( + torch.take, + r""" +take(input, index) -> Tensor + +Returns a new tensor with the elements of :attr:`input` at the given indices. +The input tensor is treated as if it were viewed as a 1-D tensor. The result +takes the same shape as the indices. + +Args: + {input} + index (LongTensor): the indices into tensor + +Example:: + + >>> src = torch.tensor([[4, 3, 5], + ... [6, 7, 8]]) + >>> torch.take(src, torch.tensor([0, 2, 5])) + tensor([ 4, 5, 8]) +""".format(**common_args), +) + +add_docstr( + torch.take_along_dim, + r""" +take_along_dim(input, indices, dim=None, *, out=None) -> Tensor + +Selects values from :attr:`input` at the 1-dimensional indices from :attr:`indices` along the given :attr:`dim`. + +If :attr:`dim` is None, the input array is treated as if it has been flattened to 1d. + +Functions that return indices along a dimension, like :func:`torch.argmax` and :func:`torch.argsort`, +are designed to work with this function. See the examples below. + +.. note:: + This function is similar to NumPy's `take_along_axis`. + See also :func:`torch.gather`. + +Args: + {input} + indices (LongTensor): the indices into :attr:`input`. Must have long dtype. + dim (int, optional): dimension to select along. Default: 0 + +Keyword args: + {out} + +Example:: + + >>> t = torch.tensor([[10, 30, 20], [60, 40, 50]]) + >>> max_idx = torch.argmax(t) + >>> torch.take_along_dim(t, max_idx) + tensor([60]) + >>> sorted_idx = torch.argsort(t, dim=1) + >>> torch.take_along_dim(t, sorted_idx, dim=1) + tensor([[10, 20, 30], + [40, 50, 60]]) +""".format(**common_args), +) + +add_docstr( + torch.tan, + r""" +tan(input, *, out=None) -> Tensor + +Returns a new tensor with the tangent of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \tan(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([-1.2027, -1.7687, 0.4412, -1.3856]) + >>> torch.tan(a) + tensor([-2.5930, 4.9859, 0.4722, -5.3366]) +""".format(**common_args), +) + +add_docstr( + torch.tanh, + r""" +tanh(input, *, out=None) -> Tensor + +Returns a new tensor with the hyperbolic tangent of the elements +of :attr:`input`. + +.. math:: + \text{out}_{i} = \tanh(\text{input}_{i}) +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 0.8986, -0.7279, 1.1745, 0.2611]) + >>> torch.tanh(a) + tensor([ 0.7156, -0.6218, 0.8257, 0.2553]) +""".format(**common_args), +) + +add_docstr( + # torch.softmax doc str. Point this to torch.nn.functional.softmax + torch.softmax, + r""" +softmax(input, dim, *, dtype=None) -> Tensor + +Alias for :func:`torch.nn.functional.softmax`. +""", +) + +add_docstr( + torch.topk, + r""" +topk(input, k, dim=None, largest=True, sorted=True, *, out=None) -> (Tensor, LongTensor) + +Returns the :attr:`k` largest elements of the given :attr:`input` tensor along +a given dimension. + +If :attr:`dim` is not given, the last dimension of the `input` is chosen. + +If :attr:`largest` is ``False`` then the `k` smallest elements are returned. + +A namedtuple of `(values, indices)` is returned with the `values` and +`indices` of the largest `k` elements of each row of the `input` tensor in the +given dimension `dim`. + +The boolean option :attr:`sorted` if ``True``, will make sure that the returned +`k` elements are themselves sorted + +.. note:: + When using `torch.topk`, the indices of tied elements are not guaranteed to be stable + and may vary across different invocations. + +Args: + {input} + k (int): the k in "top-k" + dim (int, optional): the dimension to sort along + largest (bool, optional): controls whether to return largest or + smallest elements + sorted (bool, optional): controls whether to return the elements + in sorted order + +Keyword args: + out (tuple, optional): the output tuple of (Tensor, LongTensor) that can be + optionally given to be used as output buffers + +Example:: + + >>> x = torch.arange(1., 6.) + >>> x + tensor([ 1., 2., 3., 4., 5.]) + >>> torch.topk(x, 3) + torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2])) +""".format(**common_args), +) + +add_docstr( + torch.trace, + r""" +trace(input) -> Tensor + +Returns the sum of the elements of the diagonal of the input 2-D matrix. + +Example:: + + >>> x = torch.arange(1., 10.).view(3, 3) + >>> x + tensor([[ 1., 2., 3.], + [ 4., 5., 6.], + [ 7., 8., 9.]]) + >>> torch.trace(x) + tensor(15.) +""", +) + +add_docstr( + torch.transpose, + r""" +transpose(input, dim0, dim1) -> Tensor + +Returns a tensor that is a transposed version of :attr:`input`. +The given dimensions :attr:`dim0` and :attr:`dim1` are swapped. + +If :attr:`input` is a strided tensor then the resulting :attr:`out` +tensor shares its underlying storage with the :attr:`input` tensor, so +changing the content of one would change the content of the other. + +If :attr:`input` is a :ref:`sparse tensor ` then the +resulting :attr:`out` tensor *does not* share the underlying storage +with the :attr:`input` tensor. + +If :attr:`input` is a :ref:`sparse tensor ` with compressed +layout (SparseCSR, SparseBSR, SparseCSC or SparseBSC) the arguments +:attr:`dim0` and :attr:`dim1` must be both batch dimensions, or must +both be sparse dimensions. The batch dimensions of a sparse tensor are the +dimensions preceding the sparse dimensions. + +.. note:: + Transpositions which interchange the sparse dimensions of a `SparseCSR` + or `SparseCSC` layout tensor will result in the layout changing between + the two options. Transposition of the sparse dimensions of a ` SparseBSR` + or `SparseBSC` layout tensor will likewise generate a result with the + opposite layout. + + +Args: + {input} + dim0 (int): the first dimension to be transposed + dim1 (int): the second dimension to be transposed + +Example:: + + >>> x = torch.randn(2, 3) + >>> x + tensor([[ 1.0028, -0.9893, 0.5809], + [-0.1669, 0.7299, 0.4942]]) + >>> torch.transpose(x, 0, 1) + tensor([[ 1.0028, -0.1669], + [-0.9893, 0.7299], + [ 0.5809, 0.4942]]) + +See also :func:`torch.t`. +""".format(**common_args), +) + +add_docstr( + torch.triangular_solve, + r""" +triangular_solve(b, A, upper=True, transpose=False, unitriangular=False, *, out=None) -> (Tensor, Tensor) + +Solves a system of equations with a square upper or lower triangular invertible matrix :math:`A` +and multiple right-hand sides :math:`b`. + +In symbols, it solves :math:`AX = b` and assumes :math:`A` is square upper-triangular +(or lower-triangular if :attr:`upper`\ `= False`) and does not have zeros on the diagonal. + +`torch.triangular_solve(b, A)` can take in 2D inputs `b, A` or inputs that are +batches of 2D matrices. If the inputs are batches, then returns +batched outputs `X` + +If the diagonal of :attr:`A` contains zeros or elements that are very close to zero and +:attr:`unitriangular`\ `= False` (default) or if the input matrix is badly conditioned, +the result may contain `NaN` s. + +Supports input of float, double, cfloat and cdouble data types. + +.. warning:: + + :func:`torch.triangular_solve` is deprecated in favor of :func:`torch.linalg.solve_triangular` + and will be removed in a future PyTorch release. + :func:`torch.linalg.solve_triangular` has its arguments reversed and does not return a + copy of one of the inputs. + + ``X = torch.triangular_solve(B, A).solution`` should be replaced with + + .. code:: python + + X = torch.linalg.solve_triangular(A, B) + +Args: + b (Tensor): multiple right-hand sides of size :math:`(*, m, k)` where + :math:`*` is zero of more batch dimensions + A (Tensor): the input triangular coefficient matrix of size :math:`(*, m, m)` + where :math:`*` is zero or more batch dimensions + upper (bool, optional): whether :math:`A` is upper or lower triangular. Default: ``True``. + transpose (bool, optional): solves `op(A)X = b` where `op(A) = A^T` if this flag is ``True``, + and `op(A) = A` if it is ``False``. Default: ``False``. + unitriangular (bool, optional): whether :math:`A` is unit triangular. + If True, the diagonal elements of :math:`A` are assumed to be + 1 and not referenced from :math:`A`. Default: ``False``. + +Keyword args: + out ((Tensor, Tensor), optional): tuple of two tensors to write + the output to. Ignored if `None`. Default: `None`. + +Returns: + A namedtuple `(solution, cloned_coefficient)` where `cloned_coefficient` + is a clone of :math:`A` and `solution` is the solution :math:`X` to :math:`AX = b` + (or whatever variant of the system of equations, depending on the keyword arguments.) + +Examples:: + + >>> A = torch.randn(2, 2).triu() + >>> A + tensor([[ 1.1527, -1.0753], + [ 0.0000, 0.7986]]) + >>> b = torch.randn(2, 3) + >>> b + tensor([[-0.0210, 2.3513, -1.5492], + [ 1.5429, 0.7403, -1.0243]]) + >>> torch.triangular_solve(b, A) + torch.return_types.triangular_solve( + solution=tensor([[ 1.7841, 2.9046, -2.5405], + [ 1.9320, 0.9270, -1.2826]]), + cloned_coefficient=tensor([[ 1.1527, -1.0753], + [ 0.0000, 0.7986]])) +""", +) + +add_docstr( + torch.tril, + r""" +tril(input, diagonal=0, *, out=None) -> Tensor + +Returns the lower triangular part of the matrix (2-D tensor) or batch of matrices +:attr:`input`, the other elements of the result tensor :attr:`out` are set to 0. + +The lower triangular part of the matrix is defined as the elements on and +below the diagonal. + +The argument :attr:`diagonal` controls which diagonal to consider. If +:attr:`diagonal` = 0, all elements on and below the main diagonal are +retained. A positive value includes just as many diagonals above the main +diagonal, and similarly a negative value excludes just as many diagonals below +the main diagonal. The main diagonal are the set of indices +:math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` where +:math:`d_{1}, d_{2}` are the dimensions of the matrix. +""" + + r""" +Args: + {input} + diagonal (int, optional): the diagonal to consider + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[-1.0813, -0.8619, 0.7105], + [ 0.0935, 0.1380, 2.2112], + [-0.3409, -0.9828, 0.0289]]) + >>> torch.tril(a) + tensor([[-1.0813, 0.0000, 0.0000], + [ 0.0935, 0.1380, 0.0000], + [-0.3409, -0.9828, 0.0289]]) + + >>> b = torch.randn(4, 6) + >>> b + tensor([[ 1.2219, 0.5653, -0.2521, -0.2345, 1.2544, 0.3461], + [ 0.4785, -0.4477, 0.6049, 0.6368, 0.8775, 0.7145], + [ 1.1502, 3.2716, -1.1243, -0.5413, 0.3615, 0.6864], + [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024, 0.0978]]) + >>> torch.tril(b, diagonal=1) + tensor([[ 1.2219, 0.5653, 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.4785, -0.4477, 0.6049, 0.0000, 0.0000, 0.0000], + [ 1.1502, 3.2716, -1.1243, -0.5413, 0.0000, 0.0000], + [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024, 0.0000]]) + >>> torch.tril(b, diagonal=-1) + tensor([[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [ 0.4785, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [ 1.1502, 3.2716, 0.0000, 0.0000, 0.0000, 0.0000], + [-0.0614, -0.7344, -1.3164, 0.0000, 0.0000, 0.0000]]) +""".format(**common_args), +) + +# docstr is split in two parts to avoid format mis-captureing :math: braces '{}' +# as common args. +add_docstr( + torch.tril_indices, + r""" +tril_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) -> Tensor + +Returns the indices of the lower triangular part of a :attr:`row`-by- +:attr:`col` matrix in a 2-by-N Tensor, where the first row contains row +coordinates of all indices and the second row contains column coordinates. +Indices are ordered based on rows and then columns. + +The lower triangular part of the matrix is defined as the elements on and +below the diagonal. + +The argument :attr:`offset` controls which diagonal to consider. If +:attr:`offset` = 0, all elements on and below the main diagonal are +retained. A positive value includes just as many diagonals above the main +diagonal, and similarly a negative value excludes just as many diagonals below +the main diagonal. The main diagonal are the set of indices +:math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` +where :math:`d_{1}, d_{2}` are the dimensions of the matrix. + +.. note:: + When running on CUDA, ``row * col`` must be less than :math:`2^{59}` to + prevent overflow during calculation. +""" + + r""" +Args: + row (``int``): number of rows in the 2-D matrix. + col (``int``): number of columns in the 2-D matrix. + offset (``int``): diagonal offset from the main diagonal. + Default: if not provided, 0. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor, + only support ``torch.int``, ``torch.long``. Default: if ``None``, ``torch.long``. + {device} + layout (:class:`torch.layout`, optional): currently only support ``torch.strided``. + +Example:: + + >>> a = torch.tril_indices(3, 3) + >>> a + tensor([[0, 1, 1, 2, 2, 2], + [0, 0, 1, 0, 1, 2]]) + + >>> a = torch.tril_indices(4, 3, -1) + >>> a + tensor([[1, 2, 2, 3, 3, 3], + [0, 0, 1, 0, 1, 2]]) + + >>> a = torch.tril_indices(4, 3, 1) + >>> a + tensor([[0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], + [0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.triu, + r""" +triu(input, diagonal=0, *, out=None) -> Tensor + +Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices +:attr:`input`, the other elements of the result tensor :attr:`out` are set to 0. + +The upper triangular part of the matrix is defined as the elements on and +above the diagonal. + +The argument :attr:`diagonal` controls which diagonal to consider. If +:attr:`diagonal` = 0, all elements on and above the main diagonal are +retained. A positive value excludes just as many diagonals above the main +diagonal, and similarly a negative value includes just as many diagonals below +the main diagonal. The main diagonal are the set of indices +:math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` where +:math:`d_{1}, d_{2}` are the dimensions of the matrix. +""" + + r""" +Args: + {input} + diagonal (int, optional): the diagonal to consider + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(3, 3) + >>> a + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.2072, -1.0680, 0.6602], + [ 0.3480, -0.5211, -0.4573]]) + >>> torch.triu(a) + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.0000, -1.0680, 0.6602], + [ 0.0000, 0.0000, -0.4573]]) + >>> torch.triu(a, diagonal=1) + tensor([[ 0.0000, 0.5207, 2.0049], + [ 0.0000, 0.0000, 0.6602], + [ 0.0000, 0.0000, 0.0000]]) + >>> torch.triu(a, diagonal=-1) + tensor([[ 0.2309, 0.5207, 2.0049], + [ 0.2072, -1.0680, 0.6602], + [ 0.0000, -0.5211, -0.4573]]) + + >>> b = torch.randn(4, 6) + >>> b + tensor([[ 0.5876, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [-0.2447, 0.9556, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.4333, 0.3146, 0.6576, -1.0432, 0.9348, -0.4410], + [-0.9888, 1.0679, -1.3337, -1.6556, 0.4798, 0.2830]]) + >>> torch.triu(b, diagonal=1) + tensor([[ 0.0000, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [ 0.0000, 0.0000, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.0000, 0.0000, 0.0000, -1.0432, 0.9348, -0.4410], + [ 0.0000, 0.0000, 0.0000, 0.0000, 0.4798, 0.2830]]) + >>> torch.triu(b, diagonal=-1) + tensor([[ 0.5876, -0.0794, -1.8373, 0.6654, 0.2604, 1.5235], + [-0.2447, 0.9556, -1.2919, 1.3378, -0.1768, -1.0857], + [ 0.0000, 0.3146, 0.6576, -1.0432, 0.9348, -0.4410], + [ 0.0000, 0.0000, -1.3337, -1.6556, 0.4798, 0.2830]]) +""".format(**common_args), +) + +# docstr is split in two parts to avoid format mis-capturing :math: braces '{}' +# as common args. +add_docstr( + torch.triu_indices, + r""" +triu_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) -> Tensor + +Returns the indices of the upper triangular part of a :attr:`row` by +:attr:`col` matrix in a 2-by-N Tensor, where the first row contains row +coordinates of all indices and the second row contains column coordinates. +Indices are ordered based on rows and then columns. + +The upper triangular part of the matrix is defined as the elements on and +above the diagonal. + +The argument :attr:`offset` controls which diagonal to consider. If +:attr:`offset` = 0, all elements on and above the main diagonal are +retained. A positive value excludes just as many diagonals above the main +diagonal, and similarly a negative value includes just as many diagonals below +the main diagonal. The main diagonal are the set of indices +:math:`\lbrace (i, i) \rbrace` for :math:`i \in [0, \min\{d_{1}, d_{2}\} - 1]` +where :math:`d_{1}, d_{2}` are the dimensions of the matrix. + +.. note:: + When running on CUDA, ``row * col`` must be less than :math:`2^{59}` to + prevent overflow during calculation. +""" + + r""" +Args: + row (``int``): number of rows in the 2-D matrix. + col (``int``): number of columns in the 2-D matrix. + offset (``int``): diagonal offset from the main diagonal. + Default: if not provided, 0. + +Keyword args: + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor, + only support ``torch.int``, ``torch.long``. Default: if ``None``, ``torch.long``. + {device} + layout (:class:`torch.layout`, optional): currently only support ``torch.strided``. + +Example:: + + >>> a = torch.triu_indices(3, 3) + >>> a + tensor([[0, 0, 0, 1, 1, 2], + [0, 1, 2, 1, 2, 2]]) + + >>> a = torch.triu_indices(4, 3, -1) + >>> a + tensor([[0, 0, 0, 1, 1, 1, 2, 2, 3], + [0, 1, 2, 0, 1, 2, 1, 2, 2]]) + + >>> a = torch.triu_indices(4, 3, 1) + >>> a + tensor([[0, 0, 1], + [1, 2, 2]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.true_divide, + r""" +true_divide(dividend, divisor, *, out) -> Tensor + +Alias for :func:`torch.div` with ``rounding_mode=None``. +""", +) + +add_docstr( + torch.trunc, + r""" +trunc(input, *, out=None) -> Tensor + +Returns a new tensor with the truncated integer values of +the elements of :attr:`input`. + +For integer inputs, follows the array-api convention of returning a +copy of the input tensor. + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.randn(4) + >>> a + tensor([ 3.4742, 0.5466, -0.8008, -0.9079]) + >>> torch.trunc(a) + tensor([ 3., 0., -0., -0.]) +""".format(**common_args), +) + +add_docstr( + torch.fake_quantize_per_tensor_affine, + r""" +fake_quantize_per_tensor_affine(input, scale, zero_point, quant_min, quant_max) -> Tensor + +Returns a new tensor with the data in :attr:`input` fake quantized using :attr:`scale`, +:attr:`zero_point`, :attr:`quant_min` and :attr:`quant_max`. + +.. math:: + \text{output} = ( + min( + \text{quant\_max}, + max( + \text{quant\_min}, + \text{std::nearby\_int}(\text{input} / \text{scale}) + \text{zero\_point} + ) + ) - \text{zero\_point} + ) \times \text{scale} + +Args: + input (Tensor): the input value(s), ``torch.float32`` tensor + scale (double scalar or ``float32`` Tensor): quantization scale + zero_point (int64 scalar or ``int32`` Tensor): quantization zero_point + quant_min (int64): lower bound of the quantized domain + quant_max (int64): upper bound of the quantized domain + +Returns: + Tensor: A newly fake_quantized ``torch.float32`` tensor + +Example:: + + >>> x = torch.randn(4) + >>> x + tensor([ 0.0552, 0.9730, 0.3973, -1.0780]) + >>> torch.fake_quantize_per_tensor_affine(x, 0.1, 0, 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) + >>> torch.fake_quantize_per_tensor_affine(x, torch.tensor(0.1), torch.tensor(0), 0, 255) + tensor([0.1000, 1.0000, 0.4000, 0.0000]) +""", +) + +add_docstr( + torch.fake_quantize_per_channel_affine, + r""" +fake_quantize_per_channel_affine(input, scale, zero_point, axis, quant_min, quant_max) -> Tensor + +Returns a new tensor with the data in :attr:`input` fake quantized per channel using :attr:`scale`, +:attr:`zero_point`, :attr:`quant_min` and :attr:`quant_max`, across the channel specified by :attr:`axis`. + +.. math:: + \text{output} = ( + min( + \text{quant\_max}, + max( + \text{quant\_min}, + \text{std::nearby\_int}(\text{input} / \text{scale}) + \text{zero\_point} + ) + ) - \text{zero\_point} + ) \times \text{scale} + +Args: + input (Tensor): the input value(s), in ``torch.float32`` + scale (Tensor): quantization scale, per channel in ``torch.float32`` + zero_point (Tensor): quantization zero_point, per channel in ``torch.int32`` or ``torch.half`` or ``torch.float32`` + axis (int32): channel axis + quant_min (int64): lower bound of the quantized domain + quant_max (int64): upper bound of the quantized domain + +Returns: + Tensor: A newly fake_quantized per channel ``torch.float32`` tensor + +Example:: + + >>> x = torch.randn(2, 2, 2) + >>> x + tensor([[[-0.2525, -0.0466], + [ 0.3491, -0.2168]], + + [[-0.5906, 1.6258], + [ 0.6444, -0.0542]]]) + >>> scales = (torch.randn(2) + 1) * 0.05 + >>> scales + tensor([0.0475, 0.0486]) + >>> zero_points = torch.zeros(2).to(torch.int32) + >>> zero_points + tensor([0, 0]) + >>> torch.fake_quantize_per_channel_affine(x, scales, zero_points, 1, 0, 255) + tensor([[[0.0000, 0.0000], + [0.3405, 0.0000]], + + [[0.0000, 1.6134], + [0.6323, 0.0000]]]) +""", +) + +add_docstr( + torch.fix, + r""" +fix(input, *, out=None) -> Tensor + +Alias for :func:`torch.trunc` +""", +) + +add_docstr( + torch.unsqueeze, + r""" +unsqueeze(input, dim) -> Tensor + +Returns a new tensor with a dimension of size one inserted at the +specified position. + +The returned tensor shares the same underlying data with this tensor. + +A :attr:`dim` value within the range ``[-input.dim() - 1, input.dim() + 1)`` +can be used. Negative :attr:`dim` will correspond to :meth:`unsqueeze` +applied at :attr:`dim` = ``dim + input.dim() + 1``. + +Args: + {input} + dim (int): the index at which to insert the singleton dimension + +Example:: + + >>> x = torch.tensor([1, 2, 3, 4]) + >>> torch.unsqueeze(x, 0) + tensor([[ 1, 2, 3, 4]]) + >>> torch.unsqueeze(x, 1) + tensor([[ 1], + [ 2], + [ 3], + [ 4]]) +""".format(**common_args), +) + +add_docstr( + torch.var, + r""" +var(input, dim=None, *, correction=1, keepdim=False, out=None) -> Tensor + +Calculates the variance over the dimensions specified by :attr:`dim`. :attr:`dim` +can be a single dimension, list of dimensions, or ``None`` to reduce over all +dimensions. + +The variance (:math:`\sigma^2`) is calculated as + +.. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + +where :math:`x` is the sample set of elements, :math:`\bar{x}` is the +sample mean, :math:`N` is the number of samples and :math:`\delta N` is +the :attr:`correction`. +""" + + r""" + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + +Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + {opt_keepdim} + {out} + +Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var(a, dim=1, keepdim=True) + tensor([[1.0631], + [0.5590], + [1.4893], + [0.8258]]) + +.. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + +""".format(**multi_dim_common), +) + +add_docstr( + torch.var_mean, + r""" +var_mean(input, dim=None, *, correction=1, keepdim=False, out=None) -> (Tensor, Tensor) + +Calculates the variance and mean over the dimensions specified by :attr:`dim`. +:attr:`dim` can be a single dimension, list of dimensions, or ``None`` to +reduce over all dimensions. + +The variance (:math:`\sigma^2`) is calculated as + +.. math:: \sigma^2 = \frac{1}{\max(0,~N - \delta N)}\sum_{i=0}^{N-1}(x_i-\bar{x})^2 + +where :math:`x` is the sample set of elements, :math:`\bar{x}` is the +sample mean, :math:`N` is the number of samples and :math:`\delta N` is +the :attr:`correction`. +""" + + r""" + +{keepdim_details} + +Args: + {input} + {opt_dim_all_reduce} + +Keyword args: + correction (int): difference between the sample size and sample degrees of freedom. + Defaults to `Bessel's correction`_, ``correction=1``. + + .. versionchanged:: 2.0 + Previously this argument was called ``unbiased`` and was a boolean + with ``True`` corresponding to ``correction=1`` and ``False`` being + ``correction=0``. + {opt_keepdim} + {out} + +Returns: + A tuple (var, mean) containing the variance and mean. + +Example: + + >>> a = torch.tensor( + ... [[ 0.2035, 1.2959, 1.8101, -0.4644], + ... [ 1.5027, -0.3270, 0.5905, 0.6538], + ... [-1.5745, 1.3330, -0.5596, -0.6548], + ... [ 0.1264, -0.5080, 1.6420, 0.1992]] + ... ) # fmt: skip + >>> torch.var_mean(a, dim=0, keepdim=True) + (tensor([[1.5926, 1.0056, 1.2005, 0.3646]]), + tensor([[ 0.0645, 0.4485, 0.8707, -0.0665]])) + +.. _Bessel's correction: https://en.wikipedia.org/wiki/Bessel%27s_correction + +""".format(**multi_dim_common), +) + +add_docstr( + torch.zeros, + r""" +zeros(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Returns a tensor filled with the scalar value `0`, with the shape defined +by the variable argument :attr:`size`. + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + +Keyword args: + {out} + {dtype} + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.zeros(2, 3) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) + + >>> torch.zeros(5) + tensor([ 0., 0., 0., 0., 0.]) +""".format(**factory_common_args), +) + +add_docstr( + torch.zeros_like, + r""" +zeros_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + +Returns a tensor filled with the scalar value `0`, with the same size as +:attr:`input`. ``torch.zeros_like(input)`` is equivalent to +``torch.zeros(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + +.. warning:: + As of 0.4, this function does not support an :attr:`out` keyword. As an alternative, + the old ``torch.zeros_like(input, out=output)`` is equivalent to + ``torch.zeros(input.size(), out=output)``. + +Args: + {input} + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +Example:: + + >>> input = torch.empty(2, 3) + >>> torch.zeros_like(input) + tensor([[ 0., 0., 0.], + [ 0., 0., 0.]]) +""".format(**factory_like_common_args), +) + +add_docstr( + torch.empty, + """ +empty(*size, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False, pin_memory=False, \ +memory_format=torch.contiguous_format) -> Tensor + +Returns a tensor filled with uninitialized data. The shape of the tensor is +defined by the variable argument :attr:`size`. + +.. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + +Args: + size (int...): a sequence of integers defining the shape of the output tensor. + Can be a variable number of arguments or a collection like a list or tuple. + +Keyword args: + {out} + {dtype} + {layout} + {device} + {requires_grad} + {pin_memory} + {memory_format} + +Example:: + + >>> torch.empty((2,3), dtype=torch.int64) + tensor([[ 9.4064e+13, 2.8000e+01, 9.3493e+13], + [ 7.5751e+18, 7.1428e+18, 7.5955e+18]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.empty_like, + r""" +empty_like(input, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) -> Tensor + +Returns an uninitialized tensor with the same size as :attr:`input`. +``torch.empty_like(input)`` is equivalent to +``torch.empty(input.size(), dtype=input.dtype, layout=input.layout, device=input.device)``. + +.. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + + When ``torch.preserve_format`` is used: + If the input tensor is dense (i.e., non-overlapping strided), + its memory format (including strides) is retained. + Otherwise (e.g., a non-dense view like a stepped slice), + the output is converted to the dense format. + +Args: + {input} + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +Example:: + + >>> a=torch.empty((2,3), dtype=torch.int32, device = 'cuda') + >>> torch.empty_like(a) + tensor([[0, 0, 0], + [0, 0, 0]], device='cuda:0', dtype=torch.int32) +""".format(**factory_like_common_args), +) + +add_docstr( + torch.empty_strided, + r""" +empty_strided(size, stride, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + +Creates a tensor with the specified :attr:`size` and :attr:`stride` and filled with undefined data. + +.. warning:: + If the constructed tensor is "overlapped" (with multiple indices referring to the same element + in memory) its behavior is undefined. + +.. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + +Args: + size (tuple of int): the shape of the output tensor + stride (tuple of int): the strides of the output tensor + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {pin_memory} + +Example:: + + >>> a = torch.empty_strided((2, 3), (1, 2)) + >>> a + tensor([[8.9683e-44, 4.4842e-44, 5.1239e+07], + [0.0000e+00, 0.0000e+00, 3.0705e-41]]) + >>> a.stride() + (1, 2) + >>> a.size() + torch.Size([2, 3]) +""".format(**factory_common_args), +) + +add_docstr( + torch.empty_permuted, + r""" +empty_permuted(size, physical_layout, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) -> Tensor + +Creates an uninitialized, non-overlapping and dense tensor with the +specified :attr:`size`, with :attr:`physical_layout` specifying how the +dimensions are physically laid out in memory (each logical dimension is listed +from outermost to innermost). :attr:`physical_layout` is a generalization +of NCHW/NHWC notation: if each dimension is assigned a number according to +what order they occur in size (N=0, C=1, H=2, W=3), then NCHW is ``(0, 1, 2, 3)`` +while NHWC is ``(0, 2, 3, 1)``. Equivalently, the strides of the output +tensor ``t`` are such that ``t.stride(physical_layout[i]) == contiguous_strides[i]`` +(notably, this function is *not* equivalent to ``torch.empty(size).permute(physical_layout)``). + +Unlike :func:`torch.empty_strided`, this is guaranteed to produce a dense +tensor with no overlaps. If possible, prefer using this function over +:func:`torch.empty_strided` or manual use of :func:`torch.as_strided`. + +.. note:: + If :func:`torch.use_deterministic_algorithms()` and + :attr:`torch.utils.deterministic.fill_uninitialized_memory` are both set to + ``True``, the output tensor is initialized to prevent any possible + nondeterministic behavior from using the data as an input to an operation. + Floating point and complex tensors are filled with NaN, and integer tensors + are filled with the maximum value. + +Args: + size (tuple of int): the shape of the output tensor + physical_layout (tuple of int): the ordering of dimensions physically in memory + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {pin_memory} + +Examples: + + >>> torch.empty((2, 3, 5, 7)).stride() + (105, 35, 7, 1) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 1, 2, 3)).stride() + (105, 35, 7, 1) + >>> torch.empty((2, 3, 5, 7), memory_format=torch.channels_last).stride() + (105, 1, 21, 3) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 2, 3, 1)).stride() + (105, 1, 21, 3) + >>> torch.empty_permuted((2, 3, 5, 7), (0, 2, 3, 1)).dim_order() + (0, 2, 3, 1) +""".format(**factory_common_args), +) + +add_docstr( + torch.full, + r""" +full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor + +Creates a tensor of size :attr:`size` filled with :attr:`fill_value`. The +tensor's dtype is inferred from :attr:`fill_value`. + +Args: + size (int...): a list, tuple, or :class:`torch.Size` of integers defining the + shape of the output tensor. + fill_value (Scalar): the value to fill the output tensor with. + +Keyword args: + {out} + {dtype} + {layout} + {device} + {requires_grad} + +Example:: + + >>> torch.full((2, 3), 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) +""".format(**factory_common_args), +) + +add_docstr( + torch.full_like, + """ +full_like(input, fill_value, \\*, dtype=None, layout=torch.strided, device=None, requires_grad=False, \ +memory_format=torch.preserve_format) -> Tensor + +Returns a tensor with the same size as :attr:`input` filled with :attr:`fill_value`. +``torch.full_like(input, fill_value)`` is equivalent to +``torch.full(input.size(), fill_value, dtype=input.dtype, layout=input.layout, device=input.device)``. + +Args: + {input} + fill_value: the number to fill the output tensor with. + +Keyword args: + {dtype} + {layout} + {device} + {requires_grad} + {memory_format} + +Example:: + + >>> x = torch.ones(2, 3) + >>> torch.full_like(x, 3.141592) + tensor([[ 3.1416, 3.1416, 3.1416], + [ 3.1416, 3.1416, 3.1416]]) + >>> torch.full_like(x, 7) + tensor([[7., 7., 7.], + [7., 7., 7.]]) + >>> torch.full_like(x, 0.5, dtype=torch.int32) + tensor([[0, 0, 0], + [0, 0, 0]], dtype=torch.int32) + >>> y = torch.randn(3, 4, dtype=torch.float64) + >>> torch.full_like(y, -1.0) + tensor([[-1., -1., -1., -1.], + [-1., -1., -1., -1.], + [-1., -1., -1., -1.]], dtype=torch.float64) +""".format(**factory_like_common_args), +) + +add_docstr( + torch.det, + r""" +det(input) -> Tensor + +Alias for :func:`torch.linalg.det` +""", +) + +add_docstr( + torch.where, + r""" +where(condition, input, other, *, out=None) -> Tensor + +Return a tensor of elements selected from either :attr:`input` or :attr:`other`, depending on :attr:`condition`. + +The operation is defined as: + +.. math:: + \text{out}_i = \begin{cases} + \text{input}_i & \text{if } \text{condition}_i \\ + \text{other}_i & \text{otherwise} \\ + \end{cases} +""" + + r""" +.. note:: + The tensors :attr:`condition`, :attr:`input`, :attr:`other` must be :ref:`broadcastable `. + +Arguments: + condition (BoolTensor): When True (nonzero), yield input, otherwise yield other + input (Tensor or Scalar): value (if :attr:`input` is a scalar) or values selected at indices + where :attr:`condition` is ``True`` + other (Tensor or Scalar): value (if :attr:`other` is a scalar) or values selected at indices + where :attr:`condition` is ``False`` + +Keyword args: + {out} + +Returns: + Tensor: A tensor of shape equal to the broadcasted shape of :attr:`condition`, :attr:`input`, :attr:`other` + +Example:: + + >>> x = torch.randn(3, 2) + >>> y = torch.ones(3, 2) + >>> x + tensor([[-0.4620, 0.3139], + [ 0.3898, -0.7197], + [ 0.0478, -0.1657]]) + >>> torch.where(x > 0, 1.0, 0.0) + tensor([[0., 1.], + [1., 0.], + [1., 0.]]) + >>> torch.where(x > 0, x, y) + tensor([[ 1.0000, 0.3139], + [ 0.3898, 1.0000], + [ 0.0478, 1.0000]]) + >>> x = torch.randn(2, 2, dtype=torch.double) + >>> x + tensor([[ 1.0779, 0.0383], + [-0.8785, -1.1089]], dtype=torch.float64) + >>> torch.where(x > 0, x, 0.) + tensor([[1.0779, 0.0383], + [0.0000, 0.0000]], dtype=torch.float64) + +.. function:: where(condition) -> tuple of LongTensor + :noindex: + +``torch.where(condition)`` is identical to +``torch.nonzero(condition, as_tuple=True)``. + +.. note:: + See also :func:`torch.nonzero`. +""".format(**common_args), +) + +add_docstr( + torch.logdet, + r""" +logdet(input) -> Tensor + +Calculates log determinant of a square matrix or batches of square matrices. + +It returns ``-inf`` if the input has a determinant of zero, and ``NaN`` if it has +a negative determinant. + +.. note:: + Backward through :meth:`logdet` internally uses SVD results when :attr:`input` + is not invertible. In this case, double backward through :meth:`logdet` will + be unstable in when :attr:`input` doesn't have distinct singular values. See + :func:`torch.linalg.svd` for details. + +.. seealso:: + + :func:`torch.linalg.slogdet` computes the sign (resp. angle) and natural logarithm of the + absolute value of the determinant of real-valued (resp. complex) square matrices. + +Arguments: + input (Tensor): the input tensor of size ``(*, n, n)`` where ``*`` is zero or more + batch dimensions. + +Example:: + + >>> A = torch.randn(3, 3) + >>> torch.det(A) + tensor(0.2611) + >>> torch.logdet(A) + tensor(-1.3430) + >>> A + tensor([[[ 0.9254, -0.6213], + [-0.5787, 1.6843]], + + [[ 0.3242, -0.9665], + [ 0.4539, -0.0887]], + + [[ 1.1336, -0.4025], + [-0.7089, 0.9032]]]) + >>> A.det() + tensor([1.1990, 0.4099, 0.7386]) + >>> A.det().log() + tensor([ 0.1815, -0.8917, -0.3031]) +""", +) + +add_docstr( + torch.slogdet, + r""" +slogdet(input) -> (Tensor, Tensor) + +Alias for :func:`torch.linalg.slogdet` +""", +) + +add_docstr( + torch.pinverse, + r""" +pinverse(input, rcond=1e-15) -> Tensor + +Alias for :func:`torch.linalg.pinv` +""", +) + +add_docstr( + torch.hann_window, + """ +hann_window(window_length, periodic=True, *, dtype=None, \ +layout=torch.strided, device=None, requires_grad=False) -> Tensor +""" + + r""" +Hann window function. + +.. math:: + w[n] = \frac{1}{2}\ \left[1 - \cos \left( \frac{2 \pi n}{N - 1} \right)\right] = + \sin^2 \left( \frac{\pi n}{N - 1} \right), + +where :math:`N` is the full window size. + +The input :attr:`window_length` is a positive integer controlling the +returned window size. :attr:`periodic` flag determines whether the returned +window trims off the last duplicate value from the symmetric window and is +ready to be used as a periodic window with functions like +:meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in +above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have +``torch.hann_window(L, periodic=True)`` equal to +``torch.hann_window(L + 1, periodic=False)[:-1])``. + +.. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. +""" + + r""" +Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window + +""".format(**factory_common_args), +) + + +add_docstr( + torch.hamming_window, + """ +hamming_window(window_length, *, dtype=None, layout=None, device=None, pin_memory=False, \ +requires_grad=False) -> Tensor +""" + + r""" +Hamming window function. + +.. math:: + w[n] = \alpha - \beta\ \cos \left( \frac{2 \pi n}{N - 1} \right), + +where :math:`N` is the full window size. + +The input :attr:`window_length` is a positive integer controlling the +returned window size. :attr:`periodic` flag determines whether the returned +window trims off the last duplicate value from the symmetric window and is +ready to be used as a periodic window with functions like +:meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in +above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have +``torch.hamming_window(L, periodic=True)`` equal to +``torch.hamming_window(L + 1, periodic=False)[:-1])``. + +.. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. + +.. note:: + This is a generalized version of :meth:`torch.hann_window`. +""" + + r""" +Arguments: + window_length (int): the size of returned window + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {pin_memory} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window. + +.. function:: hamming_window(window_length, periodic, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + +Hamming window function with periodic specified. + +Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {pin_memory} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window. + +.. function:: hamming_window(window_length, periodic, float alpha, *, dtype=None, layout=None, device=None, \ + pin_memory=False, requires_grad=False) -> Tensor + :noindex: + +Hamming window function with periodic and alpha specified. + +Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {pin_memory} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window. + +.. function:: hamming_window(window_length, periodic, float alpha, float beta, *, dtype=None, layout=None, \ + device=None, pin_memory=False, requires_grad=False) -> Tensor + :noindex: + +Hamming window function with periodic, alpha and beta specified. + +Arguments: + window_length (int): the size of returned window + periodic (bool): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + alpha (float): The coefficient :math:`\alpha` in the equation above + beta (float): The coefficient :math:`\beta` in the equation above + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {pin_memory} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window. + +""".format(**factory_common_args), +) + + +add_docstr( + torch.bartlett_window, + """ +bartlett_window(window_length, periodic=True, *, dtype=None, \ +layout=torch.strided, device=None, requires_grad=False) -> Tensor +""" + + r""" +Bartlett window function. + +.. math:: + w[n] = 1 - \left| \frac{2n}{N-1} - 1 \right| = \begin{cases} + \frac{2n}{N - 1} & \text{if } 0 \leq n \leq \frac{N - 1}{2} \\ + 2 - \frac{2n}{N - 1} & \text{if } \frac{N - 1}{2} < n < N \\ + \end{cases}, + +where :math:`N` is the full window size. + +The input :attr:`window_length` is a positive integer controlling the +returned window size. :attr:`periodic` flag determines whether the returned +window trims off the last duplicate value from the symmetric window and is +ready to be used as a periodic window with functions like +:meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in +above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have +``torch.bartlett_window(L, periodic=True)`` equal to +``torch.bartlett_window(L + 1, periodic=False)[:-1])``. + +.. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. +""" + + r""" +Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window + +""".format(**factory_common_args), +) + + +add_docstr( + torch.blackman_window, + """ +blackman_window(window_length, periodic=True, *, dtype=None, \ +layout=torch.strided, device=None, requires_grad=False) -> Tensor +""" + + r""" +Blackman window function. + +.. math:: + w[n] = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{N - 1} \right) + 0.08 \cos \left( \frac{4 \pi n}{N - 1} \right) + +where :math:`N` is the full window size. + +The input :attr:`window_length` is a positive integer controlling the +returned window size. :attr:`periodic` flag determines whether the returned +window trims off the last duplicate value from the symmetric window and is +ready to be used as a periodic window with functions like +:meth:`torch.stft`. Therefore, if :attr:`periodic` is true, the :math:`N` in +above formula is in fact :math:`\text{window\_length} + 1`. Also, we always have +``torch.blackman_window(L, periodic=True)`` equal to +``torch.blackman_window(L + 1, periodic=False)[:-1]``. + +.. note:: + If :attr:`window_length` :math:`=1`, the returned window contains a single value 1. +""" + + r""" +Arguments: + window_length (int): the size of returned window + periodic (bool, optional): If True, returns a window to be used as periodic + function. If False, return a symmetric window. + +Keyword args: + {dtype} Only floating point types are supported. + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {requires_grad} + +Returns: + Tensor: A 1-D tensor of size :math:`(\text{{window\_length}},)` containing the window + +""".format(**factory_common_args), +) + + +add_docstr( + torch.kaiser_window, + """ +kaiser_window(window_length, periodic=True, beta=12.0, *, dtype=None, \ +layout=torch.strided, device=None, requires_grad=False) -> Tensor +""" + + r""" +Computes the Kaiser window with window length :attr:`window_length` and shape parameter :attr:`beta`. + +Let I_0 be the zeroth order modified Bessel function of the first kind (see :func:`torch.i0`) and +``N = L - 1`` if :attr:`periodic` is False and ``L`` if :attr:`periodic` is True, +where ``L`` is the :attr:`window_length`. This function computes: + +.. math:: + out_i = I_0 \left( \beta \sqrt{1 - \left( {\frac{i - N/2}{N/2}} \right) ^2 } \right) / I_0( \beta ) + +Calling ``torch.kaiser_window(L, B, periodic=True)`` is equivalent to calling +``torch.kaiser_window(L + 1, B, periodic=False)[:-1])``. +The :attr:`periodic` argument is intended as a helpful shorthand +to produce a periodic window as input to functions like :func:`torch.stft`. + +.. note:: + If :attr:`window_length` is one, then the returned window is a single element tensor containing a one. + +""" + + r""" +Args: + window_length (int): length of the window. + periodic (bool, optional): If True, returns a periodic window suitable for use in spectral analysis. + If False, returns a symmetric window suitable for use in filter design. + beta (float, optional): shape parameter for the window. + +Keyword args: + {dtype} + layout (:class:`torch.layout`, optional): the desired layout of returned window tensor. Only + ``torch.strided`` (dense layout) is supported. + {device} + {requires_grad} + +""".format(**factory_common_args), +) + + +add_docstr( + torch.vander, + """ +vander(x, N=None, increasing=False) -> Tensor +""" + + r""" +Generates a Vandermonde matrix. + +The columns of the output matrix are elementwise powers of the input vector :math:`x^{{(N-1)}}, x^{{(N-2)}}, ..., x^0`. +If increasing is True, the order of the columns is reversed :math:`x^0, x^1, ..., x^{{(N-1)}}`. Such a +matrix with a geometric progression in each row is named for Alexandre-Theophile Vandermonde. + +Arguments: + x (Tensor): 1-D input tensor. + N (int, optional): Number of columns in the output. If N is not specified, + a square array is returned :math:`(N = len(x))`. + increasing (bool, optional): Order of the powers of the columns. If True, + the powers increase from left to right, if False (the default) they are reversed. + +Returns: + Tensor: Vandermonde matrix. If increasing is False, the first column is :math:`x^{{(N-1)}}`, + the second :math:`x^{{(N-2)}}` and so forth. If increasing is True, the columns + are :math:`x^0, x^1, ..., x^{{(N-1)}}`. + +Example:: + + >>> x = torch.tensor([1, 2, 3, 5]) + >>> torch.vander(x) + tensor([[ 1, 1, 1, 1], + [ 8, 4, 2, 1], + [ 27, 9, 3, 1], + [125, 25, 5, 1]]) + >>> torch.vander(x, N=3) + tensor([[ 1, 1, 1], + [ 4, 2, 1], + [ 9, 3, 1], + [25, 5, 1]]) + >>> torch.vander(x, N=3, increasing=True) + tensor([[ 1, 1, 1], + [ 1, 2, 4], + [ 1, 3, 9], + [ 1, 5, 25]]) + +""".format(**factory_common_args), +) + + +add_docstr( + torch.unbind, + r""" +unbind(input, dim=0) -> seq + +Removes a tensor dimension. + +Returns a tuple of all slices along a given dimension, already without it. + +Arguments: + input (Tensor): the tensor to unbind + dim (int): dimension to remove + +Example:: + + >>> torch.unbind(torch.tensor([[1, 2, 3], + >>> [4, 5, 6], + >>> [7, 8, 9]])) + (tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9])) +""", +) + + +add_docstr( + torch.combinations, + r""" +combinations(input: Tensor, r: int = 2, with_replacement: bool = False) -> seq + +Compute combinations of length :math:`r` of the given tensor. The behavior is similar to +python's `itertools.combinations` when `with_replacement` is set to `False`, and +`itertools.combinations_with_replacement` when `with_replacement` is set to `True`. + +Arguments: + input (Tensor): 1D vector. + r (int, optional): number of elements to combine + with_replacement (bool, optional): whether to allow duplication in combination + +Returns: + Tensor: A tensor equivalent to converting all the input tensors into lists, do + `itertools.combinations` or `itertools.combinations_with_replacement` on these + lists, and finally convert the resulting list into tensor. + +Example:: + + >>> a = [1, 2, 3] + >>> list(itertools.combinations(a, r=2)) + [(1, 2), (1, 3), (2, 3)] + >>> list(itertools.combinations(a, r=3)) + [(1, 2, 3)] + >>> list(itertools.combinations_with_replacement(a, r=2)) + [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)] + >>> tensor_a = torch.tensor(a) + >>> torch.combinations(tensor_a) + tensor([[1, 2], + [1, 3], + [2, 3]]) + >>> torch.combinations(tensor_a, r=3) + tensor([[1, 2, 3]]) + >>> torch.combinations(tensor_a, with_replacement=True) + tensor([[1, 1], + [1, 2], + [1, 3], + [2, 2], + [2, 3], + [3, 3]]) + +""", +) + +add_docstr( + torch.trapezoid, + r""" +trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + +Computes the `trapezoidal rule `_ along +:attr:`dim`. By default the spacing between elements is assumed to be 1, but +:attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be +used to specify arbitrary spacing along :attr:`dim`. Only one of :attr:`x` or :attr:`dx` should be specified. + + +Assuming :attr:`y` is a one-dimensional tensor with elements :math:`{y_0, y_1, ..., y_n}`, +the default computation is + +.. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{1}{2} (y_i + y_{i-1}) + \end{aligned} + +When :attr:`dx` is specified the computation becomes + +.. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{\Delta x}{2} (y_i + y_{i-1}) + \end{aligned} + +effectively multiplying the result by :attr:`dx`. When :attr:`x` is specified, +assuming :attr:`x` is also a one-dimensional tensor with +elements :math:`{x_0, x_1, ..., x_n}`, the computation becomes + +.. math:: + \begin{aligned} + \sum_{i = 1}^{n} \frac{(x_i - x_{i-1})}{2} (y_i + y_{i-1}) + \end{aligned} + +When :attr:`x` and :attr:`y` have the same size, the computation is as described above and no broadcasting is needed. +The broadcasting behavior of this function is as follows when their sizes are different. For both :attr:`x` +and :attr:`y`, the function computes the difference between consecutive elements along +dimension :attr:`dim`. This effectively creates two tensors, `x_diff` and `y_diff`, that have +the same shape as the original tensors except their lengths along the dimension :attr:`dim` is reduced by 1. +After that, those two tensors are broadcast together to compute final output as part of the trapezoidal rule. +See the examples below for details. + +.. note:: + The trapezoidal rule is a technique for approximating the definite integral of a function + by averaging its left and right Riemann sums. The approximation becomes more accurate as + the resolution of the partition increases. + +Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + +Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + +Examples:: + + >>> # Computes the trapezoidal rule in 1D, spacing is implicitly 1 + >>> y = torch.tensor([1, 5, 10]) + >>> torch.trapezoid(y) + tensor(10.5) + + >>> # Computes the same trapezoidal rule directly to verify + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.trapezoid(y, dx=2) + 21.0 + + >>> # Computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + 28.5 + + >>> # Computes the same trapezoidal rule directly to verify + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.trapezoid(y) + tensor([ 2., 8., 14.]) + + >>> # Computes the trapezoidal rule for each column of the matrix + >>> torch.trapezoid(y, dim=0) + tensor([ 6., 8., 10.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.trapezoid(y, x) + array([5., 5., 5.]) + + >>> # Computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.trapezoid(y, x) + array([2., 4., 6.]) +""", +) + +add_docstr( + torch.trapz, + r""" +trapz(y, x, *, dim=-1) -> Tensor + +Alias for :func:`torch.trapezoid`. +""", +) + +add_docstr( + torch.cumulative_trapezoid, + r""" +cumulative_trapezoid(y, x=None, *, dx=None, dim=-1) -> Tensor + +Cumulatively computes the `trapezoidal rule `_ +along :attr:`dim`. By default the spacing between elements is assumed to be 1, but +:attr:`dx` can be used to specify a different constant spacing, and :attr:`x` can be +used to specify arbitrary spacing along :attr:`dim`. + +For more details, please read :func:`torch.trapezoid`. The difference between :func:`torch.trapezoid` +and this function is that, :func:`torch.trapezoid` returns a value for each integration, +where as this function returns a cumulative value for every spacing within the integration. This +is analogous to how `.sum` returns a value and `.cumsum` returns a cumulative sum. + +Arguments: + y (Tensor): Values to use when computing the trapezoidal rule. + x (Tensor): If specified, defines spacing between values as specified above. + +Keyword arguments: + dx (float): constant spacing between values. If neither :attr:`x` or :attr:`dx` + are specified then this defaults to 1. Effectively multiplies the result by its value. + dim (int): The dimension along which to compute the trapezoidal rule. + The last (inner-most) dimension by default. + +Examples:: + + >>> # Cumulatively computes the trapezoidal rule in 1D, spacing is implicitly 1. + >>> y = torch.tensor([1, 5, 10]) + >>> torch.cumulative_trapezoid(y) + tensor([3., 10.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> (1 + 5) / 2 + 3.0 + >>> (1 + 10 + 10) / 2 + 10.5 + + >>> # Cumulatively computes the trapezoidal rule in 1D with constant spacing of 2 + >>> # NOTE: the result is the same as before, but multiplied by 2 + >>> torch.cumulative_trapezoid(y, dx=2) + tensor([6., 21.]) + + >>> # Cumulatively computes the trapezoidal rule in 1D with arbitrary spacing + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([6., 28.5]) + + >>> # Computes the same trapezoidal rule directly up to each element to verify + >>> ((3 - 1) * (1 + 5)) / 2 + 6.0 + >>> ((3 - 1) * (1 + 5) + (6 - 3) * (5 + 10)) / 2 + 28.5 + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 matrix + >>> y = torch.arange(9).reshape(3, 3) + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> torch.cumulative_trapezoid(y) + tensor([[ 0.5, 2.], + [ 3.5, 8.], + [ 6.5, 14.]]) + + >>> # Cumulatively computes the trapezoidal rule for each column of the matrix + >>> torch.cumulative_trapezoid(y, dim=0) + tensor([[ 1.5, 2.5, 3.5], + [ 6.0, 8.0, 10.0]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with the same arbitrary spacing + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([1, 3, 6]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[2., 5.], + [2., 5.], + [2., 5.]]) + + >>> # Cumulatively computes the trapezoidal rule for each row of a 3x3 ones matrix + >>> # with different arbitrary spacing per row + >>> y = torch.ones(3, 3) + >>> x = torch.tensor([[1, 2, 3], [1, 3, 5], [1, 4, 7]]) + >>> torch.cumulative_trapezoid(y, x) + tensor([[1., 2.], + [2., 4.], + [3., 6.]]) +""", +) + +add_docstr( + torch.repeat_interleave, + r""" +repeat_interleave(input, repeats, dim=None, *, output_size=None) -> Tensor + +Repeat elements of a tensor. + +.. warning:: + + This is different from :meth:`torch.Tensor.repeat` but similar to ``numpy.repeat``. + +Args: + {input} + repeats (Tensor or int): The number of repetitions for each element. + repeats is broadcasted to fit the shape of the given axis. + dim (int, optional): The dimension along which to repeat values. + By default, use the flattened input array, and return a flat output + array. + +Keyword args: + output_size (int, optional): Total output size for the given axis + ( e.g. sum of repeats). If given, it will avoid stream synchronization + needed to calculate output shape of the tensor. + +Returns: + Tensor: Repeated tensor which has the same shape as input, except along the given axis. + +Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.repeat_interleave(2) + tensor([1, 1, 2, 2, 3, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.repeat_interleave(y, 2) + tensor([1, 1, 2, 2, 3, 3, 4, 4]) + >>> torch.repeat_interleave(y, 3, dim=1) + tensor([[1, 1, 1, 2, 2, 2], + [3, 3, 3, 4, 4, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0) + tensor([[1, 2], + [3, 4], + [3, 4]]) + >>> torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0, output_size=3) + tensor([[1, 2], + [3, 4], + [3, 4]]) + +If the `repeats` is `tensor([n1, n2, n3, ...])`, then the output will be +`tensor([0, 0, ..., 1, 1, ..., 2, 2, ..., ...])` where `0` appears `n1` times, +`1` appears `n2` times, `2` appears `n3` times, etc. + +.. function:: repeat_interleave(repeats, *) -> Tensor + :noindex: + +Repeats 0 repeats[0] times, 1 repeats[1] times, 2 repeats[2] times, etc. + +Args: + repeats (Tensor): The number of repetitions for each element. + +Returns: + Tensor: Repeated tensor of size `sum(repeats)`. + +Example:: + + >>> torch.repeat_interleave(torch.tensor([1, 2, 3])) + tensor([0, 1, 1, 2, 2, 2]) + +""".format(**common_args), +) + +add_docstr( + torch.tile, + r""" +tile(input, dims) -> Tensor + +Constructs a tensor by repeating the elements of :attr:`input`. +The :attr:`dims` argument specifies the number of repetitions +in each dimension. + +If :attr:`dims` specifies fewer dimensions than :attr:`input` has, then +ones are prepended to :attr:`dims` until all dimensions are specified. +For example, if :attr:`input` has shape (8, 6, 4, 2) and :attr:`dims` +is (2, 2), then :attr:`dims` is treated as (1, 1, 2, 2). + +Analogously, if :attr:`input` has fewer dimensions than :attr:`dims` +specifies, then :attr:`input` is treated as if it were unsqueezed at +dimension zero until it has as many dimensions as :attr:`dims` specifies. +For example, if :attr:`input` has shape (4, 2) and :attr:`dims` +is (3, 3, 2, 2), then :attr:`input` is treated as if it had the +shape (1, 1, 4, 2). + +.. note:: + + This function is similar to NumPy's tile function. + +Args: + input (Tensor): the tensor whose elements to repeat. + dims (tuple): the number of repetitions per dimension. + +Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> x.tile((2,)) + tensor([1, 2, 3, 1, 2, 3]) + >>> y = torch.tensor([[1, 2], [3, 4]]) + >>> torch.tile(y, (2, 2)) + tensor([[1, 2, 1, 2], + [3, 4, 3, 4], + [1, 2, 1, 2], + [3, 4, 3, 4]]) +""", +) + +add_docstr( + torch.quantize_per_tensor, + r""" +quantize_per_tensor(input, scale, zero_point, dtype) -> Tensor + +Converts a float tensor to a quantized tensor with given scale and zero point. + +Arguments: + input (Tensor): float tensor or list of tensors to quantize + scale (float or Tensor): scale to apply in quantization formula + zero_point (int or Tensor): offset in integer value that maps to float zero + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + +Returns: + Tensor: A newly quantized tensor or list of quantized tensors. + +Example:: + + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), 0.1, 10, torch.quint8).int_repr() + tensor([ 0, 10, 20, 30], dtype=torch.uint8) + >>> torch.quantize_per_tensor([torch.tensor([-1.0, 0.0]), torch.tensor([-2.0, 2.0])], + >>> torch.tensor([0.1, 0.2]), torch.tensor([10, 20]), torch.quint8) + (tensor([-1., 0.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=10), + tensor([-2., 2.], size=(2,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=20)) + >>> torch.quantize_per_tensor(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.tensor(0.1), torch.tensor(10), torch.quint8) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.10, zero_point=10) +""", +) + +add_docstr( + torch.quantize_per_tensor_dynamic, + r""" +quantize_per_tensor_dynamic(input, dtype, reduce_range) -> Tensor + +Converts a float tensor to a quantized tensor with scale and zero_point calculated +dynamically based on the input. + +Arguments: + input (Tensor): float tensor or list of tensors to quantize + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8`` + reduce_range (bool): a flag to indicate whether to reduce the range of quantized + data by 1 bit, it's required to avoid instruction overflow for some hardwares + +Returns: + Tensor: A newly (dynamically) quantized tensor + +Example:: + + >>> t = torch.quantize_per_tensor_dynamic(torch.tensor([-1.0, 0.0, 1.0, 2.0]), torch.quint8, False) + >>> print(t) + tensor([-1., 0., 1., 2.], size=(4,), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.011764705882352941, + zero_point=85) + >>> t.int_repr() + tensor([ 0, 85, 170, 255], dtype=torch.uint8) +""", +) + +add_docstr( + torch.quantize_per_channel, + r""" +quantize_per_channel(input, scales, zero_points, axis, dtype) -> Tensor + +Converts a float tensor to a per-channel quantized tensor with given scales and zero points. + +Arguments: + input (Tensor): float tensor to quantize + scales (Tensor): float 1D tensor of scales to use, size should match ``input.size(axis)`` + zero_points (int): integer 1D tensor of offset to use, size should match ``input.size(axis)`` + axis (int): dimension on which apply per-channel quantization + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + Has to be one of the quantized dtypes: ``torch.quint8``, ``torch.qint8``, ``torch.qint32`` + +Returns: + Tensor: A newly quantized tensor + +Example:: + + >>> x = torch.tensor([[-1.0, 0.0], [1.0, 2.0]]) + >>> torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8) + tensor([[-1., 0.], + [ 1., 2.]], size=(2, 2), dtype=torch.quint8, + quantization_scheme=torch.per_channel_affine, + scale=tensor([0.1000, 0.0100], dtype=torch.float64), + zero_point=tensor([10, 0]), axis=0) + >>> torch.quantize_per_channel(x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8).int_repr() + tensor([[ 0, 10], + [100, 200]], dtype=torch.uint8) +""", +) + + +add_docstr( + torch.quantized_batch_norm, + r""" +quantized_batch_norm(input, weight=None, bias=None, mean, var, eps, output_scale, output_zero_point) -> Tensor + +Applies batch normalization on a 4D (NCHW) quantized tensor. + +.. math:: + + y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta + +Arguments: + input (Tensor): quantized tensor + weight (Tensor): float tensor that corresponds to the gamma, size C + bias (Tensor): float tensor that corresponds to the beta, size C + mean (Tensor): float mean value in batch normalization, size C + var (Tensor): float tensor for variance, size C + eps (float): a value added to the denominator for numerical stability. + output_scale (float): output quantized tensor scale + output_zero_point (int): output quantized tensor zero_point + +Returns: + Tensor: A quantized tensor with batch normalization applied. + +Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2, 2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_batch_norm(qx, torch.ones(2), torch.zeros(2), torch.rand(2), torch.rand(2), 0.00001, 0.2, 2) + tensor([[[[-0.2000, -0.2000], + [ 1.6000, -0.2000]], + + [[-0.4000, -0.4000], + [-0.4000, 0.6000]]], + + + [[[-0.2000, -0.2000], + [-0.2000, -0.2000]], + + [[ 0.6000, -0.4000], + [ 0.6000, -0.4000]]]], size=(2, 2, 2, 2), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=0.2, zero_point=2) +""", +) + + +add_docstr( + torch.quantized_max_pool1d, + r""" +quantized_max_pool1d(input, kernel_size, stride=[], padding=0, dilation=1, ceil_mode=False) -> Tensor + +Applies a 1D max pooling over an input quantized tensor composed of several input planes. + +Arguments: + input (Tensor): quantized tensor + kernel_size (list of int): the size of the sliding window + stride (``list of int``, optional): the stride of the sliding window + padding (``list of int``, optional): padding to be added on both sides, must be >= 0 and <= kernel_size / 2 + dilation (``list of int``, optional): The stride between elements within a sliding window, must be > 0. Default 1 + ceil_mode (bool, optional): If True, will use ceil instead of floor to compute the output shape. + Defaults to False. + + +Returns: + Tensor: A quantized tensor with max_pool1d applied. + +Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_max_pool1d(qx, [2]) + tensor([[0.0000], + [1.5000]], size=(2, 1), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=1.5, zero_point=3) +""", +) + + +add_docstr( + torch.quantized_max_pool2d, + r""" +quantized_max_pool2d(input, kernel_size, stride=[], padding=0, dilation=1, ceil_mode=False) -> Tensor + +Applies a 2D max pooling over an input quantized tensor composed of several input planes. + +Arguments: + input (Tensor): quantized tensor + kernel_size (``list of int``): the size of the sliding window + stride (``list of int``, optional): the stride of the sliding window + padding (``list of int``, optional): padding to be added on both sides, must be >= 0 and <= kernel_size / 2 + dilation (``list of int``, optional): The stride between elements within a sliding window, must be > 0. Default 1 + ceil_mode (bool, optional): If True, will use ceil instead of floor to compute the output shape. + Defaults to False. + + +Returns: + Tensor: A quantized tensor with max_pool2d applied. + +Example:: + + >>> qx = torch.quantize_per_tensor(torch.rand(2, 2, 2, 2), 1.5, 3, torch.quint8) + >>> torch.quantized_max_pool2d(qx, [2,2]) + tensor([[[[1.5000]], + + [[1.5000]]], + + + [[[0.0000]], + + [[0.0000]]]], size=(2, 2, 1, 1), dtype=torch.quint8, + quantization_scheme=torch.per_tensor_affine, scale=1.5, zero_point=3) +""", +) + + +add_docstr( + torch.Stream, + r""" +Stream(device, *, priority) -> Stream + +An in-order queue of executing the respective tasks asynchronously in first in first out (FIFO) order. +It can control or synchronize the execution of other Stream or block the current host thread to ensure +the correct task sequencing. It supports with statement as a context manager to ensure the operators +within the with block are running on the corresponding stream. + +See in-depth description of the CUDA behavior at :ref:`cuda-semantics` for details +on the exact semantic that applies to all devices. + +Arguments: + device (:class:`torch.device`, optional): the desired device for the Stream. + If not given, the current :ref:`accelerator` type will be used. + priority (int, optional): priority of the stream, should be 0 or negative, where negative + numbers indicate higher priority. By default, streams have priority 0. + +Returns: + Stream: An torch.Stream object. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> with torch.Stream(device='cuda') as s_cuda: + >>> a = torch.randn(10, 5, device='cuda') + >>> b = torch.randn(5, 10, device='cuda') + >>> c = torch.mm(a, b) +""", +) + + +add_docstr( + torch.Stream.query, + r""" +Stream.query() -> bool + +Check if all the work submitted has been completed. + +Returns: + bool: A boolean indicating if all kernels in this stream are completed. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> s_cuda.query() + True +""", +) + + +add_docstr( + torch.Stream.record_event, + r""" +Stream.record_event(event) -> Event + +Record an event. En-queuing it into the Stream to allow further synchronization from the current point in the FIFO queue. + +Arguments: + event (:class:`torch.Event`, optional): event to record. If not given, a new one will be allocated. + +Returns: + Event: Recorded event. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> e_cuda = s_cuda.record_event() +""", +) + + +add_docstr( + torch.Stream.synchronize, + r""" +Stream.synchronize() -> None + +Wait for all the kernels in this stream to complete. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> s_cuda.synchronize() +""", +) + + +add_docstr( + torch.Stream.wait_event, + r""" +Stream.wait_event(event) -> None + +Make all future work submitted to the stream wait for an event. + +Arguments: + event (:class:`torch.Event`): an event to wait for. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s1_cuda = torch.Stream(device='cuda') + >>> s2_cuda = torch.Stream(device='cuda') + >>> e_cuda = s1_cuda.record_event() + >>> s2_cuda.wait_event(e_cuda) +""", +) + + +add_docstr( + torch.Stream.wait_stream, + r""" +Stream.wait_stream(stream) -> None + +Synchronize with another stream. All future work submitted to this stream will wait until all kernels +already submitted to the given stream are completed. + +Arguments: + stream (:class:`torch.Stream`): a stream to synchronize. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s1_cuda = torch.Stream(device='cuda') + >>> s2_cuda = torch.Stream(device='cuda') + >>> s2_cuda.wait_stream(s1_cuda) +""", +) + + +add_docstr( + torch.Event, + r""" +Event(device=None, *, enable_timing=False, blocking=False, interprocess=False) + +Query and record Stream status to identify or control dependencies across Stream and measure timing. + +Arguments: + device (:class:`torch.device`, optional): the desired device for the Event. + If not given, the current :ref:`accelerator` type will be used. + enable_timing (bool, optional): indicates if the event should measure time (default: ``False``) + blocking (bool, optional): if ``True``, :meth:`wait` will be blocking (default: ``False``) + interprocess (bool): if ``True``, the event can be shared between processes (default: ``False``) + +.. warning:: + + Both blocking and interprocess are not supported right now and are noops. + +Returns: + Event: An torch.Event object. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> event = torch.Event() + >>> e_cuda = torch.Event(device='cuda') +""", +) + + +add_docstr( + torch.Event.elapsed_time, + r""" +Event.elapsed_time(end_event) -> float + +Returns the elapsed time in milliseconds between when this event and the :attr:`end_event` are +each recorded via :func:`torch.Stream.record_event`. + +Arguments: + end_event (:class:`torch.Event`): The ending event has been recorded. + +Returns: + float: Time between starting and ending event in milliseconds. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> e1_cuda = s_cuda.record_event() + >>> e2_cuda = s_cuda.record_event() + >>> ms = e1_cuda.elapsed_time(e2_cuda) +""", +) + + +add_docstr( + torch.Event.query, + r""" +Event.query() -> bool + +Check if the stream where this event was recorded already moved past the point where the event was recorded. +Always returns ``True`` if the Event was not recorded. + +Returns: + bool: A boolean indicating if all work currently captured by event has completed. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> e_cuda = s_cuda.record_event() + >>> e_cuda.query() + True +""", +) + + +add_docstr( + torch.Event.record, + r""" +Event.record(stream=None) -> None + +Record the event in a given stream. The stream's device must match the event's device. +This function is equivalent to ``stream.record_event(self)``. + +Arguments: + stream (:class:`torch.Stream`, optional): A stream to be recorded. + If not given, the current stream will be used. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> e_cuda = torch.Event(device='cuda') + >>> e_cuda.record() +""", +) + + +add_docstr( + torch.Event.synchronize, + r""" +Event.synchronize() -> None + +Wait for the event to complete. This prevents the CPU thread from proceeding until the event completes. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s_cuda = torch.Stream(device='cuda') + >>> e_cuda = s_cuda.record_event() + >>> e_cuda.synchronize() +""", +) + + +add_docstr( + torch.Event.wait, + r""" +Event.wait(stream=None) -> None + +Make all future work submitted to the given stream wait for this event. + +Arguments: + stream (:class:`torch.Stream`, optional): A stream to synchronize. + If not given, the current stream will be used. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> s1_cuda = torch.Stream(device='cuda') + >>> s2_cuda = torch.Stream(device='cuda') + >>> e_cuda = s1_cuda.record() + >>> e_cuda.wait(s2) +""", +) + + +add_docstr( + torch.Generator, + r""" +Generator(device='cpu') -> Generator + +Creates and returns a generator object that manages the state of the algorithm which +produces pseudo random numbers. Used as a keyword argument in many :ref:`inplace-random-sampling` +functions. + +Arguments: + device (:class:`torch.device`, optional): the desired device for the generator. + +Returns: + Generator: An torch.Generator object. + +Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> g_cpu = torch.Generator() + >>> g_cuda = torch.Generator(device='cuda') +""", +) + + +add_docstr( + torch.Generator.set_state, + r""" +Generator.set_state(new_state) -> void + +Sets the Generator state. + +Arguments: + new_state (torch.ByteTensor): The desired state. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu_other = torch.Generator() + >>> g_cpu.set_state(g_cpu_other.get_state()) +""", +) + + +add_docstr( + torch.Generator.get_state, + r""" +Generator.get_state() -> Tensor + +Returns the Generator state as a ``torch.ByteTensor``. + +Returns: + Tensor: A ``torch.ByteTensor`` which contains all the necessary bits + to restore a Generator to a specific point in time. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu.get_state() +""", +) + +add_docstr( + torch.Generator.graphsafe_set_state, + r""" +Generator.graphsafe_set_state(state) -> None + +Sets the state of the generator to the specified state in a manner that is safe for use in graph capture. +This method is crucial for ensuring that the generator's state can be captured in the CUDA graph. + +Arguments: + state (torch.Generator): A Generator point to the new state for the generator, typically obtained from `graphsafe_get_state`. + +Example: + >>> g_cuda = torch.Generator(device='cuda') + >>> g_cuda_other = torch.Generator(device='cuda') + >>> current_state = g_cuda_other.graphsafe_get_state() + >>> g_cuda.graphsafe_set_state(current_state) +""", +) + +add_docstr( + torch.Generator.graphsafe_get_state, + r""" +Generator.graphsafe_get_state() -> torch.Generator + +Retrieves the current state of the generator in a manner that is safe for graph capture. +This method is crucial for ensuring that the generator's state can be captured in the CUDA graph. + +Returns: + torch.Generator: A Generator point to the current state of the generator + +Example: + >>> g_cuda = torch.Generator(device='cuda') + >>> current_state = g_cuda.graphsafe_get_state() +""", +) + +add_docstr( + torch.Generator.clone_state, + r""" +Generator.clone_state() -> torch.Generator + +Clones the current state of the generator and returns a new generator pointing to this cloned state. +This method is beneficial for preserving a particular state of a generator to restore at a later point. + +Returns: + torch.Generator: A Generator pointing to the newly cloned state. + +Example: + >>> g_cuda = torch.Generator(device='cuda') + >>> cloned_state = g_cuda.clone_state() +""", +) + +add_docstr( + torch.Generator.manual_seed, + r""" +Generator.manual_seed(seed) -> Generator + +Sets the seed for generating random numbers. Returns a `torch.Generator` object. Any 32-bit integer is a valid seed. + +Arguments: + seed (int): The desired seed. Value must be within the inclusive range + `[-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]`. Otherwise, a RuntimeError + is raised. Negative inputs are remapped to positive values with the formula + `0xffff_ffff_ffff_ffff + seed`. + +Returns: + Generator: An torch.Generator object. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu.manual_seed(2147483647) +""", +) + + +add_docstr( + torch.Generator.initial_seed, + r""" +Generator.initial_seed() -> int + +Returns the initial seed for generating random numbers. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu.initial_seed() + 2147483647 +""", +) + + +add_docstr( + torch.Generator.seed, + r""" +Generator.seed() -> int + +Gets a non-deterministic random number from std::random_device or the current +time and uses it to seed a Generator. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu.seed() + 1516516984916 +""", +) + + +add_docstr( + torch.Generator.device, + r""" +Generator.device -> device + +Gets the current device of the generator. + +Example:: + + >>> g_cpu = torch.Generator() + >>> g_cpu.device + device(type='cpu') +""", +) + +add_docstr( + torch._assert_async, + r""" +_assert_async(tensor) -> void + +Asynchronously assert that the contents of tensor are nonzero. For CPU tensors, +this is equivalent to ``assert tensor`` or ``assert tensor.is_nonzero()``; for +CUDA tensors, we DO NOT synchronize and you may only find out the assertion +failed at a later CUDA kernel launch. Asynchronous assertion can be helpful for +testing invariants in CUDA tensors without giving up performance. This function +is NOT intended to be used for regular error checking, as it will trash your CUDA +context if the assert fails (forcing you to restart your PyTorch process.) + +Args: + tensor (Tensor): a one element tensor to test to see if it is nonzero. Zero + elements (including False for boolean tensors) cause an assertion failure + to be raised. +""", +) + +add_docstr( + torch.searchsorted, + r""" +searchsorted(sorted_sequence, values, *, out_int32=False, right=False, side=None, out=None, sorter=None) -> Tensor + +Find the indices from the *innermost* dimension of :attr:`sorted_sequence` such that, if the +corresponding values in :attr:`values` were inserted before the indices, when sorted, the order +of the corresponding *innermost* dimension within :attr:`sorted_sequence` would be preserved. +Return a new tensor with the same size as :attr:`values`. More formally, +the returned index satisfies the following rules: + +.. list-table:: + :widths: 12 10 78 + :header-rows: 1 + + * - :attr:`sorted_sequence` + - :attr:`right` + - *returned index satisfies* + * - 1-D + - False + - ``sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]`` + * - 1-D + - True + - ``sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]`` + * - N-D + - False + - ``sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]`` + * - N-D + - True + - ``sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]`` + +Args: + sorted_sequence (Tensor): N-D or 1-D tensor, containing monotonically increasing sequence on the *innermost* + dimension unless :attr:`sorter` is provided, in which case the sequence does not + need to be sorted + values (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + +Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): if False, return the first suitable location that is found. If True, return the + last such index. If no suitable index found, return 0 for non-numerical value + (eg. nan, inf) or the size of *innermost* dimension within :attr:`sorted_sequence` + (one pass the last index of the *innermost* dimension). In other words, if False, + gets the lower bound index for each value in :attr:`values` on the corresponding + *innermost* dimension of the :attr:`sorted_sequence`. If True, gets the upper + bound index instead. Default value is False. :attr:`side` does the same and is + preferred. It will error if :attr:`side` is set to "left" while this is True. + side (str, optional): the same as :attr:`right` but preferred. "left" corresponds to False for :attr:`right` + and "right" corresponds to True for :attr:`right`. It will error if this is set to + "left" while :attr:`right` is True. Default value is None. + out (Tensor, optional): the output tensor, must be the same size as :attr:`values` if provided. + sorter (LongTensor, optional): if provided, a tensor matching the shape of the unsorted + :attr:`sorted_sequence` containing a sequence of indices that sort it in the + ascending order on the innermost dimension + + +Example:: + + >>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]) + >>> sorted_sequence + tensor([[ 1, 3, 5, 7, 9], + [ 2, 4, 6, 8, 10]]) + >>> values = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> values + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.searchsorted(sorted_sequence, values) + tensor([[1, 3, 4], + [1, 2, 4]]) + >>> torch.searchsorted(sorted_sequence, values, side='right') + tensor([[2, 3, 5], + [1, 3, 4]]) + + >>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9]) + >>> sorted_sequence_1d + tensor([1, 3, 5, 7, 9]) + >>> torch.searchsorted(sorted_sequence_1d, values) + tensor([[1, 3, 4], + [1, 3, 4]]) +""", +) + +add_docstr( + torch.bucketize, + r""" +bucketize(input, boundaries, *, out_int32=False, right=False, out=None) -> Tensor + +Returns the indices of the buckets to which each value in the :attr:`input` belongs, where the +boundaries of the buckets are set by :attr:`boundaries`. Return a new tensor with the same size +as :attr:`input`. If :attr:`right` is False (default), then the left boundary is open. Note that +this behavior is opposite the behavior of +`numpy.digitize `_. +More formally, the returned index satisfies the following rules: + +.. list-table:: + :widths: 15 85 + :header-rows: 1 + + * - :attr:`right` + - *returned index satisfies* + * - False + - ``boundaries[i-1] < input[m][n]...[l][x] <= boundaries[i]`` + * - True + - ``boundaries[i-1] <= input[m][n]...[l][x] < boundaries[i]`` + +Args: + input (Tensor or Scalar): N-D tensor or a Scalar containing the search value(s). + boundaries (Tensor): 1-D tensor, must contain a strictly increasing sequence, or the return value is undefined. + +Keyword args: + out_int32 (bool, optional): indicate the output data type. torch.int32 if True, torch.int64 otherwise. + Default value is False, i.e. default output data type is torch.int64. + right (bool, optional): determines the behavior for values in :attr:`boundaries`. See the table above. + out (Tensor, optional): the output tensor, must be the same size as :attr:`input` if provided. + + +Example:: + + >>> boundaries = torch.tensor([1, 3, 5, 7, 9]) + >>> boundaries + tensor([1, 3, 5, 7, 9]) + >>> v = torch.tensor([[3, 6, 9], [3, 6, 9]]) + >>> v + tensor([[3, 6, 9], + [3, 6, 9]]) + >>> torch.bucketize(v, boundaries) + tensor([[1, 3, 4], + [1, 3, 4]]) + >>> torch.bucketize(v, boundaries, right=True) + tensor([[2, 3, 5], + [2, 3, 5]]) +""", +) + +add_docstr( + torch.view_as_real_copy, + r""" +Performs the same operation as :func:`torch.view_as_real`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.view_as_complex_copy, + r""" +Performs the same operation as :func:`torch.view_as_complex`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.as_strided_copy, + r""" +Performs the same operation as :func:`torch.as_strided`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.diagonal_copy, + r""" +Performs the same operation as :func:`torch.diagonal`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.expand_copy, + r""" +Performs the same operation as :func:`torch.Tensor.expand`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.permute_copy, + r""" +Performs the same operation as :func:`torch.permute`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.select_copy, + r""" +Performs the same operation as :func:`torch.select`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.detach_copy, + r""" +Performs the same operation as :func:`torch.detach`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.slice_copy, + r""" +Performs the same operation as :func:`torch.slice`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.split_copy, + r""" +Performs the same operation as :func:`torch.split`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.split_with_sizes_copy, + r""" +Performs the same operation as :func:`torch.split_with_sizes`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.squeeze_copy, + r""" +Performs the same operation as :func:`torch.squeeze`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.t_copy, + r""" +Performs the same operation as :func:`torch.t`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.transpose_copy, + r""" +Performs the same operation as :func:`torch.transpose`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.unsqueeze_copy, + r""" +Performs the same operation as :func:`torch.unsqueeze`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.indices_copy, + r""" +Performs the same operation as :func:`torch.indices`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.values_copy, + r""" +Performs the same operation as :func:`torch.values`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.crow_indices_copy, + r""" +Performs the same operation as :func:`torch.crow_indices`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.col_indices_copy, + r""" +Performs the same operation as :func:`torch.col_indices`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.unbind_copy, + r""" +Performs the same operation as :func:`torch.unbind`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.view_copy, + r""" +Performs the same operation as :func:`torch.view`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.unfold_copy, + r""" +Performs the same operation as :func:`torch.unfold`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +add_docstr( + torch.alias_copy, + r""" +Performs the same operation as :func:`torch.alias`, but all output tensors +are freshly created instead of aliasing the input. +""", +) + +for unary_base_func_name in ( + "exp", + "sqrt", + "abs", + "acos", + "asin", + "atan", + "ceil", + "cos", + "cosh", + "erf", + "erfc", + "expm1", + "floor", + "log", + "log10", + "log1p", + "log2", + "neg", + "tan", + "tanh", + "sin", + "sinh", + "round", + "lgamma", + "frac", + "reciprocal", + "sigmoid", + "trunc", + "zero", +): + unary_foreach_func_name = f"_foreach_{unary_base_func_name}" + if hasattr(torch, unary_foreach_func_name): + add_docstr( + getattr(torch, unary_foreach_func_name), + rf""" +{unary_foreach_func_name}(self: List[Tensor]) -> List[Tensor] + +Apply :func:`torch.{unary_base_func_name}` to each Tensor of the input list. + """, + ) + unary_inplace_foreach_func_name = f"{unary_foreach_func_name}_" + if hasattr(torch, unary_inplace_foreach_func_name): + add_docstr( + getattr(torch, unary_inplace_foreach_func_name), + rf""" +{unary_inplace_foreach_func_name}(self: List[Tensor]) -> None + +Apply :func:`torch.{unary_base_func_name}` to each Tensor of the input list. + """, + ) diff --git a/venv/lib/python3.10/site-packages/torch/_utils.py b/venv/lib/python3.10/site-packages/torch/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9bd062cb5cec66888eb6edf429a209eb38b2fe91 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_utils.py @@ -0,0 +1,1111 @@ +# mypy: allow-untyped-defs +import copyreg +import functools +import importlib +import logging +import sys +import traceback +import warnings +from collections import defaultdict +from types import ModuleType +from typing import Any, Callable, Generic, Optional, TYPE_CHECKING +from typing_extensions import deprecated, ParamSpec + +import torch + + +def _type(self, dtype=None, non_blocking=False, **kwargs): + """Returns the type if `dtype` is not provided, else casts this object to + the specified type. + + If this is already of the correct type, no copy is performed and the + original object is returned. + + Args: + dtype (type or string): The desired type + non_blocking (bool): If ``True``, and the source is in pinned memory + and destination is on the GPU or vice versa, the copy is performed + asynchronously with respect to the host. Otherwise, the argument + has no effect. + **kwargs: For compatibility, may contain the key ``async`` in place of + the ``non_blocking`` argument. The ``async`` arg is deprecated. + """ + non_blocking = _get_async_or_non_blocking("type", non_blocking, kwargs) + if dtype is None: + return self.__module__ + "." + self.__class__.__name__ + + if isinstance(dtype, str): + dtype = _import_dotted_name(dtype) + if dtype == type(self): + return self + if self.is_sparse: + if not dtype.is_sparse: + raise RuntimeError("Cannot cast sparse tensor to dense tensor") + new_module_name = dtype.__module__.replace(".sparse", "") + new_values_type_name = new_module_name + "." + dtype.__name__ + new_values = torch.Tensor._values(self).type(new_values_type_name, non_blocking) + new_indices_type_name = new_module_name + ".LongTensor" + new_indices = torch.Tensor._indices(self).type( + new_indices_type_name, non_blocking + ) + return dtype(new_indices, new_values, self.size()) + if dtype.is_sparse: + raise RuntimeError("Cannot cast dense tensor to sparse tensor") + return dtype(self.size()).copy_(self, non_blocking) + + +def _to(self, device, non_blocking=False): + """Returns a copy of this object in device memory. + + If this object is already on the correct device, then no copy is performed + and the original object is returned. + + Args: + device (int): The destination device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + if self.device == device: + return self + + if device.type == "cpu": + pin_memory = non_blocking and self.device.type in ( + "cuda", + torch._C._get_privateuse1_backend_name(), + ) + untyped_storage = torch.empty( + self.nbytes(), dtype=torch.uint8, device=device, pin_memory=pin_memory + ).untyped_storage() + untyped_storage.copy_(self, non_blocking) + return untyped_storage + + device_module = getattr(torch, device.type, None) + assert device_module is not None, ( + f"{device.type.upper()} device module is not loaded" + ) + with device_module.device(device): + if self.is_sparse and hasattr(device_module, "sparse"): + new_type = getattr(device_module.sparse, self.__class__.__name__) + indices = getattr(torch.Tensor._indices(self), device.type)( + device, non_blocking + ) + values = getattr(torch.Tensor._values(self), device.type)( + device, non_blocking + ) + return new_type(indices, values, self.size()) + else: + assert not self.is_sparse, ( + f"sparse storage is not supported for {device.type.upper()} tensors" + ) + untyped_storage = torch.UntypedStorage(self.size(), device=device) + untyped_storage.copy_(self, non_blocking) + return untyped_storage + + +def _get_async_or_non_blocking(function_name, non_blocking, kwargs): + """Return the non-blocking flag given the function name and kwargs. + + Args: + function_name (str): the name of the function being used. + non_blocking (bool): the default value. + **kwargs (dict): the kwargs passed to the function. + """ + if not kwargs: + return non_blocking + if len(kwargs) != 1 or "async" not in kwargs: + message = "{}() got an unexpected keyword argument '{}'" + argument = list(kwargs.keys()).pop() + raise TypeError(message.format(function_name, argument)) + warnings.warn("'async' is deprecated; use 'non_blocking'") + return kwargs["async"] + + +def _get_restore_location(device): + """Return the map_location location. + + Used for rebuild functions where the tensor device is distinct from the storage + """ + + map_location = torch.serialization._serialization_tls.map_location + if map_location is None: + return device + else: + if isinstance(map_location, dict): + return map_location.get(device, device) + elif isinstance(map_location, (str, torch.device)): + return map_location + else: + assert callable(map_location) + raise RuntimeError( + "Callable map_location not supported with _rebuild_wrapper_subclass " + "or _rebuild_device_tensor_from_numpy" + ) + + +# Note [Don't serialize hooks] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# Since time immemorial, we have serialized the backward hooks associated with +# variables. This kind of half-worked--Python can pickle global functions +# (but not closures!)--but there were problems. +# +# - It's fragile. If you serialize a backward hook into a saved +# model, and then you rename the function associated with the hook, +# now your saved model is broken and you can't load it anymore. +# +# - It's not actually used. The standard recommendation is to +# serialize the *state_dict* of a model, not the model itself +# (since this is more stable to code changes affecting the model +# serialization), and the state dict saves "data" only, thus +# stripping the backward hooks. In some cases, hooks are +# essential to the well-functioning of a model (e.g., DDP), +# but DDP already manages re-adding the hooks! +# +# - We didn't serialize them in many cases. Prior to #10220, we +# were dropping backward hooks in ForkingPickler. We "fixed" this +# to be convenient with other serialization sites, but lack of +# serializing backward hooks wasn't actually the root cause of +# the bug. +# +# With these cases in mind, we have decided that a better strategy +# is to just NOT serialize hooks at all. +# +# Since this is a BC-breaking change, we should warn when we previously +# serialized a hook, but no longer do so. This will be done by adding a special +# sentinel property to hooks will be used to suppress this warning. If a hook +# has the property _torch_serialize_ignore, we will not emit a warning if we +# attempt to serialize a Tensor with this hook attached to it. +# +# By the way, when _backward_hooks is skipped, we must give an EMPTY +# OrderedDict(), if you pass a None you'll run afoul #12219. + + +# TODO: Once we decide to break serialization FC, `storage` no longer needs to +# be a TypedStorage +def _rebuild_tensor(storage, storage_offset, size, stride): + # first construct a tensor with the correct dtype/device + t = torch.empty((0,), dtype=storage.dtype, device=storage._untyped_storage.device) + return t.set_(storage._untyped_storage, storage_offset, size, stride) + + +def get_tensor_metadata(tensor): + # Tensor's Metadata for serializing. + # Currently, this only returns a dict[string, bool] specifying whether + # `conj` or `neg` bit is set. + assert isinstance(tensor, torch.Tensor) + return torch._C._get_tensor_metadata(tensor) # type: ignore[attr-defined] + + +def set_tensor_metadata(tensor, metadata): + # See `get_tensor_metadata` above + assert isinstance(metadata, dict) + assert isinstance(tensor, torch.Tensor) + torch._C._set_tensor_metadata(tensor, metadata) # type: ignore[attr-defined] + + +def _restore_device_fake_mode(tensor): + if torch._guards.detect_fake_mode(None) is not None: + if tensor.untyped_storage()._fake_device is not None: + device = _get_restore_location(tensor.untyped_storage()._fake_device) + if not isinstance(device, torch.device): + device = torch.device(device) + tensor.fake_device = torch.device(device) + return tensor + + +def _rebuild_tensor_v2( + storage, + storage_offset, + size, + stride, + requires_grad, + backward_hooks, + metadata=None, +): + tensor = _rebuild_tensor(storage, storage_offset, size, stride) + tensor.requires_grad = requires_grad + if metadata: + set_tensor_metadata(tensor, metadata) + + # NB: This line exists only for backwards compatibility; the + # general expectation is that backward_hooks is an empty + # OrderedDict. See Note [Don't serialize hooks] + tensor._backward_hooks = backward_hooks + + tensor = _restore_device_fake_mode(tensor) + return tensor + + +def _rebuild_tensor_v3( + storage, + storage_offset, + size, + stride, + requires_grad, + backward_hooks, + dtype, + metadata=None, +): + t = torch.empty( + (0,), + dtype=dtype, + device=storage._untyped_storage.device, + requires_grad=requires_grad, + ) + t.set_(storage._untyped_storage, storage_offset, size, stride) + if metadata: + set_tensor_metadata(t, metadata) + t._backward_hooks = backward_hooks + t = _restore_device_fake_mode(t) + return t + + +_sparse_tensors_to_validate: list["torch.Tensor"] = [] + + +# In _legacy_load() in serialization.py we unpickle storages after the sparse +# tensors have been already unpickled. Those storages contain data necessary for +# validating sparse tensors: indices and values. That's why sparse tensors are +# first unpickled without any validation, and then this function is called just +# before _legacy_load() returns, so that all the sparse tensors can be validated +# in bulk. +# +# The same procedure must be followed by _load() in serialization.py because due +# to Pickler semantics, we have to use the same (non-validating) function for +# unpickling sparse tensors, regardless of the caller. +def _validate_loaded_sparse_tensors(): + if not torch.sparse.check_sparse_tensor_invariants().is_enabled(): + # Skip sparse tensor invariants validation for better + # performance. See check_sparse_tensor_invariants + # documentation for how to control sparse tensor invariants + # checking. + _sparse_tensors_to_validate.clear() + return + try: + # We disable pinning check (see check_pinning=False below) to + # avoid gh-153143. In fact, pinning check is unnecessary + # anywhy when loading sparse data from external sources. + for t in _sparse_tensors_to_validate: + if t.layout is torch.sparse_coo: + torch._validate_sparse_coo_tensor_args( + t._indices(), + t._values(), + t.size(), + t.is_coalesced(), + check_pinning=False, + ) + elif t.layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + # TODO: Validation currently involves an expensive traversal + # on CPU, which may include a device transfer. + if t.layout in {torch.sparse_csr, torch.sparse_bsr}: + compressed_indices, plain_indices = ( + t.crow_indices(), + t.col_indices(), + ) + else: + compressed_indices, plain_indices = ( + t.ccol_indices(), + t.row_indices(), + ) + torch._validate_sparse_compressed_tensor_args( + compressed_indices, + plain_indices, + t.values(), + t.size(), + t.layout, + check_pinning=False, + ) + else: + raise NotImplementedError( + f"_validate_loaded_sparse_tensors for layout `{t.layout}`" + ) + + finally: + _sparse_tensors_to_validate.clear() + + +def _rebuild_sparse_tensor(layout, data): + """ + Rebuilds a sparse tensor from its sparse storage representation. + + Args: + layout (str): The sparse storage layout of the tensor. + data (tuple): The tensor's sparse storage representation. + """ + if layout == torch.sparse_coo: + if len(data) == 3: + # For BC: + indices, values, size = data + is_coalesced = None + else: + indices, values, size, is_coalesced = data + result = torch.sparse_coo_tensor( + indices, values, size, check_invariants=False, is_coalesced=is_coalesced + ) + _sparse_tensors_to_validate.append(result) + return result + + elif layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + compressed_indices, plain_indices, values, size = data + result = torch.sparse_compressed_tensor( + compressed_indices, + plain_indices, + values, + size, + layout=layout, + check_invariants=False, + ) + _sparse_tensors_to_validate.append(result) + return result + + raise NotImplementedError(f"rebuilding sparse tensor for layout {layout}") + + +def _rebuild_nested_tensor(buffer, sizes, strides, storage_offsets): + return torch._nested_view_from_buffer(buffer, sizes, strides, storage_offsets) + + +def _rebuild_device_tensor_from_cpu_tensor(data, dtype, device, requires_grad): + device = _get_restore_location(device) + tensor = data.to(dtype=dtype, device=device) + tensor.requires_grad = requires_grad + return tensor + + +def _rebuild_device_tensor_from_numpy(data, dtype, device, requires_grad): + device = _get_restore_location(device) + tensor = torch.from_numpy(data).to(dtype=dtype, device=device) + tensor.requires_grad = requires_grad + return tensor + + +# Should not be used, only here to be able to load Tensors serialized with older versions of pytorch +_rebuild_xla_tensor = _rebuild_device_tensor_from_numpy + + +def _rebuild_meta_tensor_no_storage(dtype, size, stride, requires_grad): + return torch.empty_strided( + size, stride, dtype=dtype, device="meta", requires_grad=requires_grad + ) + + +def _rebuild_wrapper_subclass( + cls, + dtype, + size, + stride, + storage_offset, + layout, + device, + requires_grad, +): + device = _get_restore_location(device) + return torch.Tensor._make_wrapper_subclass( + cls, + size, + strides=stride, + dtype=dtype, + storage_offset=storage_offset, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + +# TODO: Once we decide to break serialization FC, `storage` no longer needs to +# be a TypedStorage +def _rebuild_qtensor( + storage, + storage_offset, + size, + stride, + quantizer_params, + requires_grad, + backward_hooks, +): + qscheme = quantizer_params[0] + if qscheme == torch.per_tensor_affine: + _, scale, zero_point = quantizer_params + tensor = torch._empty_affine_quantized( + size, + scale=scale, + zero_point=zero_point, + dtype=storage.dtype, + device=storage.device, + ) + elif qscheme in (torch.per_channel_affine, torch.per_channel_affine_float_qparams): + _, scales, zero_points, axis = quantizer_params + if type(scales) is list and type(zero_points) is list: + if qscheme == torch.per_channel_affine: + scales = torch.tensor(scales, dtype=torch.double, device=storage.device) + zero_points = torch.tensor( + zero_points, dtype=torch.long, device=storage.device + ) + else: + scales = torch.tensor(scales, dtype=torch.float, device=storage.device) + zero_points = torch.tensor( + zero_points, dtype=torch.float, device=storage.device + ) + tensor = torch._empty_per_channel_affine_quantized( + size, + scales=scales, + zero_points=zero_points, + axis=axis, + dtype=storage.dtype, + device=storage.device, + ) + else: + raise RuntimeError(f"Can't deserialize quantized tensor with qscheme {qscheme}") + tensor.set_(storage, storage_offset, size, stride) + tensor.requires_grad = requires_grad + # NB: This line exists only for backwards compatibility; the + # general expectation is that backward_hooks is an empty + # OrderedDict. See Note [Don't serialize hooks] + tensor._backward_hooks = backward_hooks + return tensor + + +def _rebuild_parameter(data, requires_grad, backward_hooks): + param = torch.nn.Parameter(data, requires_grad) + # NB: This line exists only for backwards compatibility; the + # general expectation is that backward_hooks is an empty + # OrderedDict. See Note [Don't serialize hooks] + param._backward_hooks = backward_hooks + + return param + + +def _rebuild_parameter_with_state(data, requires_grad, backward_hooks, state): + param = torch.nn.Parameter(data, requires_grad) + # NB: This line exists only for backwards compatibility; the + # general expectation is that backward_hooks is an empty + # OrderedDict. See Note [Don't serialize hooks] + param._backward_hooks = backward_hooks + + # Restore state on Parameter like python attr. + param = _set_obj_state(param, state) + return param + + +def _get_obj_state(obj): + # Get the state of the python subclass + # This loosely mimics the function on the object class but since Tensor do not inherit + # from it, we cannot call that function directly + # https://github.com/python/cpython/blob/c83919bd635f4433f1c6ae8504996a9fe3c215e5/Objects/typeobject.c#L4891 + # Note that starting with Python 3.11, this `__getstate__` is always defined and thus + # the else branch will never be taken. + getstate_fn = getattr(obj, "__getstate__", None) + if getstate_fn: + state = getstate_fn() + else: + slots_to_save = copyreg._slotnames(obj.__class__) # type: ignore[attr-defined] + if slots_to_save: + state = ( + obj.__dict__, + { + name: getattr(obj, name) + for name in slots_to_save + if hasattr(obj, name) + }, + ) + else: + state = obj.__dict__ + + return state + + +def _set_obj_state(obj, state): + if isinstance(state, tuple): + if not len(state) == 2: + raise RuntimeError(f"Invalid serialized state: {state}") + dict_state = state[0] + slots_state = state[1] + else: + dict_state = state + slots_state = None + + # Starting with Python 3.11, the __dict__ attribute is lazily created + # and is serialized as None when not needed. + if dict_state: + for k, v in dict_state.items(): + setattr(obj, k, v) + + if slots_state: + for k, v in slots_state.items(): + setattr(obj, k, v) + return obj + + +def _import_dotted_name(name): + components = name.split(".") + obj = __import__(components[0]) + for component in components[1:]: + obj = getattr(obj, component) + return obj + + +def _flatten_dense_tensors(tensors): + """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of + same dense type. + + Since inputs are dense, the resulting tensor will be a concatenated 1D + buffer. Element-wise operation on this buffer will be equivalent to + operating individually. + + Args: + tensors (Iterable[Tensor]): dense tensors to flatten. + + Returns: + A contiguous 1D buffer containing input tensors. + """ + return torch._C._nn.flatten_dense_tensors(tensors) + + +def _flatten_sparse_tensors(tensors): + """Flatten sparse tensors into two contiguous 1D buffers, one of indices and + one of values. Assume tensors are of same sparse type. + + Args: + tensors (Iterable[Tensor]): sparse tensors to flatten. + + Returns: + A tuple of two contiguous 1D buffers, one containing input tensors' + indices and the other containing the values. + """ + flat_indices = torch._C._nn.flatten_dense_tensors( + [torch.Tensor._indices(t) for t in tensors] + ) + flat_values = torch._C._nn.flatten_dense_tensors( + [torch.Tensor._values(t) for t in tensors] + ) + return flat_indices, flat_values + + +def _unflatten_dense_tensors(flat, tensors): + """View a flat buffer using the sizes of tensors. Assume that tensors are of + same dense type, and that flat is given by _flatten_dense_tensors. + + Args: + flat (Tensor): flattened dense tensors to unflatten. + tensors (Iterable[Tensor]): dense tensors whose sizes will be used to + unflatten flat. + + Returns: + Unflattened dense tensors with sizes same as tensors and values from + flat. + """ + return torch._C._nn.unflatten_dense_tensors(flat, tensors) + + +def _unflatten_sparse_tensors(flat, tensors): + """View flat buffer (containing indices and values) using the sizes of + tensors. Assume that tensors are of same sparse type, and that flat is given + by _flatten_sparse_tensors. + + Args: + flat (tuple(Tensor, Tensor)): flattened indices and values of sparse + tensors to unflatten. + tensors (Iterable[Tensor]): sparse tensors whose sizes will be used to + unflatten flat. + + Returns: + Unflattened sparse tensors with sizes same as tensors and values from + flat. + """ + flat_indices, flat_values = flat + indices = torch._C._nn.unflatten_dense_tensors( + flat_indices, [torch.Tensor._indices(t) for t in tensors] + ) + values = torch._C._nn.unflatten_dense_tensors( + flat_values, [torch.Tensor._values(t) for t in tensors] + ) + outputs = [] + for t, i, v in zip(tensors, indices, values): + outputs.append(t.new(i, v, t.size())) + return tuple(outputs) + + +def _reorder_tensors_as(tensors, ordered_tensors): + """Assume that tensors are of same order as ordered_tensors within their + types, e.g., from _take_tensors. Reorder them to be of same order as + ordered_tensors. + + Args: + tensors (Iterable[Tensor]): tensors to be reordered. They should be of + the same order as ordered_tensors within their own types. + ordered_tensors (Iterable[Tensor]): tensors whose order will be the + reference. + + Returns: + Ordered tuple of tensors with contents from tensors and order of + ordered_tensors. + """ + type_dict = defaultdict(list) + for tensor in tensors: + type_dict[tensor.type()].append(tensor) + type_dict_ = {t: iter(coll) for t, coll in type_dict.items()} + return tuple(next(type_dict_[tensor.type()]) for tensor in ordered_tensors) + + +def _take_tensors(tensors, size_limit): + """Group tensors into chunks. This generator yields a chunk at each time, + each containing tensors of same type up to certain byte limit in total size. + + Args: + tensors (Sequence): A sequence of tensors to be separated into chunks. + size_limit (int): The limit of each chunk in bytes. + + Yields: + Blocks of tensors of same type and within size_limit. The yielded + tensors are only ordered as the original sequence within its types. + """ + buf_dict: defaultdict[str, list] = defaultdict(lambda: [[], 0]) + for tensor in tensors: + t = tensor.type() + if tensor.is_sparse: + indices = torch.Tensor._indices(tensor) + values = torch.Tensor._values(tensor) + size = ( + indices.numel() * indices.element_size() + + values.numel() * values.element_size() + ) + else: + size = tensor.numel() * tensor.element_size() + buf_and_size = buf_dict[t] + if buf_and_size[1] + size > size_limit and buf_and_size[1] > 0: + yield buf_and_size[0] + buf_and_size = buf_dict[t] = [[], 0] + buf_and_size[0].append(tensor) + buf_and_size[1] += size + for buf, _ in buf_dict.values(): + if len(buf) > 0: + yield buf + + +# annotation decorator to get annotations in a way that is compatible +# with both Python 2 and 3 +def annotate(ret, **kwargs): + def dec(fun): + fun.__annotations__ = dict(kwargs) + fun.__annotations__["return"] = ret + return fun + + return dec + + +def render_call(fn, args, kwargs): + str_fn = torch.overrides.resolve_name(fn) + if str_fn is None: + str_fn = str(fn) + + str_args: list[str] = [] + with torch._tensor_str.printoptions(threshold=0, edgeitems=0): + str_args.extend(repr(a) for a in args) + str_args.extend(f"{k}={repr(v)}" for k, v in kwargs.items()) + r = f"{str_fn}({', '.join(str_args)})" + return r + + +# NOTE [ Python Traceback Reference Cycle Problem ] +# +# When using sys.exc_info(), it is important to **not** store the exc_info[2], +# which is the traceback, because otherwise you will run into the traceback +# reference cycle problem, i.e., the traceback holding reference to the frame, +# and the frame (which holds reference to all the object in its temporary scope) +# holding reference the traceback. + + +class KeyErrorMessage(str): + r"""str subclass that returns itself in repr""" + + __slots__ = () + + def __repr__(self): + return self + + +class ExceptionWrapper: + r"""Wraps an exception plus traceback to communicate across threads""" + + def __init__(self, exc_info=None, where="in background"): + # It is important that we don't store exc_info, see + # NOTE [ Python Traceback Reference Cycle Problem ] + if exc_info is None: + exc_info = sys.exc_info() + self.exc_type = exc_info[0] + self.exc_msg = "".join(traceback.format_exception(*exc_info)) + self.where = where + + def reraise(self): + r"""Reraises the wrapped exception in the current thread""" + # Format a message such as: "Caught ValueError in DataLoader worker + # process 2. Original Traceback:", followed by the traceback. + msg = f"Caught {self.exc_type.__name__} {self.where}.\nOriginal {self.exc_msg}" + if self.exc_type == KeyError: + # KeyError calls repr() on its argument (usually a dict key). This + # makes stack traces unreadable. It will not be changed in Python + # (https://bugs.python.org/issue2651), so we work around it. + msg = KeyErrorMessage(msg) + elif getattr(self.exc_type, "message", None): + # Some exceptions have first argument as non-str but explicitly + # have message field + raise self.exc_type(message=msg) + try: + exception = self.exc_type(msg) + except Exception: + # If the exception takes multiple arguments or otherwise can't + # be constructed, don't try to instantiate since we don't know how to + raise RuntimeError(msg) from None + raise exception + + +def _get_available_device_type(): + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + if hasattr(torch, "xpu") and torch.xpu.is_available(): # type: ignore[attr-defined] + return "xpu" + if hasattr(torch, "mtia") and torch.mtia.is_available(): + return "mtia" + custom_backend_name = torch._C._get_privateuse1_backend_name() + custom_device_mod = getattr(torch, custom_backend_name, None) + if custom_device_mod and custom_device_mod.is_available(): + return custom_backend_name + # add more available device types here + return None + + +def _get_device_attr(get_member): + device_type = _get_available_device_type() + if device_type and device_type.lower() == "cuda": + return get_member(torch.cuda) + if device_type and device_type.lower() == "mps": + return get_member(torch.mps) + if device_type and device_type.lower() == "xpu": + return get_member(torch.xpu) # type: ignore[attr-defined] + if device_type and device_type.lower() == "mtia": + return get_member(torch.mtia) + if device_type == torch._C._get_privateuse1_backend_name(): + return get_member(getattr(torch, device_type)) + # add more available device types here + return None + + +def _get_current_device_index(): + # current device index + return _get_device_attr(lambda m: m.current_device()) + + +def _get_all_device_indices(): + # all device index + return _get_device_attr(lambda m: list(range(m.device_count()))) + + +def _get_devices_properties(device_ids): + # all device properties + return [_get_device_attr(lambda m: m.get_device_properties(i)) for i in device_ids] + + +def get_current_device_index() -> int: + r"""Checks if there are CUDA devices available and + returns the device index of the current default CUDA device. + Returns -1 in case there are no CUDA devices available. + Arguments: ``None`` + """ + if torch.cuda.device_count() > 0: + return torch.cuda.current_device() + return -1 + + +def _get_device_index( + device: Any, + optional: bool = False, + allow_cpu: bool = False, +) -> int: + r"""Gets the device index from :attr:`device`, which can be a torch.device + object, a Python integer, or ``None``. + + If :attr:`device` is a torch.device object, returns the device index if it + has index. Note that for a device without a specified index, + i.e., ``torch.device('xxx')``, this will return the current default + device of that type if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``, + CPU devices will be accepted and ``-1`` will be returned in this case. + + If :attr:`device` is a Python integer, it is returned as is. + + If :attr:`device` is ``None``, this will return the current default + device of the supported runtime platform if :attr:`optional` is ``True``. + i.e., the current default CUDA device will be returned if CUDA runtime is supported. + """ + if isinstance(device, str): + device = torch.device(device) + device_idx: Optional[int] = None + if isinstance(device, torch.device): + if not allow_cpu and device.type == "cpu": + raise ValueError(f"Expected a non cpu device, but got: {device}") + device_idx = -1 if device.type == "cpu" else device.index + if isinstance(device, int): + device_idx = device + if device_idx is None: + if optional: + # The eager API _get_current_device_index uses `lambda` functions which are + # not supported in JIT and hence not scriptable. The JIT equivalent API to get + # the current device index is `get_current_device_index()` which can + # be scripted. We use is_scripting to check the mode we are in and call the + # appropriate API. + if torch.jit.is_scripting(): + device_idx = get_current_device_index() + else: + device_idx = _get_current_device_index() + else: + raise ValueError( + f"Expected a torch.device with a specified index or an integer, but got:{device}" + ) + return device_idx + + +def _handle_complex(tensor): + """ + Returns a real view of a tensor if complex dtype else just the tensor + need to check if a UninitializedParameter because otherwise checking is_complex is an error for a LazyModule + """ + return ( + torch.view_as_real(tensor) + if not isinstance(tensor, torch.nn.UninitializedParameter) + and tensor.is_complex() + else tensor + ) + + +def _element_size(dtype): + """ + Returns the element size for a dtype, in bytes + """ + if not isinstance(dtype, torch.dtype): + raise RuntimeError(f"expected torch.dtype, but got {type(dtype)}") + + if dtype.is_complex: + return torch.finfo(dtype).bits >> 2 + elif dtype.is_floating_point: + return torch.finfo(dtype).bits >> 3 + elif dtype == torch.bool: + # NOTE: torch.bool is not supported in torch.iinfo() + return 1 + else: + return torch.iinfo(dtype).bits >> 3 + + +class _ClassPropertyDescriptor: + def __init__(self, fget, fset=None): + self.fget = fget + + def __get__(self, instance, owner=None): + if owner is None: + owner = type(instance) + return self.fget.__get__(instance, owner)() + + +def classproperty(func): + if not isinstance(func, (classmethod, staticmethod)): + func = classmethod(func) + return _ClassPropertyDescriptor(func) + + +if TYPE_CHECKING: + # TorchScript does not support `@deprecated` + # This is a workaround to avoid breaking TorchScript + @deprecated( + "`torch._utils.is_compiling` is deprecated. Use `torch.compiler.is_compiling` instead.", + category=FutureWarning, + ) + def is_compiling() -> bool: + return torch.compiler.is_compiling() + +else: + + def is_compiling() -> bool: + """ + Indicates whether we are tracing/compiling with torch.compile() or torch.export(). + """ + warnings.warn( # use `warnings.warn` instead of `@deprecated` + "`torch._utils.is_compiling` is deprecated. Use `torch.compiler.is_compiling` instead.", + # FutureWarning, # TorchScript does not support Warning type + stacklevel=2, + ) + return torch.compiler.is_compiling() + + +def _functionalize_sync(t): + # This code lives in python instead of C++ since conditioning on a certain python subclass + # is much more of a pain in C++. + from torch._subclasses.functional_tensor import FunctionalTensor + + if isinstance(t, FunctionalTensor): + # If a FunctionalTensorMode is active while syncing, we don't want it to intercept any ops that get called + # when we sync our inner tensor. + # Why? + # (1) If there are input mutations in the graph, then they will be re-applied during + # AOTAutograd when we call _sync() from inside of our functionalization kernels. + # (2) _sync() causes us to regenerate our updated the tensor from the updated base, + # which dispatches to a bunch of view ops + # (3) The input to these view ops is our inner FunctionalTensorWrapper + # (since the sync was called from C++), not the python FunctionalTensor + # (4) if a python FunctionalTensorMode is active, it will complain when it intercepts + # the view op, since it will see an input that is a C++ FunctionalTensorWrapper + # (aka a normal torch.Tensor) instead of a python `FunctionalTensor). + maybe_functional_mode = torch._C._unset_dispatch_mode( + torch._C._TorchDispatchModeKey.FUNCTIONAL + ) + try: + torch._functionalize_sync(t.elem) # type: ignore[attr-defined] + finally: + if maybe_functional_mode is not None: + torch._C._set_dispatch_mode(maybe_functional_mode) + else: + torch._functionalize_sync(t) # type: ignore[attr-defined] + + +@functools.lru_cache(2) +def _get_device_module(device_type: str): + device_module = getattr(torch, device_type, None) + if device_module is None: + raise RuntimeError( + f"Device '{device_type}' does not have a corresponding module registered as 'torch.{device_type}'." + ) + return device_module + + +def _dummy_type(name: str) -> type: + def get_err_fn(is_init: bool): + def err_fn(obj, *args, **kwargs): + if is_init: + class_name = obj.__class__.__name__ + else: + class_name = obj.__name__ + raise RuntimeError(f"Tried to instantiate dummy base class {class_name}") + + return err_fn + + return type( + name, (object,), {"__init__": get_err_fn(True), "__new__": get_err_fn(False)} + ) + + +class _LazySeedTracker: + # Since seeding is memory-less, only track the latest seed. + # Note: `manual_seed_all` followed by `manual_seed` overwrites + # the seed on current device. We track the order of **latest** + # calls between these two API. + def __init__(self): + self.manual_seed_all_cb = None + self.manual_seed_cb = None + self.call_order = [] + + def queue_seed_all(self, cb, traceback): + self.manual_seed_all_cb = (cb, traceback) + # update seed_all to be latest + self.call_order = [self.manual_seed_cb, self.manual_seed_all_cb] + + def queue_seed(self, cb, traceback): + self.manual_seed_cb = (cb, traceback) + # update seed to be latest + self.call_order = [self.manual_seed_all_cb, self.manual_seed_cb] + + def get_calls(self) -> list: + return self.call_order + + +logger = logging.getLogger(__name__) +P = ParamSpec("P") + + +class CallbackRegistry(Generic[P]): + def __init__(self, name: str): + self.name = name + self.callback_list: list[Callable[P, None]] = [] + + def add_callback(self, cb: Callable[P, None]) -> None: + self.callback_list.append(cb) + + def fire_callbacks(self, *args: P.args, **kwargs: P.kwargs) -> None: + for cb in self.callback_list: + try: + cb(*args, **kwargs) + except Exception: + logger.exception( + "Exception in callback for %s registered with gpu trace", self.name + ) + + +def try_import(module_name: str) -> Optional[ModuleType]: + # Implementation based on + # https://docs.python.org/3/library/importlib.html#checking-if-a-module-can-be-imported + if (module := sys.modules.get(module_name, None)) is not None: + return module + + if (spec := importlib.util.find_spec(module_name)) is not None: + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + + # https://docs.python.org/3/library/importlib.html#importlib.machinery.ModuleSpec.loader + # "The finder should always set this attribute" + assert spec.loader is not None, "The loader attribute should always be set" + spec.loader.exec_module(module) + return module + + return None + + +# IMPORT_MAPPING and NAME_MAPPING are adapted from https://github.com/python/cpython/blob/main/Lib/_compat_pickle.py +# for use in the weights_only Unpickler. + +IMPORT_MAPPING = { + "__builtin__": "builtins", + "copy_reg": "copyreg", + "Queue": "queue", + "repr": "reprlib", + "_abcoll": "collections.abc", + # Non-mutual mappings. + "UserDict": "collections", + "UserList": "collections", + "UserString": "collections", + "whichdb": "dbm", + "StringIO": "io", + "cStringIO": "io", +} + + +# This contains rename rules that are easy to handle. We ignore the more +# complex stuff (e.g. mapping the names in the urllib and types modules). +# These rules should be run before import names are fixed. +NAME_MAPPING = { + ("__builtin__", "xrange"): ("builtins", "range"), + ("__builtin__", "reduce"): ("functools", "reduce"), + ("__builtin__", "intern"): ("sys", "intern"), + ("__builtin__", "unichr"): ("builtins", "chr"), + ("__builtin__", "unicode"): ("builtins", "str"), + ("__builtin__", "long"): ("builtins", "int"), + ("itertools", "izip"): ("builtins", "zip"), + ("itertools", "imap"): ("builtins", "map"), + ("itertools", "ifilter"): ("builtins", "filter"), + ("itertools", "ifilterfalse"): ("itertools", "filterfalse"), + ("itertools", "izip_longest"): ("itertools", "zip_longest"), + ("UserDict", "IterableUserDict"): ("collections", "UserDict"), + ("UserList", "UserList"): ("collections", "UserList"), + ("UserString", "UserString"): ("collections", "UserString"), + # Non-mutual mappings. + ("__builtin__", "basestring"): ("builtins", "str"), + ("exceptions", "StandardError"): ("builtins", "Exception"), + ("UserDict", "UserDict"): ("collections", "UserDict"), +} diff --git a/venv/lib/python3.10/site-packages/torch/_utils_internal.py b/venv/lib/python3.10/site-packages/torch/_utils_internal.py new file mode 100644 index 0000000000000000000000000000000000000000..f20a88ce85402a7531a8622b2f6a68aa8fada9e1 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_utils_internal.py @@ -0,0 +1,367 @@ +# mypy: allow-untyped-defs +import functools +import logging +import os +import sys +import tempfile +import typing_extensions +from typing import Any, Callable, Optional, TypeVar +from typing_extensions import ParamSpec + +import torch +from torch._strobelight.compile_time_profiler import StrobelightCompileTimeProfiler + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +log = logging.getLogger(__name__) + +if os.environ.get("TORCH_COMPILE_STROBELIGHT", False): + import shutil + + if not shutil.which("strobeclient"): + log.info( + "TORCH_COMPILE_STROBELIGHT is true, but seems like you are not on a FB machine." + ) + else: + log.info("Strobelight profiler is enabled via environment variable") + StrobelightCompileTimeProfiler.enable() + +# this arbitrary-looking assortment of functionality is provided here +# to have a central place for overridable behavior. The motivating +# use is the FB build environment, where this source file is replaced +# by an equivalent. + +if os.path.basename(os.path.dirname(__file__)) == "shared": + torch_parent = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) +else: + torch_parent = os.path.dirname(os.path.dirname(__file__)) + + +def get_file_path(*path_components: str) -> str: + return os.path.join(torch_parent, *path_components) + + +def get_file_path_2(*path_components: str) -> str: + return os.path.join(*path_components) + + +def get_writable_path(path: str) -> str: + if os.access(path, os.W_OK): + return path + return tempfile.mkdtemp(suffix=os.path.basename(path)) + + +def prepare_multiprocessing_environment(path: str) -> None: + pass + + +def resolve_library_path(path: str) -> str: + return os.path.realpath(path) + + +def throw_abstract_impl_not_imported_error(opname, module, context): + if module in sys.modules: + raise NotImplementedError( + f"{opname}: We could not find the fake impl for this operator. " + ) + else: + raise NotImplementedError( + f"{opname}: We could not find the fake impl for this operator. " + f"The operator specified that you may need to import the '{module}' " + f"Python module to load the fake impl. {context}" + ) + + +# NB! This treats "skip" kwarg specially!! +def compile_time_strobelight_meta( + phase_name: str, +) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + def compile_time_strobelight_meta_inner( + function: Callable[_P, _T], + ) -> Callable[_P, _T]: + @functools.wraps(function) + def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> _T: + if "skip" in kwargs and isinstance(skip := kwargs["skip"], int): + kwargs["skip"] = skip + 1 + + # This is not needed but we have it here to avoid having profile_compile_time + # in stack traces when profiling is not enabled. + if not StrobelightCompileTimeProfiler.enabled: + return function(*args, **kwargs) + + return StrobelightCompileTimeProfiler.profile_compile_time( + function, phase_name, *args, **kwargs + ) + + return wrapper_function + + return compile_time_strobelight_meta_inner + + +# Meta only, see +# https://www.internalfb.com/intern/wiki/ML_Workflow_Observability/User_Guides/Adding_instrumentation_to_your_code/ +# +# This will cause an event to get logged to Scuba via the signposts API. You +# can view samples on the API at https://fburl.com/scuba/workflow_signpost/zh9wmpqs +# we log to subsystem "torch", and the category and name you provide here. +# Each of the arguments translate into a Scuba column. We're still figuring +# out local conventions in PyTorch, but category should be something like +# "dynamo" or "inductor", and name should be a specific string describing what +# kind of event happened. +# +# Killswitch is at +# https://www.internalfb.com/intern/justknobs/?name=pytorch%2Fsignpost#event +def signpost_event(category: str, name: str, parameters: dict[str, Any]): + log.info("%s %s: %r", category, name, parameters) + + +def add_mlhub_insight(category: str, insight: str, insight_description: str): + pass + + +def log_compilation_event(metrics): + log.info("%s", metrics) + + +def upload_graph(graph): + pass + + +def set_pytorch_distributed_envs_from_justknobs(): + pass + + +def log_export_usage(**kwargs): + pass + + +def log_draft_export_usage(**kwargs): + pass + + +def log_trace_structured_event(*args, **kwargs) -> None: + pass + + +def log_cache_bypass(*args, **kwargs) -> None: + pass + + +def log_torchscript_usage(api: str, **kwargs): + _ = api + return + + +def check_if_torch_exportable(): + return False + + +def export_training_ir_rollout_check() -> bool: + return True + + +def full_aoti_runtime_assert() -> bool: + return True + + +def log_torch_jit_trace_exportability( + api: str, + type_of_export: str, + export_outcome: str, + result: str, +): + _, _, _, _ = api, type_of_export, export_outcome, result + return + + +def justknobs_check(name: str, default: bool = True) -> bool: + """ + This function can be used to killswitch functionality in FB prod, + where you can toggle this value to False in JK without having to + do a code push. In OSS, we always have everything turned on all + the time, because downstream users can simply choose to not update + PyTorch. (If more fine-grained enable/disable is needed, we could + potentially have a map we lookup name in to toggle behavior. But + the point is that it's all tied to source code in OSS, since there's + no live server to query.) + + This is the bare minimum functionality I needed to do some killswitches. + We have a more detailed plan at + https://docs.google.com/document/d/1Ukerh9_42SeGh89J-tGtecpHBPwGlkQ043pddkKb3PU/edit + In particular, in some circumstances it may be necessary to read in + a knob once at process start, and then use it consistently for the + rest of the process. Future functionality will codify these patterns + into a better high level API. + + WARNING: Do NOT call this function at module import time, JK is not + fork safe and you will break anyone who forks the process and then + hits JK again. + """ + return default + + +def justknobs_getval_int(name: str) -> int: + """ + Read warning on justknobs_check + """ + return 0 + + +def is_fb_unit_test() -> bool: + return False + + +@functools.cache +def max_clock_rate(): + """ + unit: MHz + """ + if not torch.version.hip: + from triton.testing import nvsmi + + return nvsmi(["clocks.max.sm"])[0] + else: + # Manually set max-clock speeds on ROCm until equivalent nvmsi + # functionality in triton.testing or via pyamdsmi enablement. Required + # for test_snode_runtime unit tests. + gcn_arch = str(torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0]) + if "gfx94" in gcn_arch: + return 1700 + elif "gfx90a" in gcn_arch: + return 1700 + elif "gfx908" in gcn_arch: + return 1502 + elif "gfx12" in gcn_arch: + return 1700 + elif "gfx11" in gcn_arch: + return 1700 + elif "gfx103" in gcn_arch: + return 1967 + elif "gfx101" in gcn_arch: + return 1144 + elif "gfx95" in gcn_arch: + return 1700 # TODO: placeholder, get actual value + else: + return 1100 + + +def get_mast_job_name_version() -> Optional[tuple[str, int]]: + return None + + +TEST_MASTER_ADDR = "127.0.0.1" +TEST_MASTER_PORT = 29500 +# USE_GLOBAL_DEPS controls whether __init__.py tries to load +# libtorch_global_deps, see Note [Global dependencies] +USE_GLOBAL_DEPS = True +# USE_RTLD_GLOBAL_WITH_LIBTORCH controls whether __init__.py tries to load +# _C.so with RTLD_GLOBAL during the call to dlopen. +USE_RTLD_GLOBAL_WITH_LIBTORCH = False +# If an op was defined in C++ and extended from Python using the +# torch.library.register_fake, returns if we require that there be a +# m.set_python_module("mylib.ops") call from C++ that associates +# the C++ op with a python module. +REQUIRES_SET_PYTHON_MODULE = False + + +def maybe_upload_prof_stats_to_manifold(profile_path: str) -> Optional[str]: + print("Uploading profile stats (fb-only otherwise no-op)") + return None + + +def log_chromium_event_internal( + event: dict[str, Any], + stack: list[str], + logger_uuid: str, + start_time_ns: int, +): + return None + + +def record_chromium_event_internal( + event: dict[str, Any], +): + return None + + +def profiler_allow_cudagraph_cupti_lazy_reinit_cuda12(): + return True + + +def deprecated(): + """ + When we deprecate a function that might still be in use, we make it internal + by adding a leading underscore. This decorator is used with a private function, + and creates a public alias without the leading underscore, but has a deprecation + warning. This tells users "THIS FUNCTION IS DEPRECATED, please use something else" + without breaking them, however, if they still really really want to use the + deprecated function without the warning, they can do so by using the internal + function name. + """ + + def decorator(func: Callable[_P, _T]) -> Callable[_P, _T]: + # Validate naming convention – single leading underscore, not dunder + if not (func.__name__.startswith("_")): + raise ValueError( + "@deprecate must decorate a function whose name " + "starts with a single leading underscore (e.g. '_foo') as the api should be considered internal for deprecation." + ) + + public_name = func.__name__[1:] # drop exactly one leading underscore + module = sys.modules[func.__module__] + + # Don't clobber an existing symbol accidentally. + if hasattr(module, public_name): + raise RuntimeError( + f"Cannot create alias '{public_name}' -> symbol already exists in {module.__name__}. \ + Please rename it or consult a pytorch developer on what to do" + ) + + warning_msg = f"{func.__name__[1:]} is DEPRECATED, please consider using an alternative API(s). " + + # public deprecated alias + alias = typing_extensions.deprecated( + warning_msg, category=UserWarning, stacklevel=1 + )(func) + + alias.__name__ = public_name + + # Adjust qualname if nested inside a class or another function + if "." in func.__qualname__: + alias.__qualname__ = func.__qualname__.rsplit(".", 1)[0] + "." + public_name + else: + alias.__qualname__ = public_name + + setattr(module, public_name, alias) + + return func + + return decorator + + +def get_default_numa_options(): + """ + When using elastic agent, if no numa options are provided, we will use these + as the default. + + For external use cases, we return None, i.e. no numa binding. If you would like + to use torch's automatic numa binding capabilities, you should provide + NumaOptions to your launch config directly or use the numa binding option + available in torchrun. + + Must return None or NumaOptions, but not specifying to avoid circular import. + """ + return None + + +def log_triton_builds(fail: Optional[str]): + pass + + +def find_compile_subproc_binary() -> Optional[str]: + """ + Allows overriding the binary used for subprocesses + """ + return None diff --git a/venv/lib/python3.10/site-packages/torch/_vmap_internals.py b/venv/lib/python3.10/site-packages/torch/_vmap_internals.py new file mode 100644 index 0000000000000000000000000000000000000000..6baee77ade54d2602be2eff99fa300d65661fd31 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_vmap_internals.py @@ -0,0 +1,245 @@ +# mypy: allow-untyped-defs +import functools +from typing import Any, Callable, Optional, Union +from typing_extensions import deprecated + +import torch +from torch import Tensor +from torch.utils._pytree import _broadcast_to_and_flatten, tree_flatten, tree_unflatten + + +in_dims_t = Union[int, tuple] +out_dims_t = Union[int, tuple[int, ...]] + + +# Checks that all args-to-be-batched have the same batch dim size +def _validate_and_get_batch_size( + flat_in_dims: list[Optional[int]], + flat_args: list, +) -> int: + batch_sizes = [ + arg.size(in_dim) + for in_dim, arg in zip(flat_in_dims, flat_args) + if in_dim is not None + ] + if batch_sizes and any(size != batch_sizes[0] for size in batch_sizes): + raise ValueError( + f"vmap: Expected all tensors to have the same size in the mapped " + f"dimension, got sizes {batch_sizes} for the mapped dimension" + ) + return batch_sizes[0] + + +def _num_outputs(batched_outputs: Union[Tensor, tuple[Tensor, ...]]) -> int: + if isinstance(batched_outputs, tuple): + return len(batched_outputs) + return 1 + + +# If value is a tuple, check it has length `num_elements`. +# If value is not a tuple, make a tuple with `value` repeated `num_elements` times +def _as_tuple( + value: Any, + num_elements: int, + error_message_lambda: Callable[[], str], +) -> tuple: + if not isinstance(value, tuple): + return (value,) * num_elements + if len(value) != num_elements: + raise ValueError(error_message_lambda()) + return value + + +# Creates BatchedTensors for every Tensor in arg that should be batched. +# Returns the (potentially) batched arguments and the batch_size. +def _create_batched_inputs( + in_dims: in_dims_t, + args: tuple, + vmap_level: int, + func: Callable, +) -> tuple[tuple, int]: + if not isinstance(in_dims, int) and not isinstance(in_dims, tuple): + raise ValueError( + f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(): " + f"expected `in_dims` to be int or a (potentially nested) tuple " + f"matching the structure of inputs, got: {type(in_dims)}." + ) + if len(args) == 0: + raise ValueError( + f"vmap({_get_name(func)})(): got no inputs. Maybe you forgot to add " + f"inputs, or you are trying to vmap over a function with no inputs. " + f"The latter is unsupported." + ) + + flat_args, args_spec = tree_flatten(args) + flat_in_dims = _broadcast_to_and_flatten(in_dims, args_spec) + if flat_in_dims is None: + raise ValueError( + f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(): " + f"in_dims is not compatible with the structure of `inputs`. " + f"in_dims has structure {tree_flatten(in_dims)[1]} but inputs " + f"has structure {args_spec}." + ) + + for arg, in_dim in zip(flat_args, flat_in_dims): + if not isinstance(in_dim, int) and in_dim is not None: + raise ValueError( + f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(): " + f"Got in_dim={in_dim} for an input but in_dim must be either " + f"an integer dimension or None." + ) + if isinstance(in_dim, int) and not isinstance(arg, Tensor): + raise ValueError( + f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(): " + f"Got in_dim={in_dim} for an input but the input is of type " + f"{type(arg)}. We cannot vmap over non-Tensor arguments, " + f"please use None as the respective in_dim" + ) + if in_dim is not None and (in_dim < 0 or in_dim >= arg.dim()): + raise ValueError( + f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(): " + f"Got in_dim={in_dim} for some input, but that input is a Tensor " + f"of dimensionality {arg.dim()} so expected in_dim to satisfy " + f"0 <= in_dim < {arg.dim()}." + ) + + batch_size = _validate_and_get_batch_size(flat_in_dims, flat_args) + # See NOTE [Ignored _remove_batch_dim, _add_batch_dim] + batched_inputs = [ + arg if in_dim is None else torch._add_batch_dim(arg, in_dim, vmap_level) + for in_dim, arg in zip(flat_in_dims, flat_args) + ] + return tree_unflatten(batched_inputs, args_spec), batch_size + + +# Undos the batching (and any batch dimensions) associated with the `vmap_level`. +def _unwrap_batched( + batched_outputs: Union[Tensor, tuple[Tensor, ...]], + out_dims: out_dims_t, + vmap_level: int, + batch_size: int, + func: Callable, + allow_none_pass_through: bool = False, +) -> tuple: + num_outputs = _num_outputs(batched_outputs) + out_dims_as_tuple = _as_tuple( + out_dims, + num_outputs, + lambda: f"vmap({_get_name(func)}, ..., out_dims={out_dims}): `out_dims` must " + f"have one dim per output (got {num_outputs} outputs) of {_get_name(func)}.", + ) + + # NOTE [Ignored _remove_batch_dim, _add_batch_dim] + # There is something wrong with our type bindings for functions that begin + # with '_', see #40397. + if isinstance(batched_outputs, Tensor): + out_dim = out_dims_as_tuple[0] + return torch._remove_batch_dim(batched_outputs, vmap_level, batch_size, out_dim) # type: ignore[return-value] + if allow_none_pass_through: + return tuple( + ( + torch._remove_batch_dim(out, vmap_level, batch_size, out_dim) + if out is not None + else None + ) + for out, out_dim in zip(batched_outputs, out_dims_as_tuple) + ) + else: + return tuple( + torch._remove_batch_dim(out, vmap_level, batch_size, out_dim) + for out, out_dim in zip(batched_outputs, out_dims_as_tuple) + ) + + +# Checks that `fn` returned one or more Tensors and nothing else. +# NB: A python function that return multiple arguments returns a single tuple, +# so we are effectively checking that `outputs` is a single Tensor or a tuple of +# Tensors. +def _validate_outputs(outputs: Any, func: Callable) -> None: + if isinstance(outputs, Tensor): + return + if not isinstance(outputs, tuple): + raise ValueError( + f"vmap({_get_name(func)}, ...): `{_get_name(func)}` must only return " + f"Tensors, got type {type(outputs)} as the return." + ) + for idx, output in enumerate(outputs): + if isinstance(output, Tensor): + continue + raise ValueError( + f"vmap({_get_name(func)}, ...): `{_get_name(func)}` must only return " + f"Tensors, got type {type(output)} for return {idx}." + ) + + +def _check_out_dims_is_int_or_int_tuple(out_dims: out_dims_t, func: Callable) -> None: + if isinstance(out_dims, int): + return + if not isinstance(out_dims, tuple) or not all( + isinstance(out_dim, int) for out_dim in out_dims + ): + raise ValueError( + f"vmap({_get_name(func)}, ..., out_dims={out_dims}): `out_dims` must be " + f"an int or a tuple of int representing where in the outputs the " + f"vmapped dimension should appear." + ) + + +def _get_name(func: Callable): + if hasattr(func, "__name__"): + return func.__name__ + + # Not all callables have __name__, in fact, only static functions/methods do. + # A callable created via functools.partial or an nn.Module, to name some + # examples, don't have a __name__. + return repr(func) + + +# vmap(func)(inputs) wraps all Tensor inputs to be batched in BatchedTensors, +# sends those into func, and then unwraps the output BatchedTensors. Operations +# on BatchedTensors perform the batched operations that the user is asking for. +@deprecated( + "Please use `torch.vmap` instead of `torch._vmap_internals.vmap`.", + category=FutureWarning, +) +def vmap(func: Callable, in_dims: in_dims_t = 0, out_dims: out_dims_t = 0) -> Callable: + """ + Please use torch.vmap instead of this API. + """ + return _vmap(func, in_dims, out_dims) + + +# A version of vmap but without the initial "experimental prototype" warning +def _vmap( + func: Callable, + in_dims: in_dims_t = 0, + out_dims: out_dims_t = 0, + allow_none_pass_through: bool = False, +) -> Callable: + # The `allow_none_pass_through` argument is a temporary workaround may be removed. + # Currently it enables us to wrap the call in `autograd.grad` to the autograd engine, + # which may return None if any of the inputs are unused. See the issue discussing this: + # https://github.com/pytorch/functorch/issues/159. + @functools.wraps(func) + def wrapped(*args): + _check_out_dims_is_int_or_int_tuple(out_dims, func) + vmap_level = torch._C._vmapmode_increment_nesting() + try: + batched_inputs, batch_size = _create_batched_inputs( + in_dims, args, vmap_level, func + ) + batched_outputs = func(*batched_inputs) + if not allow_none_pass_through: + _validate_outputs(batched_outputs, func) + return _unwrap_batched( + batched_outputs, + out_dims, + vmap_level, + batch_size, + func, + allow_none_pass_through=allow_none_pass_through, + ) + finally: + torch._C._vmapmode_decrement_nesting() + + return wrapped diff --git a/venv/lib/python3.10/site-packages/torch/_weights_only_unpickler.py b/venv/lib/python3.10/site-packages/torch/_weights_only_unpickler.py new file mode 100644 index 0000000000000000000000000000000000000000..9382a5500e0eea648aea60b2f97af23a2117252f --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/_weights_only_unpickler.py @@ -0,0 +1,576 @@ +# mypy: allow-untyped-defs +# Unpickler restricted to loading only state dicts +# Restrict constructing types to a list defined in _get_allowed_globals() +# Restrict BUILD operation to `Tensor`, `Parameter` and `OrderedDict` types only +# Restrict APPEND/APPENDS to `list` +# In `GLOBALS` operation do not do class lookup by name, but rather rely on dictionary +# defined by `_get_allowed_globals()` method, that contains: +# - torch types (Storage, dtypes, Tensor, `torch.Size`), +# - `torch._utils._rebuild` functions. +# - `torch.nn.Parameter` +# - `collections.Counter` +# - `collections.OrderedDict` +# Additionally, users can use an allowlist for adding classes they have deemed as safe using +# `_add_safe_globals()` (`torch.serialization.add_safe_globals`) +# `_clear_safe_globals()` (`torch.serialization.clear_safe_globals`) +# `_get_safe_globals()` (`torch.serialization.get_safe_globals`) + +# Based of https://github.com/python/cpython/blob/main/Lib/pickle.py +# Expected to be useful for loading PyTorch model weights +# For example: +# data = urllib.request.urlopen('https://download.pytorch.org/models/resnet50-0676ba61.pth').read() +# buf = io.BytesIO(data) +# weights = torch.load(buf, weights_only = True) + +import functools as _functools +import warnings + +from _codecs import encode +from collections import Counter, OrderedDict +from pickle import ( + APPEND, + APPENDS, + BINFLOAT, + BINGET, + BININT, + BININT1, + BININT2, + BINPERSID, + BINPUT, + BINUNICODE, + BUILD, + bytes_types, + decode_long, + EMPTY_DICT, + EMPTY_LIST, + EMPTY_SET, + EMPTY_TUPLE, + GLOBAL, + LONG1, + LONG_BINGET, + LONG_BINPUT, + MARK, + NEWFALSE, + NEWOBJ, + NEWTRUE, + NONE, + PROTO, + REDUCE, + SETITEM, + SETITEMS, + SHORT_BINSTRING, + STOP, + TUPLE, + TUPLE1, + TUPLE2, + TUPLE3, + UnpicklingError, +) +from struct import unpack +from sys import maxsize +from typing import Any, Callable, Union + +import torch +from torch._utils import _sparse_tensors_to_validate, IMPORT_MAPPING, NAME_MAPPING + + +# modules in this list are never allowed, even if the user attempts to allowlist +# functions/classes from them +_blocklisted_modules = [ + "sys", + "os", + "posix", + "nt", +] + +_marked_safe_globals_set: set[Union[Callable, tuple[Callable, str]]] = set() + + +def _add_safe_globals(safe_globals: list[Union[Callable, tuple[Callable, str]]]): + global _marked_safe_globals_set + _marked_safe_globals_set = _marked_safe_globals_set.union(set(safe_globals)) + + +def _get_safe_globals() -> list[Union[Callable, tuple[Callable, str]]]: + global _marked_safe_globals_set + return list(_marked_safe_globals_set) + + +def _clear_safe_globals(): + global _marked_safe_globals_set + _marked_safe_globals_set = set() + + +def _remove_safe_globals( + globals_to_remove: list[Union[Callable, tuple[Callable, str]]], +): + global _marked_safe_globals_set + _marked_safe_globals_set = _marked_safe_globals_set - set(globals_to_remove) + + +class _safe_globals: + def __init__(self, safe_globals: list[Union[Callable, tuple[Callable, str]]]): + self.safe_globals = safe_globals + + def __enter__(self): + _add_safe_globals(self.safe_globals) + + def __exit__(self, type, value, tb): + _remove_safe_globals(self.safe_globals) + + +# Separate from _get_allowed_globals because of the lru_cache on _get_allowed_globals +# For example if user had a script like +# torch.load(file_a) +# torch.serialization._add_safe_globals([torch.foo]) +# torch.load(file_b) +# the dynamic additions to safe_globals would not be picked up by +# _get_allowed_globals due to the lru_cache +def _get_user_allowed_globals(): + rc: dict[str, Any] = {} + for f in _marked_safe_globals_set: + if isinstance(f, tuple): + if len(f) != 2: + raise ValueError( + f"Expected tuple of length 2 (global, str of callable full path), but got tuple of length: {len(f)}" + ) + if type(f[1]) is not str: + raise TypeError( + f"Expected second item in tuple to be str of callable full path, but got: {type(f[1])}" + ) + f, name = f + rc[name] = f + else: + module, name = f.__module__, f.__qualname__ + rc[f"{module}.{name}"] = f + return rc + + +def _tensor_rebuild_functions(): + return { + torch._utils._rebuild_parameter, + torch._utils._rebuild_parameter_with_state, + torch._utils._rebuild_qtensor, + torch._utils._rebuild_tensor, + torch._utils._rebuild_tensor_v2, + torch._utils._rebuild_tensor_v3, + torch._utils._rebuild_sparse_tensor, + torch._utils._rebuild_meta_tensor_no_storage, + torch._utils._rebuild_nested_tensor, + torch._utils._rebuild_wrapper_subclass, + # Allowlisting this, but not allowlisting the numpy functions by default + # Reasoning is that we don't have control over the numpy functions, but + # this utility is provided by pytorch + torch._utils._rebuild_device_tensor_from_numpy, + # In 2.6, we should no longer have a dependency on numpy and the above + # _rebuild_device_tensor_from_numpy function. + torch._utils._rebuild_device_tensor_from_cpu_tensor, + } + + +# Unpickling machinery +@_functools.lru_cache(maxsize=1) +def _get_allowed_globals(): + rc: dict[str, Any] = { + "collections.OrderedDict": OrderedDict, + "collections.Counter": Counter, + "torch.nn.parameter.Parameter": torch.nn.Parameter, + "torch.serialization._get_layout": torch.serialization._get_layout, + "torch.Size": torch.Size, + "torch.Tensor": torch.Tensor, + "torch.device": torch.device, + "_codecs.encode": encode, # for bytes + "builtins.bytearray": bytearray, # for bytearray + "builtins.set": set, # for set + "builtins.complex": complex, # for complex + } + + # dtype + for t in torch.storage._dtype_to_storage_type_map().keys(): + rc[str(t)] = t + for t in torch.storage._new_dtypes(): + rc[str(t)] = t + for t in [getattr(torch, f"uint{x}") for x in range(1, 8)]: + rc[str(t)] = t + for t in [getattr(torch, f"int{x}") for x in range(1, 8)]: + rc[str(t)] = t + + # Tensor classes + for tt in torch._tensor_classes: + rc[f"{tt.__module__}.{tt.__name__}"] = tt + # Storage classes + for ts in torch._storage_classes: + if ts not in (torch.storage.TypedStorage, torch.storage.UntypedStorage): + # Wrap legacy storage types in a dummy class + rc[f"{ts.__module__}.{ts.__name__}"] = torch.serialization.StorageType( + ts.__name__ + ) + else: + rc[f"{ts.__module__}.{ts.__name__}"] = ts + # Quantization specific + for qt in [ + torch.per_tensor_affine, + torch.per_tensor_symmetric, + torch.per_channel_affine, + torch.per_channel_symmetric, + torch.per_channel_affine_float_qparams, + ]: + rc[str(qt)] = qt + # Rebuild functions + for f in _tensor_rebuild_functions(): + rc[f"torch._utils.{f.__name__}"] = f + + # Handles Tensor Subclasses, Tensor's with attributes. + # NOTE: It calls into above rebuild functions for regular Tensor types. + rc["torch._tensor._rebuild_from_type_v2"] = torch._tensor._rebuild_from_type_v2 + return rc + + +def _read_global_instruction(readline: Callable) -> tuple[str, str]: + module = readline()[:-1].decode("utf-8") + name = readline()[:-1].decode("utf-8") + # Patch since torch.save default protocol is 2 + # users will be running this code in python > 3 + if (module, name) in NAME_MAPPING: + module, name = NAME_MAPPING[(module, name)] + elif module in IMPORT_MAPPING: + module = IMPORT_MAPPING[module] + return module, name + + +def get_globals_in_pkl(file) -> set[str]: + globals_in_checkpoint = set() + read = file.read + readline = file.readline + op_to_bytes_to_read = { + NEWOBJ[0]: 0, + REDUCE[0]: 0, + BUILD[0]: 0, + APPEND[0]: 0, + APPENDS[0]: 0, + SETITEM[0]: 0, + SETITEMS[0]: 0, + MARK[0]: 0, + TUPLE[0]: 0, + TUPLE1[0]: 0, + TUPLE2[0]: 0, + TUPLE3[0]: 0, + NONE[0]: 0, + NEWFALSE[0]: 0, + NEWTRUE[0]: 0, + EMPTY_TUPLE[0]: 0, + EMPTY_LIST[0]: 0, + EMPTY_DICT[0]: 0, + EMPTY_SET[0]: 0, + BINPERSID[0]: 0, + BININT[0]: 4, + BININT1[0]: 1, + BININT2[0]: 2, + BINFLOAT[0]: 8, + BINGET[0]: 1, + LONG_BINGET[0]: 4, + BINPUT[0]: 1, + LONG_BINPUT[0]: 4, + } + while True: + key = read(1) + if not key: + raise EOFError + assert isinstance(key, bytes_types) + if key[0] == GLOBAL[0]: + module, name = _read_global_instruction(readline) + globals_in_checkpoint.add(f"{module}.{name}") + elif key[0] in op_to_bytes_to_read: + bytes_to_read = op_to_bytes_to_read[key[0]] + if bytes_to_read: + read(bytes_to_read) + # ops where bytes to read depends on the data + elif key[0] == BINUNICODE[0]: + strlen = unpack(" maxsize: + raise UnpicklingError("String is too long") + read(strlen) + elif key[0] in {SHORT_BINSTRING[0], LONG1[0]}: + strlen = read(1)[0] + read(strlen) + # first and last op + elif key[0] == PROTO[0]: + read(1)[0] + elif key[0] == STOP[0]: + return globals_in_checkpoint + else: + raise UnpicklingError(f"Unsupported operand {key[0]}") + + +class Unpickler: + def __init__(self, file, *, encoding: str = "bytes"): + self.encoding = encoding + self.readline = file.readline + self.read = file.read + self.memo: dict[int, Any] = {} + self.proto: int = -1 + + def load(self): + """Read a pickled object representation from the open file. + + Return the reconstituted object hierarchy specified in the file. + """ + self.metastack = [] + self.stack: list[Any] = [] + self.append = self.stack.append + read = self.read + while True: + key = read(1) + if not key: + raise EOFError + assert isinstance(key, bytes_types) + # Risky operators + if key[0] == GLOBAL[0]: + module, name = _read_global_instruction(self.readline) + full_path = f"{module}.{name}" + if module in _blocklisted_modules: + raise UnpicklingError( + f"Trying to load unsupported GLOBAL {full_path} whose module {module} is blocked." + ) + if full_path in _get_allowed_globals(): + self.append(_get_allowed_globals()[full_path]) + elif full_path in _get_user_allowed_globals(): + self.append(_get_user_allowed_globals()[full_path]) + elif full_path in ( + [ + "torch.nested._internal.nested_tensor.NestedTensor", + "torch.nested._internal.nested_tensor._rebuild_njt", + "torch._dynamo.decorators._DimRange", + ] + ): + raise UnpicklingError( + "``torch.nested`` and ``torch._dynamo`` must be imported to load nested jagged tensors (NJTs)" + ) + elif full_path in ( + [ + "torch.distributed.device_mesh.DeviceMesh", + "torch.distributed.tensor._dtensor_spec.DTensorSpec", + "torch.distributed.tensor._dtensor_spec.TensorMeta", + "torch.distributed.tensor.DTensor", + "torch.distributed.tensor.placement_types.Partial", + "torch.distributed.tensor.placement_types.Replicate", + "torch.distributed.tensor.placement_types.Shard", + ] + ): + raise UnpicklingError( + "``torch.distributed.tensor`` must be imported to load DTensors" + ) + else: + builtins_name = "builtins" + if ( + builtins_name in full_path + and builtins_name == full_path[: len(builtins_name)] + ): + full_path = full_path[len(builtins_name) :] + full_path = ( + full_path[1:] + if len(full_path) > 0 and full_path[0] == "." + else builtins_name + full_path + ) + raise UnpicklingError( + f"Unsupported global: GLOBAL {full_path} was not an allowed global by default. " + f"Please use `torch.serialization.add_safe_globals([{full_path}])` or the " + f"`torch.serialization.safe_globals([{full_path}])` context manager to allowlist this global " + "if you trust this class/function." + ) + elif key[0] == NEWOBJ[0]: + args = self.stack.pop() + cls = self.stack.pop() + if cls is torch.nn.Parameter: + self.append(torch.nn.Parameter(*args)) + elif ( + cls in _get_user_allowed_globals().values() + or cls in _get_allowed_globals().values() + ): + result = cls.__new__(cls, *args) + if cls in torch._tensor_classes and "sparse" in cls.__module__: + _sparse_tensors_to_validate.append(result) + self.append(result) + else: + raise UnpicklingError( + "Can only create new object for nn.Parameter or classes allowlisted " + f"via `add_safe_globals` but got {cls}" + ) + elif key[0] == REDUCE[0]: + args = self.stack.pop() + func = self.stack[-1] + if ( + func not in _get_allowed_globals().values() + and func not in _get_user_allowed_globals().values() + ): + error_msg = ( + f"Trying to call reduce for unrecognized function {func}" + ) + if hasattr(func, "__self__"): + error_msg += f" which belongs to {func.__self__}" + raise UnpicklingError(error_msg) + result = func(*args) + if func in torch._tensor_classes and "sparse" in func.__module__: + _sparse_tensors_to_validate.append(result) + self.stack[-1] = result + elif key[0] == BUILD[0]: + state = self.stack.pop() + inst = self.stack[-1] + if type(inst) is torch.Tensor: + # Legacy unpickling + inst.set_(*state) + elif type(inst) is torch.nn.Parameter: + inst.__setstate__(state) + elif type(inst) is OrderedDict: + inst.__dict__.update(state) + elif ( + type(inst) in _get_user_allowed_globals().values() + or type(inst) in _get_allowed_globals().values() + ): + if hasattr(inst, "__setstate__"): + inst.__setstate__(state) + else: + # mimics load_build in pickle + # https://github.com/python/cpython/blob/f0c6fccd08904787a39269367f09f263d496114c/Lib/pickle.py#L1854-L1867 + slotstate = None + if isinstance(state, tuple) and len(state) == 2: + state, slotstate = state + if state: + inst.__dict__.update(state) + if slotstate: + for k, v in slotstate.items(): + setattr(inst, k, v) + else: + raise UnpicklingError( + "Can only build Tensor, Parameter, OrderedDict or types allowlisted " + f"via `add_safe_globals`, but got {type(inst)}" + ) + # Stack manipulation + elif key[0] == APPEND[0]: + item = self.stack.pop() + list_obj = self.stack[-1] + if type(list_obj) is not list: + raise UnpicklingError( + f"Can only append to lists, but got {type(list_obj)}" + ) + list_obj.append(item) + elif key[0] == APPENDS[0]: + items = self.pop_mark() + list_obj = self.stack[-1] + if type(list_obj) is not list: + raise UnpicklingError( + f"Can only extend lists, but got {type(list_obj)}" + ) + list_obj.extend(items) + elif key[0] == SETITEM[0]: + (v, k) = (self.stack.pop(), self.stack.pop()) + self.stack[-1][k] = v + elif key[0] == SETITEMS[0]: + items = self.pop_mark() + for i in range(0, len(items), 2): + self.stack[-1][items[i]] = items[i + 1] + elif key[0] == MARK[0]: + self.metastack.append(self.stack) + self.stack = [] + self.append = self.stack.append + elif key[0] == TUPLE[0]: + items = self.pop_mark() + self.append(tuple(items)) + elif key[0] == TUPLE1[0]: + self.stack[-1] = (self.stack[-1],) + elif key[0] == TUPLE2[0]: + self.stack[-2:] = [(self.stack[-2], self.stack[-1])] + elif key[0] == TUPLE3[0]: + self.stack[-3:] = [(self.stack[-3], self.stack[-2], self.stack[-1])] + # Basic types construction + elif key[0] == NONE[0]: + self.append(None) + elif key[0] == NEWFALSE[0]: + self.append(False) + elif key[0] == NEWTRUE[0]: + self.append(True) + elif key[0] == EMPTY_TUPLE[0]: + self.append(()) + elif key[0] == EMPTY_LIST[0]: + self.append([]) + elif key[0] == EMPTY_DICT[0]: + self.append({}) + elif key[0] == EMPTY_SET[0]: + self.append(set()) + elif key[0] == BININT[0]: + self.append(unpack("d", self.read(8))[0]) + elif key[0] == BINUNICODE[0]: + strlen = unpack(" maxsize: + raise UnpicklingError("String is too long") + strval = str(read(strlen), "utf-8", "surrogatepass") + self.append(strval) + elif key[0] == SHORT_BINSTRING[0]: + strlen = read(1)[0] + strdata = read(strlen) + if self.encoding != "bytes": + strdata = strdata.decode(self.encoding, "strict") + self.append(strdata) + elif key[0] == BINPERSID[0]: + pid = self.stack.pop() + # Only allow persistent load of storage + if type(pid) is not tuple and type(pid) is not int: + raise UnpicklingError( + f"persistent_load id must be tuple or int, but got {type(pid)}" + ) + if ( + type(pid) is tuple + and len(pid) > 0 + and torch.serialization._maybe_decode_ascii(pid[0]) != "storage" + ): + raise UnpicklingError( + f"Only persistent_load of storage is allowed, but got {pid[0]}" + ) + self.append(self.persistent_load(pid)) + elif key[0] in [BINGET[0], LONG_BINGET[0]]: + idx = (read(1) if key[0] == BINGET[0] else unpack(" List of Tensors + + Broadcasts the given tensors according to :ref:`broadcasting-semantics`. + + Args: + *tensors: any number of tensors of the same type + + .. warning:: + + More than one element of a broadcasted tensor may refer to a single + memory location. As a result, in-place operations (especially ones that + are vectorized) may result in incorrect behavior. If you need to write + to the tensors, please clone them first. + + Example:: + + >>> x = torch.arange(3).view(1, 3) + >>> y = torch.arange(2).view(2, 1) + >>> a, b = torch.broadcast_tensors(x, y) + >>> a.size() + torch.Size([2, 3]) + >>> a + tensor([[0, 1, 2], + [0, 1, 2]]) + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(broadcast_tensors, tensors, *tensors) + return _VF.broadcast_tensors(tensors) # type: ignore[attr-defined] + + +def broadcast_shapes(*shapes): + r"""broadcast_shapes(*shapes) -> Size + + Similar to :func:`broadcast_tensors` but for shapes. + + This is equivalent to + ``torch.broadcast_tensors(*map(torch.empty, shapes))[0].shape`` + but avoids the need create to intermediate tensors. This is useful for + broadcasting tensors of common batch shape but different rightmost shape, + e.g. to broadcast mean vectors with covariance matrices. + + Example:: + + >>> torch.broadcast_shapes((2,), (3, 1), (1, 1, 1)) + torch.Size([1, 3, 2]) + + Args: + \*shapes (torch.Size): Shapes of tensors. + + Returns: + shape (torch.Size): A shape compatible with all input shapes. + + Raises: + RuntimeError: If shapes are incompatible. + """ + # This wrapper exists to support variadic args. + # TODO Move this to C++ once the jit has better support for torch.Size. + if not torch.jit.is_tracing(): + result = torch._refs._broadcast_shapes(*shapes) + if result is None: + return torch.Size([]) + return torch.Size(result) + else: + # with implementation above, torch.jit.trace hardcodes the sizes which makes subsequent replays fail + with torch.no_grad(): + scalar = torch.zeros((), device="cpu") + tensors = [scalar.expand(shape) for shape in shapes] + tensors = broadcast_tensors(*tensors) + return tensors[0].shape + + +def split( + tensor: Tensor, + split_size_or_sections: Union[int, list[int]], + dim: int = 0, +) -> tuple[Tensor, ...]: + r"""Splits the tensor into chunks. Each chunk is a view of the original tensor. + + If :attr:`split_size_or_sections` is an integer type, then :attr:`tensor` will + be split into equally sized chunks (if possible). Last chunk will be smaller if + the tensor size along the given dimension :attr:`dim` is not divisible by + :attr:`split_size`. + + If :attr:`split_size_or_sections` is a list, then :attr:`tensor` will be split + into ``len(split_size_or_sections)`` chunks with sizes in :attr:`dim` according + to :attr:`split_size_or_sections`. + + Args: + tensor (Tensor): tensor to split. + split_size_or_sections (int) or (list(int)): size of a single chunk or + list of sizes for each chunk + dim (int): dimension along which to split the tensor. + + Example:: + + >>> a = torch.arange(10).reshape(5, 2) + >>> a + tensor([[0, 1], + [2, 3], + [4, 5], + [6, 7], + [8, 9]]) + >>> torch.split(a, 2) + (tensor([[0, 1], + [2, 3]]), + tensor([[4, 5], + [6, 7]]), + tensor([[8, 9]])) + >>> torch.split(a, [1, 4]) + (tensor([[0, 1]]), + tensor([[2, 3], + [4, 5], + [6, 7], + [8, 9]])) + """ + if has_torch_function_unary(tensor): + return handle_torch_function( + split, (tensor,), tensor, split_size_or_sections, dim=dim + ) + # Overwriting reason: + # This dispatches to two ATen functions depending on the type of + # split_size_or_sections. The branching code is in _tensor.py, which we + # call here. + return tensor.split(split_size_or_sections, dim) + + +def einsum(*args: Any) -> Tensor: + r"""einsum(equation, *operands) -> Tensor + + Sums the product of the elements of the input :attr:`operands` along dimensions specified using a notation + based on the Einstein summation convention. + + Einsum allows computing many common multi-dimensional linear algebraic array operations by representing them + in a short-hand format based on the Einstein summation convention, given by :attr:`equation`. The details of + this format are described below, but the general idea is to label every dimension of the input :attr:`operands` + with some subscript and define which subscripts are part of the output. The output is then computed by summing + the product of the elements of the :attr:`operands` along the dimensions whose subscripts are not part of the + output. For example, matrix multiplication can be computed using einsum as `torch.einsum("ij,jk->ik", A, B)`. + Here, j is the summation subscript and i and k the output subscripts (see section below for more details on why). + + Equation: + + The :attr:`equation` string specifies the subscripts (letters in `[a-zA-Z]`) for each dimension of + the input :attr:`operands` in the same order as the dimensions, separating subscripts for each operand by a + comma (','), e.g. `'ij,jk'` specify subscripts for two 2D operands. The dimensions labeled with the same subscript + must be broadcastable, that is, their size must either match or be `1`. The exception is if a subscript is + repeated for the same input operand, in which case the dimensions labeled with this subscript for this operand + must match in size and the operand will be replaced by its diagonal along these dimensions. The subscripts that + appear exactly once in the :attr:`equation` will be part of the output, sorted in increasing alphabetical order. + The output is computed by multiplying the input :attr:`operands` element-wise, with their dimensions aligned based + on the subscripts, and then summing out the dimensions whose subscripts are not part of the output. + + Optionally, the output subscripts can be explicitly defined by adding an arrow ('->') at the end of the equation + followed by the subscripts for the output. For instance, the following equation computes the transpose of a + matrix multiplication: 'ij,jk->ki'. The output subscripts must appear at least once for some input operand and + at most once for the output. + + Ellipsis ('...') can be used in place of subscripts to broadcast the dimensions covered by the ellipsis. + Each input operand may contain at most one ellipsis which will cover the dimensions not covered by subscripts, + e.g. for an input operand with 5 dimensions, the ellipsis in the equation `'ab...c'` cover the third and fourth + dimensions. The ellipsis does not need to cover the same number of dimensions across the :attr:`operands` but the + 'shape' of the ellipsis (the size of the dimensions covered by them) must broadcast together. If the output is not + explicitly defined with the arrow ('->') notation, the ellipsis will come first in the output (left-most dimensions), + before the subscript labels that appear exactly once for the input operands. e.g. the following equation implements + batch matrix multiplication `'...ij,...jk'`. + + A few final notes: the equation may contain whitespaces between the different elements (subscripts, ellipsis, + arrow and comma) but something like `'. . .'` is not valid. An empty string `''` is valid for scalar operands. + + .. note:: + + ``torch.einsum`` handles ellipsis ('...') differently from NumPy in that it allows dimensions + covered by the ellipsis to be summed over, that is, ellipsis are not required to be part of the output. + + .. note:: + + Please install opt-einsum (https://optimized-einsum.readthedocs.io/en/stable/) in order to enroll into a more + performant einsum. You can install when installing torch like so: `pip install torch[opt-einsum]` or by itself + with `pip install opt-einsum`. + + If opt-einsum is available, this function will automatically speed up computation and/or consume less memory + by optimizing contraction order through our opt_einsum backend :mod:`torch.backends.opt_einsum` (The _ vs - is + confusing, I know). This optimization occurs when there are at least three inputs, since the order does not matter + otherwise. Note that finding `the` optimal path is an NP-hard problem, thus, opt-einsum relies on different + heuristics to achieve near-optimal results. If opt-einsum is not available, the default order is to contract + from left to right. + + To bypass this default behavior, add the following to disable opt_einsum and skip path calculation: + ``torch.backends.opt_einsum.enabled = False`` + + To specify which strategy you'd like for opt_einsum to compute the contraction path, add the following line: + ``torch.backends.opt_einsum.strategy = 'auto'``. The default strategy is 'auto', and we also support 'greedy' and + 'optimal'. Disclaimer that the runtime of 'optimal' is factorial in the number of inputs! See more details in + the opt_einsum documentation (https://optimized-einsum.readthedocs.io/en/stable/path_finding.html). + + .. note:: + + As of PyTorch 1.10 :func:`torch.einsum` also supports the sublist format (see examples below). In this format, + subscripts for each operand are specified by sublists, list of integers in the range [0, 52). These sublists + follow their operands, and an extra sublist can appear at the end of the input to specify the output's + subscripts., e.g. `torch.einsum(op1, sublist1, op2, sublist2, ..., [subslist_out])`. Python's `Ellipsis` object + may be provided in a sublist to enable broadcasting as described in the Equation section above. + + Args: + equation (str): The subscripts for the Einstein summation. + operands (List[Tensor]): The tensors to compute the Einstein summation of. + + Examples:: + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> # trace + >>> torch.einsum('ii', torch.randn(4, 4)) + tensor(-1.2104) + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> # diagonal + >>> torch.einsum('ii->i', torch.randn(4, 4)) + tensor([-0.1034, 0.7952, -0.2433, 0.4545]) + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> # outer product + >>> x = torch.randn(5) + >>> y = torch.randn(4) + >>> torch.einsum('i,j->ij', x, y) + tensor([[ 0.1156, -0.2897, -0.3918, 0.4963], + [-0.3744, 0.9381, 1.2685, -1.6070], + [ 0.7208, -1.8058, -2.4419, 3.0936], + [ 0.1713, -0.4291, -0.5802, 0.7350], + [ 0.5704, -1.4290, -1.9323, 2.4480]]) + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> # batch matrix multiplication + >>> As = torch.randn(3, 2, 5) + >>> Bs = torch.randn(3, 5, 4) + >>> torch.einsum('bij,bjk->bik', As, Bs) + tensor([[[-1.0564, -1.5904, 3.2023, 3.1271], + [-1.6706, -0.8097, -0.8025, -2.1183]], + + [[ 4.2239, 0.3107, -0.5756, -0.2354], + [-1.4558, -0.3460, 1.5087, -0.8530]], + + [[ 2.8153, 1.8787, -4.3839, -1.2112], + [ 0.3728, -2.1131, 0.0921, 0.8305]]]) + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> # with sublist format and ellipsis + >>> torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2]) + tensor([[[-1.0564, -1.5904, 3.2023, 3.1271], + [-1.6706, -0.8097, -0.8025, -2.1183]], + + [[ 4.2239, 0.3107, -0.5756, -0.2354], + [-1.4558, -0.3460, 1.5087, -0.8530]], + + [[ 2.8153, 1.8787, -4.3839, -1.2112], + [ 0.3728, -2.1131, 0.0921, 0.8305]]]) + + >>> # batch permute + >>> A = torch.randn(2, 3, 4, 5) + >>> torch.einsum('...ij->...ji', A).shape + torch.Size([2, 3, 5, 4]) + + >>> # equivalent to torch.nn.functional.bilinear + >>> A = torch.randn(3, 5, 4) + >>> l = torch.randn(2, 5) + >>> r = torch.randn(2, 4) + >>> torch.einsum('bn,anm,bm->ba', l, A, r) + tensor([[-0.3430, -5.2405, 0.4494], + [ 0.3311, 5.5201, -3.0356]]) + """ + import torch.backends.opt_einsum as opt_einsum + + # This wrapper exists to support variadic args. + if len(args) < 2: + raise ValueError( + "einsum(): must specify the equation string and at least one operand, " + "or at least one operand and its subscripts list" + ) + + equation = None + operands = None + + if isinstance(args[0], torch.Tensor): + # Convert the subscript list format which is an interleaving of operand and its subscripts + # list with an optional output subscripts list at the end (see documentation for more details on this) + # to the equation string format by creating the equation string from the subscripts list and grouping the + # input operands into a tensorlist (List[Tensor]). + def parse_subscript(n: int) -> str: + if n == Ellipsis: + return "..." + if n >= 0 and n < 26: + return chr(ord("A") + n) + if n >= 26 and n < 52: + return chr(ord("a") + n - 26) + raise ValueError( + "einsum(): subscript in subscript list is not within the valid range [0, 52)" + ) + + # Parse subscripts for input operands + equation = ",".join("".join(parse_subscript(s) for s in l) for l in args[1::2]) + + # Parse optional output subscripts (provided when the number of arguments is odd) + if len(args) % 2 == 1: + equation += "->" + "".join(parse_subscript(s) for s in args[-1]) + operands = args[:-1:2] + else: + operands = args[::2] + else: + equation = args[0] + operands = args[1:] + + if has_torch_function(operands): + return handle_torch_function(einsum, operands, equation, *operands) + + if len(operands) == 1 and isinstance(operands[0], (list, tuple)): + # the old interface of passing the operands as one list argument + _operands = operands[0] + # recurse in case operands contains value that has torch function + # in the original implementation this line is omitted + return einsum(equation, *_operands) + + if len(operands) <= 2 or not opt_einsum.enabled: + # the path for contracting 0 or 1 time(s) is already optimized + # or the user has disabled using opt_einsum + return _VF.einsum(equation, operands) # type: ignore[attr-defined] + + path = None + if opt_einsum.is_available(): + _opt_einsum = opt_einsum.get_opt_einsum() + tupled_path = _opt_einsum.contract_path( + equation, *operands, optimize=opt_einsum.strategy + )[0] + # flatten path for dispatching to C++ + path = [*itertools.chain.from_iterable(tupled_path)] + return _VF.einsum(equation, operands, path=path) # type: ignore[attr-defined] + + +# This wrapper exists to support variadic args. +if TYPE_CHECKING: + # The JIT doesn't understand Union, so only add type annotation for mypy + def meshgrid( + *tensors: Union[Tensor, list[Tensor]], indexing: Optional[str] = None + ) -> tuple[Tensor, ...]: + return _meshgrid(*tensors, indexing=indexing) + +else: + + def meshgrid(*tensors, indexing: Optional[str] = None) -> tuple[Tensor, ...]: + r"""Creates grids of coordinates specified by the 1D inputs in `attr`:tensors. + + This is helpful when you want to visualize data over some + range of inputs. See below for a plotting example. + + Given :math:`N` 1D tensors :math:`T_0 \ldots T_{N-1}` as + inputs with corresponding sizes :math:`S_0 \ldots S_{N-1}`, + this creates :math:`N` N-dimensional tensors :math:`G_0 \ldots + G_{N-1}`, each with shape :math:`(S_0, ..., S_{N-1})` where + the output :math:`G_i` is constructed by expanding :math:`T_i` + to the result shape. + + .. note:: + 0D inputs are treated equivalently to 1D inputs of a + single element. + + .. warning:: + `torch.meshgrid(*tensors)` currently has the same behavior + as calling `numpy.meshgrid(*arrays, indexing='ij')`. + + In the future `torch.meshgrid` will transition to + `indexing='xy'` as the default. + + https://github.com/pytorch/pytorch/issues/50276 tracks + this issue with the goal of migrating to NumPy's behavior. + + .. seealso:: + + :func:`torch.cartesian_prod` has the same effect but it + collects the data in a tensor of vectors. + + Args: + tensors (list of Tensor): list of scalars or 1 dimensional tensors. Scalars will be + treated as tensors of size :math:`(1,)` automatically + + indexing: (str, optional): the indexing mode, either "xy" + or "ij", defaults to "ij". See warning for future changes. + + If "xy" is selected, the first dimension corresponds + to the cardinality of the second input and the second + dimension corresponds to the cardinality of the first + input. + + If "ij" is selected, the dimensions are in the same + order as the cardinality of the inputs. + + Returns: + seq (sequence of Tensors): If the input has :math:`N` + tensors of size :math:`S_0 \ldots S_{N-1}``, then the + output will also have :math:`N` tensors, where each tensor + is of shape :math:`(S_0, ..., S_{N-1})`. + + Example:: + + >>> x = torch.tensor([1, 2, 3]) + >>> y = torch.tensor([4, 5, 6]) + + Observe the element-wise pairings across the grid, (1, 4), + (1, 5), ..., (3, 6). This is the same thing as the + cartesian product. + >>> grid_x, grid_y = torch.meshgrid(x, y, indexing='ij') + >>> grid_x + tensor([[1, 1, 1], + [2, 2, 2], + [3, 3, 3]]) + >>> grid_y + tensor([[4, 5, 6], + [4, 5, 6], + [4, 5, 6]]) + + This correspondence can be seen when these grids are + stacked properly. + >>> torch.equal(torch.cat(tuple(torch.dstack([grid_x, grid_y]))), + ... torch.cartesian_prod(x, y)) + True + + `torch.meshgrid` is commonly used to produce a grid for + plotting. + >>> # xdoctest: +REQUIRES(module:matplotlib) + >>> # xdoctest: +REQUIRES(env:DOCTEST_SHOW) + >>> import matplotlib.pyplot as plt + >>> xs = torch.linspace(-5, 5, steps=100) + >>> ys = torch.linspace(-5, 5, steps=100) + >>> x, y = torch.meshgrid(xs, ys, indexing='xy') + >>> z = torch.sin(torch.sqrt(x * x + y * y)) + >>> ax = plt.axes(projection='3d') + >>> ax.plot_surface(x.numpy(), y.numpy(), z.numpy()) + >>> plt.show() + + .. image:: ../_static/img/meshgrid.png + :width: 512 + + """ + return _meshgrid(*tensors, indexing=indexing) + + +def _meshgrid(*tensors, indexing: Optional[str]): + if has_torch_function(tensors): + return handle_torch_function(meshgrid, tensors, *tensors, indexing=indexing) + if len(tensors) == 1 and isinstance(tensors[0], (list, tuple)): + # the old interface of passing the operands as one list argument + tensors = tensors[0] # type: ignore[assignment] + + # Continue allowing call of old method that takes no indexing + # kwarg for forward compatibility reasons. + # + # Remove this two weeks after landing. + kwargs = {} if indexing is None else {"indexing": indexing} + return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined] + + +def stft( + input: Tensor, + n_fft: int, + hop_length: Optional[int] = None, + win_length: Optional[int] = None, + window: Optional[Tensor] = None, + center: bool = True, + pad_mode: str = "reflect", + normalized: bool = False, + onesided: Optional[bool] = None, + return_complex: Optional[bool] = None, + align_to_window: Optional[bool] = None, +) -> Tensor: + r"""Short-time Fourier transform (STFT). + + .. warning:: + From version 1.8.0, :attr:`return_complex` must always be given + explicitly for real inputs and `return_complex=False` has been + deprecated. Strongly prefer `return_complex=True` as in a future + pytorch release, this function will only return complex tensors. + + Note that :func:`torch.view_as_real` can be used to recover a real + tensor with an extra last dimension for real and imaginary components. + + .. warning:: + From version 2.1, a warning will be provided if a :attr:`window` is + not specified. In a future release, this attribute will be required. + Not providing a window currently defaults to using a rectangular window, + which may result in undesirable artifacts. Consider using tapered windows, + such as :func:`torch.hann_window`. + + The STFT computes the Fourier transform of short overlapping windows of the + input. This giving frequency components of the signal as they change over + time. The interface of this function is modeled after (but *not* a drop-in + replacement for) librosa_ stft function. + + .. _librosa: https://librosa.org/doc/latest/generated/librosa.stft.html + + Ignoring the optional batch dimension, this method computes the following + expression: + + .. math:: + X[\omega, m] = \sum_{k = 0}^{\text{win\_length-1}}% + \text{window}[k]\ \text{input}[m \times \text{hop\_length} + k]\ % + \exp\left(- j \frac{2 \pi \cdot \omega k}{\text{n\_fft}}\right), + + where :math:`m` is the index of the sliding window, and :math:`\omega` is + the frequency :math:`0 \leq \omega < \text{n\_fft}` for ``onesided=False``, + or :math:`0 \leq \omega < \lfloor \text{n\_fft} / 2 \rfloor + 1` for ``onesided=True``. + + * :attr:`input` must be either a 1-D time sequence or a 2-D batch of time + sequences. + + * If :attr:`hop_length` is ``None`` (default), it is treated as equal to + ``floor(n_fft / 4)``. + + * If :attr:`win_length` is ``None`` (default), it is treated as equal to + :attr:`n_fft`. + + * :attr:`window` can be a 1-D tensor of size :attr:`win_length`, e.g., from + :meth:`torch.hann_window`. If :attr:`window` is ``None`` (default), it is + treated as if having :math:`1` everywhere in the window. If + :math:`\text{win\_length} < \text{n\_fft}`, :attr:`window` will be padded on + both sides to length :attr:`n_fft` before being applied. + + * If :attr:`center` is ``True`` (default), :attr:`input` will be padded on + both sides so that the :math:`t`-th frame is centered at time + :math:`t \times \text{hop\_length}`. Otherwise, the :math:`t`-th frame + begins at time :math:`t \times \text{hop\_length}`. + + * :attr:`pad_mode` determines the padding method used on :attr:`input` when + :attr:`center` is ``True``. See :meth:`torch.nn.functional.pad` for + all available options. Default is ``"reflect"``. + + * If :attr:`onesided` is ``True`` (default for real input), only values for + :math:`\omega` in :math:`\left[0, 1, 2, \dots, \left\lfloor + \frac{\text{n\_fft}}{2} \right\rfloor + 1\right]` are returned because + the real-to-complex Fourier transform satisfies the conjugate symmetry, + i.e., :math:`X[m, \omega] = X[m, \text{n\_fft} - \omega]^*`. + Note if the input or window tensors are complex, then :attr:`onesided` + output is not possible. + + * If :attr:`normalized` is ``True`` (default is ``False``), the function + returns the normalized STFT results, i.e., multiplied by :math:`(\text{frame\_length})^{-0.5}`. + + * If :attr:`return_complex` is ``True`` (default if input is complex), the + return is a ``input.dim() + 1`` dimensional complex tensor. If ``False``, + the output is a ``input.dim() + 2`` dimensional real tensor where the last + dimension represents the real and imaginary components. + + Returns either a complex tensor of size :math:`(* \times N \times T)` if + :attr:`return_complex` is true, or a real tensor of size :math:`(* \times N + \times T \times 2)`. Where :math:`*` is the optional batch size of + :attr:`input`, :math:`N` is the number of frequencies where STFT is applied + and :math:`T` is the total number of frames used. + + .. warning:: + This function changed signature at version 0.4.1. Calling with the + previous signature may cause error or return incorrect result. + + Args: + input (Tensor): the input tensor of shape `(B?, L)` where `B?` is an optional + batch dimension + n_fft (int): size of Fourier transform + hop_length (int, optional): the distance between neighboring sliding window + frames. Default: ``None`` (treated as equal to ``floor(n_fft / 4)``) + win_length (int, optional): the size of window frame and STFT filter. + Default: ``None`` (treated as equal to :attr:`n_fft`) + window (Tensor, optional): the optional window function. + Shape must be 1d and `<= n_fft` + Default: ``None`` (treated as window of all :math:`1` s) + center (bool, optional): whether to pad :attr:`input` on both sides so + that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`. + Default: ``True`` + pad_mode (str, optional): controls the padding method used when + :attr:`center` is ``True``. Default: ``"reflect"`` + normalized (bool, optional): controls whether to return the normalized STFT results + Default: ``False`` + onesided (bool, optional): controls whether to return half of results to + avoid redundancy for real inputs. + Default: ``True`` for real :attr:`input` and :attr:`window`, ``False`` otherwise. + return_complex (bool, optional): whether to return a complex tensor, or + a real tensor with an extra last dimension for the real and + imaginary components. + + .. versionchanged:: 2.0 + ``return_complex`` is now a required argument for real inputs, + as the default is being transitioned to ``True``. + + .. deprecated:: 2.0 + ``return_complex=False`` is deprecated, instead use ``return_complex=True`` + Note that calling :func:`torch.view_as_real` on the output will + recover the deprecated output format. + + Returns: + Tensor: A tensor containing the STFT result with shape `(B?, N, T, C?)` where + - `B?` is an optional batch dimension from the input. + - `N` is the number of frequency samples, `(n_fft // 2) + 1` for + `onesided=True`, or otherwise `n_fft`. + - `T` is the number of frames, `1 + L // hop_length` + for `center=True`, or `1 + (L - n_fft) // hop_length` otherwise. + - `C?` is an optional length-2 dimension of real and imaginary + components, present when `return_complex=False`. + + """ + if has_torch_function_unary(input): + return handle_torch_function( + stft, + (input,), + input, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window, + center=center, + pad_mode=pad_mode, + normalized=normalized, + onesided=onesided, + return_complex=return_complex, + align_to_window=align_to_window, + ) + if center and align_to_window is not None: + raise RuntimeError( + "stft align_to_window should only be set when center = false" + ) + # NOTE: Do not edit. This code will be removed once the forward-compatibility + # period is over for PR #73432 + if center: + signal_dim = input.dim() + extended_shape = [1] * (3 - signal_dim) + list(input.size()) + pad = int(n_fft // 2) + input = F.pad(input.view(extended_shape), [pad, pad], pad_mode) + input = input.view(input.shape[-signal_dim:]) + return _VF.stft( # type: ignore[attr-defined] + input, + n_fft, + hop_length, + win_length, + window, + normalized, + onesided, + return_complex, + align_to_window, + ) + + +istft = _add_docstr( + torch.istft, + "istft(input, n_fft, hop_length=None, win_length=None, window=None, center=True, " + "normalized=False, onesided=None, length=None, return_complex=False) -> Tensor:\n" + r""" +Inverse short time Fourier Transform. This is expected to be the inverse of :func:`~torch.stft`. + +.. warning:: + From version 2.1, a warning will be provided if a :attr:`window` is + not specified. In a future release, this attribute will be required. + Please provide the same window used in the stft call. + +It has the same parameters (+ additional optional parameter of :attr:`length`) and it should return the +least squares estimation of the original signal. The algorithm will check using the NOLA condition ( +nonzero overlap). + +Important consideration in the parameters :attr:`window` and :attr:`center` so that the envelope +created by the summation of all the windows is never zero at certain point in time. Specifically, +:math:`\sum_{t=-\infty}^{\infty} |w|^2[n-t\times hop\_length] \cancel{=} 0`. + +Since :func:`~torch.stft` discards elements at the end of the signal if they do not fit in a frame, +``istft`` may return a shorter signal than the original signal (can occur if :attr:`center` is False +since the signal isn't padded). If `length` is given in the arguments and is longer than expected, +``istft`` will pad zeros to the end of the returned signal. + +If :attr:`center` is ``True``, then there will be padding e.g. ``'constant'``, ``'reflect'``, etc. +Left padding can be trimmed off exactly because they can be calculated but right padding cannot be +calculated without additional information. + +Example: Suppose the last window is: +``[17, 18, 0, 0, 0]`` vs ``[18, 0, 0, 0, 0]`` + +The :attr:`n_fft`, :attr:`hop_length`, :attr:`win_length` are all the same which prevents the calculation +of right padding. These additional values could be zeros or a reflection of the signal so providing +:attr:`length` could be useful. If :attr:`length` is ``None`` then padding will be aggressively removed +(some loss of signal). + +[1] D. W. Griffin and J. S. Lim, "Signal estimation from modified short-time Fourier transform," +IEEE Trans. ASSP, vol.32, no.2, pp.236-243, Apr. 1984. + +Args: + input (Tensor): The input tensor. Expected to be in the format of :func:`~torch.stft`, + output. That is a complex tensor of shape `(B?, N, T)` where + + - `B?` is an optional batch dimension + - `N` is the number of frequency samples, `(n_fft // 2) + 1` + for onesided input, or otherwise `n_fft`. + - `T` is the number of frames, `1 + length // hop_length` for centered stft, + or `1 + (length - n_fft) // hop_length` otherwise. + + .. versionchanged:: 2.0 + Real datatype inputs are no longer supported. Input must now have a + complex datatype, as returned by ``stft(..., return_complex=True)``. + n_fft (int): Size of Fourier transform + hop_length (Optional[int]): The distance between neighboring sliding window frames. + (Default: ``n_fft // 4``) + win_length (Optional[int]): The size of window frame and STFT filter. (Default: ``n_fft``) + window (Optional[torch.Tensor]): The optional window function. + Shape must be 1d and `<= n_fft` + (Default: ``torch.ones(win_length)``) + center (bool): Whether :attr:`input` was padded on both sides so that the :math:`t`-th frame is + centered at time :math:`t \times \text{hop\_length}`. + (Default: ``True``) + normalized (bool): Whether the STFT was normalized. (Default: ``False``) + onesided (Optional[bool]): Whether the STFT was onesided. + (Default: ``True`` if `n_fft != fft_size` in the input size) + length (Optional[int]): The amount to trim the signal by (i.e. the + original signal length). Defaults to `(T - 1) * hop_length` for + centered stft, or `n_fft + (T - 1) * hop_length` otherwise, where `T` + is the number of input frames. + return_complex (Optional[bool]): + Whether the output should be complex, or if the input should be + assumed to derive from a real signal and window. + Note that this is incompatible with ``onesided=True``. + (Default: ``False``) + +Returns: + Tensor: Least squares estimation of the original signal of shape `(B?, length)` where + `B?` is an optional batch dimension from the input tensor. +""", +) + + +if TYPE_CHECKING: + # These _impl functions return a variable number of tensors as output with + # __torch_function__; tuple unpacking is done already rather than being + # done by the caller of the _impl function + _unique_impl_out = Any +else: + _unique_impl_out = tuple[Tensor, Tensor, Tensor] + + +def _unique_impl( + input: Tensor, + sorted: bool = True, + return_inverse: bool = False, + return_counts: bool = False, + dim: Optional[int] = None, +) -> _unique_impl_out: + r"""unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None) -> tuple[Tensor, Tensor, Tensor] + + Returns the unique elements of the input tensor. + + .. note:: This function is different from :func:`torch.unique_consecutive` in the sense that + this function also eliminates non-consecutive duplicate values. + + .. note:: Currently in the CUDA implementation and the CPU implementation, + `torch.unique` always sort the tensor at the beginning regardless of the `sort` argument. + Sorting could be slow, so if your input tensor is already sorted, it is recommended to use + :func:`torch.unique_consecutive` which avoids the sorting. + + Args: + input (Tensor): the input tensor + sorted (bool): Whether to sort the unique elements in ascending order + before returning as output. + return_inverse (bool): Whether to also return the indices for where + elements in the original input ended up in the returned unique list. + return_counts (bool): Whether to also return the counts for each unique + element. + dim (int, optional): the dimension to operate upon. If ``None``, the + unique of the flattened input is returned. Otherwise, each of the + tensors indexed by the given dimension is treated as one of the + elements to apply the unique operation upon. See examples for more + details. Default: ``None`` + + Returns: + (Tensor, Tensor (optional), Tensor (optional)): A tensor or a tuple of tensors containing + + - **output** (*Tensor*): the output list of unique scalar elements. + - **inverse_indices** (*Tensor*): (optional) if + :attr:`return_inverse` is True, there will be an additional + returned tensor (same shape as input) representing the indices + for where elements in the original input map to in the output; + otherwise, this function will only return a single tensor. + - **counts** (*Tensor*): (optional) if + :attr:`return_counts` is True, there will be an additional + returned tensor (same shape as output or output.size(dim), + if dim was specified) representing the number of occurrences + for each unique value or tensor. + + Example:: + + >>> output = torch.unique(torch.tensor([1, 3, 2, 3], dtype=torch.long)) + >>> output + tensor([1, 2, 3]) + + >>> output, inverse_indices = torch.unique( + ... torch.tensor([1, 3, 2, 3], dtype=torch.long), sorted=True, return_inverse=True) + >>> output + tensor([1, 2, 3]) + >>> inverse_indices + tensor([0, 2, 1, 2]) + + >>> output, inverse_indices = torch.unique( + ... torch.tensor([[1, 3], [2, 3]], dtype=torch.long), sorted=True, return_inverse=True) + >>> output + tensor([1, 2, 3]) + >>> inverse_indices + tensor([[0, 2], + [1, 2]]) + + >>> a = torch.tensor([ + ... [ + ... [1, 1, 0, 0], + ... [1, 1, 0, 0], + ... [0, 0, 1, 1], + ... ], + ... [ + ... [0, 0, 1, 1], + ... [0, 0, 1, 1], + ... [1, 1, 1, 1], + ... ], + ... [ + ... [1, 1, 0, 0], + ... [1, 1, 0, 0], + ... [0, 0, 1, 1], + ... ], + ... ]) + + >>> # If we call `torch.unique(a, dim=0)`, each of the tensors `a[idx, :, :]` + >>> # will be compared. We can see that `a[0, :, :]` and `a[2, :, :]` match + >>> # each other, so one of them will be removed. + >>> (a[0, :, :] == a[2, :, :]).all() + tensor(True) + >>> a_unique_dim0 = torch.unique(a, dim=0) + >>> a_unique_dim0 + tensor([[[0, 0, 1, 1], + [0, 0, 1, 1], + [1, 1, 1, 1]], + [[1, 1, 0, 0], + [1, 1, 0, 0], + [0, 0, 1, 1]]]) + + >>> # Notice which sub-tensors from `a` match with the sub-tensors from + >>> # `a_unique_dim0`: + >>> (a_unique_dim0[0, :, :] == a[1, :, :]).all() + tensor(True) + >>> (a_unique_dim0[1, :, :] == a[0, :, :]).all() + tensor(True) + + >>> # For `torch.unique(a, dim=1)`, each of the tensors `a[:, idx, :]` are + >>> # compared. `a[:, 0, :]` and `a[:, 1, :]` match each other, so one of + >>> # them will be removed. + >>> (a[:, 0, :] == a[:, 1, :]).all() + tensor(True) + >>> torch.unique(a, dim=1) + tensor([[[0, 0, 1, 1], + [1, 1, 0, 0]], + [[1, 1, 1, 1], + [0, 0, 1, 1]], + [[0, 0, 1, 1], + [1, 1, 0, 0]]]) + + >>> # For `torch.unique(a, dim=2)`, the tensors `a[:, :, idx]` are compared. + >>> # `a[:, :, 0]` and `a[:, :, 1]` match each other. Also, `a[:, :, 2]` and + >>> # `a[:, :, 3]` match each other as well. So in this case, two of the + >>> # sub-tensors will be removed. + >>> (a[:, :, 0] == a[:, :, 1]).all() + tensor(True) + >>> (a[:, :, 2] == a[:, :, 3]).all() + tensor(True) + >>> torch.unique(a, dim=2) + tensor([[[0, 1], + [0, 1], + [1, 0]], + [[1, 0], + [1, 0], + [1, 1]], + [[0, 1], + [0, 1], + [1, 0]]]) + """ + if has_torch_function_unary(input): + return handle_torch_function( + unique, + (input,), + input, + sorted=sorted, + return_inverse=return_inverse, + return_counts=return_counts, + dim=dim, + ) + + if dim is not None: + output, inverse_indices, counts = _VF.unique_dim( + input, + dim, + sorted=sorted, + return_inverse=return_inverse, + return_counts=return_counts, + ) + else: + output, inverse_indices, counts = torch._unique2( + input, + sorted=sorted, + return_inverse=return_inverse, + return_counts=return_counts, + ) + return output, inverse_indices, counts + + +def _unique_consecutive_impl( + input: Tensor, + return_inverse: bool = False, + return_counts: bool = False, + dim: Optional[int] = None, +) -> _unique_impl_out: + r"""Eliminates all but the first element from every consecutive group of equivalent elements. + + .. note:: This function is different from :func:`torch.unique` in the sense that this function + only eliminates consecutive duplicate values. This semantics is similar to `std::unique` + in C++. + + Args: + input (Tensor): the input tensor + return_inverse (bool): Whether to also return the indices for where + elements in the original input ended up in the returned unique list. + return_counts (bool): Whether to also return the counts for each unique + element. + dim (int): the dimension to apply unique. If ``None``, the unique of the + flattened input is returned. default: ``None`` + + Returns: + (Tensor, Tensor (optional), Tensor (optional)): A tensor or a tuple of tensors containing + + - **output** (*Tensor*): the output list of unique scalar elements. + - **inverse_indices** (*Tensor*): (optional) if + :attr:`return_inverse` is True, there will be an additional + returned tensor (same shape as input) representing the indices + for where elements in the original input map to in the output; + otherwise, this function will only return a single tensor. + - **counts** (*Tensor*): (optional) if + :attr:`return_counts` is True, there will be an additional + returned tensor (same shape as output or output.size(dim), + if dim was specified) representing the number of occurrences + for each unique value or tensor. + + Example:: + + >>> x = torch.tensor([1, 1, 2, 2, 3, 1, 1, 2]) + >>> output = torch.unique_consecutive(x) + >>> output + tensor([1, 2, 3, 1, 2]) + + >>> output, inverse_indices = torch.unique_consecutive(x, return_inverse=True) + >>> output + tensor([1, 2, 3, 1, 2]) + >>> inverse_indices + tensor([0, 0, 1, 1, 2, 3, 3, 4]) + + >>> output, counts = torch.unique_consecutive(x, return_counts=True) + >>> output + tensor([1, 2, 3, 1, 2]) + >>> counts + tensor([2, 2, 1, 2, 1]) + """ + if has_torch_function_unary(input): + return handle_torch_function( + unique_consecutive, + (input,), + input, + return_inverse=return_inverse, + return_counts=return_counts, + dim=dim, + ) + output, inverse_indices, counts = _VF.unique_consecutive( # type: ignore[attr-defined] + input, return_inverse=return_inverse, return_counts=return_counts, dim=dim + ) + return output, inverse_indices, counts + + +def _return_counts( + input, + sorted=True, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, bool, Optional[int]) -> tuple[Tensor, Tensor] + + if has_torch_function_unary(input): + return _unique_impl(input, sorted, return_inverse, return_counts, dim) + + output, _, counts = _unique_impl(input, sorted, return_inverse, return_counts, dim) + return output, counts + + +def _return_output( + input, + sorted=True, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, bool, Optional[int]) -> Tensor + + if has_torch_function_unary(input): + return _unique_impl(input, sorted, return_inverse, return_counts, dim) + + output, _, _ = _unique_impl(input, sorted, return_inverse, return_counts, dim) + return output + + +def _return_inverse( + input, + sorted=True, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, bool, Optional[int]) -> tuple[Tensor, Tensor] + + if has_torch_function_unary(input): + return _unique_impl(input, sorted, return_inverse, return_counts, dim) + + output, inverse_indices, _ = _unique_impl( + input, sorted, return_inverse, return_counts, dim + ) + return output, inverse_indices + + +_return_inverse_false = boolean_dispatch( + arg_name="return_counts", + arg_index=3, + default=False, + if_true=_return_counts, + if_false=_return_output, + module_name=__name__, + func_name="unique", +) + +_return_inverse_true = boolean_dispatch( + arg_name="return_counts", + arg_index=3, + default=False, + if_true=_unique_impl, + if_false=_return_inverse, + module_name=__name__, + func_name="unique", +) + +# The return type of unique depends on `return_inverse`, and `return_counts` so in order to +# resolve the output type in TorchScript we need to statically know the value of both parameters + +unique = boolean_dispatch( + arg_name="return_inverse", + arg_index=2, + default=False, + if_true=_return_inverse_true, + if_false=_return_inverse_false, + module_name=__name__, + func_name="unique", +) +unique.__doc__ = _unique_impl.__doc__ + + +def _consecutive_return_counts( + input, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, Optional[int]) -> tuple[Tensor, Tensor] + + if has_torch_function_unary(input): + return _unique_consecutive_impl(input, return_inverse, return_counts, dim) + + output, _, counts = _unique_consecutive_impl( + input, return_inverse, return_counts, dim + ) + return output, counts + + +def _consecutive_return_output( + input, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, Optional[int]) -> Tensor + + if has_torch_function_unary(input): + return _unique_consecutive_impl(input, return_inverse, return_counts, dim) + + output, _, _ = _unique_consecutive_impl(input, return_inverse, return_counts, dim) + return output + + +def _consecutive_return_inverse( + input, + return_inverse=False, + return_counts=False, + dim=None, +): + # type: (Tensor, bool, bool, Optional[int]) -> tuple[Tensor, Tensor] + + if has_torch_function_unary(input): + return _unique_consecutive_impl(input, return_inverse, return_counts, dim) + + output, inverse_indices, _ = _unique_consecutive_impl( + input, return_inverse, return_counts, dim + ) + return output, inverse_indices + + +_consecutive_return_inverse_false = boolean_dispatch( + arg_name="return_counts", + arg_index=1, + default=False, + if_true=_consecutive_return_counts, + if_false=_consecutive_return_output, + module_name=__name__, + func_name="unique_consecutive", +) + +_consecutive_return_inverse_true = boolean_dispatch( + arg_name="return_counts", + arg_index=1, + default=False, + if_true=_unique_consecutive_impl, + if_false=_consecutive_return_inverse, + module_name=__name__, + func_name="unique_consecutive", +) + +# The return type of unique depends on `return_inverse`, and `return_counts` so in order to +# resolve the output type in TorchScript we need to statically know the value of both parameters + +unique_consecutive = boolean_dispatch( + arg_name="return_inverse", + arg_index=2, + default=False, + if_true=_consecutive_return_inverse_true, + if_false=_consecutive_return_inverse_false, + module_name=__name__, + func_name="unique_consecutive", +) +unique_consecutive.__doc__ = _unique_consecutive_impl.__doc__ + +if TYPE_CHECKING: + pass + # There's no good way to use this type annotation without breaking JIT + # overloads. So leave untyped for mypy for now. +else: + + @overload + def tensordot( + a, + b, + dims: int = 2, + out: Optional[torch.Tensor] = None, + ): + pass + + @overload + def tensordot( # noqa: F811 + a, + b, + dims: tuple[list[int], list[int]], + out: Optional[torch.Tensor] = None, + ): + pass + + @overload + def tensordot( # noqa: F811 + a, + b, + dims: list[list[int]], + out: Optional[torch.Tensor] = None, + ): + pass + + @overload + def tensordot( # noqa: F811 + a, + b, + dims: torch.Tensor, + out: Optional[torch.Tensor] = None, + ): + pass + + +def tensordot( # noqa: F811 + a, + b, + dims=2, + out: Optional[torch.Tensor] = None, +): + r"""Returns a contraction of a and b over multiple dimensions. + + :attr:`tensordot` implements a generalized matrix product. + + Args: + a (Tensor): Left tensor to contract + b (Tensor): Right tensor to contract + dims (int or Tuple[List[int], List[int]] or List[List[int]] containing two lists or Tensor): number of dimensions to + contract or explicit lists of dimensions for :attr:`a` and + :attr:`b` respectively + + When called with a non-negative integer argument :attr:`dims` = :math:`d`, and + the number of dimensions of :attr:`a` and :attr:`b` is :math:`m` and :math:`n`, + respectively, :func:`~torch.tensordot` computes + + .. math:: + r_{i_0,...,i_{m-d}, i_d,...,i_n} + = \sum_{k_0,...,k_{d-1}} a_{i_0,...,i_{m-d},k_0,...,k_{d-1}} \times b_{k_0,...,k_{d-1}, i_d,...,i_n}. + + When called with :attr:`dims` of the list form, the given dimensions will be contracted + in place of the last :math:`d` of :attr:`a` and the first :math:`d` of :math:`b`. The sizes + in these dimensions must match, but :func:`~torch.tensordot` will deal with broadcasted + dimensions. + + Examples:: + + >>> a = torch.arange(60.).reshape(3, 4, 5) + >>> b = torch.arange(24.).reshape(4, 3, 2) + >>> torch.tensordot(a, b, dims=([1, 0], [0, 1])) + tensor([[4400., 4730.], + [4532., 4874.], + [4664., 5018.], + [4796., 5162.], + [4928., 5306.]]) + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> a = torch.randn(3, 4, 5, device='cuda') + >>> b = torch.randn(4, 5, 6, device='cuda') + >>> c = torch.tensordot(a, b, dims=2).cpu() + tensor([[ 8.3504, -2.5436, 6.2922, 2.7556, -1.0732, 3.2741], + [ 3.3161, 0.0704, 5.0187, -0.4079, -4.3126, 4.8744], + [ 0.8223, 3.9445, 3.2168, -0.2400, 3.4117, 1.7780]]) + + >>> a = torch.randn(3, 5, 4, 6) + >>> b = torch.randn(6, 4, 5, 3) + >>> torch.tensordot(a, b, dims=([2, 1, 3], [1, 2, 0])) + tensor([[ 7.7193, -2.4867, -10.3204], + [ 1.5513, -14.4737, -6.5113], + [ -0.2850, 4.2573, -3.5997]]) + """ + if has_torch_function_variadic(a, b): + return handle_torch_function(tensordot, (a, b), a, b, dims=dims, out=out) + + if not isinstance(dims, (tuple, list, torch.Tensor, int, torch.SymInt)): + raise RuntimeError( + "tensordot expects dims to be int or " + + "tuple[list[int], list[int]] or " + + "list[list[int]] containing two lists, but got " + + f"dims={dims}" + ) + + dims_a: list[int] = [] + dims_b: list[int] = [] + + if isinstance(dims, (tuple, list)): + dims_a, dims_b = dims + + if isinstance(dims, torch.Tensor): + num_elements = dims.numel() + if num_elements > 1: + assert dims.size()[0] == 2 + dims_a = torch.jit.annotate(list[int], dims[0].tolist()) + dims_b = torch.jit.annotate(list[int], dims[1].tolist()) + else: + dims_val = int(dims.item()) + if dims_val < 0: + raise RuntimeError(f"tensordot expects dims >= 0, but got dims={dims}") + dims_a = list(range(-dims_val, 0)) + dims_b = list(range(dims_val)) + + if isinstance(dims, (int, torch.SymInt)): + if dims < 0: + raise RuntimeError(f"tensordot expects dims >= 0, but got dims={dims}") + if dims > min(a.dim(), b.dim()): + raise RuntimeError( + f"tensordot expects dims < ndim_a or ndim_b, but got dims={dims}" + ) + dims_a = list(range(-dims, 0)) + dims_b = list(range(dims)) + + if out is None: + return _VF.tensordot(a, b, dims_a, dims_b) # type: ignore[attr-defined] + else: + return _VF.tensordot(a, b, dims_a, dims_b, out=out) # type: ignore[attr-defined] + + +def cartesian_prod(*tensors: Tensor) -> Tensor: + """Do cartesian product of the given sequence of tensors. The behavior is similar to + python's `itertools.product`. + + Args: + *tensors: any number of 1 dimensional tensors. + + Returns: + Tensor: A tensor equivalent to converting all the input tensors into lists, + do `itertools.product` on these lists, and finally convert the resulting list + into tensor. + + Example:: + + >>> import itertools + >>> a = [1, 2, 3] + >>> b = [4, 5] + >>> list(itertools.product(a, b)) + [(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)] + >>> tensor_a = torch.tensor(a) + >>> tensor_b = torch.tensor(b) + >>> torch.cartesian_prod(tensor_a, tensor_b) + tensor([[1, 4], + [1, 5], + [2, 4], + [2, 5], + [3, 4], + [3, 5]]) + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(cartesian_prod, tensors, *tensors) + return _VF.cartesian_prod(tensors) # type: ignore[attr-defined] + + +def block_diag(*tensors): + """Create a block diagonal matrix from provided tensors. + + Args: + *tensors: One or more tensors with 0, 1, or 2 dimensions. + + Returns: + Tensor: A 2 dimensional tensor with all the input tensors arranged in + order such that their upper left and lower right corners are + diagonally adjacent. All other elements are set to 0. + + Example:: + + >>> import torch + >>> A = torch.tensor([[0, 1], [1, 0]]) + >>> B = torch.tensor([[3, 4, 5], [6, 7, 8]]) + >>> C = torch.tensor(7) + >>> D = torch.tensor([1, 2, 3]) + >>> E = torch.tensor([[4], [5], [6]]) + >>> torch.block_diag(A, B, C, D, E) + tensor([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 3, 4, 5, 0, 0, 0, 0, 0], + [0, 0, 6, 7, 8, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 7, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 2, 3, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 4], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 5], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 6]]) + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(block_diag, tensors, *tensors) + return torch._C._VariableFunctions.block_diag(tensors) # type: ignore[attr-defined] + + +def cdist(x1, x2, p=2.0, compute_mode="use_mm_for_euclid_dist_if_necessary"): + # type: (Tensor, Tensor, float, str) -> (Tensor) + r"""Computes batched the p-norm distance between each pair of the two collections of row vectors. + + Args: + x1 (Tensor): input tensor where the last two dimensions represent the points and the feature dimension respectively. + The shape can be :math:`D_1 \times D_2 \times \cdots \times D_n \times P \times M`, + where :math:`P` is the number of points and :math:`M` is the feature dimension. + x2 (Tensor): input tensor where the last two dimensions also represent the points and the feature dimension respectively. + The shape can be :math:`D_1' \times D_2' \times \cdots \times D_m' \times R \times M`, + where :math:`R` is the number of points and :math:`M` is the feature dimension, + which should match the feature dimension of `x1`. + p: p value for the p-norm distance to calculate between each vector pair + :math:`\in [0, \infty]`. + compute_mode: + 'use_mm_for_euclid_dist_if_necessary' - will use matrix multiplication approach to calculate + euclidean distance (p = 2) if P > 25 or R > 25 + 'use_mm_for_euclid_dist' - will always use matrix multiplication approach to calculate + euclidean distance (p = 2) + 'donot_use_mm_for_euclid_dist' - will never use matrix multiplication approach to calculate + euclidean distance (p = 2) + Default: use_mm_for_euclid_dist_if_necessary. + + If x1 has shape :math:`B \times P \times M` and x2 has shape :math:`B \times R \times M` then the + output will have shape :math:`B \times P \times R`. + + This function is equivalent to `scipy.spatial.distance.cdist(input,'minkowski', p=p)` + if :math:`p \in (0, \infty)`. When :math:`p = 0` it is equivalent to + `scipy.spatial.distance.cdist(input, 'hamming') * M`. When :math:`p = \infty`, the closest + scipy function is `scipy.spatial.distance.cdist(xn, lambda x, y: np.abs(x - y).max())`. + + Example: + + >>> a = torch.tensor([[0.9041, 0.0196], [-0.3108, -2.4423], [-0.4821, 1.059]]) + >>> a + tensor([[ 0.9041, 0.0196], + [-0.3108, -2.4423], + [-0.4821, 1.0590]]) + >>> b = torch.tensor([[-2.1763, -0.4713], [-0.6986, 1.3702]]) + >>> b + tensor([[-2.1763, -0.4713], + [-0.6986, 1.3702]]) + >>> torch.cdist(a, b, p=2) + tensor([[3.1193, 2.0959], + [2.7138, 3.8322], + [2.2830, 0.3791]]) + """ + if has_torch_function_variadic(x1, x2): + return handle_torch_function( + cdist, (x1, x2), x1, x2, p=p, compute_mode=compute_mode + ) + if compute_mode == "use_mm_for_euclid_dist_if_necessary": + return _VF.cdist(x1, x2, p, None) # type: ignore[attr-defined] + elif compute_mode == "use_mm_for_euclid_dist": + return _VF.cdist(x1, x2, p, 1) # type: ignore[attr-defined] + elif compute_mode == "donot_use_mm_for_euclid_dist": + return _VF.cdist(x1, x2, p, 2) # type: ignore[attr-defined] + else: + raise ValueError(f"{compute_mode} is not a valid value for compute_mode") + + +def atleast_1d(*tensors): + r""" + Returns a 1-dimensional view of each input tensor with zero dimensions. + Input tensors with one or more dimensions are returned as-is. + + Args: + input (Tensor or sequence of Tensors): tensor(s) to be converted to at least 1-dimensional. + + Returns: + output (Tensor or tuple of Tensors) + + Example:: + + >>> x = torch.arange(2) + >>> x + tensor([0, 1]) + >>> torch.atleast_1d(x) + tensor([0, 1]) + >>> x = torch.tensor(1.) + >>> x + tensor(1.) + >>> torch.atleast_1d(x) + tensor([1.]) + >>> x = torch.tensor(0.5) + >>> y = torch.tensor(1.) + >>> torch.atleast_1d((x, y)) + (tensor([0.5000]), tensor([1.])) + >>> torch.atleast_1d() + () + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(atleast_1d, tensors, *tensors) + if len(tensors) == 1: + tensors = tensors[0] + return _VF.atleast_1d(tensors) # type: ignore[attr-defined] + + +def atleast_2d(*tensors): + r""" + Returns a 2-dimensional view of each input tensor with zero dimensions. + Input tensors with two or more dimensions are returned as-is. + + Args: + input (Tensor or sequence of Tensors): tensor(s) to be converted to at least 2-dimensional. + + Returns: + output (Tensor or tuple of Tensors) + + Example:: + + >>> x = torch.tensor(1.) + >>> x + tensor(1.) + >>> torch.atleast_2d(x) + tensor([[1.]]) + >>> x = torch.arange(4).view(2, 2) + >>> x + tensor([[0, 1], + [2, 3]]) + >>> torch.atleast_2d(x) + tensor([[0, 1], + [2, 3]]) + >>> x = torch.tensor(0.5) + >>> y = torch.tensor(1.) + >>> torch.atleast_2d((x, y)) + (tensor([[0.5000]]), tensor([[1.]])) + >>> torch.atleast_2d() + () + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(atleast_2d, tensors, *tensors) + if len(tensors) == 1: + tensors = tensors[0] + return _VF.atleast_2d(tensors) # type: ignore[attr-defined] + + +def atleast_3d(*tensors): + r""" + Returns a 3-dimensional view of each input tensor with zero dimensions. + Input tensors with three or more dimensions are returned as-is. + + Args: + input (Tensor or sequence of Tensors): tensor(s) to be converted to at least 3-dimensional. + + Returns: + output (Tensor or tuple of Tensors) + + Example: + + >>> x = torch.tensor(0.5) + >>> x + tensor(0.5000) + >>> torch.atleast_3d(x) + tensor([[[0.5000]]]) + >>> y = torch.arange(4).view(2, 2) + >>> y + tensor([[0, 1], + [2, 3]]) + >>> torch.atleast_3d(y) + tensor([[[0], + [1]], + + [[2], + [3]]]) + >>> x = torch.tensor(1).view(1, 1, 1) + >>> x + tensor([[[1]]]) + >>> torch.atleast_3d(x) + tensor([[[1]]]) + >>> x = torch.tensor(0.5) + >>> y = torch.tensor(1.0) + >>> torch.atleast_3d((x, y)) + (tensor([[[0.5000]]]), tensor([[[1.]]])) + >>> torch.atleast_3d() + () + """ + # This wrapper exists to support variadic args. + if has_torch_function(tensors): + return handle_torch_function(atleast_3d, tensors, *tensors) + if len(tensors) == 1: + tensors = tensors[0] + return _VF.atleast_3d(tensors) # type: ignore[attr-defined] + + +if TYPE_CHECKING: + pass + # There's no good way to use this type annotation; cannot rename norm() to + # _norm_impl() in a way that doesn't break JIT overloads. So leave untyped + # for mypy for now. + # def norm(input: Tensor, + # p: Optional[Union[str, Number]] = "fro", + # dim: Optional[Union[int, List[int]]] = None, + # keepdim: bool = False, + # out: Optional[Tensor] = None, + # dtype: _dtype = None) -> Tensor: + # return _norm_impl(input, p, dim, keepdim, out, dtype) +else: + # TODO: type dim as BroadcastingList when + # https://github.com/pytorch/pytorch/issues/33782 is fixed + @overload + def norm( + input, + p="fro", + dim=None, + keepdim=False, + out=None, + dtype=None, + ): + # type: (Tensor, str, Optional[List[int]], bool, Optional[Tensor], Optional[int]) -> Tensor + pass + + @overload + def norm( # noqa: F811 + input, + p="fro", + dim=None, + keepdim=False, + out=None, + dtype=None, + ): + # type: (Tensor, Optional[number], Optional[List[int]], bool, Optional[Tensor], Optional[int]) -> Tensor + pass + + @overload + def norm( # noqa: F811 + input, + p="fro", + dim=None, + keepdim=False, + out=None, + dtype=None, + ): + # type: (Tensor, Optional[number], Optional[int], bool, Optional[Tensor], Optional[int]) -> Tensor + pass + + @overload + def norm( # noqa: F811 + input, + p="fro", + dim=None, + keepdim=False, + out=None, + dtype=None, + ): + # type: (Tensor, str, Optional[int], bool, Optional[Tensor], Optional[int]) -> Tensor + pass + + +def norm( # noqa: F811 + input, + p: Optional[Union[float, str]] = "fro", + dim=None, + keepdim=False, + out=None, + dtype=None, +): + r"""Returns the matrix norm or vector norm of a given tensor. + + .. warning:: + + torch.norm is deprecated and may be removed in a future PyTorch release. + Its documentation and behavior may be incorrect, and it is no longer + actively maintained. + + Use :func:`torch.linalg.vector_norm` when computing vector norms and + :func:`torch.linalg.matrix_norm` when computing matrix norms. + For a function with a similar behavior as this one see :func:`torch.linalg.norm`. + Note, however, the signature for these functions is slightly different than the + signature for ``torch.norm``. + + Args: + input (Tensor): The input tensor. Its data type must be either a floating + point or complex type. For complex inputs, the norm is calculated using the + absolute value of each element. If the input is complex and neither + :attr:`dtype` nor :attr:`out` is specified, the result's data type will + be the corresponding floating point type (e.g. float if :attr:`input` is + complexfloat). + + p (int, float, inf, -inf, 'fro', 'nuc', optional): the order of norm. Default: ``'fro'`` + The following norms can be calculated: + + ====== ============== ========================== + ord matrix norm vector norm + ====== ============== ========================== + 'fro' Frobenius norm -- + 'nuc' nuclear norm -- + Number -- sum(abs(x)**ord)**(1./ord) + ====== ============== ========================== + + The vector norm can be calculated across any number of dimensions. + The corresponding dimensions of :attr:`input` are flattened into + one dimension, and the norm is calculated on the flattened + dimension. + + Frobenius norm produces the same result as ``p=2`` in all cases + except when :attr:`dim` is a list of three or more dims, in which + case Frobenius norm throws an error. + + Nuclear norm can only be calculated across exactly two dimensions. + + dim (int, tuple of ints, list of ints, optional): + Specifies which dimension or dimensions of :attr:`input` to + calculate the norm across. If :attr:`dim` is ``None``, the norm will + be calculated across all dimensions of :attr:`input`. If the norm + type indicated by :attr:`p` does not support the specified number of + dimensions, an error will occur. + keepdim (bool, optional): whether the output tensors have :attr:`dim` + retained or not. Ignored if :attr:`dim` = ``None`` and + :attr:`out` = ``None``. Default: ``False`` + out (Tensor, optional): the output tensor. Ignored if + :attr:`dim` = ``None`` and :attr:`out` = ``None``. + dtype (:class:`torch.dtype`, optional): the desired data type of + returned tensor. If specified, the input tensor is casted to + :attr:`dtype` while performing the operation. Default: None. + + .. note:: + Even though ``p='fro'`` supports any number of dimensions, the true + mathematical definition of Frobenius norm only applies to tensors with + exactly two dimensions. :func:`torch.linalg.matrix_norm` with ``ord='fro'`` + aligns with the mathematical definition, since it can only be applied across + exactly two dimensions. + + Example:: + + >>> import torch + >>> a = torch.arange(9, dtype= torch.float) - 4 + >>> b = a.reshape((3, 3)) + >>> torch.norm(a) + tensor(7.7460) + >>> torch.norm(b) + tensor(7.7460) + >>> torch.norm(a, float('inf')) + tensor(4.) + >>> torch.norm(b, float('inf')) + tensor(4.) + >>> c = torch.tensor([[ 1, 2, 3], [-1, 1, 4]] , dtype=torch.float) + >>> torch.norm(c, dim=0) + tensor([1.4142, 2.2361, 5.0000]) + >>> torch.norm(c, dim=1) + tensor([3.7417, 4.2426]) + >>> torch.norm(c, p=1, dim=1) + tensor([6., 6.]) + >>> d = torch.arange(8, dtype=torch.float).reshape(2, 2, 2) + >>> torch.norm(d, dim=(1, 2)) + tensor([ 3.7417, 11.2250]) + >>> torch.norm(d[0, :, :]), torch.norm(d[1, :, :]) + (tensor(3.7417), tensor(11.2250)) + """ + + if has_torch_function_unary(input): + return handle_torch_function( + norm, (input,), input, p=p, dim=dim, keepdim=keepdim, out=out, dtype=dtype + ) + + # NB. All the repeated code and weird python is to please TorchScript. + # For a more compact implementation see the relevant function in `_refs/__init__.py` + + # We don't do this for MPS or sparse tensors + if input.layout == torch.strided and input.device.type in ( + "cpu", + "cuda", + "xpu", + "meta", + torch.utils.backend_registration._privateuse1_backend_name, + ): + if dim is not None: + if isinstance(dim, (int, torch.SymInt)): + _dim = [dim] + else: + _dim = dim + else: + _dim = None # type: ignore[assignment] + + if isinstance(p, str): + if p == "fro" and ( + dim is None or isinstance(dim, (int, torch.SymInt)) or len(dim) <= 2 + ): + if out is None: + return torch.linalg.vector_norm( + input, 2, _dim, keepdim, dtype=dtype + ) + else: + return torch.linalg.vector_norm( + input, 2, _dim, keepdim, dtype=dtype, out=out + ) + + # Here we either call the nuclear norm, or we call matrix_norm with some arguments + # that will throw an error + if _dim is None: + _dim = list(range(input.ndim)) + if out is None: + return torch.linalg.matrix_norm(input, p, _dim, keepdim, dtype=dtype) + else: + return torch.linalg.matrix_norm( + input, p, _dim, keepdim, dtype=dtype, out=out + ) + else: + # NB. p should be Union[str, number], not Optional! + _p = 2.0 if p is None else p + if out is None: + return torch.linalg.vector_norm(input, _p, _dim, keepdim, dtype=dtype) + else: + return torch.linalg.vector_norm( + input, _p, _dim, keepdim, dtype=dtype, out=out + ) + + ndim = input.dim() + + # catch default case + if dim is None and out is None and dtype is None and p is not None: + if isinstance(p, str): + if p == "fro": + return _VF.frobenius_norm(input, dim=(), keepdim=keepdim) + if not isinstance(p, str): + _dim = list(range(ndim)) + return _VF.norm(input, p, dim=_dim, keepdim=keepdim) # type: ignore[attr-defined] + + # TODO: when https://github.com/pytorch/pytorch/issues/33782 is fixed + # remove the overloads where dim is an int and replace with BraodcastingList1 + # and remove next four lines, replace _dim with dim + if dim is not None: + if isinstance(dim, (int, torch.SymInt)): + _dim = [dim] + else: + _dim = dim + else: + _dim = None # type: ignore[assignment] + + if isinstance(p, str): + if p == "fro": + if dtype is not None: + raise ValueError("dtype argument is not supported in frobenius norm") + + if _dim is None: + _dim = list(range(ndim)) + if out is None: + return _VF.frobenius_norm(input, _dim, keepdim=keepdim) # type: ignore[arg-type] + else: + return _VF.frobenius_norm(input, _dim, keepdim=keepdim, out=out) # type: ignore[arg-type] + elif p == "nuc": + if dtype is not None: + raise ValueError("dtype argument is not supported in nuclear norm") + if _dim is None: + if out is None: + return _VF.nuclear_norm(input, keepdim=keepdim) # type: ignore[arg-type] + else: + return _VF.nuclear_norm(input, keepdim=keepdim, out=out) # type: ignore[arg-type] + else: + if out is None: + return _VF.nuclear_norm(input, _dim, keepdim=keepdim) # type: ignore[arg-type] + else: + return _VF.nuclear_norm(input, _dim, keepdim=keepdim, out=out) # type: ignore[arg-type] + raise RuntimeError(f"only valid string values are 'fro' and 'nuc', found {p}") + else: + if _dim is None: + _dim = list(range(ndim)) + + if out is None: + if dtype is None: + return _VF.norm(input, p, _dim, keepdim=keepdim) # type: ignore[attr-defined] + else: + return _VF.norm(input, p, _dim, keepdim=keepdim, dtype=dtype) # type: ignore[attr-defined] + else: + if dtype is None: + return _VF.norm(input, p, _dim, keepdim=keepdim, out=out) # type: ignore[attr-defined] + else: + return _VF.norm(input, p, _dim, keepdim=keepdim, dtype=dtype, out=out) # type: ignore[attr-defined] + + +def unravel_index( + indices: Tensor, + shape: Union[int, Sequence[int], torch.Size], +) -> tuple[Tensor, ...]: + r"""Converts a tensor of flat indices into a tuple of coordinate tensors that + index into an arbitrary tensor of the specified shape. + + Args: + indices (Tensor): An integer tensor containing indices into the + flattened version of an arbitrary tensor of shape :attr:`shape`. + All elements must be in the range ``[0, prod(shape) - 1]``. + + shape (int, sequence of ints, or torch.Size): The shape of the arbitrary + tensor. All elements must be non-negative. + + Returns: + tuple of Tensors: Each ``i``-th tensor in the output corresponds with + dimension ``i`` of :attr:`shape`. Each tensor has the same shape as + ``indices`` and contains one index into dimension ``i`` for each of the + flat indices given by ``indices``. + + Example:: + + >>> import torch + >>> torch.unravel_index(torch.tensor(4), (3, 2)) + (tensor(2), + tensor(0)) + + >>> torch.unravel_index(torch.tensor([4, 1]), (3, 2)) + (tensor([2, 0]), + tensor([0, 1])) + + >>> torch.unravel_index(torch.tensor([0, 1, 2, 3, 4, 5]), (3, 2)) + (tensor([0, 0, 1, 1, 2, 2]), + tensor([0, 1, 0, 1, 0, 1])) + + >>> torch.unravel_index(torch.tensor([1234, 5678]), (10, 10, 10, 10)) + (tensor([1, 5]), + tensor([2, 6]), + tensor([3, 7]), + tensor([4, 8])) + + >>> torch.unravel_index(torch.tensor([[1234], [5678]]), (10, 10, 10, 10)) + (tensor([[1], [5]]), + tensor([[2], [6]]), + tensor([[3], [7]]), + tensor([[4], [8]])) + + >>> torch.unravel_index(torch.tensor([[1234], [5678]]), (100, 100)) + (tensor([[12], [56]]), + tensor([[34], [78]])) + """ + if has_torch_function_unary(indices): + return handle_torch_function(unravel_index, (indices,), indices, shape=shape) + res_tensor = _unravel_index(indices, shape) + return res_tensor.unbind(-1) + + +def _unravel_index(indices: Tensor, shape: Union[int, Sequence[int]]) -> Tensor: + torch._check_type( + not indices.is_complex() + and not indices.is_floating_point() + and not indices.dtype == torch.bool, + lambda: f"expected 'indices' to be integer dtype, but got {indices.dtype}", + ) + + torch._check_type( + isinstance(shape, (int, torch.SymInt, Sequence)), + lambda: f"expected 'shape' to be int or sequence of ints, but got {type(shape)}", + ) + + if isinstance(shape, (int, torch.SymInt)): + shape = torch.Size([shape]) + else: + for dim in shape: + torch._check_type( + isinstance(dim, (int, torch.SymInt)), + lambda: f"expected 'shape' sequence to only contain ints, but got {type(dim)}", + ) + shape = torch.Size(shape) + + torch._check_value( + all(dim >= 0 for dim in shape), + lambda: f"'shape' cannot have negative values, but got {tuple(shape)}", + ) + + coefs = list( + reversed( + list( + itertools.accumulate( + reversed(shape[1:] + torch.Size([1])), func=operator.mul + ) + ) + ) + ) + return indices.unsqueeze(-1).floor_divide( + torch.tensor(coefs, device=indices.device, dtype=torch.int64) + ) % torch.tensor(shape, device=indices.device, dtype=torch.int64) + + +def chain_matmul(*matrices, out=None): + r"""Returns the matrix product of the :math:`N` 2-D tensors. This product is efficiently computed + using the matrix chain order algorithm which selects the order in which incurs the lowest cost in terms + of arithmetic operations (`[CLRS]`_). Note that since this is a function to compute the product, :math:`N` + needs to be greater than or equal to 2; if equal to 2 then a trivial matrix-matrix product is returned. + If :math:`N` is 1, then this is a no-op - the original matrix is returned as is. + + .. warning:: + + :func:`torch.chain_matmul` is deprecated and will be removed in a future PyTorch release. + Use :func:`torch.linalg.multi_dot` instead, which accepts a list of two or more tensors + rather than multiple arguments. + + Args: + matrices (Tensors...): a sequence of 2 or more 2-D tensors whose product is to be determined. + out (Tensor, optional): the output tensor. Ignored if :attr:`out` = ``None``. + + Returns: + Tensor: if the :math:`i^{th}` tensor was of dimensions :math:`p_{i} \times p_{i + 1}`, then the product + would be of dimensions :math:`p_{1} \times p_{N + 1}`. + + Example:: + + >>> # xdoctest: +SKIP + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> a = torch.randn(3, 4) + >>> b = torch.randn(4, 5) + >>> c = torch.randn(5, 6) + >>> d = torch.randn(6, 7) + >>> # will raise a deprecation warning + >>> torch.chain_matmul(a, b, c, d) + tensor([[ -2.3375, -3.9790, -4.1119, -6.6577, 9.5609, -11.5095, -3.2614], + [ 21.4038, 3.3378, -8.4982, -5.2457, -10.2561, -2.4684, 2.7163], + [ -0.9647, -5.8917, -2.3213, -5.2284, 12.8615, -12.2816, -2.5095]]) + + .. _`[CLRS]`: https://mitpress.mit.edu/books/introduction-algorithms-third-edition + """ + # This wrapper exists to support variadic args. + if has_torch_function(matrices): + return handle_torch_function(chain_matmul, matrices, *matrices) + + if out is None: + return _VF.chain_matmul(matrices) # type: ignore[attr-defined] + else: + return _VF.chain_matmul(matrices, out=out) # type: ignore[attr-defined] + + +def _lu_impl(A, pivot=True, get_infos=False, out=None): + # type: (Tensor, bool, bool, Any) -> tuple[Tensor, Tensor, Tensor] + r"""Computes the LU factorization of a matrix or batches of matrices + :attr:`A`. Returns a tuple containing the LU factorization and + pivots of :attr:`A`. Pivoting is done if :attr:`pivot` is set to + ``True``. + + .. warning:: + + :func:`torch.lu` is deprecated in favor of :func:`torch.linalg.lu_factor` + and :func:`torch.linalg.lu_factor_ex`. :func:`torch.lu` will be removed in a + future PyTorch release. + ``LU, pivots, info = torch.lu(A, compute_pivots)`` should be replaced with + + .. code:: python + + LU, pivots = torch.linalg.lu_factor(A, compute_pivots) + + ``LU, pivots, info = torch.lu(A, compute_pivots, get_infos=True)`` should be replaced with + + .. code:: python + + LU, pivots, info = torch.linalg.lu_factor_ex(A, compute_pivots) + + .. note:: + * The returned permutation matrix for every matrix in the batch is + represented by a 1-indexed vector of size ``min(A.shape[-2], A.shape[-1])``. + ``pivots[i] == j`` represents that in the ``i``-th step of the algorithm, + the ``i``-th row was permuted with the ``j-1``-th row. + * LU factorization with :attr:`pivot` = ``False`` is not available + for CPU, and attempting to do so will throw an error. However, + LU factorization with :attr:`pivot` = ``False`` is available for + CUDA. + * This function does not check if the factorization was successful + or not if :attr:`get_infos` is ``True`` since the status of the + factorization is present in the third element of the return tuple. + * In the case of batches of square matrices with size less or equal + to 32 on a CUDA device, the LU factorization is repeated for + singular matrices due to the bug in the MAGMA library + (see magma issue 13). + * ``L``, ``U``, and ``P`` can be derived using :func:`torch.lu_unpack`. + + .. warning:: + The gradients of this function will only be finite when :attr:`A` is full rank. + This is because the LU decomposition is just differentiable at full rank matrices. + Furthermore, if :attr:`A` is close to not being full rank, + the gradient will be numerically unstable as it depends on the computation of :math:`L^{-1}` and :math:`U^{-1}`. + + Args: + A (Tensor): the tensor to factor of size :math:`(*, m, n)` + pivot (bool, optional): Whether to compute the LU decomposition with partial pivoting, or the regular LU + decomposition. :attr:`pivot`\ `= False` not supported on CPU. Default: `True`. + get_infos (bool, optional): if set to ``True``, returns an info IntTensor. + Default: ``False`` + out (tuple, optional): optional output tuple. If :attr:`get_infos` is ``True``, + then the elements in the tuple are Tensor, IntTensor, + and IntTensor. If :attr:`get_infos` is ``False``, then the + elements in the tuple are Tensor, IntTensor. Default: ``None`` + + Returns: + (Tensor, IntTensor, IntTensor (optional)): A tuple of tensors containing + + - **factorization** (*Tensor*): the factorization of size :math:`(*, m, n)` + + - **pivots** (*IntTensor*): the pivots of size :math:`(*, \text{min}(m, n))`. + ``pivots`` stores all the intermediate transpositions of rows. + The final permutation ``perm`` could be reconstructed by + applying ``swap(perm[i], perm[pivots[i] - 1])`` for ``i = 0, ..., pivots.size(-1) - 1``, + where ``perm`` is initially the identity permutation of :math:`m` elements + (essentially this is what :func:`torch.lu_unpack` is doing). + + - **infos** (*IntTensor*, *optional*): if :attr:`get_infos` is ``True``, this is a tensor of + size :math:`(*)` where non-zero values indicate whether factorization for the matrix or + each minibatch has succeeded or failed + + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> A = torch.randn(2, 3, 3) + >>> A_LU, pivots = torch.lu(A) + >>> A_LU + tensor([[[ 1.3506, 2.5558, -0.0816], + [ 0.1684, 1.1551, 0.1940], + [ 0.1193, 0.6189, -0.5497]], + + [[ 0.4526, 1.2526, -0.3285], + [-0.7988, 0.7175, -0.9701], + [ 0.2634, -0.9255, -0.3459]]]) + >>> pivots + tensor([[ 3, 3, 3], + [ 3, 3, 3]], dtype=torch.int32) + >>> A_LU, pivots, info = torch.lu(A, get_infos=True) + >>> if info.nonzero().size(0) == 0: + ... print('LU factorization succeeded for all samples!') + LU factorization succeeded for all samples! + """ + # If get_infos is True, then we don't need to check for errors and vice versa + return torch._lu_with_info(A, pivot=pivot, check_errors=(not get_infos)) + + +if TYPE_CHECKING: + _ListOrSeq = Sequence[Tensor] +else: + _ListOrSeq = list[Tensor] + + +def _check_list_size(out_len: int, get_infos: bool, out: _ListOrSeq) -> None: + get_infos_int = 1 if get_infos else 0 + if out_len - get_infos_int != 2: + raise TypeError( + f"expected tuple of {2 + int(get_infos)} elements but got {out_len}" + ) + if not isinstance(out, (tuple, list)): + raise TypeError( + f"argument 'out' must be tuple of Tensors, not {type(out).__name__}" + ) + + +def _lu_with_infos(A, pivot=True, get_infos=False, out=None): + # type: (Tensor, bool, bool, Optional[tuple[Tensor, Tensor, Tensor]]) -> tuple[Tensor, Tensor, Tensor] + if has_torch_function_unary(A): + return handle_torch_function( + lu, (A,), A, pivot=pivot, get_infos=get_infos, out=out + ) + result = _lu_impl(A, pivot, get_infos, out) + if out is not None: + _check_list_size(len(out), get_infos, out) + for i in range(len(out)): + out[i].resize_as_(result[i]).copy_(result[i]) + return out + else: + return result # A_LU, pivots, infos + + +def _lu_no_infos(A, pivot=True, get_infos=False, out=None): + # type: (Tensor, bool, bool, Optional[tuple[Tensor, Tensor]]) -> tuple[Tensor, Tensor] + # need to check for torch_function here so that we exit if + if has_torch_function_unary(A): + return handle_torch_function( + lu, (A,), A, pivot=pivot, get_infos=get_infos, out=out + ) + result = _lu_impl(A, pivot, get_infos, out) + if out is not None: + _check_list_size(len(out), get_infos, out) + for i in range(len(out)): + out[i].resize_as_(result[i]).copy_(result[i]) + return out + else: + return result[0], result[1] # A_LU, pivots + + +# The return type of lu depends on `get_infos`, so in order to resolve the output type +# of lu in TorchScript we need to statically know the value of `get_infos` +lu = boolean_dispatch( + arg_name="get_infos", + arg_index=2, + default=False, + if_true=_lu_with_infos, + if_false=_lu_no_infos, + module_name=__name__, + func_name="lu", +) +lu.__doc__ = _lu_impl.__doc__ + + +def align_tensors(*tensors): + raise RuntimeError("`align_tensors` not yet implemented.") diff --git a/venv/lib/python3.10/site-packages/torch/hub.py b/venv/lib/python3.10/site-packages/torch/hub.py new file mode 100644 index 0000000000000000000000000000000000000000..fc943a4dd004d7629332c59344e8f7c2174451a3 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/hub.py @@ -0,0 +1,880 @@ +# mypy: allow-untyped-defs +import contextlib +import errno +import hashlib +import json +import os +import re +import shutil +import sys +import tempfile +import uuid +import warnings +import zipfile +from pathlib import Path +from typing import Any, Optional, Union +from typing_extensions import deprecated +from urllib.error import HTTPError, URLError +from urllib.parse import urlparse # noqa: F401 +from urllib.request import Request, urlopen + +import torch +from torch.serialization import MAP_LOCATION + + +class _Faketqdm: # type: ignore[no-redef] + def __init__(self, total=None, disable=False, unit=None, *args, **kwargs): + self.total = total + self.disable = disable + self.n = 0 + # Ignore all extra *args and **kwargs lest you want to reinvent tqdm + + def update(self, n): + if self.disable: + return + + self.n += n + if self.total is None: + sys.stderr.write(f"\r{self.n:.1f} bytes") + else: + sys.stderr.write(f"\r{100 * self.n / float(self.total):.1f}%") + sys.stderr.flush() + + # Don't bother implementing; use real tqdm if you want + def set_description(self, *args, **kwargs): + pass + + def write(self, s): + sys.stderr.write(f"{s}\n") + + def close(self): + self.disable = True + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.disable: + return + + sys.stderr.write("\n") + + +try: + from tqdm import tqdm # If tqdm is installed use it, otherwise use the fake wrapper +except ImportError: + tqdm = _Faketqdm + +__all__ = [ + "download_url_to_file", + "get_dir", + "help", + "list", + "load", + "load_state_dict_from_url", + "set_dir", +] + +# matches bfd8deac from resnet18-bfd8deac.pth +HASH_REGEX = re.compile(r"-([a-f0-9]*)\.") + +_TRUSTED_REPO_OWNERS = ( + "facebookresearch", + "facebookincubator", + "pytorch", + "fairinternal", +) +ENV_GITHUB_TOKEN = "GITHUB_TOKEN" +ENV_TORCH_HOME = "TORCH_HOME" +ENV_XDG_CACHE_HOME = "XDG_CACHE_HOME" +DEFAULT_CACHE_DIR = "~/.cache" +VAR_DEPENDENCY = "dependencies" +MODULE_HUBCONF = "hubconf.py" +READ_DATA_CHUNK = 128 * 1024 +_hub_dir: Optional[str] = None + + +@contextlib.contextmanager +def _add_to_sys_path(path): + sys.path.insert(0, path) + try: + yield + finally: + sys.path.remove(path) + + +# Copied from tools/shared/module_loader to be included in torch package +def _import_module(name, path): + import importlib.util + from importlib.abc import Loader + + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert isinstance(spec.loader, Loader) + spec.loader.exec_module(module) + return module + + +def _remove_if_exists(path): + if os.path.exists(path): + if os.path.isfile(path): + os.remove(path) + else: + shutil.rmtree(path) + + +def _git_archive_link(repo_owner, repo_name, ref): + # See https://docs.github.com/en/rest/reference/repos#download-a-repository-archive-zip + return f"https://github.com/{repo_owner}/{repo_name}/zipball/{ref}" + + +def _load_attr_from_module(module, func_name): + # Check if callable is defined in the module + if func_name not in dir(module): + return None + return getattr(module, func_name) + + +def _get_torch_home(): + torch_home = os.path.expanduser( + os.getenv( + ENV_TORCH_HOME, + os.path.join(os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), "torch"), + ) + ) + return torch_home + + +def _parse_repo_info(github): + if ":" in github: + repo_info, ref = github.split(":") + else: + repo_info, ref = github, None + repo_owner, repo_name = repo_info.split("/") + + if ref is None: + # The ref wasn't specified by the user, so we need to figure out the + # default branch: main or master. Our assumption is that if main exists + # then it's the default branch, otherwise it's master. + try: + with urlopen(f"https://github.com/{repo_owner}/{repo_name}/tree/main/"): + ref = "main" + except HTTPError as e: + if e.code == 404: + ref = "master" + else: + raise + except URLError as e: + # No internet connection, need to check for cache as last resort + for possible_ref in ("main", "master"): + if os.path.exists( + f"{get_dir()}/{repo_owner}_{repo_name}_{possible_ref}" + ): + ref = possible_ref + break + if ref is None: + raise RuntimeError( + "It looks like there is no internet connection and the " + f"repo could not be found in the cache ({get_dir()})" + ) from e + return repo_owner, repo_name, ref + + +def _read_url(url): + with urlopen(url) as r: + return r.read().decode(r.headers.get_content_charset("utf-8")) + + +def _validate_not_a_forked_repo(repo_owner, repo_name, ref): + # Use urlopen to avoid depending on local git. + headers = {"Accept": "application/vnd.github.v3+json"} + token = os.environ.get(ENV_GITHUB_TOKEN) + if token is not None: + headers["Authorization"] = f"token {token}" + for url_prefix in ( + f"https://api.github.com/repos/{repo_owner}/{repo_name}/branches", + f"https://api.github.com/repos/{repo_owner}/{repo_name}/tags", + ): + page = 0 + while True: + page += 1 + url = f"{url_prefix}?per_page=100&page={page}" + try: + response = json.loads(_read_url(Request(url, headers=headers))) + except HTTPError: + # Retry without token in case it had insufficient permissions. + del headers["Authorization"] + response = json.loads(_read_url(Request(url, headers=headers))) + # Empty response means no more data to process + if not response: + break + for br in response: + if br["name"] == ref or br["commit"]["sha"].startswith(ref): + return + + raise ValueError( + f"Cannot find {ref} in https://github.com/{repo_owner}/{repo_name}. " + "If it's a commit from a forked repo, please call hub.load() with forked repo directly." + ) + + +def _get_cache_or_reload( + github, + force_reload, + trust_repo, + calling_fn, + verbose=True, + skip_validation=False, +): + # Setup hub_dir to save downloaded files + hub_dir = get_dir() + os.makedirs(hub_dir, exist_ok=True) + # Parse github repo information + repo_owner, repo_name, ref = _parse_repo_info(github) + # Github allows branch name with slash '/', + # this causes confusion with path on both Linux and Windows. + # Backslash is not allowed in Github branch name so no need to + # to worry about it. + normalized_br = ref.replace("/", "_") + # Github renames folder repo-v1.x.x to repo-1.x.x + # We don't know the repo name before downloading the zip file + # and inspect name from it. + # To check if cached repo exists, we need to normalize folder names. + owner_name_branch = "_".join([repo_owner, repo_name, normalized_br]) + repo_dir = os.path.join(hub_dir, owner_name_branch) + # Check that the repo is in the trusted list + _check_repo_is_trusted( + repo_owner, + repo_name, + owner_name_branch, + trust_repo=trust_repo, + calling_fn=calling_fn, + ) + + use_cache = (not force_reload) and os.path.exists(repo_dir) + + if use_cache: + if verbose: + sys.stderr.write(f"Using cache found in {repo_dir}\n") + else: + # Validate the tag/branch is from the original repo instead of a forked repo + if not skip_validation: + _validate_not_a_forked_repo(repo_owner, repo_name, ref) + + cached_file = os.path.join(hub_dir, normalized_br + ".zip") + _remove_if_exists(cached_file) + + try: + url = _git_archive_link(repo_owner, repo_name, ref) + sys.stdout.write(f'Downloading: "{url}" to {cached_file}\n') + download_url_to_file(url, cached_file, progress=False) + except HTTPError as err: + if err.code == 300: + # Getting a 300 Multiple Choices error likely means that the ref is both a tag and a branch + # in the repo. This can be disambiguated by explicitly using refs/heads/ or refs/tags + # See https://git-scm.com/book/en/v2/Git-Internals-Git-References + # Here, we do the same as git: we throw a warning, and assume the user wanted the branch + warnings.warn( + f"The ref {ref} is ambiguous. Perhaps it is both a tag and a branch in the repo? " + "Torchhub will now assume that it's a branch. " + "You can disambiguate tags and branches by explicitly passing refs/heads/branch_name or " + "refs/tags/tag_name as the ref. That might require using skip_validation=True." + ) + disambiguated_branch_ref = f"refs/heads/{ref}" + url = _git_archive_link( + repo_owner, repo_name, ref=disambiguated_branch_ref + ) + download_url_to_file(url, cached_file, progress=False) + else: + raise + + with zipfile.ZipFile(cached_file) as cached_zipfile: + extraced_repo_name = cached_zipfile.infolist()[0].filename + extracted_repo = os.path.join(hub_dir, extraced_repo_name) + _remove_if_exists(extracted_repo) + # Unzip the code and rename the base folder + cached_zipfile.extractall(hub_dir) + + _remove_if_exists(cached_file) + _remove_if_exists(repo_dir) + shutil.move(extracted_repo, repo_dir) # rename the repo + + return repo_dir + + +def _check_repo_is_trusted( + repo_owner, + repo_name, + owner_name_branch, + trust_repo, + calling_fn="load", +): + hub_dir = get_dir() + filepath = os.path.join(hub_dir, "trusted_list") + + if not os.path.exists(filepath): + Path(filepath).touch() + with open(filepath) as file: + trusted_repos = tuple(line.strip() for line in file) + + # To minimize friction of introducing the new trust_repo mechanism, we consider that + # if a repo was already downloaded by torchhub, then it is already trusted (even if it's not in the allowlist) + trusted_repos_legacy = next(os.walk(hub_dir))[1] + + owner_name = "_".join([repo_owner, repo_name]) + is_trusted = ( + owner_name in trusted_repos + or owner_name_branch in trusted_repos_legacy + or repo_owner in _TRUSTED_REPO_OWNERS + ) + + # TODO: Remove `None` option in 2.0 and change the default to "check" + if trust_repo is None: + if not is_trusted: + warnings.warn( + "You are about to download and run code from an untrusted repository. In a future release, this won't " + "be allowed. To add the repository to your trusted list, change the command to {calling_fn}(..., " + "trust_repo=False) and a command prompt will appear asking for an explicit confirmation of trust, " + f"or {calling_fn}(..., trust_repo=True), which will assume that the prompt is to be answered with " + f"'yes'. You can also use {calling_fn}(..., trust_repo='check') which will only prompt for " + f"confirmation if the repo is not already trusted. This will eventually be the default behaviour" + ) + return + + if (trust_repo is False) or (trust_repo == "check" and not is_trusted): + response = input( + f"The repository {owner_name} does not belong to the list of trusted repositories and as such cannot be downloaded. " + "Do you trust this repository and wish to add it to the trusted list of repositories (y/N)?" + ) + if response.lower() in ("y", "yes"): + if is_trusted: + print("The repository is already trusted.") + elif response.lower() in ("n", "no", ""): + raise Exception("Untrusted repository.") # noqa: TRY002 + else: + raise ValueError(f"Unrecognized response {response}.") + + # At this point we're sure that the user trusts the repo (or wants to trust it) + if not is_trusted: + with open(filepath, "a") as file: + file.write(owner_name + "\n") + + +def _check_module_exists(name): + import importlib.util + + return importlib.util.find_spec(name) is not None + + +def _check_dependencies(m): + dependencies = _load_attr_from_module(m, VAR_DEPENDENCY) + + if dependencies is not None: + missing_deps = [pkg for pkg in dependencies if not _check_module_exists(pkg)] + if len(missing_deps): + raise RuntimeError(f"Missing dependencies: {', '.join(missing_deps)}") + + +def _load_entry_from_hubconf(m, model): + if not isinstance(model, str): + raise ValueError("Invalid input: model should be a string of function name") + + # Note that if a missing dependency is imported at top level of hubconf, it will + # throw before this function. It's a chicken and egg situation where we have to + # load hubconf to know what're the dependencies, but to import hubconf it requires + # a missing package. This is fine, Python will throw proper error message for users. + _check_dependencies(m) + + func = _load_attr_from_module(m, model) + + if func is None or not callable(func): + raise RuntimeError(f"Cannot find callable {model} in hubconf") + + return func + + +def get_dir() -> str: + r""" + Get the Torch Hub cache directory used for storing downloaded models & weights. + + If :func:`~torch.hub.set_dir` is not called, default path is ``$TORCH_HOME/hub`` where + environment variable ``$TORCH_HOME`` defaults to ``$XDG_CACHE_HOME/torch``. + ``$XDG_CACHE_HOME`` follows the X Design Group specification of the Linux + filesystem layout, with a default value ``~/.cache`` if the environment + variable is not set. + """ + # Issue warning to move data if old env is set + if os.getenv("TORCH_HUB"): + warnings.warn("TORCH_HUB is deprecated, please use env TORCH_HOME instead") + + if _hub_dir is not None: + return _hub_dir + return os.path.join(_get_torch_home(), "hub") + + +def set_dir(d: Union[str, os.PathLike]) -> None: + r""" + Optionally set the Torch Hub directory used to save downloaded models & weights. + + Args: + d (str): path to a local folder to save downloaded models & weights. + """ + global _hub_dir + _hub_dir = os.path.expanduser(d) + + +def list( + github, + force_reload=False, + skip_validation=False, + trust_repo=None, + verbose=True, +): + r""" + List all callable entrypoints available in the repo specified by ``github``. + + Args: + github (str): a string with format "repo_owner/repo_name[:ref]" with an optional + ref (tag or branch). If ``ref`` is not specified, the default branch is assumed to be ``main`` if + it exists, and otherwise ``master``. + Example: 'pytorch/vision:0.10' + force_reload (bool, optional): whether to discard the existing cache and force a fresh download. + Default is ``False``. + skip_validation (bool, optional): if ``False``, torchhub will check that the branch or commit + specified by the ``github`` argument properly belongs to the repo owner. This will make + requests to the GitHub API; you can specify a non-default GitHub token by setting the + ``GITHUB_TOKEN`` environment variable. Default is ``False``. + trust_repo (bool, str or None): ``"check"``, ``True``, ``False`` or ``None``. + This parameter was introduced in v1.12 and helps ensuring that users + only run code from repos that they trust. + + - If ``False``, a prompt will ask the user whether the repo should + be trusted. + - If ``True``, the repo will be added to the trusted list and loaded + without requiring explicit confirmation. + - If ``"check"``, the repo will be checked against the list of + trusted repos in the cache. If it is not present in that list, the + behaviour will fall back onto the ``trust_repo=False`` option. + - If ``None``: this will raise a warning, inviting the user to set + ``trust_repo`` to either ``False``, ``True`` or ``"check"``. This + is only present for backward compatibility and will be removed in + v2.0. + + Default is ``None`` and will eventually change to ``"check"`` in v2.0. + verbose (bool, optional): If ``False``, mute messages about hitting + local caches. Note that the message about first download cannot be + muted. Default is ``True``. + + Returns: + list: The available callables entrypoint + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB) + >>> entrypoints = torch.hub.list("pytorch/vision", force_reload=True) + """ + repo_dir = _get_cache_or_reload( + github, + force_reload, + trust_repo, + "list", + verbose=verbose, + skip_validation=skip_validation, + ) + + with _add_to_sys_path(repo_dir): + hubconf_path = os.path.join(repo_dir, MODULE_HUBCONF) + hub_module = _import_module(MODULE_HUBCONF, hubconf_path) + + # We take functions starts with '_' as internal helper functions + entrypoints = [ + f + for f in dir(hub_module) + if callable(getattr(hub_module, f)) and not f.startswith("_") + ] + + return entrypoints + + +def help(github, model, force_reload=False, skip_validation=False, trust_repo=None): + r""" + Show the docstring of entrypoint ``model``. + + Args: + github (str): a string with format with an optional + ref (a tag or a branch). If ``ref`` is not specified, the default branch is assumed + to be ``main`` if it exists, and otherwise ``master``. + Example: 'pytorch/vision:0.10' + model (str): a string of entrypoint name defined in repo's ``hubconf.py`` + force_reload (bool, optional): whether to discard the existing cache and force a fresh download. + Default is ``False``. + skip_validation (bool, optional): if ``False``, torchhub will check that the ref + specified by the ``github`` argument properly belongs to the repo owner. This will make + requests to the GitHub API; you can specify a non-default GitHub token by setting the + ``GITHUB_TOKEN`` environment variable. Default is ``False``. + trust_repo (bool, str or None): ``"check"``, ``True``, ``False`` or ``None``. + This parameter was introduced in v1.12 and helps ensuring that users + only run code from repos that they trust. + + - If ``False``, a prompt will ask the user whether the repo should + be trusted. + - If ``True``, the repo will be added to the trusted list and loaded + without requiring explicit confirmation. + - If ``"check"``, the repo will be checked against the list of + trusted repos in the cache. If it is not present in that list, the + behaviour will fall back onto the ``trust_repo=False`` option. + - If ``None``: this will raise a warning, inviting the user to set + ``trust_repo`` to either ``False``, ``True`` or ``"check"``. This + is only present for backward compatibility and will be removed in + v2.0. + + Default is ``None`` and will eventually change to ``"check"`` in v2.0. + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB) + >>> print(torch.hub.help("pytorch/vision", "resnet18", force_reload=True)) + """ + repo_dir = _get_cache_or_reload( + github, + force_reload, + trust_repo, + "help", + verbose=True, + skip_validation=skip_validation, + ) + + with _add_to_sys_path(repo_dir): + hubconf_path = os.path.join(repo_dir, MODULE_HUBCONF) + hub_module = _import_module(MODULE_HUBCONF, hubconf_path) + + entry = _load_entry_from_hubconf(hub_module, model) + + return entry.__doc__ + + +def load( + repo_or_dir, + model, + *args, + source="github", + trust_repo=None, + force_reload=False, + verbose=True, + skip_validation=False, + **kwargs, +): + r""" + Load a model from a github repo or a local directory. + + Note: Loading a model is the typical use case, but this can also be used to + for loading other objects such as tokenizers, loss functions, etc. + + If ``source`` is 'github', ``repo_or_dir`` is expected to be + of the form ``repo_owner/repo_name[:ref]`` with an optional + ref (a tag or a branch). + + If ``source`` is 'local', ``repo_or_dir`` is expected to be a + path to a local directory. + + Args: + repo_or_dir (str): If ``source`` is 'github', + this should correspond to a github repo with format ``repo_owner/repo_name[:ref]`` with + an optional ref (tag or branch), for example 'pytorch/vision:0.10'. If ``ref`` is not specified, + the default branch is assumed to be ``main`` if it exists, and otherwise ``master``. + If ``source`` is 'local' then it should be a path to a local directory. + model (str): the name of a callable (entrypoint) defined in the + repo/dir's ``hubconf.py``. + *args (optional): the corresponding args for callable ``model``. + source (str, optional): 'github' or 'local'. Specifies how + ``repo_or_dir`` is to be interpreted. Default is 'github'. + trust_repo (bool, str or None): ``"check"``, ``True``, ``False`` or ``None``. + This parameter was introduced in v1.12 and helps ensuring that users + only run code from repos that they trust. + + - If ``False``, a prompt will ask the user whether the repo should + be trusted. + - If ``True``, the repo will be added to the trusted list and loaded + without requiring explicit confirmation. + - If ``"check"``, the repo will be checked against the list of + trusted repos in the cache. If it is not present in that list, the + behaviour will fall back onto the ``trust_repo=False`` option. + - If ``None``: this will raise a warning, inviting the user to set + ``trust_repo`` to either ``False``, ``True`` or ``"check"``. This + is only present for backward compatibility and will be removed in + v2.0. + + Default is ``None`` and will eventually change to ``"check"`` in v2.0. + force_reload (bool, optional): whether to force a fresh download of + the github repo unconditionally. Does not have any effect if + ``source = 'local'``. Default is ``False``. + verbose (bool, optional): If ``False``, mute messages about hitting + local caches. Note that the message about first download cannot be + muted. Does not have any effect if ``source = 'local'``. + Default is ``True``. + skip_validation (bool, optional): if ``False``, torchhub will check that the branch or commit + specified by the ``github`` argument properly belongs to the repo owner. This will make + requests to the GitHub API; you can specify a non-default GitHub token by setting the + ``GITHUB_TOKEN`` environment variable. Default is ``False``. + **kwargs (optional): the corresponding kwargs for callable ``model``. + + Returns: + The output of the ``model`` callable when called with the given + ``*args`` and ``**kwargs``. + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB) + >>> # from a github repo + >>> repo = "pytorch/vision" + >>> model = torch.hub.load( + ... repo, "resnet50", weights="ResNet50_Weights.IMAGENET1K_V1" + ... ) + >>> # from a local directory + >>> path = "/some/local/path/pytorch/vision" + >>> # xdoctest: +SKIP + >>> model = torch.hub.load(path, "resnet50", weights="ResNet50_Weights.DEFAULT") + """ + source = source.lower() + + if source not in ("github", "local"): + raise ValueError( + f'Unknown source: "{source}". Allowed values: "github" | "local".' + ) + + if source == "github": + repo_or_dir = _get_cache_or_reload( + repo_or_dir, + force_reload, + trust_repo, + "load", + verbose=verbose, + skip_validation=skip_validation, + ) + + model = _load_local(repo_or_dir, model, *args, **kwargs) + return model + + +def _load_local(hubconf_dir, model, *args, **kwargs): + r""" + Load a model from a local directory with a ``hubconf.py``. + + Args: + hubconf_dir (str): path to a local directory that contains a + ``hubconf.py``. + model (str): name of an entrypoint defined in the directory's + ``hubconf.py``. + *args (optional): the corresponding args for callable ``model``. + **kwargs (optional): the corresponding kwargs for callable ``model``. + + Returns: + a single model with corresponding pretrained weights. + + Example: + >>> # xdoctest: +SKIP("stub local path") + >>> path = "/some/local/path/pytorch/vision" + >>> model = _load_local( + ... path, + ... "resnet50", + ... weights="ResNet50_Weights.IMAGENET1K_V1", + ... ) + """ + with _add_to_sys_path(hubconf_dir): + hubconf_path = os.path.join(hubconf_dir, MODULE_HUBCONF) + hub_module = _import_module(MODULE_HUBCONF, hubconf_path) + + entry = _load_entry_from_hubconf(hub_module, model) + model = entry(*args, **kwargs) + + return model + + +def download_url_to_file( + url: str, + dst: str, + hash_prefix: Optional[str] = None, + progress: bool = True, +) -> None: + r"""Download object at the given URL to a local path. + + Args: + url (str): URL of the object to download + dst (str): Full path where object will be saved, e.g. ``/tmp/temporary_file`` + hash_prefix (str, optional): If not None, the SHA256 downloaded file should start with ``hash_prefix``. + Default: None + progress (bool, optional): whether or not to display a progress bar to stderr + Default: True + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB) + >>> # xdoctest: +REQUIRES(POSIX) + >>> torch.hub.download_url_to_file( + ... "https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth", + ... "/tmp/temporary_file", + ... ) + + """ + file_size = None + req = Request(url, headers={"User-Agent": "torch.hub"}) + u = urlopen(req) + meta = u.info() + if hasattr(meta, "getheaders"): + content_length = meta.getheaders("Content-Length") + else: + content_length = meta.get_all("Content-Length") + if content_length is not None and len(content_length) > 0: + file_size = int(content_length[0]) + + # We deliberately save it in a temp file and move it after + # download is complete. This prevents a local working checkpoint + # being overridden by a broken download. + # We deliberately do not use NamedTemporaryFile to avoid restrictive + # file permissions being applied to the downloaded file. + dst = os.path.expanduser(dst) + for _ in range(tempfile.TMP_MAX): + tmp_dst = dst + "." + uuid.uuid4().hex + ".partial" + try: + f = open(tmp_dst, "w+b") + except FileExistsError: + continue + break + else: + raise FileExistsError(errno.EEXIST, "No usable temporary file name found") + + try: + if hash_prefix is not None: + sha256 = hashlib.sha256() + with tqdm( + total=file_size, + disable=not progress, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as pbar: + while True: + buffer = u.read(READ_DATA_CHUNK) + if len(buffer) == 0: + break + f.write(buffer) # type: ignore[possibly-undefined] + if hash_prefix is not None: + sha256.update(buffer) # type: ignore[possibly-undefined] + pbar.update(len(buffer)) + + f.close() + if hash_prefix is not None: + digest = sha256.hexdigest() # type: ignore[possibly-undefined] + if digest[: len(hash_prefix)] != hash_prefix: + raise RuntimeError( + f'invalid hash value (expected "{hash_prefix}", got "{digest}")' + ) + shutil.move(f.name, dst) + finally: + f.close() + if os.path.exists(f.name): + os.remove(f.name) + + +# Hub used to support automatically extracts from zipfile manually compressed by users. +# The legacy zip format expects only one file from torch.save() < 1.6 in the zip. +# We should remove this support since zipfile is now default zipfile format for torch.save(). +def _is_legacy_zip_format(filename: str) -> bool: + if zipfile.is_zipfile(filename): + infolist = zipfile.ZipFile(filename).infolist() + return len(infolist) == 1 and not infolist[0].is_dir() + return False + + +@deprecated( + "Falling back to the old format < 1.6. This support will be " + "deprecated in favor of default zipfile format introduced in 1.6. " + "Please redo torch.save() to save it in the new zipfile format.", + category=FutureWarning, +) +def _legacy_zip_load( + filename: str, + model_dir: str, + map_location: MAP_LOCATION, + weights_only: bool, +) -> dict[str, Any]: + # Note: extractall() defaults to overwrite file if exists. No need to clean up beforehand. + # We deliberately don't handle tarfile here since our legacy serialization format was in tar. + # E.g. resnet18-5c106cde.pth which is widely used. + with zipfile.ZipFile(filename) as f: + members = f.infolist() + if len(members) != 1: + raise RuntimeError("Only one file(not dir) is allowed in the zipfile") + f.extractall(model_dir) + extraced_name = members[0].filename + extracted_file = os.path.join(model_dir, extraced_name) + return torch.load( + extracted_file, map_location=map_location, weights_only=weights_only + ) + + +def load_state_dict_from_url( + url: str, + model_dir: Optional[str] = None, + map_location: MAP_LOCATION = None, + progress: bool = True, + check_hash: bool = False, + file_name: Optional[str] = None, + weights_only: bool = False, +) -> dict[str, Any]: + r"""Loads the Torch serialized object at the given URL. + + If downloaded file is a zip file, it will be automatically + decompressed. + + If the object is already present in `model_dir`, it's deserialized and + returned. + The default value of ``model_dir`` is ``/checkpoints`` where + ``hub_dir`` is the directory returned by :func:`~torch.hub.get_dir`. + + Args: + url (str): URL of the object to download + model_dir (str, optional): directory in which to save the object + map_location (optional): a function or a dict specifying how to remap storage locations (see torch.load) + progress (bool, optional): whether or not to display a progress bar to stderr. + Default: True + check_hash(bool, optional): If True, the filename part of the URL should follow the naming convention + ``filename-.ext`` where ```` is the first eight or more + digits of the SHA256 hash of the contents of the file. The hash is used to + ensure unique names and to verify the contents of the file. + Default: False + file_name (str, optional): name for the downloaded file. Filename from ``url`` will be used if not set. + weights_only(bool, optional): If True, only weights will be loaded and no complex pickled objects. + Recommended for untrusted sources. See :func:`~torch.load` for more details. + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_HUB) + >>> state_dict = torch.hub.load_state_dict_from_url( + ... "https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth" + ... ) + + """ + # Issue warning to move data if old env is set + if os.getenv("TORCH_MODEL_ZOO"): + warnings.warn( + "TORCH_MODEL_ZOO is deprecated, please use env TORCH_HOME instead" + ) + + if model_dir is None: + hub_dir = get_dir() + model_dir = os.path.join(hub_dir, "checkpoints") + + os.makedirs(model_dir, exist_ok=True) + + parts = urlparse(url) + filename = os.path.basename(parts.path) + if file_name is not None: + filename = file_name + cached_file = os.path.join(model_dir, filename) + if not os.path.exists(cached_file): + sys.stdout.write(f'Downloading: "{url}" to {cached_file}\n') + hash_prefix = None + if check_hash: + r = HASH_REGEX.search(filename) # r is Optional[Match[str]] + hash_prefix = r.group(1) if r else None + download_url_to_file(url, cached_file, hash_prefix, progress=progress) + + if _is_legacy_zip_format(cached_file): + return _legacy_zip_load(cached_file, model_dir, map_location, weights_only) + return torch.load(cached_file, map_location=map_location, weights_only=weights_only) diff --git a/venv/lib/python3.10/site-packages/torch/library.py b/venv/lib/python3.10/site-packages/torch/library.py new file mode 100644 index 0000000000000000000000000000000000000000..d36c181581483fdd8c37a1edd99c7e60519408cd --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/library.py @@ -0,0 +1,1688 @@ +# mypy: allow-untyped-defs +import contextlib +import functools +import inspect +import re +import sys +import traceback +import weakref +from collections.abc import Sequence +from typing import ( + Any, + Callable, + Literal, + Optional, + overload, + TYPE_CHECKING, + TypeVar, + Union, +) +from typing_extensions import deprecated, ParamSpec + +import torch +import torch._library as _library +from torch._library.custom_ops import ( + _cast, + _maybe_get_opdef, + custom_op, + CustomOpDef, + device_types_t, +) +from torch._library.infer_schema import infer_schema # noqa: F401 +from torch._library.triton import triton_op, wrap_triton +from torch._ops import OpOverload +from torch.types import _dtype + + +__all__ = [ + "Library", + "impl", + "define", + "fallthrough_kernel", + "impl_abstract", + "register_autocast", + "register_fake", + "register_torch_dispatch", + "register_vmap", + "get_ctx", + "get_kernel", + "custom_op", + "triton_op", + "wrap_triton", + "infer_schema", +] + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +# Set containing the combination of (namespace, operator, DispatchKey) for which a new kernel has been registered +# The keys in the set are of the form `namespace + "/" + op_name + "/" + dispatch_key`. +# This set is maintained to ensure that two libraries don't try to override the exact same functionality to avoid +# libraries calling into kernels not intended to be called. +_impls: set[str] = set() +_defs: set[str] = set() + +# prim is reserved by TorchScript interpreter +_reserved_namespaces = ["prim"] + + +def fallthrough_kernel(): + """ + A dummy function to pass to ``Library.impl`` in order to register a fallthrough. + """ + raise NotImplementedError("fallthrough_kernel() should never be called.") + + +class Library: + """ + A class to create libraries that can be used to register new operators or + override operators in existing libraries from Python. + A user can optionally pass in a dispatch keyname if they only want to register + kernels corresponding to only one specific dispatch key. + + To create a library to override operators in an existing library (with name ns), set the kind to "IMPL". + To create a new library (with name ns) to register new operators, set the kind to "DEF". + To create a fragment of a possibly existing library to register operators (and bypass + the limitation that there is only one library for a given namespace), set the kind to + "FRAGMENT". + + Args: + ns: library name + kind: "DEF", "IMPL", "FRAGMENT" + dispatch_key: PyTorch dispatch key (default: "") + """ + + def __init__(self, ns, kind, dispatch_key=""): + from torch.fx.operator_schemas import _SCHEMA_TO_SIGNATURE_CACHE + + if kind not in ("IMPL", "DEF", "FRAGMENT"): + raise ValueError("Unsupported kind: ", kind) + + if ns in _reserved_namespaces and (kind == "DEF" or kind == "FRAGMENT"): + raise ValueError( + ns, + " is a reserved namespace. Please try creating a library with another name.", + ) + + frame = traceback.extract_stack(limit=2)[0] + filename, lineno = frame.filename, frame.lineno + self.m: Optional[Any] = torch._C._dispatch_library( + kind, ns, dispatch_key, filename, lineno + ) + self.ns = ns + self._op_defs: set[str] = set() + self._op_impls: set[str] = set() + self._registration_handles: list[torch._library.utils.RegistrationHandle] = [] + self.kind = kind + self.dispatch_key = dispatch_key + # Use a finalizer to setup the "destructor" instead of __del__. + # Python __del__ can lead to weird things (globals and locals may already + # be gone when __del__ actually gets called!). finalizers help the + # situation because it lets us capture references and keeps them alive + weakref.finalize( + self, + _del_library, + _impls, + self._op_impls, + _defs, + self._op_defs, + self._registration_handles, + self.m, + _SCHEMA_TO_SIGNATURE_CACHE, + ) + + def __repr__(self): + return f"Library(kind={self.kind}, ns={self.ns}, dispatch_key={self.dispatch_key})>" + + def define(self, schema, alias_analysis="", *, tags=()): + r"""Defines a new operator and its semantics in the ns namespace. + + Args: + schema: function schema to define a new operator. + alias_analysis (optional): Indicates if the aliasing properties of the operator arguments can be + inferred from the schema (default behavior) or not ("CONSERVATIVE"). + tags (Tag | Sequence[Tag]): one or more torch.Tag to apply to this + operator. Tagging an operator changes the operator's behavior + under various PyTorch subsystems; please read the docs for the + torch.Tag carefully before applying it. + + Returns: + name of the operator as inferred from the schema. + + Example:: + + >>> my_lib = Library("mylib", "DEF") + >>> my_lib.define("sum(Tensor self) -> Tensor") + """ + + # This is added because we also want to disallow PURE_FUNCTION alias analysis which is a valid + # AliasAnalysis type in C++ + if alias_analysis not in ["", "FROM_SCHEMA", "CONSERVATIVE"]: + raise RuntimeError(f"Invalid alias_analysis type {alias_analysis}") + assert self.m is not None + if isinstance(tags, torch.Tag): + tags = (tags,) + + name = schema.split("(")[0] + packet_name = name.split(".")[0] if "." in name else name + has_preexisting_packet = hasattr(torch.ops, self.ns) and hasattr( + getattr(torch.ops, self.ns), packet_name + ) + + result = self.m.define(schema, alias_analysis, tuple(tags)) + name = schema.split("(")[0] + qualname = self.ns + "::" + name + + # If the OpOverloadPacket exists already, then this means we're adding a + # new OpOverload for it. Refresh the packet to include the new OpOverload. + if has_preexisting_packet: + ns = getattr(torch.ops, self.ns) + packet = getattr(ns, packet_name) + torch._ops._refresh_packet(packet) + + self._op_defs.add(qualname) + _defs.add(qualname) + return result + + def _register_fake(self, op_name, fn, _stacklevel=1, *, allow_override=False): + r"""Registers the fake impl for an operator defined in the library.""" + + source = torch._library.utils.get_source(_stacklevel + 1) + frame = sys._getframe(_stacklevel) + caller_module = inspect.getmodule(frame) + # Can be none if you call register_fake from somewhere there isn't a module + # (e.g. __main__) + caller_module_name = None if caller_module is None else caller_module.__name__ + + # TODO(rzou): We're gonna need to stage this change with torchvision, + # since torchvision is github first. + if caller_module_name is not None and caller_module_name.startswith( + "torchvision." + ): + caller_module_name = None + + qualname = f"{self.ns}::{op_name}" + entry = torch._library.simple_registry.singleton.find(qualname) + if caller_module_name is not None: + func_to_register = _check_pystubs_once(fn, qualname, caller_module_name) + else: + func_to_register = fn + + handle = entry.fake_impl.register( + func_to_register, source, lib=self, allow_override=allow_override + ) + self._registration_handles.append(handle) + + def _register_torch_dispatch_rule(self, op_name, torch_dispatch_class, fn): + r"""Registers a torch_dispatch rule for the given operator and torch_dispatch_class. + + This allows for open registration to specify the behavior between the operator + and the torch_dispatch_class without needing to modify the torch_dispatch_class + or the operator directly. + + The torch_dispatch_class is either a Tensor subclass with `__torch_dispatch__` or a + TorchDispatchMode. + + If it is a Tensor subclass, we expect fn to have the following signature: + (cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any + + If it is a TorchDispatchMode, we expect fn to have the following signature: + (mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any + """ + + qualname = f"{self.ns}::{op_name}" + entry = torch._library.simple_registry.singleton.find(qualname) + handle = entry.torch_dispatch_rules.register(torch_dispatch_class, fn) + self._registration_handles.append(handle) + + def _impl_with_aoti_compile(self, op_name, dispatch_key=""): + r"""Register the operator to use the AOTI-compiled implementation. + + Args: + op_name: operator name (along with the overload) or OpOverload object. + dispatch_key: dispatch key that the input function should be registered for. By default, it uses + the dispatch key that the library was created with. + + Example:: + + >>> my_lib = Library("aten", "IMPL") + >>> my_lib._impl_with_aoti_compile("div.Tensor", "CPU") + """ + + if dispatch_key == "": + dispatch_key = self.dispatch_key + assert torch.DispatchKeySet(dispatch_key).has(torch._C.DispatchKey.Dense) + + if isinstance(op_name, str): + name = op_name + elif isinstance(op_name, OpOverload): + name = op_name._schema.name + overload_name = op_name._schema.overload_name + if overload_name != "": + name = name + "." + overload_name + else: + raise RuntimeError( + "_impl_with_aoti_compile should be passed either a name or an OpOverload object " + "as the first argument" + ) + + key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key + if key in _impls: + # TODO: in future, add more info about where the existing function is registered (this info is + # today already returned by the C++ warning when _impl_with_aoti_compile is called but we error out before that) + raise RuntimeError( + "This is not allowed since there's already a kernel registered from python overriding {}" + "'s behavior for {} dispatch key and {} namespace.".format( + name.split("::")[-1], dispatch_key, self.ns + ) + ) + + assert self.m is not None + impl_fn: Callable = self.m.impl_with_aoti_compile + impl_fn(self.ns, name.split("::")[-1], dispatch_key) + + _impls.add(key) + self._op_impls.add(key) + + def impl( + self, op_name, fn, dispatch_key="", *, with_keyset=False, allow_override=False + ): + r"""Registers the function implementation for an operator defined in the library. + + Args: + op_name: operator name (along with the overload) or OpOverload object. + fn: function that's the operator implementation for the input dispatch key or :func:`~fallthrough_kernel` + to register a fallthrough. + dispatch_key: dispatch key that the input function should be registered for. By default, it uses + the dispatch key that the library was created with. + with_keyset: flag controlling if the current dispatcher call keyset should be passed as the first argument + to :attr:`fn` when calling. This should be used to create the appropriate keyset for redispatch calls. + allow_override: Flag controlling if we want to override an + existing registered kernel implementation. This is by + default off, and will error you're trying to register a + kernel to a dispatch key with a kernel already + registered. + + Example:: + + >>> my_lib = Library("aten", "IMPL") + >>> def div_cpu(self, other): + >>> return self * (1 / other) + >>> my_lib.impl("div.Tensor", div_cpu, "CPU") + """ + + if not callable(fn): + raise TypeError( + f"Input function is required to be a callable but found type {type(fn)}" + ) + if dispatch_key == "": + dispatch_key = self.dispatch_key + + if isinstance(op_name, str): + name = op_name + elif isinstance(op_name, OpOverload): + name = op_name._schema.name + overload_name = op_name._schema.overload_name + if overload_name != "": + name = name + "." + overload_name + else: + raise RuntimeError( + "impl should be passed either a name or an OpOverload object as the first argument" + ) + + key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key + if (not allow_override) and key in _impls: + # TODO: in future, add more info about where the existing function is registered (this info is + # today already returned by the C++ warning when impl is called but we error out before that) + raise RuntimeError( + "This is not allowed since there's already a kernel registered from python overriding {}" + "'s behavior for {} dispatch key and {} namespace.".format( + name.split("::")[-1], dispatch_key, self.ns + ) + ) + + if dispatch_key == "Meta": + dispatcher_op_name = name + if "::" not in dispatcher_op_name: + dispatcher_op_name = f"{self.ns}::{dispatcher_op_name}" + + # Internally, we shouldn't be registering meta kernels for any operators that + # have CompositeImplicitAutograd kernels. + # Instead, we should be letting those decompositions run, and writing meta kernels + # only for the base operators. + if torch._C._dispatch_has_kernel_for_dispatch_key( + dispatcher_op_name, "CompositeImplicitAutograd" + ): + raise RuntimeError( + f"We should not register a meta kernel directly to the operator '{name}'," + " because it has a CompositeImplicitAutograd kernel in core." + " Instead we should let the operator decompose, and ensure that we have meta kernels" + " for the base ops that it decomposes into." + ) + + assert self.m is not None + self.m.impl( + name, + dispatch_key if dispatch_key != "" else "CompositeImplicitAutograd", + fn, + with_keyset, + ) + + _impls.add(key) + self._op_impls.add(key) + + def fallback(self, fn, dispatch_key="", *, with_keyset=False): + r"""Registers the function implementation as the fallback for the given key. + + This function only works for a library with global namespace ("_"). + + Args: + fn: function used as fallback for the given dispatch key or :func:`~fallthrough_kernel` + to register a fallthrough. + dispatch_key: dispatch key that the input function should be registered for. By default, it uses + the dispatch key that the library was created with. + with_keyset: flag controlling if the current dispatcher call keyset should be passed as the first argument + to :attr:`fn` when calling. This should be used to create the appropriate keyset for redispatch calls. + + Example:: + + >>> my_lib = Library("_", "IMPL") + >>> def fallback_kernel(op, *args, **kwargs): + >>> # Handle all autocast ops generically + >>> # ... + >>> my_lib.fallback(fallback_kernel, "Autocast") + """ + + if dispatch_key == "": + dispatch_key = self.dispatch_key + + if self.ns != "_": + raise RuntimeError( + f"""Fallback can only be registered using library fragment on the global namespace "_" but it is {self.ns}""" + ) + + assert dispatch_key != "" + assert self.m is not None + + self.m.fallback(dispatch_key, fn, with_keyset) + + def _destroy(self): + if self.m is not None: + self.m.reset() + self.m = None + for handle in self._registration_handles: + handle.destroy() + self._registration_handles.clear() + global _impls + _impls -= self._op_impls + for name in self._op_defs: + # Delete the cached torch.ops.ns.foo if it was registered. + # Otherwise, accessing it leads to a segfault. + # It's possible that we only registered an overload in this Library + # and another library owns an alive overload. + # That's OK - the next time torch.ops.ns.foo gets called, it'll be + # recomputed to point at the right collection of overloads. + ns, name_with_overload = name.split("::") + name = name_with_overload.split(".")[0] + if not hasattr(torch.ops, ns): + continue + namespace = getattr(torch.ops, ns) + if not hasattr(namespace, name): + continue + delattr(namespace, name) + namespace._dir.remove(name) + + +def _del_library( + captured_impls, + op_impls, + captured_defs, + op_defs, + registration_handles, + m, + schema_to_signature_cache, +): + for op_def in op_defs: + name = op_def + overload_name = "" + if "." in op_def: + name, overload_name = op_def.split(".") + if ( + name, + overload_name, + ) in schema_to_signature_cache: + del schema_to_signature_cache[(name, overload_name)] + + captured_impls -= op_impls + captured_defs -= op_defs + for handle in registration_handles: + handle.destroy() + + if m is not None: + m.reset() + + +@contextlib.contextmanager +def _scoped_library(*args, **kwargs): + try: + lib = Library(*args, **kwargs) + yield lib + finally: + lib._destroy() + + +_keep_alive: list[Library] = [] + + +NAMELESS_SCHEMA = re.compile(r"\(.*\) -> .*") + + +@functools.singledispatch +def define(qualname, schema, *, lib=None, tags=()): + r"""Defines a new operator. + + In PyTorch, defining an op (short for "operator") is a two step-process: + - we need to define the op (by providing an operator name and schema) + - we need to implement behavior for how the operator interacts with + various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc. + + This entrypoint defines the custom operator (the first step) + you must then perform the second step by calling various + ``impl_*`` APIs, like :func:`torch.library.impl` or + :func:`torch.library.register_fake`. + + Args: + qualname (str): The qualified name for the operator. Should be + a string that looks like "namespace::name", e.g. "aten::sin". + Operators in PyTorch need a namespace to + avoid name collisions; a given operator may only be created once. + If you are writing a Python library, we recommend the namespace to + be the name of your top-level module. + schema (str): The schema of the operator. E.g. "(Tensor x) -> Tensor" + for an op that accepts one Tensor and returns one Tensor. It does + not contain the operator name (that is passed in ``qualname``). + lib (Optional[Library]): If provided, the lifetime of this operator + will be tied to the lifetime of the Library object. + tags (Tag | Sequence[Tag]): one or more torch.Tag to apply to this + operator. Tagging an operator changes the operator's behavior + under various PyTorch subsystems; please read the docs for the + torch.Tag carefully before applying it. + + Example:: + >>> import torch + >>> import numpy as np + >>> + >>> # Define the operator + >>> torch.library.define("mylib::sin", "(Tensor x) -> Tensor") + >>> + >>> # Add implementations for the operator + >>> @torch.library.impl("mylib::sin", "cpu") + >>> def f(x): + >>> return torch.from_numpy(np.sin(x.numpy())) + >>> + >>> # Call the new operator from torch.ops. + >>> x = torch.randn(3) + >>> y = torch.ops.mylib.sin(x) + >>> assert torch.allclose(y, x.sin()) + + """ + if not isinstance(qualname, str): + raise ValueError( + f"define(qualname, schema): expected qualname " + f"to be instance of str, got {type(qualname)}" + ) + namespace, name = torch._library.utils.parse_namespace(qualname) + if lib is None: + lib = Library(namespace, "FRAGMENT") + _keep_alive.append(lib) + if not NAMELESS_SCHEMA.fullmatch(schema): + raise ValueError( + f"define(qualname, schema, ...): expected schema " + f'to look like e.g. "(Tensor x) -> Tensor" but ' + f'got "{schema}"' + ) + lib.define(name + schema, alias_analysis="", tags=tags) + + +@define.register +def _(lib: Library, schema, alias_analysis=""): + """The old torch.library.define. + We're keeping this around for BC reasons + """ + + def wrap(f): + name = lib.define(schema, alias_analysis) + lib.impl(name, f) + return f + + return wrap + + +@overload +def impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Literal[None] = None, + *, + lib: Optional[Library] = None, +) -> Callable[[Callable[..., object]], None]: ... + + +@overload +def impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Callable[..., object], + *, + lib: Optional[Library] = None, +) -> None: ... + + +# Deprecated BC API +@overload +def impl( + lib: Library, + name: str, + dispatch_key: str = "", +) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ... + + +@functools.singledispatch +def impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Optional[Callable[_P, _T]] = None, + *, + lib: Optional[Library] = None, +) -> object: + """Register an implementation for a device type for this operator. + + You may pass "default" for ``types`` to register this implementation as the + default implementation for ALL device types. + Please only use this if the implementation truly supports all device types; + for example, this is true if it is a composition of built-in PyTorch operators. + + This API may be used as a decorator. You can use nested decorators + with this API provided they return a function and are placed inside + this API (see Example 2). + + Some valid types are: "cpu", "cuda", "xla", "mps", "ipu", "xpu". + + Args: + qualname (str): Should be a string that looks like "namespace::operator_name". + types (str | Sequence[str]): The device types to register an impl to. + lib (Optional[Library]): If provided, the lifetime of this registration + will be tied to the lifetime of the Library object. + + Examples: + >>> import torch + >>> import numpy as np + >>> # Example 1: Register function. + >>> # Define the operator + >>> torch.library.define("mylib::mysin", "(Tensor x) -> Tensor") + >>> + >>> # Add implementations for the cpu device + >>> @torch.library.impl("mylib::mysin", "cpu") + >>> def f(x): + >>> return torch.from_numpy(np.sin(x.numpy())) + >>> + >>> x = torch.randn(3) + >>> y = torch.ops.mylib.mysin(x) + >>> assert torch.allclose(y, x.sin()) + >>> + >>> # Example 2: Register function with decorator. + >>> def custom_decorator(func): + >>> def wrapper(*args, **kwargs): + >>> return func(*args, **kwargs) + 1 + >>> return wrapper + >>> + >>> # Define the operator + >>> torch.library.define("mylib::sin_plus_one", "(Tensor x) -> Tensor") + >>> + >>> # Add implementations for the operator + >>> @torch.library.impl("mylib::sin_plus_one", "cpu") + >>> @custom_decorator + >>> def f(x): + >>> return torch.from_numpy(np.sin(x.numpy())) + >>> + >>> # Call the new operator from torch.ops. + >>> x = torch.randn(3) + >>> + >>> y1 = torch.ops.mylib.sin_plus_one(x) + >>> y2 = torch.sin(x) + 1 + >>> assert torch.allclose(y1, y2) + """ + return _impl(qualname, types, func, lib=lib, disable_dynamo=False) + + +if not TYPE_CHECKING: + + @impl.register + def _( + lib: Library, name: str, dispatch_key: str = "" + ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + """Legacy torch.library.impl API. Kept around for BC""" + + def wrap(f: Callable[_P, _T]) -> Callable[_P, _T]: + lib.impl(name, f, dispatch_key) + return f + + return wrap + + +@overload +def _impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Literal[None] = None, + *, + lib: Optional[Library] = None, + disable_dynamo: bool = False, +) -> Callable[[Callable[..., object]], None]: ... + + +@overload +def _impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Callable[..., object], + *, + lib: Optional[Library] = None, + disable_dynamo: bool = False, +) -> None: ... + + +def _impl( + qualname: str, + types: Union[str, Sequence[str]], + func: Optional[Callable[..., object]] = None, + *, + lib: Optional[Library] = None, + disable_dynamo: bool = False, +) -> Optional[Callable[[Callable[..., object]], None]]: + # See impl() + if isinstance(types, str): + types = (types,) + keys = set({}) + for typ in types: + is_dispatch_key = torch._C._parse_dispatch_key(typ) + if is_dispatch_key: + # We also support passing a DispatchKey to impl. Please prefer using + # the higher-level torch.library APIs and only pass DispatchKey to + # torch.library.impl with caution (or even better, don't use this + # option and file an issue on GitHub for what you need). + # We don't advertise this to users because + # it is very easy to shoot yourself in the foot. + keys.add(typ) + else: + keys.add(_device_type_to_key(typ)) + + def register_(func: Callable[..., object]) -> None: + namespace, _ = torch._library.utils.parse_namespace(qualname) + + if lib is None: + use_lib = Library(namespace, "FRAGMENT") + _keep_alive.append(use_lib) + else: + use_lib = lib + if disable_dynamo: + + @torch._disable_dynamo + def func_no_dynamo(*args, **kwargs): + return func(*args, **kwargs) + + for key in keys: + use_lib.impl(qualname, func_no_dynamo, key) + else: + for key in keys: + use_lib.impl(qualname, func, key) + + if func is None: + return register_ + else: + register_(func) + return None + + +def _device_type_to_key(device_type: str) -> str: + if device_type == "default": + # This is technically not correct, because although all device_type + # DispatchKeys are included in CompositeExplicitAutograd, + # not everything in CompositeExplicitAutograd is associated with a + # device_type. I don't really care that much about the difference. + return "CompositeExplicitAutograd" + return torch._C._dispatch_key_for_device(device_type) + + +@deprecated( + "`torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that " + "instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.", + category=FutureWarning, +) +def impl_abstract(qualname, func=None, *, lib=None, _stacklevel=1): + r"""This API was renamed to :func:`torch.library.register_fake` in PyTorch 2.4. + Please use that instead. + """ + if func is not None: + _stacklevel = _stacklevel + 1 + return register_fake(qualname, func, lib=lib, _stacklevel=_stacklevel) + + +_op_identifier = Union[ + str, "torch._ops.OpOverload", "torch._library.custom_ops.CustomOpDef" +] + + +def register_kernel( + op: _op_identifier, + device_types: device_types_t, + func: Optional[Callable] = None, + /, + *, + lib: Optional[Library] = None, +): + """Register an implementation for a device type for this operator. + + Some valid device_types are: "cpu", "cuda", "xla", "mps", "ipu", "xpu". + This API may be used as a decorator. + + Args: + op (str | OpOverload): The operator to register an impl to. + device_types (None | str | Sequence[str]): The device_types to register an impl to. + If None, we will register to all device types -- please only use + this option if your implementation is truly device-type-agnostic. + func (Callable): The function to register as the implementation for + the given device types. + lib (Optional[Library]): If provided, the lifetime of this registration + + Examples:: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> import torch + >>> from torch import Tensor + >>> from torch.library import custom_op + >>> import numpy as np + >>> + >>> # Create a custom op that works on cpu + >>> @custom_op("mylib::numpy_sin", mutates_args=(), device_types="cpu") + >>> def numpy_sin(x: Tensor) -> Tensor: + >>> x_np = x.numpy() + >>> y_np = np.sin(x_np) + >>> return torch.from_numpy(y_np) + >>> + >>> # Add implementations for the cuda device + >>> @torch.library.register_kernel("mylib::numpy_sin", "cuda") + >>> def _(x): + >>> x_np = x.cpu().numpy() + >>> y_np = np.sin(x_np) + >>> return torch.from_numpy(y_np).to(device=x.device) + >>> + >>> x_cpu = torch.randn(3) + >>> x_cuda = x_cpu.cuda() + >>> assert torch.allclose(numpy_sin(x_cpu), x_cpu.sin()) + >>> assert torch.allclose(numpy_sin(x_cuda), x_cuda.sin()) + + """ + + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError( + f"register_kernel({op}): got unexpected type for op: {type(op)}" + ) + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + return opdef.register_kernel(device_types, func) + assert isinstance(op, str) + if device_types is None: + device_types = "CompositeExplicitAutograd" + + return _impl(op, device_types, func, lib=lib, disable_dynamo=True) + + +def register_autocast( + op: _op_identifier, + device_type: str, + cast_inputs: _dtype, + /, + *, + lib: Optional[Library] = None, +): + r"""Register an autocast dispatch rule for this custom op. + + Valid `device_type` include: "cpu" and "cuda". + + Args: + op (str | OpOverload): The operator to register an autocast dispatch rule to. + device_type(str): Device type to use. 'cuda' or 'cpu'. + The type is the same as the `type` attribute of a :class:`torch.device`. + Thus, you may obtain the device type of a tensor using `Tensor.device.type`. + cast_inputs (:class:`torch.dtype`): When custom op runs in an autocast-enabled region, + casts incoming floating-point Tensors to the target dtype (non-floating-point Tensors + are not affected), then executes custom op with autocast disabled. + lib (Optional[Library]): If provided, the lifetime of this registration + + Examples:: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> import torch + >>> from torch import Tensor + >>> from torch.library import custom_op + >>> + >>> # Create a custom op that works on cuda + >>> @torch.library.custom_op("mylib::my_sin", mutates_args=()) + >>> def my_sin(x: Tensor) -> Tensor: + >>> return torch.sin(x) + >>> + >>> # Register autocast dispatch rule for the cuda device + >>> torch.library.register_autocast("mylib::my_sin", "cuda", torch.float16) + >>> + >>> x = torch.randn(3, dtype=torch.float32, device="cuda") + >>> with torch.autocast("cuda", dtype=torch.float16): + >>> y = torch.ops.mylib.my_sin(x) + >>> assert y.dtype == torch.float16 + + """ + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError( + f"register_autocast({op}): got unexpected type for op: {type(op)}" + ) + if device_type not in ["cpu", "cuda"]: + raise ValueError(f"Unknown device type: {device_type}") + + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + return opdef.register_autocast(device_type, cast_inputs) + + assert isinstance(op, str) + qualname = op + _op = torch._library.utils.lookup_op(qualname) + + namespace, opname = torch._library.utils.parse_namespace(qualname) + if lib is None: + lib = Library(namespace, "FRAGMENT") + _keep_alive.append(lib) + + def _maybe_override_py_impl(op: torch._ops.OpOverload, dispatch_key): + def inner(kernel): + if op.has_kernel_for_dispatch_key(dispatch_key): + op.py_kernels.pop(dispatch_key) + return op.py_impl(dispatch_key)(kernel) + + return inner + + @_maybe_override_py_impl(_op, torch._C.DispatchKey.AutocastCPU) + @_maybe_override_py_impl(_op, torch._C.DispatchKey.AutocastCUDA) + def _autocast_py_impl(*args, **kwargs): + assert len(kwargs) == 0, "Custom ops do not support kwargs yet." + autocast_keyset = torch._C.DispatchKeySet( + torch._C.DispatchKey.AutocastCPU + ) | torch._C.DispatchKeySet(torch._C.DispatchKey.AutocastCUDA) + with torch._C._ExcludeDispatchKeyGuard(autocast_keyset): + return _op(*_cast(args, device_type, cast_inputs)) + + def kernel(_, *args, **kwargs): + assert len(kwargs) == 0, "Custom ops do not support kwargs yet." + return _autocast_py_impl(*args, **kwargs) + + if device_type == "cuda": + return lib.impl(opname, kernel, "AutocastCUDA", with_keyset=True) + else: + # device_type is "cpu" + return lib.impl(opname, kernel, "AutocastCPU", with_keyset=True) + + +def register_fake( + op: _op_identifier, + func: Optional[Callable] = None, + /, + *, + lib: Optional[Library] = None, + _stacklevel: int = 1, + allow_override: bool = False, +): + r"""Register a FakeTensor implementation ("fake impl") for this operator. + + Also sometimes known as a "meta kernel", "abstract impl". + + An "FakeTensor implementation" specifies the behavior of this operator on + Tensors that carry no data ("FakeTensor"). Given some input Tensors with + certain properties (sizes/strides/storage_offset/device), it specifies + what the properties of the output Tensors are. + + The FakeTensor implementation has the same signature as the operator. + It is run for both FakeTensors and meta tensors. To write a FakeTensor + implementation, assume that all Tensor inputs to the operator are + regular CPU/CUDA/Meta tensors, but they do not have storage, and + you are trying to return regular CPU/CUDA/Meta tensor(s) as output. + The FakeTensor implementation must consist of only PyTorch operations + (and may not directly access the storage or data of any input or + intermediate Tensors). + + This API may be used as a decorator (see examples). + + For a detailed guide on custom ops, please see + https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html + + Args: + op_name: Operator name (along with the overload) or OpOverload object. + func: Fake tensor implementation. + lib (Optional[Library]): Library to register the fake tensor to. + allow_override: Flag controlling if we want to override an + existing registered fake impl. This is by default off, + and will error you're trying to register a fake impl to + an operator that already has a fake impl. This also only + applies if the custom operator was not created via + torch.library.custom_op, as overriding and existing fake + impl is already allowed. + + Examples: + >>> import torch + >>> import numpy as np + >>> from torch import Tensor + >>> + >>> # Example 1: an operator without data-dependent output shape + >>> @torch.library.custom_op("mylib::custom_linear", mutates_args=()) + >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor: + >>> raise NotImplementedError("Implementation goes here") + >>> + >>> @torch.library.register_fake("mylib::custom_linear") + >>> def _(x, weight, bias): + >>> assert x.dim() == 2 + >>> assert weight.dim() == 2 + >>> assert bias.dim() == 1 + >>> assert x.shape[1] == weight.shape[1] + >>> assert weight.shape[0] == bias.shape[0] + >>> assert x.device == weight.device + >>> + >>> return (x @ weight.t()) + bias + >>> + >>> with torch._subclasses.fake_tensor.FakeTensorMode(): + >>> x = torch.randn(2, 3) + >>> w = torch.randn(3, 3) + >>> b = torch.randn(3) + >>> y = torch.ops.mylib.custom_linear(x, w, b) + >>> + >>> assert y.shape == (2, 3) + >>> + >>> # Example 2: an operator with data-dependent output shape + >>> @torch.library.custom_op("mylib::custom_nonzero", mutates_args=()) + >>> def custom_nonzero(x: Tensor) -> Tensor: + >>> x_np = x.numpy(force=True) + >>> res = np.stack(np.nonzero(x_np), axis=1) + >>> return torch.tensor(res, device=x.device) + >>> + >>> @torch.library.register_fake("mylib::custom_nonzero") + >>> def _(x): + >>> # Number of nonzero-elements is data-dependent. + >>> # Since we cannot peek at the data in an fake impl, + >>> # we use the ctx object to construct a new symint that + >>> # represents the data-dependent size. + >>> ctx = torch.library.get_ctx() + >>> nnz = ctx.new_dynamic_size() + >>> shape = [nnz, x.dim()] + >>> result = x.new_empty(shape, dtype=torch.int64) + >>> return result + >>> + >>> from torch.fx.experimental.proxy_tensor import make_fx + >>> + >>> x = torch.tensor([0, 1, 2, 3, 4, 0]) + >>> trace = make_fx(torch.ops.mylib.custom_nonzero, tracing_mode="symbolic")(x) + >>> trace.print_readable() + >>> + >>> assert torch.allclose(trace(x), torch.ops.mylib.custom_nonzero(x)) + + """ + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError(f"register_fake({op}): got unexpected type for op: {type(op)}") + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + if func is None: + return opdef.register_fake + else: + return opdef.register_fake(func) + assert isinstance(op, str) + + stacklevel = _stacklevel + + def register(func): + namespace, op_name = torch._library.utils.parse_namespace(op) + if lib is None: + use_lib = Library(namespace, "FRAGMENT") + _keep_alive.append(use_lib) + else: + use_lib = lib + use_lib._register_fake( + op_name, func, _stacklevel=stacklevel + 1, allow_override=allow_override + ) + return func + + if func is None: + return register + else: + stacklevel += 1 + return register(func) + + +def register_autograd( + op: _op_identifier, + backward: Callable, + /, + *, + setup_context: Optional[Callable] = None, + lib=None, +) -> None: + r"""Register a backward formula for this custom op. + + In order for an operator to work with autograd, you need to register + a backward formula: + 1. You must tell us how to compute gradients during the backward pass + by providing us a "backward" function. + 2. If you need any values from the forward to compute gradients, you can + use `setup_context` to save values for backward. + + ``backward`` runs during the backward pass. It accepts ``(ctx, *grads)``: + - ``grads`` is one or more gradients. The number of gradients matches + the number of outputs of the operator. + The ``ctx`` object is `the same ctx object `_ used by + :class:`torch.autograd.Function`. The semantics of ``backward_fn`` are the + same as :meth:`torch.autograd.Function.backward`. + + ``setup_context(ctx, inputs, output)`` runs during the forward pass. + Please save quantities needed for backward onto the ``ctx`` object via + either :meth:`torch.autograd.function.FunctionCtx.save_for_backward` + or assigning them as attributes of ``ctx``. If your custom op has + kwarg-only arguments, we expect the signature of ``setup_context`` + to be ``setup_context(ctx, inputs, keyword_only_inputs, output)``. + + Both ``setup_context_fn`` and ``backward_fn`` must be traceable. That is, + they may not directly access :meth:`torch.Tensor.data_ptr` and they must + not depend on or mutate global state. If you need a non-traceable backward, + you can make it a separate custom_op that you call inside ``backward_fn``. + + If you need different autograd behavior on different devices, then we + recommend creating two different custom operators, one for each device + that needs different behavior, and switching between them at runtime. + + Examples: + >>> import torch + >>> import numpy as np + >>> from torch import Tensor + >>> + >>> @torch.library.custom_op("mylib::numpy_sin", mutates_args=()) + >>> def numpy_sin(x: Tensor) -> Tensor: + >>> x_np = x.cpu().numpy() + >>> y_np = np.sin(x_np) + >>> return torch.from_numpy(y_np).to(device=x.device) + >>> + >>> def setup_context(ctx, inputs, output) -> Tensor: + >>> x, = inputs + >>> ctx.save_for_backward(x) + >>> + >>> def backward(ctx, grad): + >>> x, = ctx.saved_tensors + >>> return grad * x.cos() + >>> + >>> torch.library.register_autograd( + ... "mylib::numpy_sin", backward, setup_context=setup_context + ... ) + >>> + >>> x = torch.randn(3, requires_grad=True) + >>> y = numpy_sin(x) + >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y)) + >>> assert torch.allclose(grad_x, x.cos()) + >>> + >>> # Example with a keyword-only arg + >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=()) + >>> def numpy_mul(x: Tensor, *, val: float) -> Tensor: + >>> x_np = x.cpu().numpy() + >>> y_np = x_np * val + >>> return torch.from_numpy(y_np).to(device=x.device) + >>> + >>> def setup_context(ctx, inputs, keyword_only_inputs, output) -> Tensor: + >>> ctx.val = keyword_only_inputs["val"] + >>> + >>> def backward(ctx, grad): + >>> return grad * ctx.val + >>> + >>> torch.library.register_autograd( + ... "mylib::numpy_mul", backward, setup_context=setup_context + ... ) + >>> + >>> x = torch.randn(3, requires_grad=True) + >>> y = numpy_mul(x, val=3.14) + >>> (grad_x,) = torch.autograd.grad(y, x, torch.ones_like(y)) + >>> assert torch.allclose(grad_x, torch.full_like(x, 3.14)) + + """ + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError( + f"register_autograd({op}): got unexpected type for op: {type(op)}" + ) + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + opdef.register_autograd(backward, setup_context=setup_context) + return + + assert isinstance(op, str) + qualname = op + op = torch._library.utils.lookup_op(qualname) + schema = op._schema + if not _library.utils.is_functional_schema(schema): + raise RuntimeError( + f"Cannot register autograd formula for non-functional operator " + f"{op} with schema {schema}. Please create " + f"a functional operator and register an autograd formula for that." + ) + if _library.utils.has_kwarg_only_tensors(schema): + raise NotImplementedError( + f"register_autograd with kwarg-only Tensor args. In the original " + f"definition of the op, please make your tensors not kwarg-only. " + f"Got: {schema}" + ) + + info = _library.autograd.Info(backward, setup_context) + autograd_kernel = _library.autograd.make_autograd_impl(op, info) + namespace, opname = torch._library.utils.parse_namespace(qualname) + if lib is None: + lib = Library(namespace, "FRAGMENT") + _keep_alive.append(lib) + lib.impl(opname, autograd_kernel, "Autograd", with_keyset=True) + + +def register_torch_dispatch( + op: _op_identifier, + torch_dispatch_class: Any, + func: Optional[Callable] = None, + /, + *, + lib: Optional[Library] = None, +): + r"""Registers a torch_dispatch rule for the given operator and ``torch_dispatch_class``. + + This allows for open registration to specify the behavior between the operator + and the ``torch_dispatch_class`` without needing to modify the ``torch_dispatch_class`` + or the operator directly. + + The ``torch_dispatch_class`` is either a Tensor subclass with ``__torch_dispatch__`` or a + TorchDispatchMode. + + If it is a Tensor subclass, we expect ``func`` to have the following signature: + ``(cls, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any`` + + If it is a TorchDispatchMode, we expect ``func`` to have the following signature: + ``(mode, func: OpOverload, types: Tuple[type, ...], args, kwargs) -> Any`` + + ``args`` and ``kwargs`` will have been normalized the same way they are + in ``__torch_dispatch__`` (see :ref:`torch-dispatch-calling-convention`). + + Examples: + + >>> import torch + >>> + >>> @torch.library.custom_op("mylib::foo", mutates_args={}) + >>> def foo(x: torch.Tensor) -> torch.Tensor: + >>> return x.clone() + >>> + >>> class MyMode(torch.utils._python_dispatch.TorchDispatchMode): + >>> def __torch_dispatch__(self, func, types, args=(), kwargs=None): + >>> return func(*args, **kwargs) + >>> + >>> @torch.library.register_torch_dispatch("mylib::foo", MyMode) + >>> def _(mode, func, types, args, kwargs): + >>> x, = args + >>> return x + 1 + >>> + >>> x = torch.randn(3) + >>> y = foo(x) + >>> assert torch.allclose(y, x) + >>> + >>> with MyMode(): + >>> y = foo(x) + >>> assert torch.allclose(y, x + 1) + + """ + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError( + f"register_torch_dispatch({op}): got unexpected type for op: {type(op)}" + ) + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + return opdef.register_torch_dispatch(torch_dispatch_class, func) + assert isinstance(op, str) + + def register(func): + namespace, op_name = torch._library.utils.parse_namespace(op) + if lib is None: + use_lib = Library(namespace, "FRAGMENT") + _keep_alive.append(use_lib) + else: + use_lib = lib + use_lib._register_torch_dispatch_rule(op_name, torch_dispatch_class, func) + return func + + if func is None: + return register + else: + return register(func) + + +def register_vmap( + op: _op_identifier, + func: Optional[Callable] = None, + /, + *, + lib=None, +): + r"""Register a vmap implementation to support :func:`torch.vmap` for this custom op. + + This API may be used as a decorator (see examples). + + In order for an operator to work with :func:`torch.vmap`, you may need to register a + vmap implementation in the following signature: + + ``vmap_func(info, in_dims: Tuple[Optional[int]], *args, **kwargs)``, + + where ``*args`` and ``**kwargs`` are the arguments and kwargs for ``op``. + We do not support kwarg-only Tensor args. + + It specifies how do we compute the batched version of ``op`` given inputs with an additional + dimension (specified by ``in_dims``). + + For each arg in ``args``, ``in_dims`` has a corresponding ``Optional[int]``. It is ``None`` + if the arg is not a Tensor or if the arg is not being vmapped over, otherwise, it is an integer + specifying what dimension of the Tensor is being vmapped over. + + ``info`` is a collection of additional metadata that may be helpful: + ``info.batch_size`` specifies the size of the dimension being vmapped over, while + ``info.randomness`` is the ``randomness`` option that was passed to :func:`torch.vmap`. + + The return of the function ``func`` is a tuple of ``(output, out_dims)``. Similar to ``in_dims``, + ``out_dims`` should be of the same structure as ``output`` and contain one ``out_dim`` + per output that specifies if the output has the vmapped dimension and what index it is in. + + Examples: + >>> import torch + >>> import numpy as np + >>> from torch import Tensor + >>> from typing import Tuple + >>> + >>> def to_numpy(tensor): + >>> return tensor.cpu().numpy() + >>> + >>> lib = torch.library.Library("mylib", "FRAGMENT") + >>> @torch.library.custom_op("mylib::numpy_cube", mutates_args=()) + >>> def numpy_cube(x: Tensor) -> Tuple[Tensor, Tensor]: + >>> x_np = to_numpy(x) + >>> dx = torch.tensor(3 * x_np ** 2, device=x.device) + >>> return torch.tensor(x_np ** 3, device=x.device), dx + >>> + >>> def numpy_cube_vmap(info, in_dims, x): + >>> result = numpy_cube(x) + >>> return result, (in_dims[0], in_dims[0]) + >>> + >>> torch.library.register_vmap(numpy_cube, numpy_cube_vmap) + >>> + >>> x = torch.randn(3) + >>> torch.vmap(numpy_cube)(x) + >>> + >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=()) + >>> def numpy_mul(x: Tensor, y: Tensor) -> Tensor: + >>> return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device) + >>> + >>> @torch.library.register_vmap("mylib::numpy_mul") + >>> def numpy_mul_vmap(info, in_dims, x, y): + >>> x_bdim, y_bdim = in_dims + >>> x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1) + >>> y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1) + >>> result = x * y + >>> result = result.movedim(-1, 0) + >>> return result, 0 + >>> + >>> + >>> x = torch.randn(3) + >>> y = torch.randn(3) + >>> torch.vmap(numpy_mul)(x, y) + + .. note:: + The vmap function should aim to preserve the semantics of the entire custom operator. + That is, ``grad(vmap(op))`` should be replaceable with a ``grad(map(op))``. + + If your custom operator has any custom behavior in the backward pass, please + keep this in mind. + + """ + if not isinstance( + op, (str, torch._ops.OpOverload, torch._library.custom_ops.CustomOpDef) + ): + raise ValueError(f"register_vmap({op}): got unexpected type for op: {type(op)}") + if isinstance(op, torch._ops.OpOverload): + op = op._name + opdef = _maybe_get_opdef(op) + if opdef is not None: + return opdef.register_vmap(func) + assert isinstance(op, str) + qualname = op + op = torch._library.utils.lookup_op(qualname) + schema = op._schema + if _library.utils.has_kwarg_only_tensors(schema): + raise NotImplementedError( + f"register_vmap with kwarg-only Tensor args. In the original " + f"definition of the op, please make your tensors not kwarg-only. " + f"Got: {schema}" + ) + + def register(func): + nonlocal op, lib + + namespace, opname = torch._library.utils.parse_namespace(qualname) + if lib is None: + lib = Library(namespace, "FRAGMENT") + _keep_alive.append(lib) + + from torch._functorch.autograd_function import custom_function_call_vmap_helper + from torch._functorch.pyfunctorch import retrieve_current_functorch_interpreter + + def wrapped_func(keyset, *args, **kwargs): + interpreter = retrieve_current_functorch_interpreter() + return custom_function_call_vmap_helper( + interpreter, func, op, *args, **kwargs + ) + + lib.impl(opname, wrapped_func, "FuncTorchBatched", with_keyset=True) + + if func is None: + return register + else: + return register(func) + + +# If the op was defined in C++, then we want to make sure there was an +# m.set_python_module(module, ...) call and that the module is the +# same as the module that called torch.library.register_fake. +def _check_pystubs_once(func, qualname, actual_module_name): + checked = False + + def inner(*args, **kwargs): + nonlocal checked + if checked: + return func(*args, **kwargs) + + op = torch._library.utils.lookup_op(qualname) + if op._defined_in_python: + checked = True + return func(*args, **kwargs) + + maybe_pystub = torch._C._dispatch_pystub( + op._schema.name, op._schema.overload_name + ) + if maybe_pystub is None: + if torch._library.utils.requires_set_python_module(): + namespace = op.namespace + cpp_filename = op._handle.debug() + raise RuntimeError( + f"Operator '{qualname}' was defined in C++ and has a Python " + f"fake impl. In this situation, we require there to also be a " + f'companion C++ `m.set_python_module("{actual_module_name}")` ' + f"call, but we could not find one. Please add that to " + f"to the top of the C++ TORCH_LIBRARY({namespace}, ...) block the " + f"operator was registered in ({cpp_filename})" + ) + else: + pystub_module = maybe_pystub[0] + if actual_module_name != pystub_module: + cpp_filename = op._handle.debug() + raise RuntimeError( + f"Operator '{qualname}' specified that its python fake impl " + f"is in the Python module '{pystub_module}' but it was actually found " + f"in '{actual_module_name}'. Please either move the fake impl " + f"or correct the m.set_python_module call ({cpp_filename})" + ) + checked = True + return func(*args, **kwargs) + + return inner + + +# NOTE [ctx inside the fake implementation] +# If a user has an operator with data-dependent output shape, then when writing +# a fake implementation they must query the current ctx and use methods on the +# ctx to construct a new unbacked symint. +# +# This is done via us setting the global_ctx_getter function every time a fake +# implementation is invoked. +def get_ctx() -> "torch._library.fake_impl.FakeImplCtx": + """get_ctx() returns the current AbstractImplCtx object. + + Calling ``get_ctx()`` is only valid inside of an fake impl + (see :func:`torch.library.register_fake` for more usage details. + """ + return torch._library.fake_impl.global_ctx_getter() + + +def get_kernel( + op: _op_identifier, dispatch_key: Union[str, torch.DispatchKey] +) -> torch._C._SafeKernelFunction: + """Returns the computed kernel for a given operator and dispatch key. + + This function retrieves the kernel that would be executed for a given + operator and dispatch key combination. The returned SafeKernelFunction + can be used to call the kernel in a boxed fashion. The intended use + case for this function is to retrieve the original kernel for a given + dispatch key and then register another kernel to the same dispatch key + that calls into the original kernel for certain cases. + + Args: + op: Operator name (along with the overload) or OpOverload object + Can be a string (e.g., "aten::add.Tensor"), an OpOverload, or a CustomOpDef. + dispatch_key (str | torch.DispatchKey): The dispatch key to get the kernel for. + Can be a string (e.g., "CPU", "CUDA") or a DispatchKey enum value. + + Returns: + torch._C._SafeKernelFunction: A safe kernel function that can be used to + call the kernel. + + Raises: + RuntimeError: If the operator does not exist. + + Example: + >>> # Get the CPU kernel for torch.add + >>> kernel = torch.library.get_kernel("aten::add.Tensor", "CPU") + >>> + >>> # You can also use DispatchKey enum + >>> kernel = torch.library.get_kernel("aten::add.Tensor", torch.DispatchKey.CPU) + >>> + >>> # Or use an OpOverload directly + >>> kernel = torch.library.get_kernel(torch.ops.aten.add.Tensor, "CPU") + >>> + >>> # Example: Using get_kernel in a custom op with conditional dispatch + >>> # Get the original kernel for torch.sin + >>> original_sin_kernel = torch.library.get_kernel("aten::sin", "CPU") + >>> + >>> # If input has negative values, use original sin, otherwise return zeros + >>> def conditional_sin_impl(dispatch_keys, x): + >>> if (x < 0).any(): + >>> return original_sin_kernel.call_boxed(dispatch_keys, x) + >>> else: + >>> return torch.zeros_like(x) + >>> + >>> lib = torch.library.Library("aten", "IMPL") + >>> # with_keyset=True so the first argument to the impl is the current DispatchKeySet + >>> which needs to be the first argument to ``kernel.call_boxed`` + >>> lib.impl("sin", conditional_sin_impl, "CPU", with_keyset=True) + >>> + >>> # Test the conditional behavior + >>> x_positive = torch.tensor([1.0, 2.0]) + >>> x_mixed = torch.tensor([-1.0, 2.0]) + >>> torch.sin(x_positive) + tensor([0., 0.]) + >>> torch.sin(x_mixed) + tensor([-0.8415, 0.9093]) + """ + if not isinstance(op, (str, torch._ops.OpOverload)): + raise ValueError(f"get_kernel({op}): got unexpected type for op: {type(op)}") + + if isinstance(op, torch._ops.OpOverload): + op = op._name + + if isinstance(dispatch_key, str): + try: + dispatch_key = torch._C.DispatchKey.__members__[dispatch_key] + except KeyError: + raise ValueError(f"Invalid dispatch key: {dispatch_key}") from None + + return torch._C._dispatch_get_computed_kernel_for_dispatch_key(op, dispatch_key) + + +_OPCHECK_DEFAULT_UTILS = ( + "test_schema", + "test_autograd_registration", + "test_faketensor", + "test_aot_dispatch_dynamic", +) + + +def opcheck( + op: Union[torch._ops.OpOverload, torch._ops.OpOverloadPacket, CustomOpDef], + args: tuple[Any, ...], + kwargs: Optional[dict[str, Any]] = None, + *, + test_utils: Union[str, Sequence[str]] = _OPCHECK_DEFAULT_UTILS, + raise_exception: bool = True, + atol=None, + rtol=None, +) -> dict[str, str]: + """Given an operator and some sample arguments, tests if the operator is + registered correctly. + + That is, when you use the torch.library/TORCH_LIBRARY APIs to create a + custom op, you specified metadata (e.g. mutability info) about the custom op + and these APIs require that the functions you pass them satisfy certain + properties (e.g. no data pointer access in the fake/meta/abstract kernel) + ``opcheck`` tests these metadata and properties. + + Concretely, we test the following: + + - test_schema: If the schema matches the implementation of + the operator. For example: if the schema specifies a Tensor is mutated, + then we check the implementation mutates the Tensor. If the schema + specifies that we return a new Tensor, then we check that the + implementation returns a new Tensor (instead of an existing one or + a view of an existing one). + - test_autograd_registration: If the operator supports training + (autograd): we check that its autograd formula is registered via + torch.library.register_autograd or a manual registration to one + or more DispatchKey::Autograd keys. Any other DispatchKey-based + registrations may lead to undefined behavior. + - test_faketensor: If the operator has a FakeTensor kernel + (and if it is correct). The FakeTensor kernel is necessary ( + but not sufficient) for the operator to work with PyTorch compilation + APIs (torch.compile/export/FX). We check that a FakeTensor kernel + (also sometimes known as a meta kernel) was registered for the + operator and that it is correct. This test takes the result of + running the operator on real tensors and the result of running + the operator on FakeTensors and checks that they have the same + Tensor metadata (sizes/strides/dtype/device/etc). + - test_aot_dispatch_dynamic: If the operator has correct behavior + with PyTorch compilation APIs (torch.compile/export/FX). + This checks that the outputs (and gradients, if applicable) are the + same under eager-mode PyTorch and torch.compile. + This test is a superset of ``test_faketensor`` and is an e2e test; + other things it tests are that the operator supports + functionalization and that the backward pass (if it exists) also + supports FakeTensor and functionalization. + + For best results, please call ``opcheck`` multiple times with a + representative set of inputs. If your operator supports + autograd, please use ``opcheck`` with inputs with ``requires_grad = True``; + if your operator supports multiple devices (e.g. CPU and CUDA), please + use ``opcheck`` with inputs on all supported devices. + + Args: + op: The operator. Must either be a function decorated with + :func:`torch.library.custom_op` or an OpOverload/OpOverloadPacket + found in torch.ops.* (e.g. torch.ops.aten.sin, torch.ops.mylib.foo) + args: The args to the operator + kwargs: The kwargs to the operator + test_utils: Tests that we should run. Default: all of them. + Example: ("test_schema", "test_faketensor") + raise_exception: If we should raise an exception on the first + error. If False, we will return a dict with information + on if each test passed or not. + rtol (Optional[float]): Relative tolerance for floating point comparisons. + If specified ``atol`` must also be specified. + If omitted, default values based on the ``dtype`` are selected + (see the table in :func:`torch.testing.assert_close`). + atol (Optional[float]): Absolute tolerance for floating point comparisons. + If specified ``rtol`` must also be specified. + If omitted, default values based on the ``dtype`` are selected + (see the table in :func:`torch.testing.assert_close`). + + .. warning:: + + opcheck and :func:`torch.autograd.gradcheck` test different things; + opcheck tests if your usage of torch.library APIs is correct while + :func:`torch.autograd.gradcheck` tests if your autograd formula is + mathematically correct. Use both to test custom ops that support + gradient computation. + + Example: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> @torch.library.custom_op("mylib::numpy_mul", mutates_args=()) + >>> def numpy_mul(x: Tensor, y: float) -> Tensor: + >>> x_np = x.numpy(force=True) + >>> z_np = x_np * y + >>> return torch.from_numpy(z_np).to(x.device) + >>> + >>> @numpy_mul.register_fake + >>> def _(x, y): + >>> return torch.empty_like(x) + >>> + >>> def setup_context(ctx, inputs, output): + >>> y, = inputs + >>> ctx.y = y + >>> + >>> def backward(ctx, grad): + >>> return grad * ctx.y, None + >>> + >>> numpy_mul.register_autograd(backward, setup_context=setup_context) + >>> + >>> sample_inputs = [ + >>> (torch.randn(3), 3.14), + >>> (torch.randn(2, 3, device='cuda'), 2.718), + >>> (torch.randn(1, 10, requires_grad=True), 1.234), + >>> (torch.randn(64, 64, device='cuda', requires_grad=True), 90.18), + >>> ] + >>> + >>> for args in sample_inputs: + >>> torch.library.opcheck(numpy_mul, args) + + """ + import torch.testing._internal.optests as optests + + return optests.opcheck( + op, + args, + kwargs, + test_utils=test_utils, + raise_exception=raise_exception, + rtol=rtol, + atol=atol, + ) diff --git a/venv/lib/python3.10/site-packages/torch/overrides.py b/venv/lib/python3.10/site-packages/torch/overrides.py new file mode 100644 index 0000000000000000000000000000000000000000..c8fd7c6a228994a22b8a613d5547454c8760831f --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/overrides.py @@ -0,0 +1,2120 @@ +""" +Python implementation of ``__torch_function__`` + +While most of the torch API and handling for ``__torch_function__`` happens +at the C++ level, some of the torch API is written in Python so we need +python-level handling for ``__torch_function__`` overrides as well. The main +developer-facing functionality in this file are handle_torch_function and +has_torch_function. See torch/functional.py and test/test_overrides.py +for usage examples. + +Note +---- +heavily inspired by NumPy's ``__array_function__`` (see: +https://github.com/pytorch/pytorch/issues/24015 and +https://www.numpy.org/neps/nep-0018-array-function-protocol.html +) + +If changing this file in a way that can affect ``__torch_function__`` overhead, +please report the benchmarks in ``benchmarks/overrides_benchmark``. See the +instructions in the ``README.md`` in that directory. +""" + +import __future__ # noqa: F404 + +import collections +import contextlib +import functools +import types +import warnings +from collections.abc import Iterable +from functools import wraps +from typing import Any, Callable, Optional, TypeVar +from typing_extensions import ParamSpec + +import torch +from torch._C import ( + _add_docstr, + _get_function_stack_at, + _has_torch_function, + _has_torch_function_unary, + _has_torch_function_variadic, + _is_torch_function_mode_enabled, + _len_torch_function_stack, + _pop_torch_function_stack, + _push_on_torch_function_stack, +) + + +__all__ = [ + "get_ignored_functions", + "get_overridable_functions", + "get_testing_overrides", + "handle_torch_function", + "has_torch_function", + "resolve_name", + "is_tensor_like", + "is_tensor_method_or_property", + "wrap_torch_function", + "enable_reentrant_dispatch", +] + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def _disable_user_warnings( + func: Callable[_P, _R], + regex: str = ".*is deprecated, please use.*", + module: str = "torch", +) -> Callable[_P, _R]: + """ + Decorator that temporarily disables ``UserWarning``s for the given ``module`` if the warning message matches the + given ``regex`` pattern. + + Arguments + --------- + func : function + Function to disable the warnings for. + regex : str + A regex pattern compilable by ``re.compile``. This is used to match the ``UserWarning`` message. + module : str + The python module to which the filtering should be restricted. + + Returns + ------- + function + The wrapped function. + """ + + @wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", category=UserWarning, message=regex, module=module + ) + return func(*args, **kwargs) + + return wrapper + + +@functools.cache +@_disable_user_warnings +def get_ignored_functions() -> set[Callable]: + """ + Return public functions that cannot be overridden by ``__torch_function__``. + + Returns + ------- + set[Callable] + A tuple of functions that are publicly available in the torch API but cannot + be overridden with ``__torch_function__``. Mostly this is because none of the + arguments of these functions are tensors or tensor-likes. + + Examples + -------- + >>> torch.Tensor.as_subclass in torch.overrides.get_ignored_functions() + True + >>> torch.add in torch.overrides.get_ignored_functions() + False + """ + Tensor = torch.Tensor + return { + torch.typename, + torch.is_tensor, + torch.is_storage, + torch.set_default_tensor_type, + torch.set_default_device, + torch.get_default_device, + torch.set_rng_state, + torch.get_rng_state, + torch.manual_seed, + torch.initial_seed, + torch.seed, + torch.save, + torch.load, + torch.set_printoptions, + torch.fork, + torch.get_default_dtype, + torch.get_num_interop_threads, + torch.get_num_threads, + torch.init_num_threads, + torch.import_ir_module, + torch.import_ir_module_from_buffer, + torch.is_anomaly_enabled, + torch.is_anomaly_check_nan_enabled, + torch.is_grad_enabled, + torch.merge_type_from_type_comment, + torch.parse_ir, + torch.parse_schema, + torch.parse_type_comment, + torch.set_anomaly_enabled, + torch.set_flush_denormal, + torch.set_num_interop_threads, + torch.set_num_threads, + torch.wait, + torch.as_tensor, + torch.from_numpy, + torch.tensor, + torch.default_generator, + torch.has_cuda, + torch.has_cudnn, + torch.has_lapack, + torch.device, + torch.dtype, + torch.finfo, + torch.has_mkl, + torch.has_mps, + torch.has_mkldnn, + torch.has_openmp, + torch.iinfo, + torch.memory_format, + torch.qscheme, + torch.set_grad_enabled, + torch.no_grad, + torch.enable_grad, + torch.inference_mode, + torch.is_inference_mode_enabled, + torch.layout, + torch.align_tensors, + torch.arange, + torch.as_strided, + torch.bartlett_window, + torch.blackman_window, + torch.broadcast_shapes, + torch.can_cast, + torch.compile, + torch.cudnn_affine_grid_generator, + torch.cudnn_batch_norm, + torch.cudnn_convolution, + torch.cudnn_convolution_transpose, + torch.cudnn_convolution_relu, + torch.cudnn_convolution_add_relu, + torch.cudnn_grid_sampler, + torch.cudnn_is_acceptable, + torch.empty, + torch.empty_permuted, + torch.empty_strided, + torch.empty_quantized, + torch.export.export, + torch.export.load, + torch.export.register_dataclass, + torch.export.save, + torch.eye, + torch.fft.fftfreq, + torch.fft.rfftfreq, + torch.from_file, + torch.full, + torch.fill, + torch.hamming_window, + torch.hann_window, + torch.kaiser_window, + torch.linspace, + torch.logspace, + torch.mkldnn_adaptive_avg_pool2d, + torch.mkldnn_convolution, + torch.mkldnn_max_pool2d, + torch.mkldnn_max_pool3d, + torch.mkldnn_linear_backward_weights, + torch.mkldnn_rnn_layer, + torch.normal, + torch.ones, + torch.promote_types, + torch.rand, + torch.randn, + torch.randint, + torch.randperm, + torch.range, + torch.result_type, + torch.scalar_tensor, + torch.sparse_coo_tensor, + torch.sparse_compressed_tensor, + torch.sparse_csr_tensor, + torch.sparse_csc_tensor, + torch.sparse_bsr_tensor, + torch.sparse_bsc_tensor, + torch.sym_constrain_range, + torch.sym_constrain_range_for_size, + torch.sym_fresh_size, + torch.tril_indices, + torch.triu_indices, + torch.vander, + torch.zeros, + torch._jit_internal.boolean_dispatch, + torch.nn.functional.assert_int_or_pair, + torch.nn.functional.upsample, + torch.nn.functional.upsample_bilinear, + torch.nn.functional.upsample_nearest, + torch.nn.functional.has_torch_function, + torch.nn.functional.has_torch_function_unary, + torch.nn.functional.has_torch_function_variadic, + torch.nn.functional.handle_torch_function, + torch.nn.functional.sigmoid, + torch.nn.functional.hardsigmoid, + torch.nn.functional.tanh, + torch.nn.functional._canonical_mask, + torch.nn.functional._none_or_dtype, + # Doesn't actually take or return tensor arguments + torch.nn.init.calculate_gain, + # These are deprecated; don't test them + torch.nn.init.uniform, + torch.nn.init.normal, + torch.nn.init.constant, + torch.nn.init.eye, + torch.nn.init.dirac, + torch.nn.init.xavier_uniform, + torch.nn.init.xavier_normal, + torch.nn.init.kaiming_uniform, + torch.nn.init.kaiming_normal, + torch.nn.init.orthogonal, + torch.nn.init.sparse, + torch.nested.to_padded_tensor, + has_torch_function, + handle_torch_function, + torch.set_autocast_enabled, + torch.is_autocast_enabled, + torch.set_autocast_dtype, + torch.get_autocast_dtype, + torch.clear_autocast_cache, + torch.set_autocast_cpu_enabled, + torch.is_autocast_cpu_enabled, + torch.set_autocast_xla_enabled, + torch.is_autocast_xla_enabled, + torch.set_autocast_ipu_enabled, + torch.is_autocast_ipu_enabled, + torch.set_autocast_cpu_dtype, + torch.get_autocast_cpu_dtype, + torch.set_autocast_ipu_dtype, + torch.get_autocast_ipu_dtype, + torch.get_autocast_gpu_dtype, + torch.set_autocast_gpu_dtype, + torch.get_autocast_xla_dtype, + torch.set_autocast_xla_dtype, + torch.autocast_increment_nesting, + torch.autocast_decrement_nesting, + torch.is_autocast_cache_enabled, + torch.set_autocast_cache_enabled, + torch.nn.functional.hardswish, + torch.is_vulkan_available, + torch.are_deterministic_algorithms_enabled, + torch.use_deterministic_algorithms, + torch.is_deterministic_algorithms_warn_only_enabled, + torch.set_deterministic_debug_mode, + torch.get_device_module, + torch.get_deterministic_debug_mode, + torch.set_float32_matmul_precision, + torch.get_float32_matmul_precision, + torch.unify_type_list, + torch.is_warn_always_enabled, + torch.set_warn_always, + torch.vitals_enabled, + torch.set_vital, + torch.read_vitals, + torch.vmap, + torch.cond, + torch.frombuffer, + torch.asarray, + torch._functional_sym_constrain_range, + torch._make_dep_token, + Tensor.__delitem__, + Tensor.__dir__, + Tensor.__getattribute__, + Tensor.__init__, + Tensor.__iter__, + Tensor.__init_subclass__, + Tensor.__delattr__, + Tensor.__setattr__, + Tensor.__torch_function__, + Tensor.__torch_dispatch__, + Tensor.__new__, + Tensor.__class__, + Tensor.__subclasshook__, + Tensor.__hash__, + Tensor.as_subclass, + Tensor.eig, + Tensor.lstsq, + Tensor.reinforce, + Tensor.new, + Tensor.new_tensor, + Tensor.new_empty, + Tensor.new_empty_strided, + Tensor.new_zeros, + Tensor.new_ones, + Tensor.new_full, + Tensor._make_subclass, + Tensor.solve, + Tensor.symeig, + Tensor.stride, + Tensor.unflatten, + Tensor.to_sparse_coo, + Tensor.to_sparse_csr, + Tensor.to_sparse_csc, + Tensor.to_sparse_bsr, + Tensor.to_sparse_bsc, + Tensor._to_sparse, + Tensor._to_sparse_csr, + Tensor._to_sparse_csc, + Tensor._to_sparse_bsr, + Tensor._to_sparse_bsc, + Tensor._typed_storage, + Tensor._reduce_ex_internal, + Tensor._fix_weakref, + Tensor._view_func, + Tensor._view_func_unsafe, + Tensor._rev_view_func_unsafe, + Tensor._make_dtensor, + Tensor._make_wrapper_subclass, + Tensor._python_dispatch.__get__, + Tensor._has_symbolic_sizes_strides.__get__, + Tensor._conj, + Tensor._conj_physical, + Tensor._lazy_clone, + Tensor._neg_view, + Tensor._is_zerotensor, + Tensor._is_all_true, + Tensor._is_any_true, + Tensor._addmm_activation, + Tensor.to_padded_tensor, + Tensor._use_count, + } + + +@functools.cache +def get_default_nowrap_functions() -> set[Callable]: + """ + Return public functions that do not wrap in a subclass when invoked by + the default ``Tensor.__torch_function__`` that preserves subclasses. Typically, + these functions represent field accesses (i.e., retrieving a Tensor that + is stored somewhere on the Tensor) as opposed to computation. Users of + these functions expect object identity to be preserved over multiple accesses + (e.g., ``a.grad is a.grad``) which cannot be upheld if we're wrapping on + the fly every time (furthermore, the tensor stored here might already be + the subclass, in which case wrapping really ought not to happen). + + Not ALL property accessors have this property; for example ``Tensor.T`` actually + just creates a new transposed tensor on the fly, and so we SHOULD interpose on + these calls (you need to check the implementation of the function to see if + this is the case or not). Additionally, if a property accessor doesn't return a Tensor, + it doesn't have to be on this list (though it is harmless if it is). + """ + Tensor = torch.Tensor + return { + Tensor._base.__get__, + Tensor.grad.__get__, + Tensor._grad.__get__, + } + + +@functools.cache +@_disable_user_warnings +def get_testing_overrides() -> dict[Callable, Callable]: + """Return a dict containing dummy overrides for all overridable functions + + Returns + ------- + Dict[Callable, Callable] + A dictionary that maps overridable functions in the PyTorch API to + lambda functions that have the same signature as the real function + and unconditionally return -1. These lambda functions are useful + for testing API coverage for a type that defines ``__torch_function__``. + + Examples + -------- + >>> import inspect + >>> my_add = torch.overrides.get_testing_overrides()[torch.add] + >>> inspect.signature(my_add) + + """ + # Every function in the PyTorchAPI that can be overridden needs an entry + # in this dict. + # + # Optimally we would use inspect to get the function signature and define + # the lambda function procedurally but that is blocked by generating + # function signatures for native kernels that can be consumed by inspect. + # See Issue #28233. + Tensor = torch.Tensor + ret: dict[Callable, Callable] = { + torch.abs: lambda input, out=None: -1, + torch.absolute: lambda input, out=None: -1, + torch.adaptive_avg_pool1d: lambda input, output_size: -1, + torch.adaptive_max_pool1d: lambda inputs, output_size: -1, + torch.acos: lambda input, out=None: -1, + torch.adjoint: lambda input: -1, + torch.arccos: lambda input, out=None: -1, + torch.acosh: lambda input, out=None: -1, + torch.arccosh: lambda input, out=None: -1, + torch.add: lambda input, other, out=None: -1, + torch.addbmm: lambda input, batch1, batch2, alpha=1, beta=1, out=None: -1, + torch.addcdiv: lambda input, tensor1, tensor2, value=1, out=None: -1, + torch.addcmul: lambda input, tensor1, tensor2, value=1, out=None: -1, + torch.addmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.addmv: lambda input, mat, vec, beta=1, alpha=1, out=None: -1, + torch.addr: lambda input, vec1, vec2, beta=1, alpha=1, out=None: -1, + torch.affine_grid_generator: lambda theta, size, align_corners: -1, + torch.all: lambda input, dim=None: -1, + torch.allclose: lambda input, other, trol=1e-05, atol=1e-08, equal_nan=False: -1, + torch.alpha_dropout: lambda input, p, train, inplace=False: -1, + torch.amax: lambda input, dim=None: -1, + torch.amin: lambda input, dim=None: -1, + torch.aminmax: lambda input, dim=None, keepdim=False, out=None: -1, + torch.angle: lambda input, out=None: -1, + torch.any: lambda input, dim=None, keepdim=False, out=None: -1, + torch.argmax: lambda input: -1, + torch.argmin: lambda input: -1, + torch.argsort: lambda input, dim=None: -1, + torch.asin: lambda input, out=None: -1, + torch._assert_async: lambda input, msg: -1, + torch.arcsin: lambda input, out=None: -1, + torch.asinh: lambda input, out=None: -1, + torch.arcsinh: lambda input, out=None: -1, + torch.atan: lambda input, out=None: -1, + torch.arctan: lambda input, out=None: -1, + torch.atan2: lambda input, other, out=None: -1, + torch.arctan2: lambda input, other, out=None: -1, + torch.atanh: lambda input, out=None: -1, + torch.arctanh: lambda input, out=None: -1, + torch.atleast_1d: lambda *tensors: -1, + torch.atleast_2d: lambda *tensors: -1, + torch.atleast_3d: lambda *tensors: -1, + torch.avg_pool1d: lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True: -1, + torch.baddbmm: lambda input, batch1, batch2, alpha=1, beta=1, out=None: -1, + torch.batch_norm: lambda input, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled: -1, + torch.batch_norm_backward_elemt: lambda grad_out, input, mean, invstd, weight, sum_dy, sum_dy_xmu, count_tensor: -1, + torch.batch_norm_backward_reduce: lambda grad_out, input, mean, invstd, weight, input_g, weight_g, bias_g: -1, + torch.batch_norm_elemt: lambda input, weight, bias, mean, invstd, eps: -1, + torch.batch_norm_gather_stats: lambda input, mean, invstd, running_mean, running_var, momentum, eps, count: -1, + torch.batch_norm_gather_stats_with_counts: lambda input, mean, invstd, running_mean, running_var, momentum, eps, count: -1, + torch.batch_norm_stats: lambda input, eps: -1, + torch.batch_norm_update_stats: lambda input, running_mean, running_var, momentum: -1, + torch.bernoulli: lambda input, generator=None, out=None: -1, + torch.bilinear: lambda input1, input2, weight, bias: -1, + torch.binary_cross_entropy_with_logits: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean", pos_weight=None: -1 + ), + torch.bincount: lambda input, weights=None, minlength=0: -1, + torch.binomial: lambda count, prob, generator=None: -1, + torch.bitwise_and: lambda input, other, out=None: -1, + torch.bitwise_not: lambda input, out=None: -1, + torch.bitwise_or: lambda input, other, out=None: -1, + torch.bitwise_xor: lambda input, other, out=None: -1, + torch.bitwise_left_shift: lambda input, other, out=None: -1, + torch.bitwise_right_shift: lambda input, other, out=None: -1, + torch.block_diag: lambda *tensors: -1, + torch.bmm: lambda input, mat2, out_dtype=None, out=None: -1, + torch.broadcast_tensors: lambda *tensors: -1, + torch.broadcast_to: lambda self, size: -1, + torch.bucketize: lambda input, boundaries, out_int32=False, right=False, out=None: -1, + torch.cartesian_prod: lambda *tensors: -1, + torch.cat: lambda tensors, dim=0, out=None: -1, + torch.concat: lambda tensors, dim=0, out=None: -1, # alias for torch.cat + torch.concatenate: lambda tensors, dim=0, out=None: -1, # alias for torch.concatenate + torch.cdist: lambda x1, x2, p=2.0, compute_mode="use_mm_for_euclid_dist_if_necessary": -1, + torch.ceil: lambda input, out=None: -1, + torch.celu: lambda input, alpha=1.0, inplace=False: -1, + torch.chain_matmul: lambda *matrices, out=None: -1, + torch.channel_shuffle: lambda input, groups: -1, + torch.cholesky: lambda input, upper=False, out=None: -1, + torch.linalg.cholesky: lambda input, out=None: -1, + torch.linalg.cholesky_ex: lambda input, check_errors=False, out=None: -1, + torch.cholesky_inverse: lambda input, upper=False, out=None: -1, + torch.cholesky_solve: lambda input1, input2, upper=False, out=None: -1, + torch.choose_qparams_optimized: lambda input, numel, n_bins, ratio, bit_width: -1, + torch.chunk: lambda input, chunks, dim=0: -1, + torch.clamp: lambda input, min=None, max=None, out=None: -1, + torch.clip: lambda input, min=None, max=None, out=None: -1, + torch.clamp_min: lambda input, min, out=None: -1, + torch.clamp_max: lambda input, max, out=None: -1, + torch.column_stack: lambda tensors, out=None: -1, + torch.cov: lambda input, correction=1, fweights=None, aweights=None: -1, + torch.clone: lambda input: -1, + torch.combinations: lambda input, r=2, with_replacement=False: -1, + torch.complex: lambda real, imag: -1, + torch.copysign: lambda input, other, out=None: -1, + torch.polar: lambda abs, ang: -1, + torch.linalg.cond: lambda input, ord=None: -1, + torch.conj: lambda input, out=None: -1, + torch.conj_physical: lambda input, out=None: -1, + torch.resolve_conj: lambda input, out=None: -1, + torch.resolve_neg: lambda input, out=None: -1, + torch.constant_pad_nd: lambda input, pad, value=0: -1, + torch.conv1d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.conv2d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.conv3d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.convolution: lambda input, weight, bias, stride, padding, dilation, transposed, output_adding, groups: -1, + torch.conv_tbc: lambda input, weight, bias, pad=0: -1, + torch.conv_transpose1d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.conv_transpose2d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.conv_transpose3d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.corrcoef: lambda input: -1, + torch.cos: lambda input, out=None: -1, + torch.cosine_embedding_loss: lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1, + torch.cosh: lambda input, out=None: -1, + torch.cosine_similarity: lambda x1, x2, dim=1, eps=1e-8: -1, + torch.count_nonzero: lambda input: -1, + torch.cross: lambda input, other, dim=None, out=None: -1, + torch.linalg.cross: lambda input, other, dim=-1, out=None: -1, + torch.ctc_loss: ( + lambda log_probs, targets, input_lengths, target_lengths, blank=0, reduction="mean", zero_infinity=False: -1 + ), + torch.cummax: lambda input, dim, out=None: -1, + torch.cummin: lambda input, dim, out=None: -1, + torch.cumprod: lambda input, dim, out=None, dtype=None: -1, + torch.cumsum: lambda input, dim, out=None, dtype=None: -1, + torch.cumulative_trapezoid: lambda y, x=None, dim=-1: -1, + torch.logcumsumexp: lambda input, dim, out=None: -1, + torch.deg2rad: lambda input, out=None: -1, + torch.dequantize: lambda input: -1, + torch.det: lambda input: -1, + torch.linalg.det: lambda input: -1, # alias for torch.det # type: ignore[attr-defined] + torch.detach: lambda input: -1, + torch.diag: lambda input, diagonal=0, out=None: -1, + torch.diag_embed: lambda input, diagonal=0, out=None: -1, + torch.diagflat: lambda input, offset=0: -1, + torch.diff: lambda input, n=1, dim=-1, prepend=None, append=None, out=None: -1, + torch.diagonal: lambda input, offset=0, dim1=0, dim2=1: -1, + torch.linalg.diagonal: lambda input, offset=0, dim1=-2, dim2=-1: -1, + torch.diagonal_scatter: lambda input, src, offset=0, dim1=0, dim2=1: -1, + torch.as_strided_scatter: lambda self, src, size, stride, storage_offset=None: -1, + torch.digamma: lambda input, out=None: -1, + torch.dist: lambda input, other, p=2: -1, + torch.div: lambda input, other, rounding_mode=None, out=None: -1, + torch.divide: lambda input, other, rounding_mode=None, out=None: -1, + torch.dot: lambda input, other, out=None: -1, + torch.dropout: lambda input, p, train, inplace=False: -1, + torch.dsmm: lambda input, mat2, out_dtype=None: -1, + torch.hsmm: lambda mat1, mat2: -1, + torch.dsplit: lambda input, indices_or_sections: -1, + torch.dstack: lambda tensors, out=None: -1, + torch.linalg.eig: lambda input, out=None: -1, + torch.linalg.eigvals: lambda input, out=None: -1, + torch.linalg.eigh: lambda input, UPLO="L", out=None: -1, + torch.linalg.eigvalsh: lambda input, UPLO="L", out=None: -1, + torch.einsum: lambda equation, *operands: -1, + torch.embedding: ( + lambda input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False: -1 # noqa: B950 + ), + torch.embedding_bag: ( + lambda input, weight, offsets, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode="mean", sparse=False, per_sample_weights=None, padding_idx=None: -1 # noqa: B950 + ), + torch.empty_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.eq: lambda input, other, out=None: -1, + torch.equal: lambda input, other: -1, + torch.erf: lambda input, out=None: -1, + torch.erfc: lambda input, out=None: -1, + torch.erfinv: lambda input, out=None: -1, + torch.exp: lambda input, out=None: -1, + torch.exp2: lambda input, out=None: -1, + torch.expm1: lambda input, out=None: -1, + torch.fake_quantize_per_channel_affine: lambda input, scale, zero_point, axis, quant_min, quant_max: -1, + torch.fake_quantize_per_tensor_affine: lambda input, scale, zero_point, quant_min, quant_max: -1, + torch.fused_moving_avg_obs_fake_quant: ( + lambda x, observer_on, fake_quant_on, averaging_const, running_min, running_max, scale, zero_point, quant_min, quant_max, ch_axis, per_row_fake_quant=False, symmetric_quant=False: -1 # noqa: B950 + ), + torch.fbgemm_linear_fp16_weight: lambda input, packed_weight, bias, output: -1, + torch.fbgemm_linear_fp16_weight_fp32_activation: lambda input, packed_weight, bias, output: -1, + torch.fbgemm_linear_int8_weight: lambda input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias: -1, # noqa: B950 + torch.fbgemm_linear_int8_weight_fp32_activation: ( + lambda input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias: -1 + ), + torch.fbgemm_linear_quantize_weight: lambda input: -1, + torch.fbgemm_pack_gemm_matrix_fp16: lambda input: -1, + torch.fbgemm_pack_quantized_matrix: lambda input, a, b: -1, + torch.feature_alpha_dropout: lambda input, p, train: -1, + torch.feature_dropout: lambda input, p, train: -1, + torch.fft.ifft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.rfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.irfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.hfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.ihfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.hfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.ihfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.hfftn: lambda input, s=None, dim=-1, norm=None: -1, + torch.fft.ihfftn: lambda input, s=None, dim=-1, norm=None: -1, + torch.fft.fftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.ifftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.rfftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.irfftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.fft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.ifft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.rfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.irfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.fftshift: lambda input, dim=None: -1, + torch.fft.ifftshift: lambda input, dim=None: -1, + torch.fft.fft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fix: lambda input, out=None: -1, + torch.flatten: lambda input, start_dim=0, end_dim=-1: -1, + torch.flip: lambda input, dims: -1, + torch.fliplr: lambda input: -1, + torch.flipud: lambda input: -1, + torch.frobenius_norm: lambda input, dim=None, keepdim=False, out=None: -1, + torch.floor: lambda input, out=None: -1, + torch.floor_divide: lambda input, other: -1, + torch.float_power: lambda input, exponent, out=None: -1, + torch.fmod: lambda input, other, out=None: -1, + torch.frac: lambda input, out=None: -1, + torch.frexp: lambda input, out=None: -1, + torch.full_like: lambda input, fill_value, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False: -1, # noqa: B950 + torch._functional_assert_async: lambda input, msg, dep_token: -1, + torch.lu_unpack: lambda LU_data, LU_pivots, unpack_data=True, unpack_pivots=True: -1, + torch.gather: lambda input, dim, index, out=None, sparse_grad=False: -1, + torch.gcd: lambda input, other, out=None: -1, + torch.ge: lambda input, other, out=None: -1, + torch.get_device: lambda input: -1, + torch.greater_equal: lambda input, other, out=None: -1, + torch.geqrf: lambda input, out=None: -1, + torch.i0: lambda input, out=None: -1, + torch.inner: lambda input, other, out=None: -1, + torch.outer: lambda input, vec2, out=None: -1, + torch.ger: lambda input, vec2, out=None: -1, # alias for torch.outer + torch.gradient: lambda input, spacing=None, dim=None, edge_order=1: -1, + torch.grid_sampler: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.grid_sampler_2d: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.grid_sampler_3d: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.group_norm: lambda input, num_groups, weight=None, bias=None, eps=1e-05, cudnn_enabled=True: -1, + torch.gru: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, + torch.gru_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.gt: lambda input, other, out=None: -1, + torch.greater: lambda input, other, out=None: -1, + torch.hardshrink: lambda input, lambd=0.5: -1, + torch.hash_tensor: lambda input, dim=None, keepdim=False, mode=0, out=None: -1, + torch.heaviside: lambda input, values, out=None: -1, + torch.hinge_embedding_loss: lambda input, target, margin=1.0, size_average=None, reduce=None, reduction="mean": -1, # noqa: B950 + torch.histc: lambda input, bins=100, min=0, max=0, out=None: -1, + torch.histogram: lambda input, bins=100, min=None, max=None, weight=None, density=False, out=None: -1, + torch.histogramdd: lambda input, bins, range=None, weight=None, density=False: -1, + torch.linalg.householder_product: lambda input, tau: -1, + torch.hspmm: lambda mat1, mat2, out=None: -1, + torch.hsplit: lambda input, indices_or_sections: -1, + torch.hstack: lambda tensors, out=None: -1, + torch.hypot: lambda input, other, out=None: -1, + torch.igamma: lambda input, other, out=None: -1, + torch.igammac: lambda input, other, out=None: -1, + torch.imag: lambda input, out=None: -1, + torch.index_add: lambda input, dim, index, source: -1, + torch.index_copy: lambda input, dim, index, source: -1, + torch.index_put: lambda input, indices, values, accumulate=False: -1, + torch.index_select: lambda input, dim, index, out=None: -1, + torch.index_fill: lambda input, dim, index, value: -1, + torch.index_reduce: lambda input, dim, index, source, reduce, include_input=True: -1, + torch.isfinite: lambda tensor: -1, + torch.isin: lambda e, te, assume_unique=False, invert=False: -1, + torch.isinf: lambda tensor: -1, + torch.isreal: lambda tensor: -1, + torch.isposinf: lambda input, out=None: -1, + torch.isneginf: lambda input, out=None: -1, + torch.instance_norm: ( + lambda input, running_mean, running_var, weight, bias, use_input_stats, momentum, eps, cudnn_enabled: -1 + ), + torch.int_repr: lambda input: -1, + torch.inverse: lambda input, out=None: -1, + torch.linalg.inv: lambda input, out=None: -1, + torch.linalg.inv_ex: lambda input, check_errors=False, out=None: -1, + torch.is_complex: lambda input: -1, + torch.is_conj: lambda input: -1, + torch.is_neg: lambda input: -1, + torch.is_distributed: lambda input: -1, + torch.is_inference: lambda input: -1, + torch.is_floating_point: lambda input: -1, + torch.is_nonzero: lambda input: -1, + torch.is_same_size: lambda input, other: -1, + torch.is_signed: lambda input: -1, + torch.isclose: lambda input, other, rtol=1e-05, atol=1e-08, equal_nan=False: -1, + torch.isnan: lambda input: -1, + torch.istft: ( + lambda input, n_fft, hop_length=None, win_length=None, window=None, center=True, normalized=False, onesided=None, length=None, return_complex=False: -1 # noqa: B950 + ), + torch.kl_div: lambda input, target, size_average=None, reduce=None, reduction="mean", log_target=False: -1, + torch.kron: lambda input, other: -1, + torch.kthvalue: lambda input, k, dim=None, keepdim=False, out=None: -1, + torch.linalg.ldl_factor_ex: lambda input, hermitian=False, check_errors=False, out=None: -1, + torch.linalg.ldl_factor: lambda input, hermitian=False, out=None: -1, + torch.linalg.ldl_solve: lambda LD, pivots, B, hermitian=False, out=None: -1, + torch.layer_norm: lambda input, normalized_shape, weight=None, bias=None, esp=1e-05, cudnn_enabled=True: -1, + torch.lcm: lambda input, other, out=None: -1, + torch.ldexp: lambda input, other, out=None: -1, + torch.le: lambda input, other, out=None: -1, + torch.less_equal: lambda input, other, out=None: -1, + torch.lerp: lambda input, end, weight, out=None: -1, + torch.lgamma: lambda input, out=None: -1, + torch.lobpcg: lambda input, k=None, B=None, X=None, n=None, iK=None, niter=None, tol=None, largest=None, method=None, tracker=None, ortho_iparams=None, ortho_fparams=None, ortho_bparams=None: -1, # noqa: B950 + torch.log: lambda input, out=None: -1, + torch.log_softmax: lambda input, dim, dtype=None: -1, + torch.log10: lambda input, out=None: -1, + torch.log1p: lambda input, out=None: -1, + torch.log2: lambda input, out=None: -1, + torch.logaddexp: lambda input, other, out=None: -1, + torch.logaddexp2: lambda input, other, out=None: -1, + torch.logdet: lambda input: -1, + torch.xlogy: lambda x, y, out=None: -1, + torch.logical_and: lambda input, other, out=None: -1, + torch.logical_not: lambda input, out=None: -1, + torch.logical_or: lambda input, other, out=None: -1, + torch.logical_xor: lambda input, other, out=None: -1, + torch.logit: lambda input, eps=None: -1, + torch.logsumexp: lambda input, names, keepdim=False, out=None: -1, + torch.lstm: lambda data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional: -1, + torch.lstm_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.lt: lambda input, other, out=None: -1, + torch.less: lambda input, other, out=None: -1, + torch.lu: lambda A, pivot=True, get_infos=False, out=None: -1, + torch.lu_solve: lambda b, LU_data, LU_pivots, out=None: -1, + torch.margin_ranking_loss: lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1, # type: ignore[attr-defined] # noqa: B950 + torch.masked_fill: lambda input, mask, value: -1, + torch.masked_scatter: lambda input, mask, source: -1, + torch.masked_select: lambda input, mask, out=None: -1, + torch.matmul: lambda input, other, out=None: -1, + torch.linalg.lu: lambda input, pivot=True, out=None: -1, + torch.linalg.lu_factor: lambda input, pivot=True, out=None: -1, + torch.linalg.lu_factor_ex: lambda input, pivot=True, check_errors=False, out=None: -1, + torch.linalg.lu_solve: lambda LU, pivots, B, left=True, adjoint=False, out=None: -1, + torch.linalg.matmul: lambda input, other, out=None: -1, # alias for torch.matmul + torch.matrix_power: lambda input, n: -1, + torch.linalg.matrix_power: lambda input, n, out=None: -1, + torch.linalg.matrix_rank: lambda input, tol=None, hermitian=False: -1, + torch.linalg.multi_dot: lambda tensors, out=None: -1, + torch.matrix_exp: lambda input: -1, + torch.linalg.matrix_exp: lambda input: -1, + torch.max: lambda input, out=None: -1, + torch.maximum: lambda input, other, out=None: -1, + torch.fmax: lambda input, other, out=None: -1, + torch.max_pool1d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool2d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool3d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool1d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.mean: lambda input, dim=None: -1, + torch.nanmean: lambda input, dim=None, keepdim=False, dtype=None, out=None: -1, + torch.median: lambda input, dim=None: -1, + torch.nanmedian: lambda input, dim=None: -1, + torch.meshgrid: lambda *tensors, **kwargs: -1, + torch.min: lambda input, out=None: -1, + torch.minimum: lambda input, other, out=None: -1, + torch.fmin: lambda input, other, out=None: -1, + torch.miopen_batch_norm: ( + lambda input, weight, bias, running_mean, running_var, training, exponential_average_factor, epsilon: -1 + ), + torch.miopen_convolution: lambda input, weight, bias, padding, stride, dilation, groups, benchmark, deterministic: -1, # noqa: B950 + torch.miopen_convolution_add_relu: lambda input, weight, z, alpha, bias, stride, padding, dilation, groups: -1, + torch.miopen_convolution_relu: lambda input, weight, bias, stride, padding, dilation, groups: -1, + torch.miopen_convolution_transpose: ( + lambda input, weight, bias, padding, output_padding, stride, dilation, groups, benchmark, deterministic: -1 + ), + torch.miopen_depthwise_convolution: ( + lambda input, weight, bias, padding, stride, dilation, groups, benchmark, deterministic: -1 + ), + torch.miopen_rnn: ( + lambda input, weight, weight_stride0, hx, cx, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state: -1 # noqa: B950 + ), + torch.mm: lambda input, mat2, out_dtype=None, out=None: -1, + torch.mode: lambda input, dim=-1, keepdim=False, out=None: -1, + torch.movedim: lambda input, source, destination: -1, + torch.moveaxis: lambda input, source, destination: -1, + torch.msort: lambda input, descending=False, out=None: -1, + torch.mul: lambda input, other, out=None: -1, + torch.multiply: lambda input, other, out=None: -1, + torch.multinomial: lambda input, num_samples, replacement=False, out=None: -1, + torch.mv: lambda input, vec, out=None: -1, + torch.mvlgamma: lambda input, p: -1, + torch.narrow: lambda input, dim, start, length: -1, + torch.nan_to_num: lambda input, nan=0.0, posinf=None, neginf=None, out=None: -1, + torch.native_batch_norm: lambda input, weight, bias, running_mean, running_var, training, momentum, eps: -1, + torch._native_batch_norm_legit: lambda input, weight, bias, training, momentum, eps: -1, + torch.native_dropout: lambda input, p, train: -1, + torch.native_layer_norm: lambda input, normalized_shape, weight=None, bias=None, eps=1e-05: -1, + torch._fused_rms_norm: lambda input, normalized_shape, weight=None, eps=1e-05: -1, + torch.native_group_norm: lambda input, weight, bias, N, C, HxW, group, eps: -1, + torch.native_norm: lambda input, p=2, dim=None, keepdim=False, dtype=None: -1, + torch.native_channel_shuffle: lambda input, groups: -1, + torch.ne: lambda input, other, out=None: -1, + torch.not_equal: lambda input, other, out=None: -1, + torch.neg: lambda input, out=None: -1, + torch.negative: lambda input, out=None: -1, + torch.nextafter: lambda input, other, out=None: -1, + torch.nn.functional.adaptive_avg_pool2d: lambda input, output_size: -1, + torch.nn.functional.adaptive_avg_pool3d: lambda input, output_size: -1, + torch.nn.functional.adaptive_max_pool1d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool1d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool2d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool2d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool3d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool3d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.affine_grid: lambda theta, size, align_corners=None: -1, + torch.nn.functional.alpha_dropout: lambda input, p=0.5, training=False, inplace=False: -1, + torch.nn.functional.avg_pool2d: ( + lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None: -1 # noqa: B950 + ), + torch.nn.functional.avg_pool3d: ( + lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None: -1 # noqa: B950 + ), + torch.nn.functional.batch_norm: ( + lambda input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05: -1 + ), + torch.nn.functional.bilinear: lambda input1, input2, weight, bias=None: -1, + torch.nn.functional.binary_cross_entropy: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.binary_cross_entropy_with_logits: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean", pos_weight=None: -1 + ), + torch.nn.functional.celu: lambda input, alpha=1.0, inplace=False: -1, + torch.nn.functional.cosine_embedding_loss: ( + lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.cross_entropy: ( + lambda input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean", label_smoothing=0.0: -1 # noqa: B950 + ), + torch.nn.functional.ctc_loss: ( + lambda log_probs, targets, input_lengths, target_lengths, blank=0, reduction="mean", zero_infinity=False: -1 + ), + torch.nn.functional.dropout: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout1d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout2d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout3d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.elu: lambda input, alpha=1.0, inplace=False: -1, + torch.nn.functional.embedding: ( + lambda input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False: -1 # noqa: B950 + ), + torch.nn.functional.embedding_bag: ( + lambda input, weight, offsets=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode="mean", sparse=False, per_sample_weights=None, include_last_offset=False, padding_idx=None: -1 # noqa: B950 + ), + torch.nn.functional.feature_alpha_dropout: lambda input, p=0.5, training=False, inplace=False: -1, + torch.nn.functional.fold: lambda input, output_size, kernel_size, dilation=1, padding=0, stride=1: -1, + torch.nn.functional.fractional_max_pool2d: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool2d_with_indices: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool3d: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool3d_with_indices: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.gaussian_nll_loss: lambda input, target, var, full=False, eps=1e-06, reduction="mean": -1, + torch.nn.functional.gelu: lambda input, approximate="none": -1, + torch.nn.functional.glu: lambda input, dim=-1: -1, + torch.nn.functional.grid_sample: lambda input, grid, mode="bilinear", padding_mode="zeros", align_corners=None: -1, # noqa: B950 + torch.nn.functional.group_norm: lambda input, num_groups, weight=None, bias=None, eps=1e-05: -1, + torch.nn.functional.gumbel_softmax: lambda logits, tau=1, hard=False, eps=1e-10, dim=-1: -1, + torch.nn.functional.hardshrink: lambda input, lambd=0.5: -1, + torch.nn.functional.hardtanh: lambda input, min_val=-1.0, max_val=1.0, inplace=False: -1, + torch.nn.functional.hinge_embedding_loss: ( + lambda input, target, margin=1.0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.instance_norm: ( + lambda input, running_mean=None, running_var=None, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-05: -1 # noqa: B950 + ), + torch.nn.functional.interpolate: ( + lambda input, size=None, scale_factor=None, mode="nearest", align_corners=None, recompute_scale_factor=None, antialias=False: -1 # noqa: B950 + ), + torch.nn.functional.kl_div: lambda input, target, size_average=None, reduce=None, reduction="mean", log_target=False: -1, # noqa: B950 + torch.nn.functional.l1_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", weight=None: -1, + torch.nn.functional.layer_norm: lambda input, normalized_shape, weight=None, bias=None, eps=1e-05: -1, + torch.nn.functional.leaky_relu: lambda input, negative_slope=0.01, inplace=False: -1, + torch.nn.functional.linear: lambda input, weight, bias=None: -1, + torch.nn.functional.local_response_norm: lambda input, size, alpha=0.0001, beta=0.75, k=1.0: -1, + torch.nn.functional.log_softmax: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.logsigmoid: lambda input: -1, + torch.nn.functional.lp_pool1d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.lp_pool2d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.lp_pool3d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.margin_ranking_loss: ( + lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.max_pool1d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False: -1 + ), + torch.nn.functional.max_pool1d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool2d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False: -1 + ), + torch.nn.functional.max_pool2d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool3d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool3d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_unpool1d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.max_unpool2d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.max_unpool3d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.mse_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", weight=None: -1, + torch.nn.functional.multi_head_attention_forward: ( + lambda query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight, out_proj_bias, training=True, key_padding_mask=None, need_weights=True, attn_mask=None, use_separate_proj_weight=False, q_proj_weight=None, k_proj_weight=None, v_proj_weight=None, static_k=None, static_v=None, average_attn_weights=None, is_causal=False: -1 # noqa: B950 + ), + torch.nn.functional.multi_margin_loss: ( + lambda input, target, p=1, margin=1.0, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.multilabel_margin_loss: ( + lambda input, target, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.multilabel_soft_margin_loss: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.nll_loss: ( + lambda input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.normalize: lambda input, p=2, dim=1, eps=1e-12, out=None: -1, + torch.nn.functional.one_hot: lambda tensor, num_classes=-1: -1, + torch.nn.functional.pad: lambda input, pad, mode="constant", value=0: -1, + torch.nn.functional.pairwise_distance: lambda x1, x2, p=2.0, eps=1e-06, keepdim=False: -1, + torch.nn.functional.poisson_nll_loss: ( + lambda input, target, log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.nn.functional.prelu: lambda input, weight: -1, + torch.nn.functional.relu: lambda input, inplace=False: -1, + torch.nn.functional.relu6: lambda input, inplace=False: -1, + torch.nn.functional.rms_norm: lambda input, normalized_shape, weight=None, eps=1e-6: -1, + torch.nn.functional.rrelu: lambda input, lower=0.125, upper=0.3333333333333333, training=False, inplace=False: -1, # noqa: B950 + torch.nn.functional.selu: lambda input, inplace=False: -1, + torch.nn.functional.silu: lambda input, inplace=False: -1, + torch.nn.functional.mish: lambda input, inplace=False: -1, + torch.nn.functional.scaled_dot_product_attention: lambda query, key, value, attn_mask=None, dropout_p=0.0: -1, + torch.nn.functional.smooth_l1_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", beta=1.0: -1, # noqa: B950 + torch.nn.functional.huber_loss: lambda input, target, reduction="mean", delta=1.0, weight=None: -1, + torch.nn.functional.soft_margin_loss: lambda input, target, size_average=None, reduce=None, reduction="mean": -1, # noqa: B950 + torch.nn.functional.softmax: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.softmin: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.softplus: lambda input, beta=1, threshold=20: -1, + torch.nn.functional.softshrink: lambda input, lambd=0.5: -1, + torch.nn.functional.softsign: lambda input: -1, + torch.nn.functional.tanhshrink: lambda input: -1, + torch.nn.functional.threshold: lambda input, threshold, value, inplace=False: -1, + torch.nn.functional.triplet_margin_loss: ( + lambda anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.nn.functional.triplet_margin_with_distance_loss: ( + lambda anchor, positive, negative, *, distance_function=None, margin=1.0, swap=False, reduction="mean": -1 + ), + torch.nn.functional.unfold: lambda input, kernel_size, dilation=1, padding=0, stride=1: -1, + torch.nn.init.uniform_: lambda tensor, a=0.0, b=1.0, generator=None: -1, + torch.nn.init.normal_: lambda tensor, mean=0.0, std=1.0, generator=None: -1, + torch.nn.init.constant_: lambda tensor, val: -1, + torch.nn.init.kaiming_uniform_: lambda tensor, a=0, mode="fan_in", nonlinearity="leaky_relu", generator=None: -1, # noqa: B950 + torch.nonzero: lambda input, as_tuple=False: -1, + torch.nonzero_static: lambda input, *, size, fill_value=-1: -1, + torch.argwhere: lambda input: -1, + torch.norm: lambda input, p="fro", dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.norm: lambda input, ord=None, dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.vector_norm: lambda input, ord=2, dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.matrix_norm: lambda input, ord="fro", dim=( + -2, + -1, + ), keepdim=False, out=None, dtype=None: -1, + torch.norm_except_dim: lambda v, pow=2, dim=0: -1, + torch.nuclear_norm: lambda input, p="fro", dim=None, keepdim=False, out=None, dtype=None: -1, + torch.numel: lambda input: -1, + torch.orgqr: lambda input, tau: -1, + torch.ormqr: lambda input, input2, input3, left=True, transpose=False: -1, + torch.pairwise_distance: lambda x1, x2, p=2.0, eps=1e-06, keepdim=False: -1, + torch.permute: lambda self, dim: -1, + torch.pca_lowrank: lambda input, q=None, center=True, niter=2: -1, + torch.pdist: lambda input, p=2: -1, + torch.pinverse: lambda input, rcond=1e-15: -1, + torch.linalg.pinv: lambda input, rcond=1e-15, hermitian=False: -1, + torch.pixel_shuffle: lambda input, upscale_factor: -1, + torch.pixel_unshuffle: lambda input, downscale_factor: -1, + torch.poisson: lambda input, generator=None: -1, + torch.poisson_nll_loss: lambda input, target, log_input, full, eps, reduction: -1, + torch.polygamma: lambda input, n, out=None: -1, + torch.positive: lambda input, out=None: -1, + torch.prelu: lambda input, weight: -1, + torch.ones_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.pow: lambda input, exponent, out=None: -1, + torch.prod: lambda input, dtype=None: -1, + torch.put: lambda input, index, source, accumulate=False: -1, + torch.q_per_channel_axis: lambda input: -1, + torch.q_per_channel_scales: lambda input: -1, + torch.q_per_channel_zero_points: lambda input: -1, + torch.q_scale: lambda input: -1, + torch.q_zero_point: lambda input: -1, + torch.qr: lambda input, some=True, out=None: -1, + torch.linalg.qr: lambda input, mode="reduced", out=None: -1, + torch.quantile: lambda input, q, dim=None, keepdim=False, interpolation="linear", out=None: -1, + torch.nanquantile: lambda input, q, dim=None, keepdim=False, interpolation="linear", out=None: -1, + torch.quantize_per_channel: lambda input, scales, zero_points, axis, dtype: -1, + torch.quantize_per_tensor: lambda input, scale, zero_point, dtype: -1, + torch.quantize_per_tensor_dynamic: lambda input, dtype, reduce_range: -1, + torch.quantized_batch_norm: lambda input, weight, bias, mean, var, eps, output_scale, output_zero_point: -1, + torch.quantized_gru_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_lstm_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_max_pool1d: ( + lambda input, kernel_size, stride=(), padding=(0,), dilation=( + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_max_pool2d: ( + lambda input, kernel_size, stride=(), padding=(0, 0), dilation=( + 1, + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_max_pool3d: ( + lambda input, kernel_size, stride=(), padding=(0, 0, 0), dilation=( + 1, + 1, + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_rnn_relu_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_rnn_tanh_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.rad2deg: lambda input, out=None: -1, + torch.rand_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.randint_like: lambda input, high, dtype=None, layout=torch.strided, device=None, requires_grad=False: -1, + torch.randn_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.ravel: lambda input: -1, + torch.real: lambda input, out=None: -1, + torch.vdot: lambda input, other, out=None: -1, + torch.linalg.vecdot: lambda input, other, dim=-1, out=None: -1, + torch.view_as_real: lambda input: -1, + torch.view_as_complex: lambda input: -1, + torch.reciprocal: lambda input, out=None: -1, + torch.relu: lambda input, inplace=False: -1, + torch.remainder: lambda input, other, out=None: -1, + torch.renorm: lambda input, p, dim, maxnorm, out=None: -1, + torch.repeat_interleave: lambda input, dim=None: -1, + torch.reshape: lambda input, shape: -1, + torch.rms_norm: lambda input, normalized_shape, weight=None, eps=1e-6: -1, + torch.rnn_relu: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, # noqa: B950 + torch.rnn_relu_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.rnn_tanh: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, # noqa: B950 + torch.rnn_tanh_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.roll: lambda input, shifts, dims=None: -1, + torch.rot90: lambda input, k=1, dims=(0, 1): -1, + torch.round: lambda input, out=None: -1, + torch.row_stack: lambda tensors, out=None: -1, # alias for torch.vstack + torch._rowwise_prune: (lambda weight, mask, compressed_indices_dtype: -1), + torch.rrelu: lambda input, lower=1.0 / 8, upper=1.0 / 3, training=False, inplace=False: -1, + torch.rsqrt: lambda input, out=None: -1, + torch.rsub: lambda input, other, alpha=1: -1, + torch.saddmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.scatter: lambda input, dim, index, src: -1, + torch.scatter_add: lambda input, dim, index, src: -1, + torch.scatter_reduce: lambda input, dim, index, src, reduce, include_self=True: -1, + torch.searchsorted: lambda sorted_sequence, input, out_int32=False, right=False, out=None: -1, + torch._segment_reduce: lambda data, reduce="max", lengths=None, indices=None, offsets=None, axis=0, unsafe=False: -1, # noqa: B950 + torch.select: lambda input, dim, index: -1, + torch.select_scatter: lambda input, src, dim, index: -1, + torch.slice_inverse: lambda input, src, dim=0, start=None, end=None, step=1: -1, + torch.slice_scatter: lambda input, src, dim=0, start=None, end=None, step=1: -1, + torch.selu: lambda input, inplace=False: -1, + torch.sigmoid: lambda input, out=None: -1, + torch.sign: lambda input, out=None: -1, + torch.signbit: lambda input, out=None: -1, + torch.sgn: lambda input, out=None: -1, + torch.sin: lambda input, out=None: -1, + torch.sinc: lambda input, out=None: -1, + torch.sinh: lambda input, out=None: -1, + torch.slogdet: lambda input: -1, + torch.linalg.slogdet: lambda input: -1, + torch.smm: lambda input, mat2, out_dtype=None: -1, + torch.spmm: lambda input, mat2, out_dtype=None: -1, + torch.softmax: lambda input, dim, dtype=None: -1, + torch.linalg.solve: lambda A, B, left=True, out=None: -1, + torch.linalg.solve_ex: lambda A, B, left=True, check_errors=False, out=None: -1, + torch.sort: lambda input, dim=-1, descending=False, *, stable=False, out=None: -1, + torch.split: lambda tensor, split_size_or_sections, dim=0: -1, + torch.split_with_sizes: lambda tensor, split_size_or_sections, dim=0: -1, + torch.sqrt: lambda input, out=None: -1, + torch.square: lambda input, out=None: -1, + torch.squeeze: lambda input, dim=None, out=None: -1, + torch.sspaddmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.stack: lambda tensors, dim=0, out=None: -1, + torch.std: lambda input, dim=None: -1, + torch.std_mean: lambda input, dim=None: -1, + torch.stft: ( + lambda input, n_fft, hop_length=None, win_length=None, window=None, center=True, pad_mode="reflect", normalized=False, onesided=True, return_complex=None, align_to_window=None: -1 # noqa: B950 + ), + torch.sub: lambda input, other, out=None: -1, + torch.subtract: lambda input, other, out=None: -1, + torch.sum: lambda input, dim=None: -1, + torch.sym_float: lambda input: -1, + torch.sym_int: lambda input: -1, + torch.sym_max: lambda a, b: -1, + torch.sym_min: lambda a, b: -1, + torch.sym_not: lambda input: -1, + torch.sym_ite: lambda a, b, c: -1, + torch.sym_sum: lambda args: -1, + torch._sym_sqrt: lambda input: -1, + torch._sym_cos: lambda input: -1, + torch._sym_cosh: lambda input: -1, + torch._sym_sin: lambda input: -1, + torch._sym_sinh: lambda input: -1, + torch._sym_tan: lambda input: -1, + torch._sym_tanh: lambda input: -1, + torch._sym_asin: lambda input: -1, + torch._sym_acos: lambda input: -1, + torch._sym_atan: lambda input: -1, + torch.nansum: lambda input, dim=None: -1, + torch.svd: lambda input, some=True, compute_uv=True, out=None: -1, + torch.svd_lowrank: lambda input, q=6, niter=2, M=None: -1, + torch.linalg.svd: lambda input, full_matrices=True, out=None: -1, + torch.linalg.svdvals: lambda input, out=None: -1, + torch.swapaxes: lambda input, dim0, dim1: -1, + torch.swapdims: lambda input, axis0, axis1: -1, + torch.special.airy_ai: lambda input: -1, + torch.special.bessel_j0: lambda input: -1, + torch.special.bessel_j1: lambda input: -1, + torch.special.bessel_y0: lambda input: -1, + torch.special.bessel_y1: lambda input: -1, + torch.special.chebyshev_polynomial_t: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_u: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_v: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_w: lambda input, n, out=None: -1, + torch.special.digamma: lambda input: -1, + torch.special.entr: lambda input: -1, + torch.special.erf: lambda input: -1, + torch.special.erfc: lambda input: -1, + torch.special.erfcx: lambda input: -1, + torch.special.erfinv: lambda input: -1, + torch.special.exp2: lambda input: -1, + torch.special.expit: lambda input: -1, + torch.special.expm1: lambda input: -1, + torch.special.gammainc: lambda input, other, out=None: -1, + torch.special.gammaincc: lambda input, other, out=None: -1, + torch.special.gammaln: lambda input: -1, + torch.special.hermite_polynomial_h: lambda input, n, out=None: -1, + torch.special.hermite_polynomial_he: lambda input, n, out=None: -1, + torch.special.i0: lambda input: -1, + torch.special.i0e: lambda input: -1, + torch.special.i1: lambda input: -1, + torch.special.i1e: lambda input: -1, + torch.special.laguerre_polynomial_l: lambda input, n, out=None: -1, + torch.special.legendre_polynomial_p: lambda input, n, out=None: -1, + torch.special.log1p: lambda input: -1, + torch.special.log_ndtr: lambda input: -1, + torch.special.log_softmax: lambda input, dim, dtype=None: -1, + torch.special.logit: lambda input: -1, + torch.special.logsumexp: lambda input, dim, keepdim=False, out=None: -1, + torch.special.modified_bessel_i0: lambda input: -1, + torch.special.modified_bessel_i1: lambda input: -1, + torch.special.modified_bessel_k0: lambda input: -1, + torch.special.modified_bessel_k1: lambda input: -1, + torch.special.multigammaln: lambda input, p: -1, + torch.special.ndtr: lambda input: -1, + torch.special.ndtri: lambda input: -1, + torch.special.polygamma: lambda input, n, out=None: -1, + torch.special.psi: lambda input: -1, + torch.special.round: lambda input: -1, + torch.special.scaled_modified_bessel_k0: lambda input: -1, + torch.special.scaled_modified_bessel_k1: lambda input: -1, + torch.special.shifted_chebyshev_polynomial_t: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_u: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_v: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_w: lambda input, n, out=None: -1, + torch.special.sinc: lambda input: -1, + torch.special.softmax: lambda input, dim, dtype=None: -1, + torch.special.spherical_bessel_j0: lambda input: -1, + torch.special.xlog1py: lambda input, other, out=None: -1, + torch.special.xlogy: lambda input, other, out=None: -1, + torch.special.zeta: lambda self, other, out=None: -1, + torch.t: lambda input: -1, + torch.take: lambda input, index: -1, + torch.take_along_dim: lambda input, indices, dim=None, out=None: -1, + torch.tan: lambda input, out=None: -1, + torch.tanh: lambda input, out=None: -1, + torch.linalg.tensorinv: lambda a, ind=2: -1, + torch.linalg.tensorsolve: lambda a, b, dims=None: -1, + torch.tensordot: lambda a, b, dims=2, out=None: -1, + torch.tensor_split: lambda input, indices_or_sections, dim=0: -1, + torch.threshold: lambda input, threshold, value, inplace=False: -1, + torch.tile: lambda input, dims: -1, + torch.topk: lambda input, k, dim=-1, descending=False, out=None: -1, + torch.trace: lambda input: -1, + torch.transpose: lambda input, dim0, dim1: -1, + torch.trapz: lambda y, x=None, dim=-1: -1, + torch.trapezoid: lambda y, x=None, dim=-1: -1, + torch.triangular_solve: lambda input, A, upper=True, transpose=False, unitriangular=False: -1, + torch.linalg.solve_triangular: lambda input, B, upper, left=True, unitriangular=False: -1, + torch.tril: lambda input, diagonal=0, out=None: -1, + torch.triplet_margin_loss: ( + lambda anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.triu: lambda input, diagonal=0, out=None: -1, + torch.true_divide: lambda input, other: -1, + torch.trunc: lambda input, out=None: -1, + torch.unbind: lambda input, dim=0: -1, + torch.unflatten: lambda input, dim, sizes, names: -1, + torch.unique: lambda input, sorted=True, return_inverse=False, return_counts=False, dim=None: -1, + torch.unique_consecutive: lambda input, return_inverse=False, return_counts=False, dim=None: -1, + torch.unravel_index: lambda indices, shape: -1, + torch.unsafe_chunk: lambda input, chunks, dim=0: -1, + torch.unsafe_split: lambda tensor, split_size_or_sections, dim=0: -1, + torch.unsafe_split_with_sizes: lambda tensor, split_size_or_sections, dim=0: -1, + torch.unsqueeze: lambda input, dim, out=None: -1, + torch.linalg.vander: lambda x, N=None: -1, + torch.var: lambda input, dim=None: -1, + torch.var_mean: lambda input, dim=None: -1, + torch.vsplit: lambda input, indices_or_sections: -1, + torch.vstack: lambda tensors, out=None: -1, + torch.where: lambda condition, x=None, y=None: -1, + torch._wrapped_linear_prepack: lambda weight, weight_scale, weight_zero_point, bias : -1, + torch._wrapped_quantized_linear_prepacked: ( + lambda input, input_scale, input_zero_point, prepacked, out_scale, out_zero_point, out_channel : -1 # noqa: B950 + ), + torch.zeros_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch._fw_primal_copy: lambda self, level: -1, + torch._make_dual_copy: lambda primal, tangent, level: -1, + torch.view_as_real_copy: lambda self: -1, + torch.view_as_complex_copy: lambda self: -1, + torch._conj_copy: lambda self: -1, + torch._neg_view_copy: lambda self: -1, + torch.as_strided_copy: lambda self, size, stride, storage_offset=None: -1, + torch._sparse_broadcast_to_copy: lambda self, size: -1, + torch.diagonal_copy: lambda self, offset=0, dim1=0, dim2=1: -1, + torch.expand_copy: lambda self, size, *, implicit=False: -1, + torch.narrow_copy: lambda self, dim, start, length: -1, + torch.permute_copy: lambda self, dims: -1, + torch._reshape_alias_copy: lambda self, size, stride: -1, + torch.select_copy: lambda self, dim, index: -1, + torch.detach_copy: lambda self: -1, + torch.slice_copy: lambda self, dim=0, start=None, end=None, step=1: -1, + torch.split_copy: lambda self, split_size, dim=0: -1, + torch.split_with_sizes_copy: lambda self, split_sizes, dim=0: -1, + torch.squeeze_copy: lambda self, dim: -1, + torch.t_copy: lambda self: -1, + torch.transpose_copy: lambda self, dim0, dim1: -1, + torch.unsqueeze_copy: lambda self, dim: -1, + torch._indices_copy: lambda self: -1, + torch._values_copy: lambda self: -1, + torch.indices_copy: lambda self: -1, + torch.values_copy: lambda self: -1, + torch.crow_indices_copy: lambda self: -1, + torch.col_indices_copy: lambda self: -1, + torch.ccol_indices_copy: lambda self: -1, + torch.row_indices_copy: lambda self: -1, + torch.unbind_copy: lambda self, dim=0: -1, + torch.view_copy: lambda self, dtype: -1, + torch.unfold_copy: lambda self, dimension, size, step: -1, + torch.alias_copy: lambda self: -1, + Tensor.__floordiv__: lambda self, other: -1, + Tensor.__rfloordiv__: lambda self, other: -1, + Tensor.__ifloordiv__: lambda self, other: -1, + Tensor.__truediv__: lambda self, other: -1, + Tensor.__rtruediv__: lambda self, other: -1, + Tensor.__itruediv__: lambda self, other: -1, + Tensor.__lshift__: lambda self, other: -1, + Tensor.__rlshift__: lambda self, other: -1, + Tensor.__ilshift__: lambda self, other: -1, + Tensor.__rshift__: lambda self, other: -1, + Tensor.__rrshift__: lambda self, other: -1, + Tensor.__irshift__: lambda self, other: -1, + Tensor.__and__: lambda self, other: -1, + Tensor.__or__: lambda self, other: -1, + Tensor.__xor__: lambda self, other: -1, + Tensor.__float__: lambda self: -1, + Tensor.__complex__: lambda self: -1, + Tensor.__array__: lambda self, dtype: -1, + Tensor.__bool__: lambda self: -1, + Tensor.__contains__: lambda self, other: -1, + Tensor.__neg__: lambda self: -1, + Tensor.__invert__: lambda self: -1, + Tensor.__mod__: lambda self, other: -1, + Tensor.__rmod__: lambda self, other: -1, + Tensor.__imod__: lambda self, other: -1, + Tensor.__array_wrap__: lambda self, array: -1, + Tensor.__getitem__: lambda self, idx: -1, + Tensor.__deepcopy__: lambda self, memo: -1, + Tensor.__int__: lambda self: -1, + Tensor.__long__: lambda self: -1, + Tensor.__index__: lambda self: -1, + Tensor.__len__: lambda self: -1, + Tensor.__format__: lambda self, format_spec: -1, + Tensor.__reduce_ex__: lambda self, proto: -1, + Tensor.__reversed__: lambda self: -1, + Tensor.__repr__: lambda self, *, tensor_contents=None: -1, + Tensor.__setitem__: lambda self, k, v: -1, + Tensor.__setstate__: lambda self, d: -1, + Tensor.T.__get__: lambda self: -1, + Tensor.H.__get__: lambda self: -1, + Tensor.mT.__get__: lambda self: -1, + Tensor.mH.__get__: lambda self: -1, + Tensor._backward_hooks.__get__: lambda self: -1, + Tensor._post_accumulate_grad_hooks.__get__: lambda self: -1, + Tensor._base.__get__: lambda self: -1, + Tensor._cdata.__get__: lambda self: -1, + Tensor.grad.__get__: lambda self: -1, + Tensor._grad.__get__: lambda self: -1, + Tensor._grad_fn.__get__: lambda self: -1, + Tensor.grad_fn.__get__: lambda self: -1, + Tensor._version.__get__: lambda self: -1, + Tensor._autocast_to_reduced_precision: lambda self, cuda_enabled, cpu_enabled, cuda_dtype, cpu_dtype: -1, + Tensor._autocast_to_full_precision: lambda self, cuda_enabled, cpu_enabled: -1, + Tensor._clear_non_serializable_cached_data: lambda self: -1, + Tensor.data.__get__: lambda self: -1, + Tensor.device.__get__: lambda self: -1, + Tensor.dtype.__get__: lambda self: -1, + Tensor.is_cuda.__get__: lambda self: -1, + Tensor.is_cpu.__get__: lambda self: -1, + Tensor.is_xla.__get__: lambda self: -1, + Tensor.is_xpu.__get__: lambda self: -1, + Tensor.is_ipu.__get__: lambda self: -1, + Tensor.is_leaf.__get__: lambda self: -1, + Tensor.retains_grad.__get__: lambda self: -1, + Tensor.is_meta.__get__: lambda self: -1, + Tensor.is_mps.__get__: lambda self: -1, + Tensor.is_mtia.__get__: lambda self: -1, + Tensor.is_nested.__get__: lambda self: -1, + Tensor.is_maia.__get__: lambda self: -1, + Tensor.is_mkldnn.__get__: lambda self: -1, + Tensor.is_quantized.__get__: lambda self: -1, + Tensor.is_sparse.__get__: lambda self: -1, + Tensor.is_sparse_csr.__get__: lambda self: -1, + Tensor.is_vulkan.__get__: lambda self: -1, + Tensor.itemsize.__get__: lambda self: -1, + Tensor.layout.__get__: lambda self: -1, + Tensor.name.__get__: lambda self: -1, + Tensor.names.__get__: lambda self: -1, + Tensor.nbytes.__get__: lambda self: -1, + Tensor.ndim.__get__: lambda self: -1, + Tensor.output_nr.__get__: lambda self: -1, + Tensor.requires_grad.__get__: lambda self: -1, + Tensor.shape.__get__: lambda self: -1, + Tensor.volatile.__get__: lambda self: -1, + Tensor.real.__get__: lambda self: -1, + Tensor.imag.__get__: lambda self: -1, + Tensor.__cuda_array_interface__.__get__: lambda self: -1, + Tensor.type: lambda self, dtype=None, non_blocking=False, **kwargs: -1, + Tensor._dimI: lambda self: -1, + Tensor._dimV: lambda self: -1, + Tensor._indices: lambda self: -1, + Tensor._is_view: lambda self: -1, + Tensor._nnz: lambda self: -1, + Tensor.crow_indices: lambda self: -1, + Tensor.col_indices: lambda self: -1, + Tensor.ccol_indices: lambda self: -1, + Tensor.row_indices: lambda self: -1, + Tensor._update_names: lambda self, names, inplace: -1, + Tensor._values: lambda self: -1, + Tensor.adjoint: lambda self: -1, + Tensor.align_as: lambda self, other: -1, + Tensor.align_to: lambda self, order, ellipsis_idx: -1, + Tensor.apply_: lambda self, callable: -1, + Tensor.as_strided: lambda self, size, stride: -1, + Tensor.as_strided_: lambda self, size, stride: -1, + Tensor.backward: lambda self, gradient=None, retain_graph=None, create_graph=False, inputs=None: -1, + Tensor.bfloat16: lambda self, memory_format=torch.preserve_format: -1, + Tensor.bool: lambda self, memory_format=torch.preserve_format: -1, + Tensor.byte: lambda self, memory_format=torch.preserve_format: -1, + Tensor.char: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cauchy_: lambda self, median=0, sigma=1, *, generator=None: -1, + Tensor.coalesce: lambda self: -1, + Tensor._coalesced_: lambda self, coalesced: -1, + Tensor.contiguous: lambda self, memory_format=torch.contiguous_format: -1, + Tensor.copy_: lambda self, src, non_blocking=False: -1, + Tensor.cpu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cuda: lambda self, memory_format=torch.preserve_format: -1, + Tensor.mtia: lambda self, memory_format=torch.preserve_format: -1, + Tensor.xpu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.ipu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.data_ptr: lambda self: -1, + Tensor.dense_dim: lambda self: -1, + Tensor.diagonal_scatter: lambda self, src, offset=0, dim1=0, dim2=1: -1, + Tensor.dim: lambda self: -1, + Tensor.dim_order: lambda self, ambiguity_check=False: -1, + Tensor.double: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cdouble: lambda self, memory_format=torch.preserve_format: -1, + Tensor.element_size: lambda self: -1, + Tensor.expand: lambda self, size: -1, + Tensor.expand_as: lambda self, other: -1, + Tensor.exponential_: lambda self, lambd=1, *, generator=None: -1, + Tensor.fill_: lambda self, value: -1, + Tensor.fill_diagonal_: lambda self, value: -1, + Tensor.float: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cfloat: lambda self, memory_format=torch.preserve_format: -1, + Tensor.geometric_: lambda self, p, *, generator=None: -1, + Tensor.get_device: lambda self: -1, + Tensor.half: lambda self, memory_format=torch.preserve_format: -1, + Tensor.chalf: lambda self, memory_format=torch.preserve_format: -1, + Tensor.has_names: lambda self: -1, + Tensor.indices: lambda self: -1, + Tensor.int: lambda self, memory_format=torch.preserve_format: -1, + Tensor.is_coalesced: lambda self: -1, + Tensor.is_contiguous: lambda self: -1, + Tensor.is_inference: lambda self: -1, + Tensor.is_pinned: lambda self: -1, + Tensor.is_set_to: lambda self, tensor: -1, + Tensor.is_shared: lambda self: -1, + Tensor.item: lambda self: -1, + Tensor.log_normal_: lambda self, mean=1, std=2, *, generator=None: -1, + Tensor.log_softmax: lambda self, dim: -1, + Tensor.long: lambda self, memory_format=torch.preserve_format: -1, + Tensor.map_: lambda self, tensor, callable: -1, + Tensor.map2_: lambda self, x, y, callable: -1, + Tensor.mm: lambda self, mat2, out_dtype=None: -1, + Tensor.module_load: lambda self, other, assign=False: -1, + Tensor.narrow_copy: lambda self, dimension, start, length: -1, + Tensor.ndimension: lambda self: -1, + Tensor.nelement: lambda self: -1, + Tensor._nested_tensor_size: lambda self: -1, + Tensor._nested_tensor_storage_offsets: lambda self: -1, + Tensor._nested_tensor_strides: lambda self: -1, + Tensor.normal_: lambda self: -1, + Tensor.numpy: lambda self: -1, + Tensor.permute: lambda self, dim: -1, + Tensor.pin_memory: lambda self: -1, + Tensor.put_: lambda self, indices, tensor, accumulate=False: -1, + Tensor.qscheme: lambda self: -1, + Tensor.random_: lambda self, from_=0, to=None, *, generator=None: -1, + Tensor.record_stream: lambda self, stream: -1, + Tensor.refine_names: lambda self, names: -1, + Tensor.register_hook: lambda self, hook: -1, + Tensor.register_post_accumulate_grad_hook: lambda self, hook: -1, + Tensor.rename: lambda self, name: -1, + Tensor.repeat: lambda self, *size: -1, + Tensor.requires_grad_: lambda self, requires_grad=True: -1, + Tensor.reshape_as: lambda self, other: -1, + Tensor.resize: lambda self, *size: -1, + Tensor.resize_: lambda self, size: -1, + Tensor.resize_as: lambda self, other: -1, + Tensor.resize_as_sparse_: lambda self, other: -1, + Tensor.retain_grad: lambda self: -1, + Tensor.set_: lambda self, source=None, storage_offset=0, size=None, stride=None: -1, + Tensor.select_scatter: lambda self, src, dim, index: -1, + Tensor.share_memory_: lambda self: -1, + Tensor.short: lambda self, memory_format=torch.preserve_format: -1, + Tensor.size: lambda self: -1, + Tensor.slice_scatter: lambda self, src, dim=0, start=None, end=None, step=1: -1, + Tensor.sparse_dim: lambda self: -1, + Tensor.sparse_mask: lambda self, mask: -1, + Tensor._sparse_mask_projection: lambda self, mask, accumulate_matches=False: -1, + Tensor.sparse_resize_: lambda self, size1, size2, dense_dim: -1, + Tensor.sparse_resize_and_clear_: lambda self, size1, size2, dense_dim: -1, + Tensor.sspaddmm: lambda self, mat1, mat2, beta=1, alpha=1, out=None: -1, + Tensor.storage: lambda self: -1, + Tensor.untyped_storage: lambda self: -1, + Tensor.storage_offset: lambda self: -1, + Tensor.storage_type: lambda self: -1, + Tensor.sum_to_size: lambda self, size: -1, + Tensor.tile: lambda self, *reps: -1, + Tensor.to: lambda self, dtype, non_blocking=False, copy=False, memory_format=torch.preserve_format: -1, + Tensor.to_dense: lambda self, dtype=None, *, masked_grad=None: -1, + Tensor._to_dense: lambda self, dtype=None, masked_grad=None: -1, + Tensor.to_sparse: lambda self: -1, + Tensor.tolist: lambda self: -1, + Tensor.to_mkldnn: lambda self: -1, + Tensor.type_as: lambda self, other: -1, + Tensor.unfold: lambda self, dimension, size, step: -1, + Tensor.uniform_: lambda self, from_=0, to=1: -1, + Tensor.values: lambda self: -1, + Tensor.view: lambda self, shape: -1, + Tensor.view_as: lambda self, other: -1, + Tensor.zero_: lambda self: -1, + Tensor.__dlpack__: lambda self, stream=None, max_version=None, dl_device=None, copy=None: -1, + Tensor.__dlpack_device__: lambda self: -1, + torch.linalg.lstsq: lambda self, b, cond=None, driver=None: -1, + } # fmt: skip + + privateuse1_backend_name = ( + torch.utils.backend_registration._privateuse1_backend_name + ) + if hasattr(Tensor, privateuse1_backend_name): + ret[getattr(Tensor, privateuse1_backend_name)] = ( + lambda self, device=None, non_blocking=False, **kwargs: -1 + ) + ret[getattr(Tensor, f"is_{privateuse1_backend_name}").__get__] = lambda self: -1 + + ret2 = {} + ignored = get_ignored_functions() + + for k, v in ret.items(): + # Generate methods like __add__ and add_ by default from add + names = [ + k.__name__, # Default method + k.__name__ + "_", # Inplace variant + "__" + k.__name__ + "__", # Dunder method + "__i" + k.__name__ + "__", # Inplace dunder method + "__r" + k.__name__ + "__", # Reverse dunder method + ] + + if k.__name__.startswith("bitwise_"): + # bitwise_ have dunder methods of the form ____ + # And so on. + subname = k.__name__[len("bitwise_") :] + names.extend( + ["__" + subname + "__", "__i" + subname + "__", "__r" + subname + "__"] + ) + + for name in names: + func = getattr(Tensor, name, None) + if callable(func) and func not in ret and func not in ignored: + ret2[func] = v + + ret.update(ret2) + return ret + + +def wrap_torch_function(dispatcher: Callable): + """Wraps a given function with ``__torch_function__`` -related functionality. + + Parameters + ---------- + dispatcher: Callable + A callable that returns an iterable of Tensor-likes passed into the function. + + Note + ---- + This decorator may reduce the performance of your code. Generally, it's enough to express + your code as a series of functions that, themselves, support __torch_function__. If you + find yourself in the rare situation where this is not the case, e.g. if you're wrapping a + low-level library and you also need it to work for Tensor-likes, then this function is available. + + Examples + -------- + >>> def dispatcher(a): # Must have the same signature as func + ... return (a,) + >>> @torch.overrides.wrap_torch_function(dispatcher) + >>> def func(a): # This will make func dispatchable by __torch_function__ + ... return a + 0 + """ + + def inner(func): + @functools.wraps(func) + def wrapped(*args, **kwargs): + relevant_args = dispatcher(*args, **kwargs) + if has_torch_function(relevant_args): + return handle_torch_function(wrapped, relevant_args, *args, **kwargs) + + return func(*args, **kwargs) + + return wrapped + + return inner + + +def _get_overloaded_args( + relevant_args: Iterable[Any], + get_type_fn: Optional[Callable[[Any], type]] = None, +) -> list[Any]: + """Returns a list of arguments on which to call __torch_function__. + + Checks arguments in relevant_args for __torch_function__ implementations, + storing references to the arguments and their types in overloaded_args and + overloaded_types in order of calling precedence. Only distinct types are + considered. If a type is a subclass of another type it will have higher + precedence, otherwise the precedence order is the same as the order of + arguments in relevant_args, that is, from left-to-right in the argument list. + + The precedence-determining algorithm implemented in this function is + described in `NEP-0018`_. + + See torch::append_overloaded_arg for the equivalent function in the C++ + implementation. + + Parameters + ---------- + relevant_args : iterable of array-like + Iterable of array-like arguments to check for __torch_function__ + methods. + + get_type_fn : callable, optional + Function to call on each argument in relevant_args to get its type. + + Returns + ------- + overloaded_args : list + Arguments from relevant_args on which to call __torch_function__ + methods, in the order in which they should be called. + + .. _NEP-0018: + https://numpy.org/neps/nep-0018-array-function-protocol.html + """ + if get_type_fn is None: + get_type_fn = type + + # If torch function is not enabled, there are no overloaded types + if not torch._C._is_torch_function_enabled(): + return [] + # Runtime is O(num_arguments * num_unique_types) + overloaded_types: set[type] = set() + overloaded_args: list[Any] = [] + for arg in relevant_args: + arg_type = get_type_fn(arg) + # We only collect arguments if they have a unique type, which ensures + # reasonable performance even with a long list of possibly overloaded + # arguments. + # + # NB: Important to exclude _disabled_torch_function_impl, otherwise + # https://github.com/pytorch/pytorch/issues/64687 + if ( + arg_type not in overloaded_types + and hasattr(arg_type, "__torch_function__") + and arg_type.__torch_function__ != torch._C._disabled_torch_function_impl + ): + # Create lists explicitly for the first type (usually the only one + # done) to avoid setting up the iterator for overloaded_args. + if overloaded_types: + overloaded_types.add(arg_type) + # By default, insert argument at the end, but if it is + # subclass of another argument, insert it before that argument. + # This ensures "subclasses before superclasses". + index = len(overloaded_args) + for i, old_arg in enumerate(overloaded_args): + if issubclass(arg_type, get_type_fn(old_arg)): + index = i + break + overloaded_args.insert(index, arg) + else: + overloaded_types = {arg_type} + overloaded_args = [arg] + return overloaded_args + + +def handle_torch_function( + public_api: Callable, + relevant_args: Iterable[Any], + *args, + **kwargs, +) -> Any: + """Implement a function with checks for ``__torch_function__`` overrides. + + See torch::autograd::handle_torch_function for the equivalent of this + function in the C++ implementation. + + Arguments + --------- + public_api : function + Function exposed by the public torch API originally called like + ``public_api(*args, **kwargs)`` on which arguments are now being + checked. + relevant_args : iterable + Iterable of arguments to check for __torch_function__ methods. + args : tuple + Arbitrary positional arguments originally passed into ``public_api``. + kwargs : tuple + Arbitrary keyword arguments originally passed into ``public_api``. + + Returns + ------- + object + Result from calling ``implementation`` or an ``__torch_function__`` + method, as appropriate. + + Raises + ------ + TypeError : if no implementation is found. + + Example + ------- + >>> def func(a): + ... if has_torch_function_unary(a): + ... return handle_torch_function(func, (a,), a) + ... return a + 0 + """ + # Check for __torch_function__ methods. + overloaded_args = _get_overloaded_args(relevant_args) + # overloaded_args already have unique types. + types = tuple(map(type, overloaded_args)) + + # Check for __torch_function__ mode. + if _is_torch_function_mode_enabled(): + # if we're here, the mode must be set to a TorchFunctionStackMode + # this unsets it and calls directly into TorchFunctionStackMode's torch function + with _pop_mode_temporarily() as mode: + result = mode.__torch_function__(public_api, types, args, kwargs) + if result is not NotImplemented: + return result + + # Call overrides + for overloaded_arg in overloaded_args: + # This call needs to become a classmethod call in the future. + # See https://github.com/pytorch/pytorch/issues/63767 + torch_func_method = overloaded_arg.__torch_function__ + if ( + hasattr(torch_func_method, "__self__") + and torch_func_method.__self__ is overloaded_arg + and torch_func_method is not torch._C._disabled_torch_function_impl + ): + warnings.warn( + "Defining your `__torch_function__ as a plain method is deprecated and " + "will be an error in future, please define it as a classmethod.", + DeprecationWarning, + ) + + # Use `public_api` instead of `implementation` so __torch_function__ + # implementations can do equality/identity comparisons. + result = torch_func_method(public_api, types, args, kwargs) + + if result is not NotImplemented: + return result + + func_name = f"{public_api.__module__}.{public_api.__name__}" + msg = ( + f"no implementation found for '{func_name}' on types that implement " + f"__torch_function__: {[type(arg) for arg in overloaded_args]}" + ) + if _is_torch_function_mode_enabled(): + msg += f" nor in mode {_get_current_function_mode()}" + raise TypeError(msg) + + +has_torch_function = _add_docstr( + _has_torch_function, + r"""Check for __torch_function__ implementations in the elements of an iterable + or if a __torch_function__ mode is enabled. Considers exact ``Tensor`` s + and ``Parameter`` s non-dispatchable. Use this to guard a call to + :func:`handle_torch_function`; don't use it to test if something + is Tensor-like, use :func:`is_tensor_like` instead. + Arguments + --------- + relevant_args : iterable + Iterable or arguments to check for __torch_function__ methods. + Returns + ------- + bool + True if any of the elements of relevant_args have __torch_function__ + implementations, False otherwise. + See Also + ________ + torch.is_tensor_like + Checks if something is a Tensor-like, including an exact ``Tensor``. + """, +) + +has_torch_function_unary = _add_docstr( + _has_torch_function_unary, + r"""Special case of `has_torch_function` for single inputs. + Instead of: + `has_torch_function((t,))` + call: + `has_torch_function_unary(t)` + which skips unnecessary packing and unpacking work. + """, +) + +has_torch_function_variadic = _add_docstr( + _has_torch_function_variadic, + r"""Special case of `has_torch_function` that skips tuple creation. + + This uses the METH_FASTCALL protocol introduced in Python 3.7 + + Instead of: + `has_torch_function((a, b))` + call: + `has_torch_function_variadic(a, b)` + which skips unnecessary packing and unpacking work. + """, +) + + +@functools.cache +def _get_overridable_functions() -> tuple[ + dict[Any, list[Callable]], dict[Callable, str] +]: + overridable_funcs = collections.defaultdict(list) + index = {} + tested_namespaces = [ + ("torch", torch, torch.__all__), + ("torch.functional", torch.functional, torch.functional.__all__), + ("torch.nn.functional", torch.nn.functional, dir(torch.nn.functional)), + ("torch.nn.init", torch.nn.init, dir(torch.nn.init)), + ("torch.Tensor", torch.Tensor, dir(torch.Tensor)), + ("torch.linalg", torch.linalg, dir(torch.linalg)), + ("torch.fft", torch.fft, dir(torch.fft)), + ("torch.special", torch.special, dir(torch.special)), + ] + for namespace_str, namespace, ns_funcs in tested_namespaces: + for func_name in ns_funcs: + ignore = False + # ignore private functions or functions that are deleted in torch.__init__ + if namespace is not torch.Tensor: + if func_name.startswith("__"): + continue + elif func_name.startswith("_"): + ignore = True + elif func_name.endswith("_"): + ignore = True + elif not func_name[0].islower(): + ignore = True + elif func_name == "unique_dim": + continue + else: + func = getattr(namespace, func_name) + if getattr(object, func_name, None) == func: + continue + if func_name == "__weakref__": + continue + func = getattr(namespace, func_name) + if namespace is torch.Tensor and getattr(object, func_name, None) == func: + continue + # ignore re-exported modules + if isinstance(func, types.ModuleType): + continue + # ignore __future__ imports + if isinstance(func, __future__._Feature): + continue + + if not callable(func) and hasattr(func, "__get__"): + index[func.__get__] = f"{namespace_str}.{func_name}.__get__" + index[func.__set__] = f"{namespace_str}.{func_name}.__set__" + if ignore: + continue + if func.__get__ in get_ignored_functions(): + msg = ( + "{}.{} is in the tuple returned by torch._overrides.get_ignored_functions " + "but still has an explicit override" + ) + assert func.__get__ not in get_testing_overrides(), msg.format( + namespace, func.__name__ + ) + continue + else: + overridable_funcs[func].append(func.__get__) + continue + + if not callable(func): + continue + + index[func] = f"{namespace_str}.{func_name}" + + if ignore: + continue + + # cannot be overridden by __torch_function__ + if func in get_ignored_functions(): + msg = ( + "{}.{} is in the tuple returned by torch._overrides.get_ignored_functions " + "but still has an explicit override" + ) + assert func not in get_testing_overrides(), msg.format( + namespace, func.__name__ + ) + continue + overridable_funcs[namespace].append(func) + return overridable_funcs, index + + +@_disable_user_warnings +def get_overridable_functions() -> dict[Any, list[Callable]]: + """List functions that are overridable via __torch_function__ + + Returns + ------- + Dict[Any, List[Callable]] + A dictionary that maps namespaces that contain overridable functions + to functions in that namespace that can be overridden. + """ + return _get_overridable_functions()[0] + + +@_disable_user_warnings +def resolve_name(f): + """Get a human readable string name for a function passed to + __torch_function__ + + Arguments + --------- + f : Callable + Function to resolve the name of. + + Returns + ------- + str + Name of the function; if eval'ed it should give back the input + function. + """ + if isinstance(f, (torch._ops.OpOverload, torch._ops.OpOverloadPacket)): + return str(f) + return _get_overridable_functions()[1].get(f) + + +@functools.cache +def _get_tensor_methods() -> set[Callable]: + """Returns a set of the overridable methods on ``torch.Tensor``""" + overridable_funcs = get_overridable_functions() + methods = set(overridable_funcs[torch.Tensor]) + return methods + + +@_disable_user_warnings +def is_tensor_method_or_property(func: Callable) -> bool: + """ + Returns True if the function passed in is a handler for a + method or property belonging to ``torch.Tensor``, as passed + into ``__torch_function__``. + + .. note:: + For properties, their ``__get__`` method must be passed in. + + This may be needed, in particular, for the following reasons: + + 1. Methods/properties sometimes don't contain a `__module__` slot. + 2. They require that the first passed-in argument is an instance + of ``torch.Tensor``. + + Examples + -------- + >>> is_tensor_method_or_property(torch.Tensor.add) + True + >>> is_tensor_method_or_property(torch.add) + False + """ + return func in _get_tensor_methods() or func.__name__ == "__get__" + + +def is_tensor_like(inp): + """ + Returns ``True`` if the passed-in input is a Tensor-like. + + Currently, this occurs whenever there's a ``__torch_function__`` + attribute on the type of the input. + + Examples + -------- + A subclass of tensor is generally a Tensor-like. + + >>> class SubTensor(torch.Tensor): ... + >>> is_tensor_like(SubTensor([0])) + True + + Built-in or user types aren't usually Tensor-like. + + >>> is_tensor_like(6) + False + >>> is_tensor_like(None) + False + >>> class NotATensor: ... + >>> is_tensor_like(NotATensor()) + False + + But, they can be made Tensor-like by implementing __torch_function__. + + >>> class TensorLike: + ... @classmethod + ... def __torch_function__(cls, func, types, args, kwargs): + ... return -1 + >>> is_tensor_like(TensorLike()) + True + """ + return type(inp) is torch.Tensor or hasattr(inp, "__torch_function__") + + +class TorchFunctionMode: + """ + A ``TorchFunctionMode`` allows you to override the meaning of all + ``__torch_function__`` overridable functions within a dynamic scope, + without having to actually create a tensor subclass or manually + monkey-patch functions in the PyTorch API. Some common situations + where you should use a mode: + + * You want to override the meaning of factory functions, or other + functions that do not otherwise take a tensor as an argument + (these cannot be overridden with tensor subclasses). + + * You want to override the behavior of all functions without needing + to wrap your inputs in tensor subclasses; e.g., if you are just + interested in logging intermediate computations. + + * You want to control the order of execution of various tensor + subclasses explicitly, rather than implicitly via the return of + ``NotImplemented``. + + Independent subclasses of :class:`TorchFunctionMode` are compositional: + modes can be pushed onto a stack using ``with MyMode():``. + When you call functions in the PyTorch API inside your + ``__torch_function__`` implementation, by default, they will forward on to + the next mode on the mode stack. If you want recursively call back into + your current ``__torch_function__`` implementation, either explicitly + invoke ``self.__torch_function__(...)``, or use the context manager + ``enable_torch_function_mode(self, replace=self.inner)`` to make PyTorch + API self-referential (beware of infinite loops, in this case!) + """ + + inner: "TorchFunctionMode" + + # Force metaclass to generate constructor at the base of the hierarchy + def __init__(self) -> None: + pass + + def __torch_function__(self, func, types, args=(), kwargs=None): + raise NotImplementedError + + def __enter__(self): + _push_mode(self) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + _pop_mode() + + @classmethod + def push(cls, *args, **kwargs): + warnings.warn( + "`Mode.push()` is no longer necessary and can be replaced with just `with Mode()`" + ) + instance = cls(*args, **kwargs) + return instance + + +def _get_current_function_mode(): + stack_len = _len_torch_function_stack() + return _get_function_stack_at(stack_len - 1) if stack_len > 0 else None + + +def _get_current_function_mode_stack(): + stack_len = _len_torch_function_stack() + return [_get_function_stack_at(i) for i in range(stack_len)] + + +def _push_mode(mode): + _push_on_torch_function_stack(mode) + + +def _pop_mode(): + old = _pop_torch_function_stack() + return old + + +@contextlib.contextmanager +def _pop_mode_temporarily(): + old = _pop_mode() + try: + yield old + finally: + _push_mode(old) + + +class BaseTorchFunctionMode(TorchFunctionMode): + def __torch_function__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + return func(*args, **kwargs) + + +@contextlib.contextmanager +def _enable_torch_function(): + old_state = torch._C._get_torch_function_state() + try: + torch._C._set_torch_function_state(torch._C._TorchFunctionState.ENABLED) + yield + finally: + torch._C._set_torch_function_state(old_state) + + +@contextlib.contextmanager +def enable_reentrant_dispatch(): + # NB: this can't simply be + # `enable_reentrant_dispatch = torch._C._RestorePythonTLSSnapshot` + # because: + # 1. torch._C._RestorePythonTLSSnapshot is unavailable when this file + # initially gets imported. Probably an import order thing. + # 2. enable_reentrant_dispatch is technically public API; assigning + # it the object would change the __module__ to look private. + with torch._C._RestorePythonTLSSnapshot(): + try: + yield + finally: + pass diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__init__.py b/venv/lib/python3.10/site-packages/torch/profiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..153d4560e264198ddb8f04649aef22456afbb452 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/__init__.py @@ -0,0 +1,60 @@ +r""" +PyTorch Profiler is a tool that allows the collection of performance metrics during training and inference. +Profiler's context manager API can be used to better understand what model operators are the most expensive, +examine their input shapes and stack traces, study device kernel activity and visualize the execution trace. + +.. note:: + An earlier version of the API in :mod:`torch.autograd` module is considered legacy and will be deprecated. + +""" + +import os +from typing import Any +from typing_extensions import TypeVarTuple, Unpack + +from torch._C._autograd import _supported_activities, DeviceType, kineto_available +from torch._C._profiler import _ExperimentalConfig, ProfilerActivity, RecordScope +from torch._environment import is_fbcode +from torch.autograd.profiler import KinetoStepTracker, record_function +from torch.optim.optimizer import Optimizer, register_optimizer_step_post_hook + +from .profiler import ( + _KinetoProfile, + ExecutionTraceObserver, + profile, + ProfilerAction, + schedule, + supported_activities, + tensorboard_trace_handler, +) + + +__all__ = [ + "profile", + "schedule", + "supported_activities", + "tensorboard_trace_handler", + "ProfilerAction", + "ProfilerActivity", + "kineto_available", + "DeviceType", + "record_function", + "ExecutionTraceObserver", +] + +from . import itt + + +_Ts = TypeVarTuple("_Ts") + + +def _optimizer_post_hook( + optimizer: Optimizer, args: tuple[Unpack[_Ts]], kwargs: dict[str, Any] +) -> None: + KinetoStepTracker.increment_step("Optimizer") + + +if os.environ.get("KINETO_USE_DAEMON", "") or ( + is_fbcode() and os.environ.get("KINETO_FORCE_OPTIMIZER_HOOK", "") +): + _ = register_optimizer_step_post_hook(_optimizer_post_hook) diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a639b2a8b282a9571a49482000544d3fd5325b11 Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_memory_profiler.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_memory_profiler.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..001be10b3900447862595022e726addbd355f12a Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_memory_profiler.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_pattern_matcher.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_pattern_matcher.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..10de55693331e30146914b0ec64b0fb8f32d93d2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_pattern_matcher.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d126a38cb17818cca7d1779f03e71e1b990402c7 Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/_utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/itt.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/itt.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a44574eee09e44ba5b8fe37a36e12012c90512df Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/itt.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/python_tracer.cpython-310.pyc b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/python_tracer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12eb13dd6a484b6a3dc6be20a25d25e1fbf4aed2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/torch/profiler/__pycache__/python_tracer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py b/venv/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..d9f3a917c15253e54c9f8dd7d3be724ef6ae1b69 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/_memory_profiler.py @@ -0,0 +1,1198 @@ +# mypy: allow-untyped-defs +import collections +import dataclasses +import enum +import itertools as it +import logging +from collections.abc import Iterator +from typing import Any, cast, Optional, Union +from typing_extensions import Literal + +import torch +from torch._C import FunctionSchema +from torch._C._autograd import _ProfilerResult +from torch._C._profiler import ( + _EventType, + _ExtraFields_Allocation, + _ExtraFields_TorchOp, + _ProfilerEvent, + _TensorMetadata, + RecordScope, +) +from torch._utils import _element_size +from torch.profiler import _utils + + +KeyAndID = tuple["Key", int] +TensorAndID = tuple["TensorKey", int] + +log = logging.getLogger(__name__) + + +class Category(enum.Enum): + INPUT = enum.auto() + TEMPORARY = enum.auto() + ACTIVATION = enum.auto() + GRADIENT = enum.auto() + AUTOGRAD_DETAIL = enum.auto() + PARAMETER = enum.auto() + OPTIMIZER_STATE = enum.auto() + + +_CATEGORY_TO_COLORS = { + Category.PARAMETER: "darkgreen", + Category.OPTIMIZER_STATE: "goldenrod", + Category.INPUT: "black", + Category.TEMPORARY: "mediumpurple", + Category.ACTIVATION: "red", + Category.GRADIENT: "mediumblue", + Category.AUTOGRAD_DETAIL: "royalblue", + None: "grey", +} + +_CATEGORY_TO_INDEX = {c: i for i, c in enumerate(_CATEGORY_TO_COLORS)} + + +class Action(enum.Enum): + PREEXISTING = enum.auto() + CREATE = enum.auto() + INCREMENT_VERSION = enum.auto() + DESTROY = enum.auto() + + +_ACTION_TO_INDEX = {i: i.value for i in Action} + + +@dataclasses.dataclass(eq=True, unsafe_hash=False, frozen=True) +class Key: + device: torch.device + + +@dataclasses.dataclass +class _Storage: + """Bundle storage pointer and id. + + All profiling logic should use `allocation_id`, however it is useful to + print storage pointers for debugging and unit tests sometimes look up + values using the storage data pointer of a live Tensor.""" + + ptr: int + allocation_id: int + + def __repr__(self) -> str: + return f"{hex(self.ptr):>18} ({self.allocation_id})" + + def __eq__(self, other: object) -> bool: + return isinstance(other, _Storage) and self.allocation_id == other.allocation_id + + def __hash__(self) -> int: + return hash(self.allocation_id) + + +@dataclasses.dataclass(eq=True, unsafe_hash=True, frozen=True) +class TensorKey(Key): + """Hashable identifier for a storage which has been assigned an ID. + + A detailed description of Tensor IDs and why they are needed is given in + `torch/csrc/profiler/collection.h` when `TensorID` is declared. To + summarize, multiple Storage buffers can map to the same logical Tensor. + This dataclass is used to refer to a concrete in-memory StorageImpl of + a Tensor. + """ + + id: int + storage: _Storage + + def __repr__(self) -> str: + return f"id={self.id}: {repr(self.storage):<24} ({self.device})" + + def __lt__(self, other: "TensorKey") -> bool: + return self._as_sortable < other._as_sortable + + @staticmethod + def _make( + tensor_id: Optional[int], + storage_ptr: Optional[int], + allocation_id: Optional[int], + device: torch.device, + ) -> Optional["TensorKey"]: + if ( + tensor_id is not None + and storage_ptr is not None + and allocation_id is not None + ): + return TensorKey(device, tensor_id, _Storage(storage_ptr, allocation_id)) + return None + + @classmethod + def from_allocation(cls, alloc: _ExtraFields_Allocation) -> Optional["TensorKey"]: + return cls._make(alloc.id, alloc.ptr, alloc.allocation_id, alloc.device) + + @classmethod + def from_tensor(cls, t: Optional[_TensorMetadata]) -> Optional["TensorKey"]: + if t is not None: + return cls._make(t.id, t.storage_data_ptr, t.allocation_id, t.device) + return None + + @property + def _as_sortable(self) -> tuple[int, int, str, int]: + return self.id, self.storage.allocation_id, self.device.type, self.device.index + + +def _extract_parameters_and_gradients( + node: _ProfilerEvent, +) -> Iterator[tuple[Optional[TensorKey], Optional[TensorKey]]]: + children = node.children + + # AccumulateGrad is used in the Autograd engine to handle gradient updates. + # There are two possible cases: + # 1) This is a newly created gradient Tensor. In that case there is nothing + # to accumulate, so autograd simply detaches the Tensor. + # + # 2) There is a preexisting gradient Tensor and we need to add the newly + # computed update. This is done with an in-place add (aten::add_) op. + # (The underscore suffix denotes "in-place".) + if ( + node.typed[0] == _EventType.TorchOp + and node.typed[1].scope == RecordScope.BACKWARD_FUNCTION + # TODO(robieta): Move away from load bearing names + and node.name == "torch::autograd::AccumulateGrad" + and children + and children[0].typed[0] == _EventType.TorchOp + and children[0].name in ("aten::detach", "aten::add_") + and children[0].typed[1].inputs + and isinstance(children[0].typed[1].inputs[0], _TensorMetadata) + ): + yield None, TensorKey.from_tensor(children[0].typed[1].inputs[0]) + + # We directly instrument `torch.nn.Module` and `torch.optim.Optimizer` + # NOTE: The values captured by the python tracer are cached; they can be + # used to build up labels but do not imply that a Tensor was live at + # a particular time. + elif node.typed[0] == _EventType.PyCall: + typed_fields = node.typed[1] + assert typed_fields.module is None or typed_fields.optimizer is None + if typed_fields.module is not None: + for _, p, p_grad in typed_fields.module.parameters: + yield TensorKey.from_tensor(p), TensorKey.from_tensor(p_grad) + + if typed_fields.optimizer is not None: + for p, p_grad, _ in typed_fields.optimizer.parameters: + yield TensorKey.from_tensor(p), TensorKey.from_tensor(p_grad) + + +def extract_parameters(node: _ProfilerEvent) -> Iterator[TensorKey]: + for p, _p_grad in _extract_parameters_and_gradients(node): + if p is not None: + yield p + + +def extract_gradients( + node: _ProfilerEvent, +) -> Iterator[tuple[Optional[TensorKey], TensorKey]]: + for p, p_grad in _extract_parameters_and_gradients(node): + if p_grad is not None: + yield p, p_grad + + +def get_scopes(event: Optional[_ProfilerEvent]) -> tuple[RecordScope, ...]: + scopes = [] + while event: + if event.typed[0] == _EventType.TorchOp: + scopes.append(event.typed[1].scope) + event = event.parent + return tuple(scopes) + + +class SchemaMatcher: + """Lookup operator schema based on profiled name. + + When profiling we record the operator's name but not the schema. However + some analysis requires that information. Fortunately we can look up + registered schema from the recorded name. We do not, however, record the + overload and so we must compare the profiled arguments with all overloads + to determine viable matches. + + Note: Once https://github.com/pytorch/pytorch/issues/78871 is completed + this code will be obsolete. + """ + + @classmethod + def inputs_are_mutable(cls, t: _ExtraFields_TorchOp) -> tuple[Optional[bool], ...]: + """Determine which inputs may have mutated based on function schema. + + Note that we don't need to resolve down to a single schema to perform + this analysis. An input is mutable if it is mutable in any overload. In + practice, however, it is overwhelmingly common to match a single + overload. If we cannot find any valid schema then we must be + conservative and assume all inputs are mutable. + """ + mutable: Optional[list[bool]] = None + for schema in cls.match_schemas(t): + mutable = mutable or [False for _ in schema.arguments] + for i, arg in enumerate(schema.arguments): + mutable[i] |= getattr(arg.alias_info, "is_write", False) + + return tuple(mutable or (None for _ in t.inputs)) + + @classmethod + def match_schemas(cls, t: _ExtraFields_TorchOp) -> tuple[FunctionSchema, ...]: + signature = tuple( + # Tensor + TensorKey.from_tensor(i) + if isinstance(i, _TensorMetadata) + # + # TensorList + else [TensorKey.from_tensor(j) for j in i] + if isinstance(i, list) + # + # Scalar and uncaptured inputs. + else i + for i in t.inputs + ) + + def matches(schema) -> bool: + return len(schema.arguments) == len(signature) and all( + cls._types_match(observed, schema_arg.type) + for observed, schema_arg in zip(signature, schema.arguments) + ) + + return tuple(s for s in cls.lookup_schemas(t.name) or () if matches(s)) + + @classmethod + def _types_match(cls, observed, schema_type) -> bool: + if isinstance(schema_type, torch._C.OptionalType): + schema_type = schema_type.getElementType() + return observed is None or cls._types_match(observed, schema_type) + + if isinstance(schema_type, torch._C.AnyType): + return True + + if schema_type.isSubtypeOf(torch._C.ListType.ofTensors()): + return isinstance(observed, list) and all( + isinstance(i, TensorKey) for i in observed + ) + + type_map: tuple[tuple[Any, Union[type, tuple[type, ...]]], ...] = ( + (torch._C.TensorType, TensorKey), + (torch._C.NoneType, type(None)), + (torch._C.BoolType, bool), + (torch._C.IntType, int), + (torch._C.FloatType, float), + (torch._C.ComplexType, complex), + (torch._C.NumberType, (bool, int, float, complex)), + ) + + for jit_type, py_types in type_map: + if isinstance(schema_type, jit_type): + return isinstance(observed, py_types) + + # Profiler only records a subset of possible argument types. If we + # reach this point then the schema must call for a type that profiler + # does not record. Thus, the schema can only be a match if `observed` + # is also None. + return observed is None + + @staticmethod + def lookup_schemas(name: str) -> Optional[tuple[FunctionSchema, ...]]: + # TODO(robieta): + # _jit_get_schemas_for_operator is quite expensive. (~100us / call) + # Consider adding `functools.lru_cache` if that becomes an issue. + + try: + # Schema lookup will throw if `name` is malformed. (For example, + # schemas must be namespaced and schema lookup will fail if name + # does not include "::".) We simply catch the exception and return + # `None` to denote that `name` cannot be an operator name. + # + # Note that record_function annotations also go through this path, + # so it is expected that some names will not correspond to PyTorch + # operators. + if "::" not in name: + return None + return tuple(torch._C._jit_get_schemas_for_operator(name)) + except RuntimeError: + return None + + +class OpTree: + def __init__(self, result: _ProfilerResult) -> None: + self._root_nodes = result.experimental_event_tree() + self._sorted_nodes = tuple(sorted(self.dfs(), key=lambda x: x.start_time_ns)) + + def dfs(self, *args, **kwargs) -> Iterator[_ProfilerEvent]: + yield from _utils.traverse_dfs(self._root_nodes, *args, **kwargs) + + @property + def sorted_nodes(self) -> tuple[_ProfilerEvent, ...]: + return self._sorted_nodes + + +class SizeMap: + def __init__(self, op_tree: OpTree) -> None: + self._values: dict[TensorKey, int] = {} + + for node in op_tree.sorted_nodes: + if node.typed[0] == _EventType.TorchOp: + for t in self._flat_tensor_inputs(node.typed[1]): + self._update_values(t) + + elif node.typed[0] == _EventType.PyCall: + typed_fields = node.typed[1] + assert typed_fields.module is None or typed_fields.optimizer is None + if typed_fields.module is not None: + for _, p, p_grad in typed_fields.module.parameters: + self._update_values(p) + self._update_values(p_grad) + + if typed_fields.optimizer is not None: + for p, p_grad, state in typed_fields.optimizer.parameters: + self._update_values(p) + self._update_values(p_grad) + for _, t in state: + self._update_values(t) + + allocations: dict[TensorKey, int] = {} + for node in op_tree.sorted_nodes: + if node.typed[0] == _EventType.Allocation: + alloc_fields = node.typed[1] + key = TensorKey.from_allocation(alloc_fields) + if key: + new_size = abs(alloc_fields.alloc_size) + prior_size = allocations.setdefault(key, new_size) + + # It is possible to resize Storage in PyTorch, however we + # key on data pointer so most resizes will be treated as a + # change in storage. The one corner case that cannot be + # handled is `realloc` which successfully resizes the + # storage. At time of writing this is not done anywhere in + # the core PyTorch codebase. + if prior_size != new_size: + delta = f"{prior_size} vs. {new_size}" + log.warning("Mismatch between allocation and free: %s", delta) + + self._values.update(allocations) + + def _update_values(self, t: Optional[_TensorMetadata]) -> None: + key = TensorKey.from_tensor(t) + if key is not None and t is not None and t.layout == torch.strided: + # Scalars are represented as zero dim Tensors + n = max(i[0] * i[1] for i in zip(t.sizes or [1], t.strides or [1])) + + num_bytes = n * _element_size(t.dtype) + assert num_bytes >= 0, f"{num_bytes}" + self._values[key] = max(self._values.get(key, 0), num_bytes) + + @staticmethod + def _flat_tensor_inputs(op: _ExtraFields_TorchOp) -> Iterator[_TensorMetadata]: + for i in op.inputs: + if isinstance(i, _TensorMetadata): + yield i + elif isinstance(i, list): + yield from i + + def __getitem__(self, key: TensorKey): + return self._values[key] + + +@dataclasses.dataclass() +class DataFlowEdge: + input_version: Optional[int] = None + mutated: Optional[bool] = False + + @property + def is_allocation(self) -> bool: + return self.input_version is None + + @property + def is_deletion(self) -> bool: + return self.mutated is None + + +class DataFlowNode: + def __init__(self, event: _ProfilerEvent, graph: "DataFlowGraph") -> None: + self._event = event + self._graph = graph + self._edges: dict[TensorKey, DataFlowEdge] = self._determine_edges() + + for key, edge in self._edges.items(): + if edge.mutated and not edge.is_allocation: + self._graph.bump(key) + + # Make sure the version bumping behavior matches what we expect. + versions = {k: (v, self._graph.lookup(k)) for k, v in self.outputs.items()} + assert all(i == j for i, j in versions.values()), f"{versions}, {self._edges}" + + def _determine_edges(self) -> dict[TensorKey, DataFlowEdge]: + subtree = tuple(_utils.traverse_dfs([self._event])) + + # Start by populating edges from op inputs and outputs. + mutable_by_key: dict[Optional[TensorKey], set[Optional[bool]]] = {} + for op in (i.typed[1] for i in subtree if i.typed[0] == _EventType.TorchOp): + for op_input, mutable in zip( + op.inputs, SchemaMatcher.inputs_are_mutable(op) + ): + # Tensor + if isinstance(op_input, _TensorMetadata): + key = TensorKey.from_tensor(op_input) + mutable_by_key.setdefault(key, set()).add(mutable) + + # TensorList + elif isinstance(op_input, list): + for op_input_i in op_input: + key = TensorKey.from_tensor(op_input_i) + mutable_by_key.setdefault(key, set()).add(mutable) + + edges: collections.defaultdict[Optional[TensorKey], DataFlowEdge] + edges = collections.defaultdict(DataFlowEdge) + for key, mutable_set in mutable_by_key.items(): + if key is not None: + edges[key].input_version = self._graph.lookup(key) if key else -1 + + # We consider an op to be mutated if we encounter a schema where it + # is a mutable argument OR if it is ambiguous. (We never explicitly + # see it in any schema.) + mutated = (True in mutable_set) or (tuple(mutable_set) == (None,)) + edges[key].mutated = mutated + + # Then handle deletions. Note that deleting a Tensor implicitly adds + # it as an input edge. + for i in subtree: + if i.typed[0] == _EventType.Allocation and i.typed[1].alloc_size < 0: + key = TensorKey.from_allocation(i.typed[1]) + edge = edges[key] + assert key is None or edge.mutated is not None, f"Double delete: {key}" + edge.mutated = None + edge.input_version = self._graph.lookup(key) if key else -1 + + # And finally handle allocations. This step must be last, because the + # previous two steps optimistically add input edges. + for i in subtree: + if i.typed[0] == _EventType.Allocation and i.typed[1].alloc_size > 0: + edges[TensorKey.from_allocation(i.typed[1])].input_version = None + + # We don't need to sort the inputs, but it makes debugging and unit tests nicer. + return dict(sorted((k, v) for k, v in edges.items() if k is not None)) + + @property + def inputs(self) -> dict[TensorKey, tuple[bool, int]]: + return { + # MyPy can't see through `is_allocation` to know that + # `v.input_version` is not None. + k: (bool(v.mutated), cast(int, v.input_version)) + for k, v in self._edges.items() + if not v.is_allocation + } + + @property + def outputs(self) -> dict[TensorKey, int]: + return { + k: 0 if v.input_version is None else v.input_version + 1 + for k, v in self._edges.items() + if (v.is_allocation and not v.is_deletion) or v.mutated + } + + @property + def intermediates(self) -> tuple[TensorKey, ...]: + return tuple( + k for k, v in self._edges.items() if v.is_allocation and v.is_deletion + ) + + @property + def start_time(self) -> int: + return self._event.start_time_ns + + +class DataFlowGraph: + def __init__(self, op_tree: OpTree) -> None: + self._op_tree = op_tree + self._leaf_events = self._extract_leaf_events(op_tree) + self._active_version: dict[TensorKey, Optional[int]] = {} + self._flow_nodes = [DataFlowNode(e, self) for e in self.leaf_events] + self._flow_nodes.sort(key=lambda x: x.start_time) + self.validate() + + @property + def flow_nodes(self) -> tuple[DataFlowNode, ...]: + return tuple(self._flow_nodes) + + def validate(self) -> None: + # Check that each (Tensor, version) pair has a unique creation node + outputs: set[tuple[TensorKey, int]] = set() + for node in self.flow_nodes: + node_outputs = set(node.outputs.items()) + duplicates = outputs & node_outputs + assert not duplicates, f"{node._event.name} {node._edges} {duplicates}" + outputs |= node_outputs + + # And check that `self._nodes` forms a valid topologically sorted DAG. + tensor_versions: dict[TensorKey, int] = {} + for node in self.flow_nodes: + for key, (_, version) in node.inputs.items(): + expected = tensor_versions.get(key, 0) + assert expected == version, (expected, version) + + for key, version in node.outputs.items(): + prior_version = tensor_versions.get(key, version) + assert version >= prior_version, (version, prior_version) + tensor_versions[key] = version + + @property + def leaf_events(self) -> tuple[_ProfilerEvent, ...]: + return self._leaf_events + + @staticmethod + def _extract_leaf_events(op_tree: OpTree) -> tuple[_ProfilerEvent, ...]: + """Partially traverse the op tree and extract top level ops. + + Consider the following code: + ``` + with record_function("My annotation"): + x.zero_() + y.zero_() + ``` + + The op tree (assuming no Autograd) will look like: + + TorchOp: "My annotation" + TorchOp: zero_ + TorchOp: fill_ + TorchOp: zero_ + TorchOp: fill_ + + The recursive structure of operator calls makes data flow unwieldy. + In order to simplify analysis we would like to select the highest level + ops to represent in the graph. In this case those are the `zero_` ops; + the fact that `fill_` is called is an implementation detail. We also + do not want to group everything under "My annotation" as this could + create overly coarse bundles and lose critical semantics. + + To address this issue we walk over the graph and select the topmost + torch ops ** which match at least one operator schema **. These form + the leaves of the first pass through the op tree. (As well as any + allocations or frees which do are not part of a kernel.) These events + form the logical nodes in our data flow graph. + """ + + leaf_events: list[_ProfilerEvent] = [] + + def leaf_op(e: _ProfilerEvent) -> bool: + return e.typed[0] == _EventType.TorchOp and ( + e.typed[1].scope == RecordScope.BACKWARD_FUNCTION + or bool(SchemaMatcher.match_schemas(e.typed[1])) + ) + + def children_fn(e: _ProfilerEvent): + if leaf_op(e) or e.tag == _EventType.Allocation: + leaf_events.append(e) + return [] + + return e.children + + for _ in op_tree.dfs(children_fn=children_fn): + pass + + return tuple(sorted(leaf_events, key=lambda x: x.start_time_ns)) + + def lookup(self, key: TensorKey) -> int: + version = self._active_version.setdefault(key, 0) + assert version is not None + return version + + def bump(self, key: TensorKey) -> None: + prior_version = self._active_version.get(key, None) + assert prior_version is not None + self._active_version[key] = prior_version + 1 + + def delete(self, key: TensorKey) -> None: + assert self._active_version.setdefault(key, 0) is not None + self._active_version[key] = None + + +@dataclasses.dataclass +class CategoryElement: + by_id: Optional[Category] = None + by_key: dict[TensorKey, Category] = dataclasses.field(default_factory=dict) + by_version: dict[TensorAndID, Category] = dataclasses.field(default_factory=dict) + + # Used by unit tests to check internals. (And consequently by + # MemoryProfile.lookup) This should not be used in any other capacity. + _by_id_keyset: set[TensorKey] = dataclasses.field(default_factory=set) + + +@dataclasses.dataclass +class CategoryDict: + _values: collections.defaultdict[int, CategoryElement] = dataclasses.field( + default_factory=lambda: collections.defaultdict(CategoryElement) + ) + + def set_by_id(self, key: TensorKey, category: Category) -> None: + self._values[key.id].by_id = category + self._values[key.id]._by_id_keyset.add(key) + + def set_by_key(self, key: TensorKey, category: Category) -> None: + self._values[key.id].by_key[key] = category + + def set_by_version(self, key: TensorKey, version: int, category: Category) -> None: + self._values[key.id].by_version[(key, version)] = category + + def setdefault_by_version( + self, key: TensorKey, version: int, category: Category + ) -> None: + self._values[key.id].by_version.setdefault((key, version), category) + + def get(self, key: Key, version: int) -> Optional[Category]: + if isinstance(key, Key) and not isinstance(key, TensorKey): + return None + element = self._values[key.id] + return ( + element.by_id + or element.by_key.get(key, None) + or element.by_version.get((key, version), None) + ) + + +class MemoryProfile: + def __init__(self, result: _ProfilerResult) -> None: + self._op_tree = OpTree(result) + self._data_flow_graph = DataFlowGraph(self._op_tree) + self._size_map = SizeMap(self._op_tree) + self._categories = CategoryDict() + + self._set_gradients_and_temporaries() + self._set_parameters_using_python_tracer() + self._set_inputs() + self._set_parameters_using_data_flow() + self._set_activations() + self._set_optimizer_state() + self._set_autograd_detail() + + @property + def timeline(self) -> tuple[tuple[int, Action, KeyAndID, int], ...]: + output: list[tuple[int, Action, KeyAndID, int]] = [] + allocation_times: dict[tuple[TensorKey, bool], int] = {} + live_unknown: dict[tuple[int, torch.device], Literal[True]] = {} + for event in self._op_tree.dfs(): + if event.typed[0] == _EventType.Allocation: + alloc_fields = event.typed[1] + alloc_size = alloc_fields.alloc_size + is_allocation = alloc_size > 0 + t = event.start_time_ns + + tkey = TensorKey.from_allocation(alloc_fields) + if tkey is not None: + allocation_times[(tkey, is_allocation)] = t + + else: + key = Key(alloc_fields.device) + ptr_and_device = (alloc_fields.ptr, key.device) + if is_allocation: + if ptr_and_device in live_unknown: + output.append( + (t, Action.INCREMENT_VERSION, (key, 0), alloc_size) + ) + else: + live_unknown[ptr_and_device] = True + output.append((t, Action.CREATE, (key, 0), alloc_size)) + else: + output.append((t, Action.DESTROY, (key, 0), -alloc_size)) + if not live_unknown.pop(ptr_and_device, False): + output.append( + (-1, Action.PREEXISTING, (key, 0), -alloc_size) + ) + + snapshot = self._category_snapshot() + last_version = dict(sorted(snapshot.keys())) + + events: list[tuple[int, Action, TensorAndID]] = [ + (-1, Action.PREEXISTING, (key, version)) + for key, version in snapshot.keys() + if (key, True) not in allocation_times and version == 0 + ] + + for node in self._data_flow_graph.flow_nodes: + for key, edge in node._edges.items(): + if edge.is_allocation: + t = allocation_times[(key, True)] + events.append((t, Action.CREATE, (key, 0))) + + elif edge.mutated: + t = node._event.start_time_ns + version = edge.input_version + assert version is not None + events.append((t, Action.INCREMENT_VERSION, (key, version))) + + if edge.is_deletion: + t = allocation_times[(key, False)] + events.append((t, Action.DESTROY, (key, last_version[key]))) + + output.extend( + (time, action, (key, version), self._size_map[key]) + for time, action, (key, version) in events + ) + + output.sort(key=lambda x: (x[0], x[1].value)) + return tuple(output) + + def _is_gradient(self, *args, **kwargs) -> bool: + return self._categories.get(*args, **kwargs) == Category.GRADIENT + + def _category_snapshot(self) -> dict[TensorAndID, Optional[Category]]: + all_tensor_versions: set[TensorAndID] = set() + + for node in self._data_flow_graph.flow_nodes: + all_tensor_versions.update(((k, v) for k, (_, v) in node.inputs.items())) + all_tensor_versions.update((key, 0) for key in node.intermediates) + all_tensor_versions.update(node.outputs.items()) + + for i in self._categories._values.values(): + all_tensor_versions.update((key, 0) for key in i._by_id_keyset) + + return { + (key, version): self._categories.get(key, version) + for key, version in sorted(all_tensor_versions) + } + + def _any_version_depends_on_gradient(self) -> set[int]: + """Extract IDs of Tensors which depend or will depend on a gradient. + + Note that this weakened definition of "depends" requires us to loop + over the data flow graph multiple times because it allows dependency + information to flow backward through edges and removes the guarantee + that nodes are topologically sorted. (Or indeed, even that a valid + topological order exists.) Put another way, we have converted an + acyclic data flow graph into a cyclic graph and we are attempting to + partition cycles involving a gradient from the rest of the graph. + """ + depends_on_gradient: set[int] = set() + while True: + start_size = len(depends_on_gradient) + for node in self._data_flow_graph.flow_nodes: + ids = tuple( + key.id + for key, (_, version) in node.inputs.items() + if self._categories.get(key, version) + in (Category.GRADIENT, Category.PARAMETER) + or key.id in depends_on_gradient + ) + + if ids: + depends_on_gradient.update(ids) + depends_on_gradient.update(key.id for key in node.outputs) + + # We are guaranteed to exit because there is a finite set of + # TensorAndID pairs. In practice we do not expect to loop more than + # three times: once to identify the core parameter update loop, + # once to fold the first step into that loop, and a third time + # where no new elements are added. + if len(depends_on_gradient) == start_size: + return depends_on_gradient + + def _set_gradients_and_temporaries(self) -> None: + """Mark Tensors which are unambiguous and simple to reason about.""" + + # Gradients are straightforward to detect. We directly check the + # `.grad` property in the Python tracer, and we can detect any new + # gradient Tensors from `AccumulateGrad` ops. + for event in self._op_tree.dfs(): + for _, p_grad in extract_gradients(event): + self._categories.set_by_id(p_grad, Category.GRADIENT) + + # Similarly, temporary Tensors are easy to identify and are useful to + # flag since they can make memory use "spikier" than one would + # otherwise expect. + for node in self._data_flow_graph.flow_nodes: + for i in node.intermediates: + self._categories.set_by_key(i, Category.TEMPORARY) + + def _set_parameters_using_python_tracer(self) -> None: + for event in self._op_tree.dfs(): + for p in extract_parameters(event): + if p is not None: + self._categories.set_by_id(p, Category.PARAMETER) + + def _set_inputs(self) -> None: + """Mark inputs based on which Tensors are updated using gradients. + + The process for differentiating between inputs and activations is more + involved. Most Tensors in a training loop depend on at least one + gradient: parameters depend on them through updates, and activations + and optimizer state depend on them transitively through parameters. + Critically, we do not need to know which Tensors are parameters to + apply this method; we can simply walk the data flow graph to build the + set of all values which depend on a gradient and then obtain the set + of inputs from the conjugate set. + + There is, however, one hiccup. The first time we see a parameter is + generally on the forward pass of the first step. We know from + inspection of the data flow graph that v1 of that Tensor depends on + a gradient (provided we profile an optimizer step), but not v0. To + address this problem we weaken the definition of "depends on a + gradient" to "any version of this Tensor depends on a gradient", + which in turn strengthens the criteria for the input set enough to + filter the activations in the forward pass of the first step.""" + + # All of this analysis is predicated on using at least one training + # step (or parameters from the python tracer) to partition the graph. + # Absent that we cannot determine which Tensors are inputs and which + # ones are part of the model. + depends_on_gradient = self._any_version_depends_on_gradient() + + # We only want to annotate Tensors which actually contribute to the + # model calculation. + produces_gradient: set[TensorAndID] = set() + for node in reversed(self._data_flow_graph.flow_nodes): + tensors = {(key, version) for key, (_, version) in node.inputs.items()} + tensors |= node.outputs.items() + if any( + self._categories.get(*i) in (Category.GRADIENT, Category.PARAMETER) + or i in produces_gradient + for i in tensors + ): + produces_gradient |= tensors + + # Don't include Tensors created in the backward pass, as these are + # generally Autograd implementation details rather than proper inputs. + input_candidates = produces_gradient.copy() + for node in self._data_flow_graph.flow_nodes: + if RecordScope.BACKWARD_FUNCTION in get_scopes(node._event): + input_candidates -= set(node.outputs.items()) + + for key, version in input_candidates: + if key.id not in depends_on_gradient: + self._categories.setdefault_by_version(key, version, Category.INPUT) + + def _set_parameters_using_data_flow(self) -> None: + """Deduce which Tensors are parameters. + + Consider the following code for the step of SGD with momentum + (nesterov=False), where `d_p` is the gradient of `param` and `buf` is + the momentum buffer. + ``` + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + d_p = buf + param.add_(d_p, alpha=-lr) + ``` + Both `param` and `buf` take a gradient and perform an in-place update. + + The python tracer will inspect calls to `nn.Module.forward` and + `optim.Optimizer.step` to extract parameter and optimizer state + respectively (including parameters), so this is generally a non-issue. + + However as a fallback we can also exploit several properties of + parameters to distinguish them from other model state. + + First, they are directly used in the forward pass. (At this point we + haven't established which parts of the graph correspond to the forward + pass but we can deduce enough to suffice.) Some mutable state such as + batch norm moving averages also contribute to the forward pass, but + optimizer state does not. + + Second, a parameter is by definition used to compute at least one + gradient and depends on at least one gradient. + """ + snapshot = self._category_snapshot() + + # Determine which Tensors might be parameters based on forward pass + # data flow. Note this these are only candidates; we filter nodes that + # we know are part of the backward pass but that doesn't guarantee that + # they are part of the forward pass. + candidate_parameters: set[TensorAndID] = set() + candidate_fwd_tensors: set[TensorAndID] = { + i for i, category in snapshot.items() if category == Category.INPUT + } + + for node in self._data_flow_graph.flow_nodes: + inputs = {(key, value) for key, (_, value) in node.inputs.items()} + if ( + # Don't check nodes in the backward pass. + RecordScope.BACKWARD_FUNCTION not in get_scopes(node._event) + and not any(self._is_gradient(*i) for i in inputs) + and not any(self._is_gradient(*i) for i in node.outputs.items()) + # + # and only check nodes which depend on an input. + and candidate_fwd_tensors.intersection(inputs) + ): + candidate_fwd_tensors |= node.outputs.items() + candidate_parameters |= inputs.difference(candidate_fwd_tensors) + + # Require that each parameter eventually contributes to the value of a gradient + used_for_gradient: set[TensorAndID] = set() + for node in reversed(self._data_flow_graph.flow_nodes): + if any( + self._is_gradient(*i) or i in used_for_gradient + for i in node.outputs.items() + ): + used_for_gradient.update( + (key, version) for key, (_, version) in node.inputs.items() + ) + candidate_parameters.intersection_update(used_for_gradient) + + # and depends on a gradient. + parameter_keys = {key.id for key, _ in candidate_parameters} + parameter_keys &= self._any_version_depends_on_gradient() + + for key, _ in snapshot.keys(): + if key.id in parameter_keys: + self._categories.set_by_id(key, Category.PARAMETER) + + def _set_activations(self) -> None: + """Flood the graph to identify activations.""" + + required = {Category.INPUT, Category.ACTIVATION} + also_allowed = {Category.PARAMETER, Category.TEMPORARY} + for node in self._data_flow_graph.flow_nodes: + inputs = {(key, value) for key, (_, value) in node.inputs.items()} + input_categories = {self._categories.get(*i) for i in inputs} + + if ( + (input_categories & required) + and not (input_categories - (required | also_allowed)) + # + # Stop filling when we reach the backward pass. + and RecordScope.BACKWARD_FUNCTION not in get_scopes(node._event) + ): + for i in node.outputs.items(): + self._categories.setdefault_by_version(*i, Category.ACTIVATION) + + def _set_optimizer_state(self) -> None: + for event in self._op_tree.dfs(): + if event.typed[0] == _EventType.PyCall and event.typed[1].optimizer: + parameters = event.typed[1].optimizer.parameters + for _, t in it.chain.from_iterable( + (state for _, _, state in parameters) + ): + key = TensorKey.from_tensor(t) + if key is not None: + self._categories.set_by_id(key, Category.OPTIMIZER_STATE) + + def _set_autograd_detail(self) -> None: + prior = {None, Category.AUTOGRAD_DETAIL} + for node in self._data_flow_graph.flow_nodes: + if RecordScope.BACKWARD_FUNCTION in get_scopes(node._event): + for key, version in node.outputs.items(): + if version == 0 or self._categories.get(key, version - 1) in prior: + self._categories.setdefault_by_version( + key, version, Category.AUTOGRAD_DETAIL + ) + + +class MemoryProfileTimeline: + def __init__(self, memory_profile) -> None: + """The minimum representation of the memory profile timeline + includes the memory timeline and categories. The timeline + consists of [timestamp, action, (TensorKey, version), numbytes] + elements, to denote any actions (pre-existing, create, destroy, + or increment_version) that occurred to a specific Tensor for a + chunk of memory. The categories help map each (TensorKey, + version) pair into a category.""" + self.timeline = memory_profile.timeline + self.categories = memory_profile._categories + + def _coalesce_timeline(self, device_str): + """Convert the memory timeline and categories into a memory plot + consisting of timestamps and their respective sizes by category + for a given device. + + Input: device + Output: [timestamps, sizes by category] + """ + device = torch.device(device_str) + times: list[int] = [] + sizes: list[list[int]] = [] + + def update(key, version, delta) -> None: + category = ( + self.categories.get(key, version) + if isinstance(key, TensorKey) + else None + ) + index = _CATEGORY_TO_INDEX[category] + 1 + sizes[-1][index] += int(delta) + + t_min = -1 + for t, action, (key, version), numbytes in self.timeline: + if key.device != device: + continue + + # Convert timestamps from ns to us, to match trace events. + if t != -1: + t = int(t / 1000) + + # Save the smallest timestamp to populate pre-existing allocs. + if t_min == -1 or (t < t_min and t > 0): + t_min = t + + # Handle timestep + if len(times) == 0: + times.append(t) + sizes.append([0] + [0 for _ in _CATEGORY_TO_INDEX]) + + elif t != times[-1]: + times.append(t) + sizes.append(sizes[-1].copy()) + + # Handle memory and categories + if action in (Action.PREEXISTING, Action.CREATE): + update(key, version, numbytes) + + elif action == Action.INCREMENT_VERSION: + update(key, version, -numbytes) + update(key, version + 1, numbytes) + + elif action == Action.DESTROY: + update(key, version, -numbytes) + + else: + raise ValueError(f"Unknown action: {action}") + + times = [t_min if t < 0 else t for t in times] + return times, sizes + + def export_memory_timeline(self, path, device_str) -> None: + """Saves the memory timeline as [times, sizes by category] + as a JSON formatted file to the given path for the given + device.""" + times, sizes = self._coalesce_timeline(device_str) + # TODO: Write a faster serialize (orjson not available in CI) + import json + + with open(path, "w") as f: + json.dump([times, sizes], f) + + def export_memory_timeline_raw(self, path, device_str) -> None: + """Saves the memory timeline as raw memory event tuples in the + form of (timestamp, action, numbytes, category) + as a JSON formatted file to the given path for the given + device.""" + device = torch.device(device_str) + raw_events: list[tuple[int, int, int, int]] = [] + + def get_category_index(key, version): + category = ( + self.categories.get(key, version) + if isinstance(key, TensorKey) + else None + ) + return _CATEGORY_TO_INDEX[category] + + for t, action, (key, version), numbytes in self.timeline: + if key.device != device: + continue + + if action in (Action.PREEXISTING, Action.CREATE): + raw_events.append( + ( + t, + _ACTION_TO_INDEX[action], + numbytes, + get_category_index(key, version), + ) + ) + + elif action == Action.INCREMENT_VERSION: + raw_events.append( + ( + t, + _ACTION_TO_INDEX[action], + -numbytes, + get_category_index(key, version), + ) + ) + raw_events.append( + ( + t, + _ACTION_TO_INDEX[action], + numbytes, + get_category_index(key, version + 1), + ) + ) + + elif action == Action.DESTROY: + raw_events.append( + ( + t, + _ACTION_TO_INDEX[action], + -numbytes, + get_category_index(key, version), + ) + ) + + else: + raise ValueError(f"Unknown action: {action}") + + import json + + with open(path, "w") as f: + json.dump(raw_events, f) + + def export_memory_timeline_html( + self, path, device_str, figsize=(20, 12), title=None + ) -> None: + """Exports the memory timeline as an HTML file which contains + the memory timeline plot embedded as a PNG file.""" + # Check if user has matplotlib installed, return gracefully if not. + import importlib.util + + matplotlib_spec = importlib.util.find_spec("matplotlib") + if matplotlib_spec is None: + print( + "export_memory_timeline_html failed because matplotlib was not found." + ) + return + + from base64 import b64encode + from os import remove + from tempfile import NamedTemporaryFile + + import matplotlib.pyplot as plt + import numpy as np + + mt = self._coalesce_timeline(device_str) + times, sizes = np.array(mt[0]), np.array(mt[1]) + # For this timeline, start at 0 to match Chrome traces. + t_min = min(times) + times -= t_min + stacked = np.cumsum(sizes, axis=1) / 1024**3 + device = torch.device(device_str) + max_memory_allocated = torch.cuda.max_memory_allocated(device) + max_memory_reserved = torch.cuda.max_memory_reserved(device) + + # Plot memory timeline as stacked data + fig = plt.figure(figsize=figsize, dpi=80) + axes = fig.gca() + for category, color in _CATEGORY_TO_COLORS.items(): + i = _CATEGORY_TO_INDEX[category] + axes.fill_between( + times / 1e3, stacked[:, i], stacked[:, i + 1], color=color, alpha=0.7 + ) + fig.legend(["Unknown" if i is None else i.name for i in _CATEGORY_TO_COLORS]) + # Usually training steps are in magnitude of ms. + axes.set_xlabel("Time (ms)") + axes.set_ylabel("Memory (GB)") + title = "\n\n".join( + ([title] if title else []) + + [ + f"Max memory allocated: {max_memory_allocated / (1024**3):.2f} GiB \n" + f"Max memory reserved: {max_memory_reserved / (1024**3):.2f} GiB" + ] + ) + axes.set_title(title) + + # Embed the memory timeline image into the HTML file + tmpfile = NamedTemporaryFile("wb", suffix=".png", delete=False) + tmpfile.close() + fig.savefig(tmpfile.name, format="png") + + with open(tmpfile.name, "rb") as tmp: + encoded = b64encode(tmp.read()).decode("utf-8") + html = f""" +GPU Memory Timeline HTML + + + +""" + + with open(path, "w") as f: + f.write(html) + remove(tmpfile.name) diff --git a/venv/lib/python3.10/site-packages/torch/profiler/_pattern_matcher.py b/venv/lib/python3.10/site-packages/torch/profiler/_pattern_matcher.py new file mode 100644 index 0000000000000000000000000000000000000000..cee47f28eb04a0f83dabfe62fe3dec83eeaee462 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/_pattern_matcher.py @@ -0,0 +1,667 @@ +# mypy: allow-untyped-defs +import json +import math +import os +import re +from typing import Optional + +import torch +import torch.utils.benchmark as benchmark +from torch._C._profiler import ( + _EventType, + _ExtraFields_PyCall, + _ExtraFields_PyCCall, + _ExtraFields_TorchOp, + _ProfilerEvent, +) +from torch.profiler import profile +from torch.profiler._utils import index_of_first_match, traverse_bfs, traverse_dfs + + +class Pattern: + """ + Base class for all patterns, subclass this class and implement match() + to define custom patterns. + + In subclass, define description and skip property. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + self.prof = prof + self.should_benchmark = should_benchmark + self.name = "Please specify a name for pattern" + self.description = "Please specify a description for pattern" + self.url = "" + assert prof.profiler is not None and prof.profiler.kineto_results is not None + self.event_tree = prof.profiler.kineto_results.experimental_event_tree() + self.tid_root: dict[int, list[_ProfilerEvent]] = {} + for event in self.event_tree: + self.tid_root.setdefault(event.start_tid, []).append(event) + + @property + def skip(self) -> bool: + return False + + def report(self, event: _ProfilerEvent): + msg = ( + f"{self.description}\n[Source Code Location] {source_code_location(event)}" + ) + return msg + + def eventTreeTraversal(self): + """ + Traverse the event tree and yield all events. + Override this method in subclass to customize the traversal. + """ + yield from traverse_dfs(self.event_tree) + + def summary(self, events: list[_ProfilerEvent]): + default_summary = f"{self.name}: {len(events)} events matched." + if self.should_benchmark: + # If benchmark summary is not empty, use it. + return ( + self.benchmark_summary(events) + if hasattr(self, "benchmark") # type: ignore[attr-defined] + else default_summary + ) + return default_summary + + def benchmark_summary(self, events: list[_ProfilerEvent]) -> str: + def format_time(time_ns: int) -> str: + unit_lst = ["ns", "us", "ms"] + for unit in unit_lst: + if time_ns < 1000: + return f"{time_ns:.2f} {unit}" + time_ns //= 1000 + return f"{time_ns:.2f} s" + + assert hasattr(self, "benchmark"), "Please implement benchmark()" + shapes_factor_map = self.benchmark(events) # type: ignore[attr-defined] + original_time = sum(event.duration_time_ns for event in events) + new_time = sum( + shapes_factor_map[input_shapes(event)] * event.duration_time_ns + for event in events + ) + return ( + f"{self.name}: {len(events)} events matched. " + f"Total Estimated Speedup: {format_time(original_time - new_time)} ({round(original_time / new_time, 2)}X)" + ) + + def match(self, event: _ProfilerEvent): + """ + Return True if the event matches the pattern. + This method should be overridden in subclass. + """ + raise NotImplementedError + + def matched_events(self): + if self.skip: + return [] + matched_events = [ + event for event in self.eventTreeTraversal() if self.match(event) + ] + return matched_events + + def root_of(self, event: _ProfilerEvent): + while event.parent: + event = event.parent + return event + + def siblings_of(self, event: _ProfilerEvent): + if event.parent: + children = event.parent.children + else: + children = self.tid_root[event.start_tid] + index = children.index(event) + return children[:index], children[index + 1 :] + + def next_of(self, event: _ProfilerEvent): + _, next_events = self.siblings_of(event) + return next_events[0] if next_events else None + + def prev_of(self, event: _ProfilerEvent): + prev_events, _ = self.siblings_of(event) + return prev_events[-1] if prev_events else None + + def go_up_until(self, event: _ProfilerEvent, predicate): + if not event: + return None + while event.parent and not predicate(event): + event = event.parent + return event + + +# Patterns + + +class NamePattern(Pattern): + def __init__( + self, prof: profile, name: str, should_benchmark: bool = False + ) -> None: + super().__init__(prof, should_benchmark) + self.description = f"Matched Name Event: {name}" + self.name = name + + def match(self, event: _ProfilerEvent): + return re.search(self.name, event.name) is not None + + +class ExtraCUDACopyPattern(Pattern): + """ + This pattern identifies if we creates a constant tensor on CPU and immediately moves it to GPU. + example: torch.zeros((100, 100)).to("cuda") + + Pattern: + built-in method |built-in method + ... | aten::to + aten::fill_/aten::zero_ | aten::_to_copy + + Algorithm: + We start at node aten::to, go parent events' previous events, + and check if we have a aten::fill_/aten::zero_ as we keep going down the tree. + We always select the last child in the children list when we go down the tree. + If at any step we failed, it is not a match. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Extra CUDA Copy Pattern" + self.description = "Filled a CPU tensor and immediately moved it to GPU. Please initialize it on GPU." + self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#create-tensors-directly-on-the-target-device" + self.init_ops = { + "aten::fill_", + "aten::zero_", + "aten::normal_", + "aten::uniform_", + } + + @property + def skip(self) -> bool: + return not self.prof.with_stack or not self.prof.record_shapes + + def match(self, event): + # TODO: We should also check tensor identities + if event.name != "aten::to": + return False + to_event = event + if not event.children: + return False + event = event.children[-1] + if event.name != "aten::_to_copy": + return False + if not event.children: + return False + event = event.children[-1] + if event.name != "aten::copy_": + return False + # aten::copy_ should have the first 2 args dtype the same + dtypes = input_dtypes(event) + if len(dtypes) < 2: + return False + if dtypes[0] is None or dtypes[0] != dtypes[1]: + return False + event = to_event + # Up one level + event = event.parent + if event is None: + return False + # Check if we have a aten::fill_ in previous leaf + event = self.prev_of(event) + if event is None: + return False + while event.children: + event = event.children[-1] + # aten::zero_ is a special optimization case where fill_ is not called + if event.name in self.init_ops: + return True + return event.name in self.init_ops + # TODO: Check if tensor is reused + + def benchmark(self, events: list[_ProfilerEvent]): + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + size = shape[0] + to_timer = benchmark.Timer( + stmt='torch.ones(size).to("cuda")', globals={"size": size} + ) + de_timer = benchmark.Timer( + stmt='torch.ones(size, device="cuda")', globals={"size": size} + ) + to_time = to_timer.timeit(10).mean + de_time = de_timer.timeit(10).mean + shapes_factor_map[shape] = de_time / to_time + return shapes_factor_map + + +class ForLoopIndexingPattern(Pattern): + """ + This pattern identifies if we use a for loop to index a tensor that + can be vectorized. + example: + tensor = torch.empty((100, 100)) + for i in range(100): + tensor[i] = i + + Pattern: + aten::select | ... | aten::select | ... (Repeat) + + Algorithm: + We start at node aten::select, and we check if we can find this alternating patterns. + We also keep a dictionary to avoid duplicate match in the for loop. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "For Loop Indexing Pattern" + self.description = "For loop indexing detected. Vectorization recommended." + self.visited: set[int] = set() + + def eventTreeTraversal(self): + """ + We need to use BFS traversal order to avoid duplicate match. + """ + yield from traverse_bfs(self.event_tree) + + def match(self, event: _ProfilerEvent): + if event.name != "aten::select": + return False + if event.id in self.visited: + return False + repeat_count = 1 + _, next = self.siblings_of(event) + if len(next) <= 1: + return False + + # Custom event list matching + def same_ops(list1, list2) -> bool: + if len(list1) != len(list2): + return False + for op1, op2 in zip(list1, list2): + if op1.name != op2.name: + return False + return True + + # Record the ops between two aten::select + next_select_idx = index_of_first_match(next, lambda e: e.name == "aten::select") + if next_select_idx is None: + return False + indexing_ops = [event] + next[:next_select_idx] + next = next[len(indexing_ops) - 1 :] + for i in range(0, len(next), len(indexing_ops)): + if same_ops(indexing_ops, next[i : i + len(indexing_ops)]): + repeat_count += 1 + self.visited.add(next[i].id) + else: + break + return repeat_count >= 10 + + +class FP32MatMulPattern(Pattern): + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "FP32 MatMul Pattern" + self.description = ( + "You are currently using GPU that supports TF32. " + "Please enable TF32 by setting 'torch.backends.cuda.matmul.allow_tf32 = True'" + ) + self.url = "https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + + @property + def skip(self): + if torch.version.hip is not None: + has_tf32 = False + else: + # Anything less than sm_80 is not Ampere which doesn't support TF32 + has_tf32 = all( + int(re.sub("sm_|compute_", "", arch)) >= 80 + for arch in torch.cuda.get_arch_list() + ) + return has_tf32 is False or super().skip or not self.prof.record_shapes + + def match(self, event: _ProfilerEvent) -> bool: + # If we saw this pattern once, we don't need to match it again + if event.tag != _EventType.TorchOp: + return False + assert isinstance(event.extra_fields, _ExtraFields_TorchOp) + if event.name == "aten::mm": + if event.extra_fields.allow_tf32_cublas is False: + return True + return False + + def report(self, event: _ProfilerEvent): + return self.description + + def benchmark(self, events: list[_ProfilerEvent]): + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float32) + matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float32) + fp32_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + tf32_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + setup="torch.backends.cuda.matmul.allow_tf32 = True", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + torch.backends.cuda.matmul.allow_tf32 = False + fp32_time = fp32_timer.timeit(10).mean + tf32_time = tf32_timer.timeit(10).mean + shapes_factor_map[shape] = tf32_time / fp32_time + return shapes_factor_map + + +class OptimizerSingleTensorPattern(Pattern): + """ + This pattern identifies if we are using the single-tensor version of an optimizer. + example: + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + By adding foreach=True to enable multi-tensor optimizer, we can gain speedup when + the kernels are relatively small. + + Pattern: + XXXXX: _single_tenser_ + + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Optimizer Single Tensor Pattern" + self.optimizers_with_foreach = ["adam", "sgd", "adamw"] + self.description = ( + "Detected optimizer running with single tensor implementation. " + "Please enable multi tensor implementation by passing 'foreach=True' into optimizer." + ) + self.url = "" + + def match(self, event: _ProfilerEvent) -> bool: + for optimizer in self.optimizers_with_foreach: + if event.name.endswith(f"_single_tensor_{optimizer}"): + return True + return False + + +class SynchronizedDataLoaderPattern(Pattern): + """ + This pattern identifies if we are using num_workers=0 in DataLoader. + example: + torch.utils.data.DataLoader(dataset, batch_size=batch_size) + Add num_workers=N to the arguments. N depends on system configuration. + + Pattern: + dataloader.py(...): __iter__ + dataloader.py(...): _get_iterator + NOT dataloader.py(...): check_worker_number_rationality + + Algorithm: + If we don't see check_worker_number_rationality call in the dataloader __iter__, + It is not an asynchronous dataloader. + + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Synchronized DataLoader Pattern" + self.description = ( + "Detected DataLoader running with synchronized implementation. " + "Please enable asynchronous dataloading by setting num_workers > 0 when initializing DataLoader." + ) + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#enable-async-data-loading-and-augmentation" + ) + + def match(self, event: _ProfilerEvent) -> bool: + def is_dataloader_function(name: str, function_name: str): + return name.startswith( + os.path.join("torch", "utils", "data", "dataloader.py") + ) and name.endswith(function_name) + + # TODO: fixme! Due to lifetime issues of the function name, this field might + # actually point to an already freed string when the even is a PyCall. + # Just silently skip this to unblock testing. + try: + event.name + except UnicodeDecodeError: + return False + + if not is_dataloader_function(event.name, "__iter__"): + return False + if not event.children: + return False + event = event.children[0] + if not is_dataloader_function(event.name, "_get_iterator"): + return False + if not event.children: + return False + event = event.children[0] + return not is_dataloader_function(event.name, "check_worker_number_rationality") + # TODO: We should also check if the loader is bottleneck. + + +class GradNotSetToNonePattern(Pattern): + """ + This pattern identifies if we are not setting grad to None in zero_grad. + example: + optimizer.zero_grad() + By setting set_to_none=True, we can gain speedup + + Pattern: + XXXXX: _zero_grad + NOT aten::zeros + aten::zero_ + + aten::zero_ is called on each parameter in the model. + We also want to make sure it is not called by aten::zeros. + + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Gradient Set To Zero Instead of None Pattern" + self.description = ( + "Detected gradient set to zero instead of None. " + "Please add 'set_to_none=True' when calling zero_grad()." + ) + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#disable-gradient-calculation-for-validation-or-inference" + ) + + def match(self, event: _ProfilerEvent) -> bool: + if not event.name.endswith(": zero_grad"): + return False + if not event.children: + return False + + for sub_event in traverse_dfs(event.children): + if ( + sub_event.name == "aten::zero_" + and sub_event.parent.name != "aten::zeros" + ): + return True + # TODO: We should also check if the optimizer's numerical behavior will change. + return False + + +class Conv2dBiasFollowedByBatchNorm2dPattern(Pattern): + """ + This pattern identifies if we are enabling bias in Conv2d which is followed by BatchNorm2d. + Bias doesn't do anything when followed by batchnorm. + Pattern: + nn.Module: Conv2d | nn.Module: BatchNorm2d + ... + aten::conv2d AND dtype of third argument is not null + The third argument is the bias + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Enabling Bias in Conv2d Followed By BatchNorm Pattern" + self.description = "Detected bias enabled in Conv2d that is followed by BatchNorm2d. Please set 'bias=False' in Conv2d." + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#disable-bias-for-convolutions-directly-followed-by-a-batch-norm" + ) + + @property + def skip(self): + return self.prof.record_shapes is False or super().skip + + def match(self, event: _ProfilerEvent): + if event.name != "aten::conv2d": + return False + if len(input_dtypes(event)) < 3 or input_dtypes(event)[2] is None: + return False + # This means bias=True + event = self.go_up_until( + event, lambda e: e.name.startswith("nn.Module: Conv2d") + ) + if not event: + return False + event = self.next_of(event) + if not event: + return False + return event.name.startswith("nn.Module: BatchNorm2d") + + +class MatMulDimInFP16Pattern(Pattern): + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Matrix Multiplication Dimension Not Aligned Pattern" + self.description = "Detected matmul with dimension not aligned. Please use matmul with aligned dimension." + self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#use-mixed-precision-and-amp" + + @property + def skip(self) -> bool: + return not self.prof.with_stack or not self.prof.record_shapes + + def match(self, event: _ProfilerEvent) -> bool: + def mutiple_of(shapes, multiple): + return all(dim % multiple == 0 for shape in shapes for dim in shape[-2:]) + + if event.name not in ("aten::mm", "aten::bmm", "aten::addmm"): + return False + if not input_dtypes(event): + return False + arg_dtype = input_dtypes(event)[0] + if arg_dtype in (torch.bfloat16, torch.half) and not mutiple_of( + input_shapes(event), 8 + ): + return True + return False + + def benchmark(self, events: list[_ProfilerEvent]): + def closest_multiple(shapes, multiple): + return [multiple * math.ceil(shape / multiple) for shape in shapes] + + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float16) + matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float16) + not_aligned_dim_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + matrixA = torch.randn( + closest_multiple(shape[0], 8), device="cuda", dtype=torch.float16 + ) + matrixB = torch.randn( + closest_multiple(shape[1], 8), device="cuda", dtype=torch.float16 + ) + aligned_dim_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + not_aligned_dim_time = not_aligned_dim_timer.timeit(10).mean + aligned_dim_time = aligned_dim_timer.timeit(10).mean + shapes_factor_map[shape] = aligned_dim_time / not_aligned_dim_time + return shapes_factor_map + + +def source_code_location(event: Optional[_ProfilerEvent]) -> str: + while event: + if event.tag == _EventType.PyCall or event.tag == _EventType.PyCCall: + assert isinstance( + event.extra_fields, (_ExtraFields_PyCall, _ExtraFields_PyCCall) + ) + if not event.extra_fields.caller.file_name.startswith("torch" + os.sep): + return f"{event.extra_fields.caller.file_name}:{event.extra_fields.caller.line_number}" + event = event.parent + return "No source code location found" + + +def input_shapes(event: _ProfilerEvent): + assert isinstance(event.extra_fields, _ExtraFields_TorchOp) + return tuple(tuple(getattr(i, "sizes", ())) for i in event.extra_fields.inputs) + + +def input_dtypes(event: _ProfilerEvent): + assert isinstance(event.extra_fields, _ExtraFields_TorchOp) + return tuple(getattr(i, "dtype", None) for i in event.extra_fields.inputs) + + +def report_all_anti_patterns( + prof, + should_benchmark: bool = False, + print_enable: bool = True, + json_report_dir: Optional[str] = None, +) -> None: + report_dict: dict = {} + anti_patterns = [ + ExtraCUDACopyPattern(prof, should_benchmark), + # ForLoopIndexingPattern(prof, should_benchmark), + FP32MatMulPattern(prof, should_benchmark), + OptimizerSingleTensorPattern(prof, should_benchmark), + SynchronizedDataLoaderPattern(prof, should_benchmark), + GradNotSetToNonePattern(prof, should_benchmark), + Conv2dBiasFollowedByBatchNorm2dPattern(prof, should_benchmark), + MatMulDimInFP16Pattern(prof, should_benchmark), + ] + reported = set() + summaries = [] + message_list = [f"{'-' * 40}TorchTidy Report{'-' * 40}"] + message_list.append("Matched Events:") + + for anti_pattern in anti_patterns: + matched_events = anti_pattern.matched_events() + if not matched_events: + continue + summaries.append(anti_pattern.summary(matched_events)) + for event in matched_events: + report_msg = anti_pattern.report(event) + if report_msg not in reported: + message_list.append(report_msg) + reported.add(report_msg) + src_location, line_no = source_code_location(event).split(":") + report_dict.setdefault(src_location, []).append( + { + "line_number": int(line_no), + "name": anti_pattern.name, + "url": anti_pattern.url, + "message": anti_pattern.description, + } + ) + + if json_report_dir is not None: + json_report_path = os.path.join(json_report_dir, "torchtidy_report.json") + if os.path.exists(json_report_path): + with open(json_report_path) as f: + exisiting_report = json.load(f) + exisiting_report.update(report_dict) + report_dict = exisiting_report + with open(json_report_path, "w") as f: + json.dump(report_dict, f, indent=4) + + message_list.append("Summary:") + message_list += summaries + message_list.append(f"{'-' * 40}TorchTidy Report{'-' * 40}") + if print_enable: + print("\n".join(message_list)) diff --git a/venv/lib/python3.10/site-packages/torch/profiler/_utils.py b/venv/lib/python3.10/site-packages/torch/profiler/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5b631ef743c6e77d96d70f8df739d7ef8f689763 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/_utils.py @@ -0,0 +1,401 @@ +# mypy: allow-untyped-defs +import functools +import operator +import re +from collections import deque +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torch.autograd.profiler import profile +from torch.profiler import DeviceType + + +if TYPE_CHECKING: + from torch.autograd import _KinetoEvent + + +def _traverse(tree, next_fn, children_fn=lambda x: x.children, reverse: bool = False): + order = reversed if reverse else lambda x: x + remaining = deque(order(tree)) + while remaining: + curr_event = next_fn(remaining) + yield curr_event + for child_event in order(children_fn(curr_event)): + remaining.append(child_event) + + +traverse_dfs = functools.partial(_traverse, next_fn=lambda x: x.pop(), reverse=True) +traverse_bfs = functools.partial( + _traverse, next_fn=lambda x: x.popleft(), reverse=False +) + + +@dataclass +class EventMetrics: + duration_time_ns: int = 0 + self_time_ns: int = 0 + idle_time_ns: int = 0 + queue_depth: int = 0 + + @property + def fraction_idle_time(self): + if self.duration_time_ns == 0: + return 0.0 + return self.idle_time_ns / self.duration_time_ns + + +@dataclass +class Interval: + start: int + end: int + queue_depth: int = 0 + + +class EventKey: + def __init__(self, event) -> None: + self.event = event + + def __hash__(self): + return hash(self.event.id) + + def __eq__(self, other): + return self.event.id == other.event.id + + def __repr__(self) -> str: + return f"{self.event.name}" + + def intervals_overlap(self, intervals: list[Interval]): + overlap_time = 0 + intervals = sorted(intervals, key=lambda x: x.start) + + if intervals: + overlap_start = max(self.event.start_time_ns, intervals[0].start) + overlap_end = min(self.event.end_time_ns, intervals[0].end) + + if overlap_start < overlap_end: + overlap_time += overlap_end - overlap_start + + i, j = 0, 1 + while j < len(intervals): + prev_interval = intervals[i] + curr_interval = intervals[j] + j += 1 + if prev_interval.end > curr_interval.start: + # Completely subsumed by previous interval + if prev_interval.end > curr_interval.end: + j += 1 + continue + else: + curr_interval.start = prev_interval.end + i = j + + overlap_start = max(self.event.start_time_ns, curr_interval.start) + overlap_end = min(self.event.end_time_ns, curr_interval.end) + if overlap_start < overlap_end: + overlap_time += overlap_end - overlap_start + + return overlap_time + + +class BasicEvaluation: + def __init__(self, prof: profile) -> None: + self.profile = prof + self.metrics: dict[EventKey, EventMetrics] = {} + self.compute_self_time() + self.event_keys = sorted( + (e for e in self.metrics.keys()), key=lambda x: x.event.start_time_ns + ) + self.events = [e.event for e in self.event_keys] + self.cuda_events: list[_KinetoEvent] = [] + self.queue_depth_list = self.compute_queue_depth() + self.compute_idle_time() + + def compute_self_time(self) -> None: + """ + Computes event's self time(total time - time in child ops). + """ + assert self.profile.kineto_results is not None + stack = deque(self.profile.kineto_results.experimental_event_tree()) + + # standard iterating dfs + while stack: + curr_event = stack.pop() + self_time = curr_event.duration_time_ns + for child_event in curr_event.children: + self_time -= child_event.duration_time_ns + stack.append(child_event) + assert EventKey(curr_event) not in self.metrics, ( + f"Duplicate id: {curr_event.id}, {curr_event.name}" + ) + self.metrics[EventKey(curr_event)] = EventMetrics(self_time_ns=self_time) + self.metrics[ + EventKey(curr_event) + ].duration_time_ns = curr_event.duration_time_ns + + def compute_queue_depth(self): + """ + Computes queue_depth at each event. This will calculate the queue depth data for + All the events in the tree. + This will return a list of Interval of queue depth data of cuda launch and kernels. + """ + assert self.profile.kineto_results is not None + cuda_event_list = self.profile.kineto_results.events() + + def is_cuda_launch_kernel(e): + """Check if the event is a CUDA launch kernel.""" + launch_patterns = { + "cudaLaunchKernel", # Standard CUDA + "cudaLaunchKernelExC", # Extended C + "__cudaLaunchKernel", # Internal + "cudaLaunchCooperativeKernel", # Collaborative (single-device) + "cudaLaunchCooperativeKernelMultiDevice", # Collaborative (multi-devices) + } + name = str(getattr(e, "name", e)) + return any(name.startswith(pattern) for pattern in launch_patterns) + + def is_cuda_kernel(e): + """Check if the event is a CUDA runtime kernel.""" + # Check if the kernel is CUDA + if e.device_type() != DeviceType.CUDA: + return False + + name = str(getattr(e, "name", e)).lower() + + # Exclude memory operations + exclude_patterns = {"mem", "cpy", "alloc", "free"} + + return not any(pattern in name for pattern in exclude_patterns) + + cuda_launch_events = sorted( + (e for e in cuda_event_list if is_cuda_launch_kernel(e)), + key=lambda x: x.start_ns(), + ) + cuda_kernel_events = sorted( + (e for e in cuda_event_list if is_cuda_kernel(e)), + key=lambda x: x.start_ns(), + ) + + self.cuda_events = sorted( + cuda_launch_events + cuda_kernel_events, key=lambda x: x.start_ns() + ) + + kernel_mapping: dict[_KinetoEvent, int] = {} + last_mapped_kernel = 0 + for cuda_launch_event in cuda_launch_events: + index = index_of_first_match( + cuda_kernel_events, + lambda x: x.linked_correlation_id() + == cuda_launch_event.linked_correlation_id(), + start=last_mapped_kernel, + ) + kernel_mapping[cuda_launch_event] = index + last_mapped_kernel = index if index is not None else last_mapped_kernel + + current_kernel_index = 0 + spawned_kernel_index = -1 + + all_events = cuda_launch_events + cuda_kernel_events + self.events + + def new_old_event_comparator(event): + if hasattr(event, "start_us"): + return event.start_us() * 1000 + if hasattr(event, "start_ns"): + return event.start_ns() + if hasattr(event, "start_time_ns"): + return event.start_time_ns + raise Exception("Unknown Event Type") # noqa: TRY002 + + queue_depth_list: list[Interval] = [] + all_events.sort(key=new_old_event_comparator) + for event in all_events: + # Find latest cuda kernel event + if hasattr(event, "start_us"): + start_time = event.start_us() * 1000 + end_time = (event.start_us() + event.duration_us()) * 1000 + # Find current spawned cuda kernel event + if event in kernel_mapping and kernel_mapping[event] is not None: + spawned_kernel_index = kernel_mapping[event] + if hasattr(event, "start_ns"): + start_time = event.start_ns() + end_time = event.start_ns() + event.duration_ns() + # Find current spawned cuda kernel event + if event in kernel_mapping and kernel_mapping[event] is not None: + spawned_kernel_index = kernel_mapping[event] + elif hasattr(event, "start_time_ns"): + start_time = event.start_time_ns # type: ignore[attr-defined] + end_time = event.end_time_ns # type: ignore[attr-defined] + + while ( + current_kernel_index < len(cuda_kernel_events) + and (cuda_kernel_events[current_kernel_index].start_ns()) <= start_time # type: ignore[possibly-undefined] + ): + current_kernel_index += 1 + current_queue_depth = spawned_kernel_index - current_kernel_index + 1 + current_queue_depth = max(current_queue_depth, 0) + + if hasattr(event, "start_us") or hasattr(event, "start_ns"): + queue_depth_list.append( + Interval(start_time, end_time, current_queue_depth) # type: ignore[possibly-undefined] + ) + elif hasattr(event, "start_time_ns"): + self.metrics[EventKey(event)].queue_depth = current_queue_depth + + return queue_depth_list + + def compute_idle_time(self) -> None: + """ + Computes idle time of the profile. + """ + # Based on queue_depth_list, we can calculate idle time for all the events + idle = False + idle_start = 0 + idle_intervals: list[Interval] = [] + if self.queue_depth_list and self.events: + idle_intervals += [ + Interval(self.events[0].start_time_ns, self.queue_depth_list[0].start), + Interval(self.queue_depth_list[-1].end, self.events[-1].end_time_ns), + ] + + for data_point in self.queue_depth_list: + if data_point.queue_depth == 0 and not idle: + idle_start = data_point.end + idle = True + if data_point.queue_depth > 0 and idle: + idle_intervals.append(Interval(idle_start, data_point.start)) + idle = False + + event_list = [e.event for e in self.metrics.keys()] + for event in event_list: + self.metrics[EventKey(event)].idle_time_ns = EventKey( + event + ).intervals_overlap(idle_intervals) + + def rank_events(self, length): + """ + Filter and Rank the events based on some heuristics: + 1) Events that are in the falling phase of the queue depth. + 2) Events that have a high idle_time, self_time difference. + + Parameters: + length: The number of events to return. + """ + + # Find the interval when qd is falling to 0 + import torch + + queue_depth_list = list(reversed(self.queue_depth_list)) + qd_values = [e.queue_depth for e in queue_depth_list] + + bottom_threashold = 0 + top_threashold = 4 + decrease_interval = [] + i = 0 + while i < len(qd_values): + if qd_values[i] > bottom_threashold: + i += 1 + continue + for j in range(i + 1, len(qd_values)): + # Find next zero and if the max value between them exceeds + # the threshold, then we have a falling interval + next_minimum_idx = index_of_first_match( + qd_values, lambda x: x <= bottom_threashold, start=j + ) + peak_idx = argmax(qd_values, start=j, end=next_minimum_idx) + + # if is a valid peak, we add to list and continue + if peak_idx is not None and qd_values[peak_idx] >= top_threashold: + decrease_interval.append( + Interval( + queue_depth_list[peak_idx].start, queue_depth_list[i].start + ) + ) + i = next_minimum_idx if next_minimum_idx is not None else i + break + i += 1 + # Filter out events that are not in the decrease interval + event_list = [ + event + for event in self.metrics.keys() + if event.intervals_overlap(decrease_interval) + ] + if event_list: + self_time = torch.tensor( + [self.metrics[event].self_time_ns for event in event_list], + dtype=torch.float32, + ) + idle_time = torch.tensor( + [self.metrics[event].fraction_idle_time for event in event_list], + dtype=torch.float32, + ) + normalized_gain = (idle_time - torch.mean(idle_time)) / torch.std(idle_time) + normalized_self = (self_time - torch.mean(self_time)) / torch.std(self_time) + heuristic_score_list = normalized_gain + 0.6 * normalized_self + + # Sort events by heuristic + event_list = [ + event + for _, event in sorted( + zip(heuristic_score_list, event_list), + key=operator.itemgetter(0), + reverse=True, + ) + ] + event_list = event_list[:length] + return event_list + + def get_optimizable_events(self, length: int = 1, print_enable: bool = True): + event_list = self.rank_events(length) + if not print_enable: + return event_list + output = "Optimizable events:\n" if event_list else "No events to optimize\n" + + output += "\n".join( + [ + f"""{"-" * 80} +Event: {event} +Source code location: {source_code_location(event.event)} +Percentage idle time: {self.metrics[event].fraction_idle_time * 100:.2f}% +{"-" * 80}""" + for event in event_list + ] + ) + if print_enable: + print(output) + return event_list + + +def index_of_first_match(seq, predicate, start=0, end=None): + if end is None or end >= len(seq): + end = len(seq) + for i in range(start, end): + if predicate(seq[i]): + return i + return None + + +def argmax(seq, key=lambda x: x, start=0, end=None): + seq = seq[start:end] + if len(seq) == 0: + return None + return seq.index(max(seq, key=key)) + start + + +def source_code_location(event): + while event is not None: + match = re.search(r"\.py\(.*\)", event.name) + if match is None: + event = event.parent + continue + return event.name + return "No source code location found" + + +# Provide an OSS workaround for cudagraphs + CUPTI issue +# https://github.com/pytorch/pytorch/issues/75504 +# TODO(dberard) - deprecate / remove workaround for CUDA >= 12, when +# we stop supporting older CUDA versions. +def _init_for_cuda_graphs() -> None: + from torch.autograd.profiler import profile + + with profile(): + pass diff --git a/venv/lib/python3.10/site-packages/torch/profiler/itt.py b/venv/lib/python3.10/site-packages/torch/profiler/itt.py new file mode 100644 index 0000000000000000000000000000000000000000..7b1a6eac0f0bc8d69988fe59f7acae26302572c6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/itt.py @@ -0,0 +1,81 @@ +# mypy: allow-untyped-defs +from contextlib import contextmanager +from typing import NoReturn + + +try: + from torch._C import _itt +except ImportError: + + class _ITTStub: + @staticmethod + def _fail(*args, **kwargs) -> NoReturn: + raise RuntimeError( + "ITT functions not installed. Are you sure you have a ITT build?" + ) + + @staticmethod + def is_available() -> bool: + return False + + rangePush = _fail + rangePop = _fail + mark = _fail + + _itt = _ITTStub() # type: ignore[assignment] + + +__all__ = ["is_available", "range_push", "range_pop", "mark", "range"] + + +def is_available(): + """ + Check if ITT feature is available or not + """ + return _itt.is_available() + + +def range_push(msg): + """ + Pushes a range onto a stack of nested range span. Returns zero-based + depth of the range that is started. + + Arguments: + msg (str): ASCII message to associate with range + """ + return _itt.rangePush(msg) + + +def range_pop(): + """ + Pops a range off of a stack of nested range spans. Returns the + zero-based depth of the range that is ended. + """ + return _itt.rangePop() + + +def mark(msg): + """ + Describe an instantaneous event that occurred at some point. + + Arguments: + msg (str): ASCII message to associate with the event. + """ + return _itt.mark(msg) + + +@contextmanager +def range(msg, *args, **kwargs): + """ + Context manager / decorator that pushes an ITT range at the beginning + of its scope, and pops it at the end. If extra arguments are given, + they are passed as arguments to msg.format(). + + Args: + msg (str): message to associate with the range + """ + range_push(msg.format(*args, **kwargs)) + try: + yield + finally: + range_pop() diff --git a/venv/lib/python3.10/site-packages/torch/profiler/profiler.py b/venv/lib/python3.10/site-packages/torch/profiler/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..573541799bbe65a96aba890b97c4406d06ccfcad --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/profiler.py @@ -0,0 +1,1128 @@ +# mypy: allow-untyped-defs +import gzip +import json +import os +import shutil +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Iterable +from enum import Enum +from functools import partial +from typing import Any, Callable, Optional +from typing_extensions import Self +from warnings import warn + +import torch +import torch.autograd.profiler as prof +from torch._C import _get_privateuse1_backend_name +from torch._C._profiler import ( + _add_execution_trace_observer, + _disable_execution_trace_observer, + _enable_execution_trace_observer, + _ExperimentalConfig, + _remove_execution_trace_observer, +) +from torch._environment import is_fbcode +from torch._utils_internal import profiler_allow_cudagraph_cupti_lazy_reinit_cuda12 +from torch.autograd import kineto_available, ProfilerActivity +from torch.profiler._memory_profiler import MemoryProfile, MemoryProfileTimeline + + +__all__ = [ + "supported_activities", + "ProfilerAction", + "schedule", + "tensorboard_trace_handler", + "profile", + "ExecutionTraceObserver", +] +PROFILER_STEP_NAME = "ProfilerStep" + + +class _NumpyEncoder(json.JSONEncoder): + """ + Json encoder for numpy types (np.int, np.float, np.array etc.) + Returns default encoder if numpy is not available + """ + + def default(self, obj): + """Encode NumPy types to JSON""" + try: + import numpy as np + except ImportError: + return json.JSONEncoder.default(self, obj) + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.ndarray): + return obj.tolist() + else: + return json.JSONEncoder.default(self, obj) + + +def supported_activities(): + """ + Returns a set of supported profiler tracing activities. + + Note: profiler uses CUPTI library to trace on-device CUDA kernels. + In case when CUDA is enabled but CUPTI is not available, passing + ``ProfilerActivity.CUDA`` to profiler results in using the legacy CUDA + profiling code (same as in the legacy ``torch.autograd.profiler``). + This, in turn, results in including CUDA time in the profiler table output, + but not in the JSON trace. + """ + return torch.autograd._supported_activities() + + +class _ITraceObserver(ABC): + """Abstract interface for a Trace observer. + This satisfies 3 methods: start, stop and cleanup""" + + @abstractmethod + def start(self): + pass + + @abstractmethod + def stop(self): + pass + + @abstractmethod + def cleanup(self): + pass + + +class _KinetoProfile: + """Low-level profiler wrap the autograd profile + + Args: + activities (iterable): list of activity groups (CPU, CUDA) to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA``, + ``torch.profiler.ProfilerActivity.XPU``. + Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA + or (when available) ProfilerActivity.XPU. + record_shapes (bool): save information about operator's input shapes. + profile_memory (bool): track tensor memory allocation/deallocation (see ``export_memory_timeline`` + for more details). + with_stack (bool): record source information (file and line number) for the ops. + with_flops (bool): use formula to estimate the FLOPS of specific operators + (matrix multiplication and 2D convolution). + with_modules (bool): record module hierarchy (including function names) + corresponding to the callstack of the op. e.g. If module A's forward call's + module B's forward which contains an aten::add op, + then aten::add's module hierarchy is A.B + Note that this support exist, at the moment, only for TorchScript models + and not eager mode models. + experimental_config (_ExperimentalConfig) : A set of experimental options + used by profiler libraries like Kineto. Note, backward compatibility is not guaranteed. + execution_trace_observer (ExecutionTraceObserver) : A PyTorch Execution Trace Observer object. + `PyTorch Execution Traces `__ offer a graph based + representation of AI/ML workloads and enable replay benchmarks, simulators, and emulators. + When this argument is included the observer start() and stop() will be called for the + same time window as PyTorch profiler. + acc_events (bool): Enable the accumulation of FunctionEvents across multiple profiling cycles + + + .. note:: + This API is experimental and subject to change in the future. + + Enabling shape and stack tracing results in additional overhead. + When record_shapes=True is specified, profiler will temporarily hold references to the tensors; + that may further prevent certain optimizations that depend on the reference count and introduce + extra tensor copies. + """ + + def __init__( + self, + *, + activities: Optional[Iterable[ProfilerActivity]] = None, + record_shapes: bool = False, + profile_memory: bool = False, + with_stack: bool = False, + with_flops: bool = False, + with_modules: bool = False, + experimental_config: Optional[_ExperimentalConfig] = None, + execution_trace_observer: Optional[_ITraceObserver] = None, + acc_events: bool = False, + custom_trace_id_callback: Optional[Callable[[], str]] = None, + ) -> None: + self.activities = set(activities) if activities else supported_activities() + self.record_shapes = record_shapes + self.with_flops = with_flops + self.profile_memory = profile_memory + self.with_stack = with_stack + self.with_modules = with_modules + self.experimental_config = experimental_config + self.execution_trace_observer = execution_trace_observer + self.acc_events = acc_events + self.custom_trace_id_callback = custom_trace_id_callback + self.profiler: Optional[prof.profile] = None + self.has_cudagraphs = False + self.mem_tl: Optional[MemoryProfileTimeline] = None + self.use_device = None + if ProfilerActivity.CUDA in self.activities: + self.use_device = "cuda" + elif ProfilerActivity.XPU in self.activities: + self.use_device = "xpu" + elif ProfilerActivity.MTIA in self.activities: + self.use_device = "mtia" + elif ProfilerActivity.HPU in self.activities: + self.use_device = "hpu" + elif ProfilerActivity.PrivateUse1 in self.activities: + self.use_device = _get_privateuse1_backend_name() + + # user-defined metadata to be amended to the trace + self.preset_metadata: dict[str, str] = {} + + def start(self) -> None: + self.prepare_trace() + self.start_trace() + + def stop(self) -> None: + self.stop_trace() + + def prepare_trace(self) -> None: + if hasattr(torch, "_inductor"): + import torch._inductor.config as inductor_config + + self.has_cudagraphs = inductor_config.triton.cudagraphs + if (self.profiler is None) or (not self.acc_events): + self.profiler = prof.profile( + use_cpu=(ProfilerActivity.CPU in self.activities), + use_device=self.use_device, + record_shapes=self.record_shapes, + with_flops=self.with_flops, + profile_memory=self.profile_memory, + with_stack=self.with_stack, + with_modules=self.with_modules, + use_kineto=True, + experimental_config=self.experimental_config, + acc_events=self.acc_events, + custom_trace_id_callback=self.custom_trace_id_callback, + ) + self.profiler._prepare_trace() + + def start_trace(self) -> None: + if self.execution_trace_observer: + self.execution_trace_observer.start() + assert self.profiler is not None + self.profiler._start_trace() + + if self.profile_memory: + self.add_metadata_json("profile_memory", "1") + if self.with_stack: + self.add_metadata_json("with_stack", "1") + if self.record_shapes: + self.add_metadata_json("record_shapes", "1") + if self.with_modules: + self.add_metadata_json("with_modules", "1") + if self.with_flops: + self.add_metadata_json("with_flops", "1") + + if kineto_available(): + dist_info = self._get_distributed_info() + if dist_info: + self.add_metadata_json( + "distributedInfo", json.dumps(dist_info, cls=_NumpyEncoder) + ) + + cuda_version = None + if hasattr(torch, "version"): + from torch.torch_version import TorchVersion + + cuda_version = TorchVersion(getattr(torch.version, "cuda", "0.0")) + + if self.has_cudagraphs and ( + (cuda_version and cuda_version < "12.6") + or not profiler_allow_cudagraph_cupti_lazy_reinit_cuda12() + ): + os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1" + self.add_metadata_json("DISABLE_CUPTI_LAZY_REINIT", "1") + # FIXME: CUDA Graph does not work well with CUPTI teardown. + # 1) crashes on 1st lazy CUPTI re-init after teardown (CUDA 11) + # 2) crashes on 2nd non-lazy CUPTI re-init after teardown (CUDA 12) + # Workaround: turn off CUPTI teardown when using CUDA Graphs. + os.environ["TEARDOWN_CUPTI"] = "0" + + # Insert the preset user metadata to the trace + for k, v in self.preset_metadata.items(): + self.add_metadata_json(k, v) + + def stop_trace(self) -> None: + if self.execution_trace_observer: + self.execution_trace_observer.stop() + assert self.profiler is not None + self.profiler.__exit__(None, None, None) + + def export_chrome_trace(self, path: str): + """ + Exports the collected trace in Chrome JSON format. If kineto is enabled, only + last cycle in schedule is exported. + """ + assert self.profiler + if path.endswith(".gz"): + fp = tempfile.NamedTemporaryFile("w+b", suffix=".json", delete=False) + fp.close() + retvalue = self.profiler.export_chrome_trace(fp.name) + with open(fp.name, "rb") as fin: + with gzip.open(path, "wb") as fout: + fout.writelines(fin) + os.remove(fp.name) + return retvalue + else: + return self.profiler.export_chrome_trace(path) + + def export_stacks(self, path: str, metric: str = "self_cpu_time_total"): + """Save stack traces to a file + + Args: + path (str): save stacks file to this location; + metric (str): metric to use: "self_cpu_time_total" or "self_cuda_time_total" + """ + assert self.profiler + return self.profiler.export_stacks(path, metric) + + def toggle_collection_dynamic( + self, enable: bool, activities: Iterable[ProfilerActivity] + ) -> None: + """Toggle collection of activities on/off at any point of collection. Currently supports toggling Torch Ops + (CPU) and CUDA activity supported in Kineto + + Args: + activities (iterable): list of activity groups to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA`` + Examples: + + .. code-block:: python + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as p: + code_to_profile_0() + // turn off collection of all CUDA activity + p.toggle_collection_dynamic(False, [torch.profiler.ProfilerActivity.CUDA]) + code_to_profile_1() + // turn on collection of all CUDA activity + p.toggle_collection_dynamic(True, [torch.profiler.ProfilerActivity.CUDA]) + code_to_profile_2() + print(p.key_averages().table( + sort_by="self_cuda_time_total", row_limit=-1)) + """ + if not self.profiler: + return + self.profiler.toggle_collection_dynamic(enable, activities) + + def key_averages( + self, + group_by_input_shape: bool = False, + group_by_stack_n: int = 0, + group_by_overload_name: bool = False, + ): + """Averages events, grouping them by operator name and (optionally) input shapes, stack + and overload name. + + .. note:: + To use shape/stack functionality make sure to set record_shapes/with_stack + when creating profiler context manager. + """ + assert self.profiler + return self.profiler.key_averages( + group_by_input_shape, group_by_stack_n, group_by_overload_name + ) + + def events(self): + """ + Returns the list of unaggregated profiler events, + to be used in the trace callback or after the profiling is finished + """ + assert self.profiler + return self.profiler.function_events + + def add_metadata(self, key: str, value: str) -> None: + """ + Adds a user defined metadata with a string key and a string value + into the trace file + """ + wrapped_value = '"' + value.replace('"', '\\"') + '"' + torch.autograd._add_metadata_json(key, wrapped_value) + + def add_metadata_json(self, key: str, value: str) -> None: + """ + Adds a user defined metadata with a string key and a valid json value + into the trace file + """ + torch.autograd._add_metadata_json(key, value) + + def preset_metadata_json(self, key: str, value: str) -> None: + """ + Preset a user defined metadata when the profiler is not started + and added into the trace file later. + Metadata is in the format of a string key and a valid json value + """ + self.preset_metadata[key] = value + + def _get_distributed_info(self): + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + return None + + backend = dist.get_backend() + dist_info = { + "backend": backend, + "rank": dist.get_rank(), + "world_size": dist.get_world_size(), + "pg_count": dist.get_pg_count(), + "pg_config": dist.distributed_c10d._get_all_pg_configs(), + } + if backend == "nccl": + nccl_version = torch.cuda.nccl.version() + dist_info["nccl_version"] = ".".join(str(v) for v in nccl_version) + return dist_info + + def _memory_profile(self) -> MemoryProfile: + required = ("record_shapes", "profile_memory", "with_stack") + missing = [f"{i}=True" for i in required if not getattr(self, i)] + if missing: + raise ValueError(f"{', '.join(missing)} required for memory profiling.") + + assert self.profiler is not None and self.profiler.kineto_results is not None + return MemoryProfile(self.profiler.kineto_results) + + def export_memory_timeline(self, path: str, device: Optional[str] = None) -> None: + """Export memory event information from the profiler collected + tree for a given device, and export a timeline plot. There are 3 + exportable files using ``export_memory_timeline``, each controlled by the + ``path``'s suffix. + + - For an HTML compatible plot, use the suffix ``.html``, and a memory timeline + plot will be embedded as a PNG file in the HTML file. + + - For plot points consisting of ``[times, [sizes by category]]``, where + ``times`` are timestamps and ``sizes`` are memory usage for each category. + The memory timeline plot will be saved a JSON (``.json``) or gzipped JSON + (``.json.gz``) depending on the suffix. + + - For raw memory points, use the suffix ``.raw.json.gz``. Each raw memory + event will consist of ``(timestamp, action, numbytes, category)``, where + ``action`` is one of ``[PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY]``, + and ``category`` is one of the enums from + ``torch.profiler._memory_profiler.Category``. + + Output: Memory timeline written as gzipped JSON, JSON, or HTML. + """ + # Default to device 0, if unset. Fallback on cpu. + if device is None: + if self.use_device and self.use_device != "cuda": + device = self.use_device + ":0" + else: + device = "cuda:0" if torch.cuda.is_available() else "cpu" + + # Construct the memory timeline plot data + self.mem_tl = MemoryProfileTimeline(self._memory_profile()) + + # Depending on the file suffix, save the data as json.gz or json. + # For html, we can embed the image into an HTML file. + if path.endswith(".html"): + self.mem_tl.export_memory_timeline_html(path, device) + elif path.endswith(".gz"): + fp = tempfile.NamedTemporaryFile("w+t", suffix=".json", delete=False) + fp.close() + if path.endswith("raw.json.gz"): + self.mem_tl.export_memory_timeline_raw(fp.name, device) + else: + self.mem_tl.export_memory_timeline(fp.name, device) + with open(fp.name) as fin: + with gzip.open(path, "wt") as fout: + fout.writelines(fin) + os.remove(fp.name) + else: + self.mem_tl.export_memory_timeline(path, device) + + +class ProfilerAction(Enum): + """ + Profiler actions that can be taken at the specified intervals + """ + + NONE = 0 + WARMUP = 1 + RECORD = 2 + RECORD_AND_SAVE = 3 + + +def schedule( + *, + wait: int, + warmup: int, + active: int, + repeat: int = 0, + skip_first: int = 0, + skip_first_wait: int = 0, +) -> Callable: + """ + Returns a callable that can be used as profiler ``schedule`` argument. The profiler will skip + the first ``skip_first`` steps, then wait for ``wait`` steps, then do the warmup for the next ``warmup`` steps, + then do the active recording for the next ``active`` steps and then repeat the cycle starting with ``wait`` steps. + The optional number of cycles is specified with the ``repeat`` parameter, the zero value means that + the cycles will continue until the profiling is finished. + + The ``skip_first_wait`` parameter controls whether the first ``wait`` stage should be skipped. + This can be useful if a user wants to wait longer than ``skip_first`` between cycles, but not + for the first profile. For example, if ``skip_first`` is 10 and ``wait`` is 20, the first cycle will + wait 10 + 20 = 30 steps before warmup if ``skip_first_wait`` is zero, but will wait only 10 + steps if ``skip_first_wait`` is non-zero. All subsequent cycles will then wait 20 steps between the + last active and warmup. + """ + + def schedule_fn(step: int) -> ProfilerAction: + assert step >= 0 + if step < skip_first: + return ProfilerAction.NONE + else: + step -= skip_first + # If wait >> skip_first and we want to grab profiling early, shift left by wait if skip_first_wait is True + if skip_first_wait != 0: + step += wait + num_steps = wait + warmup + active + if repeat > 0 and step / num_steps >= repeat: + return ProfilerAction.NONE + mod_step = step % num_steps + if mod_step < wait: + return ProfilerAction.NONE + elif mod_step < wait + warmup: + return ProfilerAction.WARMUP + else: + return ( + ProfilerAction.RECORD + if mod_step < num_steps - 1 + else ProfilerAction.RECORD_AND_SAVE + ) + + assert ( + wait >= 0 and warmup >= 0 and active > 0 and repeat >= 0 and skip_first >= 0 + ), "Invalid profiler schedule arguments" + if warmup == 0: + warn("Profiler won't be using warmup, this can skew profiler results") + return schedule_fn + + +def _default_schedule_fn(_: int) -> ProfilerAction: + """ + Default profiler behavior - immediately starts recording the events, + keeps doing it on every profiler step. + """ + return ProfilerAction.RECORD + + +def tensorboard_trace_handler( + dir_name: str, worker_name: Optional[str] = None, use_gzip: bool = False +): + """ + Outputs tracing files to directory of ``dir_name``, then that directory can be + directly delivered to tensorboard as logdir. + ``worker_name`` should be unique for each worker in distributed scenario, + it will be set to '[hostname]_[pid]' by default. + """ + import socket + import time + + def handler_fn(prof) -> None: + nonlocal worker_name + if not os.path.isdir(dir_name): + try: + os.makedirs(dir_name, exist_ok=True) + except Exception as e: + raise RuntimeError("Can't create directory: " + dir_name) from e + if not worker_name: + worker_name = f"{socket.gethostname()}_{os.getpid()}" + # Use nanosecond here to avoid naming clash when exporting the trace + file_name = f"{worker_name}.{time.time_ns()}.pt.trace.json" + if use_gzip: + file_name = file_name + ".gz" + prof.export_chrome_trace(os.path.join(dir_name, file_name)) + + return handler_fn + + +class profile(_KinetoProfile): + """Profiler context manager. + + Args: + activities (iterable): list of activity groups (CPU, CUDA) to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA``, + ``torch.profiler.ProfilerActivity.XPU``. + Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA + or (when available) ProfilerActivity.XPU. + schedule (Callable): callable that takes step (int) as a single parameter and returns + ``ProfilerAction`` value that specifies the profiler action to perform at each step. + on_trace_ready (Callable): callable that is called at each step when ``schedule`` + returns ``ProfilerAction.RECORD_AND_SAVE`` during the profiling. + record_shapes (bool): save information about operator's input shapes. + profile_memory (bool): track tensor memory allocation/deallocation. + with_stack (bool): record source information (file and line number) for the ops. + with_flops (bool): use formula to estimate the FLOPs (floating point operations) of specific operators + (matrix multiplication and 2D convolution). + with_modules (bool): record module hierarchy (including function names) + corresponding to the callstack of the op. e.g. If module A's forward call's + module B's forward which contains an aten::add op, + then aten::add's module hierarchy is A.B + Note that this support exist, at the moment, only for TorchScript models + and not eager mode models. + experimental_config (_ExperimentalConfig) : A set of experimental options + used for Kineto library features. Note, backward compatibility is not guaranteed. + execution_trace_observer (ExecutionTraceObserver) : A PyTorch Execution Trace Observer object. + `PyTorch Execution Traces `__ offer a graph based + representation of AI/ML workloads and enable replay benchmarks, simulators, and emulators. + When this argument is included the observer start() and stop() will be called for the + same time window as PyTorch profiler. See the examples section below for a code sample. + acc_events (bool): Enable the accumulation of FunctionEvents across multiple profiling cycles + use_cuda (bool): + .. deprecated:: 1.8.1 + use ``activities`` instead. + + .. note:: + Use :func:`~torch.profiler.schedule` to generate the callable schedule. + Non-default schedules are useful when profiling long training jobs + and allow the user to obtain multiple traces at the different iterations + of the training process. + The default schedule simply records all the events continuously for the + duration of the context manager. + + .. note:: + Use :func:`~torch.profiler.tensorboard_trace_handler` to generate result files for TensorBoard: + + ``on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name)`` + + After profiling, result files can be found in the specified directory. Use the command: + + ``tensorboard --logdir dir_name`` + + to see the results in TensorBoard. + For more information, see + `PyTorch Profiler TensorBoard Plugin `__ + + .. note:: + Enabling shape and stack tracing results in additional overhead. + When record_shapes=True is specified, profiler will temporarily hold references to the tensors; + that may further prevent certain optimizations that depend on the reference count and introduce + extra tensor copies. + + + Examples: + + .. code-block:: python + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as p: + code_to_profile() + print(p.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1)) + + Using the profiler's ``schedule``, ``on_trace_ready`` and ``step`` functions: + + .. code-block:: python + + # Non-default profiler schedule allows user to turn profiler on and off + # on different iterations of the training loop; + # trace_handler is called every time a new trace becomes available + def trace_handler(prof): + print( + prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1) + ) + # prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json") + + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + # In this example with wait=1, warmup=1, active=2, repeat=1, + # profiler will skip the first step/iteration, + # start warming up on the second, record + # the third and the forth iterations, + # after which the trace will become available + # and on_trace_ready (when set) is called; + # the cycle repeats starting with the next step + schedule=torch.profiler.schedule(wait=1, warmup=1, active=2, repeat=1), + on_trace_ready=trace_handler, + # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log') + # used when outputting for tensorboard + ) as p: + for iter in range(N): + code_iteration_to_profile(iter) + # send a signal to the profiler that the next iteration has started + p.step() + + The following sample shows how to setup up an Execution Trace Observer (`execution_trace_observer`) + + .. code-block:: python + + with torch.profiler.profile( + ... + execution_trace_observer=( + ExecutionTraceObserver().register_callback("./execution_trace.json") + ), + ) as p: + for iter in range(N): + code_iteration_to_profile(iter) + p.step() + + You can also refer to test_execution_trace_with_kineto() in tests/profiler/test_profiler.py. + Note: One can also pass any object satisfying the _ITraceObserver interface. + """ + + def __init__( + self, + *, + activities: Optional[Iterable[ProfilerActivity]] = None, + schedule: Optional[Callable[[int], ProfilerAction]] = None, + on_trace_ready: Optional[Callable[..., Any]] = None, + record_shapes: bool = False, + profile_memory: bool = False, + with_stack: bool = False, + with_flops: bool = False, + with_modules: bool = False, + experimental_config: Optional[_ExperimentalConfig] = None, + execution_trace_observer: Optional[_ITraceObserver] = None, + acc_events: bool = False, + # deprecated: + use_cuda: Optional[bool] = None, + custom_trace_id_callback: Optional[Callable[[], str]] = None, + ) -> None: + activities_set = set(activities) if activities else supported_activities() + if use_cuda is not None: + warn( + "`use_cuda` is deprecated, use `activities` argument instead", + FutureWarning, + stacklevel=2, + ) + if use_cuda: + activities_set.add(ProfilerActivity.CUDA) + elif ProfilerActivity.CUDA in activities_set: + activities_set.remove(ProfilerActivity.CUDA) + assert len(activities_set) > 0, "No valid profiler activities found" + + super().__init__( + activities=activities, + record_shapes=record_shapes, + profile_memory=profile_memory, + with_stack=with_stack, + with_flops=with_flops, + with_modules=with_modules, + experimental_config=experimental_config, + execution_trace_observer=execution_trace_observer + if execution_trace_observer + else ExecutionTraceObserver.build_execution_trace_obs_from_env(), + acc_events=acc_events, + custom_trace_id_callback=custom_trace_id_callback, + ) + + if schedule: + self.schedule = schedule + # add step markers into the trace and table view + self.record_steps = True + else: + self.schedule = _default_schedule_fn + self.record_steps = False + self.on_trace_ready = on_trace_ready + self.step_num = 0 + self.current_action = self.schedule(self.step_num) + self.step_rec_fn: Optional[prof.record_function] = None + + self.action_map: dict[ + tuple[ProfilerAction, Optional[ProfilerAction]], list[Any] + ] = { + # key is (prev_action, current_action), value is action list corresponding to the state pair. + (ProfilerAction.NONE, ProfilerAction.NONE): [], + (ProfilerAction.NONE, ProfilerAction.WARMUP): [self.prepare_trace], + (ProfilerAction.NONE, ProfilerAction.RECORD): [ + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.NONE, ProfilerAction.RECORD_AND_SAVE): [ + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.WARMUP, ProfilerAction.NONE): [ + partial(warn, "Incorrect schedule: WARMUP followed by NONE"), + self.start_trace, + self.stop_trace, + ], + (ProfilerAction.WARMUP, ProfilerAction.WARMUP): [], + (ProfilerAction.WARMUP, ProfilerAction.RECORD): [self.start_trace], + (ProfilerAction.WARMUP, ProfilerAction.RECORD_AND_SAVE): [self.start_trace], + (ProfilerAction.RECORD, ProfilerAction.NONE): [ + partial(warn, "Incorrect schedule: RECORD followed by NONE"), + self.stop_trace, + ], + (ProfilerAction.RECORD, ProfilerAction.WARMUP): [ + partial(warn, "Incorrect schedule: RECORD followed by WARMUP"), + self.stop_trace, + ], + (ProfilerAction.RECORD, ProfilerAction.RECORD): [], + (ProfilerAction.RECORD, ProfilerAction.RECORD_AND_SAVE): [], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.NONE): [ + self.stop_trace, + self._trace_ready, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.WARMUP): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.RECORD): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.RECORD_AND_SAVE): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + self.start_trace, + ], + # used for exit action + (ProfilerAction.WARMUP, None): [self.start_trace, self.stop_trace], + (ProfilerAction.RECORD, None): [self.stop_trace, self._trace_ready], + (ProfilerAction.RECORD_AND_SAVE, None): [ + self.stop_trace, + self._trace_ready, + ], + } + # Start tracking increments to profiler step, this will be used + # by Kineto + prof.KinetoStepTracker.init_step_count(PROFILER_STEP_NAME) + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() + prof.KinetoStepTracker.erase_step_count(PROFILER_STEP_NAME) + if self.execution_trace_observer: + self.execution_trace_observer.cleanup() + + def start(self) -> None: + self._transit_action(ProfilerAction.NONE, self.current_action) + if self.record_steps: + self.step_rec_fn = prof.record_function( + "ProfilerStep#" + str(self.step_num) + ) + self.step_rec_fn.__enter__() + + def stop(self) -> None: + if self.record_steps and self.step_rec_fn: + self.step_rec_fn.__exit__(None, None, None) + self._transit_action(self.current_action, None) + + def step(self) -> None: + """ + Signals the profiler that the next profiling step has started. + """ + if self.record_steps and self.step_rec_fn: + self.step_rec_fn.__exit__(None, None, None) + prev_action = self.current_action + self.step_num += 1 + self.current_action = self.schedule(self.step_num) + + self._transit_action(prev_action, self.current_action) + if os.environ.get("KINETO_USE_DAEMON", "") or ( + is_fbcode() and os.environ.get("KINETO_FORCE_STEP_HOOK", "") + ): + prof.KinetoStepTracker.increment_step(PROFILER_STEP_NAME) + + if self.record_steps: + self.step_rec_fn = prof.record_function( + "ProfilerStep#" + str(self.step_num) + ) + self.step_rec_fn.__enter__() + + def set_custom_trace_id_callback(self, callback) -> None: + """ + Sets a callback to be called when a new trace ID is generated. + """ + self.custom_trace_id_callback = callback + + def get_trace_id(self): + """ + Returns the current trace ID. + """ + if self.profiler is None: + return None + return self.profiler.trace_id + + def _trace_ready(self) -> None: + if self.on_trace_ready: + self.on_trace_ready(self) + + def _transit_action(self, prev_action, current_action) -> None: + action_list = self.action_map.get((prev_action, current_action)) + if action_list: + for action in action_list: + action() + + def _stats(self) -> Optional[prof._ProfilerStats]: + if self.profiler is None: + return None + return self.profiler._stats + + +class ExecutionTraceObserver(_ITraceObserver): + """Execution Trace Observer + + Each process can have a single ExecutionTraceObserver instance. The observer + can be added to record function callbacks via calling register_callback() + explicitly. Without calling unregister_callback(), repeated calls to + register_callback() will not add additional observers to record function + callbacks. Once an ExecutionTraceObserver is created, the start() and stop() + methods control when the event data is recorded. + + Deleting or calling unregister_callback() will remove the observer from the + record function callbacks, finalize the output file, and will stop + incurring any overheads. + """ + + def __init__(self) -> None: + """ + Initializes the default states. + """ + self._registered = False + self._execution_trace_running = False + self.extra_resources_collection = False + self.resources_dir: str = "" + self.output_file_path: str = "" + self.output_file_path_observer: str = "" + + def __del__(self) -> None: + """ + Calls unregister_callback() to make sure to finalize outputs. + """ + self.unregister_callback() + + @staticmethod + def build_execution_trace_obs_from_env() -> Optional["ExecutionTraceObserver"]: + """ + Returns an ExecutionTraceObserver instance if the environment variable + ENABLE_PYTORCH_EXECUTION_TRACE is set to 1, otherwise returns None. + + Configures the observer to also collect extra resources if the environment variable + ``ENABLE_PYTORCH_EXECUTION_TRACE_EXTRAS=1``. These are resources such as generated kernels, + index tensor data etc. that are required to make the Execution Trace replayable. + """ + if os.environ.get("ENABLE_PYTORCH_EXECUTION_TRACE", "0") == "1": + try: + fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False) + except Exception as e: + warn( + f"Execution trace will not be recorded. Exception on creating default temporary file: {e}" + ) + return None + fp.close() + et = ExecutionTraceObserver() + et.register_callback(fp.name) + # additionally, check if the env requires us to collect extra resources + if os.environ.get("ENABLE_PYTORCH_EXECUTION_TRACE_EXTRAS", "0") == "1": + et.set_extra_resource_collection(True) + else: + et.set_extra_resource_collection(False) + return et + return None + + def set_extra_resource_collection(self, val) -> None: + """ + Collects extra resources such as generated kernels, index tensor data, and any other + metadata that is required to complete the Execution Trace content. + + The caller should call this method with val=True after calling register_callback() if they want + to collect the extra resources. + """ + self.extra_resources_collection = val + if self.extra_resources_collection: + self.get_resources_dir(can_create=True) + return + + def register_callback(self, output_file_path: str) -> Self: + """ + Adds ET observer to record function callbacks. The data will be + written to output_file_path. + """ + + def get_temp_uncompressed_file() -> str: + fp = tempfile.NamedTemporaryFile("w+b", suffix=".json", delete=False) + fp.close() + return fp.name + + if not self._registered: + self.output_file_path = output_file_path + if output_file_path.endswith(".gz"): + output_file_path = get_temp_uncompressed_file() + self.output_file_path_observer = output_file_path + self._registered = _add_execution_trace_observer(output_file_path) + return self + + def get_resources_dir(self, can_create=False) -> Optional[str]: + """ + Generates the resources directory for the generated kernels, + or index tensor data or any other metadata that is required + to complete the Execution Trace content. + + The directory is created right where the ET file is being output. + + Only works if the observer has called set_extra_resource_collection(val=True). + + Returns None if the observer is not configured with extra resource collection. + """ + if not self.extra_resources_collection: + return None + if self.resources_dir: + # already created + return self.resources_dir + generated_path = ExecutionTraceObserver.get_resources_dir_for_et_path( + self.output_file_path, create_dir=can_create + ) + if not generated_path: + # could not find of create the resources dir + return None + self.resources_dir = generated_path + return self.resources_dir + + @staticmethod + def get_resources_dir_for_et_path( + trace_path, create_dir: bool = False + ) -> Optional[str]: + work_dir, file_name = os.path.split(trace_path) + resource_dir = os.path.join( + work_dir, os.path.splitext(file_name)[0] + "_resources" + ) + if not os.path.exists(resource_dir): + if create_dir: + try: + os.mkdir(resource_dir) + except Exception: + warn(f"Execution trace exception when creating {resource_dir}") + return None + else: + return None + return resource_dir + + def unregister_callback(self) -> None: + """ + Removes ET observer from record function callbacks. + """ + + def _save_triton_kernels() -> None: + try: + resource_dir = self.get_resources_dir() + except Exception as e: + warn( + f"Execution trace exception when generating resource directory: {e}" + ) + return + if not resource_dir: + return + + # Save the kernel paths for the generated kernels + from torch._inductor.codecache import PyCodeCache as PyCodeCache + + kernel_files = [ + v.__file__ + for v in PyCodeCache.modules + if getattr(v, "__file__", None) is not None + ] + + for kernel_file in kernel_files: + if kernel_file is None: + continue + name = os.path.basename(kernel_file) + dst = os.path.join(resource_dir, name) + shutil.copyfile(kernel_file, dst) + + def _save_gz_file(uncompressed_file: str, output_file: str) -> None: + print(f"Execution Trace: compressing {uncompressed_file} to {output_file}") + with open(uncompressed_file, "rb") as fin: + with gzip.open(output_file, "wb") as fout: + fout.writelines(fin) + os.remove(uncompressed_file) + + if self._registered: + self.stop() + + try: + _save_triton_kernels() + except Exception as e: + warn(f"Execution trace failed to save kernels: {e}") + + _remove_execution_trace_observer() + if self.output_file_path.endswith("gz"): + _save_gz_file(self.output_file_path_observer, self.output_file_path) + + self._registered = False + + @property + def is_registered(self): + """ + Returns True if the execution trace observer is registered, otherwise False. + """ + return self._registered + + def is_running(self): + """ + Returns True if the observer is running, otherwise False. + """ + return self._execution_trace_running + + def start(self) -> None: + """ + Starts to capture. + """ + if self._registered and not self._execution_trace_running: + _enable_execution_trace_observer() + self._execution_trace_running = True + self._record_pg_config() + + def stop(self) -> None: + """ + Stops to capture. + """ + if self._execution_trace_running: + _disable_execution_trace_observer() + self._execution_trace_running = False + + def cleanup(self) -> None: + """ + Calls unregister_callback() to make sure to finalize outputs. + """ + self.unregister_callback() + + def get_output_file_path(self) -> Optional[str]: + """ + Returns the output file name or None. + """ + if self.output_file_path: + return self.output_file_path + else: + return None + + def _record_pg_config(self) -> None: + # Records the PG config info to the trace as node: + # ## process_group:init ## + if ( + self.is_registered + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + pg_config_info = torch.distributed.distributed_c10d._world.pg_config_info + torch.autograd._record_function_with_args_enter( + "## process_group:init ##", + json.dumps(pg_config_info, cls=_NumpyEncoder), + ) diff --git a/venv/lib/python3.10/site-packages/torch/profiler/python_tracer.py b/venv/lib/python3.10/site-packages/torch/profiler/python_tracer.py new file mode 100644 index 0000000000000000000000000000000000000000..aff0fbc32ff3a8870ce81f569daa2587f598394f --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/profiler/python_tracer.py @@ -0,0 +1,19 @@ +import os +import site +import sys + +import torch + + +def _prefix_regex() -> list[str]: + raw_paths = ( + site.getsitepackages() + + sys.path + + [site.getuserbase()] + + [site.getusersitepackages()] + + [os.path.dirname(os.path.dirname(torch.__file__))] + ) + + path_prefixes = sorted({os.path.abspath(i) for i in raw_paths}, reverse=True) + assert all(isinstance(i, str) for i in path_prefixes) + return [i + os.sep for i in path_prefixes] diff --git a/venv/lib/python3.10/site-packages/torch/py.typed b/venv/lib/python3.10/site-packages/torch/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/torch/quasirandom.py b/venv/lib/python3.10/site-packages/torch/quasirandom.py new file mode 100644 index 0000000000000000000000000000000000000000..b5d4540e592f190e5e239391c40003a0bbc30146 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/quasirandom.py @@ -0,0 +1,217 @@ +# mypy: allow-untyped-defs +from typing import Optional + +import torch + + +class SobolEngine: + r""" + The :class:`torch.quasirandom.SobolEngine` is an engine for generating + (scrambled) Sobol sequences. Sobol sequences are an example of low + discrepancy quasi-random sequences. + + This implementation of an engine for Sobol sequences is capable of + sampling sequences up to a maximum dimension of 21201. It uses direction + numbers from https://web.maths.unsw.edu.au/~fkuo/sobol/ obtained using the + search criterion D(6) up to the dimension 21201. This is the recommended + choice by the authors. + + References: + - Art B. Owen. Scrambling Sobol and Niederreiter-Xing points. + Journal of Complexity, 14(4):466-489, December 1998. + + - I. M. Sobol. The distribution of points in a cube and the accurate + evaluation of integrals. + Zh. Vychisl. Mat. i Mat. Phys., 7:784-802, 1967. + + Args: + dimension (Int): The dimensionality of the sequence to be drawn + scramble (bool, optional): Setting this to ``True`` will produce + scrambled Sobol sequences. Scrambling is + capable of producing better Sobol + sequences. Default: ``False``. + seed (Int, optional): This is the seed for the scrambling. The seed + of the random number generator is set to this, + if specified. Otherwise, it uses a random seed. + Default: ``None`` + + Examples:: + + >>> # xdoctest: +SKIP("unseeded random state") + >>> soboleng = torch.quasirandom.SobolEngine(dimension=5) + >>> soboleng.draw(3) + tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [0.5000, 0.5000, 0.5000, 0.5000, 0.5000], + [0.7500, 0.2500, 0.2500, 0.2500, 0.7500]]) + """ + + MAXBIT = 30 + MAXDIM = 21201 + + def __init__(self, dimension, scramble=False, seed=None): + if dimension > self.MAXDIM or dimension < 1: + raise ValueError( + "Supported range of dimensionality " + f"for SobolEngine is [1, {self.MAXDIM}]" + ) + + self.seed = seed + self.scramble = scramble + self.dimension = dimension + + cpu = torch.device("cpu") + + self.sobolstate = torch.zeros( + dimension, self.MAXBIT, device=cpu, dtype=torch.long + ) + torch._sobol_engine_initialize_state_(self.sobolstate, self.dimension) + + if not self.scramble: + self.shift = torch.zeros(self.dimension, device=cpu, dtype=torch.long) + else: + self._scramble() + + self.quasi = self.shift.clone(memory_format=torch.contiguous_format) + self._first_point = (self.quasi / 2**self.MAXBIT).reshape(1, -1) + self.num_generated = 0 + + def draw( + self, + n: int = 1, + out: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + r""" + Function to draw a sequence of :attr:`n` points from a Sobol sequence. + Note that the samples are dependent on the previous samples. The size + of the result is :math:`(n, dimension)`. + + Args: + n (Int, optional): The length of sequence of points to draw. + Default: 1 + out (Tensor, optional): The output tensor + dtype (:class:`torch.dtype`, optional): the desired data type of the + returned tensor. + Default: ``None`` + """ + if dtype is None: + dtype = torch.get_default_dtype() + + if self.num_generated == 0: + if n == 1: + result = self._first_point.to(dtype) + else: + result, self.quasi = torch._sobol_engine_draw( + self.quasi, + n - 1, + self.sobolstate, + self.dimension, + self.num_generated, + dtype=dtype, + ) + result = torch.cat((self._first_point.to(dtype), result), dim=-2) + else: + result, self.quasi = torch._sobol_engine_draw( + self.quasi, + n, + self.sobolstate, + self.dimension, + self.num_generated - 1, + dtype=dtype, + ) + + self.num_generated += n + + if out is not None: + out.resize_as_(result).copy_(result) + return out + + return result + + def draw_base2( + self, + m: int, + out: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + r""" + Function to draw a sequence of :attr:`2**m` points from a Sobol sequence. + Note that the samples are dependent on the previous samples. The size + of the result is :math:`(2**m, dimension)`. + + Args: + m (Int): The (base2) exponent of the number of points to draw. + out (Tensor, optional): The output tensor + dtype (:class:`torch.dtype`, optional): the desired data type of the + returned tensor. + Default: ``None`` + """ + n = 2**m + total_n = self.num_generated + n + if not (total_n & (total_n - 1) == 0): + raise ValueError( + "The balance properties of Sobol' points require " + f"n to be a power of 2. {self.num_generated} points have been " + f"previously generated, then: n={self.num_generated}+2**{m}={total_n}. " + "If you still want to do this, please use " + "'SobolEngine.draw()' instead." + ) + return self.draw(n=n, out=out, dtype=dtype) + + def reset(self): + r""" + Function to reset the ``SobolEngine`` to base state. + """ + self.quasi.copy_(self.shift) + self.num_generated = 0 + return self + + def fast_forward(self, n): + r""" + Function to fast-forward the state of the ``SobolEngine`` by + :attr:`n` steps. This is equivalent to drawing :attr:`n` samples + without using the samples. + + Args: + n (Int): The number of steps to fast-forward by. + """ + if self.num_generated == 0: + torch._sobol_engine_ff_( + self.quasi, n - 1, self.sobolstate, self.dimension, self.num_generated + ) + else: + torch._sobol_engine_ff_( + self.quasi, n, self.sobolstate, self.dimension, self.num_generated - 1 + ) + self.num_generated += n + return self + + def _scramble(self): + g: Optional[torch.Generator] = None + if self.seed is not None: + g = torch.Generator() + g.manual_seed(self.seed) + + cpu = torch.device("cpu") + + # Generate shift vector + shift_ints = torch.randint( + 2, (self.dimension, self.MAXBIT), device=cpu, generator=g + ) + self.shift = torch.mv( + shift_ints, torch.pow(2, torch.arange(0, self.MAXBIT, device=cpu)) + ) + + # Generate lower triangular matrices (stacked across dimensions) + ltm_dims = (self.dimension, self.MAXBIT, self.MAXBIT) + ltm = torch.randint(2, ltm_dims, device=cpu, generator=g).tril() + + torch._sobol_engine_scramble_(self.sobolstate, ltm, self.dimension) + + def __repr__(self): + fmt_string = [f"dimension={self.dimension}"] + if self.scramble: + fmt_string += ["scramble=True"] + if self.seed is not None: + fmt_string += [f"seed={self.seed}"] + return self.__class__.__name__ + "(" + ", ".join(fmt_string) + ")" diff --git a/venv/lib/python3.10/site-packages/torch/random.py b/venv/lib/python3.10/site-packages/torch/random.py new file mode 100644 index 0000000000000000000000000000000000000000..18a1cf9a5d5de2f2385687d9fe62c847b098642d --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/random.py @@ -0,0 +1,203 @@ +# mypy: allow-untyped-defs +import contextlib +import warnings +from collections.abc import Generator + +import torch +from torch._C import default_generator + + +def set_rng_state(new_state: torch.Tensor) -> None: + r"""Sets the random number generator state. + + .. note:: This function only works for CPU. For CUDA, please use + :func:`torch.manual_seed`, which works for both CPU and CUDA. + + Args: + new_state (torch.ByteTensor): The desired state + """ + default_generator.set_state(new_state) + + +def get_rng_state() -> torch.Tensor: + r"""Returns the random number generator state as a `torch.ByteTensor`. + + .. note:: The returned state is for the default generator on CPU only. + + See also: :func:`torch.random.fork_rng`. + """ + return default_generator.get_state() + + +def manual_seed(seed) -> torch._C.Generator: + r"""Sets the seed for generating random numbers on all devices. Returns a + `torch.Generator` object. + + Args: + seed (int): The desired seed. Value must be within the inclusive range + `[-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]`. Otherwise, a RuntimeError + is raised. Negative inputs are remapped to positive values with the formula + `0xffff_ffff_ffff_ffff + seed`. + """ + seed = int(seed) + import torch.cuda + + if not torch.cuda._is_in_bad_fork(): + torch.cuda.manual_seed_all(seed) + + import torch.mps + + if not torch.mps._is_in_bad_fork(): + torch.mps.manual_seed(seed) + + import torch.xpu + + if not torch.xpu._is_in_bad_fork(): + torch.xpu.manual_seed_all(seed) + + _seed_custom_device(seed) + + return default_generator.manual_seed(seed) + + +def seed() -> int: + r"""Sets the seed for generating random numbers to a non-deterministic + random number on all devices. Returns a 64 bit number used to seed the RNG. + """ + seed = default_generator.seed() + import torch.cuda + + if not torch.cuda._is_in_bad_fork(): + torch.cuda.manual_seed_all(seed) + + import torch.mps + + if not torch.mps._is_in_bad_fork(): + torch.mps.manual_seed(seed) + + import torch.xpu + + if not torch.xpu._is_in_bad_fork(): + torch.xpu.manual_seed_all(seed) + + _seed_custom_device(seed) + + return seed + + +def _seed_custom_device(seed) -> None: + r"""Sets the seed to generate random numbers for custom device. + + Args: + seed (int): The desired seed. + + See [Note: support the custom device with privateuse1] + """ + seed = int(seed) + custom_backend_name = torch._C._get_privateuse1_backend_name() + if hasattr(torch, custom_backend_name): + custom_device_mod = getattr(torch, custom_backend_name) + _bad_fork_name = "_is_in_bad_fork" + _seed_all_name = "manual_seed_all" + if hasattr(custom_device_mod, _bad_fork_name) and hasattr( + custom_device_mod, _seed_all_name + ): + if not getattr(custom_device_mod, _bad_fork_name)(): + getattr(custom_device_mod, _seed_all_name)(seed) + else: + message = f"Set seed for `{custom_backend_name}` device does not take effect, please add API's " + message += f"`{_bad_fork_name}` and `{_seed_all_name}` to `{custom_backend_name}` device module." + warnings.warn(message, UserWarning, stacklevel=3) + + +def initial_seed() -> int: + r"""Returns the initial seed for generating random numbers as a + Python `long`. + + .. note:: The returned seed is for the default generator on CPU only. + """ + return default_generator.initial_seed() + + +_fork_rng_warned_already = False + + +@contextlib.contextmanager +def fork_rng( + devices=None, + enabled=True, + _caller="fork_rng", + _devices_kw="devices", + device_type="cuda", +) -> Generator: + """ + Forks the RNG, so that when you return, the RNG is reset + to the state that it was previously in. + + Args: + devices (iterable of Device IDs): devices for which to fork + the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates + on all devices, but will emit a warning if your machine has a lot + of devices, since this function will run very slowly in that case. + If you explicitly specify devices, this warning will be suppressed + enabled (bool): if ``False``, the RNG is not forked. This is a convenience + argument for easily disabling the context manager without having + to delete it and unindent your Python code under it. + device_type (str): device type str, default is `cuda`. As for supported device, + see details in :ref:`accelerator` + """ + + if device_type == "meta": + yield + return + + device_type = torch.device(device_type).type + device_mod = getattr(torch, device_type, None) + if device_mod is None: + raise RuntimeError( + f"torch has no module of `{device_type}`, you should register " + + "a module by `torch._register_device_module`." + ) + global _fork_rng_warned_already + + # Internal arguments: + # _caller: the function which called fork_rng, which the user used + # _devices_kw: the devices keyword of _caller + + if not enabled: + yield + return + + if devices is None: + num_devices = device_mod.device_count() + if num_devices > 1 and not _fork_rng_warned_already: + message = ( + f"{device_type.upper()} reports that you have {num_devices} available devices, and " + f"you have used {_caller} without explicitly specifying which devices are being used. " + f"For safety, we initialize *every* {device_type.upper()} device by default, which can " + f"be quite slow if you have a lot of {device_type.upper()}s. If you know that you are only" + f" making use of a few {device_type.upper()} devices, set the environment variable " + f"{device_type.upper()}_VISIBLE_DEVICES or the '{_devices_kw}' keyword argument of {_caller} " + "with the set of devices you are actually using. For example, if you are using CPU only, " + "set device.upper()_VISIBLE_DEVICES= or devices=[]; if you are using device 0 only, " + f"set {device_type.upper()}_VISIBLE_DEVICES=0 or devices=[0]. To initialize all devices " + f"and suppress this warning, set the '{_devices_kw}' keyword argument to " + f"`range(torch.{device_type}.device_count())`." + ) + warnings.warn(message) + _fork_rng_warned_already = True + devices = list(range(num_devices)) + else: + # Protect against user passing us a generator; we need to traverse this + # multiple times but a generator will be exhausted upon first traversal + devices = list(devices) + + cpu_rng_state = torch.get_rng_state() + device_rng_states = [device_mod.get_rng_state(device) for device in devices] + + try: + yield + finally: + torch.set_rng_state(cpu_rng_state) + for device, device_rng_state in zip(devices, device_rng_states): + device_mod.set_rng_state(device_rng_state, device) diff --git a/venv/lib/python3.10/site-packages/torch/return_types.py b/venv/lib/python3.10/site-packages/torch/return_types.py new file mode 100644 index 0000000000000000000000000000000000000000..d456742be4b88ebdca9f3696a415014a500cdd33 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/return_types.py @@ -0,0 +1,51 @@ +import inspect + +import torch +from torch.utils._pytree import register_pytree_node, SequenceKey + + +__all__ = ["pytree_register_structseq", "all_return_types"] + +all_return_types = [] + +# error: Module has no attribute "_return_types" +return_types = torch._C._return_types # type: ignore[attr-defined] + + +def pytree_register_structseq(cls): + def structseq_flatten(structseq): + return list(structseq), None + + def structseq_flatten_with_keys(structseq): + values, context = structseq_flatten(structseq) + return [(SequenceKey(i), v) for i, v in enumerate(values)], context + + def structseq_unflatten(values, context): + return cls(values) + + register_pytree_node( + cls, + structseq_flatten, + structseq_unflatten, + flatten_with_keys_fn=structseq_flatten_with_keys, + ) + + +for name in dir(return_types): + if name.startswith("__"): + continue + + _attr = getattr(return_types, name) + globals()[name] = _attr + + if not name.startswith("_"): + __all__.append(name) + all_return_types.append(_attr) + + # Today everything in torch.return_types is a structseq, aka a "namedtuple"-like + # thing defined by the Python C-API. We're going to need to modify this when that + # is no longer the case. + # NB: I don't know how to check that something is a "structseq" so we do a fuzzy + # check for tuple + if inspect.isclass(_attr) and issubclass(_attr, tuple): + pytree_register_structseq(_attr) diff --git a/venv/lib/python3.10/site-packages/torch/return_types.pyi b/venv/lib/python3.10/site-packages/torch/return_types.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e8dce0869e225ea2a50200240469d220ed296e34 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/return_types.pyi @@ -0,0 +1,605 @@ +# @generated by tools/pyi/gen_pyi.py from torch/_C/return_types.pyi.in +# mypy: allow-untyped-defs + +from typing import Final, NoReturn +from typing_extensions import Self + +from torch import SymInt, Tensor +from torch.types import ( # noqa: F401 + _bool, + _device, + _dtype, + _float, + _int, + _layout, + _qscheme, + _size, + Number, +) + +__all__ = [ + "pytree_register_structseq", + "all_return_types", + "_fake_quantize_per_tensor_affine_cachemask_tensor_qparams", + "_fused_moving_avg_obs_fq_helper", + "_linalg_det", + "_linalg_eigh", + "_linalg_slogdet", + "_linalg_solve_ex", + "_linalg_svd", + "_lu_with_info", + "_scaled_dot_product_cudnn_attention", + "_scaled_dot_product_efficient_attention", + "_scaled_dot_product_flash_attention", + "_scaled_dot_product_flash_attention_for_cpu", + "_unpack_dual", + "aminmax", + "cummax", + "cummin", + "frexp", + "geqrf", + "histogram", + "histogramdd", + "kthvalue", + "lu_unpack", + "max", + "median", + "min", + "mode", + "nanmedian", + "qr", + "slogdet", + "sort", + "svd", + "topk", + "triangular_solve", +] + +def pytree_register_structseq(cls: type) -> None: ... + +class _fake_quantize_per_tensor_affine_cachemask_tensor_qparams(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _fused_moving_avg_obs_fq_helper(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_det(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def result(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_eigh(tuple[Tensor, Tensor]): # fmt: skip + @property + def eigenvalues(self) -> Tensor: ... + @property + def eigenvectors(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_slogdet(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def sign(self) -> Tensor: ... + @property + def logabsdet(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_solve_ex(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def result(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + @property + def info(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_svd(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def U(self) -> Tensor: ... + @property + def S(self) -> Tensor: ... + @property + def Vh(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _lu_with_info(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + @property + def info(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_cudnn_attention(tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + @property + def cum_seq_q(self) -> Tensor: ... + @property + def cum_seq_k(self) -> Tensor: ... + @property + def max_q(self) -> _int | SymInt: ... + @property + def max_k(self) -> _int | SymInt: ... + @property + def philox_seed(self) -> Tensor: ... + @property + def philox_offset(self) -> Tensor: ... + @property + def debug_attn_mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 9 + n_sequence_fields: Final[_int] = 9 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_efficient_attention(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def log_sumexp(self) -> Tensor: ... + @property + def philox_seed(self) -> Tensor: ... + @property + def philox_offset(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_flash_attention(tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + @property + def cum_seq_q(self) -> Tensor: ... + @property + def cum_seq_k(self) -> Tensor: ... + @property + def max_q(self) -> _int | SymInt: ... + @property + def max_k(self) -> _int | SymInt: ... + @property + def rng_state(self) -> Tensor: ... + @property + def unused(self) -> Tensor: ... + @property + def debug_attn_mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 9 + n_sequence_fields: Final[_int] = 9 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_flash_attention_for_cpu(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _unpack_dual(tuple[Tensor, Tensor]): # fmt: skip + @property + def primal(self) -> Tensor: ... + @property + def tangent(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class aminmax(tuple[Tensor, Tensor]): # fmt: skip + @property + def min(self) -> Tensor: ... + @property + def max(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class cummax(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class cummin(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class frexp(tuple[Tensor, Tensor]): # fmt: skip + @property + def mantissa(self) -> Tensor: ... + @property + def exponent(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class geqrf(tuple[Tensor, Tensor]): # fmt: skip + @property + def a(self) -> Tensor: ... + @property + def tau(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class histogram(tuple[Tensor, Tensor]): # fmt: skip + @property + def hist(self) -> Tensor: ... + @property + def bin_edges(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class histogramdd(tuple[Tensor, tuple[Tensor, ...]]): # fmt: skip + @property + def hist(self) -> Tensor: ... + @property + def bin_edges(self) -> tuple[Tensor, ...]: ... + def __new__( + cls, + sequence: tuple[Tensor, tuple[Tensor, ...]], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class kthvalue(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class lu_unpack(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def P(self) -> Tensor: ... + @property + def L(self) -> Tensor: ... + @property + def U(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class max(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class median(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class min(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class mode(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class nanmedian(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class qr(tuple[Tensor, Tensor]): # fmt: skip + @property + def Q(self) -> Tensor: ... + @property + def R(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class slogdet(tuple[Tensor, Tensor]): # fmt: skip + @property + def sign(self) -> Tensor: ... + @property + def logabsdet(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class sort(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class svd(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def U(self) -> Tensor: ... + @property + def S(self) -> Tensor: ... + @property + def V(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class topk(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class triangular_solve(tuple[Tensor, Tensor]): # fmt: skip + @property + def solution(self) -> Tensor: ... + @property + def cloned_coefficient(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +all_return_types: list[type] = ... diff --git a/venv/lib/python3.10/site-packages/torch/serialization.py b/venv/lib/python3.10/site-packages/torch/serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..a6eb314fc1a82e1a829148e34151347e32c32e01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/serialization.py @@ -0,0 +1,2133 @@ +# mypy: allow-untyped-defs +import copyreg +import difflib +import functools +import io +import os +import pickle +import re +import shutil +import struct +import sys +import tarfile +import tempfile +import threading +import warnings +from contextlib import closing, contextmanager +from enum import Enum +from typing import Any, Callable, cast, Generic, IO, Optional, TypeVar, Union +from typing_extensions import TypeAlias, TypeIs + +import torch +import torch._weights_only_unpickler as _weights_only_unpickler +from torch._sources import get_source_lines_and_file +from torch._utils import _import_dotted_name +from torch.storage import _get_dtype_from_pickle_storage_type +from torch.types import FileLike, Storage + + +__all__ = [ + "SourceChangeWarning", + "mkdtemp", + "register_package", + "check_module_version_greater_or_equal", + "validate_cuda_device", + "validate_hpu_device", + "location_tag", + "default_restore_location", + "normalize_storage_type", + "storage_to_tensor_type", + "save", + "load", + "StorageType", + "LoadEndianness", + "get_crc32_options", + "set_crc32_options", + "get_default_load_endianness", + "set_default_load_endianness", + "get_default_mmap_options", + "set_default_mmap_options", + "clear_safe_globals", + "get_safe_globals", + "add_safe_globals", + "safe_globals", + "get_unsafe_globals_in_checkpoint", + "skip_data", +] + +DEFAULT_PROTOCOL = 2 + +LONG_SIZE = struct.Struct("=l").size +INT_SIZE = struct.Struct("=i").size +SHORT_SIZE = struct.Struct("=h").size + +MAGIC_NUMBER = 0x1950A86A20F9469CFC6C +PROTOCOL_VERSION = 1001 +STORAGE_KEY_SEPARATOR = "," + +MAP_LOCATION: TypeAlias = Optional[ + Union[Callable[[Storage, str], Storage], torch.device, str, dict[str, str]] +] +STORAGE: TypeAlias = Union[Storage, torch.storage.TypedStorage, torch.UntypedStorage] + +IS_WINDOWS = sys.platform == "win32" + +UNSAFE_MESSAGE = ( + "In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` " + "from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, " + "but it can result in arbitrary code execution. Do it only if you got the file from a " + "trusted source." +) + +if not IS_WINDOWS: + from mmap import MAP_PRIVATE, MAP_SHARED +else: + MAP_SHARED, MAP_PRIVATE = None, None # type: ignore[assignment] + + +def _default_to_weights_only(pickle_module): + is_fbcode = not hasattr(torch.version, "git_version") + return pickle_module is None and not is_fbcode + + +# _serialization_tls is used to store thread local state specific to serialization +# that needs to be propagated to other files, in particular we use this for +# (1) map_location (needed for wrapper subclasses/third party devices to torch._utils) +# (2) skip_data (needed for torch.Tensor.__reduce_ex__ for skip_data ctx) +# (3) materialize_fake_tensors (needed for torch.Tensor.__reduce_ex__ for skip_data ctx) +class _SerializationLocal(threading.local): + def __init__(self): + super().__init__() + self.map_location: Optional[MAP_LOCATION] = None + self.skip_data: bool = False + self.materialize_fake_tensors: bool = False + + +_serialization_tls = _SerializationLocal() + + +class SourceChangeWarning(Warning): + pass + + +@contextmanager +def mkdtemp(): + path = tempfile.mkdtemp() + try: + yield path + finally: + shutil.rmtree(path) + + +_package_registry: list[ + tuple[ + int, + Callable[[STORAGE], Optional[str]], + Callable[[STORAGE, str], Optional[STORAGE]], + ] +] = [] + + +class LoadEndianness(Enum): + NATIVE = 1 + LITTLE = 2 + BIG = 3 + + +def get_default_load_endianness() -> Optional[LoadEndianness]: + """ + Get fallback byte order for loading files + + If byteorder mark is not present in saved checkpoint, + this byte order is used as fallback. + By default, it's "native" byte order. + + Returns: + default_load_endian: Optional[LoadEndianness] + """ + from torch.utils.serialization import config + + return config.load.endianness + + +def set_default_load_endianness(endianness): + """ + Set fallback byte order for loading files + + If byteorder mark is not present in saved checkpoint, + this byte order is used as fallback. + By default, it's "native" byte order. + + Args: + endianness: the new fallback byte order + """ + if not isinstance(endianness, LoadEndianness) and endianness is not None: + raise TypeError("Invalid argument type in function set_default_load_endianness") + from torch.utils.serialization import config + + config.load.endianness = endianness + + +def get_crc32_options() -> bool: + """ + Get whether :func:`torch.save` computes and writes crc32 for each record. + + Defaults to ``True``. + """ + from torch.utils.serialization import config + + return config.save.compute_crc32 + + +def set_crc32_options(compute_crc32: bool): + """ + Set whether :func:`torch.save` computes and writes crc32 for each record. + + .. note:: + Setting this to ``False`` may make unzipping of the ``torch.save`` output + fail or warn due to corrupted CRC32. However ``torch.load`` will be + able to load the file. + + Args: + compute_crc32 (bool): set crc32 computation flag + """ + from torch.utils.serialization import config + + config.save.compute_crc32 = compute_crc32 + + +def get_default_mmap_options() -> Optional[int]: + """ + Get default mmap options for :func:`torch.load` with ``mmap=True``. + + Defaults to ``mmap.MAP_PRIVATE``. + + + Returns: + default_mmap_options: int + """ + from torch.utils.serialization import config + + return config.load.mmap_flags + + +def _get_storage_alignment() -> int: + """ + Gets alignment for storages in torch.save files/ + + Defaults to 64. + + Returns: + storage_alginment: int + """ + from torch.utils.serialization import config + + return config.save.storage_alignment + + +class set_default_mmap_options: + """ + Context manager or function to set default mmap options for :func:`torch.load` with ``mmap=True`` to flags. + + For now, only either ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED`` are supported. + Please open an issue if you need any other option to be added here. + + .. note:: + This feature is currently not supported for Windows. + + Args: + flags: ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED`` + """ + + def __init__(self, flags: int) -> None: + if IS_WINDOWS: + raise RuntimeError( + "Changing the default mmap options is currently not supported for Windows" + ) + if flags != MAP_PRIVATE and flags != MAP_SHARED: + raise ValueError( + "Invalid argument in function set_default_mmap_options, " + f"expected mmap.MAP_PRIVATE or mmap.MAP_SHARED, but got {flags}" + ) + # global config + from torch.utils.serialization import config + + self.prev = config.load.mmap_flags + config.load.mmap_flags = flags + + def __enter__(self) -> None: + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + from torch.utils.serialization import config + + config.load.mmap_flags = self.prev + + +def clear_safe_globals() -> None: + """ + Clears the list of globals that are safe for ``weights_only`` load. + """ + _weights_only_unpickler._clear_safe_globals() + + +def get_safe_globals() -> list[Union[Callable, tuple[Callable, str]]]: + """ + Returns the list of user-added globals that are safe for ``weights_only`` load. + """ + return _weights_only_unpickler._get_safe_globals() + + +def add_safe_globals(safe_globals: list[Union[Callable, tuple[Callable, str]]]) -> None: + """ + Marks the given globals as safe for ``weights_only`` load. For example, functions + added to this list can be called during unpickling, classes could be instantiated + and have state set. + + Each item in the list can either be a function/class or a tuple of the form + (function/class, string) where string is the full path of the function/class. + + Within the serialized format, each function is identified with its full + path as ``{__module__}.{__qualname__}``. When calling this API, you can provide this + full path that should match the one in the checkpoint otherwise the default + ``{fn.__module__}.{fn.__qualname__}`` will be used. + + Args: + safe_globals (List[Union[Callable, Tuple[Callable, str]]]): list of globals to mark as safe + + Example: + >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization") + >>> import tempfile + >>> class MyTensor(torch.Tensor): + ... pass + >>> t = MyTensor(torch.randn(2, 3)) + >>> with tempfile.NamedTemporaryFile() as f: + ... torch.save(t, f.name) + # Running `torch.load(f.name, weights_only=True)` will fail with + # Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default. + # Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint. + ... torch.serialization.add_safe_globals([MyTensor]) + ... torch.load(f.name, weights_only=True) + # MyTensor([[-0.5024, -1.8152, -0.5455], + # [-0.8234, 2.0500, -0.3657]]) + """ + _weights_only_unpickler._add_safe_globals(safe_globals) + + +class safe_globals(_weights_only_unpickler._safe_globals): + r"""Context-manager that adds certain globals as safe for ``weights_only`` load. + + Args: + safe_globals: List of globals for weights_only load. + + Example: + >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization") + >>> import tempfile + >>> class MyTensor(torch.Tensor): + ... pass + >>> t = MyTensor(torch.randn(2, 3)) + >>> with tempfile.NamedTemporaryFile() as f: + ... torch.save(t, f.name) + # Running `torch.load(f.name, weights_only=True)` will fail with + # Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default. + # Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint. + ... with torch.serialization.safe_globals([MyTensor]): + ... torch.load(f.name, weights_only=True) + # MyTensor([[-0.5024, -1.8152, -0.5455], + # [-0.8234, 2.0500, -0.3657]]) + >>> assert torch.serialization.get_safe_globals() == [] + """ + + +def get_unsafe_globals_in_checkpoint(f: FileLike) -> list[str]: + """Returns a list of strings of functions/classes in a ``torch.save`` object that are not safe for ``weights_only``. + + For a given function or class ``f``, the corresponding string will be of the form + ``{f.__module__}.{f.__name__}``. + + This function will return any GLOBALs in the checkpoint that are not in the set marked safe + for ``weights_only`` (either via :func:`add_safe_globals` or :class:`safe_globals` context or + allowlisted by ``torch`` by default). + + .. note:: + This function will statically disassemble the pickle file in the checkpoint. + The implication is any classes dynamically pushed onto the stack during unpickling + will not be included in the output. + + Args: + f: File-like object or string containing the checkpoint object saved via ``torch.save`` + + Returns: + A list of strings of pickle GLOBALs in the checkpoint that are not allowlisted for ``weights_only``. + """ + default_safe_globals_strings = set( + _weights_only_unpickler._get_allowed_globals().keys() + ) + user_safe_global_strings = set( + _weights_only_unpickler._get_user_allowed_globals().keys() + ) + safe_global_strings = default_safe_globals_strings.union(user_safe_global_strings) + + with _open_file_like(f, "rb") as opened_file: + if not _is_zipfile(opened_file): + raise ValueError("Expected input to be a checkpoint returned by torch.save") + with _open_zipfile_reader(opened_file) as zip_file: + if _is_torchscript_zip(zip_file): + raise ValueError( + "Expected input to be a checkpoint returned by torch.save but got a torchscript checkpoint" + ) + data_file = io.BytesIO(zip_file.get_record("data.pkl")) + all_globals = _weights_only_unpickler.get_globals_in_pkl(data_file) + return list(all_globals.difference(safe_global_strings)) + + +class skip_data: + """ + Context-manager that skips writing/reading storage bytes for ``torch.save`` / ``torch.load`` calls. + + For the save path, storages will still be saved, but the space that their bytes would usually be written to + will be empty space. The storage bytes can then be populated in a separate pass. + + For the load path, tensors will be loaded per the checkpoint but their storages will not be populated with data. + + .. warning:: + The ``skip_data`` context manager is an early prototype and is subject to change. + + Args: + materialize_fake_tensors: Whether to materialize FakeTensors during save. This is a no-op for the load path. + + Example: + >>> # xdoctest: +SKIP("NamedTemporaryFile on Windows") + >>> import tempfile + >>> t = torch.randn(2, 3) + >>> with tempfile.NamedTemporaryFile() as f: + ... with torch.serialization.skip_data(): + ... torch.save(t, f.name) + ... torch.load(f.name, weights_only=True) + tensor([[0., 0., 0.], + [0., 0., 0.]]) + """ + + def __init__(self, materialize_fake_tensors: bool = False): + self.materialize_fake_tensors = materialize_fake_tensors + + def __enter__(self): + global _serialization_tls + self._old_skip_data = _serialization_tls.skip_data + self._old_materialize_fake_tensors = _serialization_tls.materialize_fake_tensors + _serialization_tls.skip_data = True + _serialization_tls.materialize_fake_tensors = self.materialize_fake_tensors + + def __exit__(self, type, value, tb): + global _serialization_tls + _serialization_tls.skip_data = self._old_skip_data + _serialization_tls.materialize_fake_tensors = self._old_materialize_fake_tensors + + +def _is_zipfile(f) -> bool: + # This is a stricter implementation than zipfile.is_zipfile(). + # zipfile.is_zipfile() is True if the magic number appears anywhere in the + # binary. Since we expect the files here to be generated by torch.save or + # torch.jit.save, it's safe to only check the start bytes and avoid + # collisions and assume the zip has only 1 file. + # See bugs.python.org/issue28494. + + start = f.tell() + # Read the first few bytes and match against the ZIP file signature + local_header_magic_number = b"PK\x03\x04" + read_bytes = f.read(len(local_header_magic_number)) + f.seek(start) + return read_bytes == local_header_magic_number + + +def register_package( + priority: int, + tagger: Callable[[STORAGE], Optional[str]], + deserializer: Callable[[STORAGE, str], Optional[STORAGE]], +): + """ + Registers callables for tagging and deserializing storage objects with an associated priority. + Tagging associates a device with a storage object at save time while deserializing moves a + storage object to an appropriate device at load time. :attr:`tagger` and :attr:`deserializer` + are run in the order given by their :attr:`priority` until a tagger/deserializer returns a + value that is not `None`. + + To override the deserialization behavior for a device in the global registry, one can register a + tagger with a higher priority than the existing tagger. + + This function can also be used to register a tagger and deserializer for new devices. + + Args: + priority: Indicates the priority associated with the tagger and deserializer, where a lower + value indicates higher priority. + tagger: Callable that takes in a storage object and returns its tagged device as a string + or None. + deserializer: Callable that takes in storage object and a device string and returns a storage + object on the appropriate device or None. + + Returns: + `None` + + Example: + >>> def ipu_tag(obj): + >>> if obj.device.type == 'ipu': + >>> return 'ipu' + >>> def ipu_deserialize(obj, location): + >>> if location.startswith('ipu'): + >>> ipu = getattr(torch, "ipu", None) + >>> assert ipu is not None, "IPU device module is not loaded" + >>> assert torch.ipu.is_available(), "ipu is not available" + >>> return obj.ipu(location) + >>> torch.serialization.register_package(11, ipu_tag, ipu_deserialize) + """ + queue_elem = (priority, tagger, deserializer) + _package_registry.append(queue_elem) + _package_registry.sort() + + +def check_module_version_greater_or_equal( + module, + req_version_tuple, + error_if_malformed=True, +): + """ + Check if a module's version satisfies requirements + + Usually, a module's version string will be like 'x.y.z', which would be represented + as a tuple (x, y, z), but sometimes it could be an unexpected format. If the version + string does not match the given tuple's format up to the length of the tuple, then + error and exit or emit a warning. + + Args: + module: the module to check the version of + req_version_tuple: tuple (usually of ints) representing the required version + error_if_malformed: whether we should exit if module version string is malformed + + Returns: + requirement_is_met: bool + """ + try: + version_strs = module.__version__.split(".") + # Cast module version fields to match the types of the required version + module_version = tuple( + type(req_field)(version_strs[idx]) + for idx, req_field in enumerate(req_version_tuple) + ) + requirement_is_met = module_version >= req_version_tuple + + except Exception as e: + message = ( + f"'{module.__name__}' module version string is malformed '{module.__version__}' and cannot be compared" + f" with tuple {str(req_version_tuple)}" + ) + if error_if_malformed: + raise RuntimeError(message) from e + else: + warnings.warn(message + ", but continuing assuming that requirement is met") + requirement_is_met = True + + return requirement_is_met + + +def _cpu_tag(obj): + if obj.device.type == "cpu": + return "cpu" + + +def _mps_tag(obj): + if obj.device.type == "mps": + return "mps" + + +def _meta_tag(obj): + if obj.device.type == "meta": + return "meta" + + +def _backend_tag(backend_name, obj): + if backend_name == "privateuse1": + backend_name = torch._C._get_privateuse1_backend_name() + if obj.device.type == backend_name: + if obj.device.index is None: + return backend_name + else: + return backend_name + ":" + str(obj.device.index) + + +def _cpu_deserialize(obj, location): + if location == "cpu": + return obj + + +def _mps_deserialize(obj, location): + if location.startswith("mps"): + return obj.mps() + + +def _meta_deserialize(obj, location): + if location == "meta": + return torch.UntypedStorage(obj.nbytes(), device="meta") + + +def _validate_device(location, backend_name): + """ + Check whether the device index of specified backend is valid + + In case of privateuse1 backend, your must first register a device_module for + privateuse1 using torch._register_device_module. Implement the following + methods in device_module like cuda: device_module._utils._get_device_index(location, True), + device_module.device_count(). + + Args: + location: string of device + backend_name: the backend name or the name of privateuse1, which can be renamed + + Returns: + device_index: int + """ + if not hasattr(torch, backend_name): + raise RuntimeError( + f"The {backend_name.upper()} device module is not registered. " + "If you are running on a CPU-only machine, " + "please use torch.load with map_location=torch.device('cpu') " + "to map your storages to the CPU." + ) + device_module = getattr(torch, backend_name) + if hasattr(device_module, "_utils") and hasattr( + device_module._utils, "_get_device_index" + ): + device_index = device_module._utils._get_device_index(location, True) + device = torch.device(backend_name, device_index) + else: + device = torch.device(location) + device_index = device.index if device.index else 0 + if hasattr(device_module, "is_available") and not device_module.is_available(): + raise RuntimeError( + f"Attempting to deserialize object on a {backend_name.upper()} " + f"device but torch.{backend_name}.is_available() is False. " + "If you are running on a CPU-only machine, " + "please use torch.load with map_location=torch.device('cpu') " + "to map your storages to the CPU." + ) + if hasattr(device_module, "device_count"): + device_count = device_module.device_count() + if device_index >= device_count: + raise RuntimeError( + f"Attempting to deserialize object on {backend_name.upper()} device " + f"{device_index} but torch.{backend_name}.device_count() is {device_count}. " + "Please use torch.load with map_location to map your storages " + "to an existing device." + ) + return device + + +def validate_cuda_device(location): + return _validate_device(location, "cuda").index + + +def validate_hpu_device(location): + return _validate_device(location, "hpu").index + + +def _deserialize(backend_name, obj, location): + if backend_name == "privateuse1": + backend_name = torch._C._get_privateuse1_backend_name() + if location.startswith(backend_name): + device = _validate_device(location, backend_name) + return obj.to(device=device) + + +register_package(10, _cpu_tag, _cpu_deserialize) +register_package( + 20, + functools.partial(_backend_tag, "cuda"), + functools.partial(_deserialize, "cuda"), +) +register_package(21, _mps_tag, _mps_deserialize) +register_package(22, _meta_tag, _meta_deserialize) +register_package( + 23, + functools.partial(_backend_tag, "privateuse1"), + functools.partial(_deserialize, "privateuse1"), +) +register_package( + 24, + functools.partial(_backend_tag, "hpu"), + functools.partial(_deserialize, "hpu"), +) +register_package( + 25, + functools.partial(_backend_tag, "xpu"), + functools.partial(_deserialize, "xpu"), +) + + +def location_tag( + storage: Union[Storage, torch.storage.TypedStorage, torch.UntypedStorage], +): + for _, tagger, _ in _package_registry: + location = tagger(storage) + if location: + return location + raise RuntimeError( + "don't know how to determine data location of " + torch.typename(storage) + ) + + +def default_restore_location(storage, location): + """ + Restores `storage` using a deserializer function registered for the `location`. + + This function looks in the registry for deserializer functions that match the `location`. + If found, it attempts to use them, in priority order, to restore `storage` until one + returns a not `None` result. If no deserializer can be found in the registry, or all found fail + to bear a result, it raises a `RuntimeError`. + + Args: + storage (STORAGE): the storage object to restore + location (str): the location tag associated with the storage object + + Returns: + storage: Optional[STORAGE] + + Raises: + RuntimeError: If no deserializer matching `location` is found in the registry or if + all matching ones return `None`. + """ + for _, _, fn in _package_registry: + result = fn(storage, location) + if result is not None: + return result + raise RuntimeError( + "don't know how to restore data location of " + + torch.typename(storage) + + " (tagged with " + + location + + ")" + ) + + +def normalize_storage_type(storage_type): + return getattr(torch, storage_type.__name__) + + +def storage_to_tensor_type(storage): + storage_type = type(storage) + module = _import_dotted_name(storage_type.__module__) + return getattr(module, storage_type.__name__.replace("Storage", "Tensor")) + + +def _is_path(name_or_buffer: object) -> TypeIs[Union[str, os.PathLike]]: + return isinstance(name_or_buffer, (str, os.PathLike)) + + +T = TypeVar("T") + + +class _opener(Generic[T]): + def __init__(self, file_like: T) -> None: + self.file_like: T = file_like + + def __enter__(self): + return self.file_like + + def __exit__(self, *args): + pass + + +class _open_file(_opener[IO[bytes]]): + def __init__(self, name: Union[str, os.PathLike[str]], mode: str) -> None: + super().__init__(open(name, mode)) + + def __exit__(self, *args): + self.file_like.close() + + +class _open_buffer_reader(_opener[IO[bytes]]): + def __init__(self, buffer: IO[bytes]) -> None: + super().__init__(buffer) + _check_seekable(buffer) + + +class _open_buffer_writer(_opener[IO[bytes]]): + def __exit__(self, *args): + self.file_like.flush() + + +def _open_file_like(name_or_buffer: FileLike, mode: str) -> _opener[IO[bytes]]: + if _is_path(name_or_buffer): + return _open_file(name_or_buffer, mode) + else: + if "w" in mode: + return _open_buffer_writer(name_or_buffer) + elif "r" in mode: + return _open_buffer_reader(name_or_buffer) + else: + raise RuntimeError(f"Expected 'r' or 'w' in mode but got {mode}") + + +class _open_zipfile_reader(_opener[torch._C.PyTorchFileReader]): + def __init__(self, name_or_buffer: Union[str, IO[bytes]]) -> None: + super().__init__(torch._C.PyTorchFileReader(name_or_buffer)) + + +class _open_zipfile_writer_file(_opener[torch._C.PyTorchFileWriter]): + def __init__(self, name: str) -> None: + self.file_stream = None + self.name = name + try: + self.name.encode("ascii") + except UnicodeEncodeError: + # PyTorchFileWriter only supports ascii filename. + # For filenames with non-ascii characters, we rely on Python + # for writing out the file. + self.file_stream = io.FileIO(self.name, mode="w") + super().__init__( + torch._C.PyTorchFileWriter( + self.file_stream, get_crc32_options(), _get_storage_alignment() + ) + ) + else: + super().__init__( + torch._C.PyTorchFileWriter( + self.name, get_crc32_options(), _get_storage_alignment() + ) + ) + + def __exit__(self, *args) -> None: + self.file_like.write_end_of_file() + if self.file_stream is not None: + self.file_stream.close() + + +class _open_zipfile_writer_buffer(_opener[torch._C.PyTorchFileWriter]): + def __init__(self, buffer: IO[bytes]) -> None: + if not callable(getattr(buffer, "write", None)): + msg = f"Buffer of {str(type(buffer)).strip('<>')} has no callable attribute 'write'" + if not hasattr(buffer, "write"): + raise AttributeError(msg) + raise TypeError(msg) + self.buffer = buffer + super().__init__( + torch._C.PyTorchFileWriter( + buffer, get_crc32_options(), _get_storage_alignment() + ) + ) + + def __exit__(self, *args) -> None: + self.file_like.write_end_of_file() + self.buffer.flush() + + +def _open_zipfile_writer(name_or_buffer: Union[str, IO[bytes]]) -> _opener: + container: type[_opener] + if _is_path(name_or_buffer): + container = _open_zipfile_writer_file + else: + container = _open_zipfile_writer_buffer + return container(name_or_buffer) # type: ignore[arg-type] + + +def _is_compressed_file(f) -> bool: + compress_modules = ["gzip"] + try: + return f.__module__ in compress_modules + except AttributeError: + return False + + +def _should_read_directly(f): + """ + Checks if f is a file that should be read directly. It should be read + directly if it is backed by a real file (has a fileno) and is not a + a compressed file (e.g. gzip) + """ + if _is_compressed_file(f): + return False + try: + return f.fileno() >= 0 + except io.UnsupportedOperation: + return False + except AttributeError: + return False + + +def _check_seekable(f) -> bool: + def raise_err_msg(patterns, e): + for p in patterns: + if p in str(e): + msg = ( + str(e) + + ". You can only torch.load from a file that is seekable." + + " Please pre-load the data into a buffer like io.BytesIO and" + + " try to load from it instead." + ) + raise type(e)(msg) + raise e + + try: + f.seek(f.tell()) + return True + except (io.UnsupportedOperation, AttributeError) as e: + raise_err_msg(["seek", "tell"], e) + return False + + +def _check_dill_version(pickle_module) -> None: + """Checks if using dill as the pickle module, and if so, checks if it is the correct version. + If dill version is lower than 0.3.1, a ValueError is raised. + + Args: + pickle_module: module used for pickling metadata and objects + + """ + if pickle_module is not None and pickle_module.__name__ == "dill": + required_dill_version = (0, 3, 1) + if not check_module_version_greater_or_equal( + pickle_module, required_dill_version, False + ): + raise ValueError( + ( + "'torch' supports dill >= {}, but you have dill {}." + " Please upgrade dill or switch to 'pickle'" + ).format( + ".".join([str(num) for num in required_dill_version]), + pickle_module.__version__, + ) + ) + + +def _check_save_filelike(f): + if not _is_path(f) and not hasattr(f, "write"): + raise AttributeError( + "expected 'f' to be string, path, or a file-like object with " + "a 'write' attribute" + ) + + +def save( + obj: object, + f: FileLike, + pickle_module: Any = pickle, + pickle_protocol: int = DEFAULT_PROTOCOL, + _use_new_zipfile_serialization: bool = True, + _disable_byteorder_record: bool = False, +) -> None: + # Reference: https://github.com/pytorch/pytorch/issues/54354 + # The first line of this docstring overrides the one Sphinx generates for the + # documentation. We need it so that Sphinx doesn't leak `pickle`s path from + # the build environment (e.g. `>> # xdoctest: +SKIP("makes cwd dirty") + >>> # Save to file + >>> x = torch.tensor([0, 1, 2, 3, 4]) + >>> torch.save(x, "tensor.pt") + >>> # Save to io.BytesIO buffer + >>> buffer = io.BytesIO() + >>> torch.save(x, buffer) + """ + torch._C._log_api_usage_once("torch.save") + _check_dill_version(pickle_module) + _check_save_filelike(f) + + if isinstance(f, (str, os.PathLike)): + f = os.fspath(f) + + if _use_new_zipfile_serialization: + with _open_zipfile_writer(f) as opened_zipfile: + _save( + obj, + opened_zipfile, + pickle_module, + pickle_protocol, + _disable_byteorder_record, + ) + return + else: + global _serialization_tls + if _serialization_tls.skip_data: + raise RuntimeError( + "Cannot use skip_data=True with _use_new_zipfile_serialization=False" + ) + with _open_file_like(f, "wb") as opened_file: + _legacy_save(obj, opened_file, pickle_module, pickle_protocol) + + +def _legacy_save(obj, f, pickle_module, pickle_protocol) -> None: + import torch.nn as nn + + serialized_container_types = {} + serialized_storages: dict[str, tuple[torch.UntypedStorage, torch.dtype]] = {} + + # Since loading storages that view the same data with different dtypes is + # not supported, we need to keep track of the dtype associated with each + # storage data_ptr and throw an error if the dtype is ever different. + # TODO: This feature could be added in the future + storage_dtypes: dict[int, torch.dtype] = {} + + def persistent_id(obj: Any) -> Optional[tuple]: + # FIXME: the docs say that persistent_id should only return a string + # but torch store returns tuples. This works only in the binary protocol + # see + # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects + # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 + if isinstance(obj, type) and issubclass(obj, nn.Module): + if obj in serialized_container_types: + return None + serialized_container_types[obj] = True + source_file = source = None + try: + source_lines, _, source_file = get_source_lines_and_file(obj) + source = "".join(source_lines) + except ( + Exception + ): # saving the source is optional, so we can ignore any errors + warnings.warn( + "Couldn't retrieve source code for container of " + "type " + obj.__name__ + ". It won't be checked " + "for correctness upon loading." + ) + return ("module", obj, source_file, source) + + if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): + storage: torch.UntypedStorage + + if isinstance(obj, torch.storage.TypedStorage): + # TODO: Once we decide to break serialization FC, this case + # can be deleted + storage = obj._untyped_storage + storage_dtype = obj.dtype + storage_type_str = obj._pickle_storage_type() + storage_type = getattr(torch, storage_type_str) + dtype = obj.dtype + storage_numel = obj._size() + + elif isinstance(obj, torch.UntypedStorage): + storage = obj + storage_dtype = torch.uint8 + storage_type = normalize_storage_type(type(obj)) + dtype = torch.uint8 + storage_numel = storage.nbytes() + else: + raise TypeError(f"type not recognized: {type(obj)}") + + # If storage is allocated, ensure that any other saved storages + # pointing to the same data all have the same dtype. If storage is + # not allocated, don't perform this check + if storage.data_ptr() != 0: + if storage.data_ptr() in storage_dtypes: + if storage_dtype != storage_dtypes[storage.data_ptr()]: + raise RuntimeError( + "Cannot save multiple tensors or storages that " + "view the same data as different types" + ) + else: + storage_dtypes[storage.data_ptr()] = storage_dtype + + view_metadata: Optional[tuple[str, int, int]] + + # Offset is always 0, but we keep it for backwards compatibility + # with the old serialization format (which supported storage views) + offset = 0 + storage_key = str(storage._cdata) + location = location_tag(storage) + + # TODO: There's an issue here with FC. It might be impossible to + # solve, but it's worth noting. Imagine we save a list `[storage, + # tensor]`, where `tensor.storage()` is the same as `storage`, and + # `tensor.element_size() > 1`. Let's say that `tensor.dtype == + # torch.float`. The storage will be serialized with element size + # of 1, since we're choosing to serialize the first occurrence of + # a duplicate storage. Since this legacy serialization format saves + # the numel of the storage, rather than nbytes directly, we'll be + # effectively saving nbytes in this case. We'll be able to load it + # and the tensor back up with no problems in _this_ and future + # versions of pytorch, but in older versions, here's the problem: + # the storage will be loaded up as a UntypedStorage, and then the + # FloatTensor will loaded and the UntypedStorage will be assigned to + # it. Since the storage dtype does not match the tensor dtype, this + # will cause an error. If we reverse the list, like `[tensor, + # storage]`, then we will save the `tensor.storage()` as a faked + # `FloatStorage`, and the saved size will be the correct + # dtype-specific numel count that old versions expect. `tensor` + # will be able to load up properly in old versions, pointing to + # a FloatStorage. However, `storage` is still being translated to + # a UntypedStorage, and it will try to resolve to the same + # FloatStorage that `tensor` contains. This will also cause an + # error. It doesn't seem like there's any way around this. + # Probably, we just cannot maintain FC for the legacy format if the + # saved list contains both a tensor and a storage that point to the + # same data. We should still be able to maintain FC for lists of + # just tensors, as long as all views share the same dtype as the + # tensor they are viewing. + + if storage_key not in serialized_storages: + serialized_storages[storage_key] = (storage, dtype) + is_view = storage._cdata != storage._cdata + if is_view: + view_metadata = (str(storage._cdata), offset, storage.nbytes()) + else: + view_metadata = None + + res = ( + "storage", + storage_type, + storage_key, + location, + storage_numel, + view_metadata, + ) + return res + return None + + sys_info = { + "protocol_version": PROTOCOL_VERSION, + "little_endian": sys.byteorder == "little", + "type_sizes": { + "short": SHORT_SIZE, + "int": INT_SIZE, + "long": LONG_SIZE, + }, + } + + pickle_module.dump(MAGIC_NUMBER, f, protocol=pickle_protocol) + pickle_module.dump(PROTOCOL_VERSION, f, protocol=pickle_protocol) + pickle_module.dump(sys_info, f, protocol=pickle_protocol) + + class PyTorchLegacyPickler(pickle_module.Pickler): + def persistent_id(self, obj): + return persistent_id(obj) + + pickler = PyTorchLegacyPickler(f, protocol=pickle_protocol) + pickler.dump(obj) + + serialized_storage_keys = sorted(serialized_storages.keys()) + pickle_module.dump(serialized_storage_keys, f, protocol=pickle_protocol) + f.flush() + for key in serialized_storage_keys: + storage, dtype = serialized_storages[key] + storage._write_file( + f, _should_read_directly(f), True, torch._utils._element_size(dtype) + ) + + +def _save( + obj, + zip_file, + pickle_module, + pickle_protocol, + _disable_byteorder_record, +): + serialized_storages: dict[str, torch.storage.UntypedStorage] = {} + id_map: dict[int, str] = {} + + # Since loading storages that view the same data with different dtypes is + # not supported, we need to keep track of the dtype associated with each + # storage data_ptr and throw an error if the dtype is ever different. + # TODO: This feature could be added in the future + storage_dtypes: dict[int, torch.dtype] = {} + + def persistent_id(obj): + # FIXME: the docs say that persistent_id should only return a string + # but torch store returns tuples. This works only in the binary protocol + # see + # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects + # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 + if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): + if isinstance(obj, torch.storage.TypedStorage): + # TODO: Once we decide to break serialization FC, this case + # can be deleted + storage = obj._untyped_storage + storage_dtype = obj.dtype + storage_type_str = obj._pickle_storage_type() + storage_type = getattr(torch, storage_type_str) + storage_numel = obj._size() + + else: + storage = obj + storage_dtype = torch.uint8 + storage_type = normalize_storage_type(type(obj)) + storage_numel = storage.nbytes() + + # If storage is allocated, ensure that any other saved storages + # pointing to the same data all have the same dtype. If storage is + # not allocated, don't perform this check + if str(storage.device) != "meta" and storage.data_ptr() != 0: + if storage.data_ptr() in storage_dtypes: + if storage_dtype != storage_dtypes[storage.data_ptr()]: + raise RuntimeError( + "Cannot save multiple tensors or storages that " + "view the same data as different types" + ) + else: + storage_dtypes[storage.data_ptr()] = storage_dtype + + storage_key = id_map.setdefault(storage._cdata, str(len(id_map))) + if hasattr(obj, "_fake_device") and obj._fake_device is not None: + location = str(obj._fake_device) + else: + location = location_tag(storage) + serialized_storages[storage_key] = storage + + return ("storage", storage_type, storage_key, location, storage_numel) + + return None + + # Write the pickle data for `obj` + data_buf = io.BytesIO() + + class PyTorchPickler(pickle_module.Pickler): # type: ignore[name-defined] + def persistent_id(self, obj): + return persistent_id(obj) + + pickler = PyTorchPickler(data_buf, protocol=pickle_protocol) + pickler.dump(obj) + data_value = data_buf.getvalue() + zip_file.write_record("data.pkl", data_value, len(data_value)) + # .format_version is used to track + # 1. version 1 represents the order of storages being changed from + # lexicographical based on keys to numerically ordered based on keys + # 2. version 2 represents including storage_alignment as a record + # within the zipfile + zip_file.write_record(".format_version", "1", len("1")) + storage_alignment = str(_get_storage_alignment()) + zip_file.write_record( + ".storage_alignment", storage_alignment, len(storage_alignment) + ) + + # Write byte order marker + if not _disable_byteorder_record: + if sys.byteorder not in ["little", "big"]: + raise ValueError("Unknown endianness type: " + sys.byteorder) + + zip_file.write_record("byteorder", sys.byteorder, len(sys.byteorder)) + + # Write each tensor to a file named tensor/the_tensor_key in the zip archive + for key in serialized_storages.keys(): + name = f"data/{key}" + storage = serialized_storages[key] + num_bytes = storage.nbytes() + global _serialization_tls + if _serialization_tls.skip_data: + zip_file.write_record_metadata(name, num_bytes) + else: + # given that we copy things around anyway, we might use storage.cpu() + # this means to that to get tensors serialized, you need to implement + # .cpu() on the underlying Storage + if storage.device.type != "cpu": + from torch.utils.serialization import config + + if ( + config.save.use_pinned_memory_for_d2h + and ( + acc := torch.accelerator.current_accelerator( + check_available=True + ) + ) + is not None + and acc.type == storage.device.type + ): + new_storage = torch.empty( + num_bytes, dtype=torch.uint8, device="cpu", pin_memory=True + ).untyped_storage() + new_storage.copy_(storage) + torch.accelerator.current_stream(storage.device.index).synchronize() + storage = new_storage + else: + storage = storage.cpu() + # Now that it is on the CPU we can directly copy it into the zip file + zip_file.write_record(name, storage, num_bytes) + + +def load( + f: FileLike, + map_location: MAP_LOCATION = None, + pickle_module: Any = None, + *, + weights_only: Optional[bool] = None, + mmap: Optional[bool] = None, + **pickle_load_args: Any, +) -> Any: + # Reference: https://github.com/pytorch/pytorch/issues/54354 + # The first line of this docstring overrides the one Sphinx generates for the + # documentation. We need it so that Sphinx doesn't leak `pickle`s path from + # the build environment (e.g. `>> # xdoctest: +SKIP("undefined filepaths") + >>> torch.load("tensors.pt", weights_only=True) + # Load all tensors onto the CPU + >>> torch.load( + ... "tensors.pt", + ... map_location=torch.device("cpu"), + ... weights_only=True, + ... ) + # Load all tensors onto the CPU, using a function + >>> torch.load( + ... "tensors.pt", + ... map_location=lambda storage, loc: storage, + ... weights_only=True, + ... ) + # Load all tensors onto GPU 1 + >>> torch.load( + ... "tensors.pt", + ... map_location=lambda storage, loc: storage.cuda(1), + ... weights_only=True, + ... ) # type: ignore[attr-defined] + # Map tensors from GPU 1 to GPU 0 + >>> torch.load( + ... "tensors.pt", + ... map_location={"cuda:1": "cuda:0"}, + ... weights_only=True, + ... ) + # Load tensor from io.BytesIO object + # Loading from a buffer setting weights_only=False, warning this can be unsafe + >>> with open("tensor.pt", "rb") as f: + ... buffer = io.BytesIO(f.read()) + >>> torch.load(buffer, weights_only=False) + # Load a module with 'ascii' encoding for unpickling + # Loading from a module setting weights_only=False, warning this can be unsafe + >>> torch.load("module.pt", encoding="ascii", weights_only=False) + """ + torch._C._log_api_usage_once("torch.load") + DOCS_MESSAGE = ( + "\n\nCheck the documentation of torch.load to learn more about types accepted by default with " + "weights_only https://pytorch.org/docs/stable/generated/torch.load.html." + ) + + def _get_wo_message(message: str) -> str: + unsafe_global_pattern = r"GLOBAL (\S+) was not an allowed global by default." + has_unsafe_global = re.search(unsafe_global_pattern, message) is not None + blocklist_pattern = r"whose module (\S+) is blocked" + has_blocklist = re.search(blocklist_pattern, message) is not None + import_pattern = r"(\S+) must be (\S+) to load" + has_import = re.search(import_pattern, message) is not None + if has_unsafe_global: + updated_message = ( + "Weights only load failed. This file can still be loaded, to do so you have two options, " + "\033[1mdo those steps only if you trust the source of the checkpoint\033[0m. " + f"\n\t(1) {UNSAFE_MESSAGE}\n\t(2) Alternatively, to load with `weights_only=True` please check " + "the recommended steps in the following error message.\n\tWeightsUnpickler error: " + + message + ) + else: + if has_import: + return f"Weights only load failed. {message}\n {UNSAFE_MESSAGE}\n" + else: + updated_message = f"Weights only load failed. {UNSAFE_MESSAGE}\n" + if not has_blocklist: + updated_message += ( + "Please file an issue with the following so that we can make " + "`weights_only=True` compatible with your use case: WeightsUnpickler error: " + ) + updated_message += "\n\n" + message + return updated_message + DOCS_MESSAGE + + weights_only_not_set = weights_only is None + + if weights_only_not_set: + weights_only = _default_to_weights_only(pickle_module) + + true_values = ["1", "y", "yes", "true"] + # Add ability to force safe only or non-safe weight loads via environment variables + force_weights_only_load = ( + os.getenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD", "0") in true_values + ) + force_no_weights_only_load = ( + os.getenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "0") in true_values + ) + + if force_weights_only_load and force_no_weights_only_load: + raise RuntimeError( + "Only one of `TORCH_FORCE_WEIGHTS_ONLY_LOAD` or `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD` " + "should be set, but both were set." + ) + elif force_weights_only_load: + weights_only = True + elif force_no_weights_only_load: + # TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD can only override if callsite did not explicitly set weights_only + if weights_only_not_set: + warnings.warn( + "Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the" + "`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.", + UserWarning, + stacklevel=2, + ) + weights_only = False + + if weights_only: + if pickle_module is not None: + raise RuntimeError( + "Can not safely load weights when explicit pickle_module is specified" + ) + else: + if pickle_module is None: + pickle_module = pickle + + # make flipping default BC-compatible + if mmap is None: + from torch.utils.serialization import config + + mmap = config.load.mmap + + _check_dill_version(pickle_module) + + if "encoding" not in pickle_load_args.keys(): + pickle_load_args["encoding"] = "utf-8" + + with _open_file_like(f, "rb") as opened_file: + if _is_zipfile(opened_file): + # The zipfile reader is going to advance the current file position. + # If we want to actually tail call to torch.jit.load, we need to + # reset back to the original position. + orig_position = opened_file.tell() + overall_storage = None + with _open_zipfile_reader(opened_file) as opened_zipfile: + if _is_torchscript_zip(opened_zipfile): + warnings.warn( + "'torch.load' received a zip file that looks like a TorchScript archive" + " dispatching to 'torch.jit.load' (call 'torch.jit.load' directly to" + " silence this warning)", + UserWarning, + ) + if weights_only: + raise RuntimeError( + "Cannot use ``weights_only=True`` with TorchScript archives passed to " + "``torch.load``. " + UNSAFE_MESSAGE + ) + opened_file.seek(orig_position) + return torch.jit.load(opened_file, map_location=map_location) + if mmap: + if not _is_path(f): + raise ValueError( + "f must be a file path in order to use the mmap argument" + ) + size = os.path.getsize(f) + if not IS_WINDOWS: + shared = get_default_mmap_options() == MAP_SHARED + else: + shared = False + overall_storage = torch.UntypedStorage.from_file( + os.fspath(f), shared, size + ) + if weights_only: + try: + return _load( + opened_zipfile, + map_location, + _weights_only_unpickler, + overall_storage=overall_storage, + **pickle_load_args, + ) + except pickle.UnpicklingError as e: + raise pickle.UnpicklingError(_get_wo_message(str(e))) from None + return _load( + opened_zipfile, + map_location, + pickle_module, + overall_storage=overall_storage, + **pickle_load_args, + ) + if mmap: + f_name = "" if not isinstance(f, str) else f"{f}, " + raise RuntimeError( + "mmap can only be used with files saved with " + f"`torch.save({f_name}_use_new_zipfile_serialization=True), " + "please torch.save your checkpoint with this option in order to use mmap." + ) + if weights_only: + try: + return _legacy_load( + opened_file, + map_location, + _weights_only_unpickler, + **pickle_load_args, + ) + except pickle.UnpicklingError as e: + raise pickle.UnpicklingError(_get_wo_message(str(e))) from None + return _legacy_load( + opened_file, map_location, pickle_module, **pickle_load_args + ) + + +# Register pickling support for layout instances such as +# torch.sparse_coo, etc +def _get_layout(name): + """Get layout extension object from its string representation.""" + cache = _get_layout.cache # type: ignore[attr-defined] + if not cache: + for v in torch.__dict__.values(): + if isinstance(v, torch.layout): + cache[str(v)] = v + return cache[name] + + +# There are yet not good way to type annotate function attributes https://github.com/python/mypy/issues/2087 +_get_layout.cache = {} # type: ignore[attr-defined] +copyreg.pickle(torch.layout, lambda obj: (_get_layout, (str(obj),))) + + +def _legacy_load(f, map_location, pickle_module, **pickle_load_args): + deserialized_objects: dict[int, Any] = {} + + restore_location = _get_restore_location(map_location) + + class UnpicklerWrapper(pickle_module.Unpickler): # type: ignore[name-defined] + def find_class(self, mod_name, name): + if type(name) is str and "Storage" in name: + try: + return StorageType(name) + except KeyError: + pass + return super().find_class(mod_name, name) + + def _check_container_source(container_type, source_file, original_source): + try: + current_source = "".join(get_source_lines_and_file(container_type)[0]) + except Exception: # saving the source is optional, so we can ignore any errors + warnings.warn( + "Couldn't retrieve source code for container of " + "type " + container_type.__name__ + ". It won't be checked " + "for correctness upon loading." + ) + return + if original_source != current_source: + if container_type.dump_patches: + file_name = container_type.__name__ + ".patch" + diff = difflib.unified_diff( + current_source.split("\n"), + original_source.split("\n"), + source_file, + source_file, + lineterm="", + ) + lines = "\n".join(diff) + try: + with open(file_name, "a+") as f: + file_size = f.seek(0, 2) + f.seek(0) + if file_size == 0: + f.write(lines) + elif file_size != len(lines) or f.read() != lines: + raise OSError + msg = ( + "Saved a reverse patch to " + file_name + ". " + "Run `patch -p0 < " + file_name + "` to revert your " + "changes." + ) + except OSError: + msg = ( + "Tried to save a patch, but couldn't create a " + "writable file " + file_name + ". Make sure it " + "doesn't exist and your working directory is " + "writable." + ) + else: + msg = ( + "you can retrieve the original source code by " + "accessing the object's source attribute or set " + "`torch.nn.Module.dump_patches = True` and use the " + "patch tool to revert the changes." + ) + msg = f"source code of class '{torch.typename(container_type)}' has changed. {msg}" + warnings.warn(msg, SourceChangeWarning) + + def legacy_load(f): + deserialized_objects: dict[int, Any] = {} + + def persistent_load(saved_id): + if isinstance(saved_id, tuple): + # Ignore containers that don't have any sources saved + if all(saved_id[1:]): + _check_container_source(*saved_id) + return saved_id[0] + return deserialized_objects[int(saved_id)] + + with ( + closing( + tarfile.open(fileobj=f, mode="r:", format=tarfile.PAX_FORMAT) + ) as tar, + mkdtemp() as tmpdir, + ): + if pickle_module is _weights_only_unpickler: + raise RuntimeError( + "Cannot use ``weights_only=True`` with files saved in the " + "legacy .tar format. " + UNSAFE_MESSAGE + ) + tar.extract("storages", path=tmpdir) + with open(os.path.join(tmpdir, "storages"), "rb", 0) as f: + num_storages = pickle_module.load(f, **pickle_load_args) + for _ in range(num_storages): + args = pickle_module.load(f, **pickle_load_args) + key, location, storage_type = args + dtype = storage_type._dtype + obj = cast(Storage, torch.UntypedStorage)._new_with_file( + f, torch._utils._element_size(dtype) + ) + obj = restore_location(obj, location) + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + deserialized_objects[key] = torch.storage.TypedStorage( + wrap_storage=obj, dtype=dtype, _internal=True + ) + + storage_views = pickle_module.load(f, **pickle_load_args) + for target_cdata, root_cdata, offset, numel in storage_views: + root = deserialized_objects[root_cdata] + element_size = torch._utils._element_size(root.dtype) + offset_bytes = offset * element_size + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + deserialized_objects[target_cdata] = torch.storage.TypedStorage( + wrap_storage=root._untyped_storage[ + offset_bytes : offset_bytes + numel * element_size + ], + dtype=root.dtype, + _internal=True, + ) + + tar.extract("tensors", path=tmpdir) + with open(os.path.join(tmpdir, "tensors"), "rb", 0) as f: + num_tensors = pickle_module.load(f, **pickle_load_args) + for _ in range(num_tensors): + args = pickle_module.load(f, **pickle_load_args) + key, storage_id, _original_tensor_type = args + storage = deserialized_objects[storage_id] + (ndim,) = struct.unpack(" str: + # When using encoding='bytes' in Py3, some **internal** keys stored as + # strings in Py2 are loaded as bytes. This function decodes them with + # ascii encoding, one that Py3 uses by default. + # + # NOTE: This should only be used on internal keys (e.g., `typename` and + # `location` in `persistent_load` below! + if isinstance(bytes_str, bytes): + return bytes_str.decode("ascii") + return bytes_str + + +def _get_restore_location(map_location): + if map_location is None: + restore_location = default_restore_location + elif isinstance(map_location, dict): + + def restore_location(storage, location): + location = map_location.get(location, location) + return default_restore_location(storage, location) + + elif isinstance(map_location, (str, bytes)): + + def restore_location(storage, location): + return default_restore_location(storage, map_location) + + elif isinstance(map_location, torch.device): + + def restore_location(storage, location): + return default_restore_location(storage, str(map_location)) + + else: + + def restore_location(storage, location): + result = map_location(storage, location) + if result is None: + result = default_restore_location(storage, location) + return result + + return restore_location + + +class StorageType: + def __init__(self, name): + self._dtype = _get_dtype_from_pickle_storage_type(name) + + @property + def dtype(self): + return self._dtype + + def __str__(self): + return f"StorageType(dtype={self.dtype})" + + +def _load( + zip_file, + map_location, + pickle_module, + pickle_file="data.pkl", + overall_storage=None, + **pickle_load_args, +): + restore_location = _get_restore_location(map_location) + + loaded_storages = {} + + can_calculate_storage_offsets = False + if zip_file.has_record(".format_version"): + version = zip_file.get_record(".format_version") + can_calculate_storage_offsets = version >= b"1" + + # check if byteswapping is needed + byteordername = "byteorder" + byteorderdata = None + if zip_file.has_record(byteordername): + byteorderdata = zip_file.get_record(byteordername) + if byteorderdata not in [b"little", b"big"]: + raise ValueError("Unknown endianness type: " + byteorderdata.decode()) + elif ( + get_default_load_endianness() == LoadEndianness.LITTLE + or get_default_load_endianness() is None + ): + byteorderdata = b"little" + elif get_default_load_endianness() == LoadEndianness.BIG: + byteorderdata = b"big" + elif get_default_load_endianness() == LoadEndianness.NATIVE: + pass + else: + raise ValueError("Invalid load endianness type") + + storage_alignment = 64 + if zip_file.has_record(".storage_alignment"): + storage_alignment = int(zip_file.get_record(".storage_alignment")) + + if ( + not zip_file.has_record(byteordername) + and get_default_load_endianness() is None + and sys.byteorder == "big" + ): + # Default behaviour was changed + # See https://github.com/pytorch/pytorch/issues/101688 + warnings.warn( + "The default load endianness for checkpoints without a byteorder mark " + "on big endian machines was changed from 'native' to 'little' endian, " + "to avoid this behavior please use " + "torch.serialization.set_default_load_endianness to set " + "the desired default load endianness", + UserWarning, + ) + + from torch.utils.serialization import config + + calculate_storage_offsets = config.load.calculate_storage_offsets + run_debug_asserts = os.environ.get("TORCH_SERIALIZATION_DEBUG", "0") == "1" + current_offset = None + # constants from miniz.h/miniz.c + data_descripter_size64 = 24 + data_descripter_size32 = 16 + mz_uint32_max = 0xFFFFFFFF + offsets: dict[str, int] = dict() + + def _get_offset(key, name, numel): + """ + Return the offset of the storage associated with key with record name `name` and size numel. + It is expected that the zipfile header of this storage starts at current_offset. + + WARNING: This function relies on the behavior of the zipwriter in miniz.c. In particular, + the behavior of `mz_zip_writer_add_mem_ex_v2`. The behavior of this function must be kept + in sync with that of miniz! + + After reading a storage of size numel that starts at storage_offset + if it is the first time that storage was read, update nonlocal variable + current_offset to the start of the next zipfile header by incrementing + it by numel and the data descriptor size. + """ + nonlocal current_offset, offsets + if name in offsets: + storage_offset = offsets[name] + return storage_offset + + if current_offset is None: + assert key == "0" + current_offset = zip_file.get_record_offset(name) + local_header_offset = zip_file.get_record_header_offset(name) + storage_offset = current_offset + else: + storage_offset = zip_file.get_record_offset_no_read( + current_offset, name, numel, storage_alignment + ) + local_header_offset = current_offset + + # This is only actually needed for storages that have typed_storage._data_ptr() == 0 + # after being read. Otherwise persistent_load would never "re-call" load_tensor + # for a given key. + offsets[name] = storage_offset + + # Increment current_offset to offset where next zipfile header starts + current_offset = storage_offset + numel + # add size of data descriptor after payload + if numel > 0: + if local_header_offset >= mz_uint32_max or numel >= mz_uint32_max: + current_offset += data_descripter_size64 + else: + current_offset += data_descripter_size32 + + return storage_offset + + def load_tensor(dtype, numel, key, location): + name = f"data/{key}" + if torch._guards.detect_fake_mode(None) is not None: + nbytes = numel * torch._utils._element_size(dtype) + storage = torch.UntypedStorage(nbytes, device="meta") + if can_calculate_storage_offsets: + storage._checkpoint_offset = _get_offset(key, name, numel) + else: + storage._checkpoint_offset = zip_file.get_record_offset(name) + elif _serialization_tls.skip_data: + nbytes = numel * torch._utils._element_size(dtype) + storage = torch.UntypedStorage(nbytes) + elif overall_storage is not None: + if can_calculate_storage_offsets and calculate_storage_offsets: + storage_offset = _get_offset(key, name, numel) + if run_debug_asserts: + if storage_offset != zip_file.get_record_offset(name): + raise RuntimeError( + "This is a debug assert that was run as the `TORCH_SERIALIZATION_DEBUG` environment " + f"variable was set: Incorrect offset for {name}, got {storage_offset} expected " + f"{zip_file.get_record_offset(name)}" + ) + else: + storage_offset = zip_file.get_record_offset(name) + storage = overall_storage[storage_offset : storage_offset + numel] + else: + if can_calculate_storage_offsets and run_debug_asserts: + # This is debug code that we use to test the validity of + # torch.utils.serialization.config.load.calculate_storage_offsets throughout CI + storage_offset = _get_offset(key, name, numel) + if storage_offset != zip_file.get_record_offset(name): + raise RuntimeError( + "This is a debug assert that was run as the `TORCH_SERIALIZATION_DEBUG` environment " + f"variable was set: Incorrect offset for {name}, got {storage_offset} expected " + f"{zip_file.get_record_offset(name)}" + ) + storage = ( + zip_file.get_storage_from_record(name, numel, torch.UntypedStorage) + ._typed_storage() + ._untyped_storage + ) + # swap here if byteswapping is needed + if byteorderdata is not None: + if byteorderdata.decode() != sys.byteorder: + storage.byteswap(dtype) + + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + + if torch._guards.detect_fake_mode(None) is None: + wrap_storage = restore_location(storage, location) + else: + storage._fake_device = location + wrap_storage = storage + + typed_storage = torch.storage.TypedStorage( + wrap_storage=wrap_storage, + dtype=dtype, + _internal=True, + ) + + if typed_storage._data_ptr() != 0: + loaded_storages[key] = typed_storage + + return typed_storage + + def persistent_load(saved_id): + assert isinstance(saved_id, tuple) + typename = _maybe_decode_ascii(saved_id[0]) + data = saved_id[1:] + + assert typename == "storage", ( + f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'" + ) + storage_type, key, location, numel = data + if storage_type is torch.UntypedStorage: + dtype = torch.uint8 + else: + dtype = storage_type.dtype + + if key in loaded_storages: + typed_storage = loaded_storages[key] + else: + nbytes = numel * torch._utils._element_size(dtype) + typed_storage = load_tensor( + dtype, nbytes, key, _maybe_decode_ascii(location) + ) + + return typed_storage + + load_module_mapping: dict[str, str] = { + # See https://github.com/pytorch/pytorch/pull/51633 + "torch.tensor": "torch._tensor" + } + + # Need to subclass Unpickler instead of directly monkey-patching the find_class method + # because it's marked readonly in pickle. + # The type: ignore is because mypy can't statically determine the type of this class. + class UnpicklerWrapper(pickle_module.Unpickler): # type: ignore[name-defined] + # from https://stackoverflow.com/questions/13398462/unpickling-python-objects-with-a-changed-module-path/13405732 + # Lets us override the imports that pickle uses when unpickling an object. + # This is useful for maintaining BC if we change a module path that tensor instantiation relies on. + def find_class(self, mod_name, name): + if type(name) is str and "Storage" in name: + try: + return StorageType(name) + except KeyError: + pass + mod_name = load_module_mapping.get(mod_name, mod_name) + return super().find_class(mod_name, name) + + # Load the data (which may in turn use `persistent_load` to load tensors) + data_file = io.BytesIO(zip_file.get_record(pickle_file)) + + unpickler = UnpicklerWrapper(data_file, **pickle_load_args) + unpickler.persistent_load = persistent_load + # Needed for tensors where storage device and rebuild tensor device are + # not connected (wrapper subclasses and tensors rebuilt using numpy) + global _serialization_tls + _serialization_tls.map_location = map_location + result = unpickler.load() + _serialization_tls.map_location = None + + torch._utils._validate_loaded_sparse_tensors() + torch._C._log_api_usage_metadata( + "torch.load.metadata", {"serialization_id": zip_file.serialization_id()} + ) + return result + + +def _is_torchscript_zip(zip_file): + return "constants.pkl" in zip_file.get_all_records() diff --git a/venv/lib/python3.10/site-packages/torch/storage.py b/venv/lib/python3.10/site-packages/torch/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..e651bc9d16eb181bc61903da197ce7e6ce3395bd --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/storage.py @@ -0,0 +1,1551 @@ +# mypy: allow-untyped-defs + +from __future__ import annotations + +import collections +import copy +import functools +import io +import threading +import warnings +from typing import Any, cast, Optional as _Optional, TYPE_CHECKING, TypeVar, Union +from typing_extensions import Self + +import torch +from torch._utils import _to, _type +from torch.types import _bool, _int, Storage + + +if TYPE_CHECKING: + from torch._prims_common import DeviceLikeType + + +__all__ = ["TypedStorage", "UntypedStorage"] + + +try: + import numpy as np + + HAS_NUMPY = True +except ModuleNotFoundError: + HAS_NUMPY = False + np = None # type: ignore[assignment] + + +_share_memory_lock = threading.Lock() +_share_memory_map: dict[int, threading.RLock] = {} + +T = TypeVar("T", bound="Union[_StorageBase, TypedStorage]") + + +class _StorageBase: + _cdata: Any + is_sparse: _bool = False + is_sparse_csr: _bool = False + device: torch.device + # Used when + # (1) stashing FakeTensor device onto storage in torch.serialization.skip_data + # (2) stashing device onto storage to propagate to FakeTensor when torch.load under FakeTensorMode + _fake_device: _Optional[torch.device] = None + # Used when loading with FakeTensorMode to give information about offset of storage in torch.saved-file + _checkpoint_offset: _Optional[int] = None + + def __init__(self, *args, **kwargs): + pass + + def __len__(self) -> _int: + raise NotImplementedError + + def __getitem__(self, idx): + raise NotImplementedError + + def __setitem__(self, *args, **kwargs): + raise NotImplementedError + + def copy_(self, source: T, non_blocking: _Optional[_bool] = None) -> T: + raise NotImplementedError + + def new(self) -> Union[_StorageBase, TypedStorage]: + raise NotImplementedError + + def nbytes(self) -> _int: + raise NotImplementedError + + def size(self) -> _int: + return self.nbytes() + + def type( + self, dtype: _Optional[str] = None, non_blocking: _bool = False + ) -> Union[_StorageBase, TypedStorage]: + return _type(self, dtype, non_blocking) + + def cuda( + self, device=None, non_blocking=False + ) -> Union[_StorageBase, TypedStorage]: + """Returns a copy of this object in CUDA memory. + + If this object is already in CUDA memory and on the correct device, then + no copy is performed and the original object is returned. + + Args: + device (int): The destination GPU id. Defaults to the current device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + device2 = torch.device("cuda", device) if device else torch.device("cuda") + return self.to(device=device2, non_blocking=non_blocking) + + def hpu(self, device=None, non_blocking=False) -> Union[_StorageBase, TypedStorage]: + """Returns a copy of this object in HPU memory. + + If this object is already in HPU memory and on the correct device, then + no copy is performed and the original object is returned. + + Args: + device (int): The destination HPU id. Defaults to the current device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + device2 = torch.device("hpu", device) if device else torch.device("hpu") + return self.to(device=device2, non_blocking=non_blocking) + + def element_size(self) -> _int: + raise NotImplementedError + + def get_device(self) -> _int: + return self.device.index + + def data_ptr(self) -> _int: + raise NotImplementedError + + def resizable(self) -> _bool: + raise NotImplementedError + + # Defined in torch/csrc/generic/StorageSharing.cpp + def _share_filename_cpu_(self, *args, **kwargs): + raise NotImplementedError + + def _share_fd_cpu_(self, *args, **kwargs): + raise NotImplementedError + + @classmethod + def _new_using_filename_cpu(cls, size: _int) -> Self: + raise NotImplementedError + + @classmethod + def _new_using_fd_cpu(cls, size: _int) -> Self: + raise NotImplementedError + + @classmethod + def from_buffer(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + @classmethod + def _new_shared_filename_cpu( + cls, + manager, + obj, + size, + *, + device=None, + dtype=None, + ) -> Self: + raise NotImplementedError + + @classmethod + def _release_ipc_counter(cls, *args, device=None, **kwargs): + return cls._release_ipc_counter_cuda(*args, **kwargs) + + @classmethod + def _release_ipc_counter_cuda(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + @classmethod + def _new_with_weak_ptr(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + def _shared_decref(self) -> Union[_StorageBase, TypedStorage]: + raise NotImplementedError + + def _write_file(self, *args, **kwargs): + raise NotImplementedError + + def resize_(self, size: _int): + raise NotImplementedError + + def _weak_ref(self, *args, **kwargs) -> Union[_StorageBase, TypedStorage]: + raise NotImplementedError + + def _set_from_file(self, *args, **kwargs): + raise NotImplementedError + + def _set_cdata(self, *args, **kwargs): + raise NotImplementedError + + def _share_cuda_(self, *args, **kwargs): + raise NotImplementedError + + def is_shared(self) -> _bool: + raise NotImplementedError + + @classmethod + def _new_shared_cuda(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + def _shared_incref(self, *args, **kwargs): + raise NotImplementedError + + @classmethod + def _free_weak_ref(cls, *args, **kwargs): + raise NotImplementedError + + @property + def is_cuda(self): + raise NotImplementedError + + @property + def is_hpu(self): + raise NotImplementedError + + @classmethod + def from_file(cls, filename, shared, nbytes) -> Union[_StorageBase, TypedStorage]: + raise NotImplementedError + + @classmethod + def _expired(cls, *args, **kwargs) -> Union[_StorageBase, TypedStorage]: + raise NotImplementedError + + def _byteswap(self, *args, **kwargs): + raise NotImplementedError + + def _get_filename(self, *args, **kwargs) -> _Optional[str]: + raise NotImplementedError + + def __repr__(self): + info_str = f"[{torch.typename(self)}(device={self.device}) of size {len(self)}]" + if self.device.type == "meta": + return "...\n" + info_str + data_str = " " + "\n ".join(str(self[i]) for i in range(self.size())) + return data_str + "\n" + info_str + + def __iter__(self): + return iter(self[i] for i in range(self.size())) + + def __copy__(self): + return self.clone() + + def __deepcopy__(self, memo): + memo = memo.setdefault("torch", {}) + if self._cdata in memo: + return memo[self._cdata] + new_storage = self.clone() + memo[self._cdata] = new_storage + return new_storage + + def __reduce__(self): + b = io.BytesIO() + torch.save(self, b, _use_new_zipfile_serialization=False) + return (_load_from_bytes, (b.getvalue(),)) + + def __sizeof__(self): + return super().__sizeof__() + self.size() + + def clone(self): + """Return a copy of this storage.""" + return type(self)(self.nbytes(), device=self.device).copy_(self) + + def tolist(self): + """Return a list containing the elements of this storage.""" + return list(self) + + def cpu(self): + """Return a CPU copy of this storage if it's not already on the CPU.""" + if self.device.type != "cpu": + return torch.UntypedStorage(self.size()).copy_(self, False) + return self + + def mps(self): + """Return a MPS copy of this storage if it's not already on the MPS.""" + if self.device.type != "mps": + return torch.UntypedStorage(self.size(), device="mps").copy_(self, False) + return self + + def _to(self, dtype): + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Argument 'dtype' must be torch.dtype, not {type(dtype)}") + storage = ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .to(dtype) + ._typed_storage() + ) + if storage.data_ptr() == self.data_ptr(): + storage = storage.clone() + return storage + + def to(self, *, device: DeviceLikeType, non_blocking: _bool = False): + if not isinstance(device, torch.device): + device = torch.device(device) + return _to(self, device, non_blocking) + + def double(self): + """Casts this storage to double type.""" + return self._to(torch.double) + + def float(self): + """Casts this storage to float type.""" + return self._to(torch.float) + + def half(self): + """Casts this storage to half type.""" + return self._to(torch.half) + + def long(self): + """Casts this storage to long type.""" + return self._to(torch.long) + + def int(self): + """Casts this storage to int type.""" + return self._to(torch.int) + + def short(self): + """Casts this storage to short type.""" + return self._to(torch.short) + + def char(self): + """Casts this storage to char type.""" + return self._to(torch.int8) + + def byte(self): + """Casts this storage to byte type.""" + return self._to(torch.uint8) + + def bool(self): + """Casts this storage to bool type.""" + return self._to(torch.bool) + + def bfloat16(self): + """Casts this storage to bfloat16 type.""" + return self._to(torch.bfloat16) + + def complex_double(self): + """Casts this storage to complex double type.""" + return self._to(torch.cdouble) + + def complex_float(self): + """Casts this storage to complex float type.""" + return self._to(torch.cfloat) + + def float8_e5m2(self): + """Casts this storage to float8_e5m2 type""" + return self._to(torch.float8_e5m2) + + def float8_e4m3fn(self): + """Casts this storage to float8_e4m3fn type""" + return self._to(torch.float8_e4m3fn) + + def float8_e5m2fnuz(self): + """Casts this storage to float8_e5m2fnuz type""" + return self._to(torch.float8_e5m2fnuz) + + def float8_e4m3fnuz(self): + """Casts this storage to float8_e4m3fnuz type""" + return self._to(torch.float8_e4m3fnuz) + + def is_pinned(self, device: Union[str, torch.device] = "cuda"): + r"""Determine whether the CPU storage is already pinned on device. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A boolean variable. + """ + return ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .is_pinned(device) + ) + + def pin_memory(self, device: Union[str, torch.device] = "cuda"): + r"""Copy the CPU storage to pinned memory, if it's not already pinned. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A pinned CPU storage. + """ + if self.device.type != "cpu": + raise TypeError(f"cannot pin '{self.type()}' only CPU memory can be pinned") + + pinned_tensor = ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .pin_memory(device) + ) + return pinned_tensor.untyped_storage() + + def share_memory_(self): + """See :meth:`torch.UntypedStorage.share_memory_`""" + from torch.multiprocessing import get_sharing_strategy + + if self.device.type in ["cuda", torch._C._get_privateuse1_backend_name()]: + pass # CUDA or PrivateUse1 doesn't use POSIX shared memory + elif get_sharing_strategy() == "file_system": + self._share_filename_cpu_() + else: + self._share_fd_cpu_() + return self + + @classmethod + def _new_shared(cls, size, *, device="cpu"): + """Create a new storage in shared memory with the same data type.""" + from torch.multiprocessing import get_sharing_strategy + + device = torch.device(device) + if device.type in ["cuda", torch._C._get_privateuse1_backend_name(), "hpu"]: + return cls(size, device=device) + elif get_sharing_strategy() == "file_system": + return cls._new_using_filename_cpu(size) + else: + return cls._new_using_fd_cpu(size) + + def untyped(self): + return self + + def byteswap(self, dtype): + """Swap bytes in underlying data.""" + elem_size = torch._utils._element_size(dtype) + # for complex types, don't swap first and second numbers + if dtype.is_complex: + elem_size = max(int(elem_size / 2), 1) + self._byteswap(elem_size) + + +def _share_memory_lock_protected(fn): + @functools.wraps(fn) + def wrapper(self, *args, **kwargs): + to_free = None + to_wait = None + with _share_memory_lock: + key = self._cdata + if key in _share_memory_map: + to_wait = _share_memory_map[key] + else: + _share_memory_map[key] = threading.RLock() + _share_memory_map[key].acquire() + to_free = key + + # If we're already in the process of sharing the storage, wait + # for it to be done. + if to_wait is not None: + with to_wait: + pass + + try: + return fn(self, *args, **kwargs) + finally: + # If we acquired the storage lock here and we're done working on it + # we can now release it and free the entry. + if to_free is not None: + # Ensure that the cdata from the storage didn't change and only + # the data_ptr did. + assert self._cdata == to_free + with _share_memory_lock: + _share_memory_map[to_free].release() + del _share_memory_map[to_free] + + return wrapper + + +class UntypedStorage(torch._C.StorageBase, _StorageBase): + def __getitem__(self, *args, **kwargs): + if self.device.type == "meta": + raise NotImplementedError("Not available for 'meta' device type") + return super().__getitem__(*args, **kwargs) + + @property + def is_cuda(self): + return self.device.type == "cuda" + + @property + def is_hpu(self): + return self.device.type == "hpu" + + @property + def filename(self) -> _Optional[str]: + """Returns the file name associated with this storage. + + The file name will be a string if the storage is on CPU and was created via + :meth:`~torch.from_file()` with ``shared`` as ``True``. This attribute is ``None`` otherwise. + """ + return self._get_filename() + + @_share_memory_lock_protected + def share_memory_(self, *args, **kwargs): + """ + Moves the storage to shared memory. + + This is a no-op for storages already in shared memory and for CUDA + storages, which do not need to be moved for sharing across processes. + Storages in shared memory cannot be resized. + + Note that to mitigate issues like `this `_ + it is thread safe to call this function from multiple threads on the same object. + It is NOT thread safe though to call any other function on self without proper + synchronization. Please see :doc:`/notes/multiprocessing` for more details. + + .. note:: + When all references to a storage in shared memory are deleted, the associated shared memory + object will also be deleted. PyTorch has a special cleanup process to ensure that this happens + even if the current process exits unexpectedly. + + It is worth noting the difference between :meth:`share_memory_` and :meth:`from_file` with ``shared = True`` + + #. ``share_memory_`` uses `shm_open(3) `_ to create a + POSIX shared memory object while :meth:`from_file` uses + `open(2) `_ to open the filename passed by the user. + #. Both use an `mmap(2) call `_ with ``MAP_SHARED`` + to map the file/object into the current virtual address space + #. ``share_memory_`` will call ``shm_unlink(3)`` on the object after mapping it to make sure the shared memory + object is freed when no process has the object open. ``torch.from_file(shared=True)`` does not unlink the + file. This file is persistent and will remain until it is deleted by the user. + + Returns: + ``self`` + """ + return super().share_memory_(*args, **kwargs) + + @_share_memory_lock_protected + def _share_fd_cpu_(self, *args, **kwargs): + return super()._share_fd_cpu_(*args, **kwargs) + + @_share_memory_lock_protected + def _share_filename_cpu_(self, *args, **kwargs): + return super()._share_filename_cpu_(*args, **kwargs) + + +def _load_from_bytes(b): + return torch.load(io.BytesIO(b), weights_only=False) + + +@functools.cache +def _new_dtypes(): + # These are dtypes serialized as UntypedStorage unlike those in + # _dtype_to_storage_type_map + return { + torch.float8_e5m2, + torch.float8_e4m3fn, + torch.float8_e5m2fnuz, + torch.float8_e4m3fnuz, + torch.float8_e8m0fnu, + torch.float4_e2m1fn_x2, + torch.bits8, + torch.bits16, + torch.bits1x8, + torch.bits2x4, + torch.bits4x2, + torch.complex32, + torch.uint16, + torch.uint32, + torch.uint64, + } + + +@functools.cache +def _dtype_to_storage_type_map(): + # NOTE: We should no longer add dtypes to this map. This map + # is only used for BC/FC with older PyTorch versions. Going forward, + # new dtypes of TypedStorage should not translate to a legacy + # Storage class. Instead, new dtypes of TypedStorage should + # be serialized as an UntypedStorage paired with a torch.dtype + return { + torch.double: "DoubleStorage", + torch.float: "FloatStorage", + torch.half: "HalfStorage", + torch.long: "LongStorage", + torch.int: "IntStorage", + torch.int16: "ShortStorage", + torch.int8: "CharStorage", + torch.uint8: "ByteStorage", + torch.bool: "BoolStorage", + torch.bfloat16: "BFloat16Storage", + torch.cdouble: "ComplexDoubleStorage", + torch.cfloat: "ComplexFloatStorage", + torch.qint8: "QInt8Storage", + torch.qint32: "QInt32Storage", + torch.quint8: "QUInt8Storage", + torch.quint4x2: "QUInt4x2Storage", + torch.quint2x4: "QUInt2x4Storage", + } + + +@functools.cache +def _storage_type_to_dtype_map(): + dtype_map = {val: key for key, val in _dtype_to_storage_type_map().items()} + return dtype_map + + +def _get_storage_from_sequence(sequence, dtype, device): + if dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + tmp_tensor = torch.tensor( + sequence, dtype=interpret_dtypes[dtype], device=device + ) + + else: + tmp_tensor = torch.tensor(sequence, dtype=dtype, device=device) + + return tmp_tensor._typed_storage()._untyped_storage + + +def _isint(x): + if HAS_NUMPY: + return isinstance(x, (int, np.integer)) + else: + return isinstance(x, int) + + +_always_warn_typed_storage_removal = False + + +def _get_always_warn_typed_storage_removal(): + return _always_warn_typed_storage_removal + + +def _set_always_warn_typed_storage_removal(always_warn): + global _always_warn_typed_storage_removal + assert isinstance(always_warn, bool) + _always_warn_typed_storage_removal = always_warn + + +def _warn_typed_storage_removal(stacklevel=2): + global _always_warn_typed_storage_removal + + def is_first_time(): + if not hasattr(_warn_typed_storage_removal, "has_warned"): + return True + else: + return not _warn_typed_storage_removal.__dict__["has_warned"] + + if _get_always_warn_typed_storage_removal() or is_first_time(): + message = ( + "TypedStorage is deprecated. It will be removed in the future and " + "UntypedStorage will be the only storage class. This should only matter " + "to you if you are using storages directly. To access UntypedStorage " + "directly, use tensor.untyped_storage() instead of tensor.storage()" + ) + warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) + _warn_typed_storage_removal.__dict__["has_warned"] = True + + +def _reset_warn_typed_storage_removal(): + _warn_typed_storage_removal.__dict__["has_warned"] = False + + +def _get_device_from_module(module: str): + last_part = module.rsplit(".", 1)[-1] + if last_part in ["cuda", torch._C._get_privateuse1_backend_name(), "hpu"]: + return last_part + else: + return "cpu" + + +class TypedStorage: + is_sparse: _bool = False + # Used when stashing FakeTensor device onto storage in torch.save(metadata_only=True) + _fake_device: _Optional[torch.device] = None + + dtype: torch.dtype + + @property + def _dtype(self): + return self.dtype + + @property + def filename(self) -> _Optional[str]: + """Returns the file name associated with this storage if the storage was memory mapped from a file. + or ``None`` if the storage was not created by memory mapping a file.""" + return self._untyped_storage.filename + + def fill_(self, value): + _warn_typed_storage_removal() + self._setitem(slice(0, self._size()), value) + return self + + def __new__( + cls, + *args, + wrap_storage=None, + dtype=None, + device=None, + _internal=False, + ): + if not _internal: + _warn_typed_storage_removal() + + if cls == torch.storage._LegacyStorage: + raise RuntimeError( + "Only child classes of _LegacyStorage can be instantiated" + ) + + if cls == TypedStorage: + return super().__new__(cls) + + else: + arg_error_msg = ( + f"{cls}.__new__ received an invalid combination " + f"of arguments. Expected one of:\n" + " * no arguments\n" + " * (int size)\n" + " * (Sequence data)\n" + " * (*, UntypedStorage wrap_storage)" + ) + + if device is not None: + raise RuntimeError( + arg_error_msg + "\nKeyword argument 'device' cannot be specified" + ) + + if dtype is not None: + raise RuntimeError( + arg_error_msg + "\nKeyword argument 'dtype' cannot be specified" + ) + + if wrap_storage is None: + if len(args) > 1: + raise RuntimeError( + arg_error_msg + "\nToo many positional arguments" + ) + + if ( + len(args) == 1 + and not _isint(args[0]) + and not isinstance(args[0], collections.abc.Sequence) + ): + raise TypeError( + arg_error_msg + + f"\nArgument type not recognized: {type(args[0])}" + ) + + return TypedStorage( + *args, + dtype=cls._dtype, + device=_get_device_from_module(cls.__module__), + _internal=True, + ) + + else: + if len(args) != 0: + raise RuntimeError( + arg_error_msg + + "\nNo positional arguments should be given when using " + "'wrap_storage'" + ) + + if not isinstance(wrap_storage, torch.UntypedStorage): + raise TypeError( + arg_error_msg + + f"\nArgument 'wrap_storage' must be UntypedStorage, but got {type(wrap_storage)}" + ) + + cls_device = _get_device_from_module(cls.__module__) + + if wrap_storage.device.type != cls_device: + raise RuntimeError( + arg_error_msg + + f"\nDevice of 'wrap_storage' must be {cls_device}" + f", but got {wrap_storage.device.type}" + ) + + return TypedStorage( + *args, + wrap_storage=wrap_storage, + dtype=cls.dtype, + _internal=True, + ) + + def __init__( + self, + *args, + device=None, + dtype=None, + wrap_storage=None, + _internal=False, + ): + if not _internal: + _warn_typed_storage_removal() + arg_error_msg = ( + "TypedStorage.__init__ received an invalid combination " + "of arguments. Expected one of:\n" + " * (*, torch.device device, torch.dtype dtype)\n" + " * (int size, *, torch.device device, torch.dtype dtype)\n" + " * (Sequence data, *, torch.device device, torch.dtype dtype)\n" + " * (*, UntypedStorage wrap_storage, torch.dtype dtype)" + ) + + if wrap_storage is not None: + if len(args) != 0: + raise RuntimeError( + arg_error_msg + + "\nNo positional arguments should be given when using " + "'wrap_storage'" + ) + + if dtype is None: + raise RuntimeError( + arg_error_msg + "\nArgument 'dtype' must be specified" + ) + + if not isinstance(dtype, torch.dtype): + raise TypeError( + arg_error_msg + + f"\nArgument 'dtype' must be torch.dtype, not {type(dtype)}" + ) + + if device is not None: + raise RuntimeError( + arg_error_msg + + "\nArgument 'device' should not be specified when 'wrap_storage' is given" + ) + + self.dtype = dtype + + if not isinstance(wrap_storage, torch.UntypedStorage): + raise TypeError( + arg_error_msg + + f"\nArgument 'wrap_storage' must be UntypedStorage, but got {type(wrap_storage)}" + ) + + self._untyped_storage = wrap_storage + + else: + self.dtype = torch.get_default_dtype() if dtype is None else dtype + device = torch.device("cpu" if device is None else device) + + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + if device.type == "cuda": + raise RuntimeError( + "Cannot create CUDA storage with quantized dtype" + ) + + if len(args) == 0: + self._untyped_storage = torch.UntypedStorage(device=device) + + elif len(args) == 1: + if _isint(args[0]): + self._untyped_storage = torch.UntypedStorage( + int(args[0]) * self._element_size(), device=device + ) + elif isinstance(args[0], collections.abc.Sequence): + self._untyped_storage = _get_storage_from_sequence( + args[0], self.dtype, device + ) + else: + raise TypeError( + arg_error_msg + + f"\nArgument type not recognized: {type(args[0])}" + ) + + else: + raise RuntimeError(arg_error_msg + "\nToo many positional arguments") + + @property + def is_cuda(self): + _warn_typed_storage_removal() + return self._untyped_storage.device.type == "cuda" + + @property + def is_hpu(self): + _warn_typed_storage_removal() + return self._untyped_storage.device.type == "hpu" + + def untyped(self): + """Return the internal :class:`torch.UntypedStorage`.""" + _warn_typed_storage_removal() + return self._untyped_storage + + def _new_wrapped_storage(self, untyped_storage) -> Self: + assert type(untyped_storage) == torch.UntypedStorage + + if type(self) == TypedStorage: + return cast( + Self, + TypedStorage( + wrap_storage=untyped_storage, dtype=self.dtype, _internal=True + ), + ) + else: + return type(self)(wrap_storage=untyped_storage) + + def __len__(self): + _warn_typed_storage_removal() + return self._size() + + def _maybe_wrap_index(self, idx, is_stop=False): + if idx is None: + if is_stop: + return self._size() + else: + return 0 + + else: + if type(idx) != int: + raise TypeError(f"can't index a {type(self)} with {type(idx)}") + if is_stop: + if (idx > self._size()) or (idx < -self._size()): + raise IndexError( + f"index {idx} out of range for storage of size {self.size()}" + ) + if idx > 0: + return idx + else: + return idx % self._size() + else: + if (idx >= self._size()) or (idx < -self._size()): + raise IndexError( + f"index {idx} out of range for storage of size {self.size()}" + ) + return idx % self._size() + + def __setitem__(self, idx, value): + _warn_typed_storage_removal() + return self._setitem(idx, value) + + def _setitem(self, idx, value): + if not isinstance(idx, (int, slice)): + raise RuntimeError(f"can't index a {type(self)} with {type(idx)}") + if torch.is_storage(value): + raise RuntimeError(f"cannot set item with value type {type(value)}") + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + tmp_dtype = interpret_dtypes[self.dtype] + tmp_tensor = torch.tensor( + [], dtype=tmp_dtype, device=self._untyped_storage.device + ) + tmp_tensor.set_( + TypedStorage( + wrap_storage=self._untyped_storage, dtype=tmp_dtype, _internal=True + ) + ) + else: + tmp_tensor = torch.tensor( + [], dtype=self.dtype, device=self._untyped_storage.device + ).set_(self) + + tmp_tensor[idx] = value + + def __getitem__(self, idx): + _warn_typed_storage_removal() + return self._getitem(idx) + + def _getitem(self, idx): + if self._untyped_storage.device.type == "meta": + raise NotImplementedError("Not available for 'meta' device type") + + # NOTE: Before TypedStorage existed, indexing with a slice used to be + # possible for Storage objects. However, it would return + # a storage view, which would be a hassle to implement in TypedStorage, + # so it was disabled + if isinstance(idx, slice): + raise RuntimeError( + "slices are only supported in UntypedStorage.__getitem__" + ) + elif not isinstance(idx, int): + raise RuntimeError(f"can't index a {type(self)} with {type(idx)}") + + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + return TypedStorage( + wrap_storage=self._untyped_storage, + dtype=interpret_dtypes[self.dtype], + _internal=True, + )._getitem(idx) + + idx_wrapped = self._maybe_wrap_index(idx) + from torch._subclasses.fake_tensor import unset_fake_temporarily + + with unset_fake_temporarily(): + tmp_tensor = torch.tensor( + [], dtype=self.dtype, device=self._untyped_storage.device + ).set_(self) + return tmp_tensor[idx_wrapped].item() + + def copy_(self, source: T, non_blocking: _Optional[bool] = None): + _warn_typed_storage_removal() + if isinstance(source, TypedStorage): + self._untyped_storage.copy_(source._untyped_storage, non_blocking) + else: + self._untyped_storage.copy_(source, non_blocking) + return self + + def nbytes(self): + _warn_typed_storage_removal() + return self._nbytes() + + # For internal use only, to avoid deprecation warning + def _nbytes(self): + return self._untyped_storage.nbytes() + + def type( + self, + dtype: _Optional[str] = None, + non_blocking: bool = False, + ) -> Union[_StorageBase, TypedStorage, str]: + _warn_typed_storage_removal() + if dtype is None: + legacy_class = self._get_legacy_storage_class() + + if legacy_class is not None: + return legacy_class.__module__ + "." + legacy_class.__name__ + + return ".".join([self.__module__, type(self).__name__]) + + else: + return self._untyped_storage.type(dtype, non_blocking) + + def cuda(self, device=None, non_blocking=False) -> Self: + _warn_typed_storage_removal() + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError("Cannot create CUDA storage with quantized dtype") + cuda_storage = self._untyped_storage.cuda(device, non_blocking) + return self._new_wrapped_storage(cuda_storage) + + def hpu(self, device=None, non_blocking=False) -> Self: + _warn_typed_storage_removal() + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError("Cannot create HPU storage with quantized dtype") + hpu_storage = self._untyped_storage.hpu(device, non_blocking) + return self._new_wrapped_storage(hpu_storage) + + def to(self, *, device: DeviceLikeType, non_blocking: bool = False) -> Self: + _warn_typed_storage_removal() + if not isinstance(device, torch.device): + device = torch.device(device) + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError( + f"Cannot create {device.type.upper()} storage with quantized dtype" + ) + to_storage = self._untyped_storage.to(device=device, non_blocking=non_blocking) + return self._new_wrapped_storage(to_storage) + + def element_size(self): + _warn_typed_storage_removal() + return self._element_size() + + # For internal use only, to avoid deprecation warning + def _element_size(self): + return torch._utils._element_size(self.dtype) + + def get_device(self) -> _int: + _warn_typed_storage_removal() + return self._untyped_storage.get_device() + + def __str__(self): + _warn_typed_storage_removal() + info_str = ( + f"[{torch.typename(self)}(dtype={self.dtype}, " + f"device={self.device}) of size {len(self)}]" + ) + if self.device.type == "meta": + return "...\n" + info_str + else: + data_str = " " + "\n ".join(str(self[i]) for i in range(self.size())) + return data_str + "\n" + info_str + + def __repr__(self): + _warn_typed_storage_removal() + return str(self) + + def __iter__(self): + _warn_typed_storage_removal() + return iter(self[i] for i in range(self.size())) + + def __copy__(self): + _warn_typed_storage_removal() + return self._new_wrapped_storage(copy.copy(self._untyped_storage)) + + def __deepcopy__(self, memo): + _warn_typed_storage_removal() + return self._deepcopy(memo) + + # For internal use only, to avoid deprecation warning + def _deepcopy(self, memo): + return self._new_wrapped_storage(copy.deepcopy(self._untyped_storage, memo)) + + def __sizeof__(self): + _warn_typed_storage_removal() + return super().__sizeof__() + self.nbytes() + + def clone(self): + """Return a copy of this storage.""" + _warn_typed_storage_removal() + return self._new_wrapped_storage(self._untyped_storage.clone()) + + def tolist(self): + """Return a list containing the elements of this storage.""" + _warn_typed_storage_removal() + return list(self) + + def cpu(self): + """Return a CPU copy of this storage if it's not already on the CPU.""" + _warn_typed_storage_removal() + return self._new_wrapped_storage(self._untyped_storage.cpu()) + + def is_pinned(self, device: Union[str, torch.device] = "cuda"): + r"""Determine whether the CPU TypedStorage is already pinned on device. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A boolean variable. + """ + _warn_typed_storage_removal() + return self._untyped_storage.is_pinned(device) + + def pin_memory(self, device: Union[str, torch.device] = "cuda"): + r"""Copy the CPU TypedStorage to pinned memory, if it's not already pinned. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A pinned CPU storage. + """ + _warn_typed_storage_removal() + return self._new_wrapped_storage( + self._untyped_storage.pin_memory(device=device) + ) + + def share_memory_(self): + """See :meth:`torch.UntypedStorage.share_memory_`""" + _warn_typed_storage_removal() + return self._share_memory_() + + # For internal use only, to avoid deprecation warning + def _share_memory_(self): + self._untyped_storage.share_memory_() + return self + + def _new_shared(self, size, *, device=None): + """Create a new storage in shared memory with the same data type.""" + if device is None: + device = "cpu" + device = torch.device(device) + untyped_storage = torch.UntypedStorage._new_shared( + size * self._element_size(), device=device + ) + return TypedStorage( + wrap_storage=untyped_storage, dtype=self.dtype, _internal=True + ) + + @property + def _cdata(self): + return self._untyped_storage._cdata + + @property + def device(self): + _warn_typed_storage_removal() + return self._untyped_storage.device + + def size(self): + _warn_typed_storage_removal() + return self._size() + + # For internal use only, to avoid deprecation warning + def _size(self): + # NB: don't indirect through __len__, as that requires + # an int to be returned + return self._untyped_storage.nbytes() // self._element_size() + + def pickle_storage_type(self): + _warn_typed_storage_removal() + return self._pickle_storage_type() + + # For internal use only, to avoid deprecation warning + def _pickle_storage_type(self): + try: + return _dtype_to_storage_type_map()[self.dtype] + except KeyError as e: + raise KeyError(f"dtype {self.dtype} is not recognized") from e + + def __reduce__(self): + b = io.BytesIO() + torch.save(self, b, _use_new_zipfile_serialization=False) + return (_load_from_bytes, (b.getvalue(),)) + + def data_ptr(self): + _warn_typed_storage_removal() + return self._data_ptr() + + # For internal use only, to avoid deprecation warning + def _data_ptr(self): + return self._untyped_storage.data_ptr() + + def resizable(self): + _warn_typed_storage_removal() + return self._untyped_storage.resizable() + + def resize_(self, size): + _warn_typed_storage_removal() + self._resize_(size) + + # For internal use only, to avoid deprecation warning + def _resize_(self, size): + self._untyped_storage.resize_(size * self._element_size()) + + @classmethod + def _free_weak_ref(cls, *args, **kwargs): + return UntypedStorage._free_weak_ref(*args, **kwargs) + + def _weak_ref(self, *args, **kwargs): + return self._untyped_storage._weak_ref(*args, **kwargs) + + @classmethod + def from_buffer(cls, *args, **kwargs): + _warn_typed_storage_removal() + return cls._from_buffer(*args, **kwargs) + + @classmethod + def _from_buffer(cls, *args, dtype=None, device=None, **kwargs): + if cls == TypedStorage: + dtype = torch.get_default_dtype() if dtype is None else dtype + device = torch.device("cpu" if device is None else device) + if device.type != "cpu": + raise RuntimeError( + f"TypedStorage.from_buffer: Not available for device {device.type}" + ) + untyped_storage: torch.UntypedStorage = torch.UntypedStorage.from_buffer( + *args, dtype=dtype, **kwargs + ) + + else: + if dtype is not None or len(args) == 5: + raise RuntimeError( + "from_buffer: 'dtype' can only be specified in " + "UntypedStorage.from_buffer and TypedStorage.from_buffer" + ) + if device is not None: + raise RuntimeError( + "from_buffer: 'device' can only be specified in " + "UntypedStorage.from_buffer and TypedStorage.from_buffer" + ) + + dtype = cls._dtype + untyped_storage = torch.UntypedStorage.from_buffer( + *args, dtype=dtype, **kwargs + ) + + return TypedStorage(wrap_storage=untyped_storage, dtype=dtype, _internal=True) + + def _to(self, dtype): + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Argument 'dtype' must be torch.dtype, not {type(dtype)}") + storage = ( + torch.tensor([], dtype=self.dtype, device=self.device) + .set_(self) + .to(dtype) + ._typed_storage() + ) + if storage.data_ptr() == self.data_ptr(): + storage = storage.clone() + return storage + + def double(self): + """Casts this storage to double type.""" + _warn_typed_storage_removal() + return self._to(torch.double) + + def float(self): + """Casts this storage to float type.""" + _warn_typed_storage_removal() + return self._to(torch.float) + + def half(self): + """Casts this storage to half type.""" + _warn_typed_storage_removal() + return self._to(torch.half) + + def long(self): + """Casts this storage to long type.""" + _warn_typed_storage_removal() + return self._to(torch.long) + + def int(self): + """Casts this storage to int type.""" + _warn_typed_storage_removal() + return self._to(torch.int) + + def short(self): + """Casts this storage to short type.""" + _warn_typed_storage_removal() + return self._to(torch.short) + + def char(self): + """Casts this storage to char type.""" + _warn_typed_storage_removal() + return self._to(torch.int8) + + def byte(self): + """Casts this storage to byte type.""" + _warn_typed_storage_removal() + return self._to(torch.uint8) + + def bool(self): + """Casts this storage to bool type.""" + _warn_typed_storage_removal() + return self._to(torch.bool) + + def bfloat16(self): + """Casts this storage to bfloat16 type.""" + _warn_typed_storage_removal() + return self._to(torch.bfloat16) + + def complex_double(self): + """Casts this storage to complex double type.""" + _warn_typed_storage_removal() + return self._to(torch.cdouble) + + def complex_float(self): + """Casts this storage to complex float type.""" + _warn_typed_storage_removal() + return self._to(torch.cfloat) + + def float8_e5m2(self): + """Casts this storage to float8_e5m2 type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e5m2) + + def float8_e4m3fn(self): + """Casts this storage to float8_e4m3fn type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e4m3fn) + + def float8_e5m2fnuz(self): + """Casts this storage to float8_e5m2fnuz type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e5m2fnuz) + + def float8_e4m3fnuz(self): + """Casts this storage to float8_e4m3fnuz type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e4m3fnuz) + + @classmethod + def from_file(cls, filename, shared, size): + """from_file(filename, shared=False, size=0) -> Storage + + Creates a CPU storage backed by a memory-mapped file. + + If ``shared`` is ``True``, then memory is shared between all processes. + All changes are written to the file. If ``shared`` is ``False``, then the changes on + the storage do not affect the file. + + ``size`` is the number of elements in the storage. If ``shared`` is ``False``, + then the file must contain at least ``size * sizeof(Type)`` bytes + (``Type`` is the type of storage). If ``shared`` is ``True`` the file will be created if needed. + + Args: + filename (str): file name to map + shared (bool): whether to share memory (whether ``MAP_SHARED`` or ``MAP_PRIVATE`` is passed to the + underlying `mmap(2) call `_) + size (int): number of elements in the storage + """ + _warn_typed_storage_removal() + if cls == TypedStorage: + raise RuntimeError("from_file can only be called on derived classes") + untyped_storage = UntypedStorage.from_file( + filename, shared, size * torch._utils._element_size(cls.dtype) + ) + storage = cls(wrap_storage=untyped_storage) + return storage + + @classmethod + def _expired(cls, *args, **kwargs): + return UntypedStorage._expired(*args, **kwargs) + + def _write_file(self, *args, **kwargs): + return self._untyped_storage._write_file(*args, **kwargs) + + def _set_from_file(self, *args, **kwargs): + return self._untyped_storage._set_from_file(*args, **kwargs) + + def _set_cdata(self, *args, **kwargs): + return self._untyped_storage._set_cdata(*args, **kwargs) + + def _share_cuda_(self, *args, **kwargs): + return self._untyped_storage._share_cuda_(*args, **kwargs) + + def is_shared(self): + _warn_typed_storage_removal() + return self._is_shared() + + # For internal use only, to avoid deprecation warning + def _is_shared(self): + return self._untyped_storage.is_shared() + + @classmethod + def _new_shared_cuda(cls, *args, **kwargs): + return torch.UntypedStorage._new_shared_cuda(*args, **kwargs) + + def _share_filename_cpu_(self, *args, **kwargs): + ( + manager_handle, + storage_handle, + size, + ) = self._untyped_storage._share_filename_cpu_(*args, **kwargs) + return manager_handle, storage_handle, size // self._element_size() + + def _shared_decref(self): + self._untyped_storage._shared_decref() + return self + + @classmethod + def _release_ipc_counter(cls, *args, device=None, **kwargs): + return torch.UntypedStorage._release_ipc_counter_cuda(*args, **kwargs) + + def _shared_incref(self, *args, **kwargs): + return self._untyped_storage._shared_incref(*args, **kwargs) + + def _share_fd_cpu_(self, *args, **kwargs): + fd, size = self._untyped_storage._share_fd_cpu_(*args, **kwargs) + return fd, size // self._element_size() + + def _get_legacy_storage_class(self): + if self.dtype not in _dtype_to_storage_type_map(): + return None + + storage_name = _dtype_to_storage_type_map()[self.dtype] + + if self.device.type not in [ + "cpu", + "cuda", + "hpu", + torch._C._get_privateuse1_backend_name(), + ]: + return None + + module = ( + torch if self.device.type == "cpu" else getattr(torch, self.device.type) + ) + + try: + return getattr(module, storage_name) + except AttributeError: + return None + + +TypedStorage.type.__doc__ = _type.__doc__ +TypedStorage.cuda.__doc__ = _StorageBase.cuda.__doc__ +TypedStorage.hpu.__doc__ = _StorageBase.hpu.__doc__ +TypedStorage.to.__doc__ = _to.__doc__ + + +class _LegacyStorageMeta(type): + dtype: torch.dtype + + def __instancecheck__(cls, instance): + if type(instance) == TypedStorage: + cls_device = _get_device_from_module(cls.__module__) + return (cls_device == instance.device.type) and ( + cls.dtype == instance.dtype + ) + return False + + +class _LegacyStorage(TypedStorage, metaclass=_LegacyStorageMeta): + @classmethod + def _new_shared(cls, size): # type: ignore[override] + """Create a new storage in shared memory with the same data type.""" + untyped_storage = torch.UntypedStorage._new_shared(size * cls()._element_size()) + return cls(wrap_storage=untyped_storage) + + @classmethod + def _release_ipc_counter(cls, *args, **kwargs): + return torch.UntypedStorage._release_ipc_counter_cuda(*args, **kwargs) + + @classmethod + def _new_shared_filename(cls, manager, obj, size): + bytes_size = size * torch._utils._element_size(cls.dtype) + return cls( + wrap_storage=torch.UntypedStorage._new_shared_filename_cpu( + manager, obj, bytes_size + ) + ) + + +def _get_dtype_from_pickle_storage_type(pickle_storage_type: str): + try: + return _storage_type_to_dtype_map()[pickle_storage_type] + except KeyError as e: + raise KeyError( + f'pickle storage type "{pickle_storage_type}" is not recognized' + ) from e diff --git a/venv/lib/python3.10/site-packages/torch/torch_version.py b/venv/lib/python3.10/site-packages/torch/torch_version.py new file mode 100644 index 0000000000000000000000000000000000000000..0496a1b564feefe4a52280e2d7f268516f256a70 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/torch_version.py @@ -0,0 +1,66 @@ +from collections.abc import Iterable +from typing import Any + +from torch._vendor.packaging.version import InvalidVersion, Version +from torch.version import __version__ as internal_version + + +__all__ = ["TorchVersion"] + + +class TorchVersion(str): + """A string with magic powers to compare to both Version and iterables! + Prior to 1.10.0 torch.__version__ was stored as a str and so many did + comparisons against torch.__version__ as if it were a str. In order to not + break them we have TorchVersion which masquerades as a str while also + having the ability to compare against both packaging.version.Version as + well as tuples of values, eg. (1, 2, 1) + Examples: + Comparing a TorchVersion object to a Version object + TorchVersion('1.10.0a') > Version('1.10.0a') + Comparing a TorchVersion object to a Tuple object + TorchVersion('1.10.0a') > (1, 2) # 1.2 + TorchVersion('1.10.0a') > (1, 2, 1) # 1.2.1 + Comparing a TorchVersion object against a string + TorchVersion('1.10.0a') > '1.2' + TorchVersion('1.10.0a') > '1.2.1' + """ + + __slots__ = () + + # fully qualified type names here to appease mypy + def _convert_to_version(self, inp: Any) -> Any: + if isinstance(inp, Version): + return inp + elif isinstance(inp, str): + return Version(inp) + elif isinstance(inp, Iterable): + # Ideally this should work for most cases by attempting to group + # the version tuple, assuming the tuple looks (MAJOR, MINOR, ?PATCH) + # Examples: + # * (1) -> Version("1") + # * (1, 20) -> Version("1.20") + # * (1, 20, 1) -> Version("1.20.1") + return Version(".".join(str(item) for item in inp)) + else: + raise InvalidVersion(inp) + + def _cmp_wrapper(self, cmp: Any, method: str) -> bool: + try: + return getattr(Version(self), method)(self._convert_to_version(cmp)) + except BaseException as e: + if not isinstance(e, InvalidVersion): + raise + # Fall back to regular string comparison if dealing with an invalid + # version like 'parrot' + return getattr(super(), method)(cmp) + + +for cmp_method in ["__gt__", "__lt__", "__eq__", "__ge__", "__le__"]: + setattr( + TorchVersion, + cmp_method, + lambda x, y, method=cmp_method: x._cmp_wrapper(y, method), + ) + +__version__ = TorchVersion(internal_version) diff --git a/venv/lib/python3.10/site-packages/torch/types.py b/venv/lib/python3.10/site-packages/torch/types.py new file mode 100644 index 0000000000000000000000000000000000000000..ab6f4639f444153c00f4e03a63d1b45561dad786 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/types.py @@ -0,0 +1,130 @@ +# In some cases, these basic types are shadowed by corresponding +# top-level values. The underscore variants let us refer to these +# types. See https://github.com/python/mypy/issues/4146 for why these +# workarounds is necessary +import os +from builtins import ( # noqa: F401 + bool as _bool, + bytes as _bytes, + complex as _complex, + float as _float, + int as _int, + str as _str, +) +from collections.abc import Sequence +from typing import Any, IO, TYPE_CHECKING, Union +from typing_extensions import Self, TypeAlias + +# `as` imports have better static analysis support than assignment `ExposedType: TypeAlias = HiddenType` +from torch import ( # noqa: F401 + device as _device, + DispatchKey as DispatchKey, + dtype as _dtype, + layout as _layout, + qscheme as _qscheme, + Size as Size, + SymBool as SymBool, + SymFloat as SymFloat, + SymInt as SymInt, + Tensor as Tensor, +) + + +if TYPE_CHECKING: + from torch.autograd.graph import GradientEdge + + +__all__ = ["Number", "Device", "FileLike", "Storage"] + +# Convenience aliases for common composite types that we need +# to talk about in PyTorch +_TensorOrTensors: TypeAlias = Union[Tensor, Sequence[Tensor]] # noqa: PYI047 +_TensorOrTensorsOrGradEdge: TypeAlias = Union[ # noqa: PYI047 + Tensor, + Sequence[Tensor], + "GradientEdge", + Sequence["GradientEdge"], +] + +_size: TypeAlias = Union[Size, list[int], tuple[int, ...]] # noqa: PYI042,PYI047 +_symsize: TypeAlias = Union[Size, Sequence[Union[int, SymInt]]] # noqa: PYI042,PYI047 +_dispatchkey: TypeAlias = Union[str, DispatchKey] # noqa: PYI042,PYI047 + +# int or SymInt +IntLikeType: TypeAlias = Union[int, SymInt] +# float or SymFloat +FloatLikeType: TypeAlias = Union[float, SymFloat] +# bool or SymBool +BoolLikeType: TypeAlias = Union[bool, SymBool] + +py_sym_types = (SymInt, SymFloat, SymBool) # left un-annotated intentionally +PySymType: TypeAlias = Union[SymInt, SymFloat, SymBool] + +# Meta-type for "numeric" things; matches our docs +Number: TypeAlias = Union[int, float, bool] +# tuple for isinstance(x, Number) checks. +# FIXME: refactor once python 3.9 support is dropped. +_Number = (int, float, bool) + +FileLike: TypeAlias = Union[str, os.PathLike[str], IO[bytes]] + +# Meta-type for "device-like" things. Not to be confused with 'device' (a +# literal device object). This nomenclature is consistent with PythonArgParser. +# None means use the default device (typically CPU) +Device: TypeAlias = Union[_device, str, int, None] + + +# Storage protocol implemented by ${Type}StorageBase classes +class Storage: + _cdata: int + device: _device + dtype: _dtype + _torch_load_uninitialized: bool + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + raise NotImplementedError + + def _new_shared(self, size: int) -> Self: + raise NotImplementedError + + def _write_file( + self, + f: Any, + is_real_file: bool, + save_size: bool, + element_size: int, + ) -> None: + raise NotImplementedError + + def element_size(self) -> int: + raise NotImplementedError + + def is_shared(self) -> bool: + raise NotImplementedError + + def share_memory_(self) -> Self: + raise NotImplementedError + + def nbytes(self) -> int: + raise NotImplementedError + + def cpu(self) -> Self: + raise NotImplementedError + + def data_ptr(self) -> int: + raise NotImplementedError + + def from_file( + self, + filename: str, + shared: bool = False, + nbytes: int = 0, + ) -> Self: + raise NotImplementedError + + def _new_with_file( + self, + f: Any, + element_size: int, + ) -> Self: + raise NotImplementedError diff --git a/venv/lib/python3.10/site-packages/torch/version.py b/venv/lib/python3.10/site-packages/torch/version.py new file mode 100644 index 0000000000000000000000000000000000000000..355bbe1201d819af9cdf37c4d00b266907511238 --- /dev/null +++ b/venv/lib/python3.10/site-packages/torch/version.py @@ -0,0 +1,9 @@ +from typing import Optional + +__all__ = ['__version__', 'debug', 'cuda', 'git_version', 'hip', 'xpu'] +__version__ = '2.9.0+cu128' +debug = False +cuda: Optional[str] = '12.8' +git_version = '0fabc3ba44823f257e70ce397d989c8de5e362c1' +hip: Optional[str] = None +xpu: Optional[str] = None diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..b42d49f2fbe1ae4f0052ecdb6857c6a2054351eb --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/METADATA @@ -0,0 +1,46 @@ +Metadata-Version: 2.4 +Name: triton +Version: 3.5.0 +Summary: A language and compiler for custom Deep Learning operations +Home-page: https://github.com/triton-lang/triton/ +Author: Philippe Tillet +Author-email: phil@openai.com +Keywords: Compiler,Deep Learning +Classifier: Development Status :: 4 - Beta +Classifier: Intended Audience :: Developers +Classifier: Topic :: Software Development :: Build Tools +Classifier: License :: OSI Approved :: MIT License +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Requires-Python: >=3.10,<3.15 +License-File: LICENSE +Requires-Dist: importlib-metadata; python_version < "3.10" +Provides-Extra: build +Requires-Dist: cmake<4.0,>=3.20; extra == "build" +Requires-Dist: lit; extra == "build" +Provides-Extra: tests +Requires-Dist: autopep8; extra == "tests" +Requires-Dist: isort; extra == "tests" +Requires-Dist: numpy; extra == "tests" +Requires-Dist: pytest; extra == "tests" +Requires-Dist: pytest-forked; extra == "tests" +Requires-Dist: pytest-xdist; extra == "tests" +Requires-Dist: scipy>=1.7.1; extra == "tests" +Requires-Dist: llnl-hatchet; extra == "tests" +Provides-Extra: tutorials +Requires-Dist: matplotlib; extra == "tutorials" +Requires-Dist: pandas; extra == "tutorials" +Requires-Dist: tabulate; extra == "tutorials" +Dynamic: author +Dynamic: author-email +Dynamic: classifier +Dynamic: home-page +Dynamic: keywords +Dynamic: license-file +Dynamic: provides-extra +Dynamic: requires-dist +Dynamic: requires-python +Dynamic: summary diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..f62aedd17cfc407a99b8a68f17884f07a6204712 --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/RECORD @@ -0,0 +1,490 @@ +../../../bin/proton,sha256=Mo2gh0yoNuBlNh3BkOGVggNdezkIbR0t2hRe1m4s3o4,219 +../../../bin/proton-viewer,sha256=D8Q-U_SQJ6TjH6eUXCaVbfRglfE9WbcxrHZblXdGFdY,219 +triton-3.5.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +triton-3.5.0.dist-info/METADATA,sha256=y2dB0gBbNg9cyGmlGPoS9op5mG_AwmXeo-LKFPaTWbs,1663 +triton-3.5.0.dist-info/RECORD,, +triton-3.5.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton-3.5.0.dist-info/WHEEL,sha256=Obtqci3x5vy5ZivY2BiOH9GHO8-xQ0d4HFhTOtXMNhw,152 +triton-3.5.0.dist-info/entry_points.txt,sha256=js6289FRk18_JLmbNfGsmtKP77D6aYawJOSIHv3X89I,176 +triton-3.5.0.dist-info/licenses/LICENSE,sha256=kmQPuXIi_Qppj_KM4MN4LBcmI_jWxgm1V2NqgPKPuUY,1132 +triton-3.5.0.dist-info/top_level.txt,sha256=WBiIZyv6n9Y7MIh-HPHSv2w1RDk7EFL__7ZgQRrmHYs,7 +triton/FileCheck,sha256=cfnT9gRYpNgu4gzEv2s7ofOpCOyJy4y5wI-5xTiB3AU,1087840 +triton/_C/libproton.so,sha256=4hDgepHFfLFvU3Onc62msjBSeFxIRW4JFnGxA81A4aA,18356176 +triton/_C/libtriton.so,sha256=Z8hV39vZChZ2rC9938Yva6HzR37Nv2KPAgirY05KJx8,400626128 +triton/__init__.py,sha256=3k-6i-O8Fi4127sAdQdTs0yN6rNV_blAuYY2lIUQzq4,1661 +triton/__pycache__/__init__.cpython-310.pyc,, +triton/__pycache__/_filecheck.cpython-310.pyc,, +triton/__pycache__/_internal_testing.cpython-310.pyc,, +triton/__pycache__/_utils.cpython-310.pyc,, +triton/__pycache__/errors.cpython-310.pyc,, +triton/__pycache__/knobs.cpython-310.pyc,, +triton/__pycache__/testing.cpython-310.pyc,, +triton/_filecheck.py,sha256=2i88UORAasMECwAcW5NlCfJbYae7cVd2zOhAoGjI_uc,3164 +triton/_internal_testing.py,sha256=LMdOSxV1JUDQQSqV__YsfcDab0E5yrUm0KSwPLghCp8,8405 +triton/_utils.py,sha256=JjqxiYSZ4GSNTkOHWS6tJ6ME114bGahRpQfuH79Dv08,3595 +triton/backends/__init__.py,sha256=X7290kf96Fk9QnfLScsX4UDG3zPyH_-31E4A7pVOijM,1612 +triton/backends/__pycache__/__init__.cpython-310.pyc,, +triton/backends/__pycache__/compiler.cpython-310.pyc,, +triton/backends/__pycache__/driver.cpython-310.pyc,, +triton/backends/amd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/backends/amd/__pycache__/__init__.cpython-310.pyc,, +triton/backends/amd/__pycache__/compiler.cpython-310.pyc,, +triton/backends/amd/__pycache__/driver.cpython-310.pyc,, +triton/backends/amd/compiler.py,sha256=A7_8wWyda-Xqj95kXK86bjwKyJLuxJBRp51MyZtmucc,20874 +triton/backends/amd/driver.c,sha256=dEpw00lFUSMSvQ5AnNc_4XDAZIHVXvUlJ5K-LmAxkvU,11149 +triton/backends/amd/driver.py,sha256=5oUFCY7vPnwLBQy5HaJYerbBDUVcoLlB_wb6gGjhYwU,26951 +triton/backends/amd/include/hip/amd_detail/amd_channel_descriptor.h,sha256=_2myGIdBTE0plFbGKOSx8HUqGZd0UBHo-YvKe2xkpbU,11708 +triton/backends/amd/include/hip/amd_detail/amd_device_functions.h,sha256=1TwUYh4RJhoiz850vLNicO1fFxFHR-vC-16ttDGDGGw,31367 +triton/backends/amd/include/hip/amd_detail/amd_hip_atomic.h,sha256=KfiNN5FjV6mqSAOam6VJDu9tI7l8uaRUrQ5EMiP5Ymw,50830 +triton/backends/amd/include/hip/amd_detail/amd_hip_bf16.h,sha256=AWsz68vSfi9PnCsIIMOQkCStRLbEat7f6om9_J2Ahj4,61352 +triton/backends/amd/include/hip/amd_detail/amd_hip_bfloat16.h,sha256=cFJlQEELGau_9geACeuiiFHyuAWCD6-VuSqcTnqajX0,9484 +triton/backends/amd/include/hip/amd_detail/amd_hip_common.h,sha256=dzkuIzuklqTRaNJjKLqfFEm6Fh4tK_FkTjYHFsZkmCI,1370 +triton/backends/amd/include/hip/amd_detail/amd_hip_complex.h,sha256=SEygl8X_MCXDVXxNIBm5Ds0eWwa-ojVXUUW48SIgsX8,5855 +triton/backends/amd/include/hip/amd_detail/amd_hip_cooperative_groups.h,sha256=t8rqzpZS--LGLI1D-pNBNCp7TCKGS0H136ps6m6VOlI,31861 +triton/backends/amd/include/hip/amd_detail/amd_hip_fp16.h,sha256=86Nw97iaiC4QV5xBv8d3Bwc4FioMh5DQuCHj3sh_Yrw,57854 +triton/backends/amd/include/hip/amd_detail/amd_hip_fp8.h,sha256=wc1l471D-BBroLMoV1uEsPXTp-zQjSpq-jF8bqeNZX8,53765 +triton/backends/amd/include/hip/amd_detail/amd_hip_gl_interop.h,sha256=djlpeoEOqtX0gVVRgqqakQZkwNzLudnK_ixbShizU7M,3861 +triton/backends/amd/include/hip/amd_detail/amd_hip_math_constants.h,sha256=u1fIaf-AiWF70ZA1zxVkUIbRqoJLu5lrfYbgt_usySk,5890 +triton/backends/amd/include/hip/amd_detail/amd_hip_runtime.h,sha256=ZvDsQ0AiZnJ178NuAsA7AuHrySXbN3aFs5Z9m2tsIDg,13954 +triton/backends/amd/include/hip/amd_detail/amd_hip_runtime_pt_api.h,sha256=fc4mtHBkWmiSRh8m-dxIxvu9zsweLTwEgohkntYcgJw,9997 +triton/backends/amd/include/hip/amd_detail/amd_hip_unsafe_atomics.h,sha256=w9nJ1S32GRl_ejDiGacteM6Zf84iovIifAzWX8Bze0Q,24202 +triton/backends/amd/include/hip/amd_detail/amd_hip_vector_types.h,sha256=qPdmRJnzlgtjVshkafoHxdHoMLkoYS9U-ZD-TjLznr0,57088 +triton/backends/amd/include/hip/amd_detail/amd_math_functions.h,sha256=46wiaEMStCczEsHtccgHlATfw_0O5j6Z8rlFkC7bmUA,3171 +triton/backends/amd/include/hip/amd_detail/amd_surface_functions.h,sha256=rsQuylNqmNhLb7PZjBz7WbruD_6YIXtOptY2BNJDxVU,11062 +triton/backends/amd/include/hip/amd_detail/amd_warp_functions.h,sha256=-uJy_hBwSxRA9gzGp4UEZ3co8D_UHVNoMr_Rvx19qik,20042 +triton/backends/amd/include/hip/amd_detail/amd_warp_sync_functions.h,sha256=2B25tjpJ_KJGijRBaEy_a2_4HLt_TXQ6eJ57zRNoPOw,11706 +triton/backends/amd/include/hip/amd_detail/concepts.hpp,sha256=7EOkpr2w2-jclUQ115yxtFCkBWJ7btUzhBOe-mR0N0M,1252 +triton/backends/amd/include/hip/amd_detail/device_library_decls.h,sha256=4clSpgf898UVjfZFVnDkcYi75A27crPsuFtLcs1s4KU,7457 +triton/backends/amd/include/hip/amd_detail/functional_grid_launch.hpp,sha256=u7hRB9kQXX575a5C7cV3gKow55DSBUCwO0dTjIswlag,8129 +triton/backends/amd/include/hip/amd_detail/grid_launch.h,sha256=tNS7CQw9gy-z930CElH3n6c5iMvpsQ_WFZK024mNzEo,1830 +triton/backends/amd/include/hip/amd_detail/grid_launch.hpp,sha256=EuAlM3olyrArebqwW5eSxo4gfjvWCGOAGAuLLmFttgw,1370 +triton/backends/amd/include/hip/amd_detail/grid_launch_GGL.hpp,sha256=KpQAuyy1Dyt45WcPaR_x-Ex-onPGEHA01DBbla7TT-k,1219 +triton/backends/amd/include/hip/amd_detail/helpers.hpp,sha256=hi2pW1mXQnbIwvmwWt_nG6A38sqLOd-QP5S9sETTs60,5707 +triton/backends/amd/include/hip/amd_detail/hip_api_trace.hpp,sha256=UMsCaqurHFSRNiHpwjaA8DExAldhh4y0UP_KOxLa25Y,101025 +triton/backends/amd/include/hip/amd_detail/hip_assert.h,sha256=fNsG23KISuY-k5JFoX-5hZ7qGQScisXuHcdEwYlXOqw,3978 +triton/backends/amd/include/hip/amd_detail/hip_cooperative_groups_helper.h,sha256=tQ_XIvGKhvrj1h7gY-IVLmKvIPhsQa0YsBflxdhUHP8,7957 +triton/backends/amd/include/hip/amd_detail/hip_fp16_gcc.h,sha256=BtFsKmTptN4TOHocEicfNbBl2JCdZWKm_bd5mc5OzYY,6660 +triton/backends/amd/include/hip/amd_detail/hip_fp16_math_fwd.h,sha256=63tKWMPdW56qWlH_HbCaF_isVXufm514ol_SxL4YjTQ,5134 +triton/backends/amd/include/hip/amd_detail/hip_ldg.h,sha256=KAEZb9H4z4DDrkaloMOeWzahiDfI2V6c68vWT3jb5fU,3652 +triton/backends/amd/include/hip/amd_detail/hip_prof_str.h,sha256=Xb0M9ztYHRqRLjKFBumbAR2rjbUnoZulweCZKTU8ejU,648748 +triton/backends/amd/include/hip/amd_detail/hip_runtime_prof.h,sha256=eZUI5cscYKb9xC9IqnaY53Jgy1dQjkERwQaw034Bt5g,2758 +triton/backends/amd/include/hip/amd_detail/host_defines.h,sha256=ZyCvaPsp4vcFiEEIt2BhapfYVvWoVI_Q2P9rZHlFmCY,7197 +triton/backends/amd/include/hip/amd_detail/hsa_helpers.hpp,sha256=Os-sJQOFI_0Abh8Ql05s0Rtfruk4NsSMfg7BtugxMgg,3232 +triton/backends/amd/include/hip/amd_detail/macro_based_grid_launch.hpp,sha256=6ocsArNa9_R6D6XCuNy8Zq23KG-j2uYsjqNCtnMrJws,67925 +triton/backends/amd/include/hip/amd_detail/math_fwd.h,sha256=nup5YhceJnngoLJCESI8qX08dNpbZci0i78WKu-wfdI,17000 +triton/backends/amd/include/hip/amd_detail/ockl_image.h,sha256=LzRPGMb515_iIAIIcbb2uQB-bTvT4xOjY51VdARD7lc,10538 +triton/backends/amd/include/hip/amd_detail/program_state.hpp,sha256=8QE9OmB8OKTy7rBr3EYEizJI2s-_1tgXpgU7zCA2Ky0,3154 +triton/backends/amd/include/hip/amd_detail/texture_fetch_functions.h,sha256=Ex1lF2gBWJxtC3yP9pXRSFywMp3gbEmyl0Sw8iL91yM,17787 +triton/backends/amd/include/hip/amd_detail/texture_indirect_functions.h,sha256=KkW5o5gMpoVMTRwzfXHA7-kZ9ynI8OaIw6jJ1EB1s98,18447 +triton/backends/amd/include/hip/channel_descriptor.h,sha256=gTYe7SzIg-m3ThOQY2vr5Rh6-uWvUP_d37v8F4T2Q14,1773 +triton/backends/amd/include/hip/device_functions.h,sha256=vkybrdk6wyZP-T1I5PRjtfcMqGYXDeBpB5jhYj358GU,1589 +triton/backends/amd/include/hip/driver_types.h,sha256=m1HI80HC80qkTeco2Jd07woL_jTy48lz9JiDCV_8zsg,18985 +triton/backends/amd/include/hip/hip_bf16.h,sha256=lLw6K5ltb6AqSuINYTq8flxxsDkBP8Y2zbqmUjBcG9c,1571 +triton/backends/amd/include/hip/hip_bfloat16.h,sha256=Nqoy9VjfjglVx2_NJcp8hyT1sJUukXRWj8XMlidv1yA,1755 +triton/backends/amd/include/hip/hip_common.h,sha256=q5aPhG3DHW0iUJ7ayS5lfM_ZnZQNbMmLmfdHlOwbPdA,3450 +triton/backends/amd/include/hip/hip_complex.h,sha256=TmdzQP5oVPfhBVARJYcR5eyv9HInmKMFuFoQ_1ECk_I,1594 +triton/backends/amd/include/hip/hip_cooperative_groups.h,sha256=gMLvaYQ3b-f1vcoMtEwtkN0hO5__zNfP5p5oBKmv_SE,1878 +triton/backends/amd/include/hip/hip_deprecated.h,sha256=gFLuCuKn7R_xCfum_i_Q-vi3Lg8NWHKphKZKze8DwEo,6340 +triton/backends/amd/include/hip/hip_ext.h,sha256=mlzOesF-X62g9AnWdA4MP99GRu_VtwMbffOJFloLtRc,8609 +triton/backends/amd/include/hip/hip_fp16.h,sha256=vKJh-zgDWUW7NyXxtv2ho6aVLXX8BIPfzCigEQ5d6I4,1523 +triton/backends/amd/include/hip/hip_fp8.h,sha256=C4qn0im7Uhvp226VmL_QbF2b45Hhss2eokdSbqWkBcs,1433 +triton/backends/amd/include/hip/hip_gl_interop.h,sha256=-GwkSFMBneM8akFE7pqlhi0k-Ft2uz5674wGoiaU43Q,1438 +triton/backends/amd/include/hip/hip_hcc.h,sha256=RYrArDlnTEP89xKbzIpW17_bsBY5moCitq00PL-4oWI,1307 +triton/backends/amd/include/hip/hip_math_constants.h,sha256=8bSfve5E7cDuvNAUkFUeQwSLg3iJJHuqhuD4FmHNxEM,1588 +triton/backends/amd/include/hip/hip_profile.h,sha256=sjsNuduu5Jd6s7sJndZvZLlE0RZ0wN1rTVwv5nR7If0,1304 +triton/backends/amd/include/hip/hip_runtime.h,sha256=uy90l8Nep6xNUzeGcHMoDv84BT3hMpieTV-5ijkpL5A,3058 +triton/backends/amd/include/hip/hip_runtime_api.h,sha256=h5-nVJmqq2vr9P5EBimA7DwOE339eC6dt0uzk1DaC58,402475 +triton/backends/amd/include/hip/hip_texture_types.h,sha256=AhkvjG4cDjf_ZFLg5SsSTfBnXG614PBK1XVPa7irZbk,1237 +triton/backends/amd/include/hip/hip_vector_types.h,sha256=6FcBMBkP3ZN1Enalpa9hV0VopxdBJvbUCuaxISgzbTY,1630 +triton/backends/amd/include/hip/hip_version.h,sha256=3MuxlZXlLLrhLllRr-1yrIZyv7l_zknuZrT8bDDtm3g,407 +triton/backends/amd/include/hip/hiprtc.h,sha256=dBxesFd3NCvG-NQ99s61i4fUQXfoHNmYhdF6CtE9QHQ,15904 +triton/backends/amd/include/hip/library_types.h,sha256=tPOJTQedPH5qC9meawLgKpnbFrQC2WKlfo6s0rhKoZc,2370 +triton/backends/amd/include/hip/math_functions.h,sha256=frzdJ4veBG8n9ALO4EmRrdOiDguR6FP6ygLnvOnVVSM,1815 +triton/backends/amd/include/hip/surface_types.h,sha256=uQHjITphDM7k4pnuEoDEupMUxBobzvhJpSy0unpegh4,1959 +triton/backends/amd/include/hip/texture_types.h,sha256=CtmdykZfDikhnrVfdJk3w2VK5X3Af_6rEKzU-VgLu24,6687 +triton/backends/amd/include/hsa/Brig.h,sha256=5H-btCHq40qgjjpwVAoRWf3E0ccf-J6UCPEcKx_hGKw,32705 +triton/backends/amd/include/hsa/amd_hsa_common.h,sha256=q_zN0eq-dwR7FnQ84PcpV3yZyvjHsouIAjJgKltGoX8,3912 +triton/backends/amd/include/hsa/amd_hsa_elf.h,sha256=lBT57uuPT6ra01HQtM1O3RhVKW3n41AjESDSi6KVN7w,17306 +triton/backends/amd/include/hsa/amd_hsa_kernel_code.h,sha256=C55F8a480QsW16-iwN9TIT3cKnGh6GoeoEaEv3aVh4g,12659 +triton/backends/amd/include/hsa/amd_hsa_queue.h,sha256=ZJ-k5wY30heLmQnGB0VUz36XCiVHRmspg5FRNMGIk_U,4766 +triton/backends/amd/include/hsa/amd_hsa_signal.h,sha256=FDegZnWQC04GtnqHjXOBsB-AoVSaqdhNY6Mwbua5FGA,2947 +triton/backends/amd/include/hsa/hsa.h,sha256=msdMXlo7LeNQzpBQXj2yFjZqjcWerqB46vETGGB6jJU,191191 +triton/backends/amd/include/hsa/hsa_amd_tool.h,sha256=pyZSyIVl-UA5AOhte78jvn4V3hCd0dxJAIv7KeADsPs,2843 +triton/backends/amd/include/hsa/hsa_api_trace.h,sha256=GRBMFBWlv2dYJIT8BAWa1y_q-veLIgEVa1uIm-eM-Yw,29602 +triton/backends/amd/include/hsa/hsa_api_trace_version.h,sha256=yU2aATDJyhnCZF-FmgeF2wnp-EvPLE7n5zSuw8CNAqc,3248 +triton/backends/amd/include/hsa/hsa_ext_amd.h,sha256=zfiJ_FLsLVRr3_O4FYyoMQG0zYn02EolUKCh8e8vgeI,118084 +triton/backends/amd/include/hsa/hsa_ext_finalize.h,sha256=sv0AZbDM-B1wIdQ3cHTMlpUtNacQN2PkOgX90IZol_o,20227 +triton/backends/amd/include/hsa/hsa_ext_image.h,sha256=t5YJm_aw9EePCeFL1hoIfQ8ubIjBte-ptfReq6Ts-8Y,54232 +triton/backends/amd/include/hsa/hsa_ven_amd_aqlprofile.h,sha256=IR_vo3kSqcTs-gnRdpXW3JE1RHQ71FBi4S81G8GtpZ4,19823 +triton/backends/amd/include/hsa/hsa_ven_amd_loader.h,sha256=c6cxPAzAox7u6IbFzEkQZfCuRl-Kr39WhY2_w23X1R4,26146 +triton/backends/amd/include/hsa/hsa_ven_amd_pc_sampling.h,sha256=1jQvi96s94-GjOGktmvpqjVRsTGANugOsV6XmgMfivk,18767 +triton/backends/amd/include/roctracer/ext/prof_protocol.h,sha256=6FAcvVD-dNM7uulFs2B-aTxw5xOAWGy6evdD4yUaebA,3849 +triton/backends/amd/include/roctracer/hip_ostream_ops.h,sha256=Si6RFHPn0QtMpSBboKFyfkrcOMYhTiom5oUnDEqRnP0,188184 +triton/backends/amd/include/roctracer/hsa_ostream_ops.h,sha256=jSHN5Ltd_EJtecE59xJS4wL_CuLuNkbO3QMqDgMpunk,69906 +triton/backends/amd/include/roctracer/hsa_prof_str.h,sha256=f0pO8T0JBnG_stZ_Bgawx-5_B2BNS7s9ogvEldA0UFo,124187 +triton/backends/amd/include/roctracer/roctracer.h,sha256=B8sHz2DMNprP7EqNWIGwVLY1KQMpxmhfVy4UoR8dzzY,23849 +triton/backends/amd/include/roctracer/roctracer_ext.h,sha256=vLaZ8peAxSy0cwrdEalKnUApkKspfa04iw1Mr_Zcio0,2940 +triton/backends/amd/include/roctracer/roctracer_hcc.h,sha256=NlF3R8JQ9oX9lGpm0b2n-EWJ0r3y9sP9wbwnoucaCuY,1303 +triton/backends/amd/include/roctracer/roctracer_hip.h,sha256=RCzYuNw1vLR7xK4rb06TtM9TU546UYKHJ83IMHmZEm8,1432 +triton/backends/amd/include/roctracer/roctracer_hsa.h,sha256=M8APM64XNAWSslxQisM-pcmKoUQaUdTMaKvSACyt0Ag,4108 +triton/backends/amd/include/roctracer/roctracer_plugin.h,sha256=8GGE1zDbdPCVJtbmwOCYq7X0mwFjfWRtzDYKLD4cKys,4786 +triton/backends/amd/include/roctracer/roctracer_roctx.h,sha256=gBjBk5vb0l3PbBSQ7V9iFtaM_RzkIDJEW1A_PXBihBM,2014 +triton/backends/amd/include/roctracer/roctx.h,sha256=RhJXUXRhSJ5LRE_1gm7E6-bjEMrfcFBLDLuf3UxAIh8,6717 +triton/backends/amd/lib/asanrtl.bc,sha256=1xv2RlU3WvbdsghHlmhwiHewGM2B5dKts5bERM6S89o,24508 +triton/backends/amd/lib/ockl.bc,sha256=wQKCzkKukIHbu0lyjKUYlhndc7S27xto6L54J0Bn-C0,246124 +triton/backends/amd/lib/ocml.bc,sha256=UPNTXW0gCXUNB-c6orSYwb-mz9_mjUc7zny_vfFza44,205964 +triton/backends/compiler.py,sha256=MY2_cQG26p68z8VwRv2Nlj_h2DfEhwBbN-30caMgep0,2840 +triton/backends/driver.py,sha256=JcsL7pVHSIM13i1JWvc4KYum1MIPHTeUHcF3UcV6IR4,1802 +triton/backends/nvidia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/backends/nvidia/__pycache__/__init__.cpython-310.pyc,, +triton/backends/nvidia/__pycache__/compiler.cpython-310.pyc,, +triton/backends/nvidia/__pycache__/driver.cpython-310.pyc,, +triton/backends/nvidia/bin/cuobjdump,sha256=rR0PBpn0ZgNBbrWOe5_fmik-lfkxLJ33_W_aVqbDDUE,569008 +triton/backends/nvidia/bin/nvdisasm,sha256=EIC8kJowIHYf9MeKm3vZtJN725wJeHCcvlI4HeYSofI,5894744 +triton/backends/nvidia/bin/ptxas,sha256=yWCk8jixfVxdPAGtK7wevSxa7MRZy00iO_8QtF-bj8o,31912192 +triton/backends/nvidia/compiler.py,sha256=Pj0z-6twx8Bfxw4uL912PcoLhH9iFQWSuZ-IhqQZ7mQ,21822 +triton/backends/nvidia/driver.c,sha256=itdZazTLXkclZwbZbjKo-5zEA_B05aTmUnTe9DboSoc,17811 +triton/backends/nvidia/driver.py,sha256=g-9Y8jcdp62OAYIsKvyC-aLWUWwiSe4LvYhzuyBha-A,25744 +triton/backends/nvidia/include/Openacc/cupti_openacc.h,sha256=Z0OM5e_hbd3cxdXyn3SCHqBBQawLg4QORnlm57Cr2-M,3513 +triton/backends/nvidia/include/Openmp/cupti_openmp.h,sha256=E1WNmeb_7HaUSmBegtUNe4IV1i7pXeNxgzIlyKn1zrM,3491 +triton/backends/nvidia/include/Openmp/omp-tools.h,sha256=AmuC_xPC7VPu3B-W4PmXuCNufFawhY8PjNXePaQFAOg,37403 +triton/backends/nvidia/include/builtin_types.h,sha256=JxT9Vf2q2snxTBOL9ACzNmYzTWACO2VOVUu1KdFt7_g,3150 +triton/backends/nvidia/include/channel_descriptor.h,sha256=oZIDO1kdexPb9jltUx1AsXAFknvRWAAr1456925Pqig,21846 +triton/backends/nvidia/include/common_functions.h,sha256=22LTZRVcPZzEH6MJda7nNMCvMgIjSTe0OKR7sEQj6kc,3410 +triton/backends/nvidia/include/cooperative_groups.h,sha256=y2cFxa6e-saEFA9aW22ZuTwi0wud7eEHq7XN3v30LT0,60684 +triton/backends/nvidia/include/cooperative_groups/details/async.h,sha256=xsEHCZP3nuEY3l2p8SU2d1226XiXumUvDP_Gyh8PdVY,19122 +triton/backends/nvidia/include/cooperative_groups/details/coalesced_reduce.h,sha256=pBQgFY7i64V87XNATg1UEIQHVNYOItQtHjS5B4yn8pc,4257 +triton/backends/nvidia/include/cooperative_groups/details/coalesced_scan.h,sha256=DfZv5d5W0XJv-tZVhgrIdjLjs6aCx_u0oy1lDIpjo1Q,7314 +triton/backends/nvidia/include/cooperative_groups/details/driver_abi.h,sha256=v-ZUb4UgGKJk6NR2WCWHD3x_42y-togI1urFn70Gi-g,3964 +triton/backends/nvidia/include/cooperative_groups/details/functional.h,sha256=2BV8i8Bidz0kgxuYkJCAbwFxOIZRyzHgG-c_rVKhRzc,8905 +triton/backends/nvidia/include/cooperative_groups/details/helpers.h,sha256=K9jvxnXc5-6Fum1KG4EQKJJrVZ4BhHOSAJbZR4uDL0c,26476 +triton/backends/nvidia/include/cooperative_groups/details/info.h,sha256=FOrp3Ltt4PcbK2fAM5UX9jssFZtj_LqVShzLFcKiSaY,12465 +triton/backends/nvidia/include/cooperative_groups/details/invoke.h,sha256=Osq3K-tZuXHVCMQJ708PjPo-BwMhjhjApO4b0TYLFJg,8616 +triton/backends/nvidia/include/cooperative_groups/details/memory.h,sha256=hES3SfgXIBsj2MFrC_M5COXlOirSBuuhPMAJnWoI92w,5606 +triton/backends/nvidia/include/cooperative_groups/details/partitioning.h,sha256=AQz-TheqX3onqX2RmIUipzYUVB273zhLlHJw_kX9D2U,7153 +triton/backends/nvidia/include/cooperative_groups/details/reduce.h,sha256=MjqMDwT0TyWZk4JWcF3WHw8xtwMqyizA4C3zy7f8ee0,23296 +triton/backends/nvidia/include/cooperative_groups/details/scan.h,sha256=-Ttwb2AfEEY_tsmqJjR2dojkPpoRx387SoqxgvfdBtQ,17166 +triton/backends/nvidia/include/cooperative_groups/details/sync.h,sha256=Ed4K9QrPZi43ddSqZwv1X8NG_CTsXUowSQndoUv82LU,10795 +triton/backends/nvidia/include/cooperative_groups/memcpy_async.h,sha256=erOIHuObdfxRhBWfrXE3wsZF4B2GUuqwzQrsPwKPpbg,2960 +triton/backends/nvidia/include/cooperative_groups/reduce.h,sha256=B0hgDkqM-6ueqTTgb3b34A0RH4vGz8mBf5e2jT1dJ1o,2949 +triton/backends/nvidia/include/cooperative_groups/scan.h,sha256=2EU6T5cWNwftm2B7FicV31PojoI61yo5fHXGRYkGk40,2940 +triton/backends/nvidia/include/crt/common_functions.h,sha256=-U44f4yUGmwDPwd7Q_3Cz5if05xHGPSlAzz5zMylLSQ,13559 +triton/backends/nvidia/include/crt/cudacc_ext.h,sha256=KW6n0ImOZKS0VqVmBHWTXtHI816hh88YeEgUg2aYdVU,3224 +triton/backends/nvidia/include/crt/device_double_functions.h,sha256=A1vB3g0qwnNEfcpT1d9RiGDaxqPXXgYr-Vxe2oMHyxY,39938 +triton/backends/nvidia/include/crt/device_double_functions.hpp,sha256=YYIbqYhb5Qmf8c4YfcC_jytg4FRwcXPjv3TFTwhb24E,8568 +triton/backends/nvidia/include/crt/device_fp128_functions.h,sha256=3iCKrdmPp1NIjrlGR47dCZOgV9X6MmdfmjugrF6DEts,51047 +triton/backends/nvidia/include/crt/device_functions.h,sha256=T4INuRBIlLmd3JEA8H6B5XHVANGV6j8gwyUTj5h_F18,137919 +triton/backends/nvidia/include/crt/device_functions.hpp,sha256=OVHiqBjday_aUFnDKJxTeI0VDZI8ZA6JIdUKeAKR4pA,37991 +triton/backends/nvidia/include/crt/func_macro.h,sha256=EOpDlaM917bh9cwBiFBPF689DCMBw5hFarxLxFt-i74,1755 +triton/backends/nvidia/include/crt/host_config.h,sha256=1eki3w5xuY00gIkdYbmSfZO2SoI8giZWYSTRIL2uFs0,12169 +triton/backends/nvidia/include/crt/host_defines.h,sha256=fDU_0cQdtIvEdM0zZFKY0OboMANEMgHLblifPalPGLk,10102 +triton/backends/nvidia/include/crt/host_runtime.h,sha256=lOpmkxFZVkEp8dcMAGEZRITsh-19o9jy39kdSNLc3Ng,10284 +triton/backends/nvidia/include/crt/math_functions.h,sha256=LyGJ0XUthi6WEi-YVITInha9A6SFoB4fEUUU9zii7U8,238284 +triton/backends/nvidia/include/crt/math_functions.hpp,sha256=u-CGbd0R2FZWdKG-6bdmGSor9KT_wnmISj63lPQKASM,100207 +triton/backends/nvidia/include/crt/mma.h,sha256=BooWALDWATvZupJaL7-GFQRQYOotNe_Fy11I5NGh2FA,62695 +triton/backends/nvidia/include/crt/mma.hpp,sha256=spo0LX71tUCipxK517Bssj0nc-ZHf8oMWzvHoYYB_6I,66599 +triton/backends/nvidia/include/crt/nvfunctional,sha256=FDM0zqWO6bl9jpJKz9U8CMbjt6iTKh18tQalxAvRsag,16900 +triton/backends/nvidia/include/crt/sm_100_rt.h,sha256=3cBiCU11OcjGYpr85edCN1q4m2z91FaGhBtuO0is4To,8987 +triton/backends/nvidia/include/crt/sm_100_rt.hpp,sha256=vbI2CNCY08dDI7zXwp6BNceZKl0ceXG1lveq4w-VNao,6855 +triton/backends/nvidia/include/crt/sm_70_rt.h,sha256=Kf830xymA-zmF7LsunFHLSNyhhT5UiJMocgoHBQeNns,6837 +triton/backends/nvidia/include/crt/sm_70_rt.hpp,sha256=3a_rU-Y0MSB4htBDFY4PCQ_jXiWFTe7WT1ZyhMuCJOA,7837 +triton/backends/nvidia/include/crt/sm_80_rt.h,sha256=MdJHWCRzLM__nDDf1go61rDsl9ydOW3oi6SZBfjUyc8,7743 +triton/backends/nvidia/include/crt/sm_80_rt.hpp,sha256=o-rJu-jpehCeyABGgv-8dYRB7oJTCwuNdvSCq0VURdE,6705 +triton/backends/nvidia/include/crt/sm_90_rt.h,sha256=an47m0XFBaJ3pUX9MlE4-nktP1jb3eJUXhQ3ntZtzc8,11445 +triton/backends/nvidia/include/crt/sm_90_rt.hpp,sha256=YuqVygGV6rgtWtx1J9cPpEI3BXKQBII-Ez6oZFP3wrE,9228 +triton/backends/nvidia/include/crt/storage_class.h,sha256=dzcOZ16pLaN8ejqHaXw4iHbBJ6fXWxfaU-sj2QjYzzg,4791 +triton/backends/nvidia/include/cuComplex.h,sha256=WpcgpaiPhU_o9sTPMcNTEZuyXDIc8x3sz4dUWSztL2g,12186 +triton/backends/nvidia/include/cuda.h,sha256=RWjMnnoyHkdwfNZAOYDyGsLi5VFwUA0OCj9U_rA6mss,1156988 +triton/backends/nvidia/include/cudaEGL.h,sha256=iruZU9xSGAcJ29OEX4K_Uo1o4NGP9hggv2fiOZOfDQo,39955 +triton/backends/nvidia/include/cudaEGLTypedefs.h,sha256=xF_FAN1Kar9oyHJ3cCU7jztTpxX8WylpiuYyYpGGHek,5645 +triton/backends/nvidia/include/cudaGL.h,sha256=gMT1HPGa-siuji0gAsKYr4X45Lc29HKglC_ttNSGyUM,22501 +triton/backends/nvidia/include/cudaGLTypedefs.h,sha256=dClpQI-LuXgF9rPSBsj7OkIg8g_fXDjT0hLZS8TGpOg,6576 +triton/backends/nvidia/include/cudaProfilerTypedefs.h,sha256=F2aWLIKv_AhNbxNOaZVcRsxIh0kuscnV8UMWWxkBAlY,3297 +triton/backends/nvidia/include/cudaTypedefs.h,sha256=SKfAvTOj19zxsiLGKhoxXPiopKqoe5hjj5iXkR2_v6E,115169 +triton/backends/nvidia/include/cudaVDPAU.h,sha256=Np7Nc2Wjaz--hkpbhW6f9aapr-NbcPDAgkot0sJerco,12694 +triton/backends/nvidia/include/cudaVDPAUTypedefs.h,sha256=wz8nyOUdwM9mH9JO3QZW-A9dyxt-IufSX7nggSXpCNs,4144 +triton/backends/nvidia/include/cuda_awbarrier.h,sha256=3ZH-ZlXODhSiwSY9rqSni_EQwi25QMHP6Tm-zOdxBwE,9340 +triton/backends/nvidia/include/cuda_awbarrier_helpers.h,sha256=OCskCts5bCKl_RKBe9M74zKSIsVpePn44S_aJp1tFXE,12489 +triton/backends/nvidia/include/cuda_awbarrier_primitives.h,sha256=n5__E1jYYDhlgH-f3u8MQjtz57UZ7v5VshhMye1eicM,4699 +triton/backends/nvidia/include/cuda_bf16.h,sha256=TVoq2IrbF5g67wUF7W7SoGA0l8ecEDu6gskoMB6hIxA,204512 +triton/backends/nvidia/include/cuda_bf16.hpp,sha256=OukWXoN6bgRlC-p8CFbhUN0G0uAJb_zos1mCPagscnI,136544 +triton/backends/nvidia/include/cuda_device_runtime_api.h,sha256=54l66QbwerX0wPKoJC2y7qCdGP8nv1_GgdmMV8A0x4k,46986 +triton/backends/nvidia/include/cuda_egl_interop.h,sha256=awWBBEYvUFM7AURNp2mND8H7_5kGQLRswRveXYBy-3s,37509 +triton/backends/nvidia/include/cuda_fp16.h,sha256=jrFgCo4uM9QFcr_-cAGif2BGp0lJ2ANT_gLPiLJWPdo,206851 +triton/backends/nvidia/include/cuda_fp16.hpp,sha256=o1ITDmuN67N8YUGUcvTpV3IdpS-6wwlm65M_H-8LYKs,120927 +triton/backends/nvidia/include/cuda_fp4.h,sha256=pTEQf5rLfiaU_UMXgnnsS13NH5H9FtHgdeiNuW_NkHY,13823 +triton/backends/nvidia/include/cuda_fp4.hpp,sha256=YYaUu-YRgYdj9xYu4ZDh_uPVffxkDlEr0CD_bhlF8BE,35423 +triton/backends/nvidia/include/cuda_fp6.h,sha256=6xh0E4SNmjmJZD3H5_HoZe08bQ0loUE8y3cbO19-Ad4,13963 +triton/backends/nvidia/include/cuda_fp6.hpp,sha256=qa838buZeLP32xBVqbo71uFSW5RnBWx9qp5D-SR_xc0,56455 +triton/backends/nvidia/include/cuda_fp8.h,sha256=QSTMRb9l7I9mnvT1_8KXNqLO48wWaWEgG97bDjEh1ic,18072 +triton/backends/nvidia/include/cuda_fp8.hpp,sha256=6AABZPfPlK-jmQyMdbTEE0PZTlULjuLOrsQ0Z_cubBw,97230 +triton/backends/nvidia/include/cuda_gl_interop.h,sha256=VQEswFeOBF6JN6Q0pdlkvc5WT7bD1FnTfKewvANulCc,19150 +triton/backends/nvidia/include/cuda_occupancy.h,sha256=0HavrMIWXGxIujaq72iX31-73Zprx0WBYdiln3ZNP2w,71302 +triton/backends/nvidia/include/cuda_pipeline.h,sha256=0enXG49wN4JajlQi3ahbp2ei_ufTY_Mznic7zfWmKHM,8130 +triton/backends/nvidia/include/cuda_pipeline_helpers.h,sha256=bo1L7e6vCuM-K3Il8K1z4wJUja5DyXQKdo_hSWUME-E,13852 +triton/backends/nvidia/include/cuda_pipeline_primitives.h,sha256=FnJJtuV6rHr6LgL56XDwilcSbFr6W1Hj6mf1AJaMI20,8675 +triton/backends/nvidia/include/cuda_runtime.h,sha256=GqqE7SrECGrN-Qg5Dk90LSjs-xvKlHZpRLlpH7LUehM,98570 +triton/backends/nvidia/include/cuda_runtime_api.h,sha256=EWhSESFT_vV5eYZpTBEu4EvgNtE9rhmHP503XnIGHIs,655943 +triton/backends/nvidia/include/cuda_stdint.h,sha256=XbFOk9CtJjKqk7PpYNqbSVsDxAsVM8avA4rWpPi0BjQ,4093 +triton/backends/nvidia/include/cuda_surface_types.h,sha256=Mw5Lo4b8Q-f9mogOvATGyHhu9d2t2K6XOxuqtZrSh3A,3688 +triton/backends/nvidia/include/cuda_texture_types.h,sha256=ITbX-JNnP7Rm-JSgNVdJ9pq6k8FVor8RbnruDsKq6sk,3688 +triton/backends/nvidia/include/cuda_vdpau_interop.h,sha256=bXQanWc2IFXZAKWNGl2xAz9nLvFmQpWyGrsDvfeS9FA,7727 +triton/backends/nvidia/include/cudart_platform.h,sha256=YN6sKhB0b9w5tGX1IYL7ulJVPrWAiX9A44qLv4EtW5Q,2717 +triton/backends/nvidia/include/cupti.h,sha256=JkVyAGTIMYzwm62dfVqas3nMcILhgP_Wdz6fh4_NED0,4697 +triton/backends/nvidia/include/cupti_activity.h,sha256=gJKGlG2JifW36Lx-ujJcKBlnUrNOoTxar51k28GrUtU,229848 +triton/backends/nvidia/include/cupti_activity_deprecated.h,sha256=B0p4zbll2vUn1j0ImTG6QIbpp6Hiw8y-X021Zmf7flE,137602 +triton/backends/nvidia/include/cupti_callbacks.h,sha256=aZ-SE0YMFfT9R-Uh5MHboPg0ypHMjeSSAJw3zdP7OCs,29689 +triton/backends/nvidia/include/cupti_checkpoint.h,sha256=rTz8JoWxqESBXyZWUhZJGm4xeYcx4OJOtJ7Ld13T_b0,5264 +triton/backends/nvidia/include/cupti_common.h,sha256=85m74bxUgXp3tEaPQpezeazmpsNMw41PsjNSYmQdT20,3514 +triton/backends/nvidia/include/cupti_driver_cbid.h,sha256=mkBNPYkLfcExhQZFDo0iYHlaHJWGD2vOMdtzaV-lEUk,77280 +triton/backends/nvidia/include/cupti_events.h,sha256=81wcvFvvHj8RmECbbEp5FfgjJIQDoC_81FhvqznFupY,51923 +triton/backends/nvidia/include/cupti_metrics.h,sha256=zmfZxq5VkUJp6Tj7oXEkP9oycRNw1zB9VNhoQlbhiN4,32175 +triton/backends/nvidia/include/cupti_nvtx_cbid.h,sha256=_azPtR1g4qivvX7qbvHRUg0RHCWF7iEOJyHMN9qZe9E,5912 +triton/backends/nvidia/include/cupti_pcsampling.h,sha256=ycJHT36DmPIaVzHsB3xxjXkhFyEfMCJOl3LbCsHFgyA,32144 +triton/backends/nvidia/include/cupti_pcsampling_util.h,sha256=lx8CaNXowJe5Zvc06LE-u_Zry_jODs1mM6j9Q5WIX9E,12430 +triton/backends/nvidia/include/cupti_pmsampling.h,sha256=U95hKOwIkZSbGNVP11QSmMawB8qdJsljY_tUJY4vedc,20440 +triton/backends/nvidia/include/cupti_profiler_host.h,sha256=MkkfXlKBRrRL4NfaPFiuE4D4z_gpmxiBWWTBixyyMTk,22155 +triton/backends/nvidia/include/cupti_profiler_target.h,sha256=MdLutIefwdMTI7wsce0LO3NuCm3FRgFR3GxAkqadMs4,32294 +triton/backends/nvidia/include/cupti_range_profiler.h,sha256=ue5bUA-3xCwAtQGyDe5O1d5rAmRbVbcrXKfITd4xM1I,18779 +triton/backends/nvidia/include/cupti_result.h,sha256=xQqBsZRoicBSWdk1lZAE_WeZj88MLH6ClTo58oshx-8,13114 +triton/backends/nvidia/include/cupti_runtime_cbid.h,sha256=BZJnzsvf2RjRlHKEhPjCk0CdjLI9_L-nClTwe4v9NUc,48372 +triton/backends/nvidia/include/cupti_sass_metrics.h,sha256=3RW9snJuFQdOhrEn3wDJOru05q0V_zssWrqD7tvVJKw,19674 +triton/backends/nvidia/include/cupti_target.h,sha256=x4Vz1Upb6m9ixmVpmGaKQldDWYQI3OZ-ocEXGzNK0EE,1263 +triton/backends/nvidia/include/cupti_version.h,sha256=KFXmjB4o-iZGvO8la9Sf9Urg4q4srmEimnxbPCyd2N8,4506 +triton/backends/nvidia/include/device_atomic_functions.h,sha256=OR2jNSfSKzaFri74zh4Vtz5M0z9UDBU3rKeC1rYaVQs,9500 +triton/backends/nvidia/include/device_atomic_functions.hpp,sha256=0e7MOiNNUnnloXpB_r9WT5YOws5cxgzQQAzRCYvgaFA,10486 +triton/backends/nvidia/include/device_double_functions.h,sha256=KUxId5Z1fx8SWfLRTxPD7RB-zN7zslzb4n7JaJLfL3I,3452 +triton/backends/nvidia/include/device_functions.h,sha256=bWSrhTYE9NQlss7xMSMEVusvto9j2fgUDXWVH2W_cOA,3410 +triton/backends/nvidia/include/device_launch_parameters.h,sha256=H1_CC-vvAaS26ys4XsTFkMgTxUTciAjdjswjizkisvQ,3846 +triton/backends/nvidia/include/device_types.h,sha256=2LFxoZBJPoA5V0H1EbKTEaXDi3GDJPtzOPdRHDaucIQ,3588 +triton/backends/nvidia/include/driver_functions.h,sha256=cN3IjRAz2Mj2Pj35SyxJIkZNDDusnJqaqzBdMzpQKbA,4625 +triton/backends/nvidia/include/driver_types.h,sha256=mMNbiIwg5E3k7Sk685YCSvnKYmfQ3bxWv3bkEgzOtNU,200083 +triton/backends/nvidia/include/fatbinary_section.h,sha256=NnuUfy358yGJx4enq0pBnetjv17UWa-nOlgYToUitrw,1809 +triton/backends/nvidia/include/generated_cudaGL_meta.h,sha256=dfd2QuaRdEjbStOKvaQLi1Md_qrpRQh8PfyZznJ8bWY,3115 +triton/backends/nvidia/include/generated_cudaVDPAU_meta.h,sha256=fAedsoQxaU3hIAApAWDOKsa9kgcuQw4tdyf8klLm-3k,1453 +triton/backends/nvidia/include/generated_cuda_gl_interop_meta.h,sha256=LXOqvQCej0sCgAT1LUKKYZ466EFxN4hIwf9oIhXOLF0,2250 +triton/backends/nvidia/include/generated_cuda_meta.h,sha256=DDdgfW84GVtsGbr7daNJchmmZDS_xfvDHvFCm3I1OEc,98664 +triton/backends/nvidia/include/generated_cuda_runtime_api_meta.h,sha256=CuziaDwO2Mh33paCLGKqi73PQfYNmzp38wYrhAK-fng,72208 +triton/backends/nvidia/include/generated_cuda_vdpau_interop_meta.h,sha256=8OLqWN26aEYpTWUXtbHJvA5GYhVv3ybYVOTW7yK37z8,1367 +triton/backends/nvidia/include/generated_cudart_removed_meta.h,sha256=X3I5WXmhtsJNNlgY7coJ5vg4t11G5FRR6Xo7MboIeck,5172 +triton/backends/nvidia/include/generated_nvtx_meta.h,sha256=YHb_RD8g3s4m8PJn7Z0wnxvUHarl7BOAX5ADr-BL3HI,7513 +triton/backends/nvidia/include/host_config.h,sha256=BscH_GazAZbbotddVzL5RmafbQ-QjRx8f-I1O01IBW8,3380 +triton/backends/nvidia/include/host_defines.h,sha256=bBQwQF5C1N1c2qpLV56g1c-weu9Ysgz-gIf2Kn3uz_A,3386 +triton/backends/nvidia/include/library_types.h,sha256=i-GFcw92wvcixs2bQjOj4I_q26HYY_VY4DpDvHWQCjY,5156 +triton/backends/nvidia/include/math_constants.h,sha256=cV6hAyQe8X7f7MBtaKjjIJq3BycOUDp6I5cizJX5HLw,7608 +triton/backends/nvidia/include/math_functions.h,sha256=5XcC6j-fJKttvhwc4hZNoLHNw808a2ZYIOtZ7ry7yd0,3398 +triton/backends/nvidia/include/mma.h,sha256=IY_VenxuEncwGq92MhrWUb-Xswh0ekAXLy9Rbxhxa2Y,2932 +triton/backends/nvidia/include/nvPTXCompiler.h,sha256=lF1ssKlrjsNWCsp0UdhbB6GyOgFCDO1q0ZBDzOswlj0,14471 +triton/backends/nvidia/include/nvfunctional,sha256=IkFoCi_Q4OhP9nEuBI-5jWwFlR_PfG05hJH7lSMsfWc,2975 +triton/backends/nvidia/include/nvperf_common.h,sha256=ykeTJ5I6c0z8KqMQh13hlJaMaHiqqVUB60oGXOCu7Bg,17255 +triton/backends/nvidia/include/nvperf_cuda_host.h,sha256=gC0JWoUdTyAOJs8y4uoJIhie9Xq4yF4HzoumLsYNVzU,7562 +triton/backends/nvidia/include/nvperf_host.h,sha256=wCB4mR8aIHWiqT1TsxztQgWBRh_yiq5ABFm8sb3_jwg,49197 +triton/backends/nvidia/include/nvperf_target.h,sha256=jRqQtuNLTrCzPDdyeANkTrPEijSCTjLy2A1qjKu0SdM,23607 +triton/backends/nvidia/include/sm_20_atomic_functions.h,sha256=x4ycINVq__l9B4SQPD-I48jQbKxxdBmgp8Vf2GO0Qfg,4478 +triton/backends/nvidia/include/sm_20_atomic_functions.hpp,sha256=1l5NLM8DhDbqYZ_E51LoqElQJXObkbwo57d3r-4uEbE,4107 +triton/backends/nvidia/include/sm_20_intrinsics.h,sha256=axeDr7y6nT1V6LzrSWNSaHUwXgiNjPbXn1T6Uh7hlNM,57702 +triton/backends/nvidia/include/sm_20_intrinsics.hpp,sha256=mJTejRhw1prNiP_ax1OPbkYlhEqBqO4nVI3DRDXIzpo,8392 +triton/backends/nvidia/include/sm_30_intrinsics.h,sha256=b6W8Vxp9vD9OCJI6lZuGyZYXEdQ3Ei8PTAloHNkwCcQ,16978 +triton/backends/nvidia/include/sm_30_intrinsics.hpp,sha256=yX0ebd265tJ-BDhvluP2BhadPuWXpRZPI2eeQFFt5ys,24567 +triton/backends/nvidia/include/sm_32_atomic_functions.h,sha256=HGnZgQHACE2AAb6zabGUURc53IsVZelc2BSJqvs9OgY,5703 +triton/backends/nvidia/include/sm_32_atomic_functions.hpp,sha256=CQTTvOEYp-s5hqAgLvAon11vLYDrDp8cTHdel-XRzBQ,6592 +triton/backends/nvidia/include/sm_32_intrinsics.h,sha256=Xdkogdsjy1vh8u3eGu0i5xTmHxBGAjj6_vVGR-spdOE,33539 +triton/backends/nvidia/include/sm_32_intrinsics.hpp,sha256=Gl8aSLDLcit4W3pKQS19GsDG8RYcwD65HwYB_CeZe8M,70616 +triton/backends/nvidia/include/sm_35_atomic_functions.h,sha256=a3XoEsKRCEOf0Q_5Y__rMfmC4pScv4VkUggVgVJVn44,2909 +triton/backends/nvidia/include/sm_35_intrinsics.h,sha256=0mS5-LCgvZiTvL7-MG_4YwI-zWGvM-s4xyRuMkunMC8,2664 +triton/backends/nvidia/include/sm_60_atomic_functions.h,sha256=_anfNaJsvQpDEorYeUKIkbizYkwrinBcG_ZCiECtLqI,13178 +triton/backends/nvidia/include/sm_60_atomic_functions.hpp,sha256=cgIKddDn2B3QzYlzeBILAP1IRys74QCCxsH0QqaVGls,22903 +triton/backends/nvidia/include/sm_61_intrinsics.h,sha256=h_MBL1UUDxQX_qOddSImzqyFjcrhhm_63G97pGDyreU,10902 +triton/backends/nvidia/include/sm_61_intrinsics.hpp,sha256=N-nQvcBsPMT2Umy5zR69c9K1q366W-Jqe7NpoLTqTmg,6787 +triton/backends/nvidia/include/surface_functions.h,sha256=b1O82SAvEgWWxA9uZTWQcGimzZUoem2QbAET3wh3fZc,6782 +triton/backends/nvidia/include/surface_indirect_functions.h,sha256=vy9QuFVV-ezZP-x2RT9RLp2qIUgdngACOCmalSfVFPA,10877 +triton/backends/nvidia/include/surface_types.h,sha256=XkFXD1nHbeSMgajR-UJE9uQ7TByzJnjdnUL4-yGiufk,4530 +triton/backends/nvidia/include/texture_fetch_functions.h,sha256=KLCmUxf5aY5_UalX8tSFB6e4TrjA8hyUPxLOkMFltAo,12468 +triton/backends/nvidia/include/texture_indirect_functions.h,sha256=lH_y3Ni-hq4RZ0_PMFbBM0th5-OmTn3TtqtpkHHhA8w,21163 +triton/backends/nvidia/include/texture_types.h,sha256=73ntVyg8r8fzKy5VIk6yuvC45GDeWepaLIqIk-M3Ri8,6360 +triton/backends/nvidia/include/vector_functions.h,sha256=WypGkL-IDbGOlay7g_G0p3HO7OLGRE0Do__JtiFoWxY,8003 +triton/backends/nvidia/include/vector_functions.hpp,sha256=afXhNSd3LFTZo96EPtesTLfvxd4nTmLVzgkj967rTRg,10060 +triton/backends/nvidia/include/vector_types.h,sha256=6CJ4yt3KD7zQVfm1NhrgqNYYEDEIZWwaivlFx12nhNg,13396 +triton/backends/nvidia/lib/cupti/libcheckpoint.so,sha256=BrqCvQkje5NM8W2iFy7VlDxYLKo1x5tSz8_rD_cfclA,1644872 +triton/backends/nvidia/lib/cupti/libcupti.so,sha256=_M3F_MfDL4z0s-2bF_fxh6GFpvoKW8wG40NMvRaHyAg,7595792 +triton/backends/nvidia/lib/cupti/libcupti.so.12,sha256=_M3F_MfDL4z0s-2bF_fxh6GFpvoKW8wG40NMvRaHyAg,7595792 +triton/backends/nvidia/lib/cupti/libcupti.so.2025.1.1,sha256=_M3F_MfDL4z0s-2bF_fxh6GFpvoKW8wG40NMvRaHyAg,7595792 +triton/backends/nvidia/lib/cupti/libcupti_static.a,sha256=q22dskj52HUn40wsJ6-5g9AFD_aYtOSqIlNJrbRqOko,43563448 +triton/backends/nvidia/lib/cupti/libnvperf_host.so,sha256=YJ9UBjwlPa_277p5_kklL90jJDCKsP-5kcwSXT6aHqs,25825936 +triton/backends/nvidia/lib/cupti/libnvperf_host_static.a,sha256=sTLoX5YG81JXBoSKm3fLo0RAdP3HqISmgtx8L47XF-4,33480030 +triton/backends/nvidia/lib/cupti/libnvperf_target.so,sha256=LxDhEJ_DQaV0NMKOQAa0WokkeAHiD1rOE1wwLBc-B3A,5275312 +triton/backends/nvidia/lib/cupti/libpcsamplingutil.so,sha256=5vongwd5dPsnZb-py66_mAc233Tuv5lMWcdZaYiuyDs,970064 +triton/backends/nvidia/lib/libdevice.10.bc,sha256=XC-uN8huaMOjhgWpX1EtfRLV89uYYxC-R_VzBKpype4,473728 +triton/compiler/__init__.py,sha256=S0iIXHTRJL8MUqThupDtPyJVT0PA882c2dN5VwqOGxE,284 +triton/compiler/__pycache__/__init__.cpython-310.pyc,, +triton/compiler/__pycache__/code_generator.cpython-310.pyc,, +triton/compiler/__pycache__/compiler.cpython-310.pyc,, +triton/compiler/__pycache__/errors.cpython-310.pyc,, +triton/compiler/__pycache__/make_launcher.cpython-310.pyc,, +triton/compiler/code_generator.py,sha256=xNeZwy62o7kq02rX6lZF78Wy_GzVxz2tw6Z41X3SzPQ,71445 +triton/compiler/compiler.py,sha256=w3FJJ1oD0GP8HDOc12oZ2hwX49VVQQNy5srSnu3vYgI,21115 +triton/compiler/errors.py,sha256=I9Y15pDWcL9heY4SWWdLeMDtW6Iiq2pFXzKfJ6dY_C0,1732 +triton/compiler/make_launcher.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/errors.py,sha256=8WfnuRKLG578mgY6cBA3ECruVMf9ULEKFNgRcJ6IhWM,89 +triton/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/experimental/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/__init__.py,sha256=GO2iamrLyIbBszA44Z3Aa-V2v4S0kq4TOKPWrRx0FYM,186 +triton/experimental/gluon/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/__pycache__/_compiler.cpython-310.pyc,, +triton/experimental/gluon/__pycache__/_runtime.cpython-310.pyc,, +triton/experimental/gluon/_compiler.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/experimental/gluon/_runtime.py,sha256=AVoWQuzXSKD02gZLAzbZ6iHw4uccf0rH-KpWVH-7ceI,3461 +triton/experimental/gluon/language/__init__.py,sha256=HR1Tp5k-GhrKx_-7Jt0JCoTqlhQloIFMrUiHoYJIQNY,1765 +triton/experimental/gluon/language/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/__pycache__/_core.cpython-310.pyc,, +triton/experimental/gluon/language/__pycache__/_layouts.cpython-310.pyc,, +triton/experimental/gluon/language/__pycache__/_math.cpython-310.pyc,, +triton/experimental/gluon/language/__pycache__/_semantic.cpython-310.pyc,, +triton/experimental/gluon/language/__pycache__/_standard.cpython-310.pyc,, +triton/experimental/gluon/language/_core.py,sha256=1ACDrOT_X7MJZEoYAyOJ7GFANUs3FKdwO4guC9JUyKs,15695 +triton/experimental/gluon/language/_layouts.py,sha256=_PzL_sL8Db6VF5qD5s6T08CxEKplCCySyYX9JWt6rhs,22740 +triton/experimental/gluon/language/_math.py,sha256=MN9Fvt8FaQTpxEUCKhw83IGbkwBD2Whdp686av4vv04,564 +triton/experimental/gluon/language/_semantic.py,sha256=imNOS16WVaniSo20RT_fVMYk7PjZ-lWpNM1bn1YikU8,18062 +triton/experimental/gluon/language/_standard.py,sha256=9z0ZCxxa2jQjsk9pCaofNOuQT8o3WWFoW5xJGtxbio4,2718 +triton/experimental/gluon/language/amd/__init__.py,sha256=x7o_QGewXOYfxv1IiCjteHLWP_7975a-EX11naLcGcI,110 +triton/experimental/gluon/language/amd/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/amd/__pycache__/_layouts.cpython-310.pyc,, +triton/experimental/gluon/language/amd/_layouts.py,sha256=DLSoxygfOGuoaICPjF3GmRzEYeqAFIzFgjQuUTlFknw,4489 +triton/experimental/gluon/language/amd/cdna3/__init__.py,sha256=4Yl_KgH5hE35dOQdEE0M9DLUivGxRgVend4KU1E1SV0,4235 +triton/experimental/gluon/language/amd/cdna3/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/amd/cdna4/__init__.py,sha256=qumwcZI7tvee7qLRg-Ft_J78tgU9Wzth1OdIQsYXtDs,2157 +triton/experimental/gluon/language/amd/cdna4/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/amd/cdna4/__pycache__/async_copy.cpython-310.pyc,, +triton/experimental/gluon/language/amd/cdna4/async_copy.py,sha256=N_jqFSMieyVhrJEKx8TkrF5R1ELn0ghRXmv1Wt4adnw,7222 +triton/experimental/gluon/language/extra/__init__.py,sha256=FaDqvKRb8M8K47C3YceyfF3ZiZayLC8xZN248jNlxeQ,69 +triton/experimental/gluon/language/extra/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/__init__.py,sha256=SFBuACK5P2XoYcutHEnKjqgRTboU4CPDmJz0hT6dFRQ,80 +triton/experimental/gluon/language/nvidia/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/ampere/__init__.py,sha256=4hNkRbEZ4LhblmwMqyriZdI_a8SEEaNNRl__fv5p-I0,73 +triton/experimental/gluon/language/nvidia/ampere/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/ampere/__pycache__/async_copy.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/ampere/__pycache__/mbarrier.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/ampere/async_copy.py,sha256=YfB0AN1hcEKWU2WTyef4ggWtp-HBnoSHJGoWhUTMmis,2817 +triton/experimental/gluon/language/nvidia/ampere/mbarrier.py,sha256=CraFTt4ykqhz8AaWLt6x9AQbGX5RjLGGGzym7LV3jaY,2682 +triton/experimental/gluon/language/nvidia/blackwell/__init__.py,sha256=ni_x_NTkwEukCjH6iaDiEYowYDllXv29uLDIpfj7Okk,14705 +triton/experimental/gluon/language/nvidia/blackwell/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/blackwell/__pycache__/tma.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/blackwell/tma.py,sha256=f6u_qBCYBtwls7FVQDdROymUed_NefREQwAyF4k4XSs,2007 +triton/experimental/gluon/language/nvidia/hopper/__init__.py,sha256=gJpoqWWiJ2h4uCTFh5dYub90CamRsi4_NZQi47rspNw,5351 +triton/experimental/gluon/language/nvidia/hopper/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/hopper/__pycache__/mbarrier.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/hopper/__pycache__/tma.cpython-310.pyc,, +triton/experimental/gluon/language/nvidia/hopper/mbarrier.py,sha256=NF2dvXMKWzQ0mtWZQpCWKJKZzkINCZ24d8cKdqDT60Q,1339 +triton/experimental/gluon/language/nvidia/hopper/tma.py,sha256=UK2DeQqKgHaKtTWBXUAmlv2a5MS2jrd3okPsUSrheXQ,3588 +triton/experimental/gluon/nvidia/__init__.py,sha256=ISXB4RV7RcCLsU-JhcRFeA29gCBDVk8cTwO2j99ivLc,80 +triton/experimental/gluon/nvidia/__pycache__/__init__.cpython-310.pyc,, +triton/experimental/gluon/nvidia/__pycache__/blackwell.cpython-310.pyc,, +triton/experimental/gluon/nvidia/__pycache__/hopper.cpython-310.pyc,, +triton/experimental/gluon/nvidia/blackwell.py,sha256=cllwlUCE5_YKWqySQZk7wt7Fierz345E5VwztxNRGMs,69 +triton/experimental/gluon/nvidia/hopper.py,sha256=GkNztKn3bJWx-WnLv08Dkmchu1Nh3okTpzSEsCl6_Ew,1829 +triton/instrumentation/libGPUInstrumentationTestLib.so,sha256=tWzMeA3gcdgmVnnD_ctndpYt2_yCTo4JkjdJPyR_dLk,6437216 +triton/instrumentation/libPrintLoadStoreMemSpaces.so,sha256=20mXExzQ_RFuB5AKH-7rFUUTp5fUAa_AKDGgvaQLpSY,8964312 +triton/knobs.py,sha256=IK-IOVmFmd1F5fgw9sMrg6JIW1VHkpwMeJ2B0OOZWjI,16188 +triton/language/__init__.py,sha256=qlEvCv0KojmEFCNiUwv6kLnijJTRm3OnBPsHScueu74,6563 +triton/language/__pycache__/__init__.cpython-310.pyc,, +triton/language/__pycache__/core.cpython-310.pyc,, +triton/language/__pycache__/math.cpython-310.pyc,, +triton/language/__pycache__/random.cpython-310.pyc,, +triton/language/__pycache__/semantic.cpython-310.pyc,, +triton/language/__pycache__/standard.cpython-310.pyc,, +triton/language/__pycache__/target_info.cpython-310.pyc,, +triton/language/core.py,sha256=XBKBtnwNlJ2jTM8cO2iASiyqJmWVOYTYN9dT6Rr2Ef4,119331 +triton/language/extra/__init__.py,sha256=XRXFvr7416pRsh_Rh-X6qV66SiEyVDVbxp4GSAE1mfc,655 +triton/language/extra/__pycache__/__init__.cpython-310.pyc,, +triton/language/extra/__pycache__/libdevice.cpython-310.pyc,, +triton/language/extra/cuda/__init__.py,sha256=MBBu2EUYxsp6ygjiwO4Yh1X1EswMstfaiRTMSMGtbcw,407 +triton/language/extra/cuda/__pycache__/__init__.cpython-310.pyc,, +triton/language/extra/cuda/__pycache__/gdc.cpython-310.pyc,, +triton/language/extra/cuda/__pycache__/libdevice.cpython-310.pyc,, +triton/language/extra/cuda/__pycache__/utils.cpython-310.pyc,, +triton/language/extra/cuda/gdc.py,sha256=rbVOcdD_w72sQ3-8l5KXafu-Tr0D5jYcY1fyok98WVQ,2193 +triton/language/extra/cuda/libdevice.py,sha256=J7Kl0ejbAIus7-YBn2OSK71lkm3pC7G1J-5ZdHfS82U,56764 +triton/language/extra/cuda/utils.py,sha256=phDcXCFViaq3p4ThwHrO8-FtU-8A8I3nk4mZZJVvTio,4426 +triton/language/extra/hip/__init__.py,sha256=GFcuM-R0qCB5kDbysvjD5U2KUmOGNWcWzM8r8yqDI3Y,96 +triton/language/extra/hip/__pycache__/__init__.cpython-310.pyc,, +triton/language/extra/hip/__pycache__/libdevice.cpython-310.pyc,, +triton/language/extra/hip/__pycache__/utils.cpython-310.pyc,, +triton/language/extra/hip/libdevice.py,sha256=LNynTEKXgE3PSZp-nUMVq4ft_i1y9Tao_cAvKKPn2s0,17542 +triton/language/extra/hip/utils.py,sha256=tTq-k0c-s32IvDeGkgZcmosb9UtkF8sZEaE9EJGMbSA,877 +triton/language/extra/libdevice.py,sha256=9gwrstjuvhwZH0uLgZhEQU3RvOJDZFY1oJT5-SZlpY4,6350 +triton/language/math.py,sha256=CKvuIc5iMKhz7Qgx9w-VcLfOOZadv5svKK4aGZLuHMc,7399 +triton/language/random.py,sha256=jkuFmfgZ8yvKuub9EY27zPvsC6nhkJIk05xf4y-7SR8,7102 +triton/language/semantic.py,sha256=PbS-6UJ8frDiEFr_SEvazIGTV9KY6PbonDcq6cNiW98,99927 +triton/language/standard.py,sha256=Wc8pXzqrT6YrlqYnqp_sEwKVATN1DeWeVCx7TJ5bgLY,15971 +triton/language/target_info.py,sha256=UwWJDDsTPe9E4vPTQF4Psf1s540KKLISdnChlbvhHGY,1364 +triton/profiler/__init__.py,sha256=b0XtXi8lPkxUEhD9DHkS9Rzi2JvWj57RbakLGui1jA8,284 +triton/profiler/__pycache__/__init__.cpython-310.pyc,, +triton/profiler/__pycache__/context.cpython-310.pyc,, +triton/profiler/__pycache__/flags.cpython-310.pyc,, +triton/profiler/__pycache__/language.cpython-310.pyc,, +triton/profiler/__pycache__/mode.cpython-310.pyc,, +triton/profiler/__pycache__/profile.cpython-310.pyc,, +triton/profiler/__pycache__/proton.cpython-310.pyc,, +triton/profiler/__pycache__/scope.cpython-310.pyc,, +triton/profiler/__pycache__/specs.cpython-310.pyc,, +triton/profiler/__pycache__/state.cpython-310.pyc,, +triton/profiler/__pycache__/viewer.cpython-310.pyc,, +triton/profiler/context.py,sha256=LsH0trahuaG57C7RXIedizPQzllQXP1Umd4qdYBgEtU,517 +triton/profiler/flags.py,sha256=ZfaCP9qMKc8RR1IXaCyyhirK2U4JESMb_l_Z-IKCgJU,935 +triton/profiler/hooks/__init__.py,sha256=qwb8ypLQ0oOFA3h9M3RtU1OhJR_wd2EZTKJ2pFNHeuo,123 +triton/profiler/hooks/__pycache__/__init__.cpython-310.pyc,, +triton/profiler/hooks/__pycache__/hook.cpython-310.pyc,, +triton/profiler/hooks/__pycache__/instrumentation.cpython-310.pyc,, +triton/profiler/hooks/__pycache__/launch.cpython-310.pyc,, +triton/profiler/hooks/hook.py,sha256=7x5SI016VKj5a8t3LsO8tnFRR8OVN-p8E0aE_6VQnlo,4793 +triton/profiler/hooks/instrumentation.py,sha256=Z7mPpcf4er8Y8cQpv5cevB461Kb7m0GYT4IdEWUqPgs,15271 +triton/profiler/hooks/launch.py,sha256=PTKrBs85fGT-19soqcPg0DtefAYEi95T8BCEXdktOQU,1493 +triton/profiler/language.py,sha256=Y3fcIBmIYboLGfcpiQzbXbHjn96Zh7mh5Rs2bSrYSuI,2178 +triton/profiler/mode.py,sha256=-7P8jeG7SiEUOmBuAfQWv3nzTvmsHJBvgYUh-FJXD4Y,4367 +triton/profiler/profile.py,sha256=M_X_c4c_wvlWhQ4OTCP9yrP0Jxa_6atj7_gqG0aMsxY,9002 +triton/profiler/proton.py,sha256=KNPAOyUbIewyR9vU4NMrJdijrgwO7DziujTmKy7hbUw,3090 +triton/profiler/scope.py,sha256=rE-YqDftcJq2nHNBRk0uvR234Dm6WXxhi0U0PrYdDiU,3842 +triton/profiler/specs.py,sha256=3v554WIsHHREMFxofDSMEQz18i0g0W_MJZKh6Op3nqM,2341 +triton/profiler/state.py,sha256=Mp5kcU4xTwZopfXrsgjEE5bQ2HJZR3242e0KkmVw-EQ,1346 +triton/profiler/viewer.py,sha256=A2QoWxUKzQN6qBIT4a0WhY86N96-UKnrCaqxKhyJ9sY,17552 +triton/runtime/__init__.py,sha256=mKL5cqIBDUw2WO80NRCh4s1G8KYaqgM59TTAbTkPPjQ,621 +triton/runtime/__pycache__/__init__.cpython-310.pyc,, +triton/runtime/__pycache__/_allocation.cpython-310.pyc,, +triton/runtime/__pycache__/_async_compile.cpython-310.pyc,, +triton/runtime/__pycache__/autotuner.cpython-310.pyc,, +triton/runtime/__pycache__/build.cpython-310.pyc,, +triton/runtime/__pycache__/cache.cpython-310.pyc,, +triton/runtime/__pycache__/driver.cpython-310.pyc,, +triton/runtime/__pycache__/errors.cpython-310.pyc,, +triton/runtime/__pycache__/interpreter.cpython-310.pyc,, +triton/runtime/__pycache__/jit.cpython-310.pyc,, +triton/runtime/_allocation.py,sha256=jJdUXuwt3XzAcJQibYx6gB-w2Nka2kTgW1iaQ6qfkXQ,1220 +triton/runtime/_async_compile.py,sha256=57arJTL6OkUf6RHQ9SKOKHdgqWTlkNrnyT2fYR4sdjA,1703 +triton/runtime/autotuner.py,sha256=9Fhxsl-JAZDe4i5gnDnOlmcVkX4H4HYm2OKt3BmRHAc,20160 +triton/runtime/build.py,sha256=3XIgX296pjfDLH69wTbRZn7mvvJ8hzWPLVFoUBuC2Dc,3866 +triton/runtime/cache.py,sha256=A1XD40UvpFABIiRLj_-IZPIrwFxBej1vXa2p8ukvgEA,11041 +triton/runtime/driver.py,sha256=2qxsJH3SFxYDSIffj2wXMGSviEatUQLkHFcmVboeK-A,1025 +triton/runtime/errors.py,sha256=CwfJXciwel_-K3BfQfKUpLPDWrSyTnGsfJkqJojrdfQ,1052 +triton/runtime/interpreter.py,sha256=KqlylwYSzky8QEOQvZhIEd6AGhzOgyPFPLD_uTlXIcM,61789 +triton/runtime/jit.py,sha256=6oMeA-CGdWVMlKIeLmkXraS3cul_Jmws4UHnpmRe4W0,41822 +triton/testing.py,sha256=vbEQRNrOnnzRQvVVSaiZrUo8AC0XPV40GJxfvkKYLh0,20276 +triton/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +triton/tools/__pycache__/__init__.cpython-310.pyc,, +triton/tools/__pycache__/build_extern.cpython-310.pyc,, +triton/tools/__pycache__/compile.cpython-310.pyc,, +triton/tools/__pycache__/disasm.cpython-310.pyc,, +triton/tools/__pycache__/link.cpython-310.pyc,, +triton/tools/__pycache__/mxfp.cpython-310.pyc,, +triton/tools/__pycache__/ragged_tma.cpython-310.pyc,, +triton/tools/__pycache__/tensor_descriptor.cpython-310.pyc,, +triton/tools/build_extern.py,sha256=jCr-2hu3nLGBIJhCGUQ1jAyzLttughjkiPGEwRFjLR0,13673 +triton/tools/compile.py,sha256=ECNKNHnf7RZLldrv3REeXBtsGUmcjU3xqNr9tnqxeJ8,8730 +triton/tools/disasm.py,sha256=T9jiTkdK_0nI3R_4uea0zvfioYdcR-zIZwTfuucgw6g,5026 +triton/tools/extra/cuda/compile.c,sha256=TZ9V-o1seufp9PAlaiYF4miaShd1qkobJGI84rHcpYg,2163 +triton/tools/extra/cuda/compile.h,sha256=n9QKIFZTL4RSsiXtAxBP9XGSnxjyaevQQ9bBpwDsvAg,332 +triton/tools/extra/hip/compile.cpp,sha256=aJUqzL3cLuVWMP0--eATyVOT84-eLjT8j6hmrP33dwQ,1858 +triton/tools/extra/hip/compile.h,sha256=BIRh2lo4kKXWcsomk89DKx3sY3_XLMGoCyJ5HYvl3GU,348 +triton/tools/link.py,sha256=u7qtfZRLriZkAMEGNvj8YF-k1cthmLL7BwHYqBgT63E,11871 +triton/tools/mxfp.py,sha256=YQdpBrGkOVNOtnLeRjMCeVFHWkSwUubGeWsItIjO8TU,11737 +triton/tools/ragged_tma.py,sha256=2WiUvniV80LkRfA1v_uCf8uzHsLaeBy0r4ruxX8uKhc,3167 +triton/tools/tensor_descriptor.py,sha256=CFrcp_oNIwJxTt8LGHycVYzygnNwqE1AQoaJkIr7cGY,1466 diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..12e67ca29ff180ca4133d9ec4769226edb266ce4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: false +Tag: cp310-cp310-manylinux_2_27_x86_64 +Tag: cp310-cp310-manylinux_2_28_x86_64 + diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/entry_points.txt b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..d8d06486be664b9cd4a7e4e0fddbc18703fed64a --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/entry_points.txt @@ -0,0 +1,7 @@ +[console_scripts] +proton = triton.profiler.proton:main +proton-viewer = triton.profiler.viewer:main + +[triton.backends] +amd = triton.backends.amd +nvidia = triton.backends.nvidia diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/licenses/LICENSE b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..1d0238e86b1b43e123c9eb136eb0b2c0e1658d63 --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/licenses/LICENSE @@ -0,0 +1,23 @@ +/* +* Copyright 2018-2020 Philippe Tillet +* Copyright 2020-2022 OpenAI +* +* Permission is hereby granted, free of charge, to any person obtaining +* a copy of this software and associated documentation files +* (the "Software"), to deal in the Software without restriction, +* including without limitation the rights to use, copy, modify, merge, +* publish, distribute, sublicense, and/or sell copies of the Software, +* and to permit persons to whom the Software is furnished to do so, +* subject to the following conditions: +* +* The above copyright notice and this permission notice shall be +* included in all copies or substantial portions of the Software. +* +* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +*/ diff --git a/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/top_level.txt b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..a59a965090a6473278f0cf9e7fd1d3cb9cb385c9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/triton-3.5.0.dist-info/top_level.txt @@ -0,0 +1 @@ +triton diff --git a/venv/lib/python3.10/site-packages/trl/__init__.py b/venv/lib/python3.10/site-packages/trl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..44228d2092e1ec7ee2a48b99a04c6c6fe096e627 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/__init__.py @@ -0,0 +1,198 @@ +# Copyright 2020-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 sys +import warnings +from importlib.metadata import PackageNotFoundError, version +from typing import TYPE_CHECKING + +from .import_utils import _LazyModule + + +try: + __version__ = version("trl") +except PackageNotFoundError: + __version__ = "unknown" + +_import_structure = { + "scripts": ["DatasetMixtureConfig", "ScriptArguments", "TrlParser", "get_dataset", "init_zero_verbose"], + "data_utils": [ + "apply_chat_template", + "extract_prompt", + "is_conversational", + "is_conversational_from_value", + "maybe_apply_chat_template", + "maybe_convert_to_chatml", + "maybe_extract_prompt", + "maybe_unpair_preference_dataset", + "pack_dataset", + "prepare_multimodal_messages", + "prepare_multimodal_messages_vllm", + "truncate_dataset", + "unpair_preference_dataset", + ], + "models": [ + "SUPPORTED_ARCHITECTURES", + "AutoModelForCausalLMWithValueHead", + "AutoModelForSeq2SeqLMWithValueHead", + "PreTrainedModelWrapper", + "clone_chat_template", + "create_reference_model", + "setup_chat_format", + ], + "trainer": [ + "AllTrueJudge", + "BaseBinaryJudge", + "BaseJudge", + "BasePairwiseJudge", + "BaseRankJudge", + "BCOConfig", + "BCOTrainer", + "CPOConfig", + "CPOTrainer", + "DPOConfig", + "DPOTrainer", + "FDivergenceConstants", + "FDivergenceType", + "GKDConfig", + "GKDTrainer", + "GRPOConfig", + "GRPOTrainer", + "HfPairwiseJudge", + "KTOConfig", + "KTOTrainer", + "LogCompletionsCallback", + "ModelConfig", + "NashMDConfig", + "NashMDTrainer", + "OnlineDPOConfig", + "OnlineDPOTrainer", + "OpenAIPairwiseJudge", + "ORPOConfig", + "ORPOTrainer", + "PairRMJudge", + "PPOConfig", + "PPOTrainer", + "PRMConfig", + "PRMTrainer", + "RewardConfig", + "RewardTrainer", + "RLOOConfig", + "RLOOTrainer", + "SFTConfig", + "SFTTrainer", + "WinRateCallback", + "XPOConfig", + "XPOTrainer", + ], + "trainer.callbacks": [ + "BEMACallback", + "MergeModelCallback", + "RichProgressCallback", + "SyncRefModelCallback", + "WeaveCallback", + ], + "trainer.utils": ["get_kbit_device_map", "get_peft_config", "get_quantization_config"], +} + +if TYPE_CHECKING: + from .data_utils import ( + apply_chat_template, + extract_prompt, + is_conversational, + is_conversational_from_value, + maybe_apply_chat_template, + maybe_convert_to_chatml, + maybe_extract_prompt, + maybe_unpair_preference_dataset, + pack_dataset, + prepare_multimodal_messages, + prepare_multimodal_messages_vllm, + truncate_dataset, + unpair_preference_dataset, + ) + from .models import ( + SUPPORTED_ARCHITECTURES, + AutoModelForCausalLMWithValueHead, + AutoModelForSeq2SeqLMWithValueHead, + PreTrainedModelWrapper, + clone_chat_template, + create_reference_model, + setup_chat_format, + ) + from .scripts import DatasetMixtureConfig, ScriptArguments, TrlParser, get_dataset, init_zero_verbose + from .trainer import ( + AllTrueJudge, + BaseBinaryJudge, + BaseJudge, + BasePairwiseJudge, + BaseRankJudge, + BCOConfig, + BCOTrainer, + CPOConfig, + CPOTrainer, + DPOConfig, + DPOTrainer, + FDivergenceConstants, + FDivergenceType, + GKDConfig, + GKDTrainer, + GRPOConfig, + GRPOTrainer, + HfPairwiseJudge, + KTOConfig, + KTOTrainer, + LogCompletionsCallback, + ModelConfig, + NashMDConfig, + NashMDTrainer, + OnlineDPOConfig, + OnlineDPOTrainer, + OpenAIPairwiseJudge, + ORPOConfig, + ORPOTrainer, + PairRMJudge, + PPOConfig, + PPOTrainer, + PRMConfig, + PRMTrainer, + RewardConfig, + RewardTrainer, + RLOOConfig, + RLOOTrainer, + SFTConfig, + SFTTrainer, + WinRateCallback, + XPOConfig, + XPOTrainer, + ) + from .trainer.callbacks import ( + BEMACallback, + MergeModelCallback, + RichProgressCallback, + SyncRefModelCallback, + WeaveCallback, + ) + from .trainer.utils import get_kbit_device_map, get_peft_config, get_quantization_config + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + extra_objects={"__version__": __version__}, + ) diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp1.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c01b0b567bc93bf87ec136ea975b3793d273a45c --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp1.yaml @@ -0,0 +1,28 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: FSDP +downcast_bf16: 'no' +enable_cpu_affinity: false +fsdp_config: + fsdp_activation_checkpointing: false + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_backward_prefetch: BACKWARD_PRE + fsdp_cpu_ram_efficient_loading: true + fsdp_forward_prefetch: true + fsdp_offload_params: false + fsdp_reshard_after_forward: FULL_SHARD + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_sync_module_states: true + fsdp_use_orig_params: true + fsdp_version: 1 +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp2.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af498f3eced9c2434b80113f2f22d40395e0ab8a --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/fsdp2.yaml @@ -0,0 +1,25 @@ +# Requires accelerate 1.7.0 or higher +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: FSDP +downcast_bf16: 'no' +enable_cpu_affinity: false +fsdp_config: + fsdp_activation_checkpointing: false + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_cpu_ram_efficient_loading: true + fsdp_offload_params: false + fsdp_reshard_after_forward: true + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_version: 2 +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/multi_gpu.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/multi_gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..15dad9be3ba44f7c934e1ecab98a93cb83cbc79a --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/multi_gpu.yaml @@ -0,0 +1,16 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: MULTI_GPU +downcast_bf16: 'no' +gpu_ids: all +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/single_gpu.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/single_gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ebd00a067118e56f3d63ab0f24827cfea21b24b9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/single_gpu.yaml @@ -0,0 +1,16 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: "NO" +downcast_bf16: 'no' +gpu_ids: all +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero1.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d5b5f782fb30f9fcbcc8fc58262f09eaf2e10368 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero1.yaml @@ -0,0 +1,20 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + gradient_accumulation_steps: 1 + zero3_init_flag: false + zero_stage: 1 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero2.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..239b14ac3a9ae8de73122d1154bf0d71903dc15f --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero2.yaml @@ -0,0 +1,21 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: false + zero_stage: 2 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero3.yaml b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b5a1201f8a2ee8706b63f0f80c664a1fc61a7d9d --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/accelerate_configs/zero3.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: true + zero3_save_16bit_model: true + zero_stage: 3 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/venv/lib/python3.10/site-packages/trl/cli.py b/venv/lib/python3.10/site-packages/trl/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..199b1c267032d62796a80c252c5a7fa4a27c3dd6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/cli.py @@ -0,0 +1,162 @@ +# Copyright 2020-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.resources as resources +import os +import sys + +import torch +from accelerate import logging +from accelerate.commands.launch import launch_command, launch_command_parser + +from .scripts.dpo import make_parser as make_dpo_parser +from .scripts.env import print_env +from .scripts.grpo import make_parser as make_grpo_parser +from .scripts.kto import make_parser as make_kto_parser +from .scripts.reward import make_parser as make_reward_parser +from .scripts.rloo import make_parser as make_rloo_parser +from .scripts.sft import make_parser as make_sft_parser +from .scripts.utils import TrlParser +from .scripts.vllm_serve import main as vllm_serve_main +from .scripts.vllm_serve import make_parser as make_vllm_serve_parser + + +logger = logging.get_logger(__name__) + + +def main(): + parser = TrlParser(prog="TRL CLI", usage="trl", allow_abbrev=False) + + # Add the subparsers + subparsers = parser.add_subparsers(help="available commands", dest="command", parser_class=TrlParser) + + # Add the subparsers for every script + make_dpo_parser(subparsers) + subparsers.add_parser("env", help="Print the environment information") + make_grpo_parser(subparsers) + make_kto_parser(subparsers) + make_reward_parser(subparsers) + make_rloo_parser(subparsers) + make_sft_parser(subparsers) + make_vllm_serve_parser(subparsers) + + # Parse the arguments; the remaining ones (`launch_args`) are passed to the 'accelerate launch' subparser. + # Duplicates may occur if the same argument is provided in both the config file and CLI. + # For example: launch_args = `["--num_processes", "4", "--num_processes", "8"]`. + # Deduplication and precedence (CLI over config) are handled later by launch_command_parser. + args, launch_args = parser.parse_args_and_config(return_remaining_strings=True) + + # Replace `--accelerate_config foo` with `--config_file trl/accelerate_configs/foo.yaml` if it is present in the + # launch_args. It allows the user to use predefined accelerate configs from the `trl` package. + if "--accelerate_config" in launch_args: + # Get the index of the '--accelerate_config' argument and the corresponding config name + config_index = launch_args.index("--accelerate_config") + config_name = launch_args[config_index + 1] + + # If the config_name correspond to a path in the filesystem, we don't want to override it + if os.path.isfile(config_name): + accelerate_config_path = config_name + elif resources.files("trl.accelerate_configs").joinpath(f"{config_name}.yaml").exists(): + # Get the predefined accelerate config path from the package resources + accelerate_config_path = resources.files("trl.accelerate_configs").joinpath(f"{config_name}.yaml") + else: + raise ValueError( + f"Accelerate config {config_name} is neither a file nor a valid config in the `trl` package. " + "Please provide a valid config name or a path to a config file." + ) + + # Remove '--accelerate_config' and its corresponding config name + launch_args.pop(config_index) + launch_args.pop(config_index) + + # Insert '--config_file' and the absolute path to the front of the list + launch_args = ["--config_file", str(accelerate_config_path)] + launch_args + + if args.command == "dpo": + # Get the default args for the launch command + dpo_training_script = resources.files("trl.scripts").joinpath("dpo.py") + args = launch_command_parser().parse_args([str(dpo_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "dpo" + launch_command(args) # launch training + + elif args.command == "env": + print_env() + + elif args.command == "grpo": + # Get the default args for the launch command + grpo_training_script = resources.files("trl.scripts").joinpath("grpo.py") + args = launch_command_parser().parse_args([str(grpo_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "grpo" + launch_command(args) # launch training + + elif args.command == "kto": + # Get the default args for the launch command + kto_training_script = resources.files("trl.scripts").joinpath("kto.py") + args = launch_command_parser().parse_args([str(kto_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "kto" + launch_command(args) # launch training + + elif args.command == "reward": + # Get the default args for the launch command + reward_training_script = resources.files("trl.scripts").joinpath("reward.py") + args = launch_command_parser().parse_args([str(reward_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "reward" + launch_command(args) # launch training + + elif args.command == "rloo": + # Get the default args for the launch command + rloo_training_script = resources.files("trl.scripts").joinpath("rloo.py") + args = launch_command_parser().parse_args([str(rloo_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "rloo" + launch_command(args) # launch training + + elif args.command == "sft": + # Get the path to the training script + sft_training_script = resources.files("trl.scripts").joinpath("sft.py") + + # This simulates running: `accelerate launch sft.py `. + # Note that the training script args may include launch-related arguments (e.g., `--num_processes`), + # but we rely on the script to ignore any that don't apply to it. + training_script_args = sys.argv[2:] # Remove "trl" and "sft" + args = launch_command_parser().parse_args(launch_args + [str(sft_training_script)] + training_script_args) + launch_command(args) # launch training + + elif args.command == "vllm-serve": + (script_args,) = parser.parse_args_and_config() + + # Known issue: Using DeepSpeed with tensor_parallel_size=1 and data_parallel_size>1 may cause a crash when + # launched via the CLI. Suggest running the module directly. + # More information: https://github.com/vllm-project/vllm/issues/17079 + if script_args.tensor_parallel_size == 1 and script_args.data_parallel_size > 1 and torch.cuda.is_available(): + logger.warning( + "Detected configuration: tensor_parallel_size=1 and data_parallel_size>1. This setup is known to " + "cause a crash when using the `trl vllm-serve` CLI entry point. As a workaround, please run the " + "server using the module path instead: `python -m trl.scripts.vllm_serve`", + ) + + vllm_serve_main(script_args) + + +if __name__ == "__main__": + main() diff --git a/venv/lib/python3.10/site-packages/trl/data_utils.py b/venv/lib/python3.10/site-packages/trl/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bafbca3a11825b9ba744d53c638107603a2d0147 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/data_utils.py @@ -0,0 +1,937 @@ +# Copyright 2020-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 copy +from collections import defaultdict, deque +from collections.abc import Callable, Sequence +from itertools import takewhile +from typing import Any, TypeVar + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.types +from datasets import Dataset, DatasetDict +from transformers import PreTrainedTokenizerBase, ProcessorMixin + + +DatasetType = TypeVar("DatasetType", Dataset, DatasetDict) + + +def prepare_multimodal_messages(messages: list[dict[str, Any]], images: list) -> list[dict[str, Any]]: + # docstyle-ignore # because is not parsable in the code block + """ + Convert messages into a structured multimodal format and inject the provided images into the message contents. + + Args: + messages (`list[dict[str, Any]]`): + Messages with `"role"` and `"content"`. Content may be a raw string before transformation. List of messages + a `"role"` key (`"system"`, `"user"`, or `"assistant"`) and a `"content"` key containing either a string or + a list of structured blocks if already prepared. + images (`list`): + List of image objects to insert. + + Returns: + `list[dict[str, Any]]`: A deep-copied list of messages where every `"content"` value is a list of structured + content blocks, and all `"image"` placeholders are populated with the corresponding image objects. + + Notes: + - When the input `messages` isn't already in the structured format, (i.e., all `"content"` values are strings), + the function transforms them into the structured format by wrapping text in `{"type": "text", "text": ...}` + and inserting `{"type": "image"}` placeholders for the images *before* the first user message. + - When the input `messages` is already in the structured format (i.e., all `"content"` values are lists of + structured blocks), the function only fills in the actual images in the existing `{"type": "image"}` + placeholders. If the number of placeholders does not match the number of provided images, an error is raised. + + Example: + ```python + # Input + [ + {"role": "user", "content": "What's in this image?"}, + {"role": "assistant", "content": "It looks like a cat."}, + ] + + # Output, one image provided + [ + {"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What's in this image?"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "It looks like a cat."}]}, + ] + ``` + """ + + messages = copy.deepcopy(messages) # avoid modifying the original messages + + # First, convert all messages to the structured format if needed, and insert image placeholders if needed + images_included = False + for message in messages: + if message["role"] == "system": + if isinstance(message["content"], str): # if already prepared, the content will be a list + message["content"] = [{"type": "text", "text": message["content"]}] + elif message["role"] == "user": + if isinstance(message["content"], str) and not images_included: + image_entries = [{"type": "image"} for _ in range(len(images))] + message["content"] = [*image_entries, {"type": "text", "text": message["content"]}] + images_included = True + elif isinstance(message["content"], str) and images_included: + message["content"] = [{"type": "text", "text": message["content"]}] + elif message["role"] == "assistant": + if isinstance(message["content"], str): + message["content"] = [{"type": "text", "text": message["content"]}] + else: + raise ValueError(f"Invalid role in message: {message['role']}. Expected 'user', 'assistant', or 'system'.") + + # Then, check that the number of image placeholders matches the number of images provided + num_placeholders = sum(sum(1 for part in message["content"] if part["type"] == "image") for message in messages) + if num_placeholders != len(images): + raise ValueError( + f"Number of images provided ({len(images)}) does not match number of image placeholders ({num_placeholders})." + ) + + # Then, fill in the actual images in the placeholders + img_idx = 0 + for message in messages: + for part in message["content"]: + if part["type"] == "image": + part["image"] = images[img_idx] + img_idx += 1 + + return messages + + +def prepare_multimodal_messages_vllm(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + # docstyle-ignore # because is not parsable in the code block + """ + Convert structured multimodal messages into a format compatible with vLLM. Replaces `"type": "image"` blocks with + `"type": "image_pil"` blocks, and `"image": Image` with `"image_pil": Image`. + + Args: + messages (`list[dict[str, Any]]`): + Messages with `"role"` and `"content"`. Content is expected to be a list of structured blocks. + + Returns: + `list[dict[str, Any]]`: + A deep-copied list of messages compatible with vLLM's expected input format. + + Example: + ```python + # Input + [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What's in this image?"}]}] + + # Output + [{"role": "user", "content": [{"type": "image_pil", "image_pil": }, {"type": "text", "text": "What's in this image?"}]}] + ``` + """ + messages = copy.deepcopy(messages) # avoid modifying the original messages + for message in messages: + if isinstance(message["content"], list): + for part in message["content"]: + if part["type"] == "image": + part["type"] = "image_pil" # vLLM expects 'image_pil' key for images + part["image_pil"] = part.pop("image") + return messages + + +def is_conversational(example: dict[str, Any]) -> bool: + r""" + Check if the example is in a conversational format. + + Args: + example (`dict[str, Any]`): + A single data entry of a dataset. The example can have different keys depending on the dataset type. + + Returns: + `bool`: + `True` if the data is in a conversational format, `False` otherwise. + + Examples: + + ```python + >>> example = {"prompt": [{"role": "user", "content": "What color is the sky?"}]} + >>> is_conversational(example) + True + + >>> example = {"prompt": "The sky is"} + >>> is_conversational(example) + False + ``` + """ + supported_keys = ["prompt", "chosen", "rejected", "completion", "messages"] + example_keys = {key for key in example.keys() if key in supported_keys} + + # It must have one of the supported keys + if example_keys: + key = example_keys.pop() # take the first supported key + maybe_messages = example[key] + # It must be a list of messages + if isinstance(maybe_messages, list): + maybe_message = maybe_messages[0] + # Each message must a list of dictionaries with keys "role" and "content" + if isinstance(maybe_message, dict) and "role" in maybe_message and "content" in maybe_message: + return True + + return False + + +def apply_chat_template( + example: dict[str, list[dict[str, str]]], + tokenizer: PreTrainedTokenizerBase | ProcessorMixin, + tools: list[dict | Callable] | None = None, + **template_kwargs, +) -> dict[str, str]: + r""" + Apply a chat template to a conversational example along with the schema for a list of functions in `tools`. + + For more details, see [`maybe_apply_chat_template`]. + """ + # Check that the example has the correct keys + supported_keys = ["prompt", "chosen", "rejected", "completion", "messages", "label"] + example_keys = {key for key in example.keys() if key in supported_keys} + if example_keys not in [ + {"messages"}, # language modeling + {"prompt"}, # prompt-only + {"prompt", "completion"}, # prompt-completion + {"prompt", "chosen", "rejected"}, # preference + {"chosen", "rejected"}, # preference with implicit prompt + {"prompt", "completion", "label"}, # unpaired preference + ]: + raise KeyError(f"Invalid keys in the example: {example_keys}") + + # Apply the chat template to the whole conversation + if "messages" in example: + messages = tokenizer.apply_chat_template( + example["messages"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + + # Apply the chat template to the prompt, adding the generation prompt + if "prompt" in example: + last_role = example["prompt"][-1]["role"] + if last_role == "user": + add_generation_prompt = True + continue_final_message = False + elif last_role == "assistant": + add_generation_prompt = False + continue_final_message = True + else: + raise ValueError(f"Invalid role in the last message: {last_role}") + prompt = tokenizer.apply_chat_template( + example["prompt"], + tools=tools, + continue_final_message=continue_final_message, + tokenize=False, + add_generation_prompt=add_generation_prompt, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + + # Apply the chat template to the entire prompt + completion + if "prompt" in example: # explicit prompt and prompt-completion case + if "chosen" in example: + prompt_chosen = tokenizer.apply_chat_template( + example["prompt"] + example["chosen"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + # DeepSeek-R1 inserts a token when using `add_generation_prompt`, which can cause discrepancies + # between the prompt alone and the combined prompt+completion. To ensure consistency, we extract the + # common prefix between the two. In most cases, this is a no-op. + prompt = "".join(x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_chosen, strict=False))) + + chosen = prompt_chosen[len(prompt) :] + if "rejected" in example and "prompt" in example: # explicit prompt + prompt_rejected = tokenizer.apply_chat_template( + example["prompt"] + example["rejected"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + # Handle DeepSeek-R1 token, see the above comment for details + prompt = "".join( + x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_rejected, strict=False)) + ) + rejected = prompt_rejected[len(prompt) :] + if "completion" in example: + prompt_completion = tokenizer.apply_chat_template( + example["prompt"] + example["completion"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + # Handle DeepSeek-R1 token, see the above comment for details + prompt = "".join( + x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_completion, strict=False)) + ) + completion = prompt_completion[len(prompt) :] + else: # implicit prompt case + if "chosen" in example: + chosen = tokenizer.apply_chat_template( + example["chosen"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + if "rejected" in example: + rejected = tokenizer.apply_chat_template( + example["rejected"], + tools=tools, + tokenize=False, + **example.get("chat_template_kwargs", {}), + **template_kwargs, + ) + + # Extract the completion by removing the prompt part from the prompt-completion string + output = {} + if "messages" in example: + output["text"] = messages + if "prompt" in example: + output["prompt"] = prompt + if "chosen" in example: + output["chosen"] = chosen + if "rejected" in example: + output["rejected"] = rejected + if "completion" in example: + output["completion"] = completion + if "label" in example: + output["label"] = example["label"] + + return output + + +def maybe_apply_chat_template( + example: dict[str, list[dict[str, str]]], + tokenizer: PreTrainedTokenizerBase, + tools: list[dict | Callable] | None = None, + **template_kwargs: Any, +) -> dict[str, str]: + r""" + If the example is in a conversational format, apply a chat template to it. + + Args: + example (`dict[str, list[dict[str, str]]`): + Dictionary representing a single data entry of a conversational dataset. Each data entry can have different + keys depending on the dataset type. The supported dataset types are: + + - Language modeling dataset: `"messages"`. + - Prompt-only dataset: `"prompt"`. + - Prompt-completion dataset: `"prompt"` and `"completion"`. + - Preference dataset: `"prompt"`, `"chosen"`, and `"rejected"`. + - Preference dataset with implicit prompt: `"chosen"` and `"rejected"`. + - Unpaired preference dataset: `"prompt"`, `"completion"`, and `"label"`. + + For keys `"messages"`, `"prompt"`, `"chosen"`, `"rejected"`, and `"completion"`, the values are lists of + messages, where each message is a dictionary with keys `"role"` and `"content"`. Additionally, the example + may contain a `"chat_template_kwargs"` key, which is a dictionary of additional keyword arguments to pass + to the chat template renderer. + tokenizer ([`~transformers.PreTrainedTokenizerBase`]): + Tokenizer to apply the chat template with. + tools (`list[dict | Callable]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not support + function calling, this argument will have no effect. + **template_kwargs (`Any`, *optional*): + Additional kwargs to pass to the template renderer. Will be accessible by the chat template. + + Returns: + `dict[str, str]`: + Formatted example with the chat template applied. + + Notes: + - This function does not alter the keys, except for Language modeling dataset, where `"messages"` is replaced + by `"text"`. + + - In case of prompt-only data, if the last role is `"user"`, the generation prompt is added to the prompt. + Else, if the last role is `"assistant"`, the final message is continued. + + Example: + + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct") + >>> example = { + ... "prompt": [{"role": "user", "content": "What color is the sky?"}], + ... "completion": [{"role": "assistant", "content": "It is blue."}], + ... } + >>> apply_chat_template(example, tokenizer) + {'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n', 'completion': 'It is blue.<|end|>\n'} + ``` + """ + if is_conversational(example): + return apply_chat_template(example, tokenizer, tools, **template_kwargs) + else: + return example + + +def _unpair_row(examples: list[dict[str, list[dict[str, str]]]]) -> list[dict[str, list[dict[str, str]]]]: + batch_size = len(examples["chosen"]) + new_rows = { + "completion": examples["chosen"] + examples["rejected"], + "label": [True] * batch_size + [False] * batch_size, + } + if "prompt" in examples: + new_rows["prompt"] = examples["prompt"] + examples["prompt"] + return new_rows + + +def unpair_preference_dataset( + dataset: DatasetType, num_proc: int | None = None, desc: str | None = None +) -> DatasetType: + r""" + Unpair a preference dataset. + + Args: + dataset ([`~datasets.Dataset`] or [`~datasets.DatasetDict`]): + Preference dataset to unpair. The dataset must have columns `"chosen"`, `"rejected"` and optionally + `"prompt"`. + num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + desc (`str`, *optional*): + Meaningful description to be displayed alongside with the progress bar while mapping examples. + + Returns: + [`~datasets.Dataset`]: The unpaired preference dataset. + + Example: + + ```python + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "prompt": ["The sky is", "The sun is"], + ... "chosen": [" blue.", "in the sky."], + ... "rejected": [" green.", " in the sea."], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = unpair_preference_dataset(dataset) + >>> dataset + Dataset({ + features: ['prompt', 'completion', 'label'], + num_rows: 4 + }) + + >>> dataset[0] + {'prompt': 'The sky is', 'completion': ' blue.', 'label': True} + ``` + """ + return dataset.map(_unpair_row, batched=True, remove_columns=["chosen", "rejected"], num_proc=num_proc, desc=desc) + + +def maybe_unpair_preference_dataset( + dataset: DatasetType, num_proc: int | None = None, desc: str | None = None +) -> DatasetType: + r""" + Unpair a preference dataset if it is paired. + + Args: + dataset ([`~datasets.Dataset`] or [`~datasets.DatasetDict`]): + Preference dataset to unpair. The dataset must have columns `"chosen"`, `"rejected"` and optionally + `"prompt"`. + num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + desc (`str`, *optional*): + Meaningful description to be displayed alongside with the progress bar while mapping examples. + + Returns: + [`~datasets.Dataset`] or [`~datasets.DatasetDict`]: The unpaired preference dataset if it was paired, otherwise + the original dataset. + + Example: + + ```python + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "prompt": ["The sky is", "The sun is"], + ... "chosen": [" blue.", "in the sky."], + ... "rejected": [" green.", " in the sea."], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = unpair_preference_dataset(dataset) + >>> dataset + Dataset({ + features: ['prompt', 'completion', 'label'], + num_rows: 4 + }) + + >>> dataset[0] + {'prompt': 'The sky is', 'completion': ' blue.', 'label': True} + ``` + """ + if isinstance(dataset, DatasetDict): + column_names = dataset[list(dataset.keys())[0]].column_names + else: + column_names = dataset.column_names + if "chosen" in column_names and "rejected" in column_names: + return unpair_preference_dataset(dataset, num_proc=num_proc, desc=desc) + else: + return dataset + + +def extract_prompt(example: dict[str, Sequence]) -> dict[str, Sequence]: + r""" + Extracts the shared prompt from a preference data example, where the prompt is implicit within both the chosen and + rejected completions. + + For more details, see [`maybe_extract_prompt`]. + """ + for idx in range(min(len(example["chosen"]), len(example["rejected"]))): + if example["chosen"][idx] != example["rejected"][idx]: + if example["chosen"][idx - 1] == " ": # remove space before the prompt + idx -= 1 + break + return { + "prompt": example["chosen"][:idx], + "chosen": example["chosen"][idx:], + "rejected": example["rejected"][idx:], + } + + +def maybe_extract_prompt(example: dict[str, list]) -> dict[str, list]: + r""" + Extracts the shared prompt from a preference data example, where the prompt is implicit within both the chosen and + rejected completions. + + If the example already contains a `"prompt"` key, the function returns the example as is. Else, the function + identifies the longest common sequence (prefix) of conversation turns between the "chosen" and "rejected" + completions and extracts this as the prompt. It then removes this prompt from the respective "chosen" and + "rejected" completions. + + Args: + example (`dict[str, list]`): + A dictionary representing a single data entry in the preference dataset. It must contain the keys + `"chosen"` and `"rejected"`, where each value is either conversational or standard (`str`). + + Returns: + `dict[str, list]`: A dictionary containing: + - `"prompt"`: The longest common prefix between the "chosen" and "rejected" completions. + - `"chosen"`: The remainder of the "chosen" completion, with the prompt removed. + - `"rejected"`: The remainder of the "rejected" completion, with the prompt removed. + + Examples: + + ```python + >>> example = { + ... "chosen": [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is blue."}, + ... ], + ... "rejected": [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is green."}, + ... ], + ... } + >>> extract_prompt(example) + {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}], + 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}], + 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]} + ``` + + Or, with the `map` method of [`~datasets.Dataset`]: + + ```python + >>> from trl import extract_prompt + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "chosen": [ + ... [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is blue."}, + ... ], + ... [ + ... {"role": "user", "content": "Where is the sun?"}, + ... {"role": "assistant", "content": "In the sky."}, + ... ], + ... ], + ... "rejected": [ + ... [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is green."}, + ... ], + ... [ + ... {"role": "user", "content": "Where is the sun?"}, + ... {"role": "assistant", "content": "In the sea."}, + ... ], + ... ], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = dataset.map(extract_prompt) + >>> dataset[0] + {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}], + 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}], + 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]} + ``` + """ + # Some dataset add a `"prompt"` column, even though the prompt is implicit and included in the "chosen" and + # "rejected" completions. E.g.: + # {"prompt": "What color is the sky?", + # "chosen": [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}], + # "rejected": [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}]} + # That's why we check if the prompt is also conversational before deciding not to extract it. + if "chosen" not in example or "rejected" not in example: # not a preference example + return example + if "prompt" in example: + # Both conversational or both non-conversational + chosen_conv = is_conversational({"chosen": example["chosen"]}) + prompt_conv = is_conversational({"prompt": example["prompt"]}) + if (chosen_conv and prompt_conv) or (not chosen_conv and not prompt_conv): + return example + return extract_prompt({"chosen": example["chosen"], "rejected": example["rejected"]}) + + +class _SegmentTree: + """ + A segment tree data structure that, when initialized as `_SegmentTree(maxval)`, efficiently finds the next larger + value for a given input within the range [1, maxval]. + + See [Fewer Truncations Improve Language Modeling](https://arxiv.org/abs/2404.10830) for more details. + """ + + def __init__(self, maxval: int): + self.maxval = maxval + # For non-power-of-2 values, we need to round up to the next power of 2 for the tree size + self.tree_size = 1 << (maxval - 1).bit_length() + self.tree = [0] * (2 * self.tree_size) + + def add(self, val): + assert 0 < val <= self.maxval + i = self.tree_size + val - 1 + self.tree[i] = val + while i > 1: + i >>= 1 + left, right = self.tree[i << 1], self.tree[(i << 1) + 1] + # Compare the values using if-else otherwise repeated calls to `builtins.max` become the bottleneck + self.tree[i] = left if left >= right else right + + def remove(self, val): + assert 0 < val <= self.maxval + i = self.tree_size + val - 1 + self.tree[i] = 0 + while i > 1: + i >>= 1 + left, right = self.tree[i << 1], self.tree[(i << 1) + 1] + # Compare the values using if-else otherwise repeated calls to `builtins.max` become the bottleneck + self.tree[i] = left if left >= right else right + + def search(self, val): + assert 0 < val <= self.maxval + i = 1 + while i < self.tree_size: + if self.tree[i << 1] >= val: + i = i << 1 + else: + i = (i << 1) + 1 + return self.tree[i] + + +def _pack_bfd(examples: pa.Table, seq_length: int) -> pa.Table: + """Pack sequences in a pyarrow Table using Best Fit Decreasing strategy.""" + columns = [] + list_column_idx = None + for idx, column in enumerate(examples.columns): + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + column = pc.list_slice(column, 0, seq_length) + if list_column_idx is None: + list_column_idx = idx + columns.append(column) + examples = pa.Table.from_arrays(columns, names=examples.column_names) + + ids = np.arange(len(examples)) + assert list_column_idx is not None + lengths = pc.list_value_length(examples[list_column_idx]).combine_chunks() + examples = examples.append_column("seq_lengths", lengths) # Allows us to later construct `position_ids` + lengths = pc.make_struct(lengths, ids) + lengths = lengths.sort("descending", by=0) + + segment_tree = _SegmentTree(seq_length) + segment_tree.add(seq_length) # the max, `seq_length` bin is always available + space_to_bin = defaultdict(deque) + + # Bin is represented as a dict (of example ids and sum of their lengths) to allow in-place updates + bins: list[dict] = [] + for length, idx in zip(lengths.field(0).to_numpy(), lengths.field(1).to_numpy(), strict=True): + space = segment_tree.search(length) + + if space < seq_length: + # Use existing bin with exactly this amount of space + bin = space_to_bin[space].popleft() + else: + # Create a new bin + bin = {"ids": [], "length": 0} + bins.append(bin) + + bin["ids"].append(idx) + bin["length"] += length + if space < seq_length and not space_to_bin[space]: + segment_tree.remove(space) + + space = space - length + space_to_bin[space].append(bin) + if space > 0: + segment_tree.add(space) + + examples = pc.take(examples, [id_ for bin in bins for id_ in bin["ids"]]) + offsets = np.array([0] + [bin["length"] for bin in bins]) + offsets = np.cumsum(offsets) + + assert all( + column.num_chunks == 1 for column in examples.columns + ) # `pc.take` returns a ChunkedArray with a single chunk + + lengths = examples["seq_lengths"].chunks[0] + examples = examples.drop_columns("seq_lengths") + lengths = pa.ListArray.from_arrays(np.cumsum([0] + [len(bin["ids"]) for bin in bins], dtype=np.int32), lengths) + + columns = [] + for column in examples.columns: + column = column.chunks[0] + if pa.types.is_list(column.type) or pa.types.is_large_list(column.type): + dtype = column.offsets.type.to_pandas_dtype() + column = type(column).from_arrays(offsets.astype(dtype), column.values) + columns.append(column) + return pa.Table.from_arrays(columns + [lengths], names=examples.column_names + ["seq_lengths"]) + + +def _pack_wrapped(examples: pa.Table, seq_length: int) -> pa.Table: + """Pack sequences in a pyarrow Table using a wrapped strategy.""" + columns = [] + for column in examples.columns: + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + if isinstance(column, pa.ChunkedArray): + column = column.combine_chunks() + offsets, values = column.offsets, column.values + values = values[offsets[0].as_py() : offsets[-1].as_py()] + num_elements = len(values) + dtype = offsets.type.to_pandas_dtype() # np.int32 or np.int64 + offsets = np.arange(0, num_elements, seq_length, dtype=dtype) + offsets = np.concatenate((offsets, [num_elements])) + column = type(column).from_arrays(offsets, values) + columns.append(column) + return pa.Table.from_arrays(columns, names=examples.column_names) + + +def pack_dataset( + dataset: DatasetType, seq_length: int, strategy: str = "bfd", map_kwargs: dict[str, Any] | None = None +) -> DatasetType: + r""" + Pack sequences in a dataset into chunks of size `seq_length`. + + Args: + dataset ([`~datasets.Dataset`] or [`~datasets.DatasetDict`]): + Dataset to pack + seq_length (`int`): + Target sequence length to pack to. + strategy (`str`, *optional*, defaults to `"bfd"`): + Packing strategy to use. Can be either: + + - `"bfd"` (Best Fit Decreasing): Slower but preserves sequence boundaries. Sequences are never cut in the + middle. + - `"wrapped"`: Faster but more aggressive. Ignores sequence boundaries and will cut sequences in the middle + to completely fill each packed sequence with data. + map_kwargs (`dict`, *optional*): + Additional keyword arguments to pass to the dataset's map method when packing examples. + + Returns: + [`~datasets.Dataset`] or [`~datasets.DatasetDict`]: The dataset with packed sequences. The number of examples + may decrease as sequences are combined. + + Example: + ```python + >>> from datasets import Dataset + >>> from trl import pack_dataset + + >>> examples = { + ... "input_ids": [[1, 2, 3], [4, 5], [6, 7, 8], [9]], + ... "attention_mask": [[1, 1, 0], [1, 0], [1, 0, 0], [1]], + ... } + >>> dataset = Dataset.from_dict(examples) + >>> packed_dataset = pack_dataset(dataset, seq_length=4, strategy="bfd") + >>> packed_dataset[:] + {'input_ids': [[1, 2, 3, 9], [6, 7, 8], [4, 5]], + 'attention_mask': [[1, 1, 0, 1], [1, 0, 0], [1, 0]], + 'seq_lengths': [[3, 1], [3], [2]]} + ``` + """ + if map_kwargs is None: + map_kwargs = {} + # Fast packing with pyarrow + dataset = dataset.with_format("arrow") + if strategy == "bfd": + dataset = dataset.map(_pack_bfd, batched=True, fn_kwargs={"seq_length": seq_length}, **map_kwargs) + elif strategy == "wrapped": + dataset = dataset.map(_pack_wrapped, batched=True, fn_kwargs={"seq_length": seq_length}, **map_kwargs) + else: + raise ValueError(f"Invalid packing strategy: {strategy}. Use 'bfd' or 'wrapped'.") + dataset = dataset.with_format(None) + return dataset + + +def truncate_dataset(dataset: DatasetType, max_length: int, map_kwargs: dict[str, Any] | None = None) -> DatasetType: + r""" + Truncate sequences in a dataset to a specified `max_length`. + + Args: + dataset ([`~datasets.Dataset`] or [`~datasets.DatasetDict`]): + Dataset to truncate. + max_length (`int`): + Maximum sequence length to truncate to. + map_kwargs (`dict`, *optional*): + Additional keyword arguments to pass to the dataset's map method when truncating examples. + + Returns: + [`~datasets.Dataset`] or [`~datasets.DatasetDict`]: The dataset with truncated sequences. + + Example: + ```python + >>> from datasets import Dataset + + >>> examples = { + ... "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], + ... "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], + ... } + >>> dataset = Dataset.from_dict(examples) + >>> truncated_dataset = truncate_dataset(dataset, max_length=2) + >>> truncated_dataset[:] + {'input_ids': [[1, 2], [4, 5], [8]], + 'attention_mask': [[0, 1], [0, 0], [1]]} + ``` + """ + if map_kwargs is None: + map_kwargs = {} + if isinstance(dataset, Dataset): + # Fast truncation with pyarrow + def truncate(examples): + truncated_columns = [] + for column in examples.columns: + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + column = pc.list_slice(column, 0, max_length) + truncated_columns.append(column) + return pa.Table.from_arrays(truncated_columns, names=examples.column_names) + + dataset = dataset.with_format("arrow") + dataset = dataset.map(truncate, batched=True, **map_kwargs) + dataset = dataset.with_format(None) + else: + + def truncate(examples): + truncated_examples = {} + for key, column in examples.items(): + if column and isinstance(column[0], list): + column = [val[:max_length] for val in column] + truncated_examples[key] = column + return truncated_examples + + dataset = dataset.map( + truncate, + batched=True, + **map_kwargs, + ) + return dataset + + +def is_conversational_from_value(example: dict[str, Any]) -> bool: + r""" + Check if the example is in a conversational format (from/value). Note that this format isn't recommended. Prefer + the ChatML format (role/content) + + Args: + example (`dict[str, Any]`): + A single data entry of a dataset. The example can have different keys depending on the dataset type. + + Returns: + `bool`: + `True` if the data is in a conversational Chatformat, `False` otherwise. + + Examples: + + ```python + >>> example = {"conversations": [{"from": "user", "value": "What color is the sky?"}]} + >>> is_conversational_from_value(example) + True + + >>> example = {"conversations": [{"role": "user", "content": "What color is the sky?"}]} + >>> is_conversational_from_value(example) + False + + >>> example = {"conversations": "The sky is"} + >>> is_conversational_from_value(example) + False + ``` + """ + maybe_messages = example.get("conversations") + # It must be a list of messages + if isinstance(maybe_messages, list): + maybe_message = maybe_messages[0] + # Each message must a list of dictionaries with keys "from" and "value" + if isinstance(maybe_message, dict) and "from" in maybe_message and "value" in maybe_message: + return True + + return False + + +def maybe_convert_to_chatml(example: dict[str, list]) -> dict[str, list]: + """ + Convert a conversational dataset with fields `from` and `value` to ChatML format. + + This function modifies conversational data to align with OpenAI's ChatML format: + - Replaces the key `"from"` with `"role"` in message dictionaries. + - Replaces the key `"value"` with `"content"` in message dictionaries. + - Renames `"conversations"` to `"messages"` for consistency with ChatML. + + Args: + example (`dict[str, list]`): + A single data entry containing a list of messages. + + Returns: + `dict[str, list]`: + Example reformatted to ChatML style. + + Example: + ```python + >>> from trl import maybe_convert_to_chatml + + >>> example = { + ... "conversations": [ + ... {"from": "user", "value": "What color is the sky?"}, + ... {"from": "assistant", "value": "It is blue."}, + ... ] + ... } + >>> maybe_convert_to_chatml(example) + {'messages': [{'role': 'user', 'content': 'What color is the sky?'}, + {'role': 'assistant', 'content': 'It is blue.'}]} + ``` + """ + # List of possible keys containing message lists + for key in ["prompt", "completion", "chosen", "rejected", "messages", "conversations"]: + if key in example and isinstance(example[key], list): + messages = example[key] + for message in messages: + if isinstance(message, dict): + if "from" in message: + message["role"] = message.pop("from") + if "value" in message: + message["content"] = message.pop("value") + + # Rename "conversations" to "messages" + if "conversations" in example: + example["messages"] = example.pop("conversations") + + return example diff --git a/venv/lib/python3.10/site-packages/trl/experimental/__init__.py b/venv/lib/python3.10/site-packages/trl/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..28f8f42eaf28a795dc350fbeeeff58e987b86050 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/__init__.py @@ -0,0 +1,34 @@ +# Copyright 2020-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. + +""" +Experimental submodule for TRL. + +This submodule contains unstable or incubating features. Anything here may change (or be removed) in any release +without deprecation. Use at your own risk. + +To silence this notice set environment variable TRL_EXPERIMENTAL_SILENCE=1. +""" + +import os +import warnings + + +if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "You are importing from 'trl.experimental'. APIs here are unstable and may change or be removed without " + "notice. Silence this warning by setting environment variable TRL_EXPERIMENTAL_SILENCE=1.", + UserWarning, + stacklevel=2, + ) diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__init__.py b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..612a7a5452b7a99eb1c7751e6506287482811992 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2020-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 .gfpo_config import GFPOConfig +from .gfpo_trainer import GFPOTrainer diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bd8978526fcd1be0df4d0d81962652227c7c9c33 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddcc68f5cb6c7738bce0d097cbee9a2e46dd6192 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a427b43b379ef75f63afdd9c2c2eb011f9a4df1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/__pycache__/gfpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_config.py b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b67a995db16ea559e387e4606401f88e7e420684 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_config.py @@ -0,0 +1,35 @@ +# Copyright 2020-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 dataclasses import dataclass, field + +from ...trainer.grpo_config import GRPOConfig as _GRPOConfig + + +@dataclass +class GFPOConfig(_GRPOConfig): + num_remains_in_group: int | None = field( + default=None, + metadata={ + "help": "number inputs remains after group filter function, `'num_remains_in_group'` must be >=2 if given." + }, + ) + + def __post_init__(self): + super().__post_init__() + + if self.num_remains_in_group is not None and self.num_remains_in_group >= self.num_generations: + raise ValueError( + f"Number remains in Group {self.num_remains_in_group} must be less than num_generations : {self.num_generations}." + ) diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_trainer.py b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..37de2b8442b565428bc4cd621b2d1e5905f49c92 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gfpo/gfpo_trainer.py @@ -0,0 +1,390 @@ +# Copyright 2020-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 logging +from collections.abc import Callable +from typing import Any + +import torch +from accelerate.utils import gather_object + +from ...data_utils import apply_chat_template, is_conversational, prepare_multimodal_messages +from ...trainer.grpo_trainer import GRPOTrainer as _GRPOTrainer +from ...trainer.utils import nanmax, nanmin, nanstd, pad + + +logger = logging.getLogger(__name__) + +GroupFilterFunc = Callable[[list[list[Any]], list[list[Any]]], list[list[float]]] + + +class GFPOTrainer(_GRPOTrainer): + def __init__( + self, + model, + reward_funcs, + args=None, + train_dataset=None, + eval_dataset=None, + processing_class=None, + reward_processing_classes=None, + group_filter_func=None, + callbacks=None, + optimizers=(None, None), + peft_config=None, + ): + super().__init__( + model=model, + reward_funcs=reward_funcs, + args=args, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_classes=reward_processing_classes, + callbacks=callbacks, + optimizers=optimizers, + peft_config=peft_config, + ) + self.group_filter_func = group_filter_func + self.num_remains_in_group = args.num_remains_in_group + if self.group_filter_func is None and self.num_remains_in_group is not None: + raise ValueError( + f"Group filter function must not be None when num_remains_in_group ({self.num_remains_in_group}) is given." + ) + if self.group_filter_func is not None and self.num_remains_in_group is None: + logger.warning("Group filter function is not activated since num_remains_in_group is not set") + + def _generate_and_score_completions(self, inputs): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] + if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ] + + prompt_ids_list, completion_ids_list, num_items_in_batch, sampling_per_token_logps_list, extra_fields = ( + self._generate(prompts) + ) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad(prompt_ids, padding_value=self.pad_token_id, padding_side="left") + prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left") + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right") + completion_mask = pad(completion_mask, padding_value=0, padding_side="right") + if sampling_per_token_logps_list is not None: + sampling_per_token_logps = [torch.tensor(logps, device=device) for logps in sampling_per_token_logps_list] + sampling_per_token_logps = pad(sampling_per_token_logps, padding_value=0.0, padding_side="right") + else: + sampling_per_token_logps = None + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + num_images = [len(img_list) for img_list in images] if images is not None else None + + # Get forward_kwargs for models with multimodal inputs + if images is not None: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class)["prompt"] for prompt in prompts + ] + prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + + with torch.no_grad(): + # If the generation and optimization steps are misaligned—i.e., if generation does not occur at the end of + # a full optimizer step (when gradient_accumulation_steps is not a multiple of generate_every)—then the + # samples may come from an earlier version of the model. In that case, we need to track old_per_token_logps + # for importance sampling. If the steps are aligned, importance sampling isn't necessary and we set + # old_per_token_logps to None. + # When using vLLM, we always compute old_per_token_logps for importance sampling, it was shown that the + # distribution mismatch between vLLM and the training model can be large and harm the training. + generate_every = self.args.steps_per_generation * self.num_iterations # generation frequency + if self.args.gradient_accumulation_steps % generate_every != 0 or ( + self.use_vllm and self.vllm_importance_sampling_correction + ): + old_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + old_per_token_logps = None + + # Compute the importance sampling ratio when using vLLM, to correct for potential distribution mismatch + if self.use_vllm and self.vllm_importance_sampling_correction: + importance_sampling_ratio = torch.exp(old_per_token_logps - sampling_per_token_logps) + importance_sampling_ratio = torch.clamp( + importance_sampling_ratio, max=self.vllm_importance_sampling_cap + ) + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + ref_per_token_logps = None + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text, strict=True): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + num_in_group = self.num_generations + num_inputs_in_device = len(prompts) + + if self.num_remains_in_group is not None and mode == "train": + num_in_group = self.num_remains_in_group + + all_completions = gather_object(completions) + + group_filter_scores = self.group_filter_func( + group_completions=[ + all_completions[i : i + 1 * self.num_generations] + for i in range(len(all_completions) // self.num_generations) + ], + group_rewards=rewards.view(-1, self.num_generations).tolist(), + ) + group_filter_scores = torch.tensor(group_filter_scores, device=device) + + _, group_local_indices = torch.topk(group_filter_scores, self.num_remains_in_group, dim=-1) + group_row_offsets = torch.arange(0, len(all_completions), self.num_generations, device=device).unsqueeze(1) + group_global_indices = group_row_offsets + group_local_indices + group_global_indices = group_global_indices.flatten() + + rewards = rewards[group_global_indices].contiguous() + rewards_per_func = rewards_per_func[group_global_indices, :].contiguous() + + num_inputs_in_device = int(len(prompts) / self.num_generations * self.num_remains_in_group) + + # Compute grouped-wise rewards + mean_grouped_rewards = rewards.view(-1, num_in_group).mean(dim=1) + + # Normalize the rewards to compute the advantages + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(num_in_group, dim=0) + advantages = rewards - mean_grouped_rewards + + if self.scale_rewards in ["group", "none"]: + # If self.scale_rewards = "none", we'll still log group level std + std_rewards = rewards.view(-1, num_in_group).std(dim=1) + std_rewards = std_rewards.repeat_interleave(num_in_group, dim=0) + elif self.scale_rewards == "batch": + # Compute global std + std_rewards = rewards.std().expand_as(rewards) + else: + raise ValueError( + f"Invalid value for scale_rewards: {self.scale_rewards}. Must be one of 'batch', 'group', or 'none'." + ) + + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) + if self.scale_rewards != "none": + advantages = advantages / (std_rewards + 1e-4) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * num_inputs_in_device, + (self.accelerator.process_index + 1) * num_inputs_in_device, + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + if self.num_remains_in_group is not None and mode == "train": + local_input_indices_to_keep = group_global_indices[process_slice] - self.accelerator.process_index * len( + prompts + ) # step is length of prompts + + prompt_ids = prompt_ids[local_input_indices_to_keep].contiguous() + prompt_mask = prompt_mask[local_input_indices_to_keep].contiguous() + completion_ids = completion_ids[local_input_indices_to_keep].contiguous() + completion_mask = completion_mask[local_input_indices_to_keep].contiguous() + attention_mask = attention_mask[local_input_indices_to_keep].contiguous() + completion_lengths = completion_mask.sum(1) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + num_items_in_batch = agg_completion_lengths.sum() + + if sampling_per_token_logps is not None: + sampling_per_token_logps = sampling_per_token_logps[local_input_indices_to_keep].contiguous() + if old_per_token_logps is not None: + old_per_token_logps = old_per_token_logps[local_input_indices_to_keep].contiguous() + if ref_per_token_logps is not None: + ref_per_token_logps = ref_per_token_logps[local_input_indices_to_keep].contiguous() + if self.use_vllm and self.vllm_importance_sampling_correction: + importance_sampling_ratio = importance_sampling_ratio[local_input_indices_to_keep].contiguous() + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + all_prompts_text = gather_object(prompts_text) + all_completions_text = gather_object(completions_text) + all_images = gather_object(images) if images is not None else None + if self.num_remains_in_group is not None and mode == "train": + group_global_indices_list = group_global_indices.tolist() + all_prompts_text = [all_prompts_text[i] for i in group_global_indices_list] + all_completions_text = [all_completions_text[i] for i in group_global_indices_list] + if images is not None: + all_images = [all_images[i] for i in group_global_indices_list] + + self._logs["prompt"].extend(all_prompts_text) + self._logs["completion"].extend(all_completions_text) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + if images is not None: + self._logs["images"].extend(all_images) + + if self.use_vllm and self.vllm_importance_sampling_correction: + delta = torch.abs(old_per_token_logps - sampling_per_token_logps) + delta = delta[completion_mask.bool()] + mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( + self.accelerator.gather(mean_delta).mean().item() + ) + self._metrics[mode]["sampling/sampling_logp_difference/max"].append( + self.accelerator.gather(max_delta).max().item() + ) + + flat_is_ratio = importance_sampling_ratio[completion_mask.bool()] + min_importance_sampling_ratio = ( + torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + mean_importance_sampling_ratio = ( + torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + max_importance_sampling_ratio = ( + torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( + nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( + self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( + nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + ) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "advantages": advantages, + "num_items_in_batch": num_items_in_batch, + } + if old_per_token_logps is not None: + output["old_per_token_logps"] = old_per_token_logps + if self.use_vllm and self.vllm_importance_sampling_correction: + output["importance_sampling_ratio"] = importance_sampling_ratio + if ref_per_token_logps is not None: + output["ref_per_token_logps"] = ref_per_token_logps + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if images is not None: + output["num_images"] = num_images + return output diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/__init__.py b/venv/lib/python3.10/site-packages/trl/experimental/gold/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97b3b3c9f4ff631b7cd9176f72c36a667ce7e718 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gold/__init__.py @@ -0,0 +1,19 @@ +# Copyright 2020-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 .gold_config import GOLDConfig +from .gold_trainer import GOLDTrainer + + +__all__ = ["GOLDConfig", "GOLDTrainer"] diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7328c343ff20c12ee6cb355f068519b7e8a77e37 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd36d88d0921d452932ced12efa0d40105022829 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fbc0e745fd130f16e41aa2e20c6f828cd56f8a0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1db94ed4773979ee298b0bb4fe170bdc3f250828 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gold/__pycache__/gold_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/gold.py b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold.py new file mode 100644 index 0000000000000000000000000000000000000000..8620d2cbc8551fef08de0914cb8cd8cc0b422779 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold.py @@ -0,0 +1,142 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl @ git+https://github.com/huggingface/trl.git", +# "peft", +# "trackio", +# ] +# /// + +# docstyle-ignore +""" +# Full training: +python trl/experimental/gold/gold.py \ + --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \ + --teacher_model_name_or_path Qwen/Qwen2-1.5B-Instruct \ + --dataset_name trl-lib/chatbot_arena_completions \ + --learning_rate 2e-5 \ + --per_device_train_batch_size 4 \ + --gradient_accumulation_steps 8 \ + --output_dir gold-model \ + --num_train_epochs 1 \ + --push_to_hub \ + --gradient_checkpointing + +# LoRA: +python trl/experimental/gold/gold.py \ + --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \ + --teacher_model_name_or_path Qwen/Qwen2-1.5B-Instruct \ + --dataset_name trl-lib/chatbot_arena_completions \ + --learning_rate 2e-4 \ + --per_device_train_batch_size 4 \ + --gradient_accumulation_steps 8 \ + --output_dir gold-model \ + --num_train_epochs 1 \ + --push_to_hub \ + --gradient_checkpointing \ + --use_peft \ + --lora_r 64 \ + --lora_alpha 16 +""" + +from datasets import load_dataset +from transformers import AutoTokenizer, GenerationConfig + +from trl import ( + LogCompletionsCallback, + ModelConfig, + ScriptArguments, + TrlParser, + get_kbit_device_map, + get_peft_config, + get_quantization_config, +) +from trl.experimental.gold.gold_config import GOLDConfig +from trl.experimental.gold.gold_trainer import GOLDTrainer + + +if __name__ == "__main__": + parser = TrlParser((ScriptArguments, GOLDConfig, ModelConfig)) + script_args, training_args, model_args = parser.parse_args_and_config() + + ################ + # Model & Tokenizer + ################ + quantization_config = get_quantization_config(model_args) + model_kwargs = dict( + revision=training_args.student_model_revision, + trust_remote_code=model_args.trust_remote_code, + attn_implementation=model_args.attn_implementation, + torch_dtype=model_args.dtype, + use_cache=False if training_args.gradient_checkpointing else True, + device_map=get_kbit_device_map() if quantization_config is not None else None, + quantization_config=quantization_config, + ) + training_args.model_init_kwargs = model_kwargs + + if training_args.teacher_tokenizer_name_or_path is None and training_args.use_uld_loss: + training_args.teacher_tokenizer_name_or_path = training_args.teacher_model_name_or_path + teacher_model_kwargs = dict( + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + attn_implementation=model_args.attn_implementation, + torch_dtype=model_args.dtype, + use_cache=True, + device_map=get_kbit_device_map() if quantization_config is not None else None, + quantization_config=quantization_config, + ) + training_args.teacher_model_init_kwargs = teacher_model_kwargs + + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + padding_side="left", + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + ################ + # Dataset + ################ + dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config) + + ################ + # Training + ################ + trainer = GOLDTrainer( + model=model_args.model_name_or_path, + teacher_model=training_args.teacher_model_name_or_path, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=get_peft_config(model_args), + ) + + if training_args.eval_strategy != "no": + generation_config = GenerationConfig( + max_new_tokens=training_args.max_new_tokens, do_sample=True, temperature=training_args.temperature + ) + completions_callback = LogCompletionsCallback(trainer, generation_config, num_prompts=8) + trainer.add_callback(completions_callback) + + trainer.train() + + # Save and push to hub + trainer.save_model(training_args.output_dir) + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_config.py b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6c2d9c34e1245c7e3936862c4e036af388deb7fd --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_config.py @@ -0,0 +1,415 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + +from ...trainer.sft_config import SFTConfig + + +@dataclass +class GOLDConfig(SFTConfig): + r""" + Configuration class for [`GOLDTrainer`]. + + This class includes only the parameters that are specific to GOLD training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. + + Args: + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + lmbda (`float`, *optional*, defaults to `0.5`): + Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy + student-generated outputs). + beta (`float`, *optional*, defaults to `0.5`): + Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When + beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. + max_completion_length (`int`, *optional*, defaults to `128`): + Maximum number of tokens to generate per completion. + teacher_model_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being + trained. + teacher_model_init_kwargs (`dict[str, Any]]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model + from a string. + teacher_tokenizer_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Tokenizer name or path for the teacher model. If None when using ULD loss, will use the same tokenizer as + the student model (not recommended for cross-tokenizer distillation). + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + seq_kd (`bool`, *optional*, defaults to `False`): + Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on + teacher-generated output). + use_uld_loss (`bool`, *optional*, defaults to `False`): + Whether to use Universal Logit Distillation (ULD) loss instead of Generalized Jensen-Shannon Divergence + loss. + uld_crossentropy_weight (`float`, *optional*, defaults to `0.0`): + Weight for the cross-entropy loss component in ULD loss. If 0, only ULD distillation loss is used. + uld_distillation_weight (`float`, *optional*, defaults to `1.0`): + Weight for the distillation loss component in ULD loss. + uld_student_temperature (`float`, *optional*, defaults to `1.0`): + Temperature for student logits in ULD loss computation. + uld_teacher_temperature (`float`, *optional*, defaults to `1.0`): + Temperature for teacher logits in ULD loss computation. + uld_skip_student_eos (`bool`, *optional*, defaults to `True`): + Whether to skip EOS token for student in ULD loss computation. + uld_skip_teacher_eos (`bool`, *optional*, defaults to `True`): + Whether to skip EOS token for teacher in ULD loss computation. + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions from the student model. Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode for student vLLM integration. Either `"server"` (connect to a running TRL vLLM server) or `"colocate"` + (run vLLM in the same process). + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server for the student model (if `vllm_mode="server"`). + vllm_server_port (`int`, *optional*, defaults to `8001`): + Port of the vLLM server for the student model (if `vllm_mode="server"`). + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Timeout for connecting to the student vLLM server (if `vllm_mode="server"`). + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.9`): + GPU memory utilization for the colocated student vLLM engine (if `vllm_mode="colocate"`). It is recommended + to set this to a low value if the student and teacher models share the same GPU. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Tensor parallel size for the colocated student vLLM engine (if `vllm_mode="colocate"`). + vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): + Regex for vLLM guided decoding for the student model. + vllm_sync_frequency (`int`, *optional*, defaults to `1`): + Frequency (in training steps) to synchronize student model weights to vLLM engine. Set to 1 to sync after + every step. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Whether to enable sleep mode for the student vLLM engine. If set to `True`, the engine will enter sleep + mode after each training step to save resources. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["teacher_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-7, + metadata={"help": "The initial learning rate for AdamW."}, + ) + + # GOLD-specific parameters + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=0.95, + metadata={ + "help": "If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to " + "`top_p` or higher are kept for generation." + }, + ) + top_k: int = field( + default=0, + metadata={"help": "The number of highest probability vocabulary tokens to keep for top-k-filtering."}, + ) + lmbda: float = field( + default=0.5, + metadata={ + "help": "Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy " + "student-generated outputs)." + }, + ) + beta: float = field( + default=0.5, + metadata={ + "help": "Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence " + "loss. When beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL " + "Divergence." + }, + ) + max_completion_length: int = field( + default=128, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + student_model_revision: str = field( + default="main", + metadata={ + "help": "Revision of the student model to use. If not specified, the default revision of the model will be used." + }, + ) + teacher_model_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Model name or path of the teacher model. If `None`, the teacher model will be the same as the " + "model being trained." + }, + ) + teacher_model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "teacher model from a string." + }, + ) + teacher_tokenizer_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Tokenizer name or path for the teacher model. If None when using ULD loss, will use the same " + "tokenizer as the student model (not recommended for cross-tokenizer distillation)." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropouts in `model`."}, + ) + seq_kd: bool = field( + default=False, + metadata={ + "help": "Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised " + "FT on teacher-generated output)." + }, + ) + steps_per_generation: int | None = field( + default=None, + metadata={ + "help": "Number of optimization steps per generation. If `None`, it defaults to gradient_accumulation_steps." + }, + ) + + # ULD Loss parameters + use_uld_loss: bool = field( + default=False, + metadata={ + "help": "Whether to use Universal Logit Distillation (ULD) loss instead of Generalized Jensen-Shannon Divergence loss." + }, + ) + use_extended_uld: bool = field( + default=True, + metadata={ + "help": ( + "Whether to enable extended ULD alignment that uses tokenizers to align and merge token " + "probabilities across student and teacher tokenizations. When True, the trainer will compute " + "token mappings and merge probabilities for split tokens; when False, ULD will use simple " + "positional truncation like in the original ULD paper." + ) + }, + ) + uld_use_hybrid_loss: bool = field( + default=False, + metadata={ + "help": ( + "Whether to use a hybrid loss that combines ULD loss and JSD loss. When True, the final loss is a " + "a combination of JSD for known token mappings and ULD for unknown token mappings." + ) + }, + ) + uld_hybrid_matched_weight: float | None = field( + default=None, + metadata={ + "help": ( + "Weight for the matched token loss component when using hybrid ULD + JSD loss. This weight scales " + "the JSD loss computed over tokens that have a direct mapping between student and teacher " + "tokenizations. If None, uses adaptive weighting based on vocabulary overlap. Must be set together " + "with uld_hybrid_unmatched_weight (both None or both float)." + ) + }, + ) + uld_hybrid_unmatched_weight: float | None = field( + default=None, + metadata={ + "help": ( + "Weight for the unmatched token loss component when using hybrid ULD + JSD loss. This weight scales " + "the ULD loss computed over tokens that do not have a direct mapping between student and teacher " + "tokenizations. If None, uses adaptive weighting based on vocabulary overlap. Must be set together " + "with uld_hybrid_matched_weight (both None or both float)." + ) + }, + ) + uld_crossentropy_weight: float = field( + default=0.0, + metadata={"help": "Weight for the cross-entropy loss component in ULD loss."}, + ) + uld_distillation_weight: float = field( + default=1.0, + metadata={"help": "Weight for the distillation loss component in ULD loss."}, + ) + uld_student_temperature: float = field( + default=1.0, + metadata={"help": "Temperature for student logits in ULD loss computation."}, + ) + uld_teacher_temperature: float = field( + default=1.0, + metadata={"help": "Temperature for teacher logits in ULD loss computation."}, + ) + + uld_skip_student_eos: bool = field( + default=True, + metadata={"help": "Whether to skip EOS token for student in ULD loss computation."}, + ) + uld_skip_teacher_eos: bool = field( + default=True, + metadata={"help": "Whether to skip EOS token for teacher in ULD loss computation."}, + ) + + # transformers paged attention + use_transformers_paged: bool = field( + default=False, + metadata={ + "help": "Whether to use the `transformers` paged implementation for generation. If set to `True`, the " + "`transformers` paged implementation will be used for generation instead of the default padded " + "implementation." + }, + ) + + # vLLM parameters + use_vllm: bool = field( + default=False, + metadata={"help": "Whether to use vLLM for generating completions. Requires `vllm` to be installed."}, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": 'Mode for vLLM integration. Either "server" (connect to a running TRL vLLM server) or "colocate" (run vLLM in the same process).' + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": 'Host of the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_server_port: int = field( + default=8001, + metadata={"help": 'Port of the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={"help": 'Timeout (in seconds) for connecting to the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_gpu_memory_utilization: float = field( + default=0.9, + metadata={ + "help": 'GPU memory utilization for the colocated vLLM engine when `vllm_mode="colocate"`. Lower values reduce contention when sharing a device with the student/teacher models.' + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={"help": 'Tensor parallel size for the colocated vLLM engine when `vllm_mode="colocate"`.'}, + ) + vllm_guided_decoding_regex: str | None = field( + default=None, + metadata={"help": "Regex pattern used for vLLM guided decoding (optional)."}, + ) + vllm_sync_frequency: int = field( + default=1, + metadata={ + "help": "Frequency (in training steps) to synchronize model weights to the vLLM engine. Set to 1 to sync after every step." + }, + ) + vllm_enable_sleep_mode: bool = field( + default=False, + metadata={ + "help": "Whether to enable sleep mode for the colocated vLLM engine. When `True`, the engine sleeps during the optimizer step and wakes for weight sync and generation." + }, + ) + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + log_completions_steps: int = field( + default=100, + metadata={ + "help": "Number of steps between logging (prompt, completion) pairs. Only used if `log_completions` is " + "set to `True`." + }, + ) + num_completions_to_print: int | None = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_entity: str | None = field( + default=None, + metadata={"help": ("The entity to store runs under.")}, + ) + wandb_project: str | None = field( + default=None, + metadata={"help": ("The project to store runs under.")}, + ) + wandb_run_group: str | None = field( + default=None, + metadata={"help": ("The group to store runs under.")}, + ) + wandb_log_unique_prompts: bool = field( + default=True, + metadata={ + "help": ("Whether to log the unique prompts to wandb. This will create a new run for each unique prompt.") + }, + ) + callbacks: list[str] = field( + default_factory=lambda: [], + metadata={"help": "The callbacks to run during training."}, + ) + hub_model_revision: str | None = field( + default="main", metadata={"help": "The Hub model branch to push the model to."} + ) + num_completions_to_print: int = field(default=5, metadata={"help": "Number of completions to print."}) + overwrite_hub_revision: bool = field(default=False, metadata={"help": "Whether to overwrite the Hub revision."}) + push_to_hub_revision: bool = field(default=False, metadata={"help": "Whether to push to a Hub revision/branch."}) + trl_project: str = field( + default="smollm3", + metadata={ + "help": "The TRL project to use for evaluation. This is used to determine the path to the evaluation script." + }, + ) + + def __post_init__(self): + super().__post_init__() + # check lmbda and beta are in the range [0, 1] + if self.lmbda < 0.0 or self.lmbda > 1.0: + raise ValueError("lmbda must be in the range [0.0, 1.0].") + if self.beta < 0.0 or self.beta > 1.0: + raise ValueError("beta must be in the range [0.0, 1.0].") + + # Validate that max_length is sufficient for max_completion_length + if self.max_length is not None and self.max_completion_length >= self.max_length: + raise ValueError( + f"max_completion_length ({self.max_completion_length}) must be smaller than max_length ({self.max_length}) " + f"to leave room for the prompt. Consider increasing max_length or reducing max_completion_length." + ) + + if self.steps_per_generation is None: + self.steps_per_generation = self.gradient_accumulation_steps + + # Validate ULD parameters + if self.use_uld_loss: + if self.uld_crossentropy_weight < 0.0: + raise ValueError("uld_crossentropy_weight must be non-negative.") + if self.uld_distillation_weight < 0.0: + raise ValueError("uld_distillation_weight must be non-negative.") + if self.uld_student_temperature <= 0.0: + raise ValueError("uld_student_temperature must be positive.") + if self.uld_teacher_temperature <= 0.0: + raise ValueError("uld_teacher_temperature must be positive.") + + # Validate hybrid loss weights - both must be None or both must be set + if self.uld_use_hybrid_loss: + if (self.uld_hybrid_matched_weight is None) != (self.uld_hybrid_unmatched_weight is None): + raise ValueError( + "uld_hybrid_matched_weight and uld_hybrid_unmatched_weight must both be None (for adaptive " + "weighting) or both be set to numeric values. Got uld_hybrid_matched_weight=" + f"{self.uld_hybrid_matched_weight} and uld_hybrid_unmatched_weight=" + f"{self.uld_hybrid_unmatched_weight}." + ) + if self.uld_hybrid_matched_weight is not None: + if self.uld_hybrid_matched_weight < 0.0: + raise ValueError("uld_hybrid_matched_weight must be non-negative.") + if self.uld_hybrid_unmatched_weight < 0.0: + raise ValueError("uld_hybrid_unmatched_weight must be non-negative.") diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_trainer.py b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..b2697d94623fc805c247c2857515f0fcd50a222b --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gold/gold_trainer.py @@ -0,0 +1,2043 @@ +# Copyright 2020-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 os +import random +import textwrap +import warnings +from collections import defaultdict, deque +from collections.abc import Callable +from contextlib import nullcontext +from typing import Any, Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from accelerate.utils import DistributedType, broadcast_object_list, gather_object, is_peft_model +from datasets import Dataset, IterableDataset +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from transformers import AutoTokenizer +from transformers.data.data_collator import DataCollator +from transformers.feature_extraction_utils import FeatureExtractionMixin +from transformers.generation.configuration_utils import GenerationConfig +from transformers.image_processing_utils import BaseImageProcessor +from transformers.integrations.integration_utils import is_wandb_available +from transformers.modeling_utils import PreTrainedModel +from transformers.processing_utils import ProcessorMixin +from transformers.tokenization_utils_base import PreTrainedTokenizerBase +from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState +from transformers.trainer_utils import EvalPrediction +from transformers.utils import ( + is_flash_attn_2_available, + is_liger_kernel_available, + is_peft_available, + is_rich_available, +) + +from ...data_utils import is_conversational, maybe_convert_to_chatml, pack_dataset, truncate_dataset +from ...extras.profiling import profiling_decorator +from ...extras.vllm_client import VLLMClient +from ...import_utils import is_vllm_available +from ...models import prepare_deepspeed +from ...models.utils import unwrap_model_for_generation +from ...trainer.sft_trainer import SFTTrainer +from ...trainer.utils import ( + DataCollatorForChatML, + create_model_from_path, + disable_dropout_in_model, + empty_cache, + ensure_master_addr_port, + pad, +) +from .gold_config import GOLDConfig + + +if is_peft_available(): + from peft import PeftConfig + +if is_wandb_available(): + import wandb + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearJSDLoss + +if is_rich_available(): + from rich.console import Console + from rich.panel import Panel + from rich.table import Table + from rich.text import Text + + +def print_prompt_completions_sample_uld( + prompts: list[str], + completions: list[str], + step: int, + num_samples: int = None, +) -> None: + """ + Print out a sample of model completions to the console with multiple reward metrics. + + This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring model outputs + during training. It requires the `rich` library to be installed. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[str]`): + List of completions corresponding to the prompts. + rewards (`dict[str, list[float]]`): + Dictionary where keys are reward names and values are lists of rewards. + advantages (`list[float]`): + List of advantages corresponding to the prompts and completions. + step (`int`): + Current training step number, used in the output title. + num_samples (`int` or `None`, *optional*, defaults to `None`): + Number of random samples to display. If `None` (default), all items will be displayed. + + Example: + ```python + >>> from trl.trainer.utils import print_prompt_completions_sample + + >>> prompts = ["The sky is", "The sun is"] + >>> completions = [" blue.", " in the sky."] + >>> rewards = {"Correctness": [0.123, 0.456], "Format": [0.789, 0.101]} + >>> advantages = [0.987, 0.654] + >>> print_prompt_completions_sample(prompts, completions, rewards, advantages, 42) + ╭──────────────────────────── Step 42 ─────────────────────────────╮ + │ ┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━┓ │ + │ ┃ Prompt ┃ Completion ┃ Correctness ┃ Format ┃ Advantage ┃ │ + │ ┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━┩ │ + │ │ The sky is │ blue. │ 0.12 │ 0.79 │ 0.99 │ │ + │ ├────────────┼──────────────┼─────────────┼────────┼───────────┤ │ + │ │ The sun is │ in the sky. │ 0.46 │ 0.10 │ 0.65 │ │ + │ └────────────┴──────────────┴─────────────┴────────┴───────────┘ │ + ╰──────────────────────────────────────────────────────────────────╯ + ``` + """ + if not is_rich_available(): + raise ImportError( + "The function `print_prompt_completions_sample` requires the `rich` library. Please install it with " + "`pip install rich`." + ) + console = Console() + table = Table(show_header=True, header_style="bold white", expand=True) + + # Add columns + table.add_column("Prompt", style="bright_yellow") + table.add_column("Completion", style="bright_green") + + # Some basic input validation + if num_samples is not None: + if num_samples >= len(prompts): + num_samples = None + elif num_samples <= 0: + return + + # Subsample data if num_samples is specified + if num_samples is not None: + indices = random.sample(range(len(prompts)), num_samples) + prompts = [prompts[i] for i in indices] + completions = [completions[i] for i in indices] + + for i in range(len(prompts)): + table.add_row(Text(prompts[i]), Text(completions[i])) + table.add_section() # Adds a separator between rows + + panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white") + console.print(panel) + + +def build_teacher_inputs_from_texts( + tokenizer: PreTrainedTokenizerBase, + prompt_texts: list[str], + completion_texts: list[str], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Tokenize teacher prompts/completions and produce tensors ready for GOLD loss.""" + + pad_token_id = tokenizer.pad_token_id + eos_token_id = tokenizer.eos_token_id + + prompt_token_ids = tokenizer(prompt_texts, add_special_tokens=True)["input_ids"] + completion_token_ids = tokenizer(completion_texts, add_special_tokens=False)["input_ids"] + + sequences: list[torch.Tensor] = [] + attention_masks: list[torch.Tensor] = [] + labels_list: list[torch.Tensor] = [] + prompt_lengths: list[int] = [] + + for prompt_ids, completion_ids in zip(prompt_token_ids, completion_token_ids, strict=True): + # Remove trailing EOS from prompt so completions can extend cleanly + if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: + prompt_ids = prompt_ids[:-1] + + prompt_lengths.append(len(prompt_ids)) + sequence = list(prompt_ids) + sequence.extend(completion_ids) + if eos_token_id is not None: + sequence.append(eos_token_id) + + seq_tensor = torch.tensor(sequence, dtype=torch.long) + sequences.append(seq_tensor) + attention_masks.append(torch.ones_like(seq_tensor)) + + labels = seq_tensor.clone() + labels[: len(prompt_ids)] = -100 + if pad_token_id is not None: + labels[labels == pad_token_id] = -100 + labels_list.append(labels) + + teacher_input_ids = pad( + sequences, + padding_side="right", + padding_value=pad_token_id if pad_token_id is not None else 0, + ) + teacher_attention_mask = pad(attention_masks, padding_side="right", padding_value=0).bool() + teacher_labels = pad(labels_list, padding_side="right", padding_value=-100) + + if eos_token_id is not None: + for row in range(teacher_attention_mask.size(0)): + valid = ( + teacher_input_ids[row] != pad_token_id + if pad_token_id is not None + else teacher_attention_mask[row].bool() + ) + if valid.any(): + last_idx = valid.nonzero(as_tuple=True)[0][-1] + teacher_attention_mask[row, last_idx + 1 :] = False + + teacher_prompt_length = max(prompt_lengths) if prompt_lengths else 0 + + return teacher_input_ids, teacher_labels, teacher_attention_mask, teacher_prompt_length + + +class ULDLoss(nn.Module): + """ + Universal Logit Distillation Loss. + """ + + def __init__(self, config: GOLDConfig, student_tokenizer=None, teacher_tokenizer=None): + super().__init__() + self.crossentropy_weight = config.uld_crossentropy_weight + self.distillation_weight = config.uld_distillation_weight + self.student_temperature = config.uld_student_temperature + self.teacher_temperature = config.uld_teacher_temperature + self.skip_student_eos = config.uld_skip_student_eos + self.skip_teacher_eos = config.uld_skip_teacher_eos + self.use_extended_uld = config.use_extended_uld + self.ignore_index = -100 + + # Add tokenizers for enhanced alignment + self.student_tokenizer = student_tokenizer + self.teacher_tokenizer = teacher_tokenizer + + # Hybrid ULD configuration + self.use_hybrid_loss = getattr(config, "uld_use_hybrid_loss", False) + self.hybrid_matched_weight = getattr(config, "uld_hybrid_matched_weight", None) + self.hybrid_unmatched_weight = getattr(config, "uld_hybrid_unmatched_weight", None) + self.beta = getattr(config, "beta", 1.0) # For JSD loss in hybrid matched tokens + + # Initialize vocabulary mapping for hybrid loss + self._vocab_mapping = None + self._teacher_matched_ids = None + self._student_matched_ids = None + if self.use_hybrid_loss and student_tokenizer is not None and teacher_tokenizer is not None: + self._initialize_vocabulary_mapping() + + def __call__( + self, student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + ): + """ + Compute ULD loss with GKD trainer interface. + + Args: + student_logits: Student model logits [batch_size, seq_len, vocab_size] + teacher_logits: Teacher model logits [batch_size, seq_len, vocab_size] + student_labels: Student target labels [batch_size, seq_len] + teacher_labels: Teacher target labels [batch_size, seq_len] + student_input_ids: Student input token IDs [batch_size, seq_len] + teacher_input_ids: Teacher input token IDs [batch_size, seq_len] + + Returns: + Total loss (cross-entropy + distillation) + """ + # Compute cross-entropy loss for student + if self.crossentropy_weight > 0: + shift_logits = student_logits[..., :-1, :].contiguous() + shift_labels = student_labels[..., 1:].contiguous() + loss_fct = nn.CrossEntropyLoss(ignore_index=self.ignore_index) + crossentropy_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + crossentropy_loss = self.crossentropy_weight * crossentropy_loss + else: + crossentropy_loss = 0.0 + + # Compute distillation loss using ULD approximation + distillation_loss = self._compute_distillation_loss( + student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + ) + + return crossentropy_loss + distillation_loss + + def _initialize_vocabulary_mapping(self): + """Initialize vocabulary mapping for hybrid ULD loss.""" + # Computing vocabulary mapping for hybrid ULD + + student_vocab = self.student_tokenizer.get_vocab() + teacher_vocab = self.teacher_tokenizer.get_vocab() + + # Create reverse mapping for student + student_token_to_id = dict(student_vocab.items()) + + vocab_mapping = {} + teacher_matched_ids = set() + student_matched_ids = set() + + for token_str, teacher_id in teacher_vocab.items(): + if token_str in student_token_to_id: + student_id = student_token_to_id[token_str] + vocab_mapping[teacher_id] = student_id + teacher_matched_ids.add(teacher_id) + student_matched_ids.add(student_id) + + self._vocab_mapping = vocab_mapping + self._teacher_matched_ids = teacher_matched_ids + self._student_matched_ids = student_matched_ids + + def _compute_distillation_loss( + self, student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + ): + """ + Compute the Universal Logit Distillation loss with token mapping. + + This version uses actual input_ids for accurate token mapping and multiplies probabilities for split tokens. + Both student_input_ids and teacher_input_ids are required for optimal alignment. + """ + # Get answer regions (same as original) + student_answer_index, student_answer_size = self._get_start_and_size_answers(student_labels) + teacher_answer_index, teacher_answer_size = self._get_start_and_size_answers(teacher_labels) + + if self.skip_student_eos: + student_answer_size = [size - 1 for size in student_answer_size] + if self.skip_teacher_eos: + teacher_answer_size = [size - 1 for size in teacher_answer_size] + + # Handle edge case where all answer sizes are 0 + if ( + not student_answer_size + or not teacher_answer_size + or max(max(student_answer_size), max(teacher_answer_size)) <= 0 + ): + return torch.zeros(1, device=student_logits.device, requires_grad=True) * student_logits.sum() * 1e-8 + + batch_size = student_logits.size(0) + distillation_losses = [] + + for i in range(batch_size): + # Get answer regions for this batch item + student_start = student_answer_index[i] + student_size = student_answer_size[i] + teacher_start = teacher_answer_index[i] + teacher_size = teacher_answer_size[i] + + if student_size <= 0 or teacher_size <= 0: + loss_i = student_logits[i].sum() * 0.0 + distillation_losses.append(loss_i) + continue + + # Extract answer logits + student_answer_logits = student_logits[i, student_start : student_start + student_size] + teacher_answer_logits = teacher_logits[i, teacher_start : teacher_start + teacher_size] + + # Convert to probabilities + student_probs = F.softmax(student_answer_logits / self.student_temperature, dim=-1) + teacher_probs = F.softmax(teacher_answer_logits / self.teacher_temperature, dim=-1) + + # Get token IDs for mapping (always use actual input_ids) + student_token_ids = student_input_ids[i, student_start : student_start + student_size].tolist() + teacher_token_ids = teacher_input_ids[i, teacher_start : teacher_start + teacher_size].tolist() + + if self.use_extended_uld: + # Build alignment groups directly from token ids using greedy text matching + student_alignment_groups, teacher_alignment_groups = self._build_alignment_groups_from_ids( + student_token_ids, teacher_token_ids + ) + + # Merge student probabilities using student alignment groups + student_aligned = self._merge_probabilities_with_alignment_groups( + student_probs, student_alignment_groups + ) + + # Merge teacher probabilities using teacher alignment groups + teacher_aligned = self._merge_probabilities_with_alignment_groups( + teacher_probs, teacher_alignment_groups + ) + else: + min_length = min(len(student_token_ids), len(teacher_token_ids)) + student_aligned = student_probs[:min_length, :] + teacher_aligned = teacher_probs[:min_length, :] + + # Apply ULD loss computation + if self.use_hybrid_loss and self._vocab_mapping is not None: + # Use hybrid approach: direct comparison for matched tokens, sorting for unmatched + aligned_loss = self._compute_hybrid_uld_loss(student_aligned, teacher_aligned) + else: + # Original approach: sort all probabilities + student_sorted = student_aligned.sort(dim=-1, descending=True).values + teacher_sorted = teacher_aligned.sort(dim=-1, descending=True).values + + # Pad vocabularies to same size + student_vocab_size = student_sorted.size(-1) + teacher_vocab_size = teacher_sorted.size(-1) + max_vocab_size = max(student_vocab_size, teacher_vocab_size) + + if student_vocab_size < max_vocab_size: + student_sorted = F.pad(student_sorted, (0, max_vocab_size - student_vocab_size)) + if teacher_vocab_size < max_vocab_size: + teacher_sorted = F.pad(teacher_sorted, (0, max_vocab_size - teacher_vocab_size)) + + # Compute L1 distance (ULD approach) + aligned_loss = F.l1_loss(student_sorted, teacher_sorted, reduction="sum") + aligned_loss /= student_aligned.size(0) # Normalize by sequence length + distillation_losses.append(aligned_loss) + + distillation_loss = torch.stack(distillation_losses).mean() + return self.distillation_weight * distillation_loss + + def _build_alignment_groups_from_ids(self, student_token_ids, teacher_token_ids): + """ + Build alignment groups using a greedy substring-equality algorithm on decoded token pieces. + + Args: + student_token_ids: List[int] + teacher_token_ids: List[int] + + Returns: + Tuple[List[List[int]], List[List[int]]]: student and teacher alignment groups + """ + + def to_canonical_pieces(tok, ids): + pieces = [] + prev = "" + for k in range(len(ids)): + # IMPORTANT: Do NOT skip special tokens - we need to align them too + cur = tok.decode(ids[: k + 1], skip_special_tokens=False, clean_up_tokenization_spaces=False) + # Extract the incremental addition (may include spaces/ZWJ/etc.) + pieces.append(cur[len(prev) :]) + prev = cur + return pieces + + s_pieces = to_canonical_pieces(self.student_tokenizer, student_token_ids) + t_pieces = to_canonical_pieces(self.teacher_tokenizer, teacher_token_ids) + + i = j = 0 + s_buf = t_buf = "" + s_group = [] + t_group = [] + s_groups = [] + t_groups = [] + + def flush(): + if s_group and t_group: + s_groups.append(s_group.copy()) + t_groups.append(t_group.copy()) + + # Greedily accumulate pieces until substrings match, then flush + while i < len(s_pieces) or j < len(t_pieces): + if s_buf == t_buf and s_buf != "": + flush() + s_buf = t_buf = "" + s_group = [] + t_group = [] + continue + + if s_buf == "" and i < len(s_pieces): + s_buf += s_pieces[i] + s_group.append(i) + i += 1 + continue + if t_buf == "" and j < len(t_pieces): + t_buf += t_pieces[j] + t_group.append(j) + j += 1 + continue + + if len(s_buf) <= len(t_buf): + if i < len(s_pieces): + s_buf += s_pieces[i] + s_group.append(i) + i += 1 + elif j < len(t_pieces): + t_buf += t_pieces[j] + t_group.append(j) + j += 1 + else: + if j < len(t_pieces): + t_buf += t_pieces[j] + t_group.append(j) + j += 1 + elif i < len(s_pieces): + s_buf += s_pieces[i] + s_group.append(i) + i += 1 + + # Flush any remainder if both sides accumulated something + if s_buf == t_buf and s_group and t_group: + flush() + elif s_group or t_group: + # Handle remaining unmatched tokens by forcing a flush + # This ensures both sides have the same number of alignment groups + if s_group or t_group: + # Ensure both groups have content (even if empty list) + if not s_group: + s_group = [] + if not t_group: + t_group = [] + # Force flush even if buffers don't match + if s_group or t_group: + s_groups.append(s_group.copy() if s_group else []) + t_groups.append(t_group.copy() if t_group else []) + + return s_groups, t_groups + + def _merge_probabilities_with_alignment_groups(self, probs, alignment_groups): + """ + Merge probabilities based on alignment groups. + + Args: + probs: Probability tensor [seq_len, vocab_size] + alignment_groups: List of alignment groups (each group is a list of positions to merge) + + Returns: + Merged probability tensor [num_groups, vocab_size] + """ + if not alignment_groups: + return probs + + # Create aligned tensor + vocab_size = probs.size(-1) + target_len = len(alignment_groups) + aligned_probs = torch.zeros(target_len, vocab_size, device=probs.device) + + # Process each alignment group + for group_idx, group in enumerate(alignment_groups): + # Handle probability merging + if len(group) > 1: + # Multiple tokens map to this group - merge them + eps = 1e-8 + logp = torch.log(probs[group[0]].clamp_min(eps)) + for idx in group[1:]: + if idx < probs.size(0): + logp = logp + torch.log(probs[idx].clamp_min(eps)) + aligned_probs[group_idx] = torch.softmax(logp, dim=-1) + elif len(group) == 1: + aligned_probs[group_idx] = probs[group[0]] + else: + # No tokens map to this group + aligned_probs[group_idx] = torch.zeros_like(probs[0]) + + return aligned_probs + + def _compute_hybrid_uld_loss(self, student_aligned, teacher_aligned): + """ + Compute hybrid ULD loss on aligned probability distributions. This method: + 1. Directly compares probabilities for tokens with matching vocabulary entries + 2. Uses sorting approach only for tokens with different vocabulary entries + + Args: + student_aligned: Aligned student probabilities [seq_len, student_vocab_size] + teacher_aligned: Aligned teacher probabilities [seq_len, teacher_vocab_size] + Returns: + Combined hybrid loss + """ + device = student_aligned.device + # seq_len = student_aligned.size(0) # Unused variable + student_vocab_size = student_aligned.size(-1) + teacher_vocab_size = teacher_aligned.size(-1) + + # Convert sets to sorted tensors for indexing + if self._teacher_matched_ids: + teacher_matched_indices = torch.tensor(sorted(self._teacher_matched_ids), dtype=torch.long, device=device) + student_matched_indices = torch.tensor( + [self._vocab_mapping[tid.item()] for tid in teacher_matched_indices], dtype=torch.long, device=device + ) + else: + teacher_matched_indices = torch.tensor([], dtype=torch.long, device=device) + student_matched_indices = torch.tensor([], dtype=torch.long, device=device) + + # Create masks for unmatched tokens + teacher_matched_mask = torch.zeros(teacher_vocab_size, dtype=torch.bool, device=device) + student_matched_mask = torch.zeros(student_vocab_size, dtype=torch.bool, device=device) + + if len(teacher_matched_indices) > 0: + teacher_matched_mask[teacher_matched_indices] = True + student_matched_mask[student_matched_indices] = True + + # 1. JSD loss for matched vocabulary tokens (direct semantic correspondence) + matched_loss = torch.tensor(0.0, device=device) + matched_token_count = 0 + if len(teacher_matched_indices) > 0: + # Extract probabilities for matched tokens + teacher_matched_probs = teacher_aligned[:, teacher_matched_indices] # [seq_len, num_matched] + student_matched_probs = student_aligned[:, student_matched_indices] # [seq_len, num_matched] + matched_token_count = teacher_matched_probs.size(-1) + + # Use JSD loss for semantically aligned tokens + # Convert probabilities back to logits for JSD computation + + # Apply generalized JSD loss to matched tokens + matched_loss = self._compute_jsd_loss_for_matched_tokens(student_matched_probs, teacher_matched_probs) + + # 2. Sorted comparison loss for unmatched vocabulary tokens + teacher_unmatched_mask = ~teacher_matched_mask + student_unmatched_mask = ~student_matched_mask + + teacher_unmatched_probs = teacher_aligned[:, teacher_unmatched_mask] # [seq_len, num_teacher_unmatched] + student_unmatched_probs = student_aligned[:, student_unmatched_mask] # [seq_len, num_student_unmatched] + + unmatched_loss = torch.tensor(0.0, device=device) + if teacher_unmatched_probs.size(-1) > 0 and student_unmatched_probs.size(-1) > 0: + # Sort unmatched probabilities + teacher_unmatched_sorted = teacher_unmatched_probs.sort(dim=-1, descending=True).values + student_unmatched_sorted = student_unmatched_probs.sort(dim=-1, descending=True).values + + # Pad to same size if needed + teacher_unmatched_size = teacher_unmatched_sorted.size(-1) + student_unmatched_size = student_unmatched_sorted.size(-1) + max_unmatched_size = max(teacher_unmatched_size, student_unmatched_size) + + if teacher_unmatched_size < max_unmatched_size: + teacher_unmatched_sorted = F.pad( + teacher_unmatched_sorted, (0, max_unmatched_size - teacher_unmatched_size) + ) + if student_unmatched_size < max_unmatched_size: + student_unmatched_sorted = F.pad( + student_unmatched_sorted, (0, max_unmatched_size - student_unmatched_size) + ) + + # L1 loss on sorted unmatched tokens + unmatched_loss = F.l1_loss(student_unmatched_sorted, teacher_unmatched_sorted, reduction="sum") + unmatched_loss /= student_aligned.size(0) # Normalize by sequence length + + # 3. Combine losses with weights + if self.hybrid_matched_weight is None: + # Use adaptive weighting based on vocabulary overlap + hybrid_matched_weight = matched_token_count / max(1, teacher_vocab_size) + hybrid_unmatched_weight = 1.0 - hybrid_matched_weight + else: + # Use fixed weights provided in config + hybrid_matched_weight = self.hybrid_matched_weight + hybrid_unmatched_weight = self.hybrid_unmatched_weight + + total_loss = hybrid_matched_weight * matched_loss + hybrid_unmatched_weight * unmatched_loss + + # Store matched/unmatched components for logging + self.last_matched_loss = matched_loss + self.last_unmatched_loss = unmatched_loss + + return total_loss + + def _compute_jsd_loss_for_matched_tokens(self, student_logits, teacher_logits): + """ + Compute JSD loss for matched vocabulary tokens. + + Args: + student_logits: Student logits for matched tokens [seq_len, num_matched] + teacher_logits: Teacher logits for matched tokens [seq_len, num_matched] + Returns: + JSD loss for matched tokens + """ + # Reshape to [batch_size * seq_len, vocab_size] format expected by generalized_jsd_loss + batch_seq_len, num_matched = student_logits.shape + + student_logits_reshaped = student_logits.view(-1, num_matched) + teacher_logits_reshaped = teacher_logits.view(-1, num_matched) + + # Use the GOLD generalized JSD loss implementation that accepts probability inputs + jsd_loss = GOLDTrainer.generalized_jsd_loss( + student_logits_reshaped, + teacher_logits_reshaped, + labels=None, # No masking needed for matched tokens + beta=self.beta, # Standard JSD beta + temperature=1.0, # Already applied in main computation + reduction="batchmean", + logits_are_probs=True, + ) + + return jsd_loss + + def _get_start_and_size_answers(self, answer_tensors): + answers_index = [] + answers_size = [] + + for answer in answer_tensors: + answer_mask = answer.ne(self.ignore_index) + if not answer_mask.any(): + answers_index.append(0) + answers_size.append(0) + continue + + valid_indices = answer_mask.nonzero(as_tuple=True)[0] + answers_index.append(int(valid_indices[0].item())) + answers_size.append(int(answer_mask.sum().item())) + return answers_index, answers_size + + +class GOLDVLLMSyncCallback(TrainerCallback): + """Sync the model weights to vLLM after training steps when it's safe to do so.""" + + def __init__(self, trainer): + self.trainer = trainer + + def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs): + """Sync weights after training step when DeepSpeed is stable.""" + if ( + self.trainer.use_vllm + and state.global_step != self.trainer._last_vllm_sync_step + and state.global_step % self.trainer.vllm_sync_frequency == 0 + ): + # Check if this is a step where gradients are synchronized + # This happens at the end of gradient accumulation cycles + if hasattr(self.trainer.accelerator, "sync_gradients") and self.trainer.accelerator.sync_gradients: + self.trainer._move_model_to_vllm() + self.trainer._last_vllm_sync_step = state.global_step + + +class GOLDTrainer(SFTTrainer): + _tag_names = ["trl", "gold"] + _name = "GOLD" + _paper = { + "title": "Unlocking On-Policy Distillation for Any Model Family", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @misc{patino2025unlocking, + title = {{Unlocking On-Policy Distillation for Any Model Family}}, + author = {Carlos Miguel Patiño and Kashif Rasul and Quentin Gallouédec and Ben Burtenshaw and Sergio Paniego and Vaibhav Srivastav and Thibaud Frere and Ed Beeching and Lewis Tunstall and Leandro von Werra and Thomas Wolf}, + year = 2025, + url = {https://huggingface.co/spaces/HuggingFaceH4/general-on-policy-logit-distillation}, + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str | None = None, + teacher_model: PreTrainedModel | nn.Module | str = None, + args: GOLDConfig | None = None, + data_collator: DataCollator | None = None, # type: ignore + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: Optional["PeftConfig"] = None, + ): + self.model_name_or_path = model if isinstance(model, str) else model.config._name_or_path + self.model_revision = getattr(args, "student_model_revision", None) + if isinstance(model, str) and self.model_revision is not None: + args.model_init_kwargs = args.model_init_kwargs or {} + args.model_init_kwargs.setdefault("revision", self.model_revision) + + # Respect a user-provided data_collator; otherwise, provide a ChatML collator that + if data_collator is None: + data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + + # Liger fused GKD loss (JSD) + self.use_liger_gkd_loss = False + if args.use_liger_kernel: + self.liger_jsd_loss = LigerFusedLinearJSDLoss( + beta=args.beta, + ignore_index=-100, + temperature=args.temperature, + compiled=False, + ) + self.use_liger_gkd_loss = True + + if args.teacher_model_init_kwargs is None: + teacher_model_init_kwargs = {} + elif not isinstance(teacher_model, str): + raise ValueError( + "You passed teacher_model_init_kwargs to the GOLDConfig, but your teacher_model is already instantiated." + ) + else: + teacher_model_init_kwargs = args.teacher_model_init_kwargs + teacher_model_init_kwargs["torch_dtype"] = ( + teacher_model_init_kwargs["torch_dtype"] + if teacher_model_init_kwargs["torch_dtype"] in ["auto", None] + else getattr(torch, teacher_model_init_kwargs["torch_dtype"]) + ) + + if args.use_uld_loss and args.teacher_tokenizer_name_or_path is None: + if isinstance(teacher_model, str): + args.teacher_tokenizer_name_or_path = teacher_model + else: + raise ValueError( + "`teacher_tokenizer_name_or_path` must be set when using ULD loss with a pre-instantiated teacher model." + ) + + if isinstance(teacher_model, str): + init_kwargs = dict(teacher_model_init_kwargs) + if "torch_dtype" in init_kwargs and "dtype" not in init_kwargs: + init_kwargs["dtype"] = init_kwargs.pop("torch_dtype") + teacher_model = create_model_from_path(teacher_model, **init_kwargs) + self.use_uld_loss = args.use_uld_loss + self.teacher_tokenizer = None + if args.use_uld_loss and args.teacher_tokenizer_name_or_path is not None: + self.teacher_tokenizer = AutoTokenizer.from_pretrained(args.teacher_tokenizer_name_or_path) + if not hasattr(self.teacher_tokenizer, "pad_token") or self.teacher_tokenizer.pad_token is None: + self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token + + # Hybrid ULD loss configuration is handled in ULDLoss class + + super().__init__( + model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + peft_config=peft_config, + ) + + if args.disable_dropout: + disable_dropout_in_model(self.model) + if not args.use_uld_loss: + teacher_model.resize_token_embeddings(self.model.config.vocab_size) + + if self.is_deepspeed_enabled: + self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) + else: + self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) + + self.lmbda = args.lmbda + self.beta = args.beta + self.temperature = args.temperature + self.top_p = args.top_p + self.seq_kd = args.seq_kd + + # Track per-step loss statistics for on/off-policy batches (used in logging) + self._on_policy_loss_total = 0.0 + self._off_policy_loss_total = 0.0 + self._on_policy_step_equiv = 0.0 + self._off_policy_step_equiv = 0.0 + + # Hybrid ULD matched/unmatched accumulators (logged every step when ULD hybrid is used) + self._matched_sum = 0.0 + self._unmatched_sum = 0.0 + self._matched_step_eq = 0.0 + self._unmatched_step_eq = 0.0 + + self.use_transformers_paged = args.use_transformers_paged or False + + self.uld_loss_fn = None + if self.use_uld_loss: + self.uld_loss_fn = ULDLoss( + config=args, + student_tokenizer=processing_class, + teacher_tokenizer=self.teacher_tokenizer, + ) + + self.generation_config = GenerationConfig( + max_new_tokens=args.max_completion_length, + temperature=args.temperature, + top_p=args.top_p, + do_sample=True, + top_k=args.top_k, + pad_token_id=self.processing_class.pad_token_id, + ) + if ( + hasattr(self.model.generation_config, "eos_token_id") + and self.model.generation_config.eos_token_id is not None + ): + self.generation_config.eos_token_id = self.model.generation_config.eos_token_id + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.log_completion_steps = args.log_completions_steps + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # maxlen is set to the total number of forward passes per step. This value of `maxlen` ensures we log only the + # final optimization step. + maxlen = self.accelerator.num_processes * args.per_device_train_batch_size * args.steps_per_generation + self._textual_logs = { + "prompt": deque(maxlen=maxlen), + "completion": deque(maxlen=maxlen), + "rewards": defaultdict(lambda: deque(maxlen=maxlen)), + "advantages": deque(maxlen=maxlen), + } + + self.use_vllm = args.use_vllm + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and use_vllm is set to True. Please install vLLM with " + "`pip install vllm` to use it." + ) + self.vllm_mode = args.vllm_mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization + self.vllm_enable_sleep_mode = args.vllm_enable_sleep_mode + if self.vllm_mode == "server": + if self.accelerator.is_main_process: + self.vllm_client = VLLMClient( + host=args.vllm_server_host, + server_port=args.vllm_server_port, + connection_timeout=args.vllm_server_timeout, + ) + self.vllm_client.init_communicator() + elif self.vllm_mode == "colocate": + student_model_name_or_path = self.model_name_or_path + + # Make sure tensor_parallel_size divides world size evenly + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + # Create subgroups of ranks for TP + self.vllm_tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list( + range( + i * self.vllm_tensor_parallel_size, + (i + 1) * self.vllm_tensor_parallel_size, + ) + ) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + # vLLM requires the environment variables to be set for distributed training. + os.environ["RANK"] = str(self.accelerator.process_index) + os.environ["LOCAL_RANK"] = str(self.accelerator.local_process_index) + os.environ["WORLD_SIZE"] = str(self.accelerator.num_processes) + ensure_master_addr_port() + + self.vllm_engine = LLM( + model=student_model_name_or_path, + revision=self.model_revision, + tensor_parallel_size=self.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + max_num_seqs=self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps, + max_model_len=args.max_length, + distributed_executor_backend="external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + seed=self.accelerator.process_index // self.vllm_tensor_parallel_size, + enable_sleep_mode=self.vllm_enable_sleep_mode, + ) + + if self.vllm_enable_sleep_mode: + self.vllm_engine.sleep(level=2) + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + raise ValueError(f"Unknown vllm_mode: {self.vllm_mode}") + self.vllm_guided_decoding_regex = args.vllm_guided_decoding_regex + self.vllm_sync_frequency = args.vllm_sync_frequency + self._last_vllm_sync_step = -1 + + self.add_callback(GOLDVLLMSyncCallback(self)) + + def _set_signature_columns_if_needed(self): + super()._set_signature_columns_if_needed() + required_columns = [ + "prompts", + "prompt_attention_mask", + "messages", + "chat_template_kwargs", + "tools", + "original_prompt_text", + "original_completion_text", + ] + if self._signature_columns is None: + self._signature_columns = required_columns + else: + for column in required_columns: + if column not in self._signature_columns: + self._signature_columns.append(column) + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + args, + packing: bool, + formatting_func: Callable[[dict], str] | None, + dataset_name: str, + ) -> Dataset | IterableDataset: + """ + Override dataset preparation to preserve original text for cross-tokenizer distillation and ensure + attention_mask is always added for DataCollatorForChatML compatibility. + """ + # Check if dataset is already processed + column_names = list(next(iter(dataset)).keys()) + is_processed = "input_ids" in column_names + + # Use our enhanced dataset preparation for: + # 1. ULD loss with cross-tokenizer (need original text preservation) + # 2. Any unprocessed dataset (need attention_mask for DataCollatorForChatML) + if not is_processed or (self.use_uld_loss and self.teacher_tokenizer is not None): + # For unprocessed datasets, use our enhanced tokenization + return self._prepare_dataset_with_original_text( + dataset, processing_class, args, packing, formatting_func, dataset_name + ) + + # Use parent implementation for all other cases + return super()._prepare_dataset(dataset, processing_class, args, packing, formatting_func, dataset_name) + + def _prepare_dataset_with_original_text( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + args, + packing: bool, + formatting_func: Callable[[dict], str] | None, + dataset_name: str, + ) -> Dataset | IterableDataset: + """ + Prepare dataset while preserving original text for cross-tokenizer distillation. + """ + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + # Apply the formatting function if any + if formatting_func is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Applying formatting function to {dataset_name} dataset" + + def _func(example): + return {"text": formatting_func(example)} + + try: + dataset = dataset.map(_func, batched=False, **map_kwargs) + except Exception as e: + warnings.warn( + f"Failed to apply the formatting function due to the following error: {e}. This may be " + "because the function is designed for batched input. Please update it to process one example " + "at a time (i.e., accept and return a single example). For now, we will attempt to apply the " + "function in batched mode, but note that batched formatting is deprecated and will be removed " + "in version 0.21.", + DeprecationWarning, + ) + dataset = dataset.map(_func, batched=True, **map_kwargs) + + # Convert the dataset to ChatML if needed + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Converting {dataset_name} dataset to ChatML" + column_names = next(iter(dataset)).keys() + dataset = dataset.map( + maybe_convert_to_chatml, + remove_columns="conversations" if "conversations" in column_names else None, + **map_kwargs, + ) + + # Apply the chat template if needed and preserve original text + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if "text" in example and not example["text"].endswith(eos_token): # language modeling case + example["text"] = example["text"] + eos_token + elif "completion" in example and not example["completion"].endswith(eos_token): + example["completion"] = example["completion"] + eos_token + return example + + dataset = dataset.map( + add_eos, + fn_kwargs={"eos_token": processing_class.eos_token}, + remove_columns="messages" if "messages" in column_names else None, # renamed to "text" + **map_kwargs, + ) + + # Tokenize the dataset while preserving original text + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset (preserving original text)" + + def tokenize_with_original_text(example, processing_class, dataset_text_field, assistant_only_loss): + """Modified tokenization function that preserves original text.""" + result = {} + + if "prompt" in example: # prompt-completion case + # Store original text + result["original_prompt_text"] = example["prompt"] + result["original_completion_text"] = example["completion"] + + if is_conversational(example): + prompt_ids = processing_class.apply_chat_template( + example["prompt"], **example.get("chat_template_kwargs", {}) + ) + prompt_completion_ids = processing_class.apply_chat_template( + example["prompt"] + example["completion"], **example.get("chat_template_kwargs", {}) + ) + else: + prompt_ids = processing_class(text=example["prompt"]).input_ids + prompt_completion_ids = processing_class( + text=example["prompt"] + example["completion"] + ).input_ids + + # Check if the tokenized prompt starts with the tokenized prompt+completion + if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids: + warnings.warn( + "Mismatch between tokenized prompt and the start of tokenized prompt+completion. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + + # Create a completion mask + completion_mask = [0] * len(prompt_ids) + [1] * (len(prompt_completion_ids) - len(prompt_ids)) + result.update( + { + "input_ids": prompt_completion_ids, + "completion_mask": completion_mask, + "attention_mask": [1] * len(prompt_completion_ids), # Add attention mask + } + ) + + else: # language modeling or conversational case + if is_conversational(example): + # For conversational data (ChatML), extract prompt and completion properly + messages = example["messages"] + + # Extract user and assistant messages separately + user_messages = [msg for msg in messages if msg["role"] != "assistant"] + assistant_messages = [msg for msg in messages if msg["role"] == "assistant"] + + if user_messages and assistant_messages: + # Apply chat template to get the prompt (everything up to assistant) + prompt_text = processing_class.apply_chat_template( + user_messages, + tokenize=False, + add_generation_prompt=True, # Add assistant prompt + **example.get("chat_template_kwargs", {}), + ) + + # Get the full conversation with assistant response + full_text = processing_class.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + **example.get("chat_template_kwargs", {}), + ) + + # Extract completion as everything after the prompt + # This ensures we capture any extra tokens (like tags) that the template adds + if full_text.startswith(prompt_text): + completion_text = full_text[len(prompt_text) :] + else: + # Fallback: use assistant content + EOS + assistant_content = assistant_messages[0]["content"] + completion_text = ( + assistant_content + processing_class.eos_token + if hasattr(processing_class, "eos_token") + else assistant_content + ) + + # Store original text for cross-tokenizer distillation + result["original_prompt_text"] = prompt_text + result["original_completion_text"] = completion_text + else: + # Fallback: use empty prompt and full text as completion + full_text = processing_class.apply_chat_template( + messages, tokenize=False, **example.get("chat_template_kwargs", {}) + ) + result["original_prompt_text"] = "" + result["original_completion_text"] = full_text + + # Process the conversation normally + processed = processing_class.apply_chat_template( + example["messages"], + return_dict=True, + return_assistant_tokens_mask=assistant_only_loss, + **example.get("chat_template_kwargs", {}), + ) + if "assistant_masks" in processed and 1 not in processed["assistant_masks"]: + raise RuntimeError( + "You're using `assistant_only_loss=True`, but at least one example has no " + "assistant tokens. This usually means the tokenizer's chat template doesn't " + "generate assistant masks — it may be missing the `{% generation %}` tag. Please " + "check the template and ensure it's correctly configured to support assistant " + "masking." + ) + result.update({k: processed[k] for k in ("input_ids", "assistant_masks") if k in processed}) + # Add attention_mask if not already present + if "attention_mask" not in result: + result["attention_mask"] = [1] * len(result["input_ids"]) + else: + # For regular language modeling, store the full text as completion and empty prompt + result["original_prompt_text"] = "" + result["original_completion_text"] = example.get(dataset_text_field, example.get("text", "")) + + tokenized = processing_class(text=example[dataset_text_field]) + result.update( + { + "input_ids": tokenized.input_ids, + "attention_mask": getattr(tokenized, "attention_mask", [1] * len(tokenized.input_ids)), + } + ) + + return result + + dataset = dataset.map( + tokenize_with_original_text, + fn_kwargs={ + "processing_class": processing_class, + "dataset_text_field": args.dataset_text_field, + "assistant_only_loss": args.assistant_only_loss, + }, + **map_kwargs, + ) + + # Pack or truncate + if packing: + if args.max_length is None: + raise ValueError("When packing is enabled, `max_length` can't be `None`.") + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Packing {dataset_name} dataset" + + columns_to_keep = ["input_ids", "original_prompt_text", "original_completion_text"] + existing_columns = set(dataset.column_names) + columns_to_select = [col for col in columns_to_keep if col in existing_columns] + + dataset = dataset.select_columns(columns_to_select) + dataset = pack_dataset(dataset, args.max_length, args.packing_strategy, map_kwargs) + elif args.max_length is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Truncating {dataset_name} dataset" + dataset = truncate_dataset(dataset, args.max_length, map_kwargs) + + if args.use_liger_kernel: + required_columns = { + "input_ids", + "attention_mask", + "position_ids", + "completion_mask", + "assistant_masks", + "original_prompt_text", + "original_completion_text", + } + dataset = dataset.select_columns(required_columns.intersection(dataset.column_names)) + + return dataset + + @staticmethod + def generalized_jsd_loss( + student_logits, + teacher_logits, + labels=None, + beta=0.5, + temperature=1.0, + reduction="batchmean", + logits_are_probs=False, + ): + """ + Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) + of https://huggingface.co/papers/2306.13649 for the definition. + + Args: + student_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + teacher_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + labels: + Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing + loss + beta: + Interpolation coefficient between 0 and 1 (default: 0.5) + temperature: + Softmax temperature (default: 1.0) + reduction: + Specifies the reduction to apply to the output (default: 'batchmean') + + Returns: + loss: Scalar tensor with the generalized JSD loss + """ + + if logits_are_probs: + student_log_probs = torch.log(student_logits.clamp_min(1e-8)) + teacher_log_probs = torch.log(teacher_logits.clamp_min(1e-8)) + else: + # Apply temperature scaling to logits before computing probabilities + student_logits = student_logits / temperature + teacher_logits = teacher_logits / temperature + # Compute log probabilities for student and probabilities for teacher + student_log_probs = F.log_softmax(student_logits, dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) + + if beta == 0: + jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) + elif beta == 1: + jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) + else: + # Compute the log of the mixture distribution + # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture + beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device) + mixture_log_probs = torch.logsumexp( + torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]), + dim=0, + ) + + # Compute KL divergences using F.kl_div + # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. + kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) + kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) + + # Compute the Generalized Jensen-Shannon Divergence + jsd = beta * kl_teacher + (1 - beta) * kl_student + + # Masking + if labels is not None: + mask = labels != -100 + jsd = jsd[mask] + + # Apply reduction + if reduction == "batchmean": + return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0) + elif reduction == "sum": + return jsd.sum() + elif reduction == "mean": + return jsd.mean() + else: + return jsd + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if self.use_uld_loss and self.teacher_tokenizer is not None: + if "original_prompt_text" in inputs and "original_completion_text" in inputs: + prompt_texts = inputs["original_prompt_text"] + completion_texts = inputs["original_completion_text"] + full_texts = [p + c for p, c in zip(prompt_texts, completion_texts, strict=True)] + else: + # Fallback: decode student input_ids (current approach) + # WARNING: This may not work perfectly for cross-tokenizer distillation + full_sequences = inputs["input_ids"] + full_texts = self.processing_class.batch_decode(full_sequences, skip_special_tokens=False) + + # Try to split prompt/completion using original prompt length + prompt_lengths = inputs["prompts"].shape[1] + prompt_texts = self.processing_class.batch_decode(inputs["prompts"], skip_special_tokens=False) + completion_texts = [ + full.replace(prompt, "", 1) for full, prompt in zip(full_texts, prompt_texts, strict=True) + ] + + ( + teacher_input_ids, + teacher_labels, + teacher_attention_mask, + teacher_prompt_length, + ) = build_teacher_inputs_from_texts( + self.teacher_tokenizer, + prompt_texts, + completion_texts, + ) + + teacher_input_ids = teacher_input_ids.to(self.accelerator.device) + teacher_labels = teacher_labels.to(self.accelerator.device) + teacher_attention_mask = teacher_attention_mask.to(self.accelerator.device) + + outputs_student = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + + self.teacher_model.eval() + with torch.no_grad(): + outputs_teacher = self.teacher_model( + input_ids=teacher_input_ids, + attention_mask=teacher_attention_mask, + ) + + # These are not used for ULD loss but are needed if JSD loss were to be used in this branch + student_prompt_length = inputs["prompts"].shape[1] + shifted_student_logits = outputs_student.logits[:, student_prompt_length - 1 : -1, :] + shifted_teacher_logits = outputs_teacher.logits[:, teacher_prompt_length - 1 : -1, :] + shifted_labels = inputs["labels"][:, student_prompt_length:] + else: + if self.use_liger_gkd_loss: + # Forward only through the base models (avoid lm_head to save memory) + unwrapped_student = self.accelerator.unwrap_model(model) + if hasattr(unwrapped_student, "get_decoder") and unwrapped_student.get_decoder() is not None: + base_student = unwrapped_student.get_decoder() + else: + base_student = getattr( + unwrapped_student, getattr(unwrapped_student, "base_model_prefix", "model"), unwrapped_student + ) + + student_outputs = base_student( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + + self.teacher_model.eval() + unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model) + if hasattr(unwrapped_teacher, "get_decoder") and unwrapped_teacher.get_decoder() is not None: + base_teacher = unwrapped_teacher.get_decoder() + else: + base_teacher = getattr( + unwrapped_teacher, getattr(unwrapped_teacher, "base_model_prefix", "model"), unwrapped_teacher + ) + with torch.no_grad(): + teacher_outputs = base_teacher( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + + # hidden states (shifted) + student_hidden = student_outputs.last_hidden_state[:, :-1] + teacher_hidden = teacher_outputs.last_hidden_state[:, :-1] + + # Release full outputs to free memory + del student_outputs, teacher_outputs + + # labels mask and labels (shifted) + labels_mask = inputs["labels"] != -100 + masked_input_ids = torch.where( + labels_mask, inputs["input_ids"], torch.full_like(inputs["input_ids"], -100) + ) + true_labels = masked_input_ids[:, 1:].contiguous() + + # heads + student_head = unwrapped_student.get_output_embeddings() + teacher_head = unwrapped_teacher.get_output_embeddings() + + # liger fused jsd loss + loss = self.liger_jsd_loss( + student_input=student_hidden, + student_weight=student_head.weight, + teacher_input=teacher_hidden, + teacher_weight=teacher_head.weight, + true_labels=true_labels, + student_bias=getattr(student_head, "bias", None), + teacher_bias=getattr(teacher_head, "bias", None), + ) + + # Release hidden states after loss computation + del student_hidden, teacher_hidden, true_labels + else: + # Original behavior for same tokenizer or when teacher_tokenizer is not provided + outputs_student = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + self.teacher_model.eval() + with torch.no_grad(): + outputs_teacher = self.teacher_model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + prompt_lengths = inputs["prompts"].shape[1] + shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] + shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :] + shifted_labels = inputs["labels"][:, prompt_lengths:] + loss = self.generalized_jsd_loss( + student_logits=shifted_student_logits, + teacher_logits=shifted_teacher_logits, + labels=shifted_labels, + beta=self.beta, + ) + + if self.use_uld_loss: + student_input_ids = inputs["input_ids"] + + # Use the *teacher* labels created above, not the student's. + teacher_labels_for_loss = teacher_labels if "teacher_labels" in locals() else inputs["labels"] + teacher_input_ids_for_loss = teacher_input_ids if "teacher_input_ids" in locals() else inputs["input_ids"] + + # Create properly masked student labels (fixing batch size > 1 issue) + student_labels = inputs["labels"].clone() + if hasattr(self.processing_class, "pad_token_id") and self.processing_class.pad_token_id is not None: + student_labels[student_labels == self.processing_class.pad_token_id] = -100 + + # Also mask pad tokens in teacher labels for consistency + if ( + hasattr(self, "teacher_tokenizer") + and hasattr(self.teacher_tokenizer, "pad_token_id") + and self.teacher_tokenizer.pad_token_id is not None + ): + teacher_labels[teacher_labels == self.teacher_tokenizer.pad_token_id] = -100 + + loss = self.uld_loss_fn( + student_logits=outputs_student.logits, + teacher_logits=outputs_teacher.logits, + student_labels=student_labels, + teacher_labels=teacher_labels_for_loss, + student_input_ids=student_input_ids, + teacher_input_ids=teacher_input_ids_for_loss, + ) + + # If ULD hybrid mode produced per-step matched/unmatched components, accumulate them for logging. + # Use gradient_accumulation_steps to mirror Trainer's windowing behavior. + if hasattr(self.uld_loss_fn, "last_matched_loss") and hasattr(self.uld_loss_fn, "last_unmatched_loss"): + try: + ga = max(1, int(self.args.gradient_accumulation_steps)) + except Exception: + ga = 1 + step_eq = 1.0 / ga + # read scalar values for logging + matched_val = ( + self.uld_loss_fn.last_matched_loss.item() + if self.uld_loss_fn.last_matched_loss is not None + else 0.0 + ) + unmatched_val = ( + self.uld_loss_fn.last_unmatched_loss.item() + if self.uld_loss_fn.last_unmatched_loss is not None + else 0.0 + ) + + self._matched_sum += matched_val + self._unmatched_sum += unmatched_val + self._matched_step_eq += step_eq + self._unmatched_step_eq += step_eq + + empty_cache() + + return (loss, outputs_student) if return_outputs else loss + + def generate_on_policy_outputs(self, model, inputs, generation_config, pad_token_id=None): + # Generate output with respect to the prompt only + if self.use_transformers_paged: + previous_attn = self.model.config._attn_implementation + if is_flash_attn_2_available(): + model.config._attn_implementation = "paged_attention" + else: + model.config._attn_implementation = "sdpa_paged" + prompt_mask = inputs.get("prompt_attention_mask") + prompts_tensor = inputs["prompts"] + if prompt_mask is not None: + prompt_sequences = [ + row[mask.bool()].detach().cpu().tolist() + for row, mask in zip(prompts_tensor, prompt_mask, strict=True) + ] + else: + prompt_sequences = [row.detach().cpu().tolist() for row in prompts_tensor] + generated_outputs = model.generate_batch(prompt_sequences, generation_config=generation_config) + model.config._attn_implementation = previous_attn + + completion_ids = [output.generated_tokens for output in generated_outputs.values()] + generated_tokens = torch.stack([torch.tensor(ids, device=model.device) for ids in completion_ids]) + else: + generated_outputs = model.generate( + input_ids=inputs["prompts"], + attention_mask=inputs.get("prompt_attention_mask", None), + generation_config=generation_config, + return_dict_in_generate=True, + ) + # Get the generated token IDs + generated_tokens = generated_outputs.sequences + + batch_size = generated_tokens.size(0) + device = generated_tokens.device + + prompt_mask = inputs.get("prompt_attention_mask") + pad_token_id = pad_token_id if pad_token_id is not None else self.processing_class.pad_token_id + + if prompt_mask is not None: + prompt_lengths = prompt_mask.sum(dim=1).to(torch.long) + else: + if pad_token_id is not None: + prompt_lengths = (inputs["prompts"] != pad_token_id).sum(dim=1).to(torch.long) + else: + prompt_lengths = torch.full( + (batch_size,), + inputs["prompts"].shape[1], + dtype=torch.long, + device=device, + ) + + new_input_ids = generated_tokens + new_attention_mask = torch.ones_like(new_input_ids) + if pad_token_id is not None: + new_attention_mask[new_input_ids == pad_token_id] = 0 + + new_labels = torch.full_like(new_input_ids, -100) + for idx in range(batch_size): + length = int(prompt_lengths[idx].item()) + new_labels[idx, length:] = new_input_ids[idx, length:] + + if pad_token_id is not None: + new_labels[new_input_ids == pad_token_id] = -100 + + prompt_texts = [] + completion_texts = [] + for idx in range(batch_size): + length = int(prompt_lengths[idx].item()) + prompt_tokens = inputs["prompts"][idx] + if prompt_mask is not None: + prompt_tokens = prompt_tokens[prompt_mask[idx].bool()] + elif pad_token_id is not None: + prompt_tokens = prompt_tokens[prompt_tokens != pad_token_id] + prompt_texts.append( + self.processing_class.decode( + prompt_tokens.tolist(), + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ) + ) + completion_tokens = new_input_ids[idx, length:] + completion_texts.append( + self.processing_class.decode( + completion_tokens.tolist(), + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ) + ) + + return new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts + + @profiling_decorator + def _generate_on_policy_outputs_vllm(self, inputs, generation_config, pad_token_id=None): + device = self.accelerator.device + + # Decode prompts for vLLM (without special tokens - vLLM expects clean text) + prompts_text_for_vllm = self.processing_class.batch_decode( + inputs["prompts"], + skip_special_tokens=True, + # clean_up_tokenization_spaces=False # Keep this commented unless specific issues arise + ) + # Remove padding token text if it appears, as vLLM expects clean prompts + if self.processing_class.pad_token: + prompts_text_for_vllm = [p.replace(self.processing_class.pad_token, "") for p in prompts_text_for_vllm] + + # Also decode prompts WITH special tokens for ULD loss computation + prompts_text_with_special = self.processing_class.batch_decode( + inputs["prompts"], + skip_special_tokens=False, + ) + + # system_prompt = "Please reason step by step, and put your final answer within \\boxed{}." + # target_system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." + # prompts_text = [p.replace(target_system_prompt, system_prompt) for p in prompts_text] + # Add system prompt to prompts + + max_completion_length = generation_config.max_completion_length + temperature = generation_config.temperature + # vLLM uses top_k=-1 for no top_k, transformers uses 0 or None. + top_k = generation_config.top_k if generation_config.top_k and generation_config.top_k > 0 else -1 + # top_p, repetition_penalty, min_p are not directly in generation_config, get from trainer args + top_p = self.args.top_p if hasattr(self.args, "top_p") else 1.0 + repetition_penalty = self.args.repetition_penalty if hasattr(self.args, "repetition_penalty") else 1.0 + min_p = self.args.min_p if hasattr(self.args, "min_p") else 0.0 + + if self.vllm_mode == "server": + all_prompts_text = gather_object(prompts_text_for_vllm) + if self.accelerator.is_main_process: + completion_ids = self.vllm_client.generate( + prompts=all_prompts_text, + n=1, # In GKD, we generate 1 completion per prompt from student + repetition_penalty=repetition_penalty, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + max_tokens=max_completion_length, + guided_decoding_regex=self.vllm_guided_decoding_regex, + ) + else: + completion_ids = [None] * len(all_prompts_text) + completion_ids = broadcast_object_list(completion_ids, from_process=0) + process_slice = slice( + self.accelerator.process_index * len(prompts_text_for_vllm), + (self.accelerator.process_index + 1) * len(prompts_text_for_vllm), + ) + completion_ids = completion_ids[process_slice] + elif self.vllm_mode == "colocate": + if self.vllm_guided_decoding_regex: + guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.vllm_guided_decoding_regex) + else: + guided_decoding = None + sampling_params = SamplingParams( + n=1, + repetition_penalty=repetition_penalty, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + max_tokens=max_completion_length, + guided_decoding=guided_decoding, + ) + + if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts_text_for_vllm) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts_text_for_vllm, group=self.vllm_tp_group) + all_prompts_text = [p for sublist in gathered_prompts for p in sublist] + else: + all_prompts_text = prompts_text_for_vllm + + all_outputs = self.vllm_engine.generate(all_prompts_text, sampling_params=sampling_params, use_tqdm=False) + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + + if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.vllm_tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + completion_ids = completion_ids[tp_slice] + + if self.vllm_enable_sleep_mode: + self.vllm_engine.sleep(level=2) + else: + raise ValueError(f"Unknown vllm_mode: {self.vllm_mode}") + + # We need to combine prompt and completion for new_input_ids + # Tokenize prompts again to get prompt_ids on the correct device and format + # Use prompts_text_for_vllm (without special tokens) for tokenization since vLLM expects clean text + # Ensure add_special_tokens=False as vLLM typically handles prompts as raw text + # Calculate max_length for prompts, ensuring it's positive + prompt_max_length = max(1, self.args.max_length - max_completion_length) if self.args.max_length else None + prompt_tokenized = self.processing_class( + prompts_text_for_vllm, + return_tensors="pt", + padding="longest", + truncation=True if prompt_max_length else False, + max_length=prompt_max_length, + add_special_tokens=False, + ).to(device) + prompt_ids = prompt_tokenized.input_ids + + completion_ids_tensors = [torch.tensor(ids, device=device) for ids in completion_ids] + # Manually pad/truncate completions to max_completion_length length before using pad function + padded_completion_ids_list = [] + for completion_tensor in completion_ids_tensors: + if len(completion_tensor) > max_completion_length: + # Truncate if longer than max_completion_length + padded_completion_ids_list.append(completion_tensor[:max_completion_length]) + elif len(completion_tensor) < max_completion_length: + # Pad if shorter than max_completion_length + padding_needed = max_completion_length - len(completion_tensor) + padded_tensor = torch.cat( + [ + completion_tensor, + torch.full((padding_needed,), pad_token_id, device=device, dtype=completion_tensor.dtype), + ] + ) + padded_completion_ids_list.append(padded_tensor) + else: + # Already the right length + padded_completion_ids_list.append(completion_tensor) + + # Now all tensors are the same length, so we can stack them + padded_completion_ids = torch.stack(padded_completion_ids_list) + + # Ensure prompt_ids and padded_completion_ids are 2D + if prompt_ids.ndim == 1: + prompt_ids = prompt_ids.unsqueeze(0) + if padded_completion_ids.ndim == 1: + padded_completion_ids = padded_completion_ids.unsqueeze(0) + + new_input_ids = torch.cat([prompt_ids, padded_completion_ids], dim=1) + + new_attention_mask = torch.ones_like(new_input_ids, device=device) + new_labels = new_input_ids.clone() + + if pad_token_id is not None: + new_labels[new_labels == pad_token_id] = -100 + new_attention_mask[new_input_ids == pad_token_id] = 0 + + # Mask prompt tokens in labels + prompt_lengths = prompt_ids.shape[1] + new_labels[:, :prompt_lengths] = -100 + + # IMPORTANT: Preserve original text for cross-tokenizer ULD loss + # Use prompts_text_with_special (with special tokens) for ULD loss computation + # Extract completion texts from the generated completion IDs + completion_texts = [] + for comp_ids in completion_ids: + completion_text = self.processing_class.decode(comp_ids, skip_special_tokens=False) + completion_texts.append(completion_text) + + return new_input_ids, new_attention_mask, new_labels, prompts_text_with_special, completion_texts + + def _sync_fsdp_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with student vLLM.""" + if visited is None: + visited = set() + + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + # recurse into the child + self._sync_fsdp_params_to_vllm(child_module, prefix=child_prefix, visited=visited) + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + for extra in ("_fsdp_wrapped_module.", "_checkpoint_wrapped_module."): + full_name = full_name.replace(extra, "") + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.vllm_engine.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + def _move_model_to_vllm(self): + """Synchronize student model weights to vLLM engine.""" + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if self.vllm_mode == "colocate" and self.vllm_enable_sleep_mode: + empty_cache() + self.vllm_engine.wake_up(tags=["weights"]) + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + self._sync_fsdp_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = name.replace("modules_to_save.default.", "") + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.vllm_engine.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + # use memory-efficient post-order traversal for FSDP + self._sync_fsdp_params_to_vllm(self.model) + else: + # For DeepSpeed ZeRO-3, gather each parameter individually like GRPO trainer + for name, param in self.model.named_parameters(): + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.vllm_engine.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.vllm_engine.reset_prefix_cache() + + def _wake_vllm_if_needed(self): + if self.vllm_mode == "colocate" and self.vllm_enable_sleep_mode: + empty_cache() + self.vllm_engine.wake_up(tags=["kv_cache"]) + + @profiling_decorator + def training_step( + self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + ) -> torch.Tensor: + """ + Perform a training step for the General Online Logit Distillation (GOLD) model. + + This method implements the on-policy learning approach described in the GOLD blog post. With probability + `self.lmbda`, it generates new responses using the student model, which are then used for training instead of + the offline original inputs. + """ + on_policy = False + if random.random() <= self.lmbda: + on_policy = True + if self.use_vllm: + self._wake_vllm_if_needed() + result = self._generate_on_policy_outputs_vllm( + inputs, self.generation_config, self.processing_class.pad_token_id + ) + new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + else: + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model: + result = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + new_input_ids, new_attention_mask, new_labels, prompt_texts, completion_texts = result + + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + + # CRITICAL: Preserve original text for cross-tokenizer ULD loss + # This ensures both off-policy (dataset) and on-policy (generated) samples + # can use proper text-based alignment for different tokenizers + inputs["original_prompt_text"] = prompt_texts + inputs["original_completion_text"] = completion_texts + + # Log prompt and completion texts + self._textual_logs["prompt"].extend(gather_object(prompt_texts)) + self._textual_logs["completion"].extend(gather_object(completion_texts)) + + loss = super().training_step(model, inputs, num_items_in_batch) + + loss_scalar = float(loss.detach()) + ga = max(1, int(self.args.gradient_accumulation_steps)) + step_equiv = 1.0 / ga + + if on_policy: + self._on_policy_loss_total += loss_scalar + self._on_policy_step_equiv += step_equiv + else: + self._off_policy_loss_total += loss_scalar + self._off_policy_step_equiv += step_equiv + return loss + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + if mode == "train": + device = self.accelerator.device if hasattr(self.accelerator, "device") else torch.device("cpu") + # include matched/unmatched accumulators for distributed reduction + vec = torch.tensor( + [ + self._on_policy_loss_total, + self._off_policy_loss_total, + self._on_policy_step_equiv, + self._off_policy_step_equiv, + self._matched_sum, + self._unmatched_sum, + self._matched_step_eq, + self._unmatched_step_eq, + ], + dtype=torch.float64, + device=device, + ) + + # Sum across processes so we mirror Trainer's distributed reduction + if ( + getattr(self.accelerator, "distributed_type", DistributedType.NO) != DistributedType.NO + and dist.is_available() + and dist.is_initialized() + ): + dist.all_reduce(vec, op=dist.ReduceOp.SUM) + + ( + on_sum, + off_sum, + on_eq, + off_eq, + matched_sum, + unmatched_sum, + matched_eq, + unmatched_eq, + ) = vec.tolist() + + # Compute category averages over the *same window* as Trainer's logs + # (avoid div-by-zero if, e.g., no on-policy steps in the window) + if on_eq > 0: + logs["on_policy_loss"] = round(on_sum / on_eq, 4) + if off_eq > 0: + logs["off_policy_loss"] = round(off_sum / off_eq, 4) + + # matched/unmatched averaged over same logging window (if present) + if matched_eq > 0: + logs["matched_loss"] = round(matched_sum / matched_eq, 4) + if unmatched_eq > 0: + logs["unmatched_loss"] = round(unmatched_sum / unmatched_eq, 4) + + # Reset window accumulators after logging (just like Trainer resets its window) + self._on_policy_loss_total = self._off_policy_loss_total = 0.0 + self._on_policy_step_equiv = self._off_policy_step_equiv = 0.0 + self._matched_sum = self._unmatched_sum = 0.0 + self._matched_step_eq = self._unmatched_step_eq = 0.0 + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if ( + self.accelerator.is_main_process + and self.log_completions + and ((self.state.global_step % self.log_completion_steps) == 0) + ): + if is_rich_available(): + print_prompt_completions_sample_uld( + self._textual_logs["prompt"], + self._textual_logs["completion"], + self.state.global_step, + self.num_completions_to_print, + ) + + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + import pandas as pd + + table = { + "step": [str(self.state.global_step)] * len(self._textual_logs["prompt"]), + "prompt": self._textual_logs["prompt"], + "completion": self._textual_logs["completion"], + } + df = pd.DataFrame(table) + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + if self.num_completions_to_print and len(df) > 0: + df = df.sample(n=self.num_completions_to_print, random_state=42) + wandb.log({"completions": wandb.Table(dataframe=df)}) diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__init__.py b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1482c63fc9233ba137d7a510b7c1fd08596f13d6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2020-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 .grpo_trainer import GRPOTrainer diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9a77c30573ade5c87289bf58214511001357572 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/grpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/grpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed07e8f9a425940d520f579a970eb1e3e83873b0 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/__pycache__/grpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/grpo_trainer.py b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/grpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f267bbd4b622d5f74b8d1acbd28c9f923a1308f9 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/gspo_token/grpo_trainer.py @@ -0,0 +1,156 @@ +# Copyright 2020-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 torch + +from trl import GRPOTrainer as _GRPOTrainer + +from ...trainer.utils import nanmax, nanmin + + +class GRPOTrainer(_GRPOTrainer): + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + ) + + if self.top_entropy_quantile < 1.0: + entropy_mask = self.get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile) + else: + entropy_mask = None + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + + # Compute the loss + advantages = inputs["advantages"] + # When num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps, + # old_per_token_logps == per_token_logps. In this case we can skip its computation + # (see _generate_and_score_completions) and instead use per_token_logps.detach(). + # The exception is when using vLLM, where we always compute old_per_token_logps + # for importance sampling + old_per_token_logps = inputs.get("old_per_token_logps") + old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps + + log_ratio = per_token_logps - old_per_token_logps + if self.importance_sampling_level == "token": + log_importance_weights = log_ratio + elif self.importance_sampling_level == "sequence": + log_importance_weights = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + log_importance_weights = log_importance_weights.unsqueeze(-1) + elif self.importance_sampling_level == "sequence_token": + # GSPO-token: sg[si(θ)] * πθ(yi,t)/sg[πθ(yi,t)] + seq_level_log_weight = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + seq_level_log_weight = seq_level_log_weight.detach().unsqueeze(-1) # Stop gradient + log_importance_weights = per_token_logps - per_token_logps.detach() + seq_level_log_weight + else: + raise ValueError( + f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " + "and 'sequence'." + ) + # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on + # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) + + coef_1 = torch.exp(log_importance_weights) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + + # Two-sided clipping + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages.unsqueeze(1) + per_token_loss2 = coef_2 * advantages.unsqueeze(1) + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + if entropy_mask is not None: + per_token_loss = per_token_loss * entropy_mask + + if self.use_vllm and self.vllm_importance_sampling_correction: + per_token_loss = per_token_loss * inputs["importance_sampling_ratio"] + + if self.beta != 0.0: + per_token_loss = per_token_loss + self.beta * per_token_kl + + if self.loss_type == "grpo": + loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean() + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "bnpo": + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "dr_grpo": + loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.max_completion_length) + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "dapo": + normalizer = inputs["num_items_in_batch"] / self.accelerator.num_processes + loss = (per_token_loss * completion_mask).sum() / normalizer + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + # Log the metrics + mode = "train" if self.model.training else "eval" + + completion_token_count = completion_mask.sum().clamp(min=1.0) + + def masked_batch_mean(x): + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return x.mean() + else: + return (x * completion_mask).sum() / completion_token_count + + if self.beta != 0.0: + mean_kl = masked_batch_mean(per_token_kl) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + mean_entropy = masked_batch_mean(entropies) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages.unsqueeze(1) < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages.unsqueeze(1) > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = masked_batch_mean(is_low_clipped.float()) + high_clip = masked_batch_mean(is_high_clipped.float()) + clip_ratio = masked_batch_mean(is_region_clipped.float()) + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss diff --git a/venv/lib/python3.10/site-packages/trl/experimental/papo/__init__.py b/venv/lib/python3.10/site-packages/trl/experimental/papo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d976785a8e383c5718ae52abfac9bec3d894ec01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/papo/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2020-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 .papo_config import PAPOConfig +from .papo_trainer import PAPOTrainer diff --git a/venv/lib/python3.10/site-packages/trl/experimental/papo/papo_config.py b/venv/lib/python3.10/site-packages/trl/experimental/papo/papo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..9cf31abc3e50447550de4a8dc2b6e7ef68422585 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/experimental/papo/papo_config.py @@ -0,0 +1,73 @@ +# Copyright 2020-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 dataclasses import dataclass +from typing import Literal + +from ...trainer.grpo_config import GRPOConfig + + +@dataclass +class PAPOConfig(GRPOConfig): + """ + Configuration class for PAPOTrainer. + + PAPO (Perception-Aware Policy Optimization) extends GRPO/DAPO for multimodal reasoning by adding an implicit + perception loss and double entropy regularization. + + Args: + perception_loss_weight (`float`, *optional*, defaults to `0.1`): + gamma Weight coefficient for the perception loss term. This encourages the model to be sensitive to visual + changes. + + mask_ratio (`float`, *optional*, defaults to `0.3`): + Ratio of the image to mask when computing perception loss. + + mask_type (`Literal["random", "patch", "grid"]`, *optional*, defaults to `"random"`): + Type of masking strategy to use. + + der_loss_weight1 (`float`, *optional*, defaults to `0.03`): + eta1 Weight coefficient for the Double Entropy Regularization (DER) term. This term encourages confident + predictions with original images (low entropy) and uncertain predictions with masked images (high entropy). + + der_loss_weight2 (`float`, *optional*, defaults to `0.03`): + eta2 Weight coefficient for the Double Entropy Regularization (DER) term. This term encourages confident + predictions with original images (low entropy) and uncertain predictions with masked images (high entropy). + + loss_type (`Literal["grpo", "dapo"]`, inherited from GRPOConfig): + Base loss type to use. Set to "grpo" for PAPO-G or "dapo" for PAPO-D. + """ + + perception_loss_weight: float = 0.1 + mask_ratio: float = 0.3 + mask_type: Literal["random", "patch", "grid"] = "random" + + # Added for Double Entropy Regularization + der_loss_weight1: float = 0.03 + der_loss_weight2: float = 0.03 + + def __post_init__(self): + super().__post_init__() + + # Validation + if not 0.0 <= self.mask_ratio <= 1.0: + raise ValueError(f"mask_ratio must be between 0 and 1, got {self.mask_ratio}") + + if self.der_loss_weight1 < 0 or self.der_loss_weight2 < 0: + raise ValueError( + f"der_loss_weight1 and der_loss_weight2 must be non-negative, got {self.der_loss_weight1} and {self.der_loss_weight2}" + ) + + if self.mask_type not in ["random", "patch", "grid"]: + raise ValueError(f"mask_type must be one of ['random', 'patch', 'grid'], got {self.mask_type}") diff --git a/venv/lib/python3.10/site-packages/trl/extras/__init__.py b/venv/lib/python3.10/site-packages/trl/extras/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a317018578146718bb56d1b494f882305aa0535c --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/extras/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2020-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. diff --git a/venv/lib/python3.10/site-packages/trl/extras/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4737337951761fecfa23e0bb8c8019db285549bf Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fab48d4152a33a2f03b199eabfc902d9724a816 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/extras/__pycache__/profiling.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/profiling.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7bb793f7df7ecf525047dae105315fde9e066f6 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/profiling.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/extras/__pycache__/vllm_client.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/vllm_client.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..987aa87eed28ca2fe815243ddbf6feef71646744 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/extras/__pycache__/vllm_client.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/extras/dataset_formatting.py b/venv/lib/python3.10/site-packages/trl/extras/dataset_formatting.py new file mode 100644 index 0000000000000000000000000000000000000000..d6a923b36cad6ad9ae0d2aa5342ef812bbb8f12b --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/extras/dataset_formatting.py @@ -0,0 +1,158 @@ +# Copyright 2020-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 logging +import warnings +from collections.abc import Callable +from typing import Literal + +import datasets +from datasets import Dataset, Value +from packaging import version +from transformers import AutoTokenizer + + +if version.parse(datasets.__version__) >= version.parse("4.0.0"): + from datasets import List + + FORMAT_MAPPING = { + "chatml": List({"content": Value(dtype="string", id=None), "role": Value(dtype="string", id=None)}), + "instruction": {"completion": Value(dtype="string", id=None), "prompt": Value(dtype="string", id=None)}, + } +else: + FORMAT_MAPPING = { + "chatml": [{"content": Value(dtype="string", id=None), "role": Value(dtype="string", id=None)}], + "instruction": {"completion": Value(dtype="string", id=None), "prompt": Value(dtype="string", id=None)}, + } + + +def conversations_formatting_function( + tokenizer: AutoTokenizer, messages_field: Literal["messages", "conversations"], tools: list | None = None +): + r""" + return a callable function that takes in a "messages" dataset and returns a formatted dataset, based on the + tokenizer apply chat template to the dataset along with the schema of the list of functions in the tools list. + + + + `conversations_formatting_function` is deprecated and will be removed in version 0.27. Please use + `tokenizer.apply_chat_template()` directly instead. + + + """ + warnings.warn( + "`conversations_formatting_function` is deprecated and will be removed in TRL 0.27. " + "Please use `tokenizer.apply_chat_template()` directly instead.", + FutureWarning, + stacklevel=2, + ) + + def format_dataset(examples): + if isinstance(examples[messages_field][0], list): + output_texts = [] + for i in range(len(examples[messages_field])): + output_texts.append( + tokenizer.apply_chat_template(examples[messages_field][i], tokenize=False, tools=tools) + ) + return output_texts + else: + return tokenizer.apply_chat_template(examples[messages_field], tokenize=False, tools=tools) + + return format_dataset + + +def instructions_formatting_function(tokenizer: AutoTokenizer): + r""" + return a callable function that takes in an "instructions" dataset and returns a formatted dataset, based on the + tokenizer apply chat template to the dataset + + + + `instructions_formatting_function` is deprecated and will be removed in version 0.27. Please use + `tokenizer.apply_chat_template()` directly instead. + + + """ + warnings.warn( + "`instructions_formatting_function` is deprecated and will be removed in TRL 0.27. " + "Please use `tokenizer.apply_chat_template()` directly instead.", + FutureWarning, + stacklevel=2, + ) + + def format_dataset(examples): + if isinstance(examples["prompt"], list): + output_texts = [] + for i in range(len(examples["prompt"])): + converted_sample = [ + {"role": "user", "content": examples["prompt"][i]}, + {"role": "assistant", "content": examples["completion"][i]}, + ] + output_texts.append(tokenizer.apply_chat_template(converted_sample, tokenize=False)) + return output_texts + else: + converted_sample = [ + {"role": "user", "content": examples["prompt"]}, + {"role": "assistant", "content": examples["completion"]}, + ] + return tokenizer.apply_chat_template(converted_sample, tokenize=False) + + return format_dataset + + +def get_formatting_func_from_dataset( + dataset: Dataset, tokenizer: AutoTokenizer, tools: list | None = None +) -> Callable | None: + r""" + Finds the correct formatting function based on the dataset structure. Currently supported datasets are: + - `ChatML` with [{"role": str, "content": str}] + - `instruction` with [{"prompt": str, "completion": str}] + + Args: + dataset (Dataset): User dataset + tokenizer (AutoTokenizer): Tokenizer used for formatting + tools (list, *optional*): List of tools (callable functions) that will be accessible to the model. + If the template does not support function calling, this argument will have no effect. + + Returns: + Callable: Formatting function if the dataset format is supported else None + + + + `get_formatting_func_from_dataset` is deprecated and will be removed in version 0.27. Please use + `tokenizer.apply_chat_template()` directly instead. + + + """ + warnings.warn( + "`get_formatting_func_from_dataset` is deprecated and will be removed in TRL 0.27. " + "Please use `tokenizer.apply_chat_template()` directly instead.", + FutureWarning, + stacklevel=2, + ) + + if isinstance(dataset, Dataset): + if "messages" in dataset.features: + if dataset.features["messages"] == FORMAT_MAPPING["chatml"]: + logging.info("Formatting dataset with chatml format") + return conversations_formatting_function(tokenizer, "messages", tools) + if "conversations" in dataset.features: + if dataset.features["conversations"] == FORMAT_MAPPING["chatml"]: + logging.info("Formatting dataset with chatml format") + return conversations_formatting_function(tokenizer, "conversations", tools) + elif dataset.features == FORMAT_MAPPING["instruction"]: + logging.info("Formatting dataset with instruction format") + return instructions_formatting_function(tokenizer) + + return None diff --git a/venv/lib/python3.10/site-packages/trl/extras/profiling.py b/venv/lib/python3.10/site-packages/trl/extras/profiling.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc7b40b5aa35840583360fcdcfabfac9bd86112 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/extras/profiling.py @@ -0,0 +1,100 @@ +# Copyright 2020-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 contextlib +import functools +import time +from collections.abc import Callable, Generator + +from transformers import Trainer +from transformers.integrations import is_mlflow_available, is_wandb_available + + +if is_wandb_available(): + import wandb + +if is_mlflow_available(): + import mlflow + + +@contextlib.contextmanager +def profiling_context(trainer: Trainer, name: str) -> Generator[None, None, None]: + """ + A context manager function for profiling a block of code. Results are logged to Weights & Biases or MLflow + depending on the trainer's configuration. + + Args: + trainer (`~transformers.Trainer`): + Trainer object. + name (`str`): + Name of the block to be profiled. Used as a key in the logged dictionary. + + Example: + ```python + from transformers import Trainer + from trl.extras.profiling import profiling_context + + + class MyTrainer(Trainer): + def some_method(self): + A = np.random.rand(1000, 1000) + B = np.random.rand(1000, 1000) + with profiling_context(self, "matrix_multiplication"): + # Code to profile: simulate a computationally expensive operation + result = A @ B # Matrix multiplication + ``` + """ + start_time = time.perf_counter() + yield + end_time = time.perf_counter() + duration = end_time - start_time + + profiling_metrics = {f"profiling/Time taken: {trainer.__class__.__name__}.{name}": duration} + if "wandb" in trainer.args.report_to and wandb.run is not None and trainer.accelerator.is_main_process: + wandb.log(profiling_metrics) + + if "mlflow" in trainer.args.report_to and mlflow.run is not None and trainer.accelerator.is_main_process: + mlflow.log_metrics(profiling_metrics, step=trainer.state.global_step) + + +def profiling_decorator(func: Callable) -> Callable: + """ + Decorator to profile a function and log execution time using [`extras.profiling.profiling_context`]. + + Args: + func (`Callable`): + Function to be profiled. + + Example: + ```python + from transformers import Trainer + from trl.extras.profiling import profiling_decorator + + + class MyTrainer(Trainer): + @profiling_decorator + def some_method(self): + A = np.random.rand(1000, 1000) + B = np.random.rand(1000, 1000) + # Code to profile: simulate a computationally expensive operation + result = A @ B + ``` + """ + + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + with profiling_context(self, func.__name__): + return func(self, *args, **kwargs) + + return wrapper diff --git a/venv/lib/python3.10/site-packages/trl/extras/vllm_client.py b/venv/lib/python3.10/site-packages/trl/extras/vllm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..5d76d422c16c0f8d332b601b14032939d49698e8 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/extras/vllm_client.py @@ -0,0 +1,515 @@ +# Copyright 2020-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 atexit +import base64 +import copy +import logging +import socket +import time +from io import BytesIO +from urllib.parse import urlparse + +import torch +import torch.distributed.distributed_c10d as c10d +from torch import nn +from transformers import is_torch_xpu_available + +from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available + + +if is_requests_available(): + import requests + from requests import ConnectionError + + +if is_vllm_available(): + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + if is_vllm_ascend_available(): + from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator + + +logger = logging.getLogger(__name__) + + +def pil_to_base64(image): + buffer = BytesIO() + image.save(buffer, format="PNG") + img_bytes = buffer.getvalue() + return base64.b64encode(img_bytes).decode("utf-8") + + +class VLLMClient: + """ + A client class to interact with a vLLM server. + + This class provides methods to generate completions, initialize and manage weight update groups, and update model + weights in a distributed setting. Before using it, start the vLLM server with `trl vllm-serve`. + + Args: + base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `host` and `server_port` are + ignored. + host (`str`, *optional*, defaults to `"0.0.0.0"`): + IP address of the vLLM server. Ignored if `base_url` is provided. + server_port (`int`, *optional*, defaults to `8000`): + Port number of the vLLM server. Ignored if `base_url` is provided. + group_port (`int`, *optional*, defaults to `51216`): + Port number for the weight update group. + connection_timeout (`float`, *optional*, defaults to `0.0`): + Total timeout duration in seconds to wait for the server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + Examples: + Run the vLLM server with the model `Qwen/Qwen2.5-7B`: + + ``` + $ trl vllm-serve --model Qwen/Qwen2.5-7B + ... + INFO: Application startup complete. + INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) + ``` + + Use the client to generate completions and update model weights: + + ```python + >>> from trl.extras.vllm_client import VLLMClient + + >>> client = VLLMClient() + >>> client.generate(["Hello, AI!", "Tell me a joke"]) + {'prompt_ids': [[9707, 11, 15235, 0], + [40451, 752, 264, 21646]], + 'completion_ids': [[11479, 752, 5046, 279, 1465, 304, 419, 23670, 2038, 358, 2776, 4378, 369, 847, 15549, 6733], + [911, 19654, 382, 3838, 1558, 279, 16158, 1977, 979, 498, 2299, 4460, 311, 10542, 432, 518]], + 'logprobs': [[-5.193126201629639, -0.05592319369316101, -4.861808776855469, -1.673396110534668, -2.6316866874694824, -0.2861405313014984, -0.35006725788116455, -5.23351526260376, -0.1447441577911377, -5.21489953994751, -1.6022650003433228, -1.9649192094802856, -2.1338791847229004, -1.2775304317474365, -10.004860877990723, -4.171003818511963], + [-0.012896230444312096, -5.747106552124023, -1.5248860120773315, -1.9286258220672607, -2.8512537479400635, -2.8055880069732666, -3.019822835922241, -0.37132859230041504, -0.6311739087104797, -2.562908411026001, -3.1664533615112305, -2.685293436050415, -0.007259538397192955, -7.339841842651367, -1.188662052154541, -3.54781436920166]]} + + >>> from transformers import AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B", device_map="cuda") + >>> client.init_communicator(device="cuda") + >>> client.update_model_params(model) + ``` + + There are several ways to initialize the client: + + ```python + VLLMClient(base_url="http://localhost:8000") + VLLMClient(base_url="http://192.168.1.100:8000") + VLLMClient(host="localhost", server_port=8000) + VLLMClient(host="192.168.1.100", server_port=8000) + ``` + """ + + def __init__( + self, + base_url: str | None = None, + host: str = "0.0.0.0", + server_port: int = 8000, + group_port: int = 51216, + connection_timeout: float = 0.0, + ): + if not is_requests_available(): + raise ImportError("requests is not installed. Please install it with `pip install requests`.") + if not is_vllm_available(): + raise ImportError("vLLM is not installed. Please install it with `pip install trl[vllm]`.") + + self.session = requests.Session() + + if base_url is not None: + # Parse the base_url to extract host and port + parsed_url = urlparse(base_url) + self.host = socket.gethostbyname(parsed_url.hostname) + scheme = parsed_url.scheme or "http" + self.base_url = f"{scheme}://{parsed_url.netloc}{parsed_url.path}" + else: + self.host = host + self.server_port = server_port + self.base_url = f"http://{self.host}:{self.server_port}" + self.group_port = group_port + self.check_server(connection_timeout) # check server and fail after timeout + + def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): + """ + Check server availability with retries on failure, within a total timeout duration. If the server is not up + after the total timeout duration, raise a `ConnectionError`. + + Args: + retry_interval (`float`, *optional*, defaults to `2.0`): + Interval in seconds between retries. + total_timeout (`float`, *optional*, defaults to `0.0`): + Total timeout duration in seconds. + """ + url = f"{self.base_url}/health/" + start_time = time.time() # Record the start time + + while True: + try: + response = requests.get(url) + except requests.exceptions.RequestException as exc: + # Check if the total timeout duration has passed + elapsed_time = time.time() - start_time + if elapsed_time >= total_timeout: + raise ConnectionError( + f"The vLLM server can't be reached at {self.base_url} after {total_timeout} seconds. Make " + "sure the server is running by running `trl vllm-serve`." + ) from exc + else: + if response.status_code == 200: + if "X-Forwarded-For" in response.headers: + self.host = response.headers["X-Forwarded-For"] + logger.info("Server is up!") + return None + + # Retry logic: wait before trying again + logger.info(f"Server is not up yet. Retrying in {retry_interval} seconds...") + time.sleep(retry_interval) + + def generate( + self, + prompts: list[str], + images: list | None = None, + n: int = 1, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + max_tokens: int = 16, + truncate_prompt_tokens: int | None = None, + guided_decoding_regex: str | None = None, + generation_kwargs: dict | None = None, + ) -> dict[str, list[list[int]]]: + """ + Generates model completions for the provided prompts. + + Args: + prompts (`list[str]`): + List of text prompts for which the model will generate completions. + images (`list[PIL.Image]`, *optional*): + List of PIL Images to send along with the prompts. + n (`int`, *optional*, defaults to `1`): + Number of completions to generate for each prompt. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Parameter for repetition penalty. 1.0 means no penalty. + temperature (`float`, *optional*, defaults to `1.0`): + Temperature parameter for sampling. Higher values increase diversity. + top_p (`float`, *optional*, defaults to `1.0`): + Top-p sampling parameter.`1.0` means no truncation. + top_k (`int`, *optional*, defaults to `-1`): + Top-k sampling parameter. `-1` means no truncation. + min_p (`float`, *optional*, defaults to `0.0`): + Minimum probability for sampling. + max_tokens (`int`, *optional*, defaults to `16`): + Maximum number of tokens to generate for each prompt. + truncate_prompt_tokens (`int`, *optional*): + If set to `-1`, will use the truncation size supported by the model. If set to an integer k, will use + only the last k tokens from the prompt (i.e., left truncation). If set to `None`, truncation is + disabled. + guided_decoding_regex (`str`, *optional*): + Regular expression to guide the decoding process. + generation_kwargs (`dict`, *optional*): + Additional generation parameters to pass to the vLLM `SamplingParams`. This can include parameters like + `seed`, `frequency_penalty`, etc. If it contains keys that conflict with the other parameters, they + will override them. + + Returns: + `dict` with keys: + - `prompt_ids` (`list[list[int]]`): + List of lists of token IDs representing the tokenized input prompts. + - `completion_ids` (`list[list[int]]`): + List of lists of token IDs representing the model-generated completions for each prompt. + - `logprobs` (`list[list[float]]`): + List of lists of log probabilities for each generated token. + """ + url = f"{self.base_url}/generate/" + + # Convert PIL images to base64 strings + images = [pil_to_base64(img) for img in images] if images else None + + response = self.session.post( + url, + json={ + "prompts": prompts, + "images": images, + "n": n, + "repetition_penalty": repetition_penalty, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_tokens": max_tokens, + "truncate_prompt_tokens": truncate_prompt_tokens, + "guided_decoding_regex": guided_decoding_regex, + "generation_kwargs": generation_kwargs or {}, + }, + ) + if response.status_code == 200: + json_response = response.json() + return { + "prompt_ids": json_response["prompt_ids"], + "completion_ids": json_response["completion_ids"], + "logprobs": json_response["logprobs"], + } + else: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + def chat( + self, + messages: list[list[dict]], + n: int = 1, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + max_tokens: int = 16, + truncate_prompt_tokens: int | None = None, + guided_decoding_regex: str | None = None, + generation_kwargs: dict | None = None, + chat_template_kwargs: dict | None = None, + ) -> dict[str, list[list[int]]]: + """ + Generates model completions for the provided chat messages. + + Args: + messages (`list[list[dict]]`): + List of message lists for which the model will generate completions. Each message is a dictionary with + keys like "role" and "content". + n (`int`, *optional*, defaults to `1`): + Number of completions to generate for each message list. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Parameter for repetition penalty. 1.0 means no penalty. + temperature (`float`, *optional*, defaults to `1.0`): + Temperature parameter for sampling. Higher values increase diversity. + top_p (`float`, *optional*, defaults to `1.0`): + Top-p sampling parameter.`1.0` means no truncation. + top_k (`int`, *optional*, defaults to `-1`): + Top-k sampling parameter. `-1` means no truncation. + min_p (`float`, *optional*, defaults to `0.0`): + Minimum probability for sampling. + max_tokens (`int`, *optional*, defaults to `16`): + Maximum number of tokens to generate for each message list. + truncate_prompt_tokens (`int`, *optional*): + If set to `-1`, will use the truncation size supported by the model. If set to an integer k, will use + only the last k tokens from the prompt (i.e., left truncation). If set to `None`, truncation is + disabled. + guided_decoding_regex (`str`, *optional*): + Regular expression to guide the decoding process. + generation_kwargs (`dict`, *optional*): + Additional generation parameters to pass to the vLLM `SamplingParams`. This can include parameters like + `seed`, `frequency_penalty`, etc. If it contains keys that conflict with the other parameters, they + will override them. + chat_template_kwargs (`dict`, *optional*): + Additional keyword arguments to customize the chat template used by the model. + + Returns: + `dict` with keys: + - `prompt_ids` (`list[list[int]]`): + List of lists of token IDs representing the tokenized input messages. + - `completion_ids` (`list[list[int]]`): + List of lists of token IDs representing the model-generated completions for each message list. + - `logprobs` (`list[list[float]]`): + List of lists of log probabilities for each generated token. + """ + url = f"{self.base_url}/chat/" + + # Convert PIL images to base64 strings + messages = copy.deepcopy(messages) # avoid modifying the original messages + for message_list in messages: + for message in message_list: + if isinstance(message["content"], list): + for part in message["content"]: + if part["type"] == "image_pil": + part["image_pil"] = pil_to_base64(part["image_pil"]) + + response = self.session.post( + url, + json={ + "messages": messages, + "n": n, + "repetition_penalty": repetition_penalty, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_tokens": max_tokens, + "truncate_prompt_tokens": truncate_prompt_tokens, + "guided_decoding_regex": guided_decoding_regex, + "generation_kwargs": generation_kwargs or {}, + "chat_template_kwargs": chat_template_kwargs or {}, + }, + ) + if response.status_code == 200: + json_response = response.json() + return { + "prompt_ids": json_response["prompt_ids"], + "completion_ids": json_response["completion_ids"], + "logprobs": json_response["logprobs"], + } + else: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + def init_communicator(self, device: torch.device | str | int = 0): + """ + Initializes the weight update group in a distributed setup for model synchronization. + + Args: + device (`torch.device`, `str`, or `int`, *optional*, defaults to `0`): + Device of trainer main process. It's the device that will be used for the weights synchronization. Can + be a `torch.device` object, a string like `'cuda:0'`, or an integer device index. + """ + # Get the world size from the server + url = f"{self.base_url}/get_world_size/" + response = requests.get(url) + if response.status_code == 200: + vllm_world_size = response.json()["world_size"] + else: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + world_size = vllm_world_size + 1 # add the client to the world + self.rank = vllm_world_size # the client's rank is the last process + + # Initialize weight update group + url = f"{self.base_url}/init_communicator/" + # Will simplify it after torch xpu 2.9 support get uuid. + if is_torch_xpu_available(): + if hasattr(torch.xpu.get_device_properties(device), "uuid"): + client_device_uuid = str(torch.xpu.get_device_properties(device).uuid) + else: + client_device_uuid = "42" + else: + client_device_uuid = str(torch.cuda.get_device_properties(device).uuid) + + # Set the weight update group's host to "0.0.0.0" so that + # clients from different IPs can send updated weights + response = self.session.post( + url, + json={ + "host": "0.0.0.0", + "port": self.group_port, + "world_size": world_size, + "client_device_uuid": client_device_uuid, + }, + ) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + # Brief delay to allow server initialization. While not strictly required (client socket will retry on + # connection failure), this prevents log warnings like: + # [W416 23:24:57.460001114 socket.cpp:204] [c10d] The hostname of the client socket cannot be retrieved. err=-3 + time.sleep(0.1) + + # Set up the communication group for weight broadcasting + if is_torch_xpu_available(): + store = torch.distributed.TCPStore( + host_name=self.host, port=self.group_port, world_size=world_size, is_master=(self.rank == 0) + ) + prefixed_store = c10d.PrefixStore("client2server", store) + pg = c10d.ProcessGroupXCCL( + store=prefixed_store, + rank=self.rank, + size=world_size, + ) + self.communicator = pg + else: + pg = StatelessProcessGroup.create( + host=self.host, port=self.group_port, rank=self.rank, world_size=world_size + ) + self.communicator = PyNcclCommunicator(pg, device=device) + + # When the client object is deleted, close the weight update group + atexit.register(self.close_communicator) + + def update_named_param(self, name: str, weights: torch.Tensor): + """ + Updates a specific named parameter in the model and broadcasts it to other processes. + + Args: + name (`str`): + Name of the layer whose weights are being updated. + weights (`torch.Tensor`): + Tensor containing the updated weights. + """ + dtype, shape = str(weights.dtype), tuple(weights.shape) + url = f"{self.base_url}/update_named_param/" + response = self.session.post(url, json={"name": name, "dtype": dtype, "shape": shape}) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + if is_torch_xpu_available(): + # Use XCCL to broadcast the updated weights from the client (src) to all workers. + self.communicator.broadcast(weights, root=self.rank) + self.communicator.barrier() + else: + # Use NCCL to broadcast the updated weights from the client (src) to all workers. + self.communicator.broadcast(weights, src=self.rank) + self.communicator.group.barrier() + + def update_model_params(self, model: nn.Module): + """ + Updates all parameters of the given model by calling `update_named_param` for each parameter in the model. + + Args: + model (`nn.Module`): + Model whose parameters (weights/biases) are to be updated. + """ + for name, param in model.named_parameters(): + # Update each parameter individually + self.update_named_param(name, param.data) + + def reset_prefix_cache(self): + """ + Resets the prefix cache for the model. + """ + url = f"{self.base_url}/reset_prefix_cache/" + response = self.session.post(url) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + def close_communicator(self): + """ + Closes the weight update group and cleans up the communication group. + """ + url = f"{self.base_url}/close_communicator/" + + try: + response = self.session.post(url) + except ConnectionError: + # The server might be already down, so we don't need to close the communicator + pass + else: + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + +# Example usage +if __name__ == "__main__": + from vllm import SamplingParams + + device = "xpu" if is_torch_xpu_available() else "cuda" + client = VLLMClient() + client.init_communicator(device=device) + + # Generate completions + responses = client.generate(["Hello, AI!", "Tell me a joke"], n=4, max_tokens=32, sampling_params=SamplingParams()) + print("Responses:", responses) # noqa + + # Update model weights + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B").to(device) + client.update_model_params(model) diff --git a/venv/lib/python3.10/site-packages/trl/import_utils.py b/venv/lib/python3.10/site-packages/trl/import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4d8a9c84ce098ae9f22786cfbf35d60934ec3027 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/import_utils.py @@ -0,0 +1,164 @@ +# Copyright 2020-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 +import os +import warnings +from itertools import chain +from types import ModuleType +from typing import Any + +from packaging import version +from transformers.utils.import_utils import _is_package_available + + +LIGER_KERNEL_MIN_VERSION = "0.5.8" + +# Use same as transformers.utils.import_utils +_deepspeed_available = _is_package_available("deepspeed") +_fastapi_available = _is_package_available("fastapi") +_joblib_available = _is_package_available("joblib") +_liger_kernel_available, _liger_kernel_version = _is_package_available("liger_kernel", return_version=True) +_llm_blender_available = _is_package_available("llm_blender") +_math_verify_available = _is_package_available("math_verify") +_mergekit_available = _is_package_available("mergekit") +_pydantic_available = _is_package_available("pydantic") +_requests_available = _is_package_available("requests") +_unsloth_available = _is_package_available("unsloth") +_uvicorn_available = _is_package_available("uvicorn") +_vllm_available, _vllm_version = _is_package_available("vllm", return_version=True) +_vllm_ascend_available = _is_package_available("vllm_ascend") +_weave_available = _is_package_available("weave") + + +def is_deepspeed_available() -> bool: + return _deepspeed_available + + +def is_fastapi_available() -> bool: + return _fastapi_available + + +def is_joblib_available() -> bool: + return _joblib_available + + +def is_liger_kernel_available(min_version: str = LIGER_KERNEL_MIN_VERSION) -> bool: + return _liger_kernel_available and version.parse(_liger_kernel_version) >= version.parse(min_version) + + +def is_llm_blender_available() -> bool: + return _llm_blender_available + + +def is_math_verify_available() -> bool: + return _math_verify_available + + +def is_mergekit_available() -> bool: + return _mergekit_available + + +def is_pydantic_available() -> bool: + return _pydantic_available + + +def is_requests_available() -> bool: + return _requests_available + + +def is_unsloth_available() -> bool: + return _unsloth_available + + +def is_uvicorn_available() -> bool: + return _uvicorn_available + + +def is_vllm_available() -> bool: + if _vllm_available and version.parse(_vllm_version) != version.parse("0.10.2"): + warnings.warn( + f"TRL currently only supports vLLM version `0.10.2`. You have version {_vllm_version} installed. We " + "recommend to install this version to avoid compatibility issues.", + UserWarning, + ) + return _vllm_available + + +def is_vllm_ascend_available() -> bool: + return _vllm_ascend_available + + +def is_weave_available() -> bool: + return _weave_available + + +class _LazyModule(ModuleType): + """ + Module class that surfaces all objects but only performs associated imports when the objects are requested. + """ + + # Very heavily inspired by optuna.integration._IntegrationModule + # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py + def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None): + super().__init__(name) + self._modules = set(import_structure.keys()) + self._class_to_module = {} + for key, values in import_structure.items(): + for value in values: + self._class_to_module[value] = key + # Needed for autocompletion in an IDE + self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values())) + self.__file__ = module_file + self.__spec__ = module_spec + self.__path__ = [os.path.dirname(module_file)] + self._objects = {} if extra_objects is None else extra_objects + self._name = name + self._import_structure = import_structure + + # Needed for autocompletion in an IDE + def __dir__(self): + result = super().__dir__() + # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether + # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir. + for attr in self.__all__: + if attr not in result: + result.append(attr) + return result + + def __getattr__(self, name: str) -> Any: + if name in self._objects: + return self._objects[name] + if name in self._modules: + value = self._get_module(name) + elif name in self._class_to_module.keys(): + module = self._get_module(self._class_to_module[name]) + value = getattr(module, name) + else: + raise AttributeError(f"module {self.__name__} has no attribute {name}") + + setattr(self, name, value) + return value + + def _get_module(self, module_name: str): + try: + return importlib.import_module("." + module_name, self.__name__) + except Exception as e: + raise RuntimeError( + f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its" + f" traceback):\n{e}" + ) from e + + def __reduce__(self): + return (self.__class__, (self._name, self.__file__, self._import_structure)) diff --git a/venv/lib/python3.10/site-packages/trl/mergekit_utils.py b/venv/lib/python3.10/site-packages/trl/mergekit_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fc9787b8f6bffe60149c901b11c28ab1ae163007 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/mergekit_utils.py @@ -0,0 +1,282 @@ +# Copyright 2020-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 torch +from huggingface_hub import HfApi + +from trl.import_utils import is_mergekit_available + + +if is_mergekit_available(): + from mergekit.config import MergeConfiguration + from mergekit.merge import MergeOptions, run_merge + + +def upload_model_to_hf(folder_path: str, repo_id: str): + api = HfApi() + # Create the repository if it doesn't exist + repo = api.create_repo(repo_id, repo_type="model") + + # Upload the folder to the specified repository + api.upload_folder( + folder_path=folder_path, + repo_id=repo.repo_id, + repo_type=repo.repo_type, + ) + + +class MergeConfig: + r""" + Configuration class for merging two models using `mergekit`. + + This class provides a structured way to configure and generate merge configurations for various merge methods, such + as `linear`, `ties`, `dare_ties`, and `slerp`. + + Args: + method (`str`, *optional*, defaults to `"linear"`): + Merge method to use. Supported methods include: + + - `"linear"`: Linearly combines two models with specified weights. + - `"ties"`: Combines two models using the TIES method with density parameters. + - `"dare_ties"`: A variant of TIES for domain adaptation. + - `"slerp"`: Combines models using spherical linear interpolation. + + Note: + + For more details about the merge methods and how they are implemented, see the [MergeKit GitHub + repository](https://github.com/arcee-ai/mergekit?tab=readme-ov-file#merge-methods). + + Attributes: + method (`str`): The merge method to use. + policy_model_path (`str` or `None`): Path to the policy model. + target_model_path (`str` or `None`): Path to the target model. + policy_model_weight (`float`): Weight for the policy model (for `linear` and `ties` methods). + target_model_weight (`float`): Weight for the target model (for `linear` and `ties` methods). + policy_model_density (`list[float]`): Density parameters for the policy model (for `ties` and `dare_ties`). + target_model_density (`list[float]`): Density parameters for the target model (for `ties` and `dare_ties`). + normalize (`float` or `None`): Normalization factor for the TIES method. + t_values (`float` or `None`): Interpolation factor for the SLERP method. + dtype (`str`): Data type to use for merging, e.g., `"float16"`. + """ + + def __init__(self, method: str = "linear"): + if not is_mergekit_available(): + raise ImportError("MergeConfig requires the `mergekit` extra. To install, run `pip install mergekit`.") + self.method = method + self.policy_model_path = None + self.target_model_path = None + + # Initialize relevant parameters based on the method + if method == "linear": + self.policy_model_weight = 0.5 + self.target_model_weight = 0.5 + self.dtype = "float16" + elif method == "ties": + self.policy_model_weight = 1.0 + self.policy_model_density = [1.0, 0.7, 0.1] + self.target_model_weight = 1.0 + self.target_model_density = [1.0] + self.normalize = 1.0 + self.dtype = "float16" + elif method == "dare_ties": + self.policy_model_weight = 1.0 + self.policy_model_density = [1.0, 0.7, 0.1] + self.target_model_weight = 1.0 + self.target_model_density = [1.0] + self.normalize = 1.0 + self.dtype = "float16" + elif method == "slerp": + self.t_values = 0.5 + self.dtype = "float16" + else: + raise ValueError(f"Unsupported merge method: {method}") + + def create_merge_config_linear(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a linear merge of two models with specified weights. + """ + # Create the merge configuration dictionary + merge_config_dict = { + "dtype": self.dtype, + "merge_method": "linear", + "models": [ + {"model": self.policy_model_path, "parameters": {"weight": self.policy_model_weight}}, + {"model": self.target_model_path, "parameters": {"weight": self.target_model_weight}}, + ], + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_ties(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a TIES merge of two models, with specified weights and densities. + """ + # Create the TIES merge configuration dictionary + merge_config_dict = { + "merge_method": "ties", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.target_model_density, "weight": self.target_model_weight}, + }, + { + "model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.policy_model_density, "weight": self.policy_model_weight}, + }, + ], + "parameters": {"normalize": self.normalize}, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_dare_ties(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a DARE TIES merge of two models, with specified weights and densities. + """ + # Create the DARE TIES merge configuration dictionary + merge_config_dict = { + "merge_method": "dare_ties", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.target_model_density, "weight": self.target_model_weight}, + }, + { + "model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.policy_model_density, "weight": self.policy_model_weight}, + }, + ], + "parameters": {"normalize": self.normalize}, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_slerp(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a SLERP merge of a model with a base model. + """ + + # Create the SLERP merge configuration dictionary + merge_config_dict = { + "merge_method": "slerp", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": None, # No specific parameters for SLERP model + } + ], + "parameters": { + "t": self.t_values # Set the t values for SLERP + }, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create(self) -> "MergeConfiguration": + if self.method == "linear": + return self.create_merge_config_linear() + elif self.method == "ties": + return self.create_merge_config_ties() + elif self.method == "dare_ties": + return self.create_merge_config_dare_ties() + elif self.method == "slerp": + return self.create_merge_config_slerp() + + +def merge_models(config: MergeConfig, out_path: str): + """ + Merge two models using mergekit + + Args: + config ([`MergeConfig`]): The merge configuration. + out_path (`str`): The output path for the merged model. + """ + if not is_mergekit_available(): + raise ImportError("merge_models requires the `mergekit` extra. To install, run `pip install mergekit`.") + run_merge( + config, + out_path=out_path, + options=MergeOptions( + device="auto", + cuda=torch.cuda.is_available(), + copy_tokenizer=True, + lazy_unpickle=False, + low_cpu_memory=False, + ), + ) diff --git a/venv/lib/python3.10/site-packages/trl/py.typed b/venv/lib/python3.10/site-packages/trl/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/trl/rewards/__init__.py b/venv/lib/python3.10/site-packages/trl/rewards/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ce10ce99fe26c6d5819d724a241ee8d1d6864b62 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/rewards/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2020-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 sys +from typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "accuracy_rewards": ["accuracy_reward"], + "format_rewards": ["think_format_reward"], + "other_rewards": ["get_soft_overlong_punishment"], +} + + +if TYPE_CHECKING: + from .accuracy_rewards import accuracy_reward + from .format_rewards import think_format_reward + from .other_rewards import get_soft_overlong_punishment + + +else: + sys.modules[__name__] = _LazyModule(__name__, __file__, _import_structure, module_spec=__spec__) diff --git a/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..783c73ff3b9697bb2d6f9309868b7593eda25f4b Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/accuracy_rewards.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/accuracy_rewards.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c29588435381449cab762f65a77c197678f0b4e Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/accuracy_rewards.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/format_rewards.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/format_rewards.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b889396a8766b5467e409e19c2d4dc9ea6d69140 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/format_rewards.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/other_rewards.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/other_rewards.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a7bea67d2d52dfe53dce33c39635614b36348b8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/rewards/__pycache__/other_rewards.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/rewards/accuracy_rewards.py b/venv/lib/python3.10/site-packages/trl/rewards/accuracy_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..1ae7d21426dd5e06bc8fd052675b665af51a6cba --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/rewards/accuracy_rewards.py @@ -0,0 +1,91 @@ +# Copyright 2020-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 trl.import_utils import is_math_verify_available + + +if is_math_verify_available(): + from latex2sympy2_extended import NormalizationConfig + from math_verify import LatexExtractionConfig, parse, verify + + +def accuracy_reward(completions: list[list[dict[str, str]]], solution: list[str], **kwargs) -> list[float | None]: + r""" + Reward function that checks if the completion is the same as the ground truth. + - If both gold and prediction are parseable → use math verification. + - If not parseable → compare as normalized text. + + Args: + completions (`list[list[dict[str, str]]]`): + List of completions to be evaluated. Each completion must be a list of one message, i.e. a dictionary + containing the key `"content"` with the value being the text of the completion. + solution: (`list[str]`): + List of the raw-text solutions to the questions/problems/prompts. + **kwargs: + Additional keyword arguments. This function does not use them, but they are required in the function + signature to ensure compatibility with trainers like [`GRPOTrainer`]. + Example: + ```python + >>> from trl.rewards import accuracy_reward + + >>> solution = [r"\frac{1}{3}", r"\frac{1}{3}"] + >>> completion = [ + ... [{"role": "assistant", "content": r"My answer is \boxed{\frac{1}{3}}"}], + ... [{"role": "assistant", "content": r"My answer is \boxed{\frac{1}{2}}"}], + ... ] + >>> accuracy_reward(completion, solution) + [1.0, 0.0] + ``` + """ + if not is_math_verify_available(): + raise ImportError("Please install the `math_verify` package to use accuracy_reward") + + contents = [completion[0]["content"] for completion in completions] + rewards = [] + for content, sol in zip(contents, solution, strict=True): + gold_parsed = parse( + sol, + extraction_mode="first_match", + ) + if len(gold_parsed) != 0: + # We require the answer to be provided in correct latex (no malformed operators) + answer_parsed = parse( + content, + extraction_config=[ + LatexExtractionConfig( + normalization_config=NormalizationConfig( + nits=False, + malformed_operators=False, + basic_latex=True, + boxed="all", + units=True, + ), + # Ensures that boxed is tried first + boxed_match_priority=0, + try_extract_without_anchor=False, + ) + ], + extraction_mode="first_match", + ) + # Compute binary rewards if verifiable, `None` otherwise to skip this example + try: + reward = float(verify(gold_parsed, answer_parsed)) + except Exception: + reward = None + else: + # If the gold solution is not parseable, we assign `None` to skip this example + reward = float(content.strip().lower() == sol.strip().lower()) + rewards.append(reward) + + return rewards diff --git a/venv/lib/python3.10/site-packages/trl/rewards/format_rewards.py b/venv/lib/python3.10/site-packages/trl/rewards/format_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..4df6b0a3bcc539ab6cab270dfde2b1071e615260 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/rewards/format_rewards.py @@ -0,0 +1,50 @@ +# Copyright 2020-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 re + + +def think_format_reward(completions: list[list[dict[str, str]]], **kwargs) -> list[float]: + r""" + Reward function that checks if the reasoning process is enclosed within `""` and `""` tags. The + function returns a reward of 1.0 if the format is correct, otherwise 0.0. + + Args: + completions (`list[list[dict[str, str]]]`): + List of completions to be evaluated. Each completion must be a list of one message, i.e. a dictionary + containing the key `"content"` with the value being the text of the completion. + **kwargs: + Additional keyword arguments. This function does not use them, but they are required in the function + signature to ensure compatibility with trainers like [`GRPOTrainer`]. + + Returns: + `list[float]`: + A list of rewards, where each reward is 1.0 if the completion matches the expected format, otherwise 0.0. + + Example: + ```python + >>> from trl.rewards import think_format_reward + + >>> completions = [ + ... [{"content": "\nThis is my reasoning.\n\nThis is my answer."}], + ... [{"content": "\nThis is my reasoning.\nThis is my answer."}], + ... ] + >>> think_format_reward(completions) + [1.0, 0.0] + ``` + """ + pattern = r"^(?!.*)(.*?).*$" + completion_contents = [completion[0]["content"] for completion in completions] + matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE) for content in completion_contents] + return [1.0 if match else 0.0 for match in matches] diff --git a/venv/lib/python3.10/site-packages/trl/rewards/other_rewards.py b/venv/lib/python3.10/site-packages/trl/rewards/other_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..ad922b4667be3093bbd9261b4aa0cdfa4aead066 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/rewards/other_rewards.py @@ -0,0 +1,62 @@ +# Copyright 2020-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 collections.abc import Callable + + +def get_soft_overlong_punishment(max_completion_len: int, soft_punish_cache: int) -> Callable: + # docstyle-ignore + r""" + Reward function that penalizes overlong completions. It is used to penalize overlong completions, but not to reward + shorter completions. Reference: Eq. (13) from the DAPO paper (https://huggingface.co/papers/2503.14476) + + $$ + R_{\text{length}}(y) = \begin{cases} + 0, & |y| \le L_{\max} - L_{\text{cache}} \\ + \dfrac{(L_{\max} - L_{\text{cache}}) - |y|}{L_{\text{cache}}}, & L_{\max} - L_{\text{cache}} < |y| \le L_{\max} \\ + -1, & L_{\max} < |y| + \end{cases} + $$ + + Args: + max_completion_len (`int`): + Maximum length of the completion, \( L_{\max} \). + soft_punish_cache (`int`): + Minimum length of the completion, \( L_{\text{cache}} \). If set to `0`, no minimum length is applied. + + Example: + ```python + from trl.rewards import get_soft_overlong_punishment + + soft_overlong_punishment = get_soft_overlong_punishment(max_completion_len=100, soft_punish_cache=20) + completion_ids = [[1] * 90] # simulating a completion with 90 tokens. 90 is between 80 and 100. + rewards = soft_overlong_punishment(completion_ids) + print(rewards) # [-0.5] + ``` + """ + + def soft_overlong_punishment_reward(completion_ids: list[list[int]], **kwargs) -> list[float]: + """Reward function that penalizes overlong completions.""" + rewards = [] + for ids in completion_ids: + completion_length = len(ids) + if completion_length <= max_completion_len - soft_punish_cache: + rewards.append(0.0) + elif max_completion_len - soft_punish_cache < completion_length <= max_completion_len: + rewards.append((max_completion_len - soft_punish_cache - completion_length) / soft_punish_cache) + else: + rewards.append(-1.0) + return rewards + + return soft_overlong_punishment_reward diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__init__.py b/venv/lib/python3.10/site-packages/trl/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..720a91a688f6a3110250ced02fb8061f0e9b1b66 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2020-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 typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "utils": ["DatasetMixtureConfig", "ScriptArguments", "TrlParser", "get_dataset", "init_zero_verbose"], +} + +if TYPE_CHECKING: + from .utils import DatasetMixtureConfig, ScriptArguments, TrlParser, get_dataset, init_zero_verbose +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..74e0ec50690c60ae46170a0120f48523c9c07b6d Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/dpo.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/dpo.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..453823c6caad6f799283a3d1aaad3d95fec4f861 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/dpo.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/env.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/env.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..064aa0e6ce1b4c3f8c4374202eff80a6cf252b91 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/env.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/grpo.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/grpo.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa5ce2e677328a63d0a2a351048a5f590fb85b5b Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/grpo.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/kto.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/kto.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d64f8ecdc1ffcd027db39ed5414d3aa1f4ee631c Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/kto.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/reward.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/reward.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2db8b2a48e78930236ee95e756fafe04ae437fb8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/reward.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/rloo.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/rloo.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b647fa8df4444fc78a68aeb01d131a274162f6e Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/rloo.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/sft.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/sft.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95631c77ae5b7a32cc83b99a2c785b10abac33a5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/sft.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5627e5dcd2e89de57a5dcaed9977385da4cea16 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f33ed2e89da585b492222aec749c421443fd97b Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/scripts/dpo.py b/venv/lib/python3.10/site-packages/trl/scripts/dpo.py new file mode 100644 index 0000000000000000000000000000000000000000..3efafd926ea26994af7487551261d7154189324f --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/dpo.py @@ -0,0 +1,186 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +""" +# Full training +```bash +python trl/scripts/dpo.py \ + --dataset_name trl-lib/ultrafeedback_binarized \ + --model_name_or_path Qwen/Qwen2-0.5B-Instruct \ + --learning_rate 5.0e-7 \ + --num_train_epochs 1 \ + --per_device_train_batch_size 2 \ + --max_steps 1000 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eval_strategy steps \ + --eval_steps 50 \ + --output_dir Qwen2-0.5B-DPO \ + --no_remove_unused_columns +``` + +# LoRA: +```bash +python trl/scripts/dpo.py \ + --dataset_name trl-lib/ultrafeedback_binarized \ + --model_name_or_path Qwen/Qwen2-0.5B-Instruct \ + --learning_rate 5.0e-6 \ + --num_train_epochs 1 \ + --per_device_train_batch_size 2 \ + --max_steps 1000 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eval_strategy steps \ + --eval_steps 50 \ + --output_dir Qwen2-0.5B-DPO \ + --no_remove_unused_columns \ + --use_peft \ + --lora_r 32 \ + --lora_alpha 16 +``` +""" + +import argparse +import os + +import torch +from accelerate import logging +from datasets import load_dataset +from transformers import AutoModelForCausalLM + +from trl import ( + DatasetMixtureConfig, + DPOConfig, + DPOTrainer, + ModelConfig, + ScriptArguments, + TrlParser, + get_dataset, + get_kbit_device_map, + get_peft_config, + get_quantization_config, +) + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +def main(script_args, training_args, model_args, dataset_args): + ################ + # Model + ################### + dtype = model_args.dtype if model_args.dtype in ["auto", None] else getattr(torch, model_args.dtype) + model_kwargs = dict( + revision=model_args.model_revision, + attn_implementation=model_args.attn_implementation, + dtype=dtype, + ) + quantization_config = get_quantization_config(model_args) + if quantization_config is not None: + # Passing None would not be treated the same as omitting the argument, so we include it only when valid. + model_kwargs["device_map"] = get_kbit_device_map() + model_kwargs["quantization_config"] = quantization_config + + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, **model_kwargs + ) + peft_config = get_peft_config(model_args) + if peft_config is None: + ref_model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, **model_kwargs + ) + else: + ref_model = None + if script_args.ignore_bias_buffers: + # torch distributed hack + model._ddp_params_and_buffers_to_ignore = [ + name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool + ] + + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the DPO trainer + trainer = DPOTrainer( + model, + ref_model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + peft_config=peft_config, + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + if training_args.eval_strategy != "no": + metrics = trainer.evaluate() + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (ScriptArguments, DPOConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser("dpo", help="Run the DPO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/env.py b/venv/lib/python3.10/site-packages/trl/scripts/env.py new file mode 100644 index 0000000000000000000000000000000000000000..74ca4d5c9ed097d8f9401ebf03b25e67a6422ed0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/env.py @@ -0,0 +1,88 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# ] +# /// + +import os +import platform +from importlib.metadata import version + +import torch +from accelerate.commands.config import default_config_file, load_config_from_file +from transformers import is_bitsandbytes_available +from transformers.utils import is_openai_available, is_peft_available + +from trl import __version__ +from trl.import_utils import ( + is_deepspeed_available, + is_liger_kernel_available, + is_llm_blender_available, + is_vllm_available, +) +from trl.scripts.utils import get_git_commit_hash + + +def print_env(): + devices = None + if torch.cuda.is_available(): + devices = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())] + elif torch.backends.mps.is_available(): + devices = ["MPS"] + elif torch.xpu.is_available(): + devices = [torch.xpu.get_device_name(i) for i in range(torch.xpu.device_count())] + + accelerate_config = accelerate_config_str = "not found" + + # Get the default from the config file. + if os.path.isfile(default_config_file): + accelerate_config = load_config_from_file(default_config_file).to_dict() + + accelerate_config_str = ( + "\n" + "\n".join([f" - {prop}: {val}" for prop, val in accelerate_config.items()]) + if isinstance(accelerate_config, dict) + else accelerate_config + ) + + commit_hash = get_git_commit_hash("trl") + + info = { + "Platform": platform.platform(), + "Python version": platform.python_version(), + "TRL version": f"{__version__}+{commit_hash[:7]}" if commit_hash else __version__, + "PyTorch version": version("torch"), + "accelerator(s)": ", ".join(devices) if devices is not None else "cpu", + "Transformers version": version("transformers"), + "Accelerate version": version("accelerate"), + "Accelerate config": accelerate_config_str, + "Datasets version": version("datasets"), + "HF Hub version": version("huggingface_hub"), + "bitsandbytes version": version("bitsandbytes") if is_bitsandbytes_available() else "not installed", + "DeepSpeed version": version("deepspeed") if is_deepspeed_available() else "not installed", + "Liger-Kernel version": version("liger_kernel") if is_liger_kernel_available() else "not installed", + "LLM-Blender version": version("llm_blender") if is_llm_blender_available() else "not installed", + "OpenAI version": version("openai") if is_openai_available() else "not installed", + "PEFT version": version("peft") if is_peft_available() else "not installed", + "vLLM version": version("vllm") if is_vllm_available() else "not installed", + } + + info_str = "\n".join([f"- {prop}: {val}" for prop, val in info.items()]) + print(f"\nCopy-paste the following information when reporting an issue:\n\n{info_str}\n") # noqa + + +if __name__ == "__main__": + print_env() diff --git a/venv/lib/python3.10/site-packages/trl/scripts/grpo.py b/venv/lib/python3.10/site-packages/trl/scripts/grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..7e46a3858fcdec3404946c0b93523fa03d543d01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/grpo.py @@ -0,0 +1,174 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +import argparse +import importlib +import os +import sys +from dataclasses import dataclass, field + +from accelerate import logging +from datasets import load_dataset + +from trl import ( + DatasetMixtureConfig, + GRPOConfig, + GRPOTrainer, + ModelConfig, + ScriptArguments, + TrlParser, + get_dataset, + get_peft_config, +) +from trl.rewards import accuracy_reward, get_soft_overlong_punishment, think_format_reward + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +reward_funcs_registry = { + "accuracy_reward": accuracy_reward, + "think_format_reward": think_format_reward, + "get_soft_overlong_punishment": get_soft_overlong_punishment(max_completion_len=1280, soft_punish_cache=256), +} + + +@dataclass +class GRPOScriptArguments(ScriptArguments): + """ + Script arguments for the GRPO training script. + + Args: + reward_model_name_or_path (`str`, *optional*): + Reward model id of a pretrained model hosted inside a model repo on huggingface.co or local path to a + directory containing model weights saved using [`~transformers.PreTrainedModel.save_pretrained`]. + reward_funcs (`list[str]`, *optional*): + Reward functions to use. Supported values are: + + - `"accuracy_reward"` + - `"think_format_reward"` + - `"get_soft_overlong_punishment"` (used value are `max_completion_len=1280`, `soft_punish_cache=256`) + - any dotted import path " (e.g., `'my_lib.rewards.custom_reward'`). + """ + + reward_model_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Reward model id of a pretrained model hosted inside a model repo on huggingface.co or " + "local path to a directory containing model weights saved using `PreTrainedModel.save_pretrained`." + }, + ) + reward_funcs: list[str] | None = field( + default=None, + metadata={ + "help": "Reward functions to use. Supported values are: `accuracy_reward`, `think_format_reward`, " + "`get_soft_overlong_punishment` (used value are `max_completion_len=1280`, `soft_punish_cache=256`), or " + "any dotted import path (e.g., `'my_lib.rewards.custom_reward'`)." + }, + ) + + +def main(script_args, training_args, model_args, dataset_args): + # Get the reward models and functions + reward_funcs = [] + if script_args.reward_model_name_or_path: + reward_funcs.append(script_args.reward_model_name_or_path) + + if script_args.reward_funcs: + for func_name in script_args.reward_funcs: + if func_name in reward_funcs_registry: + reward_funcs.append(reward_funcs_registry[func_name]) + elif "." in func_name: + module_path, func_name = func_name.rsplit(".", 1) + sys.path.insert(0, os.getcwd()) + module = importlib.import_module(module_path) + reward_func = getattr(module, func_name) + reward_funcs.append(reward_func) + else: + raise ValueError( + f"Could not load reward function '{func_name}'. Expected one of " + f"{list(reward_funcs_registry.keys())} or a valid import path." + ) + + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the GRPO trainer + trainer = GRPOTrainer( + model=model_args.model_name_or_path, + reward_funcs=reward_funcs, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + peft_config=get_peft_config(model_args), + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (GRPOScriptArguments, GRPOConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser("grpo", help="Run the GRPO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/kto.py b/venv/lib/python3.10/site-packages/trl/scripts/kto.py new file mode 100644 index 0000000000000000000000000000000000000000..3f8b9466a6544f58c1a45d4bd3e437b314625940 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/kto.py @@ -0,0 +1,166 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +""" +Run the KTO training script with the commands below. In general, the optimal configuration for KTO will be similar to +that of DPO. + +# Full training: +```bash +python trl/scripts/kto.py \ + --dataset_name trl-lib/kto-mix-14k \ + --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \ + --per_device_train_batch_size 16 \ + --num_train_epochs 1 \ + --learning_rate 5e-7 \ + --lr_scheduler_type=cosine \ + --gradient_accumulation_steps 1 \ + --eval_steps 500 \ + --output_dir=kto-aligned-model \ + --warmup_ratio 0.1 \ + --logging_first_step +``` + +# QLoRA: +```bash +# QLoRA: +python trl/scripts/kto.py \ + --dataset_name trl-lib/kto-mix-14k \ + --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \ + --per_device_train_batch_size 8 \ + --num_train_epochs 1 \ + --learning_rate 5e-7 \ + --lr_scheduler_type=cosine \ + --gradient_accumulation_steps 1 \ + --eval_steps 500 \ + --output_dir=kto-aligned-model-lora \ + --warmup_ratio 0.1 \ + --logging_first_step \ + --use_peft \ + --load_in_4bit \ + --lora_target_modules=all-linear \ + --lora_r=16 \ + --lora_alpha=16 +``` +""" + +import argparse +import os + +from accelerate import logging +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer + +from trl import ( + DatasetMixtureConfig, + KTOConfig, + KTOTrainer, + ModelConfig, + ScriptArguments, + TrlParser, + get_dataset, + get_peft_config, +) + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +def main(script_args, training_args, model_args, dataset_args): + # Load a pretrained model + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + ref_model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the KTO trainer + trainer = KTOTrainer( + model, + ref_model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=get_peft_config(model_args), + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (ScriptArguments, KTOConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser("kto", help="Run the KTO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/reward.py b/venv/lib/python3.10/site-packages/trl/scripts/reward.py new file mode 100644 index 0000000000000000000000000000000000000000..45a7a76cda2913a3baab2a5aa4b99c8dd29ef122 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/reward.py @@ -0,0 +1,108 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +import argparse +import os + +from accelerate import logging +from datasets import load_dataset + +from trl import ( + DatasetMixtureConfig, + ModelConfig, + RewardConfig, + RewardTrainer, + ScriptArguments, + TrlParser, + get_dataset, + get_peft_config, +) + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +def main(script_args, training_args, model_args, dataset_args): + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the RewardTrainer + trainer = RewardTrainer( + model=model_args.model_name_or_path, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + peft_config=get_peft_config(model_args), + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (ScriptArguments, RewardConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser( + "reward", help="Run the reward training script", dataclass_types=dataclass_types + ) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/rloo.py b/venv/lib/python3.10/site-packages/trl/scripts/rloo.py new file mode 100644 index 0000000000000000000000000000000000000000..9545d0f11c697e1da0512591aa509c996482ffd4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/rloo.py @@ -0,0 +1,174 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +import argparse +import importlib +import os +import sys +from dataclasses import dataclass, field + +from accelerate import logging +from datasets import load_dataset + +from trl import ( + DatasetMixtureConfig, + ModelConfig, + RLOOConfig, + RLOOTrainer, + ScriptArguments, + TrlParser, + get_dataset, + get_peft_config, +) +from trl.rewards import accuracy_reward, get_soft_overlong_punishment, think_format_reward + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +reward_funcs_registry = { + "accuracy_reward": accuracy_reward, + "think_format_reward": think_format_reward, + "get_soft_overlong_punishment": get_soft_overlong_punishment(max_completion_len=1280, soft_punish_cache=256), +} + + +@dataclass +class RLOOScriptArguments(ScriptArguments): + """ + Script arguments for the RLOO training script. + + Args: + reward_model_name_or_path (`str`, *optional*): + Reward model id of a pretrained model hosted inside a model repo on huggingface.co or local path to a + directory containing model weights saved using [`~transformers.PreTrainedModel.save_pretrained`]. + reward_funcs (`list[str]`, *optional*): + Reward functions to use. Supported values are: + + - `"accuracy_reward"` + - `"think_format_reward"` + - `"get_soft_overlong_punishment"` (used value are `max_completion_len=1280`, `soft_punish_cache=256`) + - any dotted import path " (e.g., `'my_lib.rewards.custom_reward'`). + """ + + reward_model_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Reward model id of a pretrained model hosted inside a model repo on huggingface.co or " + "local path to a directory containing model weights saved using `PreTrainedModel.save_pretrained`." + }, + ) + reward_funcs: list[str] | None = field( + default=None, + metadata={ + "help": "Reward functions to use. Supported values are: `accuracy_reward`, `think_format_reward`, " + "`get_soft_overlong_punishment` (used value are `max_completion_len=1280`, `soft_punish_cache=256`), or " + "any dotted import path (e.g., `'my_lib.rewards.custom_reward'`)." + }, + ) + + +def main(script_args, training_args, model_args, dataset_args): + # Get the reward models and functions + reward_funcs = [] + if script_args.reward_model_name_or_path: + reward_funcs.append(script_args.reward_model_name_or_path) + + if script_args.reward_funcs: + for func_name in script_args.reward_funcs: + if func_name in reward_funcs_registry: + reward_funcs.append(reward_funcs_registry[func_name]) + elif "." in func_name: + module_path, func_name = func_name.rsplit(".", 1) + sys.path.insert(0, os.getcwd()) + module = importlib.import_module(module_path) + reward_func = getattr(module, func_name) + reward_funcs.append(reward_func) + else: + raise ValueError( + f"Could not load reward function '{func_name}'. Expected one of " + f"{list(reward_funcs_registry.keys())} or a valid import path." + ) + + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the RLOO trainer + trainer = RLOOTrainer( + model=model_args.model_name_or_path, + reward_funcs=reward_funcs, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + peft_config=get_peft_config(model_args), + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (RLOOScriptArguments, RLOOConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser("rloo", help="Run the RLOO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/sft.py b/venv/lib/python3.10/site-packages/trl/scripts/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..2760d3a6c970454f82f86c84d07f9acebcccb266 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/sft.py @@ -0,0 +1,177 @@ +# Copyright 2020-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. + +# /// script +# dependencies = [ +# "trl", +# "peft", +# "trackio", +# "kernels", +# ] +# /// + +""" +# Full training +``` +python trl/scripts/sft.py \ + --model_name_or_path Qwen/Qwen2-0.5B \ + --dataset_name trl-lib/Capybara \ + --learning_rate 2.0e-5 \ + --num_train_epochs 1 \ + --packing \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eos_token '<|im_end|>' \ + --eval_strategy steps \ + --eval_steps 100 \ + --output_dir Qwen2-0.5B-SFT \ + --push_to_hub +``` + +# LoRA +``` +python trl/scripts/sft.py \ + --model_name_or_path Qwen/Qwen2-0.5B \ + --dataset_name trl-lib/Capybara \ + --learning_rate 2.0e-4 \ + --num_train_epochs 1 \ + --packing \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eos_token '<|im_end|>' \ + --eval_strategy steps \ + --eval_steps 100 \ + --use_peft \ + --lora_r 32 \ + --lora_alpha 16 \ + --output_dir Qwen2-0.5B-SFT \ + --push_to_hub +``` +""" + +import argparse +import os + +from accelerate import logging +from datasets import load_dataset +from transformers import AutoConfig, AutoModelForCausalLM +from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES + +from trl import ( + DatasetMixtureConfig, + ModelConfig, + ScriptArguments, + SFTConfig, + SFTTrainer, + TrlParser, + get_dataset, + get_kbit_device_map, + get_peft_config, + get_quantization_config, +) + + +logger = logging.get_logger(__name__) + +# Enable logging in a Hugging Face Space +os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio") + + +def main(script_args, training_args, model_args, dataset_args): + ################ + # Model init kwargs + ################ + model_kwargs = dict( + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + attn_implementation=model_args.attn_implementation, + dtype=model_args.dtype, + ) + quantization_config = get_quantization_config(model_args) + if quantization_config is not None: + # Passing None would not be treated the same as omitting the argument, so we include it only when valid. + model_kwargs["device_map"] = get_kbit_device_map() + model_kwargs["quantization_config"] = quantization_config + + # Create model + config = AutoConfig.from_pretrained(model_args.model_name_or_path) + valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values() + + if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures): + from transformers import AutoModelForImageTextToText + + model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs) + + # Load the dataset + if dataset_args.datasets and script_args.dataset_name: + logger.warning( + "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the " + "dataset and `dataset_name` will be ignored." + ) + dataset = get_dataset(dataset_args) + elif dataset_args.datasets and not script_args.dataset_name: + dataset = get_dataset(dataset_args) + elif not dataset_args.datasets and script_args.dataset_name: + dataset = load_dataset( + script_args.dataset_name, name=script_args.dataset_config, streaming=script_args.dataset_streaming + ) + else: + raise ValueError("Either `datasets` or `dataset_name` must be provided.") + + # Initialize the SFT trainer + trainer = SFTTrainer( + model=model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + peft_config=get_peft_config(model_args), + ) + + # Train the model + trainer.train() + + # Log training complete + trainer.accelerator.print("✅ Training completed.") + + # Save and push to Hub + trainer.save_model(training_args.output_dir) + trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.") + + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + trainer.accelerator.print(f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}.") + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + dataclass_types = (ScriptArguments, SFTConfig, ModelConfig, DatasetMixtureConfig) + if subparsers is not None: + parser = subparsers.add_parser("sft", help="Run the SFT training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, dataset_args, _ = parser.parse_args_and_config( + return_remaining_strings=True + ) + main(script_args, training_args, model_args, dataset_args) diff --git a/venv/lib/python3.10/site-packages/trl/scripts/utils.py b/venv/lib/python3.10/site-packages/trl/scripts/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fe8386ec5e1fae89a27dfe2479d0a71780e48876 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/utils.py @@ -0,0 +1,480 @@ +# Copyright 2020-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 argparse +import importlib +import inspect +import logging +import os +import subprocess +import sys +from collections.abc import Iterable +from dataclasses import dataclass, field + +import datasets +import yaml +from datasets import DatasetDict, concatenate_datasets +from transformers import HfArgumentParser +from transformers.hf_argparser import DataClass, DataClassType +from transformers.utils import is_rich_available + + +def _ensure_transformers_parallelism_config() -> None: + """ + Ensure that ``transformers.training_args`` always defines the symbol `ParallelismConfig` so that Python's + `typing.get_type_hints` can resolve annotations on `transformers.TrainingArguments` without raising a `NameError`. + + This is needed when running with ``accelerate<1.10.1``, where the module ``accelerate.parallelism_config`` did not + exist and therefore the type alias is not imported by Transformers. + + See upstream fix PR in transformers#40818. + """ + from typing import Any + + import transformers.training_args + + if not hasattr(transformers.training_args, "ParallelismConfig"): + transformers.training_args.ParallelismConfig = Any + + +_ensure_transformers_parallelism_config() # before creating HfArgumentParser + +logger = logging.getLogger(__name__) + + +@dataclass +class DatasetConfig: + """ + Configuration for a dataset. + + This class matches the signature of [`~datasets.load_dataset`] and the arguments are used directly in the + [`~datasets.load_dataset`] function. You can refer to the [`~datasets.load_dataset`] documentation for more + details. + + Parameters: + path (`str`): + Path or name of the dataset. + name (`str`, *optional*): + Defining the name of the dataset configuration. + data_dir (`str`, *optional*): + Defining the `data_dir` of the dataset configuration. If specified for the generic builders(csv, text etc.) + or the Hub datasets and `data_files` is `None`, the behavior is equal to passing `os.path.join(data_dir, + **)` as `data_files` to reference all the files in a directory. + data_files (`str` or `Sequence` or `Mapping`, *optional*): + Path(s) to source data file(s). + split (`str`, *optional*, defaults to `"train"`): + Which split of the data to load. + columns (`list[str]`, *optional*): + List of column names to select from the dataset. If `None`, all columns are selected. + """ + + path: str + name: str | None = None + data_dir: str | None = None + data_files: str | list[str] | dict[str, str] | None = None + split: str = "train" + columns: list[str] | None = None + + +@dataclass +class DatasetMixtureConfig: + """ + Configuration class for a mixture of datasets. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + datasets (`list[DatasetConfig]`): + List of dataset configurations to include in the mixture. + streaming (`bool`, *optional*, defaults to `False`): + Whether to stream the datasets. If `True`, the datasets will be loaded in streaming mode. + test_split_size (`float`, *optional*): + Size of the test split. Refer to the `test_size` parameter in the [`~datasets.train_test_split`] function + for more details. If `None`, the dataset will not be split into train and test sets. + + Usage: + When using the CLI, you can add the following section to your YAML config file: + + ```yaml + datasets: + - path: ... + name: ... + data_dir: ... + data_files: ... + split: ... + columns: ... + - path: ... + name: ... + data_dir: ... + data_files: ... + split: ... + columns: ... + streaming: ... + test_split_size: ... + ``` + """ + + datasets: list[DatasetConfig] = field( + default_factory=list, + metadata={"help": "List of dataset configurations to include in the mixture."}, + ) + streaming: bool = field( + default=False, + metadata={"help": "Whether to stream the datasets. If True, the datasets will be loaded in streaming mode."}, + ) + test_split_size: float | None = field( + default=None, + metadata={ + "help": "Size of the test split. Refer to the `test_size` parameter in the `datasets.train_test_split` " + "function for more details. If None, the dataset will not be split into train and test sets." + }, + ) + + def __post_init__(self): + # Convert any dataset dicts (from CLI/config parsing) into DatasetConfig objects + for idx, dataset in enumerate(self.datasets): + if isinstance(dataset, dict): + # If it's a dict, convert it to DatasetConfig + self.datasets[idx] = DatasetConfig(**dataset) + + +@dataclass +class ScriptArguments: + """ + Arguments common to all scripts. + + Args: + dataset_name (`str`,, *optional*): + Path or name of the dataset to load. If `datasets` is provided, this will be ignored. + dataset_config (`str`, *optional*): + Dataset configuration name. Corresponds to the `name` argument of the [`~datasets.load_dataset`] function. + If `datasets` is provided, this will be ignored. + dataset_train_split (`str`, *optional*, defaults to `"train"`): + Dataset split to use for training. If `datasets` is provided, this will be ignored. + dataset_test_split (`str`, *optional*, defaults to `"test"`): + Dataset split to use for evaluation. If `datasets` is provided, this will be ignored. + dataset_streaming (`bool`, *optional*, defaults to `False`): + Whether to stream the dataset. If True, the dataset will be loaded in streaming mode. If `datasets` is + provided, this will be ignored. + gradient_checkpointing_use_reentrant (`bool`, *optional*, defaults to `False`): + Whether to apply `use_reentrant` for gradient checkpointing. + ignore_bias_buffers (`bool`, *optional*, defaults to `False`): + Debug argument for distributed training. Fix for DDP issues with LM bias/mask buffers - invalid scalar + type, inplace operation. See + https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992. + """ + + dataset_name: str | None = field( + default=None, + metadata={"help": "Path or name of the dataset to load. If `datasets` is provided, this will be ignored."}, + ) + dataset_config: str | None = field( + default=None, + metadata={ + "help": "Dataset configuration name. Corresponds to the `name` argument of the `datasets.load_dataset` " + "function. If `datasets` is provided, this will be ignored." + }, + ) + dataset_train_split: str = field( + default="train", + metadata={"help": "Dataset split to use for training. If `datasets` is provided, this will be ignored."}, + ) + dataset_test_split: str = field( + default="test", + metadata={"help": "Dataset split to use for evaluation. If `datasets` is provided, this will be ignored."}, + ) + dataset_streaming: bool = field( + default=False, + metadata={ + "help": "Whether to stream the dataset. If True, the dataset will be loaded in streaming mode. If " + "`datasets` is provided, this will be ignored." + }, + ) + gradient_checkpointing_use_reentrant: bool = field( + default=False, + metadata={"help": "Whether to apply `use_reentrant` for gradient checkpointing."}, + ) + ignore_bias_buffers: bool = field( + default=False, + metadata={ + "help": "Debug argument for distributed training. Fix for DDP issues with LM bias/mask buffers - invalid " + "scalar type, inplace operation. See " + "https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992." + }, + ) + + +def init_zero_verbose(): + """ + Perform zero verbose init - use this method on top of the CLI modules to make logging and warning output cleaner. + Uses Rich if available, falls back otherwise. + """ + import logging + import warnings + + FORMAT = "%(message)s" + + if is_rich_available(): + from rich.logging import RichHandler + + handler = RichHandler() + else: + handler = logging.StreamHandler() + + logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[handler], level=logging.ERROR) + + # Custom warning handler to redirect warnings to the logging system + def warning_handler(message, category, filename, lineno, file=None, line=None): + logging.warning(f"{filename}:{lineno}: {category.__name__}: {message}") + + # Add the custom warning handler - we need to do that before importing anything to make sure the loggers work well + warnings.showwarning = warning_handler + + +class TrlParser(HfArgumentParser): + """ + A subclass of [`transformers.HfArgumentParser`] designed for parsing command-line arguments with dataclass-backed + configurations, while also supporting configuration file loading and environment variable management. + + Args: + dataclass_types (`DataClassType | Iterable[DataClassType]`, *optional*): + Dataclass types to use for argument parsing. + **kwargs: + Additional keyword arguments passed to the [`transformers.HfArgumentParser`] constructor. + + Examples: + + ```yaml + # config.yaml + env: + VAR1: value1 + arg1: 23 + ``` + + ```python + # main.py + import os + from dataclasses import dataclass + from trl import TrlParser + + + @dataclass + class MyArguments: + arg1: int + arg2: str = "alpha" + + + parser = TrlParser(dataclass_types=[MyArguments]) + training_args = parser.parse_args_and_config() + + print(training_args, os.environ.get("VAR1")) + ``` + + ```bash + $ python main.py --config config.yaml + (MyArguments(arg1=23, arg2='alpha'),) value1 + + $ python main.py --arg1 5 --arg2 beta + (MyArguments(arg1=5, arg2='beta'),) None + ``` + """ + + def __init__( + self, + dataclass_types: DataClassType | Iterable[DataClassType] | None = None, + **kwargs, + ): + # Make sure dataclass_types is an iterable + if dataclass_types is None: + dataclass_types = [] + elif not isinstance(dataclass_types, Iterable): + dataclass_types = [dataclass_types] + + # Check that none of the dataclasses have the "config" field + for dataclass_type in dataclass_types: + if "config" in dataclass_type.__dataclass_fields__: + raise ValueError( + f"Dataclass {dataclass_type.__name__} has a field named 'config'. This field is reserved for the " + f"config file path and should not be used in the dataclass." + ) + + super().__init__(dataclass_types=dataclass_types, **kwargs) + + def parse_args_and_config( + self, + args: Iterable[str] | None = None, + return_remaining_strings: bool = False, + fail_with_unknown_args: bool = True, + ) -> tuple[DataClass, ...]: + """ + Parse command-line args and config file into instances of the specified dataclass types. + + This method wraps [`transformers.HfArgumentParser.parse_args_into_dataclasses`] and also parses the config file + specified with the `--config` flag. The config file (in YAML format) provides argument values that replace the + default values in the dataclasses. Command line arguments can override values set by the config file. The + method also sets any environment variables specified in the `env` field of the config file. + """ + args = list(args) if args is not None else sys.argv[1:] + if "--config" in args: + # Get the config file path from + config_index = args.index("--config") + args.pop(config_index) # remove the --config flag + config_path = args.pop(config_index) # get the path to the config file + with open(config_path) as yaml_file: + config = yaml.safe_load(yaml_file) + + # Set the environment variables specified in the config file + if "env" in config: + env_vars = config.pop("env", {}) + if not isinstance(env_vars, dict): + raise ValueError("`env` field should be a dict in the YAML file.") + for key, value in env_vars.items(): + os.environ[key] = str(value) + + # Set the defaults from the config values + config_remaining_strings = self.set_defaults_with_config(**config) + else: + config_remaining_strings = [] + + # Parse the arguments from the command line + output = self.parse_args_into_dataclasses(args=args, return_remaining_strings=return_remaining_strings) + + # Merge remaining strings from the config file with the remaining strings from the command line + if return_remaining_strings: + args_remaining_strings = output[-1] + return output[:-1] + (config_remaining_strings + args_remaining_strings,) + elif fail_with_unknown_args and config_remaining_strings: + raise ValueError( + f"Unknown arguments from config file: {config_remaining_strings}. Please remove them, add them to the " + "dataclass, or set `fail_with_unknown_args=False`." + ) + else: + return output + + def set_defaults_with_config(self, **kwargs) -> list[str]: + """ + Overrides the parser's default values with those provided via keyword arguments, including for subparsers. + + Any argument with an updated default will also be marked as not required if it was previously required. + + Returns a list of strings that were not consumed by the parser. + """ + + def apply_defaults(parser, kw): + used_keys = set() + for action in parser._actions: + # Handle subparsers recursively + if isinstance(action, argparse._SubParsersAction): + for subparser in action.choices.values(): + used_keys.update(apply_defaults(subparser, kw)) + elif action.dest in kw: + action.default = kw[action.dest] + action.required = False + used_keys.add(action.dest) + return used_keys + + used_keys = apply_defaults(self, kwargs) + # Remaining args not consumed by the parser + remaining = [ + item for key, value in kwargs.items() if key not in used_keys for item in (f"--{key}", str(value)) + ] + return remaining + + +def get_git_commit_hash(package_name): + try: + # Import the package to locate its path + package = importlib.import_module(package_name) + # Get the path to the package using inspect + package_path = os.path.dirname(inspect.getfile(package)) + + # Navigate up to the Git repository root if the package is inside a subdirectory + git_repo_path = os.path.abspath(os.path.join(package_path, "..")) + git_dir = os.path.join(git_repo_path, ".git") + + if os.path.isdir(git_dir): + # Run the git command to get the current commit hash + commit_hash = ( + subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=git_repo_path).strip().decode("utf-8") + ) + return commit_hash + else: + return None + except Exception as e: + return f"Error: {str(e)}" + + +def get_dataset(mixture_config: DatasetMixtureConfig) -> DatasetDict: + """ + Load a mixture of datasets based on the configuration. + + Args: + mixture_config ([`DatasetMixtureConfig`]): + Script arguments containing dataset configuration. + + Returns: + [`~datasets.DatasetDict`]: + Combined dataset(s) from the mixture configuration, with optional train/test split if `test_split_size` is + set. + + Example: + ```python + from trl import DatasetMixtureConfig, get_dataset + from trl.scripts.utils import DatasetConfig + + mixture_config = DatasetMixtureConfig(datasets=[DatasetConfig(path="trl-lib/tldr")]) + dataset = get_dataset(mixture_config) + print(dataset) + ``` + + ``` + DatasetDict({ + train: Dataset({ + features: ['prompt', 'completion'], + num_rows: 116722 + }) + }) + ``` + """ + logger.info(f"Creating dataset mixture with {len(mixture_config.datasets)} datasets") + datasets_list = [] + for dataset_config in mixture_config.datasets: + logger.info(f"Loading dataset for mixture: {dataset_config.path} (config name: {dataset_config.name})") + dataset = datasets.load_dataset( + path=dataset_config.path, + name=dataset_config.name, + data_dir=dataset_config.data_dir, + data_files=dataset_config.data_files, + split=dataset_config.split, + streaming=mixture_config.streaming, + ) + if dataset_config.columns is not None: + dataset = dataset.select_columns(dataset_config.columns) + datasets_list.append(dataset) + + if datasets_list: + combined_dataset = concatenate_datasets(datasets_list) + if isinstance(combined_dataset, datasets.Dataset): # IterableDataset does not have a length + logger.info(f"Created dataset mixture with {len(combined_dataset)} examples") + + if mixture_config.test_split_size is not None: + logger.info(f"Splitting dataset into train and test sets with test size: {mixture_config.test_split_size}") + combined_dataset = combined_dataset.train_test_split(test_size=mixture_config.test_split_size) + return combined_dataset + else: + return DatasetDict({"train": combined_dataset}) + else: + raise ValueError("No datasets were loaded from the mixture configuration") diff --git a/venv/lib/python3.10/site-packages/trl/scripts/vllm_serve.py b/venv/lib/python3.10/site-packages/trl/scripts/vllm_serve.py new file mode 100644 index 0000000000000000000000000000000000000000..d5884b8290a4cd8e6492f1a72d09578f82cc2a03 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/scripts/vllm_serve.py @@ -0,0 +1,853 @@ +# Copyright 2020-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 argparse +import base64 +import logging +import os +from collections.abc import Sequence +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from io import BytesIO +from itertools import chain +from multiprocessing import Pipe, Process +from multiprocessing.connection import Connection + +import torch +import torch.distributed.distributed_c10d as c10d +from transformers import is_torch_xpu_available, is_vision_available + +from trl import TrlParser +from trl.import_utils import ( + is_fastapi_available, + is_pydantic_available, + is_uvicorn_available, + is_vllm_ascend_available, + is_vllm_available, +) + + +if is_fastapi_available(): + from fastapi import FastAPI + + +if is_pydantic_available(): + from pydantic import BaseModel + + +if is_uvicorn_available(): + import uvicorn + + +if is_vision_available(): + from PIL import Image + + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.parallel_state import get_world_group + from vllm.distributed.utils import StatelessProcessGroup + from vllm.sampling_params import GuidedDecodingParams + from vllm.utils import get_open_port + + if is_vllm_ascend_available(): + from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator + + +logger = logging.getLogger(__name__) + +# We use CUDA with multiprocessing, so we must use the 'spawn' start method. Otherwise, we will get the following +# error: RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use +# the 'spawn' start method +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + + +class WeightSyncWorkerExtension: + """ + A vLLM worker extension that enables weight synchronization between a client and multiple server workers. + + This worker uses a `StatelessProcessGroup` to establish communication and a `PyNcclCommunicator` or + `ProcessGroupXCCL` to handle efficient GPU-based communication using NCCL. The primary purpose of this class is to + receive updated model weights from a client process and distribute them to all worker processes participating in + model inference. + """ + + # The following attributes are initialized when `init_communicator` method is called. + communicator = None # Communicator for weight updates + client_rank = None # Source rank for broadcasting updated weights + + def init_communicator(self, host: str, port: int, world_size: int, client_device_uuid: str) -> None: + """ + Initializes the weight update communicator using a stateless process group. + + This method creates a `StatelessProcessGroup` that allows external training processes to communicate with vLLM + workers without interfering with the global torch distributed group. + + Args: + host (`str`): + Hostname or IP address of the master node. + port (`int`): + Port number to be used for communication. + world_size (`int`): + Total number of participating processes in the update group. + client_device_uuid (`str`): + UUID of the device of client main process. Used to assert that devices are different from vllm workers + devices. + """ + if self.communicator is not None: + raise RuntimeError("Weight update group already initialized. Call close_communicator first.") + + # TODO: will remove after torch xpu 2.9 support uuid in get_device_properties + if torch.cuda.is_available() or ( + is_torch_xpu_available() and hasattr(torch.xpu.get_device_properties(self.device), "uuid") + ): + accelerator_module = torch.xpu if is_torch_xpu_available() else torch.cuda + if client_device_uuid == str(accelerator_module.get_device_properties(self.device).uuid): + raise RuntimeError( + f"Attempting to use the same CUDA device (UUID: {client_device_uuid}) for multiple distinct " + "roles/ranks within the same communicator. This setup is unsupported and will likely lead to program " + "hangs or incorrect behavior. Ensure that trainer is using different devices than vLLM server." + ) + # Get the rank of the current worker in the global world group. + rank = get_world_group().rank + + if is_torch_xpu_available(): + store = torch.distributed.TCPStore(host_name=host, port=port, world_size=world_size, is_master=(rank == 0)) + prefixed_store = c10d.PrefixStore("client2server", store) + pg = c10d.ProcessGroupXCCL( + store=prefixed_store, + rank=rank, + size=world_size, + ) + self.communicator = pg + else: + # Create a stateless process group to manage communication between training processes and vLLM workers. + # Initialize the NCCL-based communicator for weight synchronization. + pg = StatelessProcessGroup.create(host=host, port=port, rank=rank, world_size=world_size) + self.communicator = PyNcclCommunicator(pg, device=self.device) + + # The client process that sends updated weights has the highest rank (world_size - 1). + self.client_rank = world_size - 1 + + def update_named_param(self, name: str, dtype: str, shape: Sequence[int]) -> None: + """ + Receives updated weights from the client process and updates the named parameter in the model. + + Args: + name (`str`): + Name of the weight tensor being updated. + dtype (`str`): + Data type of the weight tensor as a string (e.g., `"torch.float32"`). + shape (`Sequence[int]`): + Shape of the weight tensor. + """ + if self.communicator is None: + raise RuntimeError("Communicator not initialized. Call `init_communicator` first.") + + dtype = getattr(torch, dtype.split(".")[-1]) + # Allocate memory for the incoming weight tensor on the correct device. + weight = torch.empty(shape, dtype=dtype, device=self.device) + + if is_torch_xpu_available(): + # Use XCCL to broadcast the updated weights from the client (src) to all workers. + self.communicator.broadcast(weight, root=self.client_rank) + self.communicator.barrier() + else: + # Use NCCL to broadcast the updated weights from the client (src) to all workers. + self.communicator.broadcast(weight, src=self.client_rank) + self.communicator.group.barrier() + + # Load the received weights into the model. + self.model_runner.model.load_weights(weights=[(name, weight)]) + + def close_communicator(self) -> None: + """ + Closes the communicator when weight synchronization is no longer needed. + + This method deletes the NCCL communicator to release associated resources. + """ + + if self.communicator is not None: + del self.communicator + self.communicator = None # Ensure attribute is reset to None + self.client_rank = None # Ensure attribute is reset to None + + +@dataclass +class ScriptArguments: + r""" + Arguments for the script. + + Args: + model (`str`): + Model name or path to load the model from. + revision (`str`, *optional*): + Revision to use for the model. If not specified, the default branch will be used. + tensor_parallel_size (`int`, *optional*, defaults to `1`): + Number of tensor parallel workers to use. + data_parallel_size (`int`, *optional*, defaults to `1`): + Number of data parallel workers to use. + host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host address to run the server on. + port (`int`, *optional*, defaults to `8000`): + Port to run the server on. + gpu_memory_utilization (`float`, *optional*, defaults to `0.9`): + Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache on the + device dedicated to generation powered by vLLM. Higher values will increase the KV cache size and thus + improve the model's throughput. However, if the value is too high, it may cause out-of-memory (OOM) errors + during initialization. + dtype (`str`, *optional*, defaults to `"auto"`): + Data type to use for vLLM generation. If set to `"auto"`, the data type will be automatically determined + based on the model configuration. Find the supported values in the vLLM documentation. + max_model_len (`int`, *optional*): + If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced + `vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model + context size, which might be much larger than the KV cache, leading to inefficiencies. + enable_prefix_caching (`bool`, *optional*): + Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the hardware support + this feature. + enforce_eager (`bool`, *optional*, defaults to `False`): + Whether to enforce eager execution. If set to `True`, we will disable CUDA graph and always execute the + model in eager mode. If `False` (default behavior), we will use CUDA graph and eager execution in hybrid. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + kv_cache_dtype (`str`, *optional*, defaults to `"auto"`): + Data type to use for KV cache. If set to `"auto"`, the dtype will default to the model data type. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to trust remote code when loading models. Set to `True` to allow executing code from model + repositories. This is required for some custom models but introduces security risks. + log_level (`str`, *optional*, defaults to `"info"`): + Log level for uvicorn. Possible choices: `"critical"`, `"error"`, `"warning"`, `"info"`, `"debug"`, + `"trace"`. + """ + + model: str = field( + metadata={"help": "Model name or path to load the model from."}, + ) + revision: str | None = field( + default=None, + metadata={"help": "Revision to use for the model. If not specified, the default branch will be used."}, + ) + tensor_parallel_size: int = field( + default=1, + metadata={"help": "Number of tensor parallel workers to use."}, + ) + data_parallel_size: int = field( + default=1, + metadata={"help": "Number of data parallel workers to use."}, + ) + host: str = field( + default="0.0.0.0", + metadata={"help": "Host address to run the server on."}, + ) + port: int = field( + default=8000, + metadata={"help": "Port to run the server on."}, + ) + gpu_memory_utilization: float = field( + default=0.9, + metadata={ + "help": "Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV " + "cache on the device dedicated to generation powered by vLLM. Higher values will increase the KV cache " + "size and thus improve the model's throughput. However, if the value is too high, it may cause " + "out-of-memory (OOM) errors during initialization." + }, + ) + dtype: str = field( + default="auto", + metadata={ + "help": "Data type to use for vLLM generation. If set to 'auto', the data type will be automatically " + "determined based on the model configuration. Find the supported values in the vLLM documentation." + }, + ) + max_model_len: int | None = field( + default=None, + metadata={ + "help": "If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced " + "`vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model " + "context size, which might be much larger than the KV cache, leading to inefficiencies." + }, + ) + enable_prefix_caching: bool | None = field( + default=None, + metadata={ + "help": "Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the " + "hardware support this feature." + }, + ) + enforce_eager: bool | None = field( + default=False, + metadata={ + "help": "Whether to enforce eager execution. If set to `True`, we will disable CUDA graph and always " + "execute the model in eager mode. If `False` (default behavior), we will use CUDA graph and eager " + "execution in hybrid." + }, + ) + kv_cache_dtype: str = field( + default="auto", + metadata={ + "help": "Data type to use for KV cache. If set to 'auto', the dtype will default to the model data type." + }, + ) + trust_remote_code: bool = field( + default=False, + metadata={ + "help": "Whether to trust remote code when loading models. Set to True to allow executing code from model " + "repositories. This is required for some custom models but introduces security risks." + }, + ) + log_level: str = field( + default="info", + metadata={ + "help": "Log level for uvicorn. Possible choices: 'critical', 'error', 'warning', 'info', 'debug', " + "'trace'." + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={ + "help": "Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: " + "Use the `transformers` backend for model implementation. `vllm`: Use the `vllm` library for " + "model implementation." + }, + ) + + +def llm_worker( + script_args: ScriptArguments, data_parallel_rank: int, master_port: int, connection: Connection +) -> None: + # Set required environment variables for DP to work with vLLM + os.environ["VLLM_DP_RANK"] = str(data_parallel_rank) + os.environ["VLLM_DP_RANK_LOCAL"] = str(data_parallel_rank) + os.environ["VLLM_DP_SIZE"] = str(script_args.data_parallel_size) + os.environ["VLLM_DP_MASTER_PORT"] = str(master_port) + + llm = LLM( + model=script_args.model, + revision=script_args.revision, + tensor_parallel_size=script_args.tensor_parallel_size, + gpu_memory_utilization=script_args.gpu_memory_utilization, + enforce_eager=script_args.enforce_eager, + dtype=script_args.dtype, + # Automatic Prefix Caching caches the KV cache of existing queries, so that a new query can + # directly reuse the KV cache if it shares the same prefix with one of the existing queries. + # This is particularly useful here because we generate completions from the same prompts. + enable_prefix_caching=script_args.enable_prefix_caching, + kv_cache_dtype=script_args.kv_cache_dtype, + max_model_len=script_args.max_model_len, + worker_extension_cls="trl.scripts.vllm_serve.WeightSyncWorkerExtension", + trust_remote_code=script_args.trust_remote_code, + model_impl=script_args.vllm_model_impl, + ) + + # Send ready signal to parent process + connection.send({"status": "ready"}) + + while True: + # Wait for commands from the parent process + try: + command = connection.recv() + except KeyboardInterrupt: + llm.collective_rpc(method="close_communicator") + break + + # Handle commands + if command["type"] in ["call", "fire_and_forget"]: + method_name = command["method"] + args, kwargs = command.get("args", ()), command.get("kwargs", {}) + method = getattr(llm, method_name) + result = method(*args, **kwargs) + if command["type"] == "call": + connection.send(result) + elif command["type"] == "shutdown": + break + + +def chunk_list(lst: list, n: int) -> list[list]: + """ + Split list `lst` into `n` evenly distributed sublists. + + Example: + ```python + >>> chunk_list([1, 2, 3, 4, 5, 6], 2) + [[1, 2, 3], [4, 5, 6]] + + >>> chunk_list([1, 2, 3, 4, 5, 6], 4) + [[1, 2], [3, 4], [5], [6]] + + >>> chunk_list([1, 2, 3, 4, 5, 6], 8) + [[1], [2], [3], [4], [5], [6], [], []] + ``` + """ + k, r = divmod(len(lst), n) + return [lst[i * k + min(i, r) : (i + 1) * k + min(i + 1, r)] for i in range(n)] + + +def sanitize_logprob(logprob): + import math + + value = logprob.logprob + if math.isnan(value): + logger.warning(f"Generated NaN logprob, token logprob '{logprob}' will be ignored") + return None + + return value + + +def main(script_args: ScriptArguments): + if not is_fastapi_available(): + raise ImportError( + "FastAPI is required to run the vLLM serve script. Please install it using `pip install fastapi`." + ) + + if not is_pydantic_available(): + raise ImportError( + "Pydantic is required to run the vLLM serve script. Please install it using `pip install pydantic`." + ) + + if not is_uvicorn_available(): + raise ImportError( + "Uvicorn is required to run the vLLM serve script. Please install it using `pip install uvicorn`." + ) + + if not is_vllm_available(): + raise ImportError("vLLM is required to run the vLLM serve script. Please install it using `pip install vllm`.") + + # Spawn dp workers, and setup pipes for communication + master_port = get_open_port() + connections = [] + processes = [] + for data_parallel_rank in range(script_args.data_parallel_size): + parent_connection, child_connection = Pipe() + process = Process(target=llm_worker, args=(script_args, data_parallel_rank, master_port, child_connection)) + process.start() + connections.append(parent_connection) + processes.append(process) + + @asynccontextmanager + async def lifespan(app: FastAPI): + # Wait for all workers to send "ready" + ready_connections = set() + while len(ready_connections) < script_args.data_parallel_size: + for connection in connections: + msg = connection.recv() + if isinstance(msg, dict) and msg.get("status") == "ready": + ready_connections.add(connection) + + yield + + # Wait for processes to terminate + for process in processes: + process.join(timeout=10) # Wait for 10 seconds for the process to terminate + if process.is_alive(): + logger.warning(f"Process {process} is still alive after 10 seconds, attempting to terminate...") + process.terminate() + process.join() # ensure process termination after calling terminate() + + app = FastAPI(lifespan=lifespan) + + # Define the endpoints for the model server + @app.get("/health/") + async def health(): + """ + Health check endpoint to verify that the server is running. + """ + return {"status": "ok"} + + @app.get("/get_world_size/") + async def get_world_size(): + """ + Retrieves the world size of the LLM engine, which is `tensor_parallel_size * data_parallel_size`. + + Returns: + `dict`: + A dictionary containing the world size. + + Example response: + ```json + {"world_size": 8} + ``` + """ + return {"world_size": script_args.tensor_parallel_size * script_args.data_parallel_size} + + class GenerateRequest(BaseModel): + prompts: list[str] + images: list[str] | None = None + n: int = 1 + repetition_penalty: float = 1.0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + max_tokens: int = 16 + truncate_prompt_tokens: int | None = None + guided_decoding_regex: str | None = None + generation_kwargs: dict = field(default_factory=dict) + + class GenerateResponse(BaseModel): + prompt_ids: list[list[int]] + completion_ids: list[list[int]] + logprobs: list[list[float]] + + @app.post("/generate/", response_model=GenerateResponse) + async def generate(request: GenerateRequest): + """ + Generates completions for the provided prompts. + + Args: + request (`GenerateRequest`): + - `prompts` (list of `str`): A list of prompts (text strings) for the model to generate completions. + - `images` (list of `str`, *optional*, default to `None`): A list of base64 encoded images to process + along with prompts. + - `n` (`int`, *optional*, defaults to `1`): Number of completions to generate for each prompt. + - `repetition_penalty` (`float`, *optional*, defaults to `1.0`): Repetition penalty to apply during + generation. + - `temperature` (`float`, *optional*, defaults to `1.0`): Temperature for sampling. Higher values lead + to more random outputs. + - `top_p` (`float`, *optional*, defaults to `1.0`): Top-p (nucleus) sampling parameter. It controls the + diversity of the generated text. + - `top_k` (`int`, *optional*, defaults to `-1`): Top-k sampling parameter. If set to `-1`, it disables + top-k sampling. + - `min_p` (`float`, *optional*, defaults to `0.0`): Minimum probability threshold for sampling. + - `max_tokens` (`int`, *optional*, defaults to `16`): Maximum number of tokens to generate for each + completion. + - `truncate_prompt_tokens` (`int`, *optional*): If set to `-1`, will use the truncation size supported + by the model. If set to an integer k, will use only the last k tokens from the prompt (i.e., left + truncation). If set to `None`, truncation is disabled. + - `guided_decoding_regex` (`str`, *optional*): A regex pattern for guided decoding. If provided, the + model will only generate tokens that match this regex pattern. + - `generation_kwargs` (`dict`, *optional*): Additional generation parameters to pass to the vLLM + `SamplingParams`. This can include parameters like `seed`, `frequency_penalty`, etc. If it contains + keys that conflict with the other parameters, they will override them. + + Returns: + `GenerateResponse`: + - `prompt_ids` (list of list of `int`): A list of lists of token IDs for each input prompt. + - `completion_ids` (list of list of `int`): A list of lists of token IDs for each generated completion. + - `logprobs` (list of list of `float`): A list of lists of log probabilities for each token in the + generated completions. + + Example request: + ```json + {"prompts": ["Hello world", "What is AI?"]} + ``` + + Example response: + ```json + { + "prompt_ids": [[101, 102], [201, 202]], + "completion_ids": [[103, 104, 105], [203, 204, 205]], + "logprobs": [[-0.1, -0.2, -0.3], [-0.4, -0.5, -0.6]] + } + ``` + """ + request.images = request.images or [None] * len(request.prompts) + + prompts = [] + for prompt, image in zip(request.prompts, request.images, strict=True): + row = {"prompt": prompt} + if image is not None: + row["multi_modal_data"] = {"image": Image.open(BytesIO(base64.b64decode(image)))} + prompts.append(row) + + # Guided decoding, if enabled + if request.guided_decoding_regex is not None: + guided_decoding = GuidedDecodingParams(regex=request.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": request.n, + "repetition_penalty": request.repetition_penalty, + "temperature": request.temperature, + "top_p": request.top_p, + "top_k": request.top_k, + "min_p": request.min_p, + "max_tokens": request.max_tokens, + "truncate_prompt_tokens": request.truncate_prompt_tokens, + "guided_decoding": guided_decoding, + "logprobs": 0, # enable returning log probabilities; 0 means for the sampled tokens only + } + generation_kwargs.update(request.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + # Evenly distribute prompts across DP ranks + chunked_prompts = chunk_list(prompts, script_args.data_parallel_size) + + # Send the prompts to each worker + for connection, prompts in zip(connections, chunked_prompts, strict=True): + # When the number of prompts is less than data_parallel_size, some workers will receive empty prompts. + # However, vLLM requires that we always send at least one prompt. So we send a placeholder prompt to comply + # with vLLM's requirement, and we later ignore the result. + if not prompts: + prompts = [""] + kwargs = {"prompts": prompts, "sampling_params": sampling_params} + connection.send({"type": "call", "method": "generate", "kwargs": kwargs}) + + # Receive results + all_outputs = [connection.recv() for connection in connections] + + # Handle empty prompts (see above) + all_outputs = [output for output, prompts in zip(all_outputs, chunked_prompts, strict=True) if prompts] + + # Flatten and combine all results + all_outputs = list(chain.from_iterable(all_outputs)) # from list of list to single list + prompt_ids = [output.prompt_token_ids for output in all_outputs] + completion_ids = [list(output.token_ids) for outputs in all_outputs for output in outputs.outputs] + logprobs: list[list[float]] = [ + [sanitize_logprob(next(iter(logprob.values()))) for logprob in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] + return {"prompt_ids": prompt_ids, "completion_ids": completion_ids, "logprobs": logprobs} + + class ChatRequest(BaseModel): + messages: list[list[dict]] + n: int = 1 + repetition_penalty: float = 1.0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + max_tokens: int = 16 + truncate_prompt_tokens: int | None = None + guided_decoding_regex: str | None = None + generation_kwargs: dict = field(default_factory=dict) + chat_template_kwargs: dict = field(default_factory=dict) + + class ChatResponse(BaseModel): + prompt_ids: list[list[int]] + completion_ids: list[list[int]] + logprobs: list[list[float]] + + @app.post("/chat/", response_model=ChatResponse) + async def chat(request: ChatRequest): + """ + Generates completions for the provided chat messages. + + Args: + request (`ChatRequest`): + - `messages` (list of `dict`): A list of messages (dicts with "role" and "content" keys) for the model + to generate completions. + - `n` (`int`, *optional*, defaults to `1`): Number of completions to generate for each prompt. + - `repetition_penalty` (`float`, *optional*, defaults to `1.0`): Repetition penalty to apply during + generation. + - `temperature` (`float`, *optional*, defaults to `1.0`): Temperature for sampling. Higher values lead + to more random outputs. + - `top_p` (`float`, *optional*, defaults to `1.0`): Top-p (nucleus) sampling parameter. It controls the + diversity of the generated text. + - `top_k` (`int`, *optional*, defaults to `-1`): Top-k sampling parameter. If set to `-1`, it disables + top-k sampling. + - `min_p` (`float`, *optional*, defaults to `0.0`): Minimum probability threshold for sampling. + - `max_tokens` (`int`, *optional*, defaults to `16`): Maximum number of tokens to generate for each + completion. + - `truncate_prompt_tokens` (`int`, *optional*): If set to `-1`, will use the truncation size supported + by the model. If set to an integer k, will use only the last k tokens from the prompt (i.e., left + truncation). If set to `None`, truncation is disabled. + - `guided_decoding_regex` (`str`, *optional*): A regex pattern for guided decoding. If provided, the + model will only generate tokens that match this regex pattern. + - `generation_kwargs` (`dict`, *optional*): Additional generation parameters to pass to the vLLM + `SamplingParams`. This can include parameters like `seed`, `frequency_penalty`, etc. If it contains + keys that conflict with the other parameters, they will override them. + - `chat_template_kwargs` (`dict`, *optional*): Additional keyword arguments to pass to the chat + template. + + Returns: + `ChatResponse`: + - `prompt_ids` (list of list of `int`): A list of lists of token IDs for each input prompt. + - `completion_ids` (list of list of `int`): A list of lists of token IDs for each generated completion. + - `logprobs` (list of list of `float`): A list of lists of log probabilities for each token in the + generated completions. + + Example request: + ```bash + curl -X POST 'http://0.0.0.0:8000/chat/' \ + -H 'Content-Type: application/json' \ + -d '{"messages": [[{ "role": "user", "content": "Hello!" }]]}' + ``` + + Example response: + ```json + { + "prompt_ids": [[151644, 872, 198, 9707, 0, 151645, 198, 151644, 77091, 198]], + "completion_ids":[[151667, 198, 32313, 11, 279, 1196, 1101, 1053, 330, 9707, 8958, 773, 358, 1184, 311, 5889]], + "logprobs": [[-0.00029404606902971864, -3.576278118089249e-07, -0.09024181962013245, -6.389413465512916e-05, -0.038671817630529404, -0.00013314791431184858, -0.5868351459503174, -0.09682723134756088, -0.06609706580638885, -0.00023803261865396053, -0.02242819033563137, -0.8185162544250488, -0.04954879730939865, -0.3169460594654083, -4.887569048150908e-06, -0.006023705471307039]] + } + ``` + """ + # Convert PIL images to base64 strings + for message_list in request.messages: + for message in message_list: + if isinstance(message["content"], list): + for part in message["content"]: + if part["type"] == "image_pil": + part["image_pil"] = Image.open(BytesIO(base64.b64decode(part["image_pil"]))) + + # Guided decoding, if enabled + if request.guided_decoding_regex is not None: + guided_decoding = GuidedDecodingParams(regex=request.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": request.n, + "repetition_penalty": request.repetition_penalty, + "temperature": request.temperature, + "top_p": request.top_p, + "top_k": request.top_k, + "min_p": request.min_p, + "max_tokens": request.max_tokens, + "truncate_prompt_tokens": request.truncate_prompt_tokens, + "guided_decoding": guided_decoding, + "logprobs": 0, # enable returning log probabilities; 0 means for the sampled tokens only + } + generation_kwargs.update(request.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + # Evenly distribute prompts across DP ranks + chunked_messages = chunk_list(request.messages, script_args.data_parallel_size) + + # Send the messages to each worker + for connection, messages in zip(connections, chunked_messages, strict=True): + # When the number of messages is less than data_parallel_size, some workers will receive empty messages. + # However, vLLM requires that we always send at least one prompt. So we send a placeholder prompt to comply + # with vLLM's requirement, and we later ignore the result. + if not messages: + messages = [[{"role": "user", "content": ""}]] + kwargs = { + "messages": messages, + "sampling_params": sampling_params, + "chat_template_kwargs": request.chat_template_kwargs, + } + connection.send({"type": "call", "method": "chat", "kwargs": kwargs}) + + # Receive results + all_outputs = [connection.recv() for connection in connections] + + # Handle empty prompts (see above) + all_outputs = [output for output, prompts in zip(all_outputs, chunked_messages, strict=True) if prompts] + + # Flatten and combine all results + all_outputs = list(chain.from_iterable(all_outputs)) # from list of list to single list + prompt_ids = [output.prompt_token_ids for output in all_outputs] + completion_ids = [list(output.token_ids) for outputs in all_outputs for output in outputs.outputs] + logprobs: list[list[float]] = [ + [sanitize_logprob(next(iter(logprob.values()))) for logprob in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] + return {"prompt_ids": prompt_ids, "completion_ids": completion_ids, "logprobs": logprobs} + + class InitCommunicatorRequest(BaseModel): + host: str + port: int + world_size: int + client_device_uuid: str + + @app.post("/init_communicator/") + async def init_communicator(request: InitCommunicatorRequest): + """ + Initializes the communicator for synchronizing model weights between a client and multiple server workers. + + Args: + request (`InitCommunicatorRequest`): + - `host` (`str`): Hostname or IP address of the master node. + - `port` (`int`): Port number to be used for communication. + - `world_size` (`int`): Total number of participating processes in the group. + - `client_device_uuid` (`str`): UUID of the device of client main process. Used to assert that devices + are different from vLLM workers devices. + """ + world_size = script_args.tensor_parallel_size * script_args.data_parallel_size + 1 + + # The function init_communicator is called this way: init_communicator(host, port, world_size) + # So with collective_rpc we need to call it this way: + # llm.collective_rpc(method="init_communicator", args=(host, port, world_size)) + kwargs = { + "method": "init_communicator", + "args": (request.host, request.port, world_size, request.client_device_uuid), + } + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + + return {"message": "Request received, initializing communicator"} + + class UpdateWeightsRequest(BaseModel): + name: str + dtype: str + shape: list[int] + + @app.post("/update_named_param/") + async def update_named_param(request: UpdateWeightsRequest): + """ + Updates the model weights with the provided tensor. + + Once this endpoint is called, the client process should broadcast the updated weights to all server workers. + + Args: + request (`UpdateWeightsRequest`): + - `name` (`str`): Name of the weight tensor being updated. + - `dtype` (`str`): Data type of the weight tensor (e.g., `"torch.float32"`). + - `shape` (list of `int`): Shape of the weight + + """ + # The function update_named_param is called this way: update_named_param("name", "torch.float32", (10, 10)) + # So with collective_rpc we need to call it this way: + # llm.collective_rpc("update_named_param", args=("name", "torch.float32", (10, 10))) + kwargs = {"method": "update_named_param", "args": (request.name, request.dtype, tuple(request.shape))} + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + + return {"message": "Request received, updating named parameter"} + + @app.post("/reset_prefix_cache/") + async def reset_prefix_cache(): + """ + Resets the prefix cache for the model. + """ + for connection in connections: + connection.send({"type": "call", "method": "reset_prefix_cache"}) + # Wait for and collect all results + all_outputs = [connection.recv() for connection in connections] + success = all(output for output in all_outputs) + return {"message": "Request received, resetting prefix cache status: " + str(success)} + + @app.post("/close_communicator/") + async def close_communicator(): + """ + Closes the weight update group and cleans up associated resources. + """ + kwargs = {"method": "close_communicator"} + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + return {"message": "Request received, closing communicator"} + + # Start the server + uvicorn.run(app, host=script_args.host, port=script_args.port, log_level=script_args.log_level) + + +def make_parser(subparsers: argparse._SubParsersAction | None = None): + if subparsers is not None: + parser = subparsers.add_parser("vllm-serve", help="Run the vLLM serve script", dataclass_types=ScriptArguments) + else: + parser = TrlParser(ScriptArguments) + return parser + + +if __name__ == "__main__": + parser = make_parser() + (script_args,) = parser.parse_args_and_config() + main(script_args) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__init__.py b/venv/lib/python3.10/site-packages/trl/trainer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98846bf71597fcc08ac322e6b8e5415029408ceb --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/__init__.py @@ -0,0 +1,141 @@ +# Copyright 2020-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 typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "bco_config": ["BCOConfig"], + "bco_trainer": ["BCOTrainer"], + "callbacks": [ + "BEMACallback", + "LogCompletionsCallback", + "MergeModelCallback", + "RichProgressCallback", + "SyncRefModelCallback", + "WeaveCallback", + "WinRateCallback", + ], + "cpo_config": ["CPOConfig"], + "cpo_trainer": ["CPOTrainer"], + "dpo_config": ["DPOConfig", "FDivergenceConstants", "FDivergenceType"], + "dpo_trainer": ["DPOTrainer"], + "gkd_config": ["GKDConfig"], + "gkd_trainer": ["GKDTrainer"], + "grpo_config": ["GRPOConfig"], + "grpo_trainer": ["GRPOTrainer"], + "judges": [ + "AllTrueJudge", + "BaseBinaryJudge", + "BaseJudge", + "BasePairwiseJudge", + "BaseRankJudge", + "HfPairwiseJudge", + "OpenAIPairwiseJudge", + "PairRMJudge", + ], + "kto_config": ["KTOConfig"], + "kto_trainer": ["KTOTrainer"], + "model_config": ["ModelConfig"], + "nash_md_config": ["NashMDConfig"], + "nash_md_trainer": ["NashMDTrainer"], + "online_dpo_config": ["OnlineDPOConfig"], + "online_dpo_trainer": ["OnlineDPOTrainer"], + "orpo_config": ["ORPOConfig"], + "orpo_trainer": ["ORPOTrainer"], + "ppo_config": ["PPOConfig"], + "ppo_trainer": ["PPOTrainer"], + "prm_config": ["PRMConfig"], + "prm_trainer": ["PRMTrainer"], + "reward_config": ["RewardConfig"], + "reward_trainer": ["RewardTrainer"], + "rloo_config": ["RLOOConfig"], + "rloo_trainer": ["RLOOTrainer"], + "sft_config": ["SFTConfig"], + "sft_trainer": ["SFTTrainer"], + "utils": [ + "RunningMoments", + "compute_accuracy", + "disable_dropout_in_model", + "empty_cache", + "peft_module_casting_to_bf16", + ], + "xpo_config": ["XPOConfig"], + "xpo_trainer": ["XPOTrainer"], +} + +if TYPE_CHECKING: + from .bco_config import BCOConfig + from .bco_trainer import BCOTrainer + from .callbacks import ( + BEMACallback, + LogCompletionsCallback, + MergeModelCallback, + RichProgressCallback, + SyncRefModelCallback, + WeaveCallback, + WinRateCallback, + ) + from .cpo_config import CPOConfig + from .cpo_trainer import CPOTrainer + from .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType + from .dpo_trainer import DPOTrainer + from .gkd_config import GKDConfig + from .gkd_trainer import GKDTrainer + from .grpo_config import GRPOConfig + from .grpo_trainer import GRPOTrainer + from .judges import ( + AllTrueJudge, + BaseBinaryJudge, + BaseJudge, + BasePairwiseJudge, + BaseRankJudge, + HfPairwiseJudge, + OpenAIPairwiseJudge, + PairRMJudge, + ) + from .kto_config import KTOConfig + from .kto_trainer import KTOTrainer + from .model_config import ModelConfig + from .nash_md_config import NashMDConfig + from .nash_md_trainer import NashMDTrainer + from .online_dpo_config import OnlineDPOConfig + from .online_dpo_trainer import OnlineDPOTrainer + from .orpo_config import ORPOConfig + from .orpo_trainer import ORPOTrainer + from .ppo_config import PPOConfig + from .ppo_trainer import PPOTrainer + from .prm_config import PRMConfig + from .prm_trainer import PRMTrainer + from .reward_config import RewardConfig + from .reward_trainer import RewardTrainer + from .rloo_config import RLOOConfig + from .rloo_trainer import RLOOTrainer + from .sft_config import SFTConfig + from .sft_trainer import SFTTrainer + from .utils import ( + RunningMoments, + compute_accuracy, + disable_dropout_in_model, + empty_cache, + peft_module_casting_to_bf16, + ) + from .xpo_config import XPOConfig + from .xpo_trainer import XPOTrainer +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/__init__.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6352df173537a00b22fea6c2627fddb43b7cfc0e Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/__init__.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/base_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/base_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6db235e7ec77ad78c0c5af27036ced626ba8939 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/base_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0444c3f03761d67f045edbd4828e57813c0c4590 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bd5042c258a8b87ed3e56014262040c53bd55ca Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/callbacks.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/callbacks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d096fea449a86bac633cfd52559a1191cf5b7db9 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/callbacks.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6efc911854b12c6197afc48059d7407cd5532613 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2d38898769b3cefc0f3856b74d94c8b0772d8c1 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b291c3fb574fb4166ee36e3be079d75d8abc4d2b Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7176e67dd67b12aa9354c4f87af641b0bf573f55 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdcbb0b5e58e4f8d6f675c5b6d113cbe41a99487 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0c510299e1dcc7760a1ce90bcde8cbf4c3675d3 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20439a86345a3da1cc45a8129222da081fc12c11 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..348728284d59ec83165e2f26e2c2071614a734f8 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/judges.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/judges.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..879452f4e1ddecd6c9a38dc98ae98d509b4c0d91 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/judges.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..197ee1ccd1fd1d870e7e0de24fd8f72f9eaf816a Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a139945dd93f3743d934571dc943f6e04f65132 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/model_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/model_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b430b463c7710d5b8829d14faff5eeb25a89538 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/model_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e67004c9c3ebfdcd20418e8d90f4bdd95af2d62 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db481ec86295e483c54513c9eeee4df370d64c11 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fc9d01b8a83042cfe17bf6966063bc21ddcc090 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f613d27ceffef69aa6ce82bedd398ea39664c36 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b980a8050d0b3e2c98bc00e6eb24f2622ea8338 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9172f0938474f380d9271506c86b6405af4dc65c Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc649f47943716a2422ad19dab3aa994c0bb68bc Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b021475376638bb6cab176a9e8e69db598ad46dc Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cfc20e0f232d67ce6772da561f5581fdd2fb78f Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4dbb48c162d01427de0d2801edf3d54b25313d5 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cf4541a57cb3935e16bf8a6a9efe01e253f82a2 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..000f4f018e3e1f6fa2878ab6adb5907d37cb0e5b Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8560e7b7744d520a53bdc363b8ff350d13488c16 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dc85f50543efec1dbae741bce283bf8e2f60b06 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbb12f7dce377ccdc7db68e59875368e9b3d3061 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d57727118a00817c17c183d3ed191b5e463f2f12 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/utils.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..832230bffca39b554707a8c251e1c872916f149a Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/utils.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_config.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82b3ece2df2de0e23e8d6e9fdaa1f8d38efbd991 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_config.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-310.pyc b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca0c08e7efdb75121433f902b9fe87ae0369d849 Binary files /dev/null and b/venv/lib/python3.10/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-310.pyc differ diff --git a/venv/lib/python3.10/site-packages/trl/trainer/base_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/base_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..9763fcd43f89554be7630075ea83ffc1e4876344 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/base_trainer.py @@ -0,0 +1,86 @@ +# Copyright 2020-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 os + +from transformers import Trainer, is_wandb_available + +from .utils import generate_model_card, get_comet_experiment_url, get_config_model_id + + +if is_wandb_available(): + import wandb + + +class BaseTrainer(Trainer): + _tag_names = [] + _name = "Base" + _paper = {} + _template_file = None + + def create_model_card( + self, + model_name: str | None = None, + dataset_name: str | None = None, + tags: str | list[str] | None = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str`, *optional*): + Name of the model. + dataset_name (`str`, *optional*): + Name of the dataset used for training. + tags (`str`, `list[str]`, *optional*): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + model_name_or_path = get_config_model_id(self.model.config) + if model_name_or_path and not os.path.isdir(model_name_or_path): + base_model = model_name_or_path + else: + base_model = None + + # Normalize tags + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + if "JOB_ID" in os.environ: + tags.add("hf_jobs") + tags.update(self._tag_names) + tags = list(tags) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.url if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name=self._name, + trainer_citation=self._paper.get("citation"), + template_file=self._template_file, + paper_title=self._paper.get("title"), + paper_id=self._paper.get("id"), + ) + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/bco_config.py b/venv/lib/python3.10/site-packages/trl/trainer/bco_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6e22dda6b1351186cd15e2c6a089550b98c6b5bc --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/bco_config.py @@ -0,0 +1,29 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass + +from ..experimental.bco import BCOConfig as _BCOConfig + + +@dataclass +class BCOConfig(_BCOConfig): + def __post_init__(self): + warnings.warn( + "The `BCOConfig` is now located in `trl.experimental`. Please update your imports to " + "`from trl.experimental.bco import BCOConfig`. The current import path will be removed and no longer " + "supported in TRL 0.29. For more information, see https://github.com/huggingface/trl/issues/4223." + ) + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/bco_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/bco_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..c98fc3202006c1a7a593617b40bf840752f07db5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/bco_trainer.py @@ -0,0 +1,29 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass + +from ..experimental.bco import BCOTrainer as _BCOTrainer + + +@dataclass +class BCOTrainer(_BCOTrainer): + def __init__(self, *args, **kwargs): + warnings.warn( + "The `BCOTrainer` is now located in `trl.experimental`. Please update your imports to " + "`from trl.experimental.bco import BCOTrainer`. The current import path will be removed and no longer " + "supported in TRL 0.29. For more information, see https://github.com/huggingface/trl/issues/4223." + ) + super().__init__(*args, **kwargs) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/callbacks.py b/venv/lib/python3.10/site-packages/trl/trainer/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..252f38ef31e6c130ee04d78a113ff305e459186d --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/callbacks.py @@ -0,0 +1,1023 @@ +# Copyright 2020-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 logging +import os + +import pandas as pd +import torch +from accelerate import Accelerator +from accelerate.state import AcceleratorState +from accelerate.utils import gather_object, is_wandb_available +from transformers import ( + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + Trainer, + TrainerCallback, + TrainerControl, + TrainerState, + TrainingArguments, +) +from transformers.trainer_utils import has_length +from transformers.utils import is_rich_available + +from ..data_utils import maybe_apply_chat_template +from ..import_utils import is_mergekit_available, is_weave_available +from ..mergekit_utils import MergeConfig, merge_models, upload_model_to_hf +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .utils import get_config_model_id, log_table_to_comet_experiment + + +if is_rich_available(): + from rich.columns import Columns + from rich.console import Console, Group + from rich.live import Live + from rich.panel import Panel + from rich.progress import Progress + from rich.table import Table + +if is_wandb_available(): + import wandb + +if is_weave_available(): + import weave + from weave import EvaluationLogger + from weave.trace.context import weave_client_context + + +# Logger for module-level logging +logger = logging.getLogger(__name__) + + +def _generate_completions( + prompts: list[str], + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + accelerator: Accelerator, + generation_config: GenerationConfig | None, + batch_size: int = 1, +) -> list[str]: + """ + Generates completions for a list of pre-formatted prompts from the given model. + + Args: + prompts (list[str]): A list of input prompts for which completions are to be generated. + model (PreTrainedModel): The pre-trained model to be used for generation. + tokenizer (PreTrainedTokenizerBase): The tokenizer to be used for encoding and decoding. + accelerator (Accelerator): The accelerator to be used for model execution. + generation_config (GenerationConfig): Configuration for text generation. + batch_size (int, optional): The number of prompts to process in each batch. Default is 1. + + Returns: + list[str]: A list of generated text completions corresponding to the input prompts. + """ + completions = [] + with unwrap_model_for_generation(model, accelerator) as unwrapped_model: + for idx in range(0, len(prompts), batch_size): + batch = prompts[idx : idx + batch_size] + tokenized_batch = tokenizer(batch, return_tensors="pt", padding=True, truncation=True).to(model.device) + generations = unwrapped_model.generate( + **tokenized_batch, + generation_config=generation_config, + ) + for prompt, generation in zip(tokenized_batch.input_ids, generations, strict=True): + # Remove prompt from generation + generation = generation[len(prompt) :] + completion = tokenizer.decode(generation, skip_special_tokens=True) + completions.append(completion) + return completions + + +class SyncRefModelCallback(TrainerCallback): + """ + Callback to synchronize the model with a reference model. + """ + + def __init__( + self, + ref_model: PreTrainedModel | torch.nn.Module, + accelerator: Accelerator | None, + ): + self.accelerator = accelerator + self.ref_model = ref_model + + @staticmethod + def _sync_target_model(model, target_model, alpha): + for target_param, copy_param in zip(target_model.parameters(), model.parameters(), strict=True): + target_param.data.mul_(1.0 - alpha).add_(copy_param.data, alpha=alpha) + + @staticmethod + def sync_target_model(model, target_model, alpha): + deepspeed_plugin = AcceleratorState().deepspeed_plugin + if deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3: + import deepspeed + + with deepspeed.zero.GatheredParameters( + list(model.parameters()) + list(target_model.parameters()), modifier_rank=0 + ): + if deepspeed.comm.get_rank() == 0: + SyncRefModelCallback._sync_target_model(model, target_model, alpha) + else: + SyncRefModelCallback._sync_target_model(model, target_model, alpha) + + def on_step_end(self, args, state, control, **kwargs): + model: PreTrainedModel = kwargs["model"] + + if self.ref_model is not None and state.global_step % args.ref_model_sync_steps == 0: + if self.accelerator: + model = self.accelerator.unwrap_model(model) + self.sync_target_model(model, self.ref_model, args.ref_model_mixup_alpha) + + +class RichProgressCallback(TrainerCallback): + """ + A [`TrainerCallback`] that displays the progress of training or evaluation using Rich. + """ + + def __init__(self): + if not is_rich_available(): + raise ImportError("RichProgressCallback requires the `rich` extra. To install, run `pip install rich`.") + + self.training_bar = None + self.evaluation_bar = None + self.training_task = None + self.evaluation_task = None + self.rich_group = None + self.rich_console = None + self.training_status = None + self.current_step = None + + def on_train_begin(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + self.training_bar = Progress() + self.evaluation_bar = Progress() + self.rich_console = Console() + self.training_status = self.rich_console.status("Nothing to log yet ...") + self.rich_group = Live(Panel(Group(self.training_bar, self.evaluation_bar, self.training_status))) + self.rich_group.start() + self.training_task = self.training_bar.add_task("[blue]Training ", total=state.max_steps) + self.current_step = 0 + + def on_step_end(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + self.training_bar.update(self.training_task, advance=state.global_step - self.current_step, update=True) + self.current_step = state.global_step + + def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs): + if not state.is_world_process_zero: + return + + if has_length(eval_dataloader): + if self.evaluation_task is None: + self.evaluation_task = self.evaluation_bar.add_task("[blue]Evaluation", total=len(eval_dataloader)) + self.evaluation_bar.update(self.evaluation_task, advance=1, update=True) + + def on_evaluate(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + if self.evaluation_task is not None: + self.evaluation_bar.remove_task(self.evaluation_task) + self.evaluation_task = None + + def on_predict(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + if self.evaluation_task is not None: + self.evaluation_bar.remove_task(self.evaluation_task) + self.evaluation_task = None + + def on_log(self, args, state, control, logs=None, **kwargs): + if not (state.is_world_process_zero and self.training_bar): + return + + # Group keys by top-level prefix + grouped_logs = {} + for key, value in logs.items(): + parts = key.split("/") + group = parts[0] if len(parts) > 1 else None + subkey = "/".join(parts[1:]) if len(parts) > 1 else key + grouped_logs.setdefault(group, {})[subkey] = value + + # Create a table per group + tables = [] + for group_name, metrics in grouped_logs.items(): + table = Table( + title=f"[bold blue]{group_name}[/]" if group_name else None, header_style="bold magenta", box=None + ) + table.add_column("Metric", justify="left", no_wrap=True) + table.add_column("Value", justify="right") + + for metric, val in metrics.items(): + formatted = f"{val:.3f}" if isinstance(val, (float, int)) else str(val) + table.add_row(metric, formatted) + + tables.append(Panel(table, border_style="cyan", padding=(0, 1))) + + # Arrange tables in columns using Columns + column_layout = Columns(tables, equal=False, expand=True) + self.training_status.update( + Panel(column_layout, title=f"[bold green]Step {state.global_step}[/bold green]", border_style="green") + ) + + def on_train_end(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + self.rich_group.stop() + self.training_bar = None + self.evaluation_bar = None + self.training_task = None + self.evaluation_task = None + self.rich_group = None + self.rich_console = None + self.training_status = None + self.current_step = None + + +def _win_rate_completions_df( + state: TrainerState, prompts: list[str], completions: list[str], winner_indices: list[str] +) -> pd.DataFrame: + global_step = [str(state.global_step)] * len(prompts) + data = list(zip(global_step, prompts, completions, winner_indices, strict=True)) + # Split completions from reference model and policy + split_data = [(item[0], item[1], item[2][0], item[2][1], item[3]) for item in data] + return pd.DataFrame(split_data, columns=["step", "prompt", "reference_model", "policy", "winner_index"]) + + +class WinRateCallback(TrainerCallback): + """ + A [`~transformers.TrainerCallback`] that computes the win rate of a model based on a reference. + + It generates completions using prompts from the evaluation dataset and compares the trained model's outputs against + a reference. The reference is either the initial version of the model (before training) or the reference model, if + available in the trainer. During each evaluation step, a judge determines how often the trained model's completions + win against the reference using a judge. The win rate is then logged in the trainer's logs under the key + `"eval_win_rate"`. + + Usage: + ```python + trainer = DPOTrainer(...) + judge = PairRMJudge() + win_rate_callback = WinRateCallback(judge=judge, trainer=trainer) + trainer.add_callback(win_rate_callback) + ``` + + Args: + judge ([`BasePairwiseJudge`]): + The judge to use for comparing completions. + trainer (`Trainer`): + Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `"prompt"` + column containing the prompts for generating completions. If the `Trainer` has a reference model (via the + `ref_model` attribute), it will use this reference model for generating the reference completions; + otherwise, it defaults to using the initial model. + generation_config ([`~transformers.GenerationConfig`], *optional*): + The generation config to use for generating completions. + num_prompts (`int`, *optional*): + The number of prompts to generate completions for. If not provided, defaults to the number of examples in + the evaluation dataset. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions before judging. + use_soft_judge (`bool`, *optional*, defaults to `False`): + Whether to use a soft judge that returns a win probability between 0 and 1 for the first completion vs the + second. + """ + + def __init__( + self, + judge: BasePairwiseJudge, + trainer: Trainer, + generation_config: GenerationConfig | None = None, + num_prompts: int | None = None, + shuffle_order: bool = True, + use_soft_judge: bool = False, + ): + self.judge = judge + self.trainer = trainer + self.shuffle_order = shuffle_order + self.generation_config = generation_config + self.ref_completions = [] + self.use_soft_judge = use_soft_judge + + if self.trainer.eval_dataset is None: + raise ValueError("Trainer must have an evaluation dataset to use the WinRateCallback.") + else: + self.eval_dataset = self.trainer.eval_dataset + + if num_prompts is not None: + self.eval_dataset = self.eval_dataset.select(range(num_prompts)) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # When the trainer is initialized, we generate completions for the reference model. + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + # Use the reference model if available, otherwise use the initial model + model = getattr(self.trainer, "ref_model", None) + # At this point, there are two cases where `ref_model` is None: + # 1. The method doesn't require a reference model. + # 2. The method uses a reference model, but `ref_model` is set to None. + # This occurs when using PEFT, where the reference model can be obtained by simply disabling the model's adapter. + # In theory, we should disable the adapter here, but since it's zero-initialized at the start of training, + # the model behaves identically with or without the adapter. + # Therefore, there's no need to explicitly disable it at this point. + if model is None: + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + self.ref_completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + # Compute initial win rate as a reference point + completions = list(zip(self.ref_completions, self.ref_completions, strict=True)) + if self.use_soft_judge: + ref_win_probs = self.judge.judge(prompts, completions, self.shuffle_order, return_scores=True) + winner_indices = [0 if score > 0.5 else 1 for score in ref_win_probs] + ref_win_probs = gather_object(ref_win_probs) + else: + winner_indices = self.judge.judge(prompts, completions, self.shuffle_order) + prompts = gather_object(prompts) + completions = gather_object(completions) + winner_indices = gather_object(winner_indices) + + # Logging + if self.trainer.accelerator.is_main_process: + win_rate = sum(winner_idx == 1 for winner_idx in winner_indices) / len(winner_indices) + if self.use_soft_judge: + avg_win_prob = 1.0 - sum(ref_win_probs) / len(ref_win_probs) + self.trainer.log({"eval_avg_win_prob": avg_win_prob, "eval_win_rate": win_rate}) + else: + self.trainer.log({"eval_win_rate": win_rate}) + + if "wandb" in args.report_to: + if wandb.run is not None: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + wandb.log({"win_rate_completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + log_table_to_comet_experiment( + name="win_rate_completions.csv", + table=df, + ) + + def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # At every evaluation step, we generate completions for the model and compare them with the reference + # completions that have been generated at the beginning of training. We then compute the win rate and log it to + # the trainer. + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + + completions = list(zip(self.ref_completions, completions, strict=True)) + + if self.use_soft_judge: + ref_win_probs = self.judge.judge(prompts, completions, self.shuffle_order, return_scores=True) + winner_indices = [0 if score > 0.5 else 1 for score in ref_win_probs] + ref_win_probs = gather_object(ref_win_probs) + else: + winner_indices = self.judge.judge(prompts, completions, self.shuffle_order) + prompts = gather_object(prompts) + completions = gather_object(completions) + winner_indices = gather_object(winner_indices) + + # Logging + if self.trainer.accelerator.is_main_process: + win_rate = sum(winner_idx == 1 for winner_idx in winner_indices) / len(winner_indices) + if self.use_soft_judge: + avg_win_prob = 1.0 - sum(ref_win_probs) / len(ref_win_probs) + self.trainer.log({"eval_avg_win_prob": avg_win_prob, "eval_win_rate": win_rate}) + else: + self.trainer.log({"eval_win_rate": win_rate}) + + if "wandb" in args.report_to: + if wandb.run is not None: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + wandb.log({"win_rate_completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + log_table_to_comet_experiment( + name="win_rate_completions.csv", + table=df, + ) + + +class LogCompletionsCallback(TrainerCallback): + r""" + A [`~transformers.TrainerCallback`] that logs completions to Weights & Biases and/or Comet. + + Usage: + ```python + trainer = DPOTrainer(...) + completions_callback = LogCompletionsCallback(trainer=trainer) + trainer.add_callback(completions_callback) + ``` + + Args: + trainer (`Trainer`): + Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `"prompt"` + column containing the prompts for generating completions. + generation_config ([`~transformers.GenerationConfig`], *optional*): + The generation config to use for generating completions. + num_prompts (`int`, *optional*): + The number of prompts to generate completions for. If not provided, defaults to the number of examples in + the evaluation dataset. + freq (`int`, *optional*): + The frequency at which to log completions. If not provided, defaults to the trainer's `eval_steps`. + """ + + def __init__( + self, + trainer: Trainer, + generation_config: GenerationConfig | None = None, + num_prompts: int | None = None, + freq: int | None = None, + ): + self.trainer = trainer + self.generation_config = generation_config + self.freq = freq + self.table = [] + self._last_logged_step = -1 + + if self.trainer.eval_dataset is None: + raise ValueError("Trainer must have an evaluation dataset to use the LogCompletionsCallback.") + else: + self.eval_dataset = self.trainer.eval_dataset + + if num_prompts is not None: + self.eval_dataset = self.eval_dataset.select(range(num_prompts)) + + def on_step_end(self, args, state, control, **kwargs): + # Only log once per step (this method may be called multiple times) + if state.global_step == self._last_logged_step: + return + + # Only log every `freq` steps (if no `freq` is provided, log every `eval_steps` steps) + freq = self.freq or state.eval_steps + if state.global_step % freq != 0: + return + + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + prompts = [maybe_apply_chat_template({"prompt": prompt}, tokenizer)["prompt"] for prompt in prompts] + completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + completions = gather_object(completions) + prompts = gather_object(prompts) + + # Build the data to log + if self.trainer.accelerator.is_main_process: + global_step = [str(state.global_step)] * len(prompts) + data = list(zip(global_step, prompts, completions, strict=True)) + self.table.extend(data) + table = pd.DataFrame(columns=["step", "prompt", "completion"], data=self.table) + + if "wandb" in args.report_to: + wandb.log({"completions": table}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=table, + ) + + # Save the last logged step, so we don't log the same completions multiple times + self._last_logged_step = state.global_step + + +class WeaveCallback(TrainerCallback): + r""" + A [`~transformers.TrainerCallback`] that logs traces and evaluations to W&B Weave. The callback uses + https://weave-docs.wandb.ai/guides/evaluation/evaluation_logger/ to log traces and evaluations at each evaluation + step. + + Supports two modes based on the `scorers` parameter: + - **Tracing Mode** (when scorers=None): Logs predictions for data exploration and analysis + - **Evaluation Mode** (when scorers provided): Logs predictions with scoring and summary metrics + + Both modes use Weave's EvaluationLogger for structured, consistent data logging. + + The callback logs data during evaluation phases (`on_evaluate`) rather than training steps, making it more + efficient and semantically correct. It gracefully handles missing weave installation by logging warnings and + skipping weave-specific functionality. It also checks for existing weave clients before initializing new ones. + + Usage: + ```python + # Tracing mode (just log predictions) + trainer = DPOTrainer(...) + weave_callback = WeaveTraceCallback(trainer=trainer) # project_name optional + trainer.add_callback(weave_callback) + + # Or specify a project name + weave_callback = WeaveTraceCallback(trainer=trainer, project_name="my-llm-training") + trainer.add_callback(weave_callback) + + + # Evaluation mode (log predictions + scores + summary) + def accuracy_scorer(prompt: str, completion: str) -> float: + # Your scoring logic here (metadata available via eval_attributes) + return score + + + weave_callback = WeaveTraceCallback( + trainer=trainer, + project_name="my-llm-training", # optional and needed only if weave client is not initialized + scorers={"accuracy": accuracy_scorer}, + ) + trainer.add_callback(weave_callback) + ``` + + Args: + trainer (`Trainer`): + Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `"prompt"` + column containing the prompts for generating completions. + project_name (`str`, *optional*): + Name of the Weave project where data will be logged. If not provided, will try to use existing weave client + or fall back to the active wandb run's project name. Raises an error if none of these are available. + scorers (`dict[str, Callable]`, *optional*): + Dictionary mapping scorer names to scorer functions. If `None`, operates in tracing mode (predictions + only). If provided, operates in evaluation mode (predictions + scores + summary). Scorer functions should + have signature: `scorer(prompt: str, completion: str) -> float | int` + generation_config ([`~transformers.GenerationConfig`], *optional*): + Generation config to use for generating completions. + num_prompts (`int` or `None`, *optional*): + Number of prompts to generate completions for. If not provided, defaults to the number of examples in the + evaluation dataset. + dataset_name (`str`, *optional*, defaults to `"eval_dataset"`): + Name for the dataset metadata in Weave. + model_name (`str`, *optional*): + Name for the model metadata in Weave. If not provided, attempts to extract from model config. + """ + + def __init__( + self, + trainer: Trainer, + project_name: str | None = None, + scorers: dict[str, callable] | None = None, + generation_config: GenerationConfig | None = None, + num_prompts: int | None = None, + dataset_name: str = "eval_dataset", + model_name: str | None = None, + ): + self.trainer = trainer + self.project_name = project_name + self.scorers = scorers or {} + self.generation_config = generation_config + self.dataset_name = dataset_name + self.model_name = model_name + self._last_logged_step = -1 + self._weave_initialized = False + self._eval_logger = None + + if self.trainer.eval_dataset is None: + raise ValueError("Trainer must have an evaluation dataset to use the WeaveCallback.") + else: + self.eval_dataset = self.trainer.eval_dataset + + if num_prompts is not None: + self.eval_dataset = self.eval_dataset.select(range(num_prompts)) + + def _initialize_weave(self): + """Initialize Weave and EvaluationLogger if not already initialized.""" + if not self._weave_initialized: + if not is_weave_available(): + logger.warning("Weave is not available. Please install weave to enable logging: `pip install weave`") + return + + if wc := weave_client_context.get_weave_client(): + self._weave_client = wc + else: + if self.project_name is None: + if is_wandb_available(): + if wandb.run is not None: + self.project_name = wandb.run.entity + "/" + wandb.run.project + logger.info(f"Using project name from active wandb run: {self.project_name}") + + if self.project_name is None: + raise ValueError( + "No existing Weave client found and no project_name provided. " + "Please either initialize weave with `weave.init('project-name')`, " + "provide a project_name to the `WeaveTraceCallback`, " + "or ensure an active wandb run exists." + ) + + self._weave_client = weave.init(self.project_name) + logger.info(f"Initialized Weave with project: {self.project_name}") + + if self.model_name is None: + self.model_name = getattr(self.trainer.model_wrapped.config, "_name_or_path", "unknown_model") + + self._EvaluationLogger = EvaluationLogger + + self._weave_initialized = True + + @property + def is_evaluation_mode(self) -> bool: + """True if scorers are provided (evaluation mode), False for tracing mode.""" + return bool(self.scorers) + + def on_train_begin(self, args, state, control, **kwargs): + """Initialize Weave when training begins.""" + self._initialize_weave() + + def on_evaluate(self, args, state, control, **kwargs): + if state.global_step == self._last_logged_step: + return + + self._initialize_weave() + + if not self._weave_initialized: + logger.debug("Weave not initialized, skipping logging") + return + + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + model = self.trainer.model_wrapped + + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + prompts = [maybe_apply_chat_template({"prompt": prompt}, tokenizer)["prompt"] for prompt in prompts] + + completions = _generate_completions( + prompts=prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + + all_prompts = gather_object(prompts) + all_completions = gather_object(completions) + + if self.trainer.accelerator.is_main_process: + eval_attributes = { + "training_step": state.global_step, + "model_name": self.model_name, + "generation_config": (self.generation_config.to_dict() if self.generation_config else None), + } + + eval_logger = self._EvaluationLogger( + model=self.model_name, + dataset=self.dataset_name, + eval_attributes=eval_attributes, + ) + + successful_predictions = 0 + total_score_values = {} # For summary statistics + + for prompt, completion in zip(all_prompts, all_completions, strict=True): + try: + pred_logger = eval_logger.log_prediction(inputs={"prompt": prompt}, output=completion) + + if self.is_evaluation_mode: + for scorer_name, scorer_func in self.scorers.items(): + try: + score = scorer_func(prompt, completion) + pred_logger.log_score(scorer=scorer_name, score=score) + + if scorer_name not in total_score_values: + total_score_values[scorer_name] = [] + total_score_values[scorer_name].append(score) + + except Exception as scorer_e: + logger.warning(f"Failed to apply scorer '{scorer_name}': {scorer_e}") + + pred_logger.finish() + successful_predictions += 1 + + except Exception as pred_e: + logger.warning(f"Failed to log prediction for prompt: {pred_e}") + # Continue with other predictions even if one fails + + if self.is_evaluation_mode and total_score_values: + try: + summary_stats = { + "total_predictions": len(all_prompts), + "successful_predictions": successful_predictions, + } + + for scorer_name, scores in total_score_values.items(): + if scores: # Only if we have valid scores + summary_stats[f"avg_{scorer_name}"] = sum(scores) / len(scores) + + eval_logger.log_summary(summary_stats) + + except Exception as summary_e: + logger.warning(f"Failed to log summary: {summary_e}") + else: + try: + eval_logger.finish() + except Exception as finish_e: + logger.warning(f"Failed to finish evaluation logger: {finish_e}") + + self._last_logged_step = state.global_step + + +class MergeModelCallback(TrainerCallback): + r""" + A [`~transformers.TrainerCallback`] that merges the policy model (the model being trained) with another model based + on a merge configuration. + + Args: + merge_config ([`MergeConfig`], *optional*): + Configuration used for the merging process. If not provided, the default [`MergeConfig`] is used. + merge_at_every_checkpoint (`bool`, *optional*, defaults to `False`): + Whether to merge the model at every checkpoint. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the merged model to the Hub after merging. + + Example: + + ```python + from trl.mergekit_utils import MergeConfig + from trl import MergeModelCallback + + config = MergeConfig() + merge_callback = MergeModelCallback(config) + trainer = DPOTrainer(..., callbacks=[merge_callback]) + ``` + """ + + def __init__( + self, + merge_config: "MergeConfig | None" = None, + merge_at_every_checkpoint: bool = False, + push_to_hub: bool = False, + ): + if not is_mergekit_available(): + raise ImportError( + "MergeModelCallback requires the `mergekit` extra. To install, run `pip install mergekit`." + ) + self.merge_config = merge_config or MergeConfig() + self.merge_at_every_checkpoint = merge_at_every_checkpoint + self.push_to_hub = push_to_hub + + def _merge_and_maybe_push(self, output_dir, global_step, model): + checkpoint_path = os.path.join(output_dir, f"checkpoint-{global_step}") + self.merge_config.policy_model_path = checkpoint_path + if self.merge_config.target_model_path is None: + self.merge_config.target_model_path = get_config_model_id(model.config) + merge_path = os.path.join(checkpoint_path, "merged") + + merge_models(self.merge_config.create(), merge_path) + + if self.push_to_hub: + repo_name = f"{output_dir}_checkpoint-{global_step}_merged" + upload_model_to_hf(merge_path, repo_name) + + def on_save(self, args, state, control, model=None, **kwargs): + if self.merge_at_every_checkpoint: + self._merge_and_maybe_push(args.output_dir, state.global_step, model) + + def on_train_end(self, args, state, control, model=None, **kwargs): + if not self.merge_at_every_checkpoint: + self._merge_and_maybe_push(args.output_dir, state.global_step, model) + + +class BEMACallback(TrainerCallback): + # docstyle-ignore + r""" + A [`~transformers.TrainerCallback`] that implements [BEMA](https://huggingface.co/papers/2508.00180) + (Bias-Corrected Exponential Moving Average) by [Adam Block](https://huggingface.co/abblock) and [Cyril + Zhang](https://huggingface.co/cyrilzhang). Code from https://github.com/abblock/bema under MIT license. + + BEMA computes model weights that scale like: + + $$ + \theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t + $$ + + where \\( \theta_t \\) is the current model weights, \\( \theta_0 \\) is a snapshot of the model weights at the + first `update_after` step, \\( \text{EMA}_t \\) is the exponential moving average of the model weights, and + \\( \alpha_t \\) is a scaling factor that decays with the number of steps \\( t \\) as + + $$ + \alpha_t = (\rho + \gamma \cdot t)^{-\eta}. + $$ + + The EMA is computed as: + + $$ + \text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t + $$ + + where \\( \beta_t \\) is a decay factor that decays with the number of steps \\( t \\) as + + $$ + \beta_t = (\rho + \gamma \cdot t)^{-\kappa}. + $$ + + Args: + update_freq (`int`, *optional*, defaults to `400`): + Update the BEMA weights every X steps. Denoted this as \\( \phi \\) in the paper. + ema_power (`float`, *optional*, defaults to `0.5`): + Power for the EMA decay factor. Denoted \\( \kappa \\) in the paper. To disable EMA, set this to `0.0`. + bias_power (`float`, *optional*, defaults to `0.2`): + Power for the BEMA scaling factor. Denoted \\( \eta \\) in the paper. To disable BEMA, set this to `0.0`. + lag (`int`, *optional*, defaults to `10`): + Initial offset in the weight decay schedule that controls early-stage smoothness by acting as a virtual + starting age for the updates. Denoted as \\( \rho \\) in the paper. + update_after (`int`, *optional*, defaults to `0`): + Burn-in time before starting to update the BEMA weights. Denoted \\( \tau \\) in the paper. + multiplier (`float`, *optional*, defaults to `1.0`): + Initial value for the EMA decay factor. Denoted as \\( \gamma \\) in the paper. + min_ema_multiplier (`float`, *optional*, defaults to `0.0`): + Minimum value for the EMA decay factor. + device (`str`, *optional*, defaults to `"cpu"`): + Device to use for the BEMA buffers, e.g. `"cpu"` or `"cuda"`. Note that in most cases, this device SHOULD + BE DIFFERENT from the device used for training in order to avoid OOM. + + Example: + + ```python + from trl import BEMACallback + + trainer = Trainer(..., callbacks=[BEMACallback()]) + ``` + """ + + def __init__( + self, + update_freq: int = 400, + ema_power: float = 0.5, + bias_power: float = 0.2, + lag: int = 10, + update_after: int = 0, + multiplier: float = 1.0, + min_ema_multiplier: float = 0.0, + device: str = "cpu", + ): + # User-provided hyperparams + self.update_freq = update_freq + self.ema_power = ema_power + self.bias_power = bias_power + self.lag = lag + self.update_after = update_after + self.multiplier = multiplier + self.min_ema_multiplier = min_ema_multiplier + self.device = device + + # Internal state + self.param_names = [] # references to training model param names + self.thetat_params = [] # references to training model params + self.theta0_params = [] # θ₀ buffers (on self.device) + self.ema_params = [] # EMA buffers (on self.device) + self.running_model = None # a copy of the model to run BEMA on + + @staticmethod + def _unwrap_model(model): + """ + Helper function to unwrap model from various wrappers including DataParallel, DistributedDataParallel, + DeepSpeed, and FSDP. + """ + # Handle DeepSpeed + if hasattr(model, "module") and hasattr(model, "engine"): + # DeepSpeed engine + return model.module + + # Handle FSDP + if hasattr(model, "_fsdp_wrapped_module"): + # FSDP wrapped model + return model._fsdp_wrapped_module + + # Handle DataParallel/DistributedDataParallel + if hasattr(model, "module"): + return model.module + + return model + + @torch.no_grad() + def on_train_begin( + self, args: TrainingArguments, state: TrainerState, control: TrainerControl, model: PreTrainedModel, **kwargs + ): + model = self._unwrap_model(model) + + # Create a new instance and load state_dict + self.running_model = type(model)(model.config).to(self.device) + self.running_model.load_state_dict(model.state_dict()) + + # Cache trainable parameters once in a fixed order + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + self.param_names.append(name) + self.thetat_params.append(param) + + # Clone θ₀ and EMA on the same device as model + theta0 = param.detach().clone().to(self.device) + self.theta0_params.append(theta0) + self.ema_params.append(theta0.clone()) # initialize EMA with θ₀ + + def _ema_beta(self, step: int) -> float: + """Compute the EMA decay factor βₜ = (ρ + γ·t)⁻ᵏᵃᵖᵖᵃ.""" + beta = (self.lag + self.multiplier * step) ** (-self.ema_power) + return max(beta, self.min_ema_multiplier) + + def _bema_alpha(self, step: int) -> float: + """Compute the BEMA scaling factor αₜ = (ρ + γ·t)⁻ᵉᵗᵃ.""" + return (self.lag + self.multiplier * step) ** (-self.bias_power) + + def _update_bema_weights(self, step: int): + beta = self._ema_beta(step) + alpha = self._bema_alpha(step) + + # Compute EMA + BEMA in-place and write directly to running_model + for thetat, theta0, ema, run_param in zip( + self.thetat_params, self.theta0_params, self.ema_params, self.running_model.parameters(), strict=True + ): + thetat = thetat.detach().to(self.device) + ema.mul_(1 - beta).add_(thetat, alpha=beta) # EMA update: ema = (1 - beta) * ema + beta * θₜ + run_param.copy_(ema + alpha * (thetat - theta0)) # BEMA update: run_param = ema + alpha * (θₜ - θ₀) + + @torch.no_grad() + def on_step_end( + self, args: TrainingArguments, state: TrainerState, control: TrainerControl, model: PreTrainedModel, **kwargs + ): + step = state.global_step + + # If we haven't reached the update_after step, skip the BEMA update + if step < self.update_after: + return + + # Snapshot θ₀ and EMA at first update + if step == self.update_after: + for thetat_param, theta0_param, ema_param in zip( + self.thetat_params, self.theta0_params, self.ema_params, strict=True + ): + theta0_param.copy_(thetat_param) + ema_param.copy_(thetat_param) + + # Update BEMA weights every `update_freq` steps + elif (step - self.update_after) % self.update_freq == 0: + self._update_bema_weights(step) + logger.info(f"Updated BEMA weights at step {step}") + + @torch.no_grad() + def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + if state.is_world_process_zero: + save_directory = f"{args.output_dir}/bema" + self.running_model.save_pretrained(save_directory) + logger.info(f"Saved BEMA model to {save_directory}") diff --git a/venv/lib/python3.10/site-packages/trl/trainer/cpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/cpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..426b00da1533e06a7fe2227853275822217daab0 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/cpo_config.py @@ -0,0 +1,218 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class CPOConfig(TrainingArguments): + r""" + Configuration class for the [`CPOTrainer`]. + + This class includes only the parameters that are specific to CPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int`, *optional*): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). + label_smoothing (`float`, *optional*, defaults to `0.0`): + Label smoothing factor. This argument is required if you want to use the default data collator. + loss_type (`str`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"hinge"`: hinge loss on the normalized likelihood from the + [SLiC](https://huggingface.co/papers/2305.10425) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + - `"simpo"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper. + - `"alphapo"`: AlphaPO loss from the [AlphaPO](https://huggingface.co/papers/2501.03884) paper. This + automatically sets `loss_type="simpo"` and `cpo_alpha=0.0`. + + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + cpo_alpha (`float`, *optional*, defaults to `1.0`): + Weight of the BC regularizer in CPO training. + simpo_gamma (`float`, *optional*, defaults to `0.5`): + Target reward margin for the SimPO loss, used only when the `loss_type="simpo"`. + alpha (`float`, *optional*, defaults to `0.0`): + Alpha parameter that controls reward function shape across all loss types. When alpha=0 (default), uses + standard log probability rewards. When `alpha != 0`, applies AlphaPO transformation: `r = (1 - p^(-alpha)) + / alpha` from the [AlphaPO paper](https://huggingface.co/papers/2501.03884). This parameter works with all + loss types. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int`, *optional*): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`,*optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool`, *optional*): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: int | None = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: int | None = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: int | None = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model." + }, + ) + label_smoothing: float = field( + default=0.0, + metadata={"help": "Label smoothing factor."}, + ) + loss_type: str = field( + default="sigmoid", + metadata={ + "help": "Type of loss to use.", + "choices": ["sigmoid", "hinge", "ipo", "simpo", "alphapo"], + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + cpo_alpha: float = field( + default=1.0, + metadata={"help": "Weight of the BC regularizer in CPO training."}, + ) + simpo_gamma: float = field( + default=0.5, + metadata={"help": "Target reward margin for the SimPO loss, used only when the `loss_type='simpo'`."}, + ) + alpha: float = field( + default=0.0, + metadata={ + "help": "Alpha parameter that controls reward function shape across all loss types. When alpha=0 " + "(default), uses standard log probability rewards. When `alpha != 0`, applies AlphaPO transformation: " + "`r = (1 - p^(-alpha)) / alpha` from the AlphaPO paper. This parameter works with all loss types." + }, + ) + label_pad_token_id: int = field( + default=-100, + metadata={"help": "Label pad token id."}, + ) + padding_value: int | None = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={"help": "If `True`, generates and logs completions from the model to W&B during evaluation."}, + ) + is_encoder_decoder: bool | None = field( + default=None, + metadata={"help": "Whether the model is an encoder-decoder model."}, + ) + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + # Syntactic sugar for AlphaPO: set loss_type to "simpo" and cpo_alpha to 0.0 + if self.loss_type == "alphapo": + self.loss_type = "simpo" + self.cpo_alpha = 0.0 + + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/cpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/cpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..565983a6a10f13be2a750b4ef0d5c5c2807e30b6 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/cpo_trainer.py @@ -0,0 +1,1097 @@ +# Copyright 2020-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 inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from collections.abc import Callable +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState, logging +from datasets import Dataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + is_comet_available, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_peft_available, is_torch_fx_proxy + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from .base_trainer import BaseTrainer +from .cpo_config import CPOConfig +from .utils import ( + DPODataCollatorWithPadding, + add_bos_token_if_needed, + add_eos_token_if_needed, + disable_dropout_in_model, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + + +if is_wandb_available(): + import wandb + + +logger = logging.get_logger(__name__) + + +class CPOTrainer(BaseTrainer): + r""" + Initialize CPOTrainer. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. + args ([`CPOConfig`]): + The CPO config arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + """ + + _tag_names = ["trl", "cpo"] + _name = "CPO" + _paper = { + "title": "Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation", + "id": "2401.08417", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{xu2024contrastive, + title = {{Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation}}, + author = {Haoran Xu and Amr Sharaf and Yunmo Chen and Weiting Tan and Lingfeng Shen and Benjamin Van Durme and Kenton Murray and Young Jin Kim}, + year = 2024, + booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=51iwkioZpn} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str | None = None, + args: CPOConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + model_init: Callable[[], PreTrainedModel] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: dict | None = None, + compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the CPOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + dtype = model_init_kwargs.get("dtype") + if dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(dtype, str) and dtype != "auto": + dtype = getattr(torch, dtype) + if dtype != "auto" and not isinstance(dtype, torch.dtype): + raise ValueError( + f"Invalid `dtype` passed to the CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." + ) + model_init_kwargs["dtype"] = dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a CPO dataset.") + if args.max_length is None: + logger.warning( + "`max_length` is not set in the CPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + logger.warning( + "`max_prompt_length` is not set in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if not max_prompt_length < max_length: + raise ValueError( + f"max_prompt_length ({max_prompt_length}) should be strictly less than max_length ({max_length})." + ) + + if args.max_completion_length is None and self.is_encoder_decoder: + logger.warning( + "When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + ) + max_completion_length = 128 + else: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + logger.warning( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + + if args.loss_type in ["hinge", "ipo"] and args.label_smoothing > 0: + logger.warning( + f"You are using the {args.loss_type} loss type that does not support label smoothing. The " + "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.", + ) + if args.loss_type == "kto_pair": + raise ValueError("Support for kto_pair has been removed in CPOTrainer. Please use KTOTrainer.") + + self.beta = args.beta + self.label_smoothing = args.label_smoothing + self.loss_type = args.loss_type + self.cpo_alpha = args.cpo_alpha + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + logger.warning( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + ) + + if args.loss_type == "simpo": + self.simpo_gamma = args.simpo_gamma + + # AlphaPO parameter for reward shaping + self.alpha = args.alpha + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in CPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + + # tokenize the dataset + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: PreTrainedModel | nn.Module | None = None) -> dict: + """Tokenize a single row from a CPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + a != b + for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"], strict=True) + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, list | torch.LongTensor], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: torch.device | None = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def cpo_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the CPO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the CPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. + """ + # Apply AlphaPO reward transformation if alpha != 0 + if self.alpha != 0.0: + # Compute probabilities + chosen_probs = torch.exp(policy_chosen_logps) + rejected_probs = torch.exp(policy_rejected_logps) + + # Apply AlphaPO transformation: r = (1 - p^(-alpha)) / alpha + policy_chosen_rewards = (1 - chosen_probs.pow(-self.alpha)) / self.alpha + policy_rejected_rewards = (1 - rejected_probs.pow(-self.alpha)) / self.alpha + + logits = (policy_chosen_rewards - policy_rejected_rewards).to(self.accelerator.device) + else: + # Standard log probability rewards when alpha = 0 + logits = (policy_chosen_logps - policy_rejected_logps).to(self.accelerator.device) + + # The beta is a temperature parameter for the CPO loss, typically something in the range of 0.1 to 0.5. + # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and + # calculates a conservative CPO loss. + + if self.loss_type == "simpo": + gamma_logratios = self.simpo_gamma / self.beta + logits = logits - gamma_logratios + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + elif self.loss_type == "sigmoid": + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + elif self.loss_type == "hinge": + losses = torch.relu(1 - self.beta * logits) + elif self.loss_type == "ipo": + # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper. + losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise ValueError( + f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'simpo']" + ) + + # Calculate rewards for logging + if self.alpha != 0.0: + # When using AlphaPO transformation, use the transformed rewards + chosen_rewards = self.beta * policy_chosen_rewards.to(self.accelerator.device).detach() + rejected_rewards = self.beta * policy_rejected_rewards.to(self.accelerator.device).detach() + else: + # Standard log probability rewards + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, list | torch.LongTensor] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + labels = concatenated_batch["concatenated_labels"].clone() + + if self.cpo_alpha == 0: + nll_loss = torch.tensor(0.0).to(self.accelerator.device) + else: + nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=self.loss_type in ["ipo", "simpo"], + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, list | torch.LongTensor], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the CPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards = self.cpo_loss( + policy_chosen_logps, + policy_rejected_logps, + ) + + loss = losses.mean() + self.cpo_alpha * policy_nll_loss + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() + metrics[f"{prefix}rewards/margins"] = ( + self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item() + ) + metrics[f"{prefix}logps/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean().item() + ) + metrics[f"{prefix}logps/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean().item() + ) + metrics[f"{prefix}logits/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}logits/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean().item() + + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs=False, + num_items_in_batch=None, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_output_decoded + + def prediction_step( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + prediction_loss_only: bool, + ignore_keys: list[str] | None = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: bool | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] + for prompt, pol in zip(random_batch["prompt"], policy_output_decoded, strict=True) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float`, *optional*): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/dpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/dpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..887cd95646fda457c4c03423c496866dba38db16 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/dpo_config.py @@ -0,0 +1,532 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from transformers import TrainingArguments + + +class FDivergenceType(Enum): + """Types of f-divergence functions for DPO loss regularization. + + Attributes: + REVERSE_KL: Reverse KL divergence. + JS_DIVERGENCE: Jensen-Shannon divergence. + ALPHA_DIVERGENCE: Alpha divergence. + + Examples: + ```python + >>> from trl.trainer.dpo_config import DPOConfig, FDivergenceType + + >>> config = DPOConfig( + ... f_divergence_type=FDivergenceType.ALPHA_DIVERGENCE, + ... f_alpha_divergence_coef=0.5, # used only with ALPHA_DIVERGENCE + ... ) + ``` + """ + + REVERSE_KL = "reverse_kl" + JS_DIVERGENCE = "js_divergence" + ALPHA_DIVERGENCE = "alpha_divergence" + + +class FDivergenceConstants: + """Constants for f-divergence types and their parameters. + + Attributes: + ALPHA_DIVERGENCE_COEF_KEY (`str`): Key for the alpha divergence coefficient. + ALPHA_DIVERGENCE_COEF_DEFAULT (`float`): Default value for the alpha divergence coefficient. + """ + + ALPHA_DIVERGENCE_COEF_KEY = "alpha_divergence_coef" + ALPHA_DIVERGENCE_COEF_DEFAULT = 1.0 + + +@dataclass +class DPOConfig(TrainingArguments): + r""" + Configuration class for the [`DPOTrainer`]. + + This class includes only the parameters that are specific to DPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of the + [`DPOTrainer`] is provided as a string. + ref_model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `ref_model` argument of the + [`DPOTrainer`] is provided as a string. + model_adapter_name (`str`, *optional*): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, *optional*): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + force_use_ref_model (`bool`, *optional*, defaults to `False`): + If you provide a PEFT model as the active model and wish to use a different model for the `ref_model`, set + this flag to `True`. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_logits_to_keep (`bool`, *optional*, defaults to `False`): + If `True`, only a specified number of logits are computed in the forward pass. This can be useful for + saving memory and speeding up training by not computing the logits for all tokens, especially in scenarios + when working with very long prompts where labels are ignored (-100). + + > Parameters that control the data preprocessing + + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + pad_token (`str`, *optional*): + Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that is also `None`, + it falls back to `processing_class.eos_token`. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Padding value to use for labels. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. + max_completion_length (`int`, *optional*): + Maximum length of the completion. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the full sequence (prompt + completion). + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the sequence exceeds `max_length`. Possible values are `"keep_end"` and + `"keep_start"`. + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the `flash_attention_2` attention implementation, which can efficiently handle the flattened + batch structure. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute the log probabilities from the reference model. Setting this to `True` allows + training without needing the reference model during training, which can help reduce GPU memory usage. If + set to `False` (default), the reference model will be used during training to compute log probabilities + on-the-fly. + precompute_ref_batch_size (`int`, *optional*): + Batch size to use when precomputing reference model log probabilities. This can be set higher than the + training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for + training and `per_device_eval_batch_size` for evaluation. + tools (`list[dict] | None`, *optional*): + List of tools (callable functions) that will be accessible to the model. If the template does not support + function calling, this argument will have no effect. + + > Parameters that control the training + + loss_type (`str` or `list[str]`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"hinge"`: hinge loss on the normalized likelihood from the + [SLiC](https://huggingface.co/papers/2305.10425) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + - `"exo_pair"`: pairwise EXO loss from the [EXO](https://huggingface.co/papers/2402.00856) paper. + - `"nca_pair"`: pairwise NCA loss from the [NCA](https://huggingface.co/papers/2402.05369) paper. + - `"robust"`: unbiased estimate of the DPO loss that is robust to preference noise from the [Robust + DPO](https://huggingface.co/papers/2403.00409) paper. + - `"bco_pair"`: pairwise BCO loss from the [BCO](https://huggingface.co/papers/2404.04656) paper. + - `"sppo_hard"`: SPPO loss with hard label from the [SPPO](https://huggingface.co/papers/2405.00675) + paper. + - `"aot"`: AOT loss for paired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper. + - `"aot_pair"`: AOT loss for unpaired datasets from the [AOT](https://huggingface.co/papers/2406.05882) + paper. + - `"discopop"`: DiscoPOP (a.k.a Log-Ratio Modulated Loss, LRML) loss from the + [DiscoPOP](https://huggingface.co/papers/2406.08414) paper. + - `"apo_zero"`: APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + - `"apo_down"`: APO-down loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + - `"sft"`: Negative log-likelihood loss (standard supervised fine-tuning loss). + + Multiple loss types can be combined using comma separation (e.g., `["sigmoid", "bco_pair", "sft"]` for + [MPO](https://huggingface.co/papers/2411.10442)). The `loss_weights` parameter can be used to specify + corresponding weights for each loss type. + + use_liger_loss (`bool`, *optional*): + Whether to use Liger loss. + + + + Parameter `use_liger_loss` is deprecated and will be removed in version 0.28.0. Use `use_liger_kernel` + instead. + + + + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_kernel` is + `True`. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). + f_divergence_type ([`FDivergenceType`] or `str`, *optional*, defaults to `FDivergenceType.REVERSE_KL`): + Type of f-divergence regularization function to compute divergence between policy and reference model. + f_alpha_divergence_coef (`float`, *optional*, defaults to `1.0`): + α coefficient in the α-divergence u^-α regularization function for DPO loss. + reference_free (`bool`, *optional*, defaults to `False`): + Whether to ignore the provided reference model and implicitly use a reference model that assigns equal + probability to all responses. + label_smoothing (`float`, *optional*, defaults to `0.0`): + Robust DPO label smoothing parameter from the [cDPO report](https://ericmitchell.ai/cdpo.pdf) and [Robust + DPO](https://huggingface.co/papers/2403.00409) paper that should be between `0.0` and `0.5`. + use_weighting (`bool`, *optional*, defaults to `False`): + Whether to weight the loss as done in the [WPO paper](https://huggingface.co/papers/2406.11827). + rpo_alpha (`float`, *optional*): + α parameter from the [RPO paper](https://huggingface.co/papers/2404.19733) (v3), which controls the + weighting of the NLL term in the loss. If `None`, no weighting is applied and the loss is the same as the + DPO loss. The paper recommends `rpo_alpha=1.0`. + ld_alpha (`float`, *optional*): + α parameter from the [LD-DPO paper](https://huggingface.co/papers/2409.06411), which controls the weighting + of the verbose token log-probabilities in responses. If `None`, no weighting is applied to the verbose + part, and the loss is equivalent to the standard DPO loss. The paper recommends setting `ld_alpha` between + `0.0` and `1.0`. + discopop_tau (`float`, *optional*, defaults to `0.05`): + τ/temperature parameter from the [DiscoPOP](https://huggingface.co/papers/2406.08414) paper, which controls + the shape of log ratio modulated loss. The paper recommends the default value `discopop_tau=0.05`. + loss_weights (`list[float]`, *optional*): + List of loss weights for multi-loss combinations. Used when combining multiple loss types. Example: `[0.8, + 0.2, 1.0]` for [MPO](https://huggingface.co/papers/2411.10442). If not provided, defaults to equal weights + (`1.0`) for all loss types. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + + > Parameters that control the logging + + generate_during_eval (`bool`, *optional*, defaults to `False`): + Whether to generate and log completions from both the model and the reference model to W&B or Comet during + evaluation. + + > Deprecated parameters + + padding_value: + + + + This parameter is deprecated and will be removed in version 0.26.0. Use `pad_token` (`str`) instead. + + + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs", "ref_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `DPOTrainer` is provided as a string." + }, + ) + ref_model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `ref_model` argument " + "of the `DPOTrainer` is provided as a string." + }, + ) + model_adapter_name: str | None = field( + default=None, + metadata={"help": "Name of the train target PEFT adapter, when using LoRA with multiple adapters."}, + ) + ref_adapter_name: str | None = field( + default=None, + metadata={"help": "Name of the reference PEFT adapter, when using LoRA with multiple adapters."}, + ) + force_use_ref_model: bool = field( + default=False, + metadata={ + "help": "If you provide a PEFT model as the active model and wish to use a different model for the " + "`ref_model`, set this flag to `True`." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + use_logits_to_keep: bool = field( + default=False, + metadata={ + "help": "If `True`, only a specified number of logits are computed in the forward pass. This can be " + "useful for saving memory and speeding up training by not computing the logits for all tokens, especially " + "in scenarios when working with very long prompts where labels are ignored (-100)." + }, + ) + + # Parameters that control the data preprocessing + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + pad_token: str | None = field( + default=None, + metadata={ + "help": "Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that " + "is also `None`, it falls back to `processing_class.eos_token`." + }, + ) + label_pad_token_id: int = field( + default=-100, + metadata={"help": "Padding value to use for labels."}, + ) + max_prompt_length: int | None = field( + default=512, + metadata={"help": "Maximum length of the prompt."}, + ) + max_completion_length: int | None = field( + default=None, + metadata={"help": "Maximum length of the completion."}, + ) + max_length: int | None = field( + default=1024, + metadata={"help": "Maximum length of the full sequence (prompt + completion)."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the sequence exceeds `max_length`. Possible values are `'keep_end'` " + "and `'keep_start'`.", + "choices": ["keep_end", "keep_start"], + }, + ) + padding_free: bool = field( + default=False, + metadata={ + "help": "Whether to perform forward passes without padding by flattening all sequences in the batch into " + "a single continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, " + "this is only supported with the `flash_attention_2` attention implementation, which can efficiently " + "handle the flattened batch structure." + }, + ) + precompute_ref_log_probs: bool = field( + default=False, + metadata={ + "help": "Whether to precompute the log probabilities from the reference model. Setting this to `True` " + "allows training without needing the reference model during training, which can help reduce GPU memory " + "usage. If set to `False` (default), the reference model will be used during training to compute log " + "probabilities on-the-fly." + }, + ) + precompute_ref_batch_size: int | None = field( + default=None, + metadata={ + "help": "Batch size to use when precomputing reference model log probabilities. This can be set higher " + "than the training batch size to speed up preprocessing. If `None`, defaults to " + "`per_device_train_batch_size` for training and `per_device_eval_batch_size` for evaluation." + }, + ) + tools: list[dict] | None = field( + default=None, + metadata={ + "help": "List of tools (callable functions) that will be accessible to the model. If the template does " + "not support function calling, this argument will have no effect." + }, + ) + + # Parameters that control the training + loss_type: list[str] = field( + default_factory=lambda: ["sigmoid"], + metadata={ + "help": "Type of loss to use. Possible values are: `'sigmoid'`, `'hinge'`, `'ipo'`, `'exo_pair'`, " + "`'nca_pair'`, `'robust'`, `'bco_pair'`, `'sppo_hard'`, `'aot'`, `'aot_pair'`, `'discopop'`, " + "`'apo_zero'`, `'apo_down'` and `'sft'`. Multiple loss types can be combined using comma separation " + "(e.g., `['sigmoid', 'bco_pair', 'sft']` for MPO). The `loss_weights` parameter can be used to specify " + "corresponding weights for each loss type." + }, + ) + use_liger_loss: bool = field( + default=None, + metadata={"help": "Whether to use Liger loss."}, + ) + base_model_attribute_name: str = field( + default="model", + metadata={ + "help": "Name of the attribute in the model that contains the base model. This is used to get the base " + "model from the model when the model does not have a `get_decoder` method in the case when " + "`use_liger_kernel` is `True`." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. " + "Higher β means less deviation from the reference model." + }, + ) + f_divergence_type: FDivergenceType | str = field( + default=FDivergenceType.REVERSE_KL, + metadata={ + "help": "Type of f-divergence regularization function to compute divergence between policy and reference " + "model." + }, + ) + f_alpha_divergence_coef: float = field( + default=1.0, + metadata={"help": "α coefficient in the α-divergence u^-α regularization function for DPO loss."}, + ) + reference_free: bool = field( + default=False, + metadata={ + "help": "Whether to ignore the provided reference model and implicitly use a reference model that assigns " + "equal probability to all responses." + }, + ) + label_smoothing: float = field( + default=0.0, + metadata={ + "help": "Robust DPO label smoothing parameter from the cDPO report and Robust DPO paper that should " + "be between `0.0` and `0.5`." + }, + ) + use_weighting: bool = field( + default=False, + metadata={"help": "Whether to weight the loss as done in the WPO paper."}, + ) + rpo_alpha: float | None = field( + default=None, + metadata={ + "help": "α parameter from the RPO paper (v3), which controls the weighting of the NLL term in the loss. " + "If `None`, no weighting is applied and the loss is the same as the DPO loss. The paper recommends " + "`rpo_alpha=1.0`." + }, + ) + ld_alpha: float | None = field( + default=None, + metadata={ + "help": "α parameter from the LD-DPO paper, which controls the weighting of the verbose token " + "log-probabilities in responses. If `None`, no weighting is applied to the verbose part, and the loss is " + "equivalent to the standard DPO loss. The paper recommends setting `ld_alpha` between `0.0` and `1.0`.", + }, + ) + discopop_tau: float = field( + default=0.05, + metadata={ + "help": "τ/temperature parameter from the DiscoPOP paper, which controls the shape of log ratio modulated " + "loss. The paper recommends the default value `discopop_tau=0.05`." + }, + ) + loss_weights: list[float] | None = field( + default=None, + metadata={ + "help": "List of loss weights for multi-loss combinations. Used when combining multiple loss types. " + "Example: `[0.8, 0.2, 1.0]` for MPO. If not provided, defaults to equal weights (`1.0`) for all loss " + "types." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + + # Parameters that control the logging + generate_during_eval: bool = field( + default=False, + metadata={ + "help": "Whether to generate and log completions from both the model and the reference model to W&B, MLFLow " + "or Comet during evaluation." + }, + ) + + # Deprecated arguments + padding_value: int | None = field( + default=None, + metadata={"help": "Deprecated, use `pad_token` (str) instead."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + self.f_divergence_type = FDivergenceType(self.f_divergence_type) + + # Normalize loss_type to string format for internal use + if hasattr(self.loss_type, "__len__") and len(self.loss_type) == 1: + self.loss_type = self.loss_type[0] + + # Validate loss_type + if self.loss_weights is not None: + loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type] + if len(self.loss_weights) != len(loss_types): + raise ValueError( + f"Length of loss_weights list ({self.loss_weights}) must match number of loss types " + f"({loss_types})." + ) + + if self.use_liger_loss is not None: + warnings.warn( + "The `use_liger_loss` argument is deprecated and will be removed in version 0.28.0. Please use " + "`use_liger_kernel` instead.", + FutureWarning, + stacklevel=2, + ) + self.use_liger_kernel = self.use_liger_loss + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/dpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/dpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..82cd59c2c047b53f87a1c3e8e1e87befe8e75bad --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/dpo_trainer.py @@ -0,0 +1,2008 @@ +# Copyright 2020-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 inspect +import random +import textwrap +import warnings +from collections import defaultdict +from collections.abc import Callable +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState, logging +from accelerate.utils import tqdm +from datasets import Dataset, IterableDataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoProcessor, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, +) +from transformers.data.data_collator import DataCollatorMixin +from transformers.integrations import ( + is_comet_available, + is_mlflow_available, + is_wandb_available, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_liger_kernel_available, is_peft_available + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from ..models import create_reference_model, prepare_deepspeed +from ..models.utils import prepare_fsdp +from .base_trainer import BaseTrainer +from .callbacks import SyncRefModelCallback +from .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType +from .utils import ( + RunningMoments, + cap_exp, + create_model_from_path, + disable_dropout_in_model, + empty_cache, + flush_left, + flush_right, + get_config_model_id, + log_table_to_comet_experiment, + pad, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import ( + PeftConfig, + PeftModel, + get_peft_model, + prepare_model_for_kbit_training, + ) + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearDPOLoss + + +if is_wandb_available(): + import wandb + +if is_mlflow_available(): + import mlflow + + +logger = logging.get_logger(__name__) + + +def shift_tokens_right(input_ids: torch.Tensor, decoder_start_token_id: int) -> torch.Tensor: + """Shift input ids one token to the right, and pad with pad_token_id""" + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + return shifted_input_ids + + +@dataclass +class DataCollatorForPreference(DataCollatorMixin): + """ + Data collator used for preference data. Inputs are dynamically padded to the maximum length of a batch if they are + not all of the same length. + + Args: + pad_token_id (`int`): + Token ID to use for padding. + return_tensors (`str`, *optional*, defaults to `"pt"`): + Type of Tensor to return. Only `"pt"` is currently supported. + + Examples: + ```python + >>> from trl import DataCollatorForPreference + + >>> collator = DataCollatorForPreference(pad_token_id=0) + >>> examples = [ + ... {"prompt_input_ids": [1, 2, 3], "chosen_input_ids": [4, 5], "rejected_input_ids": [6]}, + ... {"prompt_input_ids": [7, 8], "chosen_input_ids": [9, 10], "rejected_input_ids": [11, 12, 13]}, + ... ] + >>> collator(examples) + {'prompt_input_ids': tensor([[1, 2, 3], + [0, 7, 8]]), + 'prompt_attention_mask': tensor([[1, 1, 1], + [0, 1, 1]]), + 'chosen_input_ids': tensor([[ 4, 5], + [ 9, 10]]), + 'chosen_attention_mask': tensor([[1, 1], + [1, 1]]), + 'rejected_input_ids': tensor([[ 6, 0, 0], + [11, 12, 13]]), + 'rejected_attention_mask': tensor([[1, 0, 0], + [1, 1, 1]]) + } + ``` + """ + + pad_token_id: int + return_tensors: str = "pt" + + def torch_call(self, examples: list[list[int] | Any | dict[str, Any]]) -> dict[str, Any]: + # Convert to tensor + prompt_input_ids = [torch.tensor(example["prompt_input_ids"]) for example in examples] + prompt_attention_mask = [torch.ones_like(input_ids) for input_ids in prompt_input_ids] + chosen_input_ids = [torch.tensor(example["chosen_input_ids"]) for example in examples] + chosen_attention_mask = [torch.ones_like(input_ids) for input_ids in chosen_input_ids] + rejected_input_ids = [torch.tensor(example["rejected_input_ids"]) for example in examples] + rejected_attention_mask = [torch.ones_like(input_ids) for input_ids in rejected_input_ids] + if "pixel_values" in examples[0]: + pixel_values = [torch.tensor(example["pixel_values"]) for example in examples] + if "pixel_attention_mask" in examples[0]: + pixel_attention_mask = [torch.tensor(example["pixel_attention_mask"]) for example in examples] + if "ref_chosen_logps" in examples[0] and "ref_rejected_logps" in examples[0]: + ref_chosen_logps = torch.tensor([example["ref_chosen_logps"] for example in examples]) + ref_rejected_logps = torch.tensor([example["ref_rejected_logps"] for example in examples]) + + # Pad + output = {} + output["prompt_input_ids"] = pad(prompt_input_ids, padding_value=self.pad_token_id, padding_side="left") + output["prompt_attention_mask"] = pad(prompt_attention_mask, padding_value=0, padding_side="left") + output["chosen_input_ids"] = pad(chosen_input_ids, padding_value=self.pad_token_id) + output["chosen_attention_mask"] = pad(chosen_attention_mask, padding_value=0) + output["rejected_input_ids"] = pad(rejected_input_ids, padding_value=self.pad_token_id) + output["rejected_attention_mask"] = pad(rejected_attention_mask, padding_value=0) + if "pixel_values" in examples[0]: + output["pixel_values"] = pad(pixel_values, padding_value=0.0) + if "pixel_attention_mask" in examples[0]: + output["pixel_attention_mask"] = pad(pixel_attention_mask, padding_value=0) + if "image_sizes" in examples[0]: + output["image_sizes"] = torch.tensor([example["image_sizes"] for example in examples]) + if "ref_chosen_logps" in examples[0] and "ref_rejected_logps" in examples[0]: + output["ref_chosen_logps"] = ref_chosen_logps + output["ref_rejected_logps"] = ref_rejected_logps + if "token_type_ids" in examples[0]: + token_type_ids = [torch.tensor(example["token_type_ids"]) for example in examples] + output["token_type_ids"] = pad(token_type_ids, padding_value=0, padding_side="left") + + return output + + +class DPOTrainer(BaseTrainer): + """ + Trainer for Direct Preference Optimization (DPO) method. + + This class is a wrapper around the [`transformers.Trainer`] class and inherits all of its attributes and methods. + + Args: + model (`str | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + ref_model ([`PreTrainedModelWrapper`]): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + args ([`DPOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`DataCollatorForPreference`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. DPO supports [preference](#preference) type and. The format of the samples can + be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return + a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to + `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered + after the last eval batch to signal that the function needs to calculate and return the global summary + statistics rather than accumulating the batch-level statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + """ + + _tag_names = ["trl", "dpo"] + _name = "DPO" + _paper = { + "title": "Direct Preference Optimization: Your Language Model is Secretly a Reward Model", + "id": "2305.18290", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{rafailov2023direct, + title = {{Direct Preference Optimization: Your Language Model is Secretly a Reward Model}}, + author = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Christopher D. Manning and Stefano Ermon and Chelsea Finn}, + year = 2023, + booktitle = {Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023}, + url = {http://papers.nips.cc/paper_files/paper/2023/hash/a85b405ed65c6477a4fe8302b5e06ce7-Abstract-Conference.html}, + editor = {Alice Oh and Tristan Naumann and Amir Globerson and Kate Saenko and Moritz Hardt and Sergey Levine}, + }"""), + } + + def __init__( + self, + model: str | nn.Module | PreTrainedModel, + ref_model: PreTrainedModel | nn.Module | str | None = None, + args: DPOConfig | None = None, + data_collator: DataCollator | None = None, # type: ignore + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = DPOConfig(f"{model_name}-DPO") + + # Model and reference model + if isinstance(model, str): + model = create_model_from_path(model, **args.model_init_kwargs or {}) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `DPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + model_id = get_config_model_id(model.config) + if isinstance(ref_model, str): + ref_model = create_model_from_path(ref_model, **args.ref_model_init_kwargs or {}) + else: + if args.ref_model_init_kwargs is not None: + logger.warning( + "You passed `ref_model_init_kwargs` to the `DPOConfig`, but your model is already instantiated. " + "The `ref_model_init_kwargs` will be ignored." + ) + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you can simply omit the `ref_model` argument and it will be created for you." + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(model_id) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + if args.padding_value is not None: # deprecated, will be removed in 0.26.0. + warnings.warn( + "The `padding_value` argument is deprecated and will be removed in version 0.26.0. Please use " + "`pad_token` (str) instead." + ) + self.pad_token_id = args.padding_value + else: + pad_token = args.pad_token or tokenizer.pad_token or tokenizer.eos_token + self.pad_token_id = tokenizer.convert_tokens_to_ids(pad_token) + if self.pad_token_id is None: + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + + # PEFT configuration and model wrapping + model = self._prepare_peft_model(model, ref_model, peft_config, args) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available() or is_mlflow_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases, MLFlow or Comet to be installed." + " Please install `wandb`, `mlflow` or `comet-ml` to resolve." + ) + + self.is_encoder_decoder = model.config.is_encoder_decoder + self.is_vision_model = model.config.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.keys() + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = args.model_adapter_name + self.ref_adapter_name = args.ref_adapter_name + self.reference_free = args.reference_free + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model or args.precompute_ref_log_probs: + # The `model` with adapters turned off will be used as the reference model + self.ref_model = None + else: + self.ref_model = create_reference_model(model) + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Liger kernel + if args.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_kernel=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if args.loss_type not in ["sigmoid", "apo_zero", "apo_down", "sppo_hard", "nca_pair"]: + raise ValueError( + "You set `use_liger_kernel=True` but the loss type is not from `[sigmoid, apo_zero, apo_down, sppo_hard, nca_pair`. " + "Please set `loss_type='[sigmoid | apo_zero | apo_down | sppo_hard | nca_pair]'` to use the liger kernel." + ) + self.dpo_loss_fn = LigerFusedLinearDPOLoss( + ignore_index=args.label_pad_token_id, + beta=args.beta, + use_ref_model=not args.reference_free, + average_log_prob=False, + loss_type=args.loss_type, + ) + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in DPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Data collator + if data_collator is None: + data_collator = DataCollatorForPreference(pad_token_id=self.pad_token_id) + + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length + self.max_length = args.max_length + self.truncation_mode = args.truncation_mode + self.precompute_ref_log_probs = args.precompute_ref_log_probs + self.use_logits_to_keep = args.use_logits_to_keep + + if args.padding_free: + if model.config._attn_implementation != "flash_attention_2": + logger.warning( + "Padding-free training is enabled, but the attention implementation is not set to " + "'flash_attention_2'. Padding-free training flattens batches into a single sequence, and " + "'flash_attention_2' is the only known attention mechanism that reliably supports this. Using " + "other implementations may lead to unexpected behavior. To ensure compatibility, set " + "`attn_implementation='flash_attention_2'` in the model configuration, or verify that your " + "attention mechanism can handle flattened sequences." + ) + if args.per_device_train_batch_size == 1: + logger.warning( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + self.padding_free = args.padding_free + + # Since ref_logs are precomputed on the first call to get_train/eval_dataloader + # keep track of first called to avoid computation of future calls + self._precomputed_train_ref_log_probs = False + self._precomputed_eval_ref_log_probs = False + + self.beta = args.beta + self.label_smoothing = args.label_smoothing + self.loss_type = args.loss_type if isinstance(args.loss_type, list) else [args.loss_type] + self.loss_weights = args.loss_weights + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.use_weighting = args.use_weighting + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + logger.warning( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + ) + for loss_type in self.loss_type: + if ( + loss_type in ["hinge", "ipo", "bco_pair", "sppo_hard", "nca_pair", "apo_zero", "apo_down"] + and args.label_smoothing > 0 + ): + logger.warning( + f"You are using the {loss_type} loss type that does not support label smoothing. The " + "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this " + "warning.", + ) + if loss_type == "kto_pair": + raise ValueError("Support for kto_pair has been removed in DPOTrainer. Please use KTOTrainer.") + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + self.f_divergence_type = args.f_divergence_type + self.f_divergence_params = {FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY: args.f_alpha_divergence_coef} + self.dataset_num_proc = args.dataset_num_proc + + # Dataset preparation + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + # Deepspeed Zero-3 does not support precompute_ref_log_probs + if self.is_deepspeed_enabled: + if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." + ) + + if self.ref_model is None: + if not (self.is_peft_model or self.precompute_ref_log_probs): + raise ValueError( + "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" + ) + if args.sync_ref_model: + raise ValueError( + "You currently cannot use `ref_model=None` with TR-DPO method. Please provide `ref_model`." + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with TR-DPO method. Please set `precompute_ref_log_probs=False`." + ) + + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + if "bco_pair" in self.loss_type: + self.running = RunningMoments(self.accelerator) + + @property + def padding_value(self): + warnings.warn( + "The `padding_value` property is deprecated and will be removed in version 0.26.0. Please use " + "`pad_token_id` instead.", + ) + return self.pad_token_id + + @padding_value.setter + def padding_value(self, value): + warnings.warn( + "The `padding_value` property is deprecated and will be removed in version 0.26.0. Please use " + "`pad_token_id` instead.", + ) + self.pad_token_id = value + + def _prepare_peft_model( + self, model: PreTrainedModel, ref_model: PreTrainedModel, peft_config: Any, args: DPOConfig + ) -> PreTrainedModel: + """Prepares a model for PEFT training.""" + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if ref_model is not None and not args.force_use_ref_model: + raise ValueError( + "You passed both a ref_model and a peft_config. For training PEFT adapters with DPO there is no need to pass a reference" + " model. Please pass `ref_model=None` in case you want to train PEFT adapters, or pass a ref_model with `force_use_ref_model=True` in DPOTrainer's init." + " if you want to use a different ref_model." + ) + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + else: + model = self._prepare_gradient_checkpointing(model, args) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + else: + model = self._prepare_gradient_checkpointing(model, args) + + return model + + def _prepare_gradient_checkpointing(self, model: PreTrainedModel, args: DPOConfig): + """Prepare the gradienting checkpointing for the model.""" + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + if args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + return model + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + args: DPOConfig, + dataset_name: str, + ) -> Dataset | IterableDataset: + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc nor writer_batch_size + map_kwargs["num_proc"] = args.dataset_num_proc + map_kwargs["writer_batch_size"] = 10 + + with PartialState().main_process_first(): + # Extract prompt if needed + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Extracting prompt in {dataset_name} dataset" + dataset = dataset.map(maybe_extract_prompt, **map_kwargs) + + # Apply the chat template if needed + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Applying chat template to {dataset_name} dataset" + dataset = dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class, "tools": args.tools}, **map_kwargs + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + dataset = dataset.map( + self.tokenize_row if not self.is_vision_model else self.process_row, + remove_columns=["chosen", "rejected"], + fn_kwargs={ + "processing_class": processing_class, + "max_prompt_length": args.max_prompt_length, + "max_completion_length": args.max_completion_length, + # for enc-dec, we add the special tokens ([bos_token] + prompt + [eos_token]; completion + [eos_token]) + "add_special_tokens": False, + }, + **map_kwargs, + ) + + return dataset + + @staticmethod + def tokenize_row( + features: dict[str, str], + processing_class: PreTrainedTokenizerBase, + max_prompt_length: int | None = None, + max_completion_length: int | None = None, + add_special_tokens: bool = True, + ) -> dict[str, list[int]]: + """ + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"chosen"`, and `"rejected"`. + processing_class ([`~transformers.PreTrainedTokenizerBase`]): + Processing class used to process the data. + max_prompt_length (`int` or `None`): + Maximum length of the prompt sequence. If `None`, the prompt sequence is not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + add_special_tokens (`bool`): + Whether to add special tokens to the sequences. Typically used for encoder-decoder models. If `True`, + the prompt sequence will have a bos token prepended and an eos token appended. In any case, the + completion sequences will have an eos token appended. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"prompt_input_ids"`, `"chosen_input_ids"`, and + `"rejected_input_ids". + + Example: + ```python + >>> from transformers import GPT2Tokenizer + + >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + >>> features = {"prompt": "The sky is", "chosen": " blue", "rejected": " green"} + >>> DPOTrainer.tokenize_row( + ... features, tokenizer, max_prompt_length=3, max_completion_length=3, add_special_tokens=False + ... ) + {'prompt_input_ids': [464, 6766, 318], 'chosen_input_ids': [4171, 50256], 'rejected_input_ids': [4077, 50256]} + ``` + """ + tokenizer = processing_class # the processing class is a tokenizer + prompt_input_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"] + + # Add special tokens (typically for encoder-decoder models) + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + return { + "prompt_input_ids": prompt_input_ids, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + + @staticmethod + def process_row( + features: dict[str, str], + processing_class: PreTrainedTokenizerBase, + max_prompt_length: int | None = None, + max_completion_length: int | None = None, + add_special_tokens: bool = True, + ) -> dict[str, list[int]]: + """ + Same as `tokenize_row` but for vision models. Please refer to `tokenize_row` for more information. + """ + processor, tokenizer = processing_class, processing_class.tokenizer # the processing class is a processor + processed_features = processor(images=features["images"], text=features["prompt"], add_special_tokens=False) + + prompt_input_ids = processed_features["input_ids"][0] + pixel_values = processed_features["pixel_values"][0] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"] + + # Add special tokens (typically for encoder-decoder models) + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + output = { + "prompt_input_ids": prompt_input_ids, + "pixel_values": pixel_values, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + if "token_type_ids" in processed_features: + output["token_type_ids"] = processed_features["token_type_ids"][0] + + return output + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In DPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by `DataCollatorForPreference`, hence the override. + if self._signature_columns is None: + self._signature_columns = [ + "prompt_input_ids", + "chosen_input_ids", + "rejected_input_ids", + "image_sizes", + "token_type_ids", + "ref_chosen_logps", + "ref_rejected_logps", + ] + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. + """ + + if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size + dataloader_params = { + "batch_size": batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) + + ref_chosen_logps = [] + ref_rejected_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch) + ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( + (ref_chosen_logp, ref_rejected_logp) + ) + ref_chosen_logps.append(ref_chosen_logp.cpu()) + ref_rejected_logps.append(ref_rejected_logp.cpu()) + + # Unnecessary cache clearing to avoid OOM + empty_cache() + self.accelerator.free_memory() + + all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy() + all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy() + + self.train_dataset = self.train_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps) + self.train_dataset = self.train_dataset.add_column( + name="ref_rejected_logps", column=all_ref_rejected_logps + ) + + self._precomputed_train_ref_log_probs = True + + return super().get_train_dataloader() + + def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. + + Args: + eval_dataset (`torch.utils.data.Dataset`, *optional*): + If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted + by the `model.forward()` method are automatically removed. It must implement `__len__`. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset + + if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + dataloader_params = { + "batch_size": batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) + + ref_chosen_logps = [] + ref_rejected_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch) + ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( + (ref_chosen_logp, ref_rejected_logp) + ) + ref_chosen_logps.append(ref_chosen_logp.cpu()) + ref_rejected_logps.append(ref_rejected_logp.cpu()) + + all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy() + all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy() + + eval_dataset = eval_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps) + eval_dataset = eval_dataset.add_column(name="ref_rejected_logps", column=all_ref_rejected_logps) + + # Save calculated ref_chosen_logps and ref_rejected_logps to the eval_dataset for subsequent runs + if self.eval_dataset is not None: + self.eval_dataset = eval_dataset + self._precomputed_eval_ref_log_probs = True + + return super().get_eval_dataloader(eval_dataset=eval_dataset) + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.set_adapter(self.model_adapter_name or "default") + + def compute_ref_log_probs(self, batch: dict[str, torch.LongTensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Computes log probabilities of the reference model for a single padded batch of a DPO specific dataset.""" + compte_ref_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with torch.no_grad(), compte_ref_context_manager: + if self.ref_model is None: + with self.null_ref_context(): + ref_model_output = self.concatenated_forward(self.model, batch, is_ref_model=True) + else: + ref_model_output = self.concatenated_forward(self.ref_model, batch, is_ref_model=True) + return ref_model_output["chosen_logps"], ref_model_output["rejected_logps"] + + @staticmethod + def concatenated_inputs( + batch: dict[str, list | torch.LongTensor], padding_value: int + ) -> dict[str, torch.LongTensor]: + """ + Concatenate the `chosen` and `rejected` inputs from the batch into a single tensor for both the prompt and + completion sequences. + + Args: + batch (`dict[str, list | torch.LongTensor]`): + A batch of input data. The batch must contain the following keys: + + - `"prompt_input_ids"`: Tensor of shape `(batch_size, prompt_length)` representing the prompt input + IDs. + - `"chosen_input_ids"`: Tensor of shape `(batch_size, chosen_length)` representing the chosen + completion input IDs. + - `"rejected_input_ids"`: Tensor of shape `(batch_size, rejected_length)` representing the rejected + completion input IDs. + - `"prompt_pixel_values"` (optional): Tensor for pixel values, if available. + - `"prompt_pixel_attention_mask"` (optional): Tensor for pixel attention masks, if available. + + padding_value (`int`): + The padding value to use for the concatenated completion sequences (`chosen_input_ids` and + `rejected_input_ids`). + + Returns: + `dict[str, torch.LongTensor]`: A dictionary containing: + + - `"prompt_input_ids"`: Concatenated prompt input IDs of shape `(2 * batch_size, prompt_length)`. + - `"completion_input_ids"`: Concatenated chosen and rejected completion input IDs of shape `(2 * + batch_size, max_completion_length)`. + - `"prompt_attention_mask"`: Concatenated prompt attention masks of shape `(2 * batch_size, + prompt_length)`. + - `"completion_attention_mask"`: Concatenated chosen and rejected attention masks of shape `(2 * + batch_size, max_completion_length)`. + - `"pixel_values"` (optional): Concatenated pixel values if `"prompt_pixel_values"` are present. + - `"pixel_attention_mask"` (optional): Concatenated pixel attention masks if + `"prompt_pixel_attention_mask"` are present. + + Notes: + The completion input IDs and attention masks are padded to the maximum completion length of the chosen or + rejected sequences. + """ + output = {} + + # For the prompt, the input_ids are the same for both the chosen and rejected responses + output["prompt_input_ids"] = torch.cat([batch["prompt_input_ids"], batch["prompt_input_ids"]], dim=0) + output["prompt_attention_mask"] = torch.cat( + [batch["prompt_attention_mask"], batch["prompt_attention_mask"]], dim=0 + ) + if "pixel_values" in batch: + output["pixel_values"] = torch.cat([batch["pixel_values"], batch["pixel_values"]], dim=0) + + if "pixel_attention_mask" in batch: + output["pixel_attention_mask"] = torch.cat( + [batch["pixel_attention_mask"], batch["pixel_attention_mask"]], dim=0 + ) + if "image_sizes" in batch: + output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0) + if "token_type_ids" in batch: + output["token_type_ids"] = torch.cat((batch["token_type_ids"], batch["token_type_ids"])) + + # Concatenate the chosen and rejected completions + max_completion_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + output["completion_input_ids"] = torch.cat( + ( + pad_to_length(batch["chosen_input_ids"], max_completion_length, pad_value=padding_value), + pad_to_length(batch["rejected_input_ids"], max_completion_length, pad_value=padding_value), + ), + ) + output["completion_attention_mask"] = torch.cat( + ( + pad_to_length(batch["chosen_attention_mask"], max_completion_length, pad_value=0), + pad_to_length(batch["rejected_attention_mask"], max_completion_length, pad_value=0), + ), + ) + + return output + + def dpo_loss( + self, + chosen_logps: torch.FloatTensor, + rejected_logps: torch.FloatTensor, + ref_chosen_logps: torch.FloatTensor, + ref_rejected_logps: torch.FloatTensor, + loss_type: str = "sigmoid", + model_output: dict[str, torch.FloatTensor] = None, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """ + Compute the DPO loss for a batch of policy and reference model log probabilities. + + Args: + chosen_logps (`torch.FloatTensor`): + Log probabilities of the model for the chosen responses. Shape: `(batch_size,)`. + rejected_logps (`torch.FloatTensor`): + Log probabilities of the model for the rejected responses. Shape: `(batch_size,)`. + ref_chosen_logps (`torch.FloatTensor`): + Log probabilities of the reference model for the chosen responses. Shape: `(batch_size,)`. + ref_rejected_logps (`torch.FloatTensor`): + Log probabilities of the reference model for the rejected responses. Shape: `(batch_size,)`. + loss_type (`str`, defaults to `"sigmoid"`): + The type of loss to compute. One of: + - `"sigmoid"`: Sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"hinge"`: Hinge loss on the normalized likelihood from the + [SLiC](https://huggingface.co/papers/2305.10425) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + - `"exo_pair"`: Pairwise EXO loss from the [EXO](https://huggingface.co/papers/2402.00856) paper. + - `"nca_pair"`: Pairwise NCA loss from the [NCA](https://huggingface.co/papers/2402.05369) paper. + - `"robust"`: Unbiased estimate of the DPO loss that is robust to preference noise from the [Robust + DPO](https://huggingface.co/papers/2403.00409) paper. + - `"bco_pair"`: Pairwise BCO loss from the [BCO](https://huggingface.co/papers/2404.04656) paper. + - `"sppo_hard"`: SPPO loss with hard label from the [SPPO](https://huggingface.co/papers/2405.00675) + paper. + - `"aot"`: AOT loss for paired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper. + - `"aot_pair"`: AOT loss for unpaired datasets from the [AOT](https://huggingface.co/papers/2406.05882) + paper. + - `"discopop"`: DiscoPOP (a.k.a Log-Ratio Modulated Loss, LRML) loss from the + [DiscoPOP](https://huggingface.co/papers/2406.08414) paper. + - `"apo_zero"`: APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + - `"apo_down"`: APO-down loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + - `"sft"`: Negative log-likelihood loss (standard supervised fine-tuning loss). + model_output (`dict[str, torch.FloatTensor]`, *optional*): + The output of the model's forward pass. This is used to compute auxiliary losses if enabled. + + Returns: + A tuple of three tensors: `(losses, chosen_rewards, rejected_rewards)`. The losses tensor contains the DPO + loss for each example in the batch. The `chosen_rewards` and `rejected_rewards` tensors contain the rewards + for the chosen and rejected responses, respectively. + """ + device = self.accelerator.device + + # Get the log ratios for the chosen and rejected responses + chosen_logratios = chosen_logps.to(device) - (not self.reference_free) * ref_chosen_logps.to(device) + rejected_logratios = rejected_logps.to(device) - (not self.reference_free) * ref_rejected_logps.to(device) + + if self.f_divergence_type == FDivergenceType.ALPHA_DIVERGENCE: + # The alpha-divergence formula: (1 - u^-alpha) / alpha + # The divergence difference between the chosen and rejected sample is: + # (1 - u[w]^-alpha) / alpha - (1 - u[l]^-alpha) / alpha + # = (u[l]^-alpha - u[w]^-alpha) / alpha + # where u[w] and u[l] are the policy/reference probability ratios + # for the chosen and rejected samples, respectively. + alpha_coef = FDivergenceConstants.ALPHA_DIVERGENCE_COEF_DEFAULT + if self.f_divergence_params and FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY in self.f_divergence_params: + alpha_coef = float(self.f_divergence_params[FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY]) + logits = (cap_exp(rejected_logratios * -alpha_coef) - cap_exp(chosen_logratios * -alpha_coef)) / alpha_coef + else: + logratios = chosen_logps - rejected_logps + if self.reference_free: + ref_logratios = torch.tensor([0], dtype=logratios.dtype, device=logratios.device) + else: + ref_logratios = ref_chosen_logps - ref_rejected_logps + + logratios = logratios.to(self.accelerator.device) + ref_logratios = ref_logratios.to(self.accelerator.device) + logits = logratios - ref_logratios + + if self.f_divergence_type == FDivergenceType.JS_DIVERGENCE: + # The js-divergence formula: log(2 * u / (1 + u)) + # The divergence difference between the chosen and rejected sample is: + # log(2 * u[w] / (1 + u[w])) - log(2 * u[l] / (1 + u[l])) + # = log(u[w]) - log(u[l]) - (log(1 + u[w]) - log(1 + u[l])) + # where u[w] and u[l] are the policy/reference probability ratios + # for the chosen and rejected samples, respectively. + logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios) + + # The beta is a temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. + # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the + # labels and calculates a conservative DPO loss. + if loss_type == "sigmoid": + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + + elif loss_type == "robust": + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + + F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) / (1 - 2 * self.label_smoothing) + + elif loss_type == "exo_pair": + # eqn (16) of the EXO paper: https://huggingface.co/papers/2402.00856 + import math + + if self.label_smoothing == 0: + self.label_smoothing = 1e-3 + losses = (self.beta * logits).sigmoid() * ( + F.logsigmoid(self.beta * logits) - math.log(1 - self.label_smoothing) + ) + (-self.beta * logits).sigmoid() * (F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing)) + + elif loss_type == "hinge": + losses = torch.relu(1 - self.beta * logits) + + elif loss_type == "ipo": + # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper. + losses = (logits - 1 / (2 * self.beta)) ** 2 + + elif loss_type == "bco_pair": + chosen_logratios = chosen_logps - ref_chosen_logps + rejected_logratios = rejected_logps - ref_rejected_logps + chosen_rewards = self.beta * chosen_logratios + rejected_rewards = self.beta * rejected_logratios + rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach() + self.running.update(rewards) + delta = self.running.mean + losses = -F.logsigmoid((self.beta * chosen_logratios) - delta) - F.logsigmoid( + -(self.beta * rejected_logratios - delta) + ) + + elif loss_type == "sppo_hard": + # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, + # estimated using the PairRM score. The probability calculation is conducted outside of the trainer class. + # The version described here is the hard probability version, where P in Equation (4.7) of Algorithm 1 is + # set to 1 for the winner and 0 for the loser. + a = chosen_logps - ref_chosen_logps + b = rejected_logps - ref_rejected_logps + losses = (a - 0.5 / self.beta) ** 2 + (b + 0.5 / self.beta) ** 2 + + elif loss_type == "nca_pair": + chosen_rewards = (chosen_logps - ref_chosen_logps) * self.beta + rejected_rewards = (rejected_logps - ref_rejected_logps) * self.beta + losses = ( + -F.logsigmoid(chosen_rewards) + - 0.5 * F.logsigmoid(-chosen_rewards) + - 0.5 * F.logsigmoid(-rejected_rewards) + ) + + elif loss_type == "aot_pair": + chosen_logratios = chosen_logps - ref_chosen_logps + rejected_logratios = rejected_logps - ref_rejected_logps + chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0) + rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0) + delta = chosen_logratios_sorted - rejected_logratios_sorted + losses = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif loss_type == "aot": + logratios = chosen_logps - rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logratios_sorted, _ = torch.sort(logratios, dim=0) + ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0) + delta = logratios_sorted - ref_logratios_sorted + losses = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif loss_type == "apo_zero": + # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + losses_chosen = 1 - F.sigmoid(self.beta * chosen_logratios) # Increase chosen likelihood + losses_rejected = F.sigmoid(self.beta * rejected_logratios) # Decrease rejected likelihood + losses = losses_chosen + losses_rejected + + elif loss_type == "apo_down": + # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are worse than your model's default output. + # Decrease chosen likelihood and decrease rejected likelihood more + losses_chosen = F.sigmoid(self.beta * chosen_logratios) + losses_rejected = 1 - F.sigmoid(self.beta * (chosen_logratios - rejected_logratios)) + losses = losses_chosen + losses_rejected + + elif loss_type == "discopop": + # Eqn (5) of the DiscoPOP paper (https://huggingface.co/papers/2406.08414) + # This loss was discovered with LLM discovery + logratios = chosen_logps - rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logits = logratios - ref_logratios + logits = logits * self.beta + # Modulate the mixing coefficient based on the log ratio magnitudes + log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau) + logistic_component = -F.logsigmoid(logits) + exp_component = torch.exp(-logits) + # Blend between logistic and exponential component based on log ratio modulation + losses = logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation + + elif loss_type == "sft": + # SFT loss is the negative log likelihood loss on chosen responses + # This acts as the generation loss component in MPO + sft_loss = model_output["nll_loss"] + # Create losses tensor with same shape as other losses (per-sample) + batch_size = chosen_logps.shape[0] + losses = sft_loss.expand(batch_size) + # For SFT, we don't have preference rewards, so use zeros + chosen_rewards = torch.zeros_like(chosen_logps) + rejected_rewards = torch.zeros_like(rejected_logps) + + else: + raise ValueError( + f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', " + "'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_pair', 'discopop', 'apo_zero', " + "'apo_down', 'sft']" + ) + + chosen_rewards = self.beta * (chosen_logps.to(device) - ref_chosen_logps.to(device)).detach() + rejected_rewards = self.beta * (rejected_logps.to(device) - ref_rejected_logps.to(device)).detach() + + return losses, chosen_rewards, rejected_rewards + + def _compute_loss_liger( + self, model: nn.Module, batch: dict[str, list | torch.LongTensor] + ) -> dict[str, torch.Tensor]: + unwrapped_model = self.accelerator.unwrap_model(model) + concatenated_batch = self.concatenated_inputs(batch, padding_value=self.pad_token_id) + + model_kwargs = {} + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + # Add the pixel values and attention masks for vision models + if "pixel_values" in concatenated_batch: + model_kwargs["pixel_values"] = concatenated_batch["pixel_values"] + if "pixel_attention_mask" in concatenated_batch: + model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"] + if "image_sizes" in concatenated_batch: + model_kwargs["image_sizes"] = concatenated_batch["image_sizes"] + + prompt_attention_mask = concatenated_batch["prompt_attention_mask"] + completion_attention_mask = concatenated_batch["completion_attention_mask"] + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + # 2. Prepare decoder inputs + decoder_input_ids = shift_tokens_right( + concatenated_batch["completion_input_ids"], + unwrapped_model.config.decoder_start_token_id, + ) + # 3. Get decoder outputs + decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + hidden_states = decoder_outputs.last_hidden_state + + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_encoder_outputs = unwrapped_ref_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_ref_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + elif not self.reference_free: + with self.null_ref_context(): + ref_encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + + labels = concatenated_batch["completion_input_ids"] + loss_mask = completion_attention_mask.bool() + else: + # For decoder-only models + input_ids = torch.cat( + (concatenated_batch["prompt_input_ids"], concatenated_batch["completion_input_ids"]), dim=1 + ) + attention_mask = torch.cat( + (concatenated_batch["prompt_attention_mask"], concatenated_batch["completion_attention_mask"]), + dim=1, + ) + # Mask the prompt but not the completion for the loss + loss_mask = torch.cat( + (torch.zeros_like(prompt_attention_mask), completion_attention_mask), + dim=1, + ) + + # Flush and truncate + if self.max_length is not None and self.max_length < attention_mask.size(1): + if self.truncation_mode == "keep_start": + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + attention_mask = attention_mask[:, : self.max_length] + input_ids = input_ids[:, : self.max_length] + loss_mask = loss_mask[:, : self.max_length] + elif self.truncation_mode == "keep_end": + # Flush right before truncating left, then flush left + # [[0, 0, x, x, x, x], -> [[0, 0, x, x], + # [0, x, x, x, 0, 0]] [0, x, x, x]] + attention_mask, input_ids, loss_mask = flush_right(attention_mask, input_ids, loss_mask) + input_ids = input_ids[:, -self.max_length :] + attention_mask = attention_mask[:, -self.max_length :] + loss_mask = loss_mask[:, -self.max_length :] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + else: + raise ValueError( + f"Unknown truncation mode: '{self.truncation_mode}'. Should be one of ['keep_end', " + "'keep_start']." + ) + else: + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + + # Add logits_to_keep optimization + if self.use_logits_to_keep: + first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min() + logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 + model_kwargs["logits_to_keep"] = logits_to_keep + + model_kwargs["output_hidden_states"] = True + + # Add padding-free training support + if self.padding_free: + input_ids = input_ids[attention_mask.bool()].unsqueeze(0) + loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0) + position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1 + model_kwargs["position_ids"] = position_ids + else: + model_kwargs["attention_mask"] = attention_mask + + # Get the base model outputs (before LM head) + if hasattr(unwrapped_model, "get_decoder") and unwrapped_model.get_decoder() is not None: + base_model = unwrapped_model.get_decoder() + else: + base_attr = getattr(unwrapped_model, "base_model_prefix", self.args.base_model_attribute_name) + base_model = getattr(unwrapped_model, base_attr, unwrapped_model) + + outputs = base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + hidden_states = outputs.last_hidden_state[:, :-1] + + # Get reference hidden states if needed + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + if hasattr(unwrapped_ref_model, "get_decoder") and unwrapped_ref_model.get_decoder() is not None: + ref_base_model = unwrapped_ref_model.get_decoder() + else: + ref_attr = getattr(unwrapped_ref_model, "base_model_prefix", self.args.base_model_attribute_name) + ref_base_model = getattr(unwrapped_ref_model, ref_attr, unwrapped_ref_model) + + ref_outputs = ref_base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + elif not self.reference_free: + if hasattr(unwrapped_model, "get_decoder") and unwrapped_model.get_decoder() is not None: + ref_base_model = unwrapped_model.get_decoder() + else: + ref_attr = getattr(unwrapped_model, "base_model_prefix", self.args.base_model_attribute_name) + ref_base_model = getattr(unwrapped_model, ref_attr, unwrapped_model) + with self.null_ref_context(): + ref_outputs = ref_base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + + masked_input_ids = torch.where(loss_mask != 0, input_ids, self.label_pad_token_id) + labels = masked_input_ids[:, 1:] # Shift right for casual LM + + # Get the LM head + lm_head = unwrapped_model.get_output_embeddings() + + # Get reference model weights if needed + ref_weight = None + ref_bias = None + if not self.reference_free: + if self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_lm_head = unwrapped_ref_model.get_output_embeddings() + else: + with self.null_ref_context(): + ref_lm_head = unwrapped_model.get_output_embeddings() + ref_weight = ref_lm_head.weight + ref_bias = ref_lm_head.bias if hasattr(ref_lm_head, "bias") else None + + # Compute loss using Liger kernel + loss_output = self.dpo_loss_fn( + lm_head.weight, + hidden_states, + labels, + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + ref_input=ref_hidden_states if not self.reference_free else None, + ref_weight=ref_weight if not self.reference_free else None, + ref_bias=ref_bias if not self.reference_free else None, + ) + ( + loss, + (chosen_logps, rejected_logps, chosen_logits_mean, rejected_logits_mean, nll_loss, *aux_outputs), + ) = loss_output + + output = { + "loss": loss, + "chosen_logps": chosen_logps, + "rejected_logps": rejected_logps, + "mean_chosen_logits": chosen_logits_mean, + "mean_rejected_logits": rejected_logits_mean, + "nll_loss": nll_loss, + "chosen_rewards": aux_outputs[0], + "rejected_rewards": aux_outputs[1], + } + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, list | torch.LongTensor], is_ref_model: bool = False + ) -> dict[str, torch.Tensor]: + """ + Runs the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + + Args: + model: + Model to run the forward pass on. + batch: + Batch of input data. + is_ref_model: + Whether this method is being called for the reference model. If `True`, length desensitization is not + applied. + """ + num_examples = batch["prompt_input_ids"].shape[0] + + concatenated_batch = self.concatenated_inputs(batch, padding_value=self.pad_token_id) + + model_kwargs = {"use_cache": False} + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + # Add the pixel values and attention masks for vision models + if "pixel_values" in concatenated_batch: + model_kwargs["pixel_values"] = concatenated_batch["pixel_values"] + if "pixel_attention_mask" in concatenated_batch: + model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"] + if "image_sizes" in concatenated_batch: + model_kwargs["image_sizes"] = concatenated_batch["image_sizes"] + + prompt_input_ids = concatenated_batch["prompt_input_ids"] + prompt_attention_mask = concatenated_batch["prompt_attention_mask"] + completion_input_ids = concatenated_batch["completion_input_ids"] + completion_attention_mask = concatenated_batch["completion_attention_mask"] + if self.is_encoder_decoder: + labels = completion_input_ids + labels[completion_attention_mask == 0] = self.label_pad_token_id + outputs = model( + input_ids=prompt_input_ids, + attention_mask=prompt_attention_mask, + labels=labels, # we need the labels for the logits to be returned + **model_kwargs, + ) + logits = outputs.logits + loss_mask = completion_attention_mask.bool() + else: + # Concatenate the prompt and completion inputs + input_ids = torch.cat((prompt_input_ids, completion_input_ids), dim=1) + attention_mask = torch.cat((prompt_attention_mask, completion_attention_mask), dim=1) + if "token_type_ids" in concatenated_batch: + prompt_token_type_ids = concatenated_batch["token_type_ids"] + token_type_ids = pad_to_length(prompt_token_type_ids, input_ids.shape[1], 0) + # Mask the prompt but not the completion for the loss + loss_mask = torch.cat( + (torch.zeros_like(prompt_attention_mask), completion_attention_mask), + dim=1, + ) + + # Flush and truncate + if self.max_length is not None and self.max_length < attention_mask.size(1): + if self.truncation_mode == "keep_start": + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + if "token_type_ids" in concatenated_batch: + attention_mask, input_ids, loss_mask, token_type_ids = flush_left( + attention_mask, input_ids, loss_mask, token_type_ids + ) + else: + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + attention_mask = attention_mask[:, : self.max_length] + input_ids = input_ids[:, : self.max_length] + loss_mask = loss_mask[:, : self.max_length] + elif self.truncation_mode == "keep_end": + # Flush right before truncating left, then flush left + # [[0, 0, x, x, x, x], -> [[0, 0, x, x], + # [0, x, x, x, 0, 0]] [0, x, x, x]] + if "token_type_ids" in concatenated_batch: + attention_mask, input_ids, loss_mask, token_type_ids = flush_left( + attention_mask, input_ids, loss_mask, token_type_ids + ) + token_type_ids = token_type_ids[:, -self.max_length :] + else: + attention_mask, input_ids, loss_mask = flush_right(attention_mask, input_ids, loss_mask) + input_ids = input_ids[:, -self.max_length :] + attention_mask = attention_mask[:, -self.max_length :] + loss_mask = loss_mask[:, -self.max_length :] + if "token_type_ids" in concatenated_batch: + attention_mask, input_ids, loss_mask, token_type_ids = flush_left( + attention_mask, input_ids, loss_mask, token_type_ids + ) + else: + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + else: + raise ValueError( + f"Unknown truncation mode: '{self.truncation_mode}'. Should be one of ['keep_end', " + "'keep_start']." + ) + else: + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + if "token_type_ids" in concatenated_batch: + attention_mask, input_ids, loss_mask, token_type_ids = flush_left( + attention_mask, input_ids, loss_mask, token_type_ids + ) + else: + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + + if "token_type_ids" in concatenated_batch: + model_kwargs["token_type_ids"] = token_type_ids + + if self.use_logits_to_keep: + # Compute logits_to_keep based on loss_mask pattern: + # [[0, 0, 0, x, x, x, x], + # [0, 0, 0, x, x, x, 0]] + # ^ start computing logits from here ([:, -(7-3+1):]) + first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min() + logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 # +1 for the first label + model_kwargs["logits_to_keep"] = logits_to_keep + + model_kwargs["output_hidden_states"] = True + + if self.padding_free: + # Flatten the input_ids, position_ids, and loss_mask + # input_ids = [[a, b, c, 0], -> input_ids = [[a, b, c, d, e, f, g]] + # [d, e, f, g]] position_ids = [[0, 1, 2, 0, 1, 2, 3]] + input_ids = input_ids[attention_mask.bool()].unsqueeze(0) + loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0) + position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1 + model_kwargs["position_ids"] = position_ids + else: + model_kwargs["attention_mask"] = attention_mask + + outputs = model(input_ids, **model_kwargs) + logits = outputs.logits + + # Offset the logits by one to align with the labels + labels = torch.roll(input_ids, shifts=-1, dims=1) + loss_mask = torch.roll(loss_mask, shifts=-1, dims=1).bool() + + if self.use_logits_to_keep: + # Align labels with logits + # logits: -, -, [x2, x3, x4, x5, x6] + # ^ --------- ^ after logits[:, :-1, :] + # labels: [y0, y1, y2, y3, y4, y5, y6] + # ^ --------- ^ with logits_to_keep=4, [:, -4:] + # loss_mask: [0, 0, 0, 1, 1, 1, 1] + labels = labels[:, -logits_to_keep:] + loss_mask = loss_mask[:, -logits_to_keep:] + + if logits.shape[:2] != labels.shape[:2]: + # for LLaVA, the returned logits include the image tokens (placed before the text tokens) + seq_len = labels.shape[1] + logits = logits[:, -seq_len:] + + # Compute the log probabilities of the labels + labels[~loss_mask] = 0 # dummy token; we'll ignore the losses on these tokens later + per_token_logps = selective_log_softmax(logits, labels) + per_token_logps[~loss_mask] = 0 + per_token_logps = torch.roll(per_token_logps, shifts=1, dims=1) + + if self.padding_free: + # Unflatten the per_token_logps (shape: [1, sum_seq_len] -> [batch_size, seq_len]) + batch_size, seq_len = attention_mask.shape + per_token_logps_ = torch.zeros( + batch_size, seq_len, device=outputs.logits.device, dtype=outputs.logits.dtype + ) + per_token_logps_[attention_mask.bool()] = per_token_logps + per_token_logps = per_token_logps_ + + all_logps = per_token_logps[:, 1:].sum(-1) + + output = {} + + if self.use_weighting: + with torch.no_grad(): + # Eq (2) of the WPO paper: https://huggingface.co/papers/2406.11827 + logprobs = F.log_softmax(logits, dim=-1) + weights_adjustment_factor = torch.logsumexp(2 * logprobs, dim=-1) # same as sum(probs**2) in log space + per_token_logps_adjusted = per_token_logps - weights_adjustment_factor + all_weights = (per_token_logps_adjusted * loss_mask).sum(-1) / loss_mask.sum(-1) + chosen_weights = all_weights[:num_examples] + rejected_weights = all_weights[num_examples:] + output["policy_weights"] = torch.clamp(torch.exp(chosen_weights + rejected_weights), max=1) + + if self.args.rpo_alpha is not None or "sft" in self.loss_type: + # Only use the chosen logits for the RPO loss or SFT loss + chosen_logits = logits[:num_examples, :-1] if not self.is_encoder_decoder else logits[:num_examples] + chosen_labels = labels[:num_examples, :-1] if not self.is_encoder_decoder else labels[:num_examples] + + # Compute the log probabilities of the labels + output["nll_loss"] = F.cross_entropy( + torch.flatten(chosen_logits, end_dim=1), torch.flatten(chosen_labels, end_dim=1), ignore_index=0 + ) + + if "ipo" in self.loss_type: + all_logps = all_logps / loss_mask.sum(-1) + + if self.args.ld_alpha is not None and not is_ref_model: + # Compute response lengths based on loss_mask + completion_lengths = loss_mask.sum(dim=1) + + chosen_lengths = completion_lengths[:num_examples] + rejected_lengths = completion_lengths[num_examples:] + public_lengths = torch.min(chosen_lengths, rejected_lengths) # l_p in the paper + public_lengths = torch.cat([public_lengths, public_lengths], dim=0) + + seq_len = per_token_logps.size(1) + position_ids = torch.arange(seq_len, device=per_token_logps.device).expand_as(per_token_logps) + + ld_mask = position_ids < public_lengths.unsqueeze(1) + mask = position_ids < completion_lengths.unsqueeze(1) + + front_mask = (ld_mask & mask).float() + rear_mask = (~ld_mask & mask).float() + front_logps = (per_token_logps * front_mask).sum(dim=1) + rear_logps = (per_token_logps * rear_mask).sum(dim=1) + + all_logps = front_logps + self.args.ld_alpha * rear_logps + + output["chosen_logps"] = all_logps[:num_examples] + output["rejected_logps"] = all_logps[num_examples:] + + # Compute the mean logits + if self.padding_free: + # position_ids contains a sequence of range identifiers (e.g., [[0, 1, 2, 0, 1, 2, 3, ...]]). + # There are 2*num_examples ranges in total: the first half corresponds to the chosen tokens, + # and the second half to the rejected tokens. + # To find the start of the rejected tokens, we look for the num_examples+1-th zero in pos_id. + split_idx = (position_ids == 0).nonzero(as_tuple=True)[1][num_examples] + mean_chosen_logits = logits[0, :split_idx][loss_mask[0, :split_idx]].mean() + mean_rejected_logits = logits[0, split_idx:][loss_mask[0, split_idx:]].mean() + else: + mean_chosen_logits = logits[:num_examples][loss_mask[:num_examples]].mean() + mean_rejected_logits = logits[num_examples:][loss_mask[num_examples:]].mean() + + output["mean_chosen_logits"] = mean_chosen_logits + output["mean_rejected_logits"] = mean_rejected_logits + + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def get_batch_loss_metrics( + self, + model: PreTrainedModel | nn.Module, + batch: dict[str, list | torch.LongTensor], + train_eval: Literal["train", "eval"] = "train", + ) -> tuple[torch.Tensor, dict[str, float]]: + """Compute the DPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + if self.args.use_liger_kernel: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + chosen_rewards = model_output["chosen_rewards"] + rejected_rewards = model_output["rejected_rewards"] + else: + model_output = self.concatenated_forward(model, batch) + + # if ref_chosen_logps and ref_rejected_logps in batch use them, otherwise use the reference model + if "ref_chosen_logps" in batch and "ref_rejected_logps" in batch: + ref_chosen_logps = batch["ref_chosen_logps"] + ref_rejected_logps = batch["ref_rejected_logps"] + else: + ref_chosen_logps, ref_rejected_logps = self.compute_ref_log_probs(batch) + + # Initialize combined losses + losses = 0 + chosen_rewards = 0 + rejected_rewards = 0 + + # Compute losses for each loss type + for idx, loss_type in enumerate(self.loss_type): + # Compute individual loss using standard DPO loss function + _losses, _chosen_rewards, _rejected_rewards = self.dpo_loss( + model_output["chosen_logps"], + model_output["rejected_logps"], + ref_chosen_logps, + ref_rejected_logps, + loss_type, + model_output, + ) + + # Add weighted contributions + weight = self.loss_weights[idx] if self.loss_weights else 1.0 + losses = losses + _losses * weight + chosen_rewards = chosen_rewards + _chosen_rewards * weight + rejected_rewards = rejected_rewards + _rejected_rewards * weight + + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + if self.args.rpo_alpha is not None: + losses = losses + self.args.rpo_alpha * model_output["nll_loss"] # RPO loss from V3 of the paper + + if self.use_weighting: + losses = losses * model_output["policy_weights"] + + if self.aux_loss_enabled: + losses = losses + self.aux_loss_coef * model_output["aux_loss"] + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() + metrics[f"{prefix}rewards/margins"] = ( + self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item() + ) + metrics[f"{prefix}logps/chosen"] = ( + self.accelerator.gather_for_metrics(model_output["chosen_logps"]).detach().mean().item() + ) + metrics[f"{prefix}logps/rejected"] = ( + self.accelerator.gather_for_metrics(model_output["rejected_logps"]).detach().mean().item() + ) + metrics[f"{prefix}logits/chosen"] = ( + self.accelerator.gather_for_metrics(model_output["mean_chosen_logits"]).detach().mean().item() + ) + metrics[f"{prefix}logits/rejected"] = ( + self.accelerator.gather_for_metrics(model_output["mean_rejected_logits"]).detach().mean().item() + ) + if self.args.rpo_alpha is not None or "sft" in self.loss_type: + metrics[f"{prefix}nll_loss"] = ( + self.accelerator.gather_for_metrics(model_output["nll_loss"]).detach().mean().item() + ) + if self.aux_loss_enabled: + metrics[f"{prefix}aux_loss"] = ( + self.accelerator.gather_for_metrics(model_output["aux_loss"]).detach().mean().item() + ) + + return losses.mean(), metrics + + def compute_loss( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs=False, + num_items_in_batch=None, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return loss, metrics + + return loss + + def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.pad_token_id, + ) + + # if ref_output in batch use that otherwise use the reference model + if "ref_output" in batch: + ref_output = batch["ref_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + ref_output = self.model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.pad_token_id, + ) + else: + ref_output = self.ref_model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + ref_output = pad_to_length(ref_output, self.max_length, self.pad_token_id) + ref_output_decoded = self.processing_class.batch_decode(ref_output, skip_special_tokens=True) + + return policy_output_decoded, ref_output_decoded + + def prediction_step( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + prediction_loss_only: bool, + ignore_keys: list[str] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return loss.detach(), None, None + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: bool | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip( + random_batch_dataset["prompt"], policy_output_decoded, ref_output_decoded, strict=True + ) + ], + ) + if "wandb" in self.args.report_to and self.accelerator.is_main_process: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + if "mlflow" in self.args.report_to and self.accelerator.is_main_process: + mlflow.log_table(data=table, artifact_file="game_log.json") + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float`, *optional*): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/gkd_config.py b/venv/lib/python3.10/site-packages/trl/trainer/gkd_config.py new file mode 100644 index 0000000000000000000000000000000000000000..f5a047945fb978a8da4e6d93541419ceffab4db4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/gkd_config.py @@ -0,0 +1,112 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + +from .sft_config import SFTConfig + + +@dataclass +class GKDConfig(SFTConfig): + """ + Configuration class for [`GKDTrainer`]. + + This class includes only the parameters that are specific to GKD training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. + + Args: + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + lmbda (`float`, *optional*, defaults to `0.5`): + Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy + student-generated outputs). + beta (`float`, *optional*, defaults to `0.5`): + Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When + beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. + max_new_tokens (`int`, *optional*, defaults to `128`): + Maximum number of tokens to generate per completion. + teacher_model_name_or_path (`str`, *optional*): + Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being + trained. + teacher_model_init_kwargs (`dict[str, Any]]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model + from a string. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + seq_kd (`bool`, *optional*, defaults to `False`): + Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on + teacher-generated output). + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["teacher_model_init_kwargs"] + + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + lmbda: float = field( + default=0.5, + metadata={ + "help": "Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy " + "student-generated outputs)." + }, + ) + beta: float = field( + default=0.5, + metadata={ + "help": "Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence " + "loss. When beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL " + "Divergence." + }, + ) + max_new_tokens: int = field( + default=128, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + teacher_model_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Model name or path of the teacher model. If `None`, the teacher model will be the same as the " + "model being trained." + }, + ) + teacher_model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "teacher model from a string." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropouts in `model`."}, + ) + seq_kd: bool = field( + default=False, + metadata={ + "help": "Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised " + "FT on teacher-generated output)." + }, + ) + + def __post_init__(self): + super().__post_init__() + # check lmbda and beta are in the range [0, 1] + if self.lmbda < 0.0 or self.lmbda > 1.0: + raise ValueError("lmbda must be in the range [0.0, 1.0].") + if self.beta < 0.0 or self.beta > 1.0: + raise ValueError("beta must be in the range [0.0, 1.0].") diff --git a/venv/lib/python3.10/site-packages/trl/trainer/gkd_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/gkd_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..8db258cfbb815414a267f4ac6137974a34ed255a --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/gkd_trainer.py @@ -0,0 +1,449 @@ +# Copyright 2020-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 os +import random +import textwrap +import warnings +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_liger_kernel_available, is_peft_available + +from ..models import prepare_deepspeed +from ..models.utils import unwrap_model_for_generation +from .gkd_config import GKDConfig +from .sft_trainer import SFTTrainer +from .utils import DataCollatorForChatML, disable_dropout_in_model, empty_cache + + +if is_peft_available(): + from peft import PeftConfig + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearJSDLoss + + +class GKDTrainer(SFTTrainer): + """Trainer for Generalized Knowledge Distillation (GKD) of language models. + + For details on GKD, see the paper: [On-Policy Distillation of Language Models: Learning from Self-Generated + Mistakes](https://huggingface.co/papers/2306.13649). + + Args: + model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `str`, *optional*): + Model to be trained, or the string identifier of the model to be instantiated from a pretrained model. + teacher_model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `str`, *optional*): + Teacher model for knowledge distillation, or the string identifier of the model to be instantiated from a + pretrained model. + args ([`GKDConfig`], *optional*): + Training arguments. + data_collator ([`~transformers.DataCollator`], *optional*): + Data collator to batch samples from the dataset. It defaults to a [`DataCollatorForChatML`] using the + `processing_class`. + train_dataset ([`~datasets.Dataset`], *optional*): + Dataset for training. + eval_dataset ([`~datasets.Dataset`] or `dict` of [`~datasets.Dataset`], *optional*): + Dataset for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Class to process the data. + compute_metrics (`Callable`, *optional*): + Function to compute metrics at evaluation. Must take in an [`~transformers.EvalPrediction`] and return a + dictionary string to float. + callbacks (`list` of [`~transformers.TrainerCallback`], *optional*): + Callbacks to use during training. + optimizers (`tuple` of `torch.optim.Optimizer` and `torch.optim.lr_scheduler.LambdaLR`, *optional*, defaults to `(None, None)`): + Tuple containing the optimizer and the learning rate scheduler to use for training. + preprocess_logits_for_metrics (`Callable`, *optional*): + Function to preprocess the logits before computing the metrics. Must take in the `logits` and `labels` and + return the logits to be used for metrics computation. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration to use PEFT for training. If `None`, PEFT is not used. If provided, the `model` will be + wrapped with the specified PEFT adapter. + formatting_func (`Callable`, *optional*): + Function to format the dataset. Must take in an example and return an example. + """ + + _tag_names = ["trl", "gkd"] + _name = "GKD" + _paper = { + "title": "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes", + "id": "2306.13649", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{agarwal2024on-policy, + title = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}}, + author = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem}, + year = 2024, + booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=3zKtaqxLhW}, + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str | None = None, + teacher_model: PreTrainedModel | nn.Module | str = None, + args: GKDConfig | None = None, + data_collator: DataCollator | None = None, # type: ignore + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + formatting_func: Callable | None = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + # Ensure Trainer does not drop non-signature columns used by the collator (e.g., "prompts") + args.remove_unused_columns = False + # Respect a user-provided data_collator; otherwise, provide a ChatML collator that + if data_collator is None: + data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + + # Ensure SFTTrainer does not pre-process the dataset when using a ChatML collator, + # so that raw conversational fields (e.g., "messages") remain available to the collator. + if args.dataset_kwargs is None: + args.dataset_kwargs = {"skip_prepare_dataset": True} + else: + args.dataset_kwargs["skip_prepare_dataset"] = True + + # Liger fused GKD loss (JSD) + self.use_liger_gkd_loss = False + if args.use_liger_kernel: + self.liger_jsd_loss = LigerFusedLinearJSDLoss( + beta=args.beta, + ignore_index=-100, + temperature=args.temperature, + compiled=False, + ) + self.use_liger_gkd_loss = True + + super().__init__( + model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + peft_config=peft_config, + formatting_func=formatting_func, + ) + + if args.teacher_model_init_kwargs is None: + teacher_model_init_kwargs = {} + elif not isinstance(teacher_model, str): + raise ValueError( + "You passed teacher_model_init_kwargs to the GKDConfig, but your teacher_model is already instantiated." + ) + else: + teacher_model_init_kwargs = args.teacher_model_init_kwargs + teacher_model_init_kwargs["dtype"] = ( + teacher_model_init_kwargs["dtype"] + if teacher_model_init_kwargs["dtype"] in ["auto", None] + else getattr(torch, teacher_model_init_kwargs["dtype"]) + ) + + if isinstance(teacher_model, str): + teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(self.model) + + if self.is_deepspeed_enabled: + self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) + else: + self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) + + self.lmbda = args.lmbda + self.beta = args.beta + self.temperature = args.temperature + self.seq_kd = args.seq_kd + + self.generation_config = GenerationConfig( + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + do_sample=True, + top_k=0, + use_cache=False if args.gradient_checkpointing else True, + pad_token_id=self.processing_class.pad_token_id, + ) + # Set custom EOS tokens if they are specified by the model's generation + # config. This is important for models with the Llama 3 chat template, + # which use special tokens <|eot_id|> and <|eom_id|> to mark the end of + # turns or messages. + if ( + hasattr(self.model.generation_config, "eos_token_id") + and self.model.generation_config.eos_token_id is not None + ): + self.generation_config.eos_token_id = self.model.generation_config.eos_token_id + + @staticmethod + def generalized_jsd_loss( + student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction="batchmean" + ): + """ + Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) + of https://huggingface.co/papers/2306.13649 for the definition. + + Args: + student_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + teacher_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + labels: + Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing + loss + beta: + Interpolation coefficient between 0 and 1 (default: 0.5) + temperature: + Softmax temperature (default: 1.0) + reduction: + Specifies the reduction to apply to the output (default: 'batchmean') + + Returns: + loss: Scalar tensor with the generalized JSD loss + """ + + # Apply temperature scaling + student_logits = student_logits / temperature + teacher_logits = teacher_logits / temperature + + # Compute log probabilities for student and probabilities for teacher + student_log_probs = F.log_softmax(student_logits, dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) + + if beta == 0: + jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) + elif beta == 1: + jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) + else: + # Compute the log of the mixture distribution + # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture + beta = torch.tensor(beta, dtype=student_log_probs.dtype) + mixture_log_probs = torch.logsumexp( + torch.stack([student_log_probs + torch.log(1 - beta), teacher_log_probs + torch.log(beta)]), + dim=0, + ) + + # Compute KL divergences using F.kl_div + # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. + kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) + kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) + + # Compute the Generalized Jensen-Shannon Divergence + jsd = beta * kl_teacher + (1 - beta) * kl_student + + # Masking + if labels is not None: + mask = labels != -100 + jsd = jsd[mask] + + # Apply reduction + if reduction == "batchmean": + return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0) + elif reduction == "sum": + return jsd.sum() + elif reduction == "mean": + return jsd.mean() + else: + return jsd + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if self.use_liger_gkd_loss: + # Forward only through the base models (avoid lm_head to save memory) + unwrapped_student = self.accelerator.unwrap_model(model) + if hasattr(unwrapped_student, "get_decoder") and unwrapped_student.get_decoder() is not None: + base_student = unwrapped_student.get_decoder() + else: + base_student = getattr( + unwrapped_student, getattr(unwrapped_student, "base_model_prefix", "model"), unwrapped_student + ) + + student_outputs = base_student( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + + self.teacher_model.eval() + unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model) + if hasattr(unwrapped_teacher, "get_decoder") and unwrapped_teacher.get_decoder() is not None: + base_teacher = unwrapped_teacher.get_decoder() + else: + base_teacher = getattr( + unwrapped_teacher, getattr(unwrapped_teacher, "base_model_prefix", "model"), unwrapped_teacher + ) + with torch.no_grad(): + teacher_outputs = base_teacher( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + + # hidden states (shifted) + student_hidden = student_outputs.last_hidden_state[:, :-1] + teacher_hidden = teacher_outputs.last_hidden_state[:, :-1] + + # Release full outputs to free memory + del student_outputs, teacher_outputs + + # labels mask and labels (shifted) + labels_mask = inputs["labels"] != -100 + masked_input_ids = torch.where( + labels_mask, inputs["input_ids"], torch.full_like(inputs["input_ids"], -100) + ) + true_labels = masked_input_ids[:, 1:].contiguous() + + # Release intermediate tensors + del labels_mask, masked_input_ids + + # heads + student_head = unwrapped_student.get_output_embeddings() + teacher_head = unwrapped_teacher.get_output_embeddings() + + # liger fused jsd loss + loss = self.liger_jsd_loss( + student_input=student_hidden, + student_weight=student_head.weight, + teacher_input=teacher_hidden, + teacher_weight=teacher_head.weight, + true_labels=true_labels, + student_bias=getattr(student_head, "bias", None), + teacher_bias=getattr(teacher_head, "bias", None), + ) + + # Release hidden states after loss computation + del student_hidden, teacher_hidden, true_labels + else: + # compute student output + student_outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # compute teacher output in eval mode + self.teacher_model.eval() + with torch.no_grad(): + teacher_outputs = self.teacher_model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # slice the logits for the generated tokens using the inputs["prompts"] lengths + prompt_lengths = inputs["prompts"].shape[1] + shifted_student_logits = student_outputs.logits[:, prompt_lengths - 1 : -1, :] + shifted_teacher_logits = teacher_outputs.logits[:, prompt_lengths - 1 : -1, :] + shifted_labels = inputs["labels"][:, prompt_lengths:] + + # compute loss + loss = self.generalized_jsd_loss( + student_logits=shifted_student_logits, + teacher_logits=shifted_teacher_logits, + labels=shifted_labels, + beta=self.beta, + ) + + # empty cache + empty_cache() + + # Return loss + return (loss, student_outputs) if return_outputs else loss + + @staticmethod + def generate_on_policy_outputs(model, inputs, generation_config, pad_token_id=None): + # Generate output with respect to the prompt-only + generated_outputs = model.generate( + input_ids=inputs["prompts"], + attention_mask=inputs.get("prompt_attention_mask", None), + generation_config=generation_config, + return_dict_in_generate=True, + ) + + # Get the generated token IDs + generated_tokens = generated_outputs.sequences + # Calculate new attention mask + new_attention_mask = torch.ones_like(generated_tokens) + new_labels = generated_tokens.clone() + + # If there's pad_token_id, set attention mask to 0 for padding tokens + if pad_token_id is not None: + new_labels[new_labels == pad_token_id] = -100 + new_attention_mask[generated_tokens == pad_token_id] = 0 + + return generated_tokens, new_attention_mask, new_labels + + def training_step( + self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + ) -> torch.Tensor: + """ + Perform a training step for the Generalized Knowledge Distillation (GKD) model. + + This method implements the on-policy learning approach described in the GKD paper. With probability + `self.lmbda`, it generates new responses using the student model, which are then used for training instead of + the original inputs. + """ + if self.seq_kd: + with unwrap_model_for_generation(self.teacher_model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + if random.random() <= self.lmbda: + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + + loss = super().training_step(model, inputs, num_items_in_batch) + return loss diff --git a/venv/lib/python3.10/site-packages/trl/trainer/grpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/grpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6001a8dc5240ab605aff45d1997c25e3b928f241 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/grpo_config.py @@ -0,0 +1,730 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass, field + +from transformers import TrainingArguments + + +@dataclass +class GRPOConfig(TrainingArguments): + r""" + Configuration class for the [`GRPOTrainer`]. + + This class includes only the parameters that are specific to GRPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`GRPOTrainer`] is provided as a string. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + cast_lm_head_to_fp32 (`bool`, *optional*, defaults to `False`): + Whether to cast the language modeling head of the policy and reference models to float32. As recommended by + the [ScaleRL](https://huggingface.co/papers/2510.13786) recipe. This flag is only supported when the model + has untied word embedding and language modeling head layers i.e. `tie_word_embeddings` in the model config + is False. + + > Parameters that control the data preprocessing + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left. + num_generations (`int` or `None`, *optional*, defaults to `8`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + + > Parameters that control generation + + generation_batch_size: (`int`, *optional*): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one + generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. + steps_per_generation: (`int`, *optional*): + Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive + with `generation_batch_size`. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + chat_template_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + use_transformers_paged (`bool`, *optional*, defaults to `False`): + Whether to use the `transformers` paged implementation for generation. If set to `True`, the `transformers` + paged implementation will be used for generation instead of the default padded implementation. This + parameter is only effective when `use_vllm` is set to `False`. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_guided_decoding_regex (`str`, *optional*): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step and woken + for weight sync and generation. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.0`): + KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving + training speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + delta (`float`, *optional*): + Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard + GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in + the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). + epsilon_high (`float`, *optional*): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + importance_sampling_level (`str`, *optional*, defaults to `"token"`): + Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"` + keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the + log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO + paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more + stable training and better alignment with sequence-level rewards. + reward_weights (`list[float]`, *optional*): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`): + Specifies the scaling strategy for rewards. Supported values are: + + - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring + unit variance within a group. + - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the + [PPO Lite paper](https://huggingface.co/papers/2508.08221). + - `False` or `"none"`: no scaling is applied. The [Dr. GRPO + paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the + standard deviation introduces a question-level difficulty bias. + loss_type (`str`, *optional*, defaults to `"dapo"`): + Specifies the loss formulation to use. Supported values are: + + - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to + length bias—this approach tends to prefer shorter completions with positive advantages and longer ones + with negative advantages. + - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was + introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias. + The value of the constant corresponds to `max_completion_length`. + - `"dapo"` (default): Aggregates token-level losses by normalizing with the number of active token in the + global accumulated batch. This method was introduced in the [DAPO + paper](https://huggingface.co/papers/2503.14476) to eliminate length bias. + - `"bnpo"`: Aggregates token-level losses by normalizing with the number of active token in the local + batch. Note that normalization is performed over the local batch only, so results may slightly vary + depending on the local batch size, despite a constant effective batch size. When using + `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + top_entropy_quantile (`float`, *optional*, defaults to `1.0`): + ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy + loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence + position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token; + `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with + `mask_truncated_completions=True`, only tokens from non-truncated completions are considered. + use_liger_loss (`bool`, *optional*): + Whether to use Liger loss. + + + + Parameter `use_liger_loss` is deprecated and will be removed in version 0.28.0. Use `use_liger_kernel` + instead. + + + + vllm_importance_sampling_correction (`bool`, *optional*, defaults to `True`): + Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and recomputed + logprobs. [Your Efficient RL Framework Secretly Brings You Off-Policy RL + Training](https://fengyao.notion.site/off-policy-rl) highlights that using a separate generation framework + (such as vLLM) can introduce off-policy effects due to subtle implementation differences between generation + and training backends. TIS is proposed as a remedy for this issue. + vllm_importance_sampling_cap (`float`, *optional*, defaults to `2.0`): + Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the importance + sampling ratio, improving training stability. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or + `trackio`. + num_completions_to_print (`int`, *optional*): + Number of completions to print with `rich`. If `None`, all completions are logged. + wandb_log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, all prompts + are logged. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: dict | str | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` " + "argument of the `GRPOTrainer` is provided as a string." + }, + ) + disable_dropout: bool = field( + default=False, + metadata={ + "help": "Whether to disable dropout in the model. This is useful for training with a reference model, as " + "it prevents the model from generating different logprobs for the same input." + }, + ) + cast_lm_head_to_fp32: bool = field( + default=False, + metadata={ + "help": "Whether to cast the language modeling head of the policy and reference, models to float32." + "As recommended by the [ScaleRL](https://huggingface.co/papers/2510.13786) recipe. This flag is only supported when the model" + " has untied word embedding and language modeling head layers i.e. `tie_word_embeddings` in the model config is False." + }, + ) + + # Parameters that control the data preprocessing + # The default value remove_unused_columns is overwritten from the parent class, because in GRPO we usually rely on + # additional columns to compute the reward + remove_unused_columns: bool | None = field( + default=False, + metadata={ + "help": "Whether to only keep the column 'prompt' in the dataset. If you use a custom reward function " + "that requires any column other than 'prompts' and 'completions', you should keep this to `False`." + }, + ) + max_prompt_length: int | None = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left." + }, + ) + num_generations: int | None = field( + default=8, + metadata={ + "help": "Number of generations to sample. The effective batch size (num_processes * per_device_batch_size " + "* gradient_accumulation_steps) must be evenly divisible by this value." + }, + ) + max_completion_length: int | None = field( + default=256, + metadata={"help": "Maximum length of the generated completion."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option " + "is not compatible with vLLM generation." + }, + ) + shuffle_dataset: bool | None = field( + default=True, + metadata={"help": "Whether to shuffle the training dataset."}, + ) + + # Parameters that control generation + generation_batch_size: int | None = field( + default=None, + metadata={ + "help": "Batch size to use for generation. If `None`, it defaults to the effective training batch size: " + "`per_device_train_batch_size * num_processes * steps_per_generation`." + }, + ) + steps_per_generation: int | None = field( + default=None, + metadata={"help": "Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`."}, + ) + temperature: float = field( + default=1.0, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=1.0, + metadata={ + "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. " + "Set to 1.0 to consider all tokens." + }, + ) + top_k: int | None = field( + default=None, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + min_p: float | None = field( + default=None, + metadata={ + "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It " + "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range." + }, + ) + generation_kwargs: dict | None = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or " + "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the " + "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that " + "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them." + }, + ) + chat_template_kwargs: dict | None = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to the `apply_chat_template` function when generating " + "completions." + }, + ) + repetition_penalty: float = field( + default=1.0, + metadata={ + "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated " + "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model " + "to repeat tokens." + }, + ) + use_transformers_paged: bool = field( + default=False, + metadata={ + "help": "Whether to use the `transformers` paged implementation for generation. If set to `True`, the " + "`transformers` paged implementation will be used for generation instead of the default padded " + "implementation. This parameter is only effective when `use_vllm` is set to `False`." + }, + ) + cache_implementation: str | None = field( + default=None, + metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."}, + ) + + # Parameters that control generation acceleration powered by vLLM + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for " + "generation instead of the default model.generate(). Requires `vllm` to be installed." + }, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `'server'` or " + "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure " + "a TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same " + "process and share the training GPUs. This avoids the need for a separate server but may cause resource " + "contention with training." + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={ + "help": "Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: " + "Use the `transformers` backend for model implementation. `vllm`: Use the `vllm` library for " + "model implementation." + }, + ) + vllm_enable_sleep_mode: bool = field( + default=False, + metadata={ + "help": "Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step " + "and woken for weight sync and generation." + }, + ) + vllm_guided_decoding_regex: str | None = field( + default=None, + metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, + ) + + # Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_base_url: str | None = field( + default=None, + metadata={ + "help": "Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` " + "and `vllm_server_port` are ignored." + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": "Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_port: int = field( + default=8000, + metadata={"help": "Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={ + "help": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up " + "after the timeout, a `ConnectionError` is raised." + }, + ) + + # Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + vllm_gpu_memory_utilization: float = field( + default=0.3, + metadata={ + "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag." + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={ + "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_tensor_parallel_size` flag." + }, + ) + + # Parameters that control the training + beta: float = field( + default=0.0, + metadata={ + "help": "KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and " + "improving training speed." + }, + ) + num_iterations: int = field( + default=1, + metadata={"help": "Number of iterations per batch (denoted as μ in the algorithm)."}, + ) + epsilon: float = field( + default=0.2, + metadata={"help": "Epsilon value for clipping."}, + ) + delta: float | None = field( + default=None, + metadata={ + "help": "Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` " + "(default), standard GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This " + "method is introduced in the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291)." + }, + ) + epsilon_high: float | None = field( + default=None, + metadata={ + "help": "Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the " + "lower-bound specified in argument `epsilon`. Paper DAPO recommends `0.28`." + }, + ) + importance_sampling_level: str = field( + default="token", + metadata={ + "help": "Controls whether importance sampling ratios are computed at the `'token'` or `'sequence'` level. " + "`'token'` keeps the raw per-token log-probability ratios (one weight per token). `'sequence'` averages " + "the log-probability ratios across valid tokens to produce a single ratio per sequence. The GSPO paper " + "shows that sequence-level sampling often yields more stable training and better alignment with " + "sequence-level rewards." + }, + ) + reward_weights: list[float] | None = field( + default=None, + metadata={ + "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all " + "rewards are weighted equally with weight `1.0`." + }, + ) + scale_rewards: str = field( + default="group", + metadata={ + "help": "Specifies the scaling strategy for rewards. Supported values are: " + "`True` or `group'` (default): rewards are scaled by the standard deviation within each group, ensuring " + "unit variance within a group. " + "`'batch'`: rewards are scaled by the standard deviation across the entire batch, as recommended in the " + "PPO Lite paper. " + "`False` or `'none'`: no scaling is applied. The Dr. GRPO paper recommends not scaling rewards, as " + "scaling by the standard deviation introduces a question-level difficulty bias." + }, + ) + loss_type: str = field( + default="dapo", + metadata={ + "help": "Specifies the loss formulation to use. Supported values are 'grpo', 'dapo', 'bnpo', and " + "'dr_grpo'. " + "'grpo': Aggregates token-level losses by normalizing over sequence length. Not recommended due to length " + "bias—this approach tends to prefer shorter completions with positive advantages and longer ones with " + "negative advantages. " + "'dapo' (default): Aggregates token-level losses by normalizing with the number of active token in the " + "global accumulated batch. This method was introduced in the DAPO paper to eliminate length bias. " + "'dr_grpo': Aggregates token-level losses by normalizing with a global constant. This method was " + "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to " + "`max_completion_length`. " + "'bnpo': Aggregates token-level losses by normalizing with the number of active token in the local batch. " + "Note that normalization is performed over the local batch only, so results may slightly vary depending " + "on the local batch size, despite a constant effective batch size. When using " + "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss." + }, + ) + mask_truncated_completions: bool = field( + default=False, + metadata={ + "help": "When enabled, truncated completions are excluded from the loss calculation, preventing them from " + "being incorrectly penalized and introducing noise during training. According to the DAPO paper, this is " + "a good practice for training stability." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + top_entropy_quantile: float = field( + default=1.0, + metadata={ + "help": "ρ parameter from Beyond the 80/20 Rule. Keeps in the policy loss term only the top-ρ quantile of " + "tokens by entropy of the probability distribution at each sequence position, improving results. Range: " + "[0.0-1.0]. A value of `0.0` masks all but the highest entropy token; `1.0` keeps all tokens. The paper " + "recommends a value of `0.2`. If used with `mask_truncated_completions=True`, only tokens from " + "non-truncated completions are considered." + }, + ) + use_liger_loss: bool = field( + default=None, + metadata={"help": "Whether to use the Liger GRPO loss."}, + ) + vllm_importance_sampling_correction: bool = field( + default=True, + metadata={ + "help": "Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and " + "recomputed logprobs. Your Efficient RL Framework Secretly Brings You Off-Policy RL " + "Training highlights that using a separate generation framework (such as vLLM) can introduce off-policy " + "effects due to subtle implementation differences between generation and training backends. TIS is " + "proposed as a remedy for this issue." + }, + ) + vllm_importance_sampling_cap: float = field( + default=2.0, + metadata={ + "help": "Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the " + "importance sampling ratio, improving training stability." + }, + ) + + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + num_completions_to_print: int | None = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_log_unique_prompts: bool | None = field( + default=False, + metadata={ + "help": "Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, " + "all prompts are logged." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + self.scale_rewards = {True: "group", False: "none"}.get(self.scale_rewards, self.scale_rewards) + + num_processes = self.world_size + # The current default effective batch size + if self.generation_batch_size is None and self.steps_per_generation is None: + self.steps_per_generation = self.gradient_accumulation_steps + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + elif self.generation_batch_size is not None and self.steps_per_generation is None: + # Just ensure the value is divisible by the global batch size + if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size " + f"({self.per_device_train_batch_size * num_processes})." + ) + self.steps_per_generation = self.generation_batch_size // ( + self.per_device_train_batch_size * num_processes + ) + elif self.generation_batch_size is None and self.steps_per_generation is not None: + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + else: + raise ValueError( + "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time" + ) + + if self.do_eval and self.eval_strategy != "no": + # Just ensure the value is divisible by the global batch size + if (self.per_device_eval_batch_size * num_processes) % self.num_generations != 0: + raise ValueError( + f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be " + f"divisible by num_generations ({self.num_generations})." + ) + + # The generation batch must contain full prompt groups (no partials), so it must be divisible by + # num_generations. + if self.generation_batch_size % self.num_generations != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations " + f"({self.num_generations})." + ) + + if self.num_generations < 2: + raise ValueError( + "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided " + f"{self.num_generations}, which is less than the minimum required." + ) + + if self.use_liger_loss is not None: + warnings.warn( + "The `use_liger_loss` argument is deprecated and will be removed in version 0.28.0. Please use " + "`use_liger_kernel` instead.", + FutureWarning, + stacklevel=2, + ) + self.use_liger_kernel = self.use_liger_loss + + if self.delta is not None and self.use_liger_kernel: + raise ValueError("Liger kernel does not support two-sided GRPO loss yet.") diff --git a/venv/lib/python3.10/site-packages/trl/trainer/grpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/grpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..fe29c78b92add42af50d0c320fe6853fc8a5a99e --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/grpo_trainer.py @@ -0,0 +1,1972 @@ +# Copyright 2020-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 inspect +import os +import textwrap +import warnings +from collections import defaultdict, deque +from collections.abc import Callable +from contextlib import nullcontext +from functools import partial +from pathlib import Path +from typing import Any + +import datasets +import pandas as pd +import torch +import torch.utils.data +import transformers +from accelerate import logging +from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed +from datasets import Dataset, IterableDataset +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.utils.data import DataLoader, Sampler +from transformers import ( + AutoConfig, + AutoModelForSequenceClassification, + AutoProcessor, + AutoTokenizer, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + is_trackio_available, + is_wandb_available, +) +from transformers.trainer_utils import seed_worker +from transformers.utils import is_datasets_available, is_peft_available, is_rich_available + +from ..data_utils import ( + apply_chat_template, + is_conversational, + prepare_multimodal_messages, + prepare_multimodal_messages_vllm, +) +from ..extras.profiling import profiling_context, profiling_decorator +from ..extras.vllm_client import VLLMClient +from ..import_utils import is_liger_kernel_available, is_vllm_available +from ..models import prepare_deepspeed, prepare_fsdp, prepare_peft_model, unwrap_model_for_generation +from ..models.utils import _ForwardRedirection +from .base_trainer import BaseTrainer +from .callbacks import SyncRefModelCallback +from .grpo_config import GRPOConfig +from .utils import ( + RepeatSampler, + disable_dropout_in_model, + ensure_master_addr_port, + entropy_from_logits, + get_config_model_id, + identity, + nanmax, + nanmin, + nanstd, + pad, + print_prompt_completions_sample, + selective_log_softmax, + shuffle_sequence_dict, + split_pixel_values_by_grid, + split_tensor_dict, + unsplit_pixel_values_by_grid, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + +if is_wandb_available(): + import wandb + +if is_trackio_available(): + import trackio + + +logger = logging.get_logger(__name__) + +# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of +# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. +RewardFunc = str | PreTrainedModel | Callable[[list, list], list[float]] + +# What we call a rollout function is a callable that takes prompts (list), args (GRPOConfig), and processing_class as +# parameters and returns a dict of generation results. Those results must include "prompt_ids", "completion_ids", and +# "logprobs" fields. Any extra fields (per-completion) are forwarded to the reward functions. +RolloutFunc = Callable[[list[str], Any, Any], dict[str, Any]] + + +class GRPOTrainer(BaseTrainer): + """ + Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the + paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language + Models](https://huggingface.co/papers/2402.03300). + + Example: + + ```python + from datasets import load_dataset + from trl import GRPOTrainer + + dataset = load_dataset("trl-lib/tldr", split="train") + + + def reward_func(completions, **kwargs): + # Dummy reward function that rewards completions with more unique letters. + return [float(len(set(completion))) for completion in completions] + + + trainer = GRPOTrainer( + model="Qwen/Qwen2-0.5B-Instruct", + reward_funcs=reward_func, + train_dataset=dataset, + ) + + trainer.train() + ``` + + Args: + model (`str | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + reward_funcs (`RewardFunc | list[RewardFunc]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can also return `None` when the reward is not applicable to those samples. This is useful + for multi-task training where different reward functions apply to different types of samples. When a + reward function returns `None` for a sample, that reward function is excluded from the reward + calculation for that sample. For more details, see [Using a custom reward + function](#using-a-custom-reward-function). + + The trainer's state is also passed to the reward function. The trainer's state is an instance of + [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the + reward function's signature. + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`GRPOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + rollout_func (`RolloutFunc`, *optional*): + Function to use for generating completions. It must take prompts, args, and processing_class as parameters + and return a dict with `"prompt_ids"`, `"completion_ids"`, and `"logprobs"` fields. Any other fields that + are forwarded to the reward functions. This feature is experimental and may change or be removed at any + time without prior notice. + """ + + _tag_names = ["trl", "grpo"] + _name = "GRPO" + _paper = { + "title": "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", + "id": "2402.03300", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{shao2024deepseekmath, + title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}}, + author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo}, + year = 2024, + eprint = {arXiv:2402.03300}, + } + """), + } + + def __init__( + self, + model: str | PreTrainedModel, + reward_funcs: RewardFunc | list[RewardFunc], + args: GRPOConfig | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + peft_config: "PeftConfig | None" = None, + rollout_func: RolloutFunc | None = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = GRPOConfig(f"{model_name}-GRPO") + + # Models + # Trained model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + dtype = model_init_kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass # dtype is already a torch.dtype or "auto" or None + elif isinstance(dtype, str): # it's a str, but not "auto" + dtype = getattr(torch, dtype) + model_init_kwargs["dtype"] = dtype + else: + raise ValueError( + "Invalid `dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + model = architecture.from_pretrained(model_id, **model_init_kwargs) + else: + model_id = get_config_model_id(model.config) + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + + # Some models (SmolVLM/Idefics3) don't support `logits_to_keep` argument and error out if we pass it + # Inspect the forward method before we wrap the model with PEFT + self.model_kwarg_keys = ( + inspect.signature(model.forward).parameters.keys() + if not hasattr(model, "get_base_model") + else inspect.signature(model.get_base_model().forward).parameters.keys() + ) + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(get_config_model_id(model.config), truncation_side="left") + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + self.pad_token = tokenizer.pad_token + self.pad_token_id = tokenizer.pad_token_id + self.eos_token_id = tokenizer.eos_token_id + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " + f"reward functions ({len(reward_funcs)})." + ) + + for i, (reward_processing_class, reward_func) in enumerate( + zip(reward_processing_classes, reward_funcs, strict=True) + ): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained(get_config_model_id(reward_func.config)) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + + self.reward_processing_classes = reward_processing_classes + + # Rollout function + if rollout_func is not None and os.environ.get("TRL_EXPERIMENTAL_SILENCE", "0") != "1": + warnings.warn( + "You are importing from 'rollout_func', which is an experimental feature. This API may change or be " + "removed at any time without prior notice. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1.", + UserWarning, + stacklevel=2, + ) + self.rollout_func = rollout_func + + # Training arguments + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper + self.num_generations = args.num_generations # = G in the GRPO paper + self.chat_template_kwargs = args.chat_template_kwargs or {} + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_paged = args.use_transformers_paged + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction + self.vllm_importance_sampling_cap = args.vllm_importance_sampling_cap + self.use_liger_kernel = args.use_liger_kernel + self.loss_type = args.loss_type + self.scale_rewards = args.scale_rewards + self.importance_sampling_level = args.importance_sampling_level + self.mask_truncated_completions = args.mask_truncated_completions + self.top_entropy_quantile = args.top_entropy_quantile + if self.use_liger_kernel and self.top_entropy_quantile < 1.0: + raise NotImplementedError( + "Liger Kernels don't currently support masking token positions based on entropy." + ) + if self.use_liger_kernel and not self.importance_sampling_level == "token": + raise NotImplementedError( + "Liger Kernels currently only support token-level importance sampling. Please set" + "`importance_sampling_level` to 'token'." + ) + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in GRPOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations (forward + backward passes), including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in GRPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning: + # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To + # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True. + # This acts as a flag to indicate that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in GRPO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + # In Trainer, `training_step` scales the loss by `gradient_accumulation_steps` only if `compute_loss_func` + # is None. For DAPO, loss scaling instead depends on the total number of completions tokens across the + # global accumulated batch. To control scaling ourselves, we must disable Trainer’s built-in scaling. The + # simplest (though a bit hacky) way is to set `compute_loss_func` to any non-None value, which bypasses + # that behavior without rewriting `training_step`. + compute_loss_func="non-None value to disable scaling", + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + self.ref_model = architecture.from_pretrained(model_id, **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Cast LM Head To FP32 + if args.cast_lm_head_to_fp32: + + def _cast_lm_head_to_fp32(target_model: PreTrainedModel): + """Cast lm_head to fp32 while preserving embedding output dtype if tied.""" + + def cast_inputs_to_fp32(module, inputs): + # Preserve other positional args and kwargs untouched + if not inputs: + return inputs + return (inputs[0].to(torch.float32),) + inputs[1:] + + original_dtype_local = target_model.lm_head.weight.dtype + target_model.lm_head = target_model.lm_head.float() + target_model.lm_head.register_forward_pre_hook(cast_inputs_to_fp32) + + if target_model.config.tie_word_embeddings: + + def cast_outputs_to_original_dtype(module, args, output): + return output.to(original_dtype_local) + + # Only cast activations; weights are now fp32 (intentional for numerical stability of logits) + target_model.model.embed_tokens.register_forward_hook(cast_outputs_to_original_dtype) + + _cast_lm_head_to_fp32(model) + if self.ref_model is not None: + _cast_lm_head_to_fp32(self.ref_model) + + # Liger loss + if self.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "Liger is required to use `use_liger_kernel` as the GRPO loss. Run `pip install liger-kernel`." + ) + # redirect the model.module forward to the model forward to ensure pre-forward hooks are called + self._forward_redirection = _ForwardRedirection() + + self.liger_grpo_loss = LigerFusedLinearGRPOLoss( + beta=self.beta, + epsilon_low=self.epsilon_low, + epsilon_high=self.epsilon_high, + temperature=self.temperature, + use_ref_model=self.beta != 0.0, + loss_type=self.loss_type, + max_completion_length=self.max_completion_length, + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # Keep logs sized to the generation batch to record only outputs from the latest model update. + self._logs = { + "images": deque(maxlen=args.generation_batch_size), + "prompt": deque(maxlen=args.generation_batch_size), + "completion": deque(maxlen=args.generation_batch_size), + "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + "advantages": deque(maxlen=args.generation_batch_size), + } + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install trl[vllm]` to use it." + ) + + if self.vllm_mode == "server": + if self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator(device=torch.cuda.current_device()) + + elif self.vllm_mode == "colocate": + # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have + # the same number of ranks + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks. + # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size)) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + # vLLM requires the environment variables to be set for distributed training. + os.environ["RANK"] = str(self.accelerator.process_index) + os.environ["LOCAL_RANK"] = str(self.accelerator.local_process_index) + os.environ["WORLD_SIZE"] = str(self.accelerator.num_processes) + # Ensure distributed rendezvous variables are set without colliding across concurrent runs + ensure_master_addr_port() + + if self.max_prompt_length is not None and self.max_completion_length is not None: + max_model_len = self.max_prompt_length + self.max_completion_length + else: + max_model_len = None + self.llm = LLM( + model=model.name_or_path, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + max_num_seqs=self.args.per_device_train_batch_size + * self.vllm_tensor_parallel_size + * self.args.steps_per_generation, + max_model_len=max_model_len, + distributed_executor_backend="external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + seed=self.accelerator.process_index // self.vllm_tensor_parallel_size, + # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768) - thinking there's not enough memory + max_num_batched_tokens=4096, + model_impl=self.args.vllm_model_impl, + enable_sleep_mode=self.args.vllm_enable_sleep_mode, + # Important so temperature scaling/logit tweaking affects the TIS log probs + logprobs_mode="processed_logprobs", + ) + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=2) + else: + raise ValueError(f"vllm_mode must be either 'server' or 'colocate', got '{self.vllm_mode}'.") + + # vLLM specific sampling arguments + self.guided_decoding_regex = args.vllm_guided_decoding_regex + + self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": tokenizer.pad_token_id, + "bos_token_id": tokenizer.bos_token_id, + "eos_token_id": tokenizer.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt", "image", "images"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. As a result, some parts of the method aren't relevant to GRPO, but we keep them to stay one line + # apart from the super method, ensuring easier maintenance in the future. + def get_train_dataloader(self): + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size * self.args.steps_per_generation, # < this is the change + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = partial( + seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index + ) + + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. The first row shows the first sampled batch, the + # second row shows the second sampled batch, and so on. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations, + seed=self.args.seed, + ) + + @profiling_decorator + def _get_last_hidden_state( + self, + unwrapped_model, + input_ids, + attention_mask, + logits_to_keep, + pixel_values=None, + image_grid_thw=None, + pixel_attention_mask=None, + image_sizes=None, + ): + if is_peft_model(unwrapped_model): + unwrapped_model = unwrapped_model.base_model.model + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask} + + # For Qwen models: + if image_grid_thw is not None and pixel_values is not None: + model_inputs["image_grid_thw"] = image_grid_thw + # For Gemma, SmolVLM2, LLaVa-Next etc.: + if pixel_values is not None: + model_inputs["pixel_values"] = pixel_values + # For SmolVLM2 + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask + # For LLaVa-Next + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + last_hidden_state = unwrapped_model.model(**model_inputs).last_hidden_state + # Exclude the last value: it corresponds to the next token pred + last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + return last_hidden_state + + def get_high_entropy_mask(self, entropies: torch.Tensor, mask: torch.Tensor, threshold: float) -> torch.Tensor: + """ + Returns a binary mask identifying tokens whose entropy exceeds a given quantile threshold. + + Args: + entropies (`torch.Tensor`): + Tensor of shape (batch_size, seq_len) with per-token entropy values. + mask (`torch.Tensor`): + Binary mask of the same shape as `entropies`, where `1` indicates valid tokens and `0` padding. + threshold (`float`): + Quantile threshold between `0.0` and `1.0` to select high-entropy tokens. + + Returns: + `torch.Tensor`: + Boolean mask of shape (batch_size, seq_len), where `True` indicates tokens with entropy >= threshold + and `False` otherwise. + """ + local = entropies[mask.bool()].float() + + # Use a negative pad_value as a sentinel because entropy values are always >= 0. + # This guarantees that the sentinel cannot collide with any real entropy value. + pad_value = -1e9 + + # Pad across processes so that every rank has the same tensor length + padded = self.accelerator.pad_across_processes(local, dim=0, pad_index=pad_value) + gathered = self.accelerator.gather(padded) + + # Drop sentinel values (safe because no entropy can be negative) + gathered = gathered[gathered != pad_value] + + if gathered.numel() == 0: + return torch.zeros_like(entropies, dtype=torch.bool) + + entropy_threshold = torch.quantile(gathered, threshold) + masked_entropies = entropies * mask.float() + entropy_mask = masked_entropies >= entropy_threshold + return entropy_mask & mask.bool() # ensure padding tokens are always masked out + + @profiling_decorator + def _get_per_token_logps_and_entropies( + self, + model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=None, + compute_entropy=False, + pixel_values=None, + image_grid_thw=None, + num_images=None, + pixel_attention_mask=None, + image_sizes=None, + token_type_ids=None, + ) -> dict[str, torch.Tensor | None]: + """Compute log-probs and (optionally) entropies for each token.""" + batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak + all_logps = [] + all_entropies = [] + for start in range(0, input_ids.size(0), batch_size): + input_ids_batch = input_ids[start : start + batch_size] + attention_mask_batch = attention_mask[start : start + batch_size] + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} + if image_grid_thw is not None and pixel_values is not None: + rows_per_image = image_grid_thw.prod(dim=-1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) + row_start, row_end = cum_rows[start].item(), cum_rows[start + batch_size].item() + model_inputs["pixel_values"] = pixel_values[row_start:row_end] + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] + model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] + elif pixel_values is not None: + model_inputs["pixel_values"] = pixel_values[start : start + batch_size] + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask[start : start + batch_size] + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes[start : start + batch_size] + if token_type_ids is not None: + model_inputs["token_type_ids"] = token_type_ids[start : start + batch_size] + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + logits = model(**model_inputs).logits + # Exclude the last value: it corresponds to the next token pred + logits = logits[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + logits = logits[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + # Divide logits by sampling temperature. + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + logits = logits / self.temperature + completion_ids = input_ids_batch[:, -logits_to_keep:] + logps = selective_log_softmax(logits, completion_ids) # compute logprobs + all_logps.append(logps) + + if compute_entropy: + with torch.no_grad(): + entropies = entropy_from_logits(logits) + all_entropies.append(entropies) + + logps = torch.cat(all_logps, dim=0) + entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None + return logps, entropies + + def _fix_param_name_to_vllm(self, name, extra_prefixes: list[str] | None = None): + extra_prefixes = extra_prefixes or [] + prefixes = ["_checkpoint_wrapped_module."] + extra_prefixes + for prefix in prefixes: + name = name.replace(prefix, "") + return name + + def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + # For FSDP1, we need to recurse into children and also use summon_full_params + if visited is None: + visited = set() + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp1_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + def _sync_fsdp2_params_to_vllm(self, module: nn.Module): + # For FSDP2, module.state_dict() already covers all parameters, so no need for recursion + for name, param in module.state_dict().items(): + if param.is_cpu: + param = param.to(torch.device("cuda")) + param = param.full_tensor() + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param)]) + + @profiling_decorator + def _move_model_to_vllm(self): + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + # TODO: does this work with FSDP? + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm( + self.model + ) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + for name, param in self.model.named_parameters(): + name = self._fix_param_name_to_vllm(name) + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.llm.reset_prefix_cache() + + @profiling_decorator + def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = split_pixel_values_by_grid(generation_batch) + generation_batch = shuffle_sequence_dict(generation_batch) + generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) + self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + self._step += 1 + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + # This allows for dynamic reward shaping based on training progress. + reward_kwargs["trainer_state"] = self.state + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) + ): + with profiling_context(self, reward_func_name): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] + texts = [ + apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] + for x in messages + ] + else: + texts = [p + c for p, c in zip(prompts, completions, strict=True)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = { + key: value[nan_row_idx] for key, value in reward_kwargs.items() if key != "trainer_state" + } + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + logger.warning( + f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _generate_single_turn(self, prompts: list): + device = self.accelerator.device + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + if self.vllm_mode == "colocate" and self.args.vllm_enable_sleep_mode: + # wake up colocated vLLM instances if needed + torch.cuda.empty_cache() # required to avoid OOM in some cases + self.llm.wake_up(tags=["weights"]) + + # First, update the vLLM weights if needed + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + if is_conversational({"prompt": prompts[0]}): + prompts = [prepare_multimodal_messages_vllm(prompt) for prompt in prompts] + + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process + if self.vllm_mode == "server": + all_prompts = gather_object(prompts) + + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts[:: self.num_generations] + + sampling_params = { + "n": self.num_generations, + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "truncate_prompt_tokens": self.max_prompt_length, + "guided_decoding_regex": self.guided_decoding_regex, + "generation_kwargs": self.args.generation_kwargs, + } + with profiling_context(self, "vLLM.generate"): + if self.rollout_func is not None: + if is_conversational({"prompt": ordered_set_of_prompts[0]}): + ordered_set_of_prompts = [ + apply_chat_template( + {"prompt": p}, self.processing_class, **self.chat_template_kwargs + )["prompt"] + for p in ordered_set_of_prompts + ] + output = self.rollout_func( + ordered_set_of_prompts, + self.args, + self.processing_class, + ) + else: + if is_conversational({"prompt": ordered_set_of_prompts[0]}): + output = self.vllm_client.chat( + messages=ordered_set_of_prompts, + **sampling_params, + chat_template_kwargs=self.chat_template_kwargs, + ) + else: + output = self.vllm_client.generate(prompts=ordered_set_of_prompts, **sampling_params) + # Extract required fields and collect any extra fields for reward functions + required_keys = {"prompt_ids", "completion_ids", "logprobs"} + extra_fields = {k: v for k, v in output.items() if k not in required_keys} + payload = (output["prompt_ids"], output["completion_ids"], output["logprobs"], extra_fields) + else: + payload = None + + # Broadcast the completions from the main process to all processes, ensuring each process receives its corresponding slice. + obj_list = [payload] + broadcast_object_list(obj_list, from_process=0) + all_prompt_ids, all_completion_ids, all_logprobs, all_extra_fields = obj_list[0] + + # At this point, we only get 1 copy of each prompt, so we need to repeat them num_generations times + all_prompt_ids = [ids for ids in all_prompt_ids for _ in range(self.num_generations)] + + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + prompt_ids = all_prompt_ids[process_slice] + completion_ids = all_completion_ids[process_slice] + logprobs = all_logprobs[process_slice] + + # Slice extra fields dict-of-lists per process (extra fields are per-completion, like completion_ids) + extra_fields = {} + for key, values in all_extra_fields.items(): + if isinstance(values, list): + extra_fields[key] = values[process_slice] + else: + extra_fields[key] = values + + # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts + elif self.vllm_mode == "colocate": + if self.guided_decoding_regex: + guided_decoding = GuidedDecodingParams(regex=self.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": 1, # vLLM on each GPU generates only 1 in colocate mode + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "truncate_prompt_tokens": self.max_prompt_length, + "guided_decoding": guided_decoding, + "logprobs": 0, # enable returning log probabilities; 0 means for the sampled tokens only + } + if self.args.generation_kwargs is not None: + generation_kwargs.update(self.args.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + if self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) + all_prompts = [p for sublist in gathered_prompts for p in sublist] + else: + all_prompts = prompts + + if self.args.vllm_enable_sleep_mode: + self.llm.wake_up(tags=["kv_cache"]) + + with profiling_context(self, "vLLM.generate"): + if is_conversational({"prompt": prompts[0]}): + all_outputs = self.llm.chat(all_prompts, sampling_params=sampling_params, use_tqdm=False) + else: + all_outputs = self.llm.generate(all_prompts, sampling_params=sampling_params, use_tqdm=False) + + all_prompt_ids = [output.prompt_token_ids for output in all_outputs] + all_completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + all_logprobs = [ + [next(iter(lp.values())).logprob for lp in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] + + if self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + prompt_ids = all_prompt_ids[tp_slice] + completion_ids = all_completion_ids[tp_slice] + logprobs = all_logprobs[tp_slice] + else: + prompt_ids = all_prompt_ids + completion_ids = all_completion_ids + logprobs = all_logprobs + + extra_fields = {} # No extra fields for colocate mode + + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=2) + + elif self.use_transformers_paged: + processor_kwargs = { + "max_length": self.max_prompt_length, + "truncation": True, + "add_special_tokens": False, + } + if is_conversational({"prompt": prompts[0]}): + processor_outputs = self.processing_class.apply_chat_template( + conversation=prompts, + **processor_kwargs, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **self.chat_template_kwargs, + ) + else: + processor_outputs = self.processing_class(text=prompts, **processor_kwargs) + + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + if self.args.cast_lm_head_to_fp32: + unwrapped_model.lm_head.to(torch.float32) + with torch.inference_mode(): + # Continuous batching API expects 'inputs' arg only + all_outputs = unwrapped_model.generate_batch( + processor_outputs["input_ids"], generation_config=self.generation_config, progress_bar=False + ) + unwrapped_model.train() # restore training mode, as generate_batch forces eval mode + completion_ids = [output.generated_tokens for output in all_outputs.values()] + prompt_ids = processor_outputs["input_ids"] + logprobs = None # not used in this case + extra_fields = {} # No extra fields for paged mode + + else: + # Regular generation path + processor_kwargs = { + "return_tensors": "pt", + "padding": True, + "padding_side": "left", + "max_length": self.max_prompt_length, + "truncation": True, + "add_special_tokens": False, + } + if is_conversational({"prompt": prompts[0]}): + generate_inputs = self.processing_class.apply_chat_template( + conversation=prompts, + **processor_kwargs, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **self.chat_template_kwargs, + ) + else: + generate_inputs = self.processing_class(text=prompts, **processor_kwargs) + generate_inputs = super()._prepare_inputs(generate_inputs) + + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + prompt_completion_ids = unwrapped_model.generate( + **generate_inputs, generation_config=self.generation_config, disable_compile=True + ) + # Compute prompt length and extract completion ids + prompt_ids, prompt_mask = generate_inputs["input_ids"], generate_inputs["attention_mask"] + prompt_length = prompt_ids.size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + prompt_ids = [p[m].tolist() for p, m in zip(prompt_ids, prompt_mask.bool(), strict=True)] + completion_ids = [c[m].tolist() for c, m in zip(completion_ids, completion_mask.bool(), strict=True)] + logprobs = None # not used in this case + extra_fields = {} # No extra fields for non-rollout_func paths + + return prompt_ids, completion_ids, logprobs, extra_fields + + def _generate(self, prompts: list): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompt_ids, completion_ids, logprobs, extra_fields = self._generate_single_turn(prompts) + + # Get completion length per sequence, used for logging + prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) + completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) + agg_prompt_lengths = self.accelerator.gather(prompt_lengths) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + total_prompt_tokens = agg_prompt_lengths.sum() + total_completion_tokens = agg_completion_lengths.sum() # = num_items_in_batch, required for the DAPO loss + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) + agg_is_truncated = self.accelerator.gather(is_truncated) + self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) + term_completion_lengths = agg_completion_lengths[~agg_is_truncated] + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + return prompt_ids, completion_ids, total_completion_tokens, logprobs, extra_fields + + def _generate_and_score_completions( + self, inputs: list[dict[str, torch.Tensor | Any]] + ) -> dict[str, torch.Tensor | Any]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] + if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ] + + prompt_ids_list, completion_ids_list, num_items_in_batch, sampling_per_token_logps_list, extra_fields = ( + self._generate(prompts) + ) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad(prompt_ids, padding_value=self.pad_token_id, padding_side="left") + prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left") + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right") + completion_mask = pad(completion_mask, padding_value=0, padding_side="right") + if sampling_per_token_logps_list is not None: + sampling_per_token_logps = [torch.tensor(logps, device=device) for logps in sampling_per_token_logps_list] + sampling_per_token_logps = pad(sampling_per_token_logps, padding_value=0.0, padding_side="right") + else: + sampling_per_token_logps = None + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + num_images = [len(img_list) for img_list in images] if images is not None else None + + # Get forward_kwargs for models with multimodal inputs + if images is not None: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] + for prompt in prompts + ] + prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + + with torch.no_grad(): + # If the generation and optimization steps are misaligned—i.e., if generation does not occur at the end of + # a full optimizer step (when gradient_accumulation_steps is not a multiple of generate_every)—then the + # samples may come from an earlier version of the model. In that case, we need to track old_per_token_logps + # for importance sampling. If the steps are aligned, importance sampling isn't necessary and we set + # old_per_token_logps to None. + # When using vLLM, we always compute old_per_token_logps for importance sampling, it was shown that the + # distribution mismatch between vLLM and the training model can be large and harm the training. + generate_every = self.args.steps_per_generation * self.num_iterations # generation frequency + if self.args.gradient_accumulation_steps % generate_every != 0 or ( + self.use_vllm and self.vllm_importance_sampling_correction + ): + old_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + old_per_token_logps = None + + # Compute the importance sampling ratio when using vLLM, to correct for potential distribution mismatch + if self.use_vllm and self.vllm_importance_sampling_correction: + importance_sampling_ratio = torch.exp(old_per_token_logps - sampling_per_token_logps) + importance_sampling_ratio = torch.clamp( + importance_sampling_ratio, max=self.vllm_importance_sampling_cap + ) + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + ref_per_token_logps = None + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text, strict=True): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Merge extra_fields from rollout_func into inputs for reward functions + if extra_fields: + for i, inp in enumerate(inputs): + for key, values in extra_fields.items(): + if isinstance(values, list) and i < len(values): + inp[key] = values[i] + elif not isinstance(values, list): + inp[key] = values + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + # Compute grouped-wise rewards + mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1) + + # Normalize the rewards to compute the advantages + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + advantages = rewards - mean_grouped_rewards + + if self.scale_rewards in ["group", "none"]: + # If self.scale_rewards = "none", we'll still log group level std + std_rewards = rewards.view(-1, self.num_generations).std(dim=1) + std_rewards = std_rewards.repeat_interleave(self.num_generations, dim=0) + elif self.scale_rewards == "batch": + # Compute global std + std_rewards = rewards.std().expand_as(rewards) + else: + raise ValueError( + f"Invalid value for scale_rewards: {self.scale_rewards}. Must be one of 'batch', 'group', or 'none'." + ) + + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) + if self.scale_rewards != "none": + advantages = advantages / (std_rewards + 1e-4) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._logs["prompt"].extend(gather_object(prompts_text)) + self._logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + if images is not None: + self._logs["images"].extend(gather_object(images)) + + if self.use_vllm and self.vllm_importance_sampling_correction: + delta = torch.abs(old_per_token_logps - sampling_per_token_logps) + delta = delta[completion_mask.bool()] + mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) + self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( + self.accelerator.gather(mean_delta).mean().item() + ) + self._metrics[mode]["sampling/sampling_logp_difference/max"].append( + self.accelerator.gather(max_delta).max().item() + ) + + flat_is_ratio = importance_sampling_ratio[completion_mask.bool()] + min_importance_sampling_ratio = ( + torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + mean_importance_sampling_ratio = ( + torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + max_importance_sampling_ratio = ( + torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) + ) + self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( + nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( + self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() + ) + self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( + nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + ) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "advantages": advantages, + "num_items_in_batch": num_items_in_batch, + } + if old_per_token_logps is not None: + output["old_per_token_logps"] = old_per_token_logps + if self.use_vllm and self.vllm_importance_sampling_correction: + output["importance_sampling_ratio"] = importance_sampling_ratio + if ref_per_token_logps is not None: + output["ref_per_token_logps"] = ref_per_token_logps + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if images is not None: + output["num_images"] = num_images + return output + + def compute_liger_loss(self, unwrapped_model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Get the last hidden state of the model + last_hidden_state = self._get_last_hidden_state( + unwrapped_model, + input_ids, + attention_mask, + logits_to_keep, + inputs.get("pixel_values"), + inputs.get("image_grid_thw"), + inputs.get("pixel_attention_mask"), + inputs.get("image_sizes"), + ) + + # compute loss and metrics using liger grpo loss + loss, metrics = self.liger_grpo_loss( + _input=last_hidden_state, + lin_weight=unwrapped_model.lm_head.weight, + selected_token_ids=completion_ids, + attention_mask=completion_mask, + advantages=inputs["advantages"], + bias=unwrapped_model.lm_head.bias, + old_per_token_logps=inputs.get("old_per_token_logps"), + ref_per_token_logps=inputs.get("ref_per_token_logps"), + ) + # Extract metrics from the liger_grpo_loss output + # KL divergence is the first metric when beta is non-zero + mean_kl = metrics[0] if self.beta != 0.0 else None + clip_ratio = metrics[-1] + + mode = "train" if self.model.training else "eval" + if self.beta != 0.0: + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item()) + self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item()) + return loss / self.current_gradient_accumulation_steps + + @profiling_decorator + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + if self.use_liger_kernel: + # Compute the loss using the liger grpo loss + unwrapped_model = self.accelerator.unwrap_model(model) + return self._forward_redirection(model, unwrapped_model, self.compute_liger_loss, unwrapped_model, inputs) + else: + return self._compute_loss(model, inputs) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + ) + + if self.top_entropy_quantile < 1.0: + entropy_mask = self.get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile) + else: + entropy_mask = None + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + + # Compute the loss + advantages = inputs["advantages"] + # When num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps, + # old_per_token_logps == per_token_logps. In this case we can skip its computation + # (see _generate_and_score_completions) and instead use per_token_logps.detach(). + # The exception is when using vLLM, where we always compute old_per_token_logps + # for importance sampling + old_per_token_logps = inputs.get("old_per_token_logps") + old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps + + log_ratio = per_token_logps - old_per_token_logps + if self.importance_sampling_level == "token": + log_importance_weights = log_ratio + elif self.importance_sampling_level == "sequence": + log_importance_weights = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + log_importance_weights = log_importance_weights.unsqueeze(-1) + else: + raise ValueError( + f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " + "and 'sequence'." + ) + # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on + # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) + + coef_1 = torch.exp(log_importance_weights) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + + # Two-sided clipping + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages.unsqueeze(1) + per_token_loss2 = coef_2 * advantages.unsqueeze(1) + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + if entropy_mask is not None: + per_token_loss = per_token_loss * entropy_mask + + if self.use_vllm and self.vllm_importance_sampling_correction: + per_token_loss = per_token_loss * inputs["importance_sampling_ratio"] + + if self.beta != 0.0: + per_token_loss = per_token_loss + self.beta * per_token_kl + + if self.loss_type == "grpo": + loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean() + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "bnpo": + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "dr_grpo": + loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.max_completion_length) + loss = loss / self.current_gradient_accumulation_steps + elif self.loss_type == "dapo": + normalizer = inputs["num_items_in_batch"] / self.accelerator.num_processes + loss = (per_token_loss * completion_mask).sum() / normalizer + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + # Log the metrics + mode = "train" if self.model.training else "eval" + + completion_token_count = completion_mask.sum().clamp(min=1.0) + + def masked_batch_mean(x): + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return x.mean() + else: + return (x * completion_mask).sum() / completion_token_count + + if self.beta != 0.0: + mean_kl = masked_batch_mean(per_token_kl) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + mean_entropy = masked_batch_mean(entropies) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages.unsqueeze(1) < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages.unsqueeze(1) > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = masked_batch_mean(is_low_clipped.float()) + high_clip = masked_batch_mean(is_high_clipped.float()) + clip_ratio = masked_batch_mean(is_region_clipped.float()) + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._logs["prompt"], + self._logs["completion"], + self._logs["rewards"], + self._logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + ) + + table = { + "step": [str(self.state.global_step)] * len(self._logs["prompt"]), + "prompt": self._logs["prompt"], + "completion": self._logs["completion"], + **self._logs["rewards"], + "advantage": self._logs["advantages"], + } + + df_base = pd.DataFrame(table) + images_raw = self._logs["images"] or [] + + for logging_backend in self.args.report_to: + if logging_backend == "wandb": + if images_raw: + images = [] + for image_list in self._logs["images"]: + images.append([wandb.Image(image) for image in image_list]) + df = pd.concat([df_base, pd.Series(images, name="image")], axis=1, copy=False) + else: + df = df_base + + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if logging_backend == "trackio": + if images_raw: + # TODO: Implement once supported upstream https://github.com/gradio-app/trackio/issues/334 + logger.info("Skipping image logging for Trackio") + df = df_base + # images = [] + # for image_list in self._logs["images"]: + # images.append([trackio.Image(image) for image in image_list]) + # df = pd.concat([df_base, pd.Series(images, name="image")], axis=1, copy=False) + else: + df = df_base + + trackio.log({"completions": trackio.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/judges.py b/venv/lib/python3.10/site-packages/trl/trainer/judges.py new file mode 100644 index 0000000000000000000000000000000000000000..819e76cf6b8fac1064c76aaf1cf003332b965452 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/judges.py @@ -0,0 +1,456 @@ +# Copyright 2020-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 concurrent.futures +import logging +from abc import ABC, abstractmethod + +import numpy as np +from accelerate import Accelerator +from huggingface_hub import InferenceClient +from transformers.utils import is_openai_available + +from ..import_utils import is_llm_blender_available + + +if is_llm_blender_available(): + import llm_blender + +if is_openai_available(): + from openai import OpenAI + + +DEFAULT_PAIRWISE_SYSTEM_PROMPT = '''I require a leaderboard for various large language models. I'll provide you with prompts given to these models and their corresponding outputs. Your task is to assess these responses, and select the model that produces the best output from a human perspective. + +## Instruction + +{{ + "instruction": """{prompt}""", +}} + +## Model Outputs + +Here are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier. + +{{ + {{ + "model_identifier": "0", + "output": """{response0}""" + }}, + {{ + "model_identifier": "1", + "output": """{response1}""" + }} +}} + +## Task + +Evaluate the models on the basis of the quality and relevance of their results, and select the model that generated the best result. Reply with the identifier of the best model. Our evaluation will only take into account the first character of your answer, so make sure it contains only one of the identifiers and nothing else (no quotation marks, no spaces, no new lines, ...). +''' + + +class BaseJudge(ABC): + """ + Base class for judges. The subclasses of this class should implement the `judge` method. + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[str], shuffle_order: bool = True) -> list: + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BaseRankJudge(ABC): + """ + Base class for LLM ranking judges. + + **Example**: + ```python + class MyRankJudge(BaseRankJudge): + def judge(self, prompts, completions, shuffle_order=True): + return ... # Your ranking logic here + + + judge = MyRankJudge() + judge.judge( + prompts=["The capital of France is", "The capital of Germany is"], + completions=[[" Paris", " Marseille", "Lyon"], [" Munich", " Berlin"]], + ) # [[0, 1, 2], [1, 0]] + ``` + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[list[int]]: + """ + Judge the completion for the given prompts and return the ranks of each completion. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[list[str]]`): + List of completions list, where each element is a list of completions for the corresponding prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + `list[list[int]]`: + List of lists of idxs, where each list contains the ranks of the completions for the corresponding + prompt. E.g., `[1, 2, 0]` means that the second completion (`idx=1`) is the best, followed by the + third, and then the first. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BasePairwiseJudge(BaseJudge): + """ + Base class for pairwise judges. + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + """ + Judge the completion pairs for the given prompts. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[list[str]]`): + List of completions pairs, where each element is a pair of completions for the corresponding prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + `list[int]`: + List of idxs, where each idx is the rank of the best completion for the corresponding prompt. E.g., `1` + means that the second completion (`idx=1`) is the best. + + Note: + If the judge returns `-1` for any prompt, it indicates that the inner process used to compute the + preference has failed. For instance, this could occur if the underlying language model returned an invalid + answer. In such cases, the caller should handle these invalid indices appropriately, possibly by + implementing fallback logic or error handling. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BaseBinaryJudge(BaseJudge): + """ + Base class for binary judges. + """ + + @abstractmethod + def judge( + self, + prompts: list[str], + completions: list[str], + gold_completions: list[str] | None = None, + shuffle_order: bool = True, + ) -> list[int]: + """ + Judge the completion for a given prompt. Used to assess if a completion satisfies a constraint. + + This base class should be used to implement binary evaluations as done in section 4.1.4 of the [CGPO + paper](https://huggingface.co/papers/2409.20370). It is relevant for assessing whether a prompt-completion pair + satisfies a specific constraint. + + Args: + prompts (`list[str]`): List of prompts. + completions (`list[str]`): List of completions. + gold_completions (`list[str]`, `optional`): List of gold completions if it exists. + shuffle_order (`bool`): Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + list[int]: A list of binary labels: + - 1 indicates that the completion satisfies the evaluated constraint. + - 0 indicates that the completion does not satisfy the evaluated constraint. + + Note: + If the judge returns -1 for any prompt, it indicates that the inner process used to compute the preference + has failed. For instance, this could occur if the underlying language model or rule based constraint + returned an invalid answer. In such cases, the caller should handle these invalid indices appropriately, + possibly by implementing fallback logic or error handling. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class PairRMJudge(BasePairwiseJudge): + # docstyle-ignore + """ + LLM judge based on the PairRM model from AllenAI. + + This judge uses the PairRM model to rank pairs of completions for given prompts. It's designed for pairwise + comparison of language model outputs. The PairRM model is loaded using the llm-blender library and runs on the + default Accelerator device. + + **Attributes**: + + blender (`llm_blender.Blender`): + An instance of the Blender class from llm-blender. + + **Example**: + ```python + >>> pairrm_judge = PairRMJudge() + >>> prompts = ["Translate 'hello' to French", "What's the capital of Japan?"] + >>> completions = [["Bonjour", "Salut"], ["Kyoto", "Tokyo"]] + >>> results = pairrm_judge.judge(prompts, completions) + >>> print(results) # [0, 1] (indicating the first completion is preferred for the first prompt and the second) + ``` + + > [!TIP] + > This class requires the llm-blender library to be installed. Install it with: `pip install llm-blender`. + """ + + def __init__(self): + if not is_llm_blender_available(): + raise ValueError("llm-blender is not installed. Please install it with `pip install llm-blender`.") + self.blender = llm_blender.Blender() + self.blender.loadranker("llm-blender/PairRM", device=Accelerator().device) + + def judge( + self, + prompts: list[str], + completions: list[list[str]], + shuffle_order: bool = True, + return_scores: bool = False, + temperature: float = 1.0, + ) -> list[int | float]: + """ + Judge the completion pairs for the given prompts using the PairRM model. + + Args: + prompts (`list[str]`): + List of prompts to judge. + completions (`list[list[str]]`): + List of completion pairs for each prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + return_scores (`bool`, *optional*, defaults to `False`): + If `True`, return probability scores of the first completion instead of ranks (i.e. a *soft-judge*). + temperature (`float`, *optional*, defaults to `1.0`): + Temperature for scaling logits if `return_scores` is True. + + Returns: + `list[int | float]`: + If `return_scores` is `False`, returns a list of ranks (`0` or `1`) for each prompt, indicating which + completion is preferred. If `return_scores` is `True`, returns softmax probabilities for the first + completion. + + Raises: + `ValueError`: + If the number of completions per prompt is not exactly 2. + + Note: + Unlike llm-blender, ranks are 0-indexed (`0` means the first completion is preferred). + """ + + if len(completions[0]) != 2: + raise ValueError("PairRM judge requires exactly 2 completions per prompt.") + + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions, strict=True)] + + # Rank the completions + ranks = self.blender.rank(prompts, completions, return_scores=return_scores, disable_tqdm=True) + if not return_scores: + ranks -= 1 # PairRM rank is 1-indexed, so we subtract 1 to make it 0-indexed + else: + # scale the logits by temperature + ranks /= temperature + + # Flip back the ranks or scores to the original order if needed + if shuffle_order: + ranks[flip_mask] = ranks[flip_mask][:, ::-1] + + # Return the ranks or score probability + if return_scores: + logit_max = np.amax(ranks, axis=-1, keepdims=True) + exp_logit_shifted = np.exp(ranks - logit_max) + probs = exp_logit_shifted / np.sum(exp_logit_shifted, axis=-1, keepdims=True) + return probs[:, 0].tolist() + else: + return ranks[:, 0].tolist() + + +class HfPairwiseJudge(BasePairwiseJudge): + """ + Pairwise judge based on the Hugging Face API with chat completion. + + This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt. + + Args: + model (`str`, *optional*, defaults to `"meta-llama/Meta-Llama-3-70B-Instruct"`): + Model to use for the judge. + token (`str`, *optional*): + Hugging Face API token to use for the [`huggingface_hub.InferenceClient`]. + system_prompt (`str`, *optional*): + The system prompt to be used for the judge. If not provided, a default prompt is used. Note that the system + prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`. Also, the + inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token + response. + """ + + def __init__( + self, + model="meta-llama/Meta-Llama-3-70B-Instruct", + token: str | None = None, + system_prompt: str | None = None, + ): + self.client = InferenceClient(model=model, token=token) + self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT + + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions, strict=True)] + + # Define a function to get the rank for a single prompt, will be called concurrently + def get_rank(prompt, candidates): + content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1]) + completion = self.client.chat_completion(messages=[{"role": "user", "content": content}], max_tokens=1) + response = completion.choices[0].message.content + if response in ["0", "1"]: + return int(response) + else: + logging.debug(f"Invalid response from the judge model: '{response}'. Returning -1.") + return -1 + + # Call the completions concurrently + with concurrent.futures.ThreadPoolExecutor() as executor: + ranks = list(executor.map(get_rank, prompts, completions)) + + # Flip back the ranks to the original order if needed + if shuffle_order: + ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)] + + # Return the ranks + return ranks + + +class OpenAIPairwiseJudge(BasePairwiseJudge): + """ + Judge based on the OpenAI API. + + This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt. + + Args: + model (`str`, *optional*, defaults to `"gpt-4-turbo-preview"`): + Model to use for the judge. + system_prompt (`str`, *optional*): + System prompt to be used for the judge. If not provided, a default prompt is used. Note that the system + prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`. Also, the + inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token + response. + max_requests (`int` or `None`, *optional*, defaults to `1000`): + Maximum number of requests to make to the OpenAI API. If set to `None`, there is no limit. + """ + + def __init__( + self, model="gpt-4-turbo-preview", system_prompt: str | None = None, max_requests: int | None = 1_000 + ): + if not is_openai_available(): + raise ValueError("OpenAI client is not installed. Please install it with 'pip install openai'.") + self.client = OpenAI() + self.model = model + self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT + self.max_requests = max_requests + self.num_requests = 0 + self._warned = False + + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + # Check if the limit of requests is reached, if so, use random choice instead + if self.max_requests is not None and self.num_requests >= self.max_requests: + if not self._warned: # Print the warning only once + logging.warning( + f"Reached the maximum number of requests ({self.max_requests}). From now on, returning -1 instead. " + " To increase the limit, set `max_requests` to a higher value, or to `None` for no limit." + ) + self._warned = True + return [-1] * len(prompts) + + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions, strict=True)] + + # Define a function to get the rank for a single prompt, will be called concurrently + def get_rank(prompt, candidates): + content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1]) + messages = [{"role": "user", "content": content}] + completion = self.client.chat.completions.create(model=self.model, messages=messages, max_tokens=1) + response = completion.choices[0].message.content + if response in ["0", "1"]: + return int(response) + else: + logging.debug(f"Invalid response from the judge model: '{response}'. Returning -1.") + return -1 + + # Call the completions concurrently + with concurrent.futures.ThreadPoolExecutor() as executor: + ranks = list(executor.map(get_rank, prompts, completions)) + + # Flip back the ranks to the original order if needed + if shuffle_order: + ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)] + + # Update the number of requests + self.num_requests += len(prompts) + + # Return the ranks + return ranks + + +class AllTrueJudge(BaseBinaryJudge): + """ + Unify the decision of multiple [`BaseBinaryJudge`] instances. + + Returns `1` only if all inner binary judges return `1`. If any judge returns `0`, it returns `0`. If any judge + returns `-1`, indicating a failure in its process, this judge will also return `-1`. + + Implements the Mixture of Judges as described in the [CGPO paper](https://huggingface.co/papers/2409.20370). + + Args: + judges (`list[BaseBinaryJudge]`): A list of [`BaseBinaryJudge`] instances whose decisions will be unified. + """ + + def __init__(self, judges: list[BaseBinaryJudge]): + self.judges = judges + + def judge( + self, + prompts: list[str], + completions: list[str], + gold_completions: list[str] | None = None, + shuffle_order: bool = True, + ) -> list[int]: + all_binary_judgments = [ + judge.judge(prompts, completions, gold_completions, shuffle_order) for judge in self.judges + ] + output = [] + for binary_judgments in zip(*all_binary_judgments, strict=True): + # Check that all values are in {0, 1, -1} + if any(binary_judgment not in {0, 1, -1} for binary_judgment in binary_judgments): + raise ValueError( + f"Invalid binary judgment: {binary_judgments}, expected list of values in {{0, 1, -1}}." + ) + + # Unify the decision + if -1 in binary_judgments: + output.append(-1) + elif all(binary_judgment == 1 for binary_judgment in binary_judgments): + output.append(1) + else: + output.append(0) + return output diff --git a/venv/lib/python3.10/site-packages/trl/trainer/kto_config.py b/venv/lib/python3.10/site-packages/trl/trainer/kto_config.py new file mode 100644 index 0000000000000000000000000000000000000000..96d8e64dfc9703cbde2dfaf21cbff1ece83a0108 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/kto_config.py @@ -0,0 +1,260 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class KTOConfig(TrainingArguments): + r""" + Configuration class for the [`KTOTrainer`]. + + This class includes only the parameters that are specific to KTO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int`, *optional*): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. + loss_type (`str`, *optional*, defaults to `"kto"`): + Type of loss to use. Possible values are: + + - `"kto"`: KTO loss from the [KTO](https://huggingface.co/papers/2402.01306) paper. + - `"apo_zero_unpaired"`: Unpaired variant of APO-zero loss from the + [APO](https://huggingface.co/papers/2408.06266) paper. + + desirable_weight (`float`, *optional*, defaults to `1.0`): + Desirable losses are weighed by this factor to counter unequal number of desirable and undesirable paris. + undesirable_weight (`float`, *optional*, defaults to `1.0`): + Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int`, *optional*): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from both the model and the reference model to W&B or Comet + during evaluation. + is_encoder_decoder (`bool`, *optional*): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute reference model log probabilities for training and evaluation datasets. This is + useful when training without the reference model to reduce the total GPU memory needed. + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + ref_model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model + from a string. + dataset_num_proc: (`int`, *optional*): + Number of processes to use for processing the dataset. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_liger_loss (`bool`, *optional*): + Whether to use Liger loss. + + + + Parameter `use_liger_loss` is deprecated and will be removed in version 0.28.0. Use `use_liger_kernel` + instead. + + + + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_kernel` is + `True`. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs", "ref_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: int | None = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: int | None = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: int | None = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model." + }, + ) + loss_type: str = field( + default="kto", + metadata={ + "help": "Type of loss to use.", + "choices": ["kto", "apo_zero_unpaired"], + }, + ) + desirable_weight: float = field( + default=1.0, + metadata={ + "help": "Desirable losses are weighed by this factor to counter unequal number of desirable and " + "undesirable pairs.", + }, + ) + undesirable_weight: float = field( + default=1.0, + metadata={ + "help": "Undesirable losses are weighed by this factor to counter unequal number of desirable and " + "undesirable pairs.", + }, + ) + label_pad_token_id: int = field( + default=-100, + metadata={ + "help": "Label pad token id. This argument is required if you want to use the default data collator." + }, + ) + padding_value: int | None = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={ + "help": "If `True`, generates and logs completions from both the model and the reference model to W&B " + "during evaluation." + }, + ) + is_encoder_decoder: bool | None = field( + default=None, + metadata={ + "help": "When using the `model_init` argument (callable) to instantiate the model instead of the `model` " + "argument, you need to specify if the model returned by the callable is an encoder-decoder model." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + precompute_ref_log_probs: bool = field( + default=False, + metadata={ + "help": "Whether to precompute reference model log probabilities for training and evaluation datasets. " + "This is useful when training without the reference model to reduce the total GPU memory needed." + }, + ) + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + ref_model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "reference model from a string." + }, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + use_liger_loss: bool = field( + default=None, + metadata={"help": "Whether to use Liger loss. It requires liger-kernel to be installed."}, + ) + base_model_attribute_name: str = field( + default="model", + metadata={ + "help": "Name of the attribute in the model that contains the base model. This is used to get the base " + "model from the model when the model does not have a `get_decoder` method in the case when " + "`use_liger_kernel` is `True`." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + if self.use_liger_loss is not None: + warnings.warn( + "The `use_liger_loss` argument is deprecated and will be removed in version 0.28.0. Please use " + "`use_liger_kernel` instead.", + FutureWarning, + stacklevel=2, + ) + self.use_liger_kernel = self.use_liger_loss + + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/kto_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/kto_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..8a7bd5f9a2f68a98129609b89672a4244e7e75fd --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/kto_trainer.py @@ -0,0 +1,1705 @@ +# Copyright 2020-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 inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from collections.abc import Callable +from contextlib import contextmanager, nullcontext +from operator import itemgetter +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState, logging +from accelerate.utils import tqdm +from datasets import Dataset, concatenate_datasets +from torch import autocast +from torch.utils.data import DataLoader, SequentialSampler +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + TrainingArguments, + is_comet_available, + is_wandb_available, +) +from transformers.trainer_utils import EvalLoopOutput, has_length +from transformers.utils import is_peft_available + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt, maybe_unpair_preference_dataset +from ..import_utils import is_liger_kernel_available +from ..models import create_reference_model, prepare_deepspeed +from .base_trainer import BaseTrainer +from .kto_config import KTOConfig +from .utils import ( + DPODataCollatorWithPadding, + disable_dropout_in_model, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearKTOLoss + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizer + + +logger = logging.get_logger(__name__) + +RUNNING_NAME = "running.pt" + + +def _get_kl_dataset(batch: dict[str, list[Any]]) -> dict[str, list[Any]]: + """ + Creates mismatched pairs of prompts and completions for the KL dataset by adding a +1 offset to the order of + completions. For best results, the mismatched outputs y' used to estimate the KL term for a batch should be the + same set as the matched outputs y used to estimate the rewards in that batch, just paired with different x. + """ + batch["answer_input_ids"] = [batch["answer_input_ids"][-1]] + batch["answer_input_ids"][:-1] + batch["answer_attention_mask"] = [batch["answer_attention_mask"][-1]] + batch["answer_attention_mask"][:-1] + return batch + + +def _tokenize( + batch: dict[str, list[Any]], + tokenizer: "PreTrainedTokenizer", +) -> dict[str, list[Any]]: + """Tokenize a batch from a KTO specific dataset.""" + prompt_tokenized = tokenizer(batch["prompt"], add_special_tokens=False) + prompt_input_ids = prompt_tokenized["input_ids"] + prompt_attention_mask = prompt_tokenized["attention_mask"] + prompt_and_completion = [ + prompt + completion for prompt, completion in zip(batch["prompt"], batch["completion"], strict=True) + ] + full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False) + full_input_ids = full_tokenized["input_ids"] + full_attention_mask = full_tokenized["attention_mask"] + + answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids, strict=True)] + answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask, strict=True)] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids, strict=True)] + # Prepare input tokens for token by token comparison + full_input_ids = [np.array(f) for f in full_input_ids] + for full, concat in zip(full_input_ids, full_concat_input_ids, strict=True): + if len(full) != len(concat): + raise ValueError( + "The elements in 'full_input_ids' and 'full_concat_input_ids' must have the same pairwise length." + ) + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = [len(p) for p in prompt_input_ids] + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx, strict=True)): + if not np.array_equal(p, f[:r]): + response_token_ids_start_idx[idx] -= 1 + + prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx, strict=True)] + prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx, strict=True)] + + for p, m in zip(prompt_input_ids, prompt_attention_mask, strict=True): + if len(p) != len(m): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx, strict=True)] + answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx, strict=True)] + + output = dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + answer_input_ids=answer_input_ids, + answer_attention_mask=answer_attention_mask, + ) + + return output + + +def _process_tokens(example: dict[str, Any], model: "PreTrainedModel" = None, **kwargs) -> dict: + """Process tokens of a KTO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + completion responses is/are too long. First we truncate the prompt; if we're still too long, we truncate the + completion. + + We also create the labels for the completion responses, which are of length equal to the sum of the length of the + prompt and the completion response, with label_pad_token_id for the prompt tokens. + """ + prompt = example["prompt"] + completion = example["completion"] + + batch = { + f"{kwargs['prefix']}prompt": prompt, + f"{kwargs['prefix']}completion": completion, + f"{kwargs['prefix']}label": example["label"], + } + + if not kwargs["is_encoder_decoder"]: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + + if not isinstance(completion, str): + raise ValueError(f"completion should be an str but got {type(completion)}") + + # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer + all_tokens = { + "prompt_input_ids": example["prompt_input_ids"], + "prompt_attention_mask": example["prompt_attention_mask"], + "answer_input_ids": example["answer_input_ids"], + "answer_attention_mask": example["answer_attention_mask"], + } + + # calculate max length by checking if BOS/EOS is already there + max_length = kwargs["max_length"] + bos_token_id = kwargs["tokenizer"].bos_token_id + eos_token_id = kwargs["tokenizer"].eos_token_id + if len(all_tokens["prompt_input_ids"]) > 0 and bos_token_id != all_tokens["prompt_input_ids"][0]: + max_length -= 1 + if len(all_tokens["answer_input_ids"]) > 0 and eos_token_id != all_tokens["answer_input_ids"][-1]: + max_length -= 1 + + # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the prompt + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["prompt_input_ids", "prompt_attention_mask"]: + if kwargs["truncation_mode"] == "keep_start": + all_tokens[k] = all_tokens[k][: kwargs["max_prompt_length"]] + elif kwargs["truncation_mode"] == "keep_end": + all_tokens[k] = all_tokens[k][-kwargs["max_prompt_length"] :] + else: + raise ValueError(f"Unknown truncation mode: {kwargs['truncation_mode']}") + + # if that's still too long, truncate the response + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["answer_input_ids", "answer_attention_mask"]: + all_tokens[k] = all_tokens[k][: max_length - kwargs["max_prompt_length"]] + + # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens + batch[f"{kwargs['prefix']}prompt_input_ids"] = all_tokens["prompt_input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = all_tokens["prompt_attention_mask"] + batch[f"{kwargs['prefix']}completion_input_ids"] = ( + all_tokens["prompt_input_ids"] + all_tokens["answer_input_ids"] + ) + batch[f"{kwargs['prefix']}completion_attention_mask"] = ( + all_tokens["prompt_attention_mask"] + all_tokens["answer_attention_mask"] + ) + + # add BOS, which affects both prompt and the full completion + if bos_token_id is not None: + if len(all_tokens["prompt_input_ids"]) == 0 or bos_token_id != all_tokens["prompt_input_ids"][0]: + batch[f"{kwargs['prefix']}prompt_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}prompt_input_ids" + ] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}prompt_attention_mask" + ] + batch[f"{kwargs['prefix']}completion_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}completion_input_ids" + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + # add EOS, which affects only the full completion + if len(all_tokens["answer_input_ids"]) == 0 or eos_token_id != all_tokens["answer_input_ids"][-1]: + batch[f"{kwargs['prefix']}completion_input_ids"] = batch[f"{kwargs['prefix']}completion_input_ids"] + [ + eos_token_id + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + [1] + + batch[f"{kwargs['prefix']}completion_labels"] = batch[f"{kwargs['prefix']}completion_input_ids"][:] + batch[f"{kwargs['prefix']}completion_labels"][: len(batch[f"{kwargs['prefix']}prompt_input_ids"])] = [ + kwargs["label_pad_token_id"] + ] * len(batch[f"{kwargs['prefix']}prompt_input_ids"]) + else: + completion_tokens = kwargs["tokenizer"]( + completion, truncation=True, max_length=kwargs["max_completion_length"], add_special_tokens=True + ) + prompt_tokens = kwargs["tokenizer"]( + prompt, truncation=True, max_length=kwargs["max_prompt_length"], add_special_tokens=True + ) + + batch[f"{kwargs['prefix']}prompt_input_ids"] = prompt_tokens["input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = prompt_tokens["attention_mask"] + + batch[f"{kwargs['prefix']}completion_labels"] = completion_tokens["input_ids"] + batch[f"{kwargs['prefix']}completion_attention_mask"] = completion_tokens["attention_mask"] + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch[f"{kwargs['prefix']}completion_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["completion_labels"]) + ) + + return batch + + +class KTOTrainer(BaseTrainer): + r""" + Initialize KTOTrainer. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. + ref_model ([`PreTrainedModelWrapper`]): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + args ([`KTOConfig`]): + The arguments to use for training. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + data_collator ([`~transformers.DataCollator`], *optional*): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + model_adapter_name (`str`, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + """ + + _tag_names = ["trl", "kto"] + _name = "KTO" + _paper = { + "title": "KTO: Model Alignment as Prospect Theoretic Optimization", + "id": "2402.01306", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{ethayarajh2024kto, + title = {{KTO: Model Alignment as Prospect Theoretic Optimization}}, + author = {Kawin Ethayarajh and Winnie Xu and Niklas Muennighoff and Dan Jurafsky and Douwe Kiela}, + year = 2024, + eprint = {arXiv:2402.01306}, + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str = None, + ref_model: PreTrainedModel | nn.Module | str | None = None, + args: KTOConfig = None, + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + data_collator: DataCollator | None = None, + model_init: Callable[[], PreTrainedModel] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: dict | None = None, + compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, + model_adapter_name: str | None = None, + ref_adapter_name: str | None = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if type(args) is TrainingArguments: + raise ValueError("Please use `KTOConfig` instead TrainingArguments.") + + if not isinstance(model, str) and ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must mass a copy of it, or `None` if you use peft." + ) + + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the KTOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + dtype = model_init_kwargs.get("dtype") + if dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(dtype, str) and dtype != "auto": + dtype = getattr(torch, dtype) + if dtype != "auto" and not isinstance(dtype, torch.dtype): + raise ValueError( + f"Invalid `dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." + ) + model_init_kwargs["dtype"] = dtype + + if args.ref_model_init_kwargs is None: + ref_model_init_kwargs = {} + elif not isinstance(ref_model, str): + raise ValueError( + "You passed ref_model_kwargs to the KTOTrainer. But your ref_model is already instantiated." + ) + else: + ref_model_init_kwargs = args.ref_model_init_kwargs + dtype = ref_model_init_kwargs.get("dtype") + if dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(dtype, str) and dtype != "auto": + dtype = getattr(torch, dtype) + if dtype != "auto" and not isinstance(dtype, torch.dtype): + raise ValueError( + f"Invalid `dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." + ) + ref_model_init_kwargs["dtype"] = dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + if isinstance(ref_model, str): + ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it with `pip install peft` to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = model_adapter_name + self.ref_adapter_name = ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model or args.precompute_ref_log_probs: + # The `model` with adapters turned off will be used as the reference model + self.ref_model = None + else: + self.ref_model = create_reference_model(model) + + if processing_class is None: + raise ValueError( + "max_length or a processing_class must be specified when using the default DPODataCollatorWithPadding" + ) + if args.max_length is None: + logger.warning( + "When using DPODataCollatorWithPadding, you should set `max_length` in the KTOTrainer's init" + " it will be set to `512` by default, but you should do it yourself in the future.", + ) + max_length = 512 + if args.max_length is not None: + max_length = args.max_length + + if args.max_prompt_length is None: + logger.warning( + "When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + ) + max_prompt_length = 128 + if args.max_prompt_length is not None: + max_prompt_length = args.max_prompt_length + + max_completion_length = None + if args.max_completion_length is None and self.is_encoder_decoder: + logger.warning( + "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + ) + max_completion_length = 128 + if args.max_completion_length is not None and self.is_encoder_decoder: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + logger.warning( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your KTOConfig" + " we have set it for you, but you should do it yourself in the future.", + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + self.loss_type = args.loss_type + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + self.precompute_ref_log_probs = args.precompute_ref_log_probs + + # Not all losses require a KL calculation + self.calculate_KL = True + if self.loss_type in ["apo_zero_unpaired"]: + self.calculate_KL = False + + # Since ref_logs are precomputed on the first call to get_train/eval_dataloader + # keep track of first called to avoid computation of future calls + self._precomputed_train_ref_log_probs = False + self._precomputed_eval_ref_log_probs = False + + # metric + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # KTO parameter + self.beta = args.beta + self.desirable_weight = args.desirable_weight + self.undesirable_weight = args.undesirable_weight + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + logger.warning( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + ) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in KTO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids" and "completion_input_ids". As a result, + # the trainer issues the warning: "Could not estimate the number of tokens of the input, floating-point + # operations will not be computed." To suppress this warning, we set the "estimate_tokens" key in the model's + # "warnings_issued" dictionary to True. This acts as a flag to indicate that the warning has already been + # issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed + train_dataset = train_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from train dataset" + ) + # Unpair the dataset if needed + train_dataset = maybe_unpair_preference_dataset( + train_dataset, args.dataset_num_proc, desc="Unpairing train dataset" + ) + # Apply the chat template if needed + train_dataset = train_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to train dataset", + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from eval dataset" + ) + eval_dataset = maybe_unpair_preference_dataset( + eval_dataset, args.dataset_num_proc, desc="Unpairing eval dataset" + ) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to eval dataset", + ) + + # Tokenize and prepare the training datasets + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs={"tokenizer": self.processing_class}, + num_proc=args.dataset_num_proc, + desc="Tokenizing train dataset", + ) + + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": self.processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + + train_dataset = train_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized train dataset", + ) + + # Tokenize and prepare the eval datasets + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs={"tokenizer": self.processing_class}, + batched=True, + num_proc=args.dataset_num_proc, + desc="Tokenizing eval dataset", + ) + + eval_dataset = eval_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized eval dataset", + ) + + # Get KL datasets if needed + if self.calculate_KL: + if args.per_device_train_batch_size <= 1: + raise ValueError( + "Actual (not effective) batch size must be > 1. KTO will not work properly because the KL term will be equivalent to the implied reward." + ) + + # create pairs for estimating the KL term by flipping the matched pairs in each batch of size total_batch_size + # i.e., (x_1, y_1), ..., (x_n, y_n) --> (x_1, y_n), ..., (x_n, y_1) = (x'_1, y'_1), ..., (x'_n, y'_n) + train_kl_dataset = train_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting KL train dataset", + ) + + fn_kwargs["prefix"] = "KL_" + train_kl_dataset = train_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in train_kl_dataset.column_names if c in train_dataset.column_names], + desc="Processing tokenized train KL dataset", + ) + + # merge the datasets + train_dataset = concatenate_datasets([train_dataset, train_kl_dataset], axis=1) + + if eval_dataset is not None: + # Get KL dataset + eval_kl_dataset = eval_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting eval KL dataset", + ) + + eval_kl_dataset = eval_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in eval_kl_dataset.column_names if c in eval_dataset.column_names], + desc="Processing tokenized eval KL dataset", + ) + + # merge the datasets + eval_dataset = concatenate_datasets([eval_dataset, eval_kl_dataset], axis=1) + + # calculate dataset desirability balance + num_desirable = max(sum(train_dataset["label"]), 1) + num_undesirable = max(len(train_dataset["label"]) - num_desirable, 1) # "label" is binary + + if num_desirable != num_undesirable: + # The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306 + des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2) + des_weight_upper_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2) + und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2) + und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2) + + des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound + und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound + + if not (des_weight_in_range or und_weight_in_range): + logger.warning( + "You have different amounts of desirable/positive and undesirable/negative examples but the " + "weights on the desirable and undesirable losses don't seem to be in an ideal range. Based " + f"on your data, we recommend EITHER " + f"desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}] or " + f"undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH). " + "See the documentation on how to optimally set these weights.", + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + # Deepspeed Zero-3 does not support precompute_ref_log_probs + if self.is_deepspeed_enabled: + if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." + ) + + if self.ref_model is None: + if not (self.is_peft_model or self.precompute_ref_log_probs): + raise ValueError( + "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + # Import Liger kernel if enabled + if self.args.use_liger_kernel: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_kernel=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if self.loss_type in ["apo_zero_unpaired"]: + raise ValueError( + "You cannot set `loss_type='apo_zero_unpaired'` with liger-kernel." + "Only KTO loss is supported with liger-kernel." + ) + if self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with liger kernel. Please set " + "`precompute_ref_log_probs=False`." + ) + if self.is_peft_model or self.ref_adapter_name is not None: + raise ValueError( + "You cannot use `use_liger_kernel=True` with Peft models. Please set `use_liger_kernel=False`." + ) + self.kto_loss_fn = LigerFusedLinearKTOLoss( + ignore_index=self.label_pad_token_id, beta=self.beta, use_ref_model=(self.ref_model is not None) + ) + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.set_adapter(self.model_adapter_name or "default") + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. + """ + + if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_train_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + self.train_dataset = self.train_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + if self.calculate_KL: + self.train_dataset = self.train_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + self._precomputed_train_ref_log_probs = True + + return super().get_train_dataloader() + + def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. + + Args: + eval_dataset (`torch.utils.data.Dataset`, *optional*): + If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted + by the `model.forward()` method are automatically removed. It must implement `__len__`. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset + + if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_eval_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) + + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + eval_dataset = eval_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + if self.calculate_KL: + eval_dataset = eval_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs + if self.eval_dataset is not None: + self.eval_dataset = eval_dataset + self._precomputed_eval_ref_log_probs = True + + return super().get_eval_dataloader(eval_dataset=eval_dataset) + + def compute_reference_log_probs(self, padded_batch: dict) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a KTO specific dataset.""" + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + if self.is_encoder_decoder: + completion_logits = self.model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.model( + padded_batch["completion_input_ids"], + attention_mask=padded_batch["completion_attention_mask"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + else: + if self.is_encoder_decoder: + completion_logits = self.ref_model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.ref_model( + padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + + completion_logps = self.get_batch_logps( + completion_logits, + padded_batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if self.calculate_KL: + KL_logps = self.get_batch_logps( + KL_logits, + padded_batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + else: + KL_logps = None + + return completion_logps, KL_logps + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: + Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: + The label value to ignore when computing log probabilities. + is_encoder_decoder: + Whether the model is an encoder-decoder model. If True, the labels are not shifted and the logits are + assumed to already be aligned with the labels. If False, the labels are shifted to the right by one + position, and the logits are assumed to be aligned with the shifted labels. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + else: + # Fixes end-dec RuntimeError + labels = labels.clone() + + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def forward( + self, model: nn.Module, batch: dict[str, list | torch.LongTensor] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + KL_logps = self._compute_kl_logps(model, batch) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + **model_kwargs, + ) + completion_logits = outputs.logits + + completion_logps = self.get_batch_logps( + completion_logits, + batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] + + chosen_logps = completion_logps[chosen_idx, ...] + rejected_logps = completion_logps[rejected_idx, ...] + + chosen_logits = completion_logits[chosen_idx, ...] + rejected_logits = completion_logits[rejected_idx, ...] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps, outputs.aux_loss) + else: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps) + + def kto_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + policy_KL_logps: torch.FloatTensor, + reference_chosen_logps: torch.FloatTensor, + reference_rejected_logps: torch.FloatTensor, + reference_KL_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the KTO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) + policy_KL_logps: Log probabilities of the policy model for the KL responses. Shape: (batch_size,) + reference_chosen_logps: + Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) + reference_rejected_logps: + Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in + batch_size,) + reference_KL_logps: Log probabilities of the reference model for the KL responses. Shape: (batch_size,) + + Returns: + A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, KL). The losses tensor contains the KTO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The KL tensor contains the detached KL divergence estimate + between the policy and reference models. + """ + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(policy_chosen_logps.device) + + # Chosen losses + if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0: + chosen_logratios = policy_chosen_logps - reference_chosen_logps + + if self.loss_type == "kto": + # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306) + chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl)) + elif self.loss_type == "apo_zero_unpaired": + # Unpaired variant of Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios) + + chosen_rewards = self.beta * chosen_logratios.detach() + + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + chosen_losses = torch.Tensor([]).to(self.accelerator.device) + chosen_rewards = torch.Tensor([]).to(self.accelerator.device) + + # Rejected losses + if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0: + rejected_logratios = policy_rejected_logps - reference_rejected_logps + + if self.loss_type == "kto": + rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios)) + elif self.loss_type == "apo_zero_unpaired": + rejected_losses = F.sigmoid(self.beta * rejected_logratios) + + rejected_rewards = self.beta * rejected_logratios.detach() + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + rejected_losses = torch.Tensor([]).to(self.accelerator.device) + rejected_rewards = torch.Tensor([]).to(self.accelerator.device) + + losses = torch.cat( + (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), + 0, + ) + + return losses, chosen_rewards, rejected_rewards, kl + + def _compute_kl_logps(self, model, batch): + """Compute KL log probabilities for a given batch.""" + KL_logps = None + if self.calculate_KL: + if self.is_encoder_decoder: + KL_model_kwargs = { + "input_ids": batch["KL_prompt_input_ids"], + "attention_mask": batch["KL_prompt_attention_mask"], + "labels": batch["KL_completion_labels"], + "decoder_input_ids": batch.get("KL_completion_decoder_input_ids"), + } + else: + KL_model_kwargs = { + "input_ids": batch["KL_completion_input_ids"], + "attention_mask": batch["KL_completion_attention_mask"], + } + + with torch.no_grad(): + KL_logits = model(**KL_model_kwargs).logits + + KL_logps = self.get_batch_logps( + KL_logits, + batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + return KL_logps + + def _compute_loss_liger(self, model, batch): + """ + Compute the KTO loss using the Liger-Kernel's LigerFusedLinearKTOLoss. + + Args: + model: + The policy model used for generating log probabilities and outputs. It could be an encoder-decoder + model or a regular language model. + batch: A dictionary containing the input data and labels for the batch. + + Returns: + A dictionary containing the following keys: + - "loss": The computed KTO loss for the batch. + - "chosen_logits_sum": Sum of the logits for the chosen responses from the policy model. + - "rejected_logits_sum": Sum of the logits for the rejected responses from the policy model. + - "chosen_logps": Log probabilities of the chosen responses from the policy model. + - "rejected_logps": Log probabilities of the rejected responses from the policy model. + - "chosen_rewards": Rewards for the chosen responses. + - "rejected_rewards": Rewards for the rejected responses. + - "kl": The KL divergence between the policy and reference models (detached). + + If auxiliary loss is enabled, the dictionary will also include: + - "aux_loss": The auxiliary loss from the model outputs. + """ + policy_KL_logps = self._compute_kl_logps(model, batch) + reference_KL_logps = self._compute_kl_logps(self.ref_model, batch) + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(self.accelerator.device) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get decoder outputs + outputs = model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + # 1. Get reference encoder outputs + ref_encoder_outputs = self.ref_model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get reference decoder outputs + ref_outputs = self.ref_model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + else: + # skip the lm head and get the last hidden state + if hasattr(model, "get_decoder") and model.get_decoder() is not None: + base_model = model.get_decoder() + else: + base_attr = getattr(model, "base_model_prefix", self.args.base_model_attribute_name) + base_model = getattr(model, base_attr, model) + outputs = base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + + # reference model + if hasattr(self.ref_model, "get_decoder") and self.ref_model.get_decoder() is not None: + ref_base_model = self.ref_model.get_decoder() + else: + ref_attr = getattr(self.ref_model, "base_model_prefix", self.args.base_model_attribute_name) + ref_base_model = getattr(self.ref_model, ref_attr, self.ref_model) + ref_outputs = ref_base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + lm_head = model.get_output_embeddings() + ref_lm_head = self.ref_model.get_output_embeddings() + + ( + loss, + ( + chosen_logps_sum, + rejected_logps_sum, + chosen_logits_sum, + rejected_logits_sum, + chosen_rewards_sum, + rejected_rewards_sum, + ), + ) = self.kto_loss_fn( + _input=outputs.last_hidden_state[:, :-1] if not self.is_encoder_decoder else outputs.last_hidden_state, + lin_weight=lm_head.weight, + target=batch["completion_labels"][:, 1:], + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + preference_labels=torch.tensor(batch["label"], dtype=torch.bool).to(self.accelerator.device), + ref_input=ref_outputs.last_hidden_state[:, :-1] + if not self.is_encoder_decoder + else outputs.last_hidden_state, + ref_weight=ref_lm_head.weight, + ref_bias=ref_lm_head.bias if hasattr(lm_head, "bias") else None, + kl=kl, + ) + + output = { + "loss": loss, + "chosen_logits_sum": chosen_logits_sum, + "rejected_logits_sum": rejected_logits_sum, + "chosen_logps_sum": chosen_logps_sum, + "rejected_logps_sum": rejected_logps_sum, + "chosen_rewards_sum": chosen_rewards_sum, + "rejected_rewards_sum": rejected_rewards_sum, + "kl": kl, + } + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, list | torch.LongTensor], + ): + """Compute the KTO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} + + labels = torch.tensor(batch["label"]) + num_chosen = labels.sum().to(self.accelerator.device) + num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) + + if self.args.use_liger_kernel: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + policy_chosen_logits = model_output["chosen_logits_sum"] + policy_rejected_logits = model_output["rejected_logits_sum"] + policy_chosen_logps = model_output["chosen_logps_sum"] + policy_rejected_logps = model_output["rejected_logps_sum"] + chosen_rewards = model_output["chosen_rewards_sum"] + rejected_rewards = model_output["rejected_rewards_sum"] + kl = model_output["kl"] + if self.aux_loss_enabled: + aux_loss = model_output["aux_loss"] + else: + forward_output = self.forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_KL_logps, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + # if reference_logps in batch use them, otherwise use the reference model + if "reference_logps" in batch: + chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] + + reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] + reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] + if self.calculate_KL: + reference_KL_logps = batch["reference_KL_logps"] + else: + reference_KL_logps = None + else: + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.model, batch)[:5] + else: + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.ref_model, batch)[:5] + + losses, chosen_rewards, rejected_rewards, kl = self.kto_loss( + policy_chosen_logps, + policy_rejected_logps, + policy_KL_logps, + reference_chosen_logps, + reference_rejected_logps, + reference_KL_logps, + ) + + metrics["kl"] = kl.item() + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + metrics["rewards/chosen_sum"] = ( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() + ) + metrics["logps/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() + ) + metrics["logits/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() + ) + metrics["count/chosen"] = all_num_chosen + + if all_num_rejected > 0: + metrics["rewards/rejected_sum"] = ( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() + ) + metrics["logps/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() + ) + metrics["logits/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() + ) + metrics["count/rejected"] = all_num_rejected + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs=False, + num_items_in_batch=None, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs) + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def _get_train_sampler(self, dataset: Dataset | None = None) -> torch.utils.data.Sampler | None: + if dataset is None: + dataset = self.train_dataset + if dataset is None or not has_length(dataset): + return None + return SequentialSampler(dataset) + + def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + # if reference_output in batch use that otherwise use the reference model + if "reference_output" in batch: + reference_output = batch["reference_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + reference_output = self.model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + else: + reference_output = self.ref_model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) + reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) + + return policy_output_decoded, reference_output_decoded + + def prediction_step( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + prediction_loss_only: bool, + ignore_keys: list[str] | None = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs) + + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = {} + if "logits/chosen_sum" in metrics: + logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] + if "logits/rejected_sum" in metrics: + logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: bool | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + target_labels = torch.tensor(random_batch["label"], dtype=torch.bool, device=self.accelerator.device) + target_indices = torch.where(~target_labels)[0] + target_batch = { + "prompt_input_ids": random_batch["prompt_input_ids"][target_indices], + "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indices], + "prompt": itemgetter(*target_indices)(random_batch["prompt"]), + } + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip( + target_batch["prompt"], policy_output_decoded, ref_output_decoded, strict=True + ) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float`, *optional*): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # train metrics should have no prefix, eval should have 'eval_' + prefix = "eval_" if train_eval == "eval" else "" + # accumulate average metrics from sums and lengths + for split in ["chosen", "rejected"]: + if f"count/{split}" in self._stored_metrics[train_eval]: + count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() + for metric in ["rewards", "logps", "logits"]: + logs[f"{prefix}{metric}/{split}"] = ( + torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() + / count_sum + ) + # delete obsolete metric + del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] + del self._stored_metrics[train_eval][f"count/{split}"] + # calculate reward margin + if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: + logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/model_config.py b/venv/lib/python3.10/site-packages/trl/trainer/model_config.py new file mode 100644 index 0000000000000000000000000000000000000000..9e3d5fe802148f40d4f64e96622fd6164fdb1a57 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/model_config.py @@ -0,0 +1,200 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass, field + + +@dataclass +class ModelConfig: + """ + Configuration class for the models. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + model_name_or_path (`str`, *optional*): + Model checkpoint for weights initialization. + model_revision (`str`, *optional*, defaults to `"main"`): + Specific model version to use. It can be a branch name, a tag name, or a commit id. + dtype (`Literal["auto", "bfloat16", "float16", "float32"]`, *optional*): + Override the default `torch.dtype` and load the model under this dtype. Possible values are + + - `"bfloat16"`: `torch.bfloat16` + - `"float16"`: `torch.float16` + - `"float32"`: `torch.float32` + - `"auto"`: Automatically derive the dtype from the model's weights. + + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow for custom models defined on the Hub in their own modeling files. This option should only + be set to `True` for repositories you trust and in which you have read the code, as it will execute code + present on the Hub on your local machine. + attn_implementation (`str`, *optional*): + Which attention implementation to use. You can run `--attn_implementation=flash_attention_2`, in which case + you must install this manually by running `pip install flash-attn --no-build-isolation`. + use_peft (`bool`, *optional*, defaults to `False`): + Whether to use PEFT for training. + lora_r (`int`, *optional*, defaults to `16`): + LoRA R value. + lora_alpha (`int`, *optional*, defaults to `32`): + LoRA alpha. + lora_dropout (`float`, *optional*, defaults to `0.05`): + LoRA dropout. + lora_target_modules (`str | list[str]`, *optional*): + LoRA target modules. + lora_target_parameters (`str | list[str]`, *optional*): + List of target parameters for LoRA. + lora_modules_to_save (`list[str]`, *optional*): + Model layers to unfreeze & train. + lora_task_type (`str`, *optional*, defaults to `"CAUSAL_LM"`): + Task type to pass for LoRA (use `"SEQ_CLS"` for reward modeling). + use_rslora (`bool`, *optional*, defaults to `False`): + Whether to use Rank-Stabilized LoRA, which sets the adapter scaling factor to `lora_alpha/√r`, instead of + the original default value of `lora_alpha/r`. + use_dora (`bool`, *optional*, defaults to `False`): + Enable [Weight-Decomposed Low-Rank Adaptation (DoRA)](https://huggingface.co/papers/2402.09353). This + technique decomposes the updates of the weights into two parts, magnitude and direction. Direction is + handled by normal LoRA, whereas the magnitude is handled by a separate learnable parameter. This can + improve the performance of LoRA, especially at low ranks. Right now, DoRA only supports linear and Conv2D + layers. DoRA introduces a bigger overhead than pure LoRA, so it is recommended to merge weights for + inference. + load_in_8bit (`bool`, *optional*, defaults to `False`): + Whether to use 8 bit precision for the base model. Works only with LoRA. + load_in_4bit (`bool`, *optional*, defaults to `False`): + Whether to use 4 bit precision for the base model. Works only with LoRA. + bnb_4bit_quant_type (`str`, *optional*, defaults to `"nf4"`): + Quantization type (`"fp4"` or `"nf4"`). + use_bnb_nested_quant (`bool`, *optional*, defaults to `False`): + Whether to use nested quantization. + """ + + model_name_or_path: str | None = field( + default=None, + metadata={"help": "Model checkpoint for weights initialization."}, + ) + model_revision: str = field( + default="main", + metadata={"help": "Specific model version to use. It can be a branch name, a tag name, or a commit id."}, + ) + dtype: str | None = field( + default=None, + metadata={ + "help": "Override the default `torch.dtype` and load the model under this dtype.", + "choices": ["auto", "bfloat16", "float16", "float32"], + }, + ) + trust_remote_code: bool = field( + default=False, + metadata={ + "help": "Whether to allow for custom models defined on the Hub in their own modeling files. This option " + "should only be set to `True` for repositories you trust and in which you have read the code, as it will " + "execute code present on the Hub on your local machine." + }, + ) + attn_implementation: str | None = field( + default=None, + metadata={ + "help": "Which attention implementation to use. You can run `--attn_implementation=flash_attention_2`, in " + "which case you must install this manually by running `pip install flash-attn --no-build-isolation`." + }, + ) + use_peft: bool = field( + default=False, + metadata={"help": "Whether to use PEFT for training."}, + ) + lora_r: int = field( + default=16, + metadata={"help": "LoRA R value."}, + ) + lora_alpha: int = field( + default=32, + metadata={"help": "LoRA alpha."}, + ) + lora_dropout: float = field( + default=0.05, + metadata={"help": "LoRA dropout."}, + ) + lora_target_modules: list[str] | None = field( + default=None, + metadata={"help": "LoRA target modules."}, + ) + lora_target_parameters: list[str] | None = field( + default=None, + metadata={"help": "List of target parameters for LoRA."}, + ) + lora_modules_to_save: list[str] | None = field( + default=None, + metadata={"help": "Model layers to unfreeze & train."}, + ) + lora_task_type: str = field( + default="CAUSAL_LM", + metadata={"help": "Task type to pass for LoRA (use 'SEQ_CLS' for reward modeling)."}, + ) + use_rslora: bool = field( + default=False, + metadata={ + "help": "Whether to use Rank-Stabilized LoRA, which sets the adapter scaling factor to `lora_alpha/√r`, " + "instead of the original default value of `lora_alpha/r`." + }, + ) + use_dora: bool = field( + default=False, + metadata={ + "help": "Enable Weight-Decomposed Low-Rank Adaptation (DoRA). This technique decomposes the updates of " + "the weights into two parts, magnitude and direction. Direction is handled by normal LoRA, whereas the " + "magnitude is handled by a separate learnable parameter. This can improve the performance of LoRA, " + "especially at low ranks. Right now, DoRA only supports linear and Conv2D layers. DoRA introduces a " + "bigger overhead than pure LoRA, so it is recommended to merge weights for inference." + }, + ) + load_in_8bit: bool = field( + default=False, + metadata={"help": "Whether to use 8 bit precision for the base model. Works only with LoRA."}, + ) + load_in_4bit: bool = field( + default=False, + metadata={"help": "Whether to use 4 bit precision for the base model. Works only with LoRA."}, + ) + bnb_4bit_quant_type: str = field( + default="nf4", + metadata={"help": "Quantization type.", "choices": ["fp4", "nf4"]}, + ) + use_bnb_nested_quant: bool = field( + default=False, + metadata={"help": "Whether to use nested quantization."}, + ) + # Deprecated params + torch_dtype: str | None = field( + default=None, + metadata={ + "help": "Override the default `torch.dtype` and load the model under this dtype.", + "choices": ["auto", "bfloat16", "float16", "float32"], + }, + ) + + def __post_init__(self): + if self.load_in_8bit and self.load_in_4bit: + raise ValueError("You can't use 8 bit and 4 bit precision at the same time") + + if self.torch_dtype and not self.dtype: + warnings.warn( + "`torch_dtype` is deprecated and will be removed in version 0.27.0, please use `dtype` instead.", + FutureWarning, + ) + self.dtype = self.torch_dtype + + if hasattr(self.lora_target_modules, "__len__") and len(self.lora_target_modules) == 1: + self.lora_target_modules = self.lora_target_modules[0] diff --git a/venv/lib/python3.10/site-packages/trl/trainer/nash_md_config.py b/venv/lib/python3.10/site-packages/trl/trainer/nash_md_config.py new file mode 100644 index 0000000000000000000000000000000000000000..07d8152f4fae45c9cf16b1815b3a225a9237c695 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/nash_md_config.py @@ -0,0 +1,46 @@ +# Copyright 2020-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 dataclasses import dataclass, field + +from trl.trainer.online_dpo_config import OnlineDPOConfig + + +@dataclass +class NashMDConfig(OnlineDPOConfig): + r""" + Configuration class for the [`NashMDTrainer`]. + + Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following: + + Parameters: + mixture_coef (`float` or `list[float]`, *optional*, defaults to `0.5`): + Logit mixture coefficient for the model and reference model. If a list of floats is provided then the + mixture coefficient is selected for each new epoch and the last coefficient is used for the rest of the + epochs. + """ + + mixture_coef: list[float] = field( + default_factory=lambda: [0.5], + metadata={ + "help": "Logit mixture coefficient for the model and reference model. If a list of floats is provided " + "then the mixture coefficient is selected for each new epoch and the last coefficient is used for the " + "rest of the epochs." + }, + ) + + def __post_init__(self): + super().__post_init__() + if hasattr(self.mixture_coef, "__len__") and len(self.mixture_coef) == 1: + self.mixture_coef = self.mixture_coef[0] diff --git a/venv/lib/python3.10/site-packages/trl/trainer/nash_md_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/nash_md_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..ea9bb74cf868bbffdb148eae61a34544c06b98f7 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/nash_md_trainer.py @@ -0,0 +1,495 @@ +# Copyright 2020-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 textwrap +from collections.abc import Callable +from typing import Any + +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset, IterableDataset +from transformers import ( + BaseImageProcessor, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, +) +from transformers.trainer_utils import EvalPrediction +from transformers.training_args import OptimizerNames +from transformers.utils import is_peft_available + +from ..data_utils import is_conversational, maybe_apply_chat_template +from ..models.modeling_base import GeometricMixtureWrapper +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .nash_md_config import NashMDConfig +from .online_dpo_trainer import OnlineDPOTrainer +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + empty_cache, + get_reward, + selective_log_softmax, + truncate_right, +) + + +if is_peft_available(): + from peft import PeftModel + + +class NashMDTrainer(OnlineDPOTrainer): + """ + Trainer for the Nash-MD method. + + It is implemented as a subclass of [`OnlineDPOTrainer`]. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model ([`PreTrainedModelWrapper`]): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + reward_funcs ([`~transformers.PreTrainedModel`]): + The reward model to score completions with, preferably an + [`~transformers.AutoModelForSequenceClassification`]. + judge ([`BasePairwiseJudge`]): + The judge to use for pairwise comparison of model completions. + args ([`NashMDConfig`]): + The NashMD config arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + peft_config (`dict`): + The peft config to use for training. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "nash-md"] + _name = "Nash-MD" + _paper = { + "title": "Nash Learning from Human Feedback", + "id": "2312.00886", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{munos2024nash, + title = {{Nash Learning from Human Feedback}}, + author = {R{\'{e}}mi Munos and Michal Valko and Daniele Calandriello and Mohammad Gheshlaghi Azar and Mark Rowland and Zhaohan Daniel Guo and Yunhao Tang and Matthieu Geist and Thomas Mesnard and C{\\^{o}}me Fiegel and Andrea Michi and Marco Selvi and Sertan Girgin and Nikola Momchev and Olivier Bachem and Daniel J. Mankowitz and Doina Precup and Bilal Piot}, + year = 2024, + booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=Y5AmNYiyCQ} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module = None, + ref_model: PreTrainedModel | nn.Module = None, + reward_funcs: PreTrainedModel | nn.Module | None = None, + judge: BasePairwiseJudge | None = None, + args: NashMDConfig | None = None, + data_collator: Callable | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + peft_config: dict | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + reward_funcs=reward_funcs, + judge=judge, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_classes=processing_class, + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + self._mixture_coef = self.args.mixture_coef + + # Overwrite the stats dictionary to include NashMD specific statistics + self.stats = { + # Remove "non_score_reward", "rlhf_reward", "scores_margin" + # Add "mixture_coef" + "loss/kl": [], + "objective/entropy": [], + "loss/score": [], + "rewards/probabilities": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + "val/model_contain_eos_token": [], + "val/ref_contain_eos_token": [], + "beta": [], + "mixture_coef": [], + } + if self.reward_funcs is not None: + if len(self.reward_funcs) != 1: + raise ValueError("NashMDTrainer only supports one reward function/model.") + self.reward_funcs = self.reward_funcs[0] + self.stats["rewards/chosen"] = [] + self.stats["rewards/rejected"] = [] + + @property + def mixture_coef(self): + if isinstance(self._mixture_coef, list): + epoch = self.state.epoch + return self._mixture_coef[epoch] if epoch < len(self._mixture_coef) else self._mixture_coef[-1] + else: + return self._mixture_coef + + def _generate_completions(self, model, prompts): + # Generate completions from the policy model. + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_for_gen_ctx: + model_output = unwrapped_policy_for_gen_ctx.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + # Get the DDP/FSDP unwrapped version of the main model. + # This will be the policy model for GeometricMixtureWrapper (PEFT adapters active if PEFT is used). + policy_model_for_gmw = self.accelerator.unwrap_model(model) + + # Determine the correct reference model for GeometricMixtureWrapper. + # This also needs to be DDP/FSDP unwrapped. + ref_model_for_gmw: torch.nn.Module + if self.ref_model is None: + # No explicit ref_model is provided. + # Use the base of the main `model` if it's a PEFT model. + # policy_model_for_gmw is already DDP-unwrapped. + if is_peft_available() and isinstance(policy_model_for_gmw, PeftModel): + ref_model_for_gmw = policy_model_for_gmw.get_base_model() + else: + # Not a PEFT model (or PEFT not available), or already a base model. + # Use the DDP-unwrapped policy model itself as the reference. + ref_model_for_gmw = policy_model_for_gmw + else: + # An explicit ref_model is provided. Unwrap it for DDP/FSDP. + ref_model_for_gmw = self.accelerator.unwrap_model(self.ref_model) + + # Both models given to GeometricMixtureWrapper (policy_model_for_gmw and ref_model_for_gmw) are DDP-unwrapped. + with torch.no_grad(): # Ensure no_grad context for mixture model generation + mixture_model = GeometricMixtureWrapper( + model=policy_model_for_gmw, + ref_model=ref_model_for_gmw, + generation_config=self.generation_config, + mixture_coef=self.mixture_coef, + device=self.accelerator.device, + ) + + mixture_output = mixture_model.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + return model_output, mixture_output + + def _process_completions(self, model_output, mixture_output, prompts): + context_length = prompts["input_ids"].shape[1] + + # Process model completions + model_completion_ids = model_output[:, context_length:] + model_completion_ids, model_completion_mask = truncate_right( + model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + model_data = { + "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1), + "raw": prompts["raw"], + } + + # Process reference model completions + mixture_completion_ids = mixture_output[:, context_length:] + mixture_completion_ids, mixture_completion_mask = truncate_right( + mixture_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + mixture_data = { + "input_ids": torch.cat((prompts["input_ids"], mixture_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], mixture_completion_mask), dim=1), + "raw": prompts["raw"], + } + + return model_data, mixture_data + + def _compute_rewards(self, model_data, mixture_data, context_length): + with torch.no_grad(): + _, model_scores, _ = get_reward( + self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, mixture_scores, _ = get_reward( + self.reward_funcs, mixture_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + + # Apply EOS penalty if needed + if self.args.missing_eos_penalty is not None: + model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + mixture_contain_eos = torch.any(mixture_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + model_scores[~model_contain_eos] -= self.args.missing_eos_penalty + mixture_scores[~mixture_contain_eos] -= self.args.missing_eos_penalty + + return model_scores, mixture_scores + + def _compute_judge(self, model_data, mixture_data, context_length): + prompts = model_data["raw"] + model_data_completions = self.processing_class.batch_decode( + model_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + model_data_completions = [completion.strip() for completion in model_data_completions] + + mixture_data_completions = self.processing_class.batch_decode( + mixture_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + mixture_data_completions = [completion.strip() for completion in mixture_data_completions] + if is_conversational({"prompt": prompts[0]}): + model_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in model_data_completions + ] + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=message) for message in prompts] + model_data_completions = [template.render(messages=completion) for completion in model_data_completions] + + mixture_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in mixture_data_completions + ] + mixture_data_completions = [ + template.render(messages=completion) for completion in mixture_data_completions + ] + + probability = self.judge.judge( + prompts, + list(zip(model_data_completions, mixture_data_completions, strict=True)), + return_scores=True, + ) + return torch.tensor(probability, device=model_data["input_ids"].device) + + def _compute_logprobs(self, model, model_data, context_length): + def compute_logprobs_for_data(m, data): + output = m(data["input_ids"], attention_mask=data["attention_mask"]) + logits = output.logits[:, context_length - 1 : -1] + token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:]) + return token_logprobs + + # Compute logprobs for model completions under the model + model_logprobs_model_data = compute_logprobs_for_data(model, model_data) + + # Compute logprobs of model completions under the reference model + with torch.no_grad(): + if self.ref_model is None: + with model.disable_adapter(): + ref_logprobs_model_data = compute_logprobs_for_data(model, model_data) + else: + ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data) + + # Mask padding tokens + model_padding_mask = model_data["attention_mask"][:, context_length:] == 0 + model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + + return (model_logprobs_model_data, ref_logprobs_model_data) + + def _compute_losses( + self, + model_logprobs_model_data, + ref_logprobs_model_data, + probability, + ): + # reinforce score where 0.5 is a control variate + score = (probability - 0.5) * model_logprobs_model_data.sum(1) + + # kl divergence via reinforce + with torch.no_grad(): + log_ratio = model_logprobs_model_data - ref_logprobs_model_data + kl_div_log = log_ratio.sum(1) + kl_div_loss = (log_ratio * model_logprobs_model_data).sum(1) + + # final loss + loss = self.beta * kl_div_loss - score + + return loss.mean(), score, kl_div_log + + def _log_statistics( + self, + model_data, + mixture_data, + model_logprobs_model_data, + ref_logprobs_model_data, + probability, + score, + kl_div, + context_length, + model_scores=None, + mixture_scores=None, + ): + # Helper function to gather and compute mean + def gather_mean(tensor): + return self.accelerator.gather_for_metrics(tensor).mean().item() + + # Log score + self.stats["loss/score"].append(gather_mean(score)) + # Log KL divergence + self.stats["loss/kl"].append(gather_mean(kl_div)) + + # Log logprobs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + self.stats["logps/chosen"].append(gather_mean(model_logprobs_model_data_sum)) + self.stats["logps/rejected"].append(gather_mean(ref_logprobs_model_data_sum)) + + # Log rewards + if self.reward_funcs is not None: + self.stats["rewards/chosen"].append(gather_mean(model_scores)) + self.stats["rewards/rejected"].append(gather_mean(mixture_scores)) + + # Log probabilities + self.stats["rewards/probabilities"].append(gather_mean(probability)) + + # Calculate entropy for model data + entropy_model_data = -model_logprobs_model_data.sum(1) + self.stats["objective/entropy"].append(gather_mean(entropy_model_data)) + + # Calculate margins + margin = model_logprobs_model_data_sum - ref_logprobs_model_data_sum + self.stats["rewards/margins"].append(gather_mean(margin)) + + # Calculate accuracy + accuracy = (margin > 0).float() + self.stats["rewards/accuracies"].append(gather_mean(accuracy)) + + # Log EOS token statistics + model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + mixture_eos = (mixture_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float())) + self.stats["val/ref_contain_eos_token"].append(gather_mean(mixture_eos.float())) + + # Log beta and mixture coef + self.stats["beta"].append(self.beta) + self.stats["mixture_coef"].append(self.mixture_coef) + + def training_step( + self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + ) -> torch.Tensor: + model.train() + + # Apply chat template and tokenize the input + batch_size = len(next(iter(inputs.values()))) + prompts = inputs["prompt"] + inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # need the prompt_ only + inputs = self._prepare_inputs(inputs) + context_length = inputs["prompt_input_ids"].shape[1] + prompts = { + "input_ids": inputs["prompt_input_ids"], + "attention_mask": inputs["prompt_attention_mask"], + "raw": prompts, + } + del inputs + + # Sample completions from both the model and the reference model + model_output, mixture_output = self._generate_completions(model, prompts) + + # Process model completions + model_data, mixture_data = self._process_completions(model_output, mixture_output, prompts) + + # Compute rewards + if self.reward_funcs is not None: + model_scores, mixture_scores = self._compute_rewards(model_data, mixture_data, context_length) + # probability of the model data vs the mixture data + probability = F.sigmoid(model_scores - mixture_scores) + else: + model_scores, mixture_scores = None, None + probability = self._compute_judge(model_data, mixture_data, context_length) + + # Compute logprobs + model_logprobs_model_data, ref_logprobs_model_data = self._compute_logprobs(model, model_data, context_length) + + # Compute loss + loss, score, kl_div = self._compute_losses(model_logprobs_model_data, ref_logprobs_model_data, probability) + + # Log everything + self._log_statistics( + model_data, + mixture_data, + model_logprobs_model_data.detach(), + ref_logprobs_model_data, + probability, + score.detach(), + kl_div.detach(), + context_length, + model_scores, + mixture_scores, + ) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps diff --git a/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..680aeb922f171453bdb95adee6e81ecfca67d5a2 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_config.py @@ -0,0 +1,385 @@ +# Copyright 2020-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 warnings +from dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class OnlineDPOConfig(TrainingArguments): + r""" + Configuration class for the [`OnlineDPOTrainer`]. + + This class includes only the parameters that are specific to Online DPO training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + reward_model_path (`str`, *optional*): + Path to the reward model. Either `judge` or `reward_model_path` must be set, but not both. + judge (`str`, *optional*): + Name of the judge to use. Either `judge` or `reward_model_path` must be set, but not both. + max_new_tokens (`int`, *optional*, defaults to `64`): + Maximum number of tokens to generate per completion. + max_length (`int`, *optional*, defaults to `256`): + Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If the + sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the completion as + possible. + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + missing_eos_penalty (`float`, *optional*): + Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to + generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive + value. This parameter only works when using `reward_funcs` and not when using `judge`. + beta (`float` or `list[float]`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β is + selected for each new epoch and the last β is used for the rest of the epochs. + loss_type (`str`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + + > Parameters that control generation + + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + use_transformers_paged (`bool`, *optional*, defaults to `False`): + Whether to use the `transformers` paged implementation for generation. If set to `True`, the `transformers` + paged implementation will be used for generation instead of the default padded implementation. This + parameter is only effective when `use_vllm` is set to `False`. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_guided_decoding_regex (`str`, *optional*): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.55`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + + > Other parameters + + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + """ + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=5e-7, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + reward_model_path: str | None = field( + default=None, + metadata={ + "help": "Path to the reward model. Either `judge` or `reward_model_path` must be set, but not both." + }, + ) + judge: str | None = field( + default=None, + metadata={ + "help": "Name of the judge to use. Either `judge` or `reward_model_path` must be set, but not both." + }, + ) + max_new_tokens: int = field( + default=64, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + max_length: int = field( + default=512, + metadata={ + "help": "Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If " + "the sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the " + "completion as possible." + }, + ) + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=1.0, + metadata={ + "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. " + "Set to 1.0 to consider all tokens." + }, + ) + top_k: int | None = field( + default=None, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + min_p: float | None = field( + default=None, + metadata={ + "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It " + "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range." + }, + ) + repetition_penalty: float = field( + default=1.0, + metadata={ + "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated " + "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model " + "to repeat tokens." + }, + ) + generation_kwargs: dict | None = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or " + "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the " + "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that " + "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them." + }, + ) + use_transformers_paged: bool = field( + default=False, + metadata={ + "help": "Whether to use the `transformers` paged implementation for generation. If set to `True`, the " + "`transformers` paged implementation will be used for generation instead of the default padded " + "implementation. This parameter is only effective when `use_vllm` is set to `False`." + }, + ) + cache_implementation: str | None = field( + default=None, + metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."}, + ) + missing_eos_penalty: float | None = field( + default=None, + metadata={ + "help": "Penalty applied to the score when the model fails to generate an EOS token. This is useful to " + "encourage to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be " + "a positive value." + }, + ) + beta: list[float] = field( + default_factory=lambda: [0.1], + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model. For the IPO loss (`loss_type='ipo'`), β is the regularization parameter denoted by " + "τ in the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β " + "is selected for each new epoch and the last β is used for the rest of the epochs." + }, + ) + loss_type: str = field( + default="sigmoid", + metadata={ + "help": "Type of loss to use.", + "choices": ["sigmoid", "ipo"], + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. Requires vLLM to be installed " + "(`pip install trl[vllm]`)." + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={ + "help": "Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: " + "Use the `transformers` backend for model implementation. `vllm`: Use the `vllm` library for " + "model implementation." + }, + ) + vllm_guided_decoding_regex: str | None = field( + default=None, + metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, + ) + vllm_gpu_memory_utilization: float | None = field( + default=0.55, + metadata={ + "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag.", + }, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `'server'` or " + "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure " + "a TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same " + "process and share the training GPUs. This avoids the need for a separate server but may cause resource " + "contention with training.", + }, + ) + vllm_server_base_url: str | None = field( + default=None, + metadata={ + "help": "Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` " + "and `vllm_server_port` are ignored.", + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": "Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_port: int = field( + default=8000, + metadata={"help": "Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={ + "help": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up " + "after the timeout, a `ConnectionError` is raised.", + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={ + "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_tensor_parallel_size` flag.", + }, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option " + "is not compatible with vLLM generation." + }, + ) + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + reward_weights: list[float] | None = field( + default=None, + metadata={ + "help": "Weights for combining multiple reward functions. Must match the number of reward functions. " + "If None, all reward functions are equally weighted." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + if hasattr(self.beta, "__len__") and len(self.beta) == 1: + self.beta = self.beta[0] + + if self.max_new_tokens >= self.max_length: + warnings.warn( + f"The configuration has `max_new_tokens` ({self.max_new_tokens}) >= `max_length` ({self.max_length}). " + "This will cause prompts to be truncated or completely removed in the forward pass. " + "To preserve prompts, ensure e.g. `max_length > max_new_tokens + 512`. ", + ) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..bbf5aea469c06b0e8fbaee10430f8175746bc85e --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/online_dpo_trainer.py @@ -0,0 +1,1475 @@ +# Copyright 2020-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 os +import re +import textwrap +import warnings +from collections.abc import Callable +from contextlib import nullcontext +from functools import wraps +from pathlib import Path +from typing import Any + +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data +import transformers +from accelerate import logging +from accelerate.utils import broadcast_object_list, gather_object, is_peft_model +from datasets import Dataset +from packaging import version +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.utils.data import DataLoader, IterableDataset +from transformers import ( + AutoModelForCausalLM, + AutoModelForSequenceClassification, + AutoTokenizer, + DataCollator, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainerCallback, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES +from transformers.trainer_utils import EvalPrediction, seed_worker +from transformers.training_args import OptimizerNames +from transformers.utils import ( + is_flash_attn_2_available, + is_peft_available, + is_sagemaker_mp_enabled, +) + +from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template +from ..extras.profiling import profiling_context +from ..extras.vllm_client import VLLMClient +from ..import_utils import is_vllm_available +from ..models import ( + create_reference_model, + prepare_deepspeed, + prepare_fsdp, + prepare_peft_model, + unwrap_model_for_generation, +) +from .base_trainer import BaseTrainer +from .judges import BasePairwiseJudge +from .online_dpo_config import OnlineDPOConfig +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + DPODataCollatorWithPadding, + disable_dropout_in_model, + empty_cache, + ensure_master_addr_port, + get_config_model_id, + pad, + truncate_right, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel + + +if is_sagemaker_mp_enabled(): + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") + +else: + IS_SAGEMAKER_MP_POST_1_10 = False + + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + + +logger = logging.get_logger(__name__) + +# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of +# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. +RewardFunc = str | PreTrainedModel | Callable[[list, list], list[float]] + + +class OnlineDPOTrainer(BaseTrainer): + r""" + Initialize OnlineDPOTrainer. + + Args: + model (`str | nn.Module | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + ref_model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `None`): + The reference model to use for training. If None is specified, the reference model will be created from the + model. + judge ([`BasePairwiseJudge`]): + The judge to use for pairwise comparison of model completions. + reward_funcs (`RewardFunc | list[RewardFunc]`, *optional*): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function: Can be a string (path to model), a [`~transformers.PreTrainedModel`], or a + custom callable function. + - A list of reward functions: Must all be of compatible types. + + Note: Only one of `judge`, or `reward_funcs` should be provided. + args ([`OnlineDPOConfig`]): + The online DPO config arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + + If set to `None`, the tokenizer for each model-based reward function is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "online-dpo"] + _name = "Online DPO" + _paper = { + "title": "Direct Language Model Alignment from Online AI Feedback", + "id": "2402.04792", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{guo2024direct, + title = {{Direct Language Model Alignment from Online AI Feedback}}, + author = {Shangmin Guo and Biao Zhang and Tianlin Liu and Tianqi Liu and Misha Khalman and Felipe Llinares and Alexandre Ram{\'{e}} and Thomas Mesnard and Yao Zhao and Bilal Piot and Johan Ferret and Mathieu Blondel}, + year = 2024, + eprint = {arXiv:2402.04792} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str, + ref_model: PreTrainedModel | nn.Module | None = None, + reward_funcs: RewardFunc | list[RewardFunc] | None = None, + judge: BasePairwiseJudge | None = None, + args: OnlineDPOConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + peft_config: "PeftConfig | None" = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + ) -> None: + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, either omit the `ref_model` argument or pass `None`." + ) + + self.ref_model = ref_model + + # Validate reward configuration - must have exactly one of: judge, or reward_funcs + reward_configs = sum(x is not None for x in [judge, reward_funcs]) + if reward_configs == 0: + raise ValueError("One of `judge` or `reward_funcs` must be provided.") + elif reward_configs > 1: + if judge is not None: + logger.warning( + "Both `judge` and `reward_funcs` are provided. Using `judge` and ignoring `reward_funcs`.", + UserWarning, + ) + reward_funcs = None + self.judge = judge + + # Handle reward_funcs + if reward_funcs is not None: + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + + # Process reward functions (convert strings to models, collect names) + model_init_kwargs = args.model_init_kwargs or {} + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + # Load model from string path + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): + self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Handle reward processing classes for reward_funcs + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + else: + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + "The number of reward processing classes must match the number of reward functions." + ) + + self.reward_processing_classes = [] + for reward_processing_class_i, reward_func in zip(reward_processing_classes, reward_funcs, strict=True): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class_i is None: + reward_processing_class_i = AutoTokenizer.from_pretrained(reward_func.config._name_or_path) + if reward_processing_class_i.pad_token_id is None: + reward_processing_class_i.pad_token = reward_processing_class_i.eos_token + # Set pad token ID on reward model config + reward_func.config.pad_token_id = reward_processing_class_i.pad_token_id + self.reward_processing_classes.append(reward_processing_class_i) + else: + self.reward_funcs = None + self.reward_func_names = [] + self.reward_processing_classes = [] + + # Handle reward_weights + if reward_funcs is not None: + if args.reward_weights is not None: + if len(args.reward_weights) != len(self.reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(self.reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(self.reward_funcs), dtype=torch.float32) + else: + self.reward_weights = None + + if args.missing_eos_penalty is not None and reward_funcs is None and judge is None: + raise ValueError("`missing_eos_penalty` is only supported when `reward_funcs` is provided.") + + if args is None: + raise ValueError("`args` must be provided.") + + # Check that the processing_class is provided + if processing_class is None: + raise ValueError("`processing_class` must be provided.") + + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + + # Handle dtype in model_init_kwargs + dtype = model_init_kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass + elif isinstance(dtype, str): + dtype = getattr(torch, dtype) + model_init_kwargs["dtype"] = dtype + else: + raise ValueError( + "Invalid `dtype` passed to `OnlineDPOConfig`. Expected either 'auto' or a string " + f"representing a `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + + model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs) + else: + if args.model_init_kwargs is not None: + raise ValueError( + "You passed `model_init_kwargs` to the `OnlineDPOConfig`, but your model is already instantiated. " + "This argument can only be used when the `model` argument is a string." + ) + self.is_encoder_decoder = model.config.is_encoder_decoder + self.is_vision_model = model.config.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.keys() + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Enable gradient checkpointing if requested + if args.gradient_checkpointing: + model = self._enable_gradient_checkpointing(model, args) + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Handle the ref_model + # Usually, the user wants the ref model to be the initial version of the model. When using PEFT, it's easy to + # get the ref model, as it's just the model with a disabled adapter. When not using PEFT, we need to create + # the ref model from the model by copying it and disable the gradients and set it in evaluation mode. + if ref_model is None: # No ref model provided, the most common case + if peft_config is None: + self.ref_model = create_reference_model(model) # copy, disable gradients, set eval mode + else: + self.ref_model = None # we don't need a ref model here, we can just disable the adapter. + else: # rare case, the user provided a ref model + self.ref_model = ref_model + self.ref_model.eval() + + # Disable the gradient and set the reward model in eval mode + if reward_funcs is not None: + for reward_func in reward_funcs: + if isinstance(reward_func, PreTrainedModel): + reward_func.eval() + + self.max_length = args.max_length + + self.stats = { + "objective/kl": [], + "objective/entropy": [], + "objective/non_score_reward": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + "val/contain_eos_token": [], + "beta": [], + } + if self.reward_funcs is not None: + self.stats["objective/rlhf_reward"] = [] + self.stats["objective/scores_margin"] = [] + self.stats["objective/scores"] = [] + + # Store generation parameters for later use + self.use_vllm = args.use_vllm + self.num_generations = 2 # Generate 2 completions per prompt for Online DPO + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_paged = args.use_transformers_paged + self.vllm_mode = args.vllm_mode if args.use_vllm else None + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size + self.vllm_model_impl = args.vllm_model_impl + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + self.pad_token = tokenizer.pad_token + self.pad_token_id = tokenizer.pad_token_id + self.eos_token_id = tokenizer.eos_token_id + + # Vision tokens for VLM support + self.image_token_id = getattr(processing_class, "image_token_id", None) + self.vision_start_token_id = getattr(processing_class, "vision_start_token_id", None) + self.vision_end_token_id = getattr(processing_class, "vision_end_token_id", None) + # Get the image token string for token collapsing + self.image_token = None + if self.image_token_id is not None: + self.image_token = tokenizer.decode([self.image_token_id]) + + # Define the collator if not provided + if data_collator is None: + data_collator = DPODataCollatorWithPadding(pad_token_id=self.pad_token_id) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in Online DPO, the sampled data does not include + # the "input_ids" key. As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + self._beta = args.beta + + # Set up generation configuration and vLLM after super().__init__ + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install trl[vllm]` to use it." + ) + + if self.vllm_mode == "server": + if self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator(device=torch.cuda.current_device()) + else: + self.vllm_client = None + elif self.vllm_mode == "colocate": + # vLLM dynamically adjusts the size of the key-value cache based on available GPU memory at instantiation. + # A larger cache size improves speed, so we would expect gpu_memory_utilization=1. + # However, at this stage, the optimizer's weights are not yet loaded onto the GPU; they will be loaded + # after the first optimizer step and remain in GPU memory throughout training. So we must reserve enough + # space for them. + # Configure vLLM parameters + vllm_kwargs = { + "model": model.name_or_path, + "tensor_parallel_size": self.vllm_tensor_parallel_size, + "gpu_memory_utilization": self.vllm_gpu_memory_utilization, + "model_impl": self.vllm_model_impl, + "max_num_seqs": self.args.per_device_train_batch_size * self.vllm_tensor_parallel_size, + "max_model_len": args.max_length + args.max_new_tokens, # max_length includes prompt + completion + "distributed_executor_backend": "external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + "seed": self.accelerator.process_index // self.vllm_tensor_parallel_size, + # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768) + "max_num_batched_tokens": 4096, + } + + # vLLM requires the environment variables to be set for distributed training. + os.environ["RANK"] = str(self.accelerator.process_index) + os.environ["LOCAL_RANK"] = str(self.accelerator.local_process_index) + os.environ["WORLD_SIZE"] = str(self.accelerator.num_processes) + # Ensure distributed rendezvous variables are set without colliding across concurrent runs + ensure_master_addr_port() + + self.llm = LLM(**vllm_kwargs) + else: + raise ValueError(f"vllm_mode must be either 'server' or 'colocate', got '{self.vllm_mode}'.") + # vLLM specific sampling arguments + self.guided_decoding_regex = args.vllm_guided_decoding_regex + self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation + + # Set up vLLM generation config + generation_params = { + "n": 2, # 2 generations per prompt for Online DPO + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": args.max_new_tokens, + "detokenize": False, # to avoid vllm to decode (we don't need it) + } + if args.generation_kwargs is not None: + generation_params.update(args.generation_kwargs) + if self.guided_decoding_regex: + generation_params["guided_decoding"] = GuidedDecodingParams(regex=self.guided_decoding_regex) + self.generation_config = SamplingParams(**generation_params) + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + # Set up transformers generation config + generation_kwargs = { + "max_new_tokens": args.max_new_tokens, + "do_sample": True, + "pad_token_id": self.pad_token_id, + "bos_token_id": tokenizer.bos_token_id, + "eos_token_id": self.eos_token_id, + "temperature": self.temperature, + "top_k": self.top_k, + "top_p": self.top_p, + "repetition_penalty": self.repetition_penalty, + "use_cache": True if not self.args.gradient_checkpointing else False, + } + # Add min_p if supported + if self.min_p is not None: + generation_kwargs["min_p"] = self.min_p + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + # Remove None values + generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None} + self.generation_config = GenerationConfig(**generation_kwargs) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + if self.reward_funcs is not None: + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + @property + def beta(self): + if isinstance(self._beta, list): + epoch = self.state.epoch + return self._beta[epoch] if epoch < len(self._beta) else self._beta[-1] + else: + return self._beta + + @staticmethod + def tokenize_row(feature, is_encoder_decoder: bool, tokenizer: PreTrainedTokenizerBase) -> dict[str, Any]: + """Tokenize a single row from a DPO specific dataset.""" + if not is_encoder_decoder: + batch = tokenizer(feature["prompt"], add_special_tokens=False) + # Add BOS token to head of prompt. Avoid adding if it's already there + if tokenizer.bos_token_id is not None: + prompt_len_input_ids = len(batch["input_ids"]) + if prompt_len_input_ids == 0 or tokenizer.bos_token_id != batch["input_ids"][0]: + batch["input_ids"] = [tokenizer.bos_token_id] + batch["input_ids"] + batch["attention_mask"] = [1] + batch["attention_mask"] + else: + batch = tokenizer(feature["prompt"], add_special_tokens=True) + batch = {f"prompt_{key}": value for key, value in batch.items()} + return batch + + # Same as Trainer.get_train_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_train_dataloader) + def get_train_dataloader(self) -> DataLoader: + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + # Same as Trainer.get_eval_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_eval_dataloader) + def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> DataLoader: + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + + # If we have persistent workers, don't do a fork bomb especially as eval datasets + # don't change during training + dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): + return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) + + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) + data_collator = self.data_collator + + dataloader_params = { + "batch_size": self.args.eval_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(eval_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset) + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + # accelerator.free_memory() will destroy the references, so + # we need to store the non-prepared version + eval_dataloader = DataLoader(eval_dataset, **dataloader_params) + if self.args.dataloader_persistent_workers: + if hasattr(self, "_eval_dataloaders"): + self._eval_dataloaders[dataloader_key] = eval_dataloader + else: + self._eval_dataloaders = {dataloader_key: eval_dataloader} + + return self.accelerator.prepare(eval_dataloader) + + def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: OnlineDPOConfig) -> PreTrainedModel: + """Enables gradient checkpointing for the model.""" + # Ensure use_cache is disabled + model.config.use_cache = False + + # Enable gradient checkpointing on the base model for PEFT + if is_peft_model(model): + model.base_model.gradient_checkpointing_enable() + # Enable gradient checkpointing for non-PEFT models + else: + model.gradient_checkpointing_enable() + + gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = ( + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + ) + + if use_reentrant: + model.enable_input_require_grads() + + return model + + def _generate_vllm(self, prompts, images=None): + eos_token_id = self.eos_token_id + pad_token_id = self.pad_token_id + + # Generate completion_ids and prompt_ids based on mode + if self.vllm_mode == "server": + completion_ids, prompt_ids = self._generate_vllm_server(prompts, images) + elif self.vllm_mode == "colocate": + completion_ids, prompt_ids = self._generate_vllm_colocate(prompts, images) + + # Shared padding, masking, and tensor conversion logic + max_prompt_length = max(len(ids) for ids in prompt_ids) + prompt_mask = [[0] * (max_prompt_length - len(ids)) + [1] * len(ids) for ids in prompt_ids] + prompt_ids = [[pad_token_id] * (max_prompt_length - len(ids)) + ids for ids in prompt_ids] + max_tokens = self.generation_config.max_tokens + completion_mask = [[1] * len(ids) + [0] * (max_tokens - len(ids)) for ids in completion_ids] + completion_ids = [ + ids + [eos_token_id] if ids[-1] != eos_token_id and len(ids) < max_tokens else ids + for ids in completion_ids + ] + completion_ids = [ids + [pad_token_id] * (max_tokens - len(ids)) for ids in completion_ids] + + # Convert to tensors + prompt_ids = torch.tensor(prompt_ids, device=self.accelerator.device) + prompt_mask = torch.tensor(prompt_mask, device=self.accelerator.device) + completion_ids = torch.tensor(completion_ids, device=self.accelerator.device) + completion_mask = torch.tensor(completion_mask, device=self.accelerator.device) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _generate_vllm_server(self, prompts, images=None): + """Generate completions using vLLM server mode""" + has_images = images is not None + + # Update vLLM server weights if needed + if hasattr(self, "_last_loaded_step") and self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + elif not hasattr(self, "_last_loaded_step"): + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + # Apply chat template if conversational + if is_conversational({"prompt": prompts[0]}): + prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] + else: + prompts_text = prompts + # Gather all prompts to main process + all_prompts = gather_object(prompts_text) + if has_images: + all_images = gather_object(images) + + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts[:: self.num_generations] + if has_images: + ordered_set_of_images = all_images[:: self.num_generations] + else: + ordered_set_of_images = None + completion_ids = self.vllm_client.generate( + prompts=ordered_set_of_prompts, + images=ordered_set_of_images, + n=self.num_generations, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=-1 if self.top_k is None else self.top_k, + min_p=0.0 if self.min_p is None else self.min_p, + max_tokens=self.generation_config.max_tokens, + guided_decoding_regex=self.guided_decoding_regex if hasattr(self, "guided_decoding_regex") else None, + generation_kwargs=self.args.generation_kwargs, + ) + # Flatten: each prompt generates 2 completions + completion_ids = [[comp_id] for prompt_completions in completion_ids for comp_id in prompt_completions] + else: + completion_ids = [None] * (len(all_prompts) * 2) + + # Broadcast completions to all processes + completion_ids = broadcast_object_list(completion_ids, from_process=0) + + # Each process takes its slice + process_slice = slice( + self.accelerator.process_index * len(prompts) * 2, + (self.accelerator.process_index + 1) * len(prompts) * 2, + ) + completion_ids = completion_ids[process_slice] + + # Create prompt_ids by tokenizing locally + prompt_inputs = self.processing_class( + text=prompts_text, + return_tensors="pt", + padding=True, + padding_side="left", + add_special_tokens=False, + ) + prompt_ids = [] + for prompt_tokens in prompt_inputs["input_ids"]: + prompt_ids.extend([prompt_tokens.tolist(), prompt_tokens.tolist()]) # 2 copies for 2 completions + return completion_ids, prompt_ids + + def _generate_vllm_colocate(self, prompts, images=None): + """Generate completions using vLLM colocate mode""" + # Update model weights if needed - only after gradient accumulation completes + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + # Apply chat template if conversational + if is_conversational({"prompt": prompts[0]}): + prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] + else: + prompts_text = prompts + + # Prepare vLLM inputs with images if available + if images is not None: + vllm_inputs = [] + for prompt, image in zip(prompts_text, images, strict=True): + if image is not None: + vllm_inputs.append({"prompt": prompt, "multi_modal_data": {"image": image}}) + else: + vllm_inputs.append(prompt) + else: + vllm_inputs = prompts_text + + outputs = self.llm.generate(vllm_inputs, self.generation_config, use_tqdm=False) + + completion_ids = [list(output.outputs[i].token_ids) for i in range(2) for output in outputs] + prompt_ids = [list(output.prompt_token_ids) for _ in range(2) for output in outputs] + + return completion_ids, prompt_ids + + def _move_model_to_vllm(self): + """Synchronize model weights to vLLM server with support for PEFT, DeepSpeed, and FSDP""" + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + # TODO: does this work with FSDP? + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + # use memory-efficient post-order traversal for FSDP + self._sync_fsdp1_params_to_vllm(self.model) + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + for name, param in self.model.named_parameters(): + name = self._fix_param_name_to_vllm(name) + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.llm.reset_prefix_cache() + + def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + # For FSDP1, we need to recurse into children and also use summon_full_params + if visited is None: + visited = set() + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp1_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + def _sync_fsdp2_params_to_vllm(self, module: nn.Module): + # For FSDP2, module.state_dict() already covers all parameters, so no need for recursion + for name, param in module.state_dict().items(): + if param.is_cpu: + param = param.to(torch.device("cuda")) + param = param.full_tensor() + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param)]) + + def _fix_param_name_to_vllm(self, name, extra_prefixes: list[str] | None = None): + """Clean parameter names for vLLM compatibility""" + extra_prefixes = extra_prefixes or [] + prefixes = ["_checkpoint_wrapped_module."] + extra_prefixes + for prefix in prefixes: + name = name.replace(prefix, "") + return name + + def process_vision_row( + self, features: dict[str, list | torch.Tensor], processing_class=None + ) -> dict[str, list[int]]: + """ + Process a vision row for VLM models (adapted from DPO trainer) + """ + processor = processing_class or self.processing_class + processed_features = processor(images=[features["image"]], text=features["prompt"], add_special_tokens=False) + + prompt_input_ids = processed_features["input_ids"][0] + + # Create the output dict with required fields + output = { + "prompt_input_ids": prompt_input_ids, + "prompt_attention_mask": processed_features["attention_mask"][0], + } + + # Add vision-specific fields + if "pixel_values" in processed_features: + output["pixel_values"] = processed_features["pixel_values"][0] + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + + return output + + def _generate(self, model, prompts, images=None): + """Generate completions using the model""" + device = next(model.parameters()).device + eos_token_id = self.eos_token_id + pad_token_id = self.pad_token_id + + # Apply chat template and tokenize the input + inputs = [{"prompt": prompt} for prompt in prompts] + + # Add images if provided (VLM support) + if images is not None: + for i, image in enumerate(images): + inputs[i]["image"] = image + + # Apply chat template to get text prompts + prompts_text = [maybe_apply_chat_template(x, self.processing_class)["prompt"] for x in inputs] + + # Handle image token collapsing/removal + # The chat template sometimes inserts a single image token into the prompt text. However, when this text is + # later tokenized, the single image token string is expanded into multiple image token IDs, depending on the + # image size. We need to handle this properly. + if self.image_token is not None and images is not None: + escaped_img_token = re.escape(self.image_token) + # Search for the image token in the chat template + if hasattr(self.processing_class, "chat_template") and self.processing_class.chat_template: + if re.search(escaped_img_token, self.processing_class.chat_template): + # Collapse repeated image tokens back into a single token + prompts_text = [ + re.sub(rf"({escaped_img_token})+", self.image_token, text) for text in prompts_text + ] + else: + # If the chat template doesn't use the image token, remove all instances + if self.vision_end_token_id is not None: + escaped_eoi_token = re.escape( + self.processing_class.tokenizer.decode([self.vision_end_token_id]) + ) + prompts_text = [ + re.sub(rf"({escaped_img_token})+{escaped_eoi_token}", "", text) for text in prompts_text + ] + else: + # If vision_end_token_id is None, just remove the image tokens + prompts_text = [re.sub(rf"({escaped_img_token})+", "", text) for text in prompts_text] + + # Prepare kwargs for processing class + kwargs = {} + if images is not None: + kwargs = {"images": [[img] for img in images]} + + # Process inputs using the processing class (handles both VLM and LLM) + prompt_inputs = self.processing_class( + text=prompts_text, + return_tensors="pt", + padding=True, + padding_side="left", + add_special_tokens=False, + **kwargs, + ) + + prompt_inputs = {k: v.to(device) for k, v in prompt_inputs.items()} + # Convert vision inputs to model's dtype for proper computation + if "pixel_values" in prompt_inputs: + # Handle DataParallel wrapped models + model_dtype = getattr(model, "dtype", None) + if model_dtype is None and hasattr(model, "module"): + model_dtype = model.module.dtype + if model_dtype is not None: + prompt_inputs["pixel_values"] = prompt_inputs["pixel_values"].to(model_dtype) + + # Sample 2 completions per prompt of size `max_new_tokens` from the model + prompt_ids = prompt_inputs["input_ids"].repeat(2, 1) + prompt_mask = prompt_inputs["attention_mask"].repeat(2, 1) + + # Prepare vision inputs if available + vision_generation_kwargs = {} + if self.is_vision_model and images is not None: + if "pixel_values" in prompt_inputs: + vision_generation_kwargs["pixel_values"] = prompt_inputs["pixel_values"].repeat(2, 1, 1, 1) + if "pixel_attention_mask" in prompt_inputs: + vision_generation_kwargs["pixel_attention_mask"] = prompt_inputs["pixel_attention_mask"].repeat(2, 1) + if "image_sizes" in prompt_inputs: + vision_generation_kwargs["image_sizes"] = prompt_inputs["image_sizes"].repeat(2, 1) + if "image_grid_thw" in prompt_inputs: + vision_generation_kwargs["image_grid_thw"] = prompt_inputs["image_grid_thw"].repeat(2, 1) + + if self.use_transformers_paged: + previous_attn = self.model_wrapped.config._attn_implementation + + if version.parse(transformers.__version__).release >= version.parse("5.0.0").release: + new_attn = "paged|flash_attention_2" if is_flash_attn_2_available() else "paged|sdpa" + else: + new_attn = "paged_attention" if is_flash_attn_2_available() else "sdpa_paged" + self.model_wrapped.config._attn_implementation = new_attn + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + with torch.inference_mode(): + all_outputs = unwrapped_model.generate_batch( + prompt_ids.tolist(), + generation_config=self.generation_config, + progress_bar=False, + ) + unwrapped_model.train() # restore training mode, as generate_batch forces eval mode + completion_ids = [output.generated_tokens for output in all_outputs.values()] + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right") + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) + # Restore the original attention implementation, training mode + self.model_wrapped.config._attn_implementation = previous_attn + + # Extract completion_ids and create completion_mask + prompt_length = prompt_ids.size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + completion_ids, completion_mask = truncate_right(completion_ids, eos_token_id, pad_token_id) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + else: + # Regular generation path + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + # Setup cache implementation if specified + if self.args.cache_implementation is not None: + unwrapped_model.generation_config.cache_implementation = self.args.cache_implementation + + # Standard generation + output = unwrapped_model.generate( + input_ids=prompt_ids, + attention_mask=prompt_mask, + generation_config=self.generation_config, + **vision_generation_kwargs, + ) + + completion_ids = output[:, prompt_ids.size(1) :] + completion_ids, completion_mask = truncate_right(completion_ids, eos_token_id, pad_token_id) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _calculate_rewards_from_functions(self, prompts, completions, completion_ids_list, **reward_kwargs): + """ + Calculate rewards using reward functions + """ + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Add trainer state to reward kwargs for dynamic reward shaping + reward_kwargs["trainer_state"] = self.state + + for i, (reward_func, reward_processing_class) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, strict=True) + ): + if isinstance(reward_func, nn.Module): # Model-based reward function + # Handle conversational vs text input + if is_conversational({"prompt": prompts[0]}): + messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] + texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages] + else: + texts = [p + c for p, c in zip(prompts, completions, strict=True)] + + # Tokenize and get reward scores + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = {k: v.to(device) for k, v in reward_inputs.items()} + + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + # Custom reward function + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # Weight and sum across all reward functions + if self.reward_weights is not None: + total_rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + else: + total_rewards = rewards_per_func.nansum(dim=1) + + return total_rewards + + def _forward(self, model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs=None): + # Get the number of tokens to truncate from prompt + num_tokens_to_truncate = max(prompt_ids.size(1) + completion_ids.size(1) - self.max_length, 0) + + # Truncate left to avoid oom + prompt_ids = prompt_ids[:, num_tokens_to_truncate:] + prompt_mask = prompt_mask[:, num_tokens_to_truncate:] + + # Concat the prompt and completion + prompt_completion_ids = torch.cat((prompt_ids, completion_ids), dim=1) + prompt_completion_mask = torch.cat((prompt_mask, completion_mask), dim=1) + + # Prepare model kwargs with vision inputs if available + model_kwargs = {"attention_mask": prompt_completion_mask} + if vision_inputs is not None: + if "pixel_values" in vision_inputs: + model_kwargs["pixel_values"] = vision_inputs["pixel_values"] + if "pixel_attention_mask" in vision_inputs: + model_kwargs["pixel_attention_mask"] = vision_inputs["pixel_attention_mask"] + if "image_sizes" in vision_inputs: + model_kwargs["image_sizes"] = vision_inputs["image_sizes"] + if "image_grid_thw" in vision_inputs: + model_kwargs["image_grid_thw"] = vision_inputs["image_grid_thw"] + + # Get the logprobs of the completions from the model + output = model(prompt_completion_ids, **model_kwargs) + + # There is 1 offset, because the model predicts the next token + prompt_len = prompt_ids.size(1) + start_idx = prompt_len - 1 if prompt_len > 0 else 0 + # Only slice off the last logit when we have a prompt, otherwise we need all logits + end_idx = -1 if prompt_len > 0 else None + logits = output.logits[:, start_idx:end_idx] + + # Take the completion tokens logprob + logprobs = torch.take_along_dim(logits.log_softmax(dim=-1), completion_ids.unsqueeze(-1), dim=2).squeeze(-1) + return logprobs + + def training_step( + self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + ) -> torch.Tensor: + model.train() + + prompts = inputs["prompt"] + batch_size = len(prompts) + + # Handle images for VLM support + has_images = "image" in inputs + images = None + if has_images: + images = inputs["image"] + # Convert conversational prompts to include image tokens + for prompt in prompts: + if isinstance(prompt, list): + for message in prompt: + if not isinstance(message, dict): + continue + content = message.get("content") + role = message.get("role") + if isinstance(content, str): + if role == "user": + message["content"] = [{"type": "image"}, {"type": "text", "text": content}] + elif role == "system": + message["content"] = [{"type": "text", "text": content}] + + if self.args.use_vllm: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_vllm(prompts, images) + else: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts, images) + + contain_eos_token = torch.any(completion_ids == self.eos_token_id, dim=-1) + + # Extract vision inputs if available for VLM support + vision_inputs = None + if has_images and self.is_vision_model and not self.args.use_vllm: + # For vision models with transformers generation, we need to prepare vision inputs + # Process the images to get vision inputs that can be passed through the forward pass + vision_inputs = {} + kwargs = {"images": [[img] for img in images]} + processed = self.processing_class( + text=[""] * len(images), # Dummy text for vision processing + return_tensors="pt", + **kwargs, + ) + # Handle DataParallel wrapped models + model_device = getattr(model, "device", None) + model_dtype = getattr(model, "dtype", None) + if model_device is None and hasattr(model, "module"): + model_device = model.module.device + model_dtype = model.module.dtype + # Move vision tensors to device and convert to model dtype + # Need to duplicate for 2 completions per prompt + if "pixel_values" in processed: + vision_inputs["pixel_values"] = ( + processed["pixel_values"].to(model_device, dtype=model_dtype).repeat(2, 1, 1, 1) + ) + if "pixel_attention_mask" in processed: + vision_inputs["pixel_attention_mask"] = processed["pixel_attention_mask"].to(model_device).repeat(2, 1) + if "image_sizes" in processed: + vision_inputs["image_sizes"] = processed["image_sizes"].to(model_device).repeat(2, 1) + if "image_grid_thw" in processed: + vision_inputs["image_grid_thw"] = processed["image_grid_thw"].to(model_device).repeat(2, 1) + + logprobs = self._forward(model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs) + with torch.no_grad(): + if self.ref_model is not None: + ref_logprobs = self._forward( + self.ref_model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs + ) + else: # peft case: we just need to disable the adapter + with self.model.disable_adapter(): + ref_logprobs = self._forward( + self.model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs + ) + + # Decode the completions, and format them if the input is conversational + device = logprobs.device + completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational({"prompt": prompts[0]}): + completions = [[{"role": "assistant", "content": completion}] for completion in completions] + + # Get the reward from reward functions or judge + if self.reward_funcs is not None: + # First create completion_ids_list for custom reward functions + completion_ids_list = [completion_ids[i].tolist() for i in range(completion_ids.shape[0])] + + # Extract additional fields from inputs for reward functions + reward_kwargs = {} + keys = [key for key in inputs if key not in ["prompt"]] + for key in keys: + if isinstance(inputs[key], (list, tuple)): + # Repeat input fields to match number of completions (2 per prompt) + reward_kwargs[key] = inputs[key] * 2 + else: + reward_kwargs[key] = inputs[key] + + # Calculate rewards using reward functions + rewards = self._calculate_rewards_from_functions( + prompts=2 * prompts, completions=completions, completion_ids_list=completion_ids_list, **reward_kwargs + ) + + # Apply missing EOS penalty if configured + if self.args.missing_eos_penalty is not None: + rewards[~contain_eos_token] -= self.args.missing_eos_penalty + + # Split rewards into chosen/rejected pairs + first_half, second_half = rewards.split(batch_size) + mask = first_half >= second_half + elif self.judge is not None: + # Once formatted, conversational data may contain special tokens (such as <|im_start|>) that are not + # directly understandable by the judge and could alter its judgment. To avoid this and make the judge + # independent of the model's chat template, we use the raw conversation data, and apply our own chat + # template to it. + if is_conversational({"prompt": prompts[0]}): + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=prompt) for prompt in prompts] + completions = [template.render(messages=completion) for completion in completions] + + ranks_of_first_completion = self.judge.judge( + prompts, list(zip(completions[:batch_size], completions[batch_size:], strict=True)) + ) + + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + mask = torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=device) + + batch_range = torch.arange(batch_size, device=device) + chosen_indices = batch_range + (~mask * batch_size) + rejected_indices = batch_range + (mask * batch_size) + + # Build tensor so that the first half is the chosen examples and the second half the rejected examples + cr_indices = torch.cat((chosen_indices, rejected_indices), dim=0) # cr = chosen and rejected + cr_logprobs = logprobs[cr_indices] + cr_ref_logprobs = ref_logprobs[cr_indices] + + # mask out the padding tokens + padding_mask = ~completion_mask.bool() + cr_padding_mask = padding_mask[cr_indices] + + cr_logprobs_sum = (cr_logprobs * ~cr_padding_mask).sum(1) + cr_ref_logprobs_sum = (cr_ref_logprobs * ~cr_padding_mask).sum(1) + + # Split the chosen and rejected examples + chosen_logprobs_sum, rejected_logprobs_sum = torch.split(cr_logprobs_sum, batch_size) + chosen_ref_logprobs_sum, rejected_ref_logprobs_sum = torch.split(cr_ref_logprobs_sum, batch_size) + pi_logratios = chosen_logprobs_sum - rejected_logprobs_sum + ref_logratios = chosen_ref_logprobs_sum - rejected_ref_logprobs_sum + + logits = pi_logratios - ref_logratios + + if self.args.loss_type == "sigmoid": + losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.loss_type}") + + loss = losses.mean() + + # Log everything + if self.reward_funcs is not None: + # When using reward_funcs, we have rewards instead of scores + scores_margin = rewards[chosen_indices] - rewards[rejected_indices] + self.stats["objective/scores_margin"].append( + self.accelerator.gather_for_metrics(scores_margin.mean()).mean().item() + ) + self.stats["objective/scores"].append(self.accelerator.gather_for_metrics(rewards.mean()).mean().item()) + self.stats["val/contain_eos_token"].append(contain_eos_token.float().mean().item()) + self.stats["logps/chosen"].append(self.accelerator.gather_for_metrics(chosen_logprobs_sum).mean().item()) + self.stats["logps/rejected"].append(self.accelerator.gather_for_metrics(rejected_logprobs_sum).mean().item()) + + kl = logprobs - ref_logprobs + mean_kl = kl.sum(1).mean() + self.stats["objective/kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + non_score_reward = (-self.beta * kl).sum(1) + mean_non_score_reward = non_score_reward.mean() + self.stats["objective/non_score_reward"].append( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + if self.reward_funcs is not None: + # Calculate RLHF reward by combining rewards with non_score_reward + rlhf_reward = rewards + non_score_reward + self.stats["objective/rlhf_reward"].append(self.accelerator.gather_for_metrics(rlhf_reward).mean().item()) + + mean_entropy = -logprobs.sum(1).mean() + self.stats["objective/entropy"].append(self.accelerator.gather_for_metrics(mean_entropy).mean().item()) + chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum) + gathered_chosen_rewards = self.accelerator.gather_for_metrics(chosen_rewards) + self.stats["rewards/chosen"].append(gathered_chosen_rewards.mean().item()) + rejected_rewards = self.beta * (rejected_logprobs_sum - rejected_ref_logprobs_sum) + gathered_rejected_rewards = self.accelerator.gather_for_metrics(rejected_rewards) + self.stats["rewards/rejected"].append(gathered_rejected_rewards.mean().item()) + margin = gathered_chosen_rewards - gathered_rejected_rewards + self.stats["rewards/margins"].append(margin.mean().item()) + accuracy = margin > 0 + self.stats["rewards/accuracies"].append(accuracy.float().mean().item()) + self.stats["beta"].append(self.beta) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + # Same as Trainer._maybe_log_save_evaluate but log our metrics + def _maybe_log_save_evaluate( + self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None + ): + if self.control.should_log and self.state.global_step > self._globalstep_last_logged: + logs: dict[str, float] = {} + + # all_gather + mean() to get average loss over all processes + tr_loss_scalar = self._nested_gather(tr_loss).mean().item() + + # reset tr_loss to zero + tr_loss -= tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + if grad_norm is not None: + logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm + if learning_rate is not None: + logs["learning_rate"] = learning_rate + else: + logs["learning_rate"] = self._get_learning_rate() + + # Add our metrics + for key, val in self.stats.items(): + logs[key] = sum(val) / len(val) + self.stats = {key: [] for key in self.stats} # reset stats + + self._total_loss_scalar += tr_loss_scalar + self._globalstep_last_logged = self.state.global_step + self.store_flos() + self.log(logs, start_time) + + metrics = None + if self.control.should_evaluate: + metrics = self._evaluate(trial, ignore_keys_for_eval) + is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) + + if self.args.save_strategy == "best": + self.control.should_save = is_new_best_metric + + if self.control.should_save: + self._save_checkpoint(model, trial) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/orpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/orpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..523beeab934e4bc33794cf1b89962d60a011e883 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/orpo_config.py @@ -0,0 +1,169 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class ORPOConfig(TrainingArguments): + r""" + Configuration class for the [`ORPOTrainer`]. + + This class includes only the parameters that are specific to ORPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int`, *optional*): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the relative ratio loss weight in the ORPO loss. In the + [paper](https://huggingface.co/papers/2403.07691), it is denoted by λ. In the + [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int`, *optional*): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool`, *optional*): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: int | None = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: int | None = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: int | None = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the relative ratio loss weight in the ORPO loss. In the paper, it is " + "denoted by λ." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + label_pad_token_id: int = field( + default=-100, + metadata={ + "help": "Label pad token id. This argument is required if you want to use the default data collator." + }, + ) + padding_value: int | None = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={"help": "If `True`, generates and logs completions from the model to W&B during evaluation."}, + ) + is_encoder_decoder: bool | None = field( + default=None, + metadata={ + "help": "When using the `model_init` argument (callable) to instantiate the model instead of the `model` " + "argument, you need to specify if the model returned by the callable is an encoder-decoder model." + }, + ) + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/orpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/orpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..fb905243800068b9b683849d59700f140388ce0b --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/orpo_trainer.py @@ -0,0 +1,1057 @@ +# Copyright 2020-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 inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from collections.abc import Callable +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState, logging +from datasets import Dataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + is_comet_available, + is_torch_xla_available, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_peft_available, is_torch_fx_proxy + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from .base_trainer import BaseTrainer +from .orpo_config import ORPOConfig +from .utils import ( + DPODataCollatorWithPadding, + add_bos_token_if_needed, + add_eos_token_if_needed, + disable_dropout_in_model, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + + +if is_wandb_available(): + import wandb + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + +logger = logging.get_logger(__name__) + + +class ORPOTrainer(BaseTrainer): + r""" + Initialize ORPOTrainer. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. + args ([`ORPOConfig`]): + The ORPO config arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + """ + + _tag_names = ["trl", "orpo"] + _name = "ORPO" + _paper = { + "title": "ORPO: Monolithic Preference Optimization without Reference Model", + "id": "2403.07691", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{hong2024orpo, + title = {{ORPO: Monolithic Preference Optimization without Reference Model}}, + author = {Jiwoo Hong and Noah Lee and James Thorne}, + year = 2024, + eprint = {arXiv:2403.07691} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | str | None = None, + args: ORPOConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + model_init: Callable[[], PreTrainedModel] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: dict | None = None, + compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the ORPOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + dtype = model_init_kwargs.get("dtype") + if dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(dtype, str) and dtype != "auto": + dtype = getattr(torch, dtype) + if dtype != "auto" and not isinstance(dtype, torch.dtype): + raise ValueError( + f"Invalid `dtype` passed to the ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." + ) + model_init_kwargs["dtype"] = dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a ORPO dataset.") + if args.max_length is None: + logger.warning( + "`max_length` is not set in the ORPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + logger.warning( + "`max_prompt_length` is not set in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if args.max_completion_length is None and self.is_encoder_decoder: + logger.warning( + "When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + ) + self.max_completion_length = 128 + else: + self.max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + logger.warning( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.processing_class = processing_class + + self.beta = args.beta + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + logger.warning( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + ) + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in ORPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: PreTrainedModel | nn.Module | None = None) -> dict: + """Tokenize a single row from a ORPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + a != b + for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"], strict=True) + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + if is_torch_xla_available(): + # Pad the sequences to global max_length to avoid TorchXLA recompilation + for k in batch: + if "labels" in k or self.is_encoder_decoder: + pad_value = self.label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = self.padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k])) + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, list | torch.LongTensor], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: torch.device | None = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def odds_ratio_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the ORPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The log odds ratio of the chosen responses over the + rejected responses ratio for logging purposes. The `log(sigmoid(log_odds_chosen))` for logging purposes. + """ + + # Derived from Eqs. (4) and (7) from https://huggingface.co/papers/2403.07691 by using log identities and exp(log(P(y|x)) = P(y|x) + log_odds = (policy_chosen_logps - policy_rejected_logps) - ( + torch.log1p(-torch.exp(policy_chosen_logps)) - torch.log1p(-torch.exp(policy_rejected_logps)) + ) + ratio = F.logsigmoid(log_odds) + losses = self.beta * ratio + + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards, torch.mean(ratio), torch.mean(log_odds) + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels = torch.where(labels == label_pad_token_id, 0, labels) + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, list | torch.LongTensor] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + if self.is_encoder_decoder: + labels = concatenated_batch["concatenated_labels"].clone() + else: + labels = concatenated_batch["concatenated_input_ids"].clone() + attention_mask = concatenated_batch["concatenated_attention_mask"] + labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id) + # orpo chosen nll loss is computed over the full prompt and response + chosen_nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=True, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + if not self.is_encoder_decoder: + chosen_logits = all_logits[:len_chosen, :-1, :] + rejected_logits = all_logits[len_chosen:, :-1, :] + else: + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, list | torch.LongTensor], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the ORPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = self.odds_ratio_loss( + policy_chosen_logps, policy_rejected_logps + ) + # full ORPO loss + loss = policy_nll_loss - losses.mean() + + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean() + metrics[f"{prefix}rewards/margins"] = self.accelerator.gather_for_metrics( + chosen_rewards - rejected_rewards + ).mean() + metrics[f"{prefix}logps/rejected"] = self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean() + metrics[f"{prefix}logps/chosen"] = self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean() + metrics[f"{prefix}logits/rejected"] = self.accelerator.gather_for_metrics( + policy_rejected_logits.detach().mean() + ).mean() + metrics[f"{prefix}logits/chosen"] = self.accelerator.gather_for_metrics( + policy_chosen_logits.detach().mean() + ).mean() + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean() + metrics[f"{prefix}log_odds_ratio"] = self.accelerator.gather_for_metrics(log_odds_ratio).detach().mean() + metrics[f"{prefix}log_odds_chosen"] = self.accelerator.gather_for_metrics(log_odds_chosen).detach().mean() + if is_torch_xla_available(): + xm.mark_step() # needed because .item() calls + for k, v in metrics.items(): + metrics[k] = v.item() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs=False, + num_items_in_batch=None, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_output_decoded + + def prediction_step( + self, + model: PreTrainedModel | nn.Module, + inputs: dict[str, torch.Tensor | Any], + prediction_loss_only: bool, + ignore_keys: list[str] | None = None, + ): + if not self.use_dpo_data_collator: + logger.warning( + "prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than " + "DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator" + ) + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: bool | None = None, + ignore_keys: list[str] | None = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] + for prompt, pol in zip(random_batch["prompt"], policy_output_decoded, strict=True) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float`, *optional*): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/ppo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/ppo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..40d48b82dbf41a6e82f975d12f9fd9f167d22b6b --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/ppo_config.py @@ -0,0 +1,135 @@ +# Copyright 2020-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 os +from dataclasses import dataclass, field +from typing import Literal + +from ..trainer.utils import OnPolicyConfig + + +@dataclass +class PPOConfig(OnPolicyConfig): + r""" + Configuration class for the [`PPOTrainer`]. + + This class includes only the parameters that are specific to PPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default + values in this class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[:-3]`): + Name of this experiment. + reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the reward model. + model_adapter_name (`str`, *optional*): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, *optional*): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + num_ppo_epochs (`int`, *optional*, defaults to `4`): + Number of epochs to train. + whiten_rewards (`bool`, *optional*, defaults to `False`): + Whether to whiten the rewards. + kl_coef (`float`, *optional*, defaults to `0.05`): + KL coefficient. + kl_estimator (`Literal["k1", "k3"]`, *optional*, defaults to `"k1"`): + Which estimator for KL-Divergence to use from [Approximating KL + Divergence](http://joschu.net/blog/kl-approx.html). Defaults to "k1", a straightforward, unbiased + estimator. Can be set to "k3", an unbiased estimator with lower variance which "appears to be a strictly + better estimator". Cannot be set to "k2", as it is used for logging purposes. + cliprange (`float`, *optional*, defaults to `0.2`): + Clip range. + vf_coef (`float`, *optional*, defaults to `0.1`): + Value function coefficient. + cliprange_value (`float`, *optional*, defaults to `0.2`): + Clip range for the value function. + gamma (`float`, *optional*, defaults to `1.0`): + Discount factor. + lam (`float`, *optional*, defaults to `0.95`): + Lambda value for GAE. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + """ + + exp_name: str = field( + default=os.path.basename(__file__)[:-3], + metadata={"help": "Name of this experiment."}, + ) + reward_model_path: str = field( + default="EleutherAI/pythia-160m", + metadata={"help": "Path to the reward model."}, + ) + model_adapter_name: str | None = field( + default=None, + metadata={"help": "Name of the train target PEFT adapter, when using LoRA with multiple adapters."}, + ) + ref_adapter_name: str | None = field( + default=None, + metadata={"help": "Name of the reference PEFT adapter, when using LoRA with multiple adapters."}, + ) + num_ppo_epochs: int = field( + default=4, + metadata={"help": "Number of epochs to train."}, + ) + whiten_rewards: bool = field( + default=False, + metadata={"help": "Whether to whiten the rewards."}, + ) + kl_coef: float = field( + default=0.05, + metadata={"help": "KL coefficient."}, + ) + kl_estimator: Literal["k1", "k3"] = field( + default="k1", + metadata={ + "help": "Which estimator for KL-Divergence to use from Approximating KL Divergence " + "(http://joschu.net/blog/kl-approx.html). Defaults to 'k1', a straightforward, unbiased estimator. Can be " + "set to 'k3', an unbiased estimator with lower variance which 'appears to be a strictly better " + "estimator'. Cannot be set to 'k2', as it is used for logging purposes." + }, + ) + cliprange: float = field( + default=0.2, + metadata={"help": "Clip range."}, + ) + vf_coef: float = field( + default=0.1, + metadata={"help": "Value function coefficient."}, + ) + cliprange_value: float = field( + default=0.2, + metadata={"help": "Clip range for the value function."}, + ) + gamma: float = field( + default=1.0, + metadata={"help": "Discount factor."}, + ) + lam: float = field( + default=0.95, + metadata={"help": "Lambda value for GAE."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation." + }, + ) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/ppo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/ppo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e5182d48f2949e99eeed0fa1632a46f9506d01 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/ppo_trainer.py @@ -0,0 +1,844 @@ +# Copyright 2020-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 gc +import math +import os +import textwrap +import time +import warnings +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +from accelerate import Accelerator, logging +from accelerate.utils import broadcast, gather_object +from datasets import Dataset +from torch.utils.data import DataLoader +from transformers import ( + BaseImageProcessor, + DataCollatorWithPadding, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + TrainerControl, +) +from transformers.integrations import get_reporting_integration_callbacks +from transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK +from transformers.trainer_callback import CallbackHandler, ExportableState, PrinterCallback +from transformers.utils import is_peft_available, is_rich_available + +from ..models import create_reference_model +from ..models.utils import unwrap_model_for_generation +from .base_trainer import BaseTrainer +from .ppo_config import PPOConfig +from .utils import ( + OnlineTrainerState, + batch_generation, + disable_dropout_in_model, + empty_cache, + exact_div, + first_true_indices, + forward, + get_reward, + log_table_to_comet_experiment, + peft_module_casting_to_bf16, + prepare_deepspeed, + print_rich_table, + selective_log_softmax, + truncate_response, +) + + +logger = logging.get_logger(__name__) + +if is_peft_available(): + from peft import PeftConfig, PeftModel, get_peft_model + + +INVALID_LOGPROB = 1.0 + + +def masked_mean(values: torch.Tensor, mask: torch.Tensor, axis: bool | None = None) -> torch.Tensor: + """Compute mean of tensor with a masked values.""" + if axis is not None: + return (values * mask).sum(axis=axis) / mask.sum(axis=axis) + else: + return (values * mask).sum() / mask.sum() + + +def masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) -> torch.Tensor: + """Compute variance of tensor with masked values.""" + mean = masked_mean(values, mask) + centered_values = values - mean + variance = masked_mean(centered_values**2, mask) + if unbiased: + mask_sum = mask.sum() + if mask_sum == 0: + raise ValueError( + "The sum of the mask is zero, which can happen when `mini_batch_size=1`;" + "try increase the `mini_batch_size` or `gradient_accumulation_steps`" + ) + # note that if mask_sum == 1, then there is a division by zero issue + # to avoid it you just need to use a larger minibatch_size + bessel_correction = mask_sum / (mask_sum - 1) + variance = variance * bessel_correction + return variance + + +def masked_whiten(values: torch.Tensor, mask: torch.Tensor, shift_mean: bool = True) -> torch.Tensor: + """Whiten values with masked values.""" + mean, var = masked_mean(values, mask), masked_var(values, mask) + whitened = (values - mean) * torch.rsqrt(var + 1e-8) + if not shift_mean: + whitened += mean + return whitened + + +# taken from https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/ppo/ppo_trainer.py#L29 +# we did this we can do a single `model = accelerator.prepare(model)` +class PolicyAndValueWrapper(nn.Module): + def __init__(self, policy, value_model) -> None: + super().__init__() + self.policy = policy + self.value_model = value_model + self.critic_backbone = getattr(value_model, value_model.base_model_prefix) + self.is_gradient_checkpointing = policy.is_gradient_checkpointing + + def forward(self, **kwargs): + output = self.critic_backbone(**kwargs) + logits = self.value_model.score(output.hidden_states[-1]) + return self.policy(**kwargs), logits + + +class PPOTrainer(BaseTrainer): + """Trainer for Proximal Policy Optimization (PPO). + + For details on PPO, see the paper: [Proximal Policy Optimization + Algorithms](https://huggingface.co/papers/1707.06347). + + Args: + args ([`PPOConfig`]): + Training arguments. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`]): + Class to process the data. + model (`torch.nn.Module`): + Model to be trained. This is the policy model. + ref_model (`torch.nn.Module`, *optional*): + Reference model used to compute the KL divergence. If `None`, a copy of the policy model is created. + reward_model (`torch.nn.Module`): + Reward model used to compute the rewards. + train_dataset ([`~datasets.Dataset`]): + Dataset for training. + value_model (`torch.nn.Module`): + Value model used to predict the value of a state. + data_collator ([`~transformers.DataCollatorWithPadding`], *optional*): + Data collator to batch and pad samples from the dataset. If `None`, a default data collator is created + using the `processing_class`. + eval_dataset ([`~datasets.Dataset`] or `dict` of [`~datasets.Dataset`], *optional*): + Dataset for evaluation. + optimizers (`tuple` of `torch.optim.Optimizer` and `torch.optim.lr_scheduler.LambdaLR`, *optional*, defaults to `(None, None)`): + Tuple containing the optimizer and the learning rate scheduler to use for training. If `None`, the + optimizer and the learning rate scheduler are created using the + [`~transformers.Trainer.create_optimizer_and_scheduler`] method. + callbacks (`list` of [`~transformers.TrainerCallback`], *optional*): + Callbacks to use during training. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration to use PEFT for training. If `None`, PEFT is not used. If provided, the policy `model` + will be wrapped with the specified PEFT adapter. + """ + + _tag_names = ["trl", "ppo"] + _name = "PPO" + _paper = { + "title": "Fine-Tuning Language Models from Human Preferences", + "id": "1909.08593", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{mziegler2019fine-tuning, + title = {{Fine-Tuning Language Models from Human Preferences}}, + author = {Daniel M. Ziegler and Nisan Stiennon and Jeffrey Wu and Tom B. Brown and Alec Radford and Dario Amodei and Paul F. Christiano and Geoffrey Irving}, + year = 2019, + eprint = {arXiv:1909.08593} + }"""), + } + + def __init__( + self, + args: PPOConfig, + processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + model: nn.Module, + ref_model: nn.Module | None, + reward_model: nn.Module, + train_dataset: Dataset, + value_model: nn.Module, + data_collator: DataCollatorWithPadding | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + # less commonly used + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + callbacks: list[TrainerCallback] | None = None, + peft_config: "PeftConfig | None" = None, + ) -> None: + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must make a copy of it, or `None` if you use peft." + ) + + self.args = args + self.processing_class = processing_class + self.policy_model = model + + # Define the collator if not provided + if data_collator is None: + data_collator = DataCollatorWithPadding(self.processing_class) + + # Handle stop token settings: update policy model's generation_config to use provided stop token + if args.stop_token and args.stop_token_id: + raise ValueError("You cannot set both `stop_token` and `stop_token_id`.") + elif args.stop_token: + if args.stop_token == "eos": + self.policy_model.generation_config.eos_token_id = self.stop_token_id = processing_class.eos_token_id + else: + raise ValueError( + f"Unknown `stop_token` {args.stop_token}. Allowed values are: `'eos'` and `None` (no stop token)." + ) + else: + self.policy_model.generation_config.eos_token_id = self.stop_token_id = args.stop_token_id # None or int + + # Check that the kl estimator is valid + if self.args.kl_estimator not in {"k1", "k3"}: + raise ValueError( + "kl_estimator must be either 'k1' (straightforward, unbiased) or 'k3' (lower variance, unbiased, " + "appears to be a strictly better estimator). See " + "[Approximating KL Divergence](http://joschu.net/blog/kl-approx.html) for details." + ) + + # peft support + if not is_peft_available() and peft_config is not None: + raise ImportError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_confg, we merge and unload it first + if isinstance(self.policy_model, PeftModel): + self.policy_model = self.policy_model.merge_and_unload() + + # get peft model with the given config + self.policy_model = get_peft_model(self.policy_model, peft_config) + if args.bf16 and getattr(self.policy_model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(self.policy_model) + + self.is_peft_model = is_peft_available() and isinstance(self.policy_model, PeftModel) + self.model_adapter_name = args.model_adapter_name + self.ref_adapter_name = args.ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model: + self.ref_model = None + else: + self.ref_model = create_reference_model(self.policy_model) + + self.reward_model = reward_model + self.train_dataset = train_dataset + self.train_dataset_len = len(train_dataset) + self.value_model = value_model + self.data_collator = data_collator + self.eval_dataset = eval_dataset + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47 + + ######### + # calculate various batch sizes + ######### + if args.total_episodes is None: # allow the users to define episodes in terms of epochs. + args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) + accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) + self.accelerator = accelerator + args.world_size = accelerator.num_processes + args.local_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps + args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) + args.batch_size = int(args.local_batch_size * args.world_size) + args.mini_batch_size = exact_div( + args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" + ) + args.local_mini_batch_size = exact_div( + args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" + ) + if args.whiten_rewards: + assert args.local_mini_batch_size >= 8, ( + f"Per-rank minibatch size {args.local_mini_batch_size} is insufficient for whitening" + ) + # `per_rank_rollout_batch_size` is our `args.local_batch_size` + # `per_rank_minibatch_size` is our `args.local_mini_batch_size` + args.num_total_batches = math.ceil( + args.total_episodes / args.batch_size + ) # we may train for more than `total_episodes` + time_tensor = torch.tensor(int(time.time()), device=accelerator.device) + time_int = broadcast(time_tensor, 0).item() # avoid different timestamps across processes + args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" + self.local_seed = args.seed + accelerator.process_index * 100003 # Prime + if args.num_sample_generations > 0: + self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) + self.local_dataloader_batch_size = args.local_batch_size + + ######### + # setup model, optimizer, and others + ######### + for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]: + if module is not None: + disable_dropout_in_model(module) + self.model = PolicyAndValueWrapper(self.policy_model, self.value_model) + self.model.config = self.policy_model.config # needed for pushing to hub + self.create_optimizer_and_scheduler( + num_training_steps=args.num_total_batches + ) # note that we are calling `self.lr_scheduler.step()` manually only at the batch level + + ######### + # trainer specifics + ######### + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + self.control = TrainerControl() + self.state = OnlineTrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + self.current_flos = 0 + self.hp_search_backend = None + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + ######### + # setup dataloader + ######### + self.dataloader = DataLoader( + self.train_dataset, + batch_size=self.local_dataloader_batch_size, + shuffle=True, + collate_fn=self.data_collator, + drop_last=True, # needed; otherwise the last batch will be of ragged shape + ) + # sync random states for DataLoader(shuffle=True) before `accelerator.prepare` + # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c + torch.manual_seed(args.seed) + self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) + torch.manual_seed(self.local_seed) # reset the local seed again + + self.eval_dataloader = DataLoader( + self.eval_dataset, + batch_size=args.per_device_eval_batch_size, + collate_fn=self.data_collator, + drop_last=True, + ) # no need to shuffle eval dataset + self.eval_dataloader = accelerator.prepare(self.eval_dataloader) + + if self.is_deepspeed_enabled: + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = prepare_deepspeed( + self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + else: + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = self.ref_model.to(self.accelerator.device) + self.reward_model = self.reward_model.to(self.accelerator.device) + + def get_train_dataloader(self) -> DataLoader: + return self.dataloader + + def get_eval_dataloader(self) -> DataLoader: + return self.eval_dataloader + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model.policy).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.policy.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.policy.set_adapter(self.model_adapter_name or "default") + + def save_model(self, output_dir: str | None = None, _internal_call: bool = False): + backup_model = self.model + self.model = self.model.policy # save only the policy + + if self.is_deepspeed_enabled: + backup_deepspeed = self.deepspeed + self.deepspeed = self.model + + super().save_model(output_dir, _internal_call) + + self.model = backup_model + + if self.is_deepspeed_enabled: + self.deepspeed = backup_deepspeed + + def train(self): + args = self.args + accelerator = self.accelerator + optimizer = self.optimizer + model = self.model + ref_policy = self.ref_model + reward_model = self.reward_model + processing_class = self.processing_class + dataloader = self.dataloader + device = accelerator.device + + def repeat_generator(): + while True: + yield from dataloader + + iter_dataloader = iter(repeat_generator()) + generation_config = GenerationConfig( + max_new_tokens=args.response_length, + temperature=(args.temperature + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + accelerator.print("===training policy===") + start_time = time.time() + stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) + approxkl_stats = torch.zeros(stats_shape, device=device) + pg_clipfrac_stats = torch.zeros(stats_shape, device=device) + pg_loss_stats = torch.zeros(stats_shape, device=device) + vf_loss_stats = torch.zeros(stats_shape, device=device) + vf_clipfrac_stats = torch.zeros(stats_shape, device=device) + entropy_stats = torch.zeros(stats_shape, device=device) + ratio_stats = torch.zeros(stats_shape, device=device) + model.train() + + # trainer state initialization + self.state.global_step = 0 + self.state.episode = 0 + self.state.max_steps = args.num_total_batches + self.state.num_train_epochs = args.total_episodes / self.train_dataset_len + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model + self.model_wrapped = self.model + + for update in range(1, args.num_total_batches + 1): + self.state.episode += 1 * args.batch_size + data = next(iter_dataloader) + with torch.no_grad(): + queries = data["input_ids"].to(device) + context_length = queries.shape[1] + responses = [] + postprocessed_responses = [] + logprobs = [] + ref_logprobs = [] + scores = [] + sequence_lengths = [] + values = [] + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + query_responses, logitss = batch_generation( + unwrapped_model.policy, + queries, + args.local_rollout_forward_batch_size, + processing_class.pad_token_id, + generation_config, + ) + + for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): + query = queries[i : i + args.local_rollout_forward_batch_size] + query_response = query_responses[i : i + args.local_rollout_forward_batch_size] + response = query_response[:, context_length:] + logits = logitss[i : i + args.local_rollout_forward_batch_size] + logprob = selective_log_softmax(logits, response) + del logits + empty_cache() + + if ref_policy is None: + with self.null_ref_context(): + ref_output = forward(model.policy, query_response, processing_class.pad_token_id) + else: + ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) + ref_logits = ref_output.logits[:, context_length - 1 : -1] + ref_logits /= args.temperature + 1e-7 + ref_logprob = selective_log_softmax(ref_logits, response) + del ref_output, ref_logits + empty_cache() + + # Response Processing 1. truncate response after the first occurrence of `stop_token_id` + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + + # Response Processing 2. run reward model on the truncated responses + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 + unwrapped_value_model = accelerator.unwrap_model(model).value_model + full_value, _, _ = get_reward( + unwrapped_value_model, query_response, processing_class.pad_token_id, context_length + ) + value = full_value[:, context_length - 1 : -1].squeeze(-1) + _, score, _ = get_reward( + reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + + responses.append(response) + postprocessed_responses.append(postprocessed_response) + logprobs.append(logprob) + ref_logprobs.append(ref_logprob) + sequence_lengths.append(sequence_length) + scores.append(score) + values.append(value) + responses = torch.cat(responses, 0) + postprocessed_responses = torch.cat(postprocessed_responses, 0) + logprobs = torch.cat(logprobs, 0) + ref_logprobs = torch.cat(ref_logprobs, 0) + sequence_lengths = torch.cat(sequence_lengths, 0) + scores = torch.cat(scores, 0) + values = torch.cat(values, 0) + del (logprob, ref_logprob, full_value, value, score, unwrapped_model) + empty_cache() + gc.collect() + + # Response Processing 3. Filter completion. Ensure that the sample contains stop_token_id + # Completions not passing that filter will receive a lower score. + contain_eos_token = torch.any(postprocessed_responses == self.processing_class.eos_token_id, dim=-1) + if self.args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + # accelerator.print(f"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}") + + # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw + response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) + padding_mask = response_idxs > sequence_lengths.unsqueeze(1) + logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) + ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) + sequence_lengths_p1 = sequence_lengths + 1 + padding_mask_p1 = response_idxs > (sequence_lengths_p1.unsqueeze(1)) + values = torch.masked_fill(values, padding_mask_p1, 0) + + # 4. compute rewards + # Formula used by http://joschu.net/blog/kl-approx.html for the k1 and k3 estimators + logr = ref_logprobs - logprobs + kl = -logr if args.kl_estimator == "k1" else (logr.exp() - 1) - logr # Else statement is k3 + non_score_reward = -args.kl_coef * kl + rewards = non_score_reward.clone() + actual_start = torch.arange(rewards.size(0), device=rewards.device) + actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths) + rewards[actual_start, actual_end] += scores + + # 5. whiten rewards + if args.whiten_rewards: + rewards = masked_whiten(rewards, mask=~padding_mask_p1, shift_mean=False) + rewards = torch.masked_fill(rewards, padding_mask_p1, 0) + + # 6. compute advantages and returns + lastgaelam = 0 + advantages_reversed = [] + gen_length = responses.shape[1] + for t in reversed(range(gen_length)): + nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0 + delta = rewards[:, t] + args.gamma * nextvalues - values[:, t] + lastgaelam = delta + args.gamma * args.lam * lastgaelam + advantages_reversed.append(lastgaelam) + advantages = torch.stack(advantages_reversed[::-1], axis=1) + returns = advantages + values + advantages = masked_whiten(advantages, ~padding_mask) + advantages = torch.masked_fill(advantages, padding_mask, 0) + empty_cache() + + # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch + for ppo_epoch_idx in range(args.num_ppo_epochs): + b_inds = np.random.permutation(args.local_batch_size) + minibatch_idx = 0 + for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): + mini_batch_end = mini_batch_start + args.local_mini_batch_size + mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] + gradient_accumulation_idx = 0 + for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): + with accelerator.accumulate(model): + micro_batch_end = micro_batch_start + args.per_device_train_batch_size + micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] + mb_advantage = advantages[micro_batch_inds] + mb_responses = responses[micro_batch_inds] + mb_query_responses = query_responses[micro_batch_inds] + mb_logprobs = logprobs[micro_batch_inds] + mb_return = returns[micro_batch_inds] + mb_values = values[micro_batch_inds] + + output, vpred_temp = forward(model, mb_query_responses, processing_class.pad_token_id) + logits = output.logits[:, context_length - 1 : -1] + logits /= args.temperature + 1e-7 + new_logprobs = selective_log_softmax(logits, mb_responses) + new_logprobs = torch.masked_fill( + new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB + ) + vpred = vpred_temp[:, context_length - 1 : -1].squeeze(-1) + vpred = torch.masked_fill(vpred, padding_mask_p1[micro_batch_inds], 0) + vpredclipped = torch.clamp( + vpred, + mb_values - args.cliprange_value, + mb_values + args.cliprange_value, + ) + vf_losses1 = torch.square(vpred - mb_return) + vf_losses2 = torch.square(vpredclipped - mb_return) + vf_loss_max = torch.max(vf_losses1, vf_losses2) + vf_loss = 0.5 * masked_mean(vf_loss_max, ~padding_mask_p1[micro_batch_inds]) + vf_clipfrac = masked_mean( + (vf_losses2 > vf_losses1).float(), ~padding_mask_p1[micro_batch_inds] + ) + logprobs_diff = new_logprobs - mb_logprobs + ratio = torch.exp(logprobs_diff) + pg_losses = -mb_advantage * ratio + pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) + pg_loss_max = torch.max(pg_losses, pg_losses2) + pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds]) + loss = pg_loss + args.vf_coef * vf_loss + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + with torch.no_grad(): + pg_clipfrac = masked_mean( + (pg_losses2 > pg_losses).float(), ~padding_mask[micro_batch_inds] + ) + prob_dist = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) + approxkl = 0.5 * (logprobs_diff**2).mean() + approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl + pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + pg_clipfrac + ) + pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss + vf_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = vf_loss + vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + vf_clipfrac + ) + entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() + ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean() + gradient_accumulation_idx += 1 + minibatch_idx += 1 + # del everything and empty cache + # fmt: off + del ( + output, vpred_temp, logits, new_logprobs, vpred, vpredclipped, + vf_losses1, vf_losses2, vf_loss, vf_clipfrac, logprobs_diff, ratio, pg_losses, pg_losses2, pg_loss_max, + pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, mb_return, + mb_advantage, mb_values, mb_responses, mb_query_responses, mb_logprobs, + ) + # fmt: on + empty_cache() + with torch.no_grad(): + mean_kl = kl.sum(1).mean() + mean_entropy = (-logprobs).sum(1).mean() + mean_non_score_reward = non_score_reward.sum(1).mean() + rlhf_reward = mean_non_score_reward + scores.mean() + eps = int(self.state.episode / (time.time() - start_time)) + metrics = {} + metrics["eps"] = eps + metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() + metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() + metrics["objective/non_score_reward"] = ( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() + metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() + metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() + metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() + metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() + metrics["loss/value_avg"] = self.accelerator.gather_for_metrics(vf_loss_stats).mean().item() + metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() + metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() + metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() + metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() + metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() + metrics["lr"] = self.lr_scheduler.get_last_lr()[0] + metrics["episode"] = self.state.episode + self.state.epoch = self.state.episode / self.train_dataset_len # used by self.log + self.state.global_step += 1 + self.log(metrics) + + self.lr_scheduler.step() + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + del kl, mean_kl, mean_entropy, mean_non_score_reward, scores, metrics, non_score_reward + empty_cache() + gc.collect() + + if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: + self.generate_completions(sampling=True) + empty_cache() + del ( + query_responses, + responses, + postprocessed_responses, + logprobs, + ref_logprobs, + values, + sequence_lengths, + contain_eos_token, + sequence_lengths_p1, + response_idxs, + padding_mask, + padding_mask_p1, + rewards, + actual_start, + actual_end, + advantages, + returns, + ) + empty_cache() + + # HF trainer specifics + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def generate_completions(self, sampling: bool = False): + args = self.args + processing_class = self.processing_class + generation_config = GenerationConfig( + max_new_tokens=self.args.response_length, + temperature=(0.01 + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + table = defaultdict(list) + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + for batch in self.eval_dataloader: + query = batch["input_ids"] + with torch.no_grad(): + context_length = query.shape[1] + query_response, _ = batch_generation( + unwrapped_model.policy, + query, + query.shape[0], + processing_class.pad_token_id, + generation_config, + ) + response = query_response[:, context_length:] + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + table["query"].extend( + gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) + ) + table["model response"].extend( + gather_object(processing_class.batch_decode(postprocessed_response)) + ) + + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + _, score, _ = get_reward( + self.reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) + + if sampling: + break + df = pd.DataFrame(table) + + if self.accelerator.is_main_process: + if is_rich_available(): + print_rich_table(df.iloc[0 : 0 + 5]) + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/prm_config.py b/venv/lib/python3.10/site-packages/trl/trainer/prm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..4d0da1f931c93c77e3362d202bd1376a26055908 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/prm_config.py @@ -0,0 +1,112 @@ +# Copyright 2020-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 dataclasses import dataclass, field + +from transformers import TrainingArguments + + +@dataclass +class PRMConfig(TrainingArguments): + r""" + Configuration class for the [`PRMTrainer`]. + + This class includes only the parameters that are specific to PRM training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) used for truncation. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt used for truncation. + max_completion_length (`int`, *optional*): + Maximum length of the completion used for truncation. The completion is the concatenation of the steps. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + step_separator (`str`, *optional*, defaults to `"\n"`): + Separator used to separate each step of the reasoning process. + train_on_last_step_only (`bool`, *optional*, defaults to `False`): + Whether to train only on the last step. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + """ + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-5, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: int | None = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) used for truncation."}, + ) + max_prompt_length: int | None = field( + default=512, + metadata={"help": "Maximum length of the prompt used for truncation."}, + ) + max_completion_length: int | None = field( + default=None, + metadata={ + "help": "Maximum length of the completion used for truncation. The completion is the concatenation of the " + "steps." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + step_separator: str = field( + default="\n", + metadata={"help": "Separator used to separate each step of the reasoning process."}, + ) + train_on_last_step_only: bool = field( + default=False, + metadata={"help": "Whether to train only on the last step."}, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/prm_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/prm_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..fc2d80058b9c36becf6b0b7dee10b2f69d60d55d --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/prm_trainer.py @@ -0,0 +1,311 @@ +# Copyright 2020-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 os +import textwrap +import warnings +from collections.abc import Callable +from itertools import chain +from pathlib import Path + +import torch +import torch.nn as nn +from accelerate import PartialState +from datasets import Dataset, features +from transformers import ( + BaseImageProcessor, + DataCollator, + DataCollatorForTokenClassification, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from ..models import prepare_peft_model +from .base_trainer import BaseTrainer +from .prm_config import PRMConfig +from .utils import compute_accuracy, disable_dropout_in_model + + +if is_peft_available(): + from peft import PeftModel + + +class PRMTrainer(BaseTrainer): + """ + Initialize PRMTrainer. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an `AutoModelForTokenClassification`. + args ([`PRMConfig`]): + The arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`~transformers.DataCollatorForTokenClassification`]) will be used which will pad the sequences to the + maximum length of the sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + compute_metrics (`Callable[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + """ + + _tag_names = ["trl", "prm"] + _name = "PRM" + _paper = { + "title": "Solving math word problems with process-and outcome-based feedback", + "id": "2211.14275", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{uesato2022solving, + title = {{Solving Math Word Problems With Process- and Outcome-Based Feedback}}, + author = {Uesato, Jonathan and Kushman, Nate and Kumar, Ramana and Song, Francis and Siegel, Noah and Wang, Lisa and Creswell, Antonia and Irving, Geoffrey and Higgins, Irina}, + year = 2022, + journal = {arXiv preprint arXiv:2211.14275} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | None = None, + args: PRMConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + model_init: Callable[[], PreTrainedModel] | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: dict | None = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default DataCollatorForTokenClassification" + ) + data_collator = DataCollatorForTokenClassification(processing_class) + + if "input_ids" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = { + "tokenizer": processing_class, + "step_separator": args.step_separator, + "max_length": args.max_length, + "max_prompt_length": args.max_prompt_length, + "max_completion_length": args.max_completion_length, + "train_on_last_step_only": args.train_on_last_step_only, + } + train_fn_kwargs = {**fn_kwargs, "is_eval": False} + train_dataset = train_dataset.map( + self.tokenize_row, + fn_kwargs=train_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=train_dataset.features, + desc="Tokenizing train dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + eval_fn_kwargs = {**fn_kwargs, "is_eval": True} + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + self.tokenize_row, + fn_kwargs=eval_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=eval_dataset.features, + desc="Tokenizing eval dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + @staticmethod + def tokenize_row( + features, + tokenizer, + step_separator, + max_length, + max_prompt_length, + max_completion_length, + train_on_last_step_only, + is_eval, + ): + r""" + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"completions"`, and `"labels"`. + tokenizer ([`~transformers.PreTrainedTokenizerBase`]): + Tokenizer used to process the data. + step_separator (`str`): + Separator between steps in the completion. + max_length (`int` or `None`): + Maximum length of the sequences (prompt + completion). If `None`, the sequences are not truncated. + max_prompt_length (`int` or `None`): + Maximum length of the prompt. If `None`, the prompt is not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + train_on_last_step_only (`bool`): + Whether to train only on the last step. If `True`, the labels are `-100` for all tokens except the last + token of the completion. + is_eval (`bool`): + Whether the function is used to tokenize samples from a training or an evaluation dataset. Used only if + `train_on_last_step_only` is set to `True`. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"input_ids"`, and `"labels". + + Example: + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B") + >>> features = { + ... "prompt": "Which number is larger, 9.8 or 9.11?", + ... "completions": ["11 is greater than 8.", "Hence, 9.11 > 9.8."], + ... "labels": [True, False], + ... } + >>> PRMTrainer.tokenize_row( + ... features, tokenizer, "\n", max_completion_length=None, train_on_last_step_only=False, is_eval=False + ... ) + {'input_ids': [23085, 1372, 374, 8131, 11, 220, 24, 13, 23, 476, 220, 24, 13, 16, 16, 30, 16, 16, 374, 7046, 1091, 220, 23, 13, 198, 39, 763, 11, 220, 24, 13, 16, 16, 861, 220, 24, 13, 23, 13, 198], + 'labels': [-100, -100, -100, -100, -100, -100, -100, -100, 1, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 0]} + ``` + """ + # Tokenize the prompt and completions + prompt_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"] + completions_ids = [ + tokenizer(completion, add_special_tokens=False)["input_ids"] for completion in features["completions"] + ] + if train_on_last_step_only and not is_eval: + labels = [-100] * (len(features["labels"]) - 1) + [int(features["labels"][-1])] + else: + labels = [int(label) for label in features["labels"]] + + # Get the ID of the separator token and add it to the completions + separator_ids = tokenizer.encode(step_separator, add_special_tokens=False) + completions_ids = [completion + separator_ids for completion in completions_ids] + + # Create the label + labels = [ + [-100] * (len(completion) - 1) + [label] for completion, label in zip(completions_ids, labels, strict=True) + ] + + # Join the completions and labels steps + completion_ids = list(chain(*completions_ids)) + labels = list(chain(*labels)) + + if tokenizer.bos_token_id is not None: + prompt_ids = [tokenizer.bos_token_id] + prompt_ids + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_ids = prompt_ids[-max_prompt_length:] + if max_completion_length is not None: + completion_ids = completion_ids[:max_completion_length] + labels = labels[:max_completion_length] + + input_ids = prompt_ids + completion_ids + labels = [-100] * len(prompt_ids) + labels + + if max_length is not None: + input_ids = input_ids[:max_length] + labels = labels[:max_length] + + return {"input_ids": input_ids, "labels": labels} + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/reward_config.py b/venv/lib/python3.10/site-packages/trl/trainer/reward_config.py new file mode 100644 index 0000000000000000000000000000000000000000..cb2a694b604885d89c4896e5b7674296eeac6937 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/reward_config.py @@ -0,0 +1,171 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class RewardConfig(TrainingArguments): + r""" + Configuration class for the [`RewardTrainer`]. + + This class includes only the parameters that are specific to Reward training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`RewardTrainer`] is provided as a string. If you're training a MoE architecture and want + to include the load balancing/auxilliary loss as a part of the final loss, remember to set + `output_router_logits=True` in this dictionary. + chat_template_path (`str`, *optional*): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + + > Parameters that control the data preprocessing + + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + eos_token (`str`, *optional*): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + pad_token (`str`, *optional*): + Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that is also `None`, + it falls back to `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Samples are filtered out if either chosen or rejected sequence + exceeds this value. If `None`, no filtering is applied. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + + > Parameters that control the training + + center_rewards_coefficient (`float`, *optional*): + Coefficient to incentivize the reward model to output mean-zero rewards (proposed by + https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`. + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-4, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `RewardTrainer` is provided as a string." + }, + ) + chat_template_path: str | None = field( + default=None, + metadata={ + "help": "If specified, sets the model's chat template. This can either be the path to a tokenizer (local " + "directory or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, " + "you must ensure that any special tokens referenced in the template are added to the tokenizer and " + "that the model's embedding layer is resized accordingly." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + + # Parameters that control the data preprocessing + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + eos_token: str | None = field( + default=None, + metadata={ + "help": "Token used to indicate the end of a turn or sequence. If `None`, it defaults to `processing_class.eos_token`." + }, + ) + pad_token: str | None = field( + default=None, + metadata={ + "help": "Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that " + "is also `None`, it falls back to `processing_class.eos_token`." + }, + ) + max_length: int | None = field( + default=1024, + metadata={ + "help": "Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from" + "the right. If `None`, no truncation is applied." + }, + ) + pad_to_multiple_of: int | None = field( + default=None, + metadata={"help": "If set, the sequences will be padded to a multiple of this value."}, + ) + + # Parameters that control the training + center_rewards_coefficient: float | None = field( + default=None, + metadata={ + "help": "Coefficient to incentivize the reward model to output mean-zero rewards (proposed by " + "https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`." + }, + ) + activation_offloading: bool = field( + default=False, + metadata={"help": "Whether to offload the activations to the CPU."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/reward_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/reward_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..60b3f5a25444c0131d3a62a47c00185fb29bac6d --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/reward_trainer.py @@ -0,0 +1,605 @@ +# Copyright 2020-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 contextlib +import logging +import os +import re +from collections import defaultdict +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +import transformers +from accelerate import PartialState +from accelerate.logging import get_logger +from datasets import Dataset, IterableDataset +from transformers import ( + AutoModelForSequenceClassification, + AutoTokenizer, + DataCollator, + PreTrainedModel, + PreTrainedTokenizerBase, +) +from transformers.data.data_collator import DataCollatorMixin +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from ..data_utils import is_conversational +from ..models import clone_chat_template, get_act_offloading_ctx_manager, prepare_peft_model +from .base_trainer import BaseTrainer +from .reward_config import RewardConfig +from .utils import disable_dropout_in_model, get_config_model_id, pad, remove_none_values + + +if is_peft_available(): + from peft import PeftConfig, PeftModel + + +logger = get_logger(__name__) + + +# AutoModelForSequenceClassification adds a new classification head when loading a CausalLM. That head is randomly +# initialized and triggers a harmless warning about uninitialized weights. We suppress just that specific warning to +# avoid confusing users. +@contextmanager +def suppress_from_pretrained_warning(logger: logging.Logger): + pattern = re.compile( + r"^Some weights of \S+ were not initialized from the model checkpoint at \S+ and are newly initialized: " + r"\[.*\]\nYou should probably TRAIN this model on a down-stream task to be able to use it for predictions and " + r"inference\.$" + ) + + class _Filter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + return not pattern.search(record.getMessage()) + + f = _Filter() + logger.addFilter(f) + try: + yield + finally: + logger.removeFilter(f) + + +@dataclass +class DataCollatorForPreference(DataCollatorMixin): + """ + Data collator used for preference data. Inputs are dynamically padded to the maximum length of a batch. + + This collator expects each example in the input list to be a dictionary containing the `"chosen_input_ids"` and + `"rejected_input_ids"` keys. The collator returns a dictionary containing the following keys: + - `"input_ids"`: Tensor of input IDs, padded to the maximum length of the batch. The first half of the batch + corresponds to the `"chosen_input_ids"` and the second half to the `"rejected_input_ids"`. + - `"attention_mask"`: Tensor of attention mask, padded to the maximum length of the batch. + + Optionally, the examples can contain a `"margin"` key, in which case the returned dictionary will also contain a + `"margin"` key with a tensor of margins. + + Args: + pad_token_id (`int`): + Token ID to use for padding. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + return_tensors (`str`, *optional*, defaults to `"pt"`): + Type of Tensor to return. Only `"pt"` is currently supported. + + Examples: + ```python + >>> from trl.trainer.reward_trainer import DataCollatorForPreference + + >>> collator = DataCollatorForPreference(pad_token_id=0) + >>> examples = [ + ... {"chosen_input_ids": [1, 2, 3], "rejected_input_ids": [4, 5]}, + ... {"chosen_input_ids": [6, 7], "rejected_input_ids": [8]}, + ... ] + >>> collator(examples) + {'input_ids': tensor([[1, 2, 3], + [6, 7, 0], + [4, 5, 0], + [8, 0, 0]]), + 'attention_mask': tensor([[1, 1, 1], + [1, 1, 0], + [1, 1, 0], + [1, 0, 0]])} + + >>> examples = [ + ... {"chosen_input_ids": [1, 2, 3], "rejected_input_ids": [4, 5], "margin": 0.5}, + ... {"chosen_input_ids": [6, 7], "rejected_input_ids": [8], "margin": 0.0}, + ... ] + >>> collator(examples) + {'input_ids': tensor([[1, 2, 3], + [6, 7, 0], + [4, 5, 0], + [8, 0, 0]]), + 'attention_mask': tensor([[1, 1, 1], + [1, 1, 0], + [1, 1, 0], + [1, 0, 0]]), + 'margin': tensor([0.5, 0.0])} + ``` + """ + + pad_token_id: int + pad_to_multiple_of: int | None = None + return_tensors: str = "pt" + + def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + # Convert to tensor + chosen_input_ids = [torch.tensor(example["chosen_input_ids"]) for example in examples] + rejected_input_ids = [torch.tensor(example["rejected_input_ids"]) for example in examples] + if "margin" in examples[0]: + margins = torch.tensor([example["margin"] for example in examples], dtype=torch.float) + input_ids = chosen_input_ids + rejected_input_ids + attention_mask = [torch.ones_like(ids) for ids in input_ids] + + output = {} + + # Pad + output["input_ids"] = pad( + input_ids, + padding_value=self.pad_token_id, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ) + output["attention_mask"] = pad( + attention_mask, + padding_value=0, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ) + if "margin" in examples[0]: + output["margin"] = margins + return output + + +class RewardTrainer(BaseTrainer): + """ + Trainer for Outcome-supervised Reward Models (ORM). + + This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + from trl import RewardTrainer + from datasets import load_dataset + + dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") + + trainer = RewardTrainer(model="Qwen/Qwen2.5-0.5B-Instruct", train_dataset=dataset) + trainer.train() + ``` + + Args: + model (`str | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `AutoModelForSequenceClassification.from_pretrained` with the keyword arguments in + `args.model_init_kwargs`. + - A sequence classification [`~transformers.PreTrainedModel`] object. + args ([`RewardConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~trainer.reward_trainer.DataCollatorForPreference`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. This trainer supports [preference](#preference) type (both implicit and + explicit prompt). The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain an `chosen_input_ids` and + `rejected_input_ids` fields. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*): + Tokenizer used to process the data. If `None`, the tokenizer is loaded from the model's name with + [`~transformers.AutoTokenizer.from_pretrained`]. A padding token, `processing_class.pad_token`, must be + set. If the processing class has not set a padding token, `processing_class.eos_token` will be used as the + default. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a + [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing + [`RewardConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a + boolean `compute_result` argument. This will be triggered after the last eval batch to signal that the + function needs to calculate and return the global summary statistics rather than accumulating the + batch-level statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. Note that if the loaded + model is a causal LM, it's highly recommended to set `modules_to_save=["score"]` in the PEFT configuration + to ensure that the reward head is properly trained. + """ + + _tag_names = ["trl", "reward-trainer"] + _name = "Reward" + _template_file = "rm_model_card.md" + + def __init__( + self, + model: str | PreTrainedModel, + args: RewardConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = RewardConfig(f"{model_name}-Reward") + + # Model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + dtype = model_init_kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass # dtype is already a torch.dtype or "auto" or None + elif isinstance(dtype, str) and dtype in ["bfloat16", "float16", "float32"]: + model_init_kwargs["dtype"] = getattr(torch, dtype) + else: + raise ValueError( + "Invalid `dtype` passed to `RewardConfig`. Expected either 'auto' or a string representing " + f"a valid `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + with suppress_from_pretrained_warning(transformers.modeling_utils.logger): + model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=1, **model_init_kwargs) + else: + model_id = get_config_model_id(model.config) + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `RewardConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + + # Processing class + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + # Handle pad token for processors or tokenizers + if args.eos_token is not None: + eos_token = args.eos_token + eos_token_id = processing_class.convert_tokens_to_ids(eos_token) + if eos_token_id is None: + raise ValueError( + f"The specified `eos_token` ('{eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + processing_class.eos_token_id = eos_token_id + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + added_tokens = [] + else: + model, processing_class, added_tokens = clone_chat_template( + model, processing_class, args.chat_template_path + ) + else: + added_tokens = [] + + # PEFT configuration and model wrapping + if peft_config is not None: + if added_tokens: + # Ensure that the added tokens are trainable + if peft_config.trainable_token_indices is None: + peft_config.trainable_token_indices = {"embed_tokens": added_tokens} + elif "embed_tokens" not in peft_config.trainable_token_indices: + peft_config.trainable_token_indices["embed_tokens"] = added_tokens + else: + peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) + + # Ensure that the lm_head is trainable + if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: + logger.warning( + "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " + "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " + "tokens, leading to degraded generation quality. To fix this, add " + "`modules_to_save=['lm_head']` to your PEFT configuration." + ) + + if peft_config.modules_to_save is None: + peft_config.modules_to_save = ["lm_head"] + else: + peft_config.modules_to_save.append("lm_head") + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + # Pad token (needed for SequenceClassification models) + # If not provided, use the one from the processing class or the eos token if the processing class does not have + # a pad token. + pad_token = args.pad_token or processing_class.pad_token or processing_class.eos_token + pad_token_id = processing_class.convert_tokens_to_ids(pad_token) + if pad_token_id is None: + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + model.config.pad_token_id = pad_token_id + processing_class.pad_token_id = pad_token_id + + # Data collator + if data_collator is None: + data_collator = DataCollatorForPreference( + pad_token_id=pad_token_id, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + # Dataset + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Initialize the Trainer. Parent class will handle: + # - DeepSpeed configuration (through create_accelerator_and_postprocess) + # - FSDP setup + # - Distributed training setup + # - Optimizer and scheduler creation + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # During evaluation, Trainer calls compute_loss() only if can_return_loss is True and label_names is empty. + self.can_return_loss = True + self.label_names = [] + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase, + args: RewardConfig, + dataset_name: str, + ) -> Dataset | IterableDataset: + # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from + # sampled data. + if isinstance(dataset, Dataset): # IterableDataset does not support `with_transform` + dataset = dataset.with_transform(remove_none_values) + + # If the dataset is already preprocessed (tokenized), skip the processing steps. + column_names = list(next(iter(dataset)).keys()) + is_processed = "chosen_input_ids" in column_names and "rejected_input_ids" in column_names + + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + if not is_processed: + # Add EOS token to the end of the sequences if needed + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if not example["chosen"].endswith(eos_token): + example["chosen"] = example["chosen"] + eos_token + if "rejected" in example and not example["rejected"].endswith(eos_token): + example["rejected"] = example["rejected"] + eos_token + return example + + dataset = dataset.map( + add_eos, + fn_kwargs={"eos_token": processing_class.eos_token}, + **map_kwargs, + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize_fn(example, processing_class): + if "prompt" in example: # explicit prompt case + example["chosen"] = example["prompt"] + example["chosen"] + example["rejected"] = example["prompt"] + example["rejected"] + + if is_conversational(example): + chosen_input_ids = processing_class.apply_chat_template( + example["chosen"], + tools=example.get("tools"), + return_dict=True, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + rejected_input_ids = processing_class.apply_chat_template( + example["rejected"], + tools=example.get("tools"), + return_dict=True, + **example.get("chat_template_kwargs", {}), + )["input_ids"] + output = {"chosen_input_ids": chosen_input_ids, "rejected_input_ids": rejected_input_ids} + else: + output = { + "chosen_input_ids": processing_class(text=example["chosen"])["input_ids"], + "rejected_input_ids": processing_class(text=example["rejected"])["input_ids"], + } + return output + + dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) + + # Filter samples that are longer than `max_length` + if args.max_length is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens" + dataset = dataset.filter( + lambda example: len(example["chosen_input_ids"]) <= args.max_length + and len(example["rejected_input_ids"]) <= args.max_length, + **map_kwargs, + ) + + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). + if self._signature_columns is None: + self._signature_columns = ["chosen_input_ids", "rejected_input_ids", "margin"] + + def compute_loss( + self, + model: nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs: bool = False, + num_items_in_batch: torch.Tensor | None = None, + ): + """ + Compute training loss and additionally compute token accuracies + """ + mode = "train" if self.model.training else "eval" + + # If not set, defaults from model config and may warn since cache isn't compatible with gradient checkpointing + inputs["use_cache"] = False + outputs = model(**inputs) + + # Split the rewards into chosen and rejected + rewards_chosen, rewards_rejected = torch.chunk(outputs.logits.squeeze(-1), chunks=2) + + # Calculate loss, optionally modulate with margin + if "margin" in inputs: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean() + else: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() + + if self.args.center_rewards_coefficient is not None: + loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) + + if mode == "train": + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + # Compute min, mean, max, accuracy and margin + with torch.no_grad(): + all_rewards = self.accelerator.gather(outputs.logits) + self._metrics[mode]["min_reward"].append(all_rewards.min().item()) + self._metrics[mode]["mean_reward"].append(all_rewards.mean().item()) + self._metrics[mode]["max_reward"].append(all_rewards.max().item()) + + mean_accuracy = (rewards_chosen > rewards_rejected).float().mean() + mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item() + self._metrics[mode]["accuracy"].append(mean_accuracy) + + mean_margin = (rewards_chosen - rewards_rejected).mean() + mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean() + self._metrics[mode]["margin"].append(mean_margin.item()) + + return (loss, outputs) if return_outputs else loss + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/rloo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/rloo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..678b4a57a369562d67ebd21ce600ca20604f3395 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/rloo_config.py @@ -0,0 +1,578 @@ +# Copyright 2020-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 dataclasses import dataclass, field + +from transformers import TrainingArguments + + +@dataclass +class RLOOConfig(TrainingArguments): + r""" + Configuration class for the [`RLOOTrainer`]. + + This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`RLOOTrainer`] is provided as a string. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left. + num_generations (`int` or `None`, *optional*, defaults to `2`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + + > Parameters that control generation + + generation_batch_size: (`int`, *optional*): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one + generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. + steps_per_generation: (`int`, *optional*): + Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive + with `generation_batch_size`. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + chat_template_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + use_transformers_paged (`bool`, *optional*, defaults to `False`): + Whether to use the `transformers` paged implementation for generation. If set to `True`, the `transformers` + paged implementation will be used for generation instead of the default padded implementation. This + parameter is only effective when `use_vllm` is set to `False`. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_guided_decoding_regex (`str`, *optional*): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step and woken + for weight sync and generation. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.05`): + KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving training + speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + epsilon_high (`float`, *optional*): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + reward_weights (`list[float]`, *optional*): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + normalize_advantages (`bool`, *optional*, defaults to `False`): + Whether to normalize advantages. Normalization is done per generation batch to have mean `0.0` and standard + deviation of `1.0`. + reward_clip_range (`tuple[float, float]`, *optional*): + Clip range for rewards as (min, max). If `None`, no clipping is applied. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or + `trackio`. + num_completions_to_print (`int`, *optional*): + Number of completions to print with `rich`. If `None`, all completions are logged. + wandb_log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, all prompts + are logged. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: dict | str | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` " + "argument of the `RLOOTrainer` is provided as a string." + }, + ) + disable_dropout: bool = field( + default=False, + metadata={ + "help": "Whether to disable dropout in the model. This is useful for training with a reference model, as " + "it prevents the model from generating different logprobs for the same input." + }, + ) + + # Parameters that control the data preprocessing + # The default value remove_unused_columns is overwritten from the parent class, because in RLOO we usually rely on + # additional columns to compute the reward + remove_unused_columns: bool | None = field( + default=False, + metadata={ + "help": "Whether to only keep the column 'prompt' in the dataset. If you use a custom reward function " + "that requires any column other than 'prompts' and 'completions', you should keep this to `False`." + }, + ) + max_prompt_length: int | None = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left." + }, + ) + num_generations: int | None = field( + default=2, + metadata={ + "help": "Number of generations to sample. The effective batch size (num_processes * per_device_batch_size " + "* gradient_accumulation_steps) must be evenly divisible by this value." + }, + ) + max_completion_length: int | None = field( + default=256, + metadata={"help": "Maximum length of the generated completion."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option " + "is not compatible with vLLM generation." + }, + ) + shuffle_dataset: bool | None = field( + default=True, + metadata={"help": "Whether to shuffle the training dataset."}, + ) + + # Parameters that control generation + generation_batch_size: int | None = field( + default=None, + metadata={ + "help": "Batch size to use for generation. If `None`, it defaults to the effective training batch size: " + "`per_device_train_batch_size * num_processes * steps_per_generation`." + }, + ) + steps_per_generation: int | None = field( + default=None, + metadata={"help": "Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`."}, + ) + temperature: float = field( + default=1.0, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=1.0, + metadata={ + "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. " + "Set to 1.0 to consider all tokens." + }, + ) + top_k: int | None = field( + default=None, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + min_p: float | None = field( + default=None, + metadata={ + "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It " + "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range." + }, + ) + generation_kwargs: dict | None = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or " + "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the " + "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that " + "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them." + }, + ) + chat_template_kwargs: dict | None = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to the `apply_chat_template` function when generating " + "completions." + }, + ) + repetition_penalty: float = field( + default=1.0, + metadata={ + "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated " + "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model " + "to repeat tokens." + }, + ) + use_transformers_paged: bool = field( + default=False, + metadata={ + "help": "Whether to use the `transformers` paged implementation for generation. If set to `True`, the " + "`transformers` paged implementation will be used for generation instead of the default padded " + "implementation. This parameter is only effective when `use_vllm` is set to `False`." + }, + ) + cache_implementation: str | None = field( + default=None, + metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."}, + ) + + # Parameters that control generation acceleration powered by vLLM + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for " + "generation instead of the default model.generate(). Requires `vllm` to be installed." + }, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `'server'` or " + "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure " + "a TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same " + "process and share the training GPUs. This avoids the need for a separate server but may cause resource " + "contention with training." + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={ + "help": "Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: " + "Use the `transformers` backend for model implementation. `vllm`: Use the `vllm` library for " + "model implementation." + }, + ) + vllm_enable_sleep_mode: bool = field( + default=False, + metadata={ + "help": "Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step " + "and woken for weight sync and generation." + }, + ) + vllm_guided_decoding_regex: str | None = field( + default=None, + metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, + ) + + # Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_base_url: str | None = field( + default=None, + metadata={ + "help": "Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` " + "and `vllm_server_port` are ignored." + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": "Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_port: int = field( + default=8000, + metadata={"help": "Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={ + "help": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up " + "after the timeout, a `ConnectionError` is raised." + }, + ) + + # Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + vllm_gpu_memory_utilization: float = field( + default=0.3, + metadata={ + "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag." + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={ + "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_tensor_parallel_size` flag." + }, + ) + + # Parameters that control the training + beta: float = field( + default=0.05, + metadata={ + "help": "KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving " + "training speed." + }, + ) + num_iterations: int = field( + default=1, + metadata={"help": "Number of iterations per batch (denoted as μ in the algorithm)."}, + ) + epsilon: float = field( + default=0.2, + metadata={"help": "Epsilon value for clipping."}, + ) + epsilon_high: float | None = field( + default=None, + metadata={ + "help": "Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the " + "lower-bound specified in argument `epsilon`. Paper DAPO recommends `0.28`." + }, + ) + reward_weights: list[float] | None = field( + default=None, + metadata={ + "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all " + "rewards are weighted equally with weight `1.0`." + }, + ) + normalize_advantages: bool = field( + default=False, + metadata={ + "help": "Whether to normalize advantages. Normalization is done per generation batch to have mean `0.0` " + "and standard deviation of `1.0`." + }, + ) + reward_clip_range: tuple[float, float] | None = field( + default=None, + metadata={"help": "Clip range for rewards as (min, max). If None, no clipping is applied."}, + ) + mask_truncated_completions: bool = field( + default=False, + metadata={ + "help": "When enabled, truncated completions are excluded from the loss calculation, preventing them from " + "being incorrectly penalized and introducing noise during training. According to the DAPO paper, this is " + "a good practice for training stability." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + num_completions_to_print: int | None = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_log_unique_prompts: bool | None = field( + default=False, + metadata={ + "help": "Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, " + "all prompts are logged." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + num_processes = self.world_size + # The current default effective batch size + if self.generation_batch_size is None and self.steps_per_generation is None: + self.steps_per_generation = self.gradient_accumulation_steps + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + elif self.generation_batch_size is not None and self.steps_per_generation is None: + # Just ensure the value is divisible by the global batch size + if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size " + f"({self.per_device_train_batch_size * num_processes})." + ) + self.steps_per_generation = self.generation_batch_size // ( + self.per_device_train_batch_size * num_processes + ) + elif self.generation_batch_size is None and self.steps_per_generation is not None: + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + else: + raise ValueError( + "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time" + ) + + if self.do_eval and self.eval_strategy != "no": + # Just ensure the value is divisible by the global batch size + if (self.per_device_eval_batch_size * num_processes) % self.num_generations != 0: + raise ValueError( + f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be " + f"divisible by num_generations ({self.num_generations})." + ) + + # The generation batch must contain full prompt groups (no partials), so it must be divisible by + # num_generations. + if self.generation_batch_size % self.num_generations != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations " + f"({self.num_generations})." + ) + + if self.num_generations < 2: + raise ValueError( + "RLOO requires at least 2 generations per prompt to calculate the advantages. You provided " + f"{self.num_generations}, which is less than the minimum required." + ) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/rloo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/rloo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..211ccc230f3898cf465485fed7ef76c0ac2b423a --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/rloo_trainer.py @@ -0,0 +1,1581 @@ +# Copyright 2020-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 inspect +import os +import textwrap +import warnings +from collections import defaultdict, deque +from collections.abc import Callable +from contextlib import nullcontext +from functools import partial +from pathlib import Path +from typing import Any + +import datasets +import pandas as pd +import torch +import torch.utils.data +import transformers +from accelerate import logging +from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed +from datasets import Dataset, IterableDataset +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.utils.data import DataLoader, Sampler +from transformers import ( + AutoConfig, + AutoModelForSequenceClassification, + AutoProcessor, + AutoTokenizer, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + is_trackio_available, + is_wandb_available, +) +from transformers.trainer_utils import seed_worker +from transformers.utils import is_datasets_available, is_peft_available, is_rich_available + +from ..data_utils import ( + apply_chat_template, + is_conversational, + prepare_multimodal_messages, + prepare_multimodal_messages_vllm, +) +from ..extras.profiling import profiling_context, profiling_decorator +from ..extras.vllm_client import VLLMClient +from ..import_utils import is_vllm_available +from ..models import prepare_deepspeed, prepare_fsdp, prepare_peft_model, unwrap_model_for_generation +from .base_trainer import BaseTrainer +from .callbacks import SyncRefModelCallback +from .rloo_config import RLOOConfig +from .utils import ( + RepeatSampler, + disable_dropout_in_model, + ensure_master_addr_port, + entropy_from_logits, + get_config_model_id, + identity, + nanmax, + nanmin, + nanstd, + pad, + print_prompt_completions_sample, + selective_log_softmax, + shuffle_sequence_dict, + split_pixel_values_by_grid, + split_tensor_dict, + unsplit_pixel_values_by_grid, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + +if is_wandb_available(): + import wandb + +if is_trackio_available(): + import trackio + + +logger = logging.get_logger(__name__) + +# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of +# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. +RewardFunc = str | PreTrainedModel | Callable[[list, list], list[float]] + + +class RLOOTrainer(BaseTrainer): + """ + Trainer for the Reinforce Leave One Out (RLOO) method. This algorithm was initially proposed in the paper [Back to + Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in + LLMs](https://huggingface.co/papers/2402.14740). + + Example: + + ```python + from datasets import load_dataset + from trl import RLOOTrainer + + dataset = load_dataset("trl-lib/tldr", split="train") + + + def reward_func(completions, **kwargs): + # Dummy reward function that rewards completions with more unique letters. + return [float(len(set(completion))) for completion in completions] + + + trainer = RLOOTrainer( + model="Qwen/Qwen2-0.5B-Instruct", + reward_funcs=reward_func, + train_dataset=dataset, + ) + + trainer.train() + ``` + + Args: + model (`str | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + reward_funcs (`RewardFunc | list[RewardFunc]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can also return `None` when the reward is not applicable to those samples. This is useful + for multi-task training where different reward functions apply to different types of samples. When a + reward function returns `None` for a sample, that reward function is excluded from the reward + calculation for that sample. For more details, see [Using a custom reward + function](#using-a-custom-reward-function). + + The trainer's state is also passed to the reward function. The trainer's state is an instance of + [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the + reward function's signature. + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`RLOOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + """ + + _tag_names = ["trl", "rloo"] + _name = "RLOO" + _paper = { + "title": "Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", + "id": "2402.14740", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @inproceedings{ahmadian2024back, + title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, + author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, + year = 2024, + booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, + pages = {12248--12267}, + publisher = {Association for Computational Linguistics}, + editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, + }"""), + } + + def __init__( + self, + model: str | PreTrainedModel = None, + reward_funcs: RewardFunc | list[RewardFunc] = None, + args: RLOOConfig | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + peft_config: "PeftConfig | None" = None, + ): + if not os.environ.get("TRL_EXPERIMENTAL_SILENCE"): + warnings.warn( + "This trainer will soon be moved to trl.experimental and is a candidate for removal. If you rely on " + "it and want it to remain, please share your comments here: " + "https://github.com/huggingface/trl/issues/4223. Silence this warning by setting environment variable " + "TRL_EXPERIMENTAL_SILENCE=1." + ) + + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = RLOOConfig(f"{model_name}-RLOO") + + # Models + # Trained model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + dtype = model_init_kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass # dtype is already a torch.dtype or "auto" or None + elif isinstance(dtype, str): # it's a str, but not "auto" + dtype = getattr(torch, dtype) + model_init_kwargs["dtype"] = dtype + else: + raise ValueError( + "Invalid `dtype` passed to `RLOOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + model = architecture.from_pretrained(model_id, **model_init_kwargs) + else: + model_id = get_config_model_id(model.config) + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `RLOOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + + # Some models (SmolVLM/Idefics3) don't support `logits_to_keep` argument and error out if we pass it + # Inspect the forward method before we wrap the model with PEFT + self.model_kwarg_keys = ( + inspect.signature(model.forward).parameters.keys() + if not hasattr(model, "get_base_model") + else inspect.signature(model.get_base_model().forward).parameters.keys() + ) + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(get_config_model_id(model.config), truncation_side="left") + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + self.pad_token = tokenizer.pad_token + self.pad_token_id = tokenizer.pad_token_id + self.eos_token_id = tokenizer.eos_token_id + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " + f"reward functions ({len(reward_funcs)})." + ) + + for i, (reward_processing_class, reward_func) in enumerate( + zip(reward_processing_classes, reward_funcs, strict=True) + ): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained(get_config_model_id(reward_func.config)) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + + self.reward_processing_classes = reward_processing_classes + + # Training arguments + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length + self.num_generations = args.num_generations + self.chat_template_kwargs = args.chat_template_kwargs or {} + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_paged = args.use_transformers_paged + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.normalize_advantages = args.normalize_advantages + self.mask_truncated_completions = args.mask_truncated_completions + self.reward_clip_range = args.reward_clip_range + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in RLOOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations (forward + backward passes), including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in RLOO, the sampled data does not include the + # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning: + # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To + # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True. + # This acts as a flag to indicate that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in RLOO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + self.ref_model = architecture.from_pretrained(model_id, **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # Keep logs sized to the generation batch to record only outputs from the latest model update. + self._logs = { + "images": deque(maxlen=args.generation_batch_size), + "prompt": deque(maxlen=args.generation_batch_size), + "completion": deque(maxlen=args.generation_batch_size), + "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + "advantages": deque(maxlen=args.generation_batch_size), + } + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install trl[vllm]` to use it." + ) + + if self.vllm_mode == "server": + if self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator(device=torch.cuda.current_device()) + + elif self.vllm_mode == "colocate": + # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have + # the same number of ranks + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks. + # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size)) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + # vLLM requires the environment variables to be set for distributed training. + os.environ["RANK"] = str(self.accelerator.process_index) + os.environ["LOCAL_RANK"] = str(self.accelerator.local_process_index) + os.environ["WORLD_SIZE"] = str(self.accelerator.num_processes) + # Ensure distributed rendezvous variables are set without colliding across concurrent runs + ensure_master_addr_port() + + if self.max_prompt_length is not None and self.max_completion_length is not None: + max_model_len = self.max_prompt_length + self.max_completion_length + else: + max_model_len = None + self.llm = LLM( + model=model.name_or_path, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + max_num_seqs=self.args.per_device_train_batch_size + * self.vllm_tensor_parallel_size + * self.args.steps_per_generation, + max_model_len=max_model_len, + distributed_executor_backend="external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + seed=self.accelerator.process_index // self.vllm_tensor_parallel_size, + # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768) - thinking there's not enough memory + max_num_batched_tokens=4096, + model_impl=self.args.vllm_model_impl, + enable_sleep_mode=self.args.vllm_enable_sleep_mode, + ) + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=2) + else: + raise ValueError(f"vllm_mode must be either 'server' or 'colocate', got '{self.vllm_mode}'.") + + # vLLM specific sampling arguments + self.guided_decoding_regex = args.vllm_guided_decoding_regex + + self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": tokenizer.pad_token_id, + "bos_token_id": tokenizer.bos_token_id, + "eos_token_id": tokenizer.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In RLOOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt", "image", "images"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. As a result, some parts of the method aren't relevant to RLOO, but we keep them to stay one line + # apart from the super method, ensuring easier maintenance in the future. + def get_train_dataloader(self): + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size * self.args.steps_per_generation, # < this is the change + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = partial( + seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index + ) + + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. The first row shows the first sampled batch, the + # second row shows the second sampled batch, and so on. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations, + seed=self.args.seed, + ) + + @profiling_decorator + def _get_per_token_logps_and_entropies( + self, + model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=None, + compute_entropy=False, + pixel_values=None, + image_grid_thw=None, + num_images=None, + pixel_attention_mask=None, + image_sizes=None, + token_type_ids=None, + ) -> dict[str, torch.Tensor | None]: + """Compute log-probs and (optionally) entropies for each token.""" + batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak + all_logps = [] + all_entropies = [] + for start in range(0, input_ids.size(0), batch_size): + input_ids_batch = input_ids[start : start + batch_size] + attention_mask_batch = attention_mask[start : start + batch_size] + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} + + if image_grid_thw is not None and pixel_values is not None: + rows_per_image = image_grid_thw.prod(dim=-1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) + row_start, row_end = cum_rows[start].item(), cum_rows[start + batch_size].item() + model_inputs["pixel_values"] = pixel_values[row_start:row_end] + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] + model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] + elif pixel_values is not None: + model_inputs["pixel_values"] = pixel_values[start : start + batch_size] + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask[start : start + batch_size] + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes[start : start + batch_size] + if token_type_ids is not None: + model_inputs["token_type_ids"] = token_type_ids[start : start + batch_size] + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + logits = model(**model_inputs).logits + # Exclude the last value: it corresponds to the next token pred + logits = logits[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + logits = logits[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + # Divide logits by sampling temperature. + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + logits = logits / self.temperature + + completion_ids = input_ids_batch[:, -logits_to_keep:] + logps = selective_log_softmax(logits, completion_ids) # compute logprobs + all_logps.append(logps) + + if compute_entropy: + with torch.no_grad(): + entropies = entropy_from_logits(logits) + all_entropies.append(entropies) + + logps = torch.cat(all_logps, dim=0) + entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None + return logps, entropies + + def _fix_param_name_to_vllm(self, name, extra_prefixes: list[str] | None = None): + extra_prefixes = extra_prefixes or [] + prefixes = ["_checkpoint_wrapped_module."] + extra_prefixes + for prefix in prefixes: + name = name.replace(prefix, "") + return name + + def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + # For FSDP1, we need to recurse into children and also use summon_full_params + if visited is None: + visited = set() + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp1_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + def _sync_fsdp2_params_to_vllm(self, module: nn.Module): + # For FSDP2, module.state_dict() already covers all parameters, so no need for recursion + for name, param in module.state_dict().items(): + if param.is_cpu: + param = param.to(torch.device("cuda")) + param = param.full_tensor() + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param)]) + + @profiling_decorator + def _move_model_to_vllm(self): + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + # TODO: does this work with FSDP? + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm( + self.model + ) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + for name, param in self.model.named_parameters(): + name = self._fix_param_name_to_vllm(name) + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.llm.reset_prefix_cache() + + @profiling_decorator + def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = split_pixel_values_by_grid(generation_batch) + generation_batch = shuffle_sequence_dict(generation_batch) + generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) + self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + self._step += 1 + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + # This allows for dynamic reward shaping based on training progress. + reward_kwargs["trainer_state"] = self.state + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) + ): + with profiling_context(self, reward_func_name): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] + texts = [ + apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] + for x in messages + ] + else: + texts = [p + c for p, c in zip(prompts, completions, strict=True)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = { + key: value[nan_row_idx] for key, value in reward_kwargs.items() if key != "trainer_state" + } + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + logger.warning( + f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _generate_single_turn(self, prompts: list): + device = self.accelerator.device + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + if self.vllm_mode == "colocate" and self.args.vllm_enable_sleep_mode: + # wake up colocated vLLM instances if needed + torch.cuda.empty_cache() # required to avoid OOM in some cases + self.llm.wake_up(tags=["weights"]) + + # First, update the vLLM weights if needed + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + prompts = [prepare_multimodal_messages_vllm(prompt) for prompt in prompts] + + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process + if self.vllm_mode == "server": + all_prompts = gather_object(prompts) + + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts[:: self.num_generations] + + sampling_params = { + "n": self.num_generations, + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "truncate_prompt_tokens": self.max_prompt_length, + "guided_decoding_regex": self.guided_decoding_regex, + "generation_kwargs": self.args.generation_kwargs, + } + with profiling_context(self, "vLLM.generate"): + if is_conversational({"prompt": ordered_set_of_prompts[0]}): + output = self.vllm_client.chat( + messages=ordered_set_of_prompts, + **sampling_params, + chat_template_kwargs=self.chat_template_kwargs, + ) + else: + output = self.vllm_client.generate(prompts=ordered_set_of_prompts, **sampling_params) + payload = (output["prompt_ids"], output["completion_ids"], output["logprobs"]) + else: + payload = None + + # Broadcast the completions from the main process to all processes, ensuring each process receives its corresponding slice. + obj_list = [payload] + broadcast_object_list(obj_list, from_process=0) + all_prompt_ids, all_completion_ids, _ = obj_list[0] + + # At this point, we only get 1 copy of each prompt, so we need to repeat them num_generations times + all_prompt_ids = [ids for ids in all_prompt_ids for _ in range(self.num_generations)] + + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + prompt_ids = all_prompt_ids[process_slice] + completion_ids = all_completion_ids[process_slice] + + # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts + elif self.vllm_mode == "colocate": + if self.guided_decoding_regex: + guided_decoding = GuidedDecodingParams(regex=self.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": 1, # vLLM on each GPU generates only 1 in colocate mode + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "truncate_prompt_tokens": self.max_prompt_length, + "guided_decoding": guided_decoding, + } + if self.args.generation_kwargs is not None: + generation_kwargs.update(self.args.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + if self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) + all_prompts = [p for sublist in gathered_prompts for p in sublist] + else: + all_prompts = prompts + + if self.args.vllm_enable_sleep_mode: + self.llm.wake_up(tags=["kv_cache"]) + + with profiling_context(self, "vLLM.generate"): + if is_conversational({"prompt": prompts[0]}): + all_outputs = self.llm.chat(all_prompts, sampling_params=sampling_params, use_tqdm=False) + else: + all_outputs = self.llm.generate(all_prompts, sampling_params=sampling_params, use_tqdm=False) + + all_prompt_ids = [output.prompt_token_ids for output in all_outputs] + all_completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + + if self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + prompt_ids = all_prompt_ids[tp_slice] + completion_ids = all_completion_ids[tp_slice] + else: + prompt_ids = all_prompt_ids + completion_ids = all_completion_ids + + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=2) + + elif self.use_transformers_paged: + processor_kwargs = { + "max_length": self.max_prompt_length, + "truncation": True, + "add_special_tokens": False, + } + if is_conversational({"prompt": prompts[0]}): + processor_outputs = self.processing_class.apply_chat_template( + conversation=prompts, + **processor_kwargs, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **self.chat_template_kwargs, + ) + else: + processor_outputs = self.processing_class(text=prompts, **processor_kwargs) + + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + with torch.inference_mode(): + # Continuous batching API expects 'inputs' arg only + all_outputs = unwrapped_model.generate_batch( + processor_outputs["input_ids"], generation_config=self.generation_config, progress_bar=False + ) + unwrapped_model.train() # restore training mode, as generate_batch forces eval mode + completion_ids = [output.generated_tokens for output in all_outputs.values()] + prompt_ids = processor_outputs["input_ids"] + + else: + # Regular generation path + processor_kwargs = { + "return_tensors": "pt", + "padding": True, + "padding_side": "left", + "max_length": self.max_prompt_length, + "truncation": True, + "add_special_tokens": False, + } + if is_conversational({"prompt": prompts[0]}): + generate_inputs = self.processing_class.apply_chat_template( + conversation=prompts, + **processor_kwargs, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **self.chat_template_kwargs, + ) + else: + generate_inputs = self.processing_class(text=prompts, **processor_kwargs) + generate_inputs = super()._prepare_inputs(generate_inputs) + + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + prompt_completion_ids = unwrapped_model.generate( + **generate_inputs, generation_config=self.generation_config, disable_compile=True + ) + # Compute prompt length and extract completion ids + prompt_ids, prompt_mask = generate_inputs["input_ids"], generate_inputs["attention_mask"] + prompt_length = prompt_ids.size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + prompt_ids = [p[m].tolist() for p, m in zip(prompt_ids, prompt_mask.bool(), strict=True)] + completion_ids = [c[m].tolist() for c, m in zip(completion_ids, completion_mask.bool(), strict=True)] + + return prompt_ids, completion_ids + + def _generate(self, prompts: list): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompt_ids, completion_ids = self._generate_single_turn(prompts) + + # Get completion length per sequence, used for logging + prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) + completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) + agg_prompt_lengths = self.accelerator.gather(prompt_lengths) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + total_prompt_tokens = agg_prompt_lengths.sum() + total_completion_tokens = agg_completion_lengths.sum() # = num_items_in_batch, required for the DAPO loss + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + agg_completion_lengths = self.accelerator.gather(completion_lengths) + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) + agg_is_truncated = self.accelerator.gather(is_truncated) + self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) + term_completion_lengths = agg_completion_lengths[~agg_is_truncated] + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + return prompt_ids, completion_ids + + def _generate_and_score_completions( + self, inputs: list[dict[str, torch.Tensor | Any]] + ) -> dict[str, torch.Tensor | Any]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] + if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ] + + prompt_ids_list, completion_ids_list = self._generate(prompts) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad(prompt_ids, padding_value=self.pad_token_id, padding_side="left") + prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left") + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right") + completion_mask = pad(completion_mask, padding_value=0, padding_side="right") + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + num_images = [len(img_list) for img_list in images] if images is not None else None + + # Get forward_kwargs for models with multimodal inputs + if images is not None: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] + for prompt in prompts + ] + prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + + with torch.no_grad(): + # Compute the per-token log probabilities for the current model + old_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + old_logps = (old_per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + prompt_completion_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + ref_per_token_logps = None + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text, strict=True): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + # Apply reward clipping if specified + if self.reward_clip_range: + rewards = rewards.clamp(min=self.reward_clip_range[0], max=self.reward_clip_range[1]) + + # Include the KL penalty in the reward + if self.beta != 0.0: + per_token_kl = old_per_token_logps - ref_per_token_logps + # Apply sequence-level KL penalty to rewards (sum KL across tokens first, then apply to each sequence) + kl = (per_token_kl * completion_mask).sum(-1) + kl = gather(kl) # rewards are gathered, so kl must be too + rewards = rewards - self.beta * kl + + grouped_rewards = rewards.view(-1, self.num_generations) + mean_grouped_rewards = grouped_rewards.mean(dim=1) + std_rewards = grouped_rewards.std(dim=1) + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) + + # RLOO advantages computation + grouped_sum = grouped_rewards.sum(dim=1, keepdim=True) # (num_prompts, 1) + baselines = (grouped_sum - grouped_rewards) / (self.num_generations - 1) # (num_prompts, num_generations) + baselines = baselines.view(-1) # Flatten back to match rewards shape + advantages = rewards - baselines + + # Normalize advantages + if self.normalize_advantages: + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-4) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Calculate and log the mean KL divergence between current and reference model + if self.beta != 0.0: + mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._logs["prompt"].extend(gather_object(prompts_text)) + self._logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + if images is not None: + self._logs["images"].extend(gather_object(images)) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "old_logps": old_logps, + "advantages": advantages, + } + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if images is not None: + output["num_images"] = num_images + return output + + @profiling_decorator + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The RLOOTrainer does not support returning outputs") + return self._compute_loss(model, inputs) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + ) + + logps = (per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS + old_logps = inputs["old_logps"] + log_ratio = logps - old_logps + + # Compute the loss + advantages = inputs["advantages"] + coef_1 = torch.exp(log_ratio) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + per_sequence_loss1 = coef_1 * advantages + per_sequence_loss2 = coef_2 * advantages + per_sequence_loss = -torch.min(per_sequence_loss1, per_sequence_loss2) + loss = per_sequence_loss.mean() + + # Log the metrics + mode = "train" if self.model.training else "eval" + + # Entropy + mean_entropy = (entropies * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) + is_region_clipped = is_low_clipped | is_high_clipped + gathered_low_clip = self.accelerator.gather(is_low_clipped.float().mean()) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(is_high_clipped.float().mean()) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float().mean()) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._logs["prompt"], + self._logs["completion"], + self._logs["rewards"], + self._logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + ) + + table = { + "step": [str(self.state.global_step)] * len(self._logs["prompt"]), + "prompt": self._logs["prompt"], + "completion": self._logs["completion"], + **self._logs["rewards"], + "advantage": self._logs["advantages"], + } + + df_base = pd.DataFrame(table) + images_raw = self._logs["images"] or [] + + for logging_backend in self.args.report_to: + if logging_backend == "wandb": + if images_raw: + images = [] + for image_list in self._logs["images"]: + images.append([wandb.Image(image) for image in image_list]) + df = pd.concat([df_base, pd.Series(images, name="image")], axis=1, copy=False) + else: + df = df_base + + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if logging_backend == "trackio": + if images_raw: + # TODO: Implement once supported upstream https://github.com/gradio-app/trackio/issues/334 + logger.info("Skipping image logging for Trackio") + df = df_base + # images = [] + # for image_list in self._logs["images"]: + # images.append([trackio.Image(image) for image in image_list]) + # df = pd.concat([df_base, pd.Series(images, name="image")], axis=1, copy=False) + else: + df = df_base + + trackio.log({"completions": trackio.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/sft_config.py b/venv/lib/python3.10/site-packages/trl/trainer/sft_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c43a7a51bf7956e359b702e23984e8a5c8564c29 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/sft_config.py @@ -0,0 +1,262 @@ +# Copyright 2020-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 dataclasses import dataclass, field +from typing import Any + +from transformers import TrainingArguments + + +@dataclass +class SFTConfig(TrainingArguments): + r""" + Configuration class for the [`SFTTrainer`]. + + This class includes only the parameters that are specific to SFT training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`SFTTrainer`] is provided as a string. If you're training a MoE architecture and want to + include the load balancing/auxilliary loss as a part of the final loss, remember to set + `output_router_logits=True` in this dictionary. + chat_template_path (`str`, *optional*): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + + > Parameters that control the data preprocessing + + dataset_text_field (`str`, *optional*, defaults to `"text"`): + Name of the column that contains text data in the dataset. + dataset_kwargs (`dict[str, Any]`, *optional*): + Dictionary of optional keyword arguments for the dataset preparation. The only supported key is + `skip_prepare_dataset`. When the model is a VLM, `skip_prepare_dataset` is automatically treated as `True` + regardless of the provided value, since preprocessing is done on the fly. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + eos_token (`str`, *optional*): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + pad_token (`str`, *optional*): + Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that is also `None`, + it falls back to `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the right. + If `None`, no truncation is applied. When packing is enabled, this value sets the sequence length. + packing (`bool`, *optional*, defaults to `False`): + Whether to group multiple sequences into fixed-length blocks to improve computational efficiency and reduce + padding. Uses `max_length` to define sequence length. + packing_strategy (`str`, *optional*, defaults to `"bfd"`): + Strategy for packing sequences. Can be either `"bfd"` (best-fit decreasing, default), or `"wrapped"`. + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch structure. When + packing is enabled with strategy `"bfd"`, padding-free is enabled, regardless of the value of this + parameter. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + eval_packing (`bool`, *optional*): + Whether to pack the eval dataset. If `None`, uses the same value as `packing`. + + > Parameters that control the training + + completion_only_loss (`bool`, *optional*): + Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is computed + only on the completion, which is supported only for [prompt-completion](#prompt-completion) datasets. If + `False`, loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: + loss is computed on the completion for [prompt-completion](#prompt-completion) datasets, and on the full + sequence for [language modeling](#language-modeling) datasets. + assistant_only_loss (`bool`, *optional*, defaults to `False`): + Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is computed only + on the assistant responses, which is supported only for [conversational](#conversational) datasets. If + `False`, loss is computed on the entire sequence. + loss_type (`str`, *optional*, defaults to `"nll"`): + Type of loss to use. Possible values are `"nll"` (negative log-likelihood, default) and `"dft"` (Dynamic + Fine-Tuning, as described in [this paper](https://huggingface.co/papers/2508.05629)). + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=2e-5, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model + model_init_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `SFTTrainer` is provided as a string. If you're training a MoE architecture and want to include the " + "load balancing/auxilliary loss as a part of the final loss, remember to set `output_router_logits=True` " + "in this dictionary." + }, + ) + chat_template_path: str | None = field( + default=None, + metadata={ + "help": "If specified, sets the model's chat template. This can either be the path to a tokenizer (local " + "directory or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, " + "you must ensure that any special tokens referenced in the template are added to the tokenizer and " + "that the model's embedding layer is resized accordingly." + }, + ) + + # Parameters that control the data preprocessing + dataset_text_field: str = field( + default="text", + metadata={"help": "Name of the column that contains text data in the dataset."}, + ) + dataset_kwargs: dict[str, Any] | None = field( + default=None, + metadata={ + "help": "Dictionary of optional keyword arguments for the dataset preparation. The only supported key is " + "`skip_prepare_dataset`. If the model is a VLM, `skip_prepare_dataset` value is ignored. When the model " + "is a VLM, `skip_prepare_dataset` is automatically treated as `True` regardless of the provided value, " + "since preprocessing is done on the fly." + }, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + eos_token: str | None = field( + default=None, + metadata={ + "help": "Token used to indicate the end of a turn or sequence. If `None`, it defaults to `processing_class.eos_token`." + }, + ) + pad_token: str | None = field( + default=None, + metadata={ + "help": "Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that " + "is also `None`, it falls back to `processing_class.eos_token`." + }, + ) + max_length: int | None = field( + default=1024, + metadata={ + "help": "Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from" + "the right. If `None`, no truncation is applied. When packing is enabled, this value sets the " + "sequence length." + }, + ) + packing: bool = field( + default=False, + metadata={ + "help": "Whether to group multiple sequences into fixed-length blocks to improve computational efficiency " + "and reduce padding. Uses `max_length` to define sequence length." + }, + ) + packing_strategy: str = field( + default="bfd", + metadata={ + "help": "Strategy for packing sequences. Can be either `'bfd'` (best-fit decreasing, default), or " + "`'wrapped'`." + }, + ) + padding_free: bool = field( + default=False, + metadata={ + "help": "Whether to perform forward passes without padding by flattening all sequences in the batch into " + "a single continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this " + "is only supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch " + "structure. When packing is enabled with strategy `'bfd'`, padding-free is enabled, regardless of the " + "value of this parameter." + }, + ) + pad_to_multiple_of: int | None = field( + default=None, + metadata={"help": "If set, the sequences will be padded to a multiple of this value."}, + ) + eval_packing: bool | None = field( + default=None, + metadata={"help": "Whether to pack the eval dataset. If `None`, uses the same value as `packing`."}, + ) + + # Parameters that control the training + completion_only_loss: bool | None = field( + default=None, + metadata={ + "help": ( + "Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is " + "computed only on the completion, which is supported only for prompt-completion datasets. If `False`, " + "loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: " + "loss is computed on the completion for prompt-completion datasets, and on the full sequence for " + "language modeling datasets." + ) + }, + ) + assistant_only_loss: bool = field( + default=False, + metadata={ + "help": ( + "Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is " + "computed only on the assistant responses, which is supported only for conversational datasets. If `False`, " + "loss is computed on the entire sequence." + ) + }, + ) + loss_type: str = field( + default="nll", + metadata={ + "help": ( + 'Type of loss to use. Possible values are `"nll"` (negative log-likelihood, default) and `"dft"` ' + "(Dynamic Fine-Tuning, as described in https://huggingface.co/papers/2508.05629)." + ) + }, + ) + activation_offloading: bool = field( + default=False, + metadata={"help": "Whether to offload the activations to the CPU."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + super().__post_init__() diff --git a/venv/lib/python3.10/site-packages/trl/trainer/sft_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/sft_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f69a5f3eddb1192a52341a07fa214b4436b48e87 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/sft_trainer.py @@ -0,0 +1,1212 @@ +# Copyright 2020-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 contextlib +import os +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from accelerate import PartialState, logging +from datasets import Dataset, IterableDataset +from transformers import ( + AutoProcessor, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainingArguments, +) +from transformers.data.data_collator import DataCollatorMixin +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from ..data_utils import ( + apply_chat_template, + is_conversational, + is_conversational_from_value, + maybe_convert_to_chatml, + pack_dataset, + prepare_multimodal_messages, + truncate_dataset, +) +from ..models import clone_chat_template, get_act_offloading_ctx_manager, prepare_peft_model +from .base_trainer import BaseTrainer +from .sft_config import SFTConfig +from .utils import ( + create_model_from_path, + entropy_from_logits, + flush_left, + get_config_model_id, + pad, + remove_none_values, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel, PeftType + + +logger = logging.get_logger(__name__) + + +FLASH_ATTENTION_VARIANTS = { + "flash_attention_2", + "flash_attention_3", + "kernels-community/flash-attn", + "kernels-community/vllm-flash-attn3", + "kernels-community/flash-attn3", +} + + +def get_dataset_column_names(dataset: Dataset | IterableDataset) -> list[str]: + return list(next(iter(dataset)).keys()) if dataset.column_names is None else dataset.column_names + + +@dataclass +class DataCollatorForLanguageModeling(DataCollatorMixin): + """ + Data collator used for language modeling data. Inputs are dynamically padded to the maximum length of a batch. + + This collator expects each example in the input list to be a dictionary containing at least the `"input_ids"` key. + If the input contains a `"completion_mask"`, it is used to set the labels to `-100` for tokens that are not in the + completion. If `"assistant_masks"` are present, they are used to set the labels to `-100` for tokens that are not + in the assistant part of the sequence. The collator returns a dictionary containing the following keys: + - `"input_ids"`: Tensor of input IDs, padded to the maximum length of the batch. + - `"labels"`: Tensor of labels, padded to the maximum length of the batch. If `completion_only_loss` is set to + `True`, tokens that are not in the completion are set to -100. If `assistant_masks` are present, tokens that are + not in the assistant part of the sequence are set to -100. If `padding_free` is set to `False`, the following key + is also returned: + - `"attention_mask"`: Tensor of attention masks, padded to the maximum length of the batch. + If `padding_free` is set to `True`, the following key is also returned: + - `"position_ids"`: Tensor of position IDs, padded to the maximum length of the batch. + + Args: + pad_token_id (`int`): + Token ID to use for padding. + completion_only_loss (`bool`, *optional*, defaults to `True`): + When the input contains a completion mask (`completion_mask`), the labels are set to -100 for the tokens + that are no in the completion. + padding_free (`bool`, *optional*, defaults to `False`): + If set to `True`, the sequences will be flattened into a single sequence, and the position IDs will be + generated accordingly and returned instead of the attention mask. + pad_to_multiple_of (`int`, *optional*): + If set, the sequences will be padded to a multiple of this value. + return_tensors (`str`, *optional*, defaults to `"pt"`): + Type of Tensor to return. Only `"pt"` is currently supported. + + Examples: + ```python + >>> from trl.trainer.sft_trainer import DataCollatorForLanguageModeling + + >>> collator = DataCollatorForLanguageModeling(pad_token_id=0) + >>> examples = [{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}] + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3], + [ 4, 5, 0]]), + 'attention_mask': tensor([[ 1, 1, 1], + [ 1, 1, 0]]), + 'labels': tensor([[ 1, 2, 3], + [ 4, 5, -100]])} + + >>> # With completion mask + >>> examples = [ + ... {"input_ids": [1, 2, 3], "completion_mask": [0, 1, 1]}, + ... {"input_ids": [4, 5], "completion_mask": [0, 1]}, + ... ] + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3], + [ 4, 5, 0]]), + 'attention_mask': tensor([[ 1, 1, 1], + [ 1, 1, 0]]), + 'labels': tensor([[-100, 2, 3], + [-100, 5, -100]])} + + >>> # With padding_free + >>> collator = DataCollatorForLanguageModeling(pad_token_id=0, padding_free=True) + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3, 4, 5]]), + 'position_ids': tensor([[0, 1, 2, 0, 1]]), + 'labels': tensor([[1, 2, 3, 4, 5]])} + ``` + """ + + pad_token_id: int + completion_only_loss: bool = True + padding_free: bool = False + pad_to_multiple_of: int | None = None + return_tensors: str = "pt" + + def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + # Convert to tensor + input_ids = [torch.tensor(example["input_ids"]) for example in examples] + if "labels" in examples[0]: + labels = [torch.tensor(example["labels"]) for example in examples] + else: + labels = [torch.tensor(example["input_ids"]) for example in examples] + + # For padding-free, we should NOT create attention_mask as it causes FlashAttention to ignore position_ids and + # compute wrong cu_seq_lens from the all-1s mask + if self.padding_free: + if "seq_lengths" in examples[0]: + position_ids = self.get_position_ids_from_packed_seq_lengths( + [example["seq_lengths"] for example in examples] + ) + else: + position_ids = [torch.arange(len(ids)) for ids in input_ids] + else: + attention_mask = [torch.ones_like(ids) for ids in input_ids] + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = [torch.tensor(example["completion_mask"]) for example in examples] + if "assistant_masks" in examples[0]: + assistant_masks = [torch.tensor(example["assistant_masks"]) for example in examples] + + # If padding_free, flatten everything into a single sequence + output = {} + if self.padding_free: + input_ids = [torch.cat(input_ids, dim=0)] + labels = [torch.cat(labels, dim=0)] + position_ids = [torch.cat(position_ids, dim=0)] + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = [torch.cat(completion_mask, dim=0)] + if "assistant_masks" in examples[0]: + assistant_masks = [torch.cat(assistant_masks, dim=0)] + + # Pad + output["input_ids"] = pad( + input_ids, + padding_value=self.pad_token_id, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ) + output["labels"] = pad( + labels, padding_value=-100, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + if self.padding_free: + output["position_ids"] = pad( + position_ids, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"][output["position_ids"] == 0] = -100 + else: + output["attention_mask"] = pad( + attention_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = pad( + completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"][completion_mask == 0] = -100 # mask everything that is not in the completion + if "assistant_masks" in examples[0]: + assistant_masks = pad( + assistant_masks, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"][assistant_masks == 0] = -100 + return output + + @staticmethod + def get_position_ids_from_packed_seq_lengths(batch_seq_lengths: list[list[int]]) -> list[torch.Tensor]: + """ + Get position IDs for packed sequences. + + Args: + batch_seq_lengths (`list[list[int]]`): + A list of lists containing the lengths of each individual document in the packed batch. + + Return: + `list[torch.Tensor]`: + A list of tensors containing the position IDs for each packed sequence. + """ + # Get lengths per row + example_lengths = [sum(seq_lengths) for seq_lengths in batch_seq_lengths] + # Flat list of lengths + batch_seq_lengths = torch.tensor( + [seq_length for seq_lengths in batch_seq_lengths for seq_length in seq_lengths] + ) + position_ids = torch.ones(sum(example_lengths), dtype=batch_seq_lengths.dtype) + position_ids[0] = 0 + # Reset position ids to 0 at the start of each sequence + position_ids[batch_seq_lengths[:-1].cumsum(0)] = -(batch_seq_lengths[:-1] - 1) + position_ids = position_ids.cumsum(0) + # Split back into one tensor per example + return list(position_ids.split(example_lengths)) + + +@dataclass +class DataCollatorForVisionLanguageModeling(DataCollatorMixin): + """ + Data collator for vision-language modeling tasks. + + Unlike text-only datasets—where the collator typically receives pre-tokenized inputs ready for batching, + vision-language data processing involves converting images into pixel values. This conversion is disk-intensive, + making upfront preprocessing of the entire dataset impractical. Therefore, this collator performs tokenization and + image processing on-the-fly to efficiently prepare batches. + + Each input example should be a dictionary containing at least: + - An `"images"` key holding the image data. + - [language modeling](#language-modeling) type: either a `"messages"` key for conversational inputs or a `"text"` + key for standard text inputs. + - [prompt-completion](#prompt-completion) type: keys `"prompt"` and `"completion"` for the prompt and completion. + + The collator outputs a dictionary including: + - `"input_ids"`: Tensor of token IDs. + - `"attention_mask"`: Tensor indicating attention mask. + - `"pixel_values"`: Tensor representing image pixel values. + - `"labels"`: Tensor for training labels. + + Additional keys may be present depending on the processor, such as `"image_grid_thw"`. + + Args: + processor ([`~transformers.ProcessorMixin`]): + The processor used to tokenize text and process images. It must be a subclass of + [`~transformers.ProcessorMixin`] and include a `tokenizer` with a defined `pad_token_id`. + max_length (`int` or `None`, optional, defaults to `None`): + Maximum sequence length for input tokens. If `None`, no truncation is applied. + completion_only_loss (`bool`, *optional*, defaults to `False`): + Whether to compute loss only on the completion part of the sequence. When `True`, the labels for the prompt + part are set to -100. It requires the dataset type to be prompt-completion. + pad_to_multiple_of (`int` or `None`, optional, defaults to `None`): + If set, the sequences will be padded to a multiple of this value. + dataset_text_field (`str`, optional, defaults to `"text"`): + Name of the column that contains text data in the dataset. This parameter is only relevant for [standard + datasets format](dataset_formats#standard). + return_tensors (`str`, optional, defaults to `"pt"`): + The tensor type to return. Currently, only `"pt"` (PyTorch tensors) is supported. + + Example: + ```python + >>> from trl.trainer.sft_trainer import DataCollatorForVisionLanguageModeling + >>> from transformers import AutoProcessor + + >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct") + >>> collator = DataCollatorForVisionLanguageModeling(processor) + >>> examples = [ + ... {"images": [Image.open("image_0.png")], "messages": [{"role": "user", "content": "What is this?"}]}, + ... {"images": [Image.open("image_1.png")], "messages": [{"role": "user", "content": "Describe this image."}]}, + ... ] + >>> collator(examples) + {'input_ids': tensor([[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, + 151644, 872, 198, 151652, 151655, 151655, 151655, 151655, 151653, 3838, 374, + 419, 30, 151645, 198], + [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, + 151644, 872, 198, 151652, 151655, 151655, 151655, 151655, 151653, 74785, 419, + 2168, 13, 151645, 198]]), + 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), + 'pixel_values': tensor([[-0.9893, 0.1785, 1.5362, ..., -0.0582, 0.8661, -0.2431], + [-0.2302, 0.9522, -1.1061, ..., 0.0555, 1.3354, -0.6412], + [ 1.2150, 0.9084, 0.7041, ..., 0.2404, -0.8403, -0.5133], + ..., + [ 0.6895, 0.2807, 0.2515, ..., -0.2004, -1.2100, 0.0555], + [ 0.8209, -0.9748, 1.5654, ..., 1.6055, -0.4706, 0.5817], + [-1.0915, 0.4559, 0.9230, ..., 0.5106, 0.0982, -0.1720]]), + 'image_grid_thw': tensor([[1, 4, 4], + [1, 4, 4]]), + 'labels': tensor([[151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, + 151644, 872, 198, 151652, 151655, 151655, 151655, 151655, 151653, 3838, 374, + 419, 30, 151645, 198], + [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, + 151644, 872, 198, 151652, 151655, 151655, 151655, 151655, 151653, 74785, 419, + 2168, 13, 151645, 198]])} + ``` + """ + + processor: ProcessorMixin + max_length: int | None = None + completion_only_loss: bool = False # default not used in practice; SFTTrainer always passes the relevant value + pad_to_multiple_of: int | None = None + dataset_text_field: str = "text" + return_tensors: str = "pt" + + def torch_call(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + if "messages" in examples[0] or self.dataset_text_field in examples[0]: + if self.completion_only_loss: + raise ValueError( + "The `completion_only_loss` argument is not supported for language modeling datasets." + ) + return self._collate_language_modeling(examples) + elif "prompt" in examples[0] and "completion" in examples[0]: + return self._collate_prompt_completion(examples) + else: + raise KeyError(f"Unexpected input keys in examples: {list(examples[0].keys())}.") + + def _collate_language_modeling(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + images = [example["images"] for example in examples] + # Transformers requires at least one image in the batch, otherwise it throws an error + if all(img_list == [] for img_list in images): + images = None + + if "messages" in examples[0]: # conversational case + messages = [prepare_multimodal_messages(example["messages"], example["images"]) for example in examples] + texts = self.processor.apply_chat_template(messages) + elif self.dataset_text_field in examples[0]: # standard case + texts = [example[self.dataset_text_field] for example in examples] + else: + raise KeyError( + "The input examples must contain either 'messages' for conversational data or 'text' for standard " + "data." + ) + + output = self.processor( + images=images, + text=texts, + padding=True, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + truncation=self.max_length is not None, + max_length=self.max_length, + return_tensors=self.return_tensors, + add_special_tokens=False, # to avoid adding the BOS, twice see https://huggingface.co/blog/qgallouedec/gotchas-in-tokenizer-behavior#7-chat-template-and-tokenization-dont-compose-due-to-special-tokens + ) + labels = output["input_ids"].clone() + labels[output["attention_mask"] == 0] = -100 + # We mask only padding tokens (-100) in the labels. Vision tokens are left unchanged because their handling in + # loss computation has to be done by the model, and masking them here would be infeasible in practice as vision + # token definitions vary across architectures. + output["labels"] = labels + return output + + def _collate_prompt_completion(self, examples: list[dict[str, Any]]) -> dict[str, Any]: + if self.pad_to_multiple_of is not None: + raise NotImplementedError( + "Padding to a multiple of a value is not yet implemented for vision-language modeling and " + "prompt-completion data yet." + ) + images = [example["images"] for example in examples] + # Transformers requires at least one image in the batch, otherwise it throws an error + if all(img_list == [] for img_list in images): + images = None + if is_conversational(examples[0]): # conversational case + for example in examples: + example["prompt"] = prepare_multimodal_messages(example["prompt"], images=example["images"]) + example["completion"] = prepare_multimodal_messages(example["completion"], images=[]) + examples = [apply_chat_template(example, self.processor) for example in examples] + + prompts = [example["prompt"] for example in examples] + completions = [example["completion"] for example in examples] + + processed_prompts = self.processor( + images=images, + text=prompts, + padding=True, + padding_side="left", + return_tensors=self.return_tensors, + add_special_tokens=False, # to avoid adding the BOS, twice see https://huggingface.co/blog/qgallouedec/gotchas-in-tokenizer-behavior#7-chat-template-and-tokenization-dont-compose-due-to-special-tokens + ) + processed_completions = self.processor( + text=completions, + padding=True, + padding_side="right", + return_tensors=self.return_tensors, + add_special_tokens=False, # to avoid adding the BOS, twice see https://huggingface.co/blog/qgallouedec/gotchas-in-tokenizer-behavior#7-chat-template-and-tokenization-dont-compose-due-to-special-tokens + ) + + # Concatenate prompts and completions + prompt_ids, completion_ids = processed_prompts["input_ids"], processed_completions["input_ids"] + prompt_mask, completion_mask = processed_prompts["attention_mask"], processed_completions["attention_mask"] + input_ids = torch.cat((prompt_ids, completion_ids), dim=1) + attention_mask = torch.cat((prompt_mask, completion_mask), dim=1) + completion_mask = torch.cat((torch.zeros_like(prompt_mask), completion_mask), dim=1) + if "token_type_ids" in processed_prompts: # special case for Gemma + prompt_token_type_ids = processed_prompts["token_type_ids"] + completion_token_type_ids = processed_completions["token_type_ids"] + token_type_ids = torch.cat((prompt_token_type_ids, completion_token_type_ids), dim=1) + + # Flush left to reduce padding + if "token_type_ids" in processed_prompts: + attention_mask, input_ids, completion_mask, token_type_ids = flush_left( + attention_mask, input_ids, completion_mask, token_type_ids + ) + else: + attention_mask, input_ids, completion_mask = flush_left(attention_mask, input_ids, completion_mask) + + # Truncate if necessary + if self.max_length is not None: + input_ids = input_ids[:, : self.max_length] + attention_mask = attention_mask[:, : self.max_length] + completion_mask = completion_mask[:, : self.max_length] + if "token_type_ids" in processed_prompts: + token_type_ids = token_type_ids[:, : self.max_length] + + # Create labels and mask padding tokens + labels = input_ids.clone() + labels[attention_mask == 0] = -100 + if self.completion_only_loss: + labels[completion_mask == 0] = -100 + + # Build the output dictionary + output = processed_prompts # we take processed_prompts because it contains the images + output["input_ids"] = input_ids + output["attention_mask"] = attention_mask + output["labels"] = labels + if "token_type_ids" in processed_prompts: + output["token_type_ids"] = token_type_ids + return output + + +def dft_loss(outputs, labels, num_items_in_batch=None): + """ + DFT loss function, as presented in [On the Generalization of SFT: A Reinforcement Learning Perspective with Reward + Rectification](https://huggingface.co/papers/2508.05629) + """ + labels = nn.functional.pad(labels, (0, 1), value=-100) + shift_labels = labels[..., 1:].contiguous() + loss_mask = shift_labels != -100 + shift_labels[~loss_mask] = 0 + logprobs = selective_log_softmax(outputs.logits, shift_labels) + per_token_loss = -logprobs.exp().detach() * logprobs + if num_items_in_batch is None: + num_items_in_batch = loss_mask.sum() + loss = (per_token_loss * loss_mask).sum() / num_items_in_batch + return loss + + +class SFTTrainer(BaseTrainer): + """ + Trainer for Supervised Fine-Tuning (SFT) method. + + This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + from datasets import load_dataset + from trl import SFTTrainer + + dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") + + trainer = SFTTrainer(model="Qwen/Qwen2-0.5B-Instruct", train_dataset=dataset) + trainer.train() + ``` + + Args: + model (`str | PreTrainedModel`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using `.from_pretrained` (where `` is derived from the model + config) with the keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. + If you're training a model with an MoE architecture and want to include the load balancing/auxilliary loss + as a part of the final loss, remember to set the `output_router_logits` config of the model to `True`. + args ([`SFTConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator ([`~transformers.DataCollator`], *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~trainer.sft_trainer.DataCollatorForLanguageModeling`] if the model is a language model + and [`~trainer.sft_trainer.DataCollatorForVisionLanguageModeling`] if the model is a vision-language model. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. SFT supports both [language modeling](#language-modeling) type and + [prompt-completion](#prompt-completion) type. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain an `input_ids` field. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoProcessor.from_pretrained`]. A padding token, `tokenizer.pad_token`, must be set. + If the processing class has not set a padding token, `tokenizer.eos_token` will be used as the default. + compute_loss_func (`Callable`, *optional*): + A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated + batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default [loss + function](https://github.com/huggingface/transformers/blob/052e652d6d53c2b26ffde87e039b723949a53493/src/transformers/trainer.py#L3618) + used by [`Trainer`]. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a + [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing + [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean + `compute_result` argument. This will be triggered after the last eval batch to signal that the function + needs to calculate and return the global summary statistics rather than accumulating the batch-level + statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your + model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + formatting_func (`Callable`, *optional*): + Formatting function applied to the dataset before tokenization. Applying the formatting function explicitly + converts the dataset into a [language modeling](#language-modeling) type. + """ + + _tag_names = ["trl", "sft"] + _name = "SFT" + + def __init__( + self, + model: str | PreTrainedModel, + args: SFTConfig | TrainingArguments | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, + compute_loss_func: Callable | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), + optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + peft_config: "PeftConfig | None" = None, + formatting_func: Callable[[dict], str] | None = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else get_config_model_id(model.config) + model_name = model_name.split("/")[-1] + args = SFTConfig(f"{model_name}-SFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, SFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + dict_args.pop("push_to_hub_token") + args = SFTConfig(**dict_args) + + # Model + if isinstance(model, str): + model = create_model_from_path(model, **args.model_init_kwargs or {}) + else: + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `SFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(get_config_model_id(model.config)) + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + self._is_vlm = True + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + self._is_vlm = False + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if args.eos_token is not None: + eos_token = args.eos_token + eos_token_id = tokenizer.convert_tokens_to_ids(eos_token) + if eos_token_id is None: + raise ValueError( + f"The specified `eos_token` ('{eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + tokenizer.eos_token_id = eos_token_id + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + added_tokens = [] + else: + model, processing_class, added_tokens = clone_chat_template( + model, processing_class, args.chat_template_path + ) + else: + added_tokens = [] + + # Catch some wrong configurations related to VLMs + if self._is_vlm and args.packing: + raise ValueError( + "Packing is not supported for vision-language models. Please set `packing=False` in the SFTConfig." + ) + if self._is_vlm and args.padding_free: + raise ValueError( + "Padding-free training is yet not supported for vision-language models. Please set " + "`padding_free=False` in the `SFTConfig`." + ) + if self._is_vlm and args.assistant_only_loss: + raise ValueError( + "Assistant-only loss is not yet supported for vision-language models. Please set " + "`assistant_only_loss=False` in the `SFTConfig`." + ) + + # PEFT configuration and model wrapping + if peft_config is not None: + if added_tokens: + # Ensure that the added tokens are trainable + if peft_config.trainable_token_indices is None: + peft_config.trainable_token_indices = {"embed_tokens": added_tokens} + elif "embed_tokens" not in peft_config.trainable_token_indices: + peft_config.trainable_token_indices["embed_tokens"] = added_tokens + else: + peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) + + # Ensure that the lm_head is trainable + if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: + logger.warning( + "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " + "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " + "tokens, leading to degraded generation quality. To fix this, add " + "`modules_to_save=['lm_head']` to your PEFT configuration." + ) + + if peft_config.modules_to_save is None: + peft_config.modules_to_save = ["lm_head"] + else: + peft_config.modules_to_save.append("lm_head") + + # In Prompt Tuning a small set of trainable virtual tokens (continuous prompt embeddings) is prepended to the + # input. We store the number of these tokens so we can account for them correctly when calculating accuracy. + self.num_virtual_tokens = 0 + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + if model.active_adapter in model.peft_config: + peft_model_config = model.peft_config[model.active_adapter] + self.num_virtual_tokens = getattr(peft_model_config, "num_virtual_tokens", 0) + + # Data collator + # BFD packing requires padding-free mode; otherwise, the collator outputs padded attention masks, causing + # FlashAttention to ignore position_ids and recompute them incorrectly from the padded attention mask. + self.padding_free = args.padding_free or (args.packing and args.packing_strategy == "bfd") + use_flash_attention = model.config._attn_implementation in FLASH_ATTENTION_VARIANTS + if self.padding_free: + if data_collator is not None: + raise ValueError("Passing a custom data collator is not supported when using padding-free.") + if args.packing and args.packing_strategy == "wrapped": + logger.warning( + "You are passing `padding_free=True` with the 'wrapped' packing strategy, which is not " + "recommended. Please refer to the documentation to understand why this is not recommended." + ) + if not use_flash_attention: + logger.warning( + "Padding-free training is enabled, but the attention implementation is not set to a supported " + "flash attention variant. Padding-free training flattens batches into a single sequence, and only " + "the following implementations are known to reliably support this: " + f"{', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. Using other implementations may lead to " + "unexpected behavior. To ensure compatibility, set `attn_implementation` in the model " + "configuration to one of these supported options or verify that your attention mechanism can " + "handle flattened sequences." + ) + + if args.per_device_train_batch_size == 1 and not args.packing: + logger.warning( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + + # Decide whether to use completion-only loss: if not specified, then it is set to True if the dataset format + # is prompt-completion, and False if the dataset format is language modeling. + dataset_sample = next(iter(train_dataset)) + if args.completion_only_loss is None: + self.completion_only_loss = "prompt" in dataset_sample and "completion" in dataset_sample + else: + self.completion_only_loss = args.completion_only_loss + + self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample + if self._is_vision_dataset and not self._is_vlm: + raise ValueError( + "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " + "model does not seem to be a vision-language model. Please check your model and dataset." + ) + + if data_collator is None and not self._is_vision_dataset: + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + pad_token = args.pad_token or tokenizer.pad_token or tokenizer.eos_token + pad_token_id = tokenizer.convert_tokens_to_ids(pad_token) + if pad_token_id is None: + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + data_collator = DataCollatorForLanguageModeling( + pad_token_id=pad_token_id, + completion_only_loss=self.completion_only_loss, + padding_free=self.padding_free, + pad_to_multiple_of=args.pad_to_multiple_of, + ) + elif data_collator is None and self._is_vision_dataset: + data_collator = DataCollatorForVisionLanguageModeling( + processor=processing_class, + max_length=args.max_length, + completion_only_loss=self.completion_only_loss, + pad_to_multiple_of=args.pad_to_multiple_of, + dataset_text_field=args.dataset_text_field, + ) + + if args.packing and args.packing_strategy == "bfd" and not use_flash_attention: + logger.warning( + "You are using packing, but the attention implementation is not set to a supported flash attention " + "variant. Packing gathers multiple samples into a single sequence, and only the following " + f"implementations are known to reliably support this: {', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. " + "Using other implementations may lead to cross-contamination between samples. To avoid this, either " + "disable packing by setting `packing=False`, or set `attn_implementation` in the model configuration " + "to one of these supported options." + ) + if args.assistant_only_loss and not is_conversational(dataset_sample): + raise ValueError( + "You set `assistant_only_loss=True`, but the dataset is not conversational. This option is only " + "supported for conversational datasets." + ) + + # Dataset + # Skip dataset preparation if `skip_prepare_dataset=True` in `dataset_kwargs`, or if it's a VLM, where + # preprocessing (e.g., image-to-pixel conversion) is too costly and done on the fly instead. + skip_prepare_dataset = ( + args.dataset_kwargs is not None + and args.dataset_kwargs.get("skip_prepare_dataset", False) + or self._is_vision_dataset + ) + if not skip_prepare_dataset: + if self.completion_only_loss and formatting_func: + raise ValueError( + "A formatting function was provided while `completion_only_loss=True`, which is incompatible. " + "Using a formatter converts the dataset to a language modeling type, conflicting with " + "completion-only loss. To resolve this, apply your formatting function before passing the " + "dataset, or disable `completion_only_loss` in `SFTConfig`." + ) + train_dataset = self._prepare_dataset( + train_dataset, processing_class, args, args.packing, formatting_func, "train" + ) + if eval_dataset is not None: + packing = args.packing if args.eval_packing is None else args.eval_packing + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, packing, formatting_func, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset( + eval_dataset, processing_class, args, packing, formatting_func, "eval" + ) + + # Loss function + if args.loss_type == "nll": + pass # use the default loss + elif args.loss_type == "dft": + if compute_loss_func is not None: + raise ValueError( + "You passed a `compute_loss_func` together with `loss_type='dft'` to the `SFTTrainer`. " + "When using `loss_type='dft'`, the loss function is internally set to the DFT loss, so passing a " + "`compute_loss_func` is not allowed." + ) + compute_loss_func = dft_loss + else: + raise ValueError(f"Invalid `loss_type` {args.loss_type} passed. Supported values are 'nll' and 'dft'.") + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Initialize the Trainer. Parent class will handle: + # - DeepSpeed configuration (through create_accelerator_and_postprocess) + # - FSDP setup + # - Distributed training setup + # - Optimizer and scheduler creation + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + + def _prepare_dataset( + self, + dataset: Dataset | IterableDataset, + processing_class: PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin, + args: SFTConfig, + packing: bool, + formatting_func: Callable[[dict], str] | None, + dataset_name: str, + ) -> Dataset | IterableDataset: + # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from + # sampled data. + if isinstance(dataset, Dataset): # IterableDataset does not support `with_transform` + dataset = dataset.with_transform(remove_none_values) + + # If the dataset is already preprocessed (tokenized), skip the processing steps. + column_names = get_dataset_column_names(dataset) + is_processed = "input_ids" in column_names + + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + # Apply the formatting function if any + if formatting_func is not None and is_processed: + logger.warning( + "You passed a dataset that is already processed (contains an `input_ids` field) together with a " + "formatting function. Therefore `formatting_func` will be ignored. Either remove the " + "`formatting_func` or pass a dataset that is not already processed.", + ) + + if formatting_func is not None and not is_processed: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Applying formatting function to {dataset_name} dataset" + + def _func(example): + return {"text": formatting_func(example)} + + dataset = dataset.map(_func, batched=False, **map_kwargs) + + if not is_processed: + # Convert the dataset to ChatML if needed + first_example = next(iter(dataset)) + if is_conversational_from_value(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Converting {dataset_name} dataset to ChatML" + column_names = get_dataset_column_names(dataset) + dataset = dataset.map( + maybe_convert_to_chatml, + remove_columns="conversations" if "conversations" in column_names else None, + **map_kwargs, + ) + + # Apply the chat template if needed + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if "text" in example and not example["text"].endswith(eos_token): # language modeling case + example["text"] = example["text"] + eos_token + elif "completion" in example and not example["completion"].endswith(eos_token): + example["completion"] = example["completion"] + eos_token + return example + + dataset = dataset.map( + add_eos, + fn_kwargs={"eos_token": processing_class.eos_token}, + remove_columns="messages" if "messages" in column_names else None, # renamed to "text" + **map_kwargs, + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize_fn(example, processing_class, dataset_text_field, assistant_only_loss): + if "prompt" in example: # prompt-completion case + output = {} + if is_conversational(example): + if self._is_vlm: + prompt = prepare_multimodal_messages(example["prompt"], images=[]) + completion = prepare_multimodal_messages(example["completion"], images=[]) + else: + prompt = example["prompt"] + completion = example["completion"] + prompt_ids = processing_class.apply_chat_template( + prompt, + tokenize=True, + add_generation_prompt=True, + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + # Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists + # even for single examples, while for LLMs it returns lists of ints. + prompt_ids = prompt_ids[0] if isinstance(prompt_ids[0], list) else prompt_ids + prompt_completion_processed = processing_class.apply_chat_template( + prompt + completion, + return_dict=True, + tokenize=True, + return_assistant_tokens_mask=assistant_only_loss, + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + # Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists + # even for single examples, while for LLMs it returns lists of ints. + prompt_completion_processed = { + k: v[0] if isinstance(v[0], list) else v + for k, v in prompt_completion_processed.items() + } + prompt_completion_ids = prompt_completion_processed["input_ids"] + if "assistant_masks" in prompt_completion_processed: + output["assistant_masks"] = prompt_completion_processed["assistant_masks"] + else: + prompt_ids = processing_class(text=example["prompt"])["input_ids"] + prompt_completion_ids = processing_class(text=example["prompt"] + example["completion"])[ + "input_ids" + ] + + # Check if the tokenized prompt starts with the tokenized prompt+completion + if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids: + logger.warning( + "Mismatch between tokenized prompt and the start of tokenized prompt+completion. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + + # Create completion mask + completion_mask = [0] * len(prompt_ids) + [1] * (len(prompt_completion_ids) - len(prompt_ids)) + output["input_ids"] = prompt_completion_ids + output["completion_mask"] = completion_mask + + else: # language modeling case + if is_conversational(example): + if self._is_vlm: + messages = prepare_multimodal_messages(example["messages"], images=[]) + else: + messages = example["messages"] + processed = processing_class.apply_chat_template( + messages, + return_dict=True, + tokenize=True, + return_assistant_tokens_mask=assistant_only_loss, + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + # Fix transformers inconsistency: for VLMs, apply_chat_template returns lists of lists + # even for single examples, while for LLMs it returns lists of ints. + processed = {k: v[0] if isinstance(v[0], list) else v for k, v in processed.items()} + output = {k: processed[k] for k in ("input_ids", "assistant_masks") if k in processed} + else: + output = {"input_ids": processing_class(text=example[dataset_text_field])["input_ids"]} + + if "assistant_masks" in output and 1 not in output["assistant_masks"]: + raise RuntimeError( + "You're using `assistant_only_loss=True`, but at least one example has no assistant " + "tokens. This usually means the tokenizer's chat template doesn't generate assistant " + "masks — it may be missing the `{% generation %}` keyword. Please check the template and " + "ensure it's correctly configured to support assistant masking." + ) + return output + + dataset = dataset.map( + tokenize_fn, + fn_kwargs={ + "processing_class": processing_class, + "dataset_text_field": args.dataset_text_field, + "assistant_only_loss": args.assistant_only_loss, + }, + **map_kwargs, + ) + + # Pack or truncate + if packing: + if args.max_length is None: + raise ValueError("When packing is enabled, `max_length` can't be `None`.") + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Packing {dataset_name} dataset" + + columns = ["input_ids"] + if "completion_mask" in get_dataset_column_names(dataset): + columns.append("completion_mask") + if "assistant_masks" in get_dataset_column_names(dataset): + columns.append("assistant_masks") + + dataset = dataset.select_columns(columns) + + # Packing adds new column "seq_lengths" needed for document aware FlashAttention + dataset = pack_dataset(dataset, args.max_length, args.packing_strategy, map_kwargs) + elif args.max_length is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Truncating {dataset_name} dataset" + dataset = truncate_dataset(dataset, args.max_length, map_kwargs) + # For Liger kernel, ensure only the essential columns + if args.use_liger_kernel: + collator_expected_keys = {"input_ids", "seq_lengths", "completion_mask", "assistant_masks"} + column_names = get_dataset_column_names(dataset) + dataset = dataset.select_columns(collator_expected_keys.intersection(column_names)) + + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). When using `train_on_completion_only` we add a "completion_mask" column to the + # dataset. So we need to override the default signature columns to include "completion_mask" as well. + if self._signature_columns is None: + if self._is_vision_dataset: + self._signature_columns = ["messages", "prompt", "completion", "images"] + else: + self._signature_columns = ["input_ids", "labels", "seq_lengths", "completion_mask", "assistant_masks"] + + def compute_loss( + self, + model: nn.Module, + inputs: dict[str, torch.Tensor | Any], + return_outputs: bool = False, + num_items_in_batch: torch.Tensor | None = None, + ): + """ + Compute training loss and additionally compute token accuracies + """ + mode = "train" if self.model.training else "eval" + + # Set aside labels as it will be dropped by super().compute_loss() if a custom `compute_loss_func` is used. + # This can be removed when this issue is fixed. + labels = inputs["labels"] + + # If not set, defaults from model config and may warn since cache isn't compatible with gradient checkpointing + inputs["use_cache"] = False + (loss, outputs) = super().compute_loss( + model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch + ) + + # Compute entropy + if not self.args.use_liger_kernel: # liger doesn't return logits + with torch.no_grad(): + per_token_entropy = entropy_from_logits(outputs.logits) + # When using Prompt Tuning, skip the virtual tokens in logits before entropy computation, since they + # do not correspond to actual input tokens. + if ( + self.num_virtual_tokens > 0 + and model.peft_config[model.active_adapter].peft_type != PeftType.PREFIX_TUNING + ): + per_token_entropy = per_token_entropy[:, self.num_virtual_tokens :] + if "attention_mask" in inputs: + attention_mask = inputs["attention_mask"] + entropy = torch.sum(per_token_entropy * attention_mask) / attention_mask.sum() + elif "position_ids" in inputs: + entropy = torch.mean(per_token_entropy) + else: + raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.") + entropy = self.accelerator.gather_for_metrics(entropy).mean().item() + self._metrics[mode]["entropy"].append(entropy) + + if mode == "train": + # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q, + # cu_seq_lens_k, and max_length_k, max_length_q and position_ids. + if "attention_mask" in inputs: + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + elif "position_ids" in inputs: + local_num_tokens = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device) + num_tokens_in_batch = self.accelerator.gather_for_metrics(local_num_tokens).sum().item() + else: + raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.") + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + # Compute token accuracy if we have labels and if the model is not using Liger (no logits) + if not self.args.use_liger_kernel: + with torch.no_grad(): + if "shift_labels" in inputs: + # When using CP, labels are pre-shifted. We must use these (and cannot manually shift) because: + # - The first discarded token from inputs["labels"] actually belongs to process n-1 + # - The last logits require the label from process n+1 + shift_logits = outputs.logits.contiguous() + shift_labels = inputs["shift_labels"] + else: + shift_logits = outputs.logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + # Prompt Tuning and P-Tuning output logits for virtual tokens but Prefix-Tuning does not. + if ( + self.num_virtual_tokens > 0 + and model.peft_config[model.active_adapter].peft_type != PeftType.PREFIX_TUNING + ): + shift_logits = shift_logits[:, self.num_virtual_tokens :, :] + + # Get predictions + predictions = shift_logits.argmax(dim=-1) + + # Create mask for non-padding tokens (assuming ignore_index is -100) + mask = shift_labels != -100 + + # Calculate accuracy only on non-padding tokens + correct_predictions = (predictions == shift_labels) & mask + total_tokens = mask.sum() + correct_tokens = correct_predictions.sum() + + # Gather the correct_tokens and total_tokens across all processes + correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) + total_tokens = self.accelerator.gather_for_metrics(total_tokens) + + # Compute the mean token accuracy and log it + total_sum = total_tokens.sum() + accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 + self._metrics[mode]["mean_token_accuracy"].append(accuracy) + if self.aux_loss_enabled: + aux_loss = outputs.aux_loss + aux_loss = self.accelerator.gather_for_metrics(aux_loss).mean().item() + self._metrics[mode]["aux_loss"].append(aux_loss) + + return (loss, outputs) if return_outputs else loss + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: float | None = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs.update(metrics) + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/venv/lib/python3.10/site-packages/trl/trainer/utils.py b/venv/lib/python3.10/site-packages/trl/trainer/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..edeb50c0323cf3563c4ae426cb61753455b315d4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/utils.py @@ -0,0 +1,2018 @@ +# Copyright 2020-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 dataclasses +import importlib.resources as pkg_resources +import json +import os +import random +import socket +import warnings +from collections.abc import Mapping, Sequence, Sized +from dataclasses import dataclass, field +from importlib.metadata import version +from itertools import accumulate +from typing import Any, Literal, TypeVar + +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +import torch.utils.data +import transformers +from accelerate import Accelerator, PartialState, logging +from accelerate.state import AcceleratorState +from huggingface_hub import ModelCard, ModelCardData +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import Sampler +from transformers import ( + AutoConfig, + BitsAndBytesConfig, + EvalPrediction, + GenerationConfig, + PretrainedConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + TrainerState, + TrainingArguments, + is_comet_available, +) +from transformers.utils import ( + ModelOutput, + is_peft_available, + is_rich_available, + is_torch_mlu_available, + is_torch_npu_available, + is_torch_xpu_available, +) + +from ..trainer.model_config import ModelConfig + + +if is_rich_available(): + from rich.console import Console + from rich.panel import Panel + from rich.table import Table + from rich.text import Text + +if is_comet_available(): + import comet_ml + +if is_peft_available(): + from peft import LoraConfig, PeftConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +class DataCollatorForChatML: + """ + Data collator for ChatML format datasets. + """ + + tokenizer: PreTrainedTokenizerBase + ignore_index: int = -100 + max_length: int = None + prompt_key: str = "prompt" + messages_key: str = "messages" + + def __post_init__(self): + if self.tokenizer.pad_token_id is None: + raise ValueError("The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.") + if self.max_length is None: + # set a sensible default + self.max_length = min(self.tokenizer.model_max_length, 1024) + + def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + input_ids = [] + attention_mask = [] + prompts_input_ids = [] + prompt_attention_mask = [] + labels = [] + + for example in examples: + formatted_prompt = example.get(self.prompt_key, None) + if formatted_prompt is None: + prompt = example[self.messages_key][:-1] + formatted_prompt = self.tokenizer.apply_chat_template( + prompt, tokenize=False, add_generation_prompt=True + ) + + if "input_ids" not in example: + message = example[self.messages_key] + formatted_message = self.tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=False + ) + + tokenized_message = self.tokenizer( + formatted_message, + truncation=False, + padding=False, + return_tensors=None, + add_special_tokens=False, + return_offsets_mapping=True, + ) + message_input_ids_full = tokenized_message["input_ids"] + offsets = tokenized_message.get("offset_mapping") + + if offsets is not None: + prompt_char_len = len(formatted_prompt) + completion_start_idx_full = next( + (idx for idx, (start, _) in enumerate(offsets) if start >= prompt_char_len), + len(message_input_ids_full), + ) + else: + tokenized_prompt_full = self.tokenizer( + formatted_prompt, + truncation=False, + padding=False, + return_tensors=None, + add_special_tokens=False, + ) + completion_start_idx_full = len(tokenized_prompt_full["input_ids"]) + + prompt_tokens_full = message_input_ids_full[:completion_start_idx_full] + completion_input_ids_full = message_input_ids_full[completion_start_idx_full:] + + if self.max_length is not None and len(message_input_ids_full) > self.max_length: + completion_ids = completion_input_ids_full + if len(completion_ids) >= self.max_length: + completion_ids = completion_ids[-self.max_length :] + prompt_ids = [] + else: + max_prompt_tokens = self.max_length - len(completion_ids) + prompt_ids = prompt_tokens_full[-max_prompt_tokens:] if max_prompt_tokens > 0 else [] + message_input_ids = prompt_ids + completion_ids + else: + message_input_ids = message_input_ids_full + prompt_ids = prompt_tokens_full + + input_ids.append(message_input_ids) + attention_mask.append([1] * len(message_input_ids)) + current_prompt_ids = prompt_ids + else: + message_input_ids = example["input_ids"] + input_ids.append(message_input_ids) + if "attention_mask" in example: + attention_mask.append(example["attention_mask"]) + else: + attention_mask.append([1] * len(message_input_ids)) + + tokenized_prompt = self.tokenizer( + formatted_prompt, + truncation=True, + max_length=len(message_input_ids), + padding=False, + return_tensors=None, + add_special_tokens=False, + ) + current_prompt_ids = tokenized_prompt["input_ids"] + + prompts_input_ids.append(current_prompt_ids) + prompt_attention_mask.append([1] * len(current_prompt_ids)) + + label = [self.ignore_index] * len(input_ids[-1]) + completion_start_idx = len(current_prompt_ids) + label[completion_start_idx:] = input_ids[-1][completion_start_idx:] + labels.append(label) + + # convert to list of tensors and pad + input_ids = [torch.tensor(ids, dtype=torch.long) for ids in input_ids] + attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in attention_mask] + labels = [torch.tensor(label, dtype=torch.long) for label in labels] + input_ids = pad(input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) + attention_mask = pad(attention_mask, padding_side="left", padding_value=0) + labels = pad(labels, padding_side="left", padding_value=self.ignore_index) + + prompts_input_ids = [torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids] + prompt_attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask] + prompts_input_ids = pad(prompts_input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) + prompt_attention_mask = pad(prompt_attention_mask, padding_side="left", padding_value=0) + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "prompts": prompts_input_ids, + "prompt_attention_mask": prompt_attention_mask, + } + + +def _is_port_free(port: int, host: str = "127.0.0.1") -> bool: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind((host, port)) + return True + except OSError: + return False + + +def _find_free_port() -> int: + candidates = (29500, 23456, 12355, 12345) + for p in candidates: + if _is_port_free(p): + return p + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def ensure_master_addr_port(addr: str | None = None, port: int | None = None) -> None: + """ + Ensure `MASTER_ADDR`/`MASTER_PORT` are set safely. + + - Respects existing environment variables. + - Defaults `MASTER_ADDR` to localhost if unset. + - Chooses a free TCP port if `MASTER_PORT` is unset to avoid collisions. + - If `MASTER_PORT` is set to `"0"` or `"auto"`, it is resolved to a free port. + """ + os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR") or addr or "localhost" + + env_port = os.environ.get("MASTER_PORT", "").strip().lower() + if port is None and env_port not in {"", "0", "auto"}: + try: + port = int(env_port) + except ValueError: + pass + + os.environ["MASTER_PORT"] = str(_find_free_port() if port in (None, 0) else port) + + +@dataclass +class RewardDataCollatorWithPadding: + # docstyle-ignore + r""" + Reward DataCollator class that pads the inputs to the maximum length of the batch. + + > [!WARNING] + > This class is deprecated and will be removed in version 0.27.0. Please use + `trl.trainer.reward_trainer.DataCollatorForPreference` instead. + + Args: + tokenizer ([`~transformers.PreTrainedTokenizerBase`]): + The tokenizer used for encoding the data. + padding (`bool | str | PaddingStrategy`, `optional`, defaults to `True`): + padding_strategy to pass to the tokenizer. + pad_to_multiple_of (`int` or `None`, `optional`, defaults to `None`): + If set will pad the sequence to a multiple of the provided value. + return_tensors (`str`, `optional`, defaults to `"pt"`): + The tensor type to use. + """ + + tokenizer: PreTrainedTokenizerBase + padding: bool | str = True + pad_to_multiple_of: int | None = None + return_tensors: str = "pt" + + def __init__(self, *args, **kwargs) -> None: + warnings.warn( + "The `RewardDataCollatorWithPadding` is deprecated and will be removed in version 0.27.0. Please use " + "`trl.trainer.reward_trainer.DataCollatorForPreference` instead.", + FutureWarning, + ) + super().__init__(*args, **kwargs) + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: + features_chosen = [] + features_rejected = [] + margin = [] + # check if we have a margin. If we do, we need to batch it as well + has_margin = "margin" in features[0] + for feature in features: + # check if the keys are named as expected + if ( + "input_ids_chosen" not in feature + or "input_ids_rejected" not in feature + or "attention_mask_chosen" not in feature + or "attention_mask_rejected" not in feature + ): + raise ValueError( + "The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`" + ) + + features_chosen.append( + { + "input_ids": feature["input_ids_chosen"], + "attention_mask": feature["attention_mask_chosen"], + } + ) + features_rejected.append( + { + "input_ids": feature["input_ids_rejected"], + "attention_mask": feature["attention_mask_rejected"], + } + ) + if has_margin: + margin.append(feature["margin"]) + batch_chosen = self.tokenizer.pad( + features_chosen, + padding=self.padding, + pad_to_multiple_of=self.pad_to_multiple_of, + return_tensors=self.return_tensors, + ) + batch_rejected = self.tokenizer.pad( + features_rejected, + padding=self.padding, + pad_to_multiple_of=self.pad_to_multiple_of, + return_tensors=self.return_tensors, + ) + batch = { + "input_ids_chosen": batch_chosen["input_ids"], + "attention_mask_chosen": batch_chosen["attention_mask"], + "input_ids_rejected": batch_rejected["input_ids"], + "attention_mask_rejected": batch_rejected["attention_mask"], + "return_loss": True, + } + if has_margin: + margin = torch.tensor(margin, dtype=torch.float) + batch["margin"] = margin + return batch + + +def pad( + tensors: list[torch.Tensor], + padding_value: int = 0, + padding_side: str = "right", + pad_to_multiple_of: int | None = None, +) -> torch.Tensor: + """ + Pads a list of tensors to the same shape along the first dimension. + + Args: + tensors (`list[torch.Tensor]`): + List of input tensors to pad. + padding_value (`int`): + Value to use for padding. Default is 0. + padding_side (`str`): + Side on which to add padding. Must be 'left' or 'right'. Default is 'right'. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + Returns: + `torch.Tensor`: + A single tensor containing the padded tensors. + + Examples: + ```python + >>> import torch + + >>> pad([torch.tensor([1, 2, 3]), torch.tensor([4, 5])]) + tensor([[1, 2, 3], + [4, 5, 0]]) + + >>> pad([torch.tensor([[1, 2], [3, 4]]), torch.tensor([[5, 6]])]) + tensor([[[1, 2], + [3, 4]], + [[5, 6], + [0, 0]]]) + ``` + """ + # Determine the maximum shape for each dimension + output_shape = np.max([t.shape for t in tensors], 0).tolist() + + # Apply pad_to_multiple_of to the first (sequence) dimension + if pad_to_multiple_of is not None: + remainder = output_shape[0] % pad_to_multiple_of + if remainder != 0: + output_shape[0] += pad_to_multiple_of - remainder + + # Create an output tensor filled with the padding value + output = torch.full((len(tensors), *output_shape), padding_value, dtype=tensors[0].dtype, device=tensors[0].device) + + for i, t in enumerate(tensors): + if padding_side == "left": + seq_start = output_shape[0] - t.shape[0] + elif padding_side == "right": + seq_start = 0 + else: + raise ValueError("padding_side must be 'left' or 'right'") + + # Define the slices + seq_slice = slice(seq_start, seq_start + t.shape[0]) + slices = (seq_slice,) + tuple(slice(0, s) for s in t.shape[1:]) + output[i][slices] = t + + return output + + +@dataclass +class DPODataCollatorWithPadding: + r""" + DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch. + + Args: + pad_token_id (`int` defaults to 0): + The tokenizer's pad_token_id. + label_pad_token_id (`int`, defaults to -100): + The label used for masking. + is_encoder_decoder (`bool` or `None`, `optional`, defaults to `None`): + Whether you model has an encoder_decoder architecture. + """ + + pad_token_id: int = 0 + label_pad_token_id: int = -100 + is_encoder_decoder: bool | None = False + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: + # first, pad everything to the same length + padded_batch = {} + for k in features[0].keys(): + if k.endswith(("_input_ids", "_attention_mask", "_labels", "_pixel_values")): + if self.is_encoder_decoder: + to_pad = [torch.LongTensor(ex[k]) for ex in features] + + if (k.startswith("prompt")) and (k.endswith("input_ids")): + if self.pad_token_id is None: + raise ValueError( + "Padding is enabled, but the tokenizer is not configured with a padding token." + " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)" + " before calling the trainer." + ) + padding_value = self.pad_token_id + elif k.endswith("_attention_mask"): + padding_value = 0 + elif k.startswith(("chosen", "rejected", "completion")) or ("decoder" in k): + padding_value = self.label_pad_token_id + else: + raise ValueError(f"Unexpected key in batch '{k}'") + padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value) + else: + # Set padding value based on the key + if k.endswith("_input_ids"): + if self.pad_token_id is None: + raise ValueError( + "Padding is enabled, but the tokenizer is not configured with a padding token." + " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)" + " before calling the trainer." + ) + padding_value = self.pad_token_id + elif k.endswith("_labels"): + padding_value = self.label_pad_token_id + elif k.endswith("_attention_mask"): + padding_value = 0 + elif k.endswith("_pixel_values"): + padding_value = 0 # TODO: check if this is correct + else: + raise ValueError(f"Unexpected key in batch '{k}'") + + # Set padding side based on the key + if k in ["prompt_input_ids", "prompt_attention_mask"]: + padding_side = "left" + else: + padding_side = "right" + + # Set the dtype + if k.endswith("_pixel_values"): + dtype = torch.float32 # will be downcasted if necessary by the Trainer + else: + dtype = torch.int64 + + # Convert to tensor and pad + to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features] + padded_batch[k] = pad(to_pad, padding_value=padding_value, padding_side=padding_side) + elif k.endswith("_logps"): + # the cached reference model logprobs + padded_batch[k] = torch.tensor([ex[k] for ex in features]) + else: + padded_batch[k] = [ex[k] for ex in features] + + return padded_batch + + +@dataclass +class RunningMoments: + """ + Calculates the running mean and standard deviation of a data stream. Reference: + https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L75 + """ + + accelerator: Accelerator + mean: float = 0 + std: float = 1 + var: float = 1 + count: float = 1e-24 + + @torch.no_grad() + def update(self, xs: torch.Tensor) -> tuple[float, float]: + """ + Updates running moments from batch's moments computed across ranks + """ + if self.accelerator.use_distributed: + xs_mean, xs_var, xs_count = get_global_statistics(self.accelerator, xs) + else: + xs_count = xs.numel() + xs_var, xs_mean = torch.var_mean(xs, unbiased=False) + xs_mean, xs_var = xs_mean.float(), xs_var.float() + + delta = xs_mean - self.mean + tot_count = self.count + xs_count + + new_sum = xs_var * xs_count + # correct old_sum deviation accounting for the new mean + old_sum = self.var * self.count + delta**2 * self.count * xs_count / tot_count + tot_sum = old_sum + new_sum + + self.mean += (delta * xs_count / tot_count).item() + new_var = tot_sum / tot_count + self.std = (new_var * tot_count / (tot_count - 1)).float().sqrt().item() + self.var = new_var.item() + self.count = tot_count + + return xs_mean.item(), (xs_var * xs_count / (xs_count - 1)).float().sqrt().item() + + def save_to_json(self, json_path: str): + """Save the content of this instance in JSON format inside `json_path`.""" + # save everything except accelerator + if self.accelerator.is_main_process: + save_dict = dataclasses.asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if k != "accelerator"}) + json_string = json.dumps(save_dict, indent=2, sort_keys=True) + "\n" + with open(json_path, "w", encoding="utf-8") as f: + f.write(json_string) + + @classmethod + def load_from_json(cls, accelerator: Accelerator, json_path: str): + """Create an instance from the content of `json_path`.""" + # load everything except accelerator + with open(json_path, encoding="utf-8") as f: + text = f.read() + return cls(accelerator=accelerator, **json.loads(text)) + + +@torch.no_grad() +def get_global_statistics( + accelerator, xs: torch.Tensor, mask=None, device="cpu" +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Computes element-wise mean and variance of the tensor across processes. Reference: + https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L57C1-L73C75 + """ + xs = xs.to(accelerator.device) + sum_and_count = torch.tensor([xs.sum(), (xs.numel() if mask is None else mask.sum())], device=xs.device) + sum_and_count = accelerator.reduce(sum_and_count) + global_sum, count = sum_and_count + global_mean = global_sum / count + + sum_var = torch.sum(((xs - global_mean) ** 2).mul(1 if mask is None else mask)) + sum_var = accelerator.reduce(sum_var) + global_var = sum_var / count + + return global_mean.to(device), global_var.to(device), count.item() + + +def compute_accuracy(eval_pred: EvalPrediction) -> dict[str, float]: + predictions, labels = eval_pred + if predictions.ndim == 3: + # Token classification task. Shapes are (batch_size, seq_len, num_labels) and (batch_size, seq_len) + # Used to compute the accuracy in the prm_trainer. + predictions = np.argmax(predictions, axis=2) + + # Flatten the predictions and labels to remove the ignored tokens. + predictions = np.array( + [ + p + for prediction, label in zip(predictions, labels, strict=True) + for (p, lbl) in zip(prediction, label, strict=True) + if lbl != -100 + ] + ) + labels = np.array([lbl for label in labels for lbl in label if lbl != -100]) + + else: + # Here, predictions is rewards_chosen and rewards_rejected. Shapes are (batch_size, 2) and (batch_size,) + # We want to see how much of the time rewards_chosen > rewards_rejected. + equal_mask = predictions[:, 0] == predictions[:, 1] + equal_predictions_count = int(equal_mask.sum()) + + if equal_predictions_count > 0: + # Before using the logger, the accelerate state must be initialized. It'susually the case when using this + # function inside a Trainer, but it may not be the case otherwise, in particular when unit testing. + PartialState() + + logger.warning( + f"There are {equal_predictions_count} out of {len(predictions[:, 0])} instances where the predictions " + "for both options are equal. These instances are ignored in the accuracy computation.", + ) + + # Filter out equal predictions + predictions = predictions[~equal_mask] + labels = labels[~equal_mask] + + # Use the remaining predictions for accuracy calculation + predictions = np.argmax(predictions, axis=1) + + accuracy = np.array(predictions == labels, dtype=float).mean().item() + return {"accuracy": accuracy} + + +def pad_to_length(tensor: torch.Tensor, length: int, pad_value: int | float, dim: int = -1) -> torch.Tensor: + if tensor.size(dim) >= length: + return tensor + else: + pad_size = list(tensor.shape) + pad_size[dim] = length - tensor.size(dim) + return torch.cat( + [ + tensor, + pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), + ], + dim=dim, + ) + + +def disable_dropout_in_model(model: torch.nn.Module) -> None: + for module in model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0 + + +def exact_div(a, b, custom_error_message=""): + q = a // b + if a != q * b: + raise ValueError(f"{custom_error_message}, inexact division: {a} / {b} = {a / b}") + return q + + +def peft_module_casting_to_bf16(model): + for name, module in model.named_modules(): + if isinstance(module, torch.nn.LayerNorm) or "norm" in name: + module = module.to(torch.float32) + elif any(x in name for x in ["lm_head", "embed_tokens", "wte", "wpe"]): + if hasattr(module, "weight"): + if module.weight.dtype == torch.float32: + module = module.to(torch.bfloat16) + + +def get_quantization_config(model_args: ModelConfig) -> BitsAndBytesConfig | None: + if model_args.load_in_4bit: + quantization_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=model_args.dtype, # For consistency with model weights, we use the same value as `dtype` + bnb_4bit_quant_type=model_args.bnb_4bit_quant_type, + bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant, + bnb_4bit_quant_storage=model_args.dtype, + ) + elif model_args.load_in_8bit: + quantization_config = BitsAndBytesConfig( + load_in_8bit=True, + ) + else: + quantization_config = None + + return quantization_config + + +def get_kbit_device_map() -> dict[str, int] | None: + if torch.cuda.is_available() or is_torch_xpu_available(): + return {"": PartialState().local_process_index} + else: + return None + + +def get_peft_config(model_args: ModelConfig) -> "PeftConfig | None": + if model_args.use_peft is False: + return None + + if not is_peft_available(): + raise ValueError( + "You need to have PEFT library installed in your environment, make sure to install `peft`. " + "Make sure to run `pip install -U peft`." + ) + + peft_config = LoraConfig( + task_type=model_args.lora_task_type, + r=model_args.lora_r, + target_modules=model_args.lora_target_modules, + lora_alpha=model_args.lora_alpha, + lora_dropout=model_args.lora_dropout, + bias="none", + use_rslora=model_args.use_rslora, + use_dora=model_args.use_dora, + modules_to_save=model_args.lora_modules_to_save, + ) + + return peft_config + + +def get_exp_cap(value, decimal=4): + """ + Get the exponent cap of a value. This is used to cap the exponent of a value to avoid overflow. The formula is : + log(value.dtype.max) E.g. for float32 data type, the maximum exponent value is 88.7228 to 4 decimal points. + + Args: + value (`torch.Tensor`): + The input tensor to obtain the data type + decimal (`int`): + The number of decimal points of the output exponent cap. eg: direct calling exp(log(torch.float32.max)) + will result in inf so we cap the exponent to 88.7228 to avoid overflow. + """ + vdtype_max = torch.zeros([1]).to(value.dtype) + torch.finfo(value.dtype).max + vdtype_log_max = torch.log(vdtype_max).to(value.device) + return torch.floor(vdtype_log_max * 10**decimal) / 10**decimal if decimal > 0 else vdtype_log_max + + +def cap_exp(value, cap=-1): + # Cap the exponent value below the upper-bound to avoid overflow, before calling torch.exp + cap = get_exp_cap(value) if cap < 0 else cap + return torch.exp(torch.clamp(value, max=cap)) + + +def print_rich_table(df: pd.DataFrame) -> None: + if not is_rich_available(): + raise ImportError( + "The function `print_rich_table` requires the `rich` library. Please install it with `pip install rich`." + ) + console = Console() + table = Table(show_lines=True) + for column in df.columns: + table.add_column(column) + for _, row in df.iterrows(): + table.add_row(*row.astype(str).tolist()) + console.print(table) + + +SIMPLE_CHAT_TEMPLATE = "{% for message in messages %}{{message['role'].capitalize() + ': ' + message['content'] + '\n\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}" + + +@dataclass +class OnlineTrainerState(TrainerState): + """ + Training state for online/on-policy trainers. + + Extends [`~transformers.TrainerState`] with an `episode` counter to track the current rollout/episode. + + Args: + episode (`int`, defaults to 0): Zero-based episode index. + """ + + episode: int = 0 + + +@dataclass +class OnPolicyConfig(TrainingArguments): + r""" + Base configuration class for on-policy trainers. + + This class includes only the parameters that are specific to some on-policy training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + run_name (`str`, *optional*): + Name of the run. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + num_mini_batches (`int`, *optional*, defaults to `1`): + Number of minibatches to split a batch into. + total_episodes (`int`, *optional*): + Total number of episodes in the dataset. + local_rollout_forward_batch_size (`int`, *optional*, defaults to `64`): + Per rank no grad forward pass in the rollout phase. + num_sample_generations (`int`, *optional*, defaults to `10`): + Number of debugging samples generations (i.e., `generate_completions` calls) throughout training. + response_length (`int`, *optional*, defaults to `53`): + Length of the response. + stop_token (`str`, *optional*): + Specifies the stop token to use for text generation. This parameter is mutually exclusive with + `stop_token_id`. + + - `None`: No stop token is applied, unless `stop_token_id` is specified. + - `'eos'`: Uses the tokenizer's `eos_token`. + + stop_token_id (`int`, *optional*): + Specifies the ID of the stop token to use for text generation. If `None`, no stop token ID is applied, + unless `stop_token` is specified. This parameter is mutually exclusive with `stop_token`. + temperature (`float`, *optional*, defaults to `0.7`): + Sampling temperature. + missing_eos_penalty (`float`, *optional*): + Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to + generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive + value. + sft_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the SFT model. + world_size (`int`, *optional*): + Number of processes (GPUs) to use for the training. + num_total_batches (`int`, *optional*): + Number of total batches to train. + micro_batch_size (`int`, *optional*): + Micro batch size across devices (HF's `per_device_train_batch_size` * `world_size`). + local_batch_size (`int`, *optional*): + Batch size per GPU (HF's `per_device_train_batch_size` * `gradient_accumulation_steps`). + batch_size (`int`, *optional*): + Batch size across devices (HF's `per_device_train_batch_size` * `world_size` * + `gradient_accumulation_steps`). + local_mini_batch_size (`int`, *optional*): + Mini batch size per GPU. + mini_batch_size (`int`, *optional*): + Mini batch size across GPUs. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the model to the Hub after training. + """ + + # Parameters whose default values are overridden from TrainingArguments + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: bool | None = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + run_name: str | None = field( + default=None, + metadata={"help": "Name of the run."}, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + num_mini_batches: int = field( + default=1, + metadata={"help": "Number of minibatches to split a batch into."}, + ) + total_episodes: int | None = field( + default=None, + metadata={"help": "Total number of episodes in the dataset."}, + ) + local_rollout_forward_batch_size: int = field( + default=64, + metadata={"help": "Per rank no grad forward pass in the rollout phase."}, + ) + num_sample_generations: int = field( + default=10, + metadata={ + "help": "Number of debugging samples generations (i.e., `generate_completions` calls) throughout training." + }, + ) + response_length: int = field( + default=53, + metadata={"help": "Length of the response."}, + ) + stop_token: Literal["eos"] | None = field( + default=None, + metadata={ + "help": "Specifies the stop token to use for text generation. This parameter is mutually exclusive with " + "`stop_token_id`." + }, + ) + stop_token_id: int | None = field( + default=None, + metadata={ + "help": "Specifies the ID of the stop token to use for text generation. If `None`, no stop token ID is " + "applied, unless `stop_token` is specified. This parameter is mutually exclusive with `stop_token`." + }, + ) + temperature: float = field( + default=0.7, + metadata={"help": "Sampling temperature."}, + ) + missing_eos_penalty: float | None = field( + default=None, + metadata={ + "help": "Penalty applied to the score when the model fails to generate an EOS token. This is useful to " + "encourage to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be " + "a positive value." + }, + ) + sft_model_path: str = field( + default="EleutherAI/pythia-160m", + metadata={"help": "Path to the SFT model."}, + ) + world_size: int | None = field( + default=None, + metadata={"help": "Number of processes (GPUs) to use for the training."}, + ) + num_total_batches: int | None = field( + default=None, + metadata={"help": "Number of total batches to train."}, + ) + micro_batch_size: int | None = field( + default=None, + metadata={"help": "Micro batch size across devices (HF's `per_device_train_batch_size` * `world_size`)."}, + ) + local_batch_size: int | None = field( + default=None, + metadata={"help": "Batch size per GPU (HF's `per_device_train_batch_size` * `gradient_accumulation_steps`)."}, + ) + batch_size: int | None = field( + default=None, + metadata={ + "help": "Batch size across devices (HF's `per_device_train_batch_size` * `world_size` * " + "`gradient_accumulation_steps`)." + }, + ) + local_mini_batch_size: int | None = field( + default=None, + metadata={"help": "Mini batch size per GPU."}, + ) + mini_batch_size: int | None = field( + default=None, + metadata={"help": "Mini batch size across GPUs."}, + ) + push_to_hub: bool = field( + default=False, + metadata={"help": "Whether to push the model to the Hub after training."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + +def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor: + """ + Takes an N-dimensional bool tensor and returns an (N-1)-dimensional tensor of integers giving the position of the + first True in each "row". + + Returns the length of the rows (bools.size(-1)) if no element is True in a given row. + + Args: + bools (`torch.Tensor`): + An N-dimensional boolean tensor. + dtype (`torch.dtype`, optional): + The desired data type of the output tensor. Defaults to `torch.long`. + + Returns: + `torch.Tensor`: + An (N-1)-dimensional tensor of integers indicating the position of the first True in each row. If no True + value is found in a row, returns the length of the row. + """ + row_len = bools.size(-1) + zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device) + return torch.min(zero_or_index, dim=-1).values + + +def get_reward( + model: torch.nn.Module, query_responses: torch.Tensor, pad_token_id: int, context_length: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Computes the reward logits and the rewards for a given model and query responses. + + Args: + model (`torch.nn.Module`): + The model used to compute the reward logits. + query_responses (`torch.Tensor`): + The tensor containing the query responses. + pad_token_id (`int`): + The token ID representing the pad token. + context_length (`int`): + The length of the context in the query responses. + + Returns: + tuple: + - `reward_logits` (`torch.Tensor`): + The logits for the reward model. + - `final_rewards` (`torch.Tensor`): + The final rewards for each query response. + - `sequence_lengths` (`torch.Tensor`): + The lengths of the sequences in the query responses. + """ + attention_mask = query_responses != pad_token_id + position_ids = attention_mask.cumsum(1) - attention_mask.long() # exclusive cumsum + lm_backbone = getattr(model, model.base_model_prefix) + input_ids = torch.masked_fill(query_responses, ~attention_mask, 0) + output = lm_backbone( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + output_hidden_states=True, + use_cache=False, # otherwise mistral-based RM would error out + ) + reward_logits = model.score(output.hidden_states[-1]) + sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length + # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454 + return ( + reward_logits, + reward_logits[ + torch.arange(reward_logits.size(0), device=reward_logits.device), + sequence_lengths, + ].squeeze(-1), + sequence_lengths, + ) + + +def forward( + model: torch.nn.Module, + query_responses: torch.Tensor, + pad_token_id: int, +) -> ModelOutput: + """ + Performs a forward pass through the model with the given query responses and pad token ID. + + Args: + model (`torch.nn.Module`): + The model to perform the forward pass. + query_responses (`torch.Tensor`): + The tensor containing the query responses. + pad_token_id (`int`): + The token ID representing the pad token. + + Returns: + `ModelOutput`: + The output of the model, including hidden states. + """ + attention_mask = query_responses != pad_token_id + position_ids = attention_mask.cumsum(1) - attention_mask.long() + input_ids = torch.masked_fill(query_responses, ~attention_mask, 0) + return model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + output_hidden_states=True, + ) + + +def prepare_deepspeed( + model: torch.nn.Module, per_device_train_batch_size: int, fp16: bool = False, bf16: bool = False +) -> torch.nn.Module: + """ + Prepares the model for training with DeepSpeed (both for stage 2 and 3), configuring the appropriate settings based + on the model and batch size. + + Args: + model (`torch.nn.Module`): + The model to be prepared for DeepSpeed training. + per_device_train_batch_size (`int`): + The training batch size per device. + fp16 (`bool`, defaults to `False`): + Whether to use FP16 precision. + bf16 (`bool`, defaults to `False`): + Whether to use BF16 precision. + + Returns: + `torch.nn.Module`: + The model initialized and configured with DeepSpeed for training. + """ + import deepspeed + + deepspeed_plugin = AcceleratorState().deepspeed_plugin + config_kwargs = deepspeed_plugin.deepspeed_config + if config_kwargs["zero_optimization"]["stage"] != 3: + config_kwargs["train_micro_batch_size_per_gpu"] = per_device_train_batch_size + config_kwargs = { + "train_micro_batch_size_per_gpu": config_kwargs["train_micro_batch_size_per_gpu"], + "prescale_gradients": False, + "wall_clock_breakdown": False, + } + if bf16: + config_kwargs["bf16"] = {"enabled": True} + elif fp16: + config_kwargs["fp16"] = {"enabled": True} + else: + if hasattr(model, "config"): + hidden_size = ( + max(model.config.hidden_sizes) + if getattr(model.config, "hidden_sizes", None) + else getattr(model.config, "hidden_size", None) + ) + if hidden_size is not None and config_kwargs["zero_optimization"]["stage"] == 3: + # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0` + # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081 + config_kwargs.update( + { + "zero_optimization.reduce_bucket_size": hidden_size * hidden_size, + "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size, + "zero_optimization.stage3_prefetch_bucket_size": 0, + } + ) + model, *_ = deepspeed.initialize(model=model, config=config_kwargs) + model.eval() + return model + + +def truncate_response(stop_token_id: int, pad_token_id: int, responses: torch.Tensor) -> torch.Tensor: + """ + Truncates the responses at the first occurrence of the stop token, filling the rest with pad tokens. + + Args: + stop_token_id (`int`): + The token ID representing the stop token where truncation occurs. + pad_token_id (`int`): + The token ID representing the pad token used to fill the truncated responses. + responses (`torch.Tensor`): + The tensor containing the responses to be truncated. + + Returns: + `torch.Tensor`: + The truncated responses tensor with pad tokens filled after the stop token. + """ + trunc_idxs = first_true_indices(responses == stop_token_id).unsqueeze(-1) + new_size = [1] * (len(responses.size()) - 1) + [responses.shape[1]] + idxs = torch.arange(responses.shape[1], device=responses.device).view(*new_size) + postprocessed_responses = torch.masked_fill(responses, idxs > trunc_idxs, pad_token_id) + return postprocessed_responses + + +def generate( + lm_backbone: torch.nn.Module, queries: torch.Tensor, pad_token_id: int, generation_config: GenerationConfig +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Generates sequences from the language model backbone in a way that does not affect padding tokens. + + Args: + lm_backbone (`torch.nn.Module`): + The language model backbone used for generation. + queries (`torch.Tensor`): + The tensor containing the input queries. + pad_token_id (`int`): + The token ID representing the pad token. + generation_config ([`~transformers.GenerationConfig`]): + The configuration for the generation process. + + Returns: + tuple: + - `generated_sequences` (`torch.Tensor`): + The concatenated tensor of input queries and generated sequences. + - `logits` (`torch.Tensor`): + The logits output from the generation process. + """ + context_length = queries.shape[1] + attention_mask = queries != pad_token_id + input_ids = torch.masked_fill(queries, ~attention_mask, 0) + output = lm_backbone.generate( + input_ids=input_ids, + attention_mask=attention_mask, + # position_ids=attention_mask.cumsum(1) - attention_mask.long(), # not needed: already adjusted in generations + # https://github.com/huggingface/transformers/blob/ac33aeeeee2a7a89b89c93c2962e6feb90daef0a/src/transformers/models/gpt2/modeling_gpt2.py#L1227-L1250 + generation_config=generation_config, + return_dict_in_generate=True, + output_scores=True, + ) + logits = torch.stack(output.scores, 1) + return torch.cat((queries, output.sequences[:, context_length:]), dim=1), logits + + +@torch.no_grad() +def batch_generation( + model: torch.nn.Module, + queries: torch.Tensor, + local_rollout_forward_batch_size: int, + pad_token_id: int, + generation_config: GenerationConfig, +): + query_responses = [] + logitss = [] + batch_size = queries.shape[0] + for i in range(0, batch_size, local_rollout_forward_batch_size): + query = queries[i : i + local_rollout_forward_batch_size] + query_response, logits = generate( + model, + query, + pad_token_id, + generation_config, + ) + query_responses.append(query_response) + logitss.append(logits) + + # padding tensors + padded_query_responses = pad(query_responses, padding_value=pad_token_id, padding_side="right") + padded_logitss = pad(logitss, padding_value=0, padding_side="right") + + # reshaping + padded_query_responses = padded_query_responses.view(-1, padded_query_responses.shape[-1])[:batch_size] + padded_logitss = padded_logitss.view(-1, *padded_logitss.shape[2:])[:batch_size] + + return padded_query_responses, padded_logitss + + +def add_bos_token_if_needed( + bos_token_id: int | None, + prompt_len_input_ids: int, + prompt_tokens: dict[str, list[int]], + chosen_prompt_len_input_ids: int, + chosen_tokens: dict[str, list[int]], + rejected_prompt_len_input_ids: int, + rejected_tokens: dict[str, list[int]], +): + if bos_token_id is not None: + if prompt_len_input_ids == 0 or bos_token_id != prompt_tokens["prompt_input_ids"][0]: + prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens["prompt_input_ids"] + prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens["prompt_attention_mask"] + if chosen_prompt_len_input_ids == 0 or bos_token_id != chosen_tokens["prompt_input_ids"][0]: + chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens["prompt_input_ids"] + chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens["prompt_attention_mask"] + if rejected_prompt_len_input_ids == 0 or bos_token_id != rejected_tokens["prompt_input_ids"][0]: + rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens["prompt_input_ids"] + rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens["prompt_attention_mask"] + return prompt_tokens, chosen_tokens, rejected_tokens + + +def add_eos_token_if_needed( + eos_token_id: int, chosen_tokens: dict[str, list[int]], rejected_tokens: dict[str, list[int]] +): + if len(chosen_tokens["input_ids"]) == 0 or eos_token_id != chosen_tokens["input_ids"][-1]: + chosen_tokens["input_ids"].append(eos_token_id) + chosen_tokens["attention_mask"].append(1) + if len(rejected_tokens["input_ids"]) == 0 or eos_token_id != rejected_tokens["input_ids"][-1]: + rejected_tokens["input_ids"].append(eos_token_id) + rejected_tokens["attention_mask"].append(1) + return chosen_tokens, rejected_tokens + + +def truncate_right( + input_ids: torch.Tensor, stop_token_id: int, pad_token_id: int +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Truncates the input tensor from the right side after the first occurrence of the stop token. + + Args: + input_ids (`torch.Tensor`): + The tensor containing the responses to be truncated + stop_token_id (`int`): + The token ID representing the stop token where truncation occurs + pad_token_id (`int`): + The token ID representing the pad token used to fill the truncated responses + + Returns: + tuple: + - `output_ids` (`torch.Tensor`): + The truncated responses tensor with pad tokens filled after the stop token + - `mask` (`torch.Tensor`): + The mask tensor to indicate the padding tokens + """ + trunc_idxs = first_true_indices(input_ids == stop_token_id).unsqueeze(-1) + new_size = [1] * (len(input_ids.size()) - 1) + [input_ids.shape[1]] + idxs = torch.arange(input_ids.shape[1], device=input_ids.device).view(*new_size) + output_ids = torch.masked_fill(input_ids, idxs > trunc_idxs, pad_token_id) + mask = torch.masked_fill(torch.ones_like(input_ids), idxs > trunc_idxs, 0) + return output_ids, mask + + +def empty_cache() -> None: + """Empties the cache of the available torch device. + + This function checks for the availability of different torch devices (XPU, MLU, NPU, CUDA) and empties the cache of + the first available device it finds. + + If none of the specific devices are available, it defaults to emptying the CUDA cache. + """ + if is_torch_xpu_available(): + torch.xpu.empty_cache() + elif is_torch_mlu_available(): + torch.mlu.empty_cache() + elif is_torch_npu_available(): + torch.npu.empty_cache() + else: + torch.cuda.empty_cache() + + +def generate_model_card( + base_model: str | None, + model_name: str, + hub_model_id: str, + dataset_name: str | None, + tags: list[str], + wandb_url: str | None, + trainer_name: str, + trainer_citation: str | None = None, + template_file: str | None = None, + paper_title: str | None = None, + paper_id: str | None = None, + comet_url: str | None = None, +) -> ModelCard: + """ + Generate a [`~huggingface_hub.ModelCard`] from a template. + + Args: + base_model (`str` or `None`): + Base model name. + model_name (`str`): + Model name. + hub_model_id (`str`): + Hub model ID as `username/model_id`. + dataset_name (`str` or `None`): + Dataset name. + tags (`list[str]`): + Tags. + wandb_url (`str` or `None`): + Weights & Biases run URL. + comet_url (`str` or `None`): + Comet experiment URL. + trainer_name (`str`): + Trainer name. + trainer_citation (`str` or `None`, defaults to `None`): + Trainer citation as a BibTeX entry. + template_file (`str` *optional*): + Template file name located in the `trl/templates` directory. Defaults to `lm_model_card.md`. + paper_title (`str` or `None`, defaults to `None`): + Paper title. + paper_id (`str` or `None`, defaults to `None`): + ArXiv paper ID as `YYMM.NNNNN`. + + Returns: + [`~huggingface_hub.ModelCard`]: + A ModelCard object. + """ + card_data = ModelCardData( + base_model=base_model, + datasets=dataset_name, + library_name="transformers", + licence="license", + model_name=model_name, + tags=["generated_from_trainer", *tags], + ) + template_file = template_file or "lm_model_card.md" + card = ModelCard.from_template( + card_data, + template_path=str(pkg_resources.files("trl").joinpath(f"templates/{template_file}")), + base_model=base_model, + model_name=model_name, + hub_model_id=hub_model_id, + dataset_name=dataset_name, + wandb_url=wandb_url, + comet_url=comet_url, + trainer_name=trainer_name, + trainer_citation=trainer_citation, + paper_title=paper_title, + paper_id=paper_id, + trl_version=version("trl"), + transformers_version=version("transformers"), + pytorch_version=version("torch"), + datasets_version=version("datasets"), + tokenizers_version=version("tokenizers"), + ) + return card + + +def get_comet_experiment_url() -> str | None: + """ + If Comet integration is enabled, return the URL of the current Comet experiment; otherwise, return `None`. + """ + if not is_comet_available(): + return None + + if comet_ml.get_running_experiment() is not None: + return comet_ml.get_running_experiment().url + + return None + + +def log_table_to_comet_experiment(name: str, table: pd.DataFrame) -> None: + """ + If Comet integration is enabled logs a table to the Comet experiment if it is currently running. + + Args: + name (`str`): + Table name. + table (`pandas.DataFrame`): + The Pandas DataFrame containing the table to log. + """ + if not is_comet_available(): + raise ModuleNotFoundError("The comet-ml is not installed. Please install it first: pip install comet-ml") + + experiment = comet_ml.get_running_experiment() + if experiment is not None: + experiment.log_table(tabular_data=table, filename=name) + + +def flush_left(mask: torch.Tensor, *tensors: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, ...]: + """ + Shift non-zero elements in the mask and corresponding tensors to the left. + + This function operates on a binary mask and any number of additional tensors with the same dimensions as the mask. + For each row, non-zero values are shifted to the leftmost positions. Then, columns that contain only zeros across + all rows are truncated from the mask and tensors. Visually, this operation can be represented as follows: + + ``` + [[0, 0, x, x, x, x], -> [[x, x, x, x], + [0, x, x, x, 0, 0]] [x, x, x, 0]] + ``` + + Args: + mask (`torch.Tensor`): + 2D tensor (binary mask) with shape `(N, M)`. + *tensors (`torch.Tensor`): + One or more 2D tensors with the same shape as `mask`. These tensors will be processed alongside `mask`, + with non-zero values shifted and excess zero columns truncated in the same manner. + + Returns: + `torch.Tensor`: + Updated binary mask with non-zero values flushed to the left and trailing zero columns removed. + `*torch.Tensor` + Updated tensors, processed in the same way as the mask. + + Example: + ```python + >>> mask = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 0, 0]]) + >>> tensor = torch.tensor([[9, 9, 2, 3, 4], [9, 5, 6, 9, 9]]) + >>> new_mask, new_tensor = flush_left(mask, tensor) + >>> print(new_mask) + tensor([[1, 1, 1], + [1, 1, 0]]) + + >>> print(new_tensor) + tensor([[2, 3, 4], + [5, 6, 0]]) + ``` + """ + _, M = mask.shape + + # Create copy of mask and tensors + mask_copy = mask.clone() + tensors = [t.clone() for t in tensors] + + # Shift non-zero values to the left + first_non_zero = mask_copy.argmax(dim=1) + pos = torch.arange(M, device=mask_copy.device).unsqueeze(0) + idx_roll = (pos + first_non_zero.unsqueeze(1)) % M + mask_roll = mask_copy.gather(1, idx_roll) + rolled_tensors = [t.gather(1, idx_roll) for t in tensors] + + # Truncate trailing columns that are all zeros in mask_roll + col_sums = mask_roll.sum(dim=0) + empty_cols = col_sums == 0 + first_empty_col = int(empty_cols.to(torch.int8).argmax()) if empty_cols.any() else M + flushed_mask = mask_roll[:, :first_empty_col] + flushed_tensors = [t[:, :first_empty_col] for t in rolled_tensors] + + if not flushed_tensors: + return flushed_mask + return flushed_mask, *flushed_tensors + + +def flush_right(mask: torch.Tensor, *tensors: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, ...]: + """ + Shift non-zero elements in the mask and corresponding tensors to the right. See `flush_left` for details. + """ + _, M = mask.shape + + # Create copy of mask and tensors + mask_copy = mask.clone() + tensors = [t.clone() for t in tensors] + + # Shift non-zero values to the right + flipped_mask = torch.fliplr(mask_copy) + first_non_zero = flipped_mask.argmax(dim=1) + pos = torch.arange(M, device=mask_copy.device).unsqueeze(0) + idx_roll = (pos - first_non_zero.unsqueeze(1)) % M + mask_roll = mask_copy.gather(1, idx_roll) + rolled_tensors = [t.gather(1, idx_roll) for t in tensors] + + # Truncate leading columns that are all zeros in mask_roll + col_sums = mask_roll.sum(dim=0) + non_empty_cols = col_sums != 0 + first_non_empty_col = int(non_empty_cols.to(torch.int8).argmax()) if non_empty_cols.any() else M + flushed_mask = mask_roll[:, first_non_empty_col:] + flushed_tensors = [t[:, first_non_empty_col:] for t in rolled_tensors] + + if not flushed_tensors: + return flushed_mask + return flushed_mask, *flushed_tensors + + +def selective_log_softmax(logits, index) -> torch.Tensor: + """ + A memory-efficient implementation of the common `log_softmax -> gather` operation. + + This function is equivalent to the following naive implementation: + ```python + logps = torch.gather(logits.log_softmax(-1), dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + ``` + + Args: + logits (`torch.Tensor`): + Logits tensor of shape `(..., num_classes)`. + index (`torch.Tensor`): + Index tensor of shape `(...)`, specifying the positions to gather from the log-softmax output. + + Returns: + `torch.Tensor`: + Gathered log probabilities with the same shape as `index`. + """ + if logits.dtype in [torch.float32, torch.float64]: + selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + # loop to reduce peak mem consumption + logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) + per_token_logps = selected_logits - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) + else: + # logsumexp approach is unstable with bfloat16, fall back to slightly less efficient approach + per_token_logps = [] + for row_logits, row_labels in zip(logits, index, strict=True): # loop to reduce peak mem consumption + row_logps = F.log_softmax(row_logits, dim=-1) + row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) + per_token_logps.append(row_per_token_logps) + per_token_logps = torch.stack(per_token_logps) + return per_token_logps + + +def entropy_from_logits(logits: torch.Tensor, chunk_size: int = 128) -> torch.Tensor: + """ + Compute the Shannon entropy (in nats) for each row of *logits* in a memory-efficient way. + + Instead of materializing the full softmax for all rows at once, the logits are flattened to shape (N, num_classes), + where N is the product of all leading dimensions. Computation is then performed in chunks of size `chunk_size` + along this flattened dimension, reducing peak memory usage. The result is reshaped back to match the input's + leading dimensions. + + Args: + logits (`torch.Tensor`): + Logits tensor of shape `(..., num_classes)`. Entropy is taken along the last axis; all leading dimensions + are preserved in the output. + chunk_size (`int`, *optional*, defaults to `128`): + Number of rows from the flattened logits to process per iteration. Smaller values reduce memory usage at + the cost of more iterations. + + Returns: + `torch.Tensor`: + Entropy values with shape `logits.shape[:-1]`. + """ + original_shape = logits.shape[:-1] # all dims except num_classes + num_classes = logits.shape[-1] + + # Flatten all leading dimensions into one + flat_logits = logits.reshape(-1, num_classes) + + entropies = [] + for chunk in flat_logits.split(chunk_size, dim=0): + logps = F.log_softmax(chunk, dim=-1) + chunk_entropy = -(torch.exp(logps) * logps).sum(-1) + entropies.append(chunk_entropy) + + entropies = torch.cat(entropies, dim=0) + return entropies.reshape(original_shape) + + +def print_prompt_completions_sample( + prompts: list, + completions: list, + rewards: dict[str, list[float]], + advantages: list[float], + step: int, + num_samples: int = None, +) -> None: + """ + Print out a sample of model completions to the console with multiple reward metrics. + + This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring model outputs + during training. It requires the `rich` library to be installed. + + Args: + prompts (`list`): + List of prompts. Can be either strings or lists of messages. + completions (`list`): + List of completions corresponding to the prompts. Can be either strings or lists of messages. + rewards (`dict[str, list[float]]`): + Dictionary where keys are reward names and values are lists of rewards. + advantages (`list[float]`): + List of advantages corresponding to the prompts and completions. + step (`int`): + Current training step number, used in the output title. + num_samples (`int`, *optional*): + Number of random samples to display. If `None` (default), all items will be displayed. + + Example: + ```python + >>> from trl.trainer.utils import print_prompt_completions_sample + + >>> prompts = ["The sky is", "The sun is"] + >>> completions = [" blue.", " in the sky."] + >>> rewards = {"Correctness": [0.123, 0.456], "Format": [0.789, 0.101]} + >>> advantages = [0.987, 0.654] + >>> print_prompt_completions_sample(prompts, completions, rewards, advantages, 42) + ╭──────────────────────────── Step 42 ─────────────────────────────╮ + │ ┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━┓ │ + │ ┃ Prompt ┃ Completion ┃ Correctness ┃ Format ┃ Advantage ┃ │ + │ ┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━┩ │ + │ │ The sky is │ blue. │ 0.12 │ 0.79 │ 0.99 │ │ + │ ├────────────┼──────────────┼─────────────┼────────┼───────────┤ │ + │ │ The sun is │ in the sky. │ 0.46 │ 0.10 │ 0.65 │ │ + │ └────────────┴──────────────┴─────────────┴────────┴───────────┘ │ + ╰──────────────────────────────────────────────────────────────────╯ + ``` + """ + if not is_rich_available(): + raise ImportError( + "The function `print_prompt_completions_sample` requires the `rich` library. Please install it with " + "`pip install rich`." + ) + console = Console() + table = Table(show_header=True, header_style="bold white", expand=True) + + # Add columns + table.add_column("Prompt", style="bright_yellow") + table.add_column("Completion", style="bright_green") + for reward_name in rewards.keys(): + table.add_column(reward_name, style="bold cyan", justify="right") + table.add_column("Advantage", style="bold magenta", justify="right") + + def format_entry(entry) -> Text: + t = Text() + if isinstance(entry, list) and all(isinstance(m, dict) for m in entry): + for j, msg in enumerate(entry): + role = msg.get("role", "") + if "content" in msg: + # Chat message + t.append(f"{role.upper()}\n", style="bold red") + t.append(msg["content"]) + elif "name" in msg and "args" in msg: + # Tool call + t.append(f"{role.upper()}\n", style="bold red") + t.append(f"{msg['name']}({msg['args']})") + else: + # Fallback + t.append(str(msg)) + if j < len(entry) - 1: + t.append("\n\n") + else: + t.append(str(entry)) + return t + + # Some basic input validation + if num_samples is not None: + if num_samples >= len(prompts): + num_samples = None + elif num_samples <= 0: + return + + # Subsample data if num_samples is specified + if num_samples is not None: + indices = random.sample(range(len(prompts)), num_samples) + prompts = [prompts[i] for i in indices] + completions = [completions[i] for i in indices] + rewards = {key: [val[i] for i in indices] for key, val in rewards.items()} + advantages = [advantages[i] for i in indices] + + for i in range(len(prompts)): + reward_values = [f"{rewards[key][i]:.2f}" for key in rewards.keys()] # 2 decimals + table.add_row( + format_entry(prompts[i]), + format_entry(completions[i]), + *reward_values, + f"{advantages[i]:.2f}", + ) + table.add_section() # Adds a separator between rows + + panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white") + console.print(panel) + + +class RepeatSampler(Sampler): + """ + Sampler that repeats the indices of a dataset in a structured manner. + + Args: + data_source (`Sized`): + Dataset to sample from. + mini_repeat_count (`int`): + Number of times to repeat each index per batch. + batch_size (`int`, *optional*, defaults to `1`): + Number of unique indices per batch. + repeat_count (`int`, *optional*, defaults to `1`): + Number of times to repeat the full sampling process. + shuffle (`bool`, *optional*, defaults to `True`): + Whether to shuffle the dataset. + seed (`int`, *optional*): + Random seed for reproducibility (only affects this sampler). + + Example: + ```python + >>> sampler = RepeatSampler(["a", "b", "c", "d", "e", "f", "g"], mini_repeat_count=2, batch_size=3, repeat_count=4) + >>> list(sampler) + [4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6] + ``` + + ```txt + mini_repeat_count = 3 + - - - + [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, | + 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, | + 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, | + repeat_count = 2 + 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, | + 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, | + 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, ...] | + --------- --------- --------- --------- + --------- --------- --------- --------- + --------- --------- --------- --------- + batch_size = 12 + ``` + """ + + def __init__( + self, + data_source: Sized, + mini_repeat_count: int, + batch_size: int = 1, + repeat_count: int = 1, + shuffle: bool = True, + seed: int | None = None, + ): + self.data_source = data_source + self.mini_repeat_count = mini_repeat_count + self.batch_size = batch_size + self.repeat_count = repeat_count + self.num_samples = len(data_source) + self.shuffle = shuffle + self.seed = seed + + if shuffle: + self.generator = torch.Generator() # Create a local random generator + if seed is not None: + self.generator.manual_seed(seed) + + def __iter__(self): + if self.shuffle: + # E.g., [2, 4, 3, 1, 0, 6, 5] (num_samples = 7) + indexes = torch.randperm(self.num_samples, generator=self.generator).tolist() + else: + indexes = list(range(self.num_samples)) + + # [2, 4, 3, 1, 0, 6, 5] + # -> [[2, 4, 3], [1, 0, 6], [5]] (batch_size = 3) + indexes = [indexes[i : i + self.batch_size] for i in range(0, len(indexes), self.batch_size)] + + # [[2, 4, 3], [1, 0, 6], [5]] + # -> [[2, 4, 3], [1, 0, 6]] + indexes = [chunk for chunk in indexes if len(chunk) == self.batch_size] + + for chunk in indexes: + for _ in range(self.repeat_count): + for index in chunk: + for _ in range(self.mini_repeat_count): + yield index + + def __len__(self) -> int: + return (self.num_samples // self.batch_size) * self.batch_size * self.mini_repeat_count * self.repeat_count + + +# torch.nanstd doesn't exist, so we define it here +def nanstd(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the standard deviation of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): + Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: + Standard deviation of the tensor, ignoring NaNs. + """ + variance = torch.nanmean((tensor - torch.nanmean(tensor, keepdim=True)) ** 2) # Compute variance ignoring NaNs + count = torch.sum(~torch.isnan(tensor)) # Count of non-NaN values + variance *= count / (count - 1) # Bessel's correction + return torch.sqrt(variance) + + +def split_tensor_dict( + tensor_dict: dict[str, torch.Tensor | None], num_chunks: int +) -> list[dict[str, torch.Tensor | None]]: + """ + Splits a dictionary of tensors along the first dimension into `num_chunks` equal parts. + + Example: + ```python + >>> x = torch.arange(12).reshape(6, 2) + >>> y = torch.arange(6).reshape(6, 1) + >>> tensor_dict = {"x": x, "y": y} + >>> split_tensor_dict(tensor_dict, 3) + [ + {"x": tensor([[0, 1], [2, 3]]), "y": tensor([[0], [1]])}, + {"x": tensor([[4, 5], [6, 7]]), "y": tensor([[2], [3]])}, + {"x": tensor([[ 8, 9], [10, 11]]), "y": tensor([[4], [5]])} + ] + ``` + """ + first_tensor = next(tensor for tensor in tensor_dict.values() if tensor is not None) + chunk_size = first_tensor.shape[0] // num_chunks + chunks = [] + for i in range(num_chunks): + chunk_dict = {} + for key, tensor in tensor_dict.items(): + if tensor is not None and (isinstance(tensor, list) or tensor.ndim > 0): + chunk_dict[key] = tensor[i * chunk_size : (i + 1) * chunk_size] + elif tensor is not None and tensor.ndim == 0: + chunk_dict[key] = tensor + else: + chunk_dict[key] = None + chunks.append(chunk_dict) + return chunks + + +def shuffle_sequence_dict(seq_dict: dict[str, Sequence | None]) -> dict[str, Sequence | None]: + """ + Shuffles all sequence-like values in a dictionary along the first dimension in unison. + + Example: + ```python + >>> x = torch.arange(6).reshape(3, 2) + >>> y = ["a", "b", "c"] + >>> seq_dict = {"x": x, "y": y} + >>> shuffle_sequence_dict(seq_dict) + {'x': tensor([[2, 3], + [0, 1], + [4, 5]]), + 'y': ['b', 'a', 'c']} + ``` + """ + # Determine batch size from the first non-None sequence + batch_size = len(next(v for v in seq_dict.values() if v is not None)) + permutation = torch.randperm(batch_size) + + def permute(v: Sequence | None) -> Sequence | None: + if v is None: + return None + if isinstance(v, torch.Tensor) and v.ndim == 0: + return v + if isinstance(v, torch.Tensor) and v.ndim >= 1: + return v[permutation] + return [v[i] for i in permutation] + + return {key: permute(val) for key, val in seq_dict.items()} + + +def nanmin(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the minimum value of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: Minimum value of the tensor, ignoring NaNs. Returns NaN if all values are NaN. + """ + if torch.isnan(tensor).all(): + return torch.tensor(float("nan"), dtype=tensor.dtype, device=tensor.device) + return torch.min(tensor[~torch.isnan(tensor)]) + + +def nanmax(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the maximum value of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: Maximum value of the tensor, ignoring NaNs. Returns NaN if all values are NaN. + """ + if torch.isnan(tensor).all(): + return torch.tensor(float("nan"), dtype=tensor.dtype, device=tensor.device) + return torch.max(tensor[~torch.isnan(tensor)]) + + +def identity(x): + """Do we really need docs for this?""" + return x + + +def split_pixel_values_by_grid(batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor | list[torch.Tensor]]: + """ + Splits `batch["pixel_values"]` into a list of tensors based on the product of each row in `batch["image_grid_thw"]` + and batch["num_images"] while keeping other entries unchanged. + """ + if "image_grid_thw" not in batch or "pixel_values" not in batch or "num_images" not in batch: + return batch + + lengths = batch["image_grid_thw"].prod(-1).tolist() # [num_images] + pixel_values = batch["pixel_values"] # [total, feature_dim] + + if sum(lengths) != pixel_values.size(0): + raise ValueError(f"Mismatch: sum(lengths) = {sum(lengths)} != pixel_values.size(0) = {pixel_values.size(0)}") + + boundaries = [0, *accumulate(batch["num_images"])] # [3, 4, 5] -> [0, 3, 7, 12] + sections = [sum(lengths[boundaries[i] : boundaries[i + 1]]) for i in range(len(batch["num_images"]))] + split_values = list(torch.split(batch["pixel_values"], sections, dim=0)) + image_grid_thw = list(torch.split(batch["image_grid_thw"], batch["num_images"], dim=0)) + return {**batch, "pixel_values": split_values, "image_grid_thw": image_grid_thw} + + +def unsplit_pixel_values_by_grid(batch: dict[str, torch.Tensor | list[torch.Tensor]]) -> dict[str, torch.Tensor]: + """ + Opposite of `split_pixel_values_by_grid`. Merges a list of tensors in `batch["pixel_values"]` back into a single + tensor along the first dimension. + """ + pixel_values = batch.get("pixel_values") + if isinstance(pixel_values, list): + merged = torch.cat(pixel_values, dim=0) + batch = {**batch, "pixel_values": merged} + + image_grid_thw = batch.get("image_grid_thw") + if isinstance(image_grid_thw, list): + merged = torch.cat(image_grid_thw, dim=0) + batch = {**batch, "image_grid_thw": merged} + + return batch + + +TListOrMapping = TypeVar("TListOrMapping", list, Mapping) + + +def remove_none_values(example: TListOrMapping) -> TListOrMapping: + """ + Recursively removes entries with `None` values from a nested structure (list or dictionary). + + Args: + example (`list` or `Mapping`): + Input nested structure (list or dictionary) from which to remove `None`. + + Example: + ```python + >>> [ + ... { + ... "a": {"aa": None, "ab": 1}, + ... "b": "my_string", + ... } + ... ] + >>> remove_none_values(example) + [{'a': {'ab': 1}, 'b': 'my_string'}] + ``` + """ + if isinstance(example, list): + return [remove_none_values(value) if isinstance(value, (dict, list)) else value for value in example] + elif isinstance(example, Mapping): + return { + key: remove_none_values(value) if isinstance(value, (dict, list)) else value + for key, value in example.items() + if value is not None + } + else: + raise TypeError("Input must be a list or a dictionary.") + + +def create_model_from_path(model_id: str, **kwargs) -> PreTrainedModel: + """ + Create a model from a given path using the specified initialization arguments. + + Args: + model_id (`str`): + Path to the model. Can be either a local directory or a model identifier from the Hugging Face Hub. + kwargs (`dict`): + Initialization keyword arguments to pass to the model's `from_pretrained` method. When `'dtype'` is + specified, it can be either a `torch.dtype` or one of the strings: `'bfloat16'`, `'float16'`, `'float32'`, + or `'auto'`. + + Returns: + [`~transformers.PreTrainedModel`]: + The instantiated model. + """ + dtype = kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass # dtype is already a torch.dtype or "auto" or None + elif isinstance(dtype, str) and dtype in ["bfloat16", "float16", "float32"]: + kwargs["dtype"] = getattr(torch, dtype) + else: + raise ValueError( + "Invalid `dtype` passed to the config. Expected either 'auto' or a string representing " + f"a valid `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + model = architecture.from_pretrained(model_id, **kwargs) + return model + + +def get_config_model_id(config: PretrainedConfig) -> str: + """ + Retrieve the model identifier from a given model configuration. + + Args: + config ([`~transformers.PreTrainedConfig`]): + Configuration from which to extract the model identifier. + + Returns: + `str`: + The model identifier associated with the model configuration. + """ + # Fall back to `config.text_config._name_or_path` if `config._name_or_path` is missing: Qwen2-VL and Qwen2.5-VL. See GH-4323 + return getattr(config, "_name_or_path", "") or getattr(getattr(config, "text_config", None), "_name_or_path", "") diff --git a/venv/lib/python3.10/site-packages/trl/trainer/xpo_config.py b/venv/lib/python3.10/site-packages/trl/trainer/xpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..78c1e562128208e0b6a2ba6e78f9becd59e8a1ac --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/xpo_config.py @@ -0,0 +1,44 @@ +# Copyright 2020-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 dataclasses import dataclass, field + +from trl.trainer.online_dpo_config import OnlineDPOConfig + + +@dataclass +class XPOConfig(OnlineDPOConfig): + r""" + Configuration class for the [`XPOTrainer`]. + + Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following: + + Parameters: + alpha (`float` or `list[float]`, *optional*, defaults to `1e-5`): + Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each new epoch + and the last alpha is used for the rest of the epochs. + """ + + alpha: list[float] = field( + default_factory=lambda: [1e-5], + metadata={ + "help": "Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each " + "new epoch and the last alpha is used for the rest of the epochs." + }, + ) + + def __post_init__(self): + super().__post_init__() + if hasattr(self.alpha, "__len__") and len(self.alpha) == 1: + self.alpha = self.alpha[0] diff --git a/venv/lib/python3.10/site-packages/trl/trainer/xpo_trainer.py b/venv/lib/python3.10/site-packages/trl/trainer/xpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..5646885e791c5aeb30636d9dfc7b038afb5b6f0c --- /dev/null +++ b/venv/lib/python3.10/site-packages/trl/trainer/xpo_trainer.py @@ -0,0 +1,538 @@ +# Copyright 2020-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 textwrap +from collections.abc import Callable +from typing import Any + +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset, IterableDataset +from transformers import ( + BaseImageProcessor, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, +) +from transformers.trainer_utils import EvalPrediction +from transformers.training_args import OptimizerNames +from transformers.utils import is_peft_available + +from ..data_utils import is_conversational, maybe_apply_chat_template +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .online_dpo_trainer import OnlineDPOTrainer +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + empty_cache, + get_reward, + selective_log_softmax, + truncate_right, +) +from .xpo_config import XPOConfig + + +if is_peft_available(): + from peft import PeftModel + + +class XPOTrainer(OnlineDPOTrainer): + """ + Trainer for Exploratory Preference Optimization (XPO). + + It is implemented as a subclass of [`OnlineDPOTrainer`]. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model ([`PreTrainedModelWrapper`]): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + reward_funcs ([`~transformers.PreTrainedModel`]): + The reward model to score completions with, preferably an + [`~transformers.AutoModelForSequenceClassification`]. + judge ([`BasePairwiseJudge`]): + The judge to use for pairwise comparison of model completions. + args ([`XPOConfig`]): + The XPO config arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`DPODataCollatorWithPadding`]) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + peft_config (`dict`): + The peft config to use for training. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "xpo"] + _name = "XPO" + _paper = { + "title": "Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF", + "id": "2405.21046", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{jung2024binary, + title = {{Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF}}, + author = {Tengyang Xie and Dylan J. Foster and Akshay Krishnamurthy and Corby Rosset and Ahmed Awadallah and Alexander Rakhlin}, + year = 2024, + eprint = {arXiv:2405.21046} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module = None, + ref_model: PreTrainedModel | nn.Module = None, + reward_funcs: nn.Module | None = None, + judge: BasePairwiseJudge | None = None, + args: XPOConfig | None = None, + data_collator: Callable | None = None, + train_dataset: Dataset | IterableDataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, + peft_config: dict | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + judge=judge, + reward_funcs=reward_funcs, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_classes=reward_processing_classes, + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + self._alpha = self.args.alpha + + # Overwrite the stats dictionary to include XPO specific statistics + self.stats = { + # Remove "non_score_reward", "rlhf_reward", "scores" + # Add "loss/dpo", "loss/xpo" + "loss/dpo": [], + "loss/xpo": [], + "objective/kl": [], + "objective/entropy": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + # Replace "contain_eos_token" by "model_contain_eos_token" and "ref_contain_eos_token" + "val/model_contain_eos_token": [], + "val/ref_contain_eos_token": [], + "alpha": [], + "beta": [], + } + if self.reward_funcs is not None: + if len(self.reward_funcs) != 1: + raise ValueError("XPOTrainer only supports one reward function/model.") + self.reward_funcs = self.reward_funcs[0] + self.stats["objective/model_scores"] = [] + self.stats["objective/ref_scores"] = [] + self.stats["objective/scores_margin"] = [] + + @property + def alpha(self): + if isinstance(self._alpha, list): + epoch = self.state.epoch + return self._alpha[epoch] if epoch < len(self._alpha) else self._alpha[-1] + else: + return self._alpha + + def _generate_completions(self, prompts, model): + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_model_for_gen: + model_output = unwrapped_policy_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + actual_model_for_ref_generation: torch.nn.Module + if self.ref_model is None: + unwrapped_main_model_for_ref_logic = self.accelerator.unwrap_model(model) + + if is_peft_available() and isinstance(unwrapped_main_model_for_ref_logic, PeftModel): + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic.get_base_model() + else: + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic + else: + actual_model_for_ref_generation = self.accelerator.unwrap_model(self.ref_model) + + with unwrap_model_for_generation(actual_model_for_ref_generation, self.accelerator) as final_ref_model_for_gen: + ref_output = final_ref_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + return model_output, ref_output + + def _process_completions(self, model_output, ref_output, prompts): + context_length = prompts["input_ids"].shape[1] + + # Process model completions + model_completion_ids = model_output[:, context_length:] + model_completion_ids, model_completion_mask = truncate_right( + model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + model_data = { + "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1), + "raw": prompts["raw"], + } + + # Process reference model completions + ref_completion_ids = ref_output[:, context_length:] + ref_completion_ids, ref_completion_mask = truncate_right( + ref_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + ref_data = { + "input_ids": torch.cat((prompts["input_ids"], ref_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], ref_completion_mask), dim=1), + "raw": prompts["raw"], + } + + return model_data, ref_data + + def _compute_rewards(self, model_data, ref_data, context_length): + with torch.no_grad(): + _, model_scores, _ = get_reward( + self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, ref_scores, _ = get_reward( + self.reward_funcs, ref_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + + # Apply EOS penalty if needed + if self.args.missing_eos_penalty is not None: + model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + ref_contain_eos = torch.any(ref_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + model_scores[~model_contain_eos] -= self.args.missing_eos_penalty + ref_scores[~ref_contain_eos] -= self.args.missing_eos_penalty + + return model_scores, ref_scores + + def _compute_judge(self, model_data, ref_data, context_length): + prompts = model_data["raw"] + model_data_completions = self.processing_class.batch_decode( + model_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + model_data_completions = [completion.strip() for completion in model_data_completions] + + ref_data_completions = self.processing_class.batch_decode( + ref_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + ref_data_completions = [completion.strip() for completion in ref_data_completions] + + if is_conversational({"prompt": prompts[0]}): + model_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in model_data_completions + ] + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=message) for message in prompts] + model_data_completions = [template.render(messages=completion) for completion in model_data_completions] + + ref_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in ref_data_completions + ] + ref_data_completions = [template.render(messages=completion) for completion in ref_data_completions] + + ranks_of_first_completion = self.judge.judge( + prompts, + list(zip(model_data_completions, ref_data_completions, strict=True)), + ) + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + return torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=model_data["input_ids"].device) + + def _compute_logprobs(self, model, model_data, ref_data, context_length): + def compute_logprobs_for_data(m, data): + output = m(data["input_ids"], attention_mask=data["attention_mask"]) + logits = output.logits[:, context_length - 1 : -1] + token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:]) + return token_logprobs + + # Compute logprobs for model completions + model_logprobs_model_data = compute_logprobs_for_data(model, model_data) + # Compute logprobs for model on reference completions (for XPO loss) + model_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + + # Compute logprobs for reference model completions + with torch.no_grad(): + if self.ref_model is None: + with model.disable_adapter(): + ref_logprobs_model_data = compute_logprobs_for_data(model, model_data) + ref_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + else: + ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data) + ref_logprobs_ref_data = compute_logprobs_for_data(self.ref_model, ref_data) + + # Mask padding tokens + model_padding_mask = model_data["attention_mask"][:, context_length:] == 0 + ref_padding_mask = ref_data["attention_mask"][:, context_length:] == 0 + model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + model_logprobs_ref_data = model_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_ref_data = ref_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + + return model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data + + def _compute_losses( + self, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ): + # Compute log probs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + # Compute logits as the difference between chosen and rejected log ratios + logits = chosen_log_ratios - rejected_log_ratios + + if self.args.loss_type == "sigmoid": + dpo_losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + dpo_losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.args.loss_type}") + + # Compute XPO specific loss + xpo_losses = self.alpha * model_logprobs_ref_data_sum + + # Total loss + loss = (dpo_losses + xpo_losses).mean() + + return loss, dpo_losses, xpo_losses + + def _log_statistics( + self, + model_data, + ref_data, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses, + xpo_losses, + context_length, + model_scores=None, + ref_scores=None, + ): + # Helper function to gather and compute mean + def gather_mean(tensor): + return self.accelerator.gather_for_metrics(tensor).mean().item() + + # Log losses + self.stats["loss/dpo"].append(gather_mean(dpo_losses)) + self.stats["loss/xpo"].append(gather_mean(xpo_losses)) + + # Log scores + if self.reward_funcs is not None: + self.stats["objective/model_scores"].append(gather_mean(model_scores)) + self.stats["objective/ref_scores"].append(gather_mean(ref_scores)) + self.stats["objective/scores_margin"].append(gather_mean(model_scores - ref_scores)) + + # Log logprobs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean())) + self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean())) + + # Log rewards + # Compute various statistics + chosen_rewards = chosen_log_ratios * self.beta + rejected_rewards = rejected_log_ratios * self.beta + self.stats["rewards/chosen"].append(gather_mean(chosen_rewards.mean())) + self.stats["rewards/rejected"].append(gather_mean(rejected_rewards.mean())) + + # Calculate KL divergence for model and ref data + kl_model_data = model_logprobs_model_data - ref_logprobs_model_data + kl_ref_data = model_logprobs_ref_data - ref_logprobs_ref_data + mean_kl = (kl_model_data.sum(1) + kl_ref_data.sum(1)).mean() / 2 + self.stats["objective/kl"].append(gather_mean(mean_kl)) + + # Calculate entropy for model and ref data + entropy_model_data = -model_logprobs_model_data.sum(1) + entropy_ref_data = -model_logprobs_ref_data.sum(1) + mean_entropy = (entropy_model_data.mean() + entropy_ref_data.mean()) / 2 + self.stats["objective/entropy"].append(gather_mean(mean_entropy)) + + # Calculate margins + margin = chosen_rewards - rejected_rewards + self.stats["rewards/margins"].append(gather_mean(margin.mean())) + + # Calculate accuracy + accuracy = (margin > 0).float() + self.stats["rewards/accuracies"].append(gather_mean(accuracy.mean())) + + # Log EOS token statistics + model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + ref_eos = (ref_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float())) + self.stats["val/ref_contain_eos_token"].append(gather_mean(ref_eos.float())) + + # Log alpha and beta + self.stats["alpha"].append(self.alpha) + self.stats["beta"].append(self.beta) + + def training_step( + self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None + ) -> torch.Tensor: + model.train() + + # Apply chat template and tokenize the input + batch_size = len(next(iter(inputs.values()))) + prompts = inputs["prompt"] + inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # need the prompt_ only + inputs = self._prepare_inputs(inputs) + context_length = inputs["prompt_input_ids"].shape[1] + prompts = { + "input_ids": inputs["prompt_input_ids"], + "attention_mask": inputs["prompt_attention_mask"], + "raw": prompts, + } + del inputs + + # Sample completions from both the model and the reference model + model_output, ref_output = self._generate_completions(prompts, model) + + # Process model completions + model_data, ref_data = self._process_completions(model_output, ref_output, prompts) + + # Compute rewards + if self.reward_funcs is not None: + model_scores, ref_scores = self._compute_rewards(model_data, ref_data, context_length) + chosen_mask = model_scores >= ref_scores + else: + model_scores, ref_scores = None, None + chosen_mask = self._compute_judge(model_data, ref_data, context_length) + + # Compute logprobs + model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data = ( + self._compute_logprobs(model, model_data, ref_data, context_length) + ) + + # Compute loss + loss, dpo_losses, xpo_losses = self._compute_losses( + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ) + + # Log everything + self._log_statistics( + model_data, + ref_data, + model_logprobs_model_data.detach(), + model_logprobs_ref_data.detach(), + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses.detach(), + xpo_losses.detach(), + context_length, + model_scores, + ref_scores, + ) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/INSTALLER b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/METADATA b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..15116c784ed1f8f0f524513d606a26bf6765e923 --- /dev/null +++ b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/METADATA @@ -0,0 +1,154 @@ +Metadata-Version: 2.4 +Name: urllib3 +Version: 2.5.0 +Summary: HTTP library with thread-safe connection pooling, file post, and more. +Project-URL: Changelog, https://github.com/urllib3/urllib3/blob/main/CHANGES.rst +Project-URL: Documentation, https://urllib3.readthedocs.io +Project-URL: Code, https://github.com/urllib3/urllib3 +Project-URL: Issue tracker, https://github.com/urllib3/urllib3/issues +Author-email: Andrey Petrov +Maintainer-email: Seth Michael Larson , Quentin Pradet , Illia Volochii +License-Expression: MIT +License-File: LICENSE.txt +Keywords: filepost,http,httplib,https,pooling,ssl,threadsafe,urllib +Classifier: Environment :: Web Environment +Classifier: Intended Audience :: Developers +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Topic :: Internet :: WWW/HTTP +Classifier: Topic :: Software Development :: Libraries +Requires-Python: >=3.9 +Provides-Extra: brotli +Requires-Dist: brotli>=1.0.9; (platform_python_implementation == 'CPython') and extra == 'brotli' +Requires-Dist: brotlicffi>=0.8.0; (platform_python_implementation != 'CPython') and extra == 'brotli' +Provides-Extra: h2 +Requires-Dist: h2<5,>=4; extra == 'h2' +Provides-Extra: socks +Requires-Dist: pysocks!=1.5.7,<2.0,>=1.5.6; extra == 'socks' +Provides-Extra: zstd +Requires-Dist: zstandard>=0.18.0; extra == 'zstd' +Description-Content-Type: text/markdown + +

+ +![urllib3](https://github.com/urllib3/urllib3/raw/main/docs/_static/banner_github.svg) + +

+ +

+ PyPI Version + Python Versions + Join our Discord + Coverage Status + Build Status on GitHub + Documentation Status
+ OpenSSF Scorecard + SLSA 3 + CII Best Practices +

+ +urllib3 is a powerful, *user-friendly* HTTP client for Python. Much of the +Python ecosystem already uses urllib3 and you should too. +urllib3 brings many critical features that are missing from the Python +standard libraries: + +- Thread safety. +- Connection pooling. +- Client-side SSL/TLS verification. +- File uploads with multipart encoding. +- Helpers for retrying requests and dealing with HTTP redirects. +- Support for gzip, deflate, brotli, and zstd encoding. +- Proxy support for HTTP and SOCKS. +- 100% test coverage. + +urllib3 is powerful and easy to use: + +```python3 +>>> import urllib3 +>>> resp = urllib3.request("GET", "http://httpbin.org/robots.txt") +>>> resp.status +200 +>>> resp.data +b"User-agent: *\nDisallow: /deny\n" +``` + +## Installing + +urllib3 can be installed with [pip](https://pip.pypa.io): + +```bash +$ python -m pip install urllib3 +``` + +Alternatively, you can grab the latest source code from [GitHub](https://github.com/urllib3/urllib3): + +```bash +$ git clone https://github.com/urllib3/urllib3.git +$ cd urllib3 +$ pip install . +``` + + +## Documentation + +urllib3 has usage and reference documentation at [urllib3.readthedocs.io](https://urllib3.readthedocs.io). + + +## Community + +urllib3 has a [community Discord channel](https://discord.gg/urllib3) for asking questions and +collaborating with other contributors. Drop by and say hello 👋 + + +## Contributing + +urllib3 happily accepts contributions. Please see our +[contributing documentation](https://urllib3.readthedocs.io/en/latest/contributing.html) +for some tips on getting started. + + +## Security Disclosures + +To report a security vulnerability, please use the +[Tidelift security contact](https://tidelift.com/security). +Tidelift will coordinate the fix and disclosure with maintainers. + + +## Maintainers + +- [@sethmlarson](https://github.com/sethmlarson) (Seth M. Larson) +- [@pquentin](https://github.com/pquentin) (Quentin Pradet) +- [@illia-v](https://github.com/illia-v) (Illia Volochii) +- [@theacodes](https://github.com/theacodes) (Thea Flowers) +- [@haikuginger](https://github.com/haikuginger) (Jess Shapiro) +- [@lukasa](https://github.com/lukasa) (Cory Benfield) +- [@sigmavirus24](https://github.com/sigmavirus24) (Ian Stapleton Cordasco) +- [@shazow](https://github.com/shazow) (Andrey Petrov) + +👋 + + +## Sponsorship + +If your company benefits from this library, please consider [sponsoring its +development](https://urllib3.readthedocs.io/en/latest/sponsors.html). + + +## For Enterprise + +Professional support for urllib3 is available as part of the [Tidelift +Subscription][1]. Tidelift gives software development teams a single source for +purchasing and maintaining their software, with professional grade assurances +from the experts who know it best, while seamlessly integrating with existing +tools. + +[1]: https://tidelift.com/subscription/pkg/pypi-urllib3?utm_source=pypi-urllib3&utm_medium=referral&utm_campaign=readme diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/RECORD b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..6c760588131100aab79a126430b69ecbbf89976c --- /dev/null +++ b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/RECORD @@ -0,0 +1,80 @@ +urllib3-2.5.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +urllib3-2.5.0.dist-info/METADATA,sha256=maYkTIZt0a-lkEC-hMZWbCBmcGZyJcYOeRk4_nuTrNc,6461 +urllib3-2.5.0.dist-info/RECORD,, +urllib3-2.5.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +urllib3-2.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87 +urllib3-2.5.0.dist-info/licenses/LICENSE.txt,sha256=Ew46ZNX91dCWp1JpRjSn2d8oRGnehuVzIQAmgEHj1oY,1093 +urllib3/__init__.py,sha256=JMo1tg1nIV1AeJ2vENC_Txfl0e5h6Gzl9DGVk1rWRbo,6979 +urllib3/__pycache__/__init__.cpython-310.pyc,, +urllib3/__pycache__/_base_connection.cpython-310.pyc,, +urllib3/__pycache__/_collections.cpython-310.pyc,, +urllib3/__pycache__/_request_methods.cpython-310.pyc,, +urllib3/__pycache__/_version.cpython-310.pyc,, +urllib3/__pycache__/connection.cpython-310.pyc,, +urllib3/__pycache__/connectionpool.cpython-310.pyc,, +urllib3/__pycache__/exceptions.cpython-310.pyc,, +urllib3/__pycache__/fields.cpython-310.pyc,, +urllib3/__pycache__/filepost.cpython-310.pyc,, +urllib3/__pycache__/poolmanager.cpython-310.pyc,, +urllib3/__pycache__/response.cpython-310.pyc,, +urllib3/_base_connection.py,sha256=T1cwH3RhzsrBh6Bz3AOGVDboRsE7veijqZPXXQTR2Rg,5568 +urllib3/_collections.py,sha256=tM7c6J1iKtWZYV_QGYb8-r7Nr1524Dehnsa0Ufh6_mU,17295 +urllib3/_request_methods.py,sha256=gCeF85SO_UU4WoPwYHIoz_tw-eM_EVOkLFp8OFsC7DA,9931 +urllib3/_version.py,sha256=ZlSUkBo_Pd90B6pM0GDO7l2vitQD3QCK3xPR_K0zFJA,511 +urllib3/connection.py,sha256=iP4pgSJtpusXyYlejzNn-gih_wWCxMU-qy6OU1kaapc,42613 +urllib3/connectionpool.py,sha256=ZEhudsa8BIubD2M0XoxBBsjxbsXwMgUScH7oQ9i-j1Y,43371 +urllib3/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +urllib3/contrib/__pycache__/__init__.cpython-310.pyc,, +urllib3/contrib/__pycache__/pyopenssl.cpython-310.pyc,, +urllib3/contrib/__pycache__/socks.cpython-310.pyc,, +urllib3/contrib/emscripten/__init__.py,sha256=u6KNgzjlFZbuAAXa_ybCR7gQ71VJESnF-IIdDA73brw,733 +urllib3/contrib/emscripten/__pycache__/__init__.cpython-310.pyc,, +urllib3/contrib/emscripten/__pycache__/connection.cpython-310.pyc,, +urllib3/contrib/emscripten/__pycache__/fetch.cpython-310.pyc,, +urllib3/contrib/emscripten/__pycache__/request.cpython-310.pyc,, +urllib3/contrib/emscripten/__pycache__/response.cpython-310.pyc,, +urllib3/contrib/emscripten/connection.py,sha256=j8DR_flE7hsoFhNfiqHLiaPaCsVbzG44jgahwvsQ52A,8771 +urllib3/contrib/emscripten/emscripten_fetch_worker.js,sha256=CDfYF_9CDobtx2lGidyJ1zjDEvwNT5F-dchmVWXDh0E,3655 +urllib3/contrib/emscripten/fetch.py,sha256=kco06lWoQ-fdFfN51-nzeTywPVBEHg89WIst33H3xcg,23484 +urllib3/contrib/emscripten/request.py,sha256=mL28szy1KvE3NJhWor5jNmarp8gwplDU-7gwGZY5g0Q,566 +urllib3/contrib/emscripten/response.py,sha256=7oVPENYZHuzEGRtG40HonpH5tAIYHsGcHPbJt2Z0U-Y,9507 +urllib3/contrib/pyopenssl.py,sha256=Xp5Ym05VgXGhHa0C4wlutvHxY8SnKSS6WLb2t5Miu0s,19720 +urllib3/contrib/socks.py,sha256=-iardc61GypsJzD6W6yuRS7KVCyfowcQrl_719H7lIM,7549 +urllib3/exceptions.py,sha256=pziumHf0Vwx3z4gvUy7ou8nlM2yIYX0N3l3znEdeF5U,9938 +urllib3/fields.py,sha256=FCf7UULSkf10cuTRUWTQESzxgl1WT8e2aCy3kfyZins,10829 +urllib3/filepost.py,sha256=U8eNZ-mpKKHhrlbHEEiTxxgK16IejhEa7uz42yqA_dI,2388 +urllib3/http2/__init__.py,sha256=xzrASH7R5ANRkPJOot5lGnATOq3KKuyXzI42rcnwmqs,1741 +urllib3/http2/__pycache__/__init__.cpython-310.pyc,, +urllib3/http2/__pycache__/connection.cpython-310.pyc,, +urllib3/http2/__pycache__/probe.cpython-310.pyc,, +urllib3/http2/connection.py,sha256=4DB0DkZEC3yIkhGjUDIHB17wrYCLaL0Ag5bDW2_mGPI,12694 +urllib3/http2/probe.py,sha256=nnAkqbhAakOiF75rz7W0udZ38Eeh_uD8fjV74N73FEI,3014 +urllib3/poolmanager.py,sha256=oKsgP1EsAI4OVgK9-9D3AYXZS5HYV8yKUSog-QbJ8Ts,23866 +urllib3/py.typed,sha256=UaCuPFa3H8UAakbt-5G8SPacldTOGvJv18pPjUJ5gDY,93 +urllib3/response.py,sha256=TVTSu6Q1U0U7hoHYMIRxxuh4zroeMo8b5EI4DOA13Eo,46480 +urllib3/util/__init__.py,sha256=-qeS0QceivazvBEKDNFCAI-6ACcdDOE4TMvo7SLNlAQ,1001 +urllib3/util/__pycache__/__init__.cpython-310.pyc,, +urllib3/util/__pycache__/connection.cpython-310.pyc,, +urllib3/util/__pycache__/proxy.cpython-310.pyc,, +urllib3/util/__pycache__/request.cpython-310.pyc,, +urllib3/util/__pycache__/response.cpython-310.pyc,, +urllib3/util/__pycache__/retry.cpython-310.pyc,, +urllib3/util/__pycache__/ssl_.cpython-310.pyc,, +urllib3/util/__pycache__/ssl_match_hostname.cpython-310.pyc,, +urllib3/util/__pycache__/ssltransport.cpython-310.pyc,, +urllib3/util/__pycache__/timeout.cpython-310.pyc,, +urllib3/util/__pycache__/url.cpython-310.pyc,, +urllib3/util/__pycache__/util.cpython-310.pyc,, +urllib3/util/__pycache__/wait.cpython-310.pyc,, +urllib3/util/connection.py,sha256=JjO722lzHlzLXPTkr9ZWBdhseXnMVjMSb1DJLVrXSnQ,4444 +urllib3/util/proxy.py,sha256=seP8-Q5B6bB0dMtwPj-YcZZQ30vHuLqRu-tI0JZ2fzs,1148 +urllib3/util/request.py,sha256=XuAsEBT58DAZYUTwpMH5Hr3A1OPoMNvNIYIunbIqbc8,8411 +urllib3/util/response.py,sha256=vQE639uoEhj1vpjEdxu5lNIhJCSUZkd7pqllUI0BZOA,3374 +urllib3/util/retry.py,sha256=bj-2YUqblxLlv8THg5fxww-DM54XCbjgZXIQ71XioCY,18459 +urllib3/util/ssl_.py,sha256=jxnQ3msYkVaokJVWqHNnAVdVtDdidrTHDeyk50gwqaQ,19786 +urllib3/util/ssl_match_hostname.py,sha256=Di7DU7zokoltapT_F0Sj21ffYxwaS_cE5apOtwueeyA,5845 +urllib3/util/ssltransport.py,sha256=Ez4O8pR_vT8dan_FvqBYS6dgDfBXEMfVfrzcdUoWfi4,8847 +urllib3/util/timeout.py,sha256=4eT1FVeZZU7h7mYD1Jq2OXNe4fxekdNvhoWUkZusRpA,10346 +urllib3/util/url.py,sha256=WRh-TMYXosmgp8m8lT4H5spoHw5yUjlcMCfU53AkoAs,15205 +urllib3/util/util.py,sha256=j3lbZK1jPyiwD34T8IgJzdWEZVT-4E-0vYIJi9UjeNA,1146 +urllib3/util/wait.py,sha256=_ph8IrUR3sqPqi0OopQgJUlH4wzkGeM5CiyA7XGGtmI,4423 diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/REQUESTED b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/WHEEL b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..12228d414b6cfed7c39d3781c85c63256a1d7fb5 --- /dev/null +++ b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: hatchling 1.27.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/licenses/LICENSE.txt b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/licenses/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..e6183d0276b26c5b87aecccf8d0d5bcd7b1148d4 --- /dev/null +++ b/venv/lib/python3.10/site-packages/urllib3-2.5.0.dist-info/licenses/LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2008-2020 Andrey Petrov and contributors. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE.